<commit_before>// +skip_license_check

/*
This file contains portions of code directly taken from the 'xenolf/lego' project.
A copy of the license for this code can be found in the file named LICENSE in
this directory.
*/

// Package clouddns implements a DNS provider for solving the DNS-01
// challenge using Google Cloud DNS.
package clouddns

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

	logf "github.com/jetstack/cert-manager/pkg/logs"

	"github.com/go-logr/logr"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/dns/v1"
	"google.golang.org/api/option"

	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
)

// DNSProvider is an implementation of the DNSProvider interface.
type DNSProvider struct {
	hostedZoneName   string
	dns01Nameservers []string
	project          string
	client           *dns.Service
	log              logr.Logger
}

func NewDNSProvider(project string, saBytes []byte, dns01Nameservers []string, ambient bool, hostedZoneName string) (*DNSProvider, error) {
	// project is a required field
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}
	// if the service account bytes are not provided, we will attempt to instantiate
	// with 'ambient credentials' (if they are allowed/enabled)
	if len(saBytes) == 0 {
		if !ambient {
			return nil, fmt.Errorf("unable to construct clouddns provider: empty credentials; perhaps you meant to enable ambient credentials?")
		}
		return NewDNSProviderCredentials(project, dns01Nameservers, hostedZoneName)
	}
	// if service account data is provided, we instantiate using that
	if len(saBytes) != 0 {
		return NewDNSProviderServiceAccountBytes(project, saBytes, dns01Nameservers, hostedZoneName)
	}
	return nil, fmt.Errorf("missing Google Cloud DNS provider credentials")
}

// NewDNSProviderEnvironment returns a DNSProvider instance configured for Google Cloud
// DNS. Project name must be passed in the environment variable: GCE_PROJECT.
// A Service Account file can be passed in the environment variable:
// GCE_SERVICE_ACCOUNT_FILE
func NewDNSProviderEnvironment(dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	project := os.Getenv("GCE_PROJECT")
	if saFile, ok := os.LookupEnv("GCE_SERVICE_ACCOUNT_FILE"); ok {
		return NewDNSProviderServiceAccount(project, saFile, dns01Nameservers, hostedZoneName)
	}
	return NewDNSProviderCredentials(project, dns01Nameservers, hostedZoneName)
}

// NewDNSProviderCredentials uses the supplied credentials to return a
// DNSProvider instance configured for Google Cloud DNS.
func NewDNSProviderCredentials(project string, dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}

	ctx := context.Background()
	client, err := google.DefaultClient(ctx, dns.NdevClouddnsReadwriteScope)
	if err != nil {
		return nil, fmt.Errorf("Unable to get Google Cloud client: %v", err)
	}
	svc, err := dns.NewService(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, fmt.Errorf("Unable to create Google Cloud DNS service: %v", err)
	}
	return &DNSProvider{
		project:          project,
		client:           svc,
		dns01Nameservers: dns01Nameservers,
		hostedZoneName:   hostedZoneName,
		log:              logf.Log.WithName("clouddns"),
	}, nil
}

// NewDNSProviderServiceAccount uses the supplied service account JSON file to
// return a DNSProvider instance configured for Google Cloud DNS.
func NewDNSProviderServiceAccount(project string, saFile string, dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}
	if saFile == "" {
		return nil, fmt.Errorf("Google Cloud Service Account file missing")
	}

	dat, err := ioutil.ReadFile(saFile)
	if err != nil {
		return nil, fmt.Errorf("Unable to read Service Account file: %v", err)
	}
	return NewDNSProviderServiceAccountBytes(project, dat, dns01Nameservers, hostedZoneName)
}

// NewDNSProviderServiceAccountBytes uses the supplied service account JSON
// file data to return a DNSProvider instance configured for Google Cloud DNS.
func NewDNSProviderServiceAccountBytes(project string, saBytes []byte, dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}
	if len(saBytes) == 0 {
		return nil, fmt.Errorf("Google Cloud Service Account data missing")
	}

	conf, err := google.JWTConfigFromJSON(saBytes, dns.NdevClouddnsReadwriteScope)
	if err != nil {
		return nil, fmt.Errorf("Unable to acquire config: %v", err)
	}

	ctx := context.Background()
	client := conf.Client(ctx)

	svc, err := dns.NewService(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, fmt.Errorf("Unable to create Google Cloud DNS service: %v", err)
	}
	return &DNSProvider{
		project:          project,
		client:           svc,
		dns01Nameservers: dns01Nameservers,
		hostedZoneName:   hostedZoneName,
	}, nil
}

// Present creates a TXT record to fulfil the dns-01 challenge.
func (c *DNSProvider) Present(domain, fqdn, value string) error {
	zone, err := c.getHostedZone(fqdn)
	if err != nil {
		return err
	}

	rec := &dns.ResourceRecordSet{
		Name:    fqdn,
		Rrdatas: []string{value},
		Ttl:     int64(60),
		Type:    "TXT",
	}
	change := &dns.Change{
		Additions: []*dns.ResourceRecordSet{rec},
	}

	// Look for existing records.
	list, err := c.client.ResourceRecordSets.List(c.project, zone).Name(fqdn).Type("TXT").Do()
	if err != nil {
		return err
	}
	if len(list.Rrsets) > 0 {
		// Attempt to delete the existing records when adding our new one.
		change.Deletions = list.Rrsets
	}

	chg, err := c.client.Changes.Create(c.project, zone, change).Do()
	if err != nil {
		return err
	}

	// wait for change to be acknowledged
	for chg.Status == "pending" {
		time.Sleep(time.Second)

		chg, err = c.client.Changes.Get(c.project, zone, chg.Id).Do()
		if err != nil {
			return err
		}
	}

	return nil
}

// CleanUp removes the TXT record matching the specified parameters.
func (c *DNSProvider) CleanUp(domain, fqdn, value string) error {
	zone, err := c.getHostedZone(fqdn)
	if err != nil {
		return err
	}

	records, err := c.findTxtRecords(zone, fqdn)
	if err != nil {
		return err
	}

	for _, rec := range records {
		change := &dns.Change{
			Deletions: []*dns.ResourceRecordSet{rec},
		}
		_, err = c.client.Changes.Create(c.project, zone, change).Do()
		if err != nil {
			return err
		}
	}
	return nil
}

// getHostedZone returns the managed-zone
func (c *DNSProvider) getHostedZone(domain string) (string, error) {
	if c.hostedZoneName != "" {
		return c.hostedZoneName, nil
	}

	authZone, err := util.FindZoneByFqdn(util.ToFqdn(domain), c.dns01Nameservers)
	if err != nil {
		return "", err
	}

	zones, err := c.client.ManagedZones.
		List(c.project).
		DnsName(authZone).
		Do()
	if err != nil {
		return "", fmt.Errorf("GoogleCloud API call failed: %v", err)
	}

	if len(zones.ManagedZones) == 0 {
		return "", fmt.Errorf("No matching GoogleCloud domain found for domain %s", authZone)
	}

	// attempt to get the first public zone
	for _, zone := range zones.ManagedZones {
		if zone.Visibility == "public" {
			return zone.Name, nil
		}
	}

	c.log.V(logf.DebugLevel).Info("No matching public GoogleCloud managed-zone for domain, falling back to a private managed-zone", authZone)
	// fall back to first available zone, if none public
	return zones.ManagedZones[0].Name, nil
}

func (c *DNSProvider) findTxtRecords(zone, fqdn string) ([]*dns.ResourceRecordSet, error) {

	recs, err := c.client.ResourceRecordSets.List(c.project, zone).Do()
	if err != nil {
		return nil, err
	}

	found := []*dns.ResourceRecordSet{}
	for _, r := range recs.Rrsets {
		if r.Type == "TXT" && r.Name == fqdn {
			found = append(found, r)
		}
	}

	return found, nil
}
<commit_msg>Fix missing log in struct<commit_after>// +skip_license_check

/*
This file contains portions of code directly taken from the 'xenolf/lego' project.
A copy of the license for this code can be found in the file named LICENSE in
this directory.
*/

// Package clouddns implements a DNS provider for solving the DNS-01
// challenge using Google Cloud DNS.
package clouddns

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

	logf "github.com/jetstack/cert-manager/pkg/logs"

	"github.com/go-logr/logr"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/dns/v1"
	"google.golang.org/api/option"

	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
)

// DNSProvider is an implementation of the DNSProvider interface.
type DNSProvider struct {
	hostedZoneName   string
	dns01Nameservers []string
	project          string
	client           *dns.Service
	log              logr.Logger
}

func NewDNSProvider(project string, saBytes []byte, dns01Nameservers []string, ambient bool, hostedZoneName string) (*DNSProvider, error) {
	// project is a required field
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}
	// if the service account bytes are not provided, we will attempt to instantiate
	// with 'ambient credentials' (if they are allowed/enabled)
	if len(saBytes) == 0 {
		if !ambient {
			return nil, fmt.Errorf("unable to construct clouddns provider: empty credentials; perhaps you meant to enable ambient credentials?")
		}
		return NewDNSProviderCredentials(project, dns01Nameservers, hostedZoneName)
	}
	// if service account data is provided, we instantiate using that
	if len(saBytes) != 0 {
		return NewDNSProviderServiceAccountBytes(project, saBytes, dns01Nameservers, hostedZoneName)
	}
	return nil, fmt.Errorf("missing Google Cloud DNS provider credentials")
}

// NewDNSProviderEnvironment returns a DNSProvider instance configured for Google Cloud
// DNS. Project name must be passed in the environment variable: GCE_PROJECT.
// A Service Account file can be passed in the environment variable:
// GCE_SERVICE_ACCOUNT_FILE
func NewDNSProviderEnvironment(dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	project := os.Getenv("GCE_PROJECT")
	if saFile, ok := os.LookupEnv("GCE_SERVICE_ACCOUNT_FILE"); ok {
		return NewDNSProviderServiceAccount(project, saFile, dns01Nameservers, hostedZoneName)
	}
	return NewDNSProviderCredentials(project, dns01Nameservers, hostedZoneName)
}

// NewDNSProviderCredentials uses the supplied credentials to return a
// DNSProvider instance configured for Google Cloud DNS.
func NewDNSProviderCredentials(project string, dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}

	ctx := context.Background()
	client, err := google.DefaultClient(ctx, dns.NdevClouddnsReadwriteScope)
	if err != nil {
		return nil, fmt.Errorf("Unable to get Google Cloud client: %v", err)
	}
	svc, err := dns.NewService(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, fmt.Errorf("Unable to create Google Cloud DNS service: %v", err)
	}
	return &DNSProvider{
		project:          project,
		client:           svc,
		dns01Nameservers: dns01Nameservers,
		hostedZoneName:   hostedZoneName,
		log:              logf.Log.WithName("clouddns"),
	}, nil
}

// NewDNSProviderServiceAccount uses the supplied service account JSON file to
// return a DNSProvider instance configured for Google Cloud DNS.
func NewDNSProviderServiceAccount(project string, saFile string, dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}
	if saFile == "" {
		return nil, fmt.Errorf("Google Cloud Service Account file missing")
	}

	dat, err := ioutil.ReadFile(saFile)
	if err != nil {
		return nil, fmt.Errorf("Unable to read Service Account file: %v", err)
	}
	return NewDNSProviderServiceAccountBytes(project, dat, dns01Nameservers, hostedZoneName)
}

// NewDNSProviderServiceAccountBytes uses the supplied service account JSON
// file data to return a DNSProvider instance configured for Google Cloud DNS.
func NewDNSProviderServiceAccountBytes(project string, saBytes []byte, dns01Nameservers []string, hostedZoneName string) (*DNSProvider, error) {
	if project == "" {
		return nil, fmt.Errorf("Google Cloud project name missing")
	}
	if len(saBytes) == 0 {
		return nil, fmt.Errorf("Google Cloud Service Account data missing")
	}

	conf, err := google.JWTConfigFromJSON(saBytes, dns.NdevClouddnsReadwriteScope)
	if err != nil {
		return nil, fmt.Errorf("Unable to acquire config: %v", err)
	}

	ctx := context.Background()
	client := conf.Client(ctx)

	svc, err := dns.NewService(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, fmt.Errorf("Unable to create Google Cloud DNS service: %v", err)
	}
	return &DNSProvider{
		project:          project,
		client:           svc,
		dns01Nameservers: dns01Nameservers,
		hostedZoneName:   hostedZoneName,
		log:              logf.Log.WithName("clouddns"),
	}, nil
}

// Present creates a TXT record to fulfil the dns-01 challenge.
func (c *DNSProvider) Present(domain, fqdn, value string) error {
	zone, err := c.getHostedZone(fqdn)
	if err != nil {
		return err
	}

	rec := &dns.ResourceRecordSet{
		Name:    fqdn,
		Rrdatas: []string{value},
		Ttl:     int64(60),
		Type:    "TXT",
	}
	change := &dns.Change{
		Additions: []*dns.ResourceRecordSet{rec},
	}

	// Look for existing records.
	list, err := c.client.ResourceRecordSets.List(c.project, zone).Name(fqdn).Type("TXT").Do()
	if err != nil {
		return err
	}
	if len(list.Rrsets) > 0 {
		// Attempt to delete the existing records when adding our new one.
		change.Deletions = list.Rrsets
	}

	chg, err := c.client.Changes.Create(c.project, zone, change).Do()
	if err != nil {
		return err
	}

	// wait for change to be acknowledged
	for chg.Status == "pending" {
		time.Sleep(time.Second)

		chg, err = c.client.Changes.Get(c.project, zone, chg.Id).Do()
		if err != nil {
			return err
		}
	}

	return nil
}

// CleanUp removes the TXT record matching the specified parameters.
func (c *DNSProvider) CleanUp(domain, fqdn, value string) error {
	zone, err := c.getHostedZone(fqdn)
	if err != nil {
		return err
	}

	records, err := c.findTxtRecords(zone, fqdn)
	if err != nil {
		return err
	}

	for _, rec := range records {
		change := &dns.Change{
			Deletions: []*dns.ResourceRecordSet{rec},
		}
		_, err = c.client.Changes.Create(c.project, zone, change).Do()
		if err != nil {
			return err
		}
	}
	return nil
}

// getHostedZone returns the managed-zone
func (c *DNSProvider) getHostedZone(domain string) (string, error) {
	if c.hostedZoneName != "" {
		return c.hostedZoneName, nil
	}

	authZone, err := util.FindZoneByFqdn(util.ToFqdn(domain), c.dns01Nameservers)
	if err != nil {
		return "", err
	}

	zones, err := c.client.ManagedZones.
		List(c.project).
		DnsName(authZone).
		Do()
	if err != nil {
		return "", fmt.Errorf("GoogleCloud API call failed: %v", err)
	}

	if len(zones.ManagedZones) == 0 {
		return "", fmt.Errorf("No matching GoogleCloud domain found for domain %s", authZone)
	}

	// attempt to get the first public zone
	for _, zone := range zones.ManagedZones {
		if zone.Visibility == "public" {
			return zone.Name, nil
		}
	}

	c.log.V(logf.DebugLevel).Info("No matching public GoogleCloud managed-zone for domain, falling back to a private managed-zone", authZone)
	// fall back to first available zone, if none public
	return zones.ManagedZones[0].Name, nil
}

func (c *DNSProvider) findTxtRecords(zone, fqdn string) ([]*dns.ResourceRecordSet, error) {

	recs, err := c.client.ResourceRecordSets.List(c.project, zone).Do()
	if err != nil {
		return nil, err
	}

	found := []*dns.ResourceRecordSet{}
	for _, r := range recs.Rrsets {
		if r.Type == "TXT" && r.Name == fqdn {
			found = append(found, r)
		}
	}

	return found, nil
}
<|endoftext|>


<commit_before>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package bundle

import (
	"reflect"
	"testing"
	"time"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/go-apiclient"
	"github.com/gojuno/minimock/v3"
	"github.com/rs/zerolog"
	"github.com/spf13/viper"
)

func TestFetchCheck(t *testing.T) {
	t.Log("Testing fetchCheck")

	zerolog.SetGlobalLevel(zerolog.Disabled)

	viper.Reset()
	viper.Set(config.KeyAPITokenKey, "foo")
	viper.Set(config.KeyAPITokenApp, "bar")
	viper.Set(config.KeyAPIURL, "baz")

	mc := minimock.NewController(t)
	client := genMockClient(mc)
	c := Bundle{client: client}

	t.Log("cid (empty)")
	{
		cid := ""
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "invalid cid (empty)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("cid (abc)")
	{
		cid := "abc"
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "invalid cid (abc)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("api error")
	{
		cid := "000"
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "unable to retrieve check bundle (/check_bundle/000): forced mock api call error" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("valid")
	{
		cid := "1234"
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err != nil {
			t.Fatalf("expected no error, got (%s)", err)
		}
	}
}

func TestFindCheck(t *testing.T) {
	t.Log("Testing findCheck")

	zerolog.SetGlobalLevel(zerolog.Disabled)

	viper.Reset()
	viper.Set(config.KeyAPITokenKey, "foo")
	viper.Set(config.KeyAPITokenApp, "bar")
	viper.Set(config.KeyAPIURL, "baz")

	mc := minimock.NewController(t)
	client := genMockClient(mc)
	c := Bundle{client: client}

	t.Log("target (empty)")
	{

		target := ""
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != -1 {
			t.Fatal("expected found == -1")
		}

		if err.Error() != "invalid check bundle target (empty)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("api error")
	{
		target := "000"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != -1 {
			t.Fatal("expected found == -1")
		}

		if err.Error() != "searching for check bundle: forced mock api call error" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("not found")
	{
		target := "not_found"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != 0 {
			t.Fatal("expected found == 0")
		}

		if err.Error() != `no check bundles matched criteria ((active:1)(type:"json:nad")(target:"not_found"))` {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("multiple")
	{
		target := "multiple"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != 2 {
			t.Fatal("expected found == 2")
		}

		if err.Error() != `more than one (2) check bundle matched criteria ((active:1)(type:"json:nad")(target:"multiple"))` {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("valid")
	{
		target := "valid"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err != nil {
			t.Fatalf("expected no error, got (%s)", err)
		}
		if found != 1 {
			t.Fatal("expected found == 1")
		}
	}
}

func TestCreateCheck(t *testing.T) {
	t.Log("Testing createCheck")

	zerolog.SetGlobalLevel(zerolog.Disabled)

	viper.Reset()
	viper.Set(config.KeyAPITokenKey, "foo")
	viper.Set(config.KeyAPITokenApp, "bar")
	viper.Set(config.KeyAPIURL, "baz")

	mc := minimock.NewController(t)
	client := genMockClient(mc)
	c := Bundle{client: client}

	t.Log("target (empty)")
	{
		target := ""
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.createCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "invalid check bundle target (empty)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}
}

func TestBundle_CID(t *testing.T) {
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	tests := []struct {
		name    string
		fields  fields
		want    string
		wantErr bool
	}{
		{"nil bundle", fields{}, "", true},
		{"valid bundle", fields{bundle: &testCheckBundle}, "/check_bundle/1234", false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.CID()
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.CID() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("Bundle.CID() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBundle_Period(t *testing.T) {
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	tests := []struct {
		name    string
		fields  fields
		want    uint
		wantErr bool
	}{
		{"nil bundle", fields{}, 0, true},
		{"valid bundle", fields{bundle: &testCheckBundle}, 60, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.Period()
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.Period() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("Bundle.Period() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBundle_Info(t *testing.T) {
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	tests := []struct {
		name    string
		fields  fields
		want    *Info
		wantErr bool
	}{
		{"nil bundle", fields{}, nil, true},
		{"valid bundle", fields{bundle: &testCheckBundle}, &Info{CID: "/check_bundle/1234", Checks: []string{"/check/1234"}}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.Info()
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.Info() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("Bundle.Info() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBundle_Refresh(t *testing.T) {
	type fields struct {
		statusActiveMetric    string
		statusActiveBroker    string
		brokerMaxResponseTime time.Duration
		brokerMaxRetries      int
		bundle                *apiclient.CheckBundle
		client                API
		lastRefresh           time.Time
		logger                zerolog.Logger
		manage                bool
		metricStates          *metricStates
		metricStateUpdate     bool
		refreshTTL            time.Duration
		stateFile             string
		statePath             string
	}
	tests := []struct {
		name    string
		fields  fields
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				statusActiveMetric:    tt.fields.statusActiveMetric,
				statusActiveBroker:    tt.fields.statusActiveBroker,
				brokerMaxResponseTime: tt.fields.brokerMaxResponseTime,
				brokerMaxRetries:      tt.fields.brokerMaxRetries,
				bundle:                tt.fields.bundle,
				client:                tt.fields.client,
				lastRefresh:           tt.fields.lastRefresh,
				logger:                tt.fields.logger,
				manage:                tt.fields.manage,
				metricStates:          tt.fields.metricStates,
				metricStateUpdate:     tt.fields.metricStateUpdate,
				refreshTTL:            tt.fields.refreshTTL,
				stateFile:             tt.fields.stateFile,
				statePath:             tt.fields.statePath,
			}
			if err := cb.Refresh(); (err != nil) != tt.wantErr {
				t.Errorf("Bundle.Refresh() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}
<commit_msg>add: check cid test<commit_after>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package bundle

import (
	"reflect"
	"testing"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/go-apiclient"
	"github.com/gojuno/minimock/v3"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
)

func TestNew(t *testing.T) {
	type args struct {
		client API
	}
	tests := []struct {
		name    string
		args    args
		want    *Bundle
		wantErr bool
	}{
		{"invalid (nil client)", args{client: nil}, nil, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := New(tt.args.client)
			if (err != nil) != tt.wantErr {
				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("New() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestFetchCheck(t *testing.T) {
	t.Log("Testing fetchCheck")

	zerolog.SetGlobalLevel(zerolog.Disabled)

	viper.Reset()
	viper.Set(config.KeyAPITokenKey, "foo")
	viper.Set(config.KeyAPITokenApp, "bar")
	viper.Set(config.KeyAPIURL, "baz")

	mc := minimock.NewController(t)
	client := genMockClient(mc)
	c := Bundle{
		client: client,
	}

	t.Log("cid (empty)")
	{
		cid := ""
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "invalid cid (empty)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("cid (abc)")
	{
		cid := "abc"
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "invalid cid (abc)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("api error")
	{
		cid := "000"
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "unable to retrieve check bundle (/check_bundle/000): forced mock api call error" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("valid")
	{
		cid := "1234"
		viper.Set(config.KeyCheckBundleID, cid)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.fetchCheckBundle(cid)
		if err != nil {
			t.Fatalf("expected no error, got (%s)", err)
		}
	}
}

func TestFindCheck(t *testing.T) {
	t.Log("Testing findCheck")

	zerolog.SetGlobalLevel(zerolog.Disabled)

	viper.Reset()
	viper.Set(config.KeyAPITokenKey, "foo")
	viper.Set(config.KeyAPITokenApp, "bar")
	viper.Set(config.KeyAPIURL, "baz")

	mc := minimock.NewController(t)
	client := genMockClient(mc)
	c := Bundle{client: client}

	t.Log("target (empty)")
	{

		target := ""
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != -1 {
			t.Fatal("expected found == -1")
		}

		if err.Error() != "invalid check bundle target (empty)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("api error")
	{
		target := "000"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != -1 {
			t.Fatal("expected found == -1")
		}

		if err.Error() != "searching for check bundle: forced mock api call error" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("not found")
	{
		target := "not_found"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != 0 {
			t.Fatal("expected found == 0")
		}

		if err.Error() != `no check bundles matched criteria ((active:1)(type:"json:nad")(target:"not_found"))` {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("multiple")
	{
		target := "multiple"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}
		if found != 2 {
			t.Fatal("expected found == 2")
		}

		if err.Error() != `more than one (2) check bundle matched criteria ((active:1)(type:"json:nad")(target:"multiple"))` {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}

	t.Log("valid")
	{
		target := "valid"
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, found, err := c.findCheckBundle()
		if err != nil {
			t.Fatalf("expected no error, got (%s)", err)
		}
		if found != 1 {
			t.Fatal("expected found == 1")
		}
	}
}

func TestCreateCheck(t *testing.T) {
	t.Log("Testing createCheck")

	zerolog.SetGlobalLevel(zerolog.Disabled)

	viper.Reset()
	viper.Set(config.KeyAPITokenKey, "foo")
	viper.Set(config.KeyAPITokenApp, "bar")
	viper.Set(config.KeyAPIURL, "baz")

	mc := minimock.NewController(t)
	client := genMockClient(mc)
	c := Bundle{client: client}

	t.Log("target (empty)")
	{
		target := ""
		viper.Set(config.KeyCheckTarget, target)
		viper.Set(config.KeyCheckEnableNewMetrics, true)

		_, err := c.createCheckBundle()
		if err == nil {
			t.Fatal("expected error")
		}

		if err.Error() != "invalid check bundle target (empty)" {
			t.Fatalf("unexpected error return (%s)", err)
		}
	}
}

func TestBundle_CID(t *testing.T) {
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	tests := []struct {
		name    string
		fields  fields
		want    string
		wantErr bool
	}{
		{"nil bundle", fields{}, "", true},
		{"valid bundle", fields{bundle: &testCheckBundle}, "/check_bundle/1234", false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.CID()
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.CID() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("Bundle.CID() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBundle_Period(t *testing.T) {
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	tests := []struct {
		name    string
		fields  fields
		want    uint
		wantErr bool
	}{
		{"nil bundle", fields{}, 0, true},
		{"valid bundle", fields{bundle: &testCheckBundle}, 60, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.Period()
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.Period() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("Bundle.Period() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBundle_Info(t *testing.T) {
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	tests := []struct {
		name    string
		fields  fields
		want    *Info
		wantErr bool
	}{
		{"nil bundle", fields{}, nil, true},
		{"valid bundle", fields{bundle: &testCheckBundle}, &Info{CID: "/check_bundle/1234", Checks: []string{"/check/1234"}}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.Info()
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.Info() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("Bundle.Info() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestBundle_Refresh(t *testing.T) {
	zerolog.SetGlobalLevel(zerolog.Disabled)
	viper.Set(config.KeyCheckBundleID, "/check_bundle/1234")
	mc := minimock.NewController(t)
	client := genMockClient(mc)
	tb := testCheckBundle
	type fields struct {
		bundle *apiclient.CheckBundle
		client API
		logger zerolog.Logger
	}
	tests := []struct {
		name    string
		fields  fields
		wantErr bool
	}{
		{"nil bundle", fields{bundle: nil, client: client, logger: log.With().Logger()}, true},
		{"valid", fields{bundle: &tb, client: client, logger: log.With().Logger()}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
				client: tt.fields.client,
				logger: tt.fields.logger,
			}
			if err := cb.Refresh(); (err != nil) != tt.wantErr {
				t.Errorf("Bundle.Refresh() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func TestBundle_CheckCID(t *testing.T) {
	tb := testCheckBundle
	type fields struct {
		bundle *apiclient.CheckBundle
	}
	type args struct {
		idx uint
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    string
		wantErr bool
	}{
		{"invalid (nil bundle)", fields{bundle: nil}, args{idx: 0}, "", true},
		{"invalid (no checks in bundle)", fields{bundle: &apiclient.CheckBundle{}}, args{idx: 0}, "", true},
		{"invalid (idx out of range)", fields{bundle: &tb}, args{idx: 10}, "", true},
		{"valid", fields{bundle: &tb}, args{idx: 0}, testCheckBundle.Checks[0], false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cb := &Bundle{
				bundle: tt.fields.bundle,
			}
			got, err := cb.CheckCID(tt.args.idx)
			if (err != nil) != tt.wantErr {
				t.Errorf("Bundle.CheckCID() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("Bundle.CheckCID() = %v, want %v", got, tt.want)
			}
		})
	}
}
<|endoftext|>


<commit_before>package fasthttp

import (
	"bufio"
	"fmt"
	"io"
	"math"
	"sync"
	"time"
	"unsafe"
)

var gmtLocation = func() *time.Location {
	x, err := time.LoadLocation("GMT")
	if err != nil {
		panic(fmt.Sprintf("cannot load GMT location: %s", err))
	}
	return x
}()

// AppendHTTPDate appends HTTP-compliant (RFC1123) representation of date
// to dst and returns dst (which may be newly allocated).
func AppendHTTPDate(dst []byte, date time.Time) []byte {
	return date.In(gmtLocation).AppendFormat(dst, time.RFC1123)
}

// AppendUint appends n to dst and returns dst (which may be newly allocated).
func AppendUint(dst []byte, n int) []byte {
	if n < 0 {
		panic("BUG: int must be positive")
	}

	var b [20]byte
	buf := b[:]
	i := len(buf) - 1
	for {
		buf[i] = '0' + byte(n%10)
		n /= 10
		if n == 0 {
			break
		}
		i--
	}

	dst = append(dst, buf[i:]...)
	return dst
}

// ParseUint parses uint from buf.
func ParseUint(buf []byte) (int, error) {
	v, n, err := parseUintBuf(buf)
	if n != len(buf) {
		return -1, fmt.Errorf("only %b bytes out of %d bytes exhausted when parsing int %q", n, len(buf), buf)
	}
	return v, err
}

func parseUintBuf(b []byte) (int, int, error) {
	n := len(b)
	if n == 0 {
		return -1, 0, fmt.Errorf("empty integer")
	}
	v := 0
	for i := 0; i < n; i++ {
		c := b[i]
		k := c - '0'
		if k > 9 {
			if i == 0 {
				return -1, i, fmt.Errorf("unexpected first char %c. Expected 0-9", c)
			}
			return v, i, nil
		}
		if i >= maxIntChars {
			return -1, i, fmt.Errorf("too long int %q", b[:i+1])
		}
		v = 10*v + int(k)
	}
	return v, n, nil
}

// ParseUfloat parses unsigned float from buf.
func ParseUfloat(buf []byte) (float64, error) {
	if len(buf) == 0 {
		return -1, fmt.Errorf("empty float number")
	}
	b := buf
	var v uint64
	var offset float64 = 1.0
	var pointFound bool
	for i, c := range b {
		if c < '0' || c > '9' {
			if c == '.' {
				if pointFound {
					return -1, fmt.Errorf("duplicate point found in %q", buf)
				}
				pointFound = true
				continue
			}
			if c == 'e' || c == 'E' {
				if i+1 >= len(b) {
					return -1, fmt.Errorf("unexpected end of float after %c. num=%q", c, buf)
				}
				b = b[i+1:]
				minus := -1
				switch b[0] {
				case '+':
					b = b[1:]
					minus = 1
				case '-':
					b = b[1:]
				default:
					minus = 1
				}
				vv, err := ParseUint(b)
				if err != nil {
					return -1, fmt.Errorf("cannot parse exponent part of %q: %s", buf, err)
				}
				return float64(v) * offset * math.Pow10(minus*int(vv)), nil
			}
			return -1, fmt.Errorf("unexpected char found %c in %q", c, buf)
		}
		v = 10*v + uint64(c-'0')
		if pointFound {
			offset /= 10
		}
	}
	return float64(v) * offset, nil
}

func readHexInt(r *bufio.Reader) (int, error) {
	n := 0
	i := 0
	var k int
	for {
		c, err := r.ReadByte()
		if err != nil {
			if err == io.EOF && i > 0 {
				return n, nil
			}
			return -1, err
		}
		k = hexbyte2int(c)
		if k < 0 {
			if i == 0 {
				return -1, fmt.Errorf("cannot read hex num from empty string")
			}
			r.UnreadByte()
			return n, nil
		}
		if i >= maxHexIntChars {
			return -1, fmt.Errorf("cannot read hex num with more than %d digits", maxHexIntChars)
		}
		n = (n << 4) | k
		i++
	}
}

var hexIntBufPool sync.Pool

func writeHexInt(w *bufio.Writer, n int) error {
	if n < 0 {
		panic("BUG: int must be positive")
	}

	v := hexIntBufPool.Get()
	if v == nil {
		v = make([]byte, maxHexIntChars+1)
	}
	buf := v.([]byte)
	i := len(buf) - 1
	for {
		buf[i] = int2hexbyte(n & 0xf)
		n >>= 4
		if n == 0 {
			break
		}
		i--
	}
	_, err := w.Write(buf[i:])
	hexIntBufPool.Put(v)
	return err
}

func int2hexbyte(n int) byte {
	if n < 10 {
		return '0' + byte(n)
	}
	return 'a' + byte(n) - 10
}

func hexCharUpper(c byte) byte {
	if c < 10 {
		return '0' + c
	}
	return c - 10 + 'A'
}

var hex2intTable = func() []byte {
	b := make([]byte, 255)
	for i := byte(0); i < 255; i++ {
		c := byte(0)
		if i >= '0' && i <= '9' {
			c = 1 + i - '0'
		} else if i >= 'a' && i <= 'f' {
			c = 1 + i - 'a' + 10
		} else if i >= 'A' && i <= 'F' {
			c = 1 + i - 'A' + 10
		}
		b[i] = c
	}
	return b
}()

func hexbyte2int(c byte) int {
	return int(hex2intTable[c]) - 1
}

const toLower = 'a' - 'A'

func uppercaseByte(p *byte) {
	c := *p
	if c >= 'a' && c <= 'z' {
		*p = c - toLower
	}
}

func lowercaseByte(p *byte) {
	c := *p
	if c >= 'A' && c <= 'Z' {
		*p = c + toLower
	}
}

func lowercaseBytes(b []byte) {
	for i, n := 0, len(b); i < n; i++ {
		lowercaseByte(&b[i])
	}
}

// unsafeBytesToStr converts byte slice to a string without memory allocation.
// See https://groups.google.com/forum/#!msg/Golang-Nuts/ENgbUzYvCuU/90yGx7GUAgAJ .
//
// Note it may break if string and/or slice header will change
// in the future go versions.
func unsafeBytesToStr(b []byte) string {
	return *(*string)(unsafe.Pointer(&b))
}

func appendQuotedArg(dst, v []byte) []byte {
	for _, c := range v {
		if c >= '0' && c <= '9' || c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c == '/' || c == '.' {
			dst = append(dst, c)
		} else {
			dst = append(dst, '%', hexCharUpper(c>>4), hexCharUpper(c&15))
		}
	}
	return dst
}

// EqualBytesStr returns true if string(b) == s.
//
// This function has no performance benefits comparing to string(b) == s.
// It is left here for backwards compatibility only.
func EqualBytesStr(b []byte, s string) bool {
	return string(b) == s
}

// AppendBytesStr appends src to dst and returns dst
// (which may be newly allocated).
//
// This function has no performance benefits comparing to append(dst, src...).
// It is left here for backwards compatibility only.
func AppendBytesStr(dst []byte, src string) []byte {
	return append(dst, src...)
}
<commit_msg>Integrated AppendUint perfromance optimization from pull request #2. Thanks to @EricLagergren<commit_after>package fasthttp

import (
	"bufio"
	"fmt"
	"io"
	"math"
	"sync"
	"time"
	"unsafe"
)

var gmtLocation = func() *time.Location {
	x, err := time.LoadLocation("GMT")
	if err != nil {
		panic(fmt.Sprintf("cannot load GMT location: %s", err))
	}
	return x
}()

// AppendHTTPDate appends HTTP-compliant (RFC1123) representation of date
// to dst and returns dst (which may be newly allocated).
func AppendHTTPDate(dst []byte, date time.Time) []byte {
	return date.In(gmtLocation).AppendFormat(dst, time.RFC1123)
}

// AppendUint appends n to dst and returns dst (which may be newly allocated).
func AppendUint(dst []byte, n int) []byte {
	if n < 0 {
		panic("BUG: int must be positive")
	}

	var b [20]byte
	buf := b[:]
	i := len(buf)
	var q int
	for n >= 10 {
		i--
		q = n / 10
		buf[i] = '0' + byte(n-q*10)
		n = q
	}
	i--
	buf[i] = '0' + byte(n)

	dst = append(dst, buf[i:]...)
	return dst
}

// ParseUint parses uint from buf.
func ParseUint(buf []byte) (int, error) {
	v, n, err := parseUintBuf(buf)
	if n != len(buf) {
		return -1, fmt.Errorf("only %b bytes out of %d bytes exhausted when parsing int %q", n, len(buf), buf)
	}
	return v, err
}

func parseUintBuf(b []byte) (int, int, error) {
	n := len(b)
	if n == 0 {
		return -1, 0, fmt.Errorf("empty integer")
	}
	v := 0
	for i := 0; i < n; i++ {
		c := b[i]
		k := c - '0'
		if k > 9 {
			if i == 0 {
				return -1, i, fmt.Errorf("unexpected first char %c. Expected 0-9", c)
			}
			return v, i, nil
		}
		if i >= maxIntChars {
			return -1, i, fmt.Errorf("too long int %q", b[:i+1])
		}
		v = 10*v + int(k)
	}
	return v, n, nil
}

// ParseUfloat parses unsigned float from buf.
func ParseUfloat(buf []byte) (float64, error) {
	if len(buf) == 0 {
		return -1, fmt.Errorf("empty float number")
	}
	b := buf
	var v uint64
	var offset float64 = 1.0
	var pointFound bool
	for i, c := range b {
		if c < '0' || c > '9' {
			if c == '.' {
				if pointFound {
					return -1, fmt.Errorf("duplicate point found in %q", buf)
				}
				pointFound = true
				continue
			}
			if c == 'e' || c == 'E' {
				if i+1 >= len(b) {
					return -1, fmt.Errorf("unexpected end of float after %c. num=%q", c, buf)
				}
				b = b[i+1:]
				minus := -1
				switch b[0] {
				case '+':
					b = b[1:]
					minus = 1
				case '-':
					b = b[1:]
				default:
					minus = 1
				}
				vv, err := ParseUint(b)
				if err != nil {
					return -1, fmt.Errorf("cannot parse exponent part of %q: %s", buf, err)
				}
				return float64(v) * offset * math.Pow10(minus*int(vv)), nil
			}
			return -1, fmt.Errorf("unexpected char found %c in %q", c, buf)
		}
		v = 10*v + uint64(c-'0')
		if pointFound {
			offset /= 10
		}
	}
	return float64(v) * offset, nil
}

func readHexInt(r *bufio.Reader) (int, error) {
	n := 0
	i := 0
	var k int
	for {
		c, err := r.ReadByte()
		if err != nil {
			if err == io.EOF && i > 0 {
				return n, nil
			}
			return -1, err
		}
		k = hexbyte2int(c)
		if k < 0 {
			if i == 0 {
				return -1, fmt.Errorf("cannot read hex num from empty string")
			}
			r.UnreadByte()
			return n, nil
		}
		if i >= maxHexIntChars {
			return -1, fmt.Errorf("cannot read hex num with more than %d digits", maxHexIntChars)
		}
		n = (n << 4) | k
		i++
	}
}

var hexIntBufPool sync.Pool

func writeHexInt(w *bufio.Writer, n int) error {
	if n < 0 {
		panic("BUG: int must be positive")
	}

	v := hexIntBufPool.Get()
	if v == nil {
		v = make([]byte, maxHexIntChars+1)
	}
	buf := v.([]byte)
	i := len(buf) - 1
	for {
		buf[i] = int2hexbyte(n & 0xf)
		n >>= 4
		if n == 0 {
			break
		}
		i--
	}
	_, err := w.Write(buf[i:])
	hexIntBufPool.Put(v)
	return err
}

func int2hexbyte(n int) byte {
	if n < 10 {
		return '0' + byte(n)
	}
	return 'a' + byte(n) - 10
}

func hexCharUpper(c byte) byte {
	if c < 10 {
		return '0' + c
	}
	return c - 10 + 'A'
}

var hex2intTable = func() []byte {
	b := make([]byte, 255)
	for i := byte(0); i < 255; i++ {
		c := byte(0)
		if i >= '0' && i <= '9' {
			c = 1 + i - '0'
		} else if i >= 'a' && i <= 'f' {
			c = 1 + i - 'a' + 10
		} else if i >= 'A' && i <= 'F' {
			c = 1 + i - 'A' + 10
		}
		b[i] = c
	}
	return b
}()

func hexbyte2int(c byte) int {
	return int(hex2intTable[c]) - 1
}

const toLower = 'a' - 'A'

func uppercaseByte(p *byte) {
	c := *p
	if c >= 'a' && c <= 'z' {
		*p = c - toLower
	}
}

func lowercaseByte(p *byte) {
	c := *p
	if c >= 'A' && c <= 'Z' {
		*p = c + toLower
	}
}

func lowercaseBytes(b []byte) {
	for i, n := 0, len(b); i < n; i++ {
		lowercaseByte(&b[i])
	}
}

// unsafeBytesToStr converts byte slice to a string without memory allocation.
// See https://groups.google.com/forum/#!msg/Golang-Nuts/ENgbUzYvCuU/90yGx7GUAgAJ .
//
// Note it may break if string and/or slice header will change
// in the future go versions.
func unsafeBytesToStr(b []byte) string {
	return *(*string)(unsafe.Pointer(&b))
}

func appendQuotedArg(dst, v []byte) []byte {
	for _, c := range v {
		if c >= '0' && c <= '9' || c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c == '/' || c == '.' {
			dst = append(dst, c)
		} else {
			dst = append(dst, '%', hexCharUpper(c>>4), hexCharUpper(c&15))
		}
	}
	return dst
}

// EqualBytesStr returns true if string(b) == s.
//
// This function has no performance benefits comparing to string(b) == s.
// It is left here for backwards compatibility only.
func EqualBytesStr(b []byte, s string) bool {
	return string(b) == s
}

// AppendBytesStr appends src to dst and returns dst
// (which may be newly allocated).
//
// This function has no performance benefits comparing to append(dst, src...).
// It is left here for backwards compatibility only.
func AppendBytesStr(dst []byte, src string) []byte {
	return append(dst, src...)
}
<|endoftext|>


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

package web

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"

	"github.com/ajg/form"
	"github.com/gorilla/mux"
	"github.com/tsuru/tsuru-autoscale/datasource"
	"github.com/tsuru/tsuru-autoscale/db"
	"github.com/tsuru/tsuru/db/dbtest"
	"gopkg.in/check.v1"
)

func server(w http.ResponseWriter, r *http.Request) {
	m := mux.NewRouter()
	Router(m)
	m.ServeHTTP(w, r)
}

func Test(t *testing.T) { check.TestingT(t) }

type S struct {
	conn *db.Storage
}

var _ = check.Suite(&S{})

func (s *S) SetUpSuite(c *check.C) {
	err := os.Setenv("MONGODB_DATABASE_NAME", "tsuru_autoscale_web")
	c.Assert(err, check.IsNil)
	s.conn, err = db.Conn()
	c.Assert(err, check.IsNil)
}

func (s *S) TearDownTest(c *check.C) {
	dbtest.ClearAllCollections(s.conn.Actions().Database)
	dbtest.ClearAllCollections(s.conn.Alarms().Database)
	dbtest.ClearAllCollections(s.conn.DataSources().Database)
	dbtest.ClearAllCollections(s.conn.Events().Database)
	dbtest.ClearAllCollections(s.conn.Instances().Database)
	dbtest.ClearAllCollections(s.conn.Wizard().Database)
}

func (s *S) TestDataSourceAdd(c *check.C) {
	recorder := httptest.NewRecorder()
	ds := datasource.DataSource{
		Name:   "new",
		URL:    "http://tsuru.io",
		Method: "GET",
	}
	v, err := form.EncodeToValues(&ds)
	c.Assert(err, check.IsNil)
	body := strings.NewReader(v.Encode())
	request, err := http.NewRequest("POST", "/datasource/add", body)
	c.Assert(err, check.IsNil)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	server(recorder, request)
	c.Assert(recorder.Code, check.Equals, http.StatusFound)
}

func (s *S) TestDataSourceAddEmptyHeader(c *check.C) {
	recorder := httptest.NewRecorder()
	ds := datasource.DataSource{
		Name:    "new",
		URL:     "http://tsuru.io",
		Method:  "GET",
		Headers: map[string]string{" ": " "},
	}
	v := url.Values{
		"key":    []string{"", "f", ""},
		"value":  []string{"", "f", ""},
		"name":   []string{"new"},
		"url":    []string{"sdfasd"},
		"method": []string{"GET"},
	}
	body := strings.NewReader(v.Encode())
	request, err := http.NewRequest("POST", "/datasource/add", body)
	c.Assert(err, check.IsNil)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	server(recorder, request)
	c.Assert(recorder.Code, check.Equals, http.StatusFound)
	r, err := datasource.Get(ds.Name)
	c.Assert(err, check.IsNil)
	c.Assert(len(r.Headers), check.Equals, 1)
}
<commit_msg>web: add test for remove data source<commit_after>// Copyright 2016 tsuru-autoscale authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package web

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"

	"github.com/ajg/form"
	"github.com/gorilla/mux"
	"github.com/tsuru/tsuru-autoscale/datasource"
	"github.com/tsuru/tsuru-autoscale/db"
	"github.com/tsuru/tsuru/db/dbtest"
	"gopkg.in/check.v1"
)

func server(w http.ResponseWriter, r *http.Request) {
	m := mux.NewRouter()
	Router(m)
	m.ServeHTTP(w, r)
}

func Test(t *testing.T) { check.TestingT(t) }

type S struct {
	conn *db.Storage
}

var _ = check.Suite(&S{})

func (s *S) SetUpSuite(c *check.C) {
	err := os.Setenv("MONGODB_DATABASE_NAME", "tsuru_autoscale_web")
	c.Assert(err, check.IsNil)
	s.conn, err = db.Conn()
	c.Assert(err, check.IsNil)
}

func (s *S) TearDownTest(c *check.C) {
	dbtest.ClearAllCollections(s.conn.Actions().Database)
	dbtest.ClearAllCollections(s.conn.Alarms().Database)
	dbtest.ClearAllCollections(s.conn.DataSources().Database)
	dbtest.ClearAllCollections(s.conn.Events().Database)
	dbtest.ClearAllCollections(s.conn.Instances().Database)
	dbtest.ClearAllCollections(s.conn.Wizard().Database)
}

func (s *S) TestDataSourceAdd(c *check.C) {
	recorder := httptest.NewRecorder()
	ds := datasource.DataSource{
		Name:   "new",
		URL:    "http://tsuru.io",
		Method: "GET",
	}
	v, err := form.EncodeToValues(&ds)
	c.Assert(err, check.IsNil)
	body := strings.NewReader(v.Encode())
	request, err := http.NewRequest("POST", "/datasource/add", body)
	c.Assert(err, check.IsNil)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	server(recorder, request)
	c.Assert(recorder.Code, check.Equals, http.StatusFound)
}

func (s *S) TestDataSourceAddEmptyHeader(c *check.C) {
	recorder := httptest.NewRecorder()
	ds := datasource.DataSource{
		Name:    "new",
		URL:     "http://tsuru.io",
		Method:  "GET",
		Headers: map[string]string{" ": " "},
	}
	v := url.Values{
		"key":    []string{"", "f", ""},
		"value":  []string{"", "f", ""},
		"name":   []string{"new"},
		"url":    []string{"sdfasd"},
		"method": []string{"GET"},
	}
	body := strings.NewReader(v.Encode())
	request, err := http.NewRequest("POST", "/datasource/add", body)
	c.Assert(err, check.IsNil)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	server(recorder, request)
	c.Assert(recorder.Code, check.Equals, http.StatusFound)
	r, err := datasource.Get(ds.Name)
	c.Assert(err, check.IsNil)
	c.Assert(len(r.Headers), check.Equals, 1)
}

func (s *S) TestDataSourceRemove(c *check.C) {
	recorder := httptest.NewRecorder()
	ds := datasource.DataSource{
		Name:   "new",
		URL:    "http://tsuru.io",
		Method: "GET",
	}
	err := datasource.New(&ds)
	c.Assert(err, check.IsNil)
	request, err := http.NewRequest("GET", "/datasource/new/delete", nil)
	c.Assert(err, check.IsNil)
	server(recorder, request)
	c.Assert(recorder.Code, check.Equals, http.StatusFound)
	_, err = datasource.Get(ds.Name)
	c.Assert(err, check.NotNil)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"time"
)

var (
	onemin  = time.Minute * 1
	fivemin = time.Minute * 5
	tenmin  = time.Minute * 10
	onehour = time.Hour * 1

	oneminTimer  = time.NewTimer(onemin)
	fiveminTimer = time.NewTimer(fivemin)
	tenminTimer  = time.NewTimer(tenmin)
	onehourTimer = time.NewTimer(onehour)
)

func StartDispatcher() {
	go func() {
		for {
			select {
			case <-oneminTimer.C:
				err := ProcessCheck(onemin)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			case <-fiveminTimer.C:
				err := ProcessCheck(fivemin)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			case <-tenminTimer.C:
				err := ProcessCheck(tenmin)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			case <-onehourTimer.C:
				err := ProcessCheck(onehour)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			}
		}
	}()
}
<commit_msg>Switch to the proper "Ticker" implementation<commit_after>package main

import (
	"log"
	"time"
)

var (
	onemin  = time.Minute * 1
	fivemin = time.Minute * 5
	tenmin  = time.Minute * 10
	onehour = time.Hour * 1

	oneminTicker  = time.NewTicker(onemin)
	fiveminTicker = time.NewTicker(fivemin)
	tenminTicker  = time.NewTicker(tenmin)
	onehourTicker = time.NewTicker(onehour)
)

func StartDispatcher() {
	go func() {
		for {
			select {
			case <-oneminTicker.C:
				err := ProcessCheck(onemin)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			case <-fiveminTicker.C:
				err := ProcessCheck(fivemin)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			case <-tenminTicker.C:
				err := ProcessCheck(tenmin)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			case <-onehourTicker.C:
				err := ProcessCheck(onehour)
				if err != nil {
					log.Println("Failed to process checks: ", err.Error())
				}
			}
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/donnpebe/go-redis-timeseries"
	"github.com/garyburd/redigo/redis"
)

type Data struct {
	Author  string `json:"author"`
	Message string `json:"mesage"`
}

func main() {
	conn, err := redis.Dial("tcp", "localhost:6379")
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	// data will be split into per 1 second key and key will have 5 day exiration per key
	// if didn't want to expire the key, set it to 0
	ts := timeseries.NewTimeSeries("dump::device1", 1800*time.Second, 5*24*time.Hour, conn)

	now := time.Now()
	fmt.Printf("Adding data points...\n\n")
	data := Data{"author 1", "message 1"}
	err = ts.Add(&data, now)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Get from %v...\n\n", now)

	dat := new(Data)
	if err = ts.Fetch(now, dat); err != nil {
		panic(err)
	}

	fmt.Println("Records")
	fmt.Println("=======")

	fmt.Printf("Author: %s, Message: %s\n", dat.Author, dat.Message)

}
<commit_msg>fix wrong info in comment<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/donnpebe/go-redis-timeseries"
	"github.com/garyburd/redigo/redis"
)

type Data struct {
	Author  string `json:"author"`
	Message string `json:"mesage"`
}

func main() {
	conn, err := redis.Dial("tcp", "localhost:6379")
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	// data will be split into per 1800 second key and key will have 5 day exiration per key
	// if didn't want to expire the key, set it to 0
	ts := timeseries.NewTimeSeries("dump::device1", 1800*time.Second, 5*24*time.Hour, conn)

	now := time.Now()
	fmt.Printf("Adding data points...\n\n")
	data := Data{"author 1", "message 1"}
	err = ts.Add(&data, now)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Get from %v...\n\n", now)

	dat := new(Data)
	if err = ts.Fetch(now, dat); err != nil {
		panic(err)
	}

	fmt.Println("Records")
	fmt.Println("=======")

	fmt.Printf("Author: %s, Message: %s\n", dat.Author, dat.Message)

}
<|endoftext|>


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

package errors

// Chain is an error and layered error annotations.
type Chain interface {
	// The error behavior of a Chain is based on the last annotation applied.
	error

	Cause() error         // Cause returns the initial error in the Chain.
	Link(error) Chain     // Link adds an annotation layer to the Chain.
	Last() (Chain, error) // Last returns the Chain and the most recent annotation.
}

// Links is an optional interface used by the Errors function.
type Links interface {
	Errors() []error // Errors returns a flat list of errors in order of annotation.
}

// New returns a new Chain based on the provided error. If the error is a Chain it
// is returned unaltered.
func New(err error) Chain {
	if c, ok := err.(Chain); ok {
		return c
	}
	return chain{err}
}

// Cause returns the initially identified cause of an error if the error is a Chain, or the error
// itself if it is not.
func Cause(err error) error {
	if c, ok := err.(Chain); ok {
		return c.Cause()
	}
	return err
}

// Link adds an annotation to an error, returning a Chain.
func Link(err, annotation error) Chain { return New(err).Link(annotation) }

// Last returns the most recent annotation of an error and the remaining chain
// after the annotation is removed or nil if no further errors remain. Last returns
// a nil Chain if the error is not a Chain.
func Last(err error) (Chain, error) {
	if c, ok := err.(Chain); ok {
		return c.Last()
	}
	return nil, err
}

// Errors returns a flat list of errors in temporal order of annotation. If the provided
// error is not a Chain a single element slice of error is returned containing the error.
// If the error implements Links, its Errors method is called and the result returned.
func Errors(err error) []error {
	if err == nil {
		return nil
	}
	switch c := err.(type) {
	case Links:
		return c.Errors()
	case Chain:
		var errs []error
		for c != nil {
			c, err = c.Last()
			errs = append(errs, err)
		}
		return reverse(errs)
	default:
		return []error{err}
	}
}

func reverse(err []error) []error {
	for i, j := 0, len(err)-1; i < j; i, j = i+1, j-1 {
		err[i], err[j] = err[j], err[i]
	}
	return err
}

// chain is the basic implementation.
type chain []error

func (c chain) Error() string {
	if len(c) > 0 {
		return c[len(c)-1].Error()
	}
	return ""
}
func (c chain) Cause() error {
	if len(c) > 0 {
		return c[0]
	}
	return nil
}
func (c chain) Link(err error) Chain { return append(c, err) }
func (c chain) Last() (Chain, error) {
	switch len(c) {
	case 0:
		return nil, nil
	case 1:
		return nil, c[0]
	default:
		return c[:len(c)-1], c[len(c)-1]
	}
}
func (c chain) Errors() []error { return c }
<commit_msg>New -> NewChain<commit_after>// Copyright ©2013 The bíogo Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package errors

// Chain is an error and layered error annotations.
type Chain interface {
	// The error behavior of a Chain is based on the last annotation applied.
	error

	Cause() error         // Cause returns the initial error in the Chain.
	Link(error) Chain     // Link adds an annotation layer to the Chain.
	Last() (Chain, error) // Last returns the Chain and the most recent annotation.
}

// Links is an optional interface used by the Errors function.
type Links interface {
	Errors() []error // Errors returns a flat list of errors in order of annotation.
}

// New returns a new Chain based on the provided error. If the error is a Chain it
// is returned unaltered.
func NewChain(err error) Chain {
	if c, ok := err.(Chain); ok {
		return c
	}
	return chain{err}
}

// Cause returns the initially identified cause of an error if the error is a Chain, or the error
// itself if it is not.
func Cause(err error) error {
	if c, ok := err.(Chain); ok {
		return c.Cause()
	}
	return err
}

// Link adds an annotation to an error, returning a Chain.
func Link(err, annotation error) Chain { return NewChain(err).Link(annotation) }

// Last returns the most recent annotation of an error and the remaining chain
// after the annotation is removed or nil if no further errors remain. Last returns
// a nil Chain if the error is not a Chain.
func Last(err error) (Chain, error) {
	if c, ok := err.(Chain); ok {
		return c.Last()
	}
	return nil, err
}

// Errors returns a flat list of errors in temporal order of annotation. If the provided
// error is not a Chain a single element slice of error is returned containing the error.
// If the error implements Links, its Errors method is called and the result returned.
func Errors(err error) []error {
	if err == nil {
		return nil
	}
	switch c := err.(type) {
	case Links:
		return c.Errors()
	case Chain:
		var errs []error
		for c != nil {
			c, err = c.Last()
			errs = append(errs, err)
		}
		return reverse(errs)
	default:
		return []error{err}
	}
}

func reverse(err []error) []error {
	for i, j := 0, len(err)-1; i < j; i, j = i+1, j-1 {
		err[i], err[j] = err[j], err[i]
	}
	return err
}

// chain is the basic implementation.
type chain []error

func (c chain) Error() string {
	if len(c) > 0 {
		return c[len(c)-1].Error()
	}
	return ""
}
func (c chain) Cause() error {
	if len(c) > 0 {
		return c[0]
	}
	return nil
}
func (c chain) Link(err error) Chain { return append(c, err) }
func (c chain) Last() (Chain, error) {
	switch len(c) {
	case 0:
		return nil, nil
	case 1:
		return nil, c[0]
	default:
		return c[:len(c)-1], c[len(c)-1]
	}
}
func (c chain) Errors() []error { return c }
<|endoftext|>


<commit_before>//go:build linux
// +build linux

/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package e2enode

import (
	"context"
	"fmt"
	"os/exec"
	"time"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/kubernetes/test/e2e/framework"
	e2eskipper "k8s.io/kubernetes/test/e2e/framework/skipper"
	testutils "k8s.io/kubernetes/test/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
)

// waitForPods waits for timeout duration, for podCount.
// If the timeout is hit, it returns the list of currently running pods.
func waitForPods(f *framework.Framework, podCount int, timeout time.Duration) (runningPods []*v1.Pod) {
	for start := time.Now(); time.Since(start) < timeout; time.Sleep(10 * time.Second) {
		podList, err := f.PodClient().List(context.TODO(), metav1.ListOptions{})
		if err != nil {
			framework.Logf("Failed to list pods on node: %v", err)
			continue
		}

		runningPods = []*v1.Pod{}
		for i := range podList.Items {
			pod := podList.Items[i]
			if r, err := testutils.PodRunningReadyOrSucceeded(&pod); err != nil || !r {
				continue
			}
			runningPods = append(runningPods, &pod)
		}
		framework.Logf("Running pod count %d", len(runningPods))
		if len(runningPods) >= podCount {
			break
		}
	}
	return runningPods
}

var _ = SIGDescribe("Restart [Serial] [Slow] [Disruptive]", func() {
	const (
		// Saturate the node. It's not necessary that all these pods enter
		// Running/Ready, because we don't know the number of cores in the
		// test node or default limits applied (if any). It's is essential
		// that no containers end up in terminated. 100 was chosen because
		// it's the max pods per node.
		podCount            = 100
		podCreationInterval = 100 * time.Millisecond
		recoverTimeout      = 5 * time.Minute
		startTimeout        = 5 * time.Minute
		// restartCount is chosen so even with minPods we exhaust the default
		// allocation of a /24.
		minPods      = 50
		restartCount = 6
	)

	f := framework.NewDefaultFramework("restart-test")
	ginkgo.Context("Container Runtime", func() {
		ginkgo.Context("Network", func() {
			ginkgo.It("should recover from ip leak", func() {
				if framework.TestContext.ContainerRuntime == "docker" {
					ginkgo.Skip("Test fails with in-tree docker. Skipping test.")
				}

				pods := newTestPods(podCount, false, imageutils.GetPauseImageName(), "restart-container-runtime-test")
				ginkgo.By(fmt.Sprintf("Trying to create %d pods on node", len(pods)))
				createBatchPodWithRateControl(f, pods, podCreationInterval)
				defer deletePodsSync(f, pods)

				// Give the node some time to stabilize, assume pods that enter RunningReady within
				// startTimeout fit on the node and the node is now saturated.
				runningPods := waitForPods(f, podCount, startTimeout)
				if len(runningPods) < minPods {
					framework.Failf("Failed to start %d pods, cannot test that restarting container runtime doesn't leak IPs", minPods)
				}

				for i := 0; i < restartCount; i++ {
					ginkgo.By(fmt.Sprintf("Killing container runtime iteration %d", i))
					// Wait for container runtime to be running
					var pid int
					gomega.Eventually(func() error {
						runtimePids, err := getPidsForProcess(framework.TestContext.ContainerRuntimeProcessName, framework.TestContext.ContainerRuntimePidFile)
						if err != nil {
							return err
						}
						if len(runtimePids) != 1 {
							return fmt.Errorf("unexpected container runtime pid list: %+v", runtimePids)
						}
						// Make sure the container runtime is running, pid got from pid file may not be running.
						pid = runtimePids[0]
						if _, err := exec.Command("sudo", "ps", "-p", fmt.Sprintf("%d", pid)).CombinedOutput(); err != nil {
							return err
						}
						return nil
					}, 1*time.Minute, 2*time.Second).Should(gomega.BeNil())
					if stdout, err := exec.Command("sudo", "kill", "-SIGKILL", fmt.Sprintf("%d", pid)).CombinedOutput(); err != nil {
						framework.Failf("Failed to kill container runtime (pid=%d): %v, stdout: %q", pid, err, string(stdout))
					}
					// Assume that container runtime will be restarted by systemd/supervisord etc.
					time.Sleep(20 * time.Second)
				}

				ginkgo.By("Checking currently Running/Ready pods")
				postRestartRunningPods := waitForPods(f, len(runningPods), recoverTimeout)
				if len(postRestartRunningPods) == 0 {
					framework.Failf("Failed to start *any* pods after container runtime restart, this might indicate an IP leak")
				}
				ginkgo.By("Confirm no containers have terminated")
				for _, pod := range postRestartRunningPods {
					if c := testutils.TerminatedContainers(pod); len(c) != 0 {
						framework.Failf("Pod %q has failed containers %+v after container runtime restart, this might indicate an IP leak", pod.Name, c)
					}
				}
				ginkgo.By(fmt.Sprintf("Container runtime restart test passed with %d pods", len(postRestartRunningPods)))
			})
		})
	})

	ginkgo.Context("Kubelet", func() {
		ginkgo.It("should correctly account for terminated pods after restart", func() {
			node := getLocalNode(f)
			cpus := node.Status.Allocatable[v1.ResourceCPU]
			numCpus := int((&cpus).Value())
			if numCpus < 1 {
				e2eskipper.Skipf("insufficient CPU available for kubelet restart test")
			}

			// create as many restartNever pods as there are allocatable CPU
			// nodes; if they are not correctly accounted for as terminated
			// later, this will fill up all node capacity
			podCountRestartNever := numCpus
			ginkgo.By(fmt.Sprintf("creating %d RestartNever pods on node", podCountRestartNever))
			restartNeverPods := newTestPods(podCountRestartNever, false, imageutils.GetE2EImage(imageutils.BusyBox), "restart-kubelet-test")
			for _, pod := range restartNeverPods {
				pod.Spec.RestartPolicy = "Never"
				pod.Spec.Containers[0].Command = []string{"echo", "hi"}
				pod.Spec.Containers[0].Resources.Limits = v1.ResourceList{
					v1.ResourceCPU: resource.MustParse("1"),
				}
			}
			createBatchPodWithRateControl(f, restartNeverPods, podCreationInterval)
			defer deletePodsSync(f, restartNeverPods)

			completedPods := waitForPods(f, podCountRestartNever, startTimeout)
			if len(completedPods) < podCountRestartNever {
				framework.Failf("Failed to run sufficient restartNever pods, got %d but expected %d", len(completedPods), podCountRestartNever)
			}

			podCountRestartAlways := (numCpus / 2) + 1
			ginkgo.By(fmt.Sprintf("creating %d RestartAlways pods on node", podCountRestartAlways))
			restartAlwaysPods := newTestPods(podCountRestartAlways, false, imageutils.GetPauseImageName(), "restart-kubelet-test")
			for _, pod := range restartAlwaysPods {
				pod.Spec.Containers[0].Resources.Limits = v1.ResourceList{
					v1.ResourceCPU: resource.MustParse("1"),
				}
			}
			createBatchPodWithRateControl(f, restartAlwaysPods, podCreationInterval)
			defer deletePodsSync(f, restartAlwaysPods)

			numAllPods := podCountRestartNever + podCountRestartAlways
			allPods := waitForPods(f, numAllPods, startTimeout)
			if len(allPods) < numAllPods {
				framework.Failf("Failed to run sufficient restartAlways pods, got %d but expected %d", len(allPods), numAllPods)
			}

			ginkgo.By("killing and restarting kubelet")
			// We want to kill the kubelet rather than a graceful restart
			startKubelet := stopKubelet()
			startKubelet()

			// If this test works correctly, each of these pods will exit
			// with no issue. But if accounting breaks, pods scheduled after
			// restart may think these old pods are consuming CPU and we
			// will get an OutOfCpu error.
			ginkgo.By("verifying restartNever pods succeed and restartAlways pods stay running")
			postRestartRunningPods := waitForPods(f, numAllPods, recoverTimeout)
			if len(postRestartRunningPods) < numAllPods {
				framework.Failf("less pods are running after node restart, got %d but expected %d", len(postRestartRunningPods), numAllPods)
			}
		})
	})
})
<commit_msg>Tweak resource requests for Kubelet restart test<commit_after>//go:build linux
// +build linux

/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package e2enode

import (
	"context"
	"fmt"
	"os/exec"
	"time"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/kubernetes/test/e2e/framework"
	e2eskipper "k8s.io/kubernetes/test/e2e/framework/skipper"
	testutils "k8s.io/kubernetes/test/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
)

// waitForPods waits for timeout duration, for podCount.
// If the timeout is hit, it returns the list of currently running pods.
func waitForPods(f *framework.Framework, podCount int, timeout time.Duration) (runningPods []*v1.Pod) {
	for start := time.Now(); time.Since(start) < timeout; time.Sleep(10 * time.Second) {
		podList, err := f.PodClient().List(context.TODO(), metav1.ListOptions{})
		if err != nil {
			framework.Logf("Failed to list pods on node: %v", err)
			continue
		}

		runningPods = []*v1.Pod{}
		for i := range podList.Items {
			pod := podList.Items[i]
			if r, err := testutils.PodRunningReadyOrSucceeded(&pod); err != nil || !r {
				continue
			}
			runningPods = append(runningPods, &pod)
		}
		framework.Logf("Running pod count %d", len(runningPods))
		if len(runningPods) >= podCount {
			break
		}
	}
	return runningPods
}

var _ = SIGDescribe("Restart [Serial] [Slow] [Disruptive]", func() {
	const (
		// Saturate the node. It's not necessary that all these pods enter
		// Running/Ready, because we don't know the number of cores in the
		// test node or default limits applied (if any). It's is essential
		// that no containers end up in terminated. 100 was chosen because
		// it's the max pods per node.
		podCount            = 100
		podCreationInterval = 100 * time.Millisecond
		recoverTimeout      = 5 * time.Minute
		startTimeout        = 5 * time.Minute
		// restartCount is chosen so even with minPods we exhaust the default
		// allocation of a /24.
		minPods      = 50
		restartCount = 6
	)

	f := framework.NewDefaultFramework("restart-test")
	ginkgo.Context("Container Runtime", func() {
		ginkgo.Context("Network", func() {
			ginkgo.It("should recover from ip leak", func() {
				if framework.TestContext.ContainerRuntime == "docker" {
					ginkgo.Skip("Test fails with in-tree docker. Skipping test.")
				}

				pods := newTestPods(podCount, false, imageutils.GetPauseImageName(), "restart-container-runtime-test")
				ginkgo.By(fmt.Sprintf("Trying to create %d pods on node", len(pods)))
				createBatchPodWithRateControl(f, pods, podCreationInterval)
				defer deletePodsSync(f, pods)

				// Give the node some time to stabilize, assume pods that enter RunningReady within
				// startTimeout fit on the node and the node is now saturated.
				runningPods := waitForPods(f, podCount, startTimeout)
				if len(runningPods) < minPods {
					framework.Failf("Failed to start %d pods, cannot test that restarting container runtime doesn't leak IPs", minPods)
				}

				for i := 0; i < restartCount; i++ {
					ginkgo.By(fmt.Sprintf("Killing container runtime iteration %d", i))
					// Wait for container runtime to be running
					var pid int
					gomega.Eventually(func() error {
						runtimePids, err := getPidsForProcess(framework.TestContext.ContainerRuntimeProcessName, framework.TestContext.ContainerRuntimePidFile)
						if err != nil {
							return err
						}
						if len(runtimePids) != 1 {
							return fmt.Errorf("unexpected container runtime pid list: %+v", runtimePids)
						}
						// Make sure the container runtime is running, pid got from pid file may not be running.
						pid = runtimePids[0]
						if _, err := exec.Command("sudo", "ps", "-p", fmt.Sprintf("%d", pid)).CombinedOutput(); err != nil {
							return err
						}
						return nil
					}, 1*time.Minute, 2*time.Second).Should(gomega.BeNil())
					if stdout, err := exec.Command("sudo", "kill", "-SIGKILL", fmt.Sprintf("%d", pid)).CombinedOutput(); err != nil {
						framework.Failf("Failed to kill container runtime (pid=%d): %v, stdout: %q", pid, err, string(stdout))
					}
					// Assume that container runtime will be restarted by systemd/supervisord etc.
					time.Sleep(20 * time.Second)
				}

				ginkgo.By("Checking currently Running/Ready pods")
				postRestartRunningPods := waitForPods(f, len(runningPods), recoverTimeout)
				if len(postRestartRunningPods) == 0 {
					framework.Failf("Failed to start *any* pods after container runtime restart, this might indicate an IP leak")
				}
				ginkgo.By("Confirm no containers have terminated")
				for _, pod := range postRestartRunningPods {
					if c := testutils.TerminatedContainers(pod); len(c) != 0 {
						framework.Failf("Pod %q has failed containers %+v after container runtime restart, this might indicate an IP leak", pod.Name, c)
					}
				}
				ginkgo.By(fmt.Sprintf("Container runtime restart test passed with %d pods", len(postRestartRunningPods)))
			})
		})
	})

	ginkgo.Context("Kubelet", func() {
		ginkgo.It("should correctly account for terminated pods after restart", func() {
			node := getLocalNode(f)
			cpus := node.Status.Allocatable[v1.ResourceCPU]
			numCpus := int((&cpus).Value())
			if numCpus < 1 {
				e2eskipper.Skipf("insufficient CPU available for kubelet restart test")
			}
			if numCpus > 18 {
				// 950m * 19 = 1805 CPUs -> not enough to block the scheduling of another 950m pod
				e2eskipper.Skipf("test will return false positives on a machine with >18 cores")
			}

			// create as many restartNever pods as there are allocatable CPU
			// nodes; if they are not correctly accounted for as terminated
			// later, this will fill up all node capacity
			podCountRestartNever := numCpus
			ginkgo.By(fmt.Sprintf("creating %d RestartNever pods on node", podCountRestartNever))
			restartNeverPods := newTestPods(podCountRestartNever, false, imageutils.GetE2EImage(imageutils.BusyBox), "restart-kubelet-test")
			for _, pod := range restartNeverPods {
				pod.Spec.RestartPolicy = "Never"
				pod.Spec.Containers[0].Command = []string{"echo", "hi"}
				pod.Spec.Containers[0].Resources.Limits = v1.ResourceList{
					v1.ResourceCPU: resource.MustParse("950m"),  // leave a little room for other workloads
				}
			}
			createBatchPodWithRateControl(f, restartNeverPods, podCreationInterval)
			defer deletePodsSync(f, restartNeverPods)

			completedPods := waitForPods(f, podCountRestartNever, startTimeout)
			if len(completedPods) < podCountRestartNever {
				framework.Failf("Failed to run sufficient restartNever pods, got %d but expected %d", len(completedPods), podCountRestartNever)
			}

			podCountRestartAlways := (numCpus / 2) + 1
			ginkgo.By(fmt.Sprintf("creating %d RestartAlways pods on node", podCountRestartAlways))
			restartAlwaysPods := newTestPods(podCountRestartAlways, false, imageutils.GetPauseImageName(), "restart-kubelet-test")
			for _, pod := range restartAlwaysPods {
				pod.Spec.Containers[0].Resources.Limits = v1.ResourceList{
					v1.ResourceCPU: resource.MustParse("1"),
				}
			}
			createBatchPodWithRateControl(f, restartAlwaysPods, podCreationInterval)
			defer deletePodsSync(f, restartAlwaysPods)

			numAllPods := podCountRestartNever + podCountRestartAlways
			allPods := waitForPods(f, numAllPods, startTimeout)
			if len(allPods) < numAllPods {
				framework.Failf("Failed to run sufficient restartAlways pods, got %d but expected %d", len(allPods), numAllPods)
			}

			ginkgo.By("killing and restarting kubelet")
			// We want to kill the kubelet rather than a graceful restart
			startKubelet := stopKubelet()
			startKubelet()

			// If this test works correctly, each of these pods will exit
			// with no issue. But if accounting breaks, pods scheduled after
			// restart may think these old pods are consuming CPU and we
			// will get an OutOfCpu error.
			ginkgo.By("verifying restartNever pods succeed and restartAlways pods stay running")
			postRestartRunningPods := waitForPods(f, numAllPods, recoverTimeout)
			if len(postRestartRunningPods) < numAllPods {
				framework.Failf("less pods are running after node restart, got %d but expected %d", len(postRestartRunningPods), numAllPods)
			}
		})
	})
})
<|endoftext|>


<commit_before>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2015 Douglas Hall
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package decode

import (
	"bytes"
	"fmt"
	"log"
	"math"
)

// PacketConfig specifies packet-specific radio configuration.
type PacketConfig struct {
	DataRate int

	BlockSize, BlockSize2       int
	SymbolLength, SymbolLength2 int
	SampleRate                  int

	PreambleSymbols, PacketSymbols int
	PreambleLength, PacketLength   int
	Preamble                       string

	BufferLength int

	CenterFreq uint32
}

func (cfg PacketConfig) Decimate(decimation int) PacketConfig {
	cfg.BlockSize /= decimation
	cfg.BlockSize2 /= decimation
	cfg.SymbolLength /= decimation
	cfg.SymbolLength2 /= decimation
	cfg.SampleRate /= decimation

	cfg.PreambleLength /= decimation
	cfg.PacketLength /= decimation

	cfg.BufferLength /= decimation

	return cfg
}

func (cfg PacketConfig) Log() {
	log.Println("BlockSize:", cfg.BlockSize)
	log.Println("CenterFreq:", cfg.CenterFreq)
	log.Println("SampleRate:", cfg.SampleRate)
	log.Println("DataRate:", cfg.DataRate)
	log.Println("SymbolLength:", cfg.SymbolLength)
	log.Println("PreambleSymbols:", cfg.PreambleSymbols)
	log.Println("PreambleLength:", cfg.PreambleLength)
	log.Println("PacketSymbols:", cfg.PacketSymbols)
	log.Println("PacketLength:", cfg.PacketLength)
	log.Println("Preamble:", cfg.Preamble)
}

// Decoder contains buffers and radio configuration.
type Decoder struct {
	Cfg PacketConfig

	Decimation int
	DecCfg     PacketConfig

	IQ        []byte
	Signal    []float64
	Filtered  []float64
	Quantized []byte

	csum  []float64
	demod Demodulator

	preamble []byte
	slices   [][]byte

	pkt []byte
}

// Create a new decoder with the given packet configuration.
func NewDecoder(cfg PacketConfig, decimation int, fastMag bool) (d Decoder) {
	d.Cfg = cfg

	d.Cfg.SymbolLength2 = d.Cfg.SymbolLength << 1
	d.Cfg.SampleRate = d.Cfg.DataRate * d.Cfg.SymbolLength

	d.Cfg.PreambleLength = d.Cfg.PreambleSymbols * d.Cfg.SymbolLength2
	d.Cfg.PacketLength = d.Cfg.PacketSymbols * d.Cfg.SymbolLength2

	d.Cfg.BlockSize = NextPowerOf2(d.Cfg.PreambleLength)
	d.Cfg.BlockSize2 = d.Cfg.BlockSize << 1

	d.Cfg.BufferLength = d.Cfg.PacketLength + d.Cfg.BlockSize

	d.Decimation = decimation
	d.DecCfg = d.Cfg.Decimate(d.Decimation)

	// Allocate necessary buffers.
	d.IQ = make([]byte, d.Cfg.BufferLength<<1)
	d.Signal = make([]float64, d.DecCfg.BufferLength)
	d.Filtered = make([]float64, d.DecCfg.BufferLength)
	d.Quantized = make([]byte, d.DecCfg.BufferLength)

	d.csum = make([]float64, (d.DecCfg.BlockSize + d.DecCfg.SymbolLength2 + 1))

	// Calculate magnitude lookup table specified by -fastmag flag.
	if fastMag {
		d.demod = NewAlphaMaxBetaMinLUT()
	} else {
		d.demod = NewSqrtMagLUT()
	}

	// Pre-calculate a byte-slice version of the preamble for searching.
	d.preamble = make([]byte, len(d.Cfg.Preamble))
	for idx := range d.Cfg.Preamble {
		if d.Cfg.Preamble[idx] == '1' {
			d.preamble[idx] = 1
		}
	}

	// Slice quantized sample buffer to make searching for the preamble more
	// memory local. Pre-allocate a flat buffer so memory is contiguous and
	// assign slices to the buffer.
	d.slices = make([][]byte, d.DecCfg.SymbolLength2)
	flat := make([]byte, d.DecCfg.BlockSize2-(d.DecCfg.BlockSize2%d.DecCfg.SymbolLength2))

	symbolsPerBlock := d.DecCfg.BlockSize2 / d.DecCfg.SymbolLength2
	for symbolOffset := range d.slices {
		lower := symbolOffset * symbolsPerBlock
		upper := (symbolOffset + 1) * symbolsPerBlock
		d.slices[symbolOffset] = flat[lower:upper]
	}

	// Signal up to the final stage is 1-bit per byte. Allocate a buffer to
	// store packed version 8-bits per byte.
	d.pkt = make([]byte, (d.DecCfg.PacketSymbols+7)>>3)

	return
}

// Decode accepts a sample block and performs various DSP techniques to extract a packet.
func (d Decoder) Decode(input []byte) []int {
	// Shift buffers to append new block.
	copy(d.IQ, d.IQ[d.Cfg.BlockSize<<1:])
	copy(d.Signal, d.Signal[d.DecCfg.BlockSize:])
	copy(d.Filtered, d.Filtered[d.DecCfg.BlockSize:])
	copy(d.Quantized, d.Quantized[d.DecCfg.BlockSize:])
	copy(d.IQ[d.Cfg.PacketLength<<1:], input[:])

	iqBlock := d.IQ[d.Cfg.PacketLength<<1:]
	signalBlock := d.Signal[d.DecCfg.PacketLength:]

	// Compute the magnitude of the new block.
	d.demod.Execute(iqBlock, signalBlock)

	signalBlock = d.Signal[d.DecCfg.PacketLength-d.DecCfg.SymbolLength2:]
	filterBlock := d.Filtered[d.DecCfg.PacketLength-d.DecCfg.SymbolLength2:]

	// Perform matched filter on new block.
	d.Filter(signalBlock, filterBlock)

	// Perform bit-decision on new block.
	Quantize(filterBlock, d.Quantized[d.DecCfg.PacketLength-d.DecCfg.SymbolLength2:])

	// Pack the quantized signal into slices for searching.
	d.Pack(d.Quantized[:d.DecCfg.BlockSize2], d.slices)

	// Return a list of indexes the preamble exists at.
	return d.Search(d.slices, d.preamble)
}

// A Demodulator knows how to demodulate an array of uint8 IQ samples into an
// array of float64 samples.
type Demodulator interface {
	Execute([]byte, []float64)
}

// Default Magnitude Lookup Table
type MagLUT []float64

// Pre-computes normalized squares with most common DC offset for rtl-sdr dongles.
func NewSqrtMagLUT() (lut MagLUT) {
	lut = make([]float64, 0x100)
	for idx := range lut {
		lut[idx] = 127.4 - float64(idx)
		lut[idx] *= lut[idx]
	}
	return
}

// Calculates complex magnitude on given IQ stream writing result to output.
func (lut MagLUT) Execute(input []byte, output []float64) {
	decIdx := 0
	dec := (len(input) / len(output))

	for idx := 0; decIdx < len(output); idx += dec {
		// fmt.Println(idx, idx/dec, len(input), len(output))
		output[decIdx] = math.Sqrt(lut[input[idx]] + lut[input[idx+1]])
		decIdx++
	}
}

// Alpha*Max + Beta*Min Magnitude Approximation Lookup Table.
type AlphaMaxBetaMinLUT []float64

// Pre-computes absolute values with most common DC offset for rtl-sdr dongles.
func NewAlphaMaxBetaMinLUT() (lut AlphaMaxBetaMinLUT) {
	lut = make([]float64, 0x100)
	for idx := range lut {
		lut[idx] = math.Abs(127.4 - float64(idx))
	}
	return
}

// Calculates complex magnitude on given IQ stream writing result to output.
func (lut AlphaMaxBetaMinLUT) Execute(input []byte, output []float64) {
	const (
		α = 0.948059448969
		ß = 0.392699081699
	)

	for idx := 0; idx < len(input); idx += 2 {
		i := lut[input[idx]]
		q := lut[input[idx+1]]
		if i > q {
			output[idx>>1] = α*i + ß*q
		} else {
			output[idx>>1] = α*q + ß*i
		}
	}
}

// Matched filter for Manchester coded signals. Output signal's sign at each
// sample determines the bit-value since Manchester symbols have odd symmetry.
func (d Decoder) Filter(input, output []float64) {
	// Computing the cumulative summation over the signal simplifies
	// filtering to the difference of a pair of subtractions.
	var sum float64
	for idx, v := range input {
		sum += v
		d.csum[idx+1] = sum
	}

	// Filter result is difference of summation of lower and upper symbols.
	lower := d.csum[d.DecCfg.SymbolLength:]
	upper := d.csum[d.DecCfg.SymbolLength2:]
	n := len(input) - d.DecCfg.SymbolLength2
	for idx := 0; idx < n; idx++ {
		output[idx] = (lower[idx] - d.csum[idx]) - (upper[idx] - lower[idx])
	}

	return
}

// Bit-value is determined by the sign of each sample after filtering.
func Quantize(input []float64, output []byte) {
	for idx, val := range input {
		output[idx] = byte(math.Float64bits(val)>>63) ^ 0x01
	}

	return
}

// Packs quantized signal into slices such that the first rank represents
// sample offsets and the second represents the value of each symbol from the
// given offset.
func (d Decoder) Pack(input []byte, slices [][]byte) {
	for symbolOffset, slice := range slices {
		for symbolIdx := range slice {
			slice[symbolIdx] = input[symbolIdx*d.DecCfg.SymbolLength2+symbolOffset]
		}
	}

	return
}

// For each sample offset look for the preamble. Return a list of indexes the
// preamble is found at. Indexes are absolute in the unsliced quantized
// buffer.
func (d Decoder) Search(slices [][]byte, preamble []byte) (indexes []int) {
	preambleLength := len(preamble)
	for symbolOffset, slice := range slices {
		for symbolIdx := range slice[:len(slice)-preambleLength] {
			if bytes.Equal(preamble, slice[symbolIdx:][:preambleLength]) {
				indexes = append(indexes, symbolIdx*d.DecCfg.SymbolLength2+symbolOffset)
			}
		}
	}

	return
}

// Given a list of indeces the preamble exists at, sample the appropriate bits
// of the signal's bit-decision. Pack bits of each index into an array of byte
// arrays and return.
func (d Decoder) Slice(indices []int) (pkts [][]byte) {
	// We will likely find multiple instances of the message so only keep
	// track of unique instances.
	seen := make(map[string]bool)

	// For each of the indices the preamble exists at.
	for _, qIdx := range indices {
		// Check that we're still within the first sample block. We'll catch
		// the message on the next sample block otherwise.
		if qIdx > d.DecCfg.BlockSize {
			continue
		}

		// Packet is 1 bit per byte, pack to 8-bits per byte.
		for pIdx := 0; pIdx < d.DecCfg.PacketSymbols; pIdx++ {
			d.pkt[pIdx>>3] <<= 1
			d.pkt[pIdx>>3] |= d.Quantized[qIdx+(pIdx*d.DecCfg.SymbolLength2)]
		}

		// Store the packet in the seen map and append to the packet list.
		pktStr := fmt.Sprintf("%02X", d.pkt)
		if !seen[pktStr] {
			seen[pktStr] = true
			pkts = append(pkts, make([]byte, len(d.pkt)))
			copy(pkts[len(pkts)-1], d.pkt)
		}
	}

	return
}

func NextPowerOf2(v int) int {
	return 1 << uint(math.Ceil(math.Log2(float64(v))))
}
<commit_msg>Implement decimation in AlphaMaxBetaMinLUT.<commit_after>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2015 Douglas Hall
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package decode

import (
	"bytes"
	"fmt"
	"log"
	"math"
)

// PacketConfig specifies packet-specific radio configuration.
type PacketConfig struct {
	DataRate int

	BlockSize, BlockSize2       int
	SymbolLength, SymbolLength2 int
	SampleRate                  int

	PreambleSymbols, PacketSymbols int
	PreambleLength, PacketLength   int
	Preamble                       string

	BufferLength int

	CenterFreq uint32
}

func (cfg PacketConfig) Decimate(decimation int) PacketConfig {
	cfg.BlockSize /= decimation
	cfg.BlockSize2 /= decimation
	cfg.SymbolLength /= decimation
	cfg.SymbolLength2 /= decimation
	cfg.SampleRate /= decimation

	cfg.PreambleLength /= decimation
	cfg.PacketLength /= decimation

	cfg.BufferLength /= decimation

	return cfg
}

func (cfg PacketConfig) Log() {
	log.Println("BlockSize:", cfg.BlockSize)
	log.Println("CenterFreq:", cfg.CenterFreq)
	log.Println("SampleRate:", cfg.SampleRate)
	log.Println("DataRate:", cfg.DataRate)
	log.Println("SymbolLength:", cfg.SymbolLength)
	log.Println("PreambleSymbols:", cfg.PreambleSymbols)
	log.Println("PreambleLength:", cfg.PreambleLength)
	log.Println("PacketSymbols:", cfg.PacketSymbols)
	log.Println("PacketLength:", cfg.PacketLength)
	log.Println("Preamble:", cfg.Preamble)
}

// Decoder contains buffers and radio configuration.
type Decoder struct {
	Cfg PacketConfig

	Decimation int
	DecCfg     PacketConfig

	IQ        []byte
	Signal    []float64
	Filtered  []float64
	Quantized []byte

	csum  []float64
	demod Demodulator

	preamble []byte
	slices   [][]byte

	pkt []byte
}

// Create a new decoder with the given packet configuration.
func NewDecoder(cfg PacketConfig, decimation int, fastMag bool) (d Decoder) {
	d.Cfg = cfg

	d.Cfg.SymbolLength2 = d.Cfg.SymbolLength << 1
	d.Cfg.SampleRate = d.Cfg.DataRate * d.Cfg.SymbolLength

	d.Cfg.PreambleLength = d.Cfg.PreambleSymbols * d.Cfg.SymbolLength2
	d.Cfg.PacketLength = d.Cfg.PacketSymbols * d.Cfg.SymbolLength2

	d.Cfg.BlockSize = NextPowerOf2(d.Cfg.PreambleLength)
	d.Cfg.BlockSize2 = d.Cfg.BlockSize << 1

	d.Cfg.BufferLength = d.Cfg.PacketLength + d.Cfg.BlockSize

	d.Decimation = decimation
	d.DecCfg = d.Cfg.Decimate(d.Decimation)

	// Allocate necessary buffers.
	d.IQ = make([]byte, d.Cfg.BufferLength<<1)
	d.Signal = make([]float64, d.DecCfg.BufferLength)
	d.Filtered = make([]float64, d.DecCfg.BufferLength)
	d.Quantized = make([]byte, d.DecCfg.BufferLength)

	d.csum = make([]float64, (d.DecCfg.BlockSize + d.DecCfg.SymbolLength2 + 1))

	// Calculate magnitude lookup table specified by -fastmag flag.
	if fastMag {
		d.demod = NewAlphaMaxBetaMinLUT()
	} else {
		d.demod = NewSqrtMagLUT()
	}

	// Pre-calculate a byte-slice version of the preamble for searching.
	d.preamble = make([]byte, len(d.Cfg.Preamble))
	for idx := range d.Cfg.Preamble {
		if d.Cfg.Preamble[idx] == '1' {
			d.preamble[idx] = 1
		}
	}

	// Slice quantized sample buffer to make searching for the preamble more
	// memory local. Pre-allocate a flat buffer so memory is contiguous and
	// assign slices to the buffer.
	d.slices = make([][]byte, d.DecCfg.SymbolLength2)
	flat := make([]byte, d.DecCfg.BlockSize2-(d.DecCfg.BlockSize2%d.DecCfg.SymbolLength2))

	symbolsPerBlock := d.DecCfg.BlockSize2 / d.DecCfg.SymbolLength2
	for symbolOffset := range d.slices {
		lower := symbolOffset * symbolsPerBlock
		upper := (symbolOffset + 1) * symbolsPerBlock
		d.slices[symbolOffset] = flat[lower:upper]
	}

	// Signal up to the final stage is 1-bit per byte. Allocate a buffer to
	// store packed version 8-bits per byte.
	d.pkt = make([]byte, (d.DecCfg.PacketSymbols+7)>>3)

	return
}

// Decode accepts a sample block and performs various DSP techniques to extract a packet.
func (d Decoder) Decode(input []byte) []int {
	// Shift buffers to append new block.
	copy(d.IQ, d.IQ[d.Cfg.BlockSize<<1:])
	copy(d.Signal, d.Signal[d.DecCfg.BlockSize:])
	copy(d.Filtered, d.Filtered[d.DecCfg.BlockSize:])
	copy(d.Quantized, d.Quantized[d.DecCfg.BlockSize:])
	copy(d.IQ[d.Cfg.PacketLength<<1:], input[:])

	iqBlock := d.IQ[d.Cfg.PacketLength<<1:]
	signalBlock := d.Signal[d.DecCfg.PacketLength:]

	// Compute the magnitude of the new block.
	d.demod.Execute(iqBlock, signalBlock)

	signalBlock = d.Signal[d.DecCfg.PacketLength-d.DecCfg.SymbolLength2:]
	filterBlock := d.Filtered[d.DecCfg.PacketLength-d.DecCfg.SymbolLength2:]

	// Perform matched filter on new block.
	d.Filter(signalBlock, filterBlock)

	// Perform bit-decision on new block.
	Quantize(filterBlock, d.Quantized[d.DecCfg.PacketLength-d.DecCfg.SymbolLength2:])

	// Pack the quantized signal into slices for searching.
	d.Pack(d.Quantized[:d.DecCfg.BlockSize2], d.slices)

	// Return a list of indexes the preamble exists at.
	return d.Search(d.slices, d.preamble)
}

// A Demodulator knows how to demodulate an array of uint8 IQ samples into an
// array of float64 samples.
type Demodulator interface {
	Execute([]byte, []float64)
}

// Default Magnitude Lookup Table
type MagLUT []float64

// Pre-computes normalized squares with most common DC offset for rtl-sdr dongles.
func NewSqrtMagLUT() (lut MagLUT) {
	lut = make([]float64, 0x100)
	for idx := range lut {
		lut[idx] = 127.4 - float64(idx)
		lut[idx] *= lut[idx]
	}
	return
}

// Calculates complex magnitude on given IQ stream writing result to output.
func (lut MagLUT) Execute(input []byte, output []float64) {
	decIdx := 0
	dec := (len(input) / len(output))

	for idx := 0; decIdx < len(output); idx += dec {
		output[decIdx] = math.Sqrt(lut[input[idx]] + lut[input[idx+1]])
		decIdx++
	}
}

// Alpha*Max + Beta*Min Magnitude Approximation Lookup Table.
type AlphaMaxBetaMinLUT []float64

// Pre-computes absolute values with most common DC offset for rtl-sdr dongles.
func NewAlphaMaxBetaMinLUT() (lut AlphaMaxBetaMinLUT) {
	lut = make([]float64, 0x100)
	for idx := range lut {
		lut[idx] = math.Abs(127.4 - float64(idx))
	}
	return
}

// Calculates complex magnitude on given IQ stream writing result to output.
func (lut AlphaMaxBetaMinLUT) Execute(input []byte, output []float64) {
	const (
		α = 0.948059448969
		ß = 0.392699081699
	)

	decIdx := 0
	dec := (len(input) / len(output))

	for idx := 0; decIdx < len(output); idx += dec {
		i := lut[input[idx]]
		q := lut[input[idx+1]]
		if i > q {
			output[decIdx] = α*i + ß*q
		} else {
			output[decIdx] = α*q + ß*i
		}
		decIdx++
	}
}

// Matched filter for Manchester coded signals. Output signal's sign at each
// sample determines the bit-value since Manchester symbols have odd symmetry.
func (d Decoder) Filter(input, output []float64) {
	// Computing the cumulative summation over the signal simplifies
	// filtering to the difference of a pair of subtractions.
	var sum float64
	for idx, v := range input {
		sum += v
		d.csum[idx+1] = sum
	}

	// Filter result is difference of summation of lower and upper symbols.
	lower := d.csum[d.DecCfg.SymbolLength:]
	upper := d.csum[d.DecCfg.SymbolLength2:]
	n := len(input) - d.DecCfg.SymbolLength2
	for idx := 0; idx < n; idx++ {
		output[idx] = (lower[idx] - d.csum[idx]) - (upper[idx] - lower[idx])
	}

	return
}

// Bit-value is determined by the sign of each sample after filtering.
func Quantize(input []float64, output []byte) {
	for idx, val := range input {
		output[idx] = byte(math.Float64bits(val)>>63) ^ 0x01
	}

	return
}

// Packs quantized signal into slices such that the first rank represents
// sample offsets and the second represents the value of each symbol from the
// given offset.
func (d Decoder) Pack(input []byte, slices [][]byte) {
	for symbolOffset, slice := range slices {
		for symbolIdx := range slice {
			slice[symbolIdx] = input[symbolIdx*d.DecCfg.SymbolLength2+symbolOffset]
		}
	}

	return
}

// For each sample offset look for the preamble. Return a list of indexes the
// preamble is found at. Indexes are absolute in the unsliced quantized
// buffer.
func (d Decoder) Search(slices [][]byte, preamble []byte) (indexes []int) {
	preambleLength := len(preamble)
	for symbolOffset, slice := range slices {
		for symbolIdx := range slice[:len(slice)-preambleLength] {
			if bytes.Equal(preamble, slice[symbolIdx:][:preambleLength]) {
				indexes = append(indexes, symbolIdx*d.DecCfg.SymbolLength2+symbolOffset)
			}
		}
	}

	return
}

// Given a list of indeces the preamble exists at, sample the appropriate bits
// of the signal's bit-decision. Pack bits of each index into an array of byte
// arrays and return.
func (d Decoder) Slice(indices []int) (pkts [][]byte) {
	// We will likely find multiple instances of the message so only keep
	// track of unique instances.
	seen := make(map[string]bool)

	// For each of the indices the preamble exists at.
	for _, qIdx := range indices {
		// Check that we're still within the first sample block. We'll catch
		// the message on the next sample block otherwise.
		if qIdx > d.DecCfg.BlockSize {
			continue
		}

		// Packet is 1 bit per byte, pack to 8-bits per byte.
		for pIdx := 0; pIdx < d.DecCfg.PacketSymbols; pIdx++ {
			d.pkt[pIdx>>3] <<= 1
			d.pkt[pIdx>>3] |= d.Quantized[qIdx+(pIdx*d.DecCfg.SymbolLength2)]
		}

		// Store the packet in the seen map and append to the packet list.
		pktStr := fmt.Sprintf("%02X", d.pkt)
		if !seen[pktStr] {
			seen[pktStr] = true
			pkts = append(pkts, make([]byte, len(d.pkt)))
			copy(pkts[len(pkts)-1], d.pkt)
		}
	}

	return
}

func NextPowerOf2(v int) int {
	return 1 << uint(math.Ceil(math.Log2(float64(v))))
}
<|endoftext|>