<commit_before>package dns

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

const (
	MaxNameLen  = 255
	MaxLabelLen = 63

	headerLen    = 12
	octetPairLen = 2
)

var (
	ErrLabelTooLong = errors.New("a label must be 63 octets or less")
	ErrNameTooLong  = errors.New("a name must be 255 octets or less")
)

type Unpacker struct {
	buffer []byte
	msg    *Msg
	i      int

	nameBuffer bytes.Buffer
}

func NewUnpacker() *Unpacker {
	return &Unpacker{}
}

func (r *Unpacker) Unpack() (msg *Msg, n int, err error) {
	err = r.unpackHeader()
	if err != nil {
		return nil, 0, err
	}

	err = r.unpackQueries()
	if err != nil {
		return nil, 0, err
	}

	return r.msg, r.i, nil
}

func (r *Unpacker) Reset(b []byte) {
	r.buffer = b
	r.msg = &Msg{Header: Header{}}
	r.i = 0
}

func (r *Unpacker) readQType() QType {
	return QType(r.unpackOctetPair())
}

func (r *Unpacker) readQClass() QClass {
	return QClass(r.unpackOctetPair())
}

func (r *Unpacker) unpackHeader() (err error) {
	if len(r.buffer) < headerLen {
		return io.ErrShortBuffer
	}
	h := Header{}
	h.ID = r.unpackOctetPair()

	flags := r.unpackOctetPair()
	h.QR = (flags & maskQR >> maskQROffset) == 1
	h.OpCode = OpCode(flags & maskOpCode >> maskOpCodeOffset)
	h.AA = (flags & maskAA >> maskAAOffset) == 1
	h.TC = (flags & maskTC >> maskTCOffset) == 1
	h.RD = (flags & maskRD >> maskRDOffset) == 1
	h.RA = (flags & maskRA >> maskRAOffset) == 1
	h.Z = byte(flags & maskZ >> maskZOffset)
	h.RCode = byte(flags & maskRCode)

	h.QDCount = r.unpackOctetPair()
	h.ANCount = r.unpackOctetPair()
	h.NSCount = r.unpackOctetPair()
	h.ARCount = r.unpackOctetPair()

	r.msg.Header = h
	return nil
}

func (r *Unpacker) unpackQueries() (err error) {
	queries := make([]Query, r.msg.Header.QDCount)
	r.msg.Queries = queries
	for i := range queries {
		qName, err := r.readName()
		if err != nil {
			return err
		}
		if r.i+octetPairLen >= len(r.buffer) {
			return io.ErrShortBuffer
		}
		queries[i].QName = qName
		queries[i].QType = r.readQType()
		queries[i].QClass = r.readQClass()
	}
	return nil
}

func (r *Unpacker) readName() (string, error) {
	r.nameBuffer.Reset()
	for {
		label, err := unpackLabel(r.buffer, r.i)
		if err == io.EOF {
			r.i++
			return r.nameBuffer.String(), nil
		} else if err != nil {
			return "", err
		}

		r.i += len(label) + 1
		r.nameBuffer.WriteString(label + ".")
		if r.nameBuffer.Len() > MaxNameLen {
			return "", ErrNameTooLong
		}
	}
}

func (r *Unpacker) unpackOctetPair() uint16 {
	ini := r.i
	r.i += octetPairLen
	return binary.BigEndian.Uint16(r.buffer[ini:r.i])
}

func (r *Unpacker) unpackByte() (byte, error) {
	if r.i >= len(r.buffer) {
		return 0, io.ErrShortBuffer
	}
	b := r.buffer[r.i]
	r.i++
	return b, nil
}

func unpackLabel(b []byte, offset int) (string, error) {
	if !checkBounds(b, offset, offset+1) {
		return "", io.ErrShortBuffer
	}
	currentByte := b[offset]
	if currentByte == 0 {
		return "", io.EOF
	}
	offset++
	switch {
	case currentByte>>6 == 3:
		// Compute offsets for the pointer.
		offset = int((currentByte&64)<<8 + b[offset])
		currentByte = b[offset]
		offset++
	}
	end := offset + int(currentByte)
	if end-offset > MaxLabelLen {
		return "", ErrLabelTooLong
	}
	if !checkBounds(b, offset, end) {
		return "", io.ErrShortBuffer
	}
	return string(b[offset:end]), nil
}

func checkBounds(b []byte, begin, end int) bool {
	return len(b[begin:]) >= end-begin
}
<commit_msg>(wip) add comments on the unpackLabel<commit_after>package dns

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

const (
	MaxNameLen  = 255
	MaxLabelLen = 63

	headerLen    = 12
	octetPairLen = 2
)

var (
	ErrLabelTooLong = errors.New("a label must be 63 octets or less")
	ErrNameTooLong  = errors.New("a name must be 255 octets or less")
)

type Unpacker struct {
	buffer []byte
	msg    *Msg
	i      int

	nameBuffer bytes.Buffer
}

func NewUnpacker() *Unpacker {
	return &Unpacker{}
}

func (r *Unpacker) Unpack() (msg *Msg, n int, err error) {
	err = r.unpackHeader()
	if err != nil {
		return nil, 0, err
	}

	err = r.unpackQueries()
	if err != nil {
		return nil, 0, err
	}

	return r.msg, r.i, nil
}

func (r *Unpacker) Reset(b []byte) {
	r.buffer = b
	r.msg = &Msg{Header: Header{}}
	r.i = 0
}

func (r *Unpacker) readQType() QType {
	return QType(r.unpackOctetPair())
}

func (r *Unpacker) readQClass() QClass {
	return QClass(r.unpackOctetPair())
}

func (r *Unpacker) unpackHeader() (err error) {
	if len(r.buffer) < headerLen {
		return io.ErrShortBuffer
	}
	h := Header{}
	h.ID = r.unpackOctetPair()

	flags := r.unpackOctetPair()
	h.QR = (flags & maskQR >> maskQROffset) == 1
	h.OpCode = OpCode(flags & maskOpCode >> maskOpCodeOffset)
	h.AA = (flags & maskAA >> maskAAOffset) == 1
	h.TC = (flags & maskTC >> maskTCOffset) == 1
	h.RD = (flags & maskRD >> maskRDOffset) == 1
	h.RA = (flags & maskRA >> maskRAOffset) == 1
	h.Z = byte(flags & maskZ >> maskZOffset)
	h.RCode = byte(flags & maskRCode)

	h.QDCount = r.unpackOctetPair()
	h.ANCount = r.unpackOctetPair()
	h.NSCount = r.unpackOctetPair()
	h.ARCount = r.unpackOctetPair()

	r.msg.Header = h
	return nil
}

func (r *Unpacker) unpackQueries() (err error) {
	queries := make([]Query, r.msg.Header.QDCount)
	r.msg.Queries = queries
	for i := range queries {
		qName, err := r.readName()
		if err != nil {
			return err
		}
		if r.i+octetPairLen >= len(r.buffer) {
			return io.ErrShortBuffer
		}
		queries[i].QName = qName
		queries[i].QType = r.readQType()
		queries[i].QClass = r.readQClass()
	}
	return nil
}

func (r *Unpacker) readName() (string, error) {
	r.nameBuffer.Reset()
	for {
		label, err := unpackLabel(r.buffer, r.i)
		if err == io.EOF {
			r.i++
			return r.nameBuffer.String(), nil
		} else if err != nil {
			return "", err
		}

		r.i += len(label) + 1
		r.nameBuffer.WriteString(label + ".")
		if r.nameBuffer.Len() > MaxNameLen {
			return "", ErrNameTooLong
		}
	}
}

func (r *Unpacker) unpackOctetPair() uint16 {
	ini := r.i
	r.i += octetPairLen
	return binary.BigEndian.Uint16(r.buffer[ini:r.i])
}

func (r *Unpacker) unpackByte() (byte, error) {
	if r.i >= len(r.buffer) {
		return 0, io.ErrShortBuffer
	}
	b := r.buffer[r.i]
	r.i++
	return b, nil
}

func unpackLabel(b []byte, offset int) (string, error) {
	if !checkBounds(b, offset, offset+1) {
		return "", io.ErrShortBuffer
	}

	// Current byte indicates the length of the label.
	// If its a null byte, label is over.
	currentByte := b[offset]
	if currentByte == 0 {
		return "", io.EOF
	}

	offset++

	// Check if its a pointer.
	if currentByte>>6 == 3 {

		// Compute the offset to the pointer.
		offset = int((currentByte&64)<<8 + b[offset])
		currentByte = b[offset]
		offset++
	}

	// Check if the label has valid length.
	endOffset := offset + int(currentByte)
	if endOffset-offset > MaxLabelLen {
		return "", ErrLabelTooLong
	}
	if !checkBounds(b, offset, endOffset) {
		return "", io.ErrShortBuffer
	}

	// Return the label.
	return string(b[offset:endOffset]), nil
}

// Check if begin and end are within bounds of a byte slice.
func checkBounds(b []byte, begin, end int) bool {
	return len(b) >= begin && len(b[begin:]) >= end-begin
}
<|endoftext|>


<commit_before>package util

import "fmt"
import "io"
import "github.com/reiver/go-whitespace"
import "bytes"
import "encoding/json"
import "os"
import "bufio"
import "log"
import (
	"sort"
	"strconv"
	"sync"
	"strings"
)

type RawInputHandler struct {
	output io.Writer
	buff   []rune
}

type SumoOperator interface {
	Process(map[string]interface{})
}

type SumoAggOperator interface {
	Process(map[string]interface{})
	Flush()
}

type ParseError string

func (e ParseError) Error() string {
	return string(e)
}

func ConnectToStdIn(operator SumoOperator) {
	fi, _ := os.Stdin.Stat() // get the FileInfo struct describing the standard input.
	if (fi.Mode() & os.ModeCharDevice) == 0 {
		ConnectToReader(operator, os.Stdin)
	} else {
		fmt.Println("No input")
		return
	}
}

func ConnectToReader(operator SumoOperator, reader io.Reader) {
	bio := bufio.NewReader(reader)
	var line, hasMoreInLine, err = bio.ReadLine()
	var buf []byte
	for err != io.EOF {
		buf = append(buf, line...)
		if !hasMoreInLine && len(buf) > 0 {
			var rawMsg interface{}
			err := json.Unmarshal(buf, &rawMsg)
			buf = []byte{}
			if err != nil {
				log.Println("Error parsing json", err)
				log.Println(string(line))
			} else {
				mapMessage, ok := rawMsg.(map[string]interface{})
				if ok {
					operator.Process(mapMessage)
				} else {
					log.Println("Unexpected JSON")
				}
			}
		}
		line, hasMoreInLine, err = bio.ReadLine()
	}
}

const Plus = "PLUS"
const StartRelation = "StartRelation"
const EndRelation = "EndRelation"
const Raw = "_raw"
const Type = "_type"
const Meta = "Meta"
const Relation = "Relation"

func IsPlus(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == Plus
}

func IsStartRelation(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == StartRelation
}

func IsEndRelation(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == EndRelation
}

func IsRelation(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == Relation
}

func IsMeta(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == Meta
}

func CreateStartRelation() map[string]interface{} {
	return map[string]interface{}{Type: StartRelation}
}

func CreateStartRelationMeta(origin string) map[string]interface{} {
	return map[string]interface{}{Type: StartRelation, Meta: origin}
}

func CreateEndRelation() map[string]interface{} {
	return map[string]interface{}{Type: EndRelation}
}

func CreateRelation(inp map[string]interface{}) map[string]interface{} {
	inp[Type] = Relation
	return inp
}

func CreateMeta(inp map[string]interface{}) map[string]interface{} {
	inp[Type] = Meta
	return inp
}

func ExtractRaw(inp map[string]interface{}) string {
	raw, ok := inp[Raw].(string)
	if ok {
		return strings.TrimRight(raw, "\n")
	} else {
		return ""
	}
}

func (handler *RawInputHandler) Process(inp []byte) {
	runes := bytes.Runes(inp)
	// If not whitespace, flush, append
	if len(runes) > 0 && !whitespace.IsWhitespace(runes[0]) {
		handler.Flush()
		handler.buff = append(handler.buff, runes...)
	} else {
		// If it is whitespace, just append with a newline
		handler.buff = append(handler.buff, '\n')
		handler.buff = append(handler.buff, runes...)
	}
}

func NewRawInputHandler(inp io.Writer) *RawInputHandler {
	return &RawInputHandler{inp, []rune{}}
}

func (handler *RawInputHandler) Flush() {
	m := make(map[string]interface{})
	m[Raw] = string(handler.buff)
	m[Type] = Plus
	handler.buff = []rune{}
	b, err := json.Marshal(m)
	if err != nil {
		fmt.Printf("ERROR!", err)
	} else {
		handler.output.Write(b)
		handler.output.Write([]byte{'\n'})
	}
}

type JsonWriter struct {
	writer io.Writer
	mu     *sync.Mutex
}

func NewJsonWriter() *JsonWriter {
	return &JsonWriter{os.Stdout, &sync.Mutex{}}
}

func (writer *JsonWriter) Write(inp map[string]interface{}) {
	jsonBytes, err := json.Marshal(inp)
	//fmt.Printf(b)
	if err != nil {
		fmt.Printf("ERROR!", err)
	} else {
		writer.mu.Lock()
		writer.writer.Write(jsonBytes)
		writer.writer.Write([]byte{'\n'})
		writer.mu.Unlock()
	}
}

func CoerceNumber(v interface{}) (float64, error) {
	return strconv.ParseFloat(fmt.Sprint(v), 64)
}

type Datum []map[string]interface{}
type By func(p1, p2 *map[string]interface{}) bool

func (a datumSorter) Len() int      { return len(a.data) }
func (a datumSorter) Swap(i, j int) { a.data[i], a.data[j] = a.data[j], a.data[i] }

// planetSorter joins a By function and a slice of Planets to be sorted.
type datumSorter struct {
	data Datum
	by   func(p1, p2 map[string]interface{}) bool // Closure used in the Less method.
}

func (a datumSorter) Less(i, j int) bool {
	return a.by(a.data[i], a.data[j])
}

func SortByField(field string, data Datum) {
	by := func(p1, p2 map[string]interface{}) bool {
		v1, err1 := CoerceNumber(p1[field])
		v2, err2 := CoerceNumber(p2[field])

		if err1 != nil || err2 != nil {
			fmt.Print(data)
			panic(err1)
		}

		return v1 < v2
	}
	sort.Sort(sort.Reverse(&datumSorter{data, by}))
}
<commit_msg>Format<commit_after>package util

import "fmt"
import "io"
import "github.com/reiver/go-whitespace"
import "bytes"
import "encoding/json"
import "os"
import "bufio"
import "log"
import (
	"sort"
	"strconv"
	"strings"
	"sync"
)

type RawInputHandler struct {
	output io.Writer
	buff   []rune
}

type SumoOperator interface {
	Process(map[string]interface{})
}

type SumoAggOperator interface {
	Process(map[string]interface{})
	Flush()
}

type ParseError string

func (e ParseError) Error() string {
	return string(e)
}

func ConnectToStdIn(operator SumoOperator) {
	fi, _ := os.Stdin.Stat() // get the FileInfo struct describing the standard input.
	if (fi.Mode() & os.ModeCharDevice) == 0 {
		ConnectToReader(operator, os.Stdin)
	} else {
		fmt.Println("No input")
		return
	}
}

func ConnectToReader(operator SumoOperator, reader io.Reader) {
	bio := bufio.NewReader(reader)
	var line, hasMoreInLine, err = bio.ReadLine()
	var buf []byte
	for err != io.EOF {
		buf = append(buf, line...)
		if !hasMoreInLine && len(buf) > 0 {
			var rawMsg interface{}
			err := json.Unmarshal(buf, &rawMsg)
			buf = []byte{}
			if err != nil {
				log.Println("Error parsing json", err)
				log.Println(string(line))
			} else {
				mapMessage, ok := rawMsg.(map[string]interface{})
				if ok {
					operator.Process(mapMessage)
				} else {
					log.Println("Unexpected JSON")
				}
			}
		}
		line, hasMoreInLine, err = bio.ReadLine()
	}
}

const Plus = "PLUS"
const StartRelation = "StartRelation"
const EndRelation = "EndRelation"
const Raw = "_raw"
const Type = "_type"
const Meta = "Meta"
const Relation = "Relation"

func IsPlus(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == Plus
}

func IsStartRelation(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == StartRelation
}

func IsEndRelation(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == EndRelation
}

func IsRelation(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == Relation
}

func IsMeta(inp map[string]interface{}) bool {
	tpe, ok := inp[Type].(string)
	return ok && tpe == Meta
}

func CreateStartRelation() map[string]interface{} {
	return map[string]interface{}{Type: StartRelation}
}

func CreateStartRelationMeta(origin string) map[string]interface{} {
	return map[string]interface{}{Type: StartRelation, Meta: origin}
}

func CreateEndRelation() map[string]interface{} {
	return map[string]interface{}{Type: EndRelation}
}

func CreateRelation(inp map[string]interface{}) map[string]interface{} {
	inp[Type] = Relation
	return inp
}

func CreateMeta(inp map[string]interface{}) map[string]interface{} {
	inp[Type] = Meta
	return inp
}

func ExtractRaw(inp map[string]interface{}) string {
	raw, ok := inp[Raw].(string)
	if ok {
		return strings.TrimRight(raw, "\n")
	} else {
		return ""
	}
}

func (handler *RawInputHandler) Process(inp []byte) {
	runes := bytes.Runes(inp)
	// If not whitespace, flush, append
	if len(runes) > 0 && !whitespace.IsWhitespace(runes[0]) {
		handler.Flush()
		handler.buff = append(handler.buff, runes...)
	} else {
		// If it is whitespace, just append with a newline
		handler.buff = append(handler.buff, '\n')
		handler.buff = append(handler.buff, runes...)
	}
}

func NewRawInputHandler(inp io.Writer) *RawInputHandler {
	return &RawInputHandler{inp, []rune{}}
}

func (handler *RawInputHandler) Flush() {
	m := make(map[string]interface{})
	m[Raw] = string(handler.buff)
	m[Type] = Plus
	handler.buff = []rune{}
	b, err := json.Marshal(m)
	if err != nil {
		fmt.Printf("ERROR!", err)
	} else {
		handler.output.Write(b)
		handler.output.Write([]byte{'\n'})
	}
}

type JsonWriter struct {
	writer io.Writer
	mu     *sync.Mutex
}

func NewJsonWriter() *JsonWriter {
	return &JsonWriter{os.Stdout, &sync.Mutex{}}
}

func (writer *JsonWriter) Write(inp map[string]interface{}) {
	jsonBytes, err := json.Marshal(inp)
	//fmt.Printf(b)
	if err != nil {
		fmt.Printf("ERROR!", err)
	} else {
		writer.mu.Lock()
		writer.writer.Write(jsonBytes)
		writer.writer.Write([]byte{'\n'})
		writer.mu.Unlock()
	}
}

func CoerceNumber(v interface{}) (float64, error) {
	return strconv.ParseFloat(fmt.Sprint(v), 64)
}

type Datum []map[string]interface{}
type By func(p1, p2 *map[string]interface{}) bool

func (a datumSorter) Len() int      { return len(a.data) }
func (a datumSorter) Swap(i, j int) { a.data[i], a.data[j] = a.data[j], a.data[i] }

// planetSorter joins a By function and a slice of Planets to be sorted.
type datumSorter struct {
	data Datum
	by   func(p1, p2 map[string]interface{}) bool // Closure used in the Less method.
}

func (a datumSorter) Less(i, j int) bool {
	return a.by(a.data[i], a.data[j])
}

func SortByField(field string, data Datum) {
	by := func(p1, p2 map[string]interface{}) bool {
		v1, err1 := CoerceNumber(p1[field])
		v2, err2 := CoerceNumber(p2[field])

		if err1 != nil || err2 != nil {
			fmt.Print(data)
			panic(err1)
		}

		return v1 < v2
	}
	sort.Sort(sort.Reverse(&datumSorter{data, by}))
}
<|endoftext|>


<commit_before>package utils

import (
	"fmt"
	"regexp"
	"strings"
)

// LineIndex returns line index of search in content.
func LineIndex(content string, search string) int {
	ret := -1
	line := ""
	for _, c := range content {
		if c == '\n' {
			line += ""
			ret++
		} else {
			line += string(c)
		}
		if strings.Index(line, search) > -1 {
			ret++
			break
		}
	}
	return ret
}

// PowerLess select titles of Power<P
func PowerLess(P int, f func(*MdTitleTree)) func(*MdTitleTree) {
	return func(n *MdTitleTree) {
		if n.Power < P && n.Title != "" {
			f(n)
		}
	}
}

// LineLess select titles of Line<P
func LineLess(P int, f func(*MdTitleTree)) func(*MdTitleTree) {
	return func(n *MdTitleTree) {
		if n.Line < P && n.Title != "" {
			f(n)
		}
	}
}

var mdTitle = regexp.MustCompile(`^([#]{1,6})\s*(.+)`)

func cntStr(in []string, what string) int {
	c := -1
	for _, i := range in {
		if i == what {
			c++
		}
	}
	if c > -1 {
		c++ // starts at 1
	}
	return c
}

// GetAllMdTitles extracts all MD titles markup.
func GetAllMdTitles(content string) []MdTitle {
	ret := []MdTitle{}
	allTitles := []string{}
	line := ""
	isInBlock := false
	isInTitle := false
	i := 0
	for _, c := range content {
		if !isInBlock && c == '\n' {
			if isInTitle {
				if mdTitle.MatchString(line) {
					got := mdTitle.FindAllStringSubmatch(line, -1)
					if len(got) > 0 {
						t := got[0][2]
						ret = append(ret, MdTitle{
							Line: i, Title: t,
							Power:     len(got[0][1]),
							Duplicate: cntStr(allTitles, t),
						})
						allTitles = append(allTitles, t)
					}
				}
			}
			i++
			isInTitle = false
			line = ""
		} else if c == '`' {
			isInBlock = !isInBlock
			line += string(c)
		} else if c == '#' && !isInBlock {
			isInTitle = true
			line += string(c)
		} else {
			if c == '\n' {
				i++
			}
			line += string(c)
		}
	}
	return ret
}

// MakeTitleTree transform a raw list of titles into a tree.
func MakeTitleTree(titles []MdTitle) *MdTitleTree {

	root := &MdTitleTree{}
	cur := root
	for _, t := range titles {

		if t.Power == 1 {
			nnew := &MdTitleTree{MdTitle: t, Parent: root}
			root.Items = append(root.Items, nnew)
			cur = nnew

		} else if t.Power > cur.Power {
			nnew := &MdTitleTree{MdTitle: t, Parent: cur}
			cur.Items = append(cur.Items, nnew)
			cur = nnew

		} else if t.Power == cur.Power {
			nnew := &MdTitleTree{MdTitle: t, Parent: cur.Parent}
			cur.Parent.Items = append(cur.Parent.Items, nnew)
			cur = nnew

		} else if t.Power < cur.Power {
			for {
				if cur.Parent.Power <= t.Power {
					break
				}
				cur = cur.Parent
			}
			cur = cur.Parent
			nnew := &MdTitleTree{MdTitle: t, Parent: cur.Parent}
			cur.Parent.Items = append(cur.Parent.Items, nnew)
			cur = nnew
		}
	}
	return root
}

// GetMdLinkHash encodes s to insert into an MD link.
func GetMdLinkHash(link string) string {
	link = strings.ToLower(link)
	link = strings.Replace(link, "/", "", -1)
	link = strings.Replace(link, "$", "", -1)
	link = strings.Replace(link, ">", "", -1)
	link = strings.Replace(link, ".", "", -1)
	link = strings.Replace(link, ";", "", -1)
	link = strings.Replace(link, "!", "", -1)
	link = strings.Replace(link, "'", "", -1)
	link = strings.Replace(link, "|", "", -1)
	link = strings.Replace(link, "[", "", -1)
	link = strings.Replace(link, "]", "", -1)
	link = strings.Replace(link, " ", "-", -1)
	return link
}

// MdTitleTree is an MdTitle with tree capabilities
type MdTitleTree struct {
	MdTitle
	Parent *MdTitleTree
	Items  []*MdTitleTree
}

// Traverse a tree
func (m *MdTitleTree) Traverse(f func(*MdTitleTree)) {
	f(m)
	for _, i := range m.Items {
		i.Traverse(f)
	}
}

// LastOf a tree
func (m *MdTitleTree) LastOf(P int) *MdTitleTree {
	var ret *MdTitleTree
	if m.Power+1 == P && len(m.Items) > 0 {
		return m.Items[len(m.Items)-1]
	} else if m.Power < P {
		for _, t := range m.Items {
			if x := t.LastOf(P); x != nil {
				ret = x
			}
		}
	}
	return ret
}

// LastOf a tree
func (m *MdTitleTree) String() string {
	x := strings.Repeat("#", m.Power)
	return fmt.Sprintf("%-5v %-15q Items:%v Line:%v", x, m.Title, len(m.Items), m.Line)
}

// MdTitle is a markdwon title.
type MdTitle struct {
	Line      int
	Power     int
	Duplicate int
	Title     string
}
<commit_msg>link: handle ,<commit_after>package utils

import (
	"fmt"
	"regexp"
	"strings"
)

// LineIndex returns line index of search in content.
func LineIndex(content string, search string) int {
	ret := -1
	line := ""
	for _, c := range content {
		if c == '\n' {
			line += ""
			ret++
		} else {
			line += string(c)
		}
		if strings.Index(line, search) > -1 {
			ret++
			break
		}
	}
	return ret
}

// PowerLess select titles of Power<P
func PowerLess(P int, f func(*MdTitleTree)) func(*MdTitleTree) {
	return func(n *MdTitleTree) {
		if n.Power < P && n.Title != "" {
			f(n)
		}
	}
}

// LineLess select titles of Line<P
func LineLess(P int, f func(*MdTitleTree)) func(*MdTitleTree) {
	return func(n *MdTitleTree) {
		if n.Line < P && n.Title != "" {
			f(n)
		}
	}
}

var mdTitle = regexp.MustCompile(`^([#]{1,6})\s*(.+)`)

func cntStr(in []string, what string) int {
	c := -1
	for _, i := range in {
		if i == what {
			c++
		}
	}
	if c > -1 {
		c++ // starts at 1
	}
	return c
}

// GetAllMdTitles extracts all MD titles markup.
func GetAllMdTitles(content string) []MdTitle {
	ret := []MdTitle{}
	allTitles := []string{}
	line := ""
	isInBlock := false
	isInTitle := false
	i := 0
	for _, c := range content {
		if !isInBlock && c == '\n' {
			if isInTitle {
				if mdTitle.MatchString(line) {
					got := mdTitle.FindAllStringSubmatch(line, -1)
					if len(got) > 0 {
						t := got[0][2]
						ret = append(ret, MdTitle{
							Line: i, Title: t,
							Power:     len(got[0][1]),
							Duplicate: cntStr(allTitles, t),
						})
						allTitles = append(allTitles, t)
					}
				}
			}
			i++
			isInTitle = false
			line = ""
		} else if c == '`' {
			isInBlock = !isInBlock
			line += string(c)
		} else if c == '#' && !isInBlock {
			isInTitle = true
			line += string(c)
		} else {
			if c == '\n' {
				i++
			}
			line += string(c)
		}
	}
	return ret
}

// MakeTitleTree transform a raw list of titles into a tree.
func MakeTitleTree(titles []MdTitle) *MdTitleTree {

	root := &MdTitleTree{}
	cur := root
	for _, t := range titles {

		if t.Power == 1 {
			nnew := &MdTitleTree{MdTitle: t, Parent: root}
			root.Items = append(root.Items, nnew)
			cur = nnew

		} else if t.Power > cur.Power {
			nnew := &MdTitleTree{MdTitle: t, Parent: cur}
			cur.Items = append(cur.Items, nnew)
			cur = nnew

		} else if t.Power == cur.Power {
			nnew := &MdTitleTree{MdTitle: t, Parent: cur.Parent}
			cur.Parent.Items = append(cur.Parent.Items, nnew)
			cur = nnew

		} else if t.Power < cur.Power {
			for {
				if cur.Parent.Power <= t.Power {
					break
				}
				cur = cur.Parent
			}
			cur = cur.Parent
			nnew := &MdTitleTree{MdTitle: t, Parent: cur.Parent}
			cur.Parent.Items = append(cur.Parent.Items, nnew)
			cur = nnew
		}
	}
	return root
}

// GetMdLinkHash encodes s to insert into an MD link.
func GetMdLinkHash(link string) string {
	link = strings.ToLower(link)
	link = strings.Replace(link, "/", "", -1)
	link = strings.Replace(link, "$", "", -1)
	link = strings.Replace(link, ">", "", -1)
	link = strings.Replace(link, ".", "", -1)
	link = strings.Replace(link, ";", "", -1)
	link = strings.Replace(link, "!", "", -1)
	link = strings.Replace(link, "'", "", -1)
	link = strings.Replace(link, "|", "", -1)
	link = strings.Replace(link, "[", "", -1)
	link = strings.Replace(link, "]", "", -1)
	link = strings.Replace(link, ",", "", -1)
	link = strings.Replace(link, " ", "-", -1)
	// should it be a regexp like /[a-z0-9-_]/i ?
	return link
}

// MdTitleTree is an MdTitle with tree capabilities
type MdTitleTree struct {
	MdTitle
	Parent *MdTitleTree
	Items  []*MdTitleTree
}

// Traverse a tree
func (m *MdTitleTree) Traverse(f func(*MdTitleTree)) {
	f(m)
	for _, i := range m.Items {
		i.Traverse(f)
	}
}

// LastOf a tree
func (m *MdTitleTree) LastOf(P int) *MdTitleTree {
	var ret *MdTitleTree
	if m.Power+1 == P && len(m.Items) > 0 {
		return m.Items[len(m.Items)-1]
	} else if m.Power < P {
		for _, t := range m.Items {
			if x := t.LastOf(P); x != nil {
				ret = x
			}
		}
	}
	return ret
}

// LastOf a tree
func (m *MdTitleTree) String() string {
	x := strings.Repeat("#", m.Power)
	return fmt.Sprintf("%-5v %-15q Items:%v Line:%v", x, m.Title, len(m.Items), m.Line)
}

// MdTitle is a markdwon title.
type MdTitle struct {
	Line      int
	Power     int
	Duplicate int
	Title     string
}
<|endoftext|>


<commit_before>package utp

import (
	"bytes"
	"crypto/rand"
	"io"
	"io/ioutil"
	"math"
	mrand "math/rand"
	"net"
	"reflect"
	"testing"
	"time"
)

func init() {
	mrand.Seed(time.Now().Unix())
}

func TestReadWrite(t *testing.T) {
	laddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:10000")
	if err != nil {
		t.Fatal(err)
	}
	saddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:11000")
	if err != nil {
		t.Fatal(err)
	}
	caddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:12000")
	if err != nil {
		t.Fatal(err)
	}

	proxy, err := newProxy(laddr, saddr, caddr, 1.0)
	if err != nil {
		t.Fatal(err)
	}
	defer proxy.close()

	usaddr, err := ResolveUTPAddr("utp", "127.0.0.1:11000")
	if err != nil {
		t.Fatal(err)
	}

	ln, err := ListenUTP("utp", usaddr)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	upaddr, err := ResolveUTPAddr("utp", "127.0.0.1:10000")
	if err != nil {
		t.Fatal(err)
	}
	ucaddr, err := ResolveUTPAddr("utp", "127.0.0.1:12000")
	if err != nil {
		t.Fatal(err)
	}

	c, err := DialUTPTimeout("utp", ucaddr, upaddr, 1000*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	err = ln.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	s, err := ln.Accept()
	if err != nil {
		t.Fatal(err)
	}

	payload := []byte("Hello!")
	_, err = c.Write(payload)
	if err != nil {
		t.Fatal(err)
	}

	err = s.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	var buf [256]byte
	l, err := s.Read(buf[:])
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(payload, buf[:l]) {
		t.Errorf("expected payload of %v; got %v", payload, buf[:l])
	}

	payload2 := []byte("World!")
	_, err = s.Write(payload2)
	if err != nil {
		t.Fatal(err)
	}

	err = c.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	l, err = c.Read(buf[:])
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(payload2, buf[:l]) {
		t.Errorf("expected payload of %v; got %v", payload2, buf[:l])
	}
}

func TestLongReadWrite(t *testing.T) {
	laddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:20000")
	if err != nil {
		t.Fatal(err)
	}
	saddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:21000")
	if err != nil {
		t.Fatal(err)
	}
	caddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:22000")
	if err != nil {
		t.Fatal(err)
	}

	proxy, err := newProxy(laddr, saddr, caddr, 1.0)
	if err != nil {
		t.Fatal(err)
	}
	defer proxy.close()

	usaddr, err := ResolveUTPAddr("utp", "127.0.0.1:21000")
	if err != nil {
		t.Fatal(err)
	}

	ln, err := ListenUTP("utp", usaddr)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	upaddr, err := ResolveUTPAddr("utp", "127.0.0.1:20000")
	if err != nil {
		t.Fatal(err)
	}
	ucaddr, err := ResolveUTPAddr("utp", "127.0.0.1:22000")
	if err != nil {
		t.Fatal(err)
	}

	c, err := DialUTPTimeout("utp", ucaddr, upaddr, 1000*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	err = ln.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	s, err := ln.Accept()
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	var payload [1048576]byte
	_, err = rand.Read(payload[:])
	if err != nil {
		t.Fatal(err)
	}

	err = s.SetDeadline(time.Now().Add(10000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	rch := make(chan []byte)
	ech := make(chan error, 2)

	go func() {
		defer c.Close()
		_, err = c.Write(payload[:])
		if err != nil {
			ech <- err
		}
	}()

	go func() {
		b, err := ioutil.ReadAll(s)
		if err != nil {
			ech <- err
			rch <- nil
		} else {
			ech <- nil
			rch <- b
		}
	}()

	err = <-ech
	if err != nil {
		t.Fatal(err)
	}

	r := <-rch
	if r == nil {
		return
	}

	if !bytes.Equal(r, payload[:]) {
		t.Errorf("expected payload of %d; got %d", len(payload[:]), len(r))
	}
}

func TestAccept(t *testing.T) {
	ln, err := Listen("utp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	c, err := DialUTPTimeout("utp", nil, ln.Addr().(*UTPAddr), 200*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	err = ln.SetDeadline(time.Now().Add(100 * time.Millisecond))
	_, err = ln.Accept()
	if err != nil {
		t.Fatal(err)
	}
}

func TestAcceptDeadline(t *testing.T) {
	ln, err := Listen("utp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	err = ln.SetDeadline(time.Now().Add(time.Millisecond))
	_, err = ln.Accept()
	if err == nil {
		t.Fatal("Accept should failed")
	}
}

func TestAcceptClosedListener(t *testing.T) {
	ln, err := Listen("utp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	err = ln.Close()
	if err != nil {
		t.Fatal(err)
	}
	_, err = ln.Accept()
	if err == nil {
		t.Fatal("Accept should failed")
	}
	_, err = ln.Accept()
	if err == nil {
		t.Fatal("Accept should failed")
	}
}

func TestPacketBinary(t *testing.T) {
	h := header{
		typ:  st_fin,
		ver:  version,
		id:   100,
		t:    50000,
		diff: 10000,
		wnd:  65535,
		seq:  100,
		ack:  200,
	}

	e := []extension{
		extension{
			typ:     ext_selective_ack,
			payload: []byte{0, 1, 0, 1},
		},
		extension{
			typ:     ext_selective_ack,
			payload: []byte{100, 0, 200, 0},
		},
	}

	p := packet{
		header:  h,
		ext:     e,
		payload: []byte("abcdefg"),
	}

	b, err := p.MarshalBinary()
	if err != nil {
		t.Fatal(err)
	}

	p2 := packet{}
	err = p2.UnmarshalBinary(b)
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(p, p2) {
		t.Errorf("expected packet of %v; got %v", p, p2)
	}
}

func TestUnmarshalShortPacket(t *testing.T) {
	b := make([]byte, 18)
	p := packet{}
	err := p.UnmarshalBinary(b)

	if err == nil {
		t.Fatal("UnmarshalBinary should fail")
	} else if err != io.EOF {
		t.Fatal(err)
	}
}

func TestPacketBuffer(t *testing.T) {
	size := 12
	b := newPacketBuffer(12, 1)

	if b.space() != size {
		t.Errorf("expected space == %d; got %d", size, b.space())
	}

	for i := 1; i <= size; i++ {
		b.push(packet{header: header{seq: uint16(i)}})
	}

	if b.space() != 0 {
		t.Errorf("expected space == 0; got %d", b.space())
	}

	err := b.push(packet{header: header{seq: 15}})
	if err == nil {
		t.Fatal("push should fail")
	}

	all := b.all()
	if len(all) != size {
		t.Errorf("expected %d packets sequence; got %d", size, len(all))
	}

	_, err = b.fetch(6)
	if err != nil {
		t.Fatal(err)
	}

	b.compact()

	err = b.push(packet{header: header{seq: 15}})
	if err != nil {
		t.Fatal(err)
	}

	err = b.push(packet{header: header{seq: 17}})
	if err != nil {
		t.Fatal(err)
	}

	for i := 7; i <= size; i++ {
		_, err := b.fetch(uint16(i))
		if err != nil {
			t.Fatal(err)
		}
	}

	all = b.all()
	if len(all) != 2 {
		t.Errorf("expected 2 packets sequence; got %d", len(all))
	}

	b.compact()
	if b.space() != 9 {
		t.Errorf("expected space == 9; got %d", b.space())
	}
}

func TestPacketBufferBoundary(t *testing.T) {
	begin := math.MaxUint16 - 3
	b := newPacketBuffer(12, begin)
	for i := begin; i != 5; i = (i + 1) % (math.MaxUint16 + 1) {
		err := b.push(packet{header: header{seq: uint16(i)}})
		if err != nil {
			t.Fatal(err)
		}
	}
}

type proxy struct {
	ln *net.UDPConn
}

func newProxy(laddr, saddr, caddr *net.UDPAddr, rate float32) (*proxy, error) {
	ln, err := net.ListenUDP("udp", laddr)
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			var buf [mtu]byte
			len, addr, err := ln.ReadFromUDP(buf[:])
			if err != nil {
				return
			}

			var p packet
			p.UnmarshalBinary(buf[:len])
			if mrand.Float32() > rate {
				continue
			}

			if addr.String() == saddr.String() {
				_, err := ln.WriteToUDP(buf[:len], caddr)
				if err != nil {
					return
				}
			} else if addr.String() == caddr.String() {
				_, err := ln.WriteToUDP(buf[:len], saddr)
				if err != nil {
					return
				}
			}
		}
	}()
	return &proxy{ln: ln}, nil
}

func (p *proxy) close() {
	p.ln.Close()
}
<commit_msg>Fix datarace in TestLongReadWrite<commit_after>package utp

import (
	"bytes"
	"crypto/rand"
	"io"
	"io/ioutil"
	"math"
	mrand "math/rand"
	"net"
	"reflect"
	"testing"
	"time"
)

func init() {
	mrand.Seed(time.Now().Unix())
}

func TestReadWrite(t *testing.T) {
	laddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:10000")
	if err != nil {
		t.Fatal(err)
	}
	saddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:11000")
	if err != nil {
		t.Fatal(err)
	}
	caddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:12000")
	if err != nil {
		t.Fatal(err)
	}

	proxy, err := newProxy(laddr, saddr, caddr, 1.0)
	if err != nil {
		t.Fatal(err)
	}
	defer proxy.close()

	usaddr, err := ResolveUTPAddr("utp", "127.0.0.1:11000")
	if err != nil {
		t.Fatal(err)
	}

	ln, err := ListenUTP("utp", usaddr)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	upaddr, err := ResolveUTPAddr("utp", "127.0.0.1:10000")
	if err != nil {
		t.Fatal(err)
	}
	ucaddr, err := ResolveUTPAddr("utp", "127.0.0.1:12000")
	if err != nil {
		t.Fatal(err)
	}

	c, err := DialUTPTimeout("utp", ucaddr, upaddr, 1000*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	err = ln.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	s, err := ln.Accept()
	if err != nil {
		t.Fatal(err)
	}

	payload := []byte("Hello!")
	_, err = c.Write(payload)
	if err != nil {
		t.Fatal(err)
	}

	err = s.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	var buf [256]byte
	l, err := s.Read(buf[:])
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(payload, buf[:l]) {
		t.Errorf("expected payload of %v; got %v", payload, buf[:l])
	}

	payload2 := []byte("World!")
	_, err = s.Write(payload2)
	if err != nil {
		t.Fatal(err)
	}

	err = c.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	l, err = c.Read(buf[:])
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(payload2, buf[:l]) {
		t.Errorf("expected payload of %v; got %v", payload2, buf[:l])
	}
}

func TestLongReadWrite(t *testing.T) {
	laddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:20000")
	if err != nil {
		t.Fatal(err)
	}
	saddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:21000")
	if err != nil {
		t.Fatal(err)
	}
	caddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:22000")
	if err != nil {
		t.Fatal(err)
	}

	proxy, err := newProxy(laddr, saddr, caddr, 1.0)
	if err != nil {
		t.Fatal(err)
	}
	defer proxy.close()

	usaddr, err := ResolveUTPAddr("utp", "127.0.0.1:21000")
	if err != nil {
		t.Fatal(err)
	}

	ln, err := ListenUTP("utp", usaddr)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	upaddr, err := ResolveUTPAddr("utp", "127.0.0.1:20000")
	if err != nil {
		t.Fatal(err)
	}
	ucaddr, err := ResolveUTPAddr("utp", "127.0.0.1:22000")
	if err != nil {
		t.Fatal(err)
	}

	c, err := DialUTPTimeout("utp", ucaddr, upaddr, 1000*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	err = ln.SetDeadline(time.Now().Add(1000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	s, err := ln.Accept()
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	var payload [1048576]byte
	_, err = rand.Read(payload[:])
	if err != nil {
		t.Fatal(err)
	}

	err = s.SetDeadline(time.Now().Add(10000 * time.Millisecond))
	if err != nil {
		t.Fatal(err)
	}

	rch := make(chan []byte)
	ech := make(chan error, 2)

	go func() {
		defer c.Close()
		_, err := c.Write(payload[:])
		if err != nil {
			ech <- err
		}
	}()

	go func() {
		b, err := ioutil.ReadAll(s)
		if err != nil {
			ech <- err
			rch <- nil
		} else {
			ech <- nil
			rch <- b
		}
	}()

	err = <-ech
	if err != nil {
		t.Fatal(err)
	}

	r := <-rch
	if r == nil {
		return
	}

	if !bytes.Equal(r, payload[:]) {
		t.Errorf("expected payload of %d; got %d", len(payload[:]), len(r))
	}
}

func TestAccept(t *testing.T) {
	ln, err := Listen("utp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	c, err := DialUTPTimeout("utp", nil, ln.Addr().(*UTPAddr), 200*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	err = ln.SetDeadline(time.Now().Add(100 * time.Millisecond))
	_, err = ln.Accept()
	if err != nil {
		t.Fatal(err)
	}
}

func TestAcceptDeadline(t *testing.T) {
	ln, err := Listen("utp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	err = ln.SetDeadline(time.Now().Add(time.Millisecond))
	_, err = ln.Accept()
	if err == nil {
		t.Fatal("Accept should failed")
	}
}

func TestAcceptClosedListener(t *testing.T) {
	ln, err := Listen("utp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	err = ln.Close()
	if err != nil {
		t.Fatal(err)
	}
	_, err = ln.Accept()
	if err == nil {
		t.Fatal("Accept should failed")
	}
	_, err = ln.Accept()
	if err == nil {
		t.Fatal("Accept should failed")
	}
}

func TestPacketBinary(t *testing.T) {
	h := header{
		typ:  st_fin,
		ver:  version,
		id:   100,
		t:    50000,
		diff: 10000,
		wnd:  65535,
		seq:  100,
		ack:  200,
	}

	e := []extension{
		extension{
			typ:     ext_selective_ack,
			payload: []byte{0, 1, 0, 1},
		},
		extension{
			typ:     ext_selective_ack,
			payload: []byte{100, 0, 200, 0},
		},
	}

	p := packet{
		header:  h,
		ext:     e,
		payload: []byte("abcdefg"),
	}

	b, err := p.MarshalBinary()
	if err != nil {
		t.Fatal(err)
	}

	p2 := packet{}
	err = p2.UnmarshalBinary(b)
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(p, p2) {
		t.Errorf("expected packet of %v; got %v", p, p2)
	}
}

func TestUnmarshalShortPacket(t *testing.T) {
	b := make([]byte, 18)
	p := packet{}
	err := p.UnmarshalBinary(b)

	if err == nil {
		t.Fatal("UnmarshalBinary should fail")
	} else if err != io.EOF {
		t.Fatal(err)
	}
}

func TestPacketBuffer(t *testing.T) {
	size := 12
	b := newPacketBuffer(12, 1)

	if b.space() != size {
		t.Errorf("expected space == %d; got %d", size, b.space())
	}

	for i := 1; i <= size; i++ {
		b.push(packet{header: header{seq: uint16(i)}})
	}

	if b.space() != 0 {
		t.Errorf("expected space == 0; got %d", b.space())
	}

	err := b.push(packet{header: header{seq: 15}})
	if err == nil {
		t.Fatal("push should fail")
	}

	all := b.all()
	if len(all) != size {
		t.Errorf("expected %d packets sequence; got %d", size, len(all))
	}

	_, err = b.fetch(6)
	if err != nil {
		t.Fatal(err)
	}

	b.compact()

	err = b.push(packet{header: header{seq: 15}})
	if err != nil {
		t.Fatal(err)
	}

	err = b.push(packet{header: header{seq: 17}})
	if err != nil {
		t.Fatal(err)
	}

	for i := 7; i <= size; i++ {
		_, err := b.fetch(uint16(i))
		if err != nil {
			t.Fatal(err)
		}
	}

	all = b.all()
	if len(all) != 2 {
		t.Errorf("expected 2 packets sequence; got %d", len(all))
	}

	b.compact()
	if b.space() != 9 {
		t.Errorf("expected space == 9; got %d", b.space())
	}
}

func TestPacketBufferBoundary(t *testing.T) {
	begin := math.MaxUint16 - 3
	b := newPacketBuffer(12, begin)
	for i := begin; i != 5; i = (i + 1) % (math.MaxUint16 + 1) {
		err := b.push(packet{header: header{seq: uint16(i)}})
		if err != nil {
			t.Fatal(err)
		}
	}
}

type proxy struct {
	ln *net.UDPConn
}

func newProxy(laddr, saddr, caddr *net.UDPAddr, rate float32) (*proxy, error) {
	ln, err := net.ListenUDP("udp", laddr)
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			var buf [mtu]byte
			len, addr, err := ln.ReadFromUDP(buf[:])
			if err != nil {
				return
			}

			var p packet
			p.UnmarshalBinary(buf[:len])
			if mrand.Float32() > rate {
				continue
			}

			if addr.String() == saddr.String() {
				_, err := ln.WriteToUDP(buf[:len], caddr)
				if err != nil {
					return
				}
			} else if addr.String() == caddr.String() {
				_, err := ln.WriteToUDP(buf[:len], saddr)
				if err != nil {
					return
				}
			}
		}
	}()
	return &proxy{ln: ln}, nil
}

func (p *proxy) close() {
	p.ln.Close()
}
<|endoftext|>


<commit_before>package kvcodec

import (
	"bufio"
	"fmt"
	"io"
	"reflect"
	"regexp"
	"strings"
	"sync"
)

const (
	tagLabel = "mapstruct"
)

type structFields map[string]fieldMeta

type fieldMeta struct {
	Key        string
	Name       string
	OmitAlways bool
	OmitEmpty  bool
}

func readField(field reflect.StructField) (meta fieldMeta) {
	meta = fieldMeta{
		Name: field.Name,
	}
	fieldTags := strings.Split(field.Tag.Get(tagLabel), ",")
	for _, tag := range fieldTags {
		if tag == "-" {
			meta.OmitAlways = true
			return meta
		}

		if tag == "omitempty" {
			meta.OmitEmpty = true
		} else if tag != "" {
			meta.Key = tag
		} else {
			meta.Key = field.Name
		}
	}

	return meta
}

var structMap = make(map[reflect.Type]structFields)
var structMapMutex sync.RWMutex

func getstructFields(rType reflect.Type) (stInfo structFields) {
	structMapMutex.RLock()
	stInfo, ok := structMap[rType]
	if !ok {
		stInfo = getFieldInfos(rType)
		structMap[rType] = stInfo
	}
	structMapMutex.RUnlock()
	return stInfo
}

func getFieldInfos(rType reflect.Type) structFields {
	fieldsCount := rType.NumField()
	fieldMap := make(structFields)

	for i := 0; i < fieldsCount; i++ {
		field := rType.Field(i)
		meta := readField(field)

		if field.PkgPath != "" {
			continue
		}

		// if the field is an embedded struct, create a fieldInfo for each of its fields
		if field.Anonymous && field.Type.Kind() == reflect.Struct {
			panic(fmt.Errorf("embedded structs not supported"))
		}

		if !meta.OmitAlways {
			fieldMap[meta.Key] = meta
		}
	}

	return fieldMap
}

func trim(s string) string {
	re := regexp.MustCompile("(\\S+|\\S+)")
	return strings.Join(re.FindAllString(s, -1), " ")
}

func Unmarshal(in io.Reader, out interface{}) error {
	outValue := reflect.ValueOf(out)
	if outValue.Kind() == reflect.Ptr {
		outValue = outValue.Elem()
	}
	outType := outValue.Type()

	fields := getstructFields(outType)

	scanner := bufio.NewScanner(in)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), ":") {
			s := strings.Split(scanner.Text(), ":")
			k, v := trim(s[0]), trim(s[1])
			if meta, ok := fields[k]; ok {
				field := outValue.FieldByName(meta.Name)
				setField(field, v)
			}
		}
	}
	return nil
}
<commit_msg>improve error return, move delimiter and tag to const<commit_after>package kvcodec

import (
	"bufio"
	"fmt"
	"io"
	"reflect"
	"regexp"
	"strings"
	"sync"
)

const (
	tagLabel = "kv"
	kvDelim  = ":"
)

type structFields map[string]FieldMeta

type FieldMeta struct {
	Key        string
	Name       string
	OmitAlways bool
	OmitEmpty  bool
}

func newFieldMeta(field reflect.StructField) (meta FieldMeta) {
	meta = FieldMeta{
		Name: field.Name,
	}
	fieldTags := strings.Split(field.Tag.Get(tagLabel), ",")
	for _, tag := range fieldTags {
		if tag == "-" {
			meta.OmitAlways = true
			return meta
		}
		if tag == "omitempty" {
			meta.OmitEmpty = true
		} else if tag != "" {
			meta.Key = tag
		} else {
			meta.Key = field.Name
		}
	}
	return meta
}

var err error
var structMap = make(map[reflect.Type]structFields)
var structMapMutex sync.RWMutex

func getStructFields(rType reflect.Type) (structFields, error) {
	structMapMutex.RLock()
	stInfo, ok := structMap[rType]
	if !ok {
		stInfo, err = newStructFields(rType)
		if err != nil {
			return nil, err
		}
		structMap[rType] = stInfo
	}
	structMapMutex.RUnlock()
	return stInfo, nil
}

func newStructFields(rType reflect.Type) (structFields, error) {
	fieldsCount := rType.NumField()
	fieldMap := make(structFields)

	for i := 0; i < fieldsCount; i++ {
		field := rType.Field(i)
		meta := newFieldMeta(field)

		if field.PkgPath != "" {
			continue
		}

		if field.Anonymous && field.Type.Kind() == reflect.Struct {
			return nil, fmt.Errorf("embedded structs not supported")
		}

		if !meta.OmitAlways {
			fieldMap[meta.Key] = meta
		}
	}

	return fieldMap, nil
}

func trim(s string) string {
	re := regexp.MustCompile("(\\S+|\\S+)")
	return strings.Join(re.FindAllString(s, -1), " ")
}

func Unmarshal(in io.Reader, out interface{}) error {
	outValue := reflect.ValueOf(out)
	if outValue.Kind() == reflect.Ptr {
		outValue = outValue.Elem()
	}

	fields, err := getStructFields(outValue.Type())
	if fields == nil {
		panic(err)
	}
	if err != nil {
		return err
	}

	scanner := bufio.NewScanner(in)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), kvDelim) {
			s := strings.Split(scanner.Text(), kvDelim)
			k, v := trim(s[0]), trim(s[1])
			if meta, ok := fields[k]; ok {
				field := outValue.FieldByName(meta.Name)
				setField(field, v)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
	"goconf.googlecode.com/hg"
	"http"
	"bufio"
	"strings"
)

func SpawnPollers(c *conf.ConfigFile) {
	sections := c.GetSections()
	for i:= range sections {
		sectionname := sections[i]
		sectiontype, _ := c.GetString(sectionname, "type")
		if sectiontype == "" {
			continue
		}
		switch sectiontype {
			case "ipbouncer":
				bouncer_url, _ := c.GetString(sectionname, "bouncer_url")
				update_url, _ := c.GetString(sectionname, "update_url")
				interval, _ := c.GetInt(sectionname, "interval")
				go IpBouncerPoller(bouncer_url, update_url, interval)
			default:
				fmt.Printf("Warning: unknown type %v\n", sectiontype)
		}
	}
}

func IpBouncerPoller(bouncer_url string, update_url string, interval int) {
	for {
		fmt.Printf("polling %s\n", bouncer_url)
		httpclient := new(http.Client)
		resp, err := httpclient.Get(bouncer_url)
		if err == nil {
			httpdata := bufio.NewReader(resp.Body)
			ipdata, _, err2 := httpdata.ReadLine()
			if err2 == nil {
				ip := string(ipdata)
				UpdateIP(update_url, ip)
			}
			resp.Body.Close()
		}
		time.Sleep(int64(interval) * int64(1000000000))
	}
}

func UpdateIP(update_url string, ip string) {
	full_url := strings.Replace(update_url, "<ip>", ip, -1)
	fmt.Printf("full_url = %v ip = %v\n", full_url, ip)
	httpclient := new(http.Client)
	resp, err := httpclient.Get(full_url)
	if err != nil {
		fmt.Printf("an error occured while updating IP: %v\n", err)
	}
	resp.Body.Close()
}
<commit_msg>added network device polling.<commit_after>package main

import (
	"fmt"
	"time"
	"goconf.googlecode.com/hg"
	"http"
	"bufio"
	"strings"
	"syscall"
	"net"
	"os"
	"unsafe"
)

func SpawnPollers(c *conf.ConfigFile) {
	sections := c.GetSections()
	for i:= range sections {
		sectionname := sections[i]
		sectiontype, _ := c.GetString(sectionname, "type")
		if sectiontype == "" {
			continue
		}
		switch sectiontype {
			case "ipbouncer":
				bouncer_url, _ := c.GetString(sectionname, "bouncer_url")
				update_url, _ := c.GetString(sectionname, "update_url")
				interval, _ := c.GetInt(sectionname, "interval")
				go IpBouncerPoller(bouncer_url, update_url, interval)
			case "device":
				device, _ := c.GetString(sectionname, "device")
				update_url, _ := c.GetString(sectionname, "update_url")
				interval, _ := c.GetInt(sectionname, "interval")
				go DevicePoller(device, update_url, interval)
			default:
				fmt.Printf("Warning: unknown type %v\n", sectiontype)
		}
	}
}

func IpBouncerPoller(bouncer_url string, update_url string, interval int) {
	for {
		fmt.Printf("polling %s\n", bouncer_url)
		httpclient := new(http.Client)
		resp, err := httpclient.Get(bouncer_url)
		if err == nil {
			httpdata := bufio.NewReader(resp.Body)
			ipdata, _, err2 := httpdata.ReadLine()
			if err2 == nil {
				ip := string(ipdata)
				UpdateIP(update_url, ip)
			}
			resp.Body.Close()
		}
		time.Sleep(int64(interval) * int64(1000000000))
	}
}

func DevicePoller(device string, update_url string, interval int) {
	for {
		fmt.Printf("polling %s\n", device)

		ip, err := GetIPFromDevice(device)
		if err == nil {
			UpdateIP(update_url, ip.String())
		} else {
			fmt.Printf("device poller error: %v\n", err)
		}

		time.Sleep(int64(interval) * int64(1000000000))
	}
}

func UpdateIP(update_url string, ip string) {
	full_url := strings.Replace(update_url, "<ip>", ip, -1)
	fmt.Printf("full_url = %v ip = %v\n", full_url, ip)
	httpclient := new(http.Client)
	resp, err := httpclient.Get(full_url)
	if err != nil {
		fmt.Printf("an error occured while updating IP: %v\n", err)
	}
	resp.Body.Close()
}

func GetIPFromDevice(device string) (ip net.IP, err os.Error) {
	var ifreqbuf [40]byte

	for i := 0 ; i < 40 ; i++ {
		ifreqbuf[i] = 0
	}

	for i := 0 ; i < len(device) ; i++ {
		ifreqbuf[i] = device[i]
	}

	socketfd, _, _ := syscall.Syscall(syscall.SYS_SOCKET, syscall.AF_INET, syscall.SOCK_DGRAM, 0)
	defer syscall.Close(int(socketfd))

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(socketfd), uintptr(syscall.SIOCGIFADDR), uintptr(unsafe.Pointer(&ifreqbuf)))
	if err = os.NewSyscallError("SYS_IOCTL", int(errno)); err != nil {
		return
	}
	ip = net.IPv4(ifreqbuf[20], ifreqbuf[21], ifreqbuf[22], ifreqbuf[23])
	return
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/keybase/go-libkb"
	"github.com/keybase/protocol/go"
	fmprpc "github.com/maxtaco/go-framed-msgpack-rpc"
	"github.com/ugorji/go/codec"
	"net"
	"net/rpc"
)

type RemoteIdentifyUI struct {
	sessionId int
	rpccli    *rpc.Client
	uicli     keybase_1.IdentifyUiClient
}

type IdentifyHandler struct {
	conn net.Conn
	cli  *rpc.Client
}

func (h *IdentifyHandler) GetRpcClient() (cli *rpc.Client) {
	if cli = h.cli; cli == nil {
		var mh codec.MsgpackHandle
		cdc := fmprpc.MsgpackSpecRpc.ClientCodec(h.conn, &mh, true)
		cli = rpc.NewClientWithCodec(cdc)
		h.cli = cli
	}
	return
}

func NewRemoteIdentifyUI(sessionId int, c *rpc.Client) *RemoteIdentifyUI {
	return &RemoteIdentifyUI{
		sessionId: sessionId,
		rpccli:    c,
		uicli:     keybase_1.IdentifyUiClient{c},
	}
}

func (h *IdentifyHandler) NewUi(sessionId int) libkb.IdentifyUI {
	return NewRemoteIdentifyUI(sessionId, h.GetRpcClient())
}

func (u *RemoteIdentifyUI) FinishWebProofCheck(p keybase_1.RemoteProof, lcr libkb.LinkCheckResult) {
	var status keybase_1.Status
	u.uicli.FinishWebProofCheck(keybase_1.FinishWebProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:       lcr.Export(),
	}, &status)
	return
}

func (u *RemoteIdentifyUI) FinishSocialProofCheck(p keybase_1.RemoteProof, lcr libkb.LinkCheckResult) {
	var status keybase_1.Status
	u.uicli.FinishSocialProofCheck(keybase_1.FinishSocialProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:     lcr.Export(),
	}, &status)
	return
}

func (u *RemoteIdentifyUI) FinishAndPrompt(io *keybase_1.IdentifyOutcome) (ret keybase_1.FinishAndPromptRes) {
	err := u.uicli.FinishAndPrompt(keybase_1.FinishAndPromptArg{SessionId: u.sessionId, Outcome: *io, }, &ret)
	if err != nil {
		ret.Status = libkb.ExportErrorAsStatus(err)
	}
	return
}

func (u *RemoteIdentifyUI) DisplayCryptocurrency(l *libkb.CryptocurrencyChainLink) {
	var status keybase_1.Status
	u.uicli.DisplayCryptocurrency(keybase_1.DisplayCryptocurrencyArg{
		SessionId: u.sessionId,
		Address:   l.ToDisplayString(),
	}, &status)
	return
}

func (u *RemoteIdentifyUI) DisplayKey(k keybase_1.FOKID, d *keybase_1.TrackDiff) {
	var status keybase_1.Status
	u.uicli.DisplayKey(keybase_1.DisplayKeyArg{SessionId: u.sessionId, Fokid: k, Diff: d}, &status)
	return
}

func (u *RemoteIdentifyUI) ReportLastTrack(t *keybase_1.TrackSummary) {
	var status keybase_1.Status
	u.uicli.ReportLastTrack(keybase_1.ReportLastTrackArg{SessionId: u.sessionId, Track: t}, &status)
	return
}

func (u *RemoteIdentifyUI) Start() {}

func (u *RemoteIdentifyUI) LaunchNetworkChecks(res *libkb.IdentifyRes) {
	var status keybase_1.Status
	u.uicli.LaunchNetworkChecks(keybase_1.LaunchNetworkChecksArg{
		SessionId: u.sessionId,
		Id:        res.ExportToUncheckedIdentity(),
	}, &status)
	return
}

func (h *IdentifyHandler) IdentifySelf(sessionId *int, res *keybase_1.Status) error {
	luarg := libkb.LoadUserArg{}
	u, err := libkb.LoadMe(luarg)
	if _, not_found := err.(libkb.NoKeyError); not_found {
		err = nil
	} else if _, not_selected := err.(libkb.NoSelectedKeyError); not_selected {
		_, err = u.IdentifySelf(h.NewUi(*sessionId))
	}
	status := libkb.ExportErrorAsStatus(err)
	res = &status
	return nil
}
<commit_msg>more u/i exim, almost done<commit_after>package main

import (
	"github.com/keybase/go-libkb"
	"github.com/keybase/protocol/go"
	fmprpc "github.com/maxtaco/go-framed-msgpack-rpc"
	"github.com/ugorji/go/codec"
	"net"
	"net/rpc"
)

type RemoteIdentifyUI struct {
	sessionId int
	rpccli    *rpc.Client
	uicli     keybase_1.IdentifyUiClient
}

type IdentifyHandler struct {
	conn net.Conn
	cli  *rpc.Client
}

func (h *IdentifyHandler) GetRpcClient() (cli *rpc.Client) {
	if cli = h.cli; cli == nil {
		var mh codec.MsgpackHandle
		cdc := fmprpc.MsgpackSpecRpc.ClientCodec(h.conn, &mh, true)
		cli = rpc.NewClientWithCodec(cdc)
		h.cli = cli
	}
	return
}

func NewRemoteIdentifyUI(sessionId int, c *rpc.Client) *RemoteIdentifyUI {
	return &RemoteIdentifyUI{
		sessionId: sessionId,
		rpccli:    c,
		uicli:     keybase_1.IdentifyUiClient{c},
	}
}

func (h *IdentifyHandler) NewUi(sessionId int) libkb.IdentifyUI {
	return NewRemoteIdentifyUI(sessionId, h.GetRpcClient())
}

func (u *RemoteIdentifyUI) FinishWebProofCheck(p keybase_1.RemoteProof, lcr keybase_1.LinkCheckResult) {
	var status keybase_1.Status
	u.uicli.FinishWebProofCheck(keybase_1.FinishWebProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:       lcr,
	}, &status)
	return
}

func (u *RemoteIdentifyUI) FinishSocialProofCheck(p keybase_1.RemoteProof, lcr keybase_1.LinkCheckResult) {
	var status keybase_1.Status
	u.uicli.FinishSocialProofCheck(keybase_1.FinishSocialProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:     lcr,
	}, &status)
	return
}

func (u *RemoteIdentifyUI) FinishAndPrompt(io *keybase_1.IdentifyOutcome) (ret keybase_1.FinishAndPromptRes) {
	err := u.uicli.FinishAndPrompt(keybase_1.FinishAndPromptArg{SessionId: u.sessionId, Outcome: *io, }, &ret)
	if err != nil {
		ret.Status = libkb.ExportErrorAsStatus(err)
	}
	return
}

func (u *RemoteIdentifyUI) DisplayCryptocurrency(l *libkb.CryptocurrencyChainLink) {
	var status keybase_1.Status
	u.uicli.DisplayCryptocurrency(keybase_1.DisplayCryptocurrencyArg{
		SessionId: u.sessionId,
		Address:   l.ToDisplayString(),
	}, &status)
	return
}

func (u *RemoteIdentifyUI) DisplayKey(k keybase_1.FOKID, d *keybase_1.TrackDiff) {
	var status keybase_1.Status
	u.uicli.DisplayKey(keybase_1.DisplayKeyArg{SessionId: u.sessionId, Fokid: k, Diff: d}, &status)
	return
}

func (u *RemoteIdentifyUI) ReportLastTrack(t *keybase_1.TrackSummary) {
	var status keybase_1.Status
	u.uicli.ReportLastTrack(keybase_1.ReportLastTrackArg{SessionId: u.sessionId, Track: t}, &status)
	return
}

func (u *RemoteIdentifyUI) Start() {}

func (u *RemoteIdentifyUI) LaunchNetworkChecks(id *keybase_1.Identity) {
	var status keybase_1.Status
	u.uicli.LaunchNetworkChecks(keybase_1.LaunchNetworkChecksArg{
		SessionId: u.sessionId,
		Id:        *id,
	}, &status)
	return
}

func (h *IdentifyHandler) IdentifySelf(sessionId *int, res *keybase_1.Status) error {
	luarg := libkb.LoadUserArg{}
	u, err := libkb.LoadMe(luarg)
	if _, not_found := err.(libkb.NoKeyError); not_found {
		err = nil
	} else if _, not_selected := err.(libkb.NoSelectedKeyError); not_selected {
		_, err = u.IdentifySelf(h.NewUi(*sessionId))
	}
	status := libkb.ExportErrorAsStatus(err)
	res = &status
	return nil
}
<|endoftext|>


<commit_before>/*
Package imdb provides types and functions for retrieving data from an IMDb
database loaded by Goim. There are types for each entity and attribute, along
with some convenience functions for loading them from the database. While there
are a lot of types---since the database is large---this package is actually
fairly minimal. It is likely that you'll find the 'search' sub-package more
useful.

The database can be queried using the 'database/sql' package, but it is
strongly recommended that you use the Open function in this package (which will
give you access to a *sql.DB value). Namely, the Open function will perform a
migration on the schema of your database to make sure it is up to date with the
version of the 'imdb' package that you're using. (If the migration fails, it
will be rolled back and your database will be left untouched.)

Also, many of the functions here require values with types in my csql package:
https://github.com/BurntSushi/csql. Mostly, these types are interfaces that
types in the 'database/sql' package already satisfy. For example, a
csql.Queryer can be a *imdb.DB, *sql.DB, *sql.Stmt or a *sql.Tx, etc.
*/
package imdb
<commit_msg>give a hint about key types since there are so many<commit_after>/*
Package imdb provides types and functions for retrieving data from an IMDb
database loaded by Goim. There are types for each entity and attribute, along
with some convenience functions for loading them from the database. While there
are a lot of types---since the database is large---this package is actually
fairly minimal. It is likely that you'll find the 'search' sub-package more
useful.

The database can be queried using the 'database/sql' package, but it is
strongly recommended that you use the Open function in this package (which will
give you access to a *sql.DB value). Namely, the Open function will perform a
migration on the schema of your database to make sure it is up to date with the
version of the 'imdb' package that you're using. (If the migration fails, it
will be rolled back and your database will be left untouched.)

Also, many of the functions here require values with types in my csql package:
https://github.com/BurntSushi/csql. Mostly, these types are interfaces that
types in the 'database/sql' package already satisfy. For example, a
csql.Queryer can be a *imdb.DB, *sql.DB, *sql.Stmt or a *sql.Tx, etc.

The central types of this package are DB, Entity, Movie, Tvshow, Episode and
Actor. Most of the other types correspond to attributes of entities.
*/
package imdb
<|endoftext|>


<commit_before>package persist

import (
	"bytes"
	"testing"

	"github.com/NebulousLabs/Sia/build"
)

func TestSaveLoad(t *testing.T) {
	var meta = Metadata{"TestSaveLoad", "0.1"}
	var saveData int = 3
	buf := new(bytes.Buffer)

	// save data to buffer
	err := Save(meta, saveData, buf)
	if err != nil {
		t.Fatal(err)
	}
	data := buf.Bytes()

	// load valid data
	var loadData int
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err != nil {
		t.Fatal(err)
	}
	if loadData != saveData {
		t.Fatalf("loaded data (%v) does not match saved data (%v)", loadData, saveData)
	}

	// load with bad metadata
	err = Load(Metadata{"BadTestSaveLoad", "0.1"}, &loadData, bytes.NewReader(data))
	if err != ErrBadHeader {
		t.Fatal("expected ErrBadHeader, got", err)
	}
	err = Load(Metadata{"TestSaveLoad", "-1"}, &loadData, bytes.NewReader(data))
	if err != ErrBadVersion {
		t.Fatal("expected ErrBadVersion, got", err)
	}

	// corrupt data, moving back to front
	data[21] = '}'
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err == nil {
		t.Fatal("expected error when loading corrupted data")
	}
	data[14] = '}'
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err == nil {
		t.Fatal("expected error when loading corrupted data")
	}
	data[0] = '}'
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err == nil {
		t.Fatal("expected error when loading corrupted data")
	}
}

func TestSaveLoadFile(t *testing.T) {
	var meta = Metadata{"TestSaveLoadFile", "0.1"}
	var saveData int = 3

	filename := build.TempDir("TestSaveLoadFile")
	err := SaveFile(meta, saveData, filename)
	if err != nil {
		t.Fatal(err)
	}

	var loadData int
	err = LoadFile(meta, &loadData, filename)
	if err != nil {
		t.Fatal(err)
	}
	if loadData != saveData {
		t.Fatalf("loaded data (%v) does not match saved data (%v)", loadData, saveData)
	}
}
<commit_msg>create persist testdir<commit_after>package persist

import (
	"bytes"
	"os"
	"testing"

	"github.com/NebulousLabs/Sia/build"
)

func TestSaveLoad(t *testing.T) {
	var meta = Metadata{"TestSaveLoad", "0.1"}
	var saveData int = 3
	buf := new(bytes.Buffer)

	// save data to buffer
	err := Save(meta, saveData, buf)
	if err != nil {
		t.Fatal(err)
	}
	data := buf.Bytes()

	// load valid data
	var loadData int
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err != nil {
		t.Fatal(err)
	}
	if loadData != saveData {
		t.Fatalf("loaded data (%v) does not match saved data (%v)", loadData, saveData)
	}

	// load with bad metadata
	err = Load(Metadata{"BadTestSaveLoad", "0.1"}, &loadData, bytes.NewReader(data))
	if err != ErrBadHeader {
		t.Fatal("expected ErrBadHeader, got", err)
	}
	err = Load(Metadata{"TestSaveLoad", "-1"}, &loadData, bytes.NewReader(data))
	if err != ErrBadVersion {
		t.Fatal("expected ErrBadVersion, got", err)
	}

	// corrupt data, moving back to front
	data[21] = '}'
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err == nil {
		t.Fatal("expected error when loading corrupted data")
	}
	data[14] = '}'
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err == nil {
		t.Fatal("expected error when loading corrupted data")
	}
	data[0] = '}'
	err = Load(meta, &loadData, bytes.NewReader(data))
	if err == nil {
		t.Fatal("expected error when loading corrupted data")
	}
}

func TestSaveLoadFile(t *testing.T) {
	var meta = Metadata{"TestSaveLoadFile", "0.1"}
	var saveData int = 3

	os.MkdirAll(build.TempDir("persist"), 0777)
	filename := build.TempDir("persist", "TestSaveLoadFile")
	err := SaveFile(meta, saveData, filename)
	if err != nil {
		t.Fatal(err)
	}

	var loadData int
	err = LoadFile(meta, &loadData, filename)
	if err != nil {
		t.Fatal(err)
	}
	if loadData != saveData {
		t.Fatalf("loaded data (%v) does not match saved data (%v)", loadData, saveData)
	}
}
<|endoftext|>


<commit_before>// main is the inbucket daemon launcher
package main

import (
	"context"
	"expvar"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/jhillyerd/inbucket/config"
	"github.com/jhillyerd/inbucket/httpd"
	"github.com/jhillyerd/inbucket/log"
	"github.com/jhillyerd/inbucket/msghub"
	"github.com/jhillyerd/inbucket/pop3d"
	"github.com/jhillyerd/inbucket/rest"
	"github.com/jhillyerd/inbucket/smtpd"
	"github.com/jhillyerd/inbucket/webui"
)

var (
	// version contains the build version number, populated during linking
	version = "undefined"

	// date contains the build date, populated during linking
	date = "undefined"

	// Command line flags
	help    = flag.Bool("help", false, "Displays this help")
	pidfile = flag.String("pidfile", "none", "Write our PID into the specified file")
	logfile = flag.String("logfile", "stderr", "Write out log into the specified file")

	// shutdownChan - close it to tell Inbucket to shut down cleanly
	shutdownChan = make(chan bool)

	// Server instances
	smtpServer *smtpd.Server
	pop3Server *pop3d.Server
)

func init() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage of inbucket [options] <conf file>:")
		flag.PrintDefaults()
	}

	// Server uptime for status page
	startTime := time.Now()
	expvar.Publish("uptime", expvar.Func(func() interface{} {
		return time.Since(startTime) / time.Second
	}))

	// Goroutine count for status page
	expvar.Publish("goroutines", expvar.Func(func() interface{} {
		return runtime.NumGoroutine()
	}))
}

func main() {
	config.Version = version
	config.BuildDate = date

	flag.Parse()
	if *help {
		flag.Usage()
		return
	}

	// Root context
	rootCtx, rootCancel := context.WithCancel(context.Background())

	// Load & Parse config
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	err := config.LoadConfig(flag.Arg(0))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to parse config: %v\n", err)
		os.Exit(1)
	}

	// Setup signal handler
	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGINT)

	// Initialize logging
	log.SetLogLevel(config.GetLogLevel())
	if err := log.Initialize(*logfile); err != nil {
		fmt.Fprintf(os.Stderr, "%v", err)
		os.Exit(1)
	}
	defer log.Close()

	log.Infof("Inbucket %v (%v) starting...", config.Version, config.BuildDate)

	// Write pidfile if requested
	if *pidfile != "none" {
		pidf, err := os.Create(*pidfile)
		if err != nil {
			log.Errorf("Failed to create %q: %v", *pidfile, err)
			os.Exit(1)
		}
		fmt.Fprintf(pidf, "%v\n", os.Getpid())
		if err := pidf.Close(); err != nil {
			log.Errorf("Failed to close PID file %q: %v", *pidfile, err)
		}
	}

	// Create message hub
	msgHub := msghub.New(rootCtx, config.GetWebConfig().MonitorHistory)

	// Grab our datastore
	ds := smtpd.DefaultFileDataStore()

	// Start HTTP server
	httpd.Initialize(config.GetWebConfig(), shutdownChan, ds, msgHub)
	webui.SetupRoutes(httpd.Router)
	rest.SetupRoutes(httpd.Router)
	go httpd.Start(rootCtx)

	// Start POP3 server
	// TODO pass datastore
	pop3Server = pop3d.New(shutdownChan)
	go pop3Server.Start(rootCtx)

	// Startup SMTP server
	smtpServer = smtpd.NewServer(config.GetSMTPConfig(), shutdownChan, ds, msgHub)
	go smtpServer.Start(rootCtx)

	// Loop forever waiting for signals or shutdown channel
signalLoop:
	for {
		select {
		case sig := <-sigChan:
			switch sig {
			case syscall.SIGHUP:
				log.Infof("Recieved SIGHUP, cycling logfile")
				log.Rotate()
			case syscall.SIGINT:
				// Shutdown requested
				log.Infof("Received SIGINT, shutting down")
				close(shutdownChan)
			case syscall.SIGTERM:
				// Shutdown requested
				log.Infof("Received SIGTERM, shutting down")
				close(shutdownChan)
			}
		case _ = <-shutdownChan:
			rootCancel()
			break signalLoop
		}
	}

	// Wait for active connections to finish
	go timedExit()
	smtpServer.Drain()
	pop3Server.Drain()

	removePIDFile()
}

// removePIDFile removes the PID file if created
func removePIDFile() {
	if *pidfile != "none" {
		if err := os.Remove(*pidfile); err != nil {
			log.Errorf("Failed to remove %q: %v", *pidfile, err)
		}
	}
}

// timedExit is called as a goroutine during shutdown, it will force an exit
// after 15 seconds
func timedExit() {
	time.Sleep(15 * time.Second)
	log.Errorf("Clean shutdown took too long, forcing exit")
	removePIDFile()
	os.Exit(0)
}
<commit_msg>Include version for final goxc release<commit_after>// main is the inbucket daemon launcher
package main

import (
	"context"
	"expvar"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/jhillyerd/inbucket/config"
	"github.com/jhillyerd/inbucket/httpd"
	"github.com/jhillyerd/inbucket/log"
	"github.com/jhillyerd/inbucket/msghub"
	"github.com/jhillyerd/inbucket/pop3d"
	"github.com/jhillyerd/inbucket/rest"
	"github.com/jhillyerd/inbucket/smtpd"
	"github.com/jhillyerd/inbucket/webui"
)

var (
	// version contains the build version number, populated during linking
	version = "1.2.0-rc2"

	// date contains the build date, populated during linking
	date = "undefined"

	// Command line flags
	help    = flag.Bool("help", false, "Displays this help")
	pidfile = flag.String("pidfile", "none", "Write our PID into the specified file")
	logfile = flag.String("logfile", "stderr", "Write out log into the specified file")

	// shutdownChan - close it to tell Inbucket to shut down cleanly
	shutdownChan = make(chan bool)

	// Server instances
	smtpServer *smtpd.Server
	pop3Server *pop3d.Server
)

func init() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage of inbucket [options] <conf file>:")
		flag.PrintDefaults()
	}

	// Server uptime for status page
	startTime := time.Now()
	expvar.Publish("uptime", expvar.Func(func() interface{} {
		return time.Since(startTime) / time.Second
	}))

	// Goroutine count for status page
	expvar.Publish("goroutines", expvar.Func(func() interface{} {
		return runtime.NumGoroutine()
	}))
}

func main() {
	config.Version = version
	config.BuildDate = date

	flag.Parse()
	if *help {
		flag.Usage()
		return
	}

	// Root context
	rootCtx, rootCancel := context.WithCancel(context.Background())

	// Load & Parse config
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	err := config.LoadConfig(flag.Arg(0))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to parse config: %v\n", err)
		os.Exit(1)
	}

	// Setup signal handler
	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGINT)

	// Initialize logging
	log.SetLogLevel(config.GetLogLevel())
	if err := log.Initialize(*logfile); err != nil {
		fmt.Fprintf(os.Stderr, "%v", err)
		os.Exit(1)
	}
	defer log.Close()

	log.Infof("Inbucket %v (%v) starting...", config.Version, config.BuildDate)

	// Write pidfile if requested
	if *pidfile != "none" {
		pidf, err := os.Create(*pidfile)
		if err != nil {
			log.Errorf("Failed to create %q: %v", *pidfile, err)
			os.Exit(1)
		}
		fmt.Fprintf(pidf, "%v\n", os.Getpid())
		if err := pidf.Close(); err != nil {
			log.Errorf("Failed to close PID file %q: %v", *pidfile, err)
		}
	}

	// Create message hub
	msgHub := msghub.New(rootCtx, config.GetWebConfig().MonitorHistory)

	// Grab our datastore
	ds := smtpd.DefaultFileDataStore()

	// Start HTTP server
	httpd.Initialize(config.GetWebConfig(), shutdownChan, ds, msgHub)
	webui.SetupRoutes(httpd.Router)
	rest.SetupRoutes(httpd.Router)
	go httpd.Start(rootCtx)

	// Start POP3 server
	// TODO pass datastore
	pop3Server = pop3d.New(shutdownChan)
	go pop3Server.Start(rootCtx)

	// Startup SMTP server
	smtpServer = smtpd.NewServer(config.GetSMTPConfig(), shutdownChan, ds, msgHub)
	go smtpServer.Start(rootCtx)

	// Loop forever waiting for signals or shutdown channel
signalLoop:
	for {
		select {
		case sig := <-sigChan:
			switch sig {
			case syscall.SIGHUP:
				log.Infof("Recieved SIGHUP, cycling logfile")
				log.Rotate()
			case syscall.SIGINT:
				// Shutdown requested
				log.Infof("Received SIGINT, shutting down")
				close(shutdownChan)
			case syscall.SIGTERM:
				// Shutdown requested
				log.Infof("Received SIGTERM, shutting down")
				close(shutdownChan)
			}
		case _ = <-shutdownChan:
			rootCancel()
			break signalLoop
		}
	}

	// Wait for active connections to finish
	go timedExit()
	smtpServer.Drain()
	pop3Server.Drain()

	removePIDFile()
}

// removePIDFile removes the PID file if created
func removePIDFile() {
	if *pidfile != "none" {
		if err := os.Remove(*pidfile); err != nil {
			log.Errorf("Failed to remove %q: %v", *pidfile, err)
		}
	}
}

// timedExit is called as a goroutine during shutdown, it will force an exit
// after 15 seconds
func timedExit() {
	time.Sleep(15 * time.Second)
	log.Errorf("Clean shutdown took too long, forcing exit")
	removePIDFile()
	os.Exit(0)
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package action

import (
	"github.com/juju/loggo"
	"github.com/juju/names"

	"github.com/juju/juju/apiserver/common"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/state"
)

var logger = loggo.GetLogger("juju.apiserver.action")

func init() {
	common.RegisterStandardFacade("Action", 0, NewActionAPI)
}

// ActionAPI implements the client API for interacting with Actions
type ActionAPI struct {
	state      *state.State
	resources  *common.Resources
	authorizer common.Authorizer
}

// NewActionAPI returns an initialized ActionAPI
func NewActionAPI(st *state.State, resources *common.Resources, authorizer common.Authorizer) (*ActionAPI, error) {
	if !authorizer.AuthClient() {
		return nil, common.ErrPerm
	}

	return &ActionAPI{
		state:      st,
		resources:  resources,
		authorizer: authorizer,
	}, nil
}

// Enqueue takes a list of Actions and queues them up to be executed by
// the designated ActionReceiver, returning the params.Action for each
// queued Action, or an error if there was a problem queueing up the
// Action.
func (a *ActionAPI) Enqueue(arg params.Actions) (params.ActionResults, error) {
	response := params.ActionResults{Results: make([]params.ActionResult, len(arg.Actions))}
	for i, action := range arg.Actions {
		current := &response.Results[i]

		receiver, err := tagToActionReceiver(a.state, action.Receiver)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		queued, err := receiver.AddAction(action.Name, action.Parameters)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}
		current.Action = &params.Action{
			Receiver:   receiver.Tag().String(),
			Tag:        queued.ActionTag().String(),
			Name:       queued.Name(),
			Parameters: queued.Parameters(),
		}
		current.Status = string(state.ActionPending)
	}
	return response, nil
}

// ListAll takes a list of Entities representing ActionReceivers and returns
// all of the Actions that have been queued or run by each of those
// Entities.
func (a *ActionAPI) ListAll(arg params.Entities) (params.ActionsByReceivers, error) {
	return a.internalList(arg, combine(actionReceiverToActions, actionReceiverToActionResults))
}

// ListPending takes a list of Entities representing ActionReceivers
// and returns all of the Actions that are queued for each of those
// Entities.
func (a *ActionAPI) ListPending(arg params.Entities) (params.ActionsByReceivers, error) {
	return a.internalList(arg, actionReceiverToActions)
}

// ListCompleted takes a list of Entities representing ActionReceivers
// and returns all of the Actions that have been run on each of those
// Entities.
func (a *ActionAPI) ListCompleted(arg params.Entities) (params.ActionsByReceivers, error) {
	return a.internalList(arg, actionReceiverToActionResults)
}

// Cancel attempts to cancel queued up Actions from running.
func (a *ActionAPI) Cancel(arg params.Entities) (params.ActionResults, error) {
	response := params.ActionResults{Results: make([]params.ActionResult, len(arg.Entities))}
	for i, entity := range arg.Entities {
		current := &response.Results[i]
		tag, err := names.ParseTag(entity.Tag)
		if err != nil {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		atag, ok := tag.(names.ActionTag)
		if !ok {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		receiver, err := tagToActionReceiver(a.state, atag.PrefixTag().String())
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		action, err := a.state.ActionByTag(atag)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		result, err := receiver.CancelAction(action)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		current.Action = &params.Action{
			Tag:        atag.String(),
			Receiver:   receiver.Tag().String(),
			Name:       result.Name(),
			Parameters: result.Parameters(),
		}
		current.Status = string(result.Status())
		output, message := result.Results()
		current.Message = message
		current.Output = output
	}
	return response, nil
}

// ServicesCharmActions returns a slice of charm Actions for a slice of services.
func (a *ActionAPI) ServicesCharmActions(args params.Entities) (params.ServicesCharmActionsResults, error) {
	result := params.ServicesCharmActionsResults{}
	for _, entity := range args.Entities {
		newResult := params.ServiceCharmActionsResult{}
		svcTag, err := names.ParseServiceTag(entity.Tag)
		if err != nil {
			newResult.Error = common.ServerError(common.ErrBadId)
			result.Results = append(result.Results, newResult)
			continue
		}
		newResult.ServiceTag = svcTag.String()
		svc, err := a.state.Service(svcTag.Id())
		if err != nil {
			newResult.Error = common.ServerError(err)
			result.Results = append(result.Results, newResult)
			continue
		}
		ch, _, err := svc.Charm()
		if err != nil {
			newResult.Error = common.ServerError(err)
			result.Results = append(result.Results, newResult)
			continue
		}
		newResult.Actions = ch.Actions()
		result.Results = append(result.Results, newResult)
	}

	return result, nil
}

// internalList takes a list of Entities representing ActionReceivers and
// returns all of the Actions the extractorFn can get out of the
// ActionReceiver.
func (a *ActionAPI) internalList(arg params.Entities, fn extractorFn) (params.ActionsByReceivers, error) {
	response := params.ActionsByReceivers{Actions: make([]params.ActionsByReceiver, len(arg.Entities))}
	for i, entity := range arg.Entities {
		current := &response.Actions[i]
		receiver, err := tagToActionReceiver(a.state, entity.Tag)
		if err != nil {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		current.Receiver = receiver.Tag().String()

		results, err := fn(receiver)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}
		current.Actions = results
	}
	return response, nil
}

// tagToActionReceiver takes a tag string and tries to convert it to an
// ActionReceiver.
func tagToActionReceiver(st *state.State, tag string) (state.ActionReceiver, error) {
	receiverTag, err := names.ParseTag(tag)
	if err != nil {
		return nil, common.ErrBadId
	}
	entity, err := st.FindEntity(receiverTag)
	if err != nil {
		return nil, common.ErrBadId
	}
	receiver, ok := entity.(state.ActionReceiver)
	if !ok {
		return nil, common.ErrBadId
	}
	return receiver, nil
}

// extractorFn is the generic signature for functions that extract
// Actions or ActionResults from an ActionReceiver, and return them as
// params.Actions.
type extractorFn func(state.ActionReceiver) ([]params.ActionResult, error)

// combine takes multiple extractorFn's and combines them into
// one function.
func combine(funcs ...extractorFn) extractorFn {
	return func(ar state.ActionReceiver) ([]params.ActionResult, error) {
		result := []params.ActionResult{}
		for _, fn := range funcs {
			items, err := fn(ar)
			if err != nil {
				return result, err
			}
			result = append(result, items...)
		}
		return result, nil
	}
}

// actionReceiverToActions iterates through the Actions() queued up for
// an ActionReceiver, and converts them to a slice of params.Action.
func actionReceiverToActions(ar state.ActionReceiver) ([]params.ActionResult, error) {
	items := []params.ActionResult{}
	actions, err := ar.PendingActions()
	if err != nil {
		return items, err
	}
	for _, action := range actions {
		if action == nil {
			continue
		}
		items = append(items, params.ActionResult{
			Action: &params.Action{
				Receiver:   ar.Tag().String(),
				Tag:        action.ActionTag().String(),
				Name:       action.Name(),
				Parameters: action.Parameters(),
			},
			Status: string(state.ActionPending),
		})
	}
	return items, nil
}

// actionReceiverToActionResults iterates through the ActionResults()
// aqueued up for n ActionReceiver, and converts them to a slice of
// aparams.Action.
func actionReceiverToActionResults(ar state.ActionReceiver) ([]params.ActionResult, error) {
	items := []params.ActionResult{}
	actions, err := ar.CompletedActions()
	if err != nil {
		return items, err
	}
	for _, action := range actions {
		if action == nil {
			continue
		}
		output, message := action.Results()
		items = append(items, params.ActionResult{
			Action: &params.Action{
				Receiver:   ar.Tag().String(),
				Tag:        action.ActionTag().String(),
				Name:       action.Name(),
				Parameters: action.Parameters(),
			},
			Status:  string(action.Status()),
			Message: message,
			Output:  output,
		})
	}
	return items, nil
}
<commit_msg>Review feedback fixes<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package action

import (
	"github.com/juju/loggo"
	"github.com/juju/names"

	"github.com/juju/juju/apiserver/common"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/state"
)

var logger = loggo.GetLogger("juju.apiserver.action")

func init() {
	common.RegisterStandardFacade("Action", 0, NewActionAPI)
}

// ActionAPI implements the client API for interacting with Actions
type ActionAPI struct {
	state      *state.State
	resources  *common.Resources
	authorizer common.Authorizer
}

// NewActionAPI returns an initialized ActionAPI
func NewActionAPI(st *state.State, resources *common.Resources, authorizer common.Authorizer) (*ActionAPI, error) {
	if !authorizer.AuthClient() {
		return nil, common.ErrPerm
	}

	return &ActionAPI{
		state:      st,
		resources:  resources,
		authorizer: authorizer,
	}, nil
}

// Enqueue takes a list of Actions and queues them up to be executed by
// the designated ActionReceiver, returning the params.Action for each
// queued Action, or an error if there was a problem queueing up the
// Action.
func (a *ActionAPI) Enqueue(arg params.Actions) (params.ActionResults, error) {
	response := params.ActionResults{Results: make([]params.ActionResult, len(arg.Actions))}
	for i, action := range arg.Actions {
		current := &response.Results[i]

		receiver, err := tagToActionReceiver(a.state, action.Receiver)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		queued, err := receiver.AddAction(action.Name, action.Parameters)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}
		current.Action = &params.Action{
			Receiver:   receiver.Tag().String(),
			Tag:        queued.ActionTag().String(),
			Name:       queued.Name(),
			Parameters: queued.Parameters(),
		}
		current.Status = string(state.ActionPending)
	}
	return response, nil
}

// ListAll takes a list of Entities representing ActionReceivers and returns
// all of the Actions that have been queued or run by each of those
// Entities.
func (a *ActionAPI) ListAll(arg params.Entities) (params.ActionsByReceivers, error) {
	return a.internalList(arg, combine(actionReceiverToActions, actionReceiverToActionResults))
}

// ListPending takes a list of Entities representing ActionReceivers
// and returns all of the Actions that are queued for each of those
// Entities.
func (a *ActionAPI) ListPending(arg params.Entities) (params.ActionsByReceivers, error) {
	return a.internalList(arg, actionReceiverToActions)
}

// ListCompleted takes a list of Entities representing ActionReceivers
// and returns all of the Actions that have been run on each of those
// Entities.
func (a *ActionAPI) ListCompleted(arg params.Entities) (params.ActionsByReceivers, error) {
	return a.internalList(arg, actionReceiverToActionResults)
}

// Cancel attempts to cancel queued up Actions from running.
func (a *ActionAPI) Cancel(arg params.Entities) (params.ActionResults, error) {
	response := params.ActionResults{Results: make([]params.ActionResult, len(arg.Entities))}
	for i, entity := range arg.Entities {
		current := &response.Results[i]
		tag, err := names.ParseTag(entity.Tag)
		if err != nil {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		actionTag, ok := tag.(names.ActionTag)
		if !ok {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		receiver, err := tagToActionReceiver(a.state, actionTag.PrefixTag().String())
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		action, err := a.state.ActionByTag(actionTag)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		result, err := receiver.CancelAction(action)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}

		current.Action = &params.Action{
			Tag:        actionTag.String(),
			Receiver:   receiver.Tag().String(),
			Name:       result.Name(),
			Parameters: result.Parameters(),
		}
		current.Status = string(result.Status())
		output, message := result.Results()
		current.Message = message
		current.Output = output
	}
	return response, nil
}

// ServicesCharmActions returns a slice of charm Actions for a slice of services.
func (a *ActionAPI) ServicesCharmActions(args params.Entities) (params.ServicesCharmActionsResults, error) {
	result := params.ServicesCharmActionsResults{Results: make([]params.ServiceCharmActionsResult, len(args.Entities))}
	for i, entity := range args.Entities {
		current := &result.Results[i]
		svcTag, err := names.ParseServiceTag(entity.Tag)
		if err != nil {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		current.ServiceTag = svcTag.String()
		svc, err := a.state.Service(svcTag.Id())
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}
		ch, _, err := svc.Charm()
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}
		current.Actions = ch.Actions()
	}
	return result, nil
}

// internalList takes a list of Entities representing ActionReceivers and
// returns all of the Actions the extractorFn can get out of the
// ActionReceiver.
func (a *ActionAPI) internalList(arg params.Entities, fn extractorFn) (params.ActionsByReceivers, error) {
	response := params.ActionsByReceivers{Actions: make([]params.ActionsByReceiver, len(arg.Entities))}
	for i, entity := range arg.Entities {
		current := &response.Actions[i]
		receiver, err := tagToActionReceiver(a.state, entity.Tag)
		if err != nil {
			current.Error = common.ServerError(common.ErrBadId)
			continue
		}
		current.Receiver = receiver.Tag().String()

		results, err := fn(receiver)
		if err != nil {
			current.Error = common.ServerError(err)
			continue
		}
		current.Actions = results
	}
	return response, nil
}

// tagToActionReceiver takes a tag string and tries to convert it to an
// ActionReceiver.
func tagToActionReceiver(st *state.State, tag string) (state.ActionReceiver, error) {
	receiverTag, err := names.ParseTag(tag)
	if err != nil {
		return nil, common.ErrBadId
	}
	entity, err := st.FindEntity(receiverTag)
	if err != nil {
		return nil, common.ErrBadId
	}
	receiver, ok := entity.(state.ActionReceiver)
	if !ok {
		return nil, common.ErrBadId
	}
	return receiver, nil
}

// extractorFn is the generic signature for functions that extract
// Actions or ActionResults from an ActionReceiver, and return them as
// params.Actions.
type extractorFn func(state.ActionReceiver) ([]params.ActionResult, error)

// combine takes multiple extractorFn's and combines them into
// one function.
func combine(funcs ...extractorFn) extractorFn {
	return func(ar state.ActionReceiver) ([]params.ActionResult, error) {
		result := []params.ActionResult{}
		for _, fn := range funcs {
			items, err := fn(ar)
			if err != nil {
				return result, err
			}
			result = append(result, items...)
		}
		return result, nil
	}
}

// actionReceiverToActions iterates through the Actions() queued up for
// an ActionReceiver, and converts them to a slice of params.Action.
func actionReceiverToActions(ar state.ActionReceiver) ([]params.ActionResult, error) {
	items := []params.ActionResult{}
	actions, err := ar.PendingActions()
	if err != nil {
		return items, err
	}
	for _, action := range actions {
		if action == nil {
			continue
		}
		items = append(items, params.ActionResult{
			Action: &params.Action{
				Receiver:   ar.Tag().String(),
				Tag:        action.ActionTag().String(),
				Name:       action.Name(),
				Parameters: action.Parameters(),
			},
			Status: string(state.ActionPending),
		})
	}
	return items, nil
}

// actionReceiverToActionResults iterates through the ActionResults()
// aqueued up for n ActionReceiver, and converts them to a slice of
// aparams.Action.
func actionReceiverToActionResults(ar state.ActionReceiver) ([]params.ActionResult, error) {
	items := []params.ActionResult{}
	actions, err := ar.CompletedActions()
	if err != nil {
		return items, err
	}
	for _, action := range actions {
		if action == nil {
			continue
		}
		output, message := action.Results()
		items = append(items, params.ActionResult{
			Action: &params.Action{
				Receiver:   ar.Tag().String(),
				Tag:        action.ActionTag().String(),
				Name:       action.Name(),
				Parameters: action.Parameters(),
			},
			Status:  string(action.Status()),
			Message: message,
			Output:  output,
		})
	}
	return items, nil
}
<|endoftext|>


<commit_before>package nats

import (
	"fmt"
	"log"
	"strings"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/common/tls"
	"github.com/influxdata/telegraf/plugins/outputs"
	"github.com/influxdata/telegraf/plugins/serializers"
	"github.com/nats-io/nats.go"
)

type NATS struct {
	Servers     []string `toml:"servers"`
	Secure      bool     `toml:"secure"`
	Name        string   `toml:"name"`
	Username    string   `toml:"username"`
	Password    string   `toml:"password"`
	Credentials string   `toml:"credentials"`
	Subject     string   `toml:"subject"`

	tls.ClientConfig

	conn       *nats.Conn
	serializer serializers.Serializer
}

var sampleConfig = `
  ## URLs of NATS servers
  servers = ["nats://localhost:4222"]

  ## Optional client name
  # name = ""

  ## Optional credentials
  # username = ""
  # password = ""

  ## Optional NATS 2.0 and NATS NGS compatible user credentials
  # credentials = "/etc/telegraf/nats.creds"

  ## NATS subject for producer messages
  subject = "telegraf"

  ## Use Transport Layer Security
  # secure = false

  ## Optional TLS Config
  # tls_ca = "/etc/telegraf/ca.pem"
  # tls_cert = "/etc/telegraf/cert.pem"
  # tls_key = "/etc/telegraf/key.pem"
  ## Use TLS but skip chain & host verification
  # insecure_skip_verify = false

  ## Data format to output.
  ## Each data format has its own unique set of configuration options, read
  ## more about them here:
  ## https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_OUTPUT.md
  data_format = "influx"
`

func (n *NATS) SetSerializer(serializer serializers.Serializer) {
	n.serializer = serializer
}

func (n *NATS) Connect() error {
	var err error

	opts := []nats.Option{
		nats.MaxReconnects(-1),
	}

	// override authentication, if any was specified
	if n.Username != "" {
		opts = append(opts, nats.UserInfo(n.Username, n.Password))
	}

	if n.Name != "" {
		opts = append(opts, nats.Name(n.Name))
	}

	if n.Secure {
		tlsConfig, err := n.ClientConfig.TLSConfig()
		if err != nil {
			return err
		}

		opts = append(opts, nats.Secure(tlsConfig))
	}

	// try and connect
	n.conn, err = nats.Connect(strings.Join(n.Servers, ","), opts...)

	return err
}

func (n *NATS) Close() error {
	n.conn.Close()
	return nil
}

func (n *NATS) SampleConfig() string {
	return sampleConfig
}

func (n *NATS) Description() string {
	return "Send telegraf measurements to NATS"
}

func (n *NATS) Write(metrics []telegraf.Metric) error {
	if len(metrics) == 0 {
		return nil
	}

	for _, metric := range metrics {
		buf, err := n.serializer.Serialize(metric)
		if err != nil {
			log.Printf("D! [outputs.nats] Could not serialize metric: %v", err)
			continue
		}

		err = n.conn.Publish(n.Subject, buf)
		if err != nil {
			return fmt.Errorf("FAILED to send NATS message: %s", err)
		}
	}
	return nil
}

func init() {
	outputs.Add("nats", func() telegraf.Output {
		return &NATS{}
	})
}
<commit_msg>nats_output: use the configured credentials file (#8986)<commit_after>package nats

import (
	"fmt"
	"log"
	"strings"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/common/tls"
	"github.com/influxdata/telegraf/plugins/outputs"
	"github.com/influxdata/telegraf/plugins/serializers"
	"github.com/nats-io/nats.go"
)

type NATS struct {
	Servers     []string `toml:"servers"`
	Secure      bool     `toml:"secure"`
	Name        string   `toml:"name"`
	Username    string   `toml:"username"`
	Password    string   `toml:"password"`
	Credentials string   `toml:"credentials"`
	Subject     string   `toml:"subject"`

	tls.ClientConfig

	conn       *nats.Conn
	serializer serializers.Serializer
}

var sampleConfig = `
  ## URLs of NATS servers
  servers = ["nats://localhost:4222"]

  ## Optional client name
  # name = ""

  ## Optional credentials
  # username = ""
  # password = ""

  ## Optional NATS 2.0 and NATS NGS compatible user credentials
  # credentials = "/etc/telegraf/nats.creds"

  ## NATS subject for producer messages
  subject = "telegraf"

  ## Use Transport Layer Security
  # secure = false

  ## Optional TLS Config
  # tls_ca = "/etc/telegraf/ca.pem"
  # tls_cert = "/etc/telegraf/cert.pem"
  # tls_key = "/etc/telegraf/key.pem"
  ## Use TLS but skip chain & host verification
  # insecure_skip_verify = false

  ## Data format to output.
  ## Each data format has its own unique set of configuration options, read
  ## more about them here:
  ## https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_OUTPUT.md
  data_format = "influx"
`

func (n *NATS) SetSerializer(serializer serializers.Serializer) {
	n.serializer = serializer
}

func (n *NATS) Connect() error {
	var err error

	opts := []nats.Option{
		nats.MaxReconnects(-1),
	}

	// override authentication, if any was specified
	if n.Username != "" && n.Password != "" {
		opts = append(opts, nats.UserInfo(n.Username, n.Password))
	}

	if n.Credentials != "" {
		opts = append(opts, nats.UserCredentials(n.Credentials))
	}

	if n.Name != "" {
		opts = append(opts, nats.Name(n.Name))
	}

	if n.Secure {
		tlsConfig, err := n.ClientConfig.TLSConfig()
		if err != nil {
			return err
		}

		opts = append(opts, nats.Secure(tlsConfig))
	}

	// try and connect
	n.conn, err = nats.Connect(strings.Join(n.Servers, ","), opts...)

	return err
}

func (n *NATS) Close() error {
	n.conn.Close()
	return nil
}

func (n *NATS) SampleConfig() string {
	return sampleConfig
}

func (n *NATS) Description() string {
	return "Send telegraf measurements to NATS"
}

func (n *NATS) Write(metrics []telegraf.Metric) error {
	if len(metrics) == 0 {
		return nil
	}

	for _, metric := range metrics {
		buf, err := n.serializer.Serialize(metric)
		if err != nil {
			log.Printf("D! [outputs.nats] Could not serialize metric: %v", err)
			continue
		}

		err = n.conn.Publish(n.Subject, buf)
		if err != nil {
			return fmt.Errorf("FAILED to send NATS message: %s", err)
		}
	}
	return nil
}

func init() {
	outputs.Add("nats", func() telegraf.Output {
		return &NATS{}
	})
}
<|endoftext|>