<commit_before>package axfrddns

/*

axfrddns -
  Fetch the zone with an AXFR request (RFC5936) to a given primary master, and
  push Dynamic DNS updates (RFC2136) to the same server.

  Both the AXFR request and the updates might be authentificated with
  a TSIG.

*/

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math"
	"math/rand"
	"strings"
	"time"

	"github.com/miekg/dns"

	"github.com/StackExchange/dnscontrol/v3/models"
	"github.com/StackExchange/dnscontrol/v3/pkg/diff"
	"github.com/StackExchange/dnscontrol/v3/providers"
)

const (
	dnsTimeout       = 30 * time.Second
	dnssecDummyLabel = "__dnssec"
	dnssecDummyTxt   = "Domain has DNSSec records, not displayed here."
)

var features = providers.DocumentationNotes{
	providers.CanUseCAA:              providers.Can(),
	providers.CanUsePTR:              providers.Can(),
	providers.CanUseNAPTR:            providers.Can(),
	providers.CanUseSRV:              providers.Can(),
	providers.CanUseSSHFP:            providers.Can(),
	providers.CanUseTLSA:             providers.Can(),
	providers.CanUseTXTMulti:         providers.Can(),
	providers.CanAutoDNSSEC:          providers.Can("Just warn when DNSSEC is requested but no RRSIG is found in the AXFR or warn when DNSSEC is not requested but RRSIG are found in the AXFR."),
	providers.CantUseNOPURGE:         providers.Cannot(),
	providers.DocCreateDomains:       providers.Cannot(),
	providers.DocDualHost:            providers.Cannot(),
	providers.DocOfficiallySupported: providers.Cannot(),
	providers.CanGetZones:            providers.Can(),
}

// AxfrDdns stores the client info for the provider.
type AxfrDdns struct {
	rand        *rand.Rand
	master      string
	nameservers []*models.Nameserver
	transferKey *Key
	updateKey   *Key
}

func initAxfrDdns(config map[string]string, providermeta json.RawMessage) (providers.DNSServiceProvider, error) {
	// config -- the key/values from creds.json
	// providermeta -- the json blob from NewReq('name', 'TYPE', providermeta)
	var err error
	api := &AxfrDdns{
		rand: rand.New(rand.NewSource(int64(time.Now().Nanosecond()))),
	}
	param := &Param{}
	if len(providermeta) != 0 {
		err := json.Unmarshal(providermeta, param)
		if err != nil {
			return nil, err
		}
	}
	var nss []string
	if config["nameservers"] != "" {
		nss = strings.Split(config["nameservers"], ",")
	}
	for _, ns := range param.DefaultNS {
		nss = append(nss, ns[0:len(ns)-1])
	}
	api.nameservers, err = models.ToNameservers(nss)
	if err != nil {
		return nil, err
	}
	if config["master"] != "" {
		api.master = config["master"]
		if !strings.Contains(api.master, ":") {
			api.master = api.master + ":53"
		}
	} else if len(api.nameservers) != 0 {
		api.master = api.nameservers[0].Name + ":53"
	} else {
		return nil, fmt.Errorf("nameservers list is empty: creds.json needs a default `nameservers` or an explicit `master`")
	}
	api.updateKey, err = readKey(config["update-key"], "update-key")
	if err != nil {
		return nil, err
	}
	api.transferKey, err = readKey(config["transfer-key"], "transfer-key")
	if err != nil {
		return nil, err
	}
	for key := range config {
		switch key {
		case "master",
			"nameservers",
			"update-key",
			"transfer-key":
			continue
		default:
			fmt.Printf("[Warning] AXFRDDNS: unknown key in `creds.json` (%s)\n", key)
		}
	}
	return api, err
}

func init() {
	providers.RegisterDomainServiceProviderType("AXFRDDNS", initAxfrDdns, features)
}

// Param is used to decode extra parameters sent to provider.
type Param struct {
	DefaultNS []string `json:"default_ns"`
}

// Key stores the individual parts of a TSIG key.
type Key struct {
	algo   string
	id     string
	secret string
}

func readKey(raw string, kind string) (*Key, error) {
	if raw == "" {
		return nil, nil
	}
	arr := strings.Split(raw, ":")
	if len(arr) != 3 {
		return nil, fmt.Errorf("invalid key format (%s) in AXFRDDNS.TSIG", kind)
	}
	var algo string
	switch arr[0] {
	case "hmac-md5", "md5":
		algo = dns.HmacMD5
	case "hmac-sha1", "sha1":
		algo = dns.HmacSHA1
	case "hmac-sha256", "sha256":
		algo = dns.HmacSHA256
	case "hmac-sha512", "sha512":
		algo = dns.HmacSHA512
	default:
		return nil, fmt.Errorf("unknown algorithm (%s) in AXFRDDNS.TSIG", kind)
	}
	_, err := base64.StdEncoding.DecodeString(arr[2])
	if err != nil {
		return nil, fmt.Errorf("cannot decode Base64 secret (%s) in AXFRDDNS.TSIG", kind)
	}
	return &Key{algo: algo, id: arr[1] + ".", secret: arr[2]}, nil
}

// GetNameservers returns the nameservers for a domain.
func (c *AxfrDdns) GetNameservers(domain string) ([]*models.Nameserver, error) {
	return c.nameservers, nil
}

// FetchZoneRecords gets the records of a zone and returns them in dns.RR format.
func (c *AxfrDdns) FetchZoneRecords(domain string) ([]dns.RR, error) {

	transfer := new(dns.Transfer)
	transfer.DialTimeout = dnsTimeout
	transfer.ReadTimeout = dnsTimeout

	request := new(dns.Msg)
	request.SetAxfr(domain + ".")

	if c.transferKey != nil {
		transfer.TsigSecret =
			map[string]string{c.transferKey.id: c.transferKey.secret}
		request.SetTsig(c.transferKey.id, c.transferKey.algo, 300, time.Now().Unix())
	}

	envelope, err := transfer.In(request, c.master)
	if err != nil {
		return nil, err
	}

	var rawRecords []dns.RR
	for msg := range envelope {
		if msg.Error != nil {
			// Fragile but more "user-friendly" error-handling
			err := msg.Error.Error()
			if err == "dns: bad xfr rcode: 9" {
				err = "NOT AUTH (9)"
			}
			return nil, fmt.Errorf("[Error] AXFRDDNS: nameserver refused to transfer the zone: %s", msg)
		}
		rawRecords = append(rawRecords, msg.RR...)
	}
	return rawRecords, nil

}

// GetZoneRecords gets the records of a zone and returns them in RecordConfig format.
func (c *AxfrDdns) GetZoneRecords(domain string) (models.Records, error) {

	rawRecords, err := c.FetchZoneRecords(domain)
	if err != nil {
		return nil, err
	}

	var foundDNSSecRecords *models.RecordConfig
	foundRecords := models.Records{}
	for _, rr := range rawRecords {
		switch rr.(type) {
		case *dns.RRSIG,
			*dns.DNSKEY,
			*dns.CDNSKEY,
			*dns.CDS,
			*dns.NSEC,
			*dns.NSEC3,
			*dns.NSEC3PARAM:
			// Ignoring DNSSec RRs, but replacing it with a single
			// "TXT" placeholder
			if foundDNSSecRecords == nil {
				foundDNSSecRecords = new(models.RecordConfig)
				foundDNSSecRecords.Type = "TXT"
				foundDNSSecRecords.SetLabel(dnssecDummyLabel, domain)
				err = foundDNSSecRecords.SetTargetTXT(dnssecDummyTxt)
				if err != nil {
					return nil, err
				}
			}
			continue
		default:
			rec := models.RRtoRC(rr, domain)
			foundRecords = append(foundRecords, &rec)
		}
	}

	if len(foundRecords) >= 1 && foundRecords[len(foundRecords)-1].Type == "SOA" {
		// The SOA is sent two times: as the first and the last record
		// See section 2.2 of RFC5936
		foundRecords = foundRecords[:len(foundRecords)-1]
	}

	if foundDNSSecRecords != nil {
		foundRecords = append(foundRecords, foundDNSSecRecords)
	}

	return foundRecords, nil

}

// GetDomainCorrections returns a list of corrections to update a domain.
func (c *AxfrDdns) GetDomainCorrections(dc *models.DomainConfig) ([]*models.Correction, error) {
	dc.Punycode()

	foundRecords, err := c.GetZoneRecords(dc.Name)
	if err != nil {
		return nil, err
	}

	if len(foundRecords) >= 1 && foundRecords[0].Type == "SOA" {
		// Ignoring the SOA, others providers  don't manage it either.
		foundRecords = foundRecords[1:]
	}

	hasDnssecRecords := false
	if len(foundRecords) >= 1 {
		last := foundRecords[len(foundRecords)-1]
		if last.Type == "TXT" &&
			last.Name == dnssecDummyLabel &&
			len(last.TxtStrings) == 1 &&
			last.TxtStrings[0] == dnssecDummyTxt {
			hasDnssecRecords = true
			foundRecords = foundRecords[0:(len(foundRecords) - 1)]
		}
	}

	if dc.AutoDNSSEC && !hasDnssecRecords {
		fmt.Printf("Warning: AUTODNSSEC is set, but no DNSKEY or RRSIG record was found in the AXFR answer!\n")
	} else if !dc.AutoDNSSEC && hasDnssecRecords {
		fmt.Printf("Warning: AUTODNSSEC is not set, but DNSKEY or RRSIG records were found in the AXFR answer!\n")
	}

	// Normalize
	models.PostProcessRecords(foundRecords)

	differ := diff.New(dc)
	_, create, del, mod, err := differ.IncrementalDiff(foundRecords)
	if err != nil {
		return nil, err
	}

	buf := &bytes.Buffer{}
	// Print a list of changes. Generate an actual change that is the zone
	changes := false
	for _, i := range create {
		changes = true
		fmt.Fprintln(buf, i)
	}
	for _, i := range del {
		changes = true
		fmt.Fprintln(buf, i)
	}
	for _, i := range mod {
		changes = true
		fmt.Fprintln(buf, i)
	}
	msg := fmt.Sprintf("DDNS UPDATES to '%s' (primary master: '%s'). Changes:\n%s", dc.Name, c.master, buf)

	corrections := []*models.Correction{}
	if changes {

		corrections = append(corrections,
			&models.Correction{
				Msg: msg,
				F: func() error {

					// An RFC2136-compliant server must silently ignore an
					// update that inserts a non-CNAME RRset when a CNAME RR
					// with the same name is present in the zone (and
					// vice-versa). Therefore we prefer to first remove records
					// and then insert new ones.
					//
					// Compliant servers must also silently ignore an update
					// that removes the last NS record of a zone. Therefore we
					// don't want to remove all NS records before inserting a
					// new one. For the particular case of NS record, we prefer
					// to insert new records before ot remove old ones.
					//
					// This remarks does not apply for "modified" NS records, as
					// updates are processed one-by-one.
					//
					// This provider does not allow modifying the TTL of an NS
					// record in a zone that defines only one NS. That would
					// would require removing the single NS record, before
					// adding the new one. But who does that anyway?

					update := new(dns.Msg)
					update.SetUpdate(dc.Name + ".")
					update.Id = uint16(c.rand.Intn(math.MaxUint16))
					for _, c := range create {
						if c.Desired.Type == "NS" {
							update.Insert([]dns.RR{c.Desired.ToRR()})
						}
					}
					for _, c := range del {
						update.Remove([]dns.RR{c.Existing.ToRR()})
					}
					for _, c := range mod {
						update.Remove([]dns.RR{c.Existing.ToRR()})
						update.Insert([]dns.RR{c.Desired.ToRR()})
					}
					for _, c := range create {
						if c.Desired.Type != "NS" {
							update.Insert([]dns.RR{c.Desired.ToRR()})
						}
					}

					client := new(dns.Client)
					client.Timeout = dnsTimeout
					if c.updateKey != nil {
						client.TsigSecret =
							map[string]string{c.updateKey.id: c.updateKey.secret}
						update.SetTsig(c.updateKey.id, c.updateKey.algo, 300, time.Now().Unix())
					}

					msg, _, err := client.Exchange(update, c.master)
					if err != nil {
						return err
					}
					if msg.MsgHdr.Rcode != 0 {
						return fmt.Errorf("[Error] AXFRDDNS: nameserver refused to update the zone: %s (%d)",
							dns.RcodeToString[msg.MsgHdr.Rcode],
							msg.MsgHdr.Rcode)
					}

					return nil
				},
			})
	}
	return corrections, nil
}
<commit_msg>AXFRDDNS: fix error message (fix #778) (#871)<commit_after>package axfrddns

/*

axfrddns -
  Fetch the zone with an AXFR request (RFC5936) to a given primary master, and
  push Dynamic DNS updates (RFC2136) to the same server.

  Both the AXFR request and the updates might be authentificated with
  a TSIG.

*/

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math"
	"math/rand"
	"strings"
	"time"

	"github.com/miekg/dns"

	"github.com/StackExchange/dnscontrol/v3/models"
	"github.com/StackExchange/dnscontrol/v3/pkg/diff"
	"github.com/StackExchange/dnscontrol/v3/providers"
)

const (
	dnsTimeout       = 30 * time.Second
	dnssecDummyLabel = "__dnssec"
	dnssecDummyTxt   = "Domain has DNSSec records, not displayed here."
)

var features = providers.DocumentationNotes{
	providers.CanUseCAA:              providers.Can(),
	providers.CanUsePTR:              providers.Can(),
	providers.CanUseNAPTR:            providers.Can(),
	providers.CanUseSRV:              providers.Can(),
	providers.CanUseSSHFP:            providers.Can(),
	providers.CanUseTLSA:             providers.Can(),
	providers.CanUseTXTMulti:         providers.Can(),
	providers.CanAutoDNSSEC:          providers.Can("Just warn when DNSSEC is requested but no RRSIG is found in the AXFR or warn when DNSSEC is not requested but RRSIG are found in the AXFR."),
	providers.CantUseNOPURGE:         providers.Cannot(),
	providers.DocCreateDomains:       providers.Cannot(),
	providers.DocDualHost:            providers.Cannot(),
	providers.DocOfficiallySupported: providers.Cannot(),
	providers.CanGetZones:            providers.Can(),
}

// AxfrDdns stores the client info for the provider.
type AxfrDdns struct {
	rand        *rand.Rand
	master      string
	nameservers []*models.Nameserver
	transferKey *Key
	updateKey   *Key
}

func initAxfrDdns(config map[string]string, providermeta json.RawMessage) (providers.DNSServiceProvider, error) {
	// config -- the key/values from creds.json
	// providermeta -- the json blob from NewReq('name', 'TYPE', providermeta)
	var err error
	api := &AxfrDdns{
		rand: rand.New(rand.NewSource(int64(time.Now().Nanosecond()))),
	}
	param := &Param{}
	if len(providermeta) != 0 {
		err := json.Unmarshal(providermeta, param)
		if err != nil {
			return nil, err
		}
	}
	var nss []string
	if config["nameservers"] != "" {
		nss = strings.Split(config["nameservers"], ",")
	}
	for _, ns := range param.DefaultNS {
		nss = append(nss, ns[0:len(ns)-1])
	}
	api.nameservers, err = models.ToNameservers(nss)
	if err != nil {
		return nil, err
	}
	if config["master"] != "" {
		api.master = config["master"]
		if !strings.Contains(api.master, ":") {
			api.master = api.master + ":53"
		}
	} else if len(api.nameservers) != 0 {
		api.master = api.nameservers[0].Name + ":53"
	} else {
		return nil, fmt.Errorf("nameservers list is empty: creds.json needs a default `nameservers` or an explicit `master`")
	}
	api.updateKey, err = readKey(config["update-key"], "update-key")
	if err != nil {
		return nil, err
	}
	api.transferKey, err = readKey(config["transfer-key"], "transfer-key")
	if err != nil {
		return nil, err
	}
	for key := range config {
		switch key {
		case "master",
			"nameservers",
			"update-key",
			"transfer-key":
			continue
		default:
			fmt.Printf("[Warning] AXFRDDNS: unknown key in `creds.json` (%s)\n", key)
		}
	}
	return api, err
}

func init() {
	providers.RegisterDomainServiceProviderType("AXFRDDNS", initAxfrDdns, features)
}

// Param is used to decode extra parameters sent to provider.
type Param struct {
	DefaultNS []string `json:"default_ns"`
}

// Key stores the individual parts of a TSIG key.
type Key struct {
	algo   string
	id     string
	secret string
}

func readKey(raw string, kind string) (*Key, error) {
	if raw == "" {
		return nil, nil
	}
	arr := strings.Split(raw, ":")
	if len(arr) != 3 {
		return nil, fmt.Errorf("invalid key format (%s) in AXFRDDNS.TSIG", kind)
	}
	var algo string
	switch arr[0] {
	case "hmac-md5", "md5":
		algo = dns.HmacMD5
	case "hmac-sha1", "sha1":
		algo = dns.HmacSHA1
	case "hmac-sha256", "sha256":
		algo = dns.HmacSHA256
	case "hmac-sha512", "sha512":
		algo = dns.HmacSHA512
	default:
		return nil, fmt.Errorf("unknown algorithm (%s) in AXFRDDNS.TSIG", kind)
	}
	_, err := base64.StdEncoding.DecodeString(arr[2])
	if err != nil {
		return nil, fmt.Errorf("cannot decode Base64 secret (%s) in AXFRDDNS.TSIG", kind)
	}
	return &Key{algo: algo, id: arr[1] + ".", secret: arr[2]}, nil
}

// GetNameservers returns the nameservers for a domain.
func (c *AxfrDdns) GetNameservers(domain string) ([]*models.Nameserver, error) {
	return c.nameservers, nil
}

// FetchZoneRecords gets the records of a zone and returns them in dns.RR format.
func (c *AxfrDdns) FetchZoneRecords(domain string) ([]dns.RR, error) {

	transfer := new(dns.Transfer)
	transfer.DialTimeout = dnsTimeout
	transfer.ReadTimeout = dnsTimeout

	request := new(dns.Msg)
	request.SetAxfr(domain + ".")

	if c.transferKey != nil {
		transfer.TsigSecret =
			map[string]string{c.transferKey.id: c.transferKey.secret}
		request.SetTsig(c.transferKey.id, c.transferKey.algo, 300, time.Now().Unix())
	}

	envelope, err := transfer.In(request, c.master)
	if err != nil {
		return nil, err
	}

	var rawRecords []dns.RR
	for msg := range envelope {
		if msg.Error != nil {
			// Fragile but more "user-friendly" error-handling
			err := msg.Error.Error()
			if err == "dns: bad xfr rcode: 9" {
				err = "NOT AUTH (9)"
			}
			return nil, fmt.Errorf("[Error] AXFRDDNS: nameserver refused to transfer the zone: %s", err)
		}
		rawRecords = append(rawRecords, msg.RR...)
	}
	return rawRecords, nil

}

// GetZoneRecords gets the records of a zone and returns them in RecordConfig format.
func (c *AxfrDdns) GetZoneRecords(domain string) (models.Records, error) {

	rawRecords, err := c.FetchZoneRecords(domain)
	if err != nil {
		return nil, err
	}

	var foundDNSSecRecords *models.RecordConfig
	foundRecords := models.Records{}
	for _, rr := range rawRecords {
		switch rr.(type) {
		case *dns.RRSIG,
			*dns.DNSKEY,
			*dns.CDNSKEY,
			*dns.CDS,
			*dns.NSEC,
			*dns.NSEC3,
			*dns.NSEC3PARAM:
			// Ignoring DNSSec RRs, but replacing it with a single
			// "TXT" placeholder
			if foundDNSSecRecords == nil {
				foundDNSSecRecords = new(models.RecordConfig)
				foundDNSSecRecords.Type = "TXT"
				foundDNSSecRecords.SetLabel(dnssecDummyLabel, domain)
				err = foundDNSSecRecords.SetTargetTXT(dnssecDummyTxt)
				if err != nil {
					return nil, err
				}
			}
			continue
		default:
			rec := models.RRtoRC(rr, domain)
			foundRecords = append(foundRecords, &rec)
		}
	}

	if len(foundRecords) >= 1 && foundRecords[len(foundRecords)-1].Type == "SOA" {
		// The SOA is sent two times: as the first and the last record
		// See section 2.2 of RFC5936
		foundRecords = foundRecords[:len(foundRecords)-1]
	}

	if foundDNSSecRecords != nil {
		foundRecords = append(foundRecords, foundDNSSecRecords)
	}

	return foundRecords, nil

}

// GetDomainCorrections returns a list of corrections to update a domain.
func (c *AxfrDdns) GetDomainCorrections(dc *models.DomainConfig) ([]*models.Correction, error) {
	dc.Punycode()

	foundRecords, err := c.GetZoneRecords(dc.Name)
	if err != nil {
		return nil, err
	}

	if len(foundRecords) >= 1 && foundRecords[0].Type == "SOA" {
		// Ignoring the SOA, others providers  don't manage it either.
		foundRecords = foundRecords[1:]
	}

	hasDnssecRecords := false
	if len(foundRecords) >= 1 {
		last := foundRecords[len(foundRecords)-1]
		if last.Type == "TXT" &&
			last.Name == dnssecDummyLabel &&
			len(last.TxtStrings) == 1 &&
			last.TxtStrings[0] == dnssecDummyTxt {
			hasDnssecRecords = true
			foundRecords = foundRecords[0:(len(foundRecords) - 1)]
		}
	}

	if dc.AutoDNSSEC && !hasDnssecRecords {
		fmt.Printf("Warning: AUTODNSSEC is set, but no DNSKEY or RRSIG record was found in the AXFR answer!\n")
	} else if !dc.AutoDNSSEC && hasDnssecRecords {
		fmt.Printf("Warning: AUTODNSSEC is not set, but DNSKEY or RRSIG records were found in the AXFR answer!\n")
	}

	// Normalize
	models.PostProcessRecords(foundRecords)

	differ := diff.New(dc)
	_, create, del, mod, err := differ.IncrementalDiff(foundRecords)
	if err != nil {
		return nil, err
	}

	buf := &bytes.Buffer{}
	// Print a list of changes. Generate an actual change that is the zone
	changes := false
	for _, i := range create {
		changes = true
		fmt.Fprintln(buf, i)
	}
	for _, i := range del {
		changes = true
		fmt.Fprintln(buf, i)
	}
	for _, i := range mod {
		changes = true
		fmt.Fprintln(buf, i)
	}
	msg := fmt.Sprintf("DDNS UPDATES to '%s' (primary master: '%s'). Changes:\n%s", dc.Name, c.master, buf)

	corrections := []*models.Correction{}
	if changes {

		corrections = append(corrections,
			&models.Correction{
				Msg: msg,
				F: func() error {

					// An RFC2136-compliant server must silently ignore an
					// update that inserts a non-CNAME RRset when a CNAME RR
					// with the same name is present in the zone (and
					// vice-versa). Therefore we prefer to first remove records
					// and then insert new ones.
					//
					// Compliant servers must also silently ignore an update
					// that removes the last NS record of a zone. Therefore we
					// don't want to remove all NS records before inserting a
					// new one. For the particular case of NS record, we prefer
					// to insert new records before ot remove old ones.
					//
					// This remarks does not apply for "modified" NS records, as
					// updates are processed one-by-one.
					//
					// This provider does not allow modifying the TTL of an NS
					// record in a zone that defines only one NS. That would
					// would require removing the single NS record, before
					// adding the new one. But who does that anyway?

					update := new(dns.Msg)
					update.SetUpdate(dc.Name + ".")
					update.Id = uint16(c.rand.Intn(math.MaxUint16))
					for _, c := range create {
						if c.Desired.Type == "NS" {
							update.Insert([]dns.RR{c.Desired.ToRR()})
						}
					}
					for _, c := range del {
						update.Remove([]dns.RR{c.Existing.ToRR()})
					}
					for _, c := range mod {
						update.Remove([]dns.RR{c.Existing.ToRR()})
						update.Insert([]dns.RR{c.Desired.ToRR()})
					}
					for _, c := range create {
						if c.Desired.Type != "NS" {
							update.Insert([]dns.RR{c.Desired.ToRR()})
						}
					}

					client := new(dns.Client)
					client.Timeout = dnsTimeout
					if c.updateKey != nil {
						client.TsigSecret =
							map[string]string{c.updateKey.id: c.updateKey.secret}
						update.SetTsig(c.updateKey.id, c.updateKey.algo, 300, time.Now().Unix())
					}

					msg, _, err := client.Exchange(update, c.master)
					if err != nil {
						return err
					}
					if msg.MsgHdr.Rcode != 0 {
						return fmt.Errorf("[Error] AXFRDDNS: nameserver refused to update the zone: %s (%d)",
							dns.RcodeToString[msg.MsgHdr.Rcode],
							msg.MsgHdr.Rcode)
					}

					return nil
				},
			})
	}
	return corrections, nil
}
<|endoftext|>


<commit_before>package teams

import (
	"testing"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
)

func memberSetup(t *testing.T) (libkb.TestContext, *kbtest.FakeUser, string) {
	tc := libkb.SetupTest(t, "team", 1)
	tc.Tp.UpgradePerUserKey = true

	u, err := kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}

	name := createTeam(tc)

	return tc, u, name
}

func TestMemberOwner(t *testing.T) {
	tc, u, name := memberSetup(t)
	defer tc.Cleanup()

	s, err := Get(context.TODO(), tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	upak, _, err := tc.G.GetUPAKLoader().Load(libkb.NewLoadUserByUIDArg(context.Background(), tc.G, u.User.GetUID()))
	if err != nil {
		t.Fatal(err)
	}
	uv := NewUserVersion(upak.Base.Username, upak.Base.EldestSeqno)
	t.Logf("user eldest seqno: %d", upak.Base.EldestSeqno)

	role, err := s.GetUserRole(uv)
	if err != nil {
		t.Fatal(err)
	}
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}
}

func TestMemberAdd(t *testing.T) {
	tc, u, name := memberSetup(t)
	defer tc.Cleanup()

	_ = u
	_ = name
}
<commit_msg>Create UserVersion better in test<commit_after>package teams

import (
	"testing"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
)

func memberSetup(t *testing.T) (libkb.TestContext, *kbtest.FakeUser, string) {
	tc := libkb.SetupTest(t, "team", 1)
	tc.Tp.UpgradePerUserKey = true

	u, err := kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}

	name := createTeam(tc)

	return tc, u, name
}

func TestMemberOwner(t *testing.T) {
	tc, u, name := memberSetup(t)
	defer tc.Cleanup()

	s, err := Get(context.TODO(), tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(tc, s, u.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}
}

func TestMemberAdd(t *testing.T) {
	tc, u, name := memberSetup(t)
	defer tc.Cleanup()

	_ = u
	_ = name
}

func uidRole(tc libkb.TestContext, state *TeamSigChainState, uid keybase1.UID) keybase1.TeamRole {
	upak, _, err := tc.G.GetUPAKLoader().Load(libkb.NewLoadUserByUIDArg(context.Background(), tc.G, uid))
	if err != nil {
		tc.T.Fatal(err)
	}

	uv := NewUserVersion(upak.Base.Username, upak.Base.EldestSeqno)

	role, err := state.GetUserRole(uv)
	if err != nil {
		tc.T.Fatal(err)
	}
	return role
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/flate"
	"encoding/binary"
	"fmt"
	"os"
)

type PDFormat struct {
	Name             [32]byte //start at byte 0
	Attributes       uint16   //32
	Version          uint16   //34
	CreationDate     uint32   //36
	ModifyDate       uint32   //40
	BackupDate       uint32   //44
	ModifyNumber     uint32   //48
	AppInfoID        uint32   //52
	SortInfoID       uint32   //56
	Type             [4]byte  //60
	Creator          [4]byte  //64
	UniqueIDSeed     uint32   //68
	NextRecordListID uint32   //72
	SectionCount     uint16   //76-78
}

type PDRecordInfoSection struct {
	DataOffset uint32  //starts at byte 0
	Attributes byte    //4
	UniqueID   [3]byte //5-8
}

type PDHeader struct {
	CompressionType uint16 //starts at byte 0
	_               uint16 //2 (always zero?)
	TextLength      uint32 //4
	RecordCount     uint16 //8
	RecordSize      uint16 //10
	CurrentPosition uint32 //12-16
}

type Mobi8Header struct {
	//note that since the values are being read
	//by bytes.Read, we can't use "_" as a name,
	//or have any unexported fields.
	//This is stupid, but whatever.
	//We use the name "SkipX" instead.
	CompressionType     uint16   //start at byte 0
	Skip1               uint16   //2 (always zero?)
	TextLength          uint32   //4
	RecordCount         uint16   //8
	RecordSize          uint16   //10
	CryptoType          uint16   //12
	Skip2               uint16   //14 filler
	Identifier          [4]byte  //16
	HeaderLength        uint32   //20
	Type                uint32   //24
	TextEncoding        uint32   //28
	UniqueID            uint32   //32
	Version             uint32   //36
	OrtographicIndex    uint32   //40
	IncflectionIndex    uint32   //44
	IndexNames          uint32   //48
	IndexKeys           uint32   //52
	Extra               [24]byte //56
	FirstNontext        uint32   //80
	TitleOffset         uint32   //84
	TitleLength         uint32   //88
	Locale              uint32   //92
	InputLanguage       uint32   //96
	OutputLanguage      uint32   //100
	MinVersion          uint32   //104
	FirstImageOffset    uint32   //108
	HuffmanRecordOffset uint32   //112
	HuffmanRecordCount  uint32   //116
	HuffmanTableOffset  uint32   //120
	HuffTableLength     uint32   //124
	ExthFlags           uint32   //128
	Skip3               [32]byte //132
	Unknown0            uint32   //164
	DrmOffset           uint32   //168
	DrmCount            uint32   //172
	DrmSize             uint32   //176
	DrmFlags            uint32   //180
	Skip4               [8]byte  //184
	FirstContentNumber  uint32   //192
	FdstFlowCount       uint32   //196
	FcisOffset          uint32   //200
	FcisCount           uint32   //204
	FlisOffset          uint32   //208
	FlisCount           uint32   //212
	Skip5               [8]byte  //216
	SrcsOffset          uint32   //224
	SrcsCount           uint32   //228
	Skip6               [8]byte  //232
	TrailDataFlags      uint16   //240
	NcxIndex            uint32   //242
	FragmentIndex       uint32   //246
	SkeletonIndex       uint32   //250
	DatpOffset          uint32   //254
	GuideIndex          uint32   //258-262
}

func GetStruct(file *os.File, hd interface{}, length int, offset int64) (rd int, err error) {
	b := make([]byte, length)
	rd, err = file.ReadAt(b, offset)
	if err != nil {
		return
	}
	buf := bytes.NewBuffer(b)
	err = binary.Read(buf, binary.BigEndian, hd)
	return
}

type ExthHeader struct {
	Identifier   uint32 //starts at byte 0
	HeaderLength uint32 //4
	RecordCount  uint32 //8-12
}

type ExthRecordInfo struct {
	RecordType   uint32 //starts at 0
	RecordLength uint32 //4-8
}

type ExthRecordData []byte

type FileHeader struct {
	Format     PDFormat
	Sections   []PDRecordInfoSection
	MobiHeader Mobi8Header
	Fcis       FcisRecord
	Flis       FlisRecord
	Eof        EofRecord
}

type FlisRecord struct {
	Identifier [4]byte //starts at 0, "FLIS"
	Skip1      uint32  //4            8
	Skip2      uint16  //8            65
	Skip3      uint16  //10           0
	Skip4      uint32  //12           0
	Skip5      uint32  //16          -1
	Skip6      uint16  //20           1
	Skip7      uint16  //22           3
	Skip8      uint32  //24           3
	Skip9      uint32  //28           1
	Skip10     uint32  //32-36       -1
}

type FcisRecord struct {
	Identifier [4]byte //starts at 0, to 4 "FCIS"
	Skip1      uint32  //4     20
	Skip2      uint32  //8     16
	Skip3      uint32  //12    1
	Skip4      uint32  //16    0
	TextLength uint32  //20    Same as PDHeader
	Skip5      uint32  //24    0
	Skip6      uint32  //28    32
	Skip7      uint32  //32    8
	Skip8      uint16  //36    1
	Skip9      uint16  //38    1
	Skip10     uint32  //40-44 0
}

type EofRecord struct {
	Skip1 uint8 //starts at 0, to 1  233
	Skip2 uint8 //1                  142
	Skip3 uint8 //2                   13
	Skip4 uint8 //3-4                 10
}

func main() {
	hd, err := GetFileHeader("file.mobi")
	check(err)
	fmt.Printf("%#v\n", hd.Format)
	fmt.Printf("%#v %#v\n", hd.Sections[0], hd.Sections[181])
	fmt.Printf("%#v\n", hd.MobiHeader)
	fmt.Printf("%#v\n", hd.Fcis)
	fmt.Printf("%#v\n", hd.Flis)
	fmt.Printf("%#v\n", hd.Eof)
}

//GetPDRecordInfoSectionList reads `count` items from `file`,
//starting at byte `offset` and placing the result in in `ris`.
//Returns the number of records read, and any error.
func GetPDRecordInfoSectionList(file *os.File, ris *[]PDRecordInfoSection, count int, start int) (ii int, err error) {
	for ii = 0; ii < count; ii++ {
		var section PDRecordInfoSection
		_, err = GetStruct(file, &section, 8, int64(start+ii*8))
		if err != nil {
			return
		}
		*ris = append(*ris, section)
	}
	return
}

//GetFileHeader reads the header information from the file at path `path`.
//Returns the FileHeader as read, and any error.
func GetFileHeader(path string) (hd FileHeader, err error) {
	file, err := os.Open(path)
	defer file.Close()
	start := 0
	if err != nil {
		return
	}

	start, err = GetStruct(file, &hd.Format, 78, 0)
	fmt.Println(start)
	if err != nil {
		return
	}

	rdr := flate.NewReader(file)
	defer rdr.Close()

	a, err := GetPDRecordInfoSectionList(file, &hd.Sections, int(hd.Format.SectionCount), start)
	fmt.Println(a)
	if err != nil {
		return
	}

	a, err = GetStruct(file, &hd.MobiHeader, 262, int64(hd.Sections[0].DataOffset))
	fmt.Println(a, err)

	if hd.MobiHeader.FcisCount > 0 {
		offset := int64(hd.Sections[hd.MobiHeader.FcisOffset].DataOffset)
		a, err = GetStruct(file, &hd.Fcis, 44, offset)
		fmt.Println("FCIS", a, err)
	}

	if hd.MobiHeader.FlisCount > 0 {
		offset := int64(hd.Sections[hd.MobiHeader.FlisOffset].DataOffset)
		a, err = GetStruct(file, &hd.Flis, 36, offset)
		fmt.Println("FLIS", a, err)
	}

	offset := int64(hd.Sections[len(hd.Sections)-1].DataOffset)
	a, err = GetStruct(file, &hd.Eof, 4, offset)
	fmt.Println("EOF", a, err)

	return
}

//check panics when there's an error.
func check(err error) {
	if err != nil {
		panic(err)
	}
}
<commit_msg>got exth bits working.<commit_after>package main

import (
	"bytes"
	"compress/flate"
	"encoding/binary"
	"fmt"
	"os"
)

type PDFormat struct {
	Name             [32]byte //start at byte 0
	Attributes       uint16   //32
	Version          uint16   //34
	CreationDate     uint32   //36
	ModifyDate       uint32   //40
	BackupDate       uint32   //44
	ModifyNumber     uint32   //48
	AppInfoID        uint32   //52
	SortInfoID       uint32   //56
	Type             [4]byte  //60
	Creator          [4]byte  //64
	UniqueIDSeed     uint32   //68
	NextRecordListID uint32   //72
	SectionCount     uint16   //76-78
}

type PDRecordInfoSection struct {
	DataOffset uint32  //starts at byte 0
	Attributes byte    //4
	UniqueID   [3]byte //5-8
}

type PDHeader struct {
	CompressionType uint16 //starts at byte 0
	_               uint16 //2 (always zero?)
	TextLength      uint32 //4
	RecordCount     uint16 //8
	RecordSize      uint16 //10
	CurrentPosition uint32 //12-16
}

type Mobi8Header struct {
	//note that since the values are being read
	//by bytes.Read, we can't use "_" as a name,
	//or have any unexported fields.
	//This is stupid, but whatever.
	//We use the name "SkipX" instead.
	CompressionType     uint16   //start at byte 0
	Skip1               uint16   //2 (always zero?)
	TextLength          uint32   //4
	RecordCount         uint16   //8
	RecordSize          uint16   //10
	CryptoType          uint16   //12
	Skip2               uint16   //14 filler
	Identifier          [4]byte  //16
	HeaderLength        uint32   //20
	Type                uint32   //24
	TextEncoding        uint32   //28
	UniqueID            uint32   //32
	Version             uint32   //36
	OrtographicIndex    uint32   //40
	IncflectionIndex    uint32   //44
	IndexNames          uint32   //48
	IndexKeys           uint32   //52
	Extra               [24]byte //56
	FirstNontext        uint32   //80
	TitleOffset         uint32   //84
	TitleLength         uint32   //88
	Locale              uint32   //92
	InputLanguage       uint32   //96
	OutputLanguage      uint32   //100
	MinVersion          uint32   //104
	FirstImageOffset    uint32   //108
	HuffmanRecordOffset uint32   //112
	HuffmanRecordCount  uint32   //116
	HuffmanTableOffset  uint32   //120
	HuffTableLength     uint32   //124
	ExthFlags           uint32   //128
	Skip3               [32]byte //132
	Unknown0            uint32   //164
	DrmOffset           uint32   //168
	DrmCount            uint32   //172
	DrmSize             uint32   //176
	DrmFlags            uint32   //180
	Skip4               [8]byte  //184
	FirstContentNumber  uint32   //192
	FdstFlowCount       uint32   //196
	FcisOffset          uint32   //200
	FcisCount           uint32   //204
	FlisOffset          uint32   //208
	FlisCount           uint32   //212
	Skip5               [8]byte  //216
	SrcsOffset          uint32   //224
	SrcsCount           uint32   //228
	Skip6               [8]byte  //232
	TrailDataFlags      uint32   //240
	NcxIndex            uint32   //244-248
}

func GetStruct(file *os.File, hd interface{}, length int, offset int64) (rd int, err error) {
	b := make([]byte, length)
	rd, err = file.ReadAt(b, offset)
	if err != nil {
		return
	}
	buf := bytes.NewBuffer(b)
	err = binary.Read(buf, binary.BigEndian, hd)
	return
}

type ExthHeader struct {
	Identifier  [4]byte //starts at byte 0
	Length      uint32  //4
	RecordCount uint32  //8-12
}

type ExthData struct {
	Header  ExthHeader
	Records []ExthRecord
}

type ExthRecordInfo struct {
	Type   uint32 //starts at 0
	Length uint32 //4-8
}

type ExthRecordData []byte

type ExthRecord struct {
	Info ExthRecordInfo
	Data ExthRecordData
}

type FileHeader struct {
	Format     PDFormat
	Sections   []PDRecordInfoSection
	MobiHeader Mobi8Header
	Fcis       FcisRecord
	Flis       FlisRecord
	Eof        EofRecord
	Exth       ExthData
}

type FlisRecord struct {
	Identifier [4]byte //starts at 0, "FLIS"
	Skip1      uint32  //4            8
	Skip2      uint16  //8            65
	Skip3      uint16  //10           0
	Skip4      uint32  //12           0
	Skip5      uint32  //16          -1
	Skip6      uint16  //20           1
	Skip7      uint16  //22           3
	Skip8      uint32  //24           3
	Skip9      uint32  //28           1
	Skip10     uint32  //32-36       -1
}

type FcisRecord struct {
	Identifier [4]byte //starts at 0, to 4 "FCIS"
	Skip1      uint32  //4     20
	Skip2      uint32  //8     16
	Skip3      uint32  //12    1
	Skip4      uint32  //16    0
	TextLength uint32  //20    Same as PDHeader
	Skip5      uint32  //24    0
	Skip6      uint32  //28    32
	Skip7      uint32  //32    8
	Skip8      uint16  //36    1
	Skip9      uint16  //38    1
	Skip10     uint32  //40-44 0
}

type EofRecord struct {
	Skip1 uint8 //starts at 0, to 1  233
	Skip2 uint8 //1                  142
	Skip3 uint8 //2                   13
	Skip4 uint8 //3-4                 10
}

func main() {
	hd, err := GetFileHeader("file.mobi")
	check(err)
	fmt.Printf("%#v\n", hd.Format)
	fmt.Printf("%#v %#v\n", hd.Sections[0], hd.Sections[181])
	fmt.Printf("%#v\n", hd.MobiHeader)
	fmt.Printf("%#v\n", hd.Fcis)
	fmt.Printf("%#v\n", hd.Flis)
	fmt.Printf("%#v\n", hd.Eof)
	fmt.Printf("%#v\n", hd.Exth.Header)
	fmt.Printf("%#v\n", hd.Exth.Header.Identifier)
	fmt.Printf("%#v %#v\n", string(hd.Exth.Records[0].Data), string(hd.Exth.Records[16].Data))
	fmt.Printf("%#v\n", string(hd.Exth.Records[len(hd.Exth.Records)-1].Data))
}

//GetPDRecordInfoSectionList reads `count` items from `file`,
//starting at byte `offset` and placing the result in in `ris`.
//Returns the number of records read, and any error.
func GetPDRecordInfoSectionList(file *os.File, ris *[]PDRecordInfoSection, count int, start int) (ii int, err error) {
	for ii = 0; ii < count; ii++ {
		var section PDRecordInfoSection
		_, err = GetStruct(file, &section, 8, int64(start+ii*8))
		if err != nil {
			return
		}
		*ris = append(*ris, section)
	}
	return
}

//GetFileHeader reads the header information from the file at path `path`.
//Returns the FileHeader as read, and any error.
func GetFileHeader(path string) (hd FileHeader, err error) {
	file, err := os.Open(path)
	defer file.Close()
	start := 0
	if err != nil {
		return
	}

	start, err = GetStruct(file, &hd.Format, 78, 0)
	fmt.Println(start)
	if err != nil {
		return
	}

	rdr := flate.NewReader(file)
	defer rdr.Close()

	bytesRead, err := GetPDRecordInfoSectionList(file, &hd.Sections, int(hd.Format.SectionCount), start)
	fmt.Println(bytesRead)
	if err != nil {
		return
	}
	offset := int64(hd.Sections[0].DataOffset)
	bytesRead, err = GetStruct(file, &hd.MobiHeader, 248, offset)
	offset += int64(bytesRead)

	if hd.MobiHeader.ExthFlags&64 == 64 {
		GetExthData(file, &hd.Exth, offset)
	}

	if hd.MobiHeader.FcisCount > 0 {
		offset := int64(hd.Sections[hd.MobiHeader.FcisOffset].DataOffset)
		bytesRead, err = GetStruct(file, &hd.Fcis, 44, offset)
		fmt.Println("FCIS", bytesRead, err)
	}

	if hd.MobiHeader.FlisCount > 0 {
		offset := int64(hd.Sections[hd.MobiHeader.FlisOffset].DataOffset)
		bytesRead, err = GetStruct(file, &hd.Flis, 36, offset)
		fmt.Println("FLIS", bytesRead, err)
	}

	offset = int64(hd.Sections[len(hd.Sections)-1].DataOffset)
	bytesRead, err = GetStruct(file, &hd.Eof, 4, offset)
	fmt.Println("EOF", bytesRead, err)

	return
}

//GetExthData reads exth data at index `offset` in `file`, and copies
//the contents into into `ex`.
//Return the total number of bytes read and the error.
func GetExthData(file *os.File, ex *ExthData, offset int64) (total int, err error) {
	read, err := GetStruct(file, &ex.Header, 12, offset)
	total += read
	for ii := uint32(0); ii < ex.Header.RecordCount; ii++ {
		var record ExthRecord
		read, err = GetStruct(file, &record.Info, 8, offset+int64(total))
		total += read
		record.Data = make([]byte, int(record.Info.Length-8))
		read, err = GetStruct(file, &record.Data, len(record.Data), offset+int64(total))
		total += read
		ex.Records = append(ex.Records, record)
	}
	return
}

//check panics when there's an error.
func check(err error) {
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package gocassa

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"

	"github.com/gocql/gocql"
	"github.com/google/btree"
)

// MockKeySpace implements the KeySpace interface and constructs in-memory tables.
type mockKeySpace struct {
	k
}

type mockOp struct {
	options Options
	funcs   []func(mockOp) error
}

func newOp(f func(mockOp) error) mockOp {
	return mockOp{
		funcs: []func(mockOp) error{f},
	}
}

func (m mockOp) Add(ops ...Op) Op {
	for _, o := range ops {
		m.funcs = append(m.funcs, o.(mockOp).funcs...)
	}
	return m
}

func (m mockOp) Run() error {
	for _, f := range m.funcs {
		err := f(m)
		if err != nil {
			return err
		}
	}
	return nil
}

func (m mockOp) WithOptions(opt Options) Op {
	return mockOp{
		options: opt,
		funcs:   m.funcs,
	}
}

func (m mockOp) RunAtomically() error {
	return m.Run()
}

func (ks *mockKeySpace) NewTable(name string, entity interface{}, fields map[string]interface{}, keys Keys) Table {
	return &MockTable{
		name:   name,
		entity: entity,
		keys:   keys,
		rows:   map[rowKey]*btree.BTree{},
	}
}

func NewMockKeySpace() KeySpace {
	ks := &mockKeySpace{}
	ks.tableFactory = ks
	return ks
}

// MockTable implements the Table interface and stores rows in-memory.
type MockTable struct {
	// rows is mapping from row key to column group key to column map
	name    string
	rows    map[rowKey]*btree.BTree
	entity  interface{}
	keys    Keys
	options Options
}

type rowKey string
type superColumn struct {
	Key     *keyPart
	Columns map[string]interface{}
}

func (c *superColumn) Less(item btree.Item) bool {
	other, ok := item.(*superColumn)
	if !ok {
		return false
	}

	return c.Key.Less(other.Key)
}

type gocqlTypeInfo struct {
	proto byte
	typ   gocql.Type
}

func (t gocqlTypeInfo) New() interface{} {
	return &gocqlTypeInfo{t.proto, t.typ}
}

func (t gocqlTypeInfo) Type() gocql.Type {
	return t.typ
}

func (t gocqlTypeInfo) Version() byte {
	return t.proto
}

func (t gocqlTypeInfo) Custom() string {
	return ""
}

type keyPart struct {
	Key   string
	Value interface{}
	Tail  *keyPart
}

func (k *keyPart) Prepend(key string, value interface{}) *keyPart {
	return &keyPart{Key: key, Value: value, Tail: k}
}

func (k *keyPart) Less(other *keyPart) bool {
	for part, other := k, other; part != nil && other != nil; part, other = part.Tail, other.Tail {
		cmp := bytes.Compare(part.Bytes(), other.Bytes())
		if cmp == 0 {
			continue
		}
		return cmp < 0
	}

	return false
}

func (k *keyPart) Bytes() []byte {
	typeInfo := &gocqlTypeInfo{
		proto: 0x03,
		typ:   cassaType(k.Value),
	}
	// FIXME handle err
	marshalled, _ := gocql.Marshal(typeInfo, k.Value)
	return marshalled
}

func (k *keyPart) RowKey() rowKey {
	buf := bytes.Buffer{}

	for part := k; part != nil; part = part.Tail {
		buf.Write(part.Bytes())
	}

	return rowKey(buf.String())
}

func (k *keyPart) ToSuperColumn() *superColumn {
	return &superColumn{Key: k}
}

func (t *MockTable) zero() interface{} {
	return reflect.New(reflect.TypeOf(t.entity)).Interface()
}

func (t *MockTable) keyFromColumnValues(columns map[string]interface{}, keys []string) (*keyPart, error) {
	var key *keyPart

	for _, column := range keys {
		value, ok := columns[column]

		if !ok {
			return nil, fmt.Errorf("Missing mandatory PRIMARY KEY part %s", column)
		}
		key = key.Prepend(column, value)
	}

	return key, nil
}

func (t *MockTable) Name() string {
	if len(t.options.TableName) > 0 {
		return t.options.TableName
	}
	return t.name
}

func (t *MockTable) getOrCreateRow(rowKey *keyPart) *btree.BTree {
	row := t.rows[rowKey.RowKey()]
	if row == nil {
		row = btree.New(2)
		t.rows[rowKey.RowKey()] = row
	}
	return row
}

func (t *MockTable) getOrCreateColumnGroup(rowKey, superColumnKey *keyPart) map[string]interface{} {
	row := t.getOrCreateRow(rowKey)
	scol := superColumnKey.ToSuperColumn()

	if row.Has(scol) {
		return row.Get(scol).(*superColumn).Columns
	}
	row.ReplaceOrInsert(scol)
	scol.Columns = map[string]interface{}{}

	return scol.Columns
}

func (t *MockTable) SetWithOptions(i interface{}, options Options) Op {
	return newOp(func(m mockOp) error {
		columns, ok := toMap(i)
		if !ok {
			return errors.New("Can't create: value not understood")
		}

		rowKey, err := t.keyFromColumnValues(columns, t.keys.PartitionKeys)
		if err != nil {
			return err
		}

		superColumnKey, err := t.keyFromColumnValues(columns, t.keys.ClusteringColumns)
		if err != nil {
			return err
		}

		superColumn := t.getOrCreateColumnGroup(rowKey, superColumnKey)

		for k, v := range columns {
			superColumn[k] = v
		}
		return nil
	})
}

func (t *MockTable) Set(i interface{}) Op {
	return t.SetWithOptions(i, t.options)
}

func (t *MockTable) Where(relations ...Relation) Filter {
	return &MockFilter{
		table:     t,
		relations: relations,
	}
}

func (t *MockTable) Create() error {
	return nil
}

func (t *MockTable) CreateStatement() (string, error) {
	return "", nil
}

func (t *MockTable) Recreate() error {
	return nil
}

func (t *MockTable) WithOptions(o Options) Table {
	return &MockTable{
		name:    t.name,
		rows:    t.rows,
		entity:  t.entity,
		keys:    t.keys,
		options: t.options.Merge(o),
	}
}

// MockFilter implements the Filter interface and works with MockTable.
type MockFilter struct {
	table     *MockTable
	relations []Relation
}

func (f *MockFilter) rowMatch(row map[string]interface{}) bool {
	for _, relation := range f.relations {
		value := row[relation.key]
		if !relation.accept(value) {
			return false
		}
	}
	return true
}

func (f *MockFilter) keyRelationMap() map[string]Relation {
	result := map[string]Relation{}

	for _, relation := range f.relations {
		result[relation.key] = relation
	}

	return result
}

func (f *MockFilter) keysFromRelations(keys []string) ([]*keyPart, error) {
	keyRelationMap := f.keyRelationMap()
	var rowKey *keyPart
	var result []*keyPart

	if len(keys) == 0 {
		return []*keyPart{nil}, nil
	}

	for i, key := range keys {
		lastKey := i == len(keys)-1
		relation, ok := keyRelationMap[key]

		if !ok {
			return nil, fmt.Errorf("Missing mandatory PRIMARY KEY part %s", key)
		}

		if relation.op != equality && !(lastKey && relation.op == in) {
			return nil, fmt.Errorf("Invalid use of PK %s", key)
		}

		if !lastKey {
			rowKey = rowKey.Prepend(key, relation.terms[0])
		} else {
			for _, term := range relation.terms {
				result = append(result, rowKey.Prepend(relation.key, term))
			}
		}
	}

	return result, nil
}

func (f *MockFilter) UpdateWithOptions(m map[string]interface{}, options Options) Op {
	return newOp(func(mock mockOp) error {
		rowKeys, err := f.keysFromRelations(f.table.keys.PartitionKeys)
		if err != nil {
			return err
		}

		for _, rowKey := range rowKeys {
			superColumnKeys, err := f.keysFromRelations(f.table.keys.ClusteringColumns)
			if err != nil {
				return err
			}

			for _, superColumnKey := range superColumnKeys {
				superColumn := f.table.getOrCreateColumnGroup(rowKey, superColumnKey)

				for _, keyPart := range []*keyPart{rowKey, superColumnKey} {
					for k := keyPart; k != nil; k = k.Tail {
						superColumn[k.Key] = k.Value
					}
				}

				for key, value := range m {
					superColumn[key] = value
				}
			}
		}

		return nil
	})
}

func (f *MockFilter) Update(m map[string]interface{}) Op {
	return f.UpdateWithOptions(m, Options{})
}

func (f *MockFilter) Delete() Op {
	return newOp(func(m mockOp) error {
		rowKeys, err := f.keysFromRelations(f.table.keys.PartitionKeys)
		if err != nil {
			return err
		}

		for _, rowKey := range rowKeys {
			row := f.table.rows[rowKey.RowKey()]
			if row == nil {
				return nil
			}

			row.Ascend(func(item btree.Item) bool {
				columns := item.(*superColumn).Columns
				if f.rowMatch(columns) {
					row.Delete(item)
				}

				return true
			})
		}

		return nil
	})
}

func (q *MockFilter) Read(out interface{}) Op {
	return newOp(func(m mockOp) error {
		rowKeys, err := q.keysFromRelations(q.table.keys.PartitionKeys)
		if err != nil {
			return err
		}

		var result []map[string]interface{}
		for _, rowKey := range rowKeys {
			row := q.table.rows[rowKey.RowKey()]
			if row == nil {
				continue
			}

			row.Ascend(func(item btree.Item) bool {
				columns := item.(*superColumn).Columns
				if q.rowMatch(columns) {
					result = append(result, columns)
				}

				return true
			})
		}
		opt := q.table.options.Merge(m.options)
		if opt.Limit > 0 && opt.Limit < len(result) {
			result = result[:opt.Limit]
		}

		return q.assignResult(result, out)
	})
}

func (q *MockFilter) assignResult(records interface{}, out interface{}) error {
	bytes, err := json.Marshal(records)
	if err != nil {
		return err
	}
	return json.Unmarshal(bytes, out)
}

func (q *MockFilter) ReadOne(out interface{}) Op {
	return newOp(func(m mockOp) error {
		slicePtrVal := reflect.New(reflect.SliceOf(reflect.ValueOf(out).Elem().Type()))

		err := q.Read(slicePtrVal.Interface()).Run()
		if err != nil {
			return err
		}

		sliceVal := slicePtrVal.Elem()
		if sliceVal.Len() < 1 {
			return RowNotFoundError{}
		}
		q.assignResult(sliceVal.Index(0).Interface(), out)
		return nil
	})
}
<commit_msg>Added locking to mock keyspace<commit_after>package gocassa

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"sync"

	"github.com/gocql/gocql"
	"github.com/google/btree"
)

// MockKeySpace implements the KeySpace interface and constructs in-memory tables.
type mockKeySpace struct {
	k
}

type mockOp struct {
	options Options
	funcs   []func(mockOp) error
}

func newOp(f func(mockOp) error) mockOp {
	return mockOp{
		funcs: []func(mockOp) error{f},
	}
}

func (m mockOp) Add(ops ...Op) Op {
	for _, o := range ops {
		m.funcs = append(m.funcs, o.(mockOp).funcs...)
	}
	return m
}

func (m mockOp) Run() error {
	for _, f := range m.funcs {
		err := f(m)
		if err != nil {
			return err
		}
	}
	return nil
}

func (m mockOp) WithOptions(opt Options) Op {
	return mockOp{
		options: opt,
		funcs:   m.funcs,
	}
}

func (m mockOp) RunAtomically() error {
	return m.Run()
}

func (ks *mockKeySpace) NewTable(name string, entity interface{}, fields map[string]interface{}, keys Keys) Table {
	return &MockTable{
		name:   name,
		entity: entity,
		keys:   keys,
		rows:   map[rowKey]*btree.BTree{},
	}
}

func NewMockKeySpace() KeySpace {
	ks := &mockKeySpace{}
	ks.tableFactory = ks
	return ks
}

// MockTable implements the Table interface and stores rows in-memory.
type MockTable struct {
	sync.RWMutex

	// rows is mapping from row key to column group key to column map
	name    string
	rows    map[rowKey]*btree.BTree
	entity  interface{}
	keys    Keys
	options Options
}

type rowKey string
type superColumn struct {
	Key     *keyPart
	Columns map[string]interface{}
}

func (c *superColumn) Less(item btree.Item) bool {
	other, ok := item.(*superColumn)
	if !ok {
		return false
	}

	return c.Key.Less(other.Key)
}

type gocqlTypeInfo struct {
	proto byte
	typ   gocql.Type
}

func (t gocqlTypeInfo) New() interface{} {
	return &gocqlTypeInfo{t.proto, t.typ}
}

func (t gocqlTypeInfo) Type() gocql.Type {
	return t.typ
}

func (t gocqlTypeInfo) Version() byte {
	return t.proto
}

func (t gocqlTypeInfo) Custom() string {
	return ""
}

type keyPart struct {
	Key   string
	Value interface{}
	Tail  *keyPart
}

func (k *keyPart) Prepend(key string, value interface{}) *keyPart {
	return &keyPart{Key: key, Value: value, Tail: k}
}

func (k *keyPart) Less(other *keyPart) bool {
	for part, other := k, other; part != nil && other != nil; part, other = part.Tail, other.Tail {
		cmp := bytes.Compare(part.Bytes(), other.Bytes())
		if cmp == 0 {
			continue
		}
		return cmp < 0
	}

	return false
}

func (k *keyPart) Bytes() []byte {
	typeInfo := &gocqlTypeInfo{
		proto: 0x03,
		typ:   cassaType(k.Value),
	}
	// FIXME handle err
	marshalled, _ := gocql.Marshal(typeInfo, k.Value)
	return marshalled
}

func (k *keyPart) RowKey() rowKey {
	buf := bytes.Buffer{}

	for part := k; part != nil; part = part.Tail {
		buf.Write(part.Bytes())
	}

	return rowKey(buf.String())
}

func (k *keyPart) ToSuperColumn() *superColumn {
	return &superColumn{Key: k}
}

func (t *MockTable) zero() interface{} {
	return reflect.New(reflect.TypeOf(t.entity)).Interface()
}

func (t *MockTable) keyFromColumnValues(columns map[string]interface{}, keys []string) (*keyPart, error) {
	var key *keyPart

	for _, column := range keys {
		value, ok := columns[column]

		if !ok {
			return nil, fmt.Errorf("Missing mandatory PRIMARY KEY part %s", column)
		}
		key = key.Prepend(column, value)
	}

	return key, nil
}

func (t *MockTable) Name() string {
	if len(t.options.TableName) > 0 {
		return t.options.TableName
	}
	return t.name
}

func (t *MockTable) getOrCreateRow(rowKey *keyPart) *btree.BTree {
	t.Lock()
	defer t.Unlock()

	row := t.rows[rowKey.RowKey()]
	if row == nil {
		row = btree.New(2)
		t.rows[rowKey.RowKey()] = row
	}
	return row
}

func (t *MockTable) getOrCreateColumnGroup(rowKey, superColumnKey *keyPart) map[string]interface{} {
	row := t.getOrCreateRow(rowKey)
	scol := superColumnKey.ToSuperColumn()

	if row.Has(scol) {
		return row.Get(scol).(*superColumn).Columns
	}
	row.ReplaceOrInsert(scol)
	scol.Columns = map[string]interface{}{}

	return scol.Columns
}

func (t *MockTable) SetWithOptions(i interface{}, options Options) Op {
	return newOp(func(m mockOp) error {
		columns, ok := toMap(i)
		if !ok {
			return errors.New("Can't create: value not understood")
		}

		rowKey, err := t.keyFromColumnValues(columns, t.keys.PartitionKeys)
		if err != nil {
			return err
		}

		superColumnKey, err := t.keyFromColumnValues(columns, t.keys.ClusteringColumns)
		if err != nil {
			return err
		}

		superColumn := t.getOrCreateColumnGroup(rowKey, superColumnKey)

		for k, v := range columns {
			superColumn[k] = v
		}
		return nil
	})
}

func (t *MockTable) Set(i interface{}) Op {
	return t.SetWithOptions(i, t.options)
}

func (t *MockTable) Where(relations ...Relation) Filter {
	return &MockFilter{
		table:     t,
		relations: relations,
	}
}

func (t *MockTable) Create() error {
	return nil
}

func (t *MockTable) CreateStatement() (string, error) {
	return "", nil
}

func (t *MockTable) Recreate() error {
	return nil
}

func (t *MockTable) WithOptions(o Options) Table {
	return &MockTable{
		name:    t.name,
		rows:    t.rows,
		entity:  t.entity,
		keys:    t.keys,
		options: t.options.Merge(o),
	}
}

// MockFilter implements the Filter interface and works with MockTable.
type MockFilter struct {
	table     *MockTable
	relations []Relation
}

func (f *MockFilter) rowMatch(row map[string]interface{}) bool {
	for _, relation := range f.relations {
		value := row[relation.key]
		if !relation.accept(value) {
			return false
		}
	}
	return true
}

func (f *MockFilter) keyRelationMap() map[string]Relation {
	result := map[string]Relation{}

	for _, relation := range f.relations {
		result[relation.key] = relation
	}

	return result
}

func (f *MockFilter) keysFromRelations(keys []string) ([]*keyPart, error) {
	keyRelationMap := f.keyRelationMap()
	var rowKey *keyPart
	var result []*keyPart

	if len(keys) == 0 {
		return []*keyPart{nil}, nil
	}

	for i, key := range keys {
		lastKey := i == len(keys)-1
		relation, ok := keyRelationMap[key]

		if !ok {
			return nil, fmt.Errorf("Missing mandatory PRIMARY KEY part %s", key)
		}

		if relation.op != equality && !(lastKey && relation.op == in) {
			return nil, fmt.Errorf("Invalid use of PK %s", key)
		}

		if !lastKey {
			rowKey = rowKey.Prepend(key, relation.terms[0])
		} else {
			for _, term := range relation.terms {
				result = append(result, rowKey.Prepend(relation.key, term))
			}
		}
	}

	return result, nil
}

func (f *MockFilter) UpdateWithOptions(m map[string]interface{}, options Options) Op {
	return newOp(func(mock mockOp) error {
		rowKeys, err := f.keysFromRelations(f.table.keys.PartitionKeys)
		if err != nil {
			return err
		}

		for _, rowKey := range rowKeys {
			superColumnKeys, err := f.keysFromRelations(f.table.keys.ClusteringColumns)
			if err != nil {
				return err
			}

			for _, superColumnKey := range superColumnKeys {
				superColumn := f.table.getOrCreateColumnGroup(rowKey, superColumnKey)

				for _, keyPart := range []*keyPart{rowKey, superColumnKey} {
					for k := keyPart; k != nil; k = k.Tail {
						superColumn[k.Key] = k.Value
					}
				}

				for key, value := range m {
					superColumn[key] = value
				}
			}
		}

		return nil
	})
}

func (f *MockFilter) Update(m map[string]interface{}) Op {
	return f.UpdateWithOptions(m, Options{})
}

func (f *MockFilter) Delete() Op {
	return newOp(func(m mockOp) error {
		f.table.Lock()
		defer f.table.Unlock()

		rowKeys, err := f.keysFromRelations(f.table.keys.PartitionKeys)
		if err != nil {
			return err
		}

		for _, rowKey := range rowKeys {
			row := f.table.rows[rowKey.RowKey()]
			if row == nil {
				return nil
			}

			row.Ascend(func(item btree.Item) bool {
				columns := item.(*superColumn).Columns
				if f.rowMatch(columns) {
					row.Delete(item)
				}

				return true
			})
		}

		return nil
	})
}

func (q *MockFilter) Read(out interface{}) Op {
	return newOp(func(m mockOp) error {
		q.table.Lock()
		defer q.table.Unlock()

		rowKeys, err := q.keysFromRelations(q.table.keys.PartitionKeys)
		if err != nil {
			return err
		}

		var result []map[string]interface{}
		for _, rowKey := range rowKeys {
			row := q.table.rows[rowKey.RowKey()]
			if row == nil {
				continue
			}

			row.Ascend(func(item btree.Item) bool {
				columns := item.(*superColumn).Columns
				if q.rowMatch(columns) {
					result = append(result, columns)
				}

				return true
			})
		}
		opt := q.table.options.Merge(m.options)
		if opt.Limit > 0 && opt.Limit < len(result) {
			result = result[:opt.Limit]
		}

		return q.assignResult(result, out)
	})
}

func (q *MockFilter) assignResult(records interface{}, out interface{}) error {
	bytes, err := json.Marshal(records)
	if err != nil {
		return err
	}
	return json.Unmarshal(bytes, out)
}

func (q *MockFilter) ReadOne(out interface{}) Op {
	return newOp(func(m mockOp) error {
		slicePtrVal := reflect.New(reflect.SliceOf(reflect.ValueOf(out).Elem().Type()))

		err := q.Read(slicePtrVal.Interface()).Run()
		if err != nil {
			return err
		}

		sliceVal := slicePtrVal.Elem()
		if sliceVal.Len() < 1 {
			return RowNotFoundError{}
		}
		q.assignResult(sliceVal.Index(0).Interface(), out)
		return nil
	})
}
<|endoftext|>


<commit_before>package checks

import (
	"fmt"
	"io"
	"os"
	"strings"
	"sync"

	"github.com/mackerelio/mackerel-agent/config"
	cli "gopkg.in/urfave/cli.v1"
	yaml "gopkg.in/yaml.v2"
)

// Command is command definition of mkr checks
var Command = cli.Command{
	Name:  "checks",
	Usage: "Utility for check plugins",
	Subcommands: []cli.Command{
		commandRun,
	},
}

var commandRun = cli.Command{
	Name:  "run",
	Usage: "run check commands in mackerel-agent.conf",
	Description: `
    Execute command of check plugins in mackerel-agent.conf all at once.
    It is used for checking setting and operation of the check plugins.
    The result is output to stdout in TAP format. If any check fails,
    it exits non-zero.
`,
	Action: doRunChecks,
}

func doRunChecks(c *cli.Context) error {
	confFile := c.GlobalString("conf")
	conf, err := config.LoadConfig(confFile)
	if err != nil {
		return err
	}
	checkers := make([]checker, len(conf.CheckPlugins))
	i := 0
	for name, p := range conf.CheckPlugins {
		checkers[i] = &checkPluginChecker{
			name: name,
			cp:   p,
		}
		i++
	}
	return runChecks(checkers, os.Stdout)
}

type result struct {
	Name     string   `yaml:"-"`
	Memo     string   `yaml:"memo,omitempty"`
	Cmd      []string `yaml:"command,flow"`
	Stdout   string   `yaml:"stdout,omitempty"`
	Stderr   string   `yaml:"stderr,omitempty"`
	ExitCode int      `yaml:"exitCode,omitempty"`
	ErrMsg   string   `yaml:"error,omitempty"`
}

func (re *result) ok() bool {
	return re.ExitCode == 0 && re.ErrMsg == ""
}

func (re *result) tapFormat(num int) string {
	okOrNot := "ok"
	if !re.ok() {
		okOrNot = "not ok"
	}
	b, _ := yaml.Marshal(re)
	// indent
	yamlStr := "  " + strings.Replace(strings.TrimSpace(string(b)), "\n", "\n  ", -1)
	return fmt.Sprintf("%s %d - %s\n  ---\n%s\n  ...",
		okOrNot, num, re.Name, yamlStr)
}

type checkPluginChecker struct {
	name string
	cp   *config.CheckPlugin
}

func (cpc *checkPluginChecker) check() *result {
	p := cpc.cp
	stdout, stderr, exitCode, err := p.Command.Run()
	cmd := p.Command.Args
	if len(cmd) == 0 {
		cmd = append(cmd, p.Command.Cmd)
	}
	errMsg := ""
	if err != nil {
		errMsg = err.Error()
	}
	return &result{
		Name:     cpc.name,
		Memo:     p.Memo,
		Cmd:      cmd,
		ExitCode: exitCode,
		Stdout:   strings.TrimSpace(stdout),
		Stderr:   strings.TrimSpace(stderr),
		ErrMsg:   errMsg,
	}
}

type checker interface {
	check() *result
}

func runChecks(checkers []checker, w io.Writer) error {
	ch := make(chan *result)
	total := len(checkers)
	go func() {
		wg := &sync.WaitGroup{}
		wg.Add(total)
		for _, c := range checkers {
			go func(c checker) {
				defer wg.Done()
				ch <- c.check()
			}(c)
		}
		wg.Wait()
		close(ch)
	}()
	fmt.Fprintln(w, "TAP version 13")
	fmt.Fprintf(w, "1..%d\n", total)
	testNum, errNum := 1, 0
	for re := range ch {
		fmt.Fprintln(w, re.tapFormat(testNum))
		testNum++
		if !re.ok() {
			errNum++
		}
	}
	if errNum > 0 {
		return fmt.Errorf("Failed %d/%d tests, %3.2f%% okay",
			errNum, total, float64(100*(total-errNum))/float64(total))
	}
	return nil
}
<commit_msg>[checks] introduce semaphore<commit_after>package checks

import (
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/mackerelio/mackerel-agent/config"
	cli "gopkg.in/urfave/cli.v1"
	yaml "gopkg.in/yaml.v2"
)

// Command is command definition of mkr checks
var Command = cli.Command{
	Name:  "checks",
	Usage: "Utility for check plugins",
	Subcommands: []cli.Command{
		commandRun,
	},
}

var commandRun = cli.Command{
	Name:  "run",
	Usage: "run check commands in mackerel-agent.conf",
	Description: `
    Execute command of check plugins in mackerel-agent.conf all at once.
    It is used for checking setting and operation of the check plugins.
    The result is output to stdout in TAP format. If any check fails,
    it exits non-zero.
`,
	Action: doRunChecks,
}

func doRunChecks(c *cli.Context) error {
	confFile := c.GlobalString("conf")
	conf, err := config.LoadConfig(confFile)
	if err != nil {
		return err
	}
	checkers := make([]checker, len(conf.CheckPlugins))
	i := 0
	for name, p := range conf.CheckPlugins {
		checkers[i] = &checkPluginChecker{
			name: name,
			cp:   p,
		}
		i++
	}
	return runChecks(checkers, os.Stdout)
}

type result struct {
	Name     string   `yaml:"-"`
	Memo     string   `yaml:"memo,omitempty"`
	Cmd      []string `yaml:"command,flow"`
	Stdout   string   `yaml:"stdout,omitempty"`
	Stderr   string   `yaml:"stderr,omitempty"`
	ExitCode int      `yaml:"exitCode,omitempty"`
	ErrMsg   string   `yaml:"error,omitempty"`
}

func (re *result) ok() bool {
	return re.ExitCode == 0 && re.ErrMsg == ""
}

func (re *result) tapFormat(num int) string {
	okOrNot := "ok"
	if !re.ok() {
		okOrNot = "not ok"
	}
	b, _ := yaml.Marshal(re)
	// indent
	yamlStr := "  " + strings.Replace(strings.TrimSpace(string(b)), "\n", "\n  ", -1)
	return fmt.Sprintf("%s %d - %s\n  ---\n%s\n  ...",
		okOrNot, num, re.Name, yamlStr)
}

type checkPluginChecker struct {
	name string
	cp   *config.CheckPlugin
}

func (cpc *checkPluginChecker) check() *result {
	p := cpc.cp
	stdout, stderr, exitCode, err := p.Command.Run()
	cmd := p.Command.Args
	if len(cmd) == 0 {
		cmd = append(cmd, p.Command.Cmd)
	}
	errMsg := ""
	if err != nil {
		errMsg = err.Error()
	}
	return &result{
		Name:     cpc.name,
		Memo:     p.Memo,
		Cmd:      cmd,
		ExitCode: exitCode,
		Stdout:   strings.TrimSpace(stdout),
		Stderr:   strings.TrimSpace(stderr),
		ErrMsg:   errMsg,
	}
}

type checker interface {
	check() *result
}

func runChecks(checkers []checker, w io.Writer) error {
	ch := make(chan *result)
	total := len(checkers)
	go func() {
		sem := make(chan struct{}, runtime.NumCPU()*2)
		wg := &sync.WaitGroup{}
		wg.Add(total)
		for _, c := range checkers {
			go func(c checker) {
				defer wg.Done()
				sem <- struct{}{}
				ch <- c.check()
				<-sem
			}(c)
		}
		wg.Wait()
		close(ch)
	}()
	fmt.Fprintln(w, "TAP version 13")
	fmt.Fprintf(w, "1..%d\n", total)
	testNum, errNum := 1, 0
	for re := range ch {
		fmt.Fprintln(w, re.tapFormat(testNum))
		testNum++
		if !re.ok() {
			errNum++
		}
	}
	if errNum > 0 {
		return fmt.Errorf("Failed %d/%d tests, %3.2f%% okay",
			errNum, total, float64(100*(total-errNum))/float64(total))
	}
	return nil
}
<|endoftext|>


<commit_before>package main
import (
	"net/http"
	"log"
	"os"
	"github.com/AtomasBots/atomas_server_hook/atomas"
	"os/exec"
)

var version string

func doOnHook() {
	exec.Command("bash", "onhook").Start()
}

func main() {
	hooks := atomas.Hooks{[]string{}}
	http.HandleFunc("/hook", atomas.CreateHookHandler(&hooks, doOnHook))
	http.HandleFunc("/hook_list", atomas.CreateHookListHandler(&hooks))
	http.HandleFunc("/health", atomas.CreateHealthService())
	http.HandleFunc("/version", atomas.CreateVersionHandler(version))
	port := os.Args[1]
	doOnHook()
	log.Fatal(http.ListenAndServe(port, nil))
}<commit_msg>reboot when no internet<commit_after>package main
import (
	"net/http"
	"log"
	"os"
	"github.com/AtomasBots/atomas_server_hook/atomas"
	"os/exec"
	"time"
)

var version string

func doOnHook() {
	exec.Command("bash", "onhook").Start()
}

func main() {
	hooks := atomas.Hooks{[]string{}}
	http.HandleFunc("/hook", atomas.CreateHookHandler(&hooks, doOnHook))
	http.HandleFunc("/hook_list", atomas.CreateHookListHandler(&hooks))
	http.HandleFunc("/health", atomas.CreateHealthService())
	http.HandleFunc("/version", atomas.CreateVersionHandler(version))
	port := os.Args[1]
	go doOnHook()
	go rebootWhenNoInternet()
	log.Fatal(http.ListenAndServe(port, nil))
}

func rebootWhenNoInternet() {
	time.Sleep(300)
	for {
		if (isWifiOk()) {
			time.Sleep(time.Second)
		} else {
			time.Sleep(time.Second * 10)
			if (!isWifiOk()) {
				exec.Command("sudo", "reboot").Start()
			}
		}
	}
}

func isWifiOk() bool {
	command := exec.Command("ping", "-c", "1", "192.168.0.1")
	command.Start()
	err := command.Wait()
	return err == nil
}<|endoftext|>


<commit_before>/*
Copyright (c) 2014 VMware, Inc. All Rights Reserved.

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 importx

import (
	"archive/tar"
	"io"
	"os"
	"path/filepath"
)

type Archive interface {
	Open(string) (io.ReadCloser, int64, error)
}

type TapeArchive struct {
	path string
}

type TapeArchiveEntry struct {
	io.Reader
	f *os.File
}

func (t *TapeArchiveEntry) Close() error {
	return t.f.Close()
}

func (t *TapeArchive) Open(name string) (io.ReadCloser, int64, error) {
	f, err := os.Open(t.path)
	if err != nil {
		return nil, 0, err
	}

	r := tar.NewReader(f)

	for {
		h, err := r.Next()

		if err == io.EOF {
			break
		}

		matched, err := filepath.Match(name, h.Name)
		if err != nil {
			return nil, 0, err
		}

		if matched {
			return &TapeArchiveEntry{r, f}, h.Size, nil
		}
	}

	_ = f.Close()

	return nil, 0, os.ErrNotExist
}

type FileArchive struct {
	path string
}

func (t *FileArchive) Open(name string) (io.ReadCloser, int64, error) {
	fpath := name

	if name != t.path {
		fpath = filepath.Join(filepath.Dir(t.path), name)
	}

	s, err := os.Stat(fpath)
	if err != nil {
		return nil, 0, err
	}

	f, err := os.Open(fpath)
	if err != nil {
		return nil, 0, err
	}

	return f, s.Size(), nil
}
<commit_msg>Modify archive.go bug<commit_after>/*
Copyright (c) 2014 VMware, Inc. All Rights Reserved.

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 importx

import (
	"archive/tar"
	"io"
	"os"
	"path/filepath"
	"path"
)

type Archive interface {
	Open(string) (io.ReadCloser, int64, error)
}

type TapeArchive struct {
	path string
}

type TapeArchiveEntry struct {
	io.Reader
	f *os.File
}

func (t *TapeArchiveEntry) Close() error {
	return t.f.Close()
}

func (t *TapeArchive) Open(name string) (io.ReadCloser, int64, error) {
	f, err := os.Open(t.path)
	if err != nil {
		return nil, 0, err
	}

	r := tar.NewReader(f)

	for {
		h, err := r.Next()

		if err == io.EOF {
			break
		}

		matched, err := filepath.Match(name, path.Base(h.Name))
		if err != nil {
			return nil, 0, err
		}

		if matched {
			return &TapeArchiveEntry{r, f}, h.Size, nil
		}
	}

	_ = f.Close()

	return nil, 0, os.ErrNotExist
}

type FileArchive struct {
	path string
}

func (t *FileArchive) Open(name string) (io.ReadCloser, int64, error) {
	fpath := name

	if name != t.path {
		fpath = filepath.Join(filepath.Dir(t.path), name)
	}

	s, err := os.Stat(fpath)
	if err != nil {
		return nil, 0, err
	}

	f, err := os.Open(fpath)
	if err != nil {
		return nil, 0, err
	}

	return f, s.Size(), nil
}
<|endoftext|>


<commit_before>package tcp

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"
	"sync/atomic"
	"time"

	"go.uber.org/zap"

	"github.com/klauspost/compress/gzip"
	"github.com/klauspost/compress/snappy"
	"github.com/lomik/go-carbon/helper"
	"github.com/lomik/go-carbon/points"
	"github.com/lomik/go-carbon/receiver"
	"github.com/lomik/go-carbon/receiver/parse"
	"github.com/lomik/graphite-pickle/framing"
	"github.com/lomik/zapwriter"
	"github.com/prometheus/client_golang/prometheus"
)

func init() {
	receiver.Register(
		"tcp",
		func() interface{} { return NewOptions() },
		func(name string, options interface{}, store func(*points.Points)) (receiver.Receiver, error) {
			return newTCP(name, options.(*Options), store)
		},
	)

	receiver.Register(
		"pickle",
		func() interface{} { return NewFramingOptions() },
		func(name string, options interface{}, store func(*points.Points)) (receiver.Receiver, error) {
			return newFraming("pickle", name, options.(*FramingOptions), store)
		},
	)

	receiver.Register(
		"protobuf",
		func() interface{} { return NewFramingOptions() },
		func(name string, options interface{}, store func(*points.Points)) (receiver.Receiver, error) {
			return newFraming("protobuf", name, options.(*FramingOptions), store)
		},
	)
}

type Options struct {
	Listen      string `toml:"listen"`
	Enabled     bool   `toml:"enabled"`
	BufferSize  int    `toml:"buffer-size"`
	Compression string `toml:"compression"`
}

func NewOptions() *Options {
	return &Options{
		Listen:     ":2003",
		Enabled:    true,
		BufferSize: 0,
	}
}

type FramingOptions struct {
	Listen         string `toml:"listen"`
	MaxMessageSize uint32 `toml:"max-message-size"`
	Enabled        bool   `toml:"enabled"`
	BufferSize     int    `toml:"buffer-size"`
}

func NewFramingOptions() *FramingOptions {
	return &FramingOptions{
		Listen:         ":2004",
		MaxMessageSize: 67108864, // 64 Mb
		Enabled:        true,
		BufferSize:     0,
	}
}

// TCP receive metrics from TCP connections
type TCP struct {
	helper.Stoppable
	out                 func(*points.Points)
	name                string // name for store metrics
	maxMessageSize      uint32
	metricsReceived     uint32
	promMetricsReceived prometheus.Gauge
	errors              uint32
	active              int32 // counter
	listener            *net.TCPListener
	isFraming           bool
	frameParser         func(body []byte) ([]*points.Points, error)
	buffer              chan *points.Points
	logger              *zap.Logger
	decompressor        decompressor
}

// Addr returns binded socket address. For bind port 0 in tests
func (rcv *TCP) Addr() net.Addr {
	if rcv.listener == nil {
		return nil
	}
	return rcv.listener.Addr()
}

func newTCP(name string, options *Options, store func(*points.Points)) (*TCP, error) {
	if !options.Enabled {
		return nil, nil
	}

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

	r := &TCP{
		out:    store,
		name:   name,
		logger: zapwriter.Logger(name),
	}

	if options.BufferSize > 0 {
		r.buffer = make(chan *points.Points, options.BufferSize)
	}

	r.decompressor = newDecompressor(options.Compression)

	err = r.Listen(addr)
	if err != nil {
		return nil, err
	}

	return r, err
}

func newFraming(parser string, name string, options *FramingOptions, store func(*points.Points)) (*TCP, error) {
	if !options.Enabled {
		return nil, nil
	}

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

	r := &TCP{
		out:            store,
		name:           name,
		logger:         zapwriter.Logger(name),
		maxMessageSize: options.MaxMessageSize,
		isFraming:      true,
	}

	switch parser {
	case "pickle":
		r.frameParser = parse.Pickle
	case "protobuf":
		r.frameParser = parse.Protobuf
	default:
		return nil, fmt.Errorf("unknown frame parser %#v", parser)
	}

	if options.BufferSize > 0 {
		r.buffer = make(chan *points.Points, options.BufferSize)
	}

	err = r.Listen(addr)
	if err != nil {
		return nil, err
	}

	return r, err
}

func (rcv *TCP) HandleConnection(conn net.Conn) {
	atomic.AddInt32(&rcv.active, 1)
	defer atomic.AddInt32(&rcv.active, -1)

	defer conn.Close()

	bconn, err := rcv.decompressor(conn)
	if err != nil {
		rcv.logger.Error("failed init decompressor", zap.Error(err))
		return
	}
	reader := bufio.NewReader(bconn)

	finished := make(chan bool)
	defer close(finished)

	rcv.Go(func(exit chan bool) {
		select {
		case <-finished:
			return
		case <-exit:
			conn.Close()
			return
		}
	})

	lastDeadline := time.Now()
	readTimeout := 2 * time.Minute
	conn.SetReadDeadline(lastDeadline.Add(readTimeout))

	for {
		now := time.Now()
		if now.Sub(lastDeadline) > (readTimeout / 4) {
			conn.SetReadDeadline(now.Add(readTimeout))
			lastDeadline = now
		}

		line, err := reader.ReadBytes('\n')

		if err != nil {
			if err == io.EOF {
				if len(line) > 0 {
					rcv.logger.Warn("unfinished line", zap.String("line", string(line)))
				}
			} else {
				atomic.AddUint32(&rcv.errors, 1)
				rcv.logger.Error("read error", zap.Error(err))
			}
			break
		}
		if len(line) > 0 { // skip empty lines
			name, value, timestamp, err := parse.PlainLine(line)
			if err != nil {
				atomic.AddUint32(&rcv.errors, 1)
				rcv.logger.Info("parse failed",
					zap.Error(err),
					zap.String("peer", conn.RemoteAddr().String()),
				)
			} else {
				atomic.AddUint32(&rcv.metricsReceived, 1)
				rcv.out(points.OnePoint(string(name), value, timestamp))
			}
		}
	}
}

func (rcv *TCP) handleFraming(conn net.Conn) {
	framedConn, _ := framing.NewConn(conn, byte(4), binary.BigEndian)
	defer func() {
		if r := recover(); r != nil {
			rcv.logger.Error("panic recovered", zap.String("traceback", fmt.Sprint(r)))
		}
	}()

	atomic.AddInt32(&rcv.active, 1)
	defer atomic.AddInt32(&rcv.active, -1)

	defer conn.Close()

	finished := make(chan bool)
	defer close(finished)

	rcv.Go(func(exit chan bool) {
		select {
		case <-finished:
			return
		case <-exit:
			conn.Close()
			return
		}
	})

	framedConn.MaxFrameSize = uint(rcv.maxMessageSize)

	for {
		conn.SetReadDeadline(time.Now().Add(2 * time.Minute))
		data, err := framedConn.ReadFrame()
		if err == io.EOF {
			return
		} else if err == framing.ErrPrefixLength {
			atomic.AddUint32(&rcv.errors, 1)
			rcv.logger.Warn("bad message size")
			return
		} else if err != nil {
			atomic.AddUint32(&rcv.errors, 1)
			rcv.logger.Warn("can't read message body", zap.Error(err))
			return
		}

		msgs, err := rcv.frameParser(data)

		if err != nil {
			atomic.AddUint32(&rcv.errors, 1)
			rcv.logger.Info("can't parse message",
				zap.String("data", string(data)),
				zap.Error(err),
			)
			return
		}

		for _, msg := range msgs {
			atomic.AddUint32(&rcv.metricsReceived, uint32(len(msg.Data)))
			rcv.out(msg)
		}
	}
}

func (rcv *TCP) Stat(send helper.StatCallback) {
	metricsReceived := atomic.LoadUint32(&rcv.metricsReceived)
	atomic.AddUint32(&rcv.metricsReceived, -metricsReceived)
	send("metricsReceived", float64(metricsReceived))
	if rcv.promMetricsReceived != nil {
		rcv.promMetricsReceived.Set(float64(metricsReceived))
	}

	active := float64(atomic.LoadInt32(&rcv.active))
	send("active", active)

	errors := atomic.LoadUint32(&rcv.errors)
	atomic.AddUint32(&rcv.errors, -errors)
	send("errors", float64(errors))

	if rcv.buffer != nil {
		send("bufferLen", float64(len(rcv.buffer)))
		send("bufferCap", float64(cap(rcv.buffer)))
	}
}

// Listen bind port. Receive messages and send to out channel
func (rcv *TCP) Listen(addr *net.TCPAddr) error {
	return rcv.StartFunc(func() error {

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

		rcv.Go(func(exit chan bool) {
			<-exit
			tcpListener.Close()
		})

		handler := rcv.HandleConnection
		if rcv.isFraming {
			handler = rcv.handleFraming
		}

		if rcv.buffer != nil {
			originalOut := rcv.out

			rcv.Go(func(exit chan bool) {
				for {
					select {
					case <-exit:
						return
					case p := <-rcv.buffer:
						originalOut(p)
					}
				}
			})

			rcv.out = func(p *points.Points) {
				rcv.buffer <- p
			}
		}

		rcv.Go(func(exit chan bool) {
			defer tcpListener.Close()

			for {

				conn, err := tcpListener.Accept()
				if err != nil {
					if strings.Contains(err.Error(), "use of closed network connection") {
						break
					}
					rcv.logger.Warn("failed to accept connection",
						zap.Error(err),
					)
					continue
				}

				rcv.Go(func(exit chan bool) {
					handler(conn)
				})
			}

		})

		rcv.listener = tcpListener

		return nil
	})
}

type decompressor func(net.Conn) (io.Reader, error)

func newDecompressor(typ string) decompressor {
	switch typ {
	case "snappy":
		return func(c net.Conn) (io.Reader, error) {
			return snappy.NewReader(c), nil
		}
	case "gzip":
		return func(c net.Conn) (io.Reader, error) {
			return gzip.NewReader(c)
		}
	default:
		return func(c net.Conn) (io.Reader, error) {
			return c, nil
		}
	}
}

func (rcv *TCP) InitPrometheus(reg prometheus.Registerer) {
	rcv.promMetricsReceived = prometheus.NewGauge(
		prometheus.GaugeOpts{
			Name: "metrics_received_tcp",
			Help: "How many metrics were received by the TCP endpoint.",
		},
	)

	reg.MustRegister(rcv.promMetricsReceived)
}
<commit_msg>Exchanged gauge for a counter.<commit_after>package tcp

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"
	"sync/atomic"
	"time"

	"go.uber.org/zap"

	"github.com/klauspost/compress/gzip"
	"github.com/klauspost/compress/snappy"
	"github.com/lomik/go-carbon/helper"
	"github.com/lomik/go-carbon/points"
	"github.com/lomik/go-carbon/receiver"
	"github.com/lomik/go-carbon/receiver/parse"
	"github.com/lomik/graphite-pickle/framing"
	"github.com/lomik/zapwriter"
	"github.com/prometheus/client_golang/prometheus"
)

func init() {
	receiver.Register(
		"tcp",
		func() interface{} { return NewOptions() },
		func(name string, options interface{}, store func(*points.Points)) (receiver.Receiver, error) {
			return newTCP(name, options.(*Options), store)
		},
	)

	receiver.Register(
		"pickle",
		func() interface{} { return NewFramingOptions() },
		func(name string, options interface{}, store func(*points.Points)) (receiver.Receiver, error) {
			return newFraming("pickle", name, options.(*FramingOptions), store)
		},
	)

	receiver.Register(
		"protobuf",
		func() interface{} { return NewFramingOptions() },
		func(name string, options interface{}, store func(*points.Points)) (receiver.Receiver, error) {
			return newFraming("protobuf", name, options.(*FramingOptions), store)
		},
	)
}

type Options struct {
	Listen      string `toml:"listen"`
	Enabled     bool   `toml:"enabled"`
	BufferSize  int    `toml:"buffer-size"`
	Compression string `toml:"compression"`
}

func NewOptions() *Options {
	return &Options{
		Listen:     ":2003",
		Enabled:    true,
		BufferSize: 0,
	}
}

type FramingOptions struct {
	Listen         string `toml:"listen"`
	MaxMessageSize uint32 `toml:"max-message-size"`
	Enabled        bool   `toml:"enabled"`
	BufferSize     int    `toml:"buffer-size"`
}

func NewFramingOptions() *FramingOptions {
	return &FramingOptions{
		Listen:         ":2004",
		MaxMessageSize: 67108864, // 64 Mb
		Enabled:        true,
		BufferSize:     0,
	}
}

// TCP receive metrics from TCP connections
type TCP struct {
	helper.Stoppable
	out                 func(*points.Points)
	name                string // name for store metrics
	maxMessageSize      uint32
	metricsReceived     uint32
	promMetricsReceived prometheus.Counter
	errors              uint32
	active              int32 // counter
	listener            *net.TCPListener
	isFraming           bool
	frameParser         func(body []byte) ([]*points.Points, error)
	buffer              chan *points.Points
	logger              *zap.Logger
	decompressor        decompressor
}

// Addr returns binded socket address. For bind port 0 in tests
func (rcv *TCP) Addr() net.Addr {
	if rcv.listener == nil {
		return nil
	}
	return rcv.listener.Addr()
}

func newTCP(name string, options *Options, store func(*points.Points)) (*TCP, error) {
	if !options.Enabled {
		return nil, nil
	}

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

	r := &TCP{
		out:    store,
		name:   name,
		logger: zapwriter.Logger(name),
	}

	if options.BufferSize > 0 {
		r.buffer = make(chan *points.Points, options.BufferSize)
	}

	r.decompressor = newDecompressor(options.Compression)

	err = r.Listen(addr)
	if err != nil {
		return nil, err
	}

	return r, err
}

func newFraming(parser string, name string, options *FramingOptions, store func(*points.Points)) (*TCP, error) {
	if !options.Enabled {
		return nil, nil
	}

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

	r := &TCP{
		out:            store,
		name:           name,
		logger:         zapwriter.Logger(name),
		maxMessageSize: options.MaxMessageSize,
		isFraming:      true,
	}

	switch parser {
	case "pickle":
		r.frameParser = parse.Pickle
	case "protobuf":
		r.frameParser = parse.Protobuf
	default:
		return nil, fmt.Errorf("unknown frame parser %#v", parser)
	}

	if options.BufferSize > 0 {
		r.buffer = make(chan *points.Points, options.BufferSize)
	}

	err = r.Listen(addr)
	if err != nil {
		return nil, err
	}

	return r, err
}

func (rcv *TCP) HandleConnection(conn net.Conn) {
	atomic.AddInt32(&rcv.active, 1)
	defer atomic.AddInt32(&rcv.active, -1)

	defer conn.Close()

	bconn, err := rcv.decompressor(conn)
	if err != nil {
		rcv.logger.Error("failed init decompressor", zap.Error(err))
		return
	}
	reader := bufio.NewReader(bconn)

	finished := make(chan bool)
	defer close(finished)

	rcv.Go(func(exit chan bool) {
		select {
		case <-finished:
			return
		case <-exit:
			conn.Close()
			return
		}
	})

	lastDeadline := time.Now()
	readTimeout := 2 * time.Minute
	conn.SetReadDeadline(lastDeadline.Add(readTimeout))

	for {
		now := time.Now()
		if now.Sub(lastDeadline) > (readTimeout / 4) {
			conn.SetReadDeadline(now.Add(readTimeout))
			lastDeadline = now
		}

		line, err := reader.ReadBytes('\n')

		if err != nil {
			if err == io.EOF {
				if len(line) > 0 {
					rcv.logger.Warn("unfinished line", zap.String("line", string(line)))
				}
			} else {
				atomic.AddUint32(&rcv.errors, 1)
				rcv.logger.Error("read error", zap.Error(err))
			}
			break
		}
		if len(line) > 0 { // skip empty lines
			name, value, timestamp, err := parse.PlainLine(line)
			if err != nil {
				atomic.AddUint32(&rcv.errors, 1)
				rcv.logger.Info("parse failed",
					zap.Error(err),
					zap.String("peer", conn.RemoteAddr().String()),
				)
			} else {
				atomic.AddUint32(&rcv.metricsReceived, 1)
				rcv.out(points.OnePoint(string(name), value, timestamp))
			}
		}
	}
}

func (rcv *TCP) handleFraming(conn net.Conn) {
	framedConn, _ := framing.NewConn(conn, byte(4), binary.BigEndian)
	defer func() {
		if r := recover(); r != nil {
			rcv.logger.Error("panic recovered", zap.String("traceback", fmt.Sprint(r)))
		}
	}()

	atomic.AddInt32(&rcv.active, 1)
	defer atomic.AddInt32(&rcv.active, -1)

	defer conn.Close()

	finished := make(chan bool)
	defer close(finished)

	rcv.Go(func(exit chan bool) {
		select {
		case <-finished:
			return
		case <-exit:
			conn.Close()
			return
		}
	})

	framedConn.MaxFrameSize = uint(rcv.maxMessageSize)

	for {
		conn.SetReadDeadline(time.Now().Add(2 * time.Minute))
		data, err := framedConn.ReadFrame()
		if err == io.EOF {
			return
		} else if err == framing.ErrPrefixLength {
			atomic.AddUint32(&rcv.errors, 1)
			rcv.logger.Warn("bad message size")
			return
		} else if err != nil {
			atomic.AddUint32(&rcv.errors, 1)
			rcv.logger.Warn("can't read message body", zap.Error(err))
			return
		}

		msgs, err := rcv.frameParser(data)

		if err != nil {
			atomic.AddUint32(&rcv.errors, 1)
			rcv.logger.Info("can't parse message",
				zap.String("data", string(data)),
				zap.Error(err),
			)
			return
		}

		for _, msg := range msgs {
			atomic.AddUint32(&rcv.metricsReceived, uint32(len(msg.Data)))
			rcv.out(msg)
		}
	}
}

func (rcv *TCP) Stat(send helper.StatCallback) {
	metricsReceived := atomic.LoadUint32(&rcv.metricsReceived)
	atomic.AddUint32(&rcv.metricsReceived, -metricsReceived)
	send("metricsReceived", float64(metricsReceived))
	if rcv.promMetricsReceived != nil {
		rcv.promMetricsReceived.Add(float64(metricsReceived))
	}

	active := float64(atomic.LoadInt32(&rcv.active))
	send("active", active)

	errors := atomic.LoadUint32(&rcv.errors)
	atomic.AddUint32(&rcv.errors, -errors)
	send("errors", float64(errors))

	if rcv.buffer != nil {
		send("bufferLen", float64(len(rcv.buffer)))
		send("bufferCap", float64(cap(rcv.buffer)))
	}
}

// Listen bind port. Receive messages and send to out channel
func (rcv *TCP) Listen(addr *net.TCPAddr) error {
	return rcv.StartFunc(func() error {

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

		rcv.Go(func(exit chan bool) {
			<-exit
			tcpListener.Close()
		})

		handler := rcv.HandleConnection
		if rcv.isFraming {
			handler = rcv.handleFraming
		}

		if rcv.buffer != nil {
			originalOut := rcv.out

			rcv.Go(func(exit chan bool) {
				for {
					select {
					case <-exit:
						return
					case p := <-rcv.buffer:
						originalOut(p)
					}
				}
			})

			rcv.out = func(p *points.Points) {
				rcv.buffer <- p
			}
		}

		rcv.Go(func(exit chan bool) {
			defer tcpListener.Close()

			for {

				conn, err := tcpListener.Accept()
				if err != nil {
					if strings.Contains(err.Error(), "use of closed network connection") {
						break
					}
					rcv.logger.Warn("failed to accept connection",
						zap.Error(err),
					)
					continue
				}

				rcv.Go(func(exit chan bool) {
					handler(conn)
				})
			}

		})

		rcv.listener = tcpListener

		return nil
	})
}

type decompressor func(net.Conn) (io.Reader, error)

func newDecompressor(typ string) decompressor {
	switch typ {
	case "snappy":
		return func(c net.Conn) (io.Reader, error) {
			return snappy.NewReader(c), nil
		}
	case "gzip":
		return func(c net.Conn) (io.Reader, error) {
			return gzip.NewReader(c)
		}
	default:
		return func(c net.Conn) (io.Reader, error) {
			return c, nil
		}
	}
}

func (rcv *TCP) InitPrometheus(reg prometheus.Registerer) {
	rcv.promMetricsReceived = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "metrics_received_tcp_total",
			Help: "Counter of metrics received via the TCP endpoint.",
		},
	)

	reg.MustRegister(rcv.promMetricsReceived)
}
<|endoftext|>


<commit_before>package ds

import (
	"strconv"

	"cloud.google.com/go/datastore"
)

// KeyGetter interface
type KeyGetter interface {
	Key() *datastore.Key
}

// KeySetter interface
type KeySetter interface {
	SetKey(*datastore.Key)
}

// KeyGetSetter interface
type KeyGetSetter interface {
	KeyGetter
	KeySetter
	NewKey(string)
}

// Model is the base model which id is int64
type Model struct {
	key *datastore.Key
	id  int64
}

// Key returns key from model
func (x *Model) Key() *datastore.Key {
	if x == nil {
		return nil
	}
	return x.key
}

// SetKey sets model key to given key
func (x *Model) SetKey(key *datastore.Key) {
	if x == nil {
		return
	}
	x.key = key
	if key == nil {
		x.id = 0
		return
	}
	x.id = key.ID
}

// SetID sets id key to model
func (x *Model) SetID(kind string, id int64) {
	x.SetKey(datastore.IDKey(kind, id, nil))
}

// ID returns id
func (x *Model) ID() int64 {
	if x == nil {
		return 0
	}
	return x.id
}

// NewKey sets incomplete key to model
func (x *Model) NewKey(kind string) {
	x.SetKey(datastore.IncompleteKey(kind, nil))
}

// StringIDModel is the base model which id is string
// but can use both id key and name key
type StringIDModel struct {
	key *datastore.Key
	id  string
}

// Key returns key from model
func (x *StringIDModel) Key() *datastore.Key {
	if x == nil {
		return nil
	}
	return x.key
}

// SetKey sets model key to given key
// if key is not name key, it will use id key
func (x *StringIDModel) SetKey(key *datastore.Key) {
	if x == nil {
		return
	}
	x.key = key
	x.id = ""
	if key == nil {
		return
	}
	if key.Name != "" {
		x.id = key.Name
		return
	}
	if key.ID != 0 {
		x.id = strconv.FormatInt(key.ID, 10)
	}
}

// SetID sets id to model
func (x *StringIDModel) SetID(kind string, id int64) {
	if id == 0 {
		// invalid key
		// if set id key to 0, datastore server will throw error, which we can not handle
		return
	}
	x.SetKey(datastore.IDKey(kind, id, nil))
}

// SetStringID sets string id to model
func (x *StringIDModel) SetStringID(kind string, id string) {
	x.SetID(kind, parseID(id))
}

// SetNameID sets name id to model
func (x *StringIDModel) SetNameID(kind string, name string) {
	x.SetKey(datastore.NameKey(kind, name, nil))
}

// NewKey sets incomplete key to model
func (x *StringIDModel) NewKey(kind string) {
	x.SetKey(datastore.IncompleteKey(kind, nil))
}

// ID return id
func (x *StringIDModel) ID() string {
	if x == nil {
		return ""
	}
	return x.id
}

// SetKey sets key to model
func SetKey(key *datastore.Key, dst interface{}) {
	if dst == nil || key == nil {
		return
	}
	if x, ok := dst.(KeySetter); ok {
		x.SetKey(key)
	}
}

// SetKeys sets keys to models
func SetKeys(keys []*datastore.Key, dst interface{}) {
	if dst == nil || keys == nil {
		return
	}
	xs := valueOf(dst)
	for i := 0; i < xs.Len(); i++ {
		x := xs.Index(i)
		if x.IsNil() {
			continue
		}
		if x, ok := x.Interface().(KeySetter); ok {
			x.SetKey(keys[i])
		}
	}
}

// SetCommitKey sets commit pending key to model
func SetCommitKey(commit *datastore.Commit, pendingKey *datastore.PendingKey, dst interface{}) {
	if dst == nil {
		return
	}
	if x, ok := dst.(KeySetter); ok {
		x.SetKey(commit.Key(pendingKey))
	}
}

// SetCommitKeys sets commit pending keys to models
func SetCommitKeys(commit *datastore.Commit, pendingKeys []*datastore.PendingKey, dst interface{}) {
	xs := valueOf(dst)
	for i := 0; i < xs.Len(); i++ {
		x := xs.Index(i)
		if x.IsNil() {
			continue
		}
		if x, ok := x.Interface().(KeySetter); ok {
			x.SetKey(commit.Key(pendingKeys[i]))
		}
	}
}

// SetID sets id to model
func SetID(kind string, id int64, dst interface{}) {
	SetKey(datastore.IDKey(kind, id, nil), dst)
}

// SetIDs sets ids to models
func SetIDs(kind string, ids []int64, dst interface{}) {
	keys := make([]*datastore.Key, len(ids))
	for i := range ids {
		keys[i] = datastore.IDKey(kind, ids[i], nil)
	}
	SetKeys(keys, dst)
}

// SetStringID sets string id to model
func SetStringID(kind string, id string, dst interface{}) {
	tid := parseID(id)
	if tid == 0 {
		return
	}
	SetKey(datastore.IDKey(kind, tid, nil), dst)
}

// SetStringIDs sets string id to models
func SetStringIDs(kind string, ids []string, dst interface{}) {
	keys := make([]*datastore.Key, len(ids))
	for i := range ids {
		keys[i] = datastore.IDKey(kind, parseID(ids[i]), nil)
	}
	SetKeys(keys, dst)
}

// SetNameIDs sets name id to models
func SetNameIDs(kind string, names []string, dst interface{}) {
	keys := make([]*datastore.Key, len(names))
	for i := range names {
		keys[i] = datastore.NameKey(kind, names[i], nil)
	}
	SetKeys(keys, dst)
}
<commit_msg>add string id to model<commit_after>package ds

import (
	"strconv"

	"cloud.google.com/go/datastore"
)

// KeyGetter interface
type KeyGetter interface {
	Key() *datastore.Key
}

// KeySetter interface
type KeySetter interface {
	SetKey(*datastore.Key)
}

// KeyGetSetter interface
type KeyGetSetter interface {
	KeyGetter
	KeySetter
	NewKey(string)
}

// Model is the base model which id is int64
type Model struct {
	key *datastore.Key
	id  int64
}

// Key returns key from model
func (x *Model) Key() *datastore.Key {
	if x == nil {
		return nil
	}
	return x.key
}

// SetKey sets model key to given key
func (x *Model) SetKey(key *datastore.Key) {
	if x == nil {
		return
	}
	x.key = key
	if key == nil {
		x.id = 0
		return
	}
	x.id = key.ID
}

// SetID sets id key to model
func (x *Model) SetID(kind string, id int64) {
	x.SetKey(datastore.IDKey(kind, id, nil))
}

// ID returns id
func (x *Model) ID() int64 {
	if x == nil {
		return 0
	}
	return x.id
}

// StringID return id in string format
func (x *Model) StringID() string {
	if x == nil {
		return ""
	}
	return strconv.FormatInt(x.id, 10)
}

// NewKey sets incomplete key to model
func (x *Model) NewKey(kind string) {
	x.SetKey(datastore.IncompleteKey(kind, nil))
}

// StringIDModel is the base model which id is string
// but can use both id key and name key
type StringIDModel struct {
	key *datastore.Key
	id  string
}

// Key returns key from model
func (x *StringIDModel) Key() *datastore.Key {
	if x == nil {
		return nil
	}
	return x.key
}

// SetKey sets model key to given key
// if key is not name key, it will use id key
func (x *StringIDModel) SetKey(key *datastore.Key) {
	if x == nil {
		return
	}
	x.key = key
	x.id = ""
	if key == nil {
		return
	}
	if key.Name != "" {
		x.id = key.Name
		return
	}
	if key.ID != 0 {
		x.id = strconv.FormatInt(key.ID, 10)
	}
}

// SetID sets id to model
func (x *StringIDModel) SetID(kind string, id int64) {
	if id == 0 {
		// invalid key
		// if set id key to 0, datastore server will throw error, which we can not handle
		return
	}
	x.SetKey(datastore.IDKey(kind, id, nil))
}

// SetStringID sets string id to model
func (x *StringIDModel) SetStringID(kind string, id string) {
	x.SetID(kind, parseID(id))
}

// SetNameID sets name id to model
func (x *StringIDModel) SetNameID(kind string, name string) {
	x.SetKey(datastore.NameKey(kind, name, nil))
}

// NewKey sets incomplete key to model
func (x *StringIDModel) NewKey(kind string) {
	x.SetKey(datastore.IncompleteKey(kind, nil))
}

// ID return id
func (x *StringIDModel) ID() string {
	if x == nil {
		return ""
	}
	return x.id
}

// SetKey sets key to model
func SetKey(key *datastore.Key, dst interface{}) {
	if dst == nil || key == nil {
		return
	}
	if x, ok := dst.(KeySetter); ok {
		x.SetKey(key)
	}
}

// SetKeys sets keys to models
func SetKeys(keys []*datastore.Key, dst interface{}) {
	if dst == nil || keys == nil {
		return
	}
	xs := valueOf(dst)
	for i := 0; i < xs.Len(); i++ {
		x := xs.Index(i)
		if x.IsNil() {
			continue
		}
		if x, ok := x.Interface().(KeySetter); ok {
			x.SetKey(keys[i])
		}
	}
}

// SetCommitKey sets commit pending key to model
func SetCommitKey(commit *datastore.Commit, pendingKey *datastore.PendingKey, dst interface{}) {
	if dst == nil {
		return
	}
	if x, ok := dst.(KeySetter); ok {
		x.SetKey(commit.Key(pendingKey))
	}
}

// SetCommitKeys sets commit pending keys to models
func SetCommitKeys(commit *datastore.Commit, pendingKeys []*datastore.PendingKey, dst interface{}) {
	xs := valueOf(dst)
	for i := 0; i < xs.Len(); i++ {
		x := xs.Index(i)
		if x.IsNil() {
			continue
		}
		if x, ok := x.Interface().(KeySetter); ok {
			x.SetKey(commit.Key(pendingKeys[i]))
		}
	}
}

// SetID sets id to model
func SetID(kind string, id int64, dst interface{}) {
	SetKey(datastore.IDKey(kind, id, nil), dst)
}

// SetIDs sets ids to models
func SetIDs(kind string, ids []int64, dst interface{}) {
	keys := make([]*datastore.Key, len(ids))
	for i := range ids {
		keys[i] = datastore.IDKey(kind, ids[i], nil)
	}
	SetKeys(keys, dst)
}

// SetStringID sets string id to model
func SetStringID(kind string, id string, dst interface{}) {
	tid := parseID(id)
	if tid == 0 {
		return
	}
	SetKey(datastore.IDKey(kind, tid, nil), dst)
}

// SetStringIDs sets string id to models
func SetStringIDs(kind string, ids []string, dst interface{}) {
	keys := make([]*datastore.Key, len(ids))
	for i := range ids {
		keys[i] = datastore.IDKey(kind, parseID(ids[i]), nil)
	}
	SetKeys(keys, dst)
}

// SetNameIDs sets name id to models
func SetNameIDs(kind string, names []string, dst interface{}) {
	keys := make([]*datastore.Key, len(names))
	for i := range names {
		keys[i] = datastore.NameKey(kind, names[i], nil)
	}
	SetKeys(keys, dst)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/jroimartin/gocui"
)

// List of all modes
const (
	modeInsert = iota
	modeOverwrite
	modeEscape
	modeConnect
)

var notInsert = map[int]bool{
	modeEscape: true,
}

var modeChars = []struct {
	c  rune
	fg gocui.Attribute
	bg gocui.Attribute
}{
	{' ', 0, gocui.ColorGreen},
	{'R', gocui.ColorWhite | gocui.AttrBold, gocui.ColorGreen},
	{' ', 0, gocui.ColorRed},
	{'c', gocui.ColorWhite | gocui.AttrBold, gocui.ColorRed},
}

func modeBox(g *gocui.Gui) {
	maxX, maxY := g.Size()

	for i := 0; i < maxX; i++ {
		g.SetRune(i, maxY-2, '─', gocui.ColorWhite, gocui.ColorBlack)
	}

	ch := modeChars[state.Mode]
	g.SetRune(0, maxY-1, ch.c, ch.fg, ch.bg)
	g.SetRune(1, maxY-1, ' ', gocui.ColorBlack, 0)
}
<commit_msg>remove fg in modeChars<commit_after>package main

import (
	"github.com/jroimartin/gocui"
)

// List of all modes
const (
	modeInsert = iota
	modeOverwrite
	modeEscape
	modeConnect
)

var notInsert = map[int]bool{
	modeEscape: true,
}

var modeChars = []struct {
	c  rune
	bg gocui.Attribute
}{
	{' ', gocui.ColorGreen},
	{'R', gocui.ColorGreen},
	{' ', gocui.ColorRed},
	{'c', gocui.ColorRed},
}

func modeBox(g *gocui.Gui) {
	maxX, maxY := g.Size()

	for i := 0; i < maxX; i++ {
		g.SetRune(i, maxY-2, '─', gocui.ColorWhite, gocui.ColorBlack)
	}

	ch := modeChars[state.Mode]
	g.SetRune(0, maxY-1, ch.c, gocui.ColorWhite|gocui.AttrBold, ch.bg)
	g.SetRune(1, maxY-1, ' ', gocui.ColorBlack, 0)
}
<|endoftext|>