<commit_before>// Package dnsprobe implements a DNS probe.
package dnsprobe // import "hkjn.me/probes/dnsprobe"

import (
	"fmt"
	"net"
	"sort"
	"strings"
	"time"

	"github.com/golang/glog"

	"hkjn.me/prober"
	"hkjn.me/probes"
)

// DnsProber probes a target host's DNS records.
type DnsProber struct {
	Target    string // host to probe
	wantMX    mxRecords
	wantA     []string
	wantNS    nsRecords
	wantCNAME string
	wantTXT   []string
}

// New returns a new instance of the DNS probe with specified options.
func New(target string, options ...func(*DnsProber)) *prober.Probe {
	return NewWithGeneric(target, []prober.Option{}, options...)
}

// NewWithGeneric returns a new instance of the DNS probe with specified options.
//
// NewWithGeneric passes through specified prober.Options, after
// applying the dnsprobe-specific options.
func NewWithGeneric(target string, genericOpts []prober.Option, options ...func(*DnsProber)) *prober.Probe {
	p := &DnsProber{Target: target}
	for _, opt := range options {
		opt(p)
	}
	return prober.NewProbe(p, fmt.Sprintf("DnsProber_%s", target),
		fmt.Sprintf("Probes DNS records of %s", target),
		append(genericOpts, prober.Interval(time.Minute*5), prober.FailurePenalty(5))...)
}

// MX applies the option that the prober wants specific MX records.
func MX(mx []*net.MX) func(*DnsProber) {
	return func(p *DnsProber) {
		wantMX := mxRecords(mx)
		sort.Sort(wantMX)
		p.wantMX = wantMX
	}
}

// A applies the option that the prober wants specific A records.
func A(a []string) func(*DnsProber) {
	return func(p *DnsProber) {
		sort.Strings(a)
		p.wantA = a
	}
}

// NS applies the option that the prober wants specific NS records.
func NS(ns []*net.NS) func(*DnsProber) {
	return func(p *DnsProber) {
		nsRec := nsRecords(ns)
		sort.Sort(nsRec)
		p.wantNS = nsRec
	}
}

// CNAME applies the option that the prober wants specific CNAME record.
func CNAME(cname string) func(*DnsProber) {
	return func(p *DnsProber) {
		p.wantCNAME = cname
	}
}

// TXT applies the option that the prober wants specific TXT records.
func TXT(txt []string) func(*DnsProber) {
	return func(p *DnsProber) {
		sort.Strings(txt)
		p.wantTXT = txt
	}
}

// Probe verifies that the target's DNS records are as expected.
func (p *DnsProber) Probe() prober.Result {
	if len(p.wantMX) > 0 {
		glog.V(1).Infof("Checking %d MX records..\n", len(p.wantMX))
		err := p.checkMX()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if len(p.wantA) > 0 {
		glog.V(1).Infof("Checking %d A records..\n", len(p.wantA))
		err := p.checkA()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if len(p.wantNS) > 0 {
		glog.V(1).Infof("Checking %d NS records..\n", len(p.wantNS))
		err := p.checkNS()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if p.wantCNAME != "" {
		glog.V(1).Infof("Checking CNAME record..\n")
		err := p.checkCNAME()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if len(p.wantTXT) > 0 {
		glog.V(1).Infof("Checking %d TXT records..\n", len(p.wantTXT))
		err := p.checkTXT()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	return prober.Passed()
}

// mxRecords is a collection of MX records, implementing sort.Interface.
//
// We need this custom order since the sort order in net.LookupMX
// randomizes records with the same preference value, but we'd like a stable
// ordering so we can compared expected vs actual records.
type mxRecords []*net.MX

func (r mxRecords) Len() int { return len(r) }

func (r mxRecords) Swap(i, j int) { r[i], r[j] = r[j], r[i] }

func (r mxRecords) Less(i, j int) bool {
	if r[i].Pref == r[j].Pref {
		return r[i].Host < r[j].Host
	}
	return r[i].Pref < r[j].Pref
}

// String returns a readable description of the MX records.
func (r mxRecords) String() string {
	s := ""
	for i, r := range r {
		if i > 0 {
			s += ", "
		}
		s += fmt.Sprintf("%s (%d)", r.Host, r.Pref)
	}
	return s
}

// checkMX verifies that the target has expected MX records.
func (p *DnsProber) checkMX() error {
	mx, err := net.LookupMX(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up MX records for %s: %v", p.Target, err)
	}
	mxRec := mxRecords(mx)
	if len(mxRec) != len(p.wantMX) {
		return fmt.Errorf("want %d MX records, got %d: %s", len(p.wantMX), len(mxRec), mxRec)
	}
	sort.Sort(mxRec)
	for i, r := range mxRec {
		if !strings.EqualFold(p.wantMX[i].Host, r.Host) {
			return fmt.Errorf("bad host %q for MX record #%d; want %q", r.Host, i, p.wantMX[i].Host)
		}
		if p.wantMX[i].Pref != r.Pref {
			return fmt.Errorf("bad prio %d for MX record #%d; want %d", i, r.Pref, p.wantMX[i].Pref)
		}
	}
	return nil
}

// checkA verifies that the target has expected A records.
func (p *DnsProber) checkA() error {
	addr, err := net.LookupHost(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up A records for %s: %v", p.Target, err)
	}
	if len(addr) != len(p.wantA) {
		return fmt.Errorf("got %d A records, want %d", len(addr), len(p.wantA))
	}
	sort.Strings(addr)
	for i, a := range addr {
		if p.wantA[i] != a {
			return fmt.Errorf("bad A record %q at #%d; want %q", a, i, p.wantA[i])
		}
	}
	return nil
}

// nsRecords is a collection of NS records, implementing sort.Interface.
type nsRecords []*net.NS

func (r nsRecords) Len() int { return len(r) }

func (r nsRecords) Swap(i, j int) { r[i], r[j] = r[j], r[i] }

func (r nsRecords) Less(i, j int) bool { return r[i].Host < r[j].Host }

// String returns a readable description of the NS records.
func (ns nsRecords) String() string {
	s := ""
	for i, r := range ns {
		if i > 0 {
			s += ", "
		}
		s += r.Host
	}
	return s
}

// checkNS verifies that the target has expected NS records.
func (p *DnsProber) checkNS() error {
	ns, err := net.LookupNS(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up NS records for %s: %v", p.Target, err)
	}
	nsRec := nsRecords(ns)
	if len(nsRec) != len(p.wantNS) {
		return fmt.Errorf("want %d NS records, got %d: %s", len(p.wantNS), len(nsRec), nsRec)
	}
	sort.Sort(nsRec)
	for i, n := range ns {
		if p.wantNS[i].Host != n.Host {
			return fmt.Errorf("bad NS record %q at #%d; want %q", n, i, p.wantNS[i])
		}
	}
	return nil
}

// checkCNAME verifies that the target has expected CNAME record.
func (p *DnsProber) checkCNAME() error {
	cname, err := net.LookupCNAME(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up CNAME record for %s: %v", p.Target, err)
	}
	if cname != p.wantCNAME {
		return fmt.Errorf("bad CNAME record %q; want %q", cname, p.wantCNAME)
	}
	return nil
}

// checkTXT verifies that the target has expected TXT records.
func (p *DnsProber) checkTXT() error {
	txt, err := net.LookupTXT(p.Target)
	if err != nil {
		return err
	}
	if len(txt) != len(p.wantTXT) {
		return fmt.Errorf("want %d TXT records, got %d: %s", len(p.wantTXT), len(txt), txt)
	}
	sort.Strings(txt)
	for i, t := range txt {
		if p.wantTXT[i] != t {
			return fmt.Errorf("bad TXT record %q at #%d; want %q", t, i, p.wantTXT[i])
		}
	}
	return nil
}

// Alert sends an alert notification via email.
func (p *DnsProber) Alert(name, desc string, badness int, records prober.Records) error {
	return probes.SendAlertEmail(name, desc, badness, records)
}
<commit_msg>Also show DNS A records we actually got on failure<commit_after>// Package dnsprobe implements a DNS probe.
package dnsprobe // import "hkjn.me/probes/dnsprobe"

import (
	"fmt"
	"net"
	"sort"
	"strings"
	"time"

	"github.com/golang/glog"

	"hkjn.me/prober"
	"hkjn.me/probes"
)

// DnsProber probes a target host's DNS records.
type DnsProber struct {
	Target    string // host to probe
	wantMX    mxRecords
	wantA     []string
	wantNS    nsRecords
	wantCNAME string
	wantTXT   []string
}

// New returns a new instance of the DNS probe with specified options.
func New(target string, options ...func(*DnsProber)) *prober.Probe {
	return NewWithGeneric(target, []prober.Option{}, options...)
}

// NewWithGeneric returns a new instance of the DNS probe with specified options.
//
// NewWithGeneric passes through specified prober.Options, after
// applying the dnsprobe-specific options.
func NewWithGeneric(target string, genericOpts []prober.Option, options ...func(*DnsProber)) *prober.Probe {
	p := &DnsProber{Target: target}
	for _, opt := range options {
		opt(p)
	}
	return prober.NewProbe(p, fmt.Sprintf("DnsProber_%s", target),
		fmt.Sprintf("Probes DNS records of %s", target),
		append(genericOpts, prober.Interval(time.Minute*5), prober.FailurePenalty(5))...)
}

// MX applies the option that the prober wants specific MX records.
func MX(mx []*net.MX) func(*DnsProber) {
	return func(p *DnsProber) {
		wantMX := mxRecords(mx)
		sort.Sort(wantMX)
		p.wantMX = wantMX
	}
}

// A applies the option that the prober wants specific A records.
func A(a []string) func(*DnsProber) {
	return func(p *DnsProber) {
		sort.Strings(a)
		p.wantA = a
	}
}

// NS applies the option that the prober wants specific NS records.
func NS(ns []*net.NS) func(*DnsProber) {
	return func(p *DnsProber) {
		nsRec := nsRecords(ns)
		sort.Sort(nsRec)
		p.wantNS = nsRec
	}
}

// CNAME applies the option that the prober wants specific CNAME record.
func CNAME(cname string) func(*DnsProber) {
	return func(p *DnsProber) {
		p.wantCNAME = cname
	}
}

// TXT applies the option that the prober wants specific TXT records.
func TXT(txt []string) func(*DnsProber) {
	return func(p *DnsProber) {
		sort.Strings(txt)
		p.wantTXT = txt
	}
}

// Probe verifies that the target's DNS records are as expected.
func (p *DnsProber) Probe() prober.Result {
	if len(p.wantMX) > 0 {
		glog.V(1).Infof("Checking %d MX records..\n", len(p.wantMX))
		err := p.checkMX()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if len(p.wantA) > 0 {
		glog.V(1).Infof("Checking %d A records..\n", len(p.wantA))
		err := p.checkA()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if len(p.wantNS) > 0 {
		glog.V(1).Infof("Checking %d NS records..\n", len(p.wantNS))
		err := p.checkNS()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if p.wantCNAME != "" {
		glog.V(1).Infof("Checking CNAME record..\n")
		err := p.checkCNAME()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	if len(p.wantTXT) > 0 {
		glog.V(1).Infof("Checking %d TXT records..\n", len(p.wantTXT))
		err := p.checkTXT()
		if err != nil {
			return prober.FailedWith(err)
		}
	}
	return prober.Passed()
}

// mxRecords is a collection of MX records, implementing sort.Interface.
//
// We need this custom order since the sort order in net.LookupMX
// randomizes records with the same preference value, but we'd like a stable
// ordering so we can compared expected vs actual records.
type mxRecords []*net.MX

func (r mxRecords) Len() int { return len(r) }

func (r mxRecords) Swap(i, j int) { r[i], r[j] = r[j], r[i] }

func (r mxRecords) Less(i, j int) bool {
	if r[i].Pref == r[j].Pref {
		return r[i].Host < r[j].Host
	}
	return r[i].Pref < r[j].Pref
}

// String returns a readable description of the MX records.
func (r mxRecords) String() string {
	s := ""
	for i, r := range r {
		if i > 0 {
			s += ", "
		}
		s += fmt.Sprintf("%s (%d)", r.Host, r.Pref)
	}
	return s
}

// checkMX verifies that the target has expected MX records.
func (p *DnsProber) checkMX() error {
	mx, err := net.LookupMX(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up MX records for %s: %v", p.Target, err)
	}
	mxRec := mxRecords(mx)
	if len(mxRec) != len(p.wantMX) {
		return fmt.Errorf("want %d MX records, got %d: %s", len(p.wantMX), len(mxRec), mxRec)
	}
	sort.Sort(mxRec)
	for i, r := range mxRec {
		if !strings.EqualFold(p.wantMX[i].Host, r.Host) {
			return fmt.Errorf("bad host %q for MX record #%d; want %q", r.Host, i, p.wantMX[i].Host)
		}
		if p.wantMX[i].Pref != r.Pref {
			return fmt.Errorf("bad prio %d for MX record #%d; want %d", i, r.Pref, p.wantMX[i].Pref)
		}
	}
	return nil
}

// checkA verifies that the target has expected A records.
func (p *DnsProber) checkA() error {
	addr, err := net.LookupHost(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up A records for %s: %v", p.Target, err)
	}
	if len(addr) != len(p.wantA) {
		return fmt.Errorf("got %d A records, want %d: %v", len(addr), len(p.wantA), addr)
	}
	sort.Strings(addr)
	for i, a := range addr {
		if p.wantA[i] != a {
			return fmt.Errorf("bad A record %q at #%d; want %q", a, i, p.wantA[i])
		}
	}
	return nil
}

// nsRecords is a collection of NS records, implementing sort.Interface.
type nsRecords []*net.NS

func (r nsRecords) Len() int { return len(r) }

func (r nsRecords) Swap(i, j int) { r[i], r[j] = r[j], r[i] }

func (r nsRecords) Less(i, j int) bool { return r[i].Host < r[j].Host }

// String returns a readable description of the NS records.
func (ns nsRecords) String() string {
	s := ""
	for i, r := range ns {
		if i > 0 {
			s += ", "
		}
		s += r.Host
	}
	return s
}

// checkNS verifies that the target has expected NS records.
func (p *DnsProber) checkNS() error {
	ns, err := net.LookupNS(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up NS records for %s: %v", p.Target, err)
	}
	nsRec := nsRecords(ns)
	if len(nsRec) != len(p.wantNS) {
		return fmt.Errorf("want %d NS records, got %d: %s", len(p.wantNS), len(nsRec), nsRec)
	}
	sort.Sort(nsRec)
	for i, n := range ns {
		if p.wantNS[i].Host != n.Host {
			return fmt.Errorf("bad NS record %q at #%d; want %q", n, i, p.wantNS[i])
		}
	}
	return nil
}

// checkCNAME verifies that the target has expected CNAME record.
func (p *DnsProber) checkCNAME() error {
	cname, err := net.LookupCNAME(p.Target)
	if err != nil {
		return fmt.Errorf("failed to look up CNAME record for %s: %v", p.Target, err)
	}
	if cname != p.wantCNAME {
		return fmt.Errorf("bad CNAME record %q; want %q", cname, p.wantCNAME)
	}
	return nil
}

// checkTXT verifies that the target has expected TXT records.
func (p *DnsProber) checkTXT() error {
	txt, err := net.LookupTXT(p.Target)
	if err != nil {
		return err
	}
	if len(txt) != len(p.wantTXT) {
		return fmt.Errorf("want %d TXT records, got %d: %s", len(p.wantTXT), len(txt), txt)
	}
	sort.Strings(txt)
	for i, t := range txt {
		if p.wantTXT[i] != t {
			return fmt.Errorf("bad TXT record %q at #%d; want %q", t, i, p.wantTXT[i])
		}
	}
	return nil
}

// Alert sends an alert notification via email.
func (p *DnsProber) Alert(name, desc string, badness int, records prober.Records) error {
	return probes.SendAlertEmail(name, desc, badness, records)
}
<|endoftext|>


<commit_before>// Copyright 2017 Authors of Cilium
//
// 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 cmd

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"text/tabwriter"

	"github.com/cilium/cilium/common"
	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/maps/policymap"
	"github.com/cilium/cilium/pkg/policy"
	"github.com/cilium/cilium/pkg/policy/api"

	"github.com/spf13/cobra"
)

var printIDs bool

// bpfPolicyListCmd represents the bpf_policy_list command
var bpfPolicyListCmd = &cobra.Command{
	Use:    "list",
	Short:  "List contents of a policy BPF map",
	PreRun: requireEndpointID,
	Run: func(cmd *cobra.Command, args []string) {
		common.RequireRootPrivilege("cilium bpf policy list")
		listMap(cmd, args)
	},
}

func init() {
	bpfPolicyCmd.AddCommand(bpfPolicyListCmd)
	bpfPolicyListCmd.Flags().BoolVarP(&printIDs, "numeric", "n", false, "Do not resolve IDs")
}

func listMap(cmd *cobra.Command, args []string) {
	lbl := args[0]

	if lbl != "" {
		if id := policy.GetReservedID(lbl); id != policy.IdentityUnknown {
			lbl = "reserved_" + strconv.FormatUint(uint64(id), 10)
		}
	} else {
		Fatalf("Need ID or label\n")
	}

	file := bpf.MapPath(policymap.MapName + lbl)
	fd, err := bpf.ObjGet(file)
	if err != nil {
		Fatalf("%s\n", err)
	}
	defer bpf.ObjClose(fd)

	m := policymap.PolicyMap{Fd: fd}
	statsMap, err := m.DumpToSlice()
	if err != nil {
		Fatalf("Error while opening bpf Map: %s\n", err)
	}

	if handleJSON(statsMap) {
		return
	}
	w := tabwriter.NewWriter(os.Stdout, 5, 0, 3, ' ', 0)
	formatMap(w, statsMap)
	w.Flush()
	if len(statsMap) == 0 {
		fmt.Printf("Policy stats empty. Perhaps the policy enforcement is disabled?\n")
	}
}

func formatMap(w io.Writer, statsMap []policymap.PolicyEntryDump) {
	const (
		labelsIDTitle  = "IDENTITY"
		labelsDesTitle = "LABELS (source:key[=value])"
		actionTitle    = "ACTION"
		bytesTitle     = "BYTES"
		packetsTitle   = "PACKETS"
	)

	labelsID := map[policy.NumericIdentity]*policy.Identity{}
	for _, stat := range statsMap {
		if !printIDs {
			id := policy.NumericIdentity(stat.Key.Identity)
			if lbls, err := client.IdentityGet(id.StringID()); err != nil {
				fmt.Fprintf(os.Stderr, "Was impossible to retrieve label ID %d: %s\n",
					id, err)
			} else {
				labelsID[id] = policy.NewIdentityFromModel(lbls)
			}
		}

	}

	if printIDs {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t\n", labelsIDTitle, actionTitle, bytesTitle, packetsTitle)
	} else {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t\n", labelsDesTitle, actionTitle, bytesTitle, packetsTitle)
	}
	for _, stat := range statsMap {
		id := policy.NumericIdentity(stat.Key.Identity)
		act := api.Decision(stat.Action)
		if printIDs {
			fmt.Fprintf(w, "%d\t%s\t%d\t%d\t\n", id, act.String(), stat.Bytes, stat.Packets)
		} else if lbls := labelsID[id]; lbls != nil {
			first := true
			for _, lbl := range lbls.Labels {
				if first {
					fmt.Fprintf(w, "%s\t%s\t%d\t%d\t\n", lbl, act.String(), stat.Bytes, stat.Packets)
					first = false
				} else {
					fmt.Fprintf(w, "%s\t\t\t\t\t\n", lbl)
				}
			}
		} else {
			fmt.Fprintf(w, "%d\t%s\t%d\t%d\t\n", id, act.String(), stat.Bytes, stat.Packets)
		}
	}
}
<commit_msg>cmd: Add `bpf policy get` alias<commit_after>// Copyright 2017 Authors of Cilium
//
// 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 cmd

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"text/tabwriter"

	"github.com/cilium/cilium/common"
	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/maps/policymap"
	"github.com/cilium/cilium/pkg/policy"
	"github.com/cilium/cilium/pkg/policy/api"

	"github.com/spf13/cobra"
)

var printIDs bool

// bpfPolicyListCmd represents the bpf_policy_list command
var bpfPolicyListCmd = &cobra.Command{
	Use:     "list",
	Aliases: []string{"get"},
	Short:   "List contents of a policy BPF map",
	PreRun:  requireEndpointID,
	Run: func(cmd *cobra.Command, args []string) {
		common.RequireRootPrivilege("cilium bpf policy list")
		listMap(cmd, args)
	},
}

func init() {
	bpfPolicyCmd.AddCommand(bpfPolicyListCmd)
	bpfPolicyListCmd.Flags().BoolVarP(&printIDs, "numeric", "n", false, "Do not resolve IDs")
}

func listMap(cmd *cobra.Command, args []string) {
	lbl := args[0]

	if lbl != "" {
		if id := policy.GetReservedID(lbl); id != policy.IdentityUnknown {
			lbl = "reserved_" + strconv.FormatUint(uint64(id), 10)
		}
	} else {
		Fatalf("Need ID or label\n")
	}

	file := bpf.MapPath(policymap.MapName + lbl)
	fd, err := bpf.ObjGet(file)
	if err != nil {
		Fatalf("%s\n", err)
	}
	defer bpf.ObjClose(fd)

	m := policymap.PolicyMap{Fd: fd}
	statsMap, err := m.DumpToSlice()
	if err != nil {
		Fatalf("Error while opening bpf Map: %s\n", err)
	}

	if handleJSON(statsMap) {
		return
	}
	w := tabwriter.NewWriter(os.Stdout, 5, 0, 3, ' ', 0)
	formatMap(w, statsMap)
	w.Flush()
	if len(statsMap) == 0 {
		fmt.Printf("Policy stats empty. Perhaps the policy enforcement is disabled?\n")
	}
}

func formatMap(w io.Writer, statsMap []policymap.PolicyEntryDump) {
	const (
		labelsIDTitle  = "IDENTITY"
		labelsDesTitle = "LABELS (source:key[=value])"
		actionTitle    = "ACTION"
		bytesTitle     = "BYTES"
		packetsTitle   = "PACKETS"
	)

	labelsID := map[policy.NumericIdentity]*policy.Identity{}
	for _, stat := range statsMap {
		if !printIDs {
			id := policy.NumericIdentity(stat.Key.Identity)
			if lbls, err := client.IdentityGet(id.StringID()); err != nil {
				fmt.Fprintf(os.Stderr, "Was impossible to retrieve label ID %d: %s\n",
					id, err)
			} else {
				labelsID[id] = policy.NewIdentityFromModel(lbls)
			}
		}

	}

	if printIDs {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t\n", labelsIDTitle, actionTitle, bytesTitle, packetsTitle)
	} else {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t\n", labelsDesTitle, actionTitle, bytesTitle, packetsTitle)
	}
	for _, stat := range statsMap {
		id := policy.NumericIdentity(stat.Key.Identity)
		act := api.Decision(stat.Action)
		if printIDs {
			fmt.Fprintf(w, "%d\t%s\t%d\t%d\t\n", id, act.String(), stat.Bytes, stat.Packets)
		} else if lbls := labelsID[id]; lbls != nil {
			first := true
			for _, lbl := range lbls.Labels {
				if first {
					fmt.Fprintf(w, "%s\t%s\t%d\t%d\t\n", lbl, act.String(), stat.Bytes, stat.Packets)
					first = false
				} else {
					fmt.Fprintf(w, "%s\t\t\t\t\t\n", lbl)
				}
			}
		} else {
			fmt.Fprintf(w, "%d\t%s\t%d\t%d\t\n", id, act.String(), stat.Bytes, stat.Packets)
		}
	}
}
<|endoftext|>


<commit_before>package clicommand

import (
	"github.com/buildkite/agent/v3/agent"
	"github.com/buildkite/agent/v3/api"
	"github.com/buildkite/agent/v3/cliconfig"
	"github.com/urfave/cli"
)

var UploadHelpDescription = `Usage:

   buildkite-agent artifact upload [options] <pattern> [destination]

Description:

   Uploads files to a job as artifacts.

   You need to ensure that the paths are surrounded by quotes otherwise the
   built-in shell path globbing will provide the files, which is currently not
   supported.

   You can specify an alternate destination on Amazon S3, Google Cloud Storage
   or Artifactory as per the examples below. This may be specified in the
   'destination' argument, or in the 'BUILDKITE_ARTIFACT_UPLOAD_DESTINATION'
   environment variable.  Otherwise, artifacts are uploaded to a
   Buildkite-managed Amazon S3 bucket.

Example:

   $ buildkite-agent artifact upload "log/**/*.log"

   You can also upload directly to Amazon S3 if you'd like to host your own artifacts:

   $ export BUILDKITE_S3_ACCESS_KEY_ID=xxx
   $ export BUILDKITE_S3_SECRET_ACCESS_KEY=yyy
   $ export BUILDKITE_S3_DEFAULT_REGION=eu-central-1 # default is us-east-1
   $ export BUILDKITE_S3_ACL=private # default is public-read
   $ buildkite-agent artifact upload "log/**/*.log" s3://name-of-your-s3-bucket/$BUILDKITE_JOB_ID

   You can use Amazon IAM assumed roles by specifying the session token:

   $ export BUILDKITE_S3_SESSION_TOKEN=zzz

   Or upload directly to Google Cloud Storage:

   $ export BUILDKITE_GS_ACL=private
   $ buildkite-agent artifact upload "log/**/*.log" gs://name-of-your-gs-bucket/$BUILDKITE_JOB_ID

   Or upload directly to Artifactory:

   $ export BUILDKITE_ARTIFACTORY_URL=http://my-artifactory-instance.com/artifactory
   $ export BUILDKITE_ARTIFACTORY_USER=carol-danvers
   $ export BUILDKITE_ARTIFACTORY_PASSWORD=xxx
   $ buildkite-agent artifact upload "log/**/*.log" rt://name-of-your-artifactory-repo/$BUILDKITE_JOB_ID`

var FollowSymlinksFlag = cli.BoolFlag{
	Name:   "follow-symlinks",
	Usage:  "Follow symbolic links while resolving globs",
	EnvVar: "BUILDKITE_AGENT_ARTIFACT_SYMLINKS",
}

type ArtifactUploadConfig struct {
	UploadPaths string `cli:"arg:0" label:"upload paths" validate:"required"`
	Destination string `cli:"arg:1" label:"destination" env:"BUILDKITE_ARTIFACT_UPLOAD_DESTINATION"`
	Job         string `cli:"job" validate:"required"`
	ContentType string `cli:"content-type"`

	// Global flags
	Debug       bool     `cli:"debug"`
	NoColor     bool     `cli:"no-color"`
	Experiments []string `cli:"experiment" normalize:"list"`
	Profile     string   `cli:"profile"`

	// API config
	DebugHTTP        bool   `cli:"debug-http"`
	AgentAccessToken string `cli:"agent-access-token" validate:"required"`
	Endpoint         string `cli:"endpoint" validate:"required"`
	NoHTTP2          bool   `cli:"no-http2"`

	// Uploader flags
	FollowSymlinks bool `cli:"follow-symlinks"`
}

var ArtifactUploadCommand = cli.Command{
	Name:        "upload",
	Usage:       "Uploads files to a job as artifacts",
	Description: UploadHelpDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "job",
			Value:  "",
			Usage:  "Which job should the artifacts be uploaded to",
			EnvVar: "BUILDKITE_JOB_ID",
		},
		cli.StringFlag{
			Name:   "content-type",
			Value:  "",
			Usage:  "A specific Content-Type to set for the artifacts (otherwise detected)",
			EnvVar: "BUILDKITE_ARTIFACT_CONTENT_TYPE",
		},

		// API Flags
		AgentAccessTokenFlag,
		EndpointFlag,
		NoHTTP2Flag,
		DebugHTTPFlag,

		// Global flags
		NoColorFlag,
		DebugFlag,
		ExperimentsFlag,
		ProfileFlag,
		FollowSymlinksFlag,
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := ArtifactUploadConfig{}

		l := CreateLogger(&cfg)

		// Load the configuration
		if err := cliconfig.Load(c, l, &cfg); err != nil {
			l.Fatal("%s", err)
		}

		// Setup any global configuration options
		done := HandleGlobalFlags(l, cfg)
		defer done()

		// Create the API client
		client := api.NewClient(l, loadAPIClientConfig(cfg, `AgentAccessToken`))

		// Setup the uploader
		uploader := agent.NewArtifactUploader(l, client, agent.ArtifactUploaderConfig{
			JobID:          cfg.Job,
			Paths:          cfg.UploadPaths,
			Destination:    cfg.Destination,
			ContentType:    cfg.ContentType,
			DebugHTTP:      cfg.DebugHTTP,
			FollowSymlinks: cfg.FollowSymlinks,
		})

		// Upload the artifacts
		if err := uploader.Upload(); err != nil {
			l.Fatal("Failed to upload artifacts: %s", err)
		}
	},
}
<commit_msg>Copy docs change from https://github.com/buildkite/docs/pull/1128<commit_after>package clicommand

import (
	"github.com/buildkite/agent/v3/agent"
	"github.com/buildkite/agent/v3/api"
	"github.com/buildkite/agent/v3/cliconfig"
	"github.com/urfave/cli"
)

var UploadHelpDescription = `Usage:

   buildkite-agent artifact upload [options] <pattern> [destination]

Description:

   Uploads files to a job as artifacts.

   You need to ensure that the paths are surrounded by quotes otherwise the
   built-in shell path globbing will provide the files, which is currently not
   supported.

   You can specify an alternate destination on Amazon S3, Google Cloud Storage
   or Artifactory as per the examples below. This may be specified in the
   'destination' argument, or in the 'BUILDKITE_ARTIFACT_UPLOAD_DESTINATION'
   environment variable.  Otherwise, artifacts are uploaded to a
   Buildkite-managed Amazon S3 bucket, where they’re retained for six months.

Example:

   $ buildkite-agent artifact upload "log/**/*.log"

   You can also upload directly to Amazon S3 if you'd like to host your own artifacts:

   $ export BUILDKITE_S3_ACCESS_KEY_ID=xxx
   $ export BUILDKITE_S3_SECRET_ACCESS_KEY=yyy
   $ export BUILDKITE_S3_DEFAULT_REGION=eu-central-1 # default is us-east-1
   $ export BUILDKITE_S3_ACL=private # default is public-read
   $ buildkite-agent artifact upload "log/**/*.log" s3://name-of-your-s3-bucket/$BUILDKITE_JOB_ID

   You can use Amazon IAM assumed roles by specifying the session token:

   $ export BUILDKITE_S3_SESSION_TOKEN=zzz

   Or upload directly to Google Cloud Storage:

   $ export BUILDKITE_GS_ACL=private
   $ buildkite-agent artifact upload "log/**/*.log" gs://name-of-your-gs-bucket/$BUILDKITE_JOB_ID

   Or upload directly to Artifactory:

   $ export BUILDKITE_ARTIFACTORY_URL=http://my-artifactory-instance.com/artifactory
   $ export BUILDKITE_ARTIFACTORY_USER=carol-danvers
   $ export BUILDKITE_ARTIFACTORY_PASSWORD=xxx
   $ buildkite-agent artifact upload "log/**/*.log" rt://name-of-your-artifactory-repo/$BUILDKITE_JOB_ID`

var FollowSymlinksFlag = cli.BoolFlag{
	Name:   "follow-symlinks",
	Usage:  "Follow symbolic links while resolving globs",
	EnvVar: "BUILDKITE_AGENT_ARTIFACT_SYMLINKS",
}

type ArtifactUploadConfig struct {
	UploadPaths string `cli:"arg:0" label:"upload paths" validate:"required"`
	Destination string `cli:"arg:1" label:"destination" env:"BUILDKITE_ARTIFACT_UPLOAD_DESTINATION"`
	Job         string `cli:"job" validate:"required"`
	ContentType string `cli:"content-type"`

	// Global flags
	Debug       bool     `cli:"debug"`
	NoColor     bool     `cli:"no-color"`
	Experiments []string `cli:"experiment" normalize:"list"`
	Profile     string   `cli:"profile"`

	// API config
	DebugHTTP        bool   `cli:"debug-http"`
	AgentAccessToken string `cli:"agent-access-token" validate:"required"`
	Endpoint         string `cli:"endpoint" validate:"required"`
	NoHTTP2          bool   `cli:"no-http2"`

	// Uploader flags
	FollowSymlinks bool `cli:"follow-symlinks"`
}

var ArtifactUploadCommand = cli.Command{
	Name:        "upload",
	Usage:       "Uploads files to a job as artifacts",
	Description: UploadHelpDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "job",
			Value:  "",
			Usage:  "Which job should the artifacts be uploaded to",
			EnvVar: "BUILDKITE_JOB_ID",
		},
		cli.StringFlag{
			Name:   "content-type",
			Value:  "",
			Usage:  "A specific Content-Type to set for the artifacts (otherwise detected)",
			EnvVar: "BUILDKITE_ARTIFACT_CONTENT_TYPE",
		},

		// API Flags
		AgentAccessTokenFlag,
		EndpointFlag,
		NoHTTP2Flag,
		DebugHTTPFlag,

		// Global flags
		NoColorFlag,
		DebugFlag,
		ExperimentsFlag,
		ProfileFlag,
		FollowSymlinksFlag,
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := ArtifactUploadConfig{}

		l := CreateLogger(&cfg)

		// Load the configuration
		if err := cliconfig.Load(c, l, &cfg); err != nil {
			l.Fatal("%s", err)
		}

		// Setup any global configuration options
		done := HandleGlobalFlags(l, cfg)
		defer done()

		// Create the API client
		client := api.NewClient(l, loadAPIClientConfig(cfg, `AgentAccessToken`))

		// Setup the uploader
		uploader := agent.NewArtifactUploader(l, client, agent.ArtifactUploaderConfig{
			JobID:          cfg.Job,
			Paths:          cfg.UploadPaths,
			Destination:    cfg.Destination,
			ContentType:    cfg.ContentType,
			DebugHTTP:      cfg.DebugHTTP,
			FollowSymlinks: cfg.FollowSymlinks,
		})

		// Upload the artifacts
		if err := uploader.Upload(); err != nil {
			l.Fatal("Failed to upload artifacts: %s", err)
		}
	},
}
<|endoftext|>


<commit_before>package mario

import (
	util "github.com/poseidon4o/go-neural/src/util"
	"math"
)

const G_CONST float64 = 9.8 * 150

var G_FORCE util.Vector = util.Vector{
	X: 0,
	Y: G_CONST,
}

const BLOCK_SIZE int = 24

var JUMP_FORCE util.Vector = util.Vector{
	X: 0,
	Y: -1000,
}

var X_ACCELERATION util.Vector = util.Vector{
	X: 20,
	Y: 0,
}

type Figure struct {
	pos     util.Vector
	vel     util.Vector
	nextPos util.Vector
	jumps   int
}

func (m *Figure) Jump() {
	if m.jumps >= 3 {
		return
	}
	m.jumps++
	m.vel = *m.vel.Add(&JUMP_FORCE)
	m.vel.Y = math.Max(m.vel.Y, JUMP_FORCE.Y)
}

func (f *Figure) Land() {
	f.jumps = 0
}

func (f *Figure) Move(dir int) {
	acc := X_ACCELERATION
	if dir < 0 {
		acc = *acc.Neg()
	}
	f.vel = *f.vel.Add(&acc)
}

type Level struct {
	size    util.Vector
	blocks  [][]util.Vector
	figures []*Figure
}

func NewLevel(w, h int) *Level {
	lvl := &Level{
		size:    *util.NewVector(float64(w), float64(h)),
		blocks:  make([][]util.Vector, 0),
		figures: make([]*Figure, 0),
	}

	// TODO generate level
	for c := 0; c < w; c += BLOCK_SIZE {
		col := make([]util.Vector, 1, 1)
		col[0] = *util.NewVector(float64(c), float64(h-BLOCK_SIZE))
		lvl.blocks = append(lvl.blocks, col)
	}

	return lvl
}

func (l *Level) FloorAt(pos *util.Vector) *util.Vector {
	idx := int(pos.X / float64(BLOCK_SIZE))

	if idx < 0 || idx >= len(l.blocks) {
		return util.NewVector(0, 0)
	}

	for c := len(l.blocks[idx]) - 1; c >= 0; c-- {
		if pos.Y < l.blocks[idx][c].Y {
			// we are above the block
			if c-1 >= 0 && pos.Y > l.blocks[idx][c-1].Y {
				// has next block and we are below it
				return &l.blocks[idx][c]
			} else if c-1 < 0 {
				// dont have next block
				return &l.blocks[idx][c]
			}
		}
	}

	return util.NewVector(0, 0)
}

func (l *Level) NewFigurePos() *util.Vector {
	return util.NewVector(1, l.FloorAt(util.NewVector(1, 1)).Y-float64(BLOCK_SIZE))
}

func (l *Level) AddFigures(count int) {
	for c := 0; c < count; c++ {

		l.figures = append(l.figures, &Figure{
			pos:     *l.NewFigurePos(),
			vel:     *util.NewVector(0, 0),
			nextPos: *util.NewVector(0, 0),
		})
	}
}

func (l *Level) Step(dt float64) {
	for c := range l.figures {
		// position += timestep * (velocity + timestep * acceleration / 2);
		// TODO not use go
		l.figures[c].nextPos = *l.figures[c].pos.Add(G_FORCE.Scale(dt / 2).Add(&l.figures[c].vel).Scale(dt))

		// velocity += timestep * acceleration;
		l.figures[c].vel = *l.figures[c].vel.Add(G_FORCE.Scale(dt))
		l.figures[c].vel.X *= 0.9
	}
}
<commit_msg>Change level rep to be 2d array of pts<commit_after>package mario

import (
	neural "github.com/poseidon4o/go-neural/src/neural"
	util "github.com/poseidon4o/go-neural/src/util"
	"math"
)

const G_CONST float64 = 9.8 * 150

var G_FORCE util.Vector = util.Vector{
	X: 0,
	Y: G_CONST,
}

const BLOCK_SIZE int = 25

var JUMP_FORCE util.Vector = util.Vector{
	X: 0,
	Y: -1000,
}

var X_ACCELERATION util.Vector = util.Vector{
	X: 20,
	Y: 0,
}

type Figure struct {
	pos     util.Vector
	vel     util.Vector
	nextPos util.Vector
	jumps   int
}

func (m *Figure) Jump() {
	if m.jumps >= 1 {
		return
	}
	m.jumps++
	m.vel = *m.vel.Add(&JUMP_FORCE)
	m.vel.Y = math.Max(m.vel.Y, JUMP_FORCE.Y)
}

func (f *Figure) Land() {
	f.jumps = 0
}

func (f *Figure) Move(dir int) {
	acc := X_ACCELERATION
	if dir < 0 {
		acc = *acc.Neg()
	}
	f.vel = *f.vel.Add(&acc)
}

type Level struct {
	size    util.Vector
	blocks  [][]*util.Vector
	figures []*Figure
}

func (l *Level) makeHole(c *int) {
	*c += 3 // int(3 + neural.RandMax(2))
}

func (l *Level) makeObstacle(c *int) {
	*c--
}

func (l *Level) makeGround(c *int) {
	r := int(l.size.Y/float64(BLOCK_SIZE)) - 1

	l.blocks[*c][r] = util.NewVector(float64(*c*BLOCK_SIZE), float64(r*BLOCK_SIZE))
}

func NewLevel(w, h int) *Level {
	blockH := h / BLOCK_SIZE
	blockW := w / BLOCK_SIZE

	lvl := &Level{
		size:    *util.NewVector(float64(w), float64(h)),
		blocks:  make([][]*util.Vector, blockW, blockW),
		figures: make([]*Figure, 0),
	}

	for c := 0; c < blockW; c++ {
		lvl.blocks[c] = make([]*util.Vector, blockH, blockH)
		for r := 0; r < blockH; r++ {
			lvl.blocks[c][r] = nil
		}
	}

	for c, obs := 0, 1; c < blockW; c, obs = c+1, obs+1 {

		pr := c
		if obs%10 == 0 {
			if neural.Chance(0.5) {
				lvl.makeHole(&c)
			} else {
				lvl.makeObstacle(&c)
			}
		} else {
			lvl.makeGround(&c)
		}
		obs += c - pr
	}

	return lvl
}

func (l *Level) FloorAt(pos *util.Vector) *util.Vector {
	wIdx := int(pos.X / float64(BLOCK_SIZE))
	hIdx := int(pos.Y / float64(BLOCK_SIZE))

	if wIdx < 0 || wIdx >= len(l.blocks) || hIdx < 0 || hIdx >= len(l.blocks[0]) {
		return nil
	}

	if hIdx+1 >= len(l.blocks[0]) {
		return nil
	}

	return l.blocks[wIdx][hIdx+1]
}

func (l *Level) NewFigurePos() *util.Vector {
	return util.NewVector(1, 1)
}

func (l *Level) AddFigures(count int) {
	for c := 0; c < count; c++ {

		l.figures = append(l.figures, &Figure{
			pos:     *l.NewFigurePos(),
			vel:     *util.NewVector(0, 0),
			nextPos: *util.NewVector(0, 0),
		})
	}
}

func (l *Level) Step(dt float64) {
	for c := range l.figures {
		// position += timestep * (velocity + timestep * acceleration / 2);
		// TODO not use go
		l.figures[c].nextPos = *l.figures[c].pos.Add(G_FORCE.Scale(dt / 2).Add(&l.figures[c].vel).Scale(dt))

		// velocity += timestep * acceleration;
		l.figures[c].vel = *l.figures[c].vel.Add(G_FORCE.Scale(dt))
		l.figures[c].vel.X *= (1 - 3*dt)
	}
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2019 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package statsd

import (
	"fmt"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/sirupsen/logrus"

	"github.com/loadimpact/k6/output"
	"github.com/loadimpact/k6/stats"
)

// New creates a new statsd connector client
func New(params output.Params) (output.Output, error) {
	return newOutput(params)
}

func newOutput(params output.Params) (*Output, error) {
	conf, err := getConsolidatedConfig(params.JSONConfig, params.Environment, params.ConfigArgument)
	if err != nil {
		return nil, err
	}
	logger := params.Logger.WithFields(logrus.Fields{"output": "statsd"})

	return &Output{
		config: conf,
		logger: logger,
	}, nil
}

var _ output.Output = &Output{}

// Output sends result data to statsd daemons with the ability to send to datadog as well
type Output struct {
	output.SampleBuffer

	periodicFlusher *output.PeriodicFlusher

	config config

	logger logrus.FieldLogger
	client *statsd.Client
}

func (o *Output) dispatch(entry stats.Sample) error {
	var tagList []string
	if o.config.EnableTags.Bool {
		tagList = processTags(o.config.TagBlocklist, entry.Tags.CloneTags())
	}

	switch entry.Metric.Type {
	case stats.Counter:
		return o.client.Count(entry.Metric.Name, int64(entry.Value), tagList, 1)
	case stats.Trend:
		return o.client.TimeInMilliseconds(entry.Metric.Name, entry.Value, tagList, 1)
	case stats.Gauge:
		return o.client.Gauge(entry.Metric.Name, entry.Value, tagList, 1)
	case stats.Rate:
		if check, ok := entry.Tags.Get("check"); ok {
			return o.client.Count(
				checkToString(check, entry.Value),
				1,
				tagList,
				1,
			)
		}
		return o.client.Count(entry.Metric.Name, int64(entry.Value), tagList, 1)
	default:
		return fmt.Errorf("unsupported metric type %s", entry.Metric.Type)
	}
}

func checkToString(check string, value float64) string {
	label := "pass"
	if value == 0 {
		label = "fail"
	}
	return "check." + check + "." + label
}

// Description returns a human-readable description of the output.
func (o *Output) Description() string {
	return fmt.Sprintf("statsd (%s)", o.config.Addr.String)
}

// Start tries to open a connection to specified statsd service and starts the goroutine for
// metric flushing.
func (o *Output) Start() error {
	o.logger.Debug("Starting...")

	var err error
	if address := o.config.Addr.String; address == "" {
		err = fmt.Errorf(
			"connection string is invalid. Received: \"%+s\"",
			address,
		)
		o.logger.Error(err)

		return err
	}

	o.client, err = statsd.NewBuffered(o.config.Addr.String, int(o.config.BufferSize.Int64))

	if err != nil {
		o.logger.Errorf("Couldn't make buffered client, %s", err)
		return err
	}

	if namespace := o.config.Namespace.String; namespace != "" {
		o.client.Namespace = namespace
	}

	pf, err := output.NewPeriodicFlusher(time.Duration(o.config.PushInterval.Duration), o.flushMetrics)
	if err != nil {
		return err
	}
	o.logger.Debug("Started!")
	o.periodicFlusher = pf

	return nil
}

// Stop flushes any remaining metrics and stops the goroutine.
func (o *Output) Stop() error {
	o.logger.Debug("Stopping...")
	defer o.logger.Debug("Stopped!")
	o.periodicFlusher.Stop()
	return o.client.Close()
}

func (o *Output) flushMetrics() {
	samples := o.GetBufferedSamples()
	start := time.Now()
	var count int
	var errorCount int
	for _, sc := range samples {
		samples := sc.GetSamples()
		count += len(samples)
		o.logger.
			WithField("samples", len(samples)).
			Debug("Pushing metrics to server")

		for _, entry := range samples {
			if err := o.dispatch(entry); err != nil {
				// No need to return error if just one metric didn't go through
				o.logger.WithError(err).Debugf("Error while sending metric %s", entry.Metric.Name)
				errorCount++
			}
		}
	}

	if count > 0 {
		if errorCount != 0 {
			o.logger.Warnf("Couldn't send %d out of %d metrics. Enable debug logging to see individual errors",
				errorCount, count)
		}
		if err := o.client.Flush(); err != nil {
			o.logger.
				WithError(err).
				Error("Couldn't flush a batch")
		}
		o.logger.WithField("t", time.Since(start)).WithField("count", count).Debug("Wrote metrics to statsd")
	}
}
<commit_msg>Fix log message in statsd output<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2019 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package statsd

import (
	"fmt"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/sirupsen/logrus"

	"github.com/loadimpact/k6/output"
	"github.com/loadimpact/k6/stats"
)

// New creates a new statsd connector client
func New(params output.Params) (output.Output, error) {
	return newOutput(params)
}

func newOutput(params output.Params) (*Output, error) {
	conf, err := getConsolidatedConfig(params.JSONConfig, params.Environment, params.ConfigArgument)
	if err != nil {
		return nil, err
	}
	logger := params.Logger.WithFields(logrus.Fields{"output": "statsd"})

	return &Output{
		config: conf,
		logger: logger,
	}, nil
}

var _ output.Output = &Output{}

// Output sends result data to statsd daemons with the ability to send to datadog as well
type Output struct {
	output.SampleBuffer

	periodicFlusher *output.PeriodicFlusher

	config config

	logger logrus.FieldLogger
	client *statsd.Client
}

func (o *Output) dispatch(entry stats.Sample) error {
	var tagList []string
	if o.config.EnableTags.Bool {
		tagList = processTags(o.config.TagBlocklist, entry.Tags.CloneTags())
	}

	switch entry.Metric.Type {
	case stats.Counter:
		return o.client.Count(entry.Metric.Name, int64(entry.Value), tagList, 1)
	case stats.Trend:
		return o.client.TimeInMilliseconds(entry.Metric.Name, entry.Value, tagList, 1)
	case stats.Gauge:
		return o.client.Gauge(entry.Metric.Name, entry.Value, tagList, 1)
	case stats.Rate:
		if check, ok := entry.Tags.Get("check"); ok {
			return o.client.Count(
				checkToString(check, entry.Value),
				1,
				tagList,
				1,
			)
		}
		return o.client.Count(entry.Metric.Name, int64(entry.Value), tagList, 1)
	default:
		return fmt.Errorf("unsupported metric type %s", entry.Metric.Type)
	}
}

func checkToString(check string, value float64) string {
	label := "pass"
	if value == 0 {
		label = "fail"
	}
	return "check." + check + "." + label
}

// Description returns a human-readable description of the output.
func (o *Output) Description() string {
	return fmt.Sprintf("statsd (%s)", o.config.Addr.String)
}

// Start tries to open a connection to specified statsd service and starts the goroutine for
// metric flushing.
func (o *Output) Start() error {
	o.logger.Debug("Starting...")

	var err error
	if address := o.config.Addr.String; address == "" {
		err = fmt.Errorf(
			"connection string is invalid. Received: \"%+s\"",
			address,
		)
		o.logger.Error(err)

		return err
	}

	o.client, err = statsd.NewBuffered(o.config.Addr.String, int(o.config.BufferSize.Int64))

	if err != nil {
		o.logger.Errorf("Couldn't make buffered client, %s", err)
		return err
	}

	if namespace := o.config.Namespace.String; namespace != "" {
		o.client.Namespace = namespace
	}

	pf, err := output.NewPeriodicFlusher(time.Duration(o.config.PushInterval.Duration), o.flushMetrics)
	if err != nil {
		return err
	}
	o.logger.Debug("Started!")
	o.periodicFlusher = pf

	return nil
}

// Stop flushes any remaining metrics and stops the goroutine.
func (o *Output) Stop() error {
	o.logger.Debug("Stopping...")
	defer o.logger.Debug("Stopped!")
	o.periodicFlusher.Stop()
	return o.client.Close()
}

func (o *Output) flushMetrics() {
	samples := o.GetBufferedSamples()
	start := time.Now()
	var count int
	var errorCount int
	for _, sc := range samples {
		samples := sc.GetSamples()
		count += len(samples)
		o.logger.
			WithField("samples", len(samples)).
			Debug("Pushing metrics to server")

		for _, entry := range samples {
			if err := o.dispatch(entry); err != nil {
				// No need to return error if just one metric didn't go through
				o.logger.WithError(err).Debugf("Error while sending metric %s", entry.Metric.Name)
				errorCount++
			}
		}
	}

	if count > 0 {
		if errorCount != 0 {
			o.logger.Warnf("Couldn't send %d out of %d metrics. Enable verbose logging with --verbose to see individual errors",
				errorCount, count)
		}
		if err := o.client.Flush(); err != nil {
			o.logger.
				WithError(err).
				Error("Couldn't flush a batch")
		}
		o.logger.WithField("t", time.Since(start)).WithField("count", count).Debug("Wrote metrics to statsd")
	}
}
<|endoftext|>


<commit_before>// +build linux

package specs

// LinuxSpec is the full specification for linux containers.
type LinuxSpec struct {
	Spec
	// Linux is platform specific configuration for linux based containers.
	Linux Linux `json:"linux"`
}

// Linux contains platform specific configuration for linux based containers.
type Linux struct {
	// Capabilities are linux capabilities that are kept for the container.
	Capabilities []string `json:"capabilities"`
	// RootfsPropagation is the rootfs mount propagation mode for the container.
	RootfsPropagation string `json:"rootfsPropagation"`
}

// User specifies linux specific user and group information for the container's
// main process.
type User struct {
	// Uid is the user id.
	Uid int32 `json:"uid"`
	// Gid is the group id.
	Gid int32 `json:"gid"`
	// AdditionalGids are additional group ids set for the container's process.
	AdditionalGids []int32 `json:"additionalGids"`
}
<commit_msg>Return golang compliant names for UID and GID in User<commit_after>// +build linux

package specs

// LinuxSpec is the full specification for linux containers.
type LinuxSpec struct {
	Spec
	// Linux is platform specific configuration for linux based containers.
	Linux Linux `json:"linux"`
}

// Linux contains platform specific configuration for linux based containers.
type Linux struct {
	// Capabilities are linux capabilities that are kept for the container.
	Capabilities []string `json:"capabilities"`
	// RootfsPropagation is the rootfs mount propagation mode for the container.
	RootfsPropagation string `json:"rootfsPropagation"`
}

// User specifies linux specific user and group information for the container's
// main process.
type User struct {
	// UID is the user id.
	UID int32 `json:"uid"`
	// GID is the group id.
	GID int32 `json:"gid"`
	// AdditionalGids are additional group ids set for the container's process.
	AdditionalGids []int32 `json:"additionalGids"`
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/grafov/autograf/client"
)

/* This is a simple example of usage of Grafana client
for copying dashboards and saving them to a disk.
Ie really useful for Grafana backups!

Usage:
  backup-dashboards http://grafana.host:3000 api-key-string-here
*/

func main() {
	var (
		boardLinks []client.FoundBoard
		rawBoard   []byte
		meta       client.BoardProperties
		err        error
	)
	if len(os.Args) != 3 {
		fmt.Fprint(os.Stderr, "Usage:  backup-dashboards http://grafana.host:3000 api-key-string-here\n")
		os.Exit(0)
	}
	c := client.New(os.Args[1], os.Args[2])
	if boardLinks, err = c.SearchDashboards("", false); err != nil {
		fmt.Fprintf(os.Stderr, fmt.Sprintf("%s\n", err))
		os.Exit(1)
	}
	for _, link := range boardLinks {
		if rawBoard, meta, err = c.GetRawDashboard(link.URI); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, link.URI))
			continue
		}
		if err = ioutil.WriteFile(fmt.Sprintf("%s.json", meta.Slug), rawBoard, os.FileMode(int(0666))); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, meta.Slug))
		}
	}
}
<commit_msg>Add example of backup datasources.<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/grafov/autograf/client"
)

/* This is a simple example of usage of Grafana client
for copying dashboards and saving them to a disk.
It really useful for Grafana backups!

Usage:
  backup-dashboards http://grafana.host:3000 api-key-string-here
*/

func main() {
	var (
		boardLinks []client.FoundBoard
		rawBoard   []byte
		meta       client.BoardProperties
		err        error
	)
	if len(os.Args) != 3 {
		fmt.Fprint(os.Stderr, "Usage:  backup-dashboards http://grafana.host:3000 api-key-string-here\n")
		os.Exit(0)
	}
	c := client.New(os.Args[1], os.Args[2])
	if boardLinks, err = c.SearchDashboards("", false); err != nil {
		fmt.Fprintf(os.Stderr, fmt.Sprintf("%s\n", err))
		os.Exit(1)
	}
	for _, link := range boardLinks {
		if rawBoard, meta, err = c.GetRawDashboard(link.URI); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, link.URI))
			continue
		}
		if err = ioutil.WriteFile(fmt.Sprintf("%s.json", meta.Slug), rawBoard, os.FileMode(int(0666))); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, meta.Slug))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"strings"
	"testing"

	"github.com/golang/dep"
	"github.com/golang/dep/internal/gps"
	"github.com/golang/dep/internal/test"
	"github.com/pkg/errors"
)

func TestRootAnalyzer_Info(t *testing.T) {
	testCases := map[bool]string{
		true:  "dep",
		false: "dep+import",
	}
	for skipTools, want := range testCases {
		a := rootAnalyzer{skipTools: skipTools}
		got := a.Info().Name
		if got != want {
			t.Errorf("Expected the name of the importer with skipTools=%t to be '%s', got '%s'", skipTools, want, got)
		}
	}
}

func TestLookupVersionForLockedProject_MatchRevisionToTag(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	rev := gps.Revision("ff2948a2ac8f538c4ecd55962e919d1e13e74baf")
	v, err := lookupVersionForLockedProject(pi, nil, rev, sm)
	h.Must(err)

	wantV := "v1.0.0"
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to be the tag paired with the manifest's pinned revision: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestLookupVersionForLockedProject_MatchRevisionToMultipleTags(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	// Both 0.8.0 and 1.0.0 use the same rev, force dep to pick the lower version
	c, _ := gps.NewSemverConstraint("<1.0.0")
	rev := gps.Revision("ff2948a2ac8f538c4ecd55962e919d1e13e74baf")
	v, err := lookupVersionForLockedProject(pi, c, rev, sm)
	h.Must(err)

	wantV := "v0.8.0"
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to satisfy the manifest's semver constraint: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestLookupVersionForLockedProject_FallbackToConstraint(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	c := gps.NewBranch("master")
	rev := gps.Revision("c575196502940c07bf89fd6d95e83b999162e051")
	v, err := lookupVersionForLockedProject(pi, c, rev, sm)
	h.Must(err)

	wantV := c.String()
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to be defaulted from the manifest's branch constraint: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestLookupVersionForLockedProject_FallbackToRevision(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	rev := gps.Revision("c575196502940c07bf89fd6d95e83b999162e051")
	v, err := lookupVersionForLockedProject(pi, nil, rev, sm)
	h.Must(err)

	wantV := rev.String()
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to be the manifest's pinned revision: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestProjectExistsInLock(t *testing.T) {
	lock := &dep.Lock{}
	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	ver := gps.NewVersion("v1.0.0")
	lock.P = append(lock.P, gps.NewLockedProject(pi, ver, nil))

	cases := []struct {
		name       string
		importPath string
		want       bool
	}{
		{
			name:       "root project",
			importPath: "github.com/sdboyer/deptest",
			want:       true,
		},
		{
			name:       "sub package",
			importPath: "github.com/sdboyer/deptest/foo",
			want:       false,
		},
		{
			name:       "nonexisting project",
			importPath: "github.com/golang/notexist",
			want:       false,
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			result := projectExistsInLock(lock, gps.ProjectRoot(c.importPath))

			if result != c.want {
				t.Fatalf("projectExistsInLock result is not as want: \n\t(GOT) %v \n\t(WNT) %v", result, c.want)
			}
		})
	}
}

func TestIsVersion(t *testing.T) {
	testcases := map[string]struct {
		wantIsVersion bool
		wantVersion   gps.Version
	}{
		"v1.0.0": {wantIsVersion: true, wantVersion: gps.NewVersion("v1.0.0").Pair("ff2948a2ac8f538c4ecd55962e919d1e13e74baf")},
		"3f4c3bea144e112a69bbe5d8d01c1b09a544253f": {wantIsVersion: false},
		"master": {wantIsVersion: false},
	}

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	for value, testcase := range testcases {
		t.Run(value, func(t *testing.T) {
			gotIsVersion, gotVersion, err := isVersion(pi, value, sm)
			h.Must(err)

			if testcase.wantIsVersion != gotIsVersion {
				t.Fatalf("unexpected isVersion result for %s: \n\t(GOT) %v \n\t(WNT) %v", value, gotIsVersion, testcase.wantIsVersion)
			}

			if testcase.wantVersion != gotVersion {
				t.Fatalf("unexpected version for %s: \n\t(GOT) %v \n\t(WNT) %v", value, testcase.wantVersion, gotVersion)
			}
		})
	}
}

// convertTestCase is a common set of validations applied to the result
// of an importer converting from an external config format to dep's.
type convertTestCase struct {
	wantConvertErr      bool
	projectRoot         gps.ProjectRoot
	wantSourceRepo      string
	wantConstraint      string
	wantRevision        gps.Revision
	wantVersion         string
	wantLockCount       int
	wantIgnoreCount     int
	wantIgnoredPackages []string
}

// validateConvertTestCase returns an error if any of the importer's
// conversion validations failed.
func validateConvertTestCase(testCase *convertTestCase, manifest *dep.Manifest, lock *dep.Lock, convertErr error) error {
	if testCase.wantConvertErr {
		if convertErr == nil {
			return errors.New("Expected the conversion to fail, but it did not return an error")
		}
		return nil
	}

	if convertErr != nil {
		return errors.Wrap(convertErr, "Expected the conversion to pass, but it returned an error")
	}

	// Ignored projects checks.
	if len(manifest.Ignored) != testCase.wantIgnoreCount {
		return errors.Errorf("Expected manifest to have %d ignored project(s), got %d",
			testCase.wantIgnoreCount,
			len(manifest.Ignored))
	}

	if !equalSlice(manifest.Ignored, testCase.wantIgnoredPackages) {
		return errors.Errorf("Expected manifest to have ignore %s, got %s",
			strings.Join(testCase.wantIgnoredPackages, ", "),
			strings.Join(manifest.Ignored, ", "))
	}

	// Constraints checks below.
	if testCase.wantConstraint != "" {
		d, ok := manifest.Constraints[testCase.projectRoot]
		if !ok {
			return errors.Errorf("Expected the manifest to have a dependency for '%s' but got none",
				testCase.projectRoot)
		}

		v := d.Constraint.String()
		if v != testCase.wantConstraint {
			return errors.Errorf("Expected manifest constraint to be %s, got %s", testCase.wantConstraint, v)
		}
	}

	// Lock checks.
	if lock != nil {
		if len(lock.P) != testCase.wantLockCount {
			return errors.Errorf("Expected lock to have %d project(s), got %d",
				testCase.wantLockCount,
				len(lock.P))
		}

		p := lock.P[0]

		if p.Ident().ProjectRoot != testCase.projectRoot {
			return errors.Errorf("Expected the lock to have a project for '%s' but got '%s'",
				testCase.projectRoot,
				p.Ident().ProjectRoot)
		}

		if p.Ident().Source != testCase.wantSourceRepo {
			return errors.Errorf("Expected locked source to be %s, got '%s'", testCase.wantSourceRepo, p.Ident().Source)
		}

		// Break down the locked "version" into a version (optional) and revision
		var gotVersion string
		var gotRevision gps.Revision
		if lpv, ok := p.Version().(gps.PairedVersion); ok {
			gotVersion = lpv.String()
			gotRevision = lpv.Revision()
		} else if lr, ok := p.Version().(gps.Revision); ok {
			gotRevision = lr
		} else {
			return errors.New("could not determine the type of the locked version")
		}

		if gotRevision != testCase.wantRevision {
			return errors.Errorf("unexpected locked revision : \n\t(GOT) %v \n\t(WNT) %v",
				testCase.wantRevision,
				gotRevision)
		}
		if gotVersion != testCase.wantVersion {
			return errors.Errorf("unexpected locked version: \n\t(GOT) %v \n\t(WNT) %v",
				testCase.wantVersion,
				gotVersion)
		}
	}
	return nil
}
<commit_msg>root_analyzer_test: correct message var placement<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"strings"
	"testing"

	"github.com/golang/dep"
	"github.com/golang/dep/internal/gps"
	"github.com/golang/dep/internal/test"
	"github.com/pkg/errors"
)

func TestRootAnalyzer_Info(t *testing.T) {
	testCases := map[bool]string{
		true:  "dep",
		false: "dep+import",
	}
	for skipTools, want := range testCases {
		a := rootAnalyzer{skipTools: skipTools}
		got := a.Info().Name
		if got != want {
			t.Errorf("Expected the name of the importer with skipTools=%t to be '%s', got '%s'", skipTools, want, got)
		}
	}
}

func TestLookupVersionForLockedProject_MatchRevisionToTag(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	rev := gps.Revision("ff2948a2ac8f538c4ecd55962e919d1e13e74baf")
	v, err := lookupVersionForLockedProject(pi, nil, rev, sm)
	h.Must(err)

	wantV := "v1.0.0"
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to be the tag paired with the manifest's pinned revision: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestLookupVersionForLockedProject_MatchRevisionToMultipleTags(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	// Both 0.8.0 and 1.0.0 use the same rev, force dep to pick the lower version
	c, _ := gps.NewSemverConstraint("<1.0.0")
	rev := gps.Revision("ff2948a2ac8f538c4ecd55962e919d1e13e74baf")
	v, err := lookupVersionForLockedProject(pi, c, rev, sm)
	h.Must(err)

	wantV := "v0.8.0"
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to satisfy the manifest's semver constraint: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestLookupVersionForLockedProject_FallbackToConstraint(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	c := gps.NewBranch("master")
	rev := gps.Revision("c575196502940c07bf89fd6d95e83b999162e051")
	v, err := lookupVersionForLockedProject(pi, c, rev, sm)
	h.Must(err)

	wantV := c.String()
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to be defaulted from the manifest's branch constraint: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestLookupVersionForLockedProject_FallbackToRevision(t *testing.T) {
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	rev := gps.Revision("c575196502940c07bf89fd6d95e83b999162e051")
	v, err := lookupVersionForLockedProject(pi, nil, rev, sm)
	h.Must(err)

	wantV := rev.String()
	gotV := v.String()
	if gotV != wantV {
		t.Fatalf("Expected the locked version to be the manifest's pinned revision: wanted '%s', got '%s'", wantV, gotV)
	}
}

func TestProjectExistsInLock(t *testing.T) {
	lock := &dep.Lock{}
	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	ver := gps.NewVersion("v1.0.0")
	lock.P = append(lock.P, gps.NewLockedProject(pi, ver, nil))

	cases := []struct {
		name       string
		importPath string
		want       bool
	}{
		{
			name:       "root project",
			importPath: "github.com/sdboyer/deptest",
			want:       true,
		},
		{
			name:       "sub package",
			importPath: "github.com/sdboyer/deptest/foo",
			want:       false,
		},
		{
			name:       "nonexisting project",
			importPath: "github.com/golang/notexist",
			want:       false,
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			result := projectExistsInLock(lock, gps.ProjectRoot(c.importPath))

			if result != c.want {
				t.Fatalf("projectExistsInLock result is not as want: \n\t(GOT) %v \n\t(WNT) %v", result, c.want)
			}
		})
	}
}

func TestIsVersion(t *testing.T) {
	testcases := map[string]struct {
		wantIsVersion bool
		wantVersion   gps.Version
	}{
		"v1.0.0": {wantIsVersion: true, wantVersion: gps.NewVersion("v1.0.0").Pair("ff2948a2ac8f538c4ecd55962e919d1e13e74baf")},
		"3f4c3bea144e112a69bbe5d8d01c1b09a544253f": {wantIsVersion: false},
		"master": {wantIsVersion: false},
	}

	pi := gps.ProjectIdentifier{ProjectRoot: gps.ProjectRoot("github.com/sdboyer/deptest")}
	h := test.NewHelper(t)
	defer h.Cleanup()

	ctx := newTestContext(h)
	sm, err := ctx.SourceManager()
	h.Must(err)
	defer sm.Release()

	for value, testcase := range testcases {
		t.Run(value, func(t *testing.T) {
			gotIsVersion, gotVersion, err := isVersion(pi, value, sm)
			h.Must(err)

			if testcase.wantIsVersion != gotIsVersion {
				t.Fatalf("unexpected isVersion result for %s: \n\t(GOT) %v \n\t(WNT) %v", value, gotIsVersion, testcase.wantIsVersion)
			}

			if testcase.wantVersion != gotVersion {
				t.Fatalf("unexpected version for %s: \n\t(GOT) %v \n\t(WNT) %v", value, gotVersion, testcase.wantVersion)
			}
		})
	}
}

// convertTestCase is a common set of validations applied to the result
// of an importer converting from an external config format to dep's.
type convertTestCase struct {
	wantConvertErr      bool
	projectRoot         gps.ProjectRoot
	wantSourceRepo      string
	wantConstraint      string
	wantRevision        gps.Revision
	wantVersion         string
	wantLockCount       int
	wantIgnoreCount     int
	wantIgnoredPackages []string
}

// validateConvertTestCase returns an error if any of the importer's
// conversion validations failed.
func validateConvertTestCase(testCase *convertTestCase, manifest *dep.Manifest, lock *dep.Lock, convertErr error) error {
	if testCase.wantConvertErr {
		if convertErr == nil {
			return errors.New("Expected the conversion to fail, but it did not return an error")
		}
		return nil
	}

	if convertErr != nil {
		return errors.Wrap(convertErr, "Expected the conversion to pass, but it returned an error")
	}

	// Ignored projects checks.
	if len(manifest.Ignored) != testCase.wantIgnoreCount {
		return errors.Errorf("Expected manifest to have %d ignored project(s), got %d",
			testCase.wantIgnoreCount,
			len(manifest.Ignored))
	}

	if !equalSlice(manifest.Ignored, testCase.wantIgnoredPackages) {
		return errors.Errorf("Expected manifest to have ignore %s, got %s",
			strings.Join(testCase.wantIgnoredPackages, ", "),
			strings.Join(manifest.Ignored, ", "))
	}

	// Constraints checks below.
	if testCase.wantConstraint != "" {
		d, ok := manifest.Constraints[testCase.projectRoot]
		if !ok {
			return errors.Errorf("Expected the manifest to have a dependency for '%s' but got none",
				testCase.projectRoot)
		}

		v := d.Constraint.String()
		if v != testCase.wantConstraint {
			return errors.Errorf("Expected manifest constraint to be %s, got %s", testCase.wantConstraint, v)
		}
	}

	// Lock checks.
	if lock != nil {
		if len(lock.P) != testCase.wantLockCount {
			return errors.Errorf("Expected lock to have %d project(s), got %d",
				testCase.wantLockCount,
				len(lock.P))
		}

		p := lock.P[0]

		if p.Ident().ProjectRoot != testCase.projectRoot {
			return errors.Errorf("Expected the lock to have a project for '%s' but got '%s'",
				testCase.projectRoot,
				p.Ident().ProjectRoot)
		}

		if p.Ident().Source != testCase.wantSourceRepo {
			return errors.Errorf("Expected locked source to be %s, got '%s'", testCase.wantSourceRepo, p.Ident().Source)
		}

		// Break down the locked "version" into a version (optional) and revision
		var gotVersion string
		var gotRevision gps.Revision
		if lpv, ok := p.Version().(gps.PairedVersion); ok {
			gotVersion = lpv.String()
			gotRevision = lpv.Revision()
		} else if lr, ok := p.Version().(gps.Revision); ok {
			gotRevision = lr
		} else {
			return errors.New("could not determine the type of the locked version")
		}

		if gotRevision != testCase.wantRevision {
			return errors.Errorf("unexpected locked revision : \n\t(GOT) %v \n\t(WNT) %v",
				gotRevision,
				testCase.wantRevision)
		}
		if gotVersion != testCase.wantVersion {
			return errors.Errorf("unexpected locked version: \n\t(GOT) %v \n\t(WNT) %v",
				gotVersion,
				testCase.wantVersion)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package manager

import (
	"errors"
)

var (
	ErrEmptyParam         = errors.New("auth with empty identity")
	ErrPermDenied         = errors.New("permission denied")
	ErrAuthenticationFail = errors.New("authentication fails")
	ErrAuthorizationFial  = errors.New("authorization fails")
)
<commit_msg>a constructive err msg<commit_after>package manager

import (
	"errors"
)

var (
	ErrEmptyParam         = errors.New("auth with empty identity")
	ErrPermDenied         = errors.New("permission denied")
	ErrAuthenticationFail = errors.New("authentication fails, retry after 5m")
	ErrAuthorizationFial  = errors.New("authorization fails, retry after 5m")
)
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"sort"

	"github.com/doneland/yquotes"
	"github.com/howeyc/ledger"
)

func portfolioHandler(w http.ResponseWriter, r *http.Request) {
	t, err := parseAssets("templates/template.portfolio.html", "templates/template.nav.html")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	trans, terr := getTransactions()
	if terr != nil {
		http.Error(w, terr.Error(), 500)
		return
	}
	balances := ledger.GetBalances(trans, []string{})

	var pData pageData
	pData.Reports = reportConfigData.Reports
	pData.Transactions = trans

	sectionTotals := make(map[string]stockInfo)

	for _, stock := range stockConfigData.Stocks {
		quote, _ := yquotes.GetPrice(stock.Ticker)
		var sprice float64
		var sclose float64
		var cprice float64
		if quote != nil {
			sprice = quote.Last
			sclose = quote.PreviousClose
		}
		si := stockInfo{Name: stock.Name,
			Section: stock.Section,
			Ticker:  stock.Ticker,
			Price:   sprice,
			Shares:  stock.Shares}
		for _, bal := range balances {
			if stock.Account == bal.Name {
				si.Cost, _ = bal.Balance.Float64()
			}
		}
		cprice = si.Cost / si.Shares
		si.MarketValue = si.Shares * si.Price
		si.GainLossOverall = si.MarketValue - si.Cost
		si.PriceChangeDay = sprice - sclose
		si.PriceChangePctDay = (si.PriceChangeDay / sclose) * 100.0
		si.PriceChangeOverall = sprice - cprice
		si.PriceChangePctOverall = (si.PriceChangeOverall / cprice) * 100.0
		si.GainLossDay = si.Shares * si.PriceChangeDay
		pData.Stocks = append(pData.Stocks, si)
	}
	stotal := stockInfo{Name: "Total", Section: "Total", Type: "Total"}
	for _, si := range pData.Stocks {
		sectionInfo := sectionTotals[si.Section]
		sectionInfo.Name = si.Section
		sectionInfo.Section = si.Section
		sectionInfo.Type = "Section Total"
		sectionInfo.Ticker = "zzz"
		sectionInfo.Cost += si.Cost
		sectionInfo.MarketValue += si.MarketValue
		sectionInfo.GainLossOverall += si.GainLossOverall
		sectionInfo.GainLossDay += si.GainLossDay
		sectionTotals[si.Section] = sectionInfo

		stotal.Cost += si.Cost
		stotal.MarketValue += si.MarketValue
		stotal.GainLossOverall += si.GainLossOverall
		stotal.GainLossDay += si.GainLossDay
	}
	stotal.PriceChangePctDay = (stotal.GainLossDay / stotal.Cost) * 100.0
	stotal.PriceChangePctOverall = (stotal.GainLossOverall / stotal.Cost) * 100.0
	pData.Stocks = append(pData.Stocks, stotal)

	for _, sectionInfo := range sectionTotals {
		sectionInfo.PriceChangePctDay = (sectionInfo.GainLossDay / sectionInfo.Cost) * 100.0
		sectionInfo.PriceChangePctOverall = (sectionInfo.GainLossOverall / sectionInfo.Cost) * 100.0
		pData.Stocks = append(pData.Stocks, sectionInfo)
	}

	sort.Slice(pData.Stocks, func(i, j int) bool {
		return pData.Stocks[i].Ticker < pData.Stocks[j].Ticker
	})
	sort.SliceStable(pData.Stocks, func(i, j int) bool {
		return pData.Stocks[i].Section < pData.Stocks[j].Section
	})

	err = t.Execute(w, pData)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
}
<commit_msg>get stock info in parallel<commit_after>package main

import (
	"net/http"
	"sort"

	"github.com/doneland/yquotes"
	"github.com/howeyc/ledger"
)

func portfolioHandler(w http.ResponseWriter, r *http.Request) {
	t, err := parseAssets("templates/template.portfolio.html", "templates/template.nav.html")
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	trans, terr := getTransactions()
	if terr != nil {
		http.Error(w, terr.Error(), 500)
		return
	}
	balances := ledger.GetBalances(trans, []string{})

	var pData pageData
	pData.Reports = reportConfigData.Reports
	pData.Transactions = trans

	sectionTotals := make(map[string]stockInfo)
	siChan := make(chan stockInfo)

	for _, stock := range stockConfigData.Stocks {
		go func(name, account, symbol, section string, shares float64) {
			quote, _ := yquotes.GetPrice(symbol)
			var sprice float64
			var sclose float64
			var cprice float64
			if quote != nil {
				sprice = quote.Last
				sclose = quote.PreviousClose
			}
			si := stockInfo{Name: name,
				Section: section,
				Ticker:  symbol,
				Price:   sprice,
				Shares:  shares}
			for _, bal := range balances {
				if account == bal.Name {
					si.Cost, _ = bal.Balance.Float64()
				}
			}
			cprice = si.Cost / si.Shares
			si.MarketValue = si.Shares * si.Price
			si.GainLossOverall = si.MarketValue - si.Cost
			si.PriceChangeDay = sprice - sclose
			si.PriceChangePctDay = (si.PriceChangeDay / sclose) * 100.0
			si.PriceChangeOverall = sprice - cprice
			si.PriceChangePctOverall = (si.PriceChangeOverall / cprice) * 100.0
			si.GainLossDay = si.Shares * si.PriceChangeDay
			siChan <- si
		}(stock.Name, stock.Account, stock.Ticker, stock.Section, stock.Shares)
	}
	for range stockConfigData.Stocks {
		pData.Stocks = append(pData.Stocks, <-siChan)
	}

	stotal := stockInfo{Name: "Total", Section: "Total", Type: "Total"}
	for _, si := range pData.Stocks {
		sectionInfo := sectionTotals[si.Section]
		sectionInfo.Name = si.Section
		sectionInfo.Section = si.Section
		sectionInfo.Type = "Section Total"
		sectionInfo.Ticker = "zzz"
		sectionInfo.Cost += si.Cost
		sectionInfo.MarketValue += si.MarketValue
		sectionInfo.GainLossOverall += si.GainLossOverall
		sectionInfo.GainLossDay += si.GainLossDay
		sectionTotals[si.Section] = sectionInfo

		stotal.Cost += si.Cost
		stotal.MarketValue += si.MarketValue
		stotal.GainLossOverall += si.GainLossOverall
		stotal.GainLossDay += si.GainLossDay
	}
	stotal.PriceChangePctDay = (stotal.GainLossDay / stotal.Cost) * 100.0
	stotal.PriceChangePctOverall = (stotal.GainLossOverall / stotal.Cost) * 100.0
	pData.Stocks = append(pData.Stocks, stotal)

	for _, sectionInfo := range sectionTotals {
		sectionInfo.PriceChangePctDay = (sectionInfo.GainLossDay / sectionInfo.Cost) * 100.0
		sectionInfo.PriceChangePctOverall = (sectionInfo.GainLossOverall / sectionInfo.Cost) * 100.0
		pData.Stocks = append(pData.Stocks, sectionInfo)
	}

	sort.Slice(pData.Stocks, func(i, j int) bool {
		return pData.Stocks[i].Ticker < pData.Stocks[j].Ticker
	})
	sort.SliceStable(pData.Stocks, func(i, j int) bool {
		return pData.Stocks[i].Section < pData.Stocks[j].Section
	})

	err = t.Execute(w, pData)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
}
<|endoftext|>


<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"database/sql/driver"
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	var data []byte

	if strings.HasPrefix(name, "Reader::") { // io.Reader
		name = name[8:]
		if handler, inMap := readerRegister[name]; inMap {
			rdr = handler()
			if rdr != nil {
				data = make([]byte, 4+mc.maxWriteSize)
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		} else {
			err = fmt.Errorf("Reader '%s' is not registered", name)
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			var file *os.File
			var fi os.FileInfo

			if file, err = os.Open(name); err == nil {
				if fi, err = file.Stat(); err == nil {
					rdr = file
					if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize {
						data = make([]byte, 4+fileSize)
					} else if fileSize <= mc.maxPacketAllowed {
						data = make([]byte, 4+mc.maxWriteSize)
					} else {
						err = fmt.Errorf("Local File '%s' too large: Size: %d, Max: %d", name, fileSize, mc.maxPacketAllowed)
					}
				}
			}
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	if rdc, ok := rdr.(io.ReadCloser); ok {
		defer func() {
			if err == nil {
				err = rdc.Close()
			} else {
				rdc.Close()
			}
		}()
	}

	// send content packets
	var ioErr error
	if err == nil {
		var n int
		for err == nil && ioErr == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				ioErr = mc.writePacket(data[:4+n])
			}
		}
		if err == io.EOF {
			err = nil
		}
		if ioErr != nil {
			errLog.Print(ioErr.Error())
			return driver.ErrBadConn
		}
	}

	// send empty packet (termination)
	ioErr = mc.writePacket([]byte{
		0x00,
		0x00,
		0x00,
		mc.sequence,
	})
	if ioErr != nil {
		errLog.Print(ioErr.Error())
		return driver.ErrBadConn
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<commit_msg>infile: refactoring<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	var data []byte

	if strings.HasPrefix(name, "Reader::") { // io.Reader
		name = name[8:]
		if handler, inMap := readerRegister[name]; inMap {
			rdr = handler()
			if rdr != nil {
				data = make([]byte, 4+mc.maxWriteSize)

				if rdc, ok := rdr.(io.ReadCloser); ok {
					defer func() {
						if err == nil {
							err = rdc.Close()
						} else {
							rdc.Close()
						}
					}()
				}
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		} else {
			err = fmt.Errorf("Reader '%s' is not registered", name)
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			var file *os.File
			var fi os.FileInfo

			if file, err = os.Open(name); err == nil {
				defer func() {
					if err == nil {
						err = file.Close()
					} else {
						file.Close()
					}
				}()

				// get file size
				if fi, err = file.Stat(); err == nil {
					rdr = file
					if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize {
						data = make([]byte, 4+fileSize)
					} else if fileSize <= mc.maxPacketAllowed {
						data = make([]byte, 4+mc.maxWriteSize)
					} else {
						err = fmt.Errorf("Local File '%s' too large: Size: %d, Max: %d", name, fileSize, mc.maxPacketAllowed)
					}
				}
			}
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	// send content packets
	if err == nil {
		var n int
		for err == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				if ioErr := mc.writePacket(data[:4+n]); ioErr != nil {
					return ioErr
				}
			}
		}
		if err == io.EOF {
			err = nil
		}
	}

	// send empty packet (termination)
	if data == nil {
		data = make([]byte, 4)
	}
	if ioErr := mc.writePacket(data[:4]); ioErr != nil {
		return ioErr
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<|endoftext|>


<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"database/sql/driver"
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	data := make([]byte, 4+mc.maxWriteSize)

	if strings.HasPrefix(name, "Reader::") { // io.Reader
		name = name[8:]
		handler, inMap := readerRegister[name]
		if handler != nil {
			rdr = handler()
		}
		if rdr == nil {
			if !inMap {
				err = fmt.Errorf("Reader '%s' is not registered", name)
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			rdr, err = os.Open(name)
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	if rdc, ok := rdr.(io.ReadCloser); ok {
		defer func() {
			if err == nil {
				err = rdc.Close()
			} else {
				rdc.Close()
			}
		}()
	}

	// send content packets
	var ioErr error
	if err == nil {
		var n int
		for err == nil && ioErr == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				ioErr = mc.writePacket(data[:4+n])
			}
		}
		if err == io.EOF {
			err = nil
		}
		if ioErr != nil {
			errLog.Print(ioErr.Error())
			return driver.ErrBadConn
		}
	}

	// send empty packet (termination)
	ioErr = mc.writePacket([]byte{
		0x00,
		0x00,
		0x00,
		mc.sequence,
	})
	if ioErr != nil {
		errLog.Print(ioErr.Error())
		return driver.ErrBadConn
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<commit_msg>infile: better buffer size calculation<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"database/sql/driver"
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	var data []byte

	if strings.HasPrefix(name, "Reader::") { // io.Reader
		name = name[8:]
		if handler, inMap := readerRegister[name]; inMap {
			rdr = handler()
			if rdr != nil {
				data = make([]byte, 4+mc.maxWriteSize)
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		} else {
			err = fmt.Errorf("Reader '%s' is not registered", name)
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			var file *os.File
			var fi os.FileInfo

			if file, err = os.Open(name); err == nil {
				if fi, err = file.Stat(); err == nil {
					rdr = file
					if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize {
						data = make([]byte, 4+fileSize)
					} else if fileSize <= mc.maxPacketAllowed {
						data = make([]byte, 4+mc.maxWriteSize)
					} else {
						err = fmt.Errorf("Local File '%s' too large: Size: %d, Max: %d", name, fileSize, mc.maxPacketAllowed)
					}
				}
			}
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	if rdc, ok := rdr.(io.ReadCloser); ok {
		defer func() {
			if err == nil {
				err = rdc.Close()
			} else {
				rdc.Close()
			}
		}()
	}

	// send content packets
	var ioErr error
	if err == nil {
		var n int
		for err == nil && ioErr == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				ioErr = mc.writePacket(data[:4+n])
			}
		}
		if err == io.EOF {
			err = nil
		}
		if ioErr != nil {
			errLog.Print(ioErr.Error())
			return driver.ErrBadConn
		}
	}

	// send empty packet (termination)
	ioErr = mc.writePacket([]byte{
		0x00,
		0x00,
		0x00,
		mc.sequence,
	})
	if ioErr != nil {
		errLog.Print(ioErr.Error())
		return driver.ErrBadConn
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<|endoftext|>