<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 Path is an object that represents an optimized version of an
// XPath-like string.
type Path struct {
	segments []segment
}

// NewPath creates an optimized version of an XPath-like string that
// can be used to query elements in an element tree.  Panics if an
// invalid path string is supplied.
func NewPath(path string) Path {
	return Path{parsePath(path)}
}

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

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

// 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)
}

// A pather is helper object used to traverse an element tree with
// a Path object.  It collects and deduplicates elements matching
// the path query.
type pather struct {
	queue      []node
	qindex     int
	results    []*Element
	inResults  map[*Element]bool
	candidates []*Element
	scratch    []*Element // used by filters
}

// 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{
		queue:      make([]node, 0),
		results:    make([]*Element, 0),
		inResults:  make(map[*Element]bool),
		candidates: make([]*Element, 0),
		scratch:    make([]*Element, 0),
	}
}

func (p *pather) add(n node) {
	p.queue = append(p.queue, n)
}

func (p *pather) remove() node {
	n := p.queue[p.qindex]
	p.qindex++
	if p.qindex > len(p.queue)/2 && p.qindex > 32 {
		p.rebalance()
	}
	return n
}

func (p *pather) rebalance() {
	count := len(p.queue) - p.qindex
	copy(p.queue[:count], p.queue[p.qindex:])
	p.queue, p.qindex = p.queue[:count], 0
}

func (p *pather) empty() bool {
	return len(p.queue) == p.qindex
}

// 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, path Path) []*Element {
	for p.add(node{e, path.segments}); !p.empty(); {
		p.eval(p.remove())
	}
	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.add(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 {
	// If path starts or ends with //, add a .
	if strings.HasPrefix(path, "//") {
		path = "." + path
	}
	if strings.HasSuffix(path, "//") {
		path = path + "."
	}

	// Paths cannot be absolute
	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)
	}

	// Filter contains [@attr='val'] or [tag='val']?
	eqindex := strings.Index(path, "='")
	if eqindex >= 0 {
		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]}
		}
	}

	// Filter contains [@attr], [N] or [tag]
	switch {
	case path[0] == '@':
		return &filterAttr{path[1:]}
	case isInteger(path):
		pos, _ := strconv.Atoi(path)
		if pos == 0 {
			pos = 1 // force to 1-based
		}
		return &filterPos{pos - 1}
	default:
		return &filterChild{path}
	}
}

// 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{}
	for _, c := range e.Child {
		if c, ok := c.(*Element); ok {
			stack.push(c)
		}
	}
	for !stack.empty() {
		e := stack.pop()
		p.candidates = append(p.candidates, e)
		for _, c := range e.Child {
			if c, ok := c.(*Element); ok {
				stack.push(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.scratch = p.scratch[:0]
	if f.index < len(p.candidates) {
		p.scratch = append(p.scratch, p.candidates[f.index])
	}
	p.candidates = p.scratch
}

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

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

// 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.scratch = p.scratch[:0]
	for _, c := range p.candidates {
		if a := c.SelectAttr(f.attr); a != nil && a.Value == f.val {
			p.scratch = append(p.scratch, c)
		}
	}
	p.candidates = p.scratch
}

// 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.scratch = p.scratch[:0]
	for _, c := range p.candidates {
		for _, cc := range c.Child {
			if cc, ok := cc.(*Element); ok && cc.Tag == f.tag {
				p.scratch = append(p.scratch, c)
			}
		}
	}
	p.candidates = p.scratch
}

// 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.scratch = p.scratch[: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.scratch = append(p.scratch, c)
			}
		}
	}
	p.candidates = p.scratch
}
<commit_msg>Bug fixes for paths.<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 Path is an object that represents an optimized version of an
// XPath-like string.
type Path struct {
	segments []segment
}

// NewPath creates an optimized version of an XPath-like string that
// can be used to query elements in an element tree.  Panics if an
// invalid path string is supplied.
func NewPath(path string) Path {
	return Path{parsePath(path)}
}

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

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

// 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)
}

// A pather is helper object used to traverse an element tree with
// a Path object.  It collects and deduplicates elements matching
// the path query.
type pather struct {
	queue      []node
	qindex     int
	results    []*Element
	inResults  map[*Element]bool
	candidates []*Element
	scratch    []*Element // used by filters
}

// 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{
		queue:      make([]node, 0),
		results:    make([]*Element, 0),
		inResults:  make(map[*Element]bool),
		candidates: make([]*Element, 0),
		scratch:    make([]*Element, 0),
	}
}

func (p *pather) add(n node) {
	p.queue = append(p.queue, n)
}

func (p *pather) remove() node {
	n := p.queue[p.qindex]
	p.qindex++
	if p.qindex > len(p.queue)/2 && p.qindex > 32 {
		p.rebalance()
	}
	return n
}

func (p *pather) rebalance() {
	count := len(p.queue) - p.qindex
	copy(p.queue[:count], p.queue[p.qindex:])
	p.queue, p.qindex = p.queue[:count], 0
}

func (p *pather) empty() bool {
	return len(p.queue) == p.qindex
}

// 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, path Path) []*Element {
	for p.add(node{e, path.segments}); !p.empty(); {
		p.eval(p.remove())
	}
	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.add(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 {
	// If path starts or ends with //, fix it
	if strings.HasPrefix(path, "//") {
		path = "." + path
	}
	if strings.HasSuffix(path, "//") {
		path = path + "*"
	}

	// Paths cannot be absolute
	if strings.HasPrefix(path, "/") {
		panic(errPath)
	}

	// Split path into segment objects
	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)
	}

	// Filter contains [@attr='val'] or [tag='val']?
	eqindex := strings.Index(path, "='")
	if eqindex >= 0 {
		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]}
		}
	}

	// Filter contains [@attr], [N] or [tag]
	switch {
	case path[0] == '@':
		return &filterAttr{path[1:]}
	case isInteger(path):
		pos, _ := strconv.Atoi(path)
		if pos == 0 {
			pos = 1 // force to 1-based
		}
		return &filterPos{pos - 1}
	default:
		return &filterChild{path}
	}
}

// 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 {
				stack.push(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.scratch = p.scratch[:0]
	if f.index < len(p.candidates) {
		p.scratch = append(p.scratch, p.candidates[f.index])
	}
	p.candidates = p.scratch
}

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

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

// 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.scratch = p.scratch[:0]
	for _, c := range p.candidates {
		if a := c.SelectAttr(f.attr); a != nil && a.Value == f.val {
			p.scratch = append(p.scratch, c)
		}
	}
	p.candidates = p.scratch
}

// 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.scratch = p.scratch[:0]
	for _, c := range p.candidates {
		for _, cc := range c.Child {
			if cc, ok := cc.(*Element); ok && cc.Tag == f.tag {
				p.scratch = append(p.scratch, c)
			}
		}
	}
	p.candidates = p.scratch
}

// 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.scratch = p.scratch[: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.scratch = append(p.scratch, c)
			}
		}
	}
	p.candidates = p.scratch
}
<|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);
}

void hack_pcap_dump(pcap_dumper_t * dumper, struct pcap_pkthdr *pkt_header,
		      u_char * pkt_data)
{
    pcap_dump((u_char *)dumper, pkt_header, pkt_data);
}
*/
import "C"

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

type Pcap struct {
	cptr *C.pcap_t
}

type PcapDumper struct {
	cptr *C.pcap_dumper_t
}

type pcapError struct{ string }

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 Version() string               { return C.GoString(C.pcap_lib_version()) }
func (p *Pcap) Datalink() int       { return int(C.pcap_datalink(p.cptr)) }
func (e *pcapError) Error() string  { return e.string }
func (p *Pcap) Geterror() error     { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

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

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

	h.cptr = C.pcap_create(dev, buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}

	C.free(unsafe.Pointer(buf))
	return
}

// Set buffer size (units in bytes) on 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
}

//  If arg p is non-zero promiscuous mode will be set on capture handle when it is activated.
func (p *Pcap) SetPromisc(promisc bool) error {
	var pro int32
	if promisc {
		pro = 1
	}

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

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

// Set read timeout (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, 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 = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Openoffline
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 = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Pcap closes a handler.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	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_ptr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	pkthdr = *pkthdr_ptr

	if nil == buf {
		return
	}
	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000) // pcap provides usec but time.Unix requires nsec
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = make([]byte, pkthdr.caplen)

	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return
}

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 (p *Pcap) SetDirection(direction string) (err error) {
	var pcap_direction C.pcap_direction_t
	if direction == "in" {
		pcap_direction = C.PCAP_D_IN
	} else if direction == "out" {
		pcap_direction = C.PCAP_D_OUT
	} else {
		pcap_direction = C.PCAP_D_INOUT
	}
	if -1 == C.pcap_setdirection(p.cptr, pcap_direction) {
		return p.Geterror()
	}
	return nil
}

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) {
		if curaddr.addr == nil {
			continue
		}
		var a IFAddress
		var err error
		if a.IP, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.netmask))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddrToIP(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
}

// Inject ...
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
}

func (p *Pcap) DumpOpen(ofile *string) (dumper *PcapDumper, err error) {
	d := new(PcapDumper)
	d.cptr = C.pcap_dump_open(p.cptr, C.CString(*ofile))
	if nil == d.cptr {
		return d, errors.New("Cannot open dumpfile")
	}
	dumper = d
	return
}

func (p *Pcap) PcapLoop(pktnum int, dumper *PcapDumper) (result int32, err error) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	var buf_ptr *C.u_char

	for i := 0; true; {
		result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr_ptr, &buf_ptr))
		switch result {
		case 0:
			continue // timeout
		case 1:
			// success : capturing packet
		case -1:
			return result, errors.New("Error in pcap next ex")
		case -2:
			return // reach EOF in offline mode
		}
		if nil == buf_ptr {
			continue
		}
		if nil != dumper {
			p.PcapDump(dumper, pkthdr_ptr, buf_ptr)
			p.PcapDumpFlush(dumper)
		}
		if pktnum > 0 {
			i++
			if i >= pktnum {
				break
			}
		}
	}
	return
}

func (p *Pcap) PcapDump(dumper *PcapDumper, pkthdr_ptr *C.struct_pcap_pkthdr, buf_ptr *C.u_char) {
	C.hack_pcap_dump(dumper.cptr, pkthdr_ptr, buf_ptr)
}

func (p *Pcap) PcapDumpFlush(dumper *PcapDumper) error {
	if -1 == C.pcap_dump_flush(dumper.cptr) {
		return p.Geterror()
	}
	return nil
}

func (p *Pcap) PcapDumpClose(dumper *PcapDumper) {
	C.pcap_dump_close(dumper.cptr)
}
<commit_msg>Statically compiling libpcap.a instead of dynamic linking<commit_after>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: /usr/lib/x86_64-linux-gnu/libpcap.a
#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);
}

void hack_pcap_dump(pcap_dumper_t * dumper, struct pcap_pkthdr *pkt_header,
		      u_char * pkt_data)
{
    pcap_dump((u_char *)dumper, pkt_header, pkt_data);
}
*/
import "C"

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

type Pcap struct {
	cptr *C.pcap_t
}

type PcapDumper struct {
	cptr *C.pcap_dumper_t
}

type pcapError struct{ string }

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 Version() string               { return C.GoString(C.pcap_lib_version()) }
func (p *Pcap) Datalink() int       { return int(C.pcap_datalink(p.cptr)) }
func (e *pcapError) Error() string  { return e.string }
func (p *Pcap) Geterror() error     { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

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

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

	h.cptr = C.pcap_create(dev, buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}

	C.free(unsafe.Pointer(buf))
	return
}

// Set buffer size (units in bytes) on 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
}

//  If arg p is non-zero promiscuous mode will be set on capture handle when it is activated.
func (p *Pcap) SetPromisc(promisc bool) error {
	var pro int32
	if promisc {
		pro = 1
	}

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

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

// Set read timeout (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, 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 = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Openoffline
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 = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Pcap closes a handler.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	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_ptr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	pkthdr = *pkthdr_ptr

	if nil == buf {
		return
	}
	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000) // pcap provides usec but time.Unix requires nsec
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = make([]byte, pkthdr.caplen)

	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return
}

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 (p *Pcap) SetDirection(direction string) (err error) {
	var pcap_direction C.pcap_direction_t
	if direction == "in" {
		pcap_direction = C.PCAP_D_IN
	} else if direction == "out" {
		pcap_direction = C.PCAP_D_OUT
	} else {
		pcap_direction = C.PCAP_D_INOUT
	}
	if -1 == C.pcap_setdirection(p.cptr, pcap_direction) {
		return p.Geterror()
	}
	return nil
}

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) {
		if curaddr.addr == nil {
			continue
		}
		var a IFAddress
		var err error
		if a.IP, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.netmask))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddrToIP(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
}

// Inject ...
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
}

func (p *Pcap) DumpOpen(ofile *string) (dumper *PcapDumper, err error) {
	d := new(PcapDumper)
	d.cptr = C.pcap_dump_open(p.cptr, C.CString(*ofile))
	if nil == d.cptr {
		return d, errors.New("Cannot open dumpfile")
	}
	dumper = d
	return
}

func (p *Pcap) PcapLoop(pktnum int, dumper *PcapDumper) (result int32, err error) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	var buf_ptr *C.u_char

	for i := 0; true; {
		result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr_ptr, &buf_ptr))
		switch result {
		case 0:
			continue // timeout
		case 1:
			// success : capturing packet
		case -1:
			return result, errors.New("Error in pcap next ex")
		case -2:
			return // reach EOF in offline mode
		}
		if nil == buf_ptr {
			continue
		}
		if nil != dumper {
			p.PcapDump(dumper, pkthdr_ptr, buf_ptr)
			p.PcapDumpFlush(dumper)
		}
		if pktnum > 0 {
			i++
			if i >= pktnum {
				break
			}
		}
	}
	return
}

func (p *Pcap) PcapDump(dumper *PcapDumper, pkthdr_ptr *C.struct_pcap_pkthdr, buf_ptr *C.u_char) {
	C.hack_pcap_dump(dumper.cptr, pkthdr_ptr, buf_ptr)
}

func (p *Pcap) PcapDumpFlush(dumper *PcapDumper) error {
	if -1 == C.pcap_dump_flush(dumper.cptr) {
		return p.Geterror()
	}
	return nil
}

func (p *Pcap) PcapDumpClose(dumper *PcapDumper) {
	C.pcap_dump_close(dumper.cptr)
}
<|endoftext|>


<commit_before>package capn_test

import (
	"bytes"

	capn "github.com/glycerine/go-capnproto"
)

func zdateFilledSegment(n int, packed bool) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	z := NewRootZ(seg)
	list := NewZdateList(seg, n)
	plist := capn.PointerList(list)

	for i := 0; i < n; i++ {
		d := NewZdate(seg)
		d.SetMonth(12)
		d.SetDay(7)
		d.SetYear(int16(2004 + i))
		plist.Set(i, capn.Object(d))
	}
	z.SetZdatevec(list)

	buf := bytes.Buffer{}
	if packed {
		seg.WriteToPacked(&buf)
	} else {
		seg.WriteTo(&buf)
	}
	return seg, buf.Bytes()
}

func zdateReader(n int, packed bool) *bytes.Reader {
	_, byteSlice := zdateFilledSegment(n, packed)
	return bytes.NewReader(byteSlice)
}

func zdataFilledSegment(n int) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	d := NewRootZdata(seg)
	b := make([]byte, n)
	for i := 0; i < len(b); i++ {
		b[i] = byte(i)
	}
	d.SetData(b)
	buf := bytes.Buffer{}
	seg.WriteTo(&buf)
	return seg, buf.Bytes()
}

func zdataReader(n int) *bytes.Reader {
	_, byteSlice := zdataFilledSegment(n)
	return bytes.NewReader(byteSlice)
}
<commit_msg>zdateFilledSegment() now fills a ZdateList with n Zdate, as intended<commit_after>package capn_test

import (
	"bytes"

	capn "github.com/glycerine/go-capnproto"
)

func zdateFilledSegment(n int, packed bool) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	z := NewRootZ(seg)
	list := NewZdateList(seg, n)
	// hand added a Set() method to messages_test.go, so plist not needed
	// plist := capn.PointerList(list)

	for i := 0; i < n; i++ {
		d := NewZdate(seg)
		d.SetMonth(12)
		d.SetDay(7)
		d.SetYear(int16(2004 + i))
		//plist.Set(i, capn.Object(d))
		list.Set(i, d)
	}
	z.SetZdatevec(list)

	buf := bytes.Buffer{}
	if packed {
		seg.WriteToPacked(&buf)
	} else {
		seg.WriteTo(&buf)
	}
	return seg, buf.Bytes()
}

func zdateReader(n int, packed bool) *bytes.Reader {
	_, byteSlice := zdateFilledSegment(n, packed)
	return bytes.NewReader(byteSlice)
}

func zdataFilledSegment(n int) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	d := NewRootZdata(seg)
	b := make([]byte, n)
	for i := 0; i < len(b); i++ {
		b[i] = byte(i)
	}
	d.SetData(b)
	buf := bytes.Buffer{}
	seg.WriteTo(&buf)
	return seg, buf.Bytes()
}

func zdataReader(n int) *bytes.Reader {
	_, byteSlice := zdataFilledSegment(n)
	return bytes.NewReader(byteSlice)
}
<|endoftext|>


<commit_before>package herd

import (
	"fmt"
	"io"
	"sort"
	"time"
)

type rDuration time.Duration

func (d rDuration) selector(sub *Sub) bool {
	if time.Since(sub.lastReachableTime) <= time.Duration(d) {
		return true
	}
	return false
}

func (herd *Herd) writeHtml(writer io.Writer) {
	if herd.updatesDisabledReason != "" {
		herd.writeDisableStatus(writer)
		fmt.Fprintln(writer, "
")
	}
	numSubs := herd.countSelectedSubs(nil)
	fmt.Fprintf(writer, "Time since current cycle start: %s
\n",
		time.Since(herd.currentScanStartTime))
	if numSubs < 1 {
		fmt.Fprintf(writer, "Duration of previous cycle: %s
\n",
			herd.previousScanDuration)
	} else {
		fmt.Fprintf(writer, "Duration of previous cycle: %s (%s/sub)
\n",
			herd.previousScanDuration,
			herd.previousScanDuration/time.Duration(numSubs))
	}
	fmt.Fprintf(writer, "Image server: %s
\n",
		herd.imageManager, herd.imageManager)
	if herd.defaultImageName != "" {
		fmt.Fprintf(writer,
			"Default image: %s
\n",
			herd.imageManager, herd.defaultImageName, herd.defaultImageName)
	}
	fmt.Fprintf(writer,
		"Number of subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectAliveSub)
	fmt.Fprintf(writer,
		"Number of alive subs: %d
\n",
		numSubs)
	fmt.Fprint(writer, "Number of reachable subs in last: ")
	herd.writeReachableSubsLink(writer, time.Minute, "1 min", "1m", true)
	herd.writeReachableSubsLink(writer, time.Minute*10, "10 min", "10m", true)
	herd.writeReachableSubsLink(writer, time.Hour, "1 hour", "1h", true)
	herd.writeReachableSubsLink(writer, time.Hour*24, "1 day", "1d", true)
	herd.writeReachableSubsLink(writer, time.Hour*24*7, "1 week", "1w", false)
	numSubs = herd.countSelectedSubs(selectDeviantSub)
	fmt.Fprintf(writer,
		"Number of deviant subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectCompliantSub)
	fmt.Fprintf(writer,
		"Number of compliant subs: %d
\n",
		numSubs)
	subs := herd.getSelectedSubs(nil)
	connectDurations := getConnectDurations(subs)
	shortPollDurations := getPollDurations(subs, false)
	fullPollDurations := getPollDurations(subs, true)
	showDurationStats(writer, connectDurations, "Connect")
	showDurationStats(writer, shortPollDurations, "Short poll")
	showDurationStats(writer, fullPollDurations, "Full poll")
	// TODO(rgooch): Figure out a way of restoring this information.
	//fmt.Fprintf(writer, "Connection slots: %d out of %d
\n",
	//len(herd.connectionSemaphore), cap(herd.connectionSemaphore))
	fmt.Fprintf(writer, "Poll slots: %d out of %d
\n",
		len(herd.pollSemaphore), cap(herd.pollSemaphore))
}

func (herd *Herd) writeDisableStatus(writer io.Writer) {
	timeString := herd.updatesDisabledTime.Format("02 Jan 2006 15:04:05.99 MST")
	if herd.updatesDisabledBy == "" {
		fmt.Fprintf(writer,
			"Updates disabled because: %s at %s\n",
			herd.updatesDisabledReason, timeString)
	} else {
		fmt.Fprintf(writer,
			"Updates disabled by: %s because: %s at %s",
			herd.updatesDisabledBy, herd.updatesDisabledReason, timeString)
	}
}

func (herd *Herd) writeReachableSubsLink(writer io.Writer,
	duration time.Duration, durationString string, query string,
	moreToCome bool) {
	numSubs := herd.countSelectedSubs(rDuration(duration).selector)
	fmt.Fprintf(writer, "%s",
		query, durationString)
	fmt.Fprintf(writer, "(%d)",
		query, numSubs)
	if moreToCome {
		fmt.Fprint(writer, ", ")
	} else {
		fmt.Fprintln(writer, "
")
	}
}

func selectAliveSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusUnknown:
		return false
	case statusConnecting:
		return false
	case statusDNSError:
		return false
	case statusConnectionRefused:
		return false
	case statusFailedToConnect:
		return false
	case statusFailedToPoll:
		return false
	}
	return true
}

func selectDeviantSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusComputingUpdate:
		return true
	case statusSendingUpdate:
		return true
	case statusUpdatesDisabled:
		return true
	case statusUpdating:
		return true
	case statusUpdateDenied:
		return true
	case statusFailedToUpdate:
		return true
	}
	return false
}

func selectCompliantSub(sub *Sub) bool {
	if sub.publishedStatus == statusSynced {
		return true
	}
	return false
}

func getConnectDurations(subs []*Sub) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		if sub.lastConnectDuration > 0 {
			durations = append(durations, sub.lastConnectDuration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

func getPollDurations(subs []*Sub, full bool) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		var duration time.Duration
		if full {
			duration = sub.lastFullPollDuration
		} else {
			duration = sub.lastShortPollDuration
		}
		if duration > 0 {
			durations = append(durations, duration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

type durationList []time.Duration

func (dl durationList) Len() int {
	return len(dl)
}

func (dl durationList) Less(i, j int) bool {
	return dl[i] < dl[j]
}

func (dl durationList) Swap(i, j int) {
	dl[i], dl[j] = dl[j], dl[i]
}

func showDurationStats(writer io.Writer, durations []time.Duration,
	durationType string) {
	if len(durations) < 1 {
		return
	}
	var avgDuration time.Duration
	for _, duration := range durations {
		avgDuration += duration
	}
	avgDuration /= time.Duration(len(durations))
	medDuration := durations[len(durations)/2]
	unit := "ns"
	scale := 1.0
	switch {
	case medDuration > 1e9:
		unit = "s"
		scale = 1e-9
	case medDuration > 1e6:
		unit = "ms"
		scale = 1e-6
	case medDuration > 1e3:
		unit = "µs"
		scale = 1e-3
	}
	fmt.Fprintf(writer,
		"%s durations: %.3f/%.3f/%.3f/%.3f %s (avg/med/min/max)
\n",
		durationType,
		float64(avgDuration)*scale, float64(medDuration)*scale,
		float64(durations[0])*scale, float64(durations[len(durations)-1])*scale,
		unit)
}
<commit_msg>Do not count unroutable or timed out subs as "alive" in subs status page.<commit_after>package herd

import (
	"fmt"
	"io"
	"sort"
	"time"
)

type rDuration time.Duration

func (d rDuration) selector(sub *Sub) bool {
	if time.Since(sub.lastReachableTime) <= time.Duration(d) {
		return true
	}
	return false
}

func (herd *Herd) writeHtml(writer io.Writer) {
	if herd.updatesDisabledReason != "" {
		herd.writeDisableStatus(writer)
		fmt.Fprintln(writer, "
")
	}
	numSubs := herd.countSelectedSubs(nil)
	fmt.Fprintf(writer, "Time since current cycle start: %s
\n",
		time.Since(herd.currentScanStartTime))
	if numSubs < 1 {
		fmt.Fprintf(writer, "Duration of previous cycle: %s
\n",
			herd.previousScanDuration)
	} else {
		fmt.Fprintf(writer, "Duration of previous cycle: %s (%s/sub)
\n",
			herd.previousScanDuration,
			herd.previousScanDuration/time.Duration(numSubs))
	}
	fmt.Fprintf(writer, "Image server: %s
\n",
		herd.imageManager, herd.imageManager)
	if herd.defaultImageName != "" {
		fmt.Fprintf(writer,
			"Default image: %s
\n",
			herd.imageManager, herd.defaultImageName, herd.defaultImageName)
	}
	fmt.Fprintf(writer,
		"Number of subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectAliveSub)
	fmt.Fprintf(writer,
		"Number of alive subs: %d
\n",
		numSubs)
	fmt.Fprint(writer, "Number of reachable subs in last: ")
	herd.writeReachableSubsLink(writer, time.Minute, "1 min", "1m", true)
	herd.writeReachableSubsLink(writer, time.Minute*10, "10 min", "10m", true)
	herd.writeReachableSubsLink(writer, time.Hour, "1 hour", "1h", true)
	herd.writeReachableSubsLink(writer, time.Hour*24, "1 day", "1d", true)
	herd.writeReachableSubsLink(writer, time.Hour*24*7, "1 week", "1w", false)
	numSubs = herd.countSelectedSubs(selectDeviantSub)
	fmt.Fprintf(writer,
		"Number of deviant subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectCompliantSub)
	fmt.Fprintf(writer,
		"Number of compliant subs: %d
\n",
		numSubs)
	subs := herd.getSelectedSubs(nil)
	connectDurations := getConnectDurations(subs)
	shortPollDurations := getPollDurations(subs, false)
	fullPollDurations := getPollDurations(subs, true)
	showDurationStats(writer, connectDurations, "Connect")
	showDurationStats(writer, shortPollDurations, "Short poll")
	showDurationStats(writer, fullPollDurations, "Full poll")
	// TODO(rgooch): Figure out a way of restoring this information.
	//fmt.Fprintf(writer, "Connection slots: %d out of %d
\n",
	//len(herd.connectionSemaphore), cap(herd.connectionSemaphore))
	fmt.Fprintf(writer, "Poll slots: %d out of %d
\n",
		len(herd.pollSemaphore), cap(herd.pollSemaphore))
}

func (herd *Herd) writeDisableStatus(writer io.Writer) {
	timeString := herd.updatesDisabledTime.Format("02 Jan 2006 15:04:05.99 MST")
	if herd.updatesDisabledBy == "" {
		fmt.Fprintf(writer,
			"Updates disabled because: %s at %s\n",
			herd.updatesDisabledReason, timeString)
	} else {
		fmt.Fprintf(writer,
			"Updates disabled by: %s because: %s at %s",
			herd.updatesDisabledBy, herd.updatesDisabledReason, timeString)
	}
}

func (herd *Herd) writeReachableSubsLink(writer io.Writer,
	duration time.Duration, durationString string, query string,
	moreToCome bool) {
	numSubs := herd.countSelectedSubs(rDuration(duration).selector)
	fmt.Fprintf(writer, "%s",
		query, durationString)
	fmt.Fprintf(writer, "(%d)",
		query, numSubs)
	if moreToCome {
		fmt.Fprint(writer, ", ")
	} else {
		fmt.Fprintln(writer, "
")
	}
}

func selectAliveSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusUnknown:
		return false
	case statusConnecting:
		return false
	case statusDNSError:
		return false
	case statusConnectionRefused:
		return false
	case statusNoRouteToHost:
		return false
	case statusConnectTimeout:
		return false
	case statusFailedToConnect:
		return false
	case statusFailedToPoll:
		return false
	}
	return true
}

func selectDeviantSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusComputingUpdate:
		return true
	case statusSendingUpdate:
		return true
	case statusUpdatesDisabled:
		return true
	case statusUpdating:
		return true
	case statusUpdateDenied:
		return true
	case statusFailedToUpdate:
		return true
	}
	return false
}

func selectCompliantSub(sub *Sub) bool {
	if sub.publishedStatus == statusSynced {
		return true
	}
	return false
}

func getConnectDurations(subs []*Sub) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		if sub.lastConnectDuration > 0 {
			durations = append(durations, sub.lastConnectDuration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

func getPollDurations(subs []*Sub, full bool) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		var duration time.Duration
		if full {
			duration = sub.lastFullPollDuration
		} else {
			duration = sub.lastShortPollDuration
		}
		if duration > 0 {
			durations = append(durations, duration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

type durationList []time.Duration

func (dl durationList) Len() int {
	return len(dl)
}

func (dl durationList) Less(i, j int) bool {
	return dl[i] < dl[j]
}

func (dl durationList) Swap(i, j int) {
	dl[i], dl[j] = dl[j], dl[i]
}

func showDurationStats(writer io.Writer, durations []time.Duration,
	durationType string) {
	if len(durations) < 1 {
		return
	}
	var avgDuration time.Duration
	for _, duration := range durations {
		avgDuration += duration
	}
	avgDuration /= time.Duration(len(durations))
	medDuration := durations[len(durations)/2]
	unit := "ns"
	scale := 1.0
	switch {
	case medDuration > 1e9:
		unit = "s"
		scale = 1e-9
	case medDuration > 1e6:
		unit = "ms"
		scale = 1e-6
	case medDuration > 1e3:
		unit = "µs"
		scale = 1e-3
	}
	fmt.Fprintf(writer,
		"%s durations: %.3f/%.3f/%.3f/%.3f %s (avg/med/min/max)
\n",
		durationType,
		float64(avgDuration)*scale, float64(medDuration)*scale,
		float64(durations[0])*scale, float64(durations[len(durations)-1])*scale,
		unit)
}
<|endoftext|>


<commit_before>package bakery_test

import (
	"fmt"

	gc "gopkg.in/check.v1"

	"gopkg.in/macaroon-bakery.v0/bakery"
	"gopkg.in/macaroon.v1"
)

type ServiceSuite struct{}

var _ = gc.Suite(&ServiceSuite{})

// TestSingleServiceFirstParty creates a single service
// with a macaroon with one first party caveat.
// It creates a request with this macaroon and checks that the service
// can verify this macaroon as valid.
func (s *ServiceSuite) TestSingleServiceFirstParty(c *gc.C) {
	p := bakery.NewServiceParams{
		Location: "loc",
		Store:    nil,
		Key:      nil,
		Locator:  nil,
	}
	service, err := bakery.NewService(p)
	c.Assert(err, gc.IsNil)

	macaroon, err := service.NewMacaroon("", nil, nil)
	c.Assert(err, gc.IsNil)
	c.Assert(macaroon.Location(), gc.Equals, "loc")
	cav := bakery.Caveat{
		Location:  "",
		Condition: "something",
	}
	err = service.AddCaveat(macaroon, cav)
	c.Assert(err, gc.IsNil)

	checker := strCompFirstPartyChecker("something")
	req := service.NewRequest(checker)

	req.AddClientMacaroon(macaroon)

	err = req.Check()
	c.Assert(err, gc.IsNil)
}

// TestMacaroonPaperFig6 implements an example flow as described in the macaroons paper:
// http://theory.stanford.edu/~ataly/Papers/macaroons.pdf
// There are three services, ts, fs, as:
// ts is a storage service which has deligated authority to a forum service fs.
// The forum service wants to require its users to be logged into to an authentication service as.
//
// The client obtains a macaroon from fs (minted by ts, with a third party caveat addressed to as).
// The client obtains a discharge macaroon from as to satisfy this caveat.
// The target service verifies the original macaroon it delegated to fs
// No direct contact between as and ts is required
func (s *ServiceSuite) TestMacaroonPaperFig6(c *gc.C) {
	ts, fs, as := setUpFig6Services(c)
	tsMacaroon := createMacaroonWithThirdPartyCaveat(c, ts, fs, bakery.Caveat{Location: "as-loc", Condition: "user==bob"})
	d, err := bakery.DischargeAll(tsMacaroon, func(firstPartyLocation string, cav macaroon.Caveat) (*macaroon.Macaroon, error) {
		c.Assert(firstPartyLocation, gc.Equals, "ts-loc")
		c.Assert(cav.Location, gc.Equals, "as-loc")
		mac, err := as.Discharge(strCompThirdPartyChecker("user==bob"), cav.Id)
		c.Assert(err, gc.IsNil)
		return mac, nil
	})
	c.Assert(err, gc.IsNil)

	// client makes request to ts
	req := ts.NewRequest(strCompFirstPartyChecker(""))
	req.AddClientMacaroon(tsMacaroon)
	// client has all the discharge macaroons. For each discharge macaroon bind it to our tsMacaroon
	// and add it to our request.
	for _, dm := range d {
		dm.Bind(tsMacaroon.Signature())
		req.AddClientMacaroon(dm)
	}

	err = req.Check()
	c.Assert(err, gc.IsNil)
}

func (s *ServiceSuite) TestMacaroonPaperFig6FailsWithoutDischarges(c *gc.C) {
	ts, fs, _ := setUpFig6Services(c)
	tsMacaroon := createMacaroonWithThirdPartyCaveat(c, ts, fs, bakery.Caveat{Location: "as-loc", Condition: "user==bob"})

	req := ts.NewRequest(strCompFirstPartyChecker(""))
	req.AddClientMacaroon(tsMacaroon)

	err := req.Check()
	c.Assert(err, gc.ErrorMatches, `verification failed: cannot find discharge macaroon for caveat ".*"`)
}

func (s *ServiceSuite) TestMacaroonPaperFig6FailsWithBindingOnTamperedSignature(c *gc.C) {
	ts, fs, as := setUpFig6Services(c)
	tsMacaroon := createMacaroonWithThirdPartyCaveat(c, ts, fs, bakery.Caveat{Location: "as-loc", Condition: "user==bob"})
	d, err := bakery.DischargeAll(tsMacaroon, func(firstPartyLocation string, cav macaroon.Caveat) (*macaroon.Macaroon, error) {
		c.Assert(firstPartyLocation, gc.Equals, "ts-loc")
		c.Assert(cav.Location, gc.Equals, "as-loc")
		mac, err := as.Discharge(strCompThirdPartyChecker("user==bob"), cav.Id)
		c.Assert(err, gc.IsNil)
		return mac, nil
	})
	c.Assert(err, gc.IsNil)

	req := ts.NewRequest(strCompFirstPartyChecker(""))
	req.AddClientMacaroon(tsMacaroon)
	for _, dm := range d {
		dm.Bind([]byte("tampered-signature"))
		req.AddClientMacaroon(dm)
	}

	err = req.Check()
	c.Assert(err, gc.ErrorMatches, "verification failed: signature mismatch after caveat verification")
}

func setUpFig6Services(c *gc.C) (*bakery.Service, *bakery.Service, *bakery.Service) {
	fsKeyPair, err := bakery.GenerateKey()
	c.Assert(err, gc.IsNil)
	asKeyPair, err := bakery.GenerateKey()
	c.Assert(err, gc.IsNil)

	publicKeyLocator := bakery.PublicKeyLocatorMap{
		"fs-loc": &fsKeyPair.Public,
		"as-loc": &asKeyPair.Public,
	}

	ts, err := bakery.NewService(bakery.NewServiceParams{Location: "ts-loc"})
	c.Assert(err, gc.IsNil)
	fs, err := bakery.NewService(bakery.NewServiceParams{
		Location: "fs-loc",
		Key:      fsKeyPair,
		Locator:  publicKeyLocator,
	})
	c.Assert(err, gc.IsNil)
	as, err := bakery.NewService(bakery.NewServiceParams{
		Location: "as-loc",
		Key:      asKeyPair,
		Locator:  publicKeyLocator,
	})
	c.Assert(err, gc.IsNil)

	return ts, fs, as
}

func createMacaroonWithThirdPartyCaveat(c *gc.C, minter, caveater *bakery.Service, cav bakery.Caveat) *macaroon.Macaroon {
	mac, err := minter.NewMacaroon("", nil, nil)
	c.Assert(err, gc.IsNil)

	err = caveater.AddCaveat(mac, cav)
	c.Assert(err, gc.IsNil)
	return mac
}

type strCompFirstPartyChecker string

func (c strCompFirstPartyChecker) CheckFirstPartyCaveat(caveat string) error {
	if caveat != string(c) {
		return fmt.Errorf("%v doesn't match %s", caveat, c)
	}
	return nil
}

type strCompThirdPartyChecker string

func (c strCompThirdPartyChecker) CheckThirdPartyCaveat(caveatId string, caveat string) ([]bakery.Caveat, error) {
	if caveat != string(c) {
		return nil, fmt.Errorf("%v doesn't match %s", caveat, c)
	}
	return nil, nil
}
<commit_msg>make the addservice function more generic, added test comments<commit_after>package bakery_test

import (
	"fmt"

	gc "gopkg.in/check.v1"

	"gopkg.in/macaroon-bakery.v0/bakery"
	"gopkg.in/macaroon.v1"
)

type ServiceSuite struct{}

var _ = gc.Suite(&ServiceSuite{})

// TestSingleServiceFirstParty creates a single service
// with a macaroon with one first party caveat.
// It creates a request with this macaroon and checks that the service
// can verify this macaroon as valid.
func (s *ServiceSuite) TestSingleServiceFirstParty(c *gc.C) {
	p := bakery.NewServiceParams{
		Location: "loc",
		Store:    nil,
		Key:      nil,
		Locator:  nil,
	}
	service, err := bakery.NewService(p)
	c.Assert(err, gc.IsNil)

	macaroon, err := service.NewMacaroon("", nil, nil)
	c.Assert(err, gc.IsNil)
	c.Assert(macaroon.Location(), gc.Equals, "loc")
	cav := bakery.Caveat{
		Location:  "",
		Condition: "something",
	}
	err = service.AddCaveat(macaroon, cav)
	c.Assert(err, gc.IsNil)

	checker := strCompFirstPartyChecker("something")
	req := service.NewRequest(checker)

	req.AddClientMacaroon(macaroon)

	err = req.Check()
	c.Assert(err, gc.IsNil)
}

// TestMacaroonPaperFig6 implements an example flow as described in the macaroons paper:
// http://theory.stanford.edu/~ataly/Papers/macaroons.pdf
// There are three services, ts, fs, as:
// ts is a storage service which has deligated authority to a forum service fs.
// The forum service wants to require its users to be logged into to an authentication service as.
//
// The client obtains a macaroon from fs (minted by ts, with a third party caveat addressed to as).
// The client obtains a discharge macaroon from as to satisfy this caveat.
// The target service verifies the original macaroon it delegated to fs
// No direct contact between as and ts is required
func (s *ServiceSuite) TestMacaroonPaperFig6(c *gc.C) {
	locator := make(bakery.PublicKeyLocatorMap)
	as := newService(c, "as-loc", locator)
	ts := newService(c, "ts-loc", locator)
	fs := newService(c, "fs-loc", locator)

	// ts creates a macaroon.
	// ts somehow sends the macaroon to fs which adds a third party caveat to be discharged by as.
	tsMacaroon := createMacaroonWithThirdPartyCaveat(c, ts, fs, bakery.Caveat{Location: "as-loc", Condition: "user==bob"})

	// client asks for a discharge macaroon for each third party caveat
	d, err := bakery.DischargeAll(tsMacaroon, func(firstPartyLocation string, cav macaroon.Caveat) (*macaroon.Macaroon, error) {
		c.Assert(firstPartyLocation, gc.Equals, "ts-loc")
		c.Assert(cav.Location, gc.Equals, "as-loc")
		mac, err := as.Discharge(strCompThirdPartyChecker("user==bob"), cav.Id)
		c.Assert(err, gc.IsNil)
		return mac, nil
	})
	c.Assert(err, gc.IsNil)

	// client makes request to ts
	req := ts.NewRequest(strCompFirstPartyChecker(""))
	req.AddClientMacaroon(tsMacaroon)
	// client has all the discharge macaroons. For each discharge macaroon bind it to our tsMacaroon
	// and add it to our request.
	for _, dm := range d {
		dm.Bind(tsMacaroon.Signature())
		req.AddClientMacaroon(dm)
	}

	err = req.Check()
	c.Assert(err, gc.IsNil)
}

// TestMacaroonPaperFig6FailsWithoutDischarges runs a similar test as TestMacaroonPaperFig6
// without the client discharging the third party caveats.
func (s *ServiceSuite) TestMacaroonPaperFig6FailsWithoutDischarges(c *gc.C) {
	locator := make(bakery.PublicKeyLocatorMap)
	ts := newService(c, "ts-loc", locator)
	fs := newService(c, "fs-loc", locator)
	_ = newService(c, "as-loc", locator)

	// ts creates a macaroon.
	// ts somehow sends the macaroon to fs which adds a third party caveat to be discharged by as.
	tsMacaroon := createMacaroonWithThirdPartyCaveat(c, ts, fs, bakery.Caveat{Location: "as-loc", Condition: "user==bob"})

	// client makes request to ts
	req := ts.NewRequest(strCompFirstPartyChecker(""))
	req.AddClientMacaroon(tsMacaroon)

	err := req.Check()
	c.Assert(err, gc.ErrorMatches, `verification failed: cannot find discharge macaroon for caveat ".*"`)
}

// TestMacaroonPaperFig6FailsWithBindingOnTamperedSignature runs a similar test as TestMacaroonPaperFig6
// with the discharge macaroon binding being done on a tampered signature.
func (s *ServiceSuite) TestMacaroonPaperFig6FailsWithBindingOnTamperedSignature(c *gc.C) {
	locator := make(bakery.PublicKeyLocatorMap)
	as := newService(c, "as-loc", locator)
	ts := newService(c, "ts-loc", locator)
	fs := newService(c, "fs-loc", locator)

	// ts creates a macaroon.
	// ts somehow sends the macaroon to fs which adds a third party caveat to be discharged by as.
	tsMacaroon := createMacaroonWithThirdPartyCaveat(c, ts, fs, bakery.Caveat{Location: "as-loc", Condition: "user==bob"})

	// client asks for a discharge macaroon for each third party caveat
	d, err := bakery.DischargeAll(tsMacaroon, func(firstPartyLocation string, cav macaroon.Caveat) (*macaroon.Macaroon, error) {
		c.Assert(firstPartyLocation, gc.Equals, "ts-loc")
		c.Assert(cav.Location, gc.Equals, "as-loc")
		mac, err := as.Discharge(strCompThirdPartyChecker("user==bob"), cav.Id)
		c.Assert(err, gc.IsNil)
		return mac, nil
	})
	c.Assert(err, gc.IsNil)

	// client makes request to ts
	req := ts.NewRequest(strCompFirstPartyChecker(""))
	req.AddClientMacaroon(tsMacaroon)

	// client has all the discharge macaroons. For each discharge macaroon bind it to our tsMacaroon
	// and add it to our request.
	for _, dm := range d {
		dm.Bind([]byte("tampered-signature")) // Bind against an incorrect signature.
		req.AddClientMacaroon(dm)
	}

	err = req.Check()
	c.Assert(err, gc.ErrorMatches, "verification failed: signature mismatch after caveat verification")
}

func newService(c *gc.C, location string, locator bakery.PublicKeyLocatorMap) *bakery.Service {
	keyPair, err := bakery.GenerateKey()
	c.Assert(err, gc.IsNil)

	svc, err := bakery.NewService(bakery.NewServiceParams{
		Location: location,
		Key:      keyPair,
		Locator:  locator,
	})
	c.Assert(err, gc.IsNil)
	if locator != nil {
		locator[location] = &keyPair.Public
	}
	return svc
}

func createMacaroonWithThirdPartyCaveat(c *gc.C, minter, caveater *bakery.Service, cav bakery.Caveat) *macaroon.Macaroon {
	mac, err := minter.NewMacaroon("", nil, nil)
	c.Assert(err, gc.IsNil)

	err = caveater.AddCaveat(mac, cav)
	c.Assert(err, gc.IsNil)
	return mac
}

type strCompFirstPartyChecker string

func (c strCompFirstPartyChecker) CheckFirstPartyCaveat(caveat string) error {
	if caveat != string(c) {
		return fmt.Errorf("%v doesn't match %s", caveat, c)
	}
	return nil
}

type strCompThirdPartyChecker string

func (c strCompThirdPartyChecker) CheckThirdPartyCaveat(caveatId string, caveat string) ([]bakery.Caveat, error) {
	if caveat != string(c) {
		return nil, fmt.Errorf("%v doesn't match %s", caveat, c)
	}
	return nil, nil
}
<|endoftext|>


<commit_before>package mysql

import (
	"fmt"
	"strings"
	"database/sql"
	"github.com/akinsella/go-playground/models"
	"github.com/akinsella/go-playground/tasks"
	_ "github.com/go-sql-driver/mysql"
)

type MySQL struct {}

type MySQLStopTimesImportTask struct {
	tasks.ImportTask
}

func (m *MySQLStopTimesImportTask) DoWork(workRoutine int) {
	m.InsertStopTimes(insertStopTimes);
}

func insertStopTimes(sts *models.StopTimes) (error)  {

	db, err := sql.Open("mysql", "gtfs:gtfs@/gtfs?charset=utf8mb4,utf8");

	if err != nil {
		panic(err.Error())
	}

	defer db.Close()

	valueStrings := make([]string, 0, len(sts.Records))
	valueArgs := make([]interface{}, 0, len(sts.Records) * 9)

	for _, st := range sts.Records {
		valueStrings = append(valueStrings, "('RATP', ?, ?, ?, ?, ?, ?, ?, ?)")
		valueArgs = append(
			valueArgs,
			st.TripId,
			st.ArrivalTime,
			st.DepartureTime,
			st.StopId,
			st.StopSequence,
			st.StopHeadSign,
			st.PickupType,
			st.DropOffType,
		)
	}

	stmt := fmt.Sprintf("INSERT INTO stop_times (agency_key, trip_id, arrival_time, departure_time, stop_id, stop_sequence, stop_head_sign, pickup_type, drop_off_type) VALUES %s", strings.Join(valueStrings, ","))

	_, err = db.Exec(stmt, valueArgs...)

	return err
}
<commit_msg>Removed unused MySQL structure<commit_after>package mysql

import (
	"fmt"
	"strings"
	"database/sql"
	"github.com/akinsella/go-playground/models"
	"github.com/akinsella/go-playground/tasks"
	_ "github.com/go-sql-driver/mysql"
)

type MySQLStopTimesImportTask struct {
	tasks.ImportTask
}

func (m *MySQLStopTimesImportTask) DoWork(workRoutine int) {
	m.InsertStopTimes(insertStopTimes);
}

func insertStopTimes(sts *models.StopTimes) (error)  {

	db, err := sql.Open("mysql", "gtfs:gtfs@/gtfs?charset=utf8mb4,utf8");

	if err != nil {
		panic(err.Error())
	}

	defer db.Close()

	valueStrings := make([]string, 0, len(sts.Records))
	valueArgs := make([]interface{}, 0, len(sts.Records) * 9)

	for _, st := range sts.Records {
		valueStrings = append(valueStrings, "('RATP', ?, ?, ?, ?, ?, ?, ?, ?)")
		valueArgs = append(
			valueArgs,
			st.TripId,
			st.ArrivalTime,
			st.DepartureTime,
			st.StopId,
			st.StopSequence,
			st.StopHeadSign,
			st.PickupType,
			st.DropOffType,
		)
	}

	stmt := fmt.Sprintf("INSERT INTO stop_times (agency_key, trip_id, arrival_time, departure_time, stop_id, stop_sequence, stop_head_sign, pickup_type, drop_off_type) VALUES %s", strings.Join(valueStrings, ","))

	_, err = db.Exec(stmt, valueArgs...)

	return err
}
<|endoftext|>


<commit_before>package providers

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/kishoreyellamraju/oauth2_proxy/cookie"
)

func (p *ProviderData) Redeem(redirectUrl, code string) (s *SessionState, err error) {
	if code == "" {
		err = errors.New("missing code")
		return
	}

	params := url.Values{}
	params.Add("redirect_uri", redirectUrl)
	params.Add("client_id", p.ClientID)
	params.Add("client_secret", p.ClientSecret)
	params.Add("code", code)
	params.Add("grant_type", "authorization_code")
	var req *http.Request
	req, err = http.NewRequest("POST", p.RedeemUrl.String(), bytes.NewBufferString(params.Encode()))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	var resp *http.Response
	resp, err = http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	var body []byte
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemUrl.String(), body)
		return
	}

	// blindly try json and x-www-form-urlencoded
	var jsonResponse struct {
		AccessToken string `json:"access_token"`
	}
	err = json.Unmarshal(body, &jsonResponse)
	if err == nil {
		s = &SessionState{
			AccessToken: jsonResponse.AccessToken,
		}
		return
	}

	var v url.Values
	v, err = url.ParseQuery(string(body))
	if err != nil {
		return
	}
	if a := v.Get("access_token"); a != "" {
		s = &SessionState{AccessToken: a}
	} else {
		err = fmt.Errorf("no access token found %s", body)
	}
	return
}

// GetLoginURL with typical oauth parameters
func (p *ProviderData) GetLoginURL(redirectURI, finalRedirect string) string {
	var a url.URL
	a = *p.LoginUrl
	params, _ := url.ParseQuery(a.RawQuery)
	params.Set("redirect_uri", redirectURI)
	params.Set("approval_prompt", "force")
	params.Add("scope", p.Scope)
	params.Set("client_id", p.ClientID)
	params.Set("response_type", "code")
	if strings.HasPrefix(finalRedirect, "/") {
		params.Add("state", finalRedirect)
	}
	a.RawQuery = params.Encode()
	return a.String()
}

// CookieForSession serializes a session state for storage in a cookie
func (p *ProviderData) CookieForSession(s *SessionState, c *cookie.Cipher) (string, error) {
	return s.EncodeSessionState(c)
}

// SessionFromCookie deserializes a session from a cookie value
func (p *ProviderData) SessionFromCookie(v string, c *cookie.Cipher) (s *SessionState, err error) {
	return DecodeSessionState(v, c)
}

func (p *ProviderData) GetEmailAddress(s *SessionState) (string, error) {
	return "", errors.New("not implemented")
}

func (p *ProviderData) ValidateSessionState(s *SessionState) bool {
	return validateToken(p, s.AccessToken, nil)
}

// RefreshSessionIfNeeded
func (p *ProviderData) RefreshSessionIfNeeded(s *SessionState) (bool, error) {
	return false, nil
}
<commit_msg>update approval_promt from force to auto<commit_after>package providers

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/kishoreyellamraju/oauth2_proxy/cookie"
)

func (p *ProviderData) Redeem(redirectUrl, code string) (s *SessionState, err error) {
	if code == "" {
		err = errors.New("missing code")
		return
	}

	params := url.Values{}
	params.Add("redirect_uri", redirectUrl)
	params.Add("client_id", p.ClientID)
	params.Add("client_secret", p.ClientSecret)
	params.Add("code", code)
	params.Add("grant_type", "authorization_code")
	var req *http.Request
	req, err = http.NewRequest("POST", p.RedeemUrl.String(), bytes.NewBufferString(params.Encode()))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	var resp *http.Response
	resp, err = http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	var body []byte
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemUrl.String(), body)
		return
	}

	// blindly try json and x-www-form-urlencoded
	var jsonResponse struct {
		AccessToken string `json:"access_token"`
	}
	err = json.Unmarshal(body, &jsonResponse)
	if err == nil {
		s = &SessionState{
			AccessToken: jsonResponse.AccessToken,
		}
		return
	}

	var v url.Values
	v, err = url.ParseQuery(string(body))
	if err != nil {
		return
	}
	if a := v.Get("access_token"); a != "" {
		s = &SessionState{AccessToken: a}
	} else {
		err = fmt.Errorf("no access token found %s", body)
	}
	return
}

// GetLoginURL with typical oauth parameters
func (p *ProviderData) GetLoginURL(redirectURI, finalRedirect string) string {
	var a url.URL
	a = *p.LoginUrl
	params, _ := url.ParseQuery(a.RawQuery)
	params.Set("redirect_uri", redirectURI)
	params.Set("approval_prompt", "auto")
	params.Add("scope", p.Scope)
	params.Set("client_id", p.ClientID)
	params.Set("response_type", "code")
	if strings.HasPrefix(finalRedirect, "/") {
		params.Add("state", finalRedirect)
	}
	a.RawQuery = params.Encode()
	return a.String()
}

// CookieForSession serializes a session state for storage in a cookie
func (p *ProviderData) CookieForSession(s *SessionState, c *cookie.Cipher) (string, error) {
	return s.EncodeSessionState(c)
}

// SessionFromCookie deserializes a session from a cookie value
func (p *ProviderData) SessionFromCookie(v string, c *cookie.Cipher) (s *SessionState, err error) {
	return DecodeSessionState(v, c)
}

func (p *ProviderData) GetEmailAddress(s *SessionState) (string, error) {
	return "", errors.New("not implemented")
}

func (p *ProviderData) ValidateSessionState(s *SessionState) bool {
	return validateToken(p, s.AccessToken, nil)
}

// RefreshSessionIfNeeded
func (p *ProviderData) RefreshSessionIfNeeded(s *SessionState) (bool, error) {
	return false, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// 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 codegen

import (
	"path"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"
	"go.uber.org/thriftrw/compile"
)

// PackageHelper manages the mapping from thrift file to generated type code and service code.
type PackageHelper struct {
	// The project package name
	packageRoot string
	// The root directory containing thrift files.
	thriftRootDir string
	// Namespace under thrift folder
	gatewayNamespace string
	// The go package name of where all the generated structs are
	genCodePackage string
	// The directory to put the generated service code.
	targetGenDir string
	// The root directory for the gateway test config files.
	testConfigsRootDir string
	// String containing copyright header to add to generated code.
	copyrightHeader string
	// The middlewares available for the endpoints
	middlewareSpecs map[string]*MiddlewareSpec
}

// NewPackageHelper creates a package helper.
func NewPackageHelper(
	packageRoot string,
	configDirName string,
	middlewareConfig string,
	thriftRootDir string,
	genCodePackage string,
	targetGenDir string,
	gatewayNamespace string,
	copyrightHeader string,
) (*PackageHelper, error) {
	genDir, err := filepath.Abs(targetGenDir)
	if err != nil {
		return nil, errors.Errorf("%s is not valid path: %s", targetGenDir, err)
	}

	genDirIndex := strings.Index(genDir, gatewayNamespace)
	if genDirIndex == -1 {
		return nil, errors.Errorf(
			"gatewayNamespace (%s) must be inside targetGenDir (%s)",
			gatewayNamespace,
			genDir,
		)
	}

	middlewareSpecs := map[string]*MiddlewareSpec{}
	middleConfig := filepath.Join(
		configDirName,
		middlewareConfig,
	)
	middlewares, err := parseMiddlewareConfig(middleConfig, configDirName)
	if err != nil {
		return nil, errors.Wrapf(
			err, "Cannot load middlewares:")
	}

	for _, mspec := range middlewares {
		middlewareSpecs[mspec.Name] = mspec
	}

	p := &PackageHelper{
		packageRoot:      packageRoot,
		thriftRootDir:    path.Clean(thriftRootDir),
		genCodePackage:   genCodePackage,
		gatewayNamespace: gatewayNamespace,
		targetGenDir:     genDir,
		copyrightHeader:  copyrightHeader,
		middlewareSpecs:  middlewareSpecs,
	}
	return p, nil
}

// PackageRoot returns the service's root package name
func (p PackageHelper) PackageRoot() string {
	return p.packageRoot
}

// MiddlewareSpecs returns a map of middlewares available
func (p PackageHelper) MiddlewareSpecs() map[string]*MiddlewareSpec {
	return p.middlewareSpecs
}

// TypeImportPath returns the Go import path for types defined in a thrift file.
func (p PackageHelper) TypeImportPath(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}

	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}
	return path.Join(
		p.genCodePackage,
		thrift[idx+len(p.thriftRootDir):len(thrift)-7],
	), nil
}

// GoGatewayPackageName returns the name of the gateway package
func (p PackageHelper) GoGatewayPackageName() string {
	nsIndex := strings.Index(p.targetGenDir, p.gatewayNamespace)
	return path.Join(
		p.gatewayNamespace,
		p.targetGenDir[nsIndex+len(p.gatewayNamespace):],
	)
}

// ThriftIDLPath returns the file path to the thrift idl folder
func (p PackageHelper) ThriftIDLPath() string {
	return p.thriftRootDir
}

// CodeGenTargetPath returns the file path where the code should
// be generated.
func (p PackageHelper) CodeGenTargetPath() string {
	return p.targetGenDir
}

// PackageGenPath returns the Go package path for generated code from a thrift file.
func (p PackageHelper) PackageGenPath(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}

	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}

	nsIndex := strings.Index(p.targetGenDir, p.gatewayNamespace)
	return path.Join(
		p.gatewayNamespace,
		p.targetGenDir[nsIndex+len(p.gatewayNamespace):],
		filepath.Dir(thrift[idx+len(p.thriftRootDir):]),
	), nil
}

// TypePackageName returns the package name that defines the type.
func (p PackageHelper) TypePackageName(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}
	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}

	// Strip the leading / and strip the .thrift on the end.
	thriftSegment := thrift[idx+len(p.thriftRootDir)+1 : len(thrift)-7]

	thriftPackageName := strings.Replace(thriftSegment, "/", "_", -1)
	return camelCase(thriftPackageName), nil
}

func (p PackageHelper) getRelativeFileName(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}
	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}
	return thrift[idx+len(p.thriftRootDir):], nil
}

// TargetClientsInitPath returns where the clients init should go
func (p PackageHelper) TargetClientsInitPath() string {
	return path.Join(p.targetGenDir, "clients", "clients.go")
}

// TargetEndpointsRegisterPath returns where the endpoints register file
// should be written to
func (p PackageHelper) TargetEndpointsRegisterPath() string {
	return path.Join(p.targetGenDir, "endpoints", "register.go")
}

// EndpointTestConfigPath returns the path for the endpoint test configs
func (p PackageHelper) EndpointTestConfigPath(
	serviceName, methodName string,
) string {
	fileName := strings.ToLower(methodName) + "_test.json"
	return path.Join(p.testConfigsRootDir, strings.ToLower(serviceName), fileName)
}

// TypeFullName returns the referred Go type name in generated code.
func (p PackageHelper) TypeFullName(typeSpec compile.TypeSpec) (string, error) {
	if typeSpec == nil {
		return "", nil
	}
	tfile := typeSpec.ThriftFile()

	if tfile == "" {
		return typeSpec.ThriftName(), nil
	}

	pkg, err := p.TypePackageName(tfile)
	if err != nil {
		return "", errors.Wrap(err, "failed to get the full type")
	}
	return pkg + "." + typeSpec.ThriftName(), nil
}
<commit_msg>codegen/package: add goGatewayNamespace field<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// 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 codegen

import (
	"path"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"
	"go.uber.org/thriftrw/compile"
)

// PackageHelper manages the mapping from thrift file to generated type code and service code.
type PackageHelper struct {
	// The project package name
	packageRoot string
	// The root directory containing thrift files.
	thriftRootDir string
	// Namespace under thrift folder
	gatewayNamespace string
	// The go package name of where all the generated structs are
	genCodePackage string
	// The directory to put the generated service code.
	targetGenDir string
	// The go package name where all the generated code is
	goGatewayNamespace string
	// The root directory for the gateway test config files.
	testConfigsRootDir string
	// String containing copyright header to add to generated code.
	copyrightHeader string
	// The middlewares available for the endpoints
	middlewareSpecs map[string]*MiddlewareSpec
}

// NewPackageHelper creates a package helper.
func NewPackageHelper(
	packageRoot string,
	configDirName string,
	middlewareConfig string,
	thriftRootDir string,
	genCodePackage string,
	targetGenDir string,
	gatewayNamespace string,
	copyrightHeader string,
) (*PackageHelper, error) {
	genDir, err := filepath.Abs(targetGenDir)
	if err != nil {
		return nil, errors.Errorf("%s is not valid path: %s", targetGenDir, err)
	}

	absConfigDir, err := filepath.Abs(configDirName)
	if err != nil {
		return nil, errors.Errorf(
			"%s is not valid path: %s", configDirName, err,
		)
	}

	relativeGenDir, err := filepath.Rel(absConfigDir, genDir)
	if err != nil {
		return nil, errors.Wrapf(
			err, "Could not compute relative dir for %q", genDir,
		)
	}

	goGatewayNamespace := filepath.Join(packageRoot, relativeGenDir)

	genDirIndex := strings.Index(genDir, gatewayNamespace)
	if genDirIndex == -1 {
		return nil, errors.Errorf(
			"gatewayNamespace (%s) must be inside targetGenDir (%s)",
			gatewayNamespace,
			genDir,
		)
	}

	middlewareSpecs := map[string]*MiddlewareSpec{}
	middleConfig := filepath.Join(
		configDirName,
		middlewareConfig,
	)
	middlewares, err := parseMiddlewareConfig(middleConfig, configDirName)
	if err != nil {
		return nil, errors.Wrapf(
			err, "Cannot load middlewares:")
	}

	for _, mspec := range middlewares {
		middlewareSpecs[mspec.Name] = mspec
	}

	p := &PackageHelper{
		packageRoot:        packageRoot,
		thriftRootDir:      path.Clean(thriftRootDir),
		genCodePackage:     genCodePackage,
		gatewayNamespace:   gatewayNamespace,
		goGatewayNamespace: goGatewayNamespace,
		targetGenDir:       genDir,
		copyrightHeader:    copyrightHeader,
		middlewareSpecs:    middlewareSpecs,
	}
	return p, nil
}

// PackageRoot returns the service's root package name
func (p PackageHelper) PackageRoot() string {
	return p.packageRoot
}

// MiddlewareSpecs returns a map of middlewares available
func (p PackageHelper) MiddlewareSpecs() map[string]*MiddlewareSpec {
	return p.middlewareSpecs
}

// TypeImportPath returns the Go import path for types defined in a thrift file.
func (p PackageHelper) TypeImportPath(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}

	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}
	return path.Join(
		p.genCodePackage,
		thrift[idx+len(p.thriftRootDir):len(thrift)-7],
	), nil
}

// GoGatewayPackageName returns the name of the gateway package
func (p PackageHelper) GoGatewayPackageName() string {
	return p.goGatewayNamespace
}

// ThriftIDLPath returns the file path to the thrift idl folder
func (p PackageHelper) ThriftIDLPath() string {
	return p.thriftRootDir
}

// CodeGenTargetPath returns the file path where the code should
// be generated.
func (p PackageHelper) CodeGenTargetPath() string {
	return p.targetGenDir
}

// PackageGenPath returns the Go package path for generated code from a thrift file.
func (p PackageHelper) PackageGenPath(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}

	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}

	nsIndex := strings.Index(p.targetGenDir, p.gatewayNamespace)
	return path.Join(
		p.gatewayNamespace,
		p.targetGenDir[nsIndex+len(p.gatewayNamespace):],
		filepath.Dir(thrift[idx+len(p.thriftRootDir):]),
	), nil
}

// TypePackageName returns the package name that defines the type.
func (p PackageHelper) TypePackageName(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}
	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}

	// Strip the leading / and strip the .thrift on the end.
	thriftSegment := thrift[idx+len(p.thriftRootDir)+1 : len(thrift)-7]

	thriftPackageName := strings.Replace(thriftSegment, "/", "_", -1)
	return camelCase(thriftPackageName), nil
}

func (p PackageHelper) getRelativeFileName(thrift string) (string, error) {
	if !strings.HasSuffix(thrift, ".thrift") {
		return "", errors.Errorf("file %s is not .thrift", thrift)
	}
	idx := strings.Index(thrift, p.thriftRootDir)
	if idx == -1 {
		return "", errors.Errorf(
			"file %s is not in thrift dir (%s)",
			thrift, p.thriftRootDir,
		)
	}
	return thrift[idx+len(p.thriftRootDir):], nil
}

// TargetClientsInitPath returns where the clients init should go
func (p PackageHelper) TargetClientsInitPath() string {
	return path.Join(p.targetGenDir, "clients", "clients.go")
}

// TargetEndpointsRegisterPath returns where the endpoints register file
// should be written to
func (p PackageHelper) TargetEndpointsRegisterPath() string {
	return path.Join(p.targetGenDir, "endpoints", "register.go")
}

// EndpointTestConfigPath returns the path for the endpoint test configs
func (p PackageHelper) EndpointTestConfigPath(
	serviceName, methodName string,
) string {
	fileName := strings.ToLower(methodName) + "_test.json"
	return path.Join(p.testConfigsRootDir, strings.ToLower(serviceName), fileName)
}

// TypeFullName returns the referred Go type name in generated code.
func (p PackageHelper) TypeFullName(typeSpec compile.TypeSpec) (string, error) {
	if typeSpec == nil {
		return "", nil
	}
	tfile := typeSpec.ThriftFile()

	if tfile == "" {
		return typeSpec.ThriftName(), nil
	}

	pkg, err := p.TypePackageName(tfile)
	if err != nil {
		return "", errors.Wrap(err, "failed to get the full type")
	}
	return pkg + "." + typeSpec.ThriftName(), nil
}
<|endoftext|>


<commit_before>package apiserver

import (
	"fmt"

	"github.com/laincloud/deployd/engine"
	"github.com/mijia/sweb/form"
	"github.com/mijia/sweb/server"
	"golang.org/x/net/context"
	"net/http"
)

type RestfulConstraints struct {
	server.BaseResource
}

func (rc RestfulConstraints) Get(ctx context.Context, r *http.Request) (int, interface{}) {
	cstType := form.ParamString(r, "type", "node")

	if constraint, ok := getEngine(ctx).GetConstraints(cstType); !ok {
		return http.StatusNotFound, fmt.Sprintf("No constraint found")
	} else {
		return http.StatusOK, constraint
	}
}

func (rc RestfulConstraints) Patch(ctx context.Context, r *http.Request) (int, interface{}) {

	cstType := form.ParamString(r, "type", "")
	cstValue := form.ParamString(r, "value", "")
	equal := form.ParamBoolean(r, "equal", false)
	soft := form.ParamBoolean(r, "soft", true)

	if cstType == "" {
		return http.StatusBadRequest, "constaint type required"
	}
	if cstValue == "" {
		return http.StatusBadRequest, "constraint value required"
	}

	constraint := engine.ConstraintSpec{cstType, equal, cstValue, soft}

	if err := getEngine(ctx).UpdateConstraints(constraint); err != nil {
		return http.StatusInternalServerError, err.Error()
	}

	urlReverser := getUrlReverser(ctx)
	return http.StatusAccepted, map[string]string{
		"message":   "Constaints will be patched",
		"check_url": urlReverser.Reverse("Get_RestfulConstraints") + "?type=" + cstType,
	}
}

func (rc RestfulConstraints) Delete(ctx context.Context, r *http.Request) (int, interface{}) {
	cstType := form.ParamString(r, "type", "node")

	if err := getEngine(ctx).DeleteConstraints(cstType); err != nil {
		if err == engine.ErrConstraintNotExists {
			return http.StatusNotFound, err.Error()
		}
		return http.StatusInternalServerError, err.Error()
	}

	urlReverser := getUrlReverser(ctx)
	return http.StatusAccepted, map[string]string{
		"message":   "Constraint will be deleted from the orc engine.",
		"check_url": urlReverser.Reverse("Get_RestfulConstraints") + "?type=" + cstType,
	}
}
<commit_msg>fix typo<commit_after>package apiserver

import (
	"fmt"

	"github.com/laincloud/deployd/engine"
	"github.com/mijia/sweb/form"
	"github.com/mijia/sweb/server"
	"golang.org/x/net/context"
	"net/http"
)

type RestfulConstraints struct {
	server.BaseResource
}

func (rc RestfulConstraints) Get(ctx context.Context, r *http.Request) (int, interface{}) {
	cstType := form.ParamString(r, "type", "node")

	if constraint, ok := getEngine(ctx).GetConstraints(cstType); !ok {
		return http.StatusNotFound, fmt.Sprintf("No constraint found")
	} else {
		return http.StatusOK, constraint
	}
}

func (rc RestfulConstraints) Patch(ctx context.Context, r *http.Request) (int, interface{}) {

	cstType := form.ParamString(r, "type", "")
	cstValue := form.ParamString(r, "value", "")
	equal := form.ParamBoolean(r, "equal", false)
	soft := form.ParamBoolean(r, "soft", true)

	if cstType == "" {
		return http.StatusBadRequest, "constraint type required"
	}
	if cstValue == "" {
		return http.StatusBadRequest, "constraint value required"
	}

	constraint := engine.ConstraintSpec{cstType, equal, cstValue, soft}

	if err := getEngine(ctx).UpdateConstraints(constraint); err != nil {
		return http.StatusInternalServerError, err.Error()
	}

	urlReverser := getUrlReverser(ctx)
	return http.StatusAccepted, map[string]string{
		"message":   "Constraints will be patched",
		"check_url": urlReverser.Reverse("Get_RestfulConstraints") + "?type=" + cstType,
	}
}

func (rc RestfulConstraints) Delete(ctx context.Context, r *http.Request) (int, interface{}) {
	cstType := form.ParamString(r, "type", "node")

	if err := getEngine(ctx).DeleteConstraints(cstType); err != nil {
		if err == engine.ErrConstraintNotExists {
			return http.StatusNotFound, err.Error()
		}
		return http.StatusInternalServerError, err.Error()
	}

	urlReverser := getUrlReverser(ctx)
	return http.StatusAccepted, map[string]string{
		"message":   "Constraint will be deleted from the orc engine.",
		"check_url": urlReverser.Reverse("Get_RestfulConstraints") + "?type=" + cstType,
	}
}
<|endoftext|>