<commit_before>package cli

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

const terraformCodeExampleOutputOnly = `
output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleGcpProvider = `
provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderEmptyOriginal = `
provider "aws" {
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderRegionOverridenExpected = `
provider "aws" {
  region = "eu-west-1"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderRegionVersionProfileOverridenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.3.0"
  profile = "foo"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderNonEmptyOriginal = `
provider "aws" {
  region  = var.aws_region
  version = "0.2.0"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderRegionOverridenVersionNotOverriddenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.2.0"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersOriginal = `
provider "aws" {
  region  = var.aws_region
  version = "0.2.0"
}

provider "aws" {
  alias   = "another"
  region  = var.aws_region
  version = "0.2.0"
}

resource "aws_instance" "example" {

}

provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

provider "aws" {
  alias  = "yet another"
  region = var.aws_region
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersRegionOverridenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.2.0"
}

provider "aws" {
  alias   = "another"
  region  = "eu-west-1"
  version = "0.2.0"
}

resource "aws_instance" "example" {

}

provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

provider "aws" {
  alias  = "yet another"
  region = "eu-west-1"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersRegionProfileVersionOverridenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.3.0"
  profile = "foo"
}

provider "aws" {
  alias   = "another"
  region  = "eu-west-1"
  version = "0.3.0"
  profile = "foo"
}

resource "aws_instance" "example" {

}

provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

provider "aws" {
  alias   = "yet another"
  region  = "eu-west-1"
  version = "0.3.0"
  profile = "foo"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal = `
# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = var.aws_region

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"
}

# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = var.aws_region

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"

  # Make sure comments are maintained
  # And don't interfere with parsing
  alias = "secondary"
}
`

const terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionOverriddenExpected = `
# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"
}

# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"

  # Make sure comments are maintained
  # And don't interfere with parsing
  alias = "secondary"
}
`

const terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionVersionProfileOverriddenExpected = `
# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.3.0"
  profile = "foo"
}

# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.3.0"

  # Make sure comments are maintained
  # And don't interfere with parsing
  alias   = "secondary"
  profile = "foo"
}
`

func TestPatchAwsProviderInTerraformCodeHappyPath(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		testName               string
		originalTerraformCode  string
		attributesToOverride   map[string]string
		expectedCodeWasUpdated bool
		expectedTerraformCode  string
	}{
		{"empty", "", nil, false, ""},
		{"empty with attributes", "", map[string]string{"region": "eu-west-1"}, false, ""},
		{"no provider", terraformCodeExampleOutputOnly, map[string]string{"region": "eu-west-1"}, false, terraformCodeExampleOutputOnly},
		{"no aws provider", terraformCodeExampleGcpProvider, map[string]string{"region": "eu-west-1"}, false, terraformCodeExampleGcpProvider},
		{"one empty aws provider, but no overrides", terraformCodeExampleAwsProviderEmptyOriginal, nil, false, terraformCodeExampleAwsProviderEmptyOriginal},
		{"one empty aws provider, with region override", terraformCodeExampleAwsProviderEmptyOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsProviderRegionOverridenExpected},
		{"one empty aws provider, with region, version, profile override", terraformCodeExampleAwsProviderEmptyOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0", "profile": "foo"}, true, terraformCodeExampleAwsProviderRegionVersionProfileOverridenExpected},
		{"one non-empty aws provider, but no overrides", terraformCodeExampleAwsProviderNonEmptyOriginal, nil, false, terraformCodeExampleAwsProviderNonEmptyOriginal},
		{"one non-empty aws provider, with region override", terraformCodeExampleAwsProviderNonEmptyOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsProviderRegionOverridenVersionNotOverriddenExpected},
		{"one non-empty aws provider, with region, version, profile override", terraformCodeExampleAwsProviderNonEmptyOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0", "profile": "foo"}, true, terraformCodeExampleAwsProviderRegionVersionProfileOverridenExpected},
		{"multiple providers, but no overrides", terraformCodeExampleAwsMultipleProvidersOriginal, nil, false, terraformCodeExampleAwsMultipleProvidersOriginal},
		{"multiple providers, with region override", terraformCodeExampleAwsMultipleProvidersOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsMultipleProvidersRegionOverridenExpected},
		{"multiple providers, with region, version, profile override", terraformCodeExampleAwsMultipleProvidersOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0", "profile": "foo"}, true, terraformCodeExampleAwsMultipleProvidersRegionProfileVersionOverridenExpected},
		{"multiple providers with comments, but no overrides", terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal, nil, false, terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal},
		{"multiple providers with comments, with region override", terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionOverriddenExpected},
		{"multiple providers with comments, with region, version, profile override", terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0", "profile": "foo"}, true, terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionVersionProfileOverriddenExpected},
	}

	for _, testCase := range testCases {
		// The following is necessary to make sure testCase's values don't
		// get updated due to concurrency within the scope of t.Run(..) below
		testCase := testCase
		t.Run(testCase.testName, func(t *testing.T) {
			t.Parallel()
			actualTerraformCode, actualCodeWasUpdated, err := patchAwsProviderInTerraformCode(testCase.originalTerraformCode, "test.tf", testCase.attributesToOverride)
			assert.NoError(t, err)
			assert.Equal(t, testCase.expectedCodeWasUpdated, actualCodeWasUpdated)
			assert.Equal(t, testCase.expectedTerraformCode, actualTerraformCode)
		})
	}
}
<commit_msg>Get rid of profile from test<commit_after>package cli

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

const terraformCodeExampleOutputOnly = `
output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleGcpProvider = `
provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderEmptyOriginal = `
provider "aws" {
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderRegionOverridenExpected = `
provider "aws" {
  region = "eu-west-1"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderRegionVersionOverridenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.3.0"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderNonEmptyOriginal = `
provider "aws" {
  region  = var.aws_region
  version = "0.2.0"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsProviderRegionOverridenVersionNotOverriddenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.2.0"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersOriginal = `
provider "aws" {
  region  = var.aws_region
  version = "0.2.0"
}

provider "aws" {
  alias   = "another"
  region  = var.aws_region
  version = "0.2.0"
}

resource "aws_instance" "example" {

}

provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

provider "aws" {
  alias  = "yet another"
  region = var.aws_region
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersRegionOverridenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.2.0"
}

provider "aws" {
  alias   = "another"
  region  = "eu-west-1"
  version = "0.2.0"
}

resource "aws_instance" "example" {

}

provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

provider "aws" {
  alias  = "yet another"
  region = "eu-west-1"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersRegionVersionOverridenExpected = `
provider "aws" {
  region  = "eu-west-1"
  version = "0.3.0"
}

provider "aws" {
  alias   = "another"
  region  = "eu-west-1"
  version = "0.3.0"
}

resource "aws_instance" "example" {

}

provider "google" {
  credentials = file("account.json")
  project     = "my-project-id"
  region      = "us-central1"
}

provider "aws" {
  alias   = "yet another"
  region  = "eu-west-1"
  version = "0.3.0"
}

output "hello" {
  value = "Hello, World"
}
`

const terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal = `
# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = var.aws_region

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"
}

# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = var.aws_region

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"

  # Make sure comments are maintained
  # And don't interfere with parsing
  alias = "secondary"
}
`

const terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionOverriddenExpected = `
# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"
}

# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.2.0"

  # Make sure comments are maintained
  # And don't interfere with parsing
  alias = "secondary"
}
`

const terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionVersionOverriddenExpected = `
# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.3.0"
}

# Make sure comments are maintained
# And don't interfere with parsing
provider "aws" {
  # Make sure comments are maintained
  # And don't interfere with parsing
  region = "eu-west-1"

  # Make sure comments are maintained
  # And don't interfere with parsing
  version = "0.3.0"

  # Make sure comments are maintained
  # And don't interfere with parsing
  alias = "secondary"
}
`

func TestPatchAwsProviderInTerraformCodeHappyPath(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		testName               string
		originalTerraformCode  string
		attributesToOverride   map[string]string
		expectedCodeWasUpdated bool
		expectedTerraformCode  string
	}{
		{"empty", "", nil, false, ""},
		{"empty with attributes", "", map[string]string{"region": "eu-west-1"}, false, ""},
		{"no provider", terraformCodeExampleOutputOnly, map[string]string{"region": "eu-west-1"}, false, terraformCodeExampleOutputOnly},
		{"no aws provider", terraformCodeExampleGcpProvider, map[string]string{"region": "eu-west-1"}, false, terraformCodeExampleGcpProvider},
		{"one empty aws provider, but no overrides", terraformCodeExampleAwsProviderEmptyOriginal, nil, false, terraformCodeExampleAwsProviderEmptyOriginal},
		{"one empty aws provider, with region override", terraformCodeExampleAwsProviderEmptyOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsProviderRegionOverridenExpected},
		{"one empty aws provider, with region, version override", terraformCodeExampleAwsProviderEmptyOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0"}, true, terraformCodeExampleAwsProviderRegionVersionOverridenExpected},
		{"one non-empty aws provider, but no overrides", terraformCodeExampleAwsProviderNonEmptyOriginal, nil, false, terraformCodeExampleAwsProviderNonEmptyOriginal},
		{"one non-empty aws provider, with region override", terraformCodeExampleAwsProviderNonEmptyOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsProviderRegionOverridenVersionNotOverriddenExpected},
		{"one non-empty aws provider, with region, version override", terraformCodeExampleAwsProviderNonEmptyOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0"}, true, terraformCodeExampleAwsProviderRegionVersionOverridenExpected},
		{"multiple providers, but no overrides", terraformCodeExampleAwsMultipleProvidersOriginal, nil, false, terraformCodeExampleAwsMultipleProvidersOriginal},
		{"multiple providers, with region override", terraformCodeExampleAwsMultipleProvidersOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsMultipleProvidersRegionOverridenExpected},
		{"multiple providers, with region, version override", terraformCodeExampleAwsMultipleProvidersOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0"}, true, terraformCodeExampleAwsMultipleProvidersRegionVersionOverridenExpected},
		{"multiple providers with comments, but no overrides", terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal, nil, false, terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal},
		{"multiple providers with comments, with region override", terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal, map[string]string{"region": "eu-west-1"}, true, terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionOverriddenExpected},
		{"multiple providers with comments, with region, version override", terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsOriginal, map[string]string{"region": "eu-west-1", "version": "0.3.0"}, true, terraformCodeExampleAwsMultipleProvidersNonEmptyWithCommentsRegionVersionOverriddenExpected},
	}

	for _, testCase := range testCases {
		// The following is necessary to make sure testCase's values don't
		// get updated due to concurrency within the scope of t.Run(..) below
		testCase := testCase
		t.Run(testCase.testName, func(t *testing.T) {
			t.Parallel()
			actualTerraformCode, actualCodeWasUpdated, err := patchAwsProviderInTerraformCode(testCase.originalTerraformCode, "test.tf", testCase.attributesToOverride)
			assert.NoError(t, err)
			assert.Equal(t, testCase.expectedCodeWasUpdated, actualCodeWasUpdated)
			assert.Equal(t, testCase.expectedTerraformCode, actualTerraformCode)
		})
	}
}
<|endoftext|>


<commit_before>package renter

import (
	"bytes"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"

	"github.com/NebulousLabs/errors"
)

// cacheData contatins the data and the timestamp for the unfinished
// download chunks
type cacheData struct {
	data       []byte
	lastAccess time.Time
}

// downloadPieceInfo contains all the information required to download and
// recover a piece of a chunk from a host. It is a value in a map where the key
// is the file contract id.
type downloadPieceInfo struct {
	index uint64
	root  crypto.Hash
}

// unfinishedDownloadChunk contains a chunk for a download that is in progress.
//
// TODO: Currently, if a standby worker is needed, all of the standby workers
// are added and the first one that is available will pick up the slack. But,
// depending on the situation, we may only want to add a handful of workers to
// make sure that a fast / optimal worker is initially able to pick up the
// slack. This could potentially be streamlined by turning the standby array
// into a standby heap, and then having some general scoring system for figuring
// out how useful a worker is, and then having some threshold that a worker
// needs to be pulled from standby to work on the download. That threshold
// should go up every time that a worker fails, to make sure that if you have
// repeated failures, you keep pulling in the fresh workers instead of getting
// stuck and always rejecting all the standby workers.
type unfinishedDownloadChunk struct {
	// Fetch + Write instructions - read only or otherwise thread safe.
	destination downloadDestination // Where to write the recovered logical chunk.
	erasureCode modules.ErasureCoder
	masterKey   crypto.TwofishKey

	// Fetch + Write instructions - read only or otherwise thread safe.
	staticChunkIndex  uint64                                     // Required for deriving the encryption keys for each piece.
	staticCacheID     string                                     // Used to uniquely identify a chunk in the chunk cache.
	staticChunkMap    map[types.FileContractID]downloadPieceInfo // Maps from file contract ids to the info for the piece associated with that contract
	staticChunkSize   uint64
	staticFetchLength uint64 // Length within the logical chunk to fetch.
	staticFetchOffset uint64 // Offset within the logical chunk that is being downloaded.
	staticPieceSize   uint64
	staticWriteOffset int64 // Offet within the writer to write the completed data.

	// Fetch + Write instructions - read only or otherwise thread safe.
	staticLatencyTarget time.Duration
	staticNeedsMemory   bool // Set to true if memory was not pre-allocated for this chunk.
	staticOverdrive     int
	staticPriority      uint64

	// Download chunk state - need mutex to access.
	failed            bool      // Indicates if the chunk has been marked as failed.
	physicalChunkData [][]byte  // Used to recover the logical data.
	pieceUsage        []bool    // Which pieces are being actively fetched.
	piecesCompleted   int       // Number of pieces that have successfully completed.
	piecesRegistered  int       // Number of pieces that workers are actively fetching.
	recoveryComplete  bool      // Whether or not the recovery has completed and the chunk memory released.
	workersRemaining  int       // Number of workers still able to fetch the chunk.
	workersStandby    []*worker // Set of workers that are able to work on this download, but are not needed unless other workers fail.

	// Memory management variables.
	memoryAllocated uint64

	// The download object, mostly to update download progress.
	download *download
	mu       sync.Mutex

	// Caching related fields
	chunkCache map[string]*cacheData
	cacheMu    *sync.Mutex
}

// fail will set the chunk status to failed. The physical chunk memory will be
// wiped and any memory allocation will be returned to the renter. The download
// as a whole will be failed as well.
func (udc *unfinishedDownloadChunk) fail(err error) {
	udc.failed = true
	udc.recoveryComplete = true
	for i := range udc.physicalChunkData {
		udc.physicalChunkData[i] = nil
	}
	udc.download.managedFail(fmt.Errorf("chunk %v failed: %v", udc.staticChunkIndex, err))
	udc.destination = nil
}

// managedCleanUp will check if the download has failed, and if not it will add
// any standby workers which need to be added. Calling managedCleanUp too many
// times is not harmful, however missing a call to managedCleanUp can lead to
// dealocks.
func (udc *unfinishedDownloadChunk) managedCleanUp() {
	// Check if the chunk is newly failed.
	udc.mu.Lock()
	if udc.workersRemaining+udc.piecesCompleted < udc.erasureCode.MinPieces() && !udc.failed {
		udc.fail(errors.New("not enough workers to continue download"))
	}
	// Return any excess memory.
	udc.returnMemory()

	// Nothing to do if the chunk has failed.
	if udc.failed {
		udc.mu.Unlock()
		return
	}

	// Check whether standby workers are required.
	chunkComplete := udc.piecesCompleted >= udc.erasureCode.MinPieces()
	desiredPiecesRegistered := udc.erasureCode.MinPieces() + udc.staticOverdrive - udc.piecesCompleted
	standbyWorkersRequired := !chunkComplete && udc.piecesRegistered < desiredPiecesRegistered
	if !standbyWorkersRequired {
		udc.mu.Unlock()
		return
	}

	// Assemble a list of standby workers, release the udc lock, and then queue
	// the chunk into the workers. The lock needs to be released early because
	// holding the udc lock and the worker lock at the same time is a deadlock
	// risk (they interact with eachother, call functions on eachother).
	var standbyWorkers []*worker
	for i := 0; i < len(udc.workersStandby); i++ {
		standbyWorkers = append(standbyWorkers, udc.workersStandby[i])
	}
	udc.workersStandby = udc.workersStandby[:0] // Workers have been taken off of standby.
	udc.mu.Unlock()
	for i := 0; i < len(standbyWorkers); i++ {
		standbyWorkers[i].managedQueueDownloadChunk(udc)
	}
}

// managedRemoveWorker will decrement a worker from the set of remaining workers
// in the udc. After a worker has been removed, the udc needs to be cleaned up.
func (udc *unfinishedDownloadChunk) managedRemoveWorker() {
	udc.mu.Lock()
	udc.workersRemaining--
	udc.mu.Unlock()
	udc.managedCleanUp()
}

// returnMemory will check on the status of all the workers and pieces, and
// determine how much memory is safe to return to the renter. This should be
// called each time a worker returns, and also after the chunk is recovered.
func (udc *unfinishedDownloadChunk) returnMemory() {
	// The maximum amount of memory is the pieces completed plus the number of
	// workers remaining.
	maxMemory := uint64(udc.workersRemaining+udc.piecesCompleted) * udc.staticPieceSize
	// If enough pieces have completed, max memory is the number of registered
	// pieces plus the number of completed pieces.
	if udc.piecesCompleted >= udc.erasureCode.MinPieces() {
		// udc.piecesRegistered is guaranteed to be at most equal to the number
		// of overdrive pieces, meaning it will be equal to or less than
		// initalMemory.
		maxMemory = uint64(udc.piecesCompleted+udc.piecesRegistered) * udc.staticPieceSize
	}
	// If the chunk recovery has completed, the maximum number of pieces is the
	// number of registered.
	if udc.recoveryComplete {
		maxMemory = uint64(udc.piecesRegistered) * udc.staticPieceSize
	}
	// Return any memory we don't need.
	if uint64(udc.memoryAllocated) > maxMemory {
		udc.download.memoryManager.Return(udc.memoryAllocated - maxMemory)
		udc.memoryAllocated = maxMemory
	}
}

// threadedRecoverLogicalData will take all of the pieces that have been
// downloaded and encode them into the logical data which is then written to the
// underlying writer for the download.
func (udc *unfinishedDownloadChunk) threadedRecoverLogicalData() error {
	// Ensure cleanup occurs after the data is recovered, whether recovery
	// succeeds or fails.
	defer udc.managedCleanUp()

	// Decrypt the chunk pieces. This doesn't need to happen under a lock,
	// because any thread potentially writing to the physicalChunkData array is
	// going to be stopped by the fact that the chunk is complete.
	for i := range udc.physicalChunkData {
		// Skip empty pieces.
		if udc.physicalChunkData[i] == nil {
			continue
		}

		key := deriveKey(udc.masterKey, udc.staticChunkIndex, uint64(i))
		decryptedPiece, err := key.DecryptBytes(udc.physicalChunkData[i])
		if err != nil {
			udc.mu.Lock()
			udc.fail(err)
			udc.mu.Unlock()
			return errors.AddContext(err, "unable to decrypt chunk")
		}
		udc.physicalChunkData[i] = decryptedPiece
	}

	// Recover the pieces into the logical chunk data.
	//
	// TODO: Might be some way to recover into the downloadDestination instead
	// of creating a buffer and then writing that.
	recoverWriter := new(bytes.Buffer)
	err := udc.erasureCode.Recover(udc.physicalChunkData, udc.staticChunkSize, recoverWriter)
	if err != nil {
		udc.mu.Lock()
		udc.fail(err)
		udc.mu.Unlock()
		return errors.AddContext(err, "unable to recover chunk")
	}
	// Clear out the physical chunk pieces, we do not need them anymore.
	for i := range udc.physicalChunkData {
		udc.physicalChunkData[i] = nil
	}

	// Get recovered data
	recoveredData := recoverWriter.Bytes()

	// Add the chunk to the cache.
	udc.addChunkToCache(recoveredData)

	// Write the bytes to the requested output.
	start := udc.staticFetchOffset
	end := udc.staticFetchOffset + udc.staticFetchLength
	_, err = udc.destination.WriteAt(recoveredData[start:end], udc.staticWriteOffset)
	if err != nil {
		udc.mu.Lock()
		udc.fail(err)
		udc.mu.Unlock()
		return errors.AddContext(err, "unable to write to download destination")
	}
	recoverWriter = nil

	// Now that the download has completed and been flushed from memory, we can
	// release the memory that was used to store the data. Call 'cleanUp' to
	// trigger the memory cleanup along with some extra checks that everything
	// is consistent.
	udc.mu.Lock()
	udc.recoveryComplete = true
	udc.mu.Unlock()

	// Update the download and signal completion of this chunk.
	udc.download.mu.Lock()
	defer udc.download.mu.Unlock()
	udc.download.chunksRemaining--
	atomic.AddUint64(&udc.download.atomicDataReceived, udc.staticFetchLength)
	if udc.download.chunksRemaining == 0 {
		// Download is complete, send out a notification and close the
		// destination writer.
		udc.download.endTime = time.Now()
		close(udc.download.completeChan)
		return udc.download.destination.Close()
	}
	return nil
}
<commit_msg>also nil destination on successful download<commit_after>package renter

import (
	"bytes"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"

	"github.com/NebulousLabs/errors"
)

// cacheData contatins the data and the timestamp for the unfinished
// download chunks
type cacheData struct {
	data       []byte
	lastAccess time.Time
}

// downloadPieceInfo contains all the information required to download and
// recover a piece of a chunk from a host. It is a value in a map where the key
// is the file contract id.
type downloadPieceInfo struct {
	index uint64
	root  crypto.Hash
}

// unfinishedDownloadChunk contains a chunk for a download that is in progress.
//
// TODO: Currently, if a standby worker is needed, all of the standby workers
// are added and the first one that is available will pick up the slack. But,
// depending on the situation, we may only want to add a handful of workers to
// make sure that a fast / optimal worker is initially able to pick up the
// slack. This could potentially be streamlined by turning the standby array
// into a standby heap, and then having some general scoring system for figuring
// out how useful a worker is, and then having some threshold that a worker
// needs to be pulled from standby to work on the download. That threshold
// should go up every time that a worker fails, to make sure that if you have
// repeated failures, you keep pulling in the fresh workers instead of getting
// stuck and always rejecting all the standby workers.
type unfinishedDownloadChunk struct {
	// Fetch + Write instructions - read only or otherwise thread safe.
	destination downloadDestination // Where to write the recovered logical chunk.
	erasureCode modules.ErasureCoder
	masterKey   crypto.TwofishKey

	// Fetch + Write instructions - read only or otherwise thread safe.
	staticChunkIndex  uint64                                     // Required for deriving the encryption keys for each piece.
	staticCacheID     string                                     // Used to uniquely identify a chunk in the chunk cache.
	staticChunkMap    map[types.FileContractID]downloadPieceInfo // Maps from file contract ids to the info for the piece associated with that contract
	staticChunkSize   uint64
	staticFetchLength uint64 // Length within the logical chunk to fetch.
	staticFetchOffset uint64 // Offset within the logical chunk that is being downloaded.
	staticPieceSize   uint64
	staticWriteOffset int64 // Offet within the writer to write the completed data.

	// Fetch + Write instructions - read only or otherwise thread safe.
	staticLatencyTarget time.Duration
	staticNeedsMemory   bool // Set to true if memory was not pre-allocated for this chunk.
	staticOverdrive     int
	staticPriority      uint64

	// Download chunk state - need mutex to access.
	failed            bool      // Indicates if the chunk has been marked as failed.
	physicalChunkData [][]byte  // Used to recover the logical data.
	pieceUsage        []bool    // Which pieces are being actively fetched.
	piecesCompleted   int       // Number of pieces that have successfully completed.
	piecesRegistered  int       // Number of pieces that workers are actively fetching.
	recoveryComplete  bool      // Whether or not the recovery has completed and the chunk memory released.
	workersRemaining  int       // Number of workers still able to fetch the chunk.
	workersStandby    []*worker // Set of workers that are able to work on this download, but are not needed unless other workers fail.

	// Memory management variables.
	memoryAllocated uint64

	// The download object, mostly to update download progress.
	download *download
	mu       sync.Mutex

	// Caching related fields
	chunkCache map[string]*cacheData
	cacheMu    *sync.Mutex
}

// fail will set the chunk status to failed. The physical chunk memory will be
// wiped and any memory allocation will be returned to the renter. The download
// as a whole will be failed as well.
func (udc *unfinishedDownloadChunk) fail(err error) {
	udc.failed = true
	udc.recoveryComplete = true
	for i := range udc.physicalChunkData {
		udc.physicalChunkData[i] = nil
	}
	udc.download.managedFail(fmt.Errorf("chunk %v failed: %v", udc.staticChunkIndex, err))
	udc.destination = nil
}

// managedCleanUp will check if the download has failed, and if not it will add
// any standby workers which need to be added. Calling managedCleanUp too many
// times is not harmful, however missing a call to managedCleanUp can lead to
// dealocks.
func (udc *unfinishedDownloadChunk) managedCleanUp() {
	// Check if the chunk is newly failed.
	udc.mu.Lock()
	if udc.workersRemaining+udc.piecesCompleted < udc.erasureCode.MinPieces() && !udc.failed {
		udc.fail(errors.New("not enough workers to continue download"))
	}
	// Return any excess memory.
	udc.returnMemory()

	// Nothing to do if the chunk has failed.
	if udc.failed {
		udc.mu.Unlock()
		return
	}

	// Check whether standby workers are required.
	chunkComplete := udc.piecesCompleted >= udc.erasureCode.MinPieces()
	desiredPiecesRegistered := udc.erasureCode.MinPieces() + udc.staticOverdrive - udc.piecesCompleted
	standbyWorkersRequired := !chunkComplete && udc.piecesRegistered < desiredPiecesRegistered
	if !standbyWorkersRequired {
		udc.mu.Unlock()
		return
	}

	// Assemble a list of standby workers, release the udc lock, and then queue
	// the chunk into the workers. The lock needs to be released early because
	// holding the udc lock and the worker lock at the same time is a deadlock
	// risk (they interact with eachother, call functions on eachother).
	var standbyWorkers []*worker
	for i := 0; i < len(udc.workersStandby); i++ {
		standbyWorkers = append(standbyWorkers, udc.workersStandby[i])
	}
	udc.workersStandby = udc.workersStandby[:0] // Workers have been taken off of standby.
	udc.mu.Unlock()
	for i := 0; i < len(standbyWorkers); i++ {
		standbyWorkers[i].managedQueueDownloadChunk(udc)
	}
}

// managedRemoveWorker will decrement a worker from the set of remaining workers
// in the udc. After a worker has been removed, the udc needs to be cleaned up.
func (udc *unfinishedDownloadChunk) managedRemoveWorker() {
	udc.mu.Lock()
	udc.workersRemaining--
	udc.mu.Unlock()
	udc.managedCleanUp()
}

// returnMemory will check on the status of all the workers and pieces, and
// determine how much memory is safe to return to the renter. This should be
// called each time a worker returns, and also after the chunk is recovered.
func (udc *unfinishedDownloadChunk) returnMemory() {
	// The maximum amount of memory is the pieces completed plus the number of
	// workers remaining.
	maxMemory := uint64(udc.workersRemaining+udc.piecesCompleted) * udc.staticPieceSize
	// If enough pieces have completed, max memory is the number of registered
	// pieces plus the number of completed pieces.
	if udc.piecesCompleted >= udc.erasureCode.MinPieces() {
		// udc.piecesRegistered is guaranteed to be at most equal to the number
		// of overdrive pieces, meaning it will be equal to or less than
		// initalMemory.
		maxMemory = uint64(udc.piecesCompleted+udc.piecesRegistered) * udc.staticPieceSize
	}
	// If the chunk recovery has completed, the maximum number of pieces is the
	// number of registered.
	if udc.recoveryComplete {
		maxMemory = uint64(udc.piecesRegistered) * udc.staticPieceSize
	}
	// Return any memory we don't need.
	if uint64(udc.memoryAllocated) > maxMemory {
		udc.download.memoryManager.Return(udc.memoryAllocated - maxMemory)
		udc.memoryAllocated = maxMemory
	}
}

// threadedRecoverLogicalData will take all of the pieces that have been
// downloaded and encode them into the logical data which is then written to the
// underlying writer for the download.
func (udc *unfinishedDownloadChunk) threadedRecoverLogicalData() error {
	// Ensure cleanup occurs after the data is recovered, whether recovery
	// succeeds or fails.
	defer udc.managedCleanUp()

	// Decrypt the chunk pieces. This doesn't need to happen under a lock,
	// because any thread potentially writing to the physicalChunkData array is
	// going to be stopped by the fact that the chunk is complete.
	for i := range udc.physicalChunkData {
		// Skip empty pieces.
		if udc.physicalChunkData[i] == nil {
			continue
		}

		key := deriveKey(udc.masterKey, udc.staticChunkIndex, uint64(i))
		decryptedPiece, err := key.DecryptBytes(udc.physicalChunkData[i])
		if err != nil {
			udc.mu.Lock()
			udc.fail(err)
			udc.mu.Unlock()
			return errors.AddContext(err, "unable to decrypt chunk")
		}
		udc.physicalChunkData[i] = decryptedPiece
	}

	// Recover the pieces into the logical chunk data.
	//
	// TODO: Might be some way to recover into the downloadDestination instead
	// of creating a buffer and then writing that.
	recoverWriter := new(bytes.Buffer)
	err := udc.erasureCode.Recover(udc.physicalChunkData, udc.staticChunkSize, recoverWriter)
	if err != nil {
		udc.mu.Lock()
		udc.fail(err)
		udc.mu.Unlock()
		return errors.AddContext(err, "unable to recover chunk")
	}
	// Clear out the physical chunk pieces, we do not need them anymore.
	for i := range udc.physicalChunkData {
		udc.physicalChunkData[i] = nil
	}

	// Get recovered data
	recoveredData := recoverWriter.Bytes()

	// Add the chunk to the cache.
	udc.addChunkToCache(recoveredData)

	// Write the bytes to the requested output.
	start := udc.staticFetchOffset
	end := udc.staticFetchOffset + udc.staticFetchLength
	_, err = udc.destination.WriteAt(recoveredData[start:end], udc.staticWriteOffset)
	if err != nil {
		udc.mu.Lock()
		udc.fail(err)
		udc.mu.Unlock()
		return errors.AddContext(err, "unable to write to download destination")
	}
	recoverWriter = nil

	// Now that the download has completed and been flushed from memory, we can
	// release the memory that was used to store the data. Call 'cleanUp' to
	// trigger the memory cleanup along with some extra checks that everything
	// is consistent.
	udc.mu.Lock()
	udc.recoveryComplete = true
	udc.mu.Unlock()

	// Update the download and signal completion of this chunk.
	udc.download.mu.Lock()
	defer udc.download.mu.Unlock()
	udc.download.chunksRemaining--
	atomic.AddUint64(&udc.download.atomicDataReceived, udc.staticFetchLength)
	if udc.download.chunksRemaining == 0 {
		// Download is complete, send out a notification and close the
		// destination writer.
		udc.download.endTime = time.Now()
		close(udc.download.completeChan)
		err := udc.download.destination.Close()
		udc.download.destination = nil
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// +build kube

/*
   Copyright 2020 Docker Compose CLI authors

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

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

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

package context

import (
	"github.com/docker/compose-cli/pkg/api"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"

	"github.com/docker/compose-cli/api/context/store"
	"github.com/docker/compose-cli/kube"
)

func init() {
	extraCommands = append(extraCommands, createKubeCommand)
	extraHelp = append(extraHelp, `
Create a Kubernetes context:
$ docker context create k8s CONTEXT [flags]
(see docker context create k8s --help)
`)
}

func createKubeCommand() *cobra.Command {
	var opts kube.ContextParams
	cmd := &cobra.Command{
		Use:   "kubernetes CONTEXT [flags]",
		Short: "Create context for a Kubernetes Cluster",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCreateKube(args[0], opts)
		},
	}

	addDescriptionFlag(cmd, &opts.Description)
	cmd.Flags().StringVar(&opts.KubeConfigPath, "kubeconfig", "", "The endpoint of the Kubernetes manager")
	cmd.Flags().StringVar(&opts.KubeContextName, "kubecontext", "", "The name of the context to use in kubeconfig")
	cmd.Flags().BoolVar(&opts.FromEnvironment, "from-env", false, "Get endpoint and creds from env vars")
	return cmd
}

func runCreateKube(contextName string, opts kube.ContextParams) error {
	if contextExists(contextName) {
		return errors.Wrapf(api.ErrAlreadyExists, "context %q", contextName)
	}

	contextData, description, err := opts.CreateContextData()
	if err != nil {
		return err
	}
	return createDockerContext(contextName, store.KubeContextType, description, contextData)
}
<commit_msg>Fix help text when for context create and kube backend<commit_after>// +build kube

/*
   Copyright 2020 Docker Compose CLI authors

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

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

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

package context

import (
	"github.com/docker/compose-cli/pkg/api"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"

	"github.com/docker/compose-cli/api/context/store"
	"github.com/docker/compose-cli/kube"
)

func init() {
	extraCommands = append(extraCommands, createKubeCommand)
	extraHelp = append(extraHelp, `
Create a Kubernetes context:
$ docker context create kubernetes CONTEXT [flags]
(see docker context create kubernetes --help)
`)
}

func createKubeCommand() *cobra.Command {
	var opts kube.ContextParams
	cmd := &cobra.Command{
		Use:   "kubernetes CONTEXT [flags]",
		Short: "Create context for a Kubernetes Cluster",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCreateKube(args[0], opts)
		},
	}

	addDescriptionFlag(cmd, &opts.Description)
	cmd.Flags().StringVar(&opts.KubeConfigPath, "kubeconfig", "", "The endpoint of the Kubernetes manager")
	cmd.Flags().StringVar(&opts.KubeContextName, "kubecontext", "", "The name of the context to use in kubeconfig")
	cmd.Flags().BoolVar(&opts.FromEnvironment, "from-env", false, "Get endpoint and creds from env vars")
	return cmd
}

func runCreateKube(contextName string, opts kube.ContextParams) error {
	if contextExists(contextName) {
		return errors.Wrapf(api.ErrAlreadyExists, "context %q", contextName)
	}

	contextData, description, err := opts.CreateContextData()
	if err != nil {
		return err
	}
	return createDockerContext(contextName, store.KubeContextType, description, contextData)
}
<|endoftext|>


<commit_before>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

package vespa

import (
	"bytes"
	"encoding/json"
	"os"
)

type VespaTlsConfig struct {
	DisableHostnameValidation bool `json:"disable-hostname-validation"`
	Files                     struct {
		PrivateKey     string `json:"private-key"`
		CaCertificates string `json:"ca-certificates"`
		Certificates   string `json:"certificates"`
	} `json:"files"`
	AuthorizedPeers []struct {
		RequiredCredentials []struct {
			Field     string `json:"field"`
			MustMatch string `json:"must-match"`
		} `json:"required-credentials"`
		Name         string   `json:"name"`
		Capabilities []string `json:"capabilities"`
	} `json:"authorized-peers"`
}

func LoadTlsConfig() (*VespaTlsConfig, error) {
	fn := os.Getenv("VESPA_TLS_CONFIG_FILE")
	if fn == "" {
		return nil, nil
	}
	contents, err := os.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	codec := json.NewDecoder(bytes.NewReader(contents))
	var parsedJson VespaTlsConfig
	err = codec.Decode(&parsedJson)
	if err != nil {
		return nil, err
	}
	return &parsedJson, nil
}

func ExportSecurityEnvToSh() {
	LoadDefaultEnv()
	cfg, _ := LoadTlsConfig()
	helper := newShellEnvExporter()
	if cfg == nil {
		helper.unsetVar("VESPA_TLS_ENABLED")
	} else {
		if fn := cfg.Files.PrivateKey; fn != "" {
			helper.overrideVar("VESPA_TLS_PRIVATE_KEY", fn)
		}
		if fn := cfg.Files.CaCertificates; fn != "" {
			helper.overrideVar("VESPA_TLS_CA_CERT", fn)
		}
		if fn := cfg.Files.Certificates; fn != "" {
			helper.overrideVar("VESPA_TLS_CERT", fn)
		}
		if cfg.DisableHostnameValidation {
			helper.overrideVar("VESPA_TLS_HOSTNAME_VALIDATION_DISABLED", "1")
		} else {
			helper.unsetVar("VESPA_TLS_HOSTNAME_VALIDATION_DISABLED")
		}
		if os.Getenv("VESPA_TLS_INSECURE_MIXED_MODE") == "" {
			helper.overrideVar("VESPA_TLS_ENABLED", "1")
		}
	}
	helper.dump()
}
<commit_msg>drop currently-unused fields<commit_after>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

package vespa

import (
	"bytes"
	"encoding/json"
	"os"
)

type VespaTlsConfig struct {
	DisableHostnameValidation bool `json:"disable-hostname-validation"`
	Files                     struct {
		PrivateKey     string `json:"private-key"`
		CaCertificates string `json:"ca-certificates"`
		Certificates   string `json:"certificates"`
	} `json:"files"`
}

func LoadTlsConfig() (*VespaTlsConfig, error) {
	fn := os.Getenv("VESPA_TLS_CONFIG_FILE")
	if fn == "" {
		return nil, nil
	}
	contents, err := os.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	codec := json.NewDecoder(bytes.NewReader(contents))
	var parsedJson VespaTlsConfig
	err = codec.Decode(&parsedJson)
	if err != nil {
		return nil, err
	}
	return &parsedJson, nil
}

func ExportSecurityEnvToSh() {
	LoadDefaultEnv()
	cfg, _ := LoadTlsConfig()
	helper := newShellEnvExporter()
	if cfg == nil {
		helper.unsetVar("VESPA_TLS_ENABLED")
	} else {
		if fn := cfg.Files.PrivateKey; fn != "" {
			helper.overrideVar("VESPA_TLS_PRIVATE_KEY", fn)
		}
		if fn := cfg.Files.CaCertificates; fn != "" {
			helper.overrideVar("VESPA_TLS_CA_CERT", fn)
		}
		if fn := cfg.Files.Certificates; fn != "" {
			helper.overrideVar("VESPA_TLS_CERT", fn)
		}
		if cfg.DisableHostnameValidation {
			helper.overrideVar("VESPA_TLS_HOSTNAME_VALIDATION_DISABLED", "1")
		} else {
			helper.unsetVar("VESPA_TLS_HOSTNAME_VALIDATION_DISABLED")
		}
		if os.Getenv("VESPA_TLS_INSECURE_MIXED_MODE") == "" {
			helper.overrideVar("VESPA_TLS_ENABLED", "1")
		}
	}
	helper.dump()
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Jean-Philippe Couture
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

func main() {
	if len(os.Args) < 3 {
		printUsage()
		return
	}
	fn := os.Args[1]
	cmd := os.Args[2]
	args := os.Args[2:]

	setEnvVars(loadEnvFile(fn))

	binary, lookErr := exec.LookPath(cmd)
	if lookErr != nil {
		binary = cmd
	}

	execErr := syscall.Exec(binary, args, os.Environ())
	if execErr != nil {
		fmt.Println("[Err] Cannot execute:", binary)
	}
}

func printUsage() {
	usage := `nv - context specific environment variables
Usage: nv <env file> <command> [arguments...]`
	fmt.Println(usage)
}

func loadEnvFile(fn string) []string {
	dat, err := ioutil.ReadFile(fn)
	if err != nil {
		panic(err)
	}
	vars := make([]string, 0)
	lines := bytes.Split(dat, []byte("\n"))
	for _, line := range lines {
		if len(line) > 0 {
			vars = append(vars, string(line))
		}
	}
	return vars
}

func setEnvVars(envVars []string) {
	for _, envVar := range envVars {
		result := strings.Split(envVar, "=")
		if len(result) >= 2 {
			name := result[0]
			value := strings.Join(result[1:], "=")
			os.Setenv(name, value)
		}
	}
}
<commit_msg>Check if the .env file exists and is readable.<commit_after>// Copyright (c) 2015 Jean-Philippe Couture
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

func main() {
	if len(os.Args) < 3 {
		printUsage()
		os.Exit(-1)
	}

	fn := os.Args[1]
	cmd := os.Args[2]
	args := os.Args[2:]

	if !fileExists(fn) {
		fmt.Printf("[Err] '%s': No such file or directory\n", fn)
		os.Exit(-1)
	}

	if !fileReadable(fn) {
		fmt.Printf("[Err] '%s': Permission denied\n", fn)
		os.Exit(-1)
	}

	setEnvVars(loadEnvFile(fn))

	binary, lookErr := exec.LookPath(cmd)
	if lookErr != nil {
		binary = cmd
	}

	execErr := syscall.Exec(binary, args, os.Environ())
	if execErr != nil {
		fmt.Println("[Err] Cannot execute:", binary)
		os.Exit(-1)
	}
}

func printUsage() {
	usage := `nv - context specific environment variables
Usage: nv <env file> <command> [arguments...]`
	fmt.Println(usage)
}

func loadEnvFile(fn string) []string {
	dat, err := ioutil.ReadFile(fn)
	if err != nil {
		panic(err)
	}
	vars := make([]string, 0)
	lines := bytes.Split(dat, []byte("\n"))
	for _, line := range lines {
		if len(line) > 0 {
			vars = append(vars, string(line))
		}
	}
	return vars
}

func setEnvVars(envVars []string) {
	for _, envVar := range envVars {
		result := strings.Split(envVar, "=")
		if len(result) >= 2 {
			name := result[0]
			value := strings.Join(result[1:], "=")
			os.Setenv(name, value)
		}
	}
}

func fileExists(filename string) bool {
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		return false
	}
	return true
}

func fileReadable(filename string) bool {
	if _, err := os.Stat(filename); err == nil {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package dockerclient

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"log"
)

type ProgressDetail struct {
	Current int64
	Total   int64
	Start   int64
}

type ErrorDetail struct {
	Code    int
	Message string
}

type BuildResponse []*JSONMessage

func (rsp BuildResponse) Last() *JSONMessage {
	if len(rsp) == 0 {
		return nil
	}
	return rsp[len(rsp)-1]
}

func (rsp BuildResponse) ImageId() string {
	last := rsp.Last()
	if last == nil {
		return ""
	}
	m := imageIdRegexp.FindStringSubmatch(last.Stream)
	if len(m) == 2 {
		return m[1]
	}
	return ""
}

func (dh *DockerHost) handleBuildImageJson(r io.Reader, f func(s *JSONMessage)) (rsp BuildResponse, e error) {
	scanner := bufio.NewReader(r)
	buf := &bytes.Buffer{}
	for {
		b, e := scanner.ReadBytes('}')
		if e == io.EOF {
			break
		} else if e != nil {
			return nil, e
		}
		buf.Write(b)
		stream := &JSONMessage{}
		e = json.Unmarshal(buf.Bytes(), stream)
		if e != nil {
			if e.Error() == "unexpected end of JSON input" {
				continue
			}
			log.Printf("ERROR: %s => %s", e.Error(), buf.String())
			return nil, e
		}
		if f != nil {
			f(stream)
		}
		rsp = append(rsp, stream)
		buf.Reset()
	}
	return rsp, nil
}

func (dh *DockerHost) handleBuildImagePlain(r io.Reader, f func(s *JSONMessage)) (rsp BuildResponse, e error) {
	reader := bufio.NewReader(r)
	for {
		b, e := reader.ReadString('\n')
		if e == io.EOF {
			break
		} else if e != nil {
			return nil, e
		}
		s := &JSONMessage{Stream: string(b)}
		rsp = append(rsp, s)
		f(s)
	}
	return rsp, nil
}
<commit_msg>add ScanJson method<commit_after>package dockerclient

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"log"
)

type ProgressDetail struct {
	Current int64
	Total   int64
	Start   int64
}

type ErrorDetail struct {
	Code    int
	Message string
}

type BuildResponse []*JSONMessage

func (rsp BuildResponse) Last() *JSONMessage {
	if len(rsp) == 0 {
		return nil
	}
	return rsp[len(rsp)-1]
}

func (rsp BuildResponse) ImageId() string {
	last := rsp.Last()
	if last == nil {
		return ""
	}
	m := imageIdRegexp.FindStringSubmatch(last.Stream)
	if len(m) == 2 {
		return m[1]
	}
	return ""
}

func ScanJson(r io.Reader, f func(b []byte) error) error {
	scanner := bufio.NewReader(r)
	buf := &bytes.Buffer{}
	for {
		b, e := scanner.ReadBytes('}')
		if e == io.EOF {
			break
		} else if e != nil {
			return e
		}
		e = f(b)
		if e != nil {
			return e
		}
		buf.Reset()
	}
	return nil
}

func (dh *DockerHost) handleBuildImageJson(r io.Reader, f func(s *JSONMessage)) (rsp BuildResponse, e error) {
	scanner := bufio.NewReader(r)
	buf := &bytes.Buffer{}
	for {
		b, e := scanner.ReadBytes('}')
		if e == io.EOF {
			break
		} else if e != nil {
			return nil, e
		}
		buf.Write(b)
		stream := &JSONMessage{}
		e = json.Unmarshal(buf.Bytes(), stream)
		if e != nil {
			if e.Error() == "unexpected end of JSON input" {
				continue
			}
			log.Printf("ERROR: %s => %s", e.Error(), buf.String())
			return nil, e
		}
		if f != nil {
			f(stream)
		}
		rsp = append(rsp, stream)
		buf.Reset()
	}
	return rsp, nil
}

func (dh *DockerHost) handleBuildImagePlain(r io.Reader, f func(s *JSONMessage)) (rsp BuildResponse, e error) {
	reader := bufio.NewReader(r)
	for {
		b, e := reader.ReadString('\n')
		if e == io.EOF {
			break
		} else if e != nil {
			return nil, e
		}
		s := &JSONMessage{Stream: string(b)}
		rsp = append(rsp, s)
		f(s)
	}
	return rsp, nil
}
<|endoftext|>


<commit_before>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	discovery "github.com/libp2p/go-libp2p-discovery"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

var DesiredRelays = 3

// AutoRelayHost is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelayHost struct {
	*basic.BasicHost
	discover discovery.Discoverer
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]pstore.PeerInfo
	addrs  []ma.Multiaddr
}

func NewAutoRelayHost(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer) *AutoRelayHost {
	autonat := autonat.NewAutoNAT(ctx, bhost)
	h := &AutoRelayHost{
		BasicHost:  bhost,
		discover:   discover,
		autonat:    autonat,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]pstore.PeerInfo),
		disconnect: make(chan struct{}, 1),
	}
	bhost.AddrsFactory = h.hostAddrs
	bhost.Network().Notify(h)
	go h.background(ctx)
	return h
}

func (h *AutoRelayHost) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	h.mx.Lock()
	defer h.mx.Unlock()
	if h.addrs != nil && h.autonat.Status() == autonat.NATStatusPrivate {
		return h.addrs
	} else {
		return h.addrsF(addrs)
	}
}

func (h *AutoRelayHost) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + 30*time.Second):
	case <-ctx.Done():
		return
	}

	for {
		wait := autonat.AutoNATRefreshInterval
		switch h.autonat.Status() {
		case autonat.NATStatusUnknown:
			wait = autonat.AutoNATRetryInterval
		case autonat.NATStatusPublic:
		case autonat.NATStatusPrivate:
			h.findRelays(ctx)
		}

		select {
		case <-h.disconnect:
			// invalidate addrs
			h.mx.Lock()
			h.addrs = nil
			h.mx.Unlock()
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (h *AutoRelayHost) findRelays(ctx context.Context) {
	h.mx.Lock()
	if len(h.relays) >= DesiredRelays {
		h.mx.Unlock()
		return
	}
	need := DesiredRelays - len(h.relays)
	h.mx.Unlock()

	limit := 20
	for ; need > limit; limit *= 2 {
	}

	dctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	pis, err := discovery.FindPeers(dctx, h.discover, "/libp2p/relay", limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())
		return
	}

	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric
	shuffleRelays(pis)

	update := 0

	for _, pi := range pis {
		h.mx.Lock()
		if _, ok := h.relays[pi.ID]; ok {
			h.mx.Unlock()
			continue
		}
		h.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 60*time.Second)
		err = h.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		h.mx.Lock()
		h.relays[pi.ID] = pi
		h.mx.Unlock()

		update++
		need--
		if need == 0 {
			break
		}
	}

	if update > 0 || h.addrs == nil {
		h.updateAddrs()
	}
}

func (h *AutoRelayHost) updateAddrs() {
	h.doUpdateAddrs()
	h.PushIdentify()
}

func (h *AutoRelayHost) doUpdateAddrs() {
	h.mx.Lock()
	defer h.mx.Unlock()

	addrs := h.addrsF(h.AllAddrs())
	raddrs := make([]ma.Multiaddr, 0, len(addrs)+len(h.relays))

	// remove our public addresses from the list and replace them by just the public IP
	for _, addr := range addrs {
		if manet.IsPublicAddr(addr) {
			ip, err := addr.ValueForProtocol(ma.P_IP4)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip4/%s", ip))
				if err != nil {
					panic(err)
				}

				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}

			ip, err = addr.ValueForProtocol(ma.P_IP6)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip6/%s", ip))
				if err != nil {
					panic(err)
				}
				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}
		} else {
			raddrs = append(raddrs, addr)
		}
	}

	circuit, err := ma.NewMultiaddr("/p2p-circuit")
	if err != nil {
		panic(err)
	}

	for _, pi := range h.relays {
		for _, addr := range pi.Addrs {
			if !manet.IsPrivateAddr(addr) {
				pub := addr.Encapsulate(circuit)
				raddrs = append(raddrs, pub)
			}
		}
	}

	h.addrs = raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

func containsAddr(lst []ma.Multiaddr, addr ma.Multiaddr) bool {
	for _, xaddr := range lst {
		if xaddr.Equal(addr) {
			return true
		}
	}
	return false
}

// notify
func (h *AutoRelayHost) Listen(inet.Network, ma.Multiaddr)      {}
func (h *AutoRelayHost) ListenClose(inet.Network, ma.Multiaddr) {}
func (h *AutoRelayHost) Connected(inet.Network, inet.Conn)      {}

func (h *AutoRelayHost) Disconnected(_ inet.Network, c inet.Conn) {
	p := c.RemotePeer()
	h.mx.Lock()
	defer h.mx.Unlock()
	if _, ok := h.relays[p]; ok {
		delete(h.relays, p)
		select {
		case h.disconnect <- struct{}{}:
		default:
		}
	}
}

func (h *AutoRelayHost) OpenedStream(inet.Network, inet.Stream) {}
func (h *AutoRelayHost) ClosedStream(inet.Network, inet.Stream) {}

var _ host.Host = (*AutoRelayHost)(nil)
<commit_msg>configurable boot delay for autorelay<commit_after>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	discovery "github.com/libp2p/go-libp2p-discovery"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

var (
	DesiredRelays = 3

	BootDelay = 90 * time.Second
)

// AutoRelayHost is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelayHost struct {
	*basic.BasicHost
	discover discovery.Discoverer
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]pstore.PeerInfo
	addrs  []ma.Multiaddr
}

func NewAutoRelayHost(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer) *AutoRelayHost {
	autonat := autonat.NewAutoNAT(ctx, bhost)
	h := &AutoRelayHost{
		BasicHost:  bhost,
		discover:   discover,
		autonat:    autonat,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]pstore.PeerInfo),
		disconnect: make(chan struct{}, 1),
	}
	bhost.AddrsFactory = h.hostAddrs
	bhost.Network().Notify(h)
	go h.background(ctx)
	return h
}

func (h *AutoRelayHost) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	h.mx.Lock()
	defer h.mx.Unlock()
	if h.addrs != nil && h.autonat.Status() == autonat.NATStatusPrivate {
		return h.addrs
	} else {
		return h.addrsF(addrs)
	}
}

func (h *AutoRelayHost) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	for {
		wait := autonat.AutoNATRefreshInterval
		switch h.autonat.Status() {
		case autonat.NATStatusUnknown:
			wait = autonat.AutoNATRetryInterval
		case autonat.NATStatusPublic:
		case autonat.NATStatusPrivate:
			h.findRelays(ctx)
		}

		select {
		case <-h.disconnect:
			// invalidate addrs
			h.mx.Lock()
			h.addrs = nil
			h.mx.Unlock()
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (h *AutoRelayHost) findRelays(ctx context.Context) {
	h.mx.Lock()
	if len(h.relays) >= DesiredRelays {
		h.mx.Unlock()
		return
	}
	need := DesiredRelays - len(h.relays)
	h.mx.Unlock()

	limit := 20
	for ; need > limit; limit *= 2 {
	}

	dctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	pis, err := discovery.FindPeers(dctx, h.discover, "/libp2p/relay", limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())
		return
	}

	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric
	shuffleRelays(pis)

	update := 0

	for _, pi := range pis {
		h.mx.Lock()
		if _, ok := h.relays[pi.ID]; ok {
			h.mx.Unlock()
			continue
		}
		h.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 60*time.Second)
		err = h.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		h.mx.Lock()
		h.relays[pi.ID] = pi
		h.mx.Unlock()

		update++
		need--
		if need == 0 {
			break
		}
	}

	if update > 0 || h.addrs == nil {
		h.updateAddrs()
	}
}

func (h *AutoRelayHost) updateAddrs() {
	h.doUpdateAddrs()
	h.PushIdentify()
}

func (h *AutoRelayHost) doUpdateAddrs() {
	h.mx.Lock()
	defer h.mx.Unlock()

	addrs := h.addrsF(h.AllAddrs())
	raddrs := make([]ma.Multiaddr, 0, len(addrs)+len(h.relays))

	// remove our public addresses from the list and replace them by just the public IP
	for _, addr := range addrs {
		if manet.IsPublicAddr(addr) {
			ip, err := addr.ValueForProtocol(ma.P_IP4)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip4/%s", ip))
				if err != nil {
					panic(err)
				}

				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}

			ip, err = addr.ValueForProtocol(ma.P_IP6)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip6/%s", ip))
				if err != nil {
					panic(err)
				}
				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}
		} else {
			raddrs = append(raddrs, addr)
		}
	}

	circuit, err := ma.NewMultiaddr("/p2p-circuit")
	if err != nil {
		panic(err)
	}

	for _, pi := range h.relays {
		for _, addr := range pi.Addrs {
			if !manet.IsPrivateAddr(addr) {
				pub := addr.Encapsulate(circuit)
				raddrs = append(raddrs, pub)
			}
		}
	}

	h.addrs = raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

func containsAddr(lst []ma.Multiaddr, addr ma.Multiaddr) bool {
	for _, xaddr := range lst {
		if xaddr.Equal(addr) {
			return true
		}
	}
	return false
}

// notify
func (h *AutoRelayHost) Listen(inet.Network, ma.Multiaddr)      {}
func (h *AutoRelayHost) ListenClose(inet.Network, ma.Multiaddr) {}
func (h *AutoRelayHost) Connected(inet.Network, inet.Conn)      {}

func (h *AutoRelayHost) Disconnected(_ inet.Network, c inet.Conn) {
	p := c.RemotePeer()
	h.mx.Lock()
	defer h.mx.Unlock()
	if _, ok := h.relays[p]; ok {
		delete(h.relays, p)
		select {
		case h.disconnect <- struct{}{}:
		default:
		}
	}
}

func (h *AutoRelayHost) OpenedStream(inet.Network, inet.Stream) {}
func (h *AutoRelayHost) ClosedStream(inet.Network, inet.Stream) {}

var _ host.Host = (*AutoRelayHost)(nil)
<|endoftext|>


<commit_before>package swarm

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"testing"
	"time"

	metrics "github.com/ipfs/go-libp2p/p2p/metrics"
	inet "github.com/ipfs/go-libp2p/p2p/net"
	peer "github.com/ipfs/go-libp2p/p2p/peer"
	testutil "github.com/ipfs/go-libp2p/testutil"

	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
	ma "gx/ipfs/QmcobAGsCjYt5DXoq9et9L8yR8er7o7Cu3DTvpaq12jYSz/go-multiaddr"
)

func EchoStreamHandler(stream inet.Stream) {
	go func() {
		defer stream.Close()

		// pull out the ipfs conn
		c := stream.Conn()
		log.Infof("%s ponging to %s", c.LocalPeer(), c.RemotePeer())

		buf := make([]byte, 4)

		for {
			if _, err := stream.Read(buf); err != nil {
				if err != io.EOF {
					log.Error("ping receive error:", err)
				}
				return
			}

			if !bytes.Equal(buf, []byte("ping")) {
				log.Errorf("ping receive error: ping != %s %v", buf, buf)
				return
			}

			if _, err := stream.Write([]byte("pong")); err != nil {
				log.Error("pond send error:", err)
				return
			}
		}
	}()
}

func makeDialOnlySwarm(ctx context.Context, t *testing.T) *Swarm {
	id := testutil.RandIdentityOrFatal(t)

	peerstore := peer.NewPeerstore()
	peerstore.AddPubKey(id.ID(), id.PublicKey())
	peerstore.AddPrivKey(id.ID(), id.PrivateKey())

	swarm, err := NewSwarm(ctx, nil, id.ID(), peerstore, metrics.NewBandwidthCounter())
	if err != nil {
		t.Fatal(err)
	}

	swarm.SetStreamHandler(EchoStreamHandler)

	return swarm
}

func makeSwarms(ctx context.Context, t *testing.T, num int) []*Swarm {
	swarms := make([]*Swarm, 0, num)

	for i := 0; i < num; i++ {
		localnp := testutil.RandPeerNetParamsOrFatal(t)

		peerstore := peer.NewPeerstore()
		peerstore.AddPubKey(localnp.ID, localnp.PubKey)
		peerstore.AddPrivKey(localnp.ID, localnp.PrivKey)

		addrs := []ma.Multiaddr{localnp.Addr}
		swarm, err := NewSwarm(ctx, addrs, localnp.ID, peerstore, metrics.NewBandwidthCounter())
		if err != nil {
			t.Fatal(err)
		}

		swarm.SetStreamHandler(EchoStreamHandler)
		swarms = append(swarms, swarm)
	}

	return swarms
}

func connectSwarms(t *testing.T, ctx context.Context, swarms []*Swarm) {

	var wg sync.WaitGroup
	connect := func(s *Swarm, dst peer.ID, addr ma.Multiaddr) {
		// TODO: make a DialAddr func.
		s.peers.AddAddr(dst, addr, peer.PermanentAddrTTL)
		if _, err := s.Dial(ctx, dst); err != nil {
			t.Fatal("error swarm dialing to peer", err)
		}
		wg.Done()
	}

	log.Info("Connecting swarms simultaneously.")
	for _, s1 := range swarms {
		for _, s2 := range swarms {
			if s2.local != s1.local { // don't connect to self.
				wg.Add(1)
				connect(s1, s2.LocalPeer(), s2.ListenAddresses()[0]) // try the first.
			}
		}
	}
	wg.Wait()

	for _, s := range swarms {
		log.Infof("%s swarm routing table: %s", s.local, s.Peers())
	}
}

func SubtestSwarm(t *testing.T, SwarmNum int, MsgNum int) {
	// t.Skip("skipping for another test")

	ctx := context.Background()
	swarms := makeSwarms(ctx, t, SwarmNum)

	// connect everyone
	connectSwarms(t, ctx, swarms)

	// ping/pong
	for _, s1 := range swarms {
		log.Debugf("-------------------------------------------------------")
		log.Debugf("%s ping pong round", s1.local)
		log.Debugf("-------------------------------------------------------")

		_, cancel := context.WithCancel(ctx)
		got := map[peer.ID]int{}
		errChan := make(chan error, MsgNum*len(swarms))
		streamChan := make(chan *Stream, MsgNum)

		// send out "ping" x MsgNum to every peer
		go func() {
			defer close(streamChan)

			var wg sync.WaitGroup
			send := func(p peer.ID) {
				defer wg.Done()

				// first, one stream per peer (nice)
				stream, err := s1.NewStreamWithPeer(ctx, p)
				if err != nil {
					errChan <- err
					return
				}

				// send out ping!
				for k := 0; k < MsgNum; k++ { // with k messages
					msg := "ping"
					log.Debugf("%s %s %s (%d)", s1.local, msg, p, k)
					if _, err := stream.Write([]byte(msg)); err != nil {
						errChan <- err
						continue
					}
				}

				// read it later
				streamChan <- stream
			}

			for _, s2 := range swarms {
				if s2.local == s1.local {
					continue // dont send to self...
				}

				wg.Add(1)
				go send(s2.local)
			}
			wg.Wait()
		}()

		// receive "pong" x MsgNum from every peer
		go func() {
			defer close(errChan)
			count := 0
			countShouldBe := MsgNum * (len(swarms) - 1)
			for stream := range streamChan { // one per peer
				defer stream.Close()

				// get peer on the other side
				p := stream.Conn().RemotePeer()

				// receive pings
				msgCount := 0
				msg := make([]byte, 4)
				for k := 0; k < MsgNum; k++ { // with k messages

					// read from the stream
					if _, err := stream.Read(msg); err != nil {
						errChan <- err
						continue
					}

					if string(msg) != "pong" {
						errChan <- fmt.Errorf("unexpected message: %s", msg)
						continue
					}

					log.Debugf("%s %s %s (%d)", s1.local, msg, p, k)
					msgCount++
				}

				got[p] = msgCount
				count += msgCount
			}

			if count != countShouldBe {
				errChan <- fmt.Errorf("count mismatch: %d != %d", count, countShouldBe)
			}
		}()

		// check any errors (blocks till consumer is done)
		for err := range errChan {
			if err != nil {
				t.Error(err.Error())
			}
		}

		log.Debugf("%s got pongs", s1.local)
		if (len(swarms) - 1) != len(got) {
			t.Errorf("got (%d) less messages than sent (%d).", len(got), len(swarms))
		}

		for p, n := range got {
			if n != MsgNum {
				t.Error("peer did not get all msgs", p, n, "/", MsgNum)
			}
		}

		cancel()
		<-time.After(10 * time.Millisecond)
	}

	for _, s := range swarms {
		s.Close()
	}
}

func TestSwarm(t *testing.T) {
	// t.Skip("skipping for another test")
	t.Parallel()

	// msgs := 1000
	msgs := 100
	swarms := 5
	SubtestSwarm(t, swarms, msgs)
}

func TestBasicSwarm(t *testing.T) {
	// t.Skip("skipping for another test")
	t.Parallel()

	msgs := 1
	swarms := 2
	SubtestSwarm(t, swarms, msgs)
}

func TestConnHandler(t *testing.T) {
	// t.Skip("skipping for another test")
	t.Parallel()

	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 5)

	gotconn := make(chan struct{}, 10)
	swarms[0].SetConnHandler(func(conn *Conn) {
		gotconn <- struct{}{}
	})

	connectSwarms(t, ctx, swarms)

	<-time.After(time.Millisecond)
	// should've gotten 5 by now.

	swarms[0].SetConnHandler(nil)

	expect := 4
	for i := 0; i < expect; i++ {
		select {
		case <-time.After(time.Second):
			t.Fatal("failed to get connections")
		case <-gotconn:
		}
	}

	select {
	case <-gotconn:
		t.Fatalf("should have connected to %d swarms", expect)
	default:
	}
}

func TestAddrBlocking(t *testing.T) {
	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 2)

	swarms[0].SetConnHandler(func(conn *Conn) {
		t.Fatal("no connections should happen!")
	})

	_, block, err := net.ParseCIDR("127.0.0.1/8")
	if err != nil {
		t.Fatal(err)
	}

	swarms[1].Filters.AddDialFilter(block)

	swarms[1].peers.AddAddr(swarms[0].LocalPeer(), swarms[0].ListenAddresses()[0], peer.PermanentAddrTTL)
	_, err = swarms[1].Dial(ctx, swarms[0].LocalPeer())
	if err == nil {
		t.Fatal("dial should have failed")
	}

	swarms[0].peers.AddAddr(swarms[1].LocalPeer(), swarms[1].ListenAddresses()[0], peer.PermanentAddrTTL)
	_, err = swarms[0].Dial(ctx, swarms[1].LocalPeer())
	if err == nil {
		t.Fatal("dial should have failed")
	}
}

func TestFilterBounds(t *testing.T) {
	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 2)

	conns := make(chan struct{}, 8)
	swarms[0].SetConnHandler(func(conn *Conn) {
		conns <- struct{}{}
	})

	// Address that we wont be dialing from
	_, block, err := net.ParseCIDR("192.0.0.1/8")
	if err != nil {
		t.Fatal(err)
	}

	// set filter on both sides, shouldnt matter
	swarms[1].Filters.AddDialFilter(block)
	swarms[0].Filters.AddDialFilter(block)

	connectSwarms(t, ctx, swarms)

	select {
	case <-time.After(time.Second):
		t.Fatal("should have gotten connection")
	case <-conns:
		t.Log("got connect")
	}
}
<commit_msg>Fix address filtering for /ip6, add tests<commit_after>package swarm

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"testing"
	"time"

	metrics "github.com/ipfs/go-libp2p/p2p/metrics"
	inet "github.com/ipfs/go-libp2p/p2p/net"
	peer "github.com/ipfs/go-libp2p/p2p/peer"
	testutil "github.com/ipfs/go-libp2p/testutil"

	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
	ma "gx/ipfs/QmcobAGsCjYt5DXoq9et9L8yR8er7o7Cu3DTvpaq12jYSz/go-multiaddr"
)

func EchoStreamHandler(stream inet.Stream) {
	go func() {
		defer stream.Close()

		// pull out the ipfs conn
		c := stream.Conn()
		log.Infof("%s ponging to %s", c.LocalPeer(), c.RemotePeer())

		buf := make([]byte, 4)

		for {
			if _, err := stream.Read(buf); err != nil {
				if err != io.EOF {
					log.Error("ping receive error:", err)
				}
				return
			}

			if !bytes.Equal(buf, []byte("ping")) {
				log.Errorf("ping receive error: ping != %s %v", buf, buf)
				return
			}

			if _, err := stream.Write([]byte("pong")); err != nil {
				log.Error("pond send error:", err)
				return
			}
		}
	}()
}

func makeDialOnlySwarm(ctx context.Context, t *testing.T) *Swarm {
	id := testutil.RandIdentityOrFatal(t)

	peerstore := peer.NewPeerstore()
	peerstore.AddPubKey(id.ID(), id.PublicKey())
	peerstore.AddPrivKey(id.ID(), id.PrivateKey())

	swarm, err := NewSwarm(ctx, nil, id.ID(), peerstore, metrics.NewBandwidthCounter())
	if err != nil {
		t.Fatal(err)
	}

	swarm.SetStreamHandler(EchoStreamHandler)

	return swarm
}

func makeSwarms(ctx context.Context, t *testing.T, num int) []*Swarm {
	swarms := make([]*Swarm, 0, num)

	for i := 0; i < num; i++ {
		localnp := testutil.RandPeerNetParamsOrFatal(t)

		peerstore := peer.NewPeerstore()
		peerstore.AddPubKey(localnp.ID, localnp.PubKey)
		peerstore.AddPrivKey(localnp.ID, localnp.PrivKey)

		addrs := []ma.Multiaddr{localnp.Addr}
		swarm, err := NewSwarm(ctx, addrs, localnp.ID, peerstore, metrics.NewBandwidthCounter())
		if err != nil {
			t.Fatal(err)
		}

		swarm.SetStreamHandler(EchoStreamHandler)
		swarms = append(swarms, swarm)
	}

	return swarms
}

func connectSwarms(t *testing.T, ctx context.Context, swarms []*Swarm) {

	var wg sync.WaitGroup
	connect := func(s *Swarm, dst peer.ID, addr ma.Multiaddr) {
		// TODO: make a DialAddr func.
		s.peers.AddAddr(dst, addr, peer.PermanentAddrTTL)
		if _, err := s.Dial(ctx, dst); err != nil {
			t.Fatal("error swarm dialing to peer", err)
		}
		wg.Done()
	}

	log.Info("Connecting swarms simultaneously.")
	for _, s1 := range swarms {
		for _, s2 := range swarms {
			if s2.local != s1.local { // don't connect to self.
				wg.Add(1)
				connect(s1, s2.LocalPeer(), s2.ListenAddresses()[0]) // try the first.
			}
		}
	}
	wg.Wait()

	for _, s := range swarms {
		log.Infof("%s swarm routing table: %s", s.local, s.Peers())
	}
}

func SubtestSwarm(t *testing.T, SwarmNum int, MsgNum int) {
	// t.Skip("skipping for another test")

	ctx := context.Background()
	swarms := makeSwarms(ctx, t, SwarmNum)

	// connect everyone
	connectSwarms(t, ctx, swarms)

	// ping/pong
	for _, s1 := range swarms {
		log.Debugf("-------------------------------------------------------")
		log.Debugf("%s ping pong round", s1.local)
		log.Debugf("-------------------------------------------------------")

		_, cancel := context.WithCancel(ctx)
		got := map[peer.ID]int{}
		errChan := make(chan error, MsgNum*len(swarms))
		streamChan := make(chan *Stream, MsgNum)

		// send out "ping" x MsgNum to every peer
		go func() {
			defer close(streamChan)

			var wg sync.WaitGroup
			send := func(p peer.ID) {
				defer wg.Done()

				// first, one stream per peer (nice)
				stream, err := s1.NewStreamWithPeer(ctx, p)
				if err != nil {
					errChan <- err
					return
				}

				// send out ping!
				for k := 0; k < MsgNum; k++ { // with k messages
					msg := "ping"
					log.Debugf("%s %s %s (%d)", s1.local, msg, p, k)
					if _, err := stream.Write([]byte(msg)); err != nil {
						errChan <- err
						continue
					}
				}

				// read it later
				streamChan <- stream
			}

			for _, s2 := range swarms {
				if s2.local == s1.local {
					continue // dont send to self...
				}

				wg.Add(1)
				go send(s2.local)
			}
			wg.Wait()
		}()

		// receive "pong" x MsgNum from every peer
		go func() {
			defer close(errChan)
			count := 0
			countShouldBe := MsgNum * (len(swarms) - 1)
			for stream := range streamChan { // one per peer
				defer stream.Close()

				// get peer on the other side
				p := stream.Conn().RemotePeer()

				// receive pings
				msgCount := 0
				msg := make([]byte, 4)
				for k := 0; k < MsgNum; k++ { // with k messages

					// read from the stream
					if _, err := stream.Read(msg); err != nil {
						errChan <- err
						continue
					}

					if string(msg) != "pong" {
						errChan <- fmt.Errorf("unexpected message: %s", msg)
						continue
					}

					log.Debugf("%s %s %s (%d)", s1.local, msg, p, k)
					msgCount++
				}

				got[p] = msgCount
				count += msgCount
			}

			if count != countShouldBe {
				errChan <- fmt.Errorf("count mismatch: %d != %d", count, countShouldBe)
			}
		}()

		// check any errors (blocks till consumer is done)
		for err := range errChan {
			if err != nil {
				t.Error(err.Error())
			}
		}

		log.Debugf("%s got pongs", s1.local)
		if (len(swarms) - 1) != len(got) {
			t.Errorf("got (%d) less messages than sent (%d).", len(got), len(swarms))
		}

		for p, n := range got {
			if n != MsgNum {
				t.Error("peer did not get all msgs", p, n, "/", MsgNum)
			}
		}

		cancel()
		<-time.After(10 * time.Millisecond)
	}

	for _, s := range swarms {
		s.Close()
	}
}

func TestSwarm(t *testing.T) {
	// t.Skip("skipping for another test")
	t.Parallel()

	// msgs := 1000
	msgs := 100
	swarms := 5
	SubtestSwarm(t, swarms, msgs)
}

func TestBasicSwarm(t *testing.T) {
	// t.Skip("skipping for another test")
	t.Parallel()

	msgs := 1
	swarms := 2
	SubtestSwarm(t, swarms, msgs)
}

func TestConnHandler(t *testing.T) {
	// t.Skip("skipping for another test")
	t.Parallel()

	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 5)

	gotconn := make(chan struct{}, 10)
	swarms[0].SetConnHandler(func(conn *Conn) {
		gotconn <- struct{}{}
	})

	connectSwarms(t, ctx, swarms)

	<-time.After(time.Millisecond)
	// should've gotten 5 by now.

	swarms[0].SetConnHandler(nil)

	expect := 4
	for i := 0; i < expect; i++ {
		select {
		case <-time.After(time.Second):
			t.Fatal("failed to get connections")
		case <-gotconn:
		}
	}

	select {
	case <-gotconn:
		t.Fatalf("should have connected to %d swarms", expect)
	default:
	}
}

func TestAddrBlocking(t *testing.T) {
	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 2)

	swarms[0].SetConnHandler(func(conn *Conn) {
		t.Fatalf("no connections should happen! -- %s", conn)
	})

	_, block, err := net.ParseCIDR("127.0.0.1/8")
	if err != nil {
		t.Fatal(err)
	}

	swarms[1].Filters.AddDialFilter(block)

	swarms[1].peers.AddAddr(swarms[0].LocalPeer(), swarms[0].ListenAddresses()[0], peer.PermanentAddrTTL)
	_, err = swarms[1].Dial(ctx, swarms[0].LocalPeer())
	if err == nil {
		t.Fatal("dial should have failed")
	}

	swarms[0].peers.AddAddr(swarms[1].LocalPeer(), swarms[1].ListenAddresses()[0], peer.PermanentAddrTTL)
	_, err = swarms[0].Dial(ctx, swarms[1].LocalPeer())
	if err == nil {
		t.Fatal("dial should have failed")
	}
}

func TestFilterBounds(t *testing.T) {
	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 2)

	conns := make(chan struct{}, 8)
	swarms[0].SetConnHandler(func(conn *Conn) {
		conns <- struct{}{}
	})

	// Address that we wont be dialing from
	_, block, err := net.ParseCIDR("192.0.0.1/8")
	if err != nil {
		t.Fatal(err)
	}

	// set filter on both sides, shouldnt matter
	swarms[1].Filters.AddDialFilter(block)
	swarms[0].Filters.AddDialFilter(block)

	connectSwarms(t, ctx, swarms)

	select {
	case <-time.After(time.Second):
		t.Fatal("should have gotten connection")
	case <-conns:
		t.Log("got connect")
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"os"

	"github.com/jacobsa/comeback/graph"
	"golang.org/x/net/context"
)

type PathAndFileInfo struct {
	// The absolute path to the file (or directory, etc.).
	Path string

	// Information about the file, as returned by os.Lstat.
	Info os.FileInfo
}

// Create a visitor that walks the directory hierarchy rooted at the given base
// path, excluding any relative path that matches any of the supplied
// exclusions, along with any of its descendents. Everything encountered and
// not excluded will be emitted to the supplied channel in an unspecified
// order. The channel will not be closed.
//
// It is expected that node names are paths relative to the supplied base path.
// In particular, to walk the entire hierarchy, use "" as the traversal root.
func NewFileSystemVisitor(
	ctx context.Context,
	basePath string,
	output chan<- PathAndFileInfo) (v graph.Visitor)
<commit_msg>Fixed a build error.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"os"

	"github.com/jacobsa/comeback/graph"
	"golang.org/x/net/context"
)

type PathAndFileInfo struct {
	// The absolute path to the file (or directory, etc.).
	Path string

	// Information about the file, as returned by os.Lstat.
	Info os.FileInfo
}

// Create a visitor that walks the directory hierarchy rooted at the given base
// path, excluding any relative path that matches any of the supplied
// exclusions, along with any of its descendents. Everything encountered and
// not excluded will be emitted to the supplied channel in an unspecified
// order. The channel will not be closed.
//
// It is expected that node names are paths relative to the supplied base path.
// In particular, to walk the entire hierarchy, use "" as the traversal root.
func NewFileSystemVisitor(
	ctx context.Context,
	basePath string,
	output chan<- PathAndFileInfo) (v graph.Visitor) {
	panic("TODO")
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"github.com/mitchellh/cli"
	"sort"
	"strings"
)

// InfoCommand is a Command implementation that queries a running
// Consul agent for various debugging statistics for operators
type InfoCommand struct {
	Ui cli.Ui
}

func (i *InfoCommand) Help() string {
	helpText := `
Usage: consul info [options]

	Provides debugging information for operators

Options:

  -rpc-addr=127.0.0.1:8400  RPC address of the Consul agent.
`
	return strings.TrimSpace(helpText)
}

func (i *InfoCommand) Run(args []string) int {
	cmdFlags := flag.NewFlagSet("info", flag.ContinueOnError)
	cmdFlags.Usage = func() { i.Ui.Output(i.Help()) }
	rpcAddr := RPCAddrFlag(cmdFlags)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	client, err := RPCClient(*rpcAddr)
	if err != nil {
		i.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
		return 1
	}
	defer client.Close()

	stats, err := client.Stats()
	if err != nil {
		i.Ui.Error(fmt.Sprintf("Error querying agent: %s", err))
		return 1
	}

	// Get the keys in sorted order
	keys := make([]string, 0, len(stats))
	for key := range stats {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	// Iterate over each top-level key
	for _, key := range keys {
		i.Ui.Output(key + ":")

		// Sort the sub-keys
		subvals := stats[key]
		subkeys := make([]string, 0, len(subvals))
		for k := range subvals {
			subkeys = append(subkeys, k)
		}
		sort.Strings(subkeys)

		// Iterate over the subkeys
		for _, subkey := range subkeys {
			val := subvals[subkey]
			i.Ui.Output(fmt.Sprintf("\t%s = %s", subkey, val))
		}
	}
	return 0
}

func (i *InfoCommand) Synopsis() string {
	return "Provides debugging information for operators"
}
<commit_msg>command/info: Warn on GOMAXPROCS = 1. Fixes #87.<commit_after>package command

import (
	"flag"
	"fmt"
	"github.com/mitchellh/cli"
	"sort"
	"strings"
)

// InfoCommand is a Command implementation that queries a running
// Consul agent for various debugging statistics for operators
type InfoCommand struct {
	Ui cli.Ui
}

func (i *InfoCommand) Help() string {
	helpText := `
Usage: consul info [options]

	Provides debugging information for operators

Options:

  -rpc-addr=127.0.0.1:8400  RPC address of the Consul agent.
`
	return strings.TrimSpace(helpText)
}

func (i *InfoCommand) Run(args []string) int {
	cmdFlags := flag.NewFlagSet("info", flag.ContinueOnError)
	cmdFlags.Usage = func() { i.Ui.Output(i.Help()) }
	rpcAddr := RPCAddrFlag(cmdFlags)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	client, err := RPCClient(*rpcAddr)
	if err != nil {
		i.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
		return 1
	}
	defer client.Close()

	stats, err := client.Stats()
	if err != nil {
		i.Ui.Error(fmt.Sprintf("Error querying agent: %s", err))
		return 1
	}

	// Get the keys in sorted order
	keys := make([]string, 0, len(stats))
	for key := range stats {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	// Iterate over each top-level key
	for _, key := range keys {
		i.Ui.Output(key + ":")

		// Sort the sub-keys
		subvals := stats[key]
		subkeys := make([]string, 0, len(subvals))
		for k := range subvals {
			subkeys = append(subkeys, k)
		}
		sort.Strings(subkeys)

		// Iterate over the subkeys
		for _, subkey := range subkeys {
			val := subvals[subkey]
			i.Ui.Output(fmt.Sprintf("\t%s = %s", subkey, val))
		}
	}

	// Check for specific warnings
	runtime, ok := stats["runtime"]
	if ok {
		if maxprocs := runtime["max_procs"]; maxprocs == "1" {
			i.Ui.Output("WARNING: It is highly recommended to set GOMAXPROCS higher than 1")
		}
	}
	return 0
}

func (i *InfoCommand) Synopsis() string {
	return "Provides debugging information for operators"
}
<|endoftext|>