<commit_before>package cmd

import (
	"fmt"
	"io"
	"math"
	//"os"
	"sort"
	"strconv"
	"strings"

	"github.com/biogo/biogo/align"
	"github.com/biogo/biogo/alphabet"
	"github.com/biogo/biogo/feat"
	"github.com/biogo/biogo/seq/linear"
	"github.com/shenwei356/bio/seqio/fastx"
)

type Query struct {
	Name      string
	Seq       string
	Strand    string
	NullScore float64
}

type Range struct {
	Start float64
	End   float64
}

func (r Range) Len() float64 {
	return r.End - r.Start
}

type Ranges []Range

type Reference struct {
	Name   string
	Seq    string
	Ranges Ranges
}

type Queries []*Query

type SeqDetector struct {
	Queries   Queries
	SearchAll bool
	Stranded  bool
	NullMode  string
	Cutoff    float64
}

func NewSeqDetector(searchAll bool, stranded bool, nullMode string, cutoff float64) *SeqDetector {
	return &SeqDetector{Queries{}, searchAll, stranded, nullMode, cutoff}
}

func (d *SeqDetector) Detect(r *Reference, rec bool) []*AlignedSeq {
	var h []*AlignedSeq
	for _, rr := range r.Ranges {
		if rec {
			h = append(h, d.detectRec(r, rr)...)
		} else {
			h = append(h, d.detectOnce(r, rr)...)
		}
	}
	return h
}

func actualRange(rr Range, l int) Range {
	s, e := rr.Start, rr.End
	if s == e {
		return rr
	}
	if math.IsNaN(rr.Start) {
		s = 0
	} else if rr.Start < 0 {
		s = float64(l) + rr.Start
	}
	if s < 0 {
		s = 0
	}
	if math.IsNaN(rr.End) {
		e = float64(l)
	} else if rr.End < 0 {
		e = float64(l) + rr.End
	}
	if e > float64(l) {
		e = float64(l)
	}
	if s > e {
		s, e = 0, float64(l)
	}

	return Range{s, e}
}

func (d *SeqDetector) detectOnce(r *Reference, rr Range) []*AlignedSeq {
	var hits []*AlignedSeq
	if rr.Len() == 0 {
		return hits
	}
	for _, q := range d.Queries {
		nr := &Reference{r.Name, r.Seq, Ranges{actualRange(rr, len(r.Seq))}}
		h := PairwiseAlignSW(nr, q)
		h.Detector = d
		if (h.Score / q.NullScore) > d.Cutoff {
			hits = append(hits, h)
		}
	}
	return bestHits(hits, -1)
}

func (d *SeqDetector) detectRec(r *Reference, rr Range) []*AlignedSeq {
	var hits []*AlignedSeq
	if rr.Len() == 0 {
		return hits
	}
	for _, q := range d.Queries {
		nr := &Reference{r.Name, r.Seq, Ranges{actualRange(rr, len(r.Seq))}}
		h := PairwiseAlignSW(nr, q)
		h.Detector = d
		if (h.Score / q.NullScore) > d.Cutoff {
			hits = append(hits, h)
		}
		if len(hits) > 0 {
			bh := bestHits(hits, 1)
			bh = append(bh, d.detectRec(r, Range{rr.Start, float64(bh[0].RefStart)})...)
			bh = append(bh, d.detectRec(r, Range{float64(bh[0].RefEnd), rr.End})...)
			hits = bh
		}
	}
	return bestHits(hits, -1)
}

type byScore []*AlignedSeq

func (a byScore) Len() int           { return len(a) }
func (a byScore) Less(i, j int) bool { return a[i].Score > a[j].Score }
func (a byScore) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func bestHits(h []*AlignedSeq, n int) []*AlignedSeq {
	if len(h) == 0 {
		return h
	}
	if n < 0 {
		n = len(h)
	}
	res := make([]*AlignedSeq, len(h))
	copy(res, h)
	sort.Stable(byScore(res))
	i := n
	if i > len(res) {
		i = len(res)
	}
	res[0].Best = true
	return res[:i]
}

func (d *SeqDetector) LoadQueries(fx string) {
	var record *fastx.Record
	var fastxReader *fastx.Reader
	var err error
	if len(d.Queries) == 0 {
		d.Queries = make(Queries, 0, 10)
	}

	fastxReader, err = fastx.NewReader(nil, fx, "")
	checkError(err)

	for {
		record, err = fastxReader.Read()
		if err != nil {
			if err == io.EOF {
				break
			}
			checkError(err)
			break
		}

		seq := string(record.Seq.Seq)
		name := strings.Split(string(record.Name), " ")[0]
		ns := d.nullScore(seq)
		d.Queries = append(d.Queries, &Query{name, seq, "+", ns})
		if !d.Stranded {
			d.Queries = append(d.Queries, &Query{name, RevCompDNA(seq), "-", ns})
		}

	}
}

func (d *SeqDetector) AddAnonQueries(qrs []string) {
	for i, q := range qrs {
		name := fmt.Sprintf("q%d", i)
		ns := d.nullScore(q)
		d.Queries = append(d.Queries, &Query{name, q, "+", ns})
		if !d.Stranded {
			d.Queries = append(d.Queries, &Query{name, RevCompDNA(q), "-", ns})
		}
	}

}

func (d *SeqDetector) nullScore(q string) float64 {
	switch d.NullMode {
	case "self":
		return PairwiseAlignSW(&Reference{Name: "Ref", Seq: q, Ranges: Ranges{Range{0, float64(len(q))}}}, &Query{Name: "Query", Seq: q}).Score
	}
	return math.NaN()
}

//Interface for getting alignment score.
type Scorer interface {
	Score() int
}

// Make a new anonymous linear.Seq.
func NewAnonLinearSeq(s string) *linear.Seq {
	return &linear.Seq{Seq: alphabet.BytesToLetters([]byte(s))}
}

//Pairwise alignment of two sequences by biogo.
func PairwiseAlignSW(r *Reference, q *Query) *AlignedSeq {
	ref := NewAnonLinearSeq(r.Seq[int(r.Ranges[0].Start):int(r.Ranges[0].End)])
	ref.Alpha = alphabet.DNAgapped
	query := NewAnonLinearSeq(q.Seq)
	query.Alpha = alphabet.DNAgapped

	smith := align.SWAffine{
		Matrix: align.Linear{
			{0, -1, -1, -1, -1},
			{-1, 1, -1, -1, -1},
			{-1, -1, 1, -1, -1},
			{-1, -1, -1, 1, -1},
			{-1, -1, -1, -1, 1},
		},
		GapOpen: -1,
	}

	aln, err := smith.Align(ref, query)
	var res *AlignedSeq
	if err == nil {
		fa := align.Format(ref, query, aln, '-')
		res = AlignInfo(r, q, aln)
		res.RefAln = fmt.Sprintf("%s", fa[0])
		res.QueryAln = fmt.Sprintf("%s", fa[1])
	} else {
		panic("Could not align sequences!")
	}
	return res
}

type AlignedSeq struct {
	Ref        *Reference
	Query      *Query
	QueryAln   string
	RefAln     string
	RefStart   int
	RefEnd     int
	QueryStart int
	QueryEnd   int
	Score      float64
	Best       bool
	Detector   *SeqDetector
}

func (a *AlignedSeq) Fields() []string {
	validFields := []string{"Ref", "RefStart", "RefEnd", "Query", "QueryStart", "QueryEnd", "Strand", "MapQual", "RawScore", "Acc", "ClipAcc", "QueryCov"}
	return validFields
}
func (a *AlignedSeq) String() string {
	validFields := []string{"Ref", "RefStart", "RefEnd", "Query", "QueryStart", "QueryEnd", "Strand", "MapQual", "RawScore", "Acc", "ClipAcc", "QueryCov"}
	fmap := make(map[string]func(*AlignedSeq) string)
	fmap["Ref"] = func(a *AlignedSeq) string {
		return a.Ref.Name
	}
	fmap["RefStart"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.RefStart)
	}
	fmap["RefEnd"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.RefEnd)
	}
	fmap["Query"] = func(a *AlignedSeq) string {
		return a.Query.Name
	}
	fmap["QueryStart"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.QueryStart)
	}
	fmap["QueryEnd"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.QueryEnd)
	}
	fmap["Strand"] = func(a *AlignedSeq) string {
		return a.Query.Strand
	}
	fmap["RawScore"] = func(a *AlignedSeq) string {
		return fmt.Sprintf("%.0f", a.Score)
	}
	fmap["MapQual"] = func(a *AlignedSeq) string {
		mq := -10 * math.Log10(1-a.Score/a.Query.NullScore)
		if math.IsInf(mq, 1) {
			mq = 60
		}
		return fmt.Sprintf("%.2f", mq)
	}
	fmap["Acc"] = func(a *AlignedSeq) string {
		diff := 0
		for i, rb := range []byte(a.RefAln) {
			if rb != a.QueryAln[i] {
				diff++
			}
		}
		length := float64(len(a.RefAln))
		acc := (length - float64(diff)) * 100 / length
		return fmt.Sprintf("%.2f", acc)
	}
	fmap["ClipAcc"] = func(a *AlignedSeq) string {
		diff := 0
		for i, rb := range []byte(a.RefAln) {
			if rb != a.QueryAln[i] {
				diff++
			}
		}
		length := float64(len(a.RefAln))
		acc := ((length - float64(diff)) * 100) / (length + float64(len(a.Query.Seq)-a.QueryEnd+a.QueryStart))
		return fmt.Sprintf("%.2f", acc)
	}
	fmap["QueryCov"] = func(a *AlignedSeq) string {
		acc := float64(a.QueryEnd-a.QueryStart) * 100 / float64(len(a.Query.Seq))
		return fmt.Sprintf("%.2f", acc)
	}

	tmp := make([]string, len(validFields))
	for i, f := range validFields {
		tmp[i] = fmap[f](a)
	}

	return strings.Join(tmp, "\t")
}

func (a *AlignedSeq) AlnString() string {
	return fmt.Sprintf("@\t%s\t+\t%d\t%d\t%s\n@\t%s\t%s\t%d\t%d\t%s", a.RefAln, a.RefStart, a.RefEnd, a.Ref.Name, a.QueryAln, a.Query.Strand, a.QueryStart, a.QueryEnd, a.Query.Name)
}

//Construct AlignedSeq structure based on raw alignment results.
func AlignInfo(r *Reference, q *Query, f []feat.Pair) *AlignedSeq {
	ref_starts := make([]int, 0)
	ref_ends := make([]int, 0)
	query_starts := make([]int, 0)
	query_ends := make([]int, 0)
	scores := make([]int, 0)

	for _, fs := range f {
		fc := fs.Features()
		fsScorer, _ := fs.(Scorer)
		scores = append(scores, fsScorer.Score())
		ref_starts = append(ref_starts, fc[0].Start())
		ref_ends = append(ref_ends, fc[0].End())
		query_starts = append(query_starts, fc[1].Start())
		query_ends = append(query_ends, fc[1].End())

	}
	res := &AlignedSeq{Ref: r, Query: q}
	res.RefStart = MinInts(ref_starts) + int(r.Ranges[0].Start)
	res.RefEnd = MaxInts(ref_ends) + int(r.Ranges[0].Start)
	res.QueryStart = MinInts(query_starts)
	res.QueryEnd = MaxInts(query_ends)
	res.Score = float64(SumInts(scores))
	return res
}
<commit_msg>fish: tweaked scoring system.<commit_after>package cmd

import (
	"fmt"
	"io"
	"math"
	//"os"
	"sort"
	"strconv"
	"strings"

	"github.com/biogo/biogo/align"
	"github.com/biogo/biogo/alphabet"
	"github.com/biogo/biogo/feat"
	"github.com/biogo/biogo/seq/linear"
	"github.com/shenwei356/bio/seqio/fastx"
)

type Query struct {
	Name      string
	Seq       string
	Strand    string
	NullScore float64
}

type Range struct {
	Start float64
	End   float64
}

func (r Range) Len() float64 {
	return r.End - r.Start
}

type Ranges []Range

type Reference struct {
	Name   string
	Seq    string
	Ranges Ranges
}

type Queries []*Query

type SeqDetector struct {
	Queries   Queries
	SearchAll bool
	Stranded  bool
	NullMode  string
	Cutoff    float64
}

func NewSeqDetector(searchAll bool, stranded bool, nullMode string, cutoff float64) *SeqDetector {
	return &SeqDetector{Queries{}, searchAll, stranded, nullMode, cutoff}
}

func (d *SeqDetector) Detect(r *Reference, rec bool) []*AlignedSeq {
	var h []*AlignedSeq
	for _, rr := range r.Ranges {
		if rec {
			h = append(h, d.detectRec(r, rr)...)
		} else {
			h = append(h, d.detectOnce(r, rr)...)
		}
	}
	return h
}

func actualRange(rr Range, l int) Range {
	s, e := rr.Start, rr.End
	if s == e {
		return rr
	}
	if math.IsNaN(rr.Start) {
		s = 0
	} else if rr.Start < 0 {
		s = float64(l) + rr.Start
	}
	if s < 0 {
		s = 0
	}
	if math.IsNaN(rr.End) {
		e = float64(l)
	} else if rr.End < 0 {
		e = float64(l) + rr.End
	}
	if e > float64(l) {
		e = float64(l)
	}
	if s > e {
		s, e = 0, float64(l)
	}

	return Range{s, e}
}

func (d *SeqDetector) detectOnce(r *Reference, rr Range) []*AlignedSeq {
	var hits []*AlignedSeq
	if rr.Len() == 0 {
		return hits
	}
	for _, q := range d.Queries {
		nr := &Reference{r.Name, r.Seq, Ranges{actualRange(rr, len(r.Seq))}}
		h := PairwiseAlignSW(nr, q)
		h.Detector = d
		if (h.Score / q.NullScore) > d.Cutoff {
			hits = append(hits, h)
		}
	}
	return bestHits(hits, -1)
}

func (d *SeqDetector) detectRec(r *Reference, rr Range) []*AlignedSeq {
	var hits []*AlignedSeq
	if rr.Len() == 0 {
		return hits
	}
	for _, q := range d.Queries {
		nr := &Reference{r.Name, r.Seq, Ranges{actualRange(rr, len(r.Seq))}}
		h := PairwiseAlignSW(nr, q)
		h.Detector = d
		if (h.Score / q.NullScore) > d.Cutoff {
			hits = append(hits, h)
		}
		if len(hits) > 0 {
			bh := bestHits(hits, 1)
			bh = append(bh, d.detectRec(r, Range{rr.Start, float64(bh[0].RefStart)})...)
			bh = append(bh, d.detectRec(r, Range{float64(bh[0].RefEnd), rr.End})...)
			hits = bh
		}
	}
	return bestHits(hits, -1)
}

type byScore []*AlignedSeq

func (a byScore) Len() int           { return len(a) }
func (a byScore) Less(i, j int) bool { return a[i].Score > a[j].Score }
func (a byScore) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func bestHits(h []*AlignedSeq, n int) []*AlignedSeq {
	if len(h) == 0 {
		return h
	}
	if n < 0 {
		n = len(h)
	}
	res := make([]*AlignedSeq, len(h))
	copy(res, h)
	sort.Stable(byScore(res))
	i := n
	if i > len(res) {
		i = len(res)
	}
	res[0].Best = true
	return res[:i]
}

func (d *SeqDetector) LoadQueries(fx string) {
	var record *fastx.Record
	var fastxReader *fastx.Reader
	var err error
	if len(d.Queries) == 0 {
		d.Queries = make(Queries, 0, 10)
	}

	fastxReader, err = fastx.NewReader(nil, fx, "")
	checkError(err)

	for {
		record, err = fastxReader.Read()
		if err != nil {
			if err == io.EOF {
				break
			}
			checkError(err)
			break
		}

		seq := string(record.Seq.Seq)
		name := strings.Split(string(record.Name), " ")[0]
		ns := d.nullScore(seq)
		d.Queries = append(d.Queries, &Query{name, seq, "+", ns})
		if !d.Stranded {
			d.Queries = append(d.Queries, &Query{name, RevCompDNA(seq), "-", ns})
		}

	}
}

func (d *SeqDetector) AddAnonQueries(qrs []string) {
	for i, q := range qrs {
		name := fmt.Sprintf("q%d", i)
		ns := d.nullScore(q)
		d.Queries = append(d.Queries, &Query{name, q, "+", ns})
		if !d.Stranded {
			d.Queries = append(d.Queries, &Query{name, RevCompDNA(q), "-", ns})
		}
	}

}

func (d *SeqDetector) nullScore(q string) float64 {
	switch d.NullMode {
	case "self":
		return PairwiseAlignSW(&Reference{Name: "Ref", Seq: q, Ranges: Ranges{Range{0, float64(len(q))}}}, &Query{Name: "Query", Seq: q}).Score
	}
	return math.NaN()
}

//Interface for getting alignment score.
type Scorer interface {
	Score() int
}

// Make a new anonymous linear.Seq.
func NewAnonLinearSeq(s string) *linear.Seq {
	return &linear.Seq{Seq: alphabet.BytesToLetters([]byte(s))}
}

//Pairwise alignment of two sequences by biogo.
func PairwiseAlignSW(r *Reference, q *Query) *AlignedSeq {
	ref := NewAnonLinearSeq(r.Seq[int(r.Ranges[0].Start):int(r.Ranges[0].End)])
	ref.Alpha = alphabet.DNAgapped
	query := NewAnonLinearSeq(q.Seq)
	query.Alpha = alphabet.DNAgapped

	smith := align.SWAffine{
		Matrix: align.Linear{
			{0, -1, -1, -1, -1},
			{-1, 4, -4, -4, -4},
			{-1, -4, 4, -4, -4},
			{-1, -4, -4, 4, -4},
			{-1, -4, -4, -4, 4},
		},
		GapOpen: -2,
	}

	aln, err := smith.Align(ref, query)
	var res *AlignedSeq
	if err == nil {
		fa := align.Format(ref, query, aln, '-')
		res = AlignInfo(r, q, aln)
		res.RefAln = fmt.Sprintf("%s", fa[0])
		res.QueryAln = fmt.Sprintf("%s", fa[1])
	} else {
		panic("Could not align sequences!")
	}
	return res
}

type AlignedSeq struct {
	Ref        *Reference
	Query      *Query
	QueryAln   string
	RefAln     string
	RefStart   int
	RefEnd     int
	QueryStart int
	QueryEnd   int
	Score      float64
	Best       bool
	Detector   *SeqDetector
}

func (a *AlignedSeq) Fields() []string {
	validFields := []string{"Ref", "RefStart", "RefEnd", "Query", "QueryStart", "QueryEnd", "Strand", "MapQual", "RawScore", "Acc", "ClipAcc", "QueryCov"}
	return validFields
}
func (a *AlignedSeq) String() string {
	validFields := []string{"Ref", "RefStart", "RefEnd", "Query", "QueryStart", "QueryEnd", "Strand", "MapQual", "RawScore", "Acc", "ClipAcc", "QueryCov"}
	fmap := make(map[string]func(*AlignedSeq) string)
	fmap["Ref"] = func(a *AlignedSeq) string {
		return a.Ref.Name
	}
	fmap["RefStart"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.RefStart)
	}
	fmap["RefEnd"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.RefEnd)
	}
	fmap["Query"] = func(a *AlignedSeq) string {
		return a.Query.Name
	}
	fmap["QueryStart"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.QueryStart)
	}
	fmap["QueryEnd"] = func(a *AlignedSeq) string {
		return strconv.Itoa(a.QueryEnd)
	}
	fmap["Strand"] = func(a *AlignedSeq) string {
		return a.Query.Strand
	}
	fmap["RawScore"] = func(a *AlignedSeq) string {
		return fmt.Sprintf("%.0f", a.Score)
	}
	fmap["MapQual"] = func(a *AlignedSeq) string {
		mq := -10 * math.Log10(1-a.Score/a.Query.NullScore)
		if math.IsInf(mq, 1) {
			mq = 60
		}
		return fmt.Sprintf("%.2f", mq)
	}
	fmap["Acc"] = func(a *AlignedSeq) string {
		diff := 0
		for i, rb := range []byte(a.RefAln) {
			if rb != a.QueryAln[i] {
				diff++
			}
		}
		length := float64(len(a.RefAln))
		acc := (length - float64(diff)) * 100 / length
		return fmt.Sprintf("%.2f", acc)
	}
	fmap["ClipAcc"] = func(a *AlignedSeq) string {
		diff := 0
		for i, rb := range []byte(a.RefAln) {
			if rb != a.QueryAln[i] {
				diff++
			}
		}
		length := float64(len(a.RefAln))
		acc := ((length - float64(diff)) * 100) / (length + float64(len(a.Query.Seq)-a.QueryEnd+a.QueryStart))
		return fmt.Sprintf("%.2f", acc)
	}
	fmap["QueryCov"] = func(a *AlignedSeq) string {
		acc := float64(a.QueryEnd-a.QueryStart) * 100 / float64(len(a.Query.Seq))
		return fmt.Sprintf("%.2f", acc)
	}

	tmp := make([]string, len(validFields))
	for i, f := range validFields {
		tmp[i] = fmap[f](a)
	}

	return strings.Join(tmp, "\t")
}

func (a *AlignedSeq) AlnString() string {
	return fmt.Sprintf("@\t%s\t+\t%d\t%d\t%s\n@\t%s\t%s\t%d\t%d\t%s", a.RefAln, a.RefStart, a.RefEnd, a.Ref.Name, a.QueryAln, a.Query.Strand, a.QueryStart, a.QueryEnd, a.Query.Name)
}

//Construct AlignedSeq structure based on raw alignment results.
func AlignInfo(r *Reference, q *Query, f []feat.Pair) *AlignedSeq {
	ref_starts := make([]int, 0)
	ref_ends := make([]int, 0)
	query_starts := make([]int, 0)
	query_ends := make([]int, 0)
	scores := make([]int, 0)

	for _, fs := range f {
		fc := fs.Features()
		fsScorer, _ := fs.(Scorer)
		scores = append(scores, fsScorer.Score())
		ref_starts = append(ref_starts, fc[0].Start())
		ref_ends = append(ref_ends, fc[0].End())
		query_starts = append(query_starts, fc[1].Start())
		query_ends = append(query_ends, fc[1].End())

	}
	res := &AlignedSeq{Ref: r, Query: q}
	res.RefStart = MinInts(ref_starts) + int(r.Ranges[0].Start)
	res.RefEnd = MaxInts(ref_ends) + int(r.Ranges[0].Start)
	res.QueryStart = MinInts(query_starts)
	res.QueryEnd = MaxInts(query_ends)
	res.Score = float64(SumInts(scores))
	return res
}
<|endoftext|>


<commit_before>// The plugin package provides the functionality to both expose a Packer
// plugin binary and to connect to an existing Packer plugin binary.
//
// Packer supports plugins in the form of self-contained external static
// Go binaries. These binaries behave in a certain way (enforced by this
// package) and are connected to in a certain way (also enforced by this
// package).
package plugin

import (
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	packrpc "github.com/mitchellh/packer/packer/rpc"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/signal"
	"runtime"
	"strconv"
)

const MagicCookieKey = "PACKER_PLUGIN_MAGIC_COOKIE"
const MagicCookieValue = "d602bf8f470bc67ca7faa0386276bbdd4330efaf76d1a219cb4d6991ca9872b2"

// This serves a single RPC connection on the given RPC server on
// a random port.
func serve(server *rpc.Server) (err error) {
	if os.Getenv(MagicCookieKey) != MagicCookieValue {
		return errors.New("Please do not execute plugins directly. Packer will execute these for you.")
	}

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	minPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MIN_PORT"), 10, 32)
	if err != nil {
		return
	}

	maxPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MAX_PORT"), 10, 32)
	if err != nil {
		return
	}

	log.Printf("Plugin minimum port: %d\n", minPort)
	log.Printf("Plugin maximum port: %d\n", maxPort)

	// Set the RPC port range
	packrpc.PortRange(int(minPort), int(maxPort))

	var address string
	var listener net.Listener
	for port := minPort; port <= maxPort; port++ {
		address = fmt.Sprintf("127.0.0.1:%d", port)
		listener, err = net.Listen("tcp", address)
		if err != nil {
			err = nil
			continue
		}

		break
	}

	defer listener.Close()

	// Output the address to stdout
	log.Printf("Plugin address: %s\n", address)
	fmt.Println(address)
	os.Stdout.Sync()

	// Accept a connection
	log.Println("Waiting for connection...")
	conn, err := listener.Accept()
	if err != nil {
		log.Printf("Error accepting connection: %s\n", err.Error())
		return
	}

	// Serve a single connection
	log.Println("Serving a plugin connection...")
	server.ServeConn(conn)
	return
}

// Registers a signal handler to "swallow" interrupts so that the
// plugin isn't killed. The main host Packer process is responsible
// for killing the plugins when interrupted.
func swallowInterrupts() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)

	go func() {
		for {
			<-ch
			log.Println("Received interrupt signal. Ignoring.")
		}
	}()
}

// Serves a builder from a plugin.
func ServeBuilder(builder packer.Builder) {
	log.Println("Preparing to serve a builder plugin...")

	server := rpc.NewServer()
	packrpc.RegisterBuilder(server, builder)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a command from a plugin.
func ServeCommand(command packer.Command) {
	log.Println("Preparing to serve a command plugin...")

	server := rpc.NewServer()
	packrpc.RegisterCommand(server, command)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a hook from a plugin.
func ServeHook(hook packer.Hook) {
	log.Println("Preparing to serve a hook plugin...")

	server := rpc.NewServer()
	packrpc.RegisterHook(server, hook)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a post-processor from a plugin.
func ServePostProcessor(p packer.PostProcessor) {
	log.Println("Preparing to serve a post-processor plugin...")

	server := rpc.NewServer()
	packrpc.RegisterPostProcessor(server, p)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a provisioner from a plugin.
func ServeProvisioner(p packer.Provisioner) {
	log.Println("Preparing to serve a provisioner plugin...")

	server := rpc.NewServer()
	packrpc.RegisterProvisioner(server, p)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}
<commit_msg>packer/plugin: Count number of interrupts atomically<commit_after>// The plugin package provides the functionality to both expose a Packer
// plugin binary and to connect to an existing Packer plugin binary.
//
// Packer supports plugins in the form of self-contained external static
// Go binaries. These binaries behave in a certain way (enforced by this
// package) and are connected to in a certain way (also enforced by this
// package).
package plugin

import (
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	packrpc "github.com/mitchellh/packer/packer/rpc"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"sync/atomic"
)

// This is a count of the number of interrupts the process has received.
// This is updated with sync/atomic whenever a SIGINT is received and can
// be checked by the plugin safely to take action.
var Interrupts int32 = 0

const MagicCookieKey = "PACKER_PLUGIN_MAGIC_COOKIE"
const MagicCookieValue = "d602bf8f470bc67ca7faa0386276bbdd4330efaf76d1a219cb4d6991ca9872b2"

// This serves a single RPC connection on the given RPC server on
// a random port.
func serve(server *rpc.Server) (err error) {
	if os.Getenv(MagicCookieKey) != MagicCookieValue {
		return errors.New("Please do not execute plugins directly. Packer will execute these for you.")
	}

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	minPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MIN_PORT"), 10, 32)
	if err != nil {
		return
	}

	maxPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MAX_PORT"), 10, 32)
	if err != nil {
		return
	}

	log.Printf("Plugin minimum port: %d\n", minPort)
	log.Printf("Plugin maximum port: %d\n", maxPort)

	// Set the RPC port range
	packrpc.PortRange(int(minPort), int(maxPort))

	var address string
	var listener net.Listener
	for port := minPort; port <= maxPort; port++ {
		address = fmt.Sprintf("127.0.0.1:%d", port)
		listener, err = net.Listen("tcp", address)
		if err != nil {
			err = nil
			continue
		}

		break
	}

	defer listener.Close()

	// Output the address to stdout
	log.Printf("Plugin address: %s\n", address)
	fmt.Println(address)
	os.Stdout.Sync()

	// Accept a connection
	log.Println("Waiting for connection...")
	conn, err := listener.Accept()
	if err != nil {
		log.Printf("Error accepting connection: %s\n", err.Error())
		return
	}

	// Serve a single connection
	log.Println("Serving a plugin connection...")
	server.ServeConn(conn)
	return
}

// Registers a signal handler to swallow and count interrupts so that the
// plugin isn't killed. The main host Packer process is responsible
// for killing the plugins when interrupted.
func countInterrupts() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)

	go func() {
		for {
			<-ch
			newCount := atomic.AddInt32(&Interrupts, 1)
			log.Printf("Received interrupt signal (count: %d). Ignoring.", newCount)
		}
	}()
}

// Serves a builder from a plugin.
func ServeBuilder(builder packer.Builder) {
	log.Println("Preparing to serve a builder plugin...")

	server := rpc.NewServer()
	packrpc.RegisterBuilder(server, builder)

	countInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a command from a plugin.
func ServeCommand(command packer.Command) {
	log.Println("Preparing to serve a command plugin...")

	server := rpc.NewServer()
	packrpc.RegisterCommand(server, command)

	countInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a hook from a plugin.
func ServeHook(hook packer.Hook) {
	log.Println("Preparing to serve a hook plugin...")

	server := rpc.NewServer()
	packrpc.RegisterHook(server, hook)

	countInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a post-processor from a plugin.
func ServePostProcessor(p packer.PostProcessor) {
	log.Println("Preparing to serve a post-processor plugin...")

	server := rpc.NewServer()
	packrpc.RegisterPostProcessor(server, p)

	countInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a provisioner from a plugin.
func ServeProvisioner(p packer.Provisioner) {
	log.Println("Preparing to serve a provisioner plugin...")

	server := rpc.NewServer()
	packrpc.RegisterProvisioner(server, p)

	countInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2018 Bitmark Inc.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package rpc

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/rpc"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/sha3"
	"golang.org/x/time/rate"

	"github.com/bitmark-inc/bitmarkd/announce"
	"github.com/bitmark-inc/bitmarkd/fault"
	"github.com/bitmark-inc/bitmarkd/reservoir"
	"github.com/bitmark-inc/bitmarkd/util"
	"github.com/bitmark-inc/listener"
	"github.com/bitmark-inc/logger"
)

type RPCConfiguration struct {
	MaximumConnections int      `gluamapper:"maximum_connections" json:"maximum_connections"`
	Bandwidth          float64  `gluamapper:"bandwidth" json:"bandwidth"`
	Listen             []string `gluamapper:"listen" json:"listen"`
	Certificate        string   `gluamapper:"certificate" json:"certificate"`
	PrivateKey         string   `gluamapper:"private_key" json:"private_key"`
	Announce           []string `gluamapper:"announce" json:"announce"`
}

type HTTPSConfiguration struct {
	MaximumConnections int                 `gluamapper:"maximum_connections" json:"maximum_connections"`
	Listen             []string            `gluamapper:"listen" json:"listen"`
	Certificate        string              `gluamapper:"certificate" json:"certificate"`
	PrivateKey         string              `gluamapper:"private_key" json:"private_key"`
	Allow              map[string][]string `gluamapper:"allow" json:"allow"`
}

// rate limiting (requests per second)
// burst limit   (total items in one request)
const (
	rateLimitAssets = 200
	rateBurstAssets = 100

	rateLimitBitmark = 200
	rateBurstBitmark = 100

	rateLimitBitmarks = 200
	rateBurstBitmarks = reservoir.MaximumIssues

	rateLimitOwner = 200
	rateBurstOwner = 100

	rateLimitNode = 200
	rateBurstNode = 100

	rateLimitTransaction = 200
	rateBurstTransaction = 100

	rateLimitBlockOwner = 200
	rateBurstBlockOwner = 100

	rateLimitShare = 200
	rateBurstShare = 100
)

// globals
type rpcData struct {
	sync.RWMutex // to allow locking

	log *logger.L // logger

	listener *listener.MultiListener

	httpServer *httpHandler

	// set once during initialise
	initialised bool
}

// global data
var globalData rpcData

// initialise peer backgrouds processes
func Initialise(rpcConfiguration *RPCConfiguration, httpsConfiguration *HTTPSConfiguration, version string) error {

	globalData.Lock()
	defer globalData.Unlock()

	// no need to start if already started
	if globalData.initialised {
		return fault.ErrAlreadyInitialised
	}

	log := logger.New("rpc")
	globalData.log = log
	log.Info("starting…")

	// servers
	err := initialiseRPC(rpcConfiguration, version)
	if nil != err {
		return err
	}
	err = initialiseHTTPS(httpsConfiguration, version)
	if nil != err {
		return err
	}

	// all data initialised
	globalData.initialised = true

	return nil
}

// finialise - stop all background tasks
func Finalise() error {

	if !globalData.initialised {
		return fault.ErrNotInitialised
	}

	globalData.log.Info("shutting down…")
	globalData.log.Flush()

	// stop background
	globalData.listener.Stop()

	// finally...
	globalData.initialised = false

	globalData.log.Info("finished")
	globalData.log.Flush()

	return nil
}

func initialiseRPC(configuration *RPCConfiguration, version string) error {
	name := "client_rpc"
	log := globalData.log

	if configuration.MaximumConnections <= 0 {
		log.Errorf("invalid %s maximum connection limit: %d", name, configuration.MaximumConnections)
		return fault.ErrMissingParameters
	}
	if configuration.Bandwidth <= 1000000 { // fail if < 1Mbps
		log.Errorf("invalid %s bandwidth: %d bps < 1Mbps", name, configuration.Bandwidth)
		return fault.ErrMissingParameters
	}

	if 0 == len(configuration.Listen) {
		log.Errorf("missing %s listen", name)
		return fault.ErrMissingParameters
	}

	// create limiter
	limiter := listener.NewBandwidthLimiter(configuration.MaximumConnections, configuration.Bandwidth)

	tlsConfiguration, fingerprint, err := getCertificate(log, name, configuration.Certificate, configuration.PrivateKey)
	if nil != err {
		return err
	}

	log.Infof("%s: SHA3-256 fingerprint: %x", name, fingerprint)

	log.Infof("multi listener for: %s", name)
	ml, err := listener.NewMultiListener(name, configuration.Listen, tlsConfiguration, limiter, Callback)
	if nil != err {
		log.Errorf("invalid %s listen addresses", name)
		return err
	}
	globalData.listener = ml

	// setup announce
	rpcs := make([]byte, 0, 100) // ***** FIX THIS: need a better default size
process_rpcs:
	for _, address := range configuration.Announce {
		if "" == address {
			continue process_rpcs
		}
		c, err := util.NewConnection(address)
		if nil != err {
			log.Errorf("invalid %s listen announce: %q  error: %s", name, address, err)
			return err
		}
		rpcs = append(rpcs, c.Pack()...)
	}
	err = announce.SetRPC(fingerprint, rpcs)
	if nil != err {
		log.Criticalf("announce.SetRPC error: %s", err)
		return err
	}

	server := createRPCServer(log, version)
	argument := &serverArgument{
		Log:    log,
		Server: server,
	}

	log.Infof("starting server: %s  with: %v", name, argument)
	globalData.listener.Start(argument)

	return nil
}

// Start server with Test instance as a service
func initialiseHTTPS(configuration *HTTPSConfiguration, version string) error {

	name := "http_rpc"
	log := globalData.log

	if 0 == len(configuration.Listen) {
		log.Infof("disable: %s", name)
		return nil
	}

	if configuration.MaximumConnections <= 0 {
		log.Errorf("invalid %s maximum connection limit: %d", name, configuration.MaximumConnections)
		return fault.ErrMissingParameters
	}

	tlsConfiguration, fingerprint, err := getCertificate(globalData.log, name, configuration.Certificate, configuration.PrivateKey)
	if nil != err {
		return err
	}

	log.Infof("%s: SHA3-256 fingerprint: %x", name, fingerprint)

	// create access control and format strings to match http.Request.RemoteAddr
	local := make(map[string]map[string]struct{})
	for path, addresses := range configuration.Allow {
		set := make(map[string]struct{})
		local[path] = set
	local_loop:
		for _, ip := range addresses {
			ip := net.ParseIP(strings.Trim(ip, " "))
			if nil == ip {
				continue local_loop
			}
			if nil != ip.To4() {
				set[ip.String()] = struct{}{}
			} else {

				set["["+ip.String()+"]"] = struct{}{}
			}
		}
	}

	server := createRPCServer(log, version)
	handler := &httpHandler{
		log:     log,
		server:  server,
		version: version,
		start:   time.Now(),
		allow:   local,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/bitmarkd/rpc", handler.rpc)
	mux.HandleFunc("/bitmarkd/details", handler.details)
	mux.HandleFunc("/bitmarkd/connections", handler.connections)
	mux.HandleFunc("/bitmarkd/peers", handler.peers)
	mux.HandleFunc("/", handler.root)

	for _, listen := range configuration.Listen {
		log.Infof("starting server: %s on: %q", name, listen)

		go ListenAndServeTLSKeyPair(listen, mux, tlsConfiguration)
	}

	return nil
}

func createRPCServer(log *logger.L, version string) *rpc.Server {

	start := time.Now().UTC()

	assets := &Assets{
		log:     log,
		limiter: rate.NewLimiter(rateLimitAssets, rateBurstAssets),
	}

	bitmark := &Bitmark{
		log:     log,
		limiter: rate.NewLimiter(rateLimitBitmark, rateBurstBitmark),
	}

	bitmarks := &Bitmarks{
		log:     log,
		limiter: rate.NewLimiter(rateLimitBitmarks, rateBurstBitmarks),
	}

	owner := &Owner{
		log:     log,
		limiter: rate.NewLimiter(rateLimitOwner, rateBurstOwner),
	}

	node := &Node{
		log:     log,
		limiter: rate.NewLimiter(rateLimitNode, rateBurstNode),
		start:   start,
		version: version,
	}

	transaction := &Transaction{
		log:     log,
		limiter: rate.NewLimiter(rateLimitTransaction, rateBurstTransaction),
		start:   start,
	}

	blockOwner := &BlockOwner{
		log:     log,
		limiter: rate.NewLimiter(rateLimitBlockOwner, rateBurstBlockOwner),
	}

	share := &Share{
		log:     log,
		limiter: rate.NewLimiter(rateLimitShare, rateBurstShare),
	}

	server := rpc.NewServer()

	server.Register(assets)
	server.Register(bitmark)
	server.Register(bitmarks)
	server.Register(owner)
	server.Register(node)
	server.Register(transaction)
	server.Register(blockOwner)
	server.Register(share)

	return server
}

// Verify that a set of listener parameters are valid
// and return the certificate
func getCertificate(log *logger.L, name, certificate, key string) (*tls.Config, [32]byte, error) {
	var fingerprint [32]byte

	keyPair, err := tls.X509KeyPair([]byte(certificate), []byte(key))
	if err != nil {
		log.Errorf("%s failed to load keypair: %v", name, err)
		return nil, fingerprint, err
	}

	tlsConfiguration := &tls.Config{
		Certificates: []tls.Certificate{
			keyPair,
		},
	}

	fingerprint = CertificateFingerprint(keyPair.Certificate[0])

	return tlsConfiguration, fingerprint, nil
}

// compute the fingerprint of a certificate
//
// FreeBSD: openssl x509 -outform DER -in bitmarkd-local-rpc.crt | sha3sum -a 256
func CertificateFingerprint(certificate []byte) [32]byte {
	return sha3.Sum256(certificate)
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// start a HTTPS server using in-memory TLS KeyPair
func ListenAndServeTLSKeyPair(addr string, handler http.Handler, cfg *tls.Config) error {
	s := &http.Server{
		Addr:           addr,
		Handler:        handler,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	cfg.NextProtos = []string{"http/1.1"}

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, cfg)

	return s.Serve(tlsListener)
}
<commit_msg>[rpc] Fix typo of Limiter<commit_after>// Copyright (c) 2014-2018 Bitmark Inc.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package rpc

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/rpc"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/sha3"
	"golang.org/x/time/rate"

	"github.com/bitmark-inc/bitmarkd/announce"
	"github.com/bitmark-inc/bitmarkd/fault"
	"github.com/bitmark-inc/bitmarkd/reservoir"
	"github.com/bitmark-inc/bitmarkd/util"
	"github.com/bitmark-inc/listener"
	"github.com/bitmark-inc/logger"
)

type RPCConfiguration struct {
	MaximumConnections int      `gluamapper:"maximum_connections" json:"maximum_connections"`
	Bandwidth          float64  `gluamapper:"bandwidth" json:"bandwidth"`
	Listen             []string `gluamapper:"listen" json:"listen"`
	Certificate        string   `gluamapper:"certificate" json:"certificate"`
	PrivateKey         string   `gluamapper:"private_key" json:"private_key"`
	Announce           []string `gluamapper:"announce" json:"announce"`
}

type HTTPSConfiguration struct {
	MaximumConnections int                 `gluamapper:"maximum_connections" json:"maximum_connections"`
	Listen             []string            `gluamapper:"listen" json:"listen"`
	Certificate        string              `gluamapper:"certificate" json:"certificate"`
	PrivateKey         string              `gluamapper:"private_key" json:"private_key"`
	Allow              map[string][]string `gluamapper:"allow" json:"allow"`
}

// rate limiting (requests per second)
// burst limit   (total items in one request)
const (
	rateLimitAssets = 200
	rateBurstAssets = 100

	rateLimitBitmark = 200
	rateBurstBitmark = 100

	rateLimitBitmarks = 200
	rateBurstBitmarks = reservoir.MaximumIssues

	rateLimitOwner = 200
	rateBurstOwner = 100

	rateLimitNode = 200
	rateBurstNode = 100

	rateLimitTransaction = 200
	rateBurstTransaction = 100

	rateLimitBlockOwner = 200
	rateBurstBlockOwner = 100

	rateLimitShare = 200
	rateBurstShare = 100
)

// globals
type rpcData struct {
	sync.RWMutex // to allow locking

	log *logger.L // logger

	listener *listener.MultiListener

	httpServer *httpHandler

	// set once during initialise
	initialised bool
}

// global data
var globalData rpcData

// initialise peer backgrouds processes
func Initialise(rpcConfiguration *RPCConfiguration, httpsConfiguration *HTTPSConfiguration, version string) error {

	globalData.Lock()
	defer globalData.Unlock()

	// no need to start if already started
	if globalData.initialised {
		return fault.ErrAlreadyInitialised
	}

	log := logger.New("rpc")
	globalData.log = log
	log.Info("starting…")

	// servers
	err := initialiseRPC(rpcConfiguration, version)
	if nil != err {
		return err
	}
	err = initialiseHTTPS(httpsConfiguration, version)
	if nil != err {
		return err
	}

	// all data initialised
	globalData.initialised = true

	return nil
}

// finialise - stop all background tasks
func Finalise() error {

	if !globalData.initialised {
		return fault.ErrNotInitialised
	}

	globalData.log.Info("shutting down…")
	globalData.log.Flush()

	// stop background
	globalData.listener.Stop()

	// finally...
	globalData.initialised = false

	globalData.log.Info("finished")
	globalData.log.Flush()

	return nil
}

func initialiseRPC(configuration *RPCConfiguration, version string) error {
	name := "client_rpc"
	log := globalData.log

	if configuration.MaximumConnections <= 0 {
		log.Errorf("invalid %s maximum connection limit: %d", name, configuration.MaximumConnections)
		return fault.ErrMissingParameters
	}
	if configuration.Bandwidth <= 1000000 { // fail if < 1Mbps
		log.Errorf("invalid %s bandwidth: %d bps < 1Mbps", name, configuration.Bandwidth)
		return fault.ErrMissingParameters
	}

	if 0 == len(configuration.Listen) {
		log.Errorf("missing %s listen", name)
		return fault.ErrMissingParameters
	}

	// create limiter
	limiter := listener.NewLimiter(configuration.MaximumConnections)

	tlsConfiguration, fingerprint, err := getCertificate(log, name, configuration.Certificate, configuration.PrivateKey)
	if nil != err {
		return err
	}

	log.Infof("%s: SHA3-256 fingerprint: %x", name, fingerprint)

	log.Infof("multi listener for: %s", name)
	ml, err := listener.NewMultiListener(name, configuration.Listen, tlsConfiguration, limiter, Callback)
	if nil != err {
		log.Errorf("invalid %s listen addresses", name)
		return err
	}
	globalData.listener = ml

	// setup announce
	rpcs := make([]byte, 0, 100) // ***** FIX THIS: need a better default size
process_rpcs:
	for _, address := range configuration.Announce {
		if "" == address {
			continue process_rpcs
		}
		c, err := util.NewConnection(address)
		if nil != err {
			log.Errorf("invalid %s listen announce: %q  error: %s", name, address, err)
			return err
		}
		rpcs = append(rpcs, c.Pack()...)
	}
	err = announce.SetRPC(fingerprint, rpcs)
	if nil != err {
		log.Criticalf("announce.SetRPC error: %s", err)
		return err
	}

	server := createRPCServer(log, version)
	argument := &serverArgument{
		Log:    log,
		Server: server,
	}

	log.Infof("starting server: %s  with: %v", name, argument)
	globalData.listener.Start(argument)

	return nil
}

// Start server with Test instance as a service
func initialiseHTTPS(configuration *HTTPSConfiguration, version string) error {

	name := "http_rpc"
	log := globalData.log

	if 0 == len(configuration.Listen) {
		log.Infof("disable: %s", name)
		return nil
	}

	if configuration.MaximumConnections <= 0 {
		log.Errorf("invalid %s maximum connection limit: %d", name, configuration.MaximumConnections)
		return fault.ErrMissingParameters
	}

	tlsConfiguration, fingerprint, err := getCertificate(globalData.log, name, configuration.Certificate, configuration.PrivateKey)
	if nil != err {
		return err
	}

	log.Infof("%s: SHA3-256 fingerprint: %x", name, fingerprint)

	// create access control and format strings to match http.Request.RemoteAddr
	local := make(map[string]map[string]struct{})
	for path, addresses := range configuration.Allow {
		set := make(map[string]struct{})
		local[path] = set
	local_loop:
		for _, ip := range addresses {
			ip := net.ParseIP(strings.Trim(ip, " "))
			if nil == ip {
				continue local_loop
			}
			if nil != ip.To4() {
				set[ip.String()] = struct{}{}
			} else {

				set["["+ip.String()+"]"] = struct{}{}
			}
		}
	}

	server := createRPCServer(log, version)
	handler := &httpHandler{
		log:     log,
		server:  server,
		version: version,
		start:   time.Now(),
		allow:   local,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/bitmarkd/rpc", handler.rpc)
	mux.HandleFunc("/bitmarkd/details", handler.details)
	mux.HandleFunc("/bitmarkd/connections", handler.connections)
	mux.HandleFunc("/bitmarkd/peers", handler.peers)
	mux.HandleFunc("/", handler.root)

	for _, listen := range configuration.Listen {
		log.Infof("starting server: %s on: %q", name, listen)

		go ListenAndServeTLSKeyPair(listen, mux, tlsConfiguration)
	}

	return nil
}

func createRPCServer(log *logger.L, version string) *rpc.Server {

	start := time.Now().UTC()

	assets := &Assets{
		log:     log,
		limiter: rate.NewLimiter(rateLimitAssets, rateBurstAssets),
	}

	bitmark := &Bitmark{
		log:     log,
		limiter: rate.NewLimiter(rateLimitBitmark, rateBurstBitmark),
	}

	bitmarks := &Bitmarks{
		log:     log,
		limiter: rate.NewLimiter(rateLimitBitmarks, rateBurstBitmarks),
	}

	owner := &Owner{
		log:     log,
		limiter: rate.NewLimiter(rateLimitOwner, rateBurstOwner),
	}

	node := &Node{
		log:     log,
		limiter: rate.NewLimiter(rateLimitNode, rateBurstNode),
		start:   start,
		version: version,
	}

	transaction := &Transaction{
		log:     log,
		limiter: rate.NewLimiter(rateLimitTransaction, rateBurstTransaction),
		start:   start,
	}

	blockOwner := &BlockOwner{
		log:     log,
		limiter: rate.NewLimiter(rateLimitBlockOwner, rateBurstBlockOwner),
	}

	share := &Share{
		log:     log,
		limiter: rate.NewLimiter(rateLimitShare, rateBurstShare),
	}

	server := rpc.NewServer()

	server.Register(assets)
	server.Register(bitmark)
	server.Register(bitmarks)
	server.Register(owner)
	server.Register(node)
	server.Register(transaction)
	server.Register(blockOwner)
	server.Register(share)

	return server
}

// Verify that a set of listener parameters are valid
// and return the certificate
func getCertificate(log *logger.L, name, certificate, key string) (*tls.Config, [32]byte, error) {
	var fingerprint [32]byte

	keyPair, err := tls.X509KeyPair([]byte(certificate), []byte(key))
	if err != nil {
		log.Errorf("%s failed to load keypair: %v", name, err)
		return nil, fingerprint, err
	}

	tlsConfiguration := &tls.Config{
		Certificates: []tls.Certificate{
			keyPair,
		},
	}

	fingerprint = CertificateFingerprint(keyPair.Certificate[0])

	return tlsConfiguration, fingerprint, nil
}

// compute the fingerprint of a certificate
//
// FreeBSD: openssl x509 -outform DER -in bitmarkd-local-rpc.crt | sha3sum -a 256
func CertificateFingerprint(certificate []byte) [32]byte {
	return sha3.Sum256(certificate)
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// start a HTTPS server using in-memory TLS KeyPair
func ListenAndServeTLSKeyPair(addr string, handler http.Handler, cfg *tls.Config) error {
	s := &http.Server{
		Addr:           addr,
		Handler:        handler,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	cfg.NextProtos = []string{"http/1.1"}

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, cfg)

	return s.Serve(tlsListener)
}
<|endoftext|>


<commit_before>package sqlgen

import (
	"context"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/samsarahq/thunder/internal/testfixtures"
	"github.com/stretchr/testify/assert"
)

func setup() (*testfixtures.TestDatabase, *DB, error) {
	testDb, err := testfixtures.NewTestDatabase()
	if err != nil {
		return nil, nil, err
	}

	if _, err = testDb.Exec(`
		CREATE TABLE users (
			id   BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY,
			name VARCHAR(255),
			uuid VARCHAR(255),
			mood VARCHAR(255)
		)
	`); err != nil {
		return nil, nil, err
	}
	schema := NewSchema()
	schema.MustRegisterType("users", AutoIncrement, User{})

	return testDb, NewDB(testDb.DB, schema), nil
}

type User struct {
	Id   int64 `sql:",primary"`
	Name string
	Uuid testfixtures.CustomType
	Mood *testfixtures.CustomType
}

type Complex struct {
	Id       int64 `sql:",primary"`
	Name     string
	Text     []byte            `sql:",string"`
	Blob     []byte            `sql:",binary"`
	Mappings map[string]string `sql:",json"`
}

func TestTagOverrides(t *testing.T) {
	schema := NewSchema()
	err := schema.RegisterType("complex", AutoIncrement, Complex{})
	assert.NoError(t, err)
}

func TestContextDeadlineEnforced(t *testing.T) {
	tdb, db, err := setup()
	if err != nil {
		t.Fatal(err)
	}
	defer tdb.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	if _, err = db.QueryExecer(ctx).ExecContext(ctx, "DO SLEEP(1)"); err == nil || err != context.DeadlineExceeded {
		t.Errorf("Expected context.DeadlineExceeded, got: %s", err)
	}
}

func TestIntegrationBasic(t *testing.T) {
	tdb, db, err := setup()
	if err != nil {
		t.Fatal(err)
	}
	defer tdb.Close()

	mood := testfixtures.CustomType{'f', 'o', 'o', 'o', 'o', 'o', 'o'}

	if _, err := db.InsertRow(context.Background(), &User{Name: "Bob", Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'}, Mood: &mood}); err != nil {
		t.Error(err)
	}

	var users []*User
	if err := db.Query(context.Background(), &users, nil, nil); err != nil {
		t.Error(err)
	}

	assert.Equal(t, []*User{
		{
			Id:   1,
			Name: "Bob",
			Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'},
			Mood: &mood,
		},
	}, users)
}

// TestContextCancelBeforeRowsScan demonstrates we don't
// always get context.Canceled back from sql library. This
// affects our error handling and we need to be aware of it.
func TestContextCancelBeforeRowsScan(t *testing.T) {
	testDb, err := testfixtures.NewTestDatabase()
	if err != nil {
		t.Fatal(err)
	}
	defer testDb.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	rows, err := testDb.QueryContext(ctx, `select "foo"`)
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()

	// When we cancel the context after rows.Next() returns true,
	// database/sql.(*Rows).initContextClose monitors the context
	// and closes rows asynchronously, and subsequent rows.Scan()
	// returns errors.New("sql: Rows are closed") instead of
	// context.Canceled.
	for rows.Next() {
		cancel()
		time.Sleep(1000 * time.Millisecond)

		var foo string
		err := rows.Scan(&foo)

		// err is not context.Canceled.
		if err == nil || err.Error() != "sql: Rows are closed" {
			t.Fatalf("expecting 'sql: Rows are closed' from rows.Scan(), got %v", err)
		}
	}
	if err := rows.Err(); err != context.Canceled {
		t.Fatalf("expecting context.Canceled from rows.Err(), got %v", err)
	}
}

func Benchmark(b *testing.B) {
	tdb, db, err := setup()
	if err != nil {
		b.Fatal(err)
	}
	defer tdb.Close()

	ctx := context.Background()

	mood := testfixtures.CustomType{'f', 'o', 'o', 'o', 'o', 'o', 'o'}
	user := &User{
		Id:   1,
		Name: "Bob",
		Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'},
		Mood: &mood,
	}

	if _, err := db.InsertRow(ctx, user); err != nil {
		b.Fatal(err)
	}

	benchmarks := []struct {
		name string
		fn   func() error
	}{
		{"Read", func() error {
			user := &User{}
			return db.QueryRow(ctx, &user, nil, nil)
		}},
		{"Read_Where", func() error {
			user := &User{}
			return db.QueryRow(ctx, &user, Filter{"name": "Bob"}, nil)
		}},
		{"Create", func() error {
			_, err := db.InsertRow(ctx, user)
			return err
		}},
		{"Update", func() error { return db.UpdateRow(ctx, user) }},
		{"Delete", func() error { return db.DeleteRow(ctx, user) }},
	}

	for _, bm := range benchmarks {
		b.Run(bm.name, func(b *testing.B) {
			b.ReportAllocs()
			for i := 0; i < b.N; i++ {
				if err := bm.fn(); err != nil {
					b.Error(err)
				}
			}
		})
	}
}
<commit_msg>go sqlgen: test batched query matches rows against filter<commit_after>package sqlgen

import (
	"context"
	"database/sql"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/samsarahq/thunder/batch"
	"github.com/samsarahq/thunder/internal/testfixtures"
	"github.com/stretchr/testify/assert"
)

func setup() (*testfixtures.TestDatabase, *DB, error) {
	testDb, err := testfixtures.NewTestDatabase()
	if err != nil {
		return nil, nil, err
	}

	if _, err = testDb.Exec(`
		CREATE TABLE users (
			id   BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY,
			name VARCHAR(255),
			uuid VARCHAR(255),
			mood VARCHAR(255)
		)
	`); err != nil {
		return nil, nil, err
	}
	schema := NewSchema()
	schema.MustRegisterType("users", AutoIncrement, User{})

	return testDb, NewDB(testDb.DB, schema), nil
}

type User struct {
	Id   int64 `sql:",primary"`
	Name string
	Uuid testfixtures.CustomType
	Mood *testfixtures.CustomType
}

type Complex struct {
	Id       int64 `sql:",primary"`
	Name     string
	Text     []byte            `sql:",string"`
	Blob     []byte            `sql:",binary"`
	Mappings map[string]string `sql:",json"`
}

func TestTagOverrides(t *testing.T) {
	schema := NewSchema()
	err := schema.RegisterType("complex", AutoIncrement, Complex{})
	assert.NoError(t, err)
}

func TestContextDeadlineEnforced(t *testing.T) {
	tdb, db, err := setup()
	if err != nil {
		t.Fatal(err)
	}
	defer tdb.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	if _, err = db.QueryExecer(ctx).ExecContext(ctx, "DO SLEEP(1)"); err == nil || err != context.DeadlineExceeded {
		t.Errorf("Expected context.DeadlineExceeded, got: %s", err)
	}
}

func TestIntegrationBasic(t *testing.T) {
	tdb, db, err := setup()
	if err != nil {
		t.Fatal(err)
	}
	defer tdb.Close()

	mood := testfixtures.CustomType{'f', 'o', 'o', 'o', 'o', 'o', 'o'}

	if _, err := db.InsertRow(context.Background(), &User{Name: "Bob", Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'}, Mood: &mood}); err != nil {
		t.Error(err)
	}

	var users []*User
	if err := db.Query(context.Background(), &users, nil, nil); err != nil {
		t.Error(err)
	}

	assert.Equal(t, []*User{
		{
			Id:   1,
			Name: "Bob",
			Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'},
			Mood: &mood,
		},
	}, users)
}

// TestContextCancelBeforeRowsScan demonstrates we don't
// always get context.Canceled back from sql library. This
// affects our error handling and we need to be aware of it.
func TestContextCancelBeforeRowsScan(t *testing.T) {
	testDb, err := testfixtures.NewTestDatabase()
	if err != nil {
		t.Fatal(err)
	}
	defer testDb.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	rows, err := testDb.QueryContext(ctx, `select "foo"`)
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()

	// When we cancel the context after rows.Next() returns true,
	// database/sql.(*Rows).initContextClose monitors the context
	// and closes rows asynchronously, and subsequent rows.Scan()
	// returns errors.New("sql: Rows are closed") instead of
	// context.Canceled.
	for rows.Next() {
		cancel()
		time.Sleep(1000 * time.Millisecond)

		var foo string
		err := rows.Scan(&foo)

		// err is not context.Canceled.
		if err == nil || err.Error() != "sql: Rows are closed" {
			t.Fatalf("expecting 'sql: Rows are closed' from rows.Scan(), got %v", err)
		}
	}
	if err := rows.Err(); err != context.Canceled {
		t.Fatalf("expecting context.Canceled from rows.Err(), got %v", err)
	}
}

// TestBatchFilter shows sqlgen batching matcher does not match if
// filter type does not exactly match column type.
func TestBatchFilter(t *testing.T) {
	tdb, db, err := setup()
	if err != nil {
		t.Fatal(err)
	}
	defer tdb.Close()

	if _, err := db.InsertRow(context.Background(), &User{
		Id:   1,
		Name: "Bob",
		Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'},
		Mood: &testfixtures.CustomType{'f', 'o', 'o', 'o', 'o', 'o', 'o'},
	}); err != nil {
		t.Fatal(err)
	}

	ctx := batch.WithBatching(context.Background())

	var user *User

	// Only int64 works because id field is int64.
	if err := db.QueryRow(ctx, &user, Filter{"id": int64(1)}, nil); err != nil {
		t.Fatal(err)
	}

	// Other int variants don't work.
	if err := db.QueryRow(ctx, &user, Filter{"id": int32(1)}, nil); err != sql.ErrNoRows {
		t.Fatalf("expecting sql.ErrNoRows, got: %v", err)
	}
	if err := db.QueryRow(ctx, &user, Filter{"id": int16(1)}, nil); err != sql.ErrNoRows {
		t.Fatalf("expecting sql.ErrNoRows, got: %v", err)
	}
	if err := db.QueryRow(ctx, &user, Filter{"id": int8(1)}, nil); err != sql.ErrNoRows {
		t.Fatalf("expecting sql.ErrNoRows, got: %v", err)
	}
	if err := db.QueryRow(ctx, &user, Filter{"id": int(1)}, nil); err != sql.ErrNoRows {
		t.Fatalf("expecting sql.ErrNoRows, got: %v", err)
	}

	// Unsigned int does not work.
	if err := db.QueryRow(ctx, &user, Filter{"id": uint(1)}, nil); err != sql.ErrNoRows {
		t.Fatalf("expecting sql.ErrNoRows, got: %v", err)
	}

	// String does not work either.
	if err := db.QueryRow(ctx, &user, Filter{"id": "1"}, nil); err != sql.ErrNoRows {
		t.Fatalf("expecting sql.ErrNoRows, got: %v", err)
	}
}

func Benchmark(b *testing.B) {
	tdb, db, err := setup()
	if err != nil {
		b.Fatal(err)
	}
	defer tdb.Close()

	ctx := context.Background()

	mood := testfixtures.CustomType{'f', 'o', 'o', 'o', 'o', 'o', 'o'}
	user := &User{
		Id:   1,
		Name: "Bob",
		Uuid: testfixtures.CustomType{'1', '1', '2', '3', '8', '4', '9', '1', '2', '9', '3'},
		Mood: &mood,
	}

	if _, err := db.InsertRow(ctx, user); err != nil {
		b.Fatal(err)
	}

	benchmarks := []struct {
		name string
		fn   func() error
	}{
		{"Read", func() error {
			user := &User{}
			return db.QueryRow(ctx, &user, nil, nil)
		}},
		{"Read_Where", func() error {
			user := &User{}
			return db.QueryRow(ctx, &user, Filter{"name": "Bob"}, nil)
		}},
		{"Create", func() error {
			_, err := db.InsertRow(ctx, user)
			return err
		}},
		{"Update", func() error { return db.UpdateRow(ctx, user) }},
		{"Delete", func() error { return db.DeleteRow(ctx, user) }},
	}

	for _, bm := range benchmarks {
		b.Run(bm.name, func(b *testing.B) {
			b.ReportAllocs()
			for i := 0; i < b.N; i++ {
				if err := bm.fn(); err != nil {
					b.Error(err)
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"time"
)

const (
	// default docker registry
	defaultRegistry = "https://index.docker.io/v1/"
)

type (
	// Daemon defines Docker daemon parameters.
	Daemon struct {
		Registry      string   // Docker registry
		Mirror        string   // Docker registry mirror
		Insecure      bool     // Docker daemon enable insecure registries
		StorageDriver string   // Docker daemon storage driver
		StoragePath   string   // Docker daemon storage path
		Disabled      bool     // DOcker daemon is disabled (already running)
		Debug         bool     // Docker daemon started in debug mode
		Bip           string   // Docker daemon network bridge IP address
		DNS           []string // Docker daemon dns server
		MTU           string   // Docker daemon mtu setting
	}

	// Login defines Docker login parameters.
	Login struct {
		Registry string // Docker registry address
		Username string // Docker registry username
		Password string // Docker registry password
		Email    string // Docker registry email
	}

	// Build defines Docker build parameters.
	Build struct {
		Name       string   // Docker build using default named tag
		Dockerfile string   // Docker build Dockerfile
		Context    string   // Docker build context
		Tags       []string // Docker build tags
		Args       []string // Docker build args
		Repo       string   // Docker build repository
	}

	// Plugin defines the Docker plugin parameters.
	Plugin struct {
		Login  Login  // Docker login configuration
		Build  Build  // Docker build configuration
		Daemon Daemon // Docker daemon configuration
		Dryrun bool   // Docker push is skipped
	}
)

// Exec executes the plugin step
func (p Plugin) Exec() error {
	// this code attempts to normalize the repository name by appending the fully
	// qualified registry name if otherwise omitted.
	if p.Login.Registry != defaultRegistry &&
		strings.HasPrefix(p.Build.Repo, defaultRegistry) {
		p.Build.Repo = p.Login.Registry + "/" + p.Build.Repo
	}

	// TODO execute code remove dangling images
	// this is problematic because we are running docker in scratch which does
	// not have bash, so we need to hack something together
	// docker images --quiet --filter=dangling=true | xargs --no-run-if-empty docker rmi

	/*
		cmd = exec.Command("docker", "images", "-q", "-f", "dangling=true")
		cmd = exec.Command("docker", append([]string{"rmi"}, images...)...)
	*/

	// start the Docker daemon server
	if !p.Daemon.Disabled {
		cmd := commandDaemon(p.Daemon)
		if p.Daemon.Debug {
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
		} else {
			cmd.Stdout = ioutil.Discard
			cmd.Stderr = ioutil.Discard
		}
		go func() {
			trace(cmd)
			cmd.Run()
		}()
	}

	// poll the docker daemon until it is started. This ensures the daemon is
	// ready to accept connections before we proceed.
	for i := 0; i < 15; i++ {
		cmd := commandInfo()
		err := cmd.Run()
		if err == nil {
			break
		}
		time.Sleep(time.Second * 1)
	}

	// login to the Docker registry
	if p.Login.Password != "" {
		cmd := commandLogin(p.Login)
		err := cmd.Run()
		if err != nil {
			return fmt.Errorf("Error authenticating: %s", err)
		}
	} else {
		fmt.Println("Registry credentials not provided. Guest mode enabled.")
	}

	var cmds []*exec.Cmd
	cmds = append(cmds, commandVersion())      // docker version
	cmds = append(cmds, commandInfo())         // docker info
	cmds = append(cmds, commandBuild(p.Build)) // docker build

	for _, tag := range p.Build.Tags {
		cmds = append(cmds, commandTag(p.Build, tag)) // docker tag

		if p.Dryrun == false {
			cmds = append(cmds, commandPush(p.Build, tag)) // docker push
		}
	}

	// execute all commands in batch mode.
	for _, cmd := range cmds {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		trace(cmd)

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

	return nil
}

const dockerExe = "/usr/local/bin/docker"

// helper function to create the docker login command.
func commandLogin(login Login) *exec.Cmd {
	if login.Email != "" {
		return commandLoginEmail(login)
	}
	return exec.Command(
		dockerExe, "login",
		"-u", login.Username,
		"-p", login.Password,
		login.Registry,
	)
}

func commandLoginEmail(login Login) *exec.Cmd {
	return exec.Command(
		dockerExe, "login",
		"-u", login.Username,
		"-p", login.Password,
		"-e", login.Email,
		login.Registry,
	)
}

// helper function to create the docker info command.
func commandVersion() *exec.Cmd {
	return exec.Command(dockerExe, "version")
}

// helper function to create the docker info command.
func commandInfo() *exec.Cmd {
	return exec.Command(dockerExe, "info")
}

// helper function to create the docker build command.
func commandBuild(build Build) *exec.Cmd {
	cmd := exec.Command(
		dockerExe, "build",
		"--pull=true",
		"--rm=true",
		"-f", build.Dockerfile,
		"-t", build.Name,
	)
	for _, arg := range build.Args {
		cmd.Args = append(cmd.Args, "--build-arg", arg)
	}
	cmd.Args = append(cmd.Args, build.Context)
	return cmd
}

// helper function to create the docker tag command.
func commandTag(build Build, tag string) *exec.Cmd {
	var (
		source = build.Name
		target = fmt.Sprintf("%s:%s", build.Repo, tag)
	)
	return exec.Command(
		dockerExe, "tag", source, target,
	)
}

// helper function to create the docker push command.
func commandPush(build Build, tag string) *exec.Cmd {
	target := fmt.Sprintf("%s:%s", build.Repo, tag)
	return exec.Command(dockerExe, "push", target)
}

// helper function to create the docker daemon command.
func commandDaemon(daemon Daemon) *exec.Cmd {
	args := []string{"daemon", "-g", daemon.StoragePath}

	if daemon.StorageDriver != "" {
		args = append(args, "-s", daemon.StorageDriver)
	}
	if daemon.Insecure && daemon.Registry != "" {
		args = append(args, "--insecure-registry", daemon.Registry)
	}
	if len(daemon.Mirror) != 0 {
		args = append(args, "--registry-mirror", daemon.Mirror)
	}
	if len(daemon.Bip) != 0 {
		args = append(args, "--bip", daemon.Bip)
	}
	for _, dns := range daemon.DNS {
		args = append(args, "--dns", dns)
	}
	if len(daemon.MTU) != 0 {
		args = append(args, "--mtu", daemon.MTU)
	}
	return exec.Command(dockerExe, args...)
}

// trace writes each command to stdout with the command wrapped in an xml
// tag so that it can be extracted and displayed in the logs.
func trace(cmd *exec.Cmd) {
	fmt.Fprintf(os.Stdout, "+ %s\n", strings.Join(cmd.Args, " "))
}
<commit_msg>fixed appending registry before the repo<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"time"
)

const (
	// default docker registry
	defaultRegistry = "https://index.docker.io/v1/"
)

type (
	// Daemon defines Docker daemon parameters.
	Daemon struct {
		Registry      string   // Docker registry
		Mirror        string   // Docker registry mirror
		Insecure      bool     // Docker daemon enable insecure registries
		StorageDriver string   // Docker daemon storage driver
		StoragePath   string   // Docker daemon storage path
		Disabled      bool     // DOcker daemon is disabled (already running)
		Debug         bool     // Docker daemon started in debug mode
		Bip           string   // Docker daemon network bridge IP address
		DNS           []string // Docker daemon dns server
		MTU           string   // Docker daemon mtu setting
	}

	// Login defines Docker login parameters.
	Login struct {
		Registry string // Docker registry address
		Username string // Docker registry username
		Password string // Docker registry password
		Email    string // Docker registry email
	}

	// Build defines Docker build parameters.
	Build struct {
		Name       string   // Docker build using default named tag
		Dockerfile string   // Docker build Dockerfile
		Context    string   // Docker build context
		Tags       []string // Docker build tags
		Args       []string // Docker build args
		Repo       string   // Docker build repository
	}

	// Plugin defines the Docker plugin parameters.
	Plugin struct {
		Login  Login  // Docker login configuration
		Build  Build  // Docker build configuration
		Daemon Daemon // Docker daemon configuration
		Dryrun bool   // Docker push is skipped
	}
)

// Exec executes the plugin step
func (p Plugin) Exec() error {
	// this code attempts to normalize the repository name by appending the fully
	// qualified registry name if otherwise omitted.
	if p.Login.Registry != defaultRegistry &&
		!strings.HasPrefix(p.Build.Repo, p.Login.Registry) {
		p.Build.Repo = p.Login.Registry + "/" + p.Build.Repo
	}

	// TODO execute code remove dangling images
	// this is problematic because we are running docker in scratch which does
	// not have bash, so we need to hack something together
	// docker images --quiet --filter=dangling=true | xargs --no-run-if-empty docker rmi

	/*
		cmd = exec.Command("docker", "images", "-q", "-f", "dangling=true")
		cmd = exec.Command("docker", append([]string{"rmi"}, images...)...)
	*/

	// start the Docker daemon server
	if !p.Daemon.Disabled {
		cmd := commandDaemon(p.Daemon)
		if p.Daemon.Debug {
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
		} else {
			cmd.Stdout = ioutil.Discard
			cmd.Stderr = ioutil.Discard
		}
		go func() {
			trace(cmd)
			cmd.Run()
		}()
	}

	// poll the docker daemon until it is started. This ensures the daemon is
	// ready to accept connections before we proceed.
	for i := 0; i < 15; i++ {
		cmd := commandInfo()
		err := cmd.Run()
		if err == nil {
			break
		}
		time.Sleep(time.Second * 1)
	}

	// login to the Docker registry
	if p.Login.Password != "" {
		cmd := commandLogin(p.Login)
		err := cmd.Run()
		if err != nil {
			return fmt.Errorf("Error authenticating: %s", err)
		}
	} else {
		fmt.Println("Registry credentials not provided. Guest mode enabled.")
	}

	var cmds []*exec.Cmd
	cmds = append(cmds, commandVersion())      // docker version
	cmds = append(cmds, commandInfo())         // docker info
	cmds = append(cmds, commandBuild(p.Build)) // docker build

	for _, tag := range p.Build.Tags {
		cmds = append(cmds, commandTag(p.Build, tag)) // docker tag

		if p.Dryrun == false {
			cmds = append(cmds, commandPush(p.Build, tag)) // docker push
		}
	}

	// execute all commands in batch mode.
	for _, cmd := range cmds {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		trace(cmd)

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

	return nil
}

const dockerExe = "/usr/local/bin/docker"

// helper function to create the docker login command.
func commandLogin(login Login) *exec.Cmd {
	if login.Email != "" {
		return commandLoginEmail(login)
	}
	return exec.Command(
		dockerExe, "login",
		"-u", login.Username,
		"-p", login.Password,
		login.Registry,
	)
}

func commandLoginEmail(login Login) *exec.Cmd {
	return exec.Command(
		dockerExe, "login",
		"-u", login.Username,
		"-p", login.Password,
		"-e", login.Email,
		login.Registry,
	)
}

// helper function to create the docker info command.
func commandVersion() *exec.Cmd {
	return exec.Command(dockerExe, "version")
}

// helper function to create the docker info command.
func commandInfo() *exec.Cmd {
	return exec.Command(dockerExe, "info")
}

// helper function to create the docker build command.
func commandBuild(build Build) *exec.Cmd {
	cmd := exec.Command(
		dockerExe, "build",
		"--pull=true",
		"--rm=true",
		"-f", build.Dockerfile,
		"-t", build.Name,
	)
	for _, arg := range build.Args {
		cmd.Args = append(cmd.Args, "--build-arg", arg)
	}
	cmd.Args = append(cmd.Args, build.Context)
	return cmd
}

// helper function to create the docker tag command.
func commandTag(build Build, tag string) *exec.Cmd {
	var (
		source = build.Name
		target = fmt.Sprintf("%s:%s", build.Repo, tag)
	)
	return exec.Command(
		dockerExe, "tag", source, target,
	)
}

// helper function to create the docker push command.
func commandPush(build Build, tag string) *exec.Cmd {
	target := fmt.Sprintf("%s:%s", build.Repo, tag)
	return exec.Command(dockerExe, "push", target)
}

// helper function to create the docker daemon command.
func commandDaemon(daemon Daemon) *exec.Cmd {
	args := []string{"daemon", "-g", daemon.StoragePath}

	if daemon.StorageDriver != "" {
		args = append(args, "-s", daemon.StorageDriver)
	}
	if daemon.Insecure && daemon.Registry != "" {
		args = append(args, "--insecure-registry", daemon.Registry)
	}
	if len(daemon.Mirror) != 0 {
		args = append(args, "--registry-mirror", daemon.Mirror)
	}
	if len(daemon.Bip) != 0 {
		args = append(args, "--bip", daemon.Bip)
	}
	for _, dns := range daemon.DNS {
		args = append(args, "--dns", dns)
	}
	if len(daemon.MTU) != 0 {
		args = append(args, "--mtu", daemon.MTU)
	}
	return exec.Command(dockerExe, args...)
}

// trace writes each command to stdout with the command wrapped in an xml
// tag so that it can be extracted and displayed in the logs.
func trace(cmd *exec.Cmd) {
	fmt.Fprintf(os.Stdout, "+ %s\n", strings.Join(cmd.Args, " "))
}
<|endoftext|>


<commit_before>package redis

// MultiCall holds data for multiple command calls.
type MultiCall struct {
	transaction bool
	c           *connection
	calls       []call
}

func newMultiCall(transaction bool, c *connection) *MultiCall {
	return &MultiCall{
		transaction: transaction,
		c:           c,
	}
}

// process calls the given multicall function, flushes the
// calls, and returns the returned Reply.
func (mc *MultiCall) process(userCalls func(*MultiCall)) *Reply {
	if mc.transaction {
		mc.Multi()
	}
	userCalls(mc)
	var r *Reply
	if !mc.transaction {
		r = mc.c.multiCall(mc.calls)
	} else {
		mc.Exec()
		r = mc.c.multiCall(mc.calls)

		execReply, err := r.At(len(r.elems) - 1)
		if err == nil {
			r.elems = execReply.elems
		} else {
			if execReply.Error != nil {
				r.Error = execReply.Error
			} else {
				r.Error = newError("unknown transaction error")
			}
		}
	}

	return r
}

func (mc *MultiCall) call(cmd Cmd, args ...interface{}) {
	mc.calls = append(mc.calls, call{cmd, args})
}

// Call queues a Redis command call for later execution.
func (mc *MultiCall) Call(cmd string, args ...interface{}) {
	mc.call(Cmd(cmd), args...)
}

// Flush sends queued command calls to the Redis server for execution and
// returns the returned Reply.
func (mc *MultiCall) Flush() (r *Reply) {
	r = mc.c.multiCall(mc.calls)
	mc.calls = nil
	return
}
<commit_msg>multicall: should be safe to ignore errors here<commit_after>package redis

// MultiCall holds data for multiple command calls.
type MultiCall struct {
	transaction bool
	c           *connection
	calls       []call
}

func newMultiCall(transaction bool, c *connection) *MultiCall {
	return &MultiCall{
		transaction: transaction,
		c:           c,
	}
}

// process calls the given multicall function, flushes the
// calls, and returns the returned Reply.
func (mc *MultiCall) process(userCalls func(*MultiCall)) *Reply {
	if mc.transaction {
		mc.Multi()
	}
	userCalls(mc)
	var r *Reply
	if !mc.transaction {
		r = mc.c.multiCall(mc.calls)
	} else {
		mc.Exec()
		r = mc.c.multiCall(mc.calls)

		execReply, _ := r.At(len(r.elems) - 1)
		if execReply.Error == nil {
			r.elems = execReply.elems
		} else {
			if execReply.Error != nil {
				r.Error = execReply.Error
			} else {
				r.Error = newError("unknown transaction error")
			}
		}
	}

	return r
}

func (mc *MultiCall) call(cmd Cmd, args ...interface{}) {
	mc.calls = append(mc.calls, call{cmd, args})
}

// Call queues a Redis command call for later execution.
func (mc *MultiCall) Call(cmd string, args ...interface{}) {
	mc.call(Cmd(cmd), args...)
}

// Flush sends queued command calls to the Redis server for execution and
// returns the returned Reply.
func (mc *MultiCall) Flush() (r *Reply) {
	r = mc.c.multiCall(mc.calls)
	mc.calls = nil
	return
}
<|endoftext|>


<commit_before>package dncomp

import "testing"

func equal(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func checkResult(t *testing.T, input []byte, err error, res, expect []string) {
	if err != nil {
		t.Error("error:", err.Error())
	} else if !equal(res, expect) {
		t.Error("decoding error for input", input,
			", expect", expect, ", got", res)
	}
}

func checkError(t *testing.T, input []byte, err error, res []string) {
	if err == nil {
		t.Error("error return failed for input", input, ", got", res)
	}
}

func TestRFC(t *testing.T) {
	input := []byte{1, 'F', 3, 'I', 'S', 'I', 4, 'A', 'R', 'P', 'A', 0,
		3, 'F', 'O', 'O', 0xc0, 0, 0xc0, 6, 0, 0}

	expect := []string{"F.ISI.ARPA", "FOO.F.ISI.ARPA", "ARPA", ""}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestEmpty(t *testing.T) {
	input := []byte{}

	expect := []string{}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestRoot(t *testing.T) {
	input := []byte{0}

	expect := []string{""}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestSimple(t *testing.T) {
	input := []byte{1, 'A'}

	expect := []string{"A"}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestInvalid1(t *testing.T) {
	input := []byte{ 1 }
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalid2(t *testing.T) {
	input := []byte{ 'A' }
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalid3(t *testing.T) {
	input := []byte{ 5, 'A', 'B' }
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalid4(t *testing.T) {
	input := []byte{ 2, 'A', 'B', 0, 1 }
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalidPointer1(t *testing.T) {
	input := []byte{ 2, 'A', 'B', 0, 0xc0, 6 }
	res, err := Decode(input)
	checkError(t, input, err, res)
}

<commit_msg>Add extra tests and fix typo in reference data<commit_after>package dncomp

import "testing"

func equal(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func checkResult(t *testing.T, input []byte, err error, res, expect []string) {
	if err != nil {
		t.Error("error:", err.Error())
	} else if !equal(res, expect) {
		t.Error("decoding error for input", input,
			", expect", expect, ", got", res)
	}
}

func checkError(t *testing.T, input []byte, err error, res []string) {
	if err == nil {
		t.Error("error return failed for input", input, ", got", res)
	}
}

func TestRFC(t *testing.T) {
	input := []byte{1, 'F', 3, 'I', 'S', 'I', 4, 'A', 'R', 'P', 'A', 0,
		3, 'F', 'O', 'O', 0xc0, 0, 0xc0, 6, 0}

	expect := []string{"F.ISI.ARPA", "FOO.F.ISI.ARPA", "ARPA", ""}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestEmpty(t *testing.T) {
	input := []byte{}

	expect := []string{}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestRoot(t *testing.T) {
	input := []byte{0}

	expect := []string{""}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestSimple(t *testing.T) {
	input := []byte{1, 'A'}

	expect := []string{"A"}

	res, err := Decode(input)
	checkResult(t, input, err, res, expect)
}

func TestInvalid1(t *testing.T) {
	input := []byte{1}
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalid2(t *testing.T) {
	input := []byte{'A'}
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalid3(t *testing.T) {
	input := []byte{5, 'A', 'B'}
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalid4(t *testing.T) {
	input := []byte{2, 'A', 'B', 0, 1}
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalidPointer1(t *testing.T) {
	input := []byte{2, 'A', 'B', 0, 0xc0, 6}
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalidPointer2(t *testing.T) {
	input := []byte{2, 'A', 'B', 0, 0xc0, 5}
	res, err := Decode(input)
	checkError(t, input, err, res)
}

func TestInvalidPointer3(t *testing.T) {
	input := []byte{2, 'A', 'B', 0, 0xc0, 4}
	res, err := Decode(input)
	checkError(t, input, err, res)
}
<|endoftext|>


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

package testing

import (
	"gopkg.in/check.v1"
)

// Contains is a checker for gocheck that asserts if an item is in a slice.
var Contains = &containsChecker{
	&check.CheckerInfo{
		Name: "Contains",
		Params: []string{
			"Container",
			"Expected to contain",
		},
	},
}

type containsChecker struct {
	*check.CheckerInfo
}

func (c *containsChecker) Check(params []interface{}, names []string) (result bool, error string) {
	slice, ok := params[0].([]interface{})
	if !ok {
		return false, "First parameter is not a []interface{}"
	}
	value, ok := params[1].(interface{})
	if !ok {
		return false, "Second parameter is not an int"
	}
	for _, v := range slice {
		if v == value {
			return true, ""
		}
	}
	return false, ""
}
<commit_msg>Add HasKey checker for maps and cleanup Contains checker<commit_after>// Copyright (c) 2014 - Max Persson <max@looplab.se>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package testing

import (
	"reflect"

	"gopkg.in/check.v1"
)

// Contains is a checker for gocheck that asserts if an item is in a slice.
var Contains = &contains{}

// HasKey is a checker for gocheck that asserts if a map has a key.
var HasKey = &hasKey{}

type contains struct{}

func (c *contains) Check(params []interface{}, names []string) (bool, string) {
	if len(params) != 2 {
		return false, "Contains takes 2 arguments: a slice and an item"
	}

	slice, ok := params[0].([]interface{})
	if !ok {
		return false, "first parameter is not a []interface{}"
	}
	value, ok := params[1].(interface{})
	if !ok {
		return false, "second parameter is not an interface"
	}

	for _, v := range slice {
		if v == value {
			return true, ""
		}
	}

	return false, ""
}

func (c *contains) Info() *check.CheckerInfo {
	return &check.CheckerInfo{
		Name:   "Contains",
		Params: []string{"slice", "item"},
	}
}

type hasKey struct{}

func (h *hasKey) Check(params []interface{}, names []string) (bool, string) {
	if len(params) != 2 {
		return false, "HasKey takes 2 arguments: a map and a key"
	}

	mapValue := reflect.ValueOf(params[0])
	if mapValue.Kind() != reflect.Map {
		return false, "first argument to HasKey must be a map"
	}

	keyValue := reflect.ValueOf(params[1])
	if !keyValue.Type().AssignableTo(mapValue.Type().Key()) {
		return false, "second argument must be assignable to the map key type"
	}

	return mapValue.MapIndex(keyValue).IsValid(), ""
}

func (h *hasKey) Info() *check.CheckerInfo {
	return &check.CheckerInfo{
		Name:   "HasKey",
		Params: []string{"map", "key"},
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package api

import (
	"errors"
	"net/http"
	"sync"
)

// RequestContextMapper keeps track of the context associated with a particular request
type RequestContextMapper interface {
	// Get returns the context associated with the given request (if any), and true if the request has an associated context, and false if it does not.
	Get(req *http.Request) (Context, bool)
	// Update maps the request to the given context. If no context was previously associated with the request, an error is returned.
	// Update should only be called with a descendant context of the previously associated context.
	// Updating to an unrelated context may return an error in the future.
	// The context associated with a request should only be updated by a limited set of callers.
	// Valid examples include the authentication layer, or an audit/tracing layer.
	Update(req *http.Request, context Context) error
}

type requestContextMap struct {
	contexts map[*http.Request]Context
	lock     sync.Mutex
}

// NewRequestContextMapper returns a new RequestContextMapper.
// The returned mapper must be added as a request filter using NewRequestContextFilter.
func NewRequestContextMapper() RequestContextMapper {
	return &requestContextMap{
		contexts: make(map[*http.Request]Context),
	}
}

// Get returns the context associated with the given request (if any), and true if the request has an associated context, and false if it does not.
// Get will only return a valid context when called from inside the filter chain set up by NewRequestContextFilter()
func (c *requestContextMap) Get(req *http.Request) (Context, bool) {
	c.lock.Lock()
	defer c.lock.Unlock()
	context, ok := c.contexts[req]
	return context, ok
}

// Update maps the request to the given context.
// If no context was previously associated with the request, an error is returned and the context is ignored.
func (c *requestContextMap) Update(req *http.Request, context Context) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	if _, ok := c.contexts[req]; !ok {
		return errors.New("No context associated")
	}
	// TODO: ensure the new context is a descendant of the existing one
	c.contexts[req] = context
	return nil
}

// init maps the request to the given context and returns true if there was no context associated with the request already.
// if a context was already associated with the request, it ignores the given context and returns false.
// init is intentionally unexported to ensure that all init calls are paired with a remove after a request is handled
func (c *requestContextMap) init(req *http.Request, context Context) bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	if _, exists := c.contexts[req]; exists {
		return false
	}
	c.contexts[req] = context
	return true
}

// remove is intentionally unexported to ensure that the context is not removed until a request is handled
func (c *requestContextMap) remove(req *http.Request) {
	c.lock.Lock()
	defer c.lock.Unlock()
	delete(c.contexts, req)
}

// NewRequestContextFilter ensures there is a Context object associated with the request before calling the passed handler.
// After the passed handler runs, the context is cleaned up.
func NewRequestContextFilter(mapper RequestContextMapper, handler http.Handler) (http.Handler, error) {
	if mapper, ok := mapper.(*requestContextMap); ok {
		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			if mapper.init(req, NewContext()) {
				// If we were the ones to successfully initialize, pair with a remove
				defer mapper.remove(req)
			}
			handler.ServeHTTP(w, req)
		}), nil
	} else {
		return handler, errors.New("Unknown RequestContextMapper implementation.")
	}

}

// IsEmpty returns true if there are no contexts registered, or an error if it could not be determined. Intended for use by tests.
func IsEmpty(requestsToContexts RequestContextMapper) (bool, error) {
	if requestsToContexts, ok := requestsToContexts.(*requestContextMap); ok {
		return len(requestsToContexts.contexts) == 0, nil
	}
	return true, errors.New("Unknown RequestContextMapper implementation")
}
<commit_msg>Cleanup handler chain<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package api

import (
	"errors"
	"net/http"
	"sync"

	"github.com/golang/glog"
)

// RequestContextMapper keeps track of the context associated with a particular request
type RequestContextMapper interface {
	// Get returns the context associated with the given request (if any), and true if the request has an associated context, and false if it does not.
	Get(req *http.Request) (Context, bool)
	// Update maps the request to the given context. If no context was previously associated with the request, an error is returned.
	// Update should only be called with a descendant context of the previously associated context.
	// Updating to an unrelated context may return an error in the future.
	// The context associated with a request should only be updated by a limited set of callers.
	// Valid examples include the authentication layer, or an audit/tracing layer.
	Update(req *http.Request, context Context) error
}

type requestContextMap struct {
	contexts map[*http.Request]Context
	lock     sync.Mutex
}

// NewRequestContextMapper returns a new RequestContextMapper.
// The returned mapper must be added as a request filter using NewRequestContextFilter.
func NewRequestContextMapper() RequestContextMapper {
	return &requestContextMap{
		contexts: make(map[*http.Request]Context),
	}
}

// Get returns the context associated with the given request (if any), and true if the request has an associated context, and false if it does not.
// Get will only return a valid context when called from inside the filter chain set up by NewRequestContextFilter()
func (c *requestContextMap) Get(req *http.Request) (Context, bool) {
	c.lock.Lock()
	defer c.lock.Unlock()
	context, ok := c.contexts[req]
	return context, ok
}

// Update maps the request to the given context.
// If no context was previously associated with the request, an error is returned and the context is ignored.
func (c *requestContextMap) Update(req *http.Request, context Context) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	if _, ok := c.contexts[req]; !ok {
		return errors.New("No context associated")
	}
	// TODO: ensure the new context is a descendant of the existing one
	c.contexts[req] = context
	return nil
}

// init maps the request to the given context and returns true if there was no context associated with the request already.
// if a context was already associated with the request, it ignores the given context and returns false.
// init is intentionally unexported to ensure that all init calls are paired with a remove after a request is handled
func (c *requestContextMap) init(req *http.Request, context Context) bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	if _, exists := c.contexts[req]; exists {
		return false
	}
	c.contexts[req] = context
	return true
}

// remove is intentionally unexported to ensure that the context is not removed until a request is handled
func (c *requestContextMap) remove(req *http.Request) {
	c.lock.Lock()
	defer c.lock.Unlock()
	delete(c.contexts, req)
}

// WithRequestContext ensures there is a Context object associated with the request before calling the passed handler.
// After the passed handler runs, the context is cleaned up.
func WithRequestContext(handler http.Handler, mapper RequestContextMapper) http.Handler {
	rcMap, ok := mapper.(*requestContextMap)
	if !ok {
		glog.Fatal("Unknown RequestContextMapper implementation.")
	}

	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if rcMap.init(req, NewContext()) {
			// If we were the ones to successfully initialize, pair with a remove
			defer rcMap.remove(req)
		}
		handler.ServeHTTP(w, req)
	})
}

// IsEmpty returns true if there are no contexts registered, or an error if it could not be determined. Intended for use by tests.
func IsEmpty(requestsToContexts RequestContextMapper) (bool, error) {
	if requestsToContexts, ok := requestsToContexts.(*requestContextMap); ok {
		return len(requestsToContexts.contexts) == 0, nil
	}
	return true, errors.New("Unknown RequestContextMapper implementation")
}
<|endoftext|>


<commit_before>package sensu

import (
	"bytes"
	"encoding/json"
	"github.com/upfluence/sensu-client-go/sensu/check"
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	MAX_FAILS = 100
	MAX_TIME  = 60 * time.Second
)

type Subscriber struct {
	Subscription string
	Client       *Client
}

func (s *Subscriber) SetClient(c *Client) error {
	s.Client = c

	return nil
}

func (s *Subscriber) Start() error {
	var output check.CheckOutput
	var b []byte

	funnel := strings.Join(
		[]string{
			s.Client.Config.Name(),
			CurrentVersion,
			strconv.Itoa(int(time.Now().Unix())),
		},
		"-",
	)

	msgChan := make(chan []byte)
	stopChan := make(chan bool)

	go s.Client.Transport.Subscribe("#", s.Subscription, funnel, msgChan, stopChan)

	log.Printf("Subscribed to %s", s.Subscription)

	failures := 0

	for {
		if failures >= MAX_FAILS {
			stopChan <- true
			msgChan = make(chan []byte)
			failures = 0
			s.Client.Transport.Close()
			time.Sleep(MAX_TIME)
			s.Client.Transport.Connect()
			go s.Client.Transport.Subscribe(
				"#",
				s.Subscription,
				funnel,
				msgChan,
				stopChan,
			)

		}

		b = <-msgChan

		payload := make(map[string]interface{})

		log.Printf("Check received : %s", bytes.NewBuffer(b).String())
		json.Unmarshal(b, &payload)

		if _, ok := payload["name"]; !ok {
			log.Printf("The name field is not filled")
			failures++
			continue
		}

		if ch, ok := check.Store[payload["name"].(string)]; ok {
			output = ch.Execute()
		} else if _, ok := payload["command"]; !ok {
			log.Printf("The command field is not filled")
			continue
		} else {
			output = (&check.ExternalCheck{payload["command"].(string)}).Execute()
		}

		p, err := json.Marshal(s.forgeCheckResponse(payload, &output))

		if err != nil {
			log.Printf("something goes wrong : %s", err.Error())
		} else {
			log.Printf("Payload sent: %s", bytes.NewBuffer(p).String())

			err = s.Client.Transport.Publish("direct", "results", "", p)
		}
	}

	return nil
}

func (s *Subscriber) forgeCheckResponse(payload map[string]interface{}, output *check.CheckOutput) map[string]interface{} {
	result := make(map[string]interface{})

	result["client"] = s.Client.Config.Name()

	formattedOuput := make(map[string]interface{})

	formattedOuput["name"] = payload["name"]
	formattedOuput["issued"] = int(payload["issued"].(float64))
	formattedOuput["output"] = output.Output
	formattedOuput["duration"] = output.Duration
	formattedOuput["status"] = output.Status
	formattedOuput["executed"] = output.Executed

	result["check"] = formattedOuput

	return result
}
<commit_msg>little bit of refactoring<commit_after>package sensu

import (
	"bytes"
	"encoding/json"
	"github.com/upfluence/sensu-client-go/sensu/check"
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	MAX_FAILS = 100
	MAX_TIME  = 60 * time.Second
)

type Subscriber struct {
	Subscription string
	Client       *Client
}

func (s *Subscriber) SetClient(c *Client) error {
	s.Client = c

	return nil
}

func (s *Subscriber) Start() error {
	var output check.CheckOutput
	var b []byte

	funnel := strings.Join(
		[]string{
			s.Client.Config.Name(),
			CurrentVersion,
			strconv.Itoa(int(time.Now().Unix())),
		},
		"-",
	)

	msgChan, stopChan := s.subscribe(funnel)
	log.Printf("Subscribed to %s", s.Subscription)

	failures := 0
	for {
		if failures >= MAX_FAILS {
			failures = 0
			stopChan <- true
			time.Sleep(MAX_TIME)
			msgChan, stopChan = s.subscribe(funnel)
		}

		b = <-msgChan

		payload := make(map[string]interface{})

		log.Printf("Check received : %s", bytes.NewBuffer(b).String())
		json.Unmarshal(b, &payload)

		if _, ok := payload["name"]; !ok {
			log.Printf("The name field is not filled")
			failures++
			continue
		}

		if ch, ok := check.Store[payload["name"].(string)]; ok {
			output = ch.Execute()
		} else if _, ok := payload["command"]; !ok {
			log.Printf("The command field is not filled")
			continue
		} else {
			output = (&check.ExternalCheck{payload["command"].(string)}).Execute()
		}

		p, err := json.Marshal(s.forgeCheckResponse(payload, &output))

		if err != nil {
			log.Printf("something goes wrong : %s", err.Error())
		} else {
			log.Printf("Payload sent: %s", bytes.NewBuffer(p).String())

			err = s.Client.Transport.Publish("direct", "results", "", p)
		}
	}

	return nil
}

func (s *Subscriber) subscribe(funnel string) (chan []byte, chan bool) {
	msgChan := make(chan []byte)
	stopChan := make(chan bool)
	go s.Client.Transport.Subscribe(
		"#",
		s.Subscription,
		funnel,
		msgChan,
		stopChan,
	)
	return msgChan, stopChan
}

func (s *Subscriber) forgeCheckResponse(payload map[string]interface{}, output *check.CheckOutput) map[string]interface{} {
	result := make(map[string]interface{})

	result["client"] = s.Client.Config.Name()

	formattedOuput := make(map[string]interface{})

	formattedOuput["name"] = payload["name"]
	formattedOuput["issued"] = int(payload["issued"].(float64))
	formattedOuput["output"] = output.Output
	formattedOuput["duration"] = output.Duration
	formattedOuput["status"] = output.Status
	formattedOuput["executed"] = output.Executed

	result["check"] = formattedOuput

	return result
}
<|endoftext|>