<commit_before>package brokers_test

import (
	"errors"
	"github.com/RichardKnop/machinery/v1"
	"github.com/RichardKnop/machinery/v1/brokers"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/retry"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

var (
	testBroker           brokers.Interface
	testAWSSQSBroker     *brokers.AWSSQSBroker
	errAWSSQSBroker      *brokers.AWSSQSBroker
	cnf                  *config.Config
	receiveMessageOutput *sqs.ReceiveMessageOutput
)

func init() {
	testAWSSQSBroker = brokers.TestAWSSQSBroker
	errAWSSQSBroker = brokers.ErrAWSSQSBroker
	cnf = brokers.TestConf
	receiveMessageOutput = brokers.ReceiveMessageOutput
	testBroker = brokers.NewAWSSQSBroker(cnf)
}

func TestNewAWSSQSBroker(t *testing.T) {
	assert.IsType(t, testAWSSQSBroker, testBroker)
}

func TestPrivateFunc_continueReceivingMessages(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	firstStep := make(chan int)
	nextStep := make(chan int)
	go func() {
		stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
		firstStep <- 1
		stopReceivingChan <- 1
	}()

	var (
		whetherContinue bool
		err             error
	)
	<-firstStep
	// Test the case that a signal was received from stopReceivingChan
	go func() {
		whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
		nextStep <- 1
	}()
	<-nextStep
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	// Test the default condition
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	// Test the error
	whetherContinue, err = errAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.NotNil(t, err)

	// Test when there is no message
	outputCopy := *receiveMessageOutput
	receiveMessageOutput.Messages = []*sqs.Message{}
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)
	// recover original value
	*receiveMessageOutput = outputCopy

}

func TestPrivateFunc_consume(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()

	// an infinite loop will be executed only when there is no error
	err = testAWSSQSBroker.ConsumeForTest(deliveries, 0, wk)
	assert.NotNil(t, err)

}

func TestPrivateFunc_consumeOne(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	err = testAWSSQSBroker.ConsumeOneForTest(receiveMessageOutput, wk)
	assert.NotNil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)

	outputCopy.Messages = []*sqs.Message{
		{
			Body: aws.String("foo message"),
		},
	}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)
}

func TestPrivateFunc_initializePool(t *testing.T) {
	concurrency := 9
	pool := make(chan struct{}, concurrency)
	testAWSSQSBroker.InitializePoolForTest(pool, concurrency)
	assert.Len(t, pool, concurrency)
}

func TestPrivateFunc_startConsuming(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("sms_worker", 0)
	retryFunc := testAWSSQSBroker.GetRetryFuncForTest()
	stopChan := testAWSSQSBroker.GetStopChanForTest()
	retryStopChan := testAWSSQSBroker.GetRetryStopChanForTest()
	assert.Nil(t, retryFunc)
	testAWSSQSBroker.StartConsumingForTest("fooTag", wk)
	assert.IsType(t, retryFunc, retry.Closure())
	assert.Equal(t, len(stopChan), 0)
	assert.Equal(t, len(retryStopChan), 0)
}

func TestPrivateFuncDefaultQueueURL(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()

	assert.EqualValues(t, *qURL, "https://sqs.foo.amazonaws.com.cn/test_queue")
}

func TestPrivateFunc_stopReceiving(t *testing.T) {
	go testAWSSQSBroker.StopReceivingForTest()
	stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
	assert.NotNil(t, <-stopReceivingChan)
}

func TestPrivateFunc_receiveMessage(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	output, err := testAWSSQSBroker.ReceiveMessageForTest(qURL)
	assert.Nil(t, err)
	assert.Equal(t, receiveMessageOutput, output)
}

func TestPrivateFunc_consumeDeliveries(t *testing.T) {
	concurrency := 0
	pool := make(chan struct{}, concurrency)
	errorsChan := make(chan error)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	go func() { deliveries <- receiveMessageOutput }()
	whetherContinue, err := testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	go func() { errorsChan <- errors.New("foo error") }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.NotNil(t, err)

	go func() { testAWSSQSBroker.GetStopChanForTest() <- 1 }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	e := <-errorsChan
	assert.True(t, whetherContinue)
	assert.NotNil(t, e)
	assert.Nil(t, err)

	go func() { deliveries <- receiveMessageOutput }()
	// wait for a while to fix racing problem
	time.Sleep(5 * time.Second)
	go func() { pool <- struct{}{} }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	p := <-pool
	assert.True(t, whetherContinue)
	assert.NotNil(t, p)
	assert.Nil(t, err)
}

func TestPrivateFunc_deleteOne(t *testing.T) {
	err := testAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.Nil(t, err)

	err = errAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.NotNil(t, err)
}
<commit_msg>use wait group to replace time waiting<commit_after>package brokers_test

import (
	"errors"
	"github.com/RichardKnop/machinery/v1"
	"github.com/RichardKnop/machinery/v1/brokers"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/retry"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/stretchr/testify/assert"
	"sync"
	"testing"
)

var (
	testBroker           brokers.Interface
	testAWSSQSBroker     *brokers.AWSSQSBroker
	errAWSSQSBroker      *brokers.AWSSQSBroker
	cnf                  *config.Config
	receiveMessageOutput *sqs.ReceiveMessageOutput
)

func init() {
	testAWSSQSBroker = brokers.TestAWSSQSBroker
	errAWSSQSBroker = brokers.ErrAWSSQSBroker
	cnf = brokers.TestConf
	receiveMessageOutput = brokers.ReceiveMessageOutput
	testBroker = brokers.NewAWSSQSBroker(cnf)
}

func TestNewAWSSQSBroker(t *testing.T) {
	assert.IsType(t, testAWSSQSBroker, testBroker)
}

func TestPrivateFunc_continueReceivingMessages(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	firstStep := make(chan int)
	nextStep := make(chan int)
	go func() {
		stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
		firstStep <- 1
		stopReceivingChan <- 1
	}()

	var (
		whetherContinue bool
		err             error
	)
	<-firstStep
	// Test the case that a signal was received from stopReceivingChan
	go func() {
		whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
		nextStep <- 1
	}()
	<-nextStep
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	// Test the default condition
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	// Test the error
	whetherContinue, err = errAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.NotNil(t, err)

	// Test when there is no message
	outputCopy := *receiveMessageOutput
	receiveMessageOutput.Messages = []*sqs.Message{}
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)
	// recover original value
	*receiveMessageOutput = outputCopy

}

func TestPrivateFunc_consume(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()

	// an infinite loop will be executed only when there is no error
	err = testAWSSQSBroker.ConsumeForTest(deliveries, 0, wk)
	assert.NotNil(t, err)

}

func TestPrivateFunc_consumeOne(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	err = testAWSSQSBroker.ConsumeOneForTest(receiveMessageOutput, wk)
	assert.NotNil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)

	outputCopy.Messages = []*sqs.Message{
		{
			Body: aws.String("foo message"),
		},
	}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)
}

func TestPrivateFunc_initializePool(t *testing.T) {
	concurrency := 9
	pool := make(chan struct{}, concurrency)
	testAWSSQSBroker.InitializePoolForTest(pool, concurrency)
	assert.Len(t, pool, concurrency)
}

func TestPrivateFunc_startConsuming(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("sms_worker", 0)
	retryFunc := testAWSSQSBroker.GetRetryFuncForTest()
	stopChan := testAWSSQSBroker.GetStopChanForTest()
	retryStopChan := testAWSSQSBroker.GetRetryStopChanForTest()
	assert.Nil(t, retryFunc)
	testAWSSQSBroker.StartConsumingForTest("fooTag", wk)
	assert.IsType(t, retryFunc, retry.Closure())
	assert.Equal(t, len(stopChan), 0)
	assert.Equal(t, len(retryStopChan), 0)
}

func TestPrivateFuncDefaultQueueURL(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()

	assert.EqualValues(t, *qURL, "https://sqs.foo.amazonaws.com.cn/test_queue")
}

func TestPrivateFunc_stopReceiving(t *testing.T) {
	go testAWSSQSBroker.StopReceivingForTest()
	stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
	assert.NotNil(t, <-stopReceivingChan)
}

func TestPrivateFunc_receiveMessage(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	output, err := testAWSSQSBroker.ReceiveMessageForTest(qURL)
	assert.Nil(t, err)
	assert.Equal(t, receiveMessageOutput, output)
}

func TestPrivateFunc_consumeDeliveries(t *testing.T) {
	concurrency := 0
	pool := make(chan struct{}, concurrency)
	errorsChan := make(chan error)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	go func() { deliveries <- receiveMessageOutput }()
	whetherContinue, err := testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	go func() { errorsChan <- errors.New("foo error") }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.NotNil(t, err)

	go func() { testAWSSQSBroker.GetStopChanForTest() <- 1 }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	e := <-errorsChan
	assert.True(t, whetherContinue)
	assert.NotNil(t, e)
	assert.Nil(t, err)

	// using a wait group and a channel to fix the racing problem
	var wg sync.WaitGroup
	wg.Add(1)
	nextStep := make(chan bool, 1)
	go func() {
		defer wg.Done()
		nextStep <- true
		deliveries <- receiveMessageOutput
	}()
	if <-nextStep {
		// <-pool will block the routine in the following steps, so pool <- struct{}{} will be executed for sure
		go func() { wg.Wait(); pool <- struct{}{} }()
	}
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	// the pool shouldn't be consumed
	p := <-pool
	assert.True(t, whetherContinue)
	assert.NotNil(t, p)
	assert.Nil(t, err)
}

func TestPrivateFunc_deleteOne(t *testing.T) {
	err := testAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.Nil(t, err)

	err = errAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.NotNil(t, err)
}
<|endoftext|>


<commit_before>package azure

import (
	"context"
	"fmt"

	"github.com/Azure/azure-sdk-for-go/services/storage/mgmt/2019-04-01/storage"
	"github.com/Azure/go-autorest/autorest"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
	"github.com/hashicorp/go-azure-helpers/authentication"
)

const (
	containerIDFormat = "https://%s.blob.core.windows.net/%s"
)

type StorageContainerGenerator struct {
	AzureService
}

func NewStorageContainerGenerator(subscriptionID string, authorizer autorest.Authorizer, rg string) *StorageContainerGenerator {
	storageContainerGenerator := new(StorageContainerGenerator)
	storageContainerGenerator.Args = map[string]interface{}{}
	storageContainerGenerator.Args["config"] = authentication.Config{SubscriptionID: subscriptionID}
	storageContainerGenerator.Args["authorizer"] = authorizer
	storageContainerGenerator.Args["resource_group"] = rg

	return storageContainerGenerator
}

func (g StorageContainerGenerator) ListBlobContainers() ([]terraformutils.Resource, error) {
	var containerResources []terraformutils.Resource
	blobContainersClient := storage.NewBlobContainersClient(g.Args["config"].(authentication.Config).SubscriptionID)
	blobContainersClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)
	ctx := context.Background()

	accounts, err := g.getStorageAccounts()
	if err != nil {
		return containerResources, err
	}

	for _, storageAccount := range accounts {
		parsedStorageAccountResourceID, err := ParseAzureResourceID(*storageAccount.ID)
		if err != nil {
			break
		}
		containerItemsIterator, err := blobContainersClient.ListComplete(ctx, parsedStorageAccountResourceID.ResourceGroup, *storageAccount.Name, "", "", "")
		if err != nil {
			return containerResources, err
		}

		for containerItemsIterator.NotDone() {
			containerItem := containerItemsIterator.Value()
			containerResources = append(containerResources,
				terraformutils.NewResource(
					fmt.Sprintf(containerIDFormat, *storageAccount.Name, *containerItem.Name),
					*containerItem.Name,
					"azurerm_storage_container",
					"azurerm",
					map[string]string{
						"storage_account_name": *storageAccount.Name,
						"name":                 *containerItem.Name,
					},
					[]string{},
					map[string]interface{}{}))

			if err := containerItemsIterator.NextWithContext(ctx); err != nil {
				return containerResources, err
			}
		}
	}

	return containerResources, nil
}

func (g *StorageContainerGenerator) getStorageAccounts() ([]storage.Account, error) {
	ctx := context.Background()
	accountsClient := storage.NewAccountsClient(g.Args["config"].(authentication.Config).SubscriptionID)

	accountsClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)
	var accounts []storage.Account
	if rg := g.Args["resource_group"].(string); rg != "" {
		accountsResult, err := accountsClient.ListByResourceGroup(ctx, rg)
		if err != nil {
			return nil, err
		}
		if paccounts := accountsResult.Value; paccounts != nil {
			for _, account := range *paccounts {
				accounts = append(accounts, account)
			}
		}
	} else {
		accountsIterator, err := accountsClient.ListComplete(ctx)
		if err != nil {
			return nil, err
		}
		for accountsIterator.NotDone() {
			account := accountsIterator.Value()
			accounts = append(accounts, account)
			if err := accountsIterator.NextWithContext(ctx); err != nil {
				return accounts, err
			}
		}
	}

	return accounts, nil
}

func (g *StorageContainerGenerator) InitResources() error {
	storageAccounts, err := g.ListBlobContainers()
	if err != nil {
		return err
	}

	g.Resources = storageAccounts

	return nil
}
<commit_msg>fix lint error<commit_after>package azure

import (
	"context"
	"fmt"

	"github.com/Azure/azure-sdk-for-go/services/storage/mgmt/2019-04-01/storage"
	"github.com/Azure/go-autorest/autorest"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
	"github.com/hashicorp/go-azure-helpers/authentication"
)

const (
	containerIDFormat = "https://%s.blob.core.windows.net/%s"
)

type StorageContainerGenerator struct {
	AzureService
}

func NewStorageContainerGenerator(subscriptionID string, authorizer autorest.Authorizer, rg string) *StorageContainerGenerator {
	storageContainerGenerator := new(StorageContainerGenerator)
	storageContainerGenerator.Args = map[string]interface{}{}
	storageContainerGenerator.Args["config"] = authentication.Config{SubscriptionID: subscriptionID}
	storageContainerGenerator.Args["authorizer"] = authorizer
	storageContainerGenerator.Args["resource_group"] = rg

	return storageContainerGenerator
}

func (g StorageContainerGenerator) ListBlobContainers() ([]terraformutils.Resource, error) {
	var containerResources []terraformutils.Resource
	blobContainersClient := storage.NewBlobContainersClient(g.Args["config"].(authentication.Config).SubscriptionID)
	blobContainersClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)
	ctx := context.Background()

	accounts, err := g.getStorageAccounts()
	if err != nil {
		return containerResources, err
	}

	for _, storageAccount := range accounts {
		parsedStorageAccountResourceID, err := ParseAzureResourceID(*storageAccount.ID)
		if err != nil {
			break
		}
		containerItemsIterator, err := blobContainersClient.ListComplete(ctx, parsedStorageAccountResourceID.ResourceGroup, *storageAccount.Name, "", "", "")
		if err != nil {
			return containerResources, err
		}

		for containerItemsIterator.NotDone() {
			containerItem := containerItemsIterator.Value()
			containerResources = append(containerResources,
				terraformutils.NewResource(
					fmt.Sprintf(containerIDFormat, *storageAccount.Name, *containerItem.Name),
					*containerItem.Name,
					"azurerm_storage_container",
					"azurerm",
					map[string]string{
						"storage_account_name": *storageAccount.Name,
						"name":                 *containerItem.Name,
					},
					[]string{},
					map[string]interface{}{}))

			if err := containerItemsIterator.NextWithContext(ctx); err != nil {
				return containerResources, err
			}
		}
	}

	return containerResources, nil
}

func (g *StorageContainerGenerator) getStorageAccounts() ([]storage.Account, error) {
	ctx := context.Background()
	accountsClient := storage.NewAccountsClient(g.Args["config"].(authentication.Config).SubscriptionID)

	accountsClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)
	var accounts []storage.Account
	if rg := g.Args["resource_group"].(string); rg != "" {
		accountsResult, err := accountsClient.ListByResourceGroup(ctx, rg)
		if err != nil {
			return nil, err
		}
		if paccounts := accountsResult.Value; paccounts != nil {
			accounts = append(accounts, *paccount...)
		}
	} else {
		accountsIterator, err := accountsClient.ListComplete(ctx)
		if err != nil {
			return nil, err
		}
		for accountsIterator.NotDone() {
			account := accountsIterator.Value()
			accounts = append(accounts, account)
			if err := accountsIterator.NextWithContext(ctx); err != nil {
				return accounts, err
			}
		}
	}

	return accounts, nil
}

func (g *StorageContainerGenerator) InitResources() error {
	storageAccounts, err := g.ListBlobContainers()
	if err != nil {
		return err
	}

	g.Resources = storageAccounts

	return nil
}
<|endoftext|>


<commit_before>// This package implements a provisioner for Packer that executes
// shell scripts within the remote machine.
package chefSolo

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"github.com/mitchellh/iochan"
	"github.com/mitchellh/mapstructure"
	"github.com/mitchellh/packer/packer"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"
	"path/filepath"
	"encoding/json"
)

const RemoteStagingPath = "/tmp/provision/chef-solo"
const RemoteFileCachePath = "/tmp/provision/chef-solo"
const RemoteCookbookPath = "/tmp/provision/chef-solo/cookbooks"
const DefaultCookbookPath = "cookbooks"

var Ui packer.Ui

type config struct {
	// An array of local paths of cookbooks to upload.
	CookbookPaths []string `mapstructure:"cookbook_paths"`

	// The local path of the cookbooks to upload.
	CookbookPath string `mapstructure:"cookbook_path"`

	// An array of recipes to run.
	RunList []string `mapstructure:"run_list"`

	// A string of JSON that will be used as the JSON attributes for the
	// Chef run.
	Json map[string]interface{}

	UseSudo bool `mapstructure:"use_sudo"`

	// If true
	SkipInstall bool `mapstructure:"skip_install"`
}

type Provisioner struct {
	config config
}

type ExecuteRecipeTemplate struct {
	SoloRbPath string
	JsonPath string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	errs := make([]error, 0)
	for _, raw := range raws {
		if err := mapstructure.Decode(raw, &p.config); err != nil {
			return err
		}
	}

	if p.config.CookbookPaths == nil {
		p.config.CookbookPaths = make([]string, 0)
	}

	if p.config.CookbookPath == "" {
		p.config.CookbookPath = DefaultCookbookPath
	}

	if len(p.config.CookbookPaths) > 0 && p.config.CookbookPath != "" {
		errs = append(errs, errors.New("Only one of cookbooks or cookbook can be specified."))
	}

	if len(p.config.CookbookPaths) == 0 {
		p.config.CookbookPaths = append(p.config.CookbookPaths, p.config.CookbookPath)
	}

	if p.config.RunList == nil {
		p.config.RunList = make([]string, 0)
	}

	if p.config.Json != nil {
		if _, err := json.Marshal(p.config.Json); err != nil {
			errs = append(errs, fmt.Errorf("Bad JSON: %s", err))
		}
	} else {
		p.config.Json = make(map[string]interface{})
	}

	for _, path := range p.config.CookbookPaths {
		pFileInfo, err := os.Stat(path)

		if err != nil || !pFileInfo.IsDir() {
			errs = append(errs, fmt.Errorf("Bad cookbook path '%s': %s", path, err))
		}
	}

	if len(errs) > 0 {
		return &packer.MultiError{errs}
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	cookbookPaths := make([]string, len(p.config.CookbookPaths))
	copy(cookbookPaths, p.config.CookbookPaths)

	Ui = ui

	// Generic setup for Chef runs
	err := InstallChefSolo(comm)
	if err != nil {
		return fmt.Errorf("Error installing Chef Solo: %s", err)
	}

	err = CreateRemoteDirectory(RemoteCookbookPath, comm)
	if err != nil {
		return fmt.Errorf("Error creating remote staging directory: %s", err)
	}

	soloRbPath, err := CreateSoloRb(p.config.CookbookPaths, comm)
	if err != nil {
		return fmt.Errorf("Error creating Chef Solo configuration file: %s", err)
	}

	jsonPath, err := CreateAttributesJson(p.config.Json, p.config.RunList, comm)
	if err != nil {
		return fmt.Errorf("Error uploading JSON attributes file: %s", err)
	}

	// Upload all cookbooks
	for _, path := range cookbookPaths {
		ui.Say(fmt.Sprintf("Copying cookbook path: %s", path))
		err = UploadLocalDirectory(path, comm)
		if err != nil {
			return fmt.Errorf("Error uploading cookbooks: %s", err)
		}
	}

	// Execute requested recipes
	ui.Say("Beginning Chef Solo run")

	// Compile the command
	var command bytes.Buffer
	t := template.Must(template.New("chef-run").Parse("sudo chef-solo --no-color -c {{.SoloRbPath}} -j {{.JsonPath}}"))
	t.Execute(&command, &ExecuteRecipeTemplate{soloRbPath, jsonPath})

	err = executeCommand(command.String(), comm)
	if err != nil {
		return fmt.Errorf("Error running Chef Solo: %s", err)
	}

	return fmt.Errorf("Die")

	return nil
}

func UploadLocalDirectory(localDir string, comm packer.Communicator) (err error) {
	visitPath := func (path string, f os.FileInfo, err error) (err2 error) {
		var remotePath = RemoteCookbookPath + "/" + path
		if f.IsDir() {
			// Make remote directory
			err = CreateRemoteDirectory(remotePath, comm)
			if err != nil {
				return err
			}
		} else {
			// Upload file to existing directory
			file, err := os.Open(path)
			if err != nil {
				return fmt.Errorf("Error opening file: %s", err)
			}

			err = comm.Upload(remotePath, file)
			if err != nil {
				return fmt.Errorf("Error uploading file: %s", err)
			}
		}
		return
	}

	err = filepath.Walk(localDir, visitPath)
	if err != nil {
		return fmt.Errorf("Error uploading cookbook %s: %s", localDir, err)
	}

	return nil
}

func CreateRemoteDirectory(path string, comm packer.Communicator) (err error) {
	//Ui.Say(fmt.Sprintf("Creating directory: %s", path))
	var copyCommand = []string{"mkdir -p", path}

	var cmd packer.RemoteCmd
	cmd.Command = strings.Join(copyCommand, " ")

	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	// Start the command
	if err := comm.Start(&cmd); err != nil {
	  return fmt.Errorf("Unable to create remote directory %s: %d", path, err)
	}

	// Wait for it to complete
	cmd.Wait()

	return
}

func CreateSoloRb(cookbookPaths []string, comm packer.Communicator) (str string, err error) {
	Ui.Say(fmt.Sprintf("Creating Chef configuration file..."))

	remotePath := RemoteStagingPath + "/solo.rb"
	tf, err := ioutil.TempFile("", "packer-chef-solo-rb")
	if err != nil {
		return "", fmt.Errorf("Error preparing Chef solo.rb: %s", err)
	}

	// Write our contents to it
	writer := bufio.NewWriter(tf)

	// Messy, messy...
	cbPathsCat := "\"" + RemoteCookbookPath + "/" + strings.Join(cookbookPaths, "\",\"" + RemoteCookbookPath + "/") + "\""
	contents := "file_cache_path \"" + RemoteFileCachePath + "\"\ncookbook_path [" + cbPathsCat + "]\n"

	if _, err := writer.WriteString(contents); err != nil {
		return "", fmt.Errorf("Error preparing solo.rb: %s", err)
	}

	if err := writer.Flush(); err != nil {
		return "", fmt.Errorf("Error preparing solo.rb: %s", err)
	}

	name := tf.Name()
	tf.Close()
	f, err := os.Open(name)
	comm.Upload(remotePath, f)

	defer os.Remove(name)

	// Upload the Chef Solo configuration file to the cookbook directory.

	if err != nil {
		return "", fmt.Errorf("Error uploading Chef Solo configuration file: %s", err)
	}
	//executeCommand("sudo cat " + remotePath, comm)
	return remotePath, nil
}

func CreateAttributesJson(jsonAttrs map[string]interface{}, runList []string, comm packer.Communicator) (str string, err error) {
	Ui.Say(fmt.Sprintf("Creating and uploading Chef attributes file"))
	remotePath := RemoteStagingPath + "/node.json"

	var formattedRunList []string
	for _, value := range runList {
		formattedRunList = append(formattedRunList, "recipe[" + value + "]")
	}

	// Add RunList to JSON
	jsonAttrs["run_list"] = formattedRunList

	// Convert to JSON string
	jsonString, err := json.MarshalIndent(jsonAttrs, "", "  ")
	if err != nil {
		return "", fmt.Errorf("Error parsing JSON attributes: %s", err)
	}

	tf, err := ioutil.TempFile("", "packer-chef-solo-json")
	if err != nil {
		return "", fmt.Errorf("Error preparing Chef attributes file: %s", err)
	}
	defer os.Remove(tf.Name())

	// Write our contents to it
	writer := bufio.NewWriter(tf)
	if _, err := writer.WriteString(string(jsonString)); err != nil {
		return "", fmt.Errorf("Error preparing Chef attributes file: %s", err)
	}

	if err := writer.Flush(); err != nil {
		return "", fmt.Errorf("Error preparing Chef attributes file: %s", err)
	}

	jsonFile := tf.Name()
	tf.Close()

	log.Printf("Opening %s for reading", jsonFile)
	f, err := os.Open(jsonFile)
	if err != nil {
		return "", fmt.Errorf("Error opening JSON attributes file: %s", err)
	}

	log.Printf("Uploading %s => %s", jsonFile, remotePath)
	err = comm.Upload(remotePath, f)
	if err != nil {
		return "", fmt.Errorf("Error uploading JSON attributes file: %s", err)
	}

	return remotePath, nil
}

func InstallChefSolo(comm packer.Communicator) (err error) {
	Ui.Say(fmt.Sprintf("Installing Chef Solo"))
	var installCommand = "curl -L https://www.opscode.com/chef/install.sh | sudo bash"
	err = executeCommand(installCommand, comm)
	if err != nil {
	  return fmt.Errorf("Unable to install Chef Solo: %d", err)
	}

	return nil
}

func executeCommand(command string, comm packer.Communicator) (err error) {
	// Setup the remote command
	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()

	var cmd packer.RemoteCmd
	cmd.Command = command
	cmd.Stdout = stdout_w
	cmd.Stderr = stderr_w

	//Ui.Say(fmt.Sprintf("Executing command: %s", cmd.Command))
	log.Printf("Executing command: %s", cmd.Command)
	err = comm.Start(&cmd)
	if err != nil {
		return fmt.Errorf("Failed executing command: %s", err)
	}

	exitChan := make(chan int, 1)
	stdoutChan := iochan.DelimReader(stdout_r, '\n')
	stderrChan := iochan.DelimReader(stderr_r, '\n')

	go func() {
		defer stdout_w.Close()
		defer stderr_w.Close()

		cmd.Wait()
		exitChan <- cmd.ExitStatus
	}()

OutputLoop:
	for {
		select {
		case output := <-stderrChan:
			Ui.Message(strings.TrimSpace(output))
		case output := <-stdoutChan:
			Ui.Message(strings.TrimSpace(output))
		case exitStatus := <-exitChan:
			log.Printf("Chef Solo provisioner exited with status %d", exitStatus)

			if exitStatus != 0 {
				return fmt.Errorf("Command exited with non-zero exit status: %d", exitStatus)
			}

			break OutputLoop
		}
	}

	// Make sure we finish off stdout/stderr because we may have gotten
	// a message from the exit channel first.
	for output := range stdoutChan {
		Ui.Message(output)
	}

	for output := range stderrChan {
		Ui.Message(output)
	}

	return nil
}
<commit_msg>Rename configuration `cookbook_paths` to `cookbooks_paths` to conform to Vagrant wording. Remove configuration `cookbook_path` to simplify some things.<commit_after>// This package implements a provisioner for Packer that executes
// shell scripts within the remote machine.
package chefSolo

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/mitchellh/iochan"
	"github.com/mitchellh/mapstructure"
	"github.com/mitchellh/packer/packer"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"
	"path/filepath"
	"encoding/json"
)

const RemoteStagingPath = "/tmp/provision/chef-solo"
const RemoteFileCachePath = "/tmp/provision/chef-solo"
const RemoteCookbookPath = "/tmp/provision/chef-solo/cookbooks"
const DefaultCookbookPath = "cookbooks"

var Ui packer.Ui

type config struct {
	// An array of local paths of cookbooks to upload.
	CookbooksPaths []string `mapstructure:"cookbooks_paths"`

	// An array of recipes to run.
	Recipes []string

	// A string of JSON that will be used as the JSON attributes for the
	// Chef run.
	Json map[string]interface{}

	UseSudo bool `mapstructure:"use_sudo"`

	// If true
	SkipInstall bool `mapstructure:"skip_install"`
}

type Provisioner struct {
	config config
}

type ExecuteRecipeTemplate struct {
	SoloRbPath string
	JsonPath string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	errs := make([]error, 0)
	for _, raw := range raws {
		if err := mapstructure.Decode(raw, &p.config); err != nil {
			return err
		}
	}

	if p.config.CookbooksPaths == nil {
		p.config.CookbooksPaths = make([]string, 0)
	}

	if p.config.Recipes == nil {
		p.config.Recipes = make([]string, 0)
	}

	if p.config.Json != nil {
		if _, err := json.Marshal(p.config.Json); err != nil {
			errs = append(errs, fmt.Errorf("Bad JSON: %s", err))
		}
	} else {
		p.config.Json = make(map[string]interface{})
	}

	for _, path := range p.config.CookbooksPaths {
		pFileInfo, err := os.Stat(path)

		if err != nil || !pFileInfo.IsDir() {
			errs = append(errs, fmt.Errorf("Bad cookbook path '%s': %s", path, err))
		}
	}

	if len(errs) > 0 {
		return &packer.MultiError{errs}
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	cookbooksPaths := make([]string, len(p.config.CookbooksPaths))
	copy(cookbooksPaths, p.config.CookbooksPaths)

	Ui = ui

	// Generic setup for Chef runs
	err := InstallChefSolo(comm)
	if err != nil {
		return fmt.Errorf("Error installing Chef Solo: %s", err)
	}

	err = CreateRemoteDirectory(RemoteCookbookPath, comm)
	if err != nil {
		return fmt.Errorf("Error creating remote staging directory: %s", err)
	}

	soloRbPath, err := CreateSoloRb(p.config.CookbooksPaths, comm)
	if err != nil {
		return fmt.Errorf("Error creating Chef Solo configuration file: %s", err)
	}

	jsonPath, err := CreateAttributesJson(p.config.Json, p.config.Recipes, comm)
	if err != nil {
		return fmt.Errorf("Error uploading JSON attributes file: %s", err)
	}

	// Upload all cookbooks
	for _, path := range cookbooksPaths {
		ui.Say(fmt.Sprintf("Copying cookbook path: %s", path))
		err = UploadLocalDirectory(path, comm)
		if err != nil {
			return fmt.Errorf("Error uploading cookbooks: %s", err)
		}
	}

	// Execute requested recipes
	ui.Say("Beginning Chef Solo run")

	// Compile the command
	var command bytes.Buffer
	t := template.Must(template.New("chef-run").Parse("sudo chef-solo --no-color -c {{.SoloRbPath}} -j {{.JsonPath}}"))
	t.Execute(&command, &ExecuteRecipeTemplate{soloRbPath, jsonPath})

	err = executeCommand(command.String(), comm)
	if err != nil {
		return fmt.Errorf("Error running Chef Solo: %s", err)
	}

	// return fmt.Errorf("Die")

	return nil
}

func UploadLocalDirectory(localDir string, comm packer.Communicator) (err error) {
	visitPath := func (path string, f os.FileInfo, err error) (err2 error) {
		var remotePath = RemoteCookbookPath + "/" + path
		if f.IsDir() {
			// Make remote directory
			err = CreateRemoteDirectory(remotePath, comm)
			if err != nil {
				return err
			}
		} else {
			// Upload file to existing directory
			file, err := os.Open(path)
			if err != nil {
				return fmt.Errorf("Error opening file: %s", err)
			}

			err = comm.Upload(remotePath, file)
			if err != nil {
				return fmt.Errorf("Error uploading file: %s", err)
			}
		}
		return
	}

	err = filepath.Walk(localDir, visitPath)
	if err != nil {
		return fmt.Errorf("Error uploading cookbook %s: %s", localDir, err)
	}

	return nil
}

func CreateRemoteDirectory(path string, comm packer.Communicator) (err error) {
	//Ui.Say(fmt.Sprintf("Creating directory: %s", path))
	var copyCommand = []string{"mkdir -p", path}

	var cmd packer.RemoteCmd
	cmd.Command = strings.Join(copyCommand, " ")

	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	// Start the command
	if err := comm.Start(&cmd); err != nil {
	  return fmt.Errorf("Unable to create remote directory %s: %d", path, err)
	}

	// Wait for it to complete
	cmd.Wait()

	return
}

func CreateSoloRb(cookbooksPaths []string, comm packer.Communicator) (str string, err error) {
	Ui.Say(fmt.Sprintf("Creating Chef configuration file..."))

	remotePath := RemoteStagingPath + "/solo.rb"
	tf, err := ioutil.TempFile("", "packer-chef-solo-rb")
	if err != nil {
		return "", fmt.Errorf("Error preparing Chef solo.rb: %s", err)
	}

	// Write our contents to it
	writer := bufio.NewWriter(tf)

	// Messy, messy...
	cbPathsCat := "\"" + RemoteCookbookPath + "/" + strings.Join(cookbooksPaths, "\",\"" + RemoteCookbookPath + "/") + "\""
	contents := "file_cache_path \"" + RemoteFileCachePath + "\"\ncookbook_path [" + cbPathsCat + "]\n"

	if _, err := writer.WriteString(contents); err != nil {
		return "", fmt.Errorf("Error preparing solo.rb: %s", err)
	}

	if err := writer.Flush(); err != nil {
		return "", fmt.Errorf("Error preparing solo.rb: %s", err)
	}

	name := tf.Name()
	tf.Close()
	f, err := os.Open(name)
	comm.Upload(remotePath, f)

	defer os.Remove(name)

	// Upload the Chef Solo configuration file to the cookbook directory.

	if err != nil {
		return "", fmt.Errorf("Error uploading Chef Solo configuration file: %s", err)
	}
	//executeCommand("sudo cat " + remotePath, comm)
	return remotePath, nil
}

func CreateAttributesJson(jsonAttrs map[string]interface{}, recipes []string, comm packer.Communicator) (str string, err error) {
	Ui.Say(fmt.Sprintf("Creating and uploading Chef attributes file"))
	remotePath := RemoteStagingPath + "/node.json"

	var formattedRecipes []string
	for _, value := range recipes {
		formattedRecipes = append(formattedRecipes, "recipe[" + value + "]")
	}

	// Add Recipes to JSON
	jsonAttrs["run_list"] = formattedRecipes

	// Convert to JSON string
	jsonString, err := json.MarshalIndent(jsonAttrs, "", "  ")
	if err != nil {
		return "", fmt.Errorf("Error parsing JSON attributes: %s", err)
	}

	tf, err := ioutil.TempFile("", "packer-chef-solo-json")
	if err != nil {
		return "", fmt.Errorf("Error preparing Chef attributes file: %s", err)
	}
	defer os.Remove(tf.Name())

	// Write our contents to it
	writer := bufio.NewWriter(tf)
	if _, err := writer.WriteString(string(jsonString)); err != nil {
		return "", fmt.Errorf("Error preparing Chef attributes file: %s", err)
	}

	if err := writer.Flush(); err != nil {
		return "", fmt.Errorf("Error preparing Chef attributes file: %s", err)
	}

	jsonFile := tf.Name()
	tf.Close()

	log.Printf("Opening %s for reading", jsonFile)
	f, err := os.Open(jsonFile)
	if err != nil {
		return "", fmt.Errorf("Error opening JSON attributes file: %s", err)
	}

	log.Printf("Uploading %s => %s", jsonFile, remotePath)
	err = comm.Upload(remotePath, f)
	if err != nil {
		return "", fmt.Errorf("Error uploading JSON attributes file: %s", err)
	}

	return remotePath, nil
}

func InstallChefSolo(comm packer.Communicator) (err error) {
	Ui.Say(fmt.Sprintf("Installing Chef Solo"))
	var installCommand = "curl -L https://www.opscode.com/chef/install.sh | sudo bash"
	err = executeCommand(installCommand, comm)
	if err != nil {
	  return fmt.Errorf("Unable to install Chef Solo: %d", err)
	}

	return nil
}

func executeCommand(command string, comm packer.Communicator) (err error) {
	// Setup the remote command
	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()

	var cmd packer.RemoteCmd
	cmd.Command = command
	cmd.Stdout = stdout_w
	cmd.Stderr = stderr_w

	//Ui.Say(fmt.Sprintf("Executing command: %s", cmd.Command))
	log.Printf("Executing command: %s", cmd.Command)
	err = comm.Start(&cmd)
	if err != nil {
		return fmt.Errorf("Failed executing command: %s", err)
	}

	exitChan := make(chan int, 1)
	stdoutChan := iochan.DelimReader(stdout_r, '\n')
	stderrChan := iochan.DelimReader(stderr_r, '\n')

	go func() {
		defer stdout_w.Close()
		defer stderr_w.Close()

		cmd.Wait()
		exitChan <- cmd.ExitStatus
	}()

OutputLoop:
	for {
		select {
		case output := <-stderrChan:
			Ui.Message(strings.TrimSpace(output))
		case output := <-stdoutChan:
			Ui.Message(strings.TrimSpace(output))
		case exitStatus := <-exitChan:
			log.Printf("Chef Solo provisioner exited with status %d", exitStatus)

			if exitStatus != 0 {
				return fmt.Errorf("Command exited with non-zero exit status: %d", exitStatus)
			}

			break OutputLoop
		}
	}

	// Make sure we finish off stdout/stderr because we may have gotten
	// a message from the exit channel first.
	for output := range stdoutChan {
		Ui.Message(output)
	}

	for output := range stderrChan {
		Ui.Message(output)
	}

	return nil
}
<|endoftext|>


<commit_before>// Library for I2Ps SAMv3 bridge (https://geti2p.com)
package sam3

import (
	"bufio"
	"bytes"
	"net"
	"errors"
    "io"
    "os"
	"strings"
)

// Used for controlling I2Ps SAMv3.
type SAM struct {
	address     string        // ipv4:port
	conn        net.Conn
    keys        *I2PKeys
}

const (
	session_OK             = "SESSION STATUS RESULT=OK DESTINATION="
	session_DUPLICATE_ID   = "SESSION STATUS RESULT=DUPLICATED_ID\n"
	session_DUPLICATE_DEST = "SESSION STATUS RESULT=DUPLICATED_DEST\n"
	session_INVALID_KEY    = "SESSION STATUS RESULT=INVALID_KEY\n"
	session_I2P_ERROR      = "SESSION STATUS RESULT=I2P_ERROR MESSAGE="
)

// Creates a new controller for the I2P routers SAM bridge.
func NewSAM(address string) (*SAM, error) {
	conn, err := net.Dial("tcp4", address)
	if err != nil {
		return nil, err
	}
	if _, err := conn.Write([]byte("HELLO VERSION MIN=3.0 MAX=3.0\n")); err != nil {
		return nil, err
	}
	buf := make([]byte, 256)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, err
	}
	if string(buf[:n]) == "HELLO REPLY RESULT=OK VERSION=3.0\n" {
		return &SAM{address, conn, nil}, nil
	} else if string(buf[:n]) == "HELLO REPLY RESULT=NOVERSION\n" {
		return nil, errors.New("That SAM bridge does not support SAMv3.")
	} else {
		return nil, errors.New(string(buf[:n]))
	}
}

// if keyfile fname does not exist
func (sam *SAM) EnsureKeyfile(fname string) (err error) {

    _, err = os.Stat(fname)
    if os.IsNotExist(err) {
        // make the keys
        var keys I2PKeys
        keys, err = sam.NewKeys()
        if err == nil {
            sam.keys = &keys
            // save keys
            var f io.WriteCloser
            f, err = os.OpenFile(fname, os.O_WRONLY, 0600)
            if err == nil {
                err = StoreKeysIncompat(keys, f)
                f.Close()
            }
        }
    } else if err == nil {
        // we haz key file
        var f *os.File
        var keys I2PKeys
        f, err = os.Open(fname)
        if err == nil {
            keys, err = LoadKeysIncompat(f)
            if err == nil {
                sam.keys = &keys
            }
        }
    }
    return
}

// Creates the I2P-equivalent of an IP address, that is unique and only the one
// who has the private keys can send messages from. The public keys are the I2P
// desination (the address) that anyone can send messages to.
func (sam *SAM) NewKeys() (I2PKeys, error) {
	if _, err := sam.conn.Write([]byte("DEST GENERATE\n")); err != nil {
		return I2PKeys{}, err
	}
	buf := make([]byte, 8192)
	n, err := sam.conn.Read(buf)
	if err != nil {
		return I2PKeys{}, err
	}
	s := bufio.NewScanner(bytes.NewReader(buf[:n]))
	s.Split(bufio.ScanWords)

	var pub, priv string
	for s.Scan() {
		text := s.Text()
		if text == "DEST" {
			continue
		} else if text == "REPLY" {
			continue
		} else if strings.HasPrefix(text, "PUB=") {
			pub = text[4:]
		} else if strings.HasPrefix(text, "PRIV=") {
			priv = text[5:]
		} else {
			return I2PKeys{}, errors.New("Failed to parse keys.")
		}
	}
	return I2PKeys{I2PAddr(pub), priv}, nil
}

// Performs a lookup, probably this order: 1) routers known addresses, cached
// addresses, 3) by asking peers in the I2P network.
func (sam *SAM) Lookup(name string) (I2PAddr, error) {
	if _, err := sam.conn.Write([]byte("NAMING LOOKUP NAME=" + name + "\n")); err != nil {
		return I2PAddr(""), err
	}
	buf := make([]byte, 4096)
	n, err := sam.conn.Read(buf)
	if err != nil {
		return I2PAddr(""), err
	}
	if n <= 13 || !strings.HasPrefix(string(buf[:n]), "NAMING REPLY ") {
		return I2PAddr(""), errors.New("Failed to parse.")
	}
	s := bufio.NewScanner(bytes.NewReader(buf[13:n]))
	s.Split(bufio.ScanWords)

	errStr := ""
	for s.Scan() {
		text := s.Text()
		if text == "RESULT=OK" {
			continue
		} else if text == "RESULT=INVALID_KEY" {
			errStr += "Invalid key."
		} else if text == "RESULT=KEY_NOT_FOUND" {
			errStr += "Unable to resolve " + name
		} else if text == "NAME=" + name {
			continue
		} else if strings.HasPrefix(text, "VALUE=") {
			return I2PAddr(text[6:]), nil
		} else if strings.HasPrefix(text, "MESSAGE=") {
			errStr += " " + text[8:]
		} else {
      continue
		}
	}
	return I2PAddr(""), errors.New(errStr)
}

// Creates a new session with the style of either "STREAM", "DATAGRAM" or "RAW",
// for a new I2P tunnel with name id, using the cypher keys specified, with the
// I2CP/streaminglib-options as specified. Extra arguments can be specified by
// setting extra to something else than []string{}.
// This sam3 instance is now a session
func (sam *SAM) newGenericSession(style, id string, keys I2PKeys, options []string, extras []string) (net.Conn, error) {

	optStr := ""
	for _, opt := range options {
		optStr += "OPTION=" + opt + " "
	}

	conn := sam.conn
	scmsg := []byte("SESSION CREATE STYLE=" + style + " ID=" + id + " DESTINATION=" + keys.String() + " " + optStr + strings.Join(extras, " ") + "\n")
	for m, i:=0, 0; m!=len(scmsg); i++ {
		if i == 15 {
			conn.Close()
			return nil, errors.New("writing to SAM failed")
		}
		n, err := conn.Write(scmsg[m:])
		if err != nil {
			conn.Close()
			return nil, err
		}
		m += n
	}
	buf := make([]byte, 4096)
	n, err := conn.Read(buf)
	if err != nil {
		conn.Close()
		return nil, err
	}
	text := string(buf[:n])
	if strings.HasPrefix(text, session_OK) {
		if keys.String() != text[len(session_OK):len(text)-1] {
			return nil, errors.New("SAMv3 created a tunnel with keys other than the ones we asked it for")
		}
		return conn, nil //&StreamSession{id, conn, keys, nil, sync.RWMutex{}, nil}, nil
	} else if text == session_DUPLICATE_ID {
		conn.Close()
		return nil, errors.New("Duplicate tunnel name")
	} else if text == session_DUPLICATE_DEST {
		conn.Close()
		return nil, errors.New("Duplicate destination")
	} else if text == session_INVALID_KEY {
		conn.Close()
		return nil, errors.New("Invalid key")
	} else if strings.HasPrefix(text, session_I2P_ERROR) {
		conn.Close()
		return nil, errors.New("I2P error " + text[len(session_I2P_ERROR):])
	} else {
		conn.Close()
		return nil, errors.New("Unable to parse SAMv3 reply: " + text)
	}
}

// close this sam session
func (sam *SAM) Close() error {
	if err := sam.conn.Close(); err != nil {
		return err
	}
	return nil
}

<commit_msg>have ensurekeyfile return the i2p keys<commit_after>// Library for I2Ps SAMv3 bridge (https://geti2p.com)
package sam3

import (
	"bufio"
	"bytes"
	"net"
	"errors"
    "io"
    "os"
	"strings"
)

// Used for controlling I2Ps SAMv3.
type SAM struct {
	address     string        // ipv4:port
	conn        net.Conn
    keys        *I2PKeys
}

const (
	session_OK             = "SESSION STATUS RESULT=OK DESTINATION="
	session_DUPLICATE_ID   = "SESSION STATUS RESULT=DUPLICATED_ID\n"
	session_DUPLICATE_DEST = "SESSION STATUS RESULT=DUPLICATED_DEST\n"
	session_INVALID_KEY    = "SESSION STATUS RESULT=INVALID_KEY\n"
	session_I2P_ERROR      = "SESSION STATUS RESULT=I2P_ERROR MESSAGE="
)

// Creates a new controller for the I2P routers SAM bridge.
func NewSAM(address string) (*SAM, error) {
	conn, err := net.Dial("tcp4", address)
	if err != nil {
		return nil, err
	}
	if _, err := conn.Write([]byte("HELLO VERSION MIN=3.0 MAX=3.0\n")); err != nil {
		return nil, err
	}
	buf := make([]byte, 256)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, err
	}
	if string(buf[:n]) == "HELLO REPLY RESULT=OK VERSION=3.0\n" {
		return &SAM{address, conn, nil}, nil
	} else if string(buf[:n]) == "HELLO REPLY RESULT=NOVERSION\n" {
		return nil, errors.New("That SAM bridge does not support SAMv3.")
	} else {
		return nil, errors.New(string(buf[:n]))
	}
}

// if keyfile fname does not exist
func (sam *SAM) EnsureKeyfile(fname string) (keys I2PKeys, err error) {

    _, err = os.Stat(fname)
    if os.IsNotExist(err) {
        // make the keys
        keys, err = sam.NewKeys()
        if err == nil {
            sam.keys = &keys
            // save keys
            var f io.WriteCloser
            f, err = os.OpenFile(fname, os.O_WRONLY, 0600)
            if err == nil {
                err = StoreKeysIncompat(keys, f)
                f.Close()
            }
        }
    } else if err == nil {
        // we haz key file
        var f *os.File
        f, err = os.Open(fname)
        if err == nil {
            keys, err = LoadKeysIncompat(f)
            if err == nil {
                sam.keys = &keys
            }
        }
    }
    return
}

// Creates the I2P-equivalent of an IP address, that is unique and only the one
// who has the private keys can send messages from. The public keys are the I2P
// desination (the address) that anyone can send messages to.
func (sam *SAM) NewKeys() (I2PKeys, error) {
	if _, err := sam.conn.Write([]byte("DEST GENERATE\n")); err != nil {
		return I2PKeys{}, err
	}
	buf := make([]byte, 8192)
	n, err := sam.conn.Read(buf)
	if err != nil {
		return I2PKeys{}, err
	}
	s := bufio.NewScanner(bytes.NewReader(buf[:n]))
	s.Split(bufio.ScanWords)

	var pub, priv string
	for s.Scan() {
		text := s.Text()
		if text == "DEST" {
			continue
		} else if text == "REPLY" {
			continue
		} else if strings.HasPrefix(text, "PUB=") {
			pub = text[4:]
		} else if strings.HasPrefix(text, "PRIV=") {
			priv = text[5:]
		} else {
			return I2PKeys{}, errors.New("Failed to parse keys.")
		}
	}
	return I2PKeys{I2PAddr(pub), priv}, nil
}

// Performs a lookup, probably this order: 1) routers known addresses, cached
// addresses, 3) by asking peers in the I2P network.
func (sam *SAM) Lookup(name string) (I2PAddr, error) {
	if _, err := sam.conn.Write([]byte("NAMING LOOKUP NAME=" + name + "\n")); err != nil {
		return I2PAddr(""), err
	}
	buf := make([]byte, 4096)
	n, err := sam.conn.Read(buf)
	if err != nil {
		return I2PAddr(""), err
	}
	if n <= 13 || !strings.HasPrefix(string(buf[:n]), "NAMING REPLY ") {
		return I2PAddr(""), errors.New("Failed to parse.")
	}
	s := bufio.NewScanner(bytes.NewReader(buf[13:n]))
	s.Split(bufio.ScanWords)

	errStr := ""
	for s.Scan() {
		text := s.Text()
		if text == "RESULT=OK" {
			continue
		} else if text == "RESULT=INVALID_KEY" {
			errStr += "Invalid key."
		} else if text == "RESULT=KEY_NOT_FOUND" {
			errStr += "Unable to resolve " + name
		} else if text == "NAME=" + name {
			continue
		} else if strings.HasPrefix(text, "VALUE=") {
			return I2PAddr(text[6:]), nil
		} else if strings.HasPrefix(text, "MESSAGE=") {
			errStr += " " + text[8:]
		} else {
      continue
		}
	}
	return I2PAddr(""), errors.New(errStr)
}

// Creates a new session with the style of either "STREAM", "DATAGRAM" or "RAW",
// for a new I2P tunnel with name id, using the cypher keys specified, with the
// I2CP/streaminglib-options as specified. Extra arguments can be specified by
// setting extra to something else than []string{}.
// This sam3 instance is now a session
func (sam *SAM) newGenericSession(style, id string, keys I2PKeys, options []string, extras []string) (net.Conn, error) {

	optStr := ""
	for _, opt := range options {
		optStr += "OPTION=" + opt + " "
	}

	conn := sam.conn
	scmsg := []byte("SESSION CREATE STYLE=" + style + " ID=" + id + " DESTINATION=" + keys.String() + " " + optStr + strings.Join(extras, " ") + "\n")
	for m, i:=0, 0; m!=len(scmsg); i++ {
		if i == 15 {
			conn.Close()
			return nil, errors.New("writing to SAM failed")
		}
		n, err := conn.Write(scmsg[m:])
		if err != nil {
			conn.Close()
			return nil, err
		}
		m += n
	}
	buf := make([]byte, 4096)
	n, err := conn.Read(buf)
	if err != nil {
		conn.Close()
		return nil, err
	}
	text := string(buf[:n])
	if strings.HasPrefix(text, session_OK) {
		if keys.String() != text[len(session_OK):len(text)-1] {
			return nil, errors.New("SAMv3 created a tunnel with keys other than the ones we asked it for")
		}
		return conn, nil //&StreamSession{id, conn, keys, nil, sync.RWMutex{}, nil}, nil
	} else if text == session_DUPLICATE_ID {
		conn.Close()
		return nil, errors.New("Duplicate tunnel name")
	} else if text == session_DUPLICATE_DEST {
		conn.Close()
		return nil, errors.New("Duplicate destination")
	} else if text == session_INVALID_KEY {
		conn.Close()
		return nil, errors.New("Invalid key")
	} else if strings.HasPrefix(text, session_I2P_ERROR) {
		conn.Close()
		return nil, errors.New("I2P error " + text[len(session_I2P_ERROR):])
	} else {
		conn.Close()
		return nil, errors.New("Unable to parse SAMv3 reply: " + text)
	}
}

// close this sam session
func (sam *SAM) Close() error {
	if err := sam.conn.Close(); err != nil {
		return err
	}
	return nil
}

<|endoftext|>


<commit_before>package junos

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/Juniper/go-netconf/netconf"
	"log"
)

// Session holds the connection information to our Junos device.
type Session struct {
	Conn *netconf.Session
}

// rollbackXML parses our rollback configuration.
type rollbackXML struct {
	XMLName xml.Name `xml:"rollback-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// RescueXML parses our rescue configuration.
type rescueXML struct {
	XMLName xml.Name `xml:"rescue-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// CommandXML parses our operational command responses.
type commandXML struct {
	Config string `xml:",innerxml"`
}

// NewSession establishes a new connection to a Junos device that we will use
// to run our commands against.
func NewSession(host, user, password string) *Session {
	s, err := netconf.DialSSH(host, netconf.SSHConfigPassword(user, password))
	if err != nil {
		log.Fatal(err)
	}

	return &Session{
		Conn: s,
	}
}

// Lock locks the candidate configuration.
func (s *Session) Lock() error {
	resp, err := s.Conn.Exec(rpcCommand["lock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// Unlock unlocks the candidate configuration.
func (s *Session) Unlock() error {
	resp, err := s.Conn.Exec(rpcCommand["unlock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// GetRollbackConfig returns the configuration of the given rollback state.
func (s *Session) GetRollbackConfig(number int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information"], number)
	reply, err := s.Conn.Exec(command)

	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// RollbackDiff compares the current active configuration to a given rollback configuration.
func (s *Session) RollbackDiff(compare int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information-compare"], compare)
	reply, err := s.Conn.Exec(command)

	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// GetRescueConfig returns the rescue configuration.
func (s *Session) GetRescueConfig() (string, error) {
	rescue := &rescueXML{}
	reply, err := s.Conn.Exec(rpcCommand["get-rescue-information"])

	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rescue)
	if err != nil {
		log.Fatal(err)
	}

	if rescue.Config == "" {
		return "No rescue configuration set.", nil
	}

	return rescue.Config, nil
}

// Command runs any operational mode command, such as "show" or "request."
// Format is either "text" or "xml".
func (s *Session) Command(cmd, format string) (string, error) {
	c := &commandXML{}
	var command string

	switch format {
	case "xml":
		command = fmt.Sprintf(rpcCommand["command-xml"], cmd)
	default:
		command = fmt.Sprintf(rpcCommand["command"], cmd)
	}
	reply, err := s.Conn.Exec(command)
	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), &c)
	if err != nil {
		log.Fatal(err)
	}

	if c.Config == "" {
		return "No output available.", nil
	}

	return c.Config, nil
}

// Close disconnects and closes the session to our Junos device.
func (s *Session) Close() {
	s.Conn.Close()
}
<commit_msg>Updated documentation<commit_after>// Package junos allows configuration of Juniper Networks' Junos devices.
package junos

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/Juniper/go-netconf/netconf"
	"log"
)

// Session holds the connection information to our Junos device.
type Session struct {
	Conn *netconf.Session
}

// rollbackXML parses our rollback configuration.
type rollbackXML struct {
	XMLName xml.Name `xml:"rollback-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// RescueXML parses our rescue configuration.
type rescueXML struct {
	XMLName xml.Name `xml:"rescue-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// CommandXML parses our operational command responses.
type commandXML struct {
	Config string `xml:",innerxml"`
}

// NewSession establishes a new connection to a Junos device that we will use
// to run our commands against.
func NewSession(host, user, password string) *Session {
	s, err := netconf.DialSSH(host, netconf.SSHConfigPassword(user, password))
	if err != nil {
		log.Fatal(err)
	}

	return &Session{
		Conn: s,
	}
}

// Lock locks the candidate configuration.
func (s *Session) Lock() error {
	resp, err := s.Conn.Exec(rpcCommand["lock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// Unlock unlocks the candidate configuration.
func (s *Session) Unlock() error {
	resp, err := s.Conn.Exec(rpcCommand["unlock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// GetRollbackConfig returns the configuration of the given rollback state.
func (s *Session) GetRollbackConfig(number int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information"], number)
	reply, err := s.Conn.Exec(command)

	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// RollbackDiff compares the current active configuration to a given rollback configuration.
func (s *Session) RollbackDiff(compare int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information-compare"], compare)
	reply, err := s.Conn.Exec(command)

	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// GetRescueConfig returns the rescue configuration.
func (s *Session) GetRescueConfig() (string, error) {
	rescue := &rescueXML{}
	reply, err := s.Conn.Exec(rpcCommand["get-rescue-information"])

	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rescue)
	if err != nil {
		log.Fatal(err)
	}

	if rescue.Config == "" {
		return "No rescue configuration set.", nil
	}

	return rescue.Config, nil
}

// Command runs any operational mode command, such as "show" or "request."
// Format is either "text" or "xml".
func (s *Session) Command(cmd, format string) (string, error) {
	c := &commandXML{}
	var command string

	switch format {
	case "xml":
		command = fmt.Sprintf(rpcCommand["command-xml"], cmd)
	default:
		command = fmt.Sprintf(rpcCommand["command"], cmd)
	}
	reply, err := s.Conn.Exec(command)
	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), &c)
	if err != nil {
		log.Fatal(err)
	}

	if c.Config == "" {
		return "No output available.", nil
	}

	return c.Config, nil
}

// Close disconnects and closes the session to our Junos device.
func (s *Session) Close() {
	s.Conn.Close()
}
<|endoftext|>


<commit_before>package sensu

import (
	"os"
	"reflect"
	"strings"
	"testing"

	"github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/check"
	stdClient "github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/client"
)

var dummyClient = &stdClient.Client{
	Name:          "test_client",
	Address:       "10.0.0.42",
	Subscriptions: strings.Split("email,messenger", ","),
}

func validateStringParameter(
	actual string,
	expected string,
	parameterName string,
	t *testing.T) {

	if actual != expected {
		t.Errorf(
			"Expected %s to be \"%s\" but got \"%s\" instead!",
			parameterName,
			expected,
			actual,
		)
	}
}

func TestRabbitMQURIDefaultValue(t *testing.T) {
	validateStringParameter(
		(&Config{}).RabbitMQURI(),
		defaultRabbitMQURI,
		"RabbitMQ URI",
		t,
	)
}

func TestRabbitMQURIFromEnvVar(t *testing.T) {
	expectedRabbitMqUri := "amqp://user:password@example.com:5672"

	os.Setenv("RABBITMQ_URI", expectedRabbitMqUri)
	defer os.Unsetenv("RABBITMQ_URI")

	validateStringParameter(
		(&Config{}).RabbitMQURI(),
		expectedRabbitMqUri,
		"RabbitMQ URI",
		t,
	)
}

func TestRabbitMQURIFromConfig(t *testing.T) {
	expectedRabbitMqUri := "amqp://user:password@example.com:5672"

	config := Config{config: &configPayload{RabbitMQURI: &expectedRabbitMqUri}}

	validateStringParameter(
		config.RabbitMQURI(),
		expectedRabbitMqUri,
		"RabbitMQ URI",
		t,
	)
}

func validateClient(actualClient *stdClient.Client, expectedClient *stdClient.Client, t *testing.T) {
	validateStringParameter(
		actualClient.Name,
		expectedClient.Name,
		"client name",
		t,
	)

	validateStringParameter(
		actualClient.Address,
		expectedClient.Address,
		"client address",
		t,
	)

	if !reflect.DeepEqual(
		actualClient.Subscriptions,
		expectedClient.Subscriptions,
	) {
		t.Errorf(
			"Expected client subscriptions to be \"%#v\" but got \"%#v\" instead!",
			expectedClient.Subscriptions,
			actualClient.Subscriptions,
		)
	}
}

func TestClientFromConfig(t *testing.T) {
	config := Config{config: &configPayload{Client: dummyClient}}

	validateClient(config.Client(), dummyClient, t)
}

func TestClientFromEnvVars(t *testing.T) {
	os.Setenv("SENSU_CLIENT_NAME", dummyClient.Name)
	defer os.Unsetenv("SENSU_CLIENT_NAME")

	os.Setenv("SENSU_CLIENT_ADDRESS", dummyClient.Address)
	defer os.Unsetenv("SENSU_CLIENT_ADDRESS")

	os.Setenv(
		"SENSU_CLIENT_SUBSCRIPTIONS",
		strings.Join(dummyClient.Subscriptions, ","),
	)
	defer os.Unsetenv("SENSU_CLIENT_SUBSCRIPTIONS")

	validateClient((&Config{}).Client(), dummyClient, t)
}

func TestClientFromEnvVarsNoSubscriptions(t *testing.T) {
	dummyClientNoSubscriptions := dummyClient
	dummyClientNoSubscriptions.Subscriptions = []string{}

	os.Setenv("SENSU_CLIENT_NAME", dummyClientNoSubscriptions.Name)
	defer os.Unsetenv("SENSU_CLIENT_NAME")

	os.Setenv("SENSU_CLIENT_ADDRESS", dummyClientNoSubscriptions.Address)
	defer os.Unsetenv("SENSU_CLIENT_ADDRESS")

	validateClient((&Config{}).Client(), dummyClientNoSubscriptions, t)
}

func TestChecksFromConfig(t *testing.T) {
	expectedCheckCount := 2
	config := Config{
		config: &configPayload{
			Checks: []*check.Check{&check.Check{}, &check.Check{}},
		},
	}

	actualCheckCount := len(config.Checks())

	if expectedCheckCount != actualCheckCount {
		t.Errorf(
			"Expected check count to be \"%d\" but got \"%d\" instead!",
			expectedCheckCount,
			actualCheckCount,
		)
	}
}

func TestNewConfigFromFile(t *testing.T) {
	if c, err := NewConfigFromFile(nil, ""); c != nil || err != errNoClientName {
		t.Errorf("Expected (nil, %v) but got (%v, %v)", errNoClientName, c, err)
	}
}

func TestRabbitMQHAConfigDefaultValue(t *testing.T) {
	haConfig := (&Config{}).RabbitMQHAConfig()

	expectedConfigCont := 1

	if len(haConfig) != expectedConfigCont {
		t.Errorf(
			"Expected the config count to be \"%d\" but got \"%d\" instead!",
			expectedConfigCont,
			len(haConfig),
		)
	}

	validateStringParameter(
		haConfig[0].GetURI(),
		defaultRabbitMQURI,
		"RabbitMQ URI",
		t,
	)
}
<commit_msg>Remove superfluous quotes from test error message<commit_after>package sensu

import (
	"os"
	"reflect"
	"strings"
	"testing"

	"github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/check"
	stdClient "github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/client"
)

var dummyClient = &stdClient.Client{
	Name:          "test_client",
	Address:       "10.0.0.42",
	Subscriptions: strings.Split("email,messenger", ","),
}

func validateStringParameter(
	actual string,
	expected string,
	parameterName string,
	t *testing.T) {

	if actual != expected {
		t.Errorf(
			"Expected %s to be \"%s\" but got \"%s\" instead!",
			parameterName,
			expected,
			actual,
		)
	}
}

func TestRabbitMQURIDefaultValue(t *testing.T) {
	validateStringParameter(
		(&Config{}).RabbitMQURI(),
		defaultRabbitMQURI,
		"RabbitMQ URI",
		t,
	)
}

func TestRabbitMQURIFromEnvVar(t *testing.T) {
	expectedRabbitMqUri := "amqp://user:password@example.com:5672"

	os.Setenv("RABBITMQ_URI", expectedRabbitMqUri)
	defer os.Unsetenv("RABBITMQ_URI")

	validateStringParameter(
		(&Config{}).RabbitMQURI(),
		expectedRabbitMqUri,
		"RabbitMQ URI",
		t,
	)
}

func TestRabbitMQURIFromConfig(t *testing.T) {
	expectedRabbitMqUri := "amqp://user:password@example.com:5672"

	config := Config{config: &configPayload{RabbitMQURI: &expectedRabbitMqUri}}

	validateStringParameter(
		config.RabbitMQURI(),
		expectedRabbitMqUri,
		"RabbitMQ URI",
		t,
	)
}

func validateClient(actualClient *stdClient.Client, expectedClient *stdClient.Client, t *testing.T) {
	validateStringParameter(
		actualClient.Name,
		expectedClient.Name,
		"client name",
		t,
	)

	validateStringParameter(
		actualClient.Address,
		expectedClient.Address,
		"client address",
		t,
	)

	if !reflect.DeepEqual(
		actualClient.Subscriptions,
		expectedClient.Subscriptions,
	) {
		t.Errorf(
			"Expected client subscriptions to be \"%#v\" but got \"%#v\" instead!",
			expectedClient.Subscriptions,
			actualClient.Subscriptions,
		)
	}
}

func TestClientFromConfig(t *testing.T) {
	config := Config{config: &configPayload{Client: dummyClient}}

	validateClient(config.Client(), dummyClient, t)
}

func TestClientFromEnvVars(t *testing.T) {
	os.Setenv("SENSU_CLIENT_NAME", dummyClient.Name)
	defer os.Unsetenv("SENSU_CLIENT_NAME")

	os.Setenv("SENSU_CLIENT_ADDRESS", dummyClient.Address)
	defer os.Unsetenv("SENSU_CLIENT_ADDRESS")

	os.Setenv(
		"SENSU_CLIENT_SUBSCRIPTIONS",
		strings.Join(dummyClient.Subscriptions, ","),
	)
	defer os.Unsetenv("SENSU_CLIENT_SUBSCRIPTIONS")

	validateClient((&Config{}).Client(), dummyClient, t)
}

func TestClientFromEnvVarsNoSubscriptions(t *testing.T) {
	dummyClientNoSubscriptions := dummyClient
	dummyClientNoSubscriptions.Subscriptions = []string{}

	os.Setenv("SENSU_CLIENT_NAME", dummyClientNoSubscriptions.Name)
	defer os.Unsetenv("SENSU_CLIENT_NAME")

	os.Setenv("SENSU_CLIENT_ADDRESS", dummyClientNoSubscriptions.Address)
	defer os.Unsetenv("SENSU_CLIENT_ADDRESS")

	validateClient((&Config{}).Client(), dummyClientNoSubscriptions, t)
}

func TestChecksFromConfig(t *testing.T) {
	expectedCheckCount := 2
	config := Config{
		config: &configPayload{
			Checks: []*check.Check{&check.Check{}, &check.Check{}},
		},
	}

	actualCheckCount := len(config.Checks())

	if expectedCheckCount != actualCheckCount {
		t.Errorf(
			"Expected check count to be \"%d\" but got \"%d\" instead!",
			expectedCheckCount,
			actualCheckCount,
		)
	}
}

func TestNewConfigFromFile(t *testing.T) {
	if c, err := NewConfigFromFile(nil, ""); c != nil || err != errNoClientName {
		t.Errorf("Expected (nil, %v) but got (%v, %v)", errNoClientName, c, err)
	}
}

func TestRabbitMQHAConfigDefaultValue(t *testing.T) {
	haConfig := (&Config{}).RabbitMQHAConfig()

	expectedConfigCont := 1

	if len(haConfig) != expectedConfigCont {
		t.Errorf(
			"Expected the config count to be %d but got %d instead!",
			expectedConfigCont,
			len(haConfig),
		)
	}

	validateStringParameter(
		haConfig[0].GetURI(),
		defaultRabbitMQURI,
		"RabbitMQ URI",
		t,
	)
}
<|endoftext|>


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

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	"github.com/golang/protobuf/proto"
	eventpb "github.com/google/schedviz/tracedata/schedviz_events_go_proto"

	"github.com/google/schedviz/analysis/sched"
	"github.com/google/schedviz/server/models"
	"github.com/google/schedviz/tracedata/trace"
)

// FsStorage is a storage service that saves collections as protos on local disk
// Implements StorageService
type FsStorage struct {
	*storageBase
	StoragePath string
}

// CreateFSStorage creates a new file system storage service that stores its files at storagePath
// and has an LRU cache of size cacheSize.
func CreateFSStorage(storagePath string, cacheSize int) (*FsStorage, error) {
	sb, err := newStorageBase(cacheSize)
	if err != nil {
		return nil, err
	}
	return &FsStorage{
		storageBase: sb,
		StoragePath: storagePath,
	}, nil
}

// DeleteCollection deletes the collection with the given name.
func (fs *FsStorage) DeleteCollection(_ context.Context, _ string, collectionUniqueName string) error {
	filePath := path.Join(fs.StoragePath, collectionUniqueName+".binproto")
	fs.mu.Lock()
	fs.lruCache.Remove(filePath)
	defer fs.mu.Unlock()
	if err := os.Remove(filePath); err != nil {
		return err
	}
	return nil
}

func (fs *FsStorage) getCollectionPath(collectionName string) string {
	return path.Join(fs.StoragePath, collectionName+".binproto")
}

func (fs *FsStorage) getCollectionNameFromFileName(fileName string) string {
	return strings.TrimSuffix(fileName, ".binproto")
}

func (fs *FsStorage) getCollectionFromDisk(collectionName string) (*eventpb.Collection, error) {
	filePath := fs.getCollectionPath(collectionName)
	bytes, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	collectionProto := &eventpb.Collection{}
	if err := proto.Unmarshal(bytes, collectionProto); err != nil {
		return nil, err
	}
	return collectionProto, nil
}

// GetCollection returns an already-saved collection with the given name.
// shouldCache controls whether or not the fetched collection will be saved in the cache to speed up
// future requests for the same collection.
func (fs *FsStorage) GetCollection(ctx context.Context, collectionName string) (*CachedCollection, error) {
	cachedCollection, ok, err := fs.getCollectionFromCache(collectionName, true /*= addCollection*/)
	if err != nil {
		return nil, err
	}
	if ok {
		if err := cachedCollection.wait(ctx); err != nil {
			return nil, err
		}
		return cachedCollection, cachedCollection.err
	}
	defer func() {
		cachedCollection.err = err
		cachedCollection.release()
	}()
	collectionProto, err := fs.getCollectionFromDisk(collectionName)
	if err != nil {
		return nil, err
	}
	collection, err := sched.NewCollection(collectionProto.EventSet,
		sched.NormalizeTimestamps(true))
	if err != nil {
		return nil, err
	}
	cachedCollection.Collection = collection
	cachedCollection.SystemTopology = convertTopologyProtoToStruct(collectionProto.Topology)
	cachedCollection.Payload = map[string]interface{}{}
	return cachedCollection, nil
}

// GetCollectionMetadata gets the metadata for the collection with the given name.
func (fs *FsStorage) GetCollectionMetadata(ctx context.Context, collectionUniqueName string) (models.Metadata, error) {
	collectionProto, err := fs.getCollectionFromDisk(collectionUniqueName)
	if err != nil {
		return models.Metadata{}, err
	}
	return convertMetadataProtoToStruct(collectionProto.Metadata), nil
}

// EditCollection edits the metadata for the collection with the given name.
func (fs *FsStorage) EditCollection(ctx context.Context, _ string, req *models.EditCollectionRequest) error {
	if len(req.CollectionName) == 0 {
		return missingFieldError("collection_name")
	}
	metadata, err := fs.GetCollectionMetadata(ctx, req.CollectionName)
	if err != nil {
		return err
	}

	oldTags := metadata.Tags
	tagSet := make(map[string]struct{})
	for _, tag := range oldTags {
		tagSet[tag] = struct{}{}
	}
	for _, tag := range req.RemoveTags {
		delete(tagSet, tag)
	}
	for _, tag := range req.AddTags {
		tagSet[tag] = struct{}{}
	}
	var newTags = []string{}
	for tag := range tagSet {
		newTags = append(newTags, tag)
	}
	metadata.Tags = newTags

	ownerSet := make(map[string]struct{})
	for _, owner := range metadata.Owners {
		ownerSet[owner] = struct{}{}
	}
	for _, owner := range req.AddOwners {
		ownerSet[owner] = struct{}{}
	}
	var newOwners = []string{}
	for owner := range ownerSet {
		newOwners = append(newOwners, owner)
	}
	metadata.Owners = newOwners
	metadata.Description = req.Description

	// Write updated metadata to file
	collectionProto, err := fs.getCollectionFromDisk(req.CollectionName)
	if err != nil {
		return err
	}
	metadataProto, err := convertMetadataStructToProto(&metadata)
	if err != nil {
		return err
	}
	collectionProto.Metadata = metadataProto

	outBytes, err := proto.Marshal(collectionProto)
	if err != nil {
		return err
	}
	filePath := fs.getCollectionPath(req.CollectionName)
	if err := ioutil.WriteFile(filePath, outBytes, 0644); err != nil {
		return err
	}
	return nil
}

// ListCollectionMetadata gets the metadata for all collections.
func (fs *FsStorage) ListCollectionMetadata(ctx context.Context, _ string, _ string) ([]models.Metadata, error) {
	files, err := ioutil.ReadDir(fs.StoragePath)
	if err != nil {
		return nil, err
	}
	// Force initialize as an empty, not nil, slice.
	// Without the curly braces, this will appear as null when empty and serialized to JSON, instead
	// of an empty array.
	var ret = []models.Metadata{}
	for _, file := range files {
		collectionName := fs.getCollectionNameFromFileName(file.Name())
		collectionProto, err := fs.getCollectionFromDisk(collectionName)
		if err != nil {
			return nil, err
		}
		metadata := convertMetadataProtoToStruct(collectionProto.Metadata)
		ret = append(ret, metadata)
	}
	return ret, nil
}

// GetCollectionParameters gets the parameters for the collection with the given name.
func (fs *FsStorage) GetCollectionParameters(ctx context.Context, collectionName string) (models.CollectionParametersResponse, error) {
	if len(collectionName) == 0 {
		return models.CollectionParametersResponse{}, missingFieldError("collection_name")
	}
	collection, err := fs.GetCollection(ctx, collectionName)
	if err != nil {
		return models.CollectionParametersResponse{}, err
	}

	startTimestamp, endTimestamp := collection.Collection.Interval()

	ftraceEvents := collection.Collection.TraceCollection.EventNames()

	ret := models.CollectionParametersResponse{
		CollectionName:   collectionName,
		CPUs:             collection.Collection.ExpandCPUs(nil),
		StartTimestampNs: int64(startTimestamp),
		EndTimestampNs:   int64(endTimestamp),
		FtraceEvents:     ftraceEvents,
	}

	return ret, nil
}

// GetFtraceEvents returns events for the specified collection
func (fs *FsStorage) GetFtraceEvents(ctx context.Context, req *models.FtraceEventsRequest) (models.FtraceEventsResponse, error) {
	if len(req.CollectionName) == 0 {
		return models.FtraceEventsResponse{}, missingFieldError("collection_name")
	}
	collection, err := fs.GetCollection(ctx, req.CollectionName)
	if err != nil {
		return models.FtraceEventsResponse{}, err
	}

	var cpus = []sched.CPUID{}
	for _, cpu := range req.Cpus {
		cpus = append(cpus, sched.CPUID(cpu))
	}

	events, err := collection.Collection.GetRawEvents(
		sched.CPUs(cpus...),
		sched.TimeRange(trace.Timestamp(req.StartTimestamp), trace.Timestamp(req.EndTimestamp)),
		sched.EventTypes(req.EventTypes...))
	if err != nil {
		return models.FtraceEventsResponse{}, err
	}

	var fTraceEventsByCPU = make(map[sched.CPUID][]*trace.Event)
	for _, ev := range events {
		cpuID := sched.CPUID(ev.CPU)
		fTraceEventsByCPU[cpuID] = append(fTraceEventsByCPU[cpuID], ev)
	}

	return models.FtraceEventsResponse{
		CollectionName: req.CollectionName,
		EventsByCPU:    fTraceEventsByCPU,
	}, nil
}

// SetFailOnUnknownEventFormat configures behavior when encountering an unknown
// event format.  If the provided bool is true, parsing fails on unknown events;
// otherwise unknown events are logged and ignored.
func (fs *FsStorage) SetFailOnUnknownEventFormat(option bool) {
	fs.failOnUnknownEventFormat = option
}

func missingFieldError(fieldName string) error {
	return fmt.Errorf("missing required field %q`", fieldName)
}
<commit_msg>Fall back to fault-tolerant loader if default event loader is unable to create collection.<commit_after>//
// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS-IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
//
package storageservice

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	eventpb "github.com/google/schedviz/tracedata/schedviz_events_go_proto"

	"github.com/google/schedviz/analysis/sched"
	"github.com/google/schedviz/server/models"
	"github.com/google/schedviz/tracedata/trace"
)

// FsStorage is a storage service that saves collections as protos on local disk
// Implements StorageService
type FsStorage struct {
	*storageBase
	StoragePath string
}

// CreateFSStorage creates a new file system storage service that stores its files at storagePath
// and has an LRU cache of size cacheSize.
func CreateFSStorage(storagePath string, cacheSize int) (*FsStorage, error) {
	sb, err := newStorageBase(cacheSize)
	if err != nil {
		return nil, err
	}
	return &FsStorage{
		storageBase: sb,
		StoragePath: storagePath,
	}, nil
}

// DeleteCollection deletes the collection with the given name.
func (fs *FsStorage) DeleteCollection(_ context.Context, _ string, collectionUniqueName string) error {
	filePath := path.Join(fs.StoragePath, collectionUniqueName+".binproto")
	fs.mu.Lock()
	fs.lruCache.Remove(filePath)
	defer fs.mu.Unlock()
	if err := os.Remove(filePath); err != nil {
		return err
	}
	return nil
}

func (fs *FsStorage) getCollectionPath(collectionName string) string {
	return path.Join(fs.StoragePath, collectionName+".binproto")
}

func (fs *FsStorage) getCollectionNameFromFileName(fileName string) string {
	return strings.TrimSuffix(fileName, ".binproto")
}

func (fs *FsStorage) getCollectionFromDisk(collectionName string) (*eventpb.Collection, error) {
	filePath := fs.getCollectionPath(collectionName)
	bytes, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	collectionProto := &eventpb.Collection{}
	if err := proto.Unmarshal(bytes, collectionProto); err != nil {
		return nil, err
	}
	return collectionProto, nil
}

// GetCollection returns an already-saved collection with the given name.
// shouldCache controls whether or not the fetched collection will be saved in the cache to speed up
// future requests for the same collection.
func (fs *FsStorage) GetCollection(ctx context.Context, collectionName string) (*CachedCollection, error) {
	cachedCollection, ok, err := fs.getCollectionFromCache(collectionName, true /*= addCollection*/)
	if err != nil {
		return nil, err
	}
	if ok {
		if err := cachedCollection.wait(ctx); err != nil {
			return nil, err
		}
		return cachedCollection, cachedCollection.err
	}
	defer func() {
		cachedCollection.err = err
		cachedCollection.release()
	}()
	collectionProto, err := fs.getCollectionFromDisk(collectionName)
	if err != nil {
		return nil, err
	}
	collection, err := createCollection(collectionProto.EventSet)
	if err != nil {
		return nil, err
	}
	cachedCollection.Collection = collection
	cachedCollection.SystemTopology = convertTopologyProtoToStruct(collectionProto.Topology)
	cachedCollection.Payload = map[string]interface{}{}
	return cachedCollection, nil
}

// GetCollectionMetadata gets the metadata for the collection with the given name.
func (fs *FsStorage) GetCollectionMetadata(ctx context.Context, collectionUniqueName string) (models.Metadata, error) {
	collectionProto, err := fs.getCollectionFromDisk(collectionUniqueName)
	if err != nil {
		return models.Metadata{}, err
	}
	return convertMetadataProtoToStruct(collectionProto.Metadata), nil
}

// EditCollection edits the metadata for the collection with the given name.
func (fs *FsStorage) EditCollection(ctx context.Context, _ string, req *models.EditCollectionRequest) error {
	if len(req.CollectionName) == 0 {
		return missingFieldError("collection_name")
	}
	metadata, err := fs.GetCollectionMetadata(ctx, req.CollectionName)
	if err != nil {
		return err
	}

	oldTags := metadata.Tags
	tagSet := make(map[string]struct{})
	for _, tag := range oldTags {
		tagSet[tag] = struct{}{}
	}
	for _, tag := range req.RemoveTags {
		delete(tagSet, tag)
	}
	for _, tag := range req.AddTags {
		tagSet[tag] = struct{}{}
	}
	var newTags = []string{}
	for tag := range tagSet {
		newTags = append(newTags, tag)
	}
	metadata.Tags = newTags

	ownerSet := make(map[string]struct{})
	for _, owner := range metadata.Owners {
		ownerSet[owner] = struct{}{}
	}
	for _, owner := range req.AddOwners {
		ownerSet[owner] = struct{}{}
	}
	var newOwners = []string{}
	for owner := range ownerSet {
		newOwners = append(newOwners, owner)
	}
	metadata.Owners = newOwners
	metadata.Description = req.Description

	// Write updated metadata to file
	collectionProto, err := fs.getCollectionFromDisk(req.CollectionName)
	if err != nil {
		return err
	}
	metadataProto, err := convertMetadataStructToProto(&metadata)
	if err != nil {
		return err
	}
	collectionProto.Metadata = metadataProto

	outBytes, err := proto.Marshal(collectionProto)
	if err != nil {
		return err
	}
	filePath := fs.getCollectionPath(req.CollectionName)
	if err := ioutil.WriteFile(filePath, outBytes, 0644); err != nil {
		return err
	}
	return nil
}

// ListCollectionMetadata gets the metadata for all collections.
func (fs *FsStorage) ListCollectionMetadata(ctx context.Context, _ string, _ string) ([]models.Metadata, error) {
	files, err := ioutil.ReadDir(fs.StoragePath)
	if err != nil {
		return nil, err
	}
	// Force initialize as an empty, not nil, slice.
	// Without the curly braces, this will appear as null when empty and serialized to JSON, instead
	// of an empty array.
	var ret = []models.Metadata{}
	for _, file := range files {
		collectionName := fs.getCollectionNameFromFileName(file.Name())
		collectionProto, err := fs.getCollectionFromDisk(collectionName)
		if err != nil {
			return nil, err
		}
		metadata := convertMetadataProtoToStruct(collectionProto.Metadata)
		ret = append(ret, metadata)
	}
	return ret, nil
}

// GetCollectionParameters gets the parameters for the collection with the given name.
func (fs *FsStorage) GetCollectionParameters(ctx context.Context, collectionName string) (models.CollectionParametersResponse, error) {
	if len(collectionName) == 0 {
		return models.CollectionParametersResponse{}, missingFieldError("collection_name")
	}
	collection, err := fs.GetCollection(ctx, collectionName)
	if err != nil {
		return models.CollectionParametersResponse{}, err
	}

	startTimestamp, endTimestamp := collection.Collection.Interval()

	ftraceEvents := collection.Collection.TraceCollection.EventNames()

	ret := models.CollectionParametersResponse{
		CollectionName:   collectionName,
		CPUs:             collection.Collection.ExpandCPUs(nil),
		StartTimestampNs: int64(startTimestamp),
		EndTimestampNs:   int64(endTimestamp),
		FtraceEvents:     ftraceEvents,
	}

	return ret, nil
}

// GetFtraceEvents returns events for the specified collection
func (fs *FsStorage) GetFtraceEvents(ctx context.Context, req *models.FtraceEventsRequest) (models.FtraceEventsResponse, error) {
	if len(req.CollectionName) == 0 {
		return models.FtraceEventsResponse{}, missingFieldError("collection_name")
	}
	collection, err := fs.GetCollection(ctx, req.CollectionName)
	if err != nil {
		return models.FtraceEventsResponse{}, err
	}

	var cpus = []sched.CPUID{}
	for _, cpu := range req.Cpus {
		cpus = append(cpus, sched.CPUID(cpu))
	}

	events, err := collection.Collection.GetRawEvents(
		sched.CPUs(cpus...),
		sched.TimeRange(trace.Timestamp(req.StartTimestamp), trace.Timestamp(req.EndTimestamp)),
		sched.EventTypes(req.EventTypes...))
	if err != nil {
		return models.FtraceEventsResponse{}, err
	}

	var fTraceEventsByCPU = make(map[sched.CPUID][]*trace.Event)
	for _, ev := range events {
		cpuID := sched.CPUID(ev.CPU)
		fTraceEventsByCPU[cpuID] = append(fTraceEventsByCPU[cpuID], ev)
	}

	return models.FtraceEventsResponse{
		CollectionName: req.CollectionName,
		EventsByCPU:    fTraceEventsByCPU,
	}, nil
}

// SetFailOnUnknownEventFormat configures behavior when encountering an unknown
// event format.  If the provided bool is true, parsing fails on unknown events;
// otherwise unknown events are logged and ignored.
func (fs *FsStorage) SetFailOnUnknownEventFormat(option bool) {
	fs.failOnUnknownEventFormat = option
}

// createCollection creates a collection with the default event loader, and
// will attempt to create a collection with the fault tolerant loader if the
// default loader failed.
func createCollection(es *eventpb.EventSet) (*sched.Collection, error) {
	coll, err := sched.NewCollection(es, sched.NormalizeTimestamps(true))
	if err == nil {
		return coll, nil
	}
	log.Warning("Failed to load collection with default loader. " +
		"Retrying with fault tolerant loader.")
	coll, err = sched.NewCollection(es,
		sched.NormalizeTimestamps(true),
		sched.UsingEventLoaders(sched.FaultTolerantEventLoaders()))
	if err != nil {
		return nil, err
	}
	return coll, nil
}

func missingFieldError(fieldName string) error {
	return fmt.Errorf("missing required field %q`", fieldName)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/all"
)

// Lidar is structure to access basic functions of LidarLite
// LidarLite_V2 blue label
// Documentation on http://lidarlite.com/docs/v2/specs_and_hardware
type Lidar struct {
	bus     embd.I2CBus
	address byte
}

// Ready - Ready status. 0 - ready for a new command, 1 - busy
//          with acquisition
// RefOverflow - Overflow detected in correlation process assotieted with
//          a reference acquisition
// SignalOverflow - Overflow detected in correlation process assotieted with a signal
//          acquisition
// SignalNotValid - Indicates that the signal correlation peak is equal to or below
//          correlation record threshold
// SecondaryReturn - Secondary return detected above correlation noise floor threshold
// Health - 1 if is good, 0 if is bad
// ErrorDetection - Process error detected/measurement invalid
// EyeSafe - This bit will go high if eye-safety protection has been activated
const (
	Ready = 1 << iota

// RefOverflow     = 1 << iota
// SigOverflow     = 1 << iota
// SignalNotValid  = 1 << iota
// SecondatyReturn = 1 << iota
// Health          = 1 << iota
// ErrorDetection  = 1 << iota
// EyeSafe         = 1 << iota
)

// NewLidar sets the configuration for the sensor
// Write 0x00 to Register 0x00 reset FPGA. Re-loads FPGA from internal Flash
// memory: all registers return to default values
func NewLidar(i2cbus, addr byte) *Lidar {

	lSensor := Lidar{bus: embd.NewI2CBus(i2cbus), address: byte(addr)}
	if e := lSensor.bus.WriteByteToReg(lSensor.address, 0x00, 0x00); e != nil {
		log.Panic("Write ", e)
	}
	time.Sleep(300 * time.Millisecond)
	return &lSensor
}

//CloseLidar closes releases the resources associated with the bus
func (ls *Lidar) CloseLidar() {
	if err := ls.bus.Close(); err == nil {
		log.Println(err)
	}
}

// GetStatus gets Mode/Status of sensor
func (ls *Lidar) GetStatus() (byte, error) {

	val, err := ls.bus.ReadByteFromReg(ls.address, 0x01)
	if err == nil {
		log.Printf("Status: %.8b\n ", val)
	} else {
		log.Println("GetStatus ", err)
		return 0, err
	}
	return val, nil

}

// Distance reads the distance from LidarLite
// stablizePreampFlag - true - take aquisition with DC stabilisation/correction.
// false - it will read faster, but you will need to stabilize DC every once in
// awhile(ex. 1 out of every 100 readings is typically good)
func (ls *Lidar) Distance(stablizePreampFlag bool) (int, error) {

	var wErr error // Write and Read errors

	if _, errSt := ls.GetStatus(); errSt == nil {
		if stablizePreampFlag {
			wErr = ls.bus.WriteByteToReg(ls.address, 0x00, 0x04)
		} else {
			wErr = ls.bus.WriteByteToReg(ls.address, 0x00, 0x03)
		}

		if wErr != nil {
			log.Println("Write ", wErr)
			return 0, wErr
		}

		time.Sleep(100 * time.Millisecond)
	} else {
		log.Println(errSt)
		return 0, errSt
	}

	_, errSt := ls.GetStatus()
	if errSt == nil {
		v1, rErr := ls.bus.ReadByteFromReg(ls.address, 0x10)
		if rErr != nil {
			log.Println("Read ", rErr)
			return 0, rErr
		}
		v2, rErr := ls.bus.ReadByteFromReg(ls.address, 0x0f)
		if rErr != nil {
			log.Println("Read", rErr)
			return 0, rErr
		}

		return ((int(v2) << 8) + int(v1)), nil

	}
	log.Println(errSt)
	return 0, errSt
}

// Velocity is measured by observing the change in distance over a fixed time
// of perion
// TODO 0x04 Check Mode Control
func (ls *Lidar) Velocity() (int, error) {
	// Write 0xa0 to 0x04 to switch on velocity mode
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x04, 0xa0); wErr != nil {
		log.Println("Write ", wErr)
		return -1, wErr
	}
	// Write 0x04 to register 0x00 to start getting distance readings
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x00, 0x04); wErr != nil {
		log.Println("Write ", wErr)
		return -1, wErr
	}
	//Read 1 byte from register 0x09 to get velocity measurement
	for {
		// Sensor is ready for reading
		stVal, err := ls.GetStatus()
		if err != nil {
			log.Println(err)
			return -1, err
		}
		if (stVal & Ready) == 0 {
			val, e := ls.bus.ReadByteFromReg(ls.address, 0x09)
			if err != nil {
				log.Println(e)
				return -1, e
			}
			return int(val), nil
		}
		time.Sleep(300 * time.Microsecond)

	}
}

func main() {
	log.SetFlags(log.Lshortfile)
	lidar := NewLidar(1, 0x62) // 0x62 the default LidarLite address
	defer lidar.CloseLidar()
	for {
		if val, err := lidar.Distance(true); err == nil {
			fmt.Println(val)
			time.Sleep(1 * time.Second)
		}
		/*if val, err := lidar.Velocity(); err == nil {
			fmt.Println(val)
		}*/
	}

}
<commit_msg>Add Continuous mode<commit_after>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/all"
)

// Lidar is structure to access basic functions of LidarLite
// LidarLite_V2 blue label
// Documentation on http://lidarlite.com/docs/v2/specs_and_hardware
type Lidar struct {
	bus     embd.I2CBus
	address byte
}

// Ready - Ready status. 0 - ready for a new command, 1 - busy
//          with acquisition
// RefOverflow - Overflow detected in correlation process assotieted with
//          a reference acquisition
// SignalOverflow - Overflow detected in correlation process assotieted with a signal
//          acquisition
// SignalNotValid - Indicates that the signal correlation peak is equal to or below
//          correlation record threshold
// SecondaryReturn - Secondary return detected above correlation noise floor threshold
// Health - 1 if is good, 0 if is bad
// ErrorDetection - Process error detected/measurement invalid
// EyeSafe - This bit will go high if eye-safety protection has been activated
const (
	Ready = 1 << iota

// RefOverflow     = 1 << iota
// SigOverflow     = 1 << iota
// SignalNotValid  = 1 << iota
// SecondatyReturn = 1 << iota
// Health          = 1 << iota
// ErrorDetection  = 1 << iota
// EyeSafe         = 1 << iota
)

// NewLidar sets the configuration for the sensor
// Write 0x00 to Register 0x00 reset FPGA. Re-loads FPGA from internal Flash
// memory: all registers return to default values
func NewLidar(i2cbus, addr byte) *Lidar {

	lSensor := Lidar{bus: embd.NewI2CBus(i2cbus), address: byte(addr)}
	if e := lSensor.bus.WriteByteToReg(lSensor.address, 0x00, 0x00); e != nil {
		log.Panic("Write ", e)
	}
	time.Sleep(300 * time.Millisecond)
	return &lSensor
}

//CloseLidar closes releases the resources associated with the bus
func (ls *Lidar) CloseLidar() {
	if err := ls.bus.Close(); err != nil {
		log.Println(err)
	}
}

// GetStatus gets Mode/Status of sensor
func (ls *Lidar) GetStatus() (byte, error) {

	val, err := ls.bus.ReadByteFromReg(ls.address, 0x01)
	if err == nil {
		log.Printf("Status: %.8b\n", val)
	} else {
		log.Println("GetStatus", err)
		return 0, err
	}
	return val, nil

}

// Distance reads the distance from LidarLite
// stablizePreampFlag - true - take aquisition with DC stabilisation/correction.
// false - it will read faster, but you will need to stabilize DC every once in
// awhile(ex. 1 out of every 100 readings is typically good)
func (ls *Lidar) Distance(stablizePreampFlag bool) (int, error) {

	var wErr error // Write and Read errors

	if _, errSt := ls.GetStatus(); errSt == nil {
		if stablizePreampFlag {
			wErr = ls.bus.WriteByteToReg(ls.address, 0x00, 0x04)
		} else {
			wErr = ls.bus.WriteByteToReg(ls.address, 0x00, 0x03)
		}

		if wErr != nil {
			log.Println("Write ", wErr)
			return 0, wErr
		}

		time.Sleep(100 * time.Millisecond)
	} else {
		log.Println(errSt)
		return 0, errSt
	}

	_, errSt := ls.GetStatus()
	if errSt == nil {
		v1, rErr := ls.bus.ReadByteFromReg(ls.address, 0x10)
		if rErr != nil {
			log.Println("Read ", rErr)
			return 0, rErr
		}
		v2, rErr := ls.bus.ReadByteFromReg(ls.address, 0x0f)
		if rErr != nil {
			log.Println("Read", rErr)
			return 0, rErr
		}

		return ((int(v2) << 8) + int(v1)), nil

	}
	log.Println(errSt)
	return 0, errSt
}

// Velocity is measured by observing the change in distance over a fixed time
// of perion
// TODO 0x04 Check Mode Control
func (ls *Lidar) Velocity() (int, error) {
	// Write 0xa0 to 0x04 to switch on velocity mode
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x04, 0xa0); wErr != nil {
		log.Println("Write ", wErr)
		return -1, wErr
	}
	// Write 0x04 to register 0x00 to start getting distance readings
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x00, 0x04); wErr != nil {
		log.Println("Write ", wErr)
		return -1, wErr
	}
	//Read 1 byte from register 0x09 to get velocity measurement
	for {
		// Sensor is ready for reading
		stVal, err := ls.GetStatus()
		if err != nil {
			log.Println(err)
			return -1, err
		}
		if (stVal & Ready) == 0 {
			val, e := ls.bus.ReadByteFromReg(ls.address, 0x09)
			if e != nil {
				log.Println(e)
				return -1, e
			}
			return int(val), nil
		}
		time.Sleep(300 * time.Microsecond)

	}
}

// BeginContinuous allows to tell the sensor to take a certain number (or
// infinite) readings allowing you to read from it at a continuous rate.
// modePinLow tells the mode pin to go low when a new reading is available.
// TODO search more about interval
func (ls *Lidar) BeginContinuous(modePinLow bool, interval, numberOfReadings byte) error {

	// Register 0x45 sets the time between measurements. Min val os 0x02
	// for proper operations.
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x45, interval); wErr != nil {
		log.Println(wErr)
		return wErr
	}

	if modePinLow {
		if wErr := ls.bus.WriteByteToReg(ls.address, 0x04, 0x21); wErr != nil {
			log.Print(wErr)
			return wErr
		}

	} else {
		// Set register 0x04 to 0x20 to look at "NON-default" value of velocity
		if wErr := ls.bus.WriteByteToReg(ls.address, 0x04, 0x20); wErr != nil {
			log.Print(wErr)
			return wErr
		}
	}
	// Set the number of readings, 0xfe = 254 readings, 0x01 = 1 reading and
	// 0xff = continuous readings
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x11, numberOfReadings); wErr != nil {
		log.Println(wErr)
		return wErr
	}

	// Initiate reading distance
	if wErr := ls.bus.WriteByteToReg(ls.address, 0x00, 0x04); wErr != nil {
		log.Println(wErr)
		return wErr
	}
	return nil
}

//DistanceContinuous reads in continuous mode
func (ls *Lidar) DistanceContinuous() (int, error) {

	val, rErr := ls.bus.ReadWordFromReg(ls.address, 0x8f)
	if rErr != nil {
		log.Println("Read", rErr)
		return -1, rErr
	}

	return int(val), nil
}

func main() {
	log.SetFlags(log.Lshortfile)
	lidar := NewLidar(1, 0x62) // 0x62 the default LidarLite address
	defer lidar.CloseLidar()
	/*for {
		if val, err := lidar.Distance(true); err == nil {
			fmt.Println(val)
			time.Sleep(1 * time.Second)
		}
		if val, err := lidar.Velocity(); err == nil {
			fmt.Println(val)
		}
	}*/

	if err := lidar.BeginContinuous(true, 0x08, 0xff); err == nil {
		for {
			val, e := lidar.DistanceContinuous()

			if e == nil {
				fmt.Println(val)
			}
		}
	}
}
<|endoftext|>