<commit_before>package benchlist

import (
	"container/list"
	"math/rand"
	"sync"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/timer"

	safemath "github.com/ava-labs/avalanchego/utils/math"
)

// QueryBenchlist ...
type QueryBenchlist interface {
	// RegisterQuery registers a sent query and returns whether the query is subject to benchlist
	RegisterQuery(ids.ShortID, uint32) bool
	// RegisterResponse registers the response to a query message
	RegisterResponse(ids.ShortID, uint32)
	// QueryFailed registers that a query did not receive a response within our synchrony bound
	QueryFailed(ids.ShortID, uint32)
}

// If a peer consistently does not respond to queries, it will
// increase latencies on the network whenever that peer is polled.
// If we cannot terminate the poll early, then the poll will wait
// the full timeout before finalizing the poll and making progress.
// This can increase network latencies to an undesirable level.

// Therefore, a benchlist is used as a heurstic to immediately fail
// queries to nodes that are consistently not responding.

type queryBenchlist struct {
	vdrs validators.Set
	// Validator ID --> Request ID --> non-empty iff
	// there is an outstanding request to this validator
	// with the corresponding requestID
	pendingQueries map[[20]byte]map[uint32]struct{}
	// Map of consecutive query failures
	consecutiveFailures map[[20]byte]int

	// Maintain benchlist
	benchlistTimes map[[20]byte]time.Time
	benchlistOrder *list.List
	benchlistSet   ids.ShortSet

	threshold    int
	halfDuration time.Duration
	maxPortion   float64

	clock timer.Clock

	metrics *metrics
	ctx     *snow.Context

	lock sync.Mutex
}

// NewQueryBenchlist ...
func NewQueryBenchlist(validators validators.Set, ctx *snow.Context, threshold int, duration time.Duration, maxPortion float64) QueryBenchlist {
	metrics := &metrics{}
	metrics.Initialize(ctx.Namespace, ctx.Metrics)

	return &queryBenchlist{
		pendingQueries:      make(map[[20]byte]map[uint32]struct{}),
		consecutiveFailures: make(map[[20]byte]int),
		benchlistTimes:      make(map[[20]byte]time.Time),
		benchlistOrder:      list.New(),
		benchlistSet:        ids.ShortSet{},
		vdrs:                validators,
		threshold:           threshold,
		halfDuration:        duration / 2,
		maxPortion:          maxPortion,
		ctx:                 ctx,
		metrics:             metrics,
	}
}

// RegisterQuery attempts to register a query from [validatorID] and returns true
// if that request should be made (not subject to benchlisting)
func (b *queryBenchlist) RegisterQuery(validatorID ids.ShortID, requestID uint32) bool {
	b.lock.Lock()
	defer b.lock.Unlock()

	key := validatorID.Key()
	if benched := b.benched(validatorID); benched {
		return false
	}
	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		validatorRequests = make(map[uint32]struct{})
		b.pendingQueries[key] = validatorRequests
	}
	validatorRequests[requestID] = struct{}{}

	return true
}

// RegisterResponse removes the query from pending
func (b *queryBenchlist) RegisterResponse(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	// Reset consecutive failures on success
	delete(b.consecutiveFailures, validatorID.Key())
}

// QueryFailed notes a failure and benchlists [validatorID] if necessary
func (b *queryBenchlist) QueryFailed(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	key := validatorID.Key()
	// Add a failure and benches [validatorID] if it has
	// passed the threshold
	b.consecutiveFailures[key]++
	if b.consecutiveFailures[key] >= b.threshold {
		b.bench(validatorID)
	}
}

func (b *queryBenchlist) bench(validatorID ids.ShortID) {
	if b.benchlistSet.Contains(validatorID) {
		return
	}

	key := validatorID.Key()

	// Add to benchlist times with randomized delay
	randomizedDuration := time.Duration(rand.Float64()*float64(b.halfDuration)) + b.halfDuration
	b.benchlistTimes[key] = b.clock.Time().Add(randomizedDuration)
	b.benchlistOrder.PushBack(validatorID)
	b.benchlistSet.Add(validatorID)
	delete(b.consecutiveFailures, key)
	b.metrics.numBenched.Inc()
	b.ctx.Log.Debug("Benching validator %s for %v after %d consecutive failed queries", validatorID, randomizedDuration, b.threshold)

	// Note: there could be a memory leak if a large number of
	// validators were added, sampled, benched, and never sampled
	// again. Due to the minimum staking amount and durations this
	// is not a realistic concern.
	b.cleanup()
}

// benched checks if [validatorID] is currently benched
// and calls cleanup if its benching period has elapsed
func (b *queryBenchlist) benched(validatorID ids.ShortID) bool {
	key := validatorID.Key()

	end, ok := b.benchlistTimes[key]
	if !ok {
		return false
	}

	if b.clock.Time().Before(end) {
		return true
	}

	// If a benched item has expired, cleanup the benchlist
	b.cleanup()
	return false
}

// cleanup ensures that we have not benched too much stake
// and removes anything from the benchlist whose time has expired
func (b *queryBenchlist) cleanup() {
	currentWeight, err := b.vdrs.SubsetWeight(b.benchlistSet)
	if err != nil {
		// Add log for this, should never happen
		b.ctx.Log.Error("Failed to calculate subset weight due to: %w. Resetting benchlist.", err)
		b.reset()
		return
	}

	numBenched := b.benchlistSet.Len()
	updatedWeight := currentWeight
	totalWeight := b.vdrs.Weight()
	maxBenchlistWeight := uint64(float64(totalWeight) * b.maxPortion)

	// Iterate over elements of the benchlist in order of expiration
	for e := b.benchlistOrder.Front(); e != nil; e = e.Next() {
		validatorID := e.Value.(ids.ShortID)
		key := validatorID.Key()
		end := b.benchlistTimes[key]
		// Remove elements with the next expiration until the next item has not
		// expired and the bench has less than the maximum weight
		// Note: this creates an edge case where benchlisting a validator
		// with a sufficient stake may clear the benchlist
		if b.clock.Time().Before(end) && currentWeight < maxBenchlistWeight {
			break
		}

		removeWeight, ok := b.vdrs.GetWeight(validatorID)
		if ok {
			newWeight, err := safemath.Sub64(currentWeight, removeWeight)
			if err != nil {
				b.ctx.Log.Error("Failed to calculate new subset weight due to: %w. Resetting benchlist.", err)
				b.reset()
				return
			}
			updatedWeight = newWeight
		}

		b.benchlistOrder.Remove(e)
		delete(b.benchlistTimes, key)
		b.benchlistSet.Remove(validatorID)
		b.metrics.numBenched.Dec()
	}

	b.ctx.Log.Debug("Benchlist weight: (%v/%v) -> (%v/%v). Benched Validators: %d -> %d",
		currentWeight,
		totalWeight,
		updatedWeight,
		totalWeight,
		numBenched,
		b.benchlistSet.Len(),
	)
	b.metrics.weightBenched.Set(float64(updatedWeight))
}

func (b *queryBenchlist) reset() {
	b.pendingQueries = make(map[[20]byte]map[uint32]struct{})
	b.consecutiveFailures = make(map[[20]byte]int)
	b.benchlistTimes = make(map[[20]byte]time.Time)
	b.benchlistOrder.Init()
	b.benchlistSet.Clear()
	b.metrics.weightBenched.Set(0)
	b.metrics.numBenched.Set(0)
}

// removeQuery returns true if the query was present
func (b *queryBenchlist) removeQuery(validatorID ids.ShortID, requestID uint32) bool {
	key := validatorID.Key()

	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		return false
	}

	_, ok = validatorRequests[requestID]
	if ok {
		delete(validatorRequests, requestID)
		if len(validatorRequests) == 0 {
			delete(b.pendingQueries, key)
		}
	}
	return ok
}
<commit_msg>Add nosec to weak random number generation in randomized delay<commit_after>package benchlist

import (
	"container/list"
	"math/rand"
	"sync"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/timer"

	safemath "github.com/ava-labs/avalanchego/utils/math"
)

// QueryBenchlist ...
type QueryBenchlist interface {
	// RegisterQuery registers a sent query and returns whether the query is subject to benchlist
	RegisterQuery(ids.ShortID, uint32) bool
	// RegisterResponse registers the response to a query message
	RegisterResponse(ids.ShortID, uint32)
	// QueryFailed registers that a query did not receive a response within our synchrony bound
	QueryFailed(ids.ShortID, uint32)
}

// If a peer consistently does not respond to queries, it will
// increase latencies on the network whenever that peer is polled.
// If we cannot terminate the poll early, then the poll will wait
// the full timeout before finalizing the poll and making progress.
// This can increase network latencies to an undesirable level.

// Therefore, a benchlist is used as a heurstic to immediately fail
// queries to nodes that are consistently not responding.

type queryBenchlist struct {
	vdrs validators.Set
	// Validator ID --> Request ID --> non-empty iff
	// there is an outstanding request to this validator
	// with the corresponding requestID
	pendingQueries map[[20]byte]map[uint32]struct{}
	// Map of consecutive query failures
	consecutiveFailures map[[20]byte]int

	// Maintain benchlist
	benchlistTimes map[[20]byte]time.Time
	benchlistOrder *list.List
	benchlistSet   ids.ShortSet

	threshold    int
	halfDuration time.Duration
	maxPortion   float64

	clock timer.Clock

	metrics *metrics
	ctx     *snow.Context

	lock sync.Mutex
}

// NewQueryBenchlist ...
func NewQueryBenchlist(validators validators.Set, ctx *snow.Context, threshold int, duration time.Duration, maxPortion float64) QueryBenchlist {
	metrics := &metrics{}
	metrics.Initialize(ctx.Namespace, ctx.Metrics)

	return &queryBenchlist{
		pendingQueries:      make(map[[20]byte]map[uint32]struct{}),
		consecutiveFailures: make(map[[20]byte]int),
		benchlistTimes:      make(map[[20]byte]time.Time),
		benchlistOrder:      list.New(),
		benchlistSet:        ids.ShortSet{},
		vdrs:                validators,
		threshold:           threshold,
		halfDuration:        duration / 2,
		maxPortion:          maxPortion,
		ctx:                 ctx,
		metrics:             metrics,
	}
}

// RegisterQuery attempts to register a query from [validatorID] and returns true
// if that request should be made (not subject to benchlisting)
func (b *queryBenchlist) RegisterQuery(validatorID ids.ShortID, requestID uint32) bool {
	b.lock.Lock()
	defer b.lock.Unlock()

	key := validatorID.Key()
	if benched := b.benched(validatorID); benched {
		return false
	}
	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		validatorRequests = make(map[uint32]struct{})
		b.pendingQueries[key] = validatorRequests
	}
	validatorRequests[requestID] = struct{}{}

	return true
}

// RegisterResponse removes the query from pending
func (b *queryBenchlist) RegisterResponse(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	// Reset consecutive failures on success
	delete(b.consecutiveFailures, validatorID.Key())
}

// QueryFailed notes a failure and benchlists [validatorID] if necessary
func (b *queryBenchlist) QueryFailed(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	key := validatorID.Key()
	// Add a failure and benches [validatorID] if it has
	// passed the threshold
	b.consecutiveFailures[key]++
	if b.consecutiveFailures[key] >= b.threshold {
		b.bench(validatorID)
	}
}

func (b *queryBenchlist) bench(validatorID ids.ShortID) {
	if b.benchlistSet.Contains(validatorID) {
		return
	}

	key := validatorID.Key()

	// Add to benchlist times with randomized delay
	randomizedDuration := time.Duration(rand.Float64()*float64(b.halfDuration)) + b.halfDuration // #nosec G404
	b.benchlistTimes[key] = b.clock.Time().Add(randomizedDuration)
	b.benchlistOrder.PushBack(validatorID)
	b.benchlistSet.Add(validatorID)
	delete(b.consecutiveFailures, key)
	b.metrics.numBenched.Inc()
	b.ctx.Log.Debug("Benching validator %s for %v after %d consecutive failed queries", validatorID, randomizedDuration, b.threshold)

	// Note: there could be a memory leak if a large number of
	// validators were added, sampled, benched, and never sampled
	// again. Due to the minimum staking amount and durations this
	// is not a realistic concern.
	b.cleanup()
}

// benched checks if [validatorID] is currently benched
// and calls cleanup if its benching period has elapsed
func (b *queryBenchlist) benched(validatorID ids.ShortID) bool {
	key := validatorID.Key()

	end, ok := b.benchlistTimes[key]
	if !ok {
		return false
	}

	if b.clock.Time().Before(end) {
		return true
	}

	// If a benched item has expired, cleanup the benchlist
	b.cleanup()
	return false
}

// cleanup ensures that we have not benched too much stake
// and removes anything from the benchlist whose time has expired
func (b *queryBenchlist) cleanup() {
	currentWeight, err := b.vdrs.SubsetWeight(b.benchlistSet)
	if err != nil {
		// Add log for this, should never happen
		b.ctx.Log.Error("Failed to calculate subset weight due to: %w. Resetting benchlist.", err)
		b.reset()
		return
	}

	numBenched := b.benchlistSet.Len()
	updatedWeight := currentWeight
	totalWeight := b.vdrs.Weight()
	maxBenchlistWeight := uint64(float64(totalWeight) * b.maxPortion)

	// Iterate over elements of the benchlist in order of expiration
	for e := b.benchlistOrder.Front(); e != nil; e = e.Next() {
		validatorID := e.Value.(ids.ShortID)
		key := validatorID.Key()
		end := b.benchlistTimes[key]
		// Remove elements with the next expiration until the next item has not
		// expired and the bench has less than the maximum weight
		// Note: this creates an edge case where benchlisting a validator
		// with a sufficient stake may clear the benchlist
		if b.clock.Time().Before(end) && currentWeight < maxBenchlistWeight {
			break
		}

		removeWeight, ok := b.vdrs.GetWeight(validatorID)
		if ok {
			newWeight, err := safemath.Sub64(currentWeight, removeWeight)
			if err != nil {
				b.ctx.Log.Error("Failed to calculate new subset weight due to: %w. Resetting benchlist.", err)
				b.reset()
				return
			}
			updatedWeight = newWeight
		}

		b.benchlistOrder.Remove(e)
		delete(b.benchlistTimes, key)
		b.benchlistSet.Remove(validatorID)
		b.metrics.numBenched.Dec()
	}

	b.ctx.Log.Debug("Benchlist weight: (%v/%v) -> (%v/%v). Benched Validators: %d -> %d",
		currentWeight,
		totalWeight,
		updatedWeight,
		totalWeight,
		numBenched,
		b.benchlistSet.Len(),
	)
	b.metrics.weightBenched.Set(float64(updatedWeight))
}

func (b *queryBenchlist) reset() {
	b.pendingQueries = make(map[[20]byte]map[uint32]struct{})
	b.consecutiveFailures = make(map[[20]byte]int)
	b.benchlistTimes = make(map[[20]byte]time.Time)
	b.benchlistOrder.Init()
	b.benchlistSet.Clear()
	b.metrics.weightBenched.Set(0)
	b.metrics.numBenched.Set(0)
}

// removeQuery returns true if the query was present
func (b *queryBenchlist) removeQuery(validatorID ids.ShortID, requestID uint32) bool {
	key := validatorID.Key()

	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		return false
	}

	_, ok = validatorRequests[requestID]
	if ok {
		delete(validatorRequests, requestID)
		if len(validatorRequests) == 0 {
			delete(b.pendingQueries, key)
		}
	}
	return ok
}
<|endoftext|>


<commit_before>package vm_test

import (
	. "github.com/onsi/ginkgo"

	. "github.com/maximilien/bosh-softlayer-cpi/softlayer/vm"

	boshlog "github.com/cloudfoundry/bosh-agent/logger"

	fakevm "github.com/maximilien/bosh-softlayer-cpi/softlayer/vm/fakes"
	fakeslclient "github.com/maximilien/softlayer-go/client/fakes"
)

var _ = Describe("SoftLayerCreator", func() {
	var (
		softLayerClient        *fakeslclient.FakeSoftLayerClient
		agentEnvServiceFactory *fakevm.FakeAgentEnvServiceFactory
		agentOptions           AgentOptions
		logger                 boshlog.Logger
		creator                SoftLayerCreator
	)

	BeforeEach(func() {
		softLayerClient = fakeslclient.NewFakeSoftLayerClient("fake-username", "fake-api-key")
		agentEnvServiceFactory = &fakevm.FakeAgentEnvServiceFactory{}
		agentOptions = AgentOptions{Mbus: "fake-mbus"}
		logger = boshlog.NewLogger(boshlog.LevelNone)

		creator = NewSoftLayerCreator(
			softLayerClient,
			agentEnvServiceFactory,
			agentOptions,
			logger,
		)
	})

	Describe("Create", func() {
	})
})
<commit_msg>added test for creator that uses fake SL client with test fixtures<commit_after>package vm_test

import (
	"os"
	"path/filepath"

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

	. "github.com/maximilien/bosh-softlayer-cpi/softlayer/vm"

	fakevm "github.com/maximilien/bosh-softlayer-cpi/softlayer/vm/fakes"
	fakeslclient "github.com/maximilien/softlayer-go/client/fakes"

	common "github.com/maximilien/bosh-softlayer-cpi/common"
	bslcstem "github.com/maximilien/bosh-softlayer-cpi/softlayer/stemcell"

	boshlog "github.com/cloudfoundry/bosh-agent/logger"

	sldatatypes "github.com/maximilien/softlayer-go/data_types"
)

var _ = Describe("SoftLayerCreator", func() {
	var (
		softLayerClient        *fakeslclient.FakeSoftLayerClient
		agentEnvServiceFactory *fakevm.FakeAgentEnvServiceFactory
		agentOptions           AgentOptions
		logger                 boshlog.Logger
		creator                SoftLayerCreator
	)

	BeforeEach(func() {
		workingDir, err := os.Getwd()
		Expect(err).ToNot(HaveOccurred())

		softLayerClient = fakeslclient.NewFakeSoftLayerClient("fake-username", "fake-api-key")
		softLayerClient.DoRawHttpRequestResponse, err = common.ReadJsonTestFixtures(filepath.Join(workingDir, "..", ".."), "softlayer", "SoftLayer_Virtual_Guest_Service_createObject.json")
		Expect(err).ToNot(HaveOccurred())

		agentEnvServiceFactory = &fakevm.FakeAgentEnvServiceFactory{}
		agentOptions = AgentOptions{Mbus: "fake-mbus"}
		logger = boshlog.NewLogger(boshlog.LevelNone)

		creator = NewSoftLayerCreator(
			softLayerClient,
			agentEnvServiceFactory,
			agentOptions,
			logger,
		)
	})

	Describe("Create", func() {
		var (
			agentID    string
			stemcell   bslcstem.FSStemcell
			cloudProps VMCloudProperties
			networks   Networks
			env        Environment
		)

		Context("valid arguments", func() {
			BeforeEach(func() {
				agentID = "fake-agent-id"
				stemcell = bslcstem.NewFSStemcell("fake-stemcell-id", logger)
				cloudProps = VMCloudProperties{
					StartCpus:  4,
					MaxMemory:  2048,
					Datacenter: sldatatypes.Datacenter{Name: "fake-datacenter"},
				}
				networks = Networks{}
				env = Environment{}
			})

			It("returns a new SoftLayerVM with correct virtual guest ID and SoftLayerClient", func() {
				vm, err := creator.Create(agentID, stemcell, cloudProps, networks, env)
				Expect(err).ToNot(HaveOccurred())
				Expect(vm.ID()).To(Equal(1234567))
			})
		})

		Context("invalid arguments", func() {
			Context("missing correct VMProperties", func() {
				BeforeEach(func() {
					agentID = "fake-agent-id"
					stemcell = bslcstem.NewFSStemcell("fake-stemcell-id", logger)
					networks = Networks{}
					env = Environment{}
				})

				It("fails when VMProperties is missing StartCpus", func() {
					cloudProps = VMCloudProperties{
						MaxMemory:  2048,
						Datacenter: sldatatypes.Datacenter{Name: "fake-datacenter"},
					}

					_, err := creator.Create(agentID, stemcell, cloudProps, networks, env)
					Expect(err).To(HaveOccurred())
				})

				It("fails when VMProperties is missing MaxMemory", func() {
					cloudProps = VMCloudProperties{
						StartCpus:  4,
						Datacenter: sldatatypes.Datacenter{Name: "fake-datacenter"},
					}

					_, err := creator.Create(agentID, stemcell, cloudProps, networks, env)
					Expect(err).To(HaveOccurred())
				})

				It("fails when VMProperties is missing MaxMemory", func() {
					cloudProps = VMCloudProperties{
						StartCpus: 4,
						MaxMemory: 1024,
					}

					_, err := creator.Create(agentID, stemcell, cloudProps, networks, env)
					Expect(err).To(HaveOccurred())
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>package command

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"sync"

	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/coreos/go-etcd/etcd"
	"github.com/coreos/etcd/store"
)

type set struct {
	key   string
	value string
	ttl   int64
}

func NewImportSnapCommand() cli.Command {
	return cli.Command{
		Name:  "import",
		Usage: "import a snapshot to a cluster",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "snap", Value: "", Usage: "Path to the vaild etcd 0.4.x snapshot."},
			cli.StringSliceFlag{Name: "hidden", Value: nil, Usage: "Hidden key space to import from snapshot"},
			cli.IntFlag{Name: "c", Value: 10, Usage: "Number of concurrent clients to import the data"},
		},
		Action: handleImportSnap,
	}
}

func handleImportSnap(c *cli.Context) {
	d, err := ioutil.ReadFile(c.String("snap"))
	if err != nil {
		if c.String("snap") == "" {
			fmt.Printf("no snapshot file provided (use --snap)")
		} else {
			fmt.Printf("cannot read snapshot file %s\n", c.String("snap"))
		}
		os.Exit(1)
	}

	st := store.New()
	err = st.Recovery(d)
	if err != nil {
		fmt.Printf("cannot recover the snapshot file: %v\n", err)
		os.Exit(1)
	}

	endpoints, err := getEndpoints(c)
	if err != nil {
		handleError(ErrorFromEtcd, err)
	}
	tr, err := getTransport(c)
	if err != nil {
		handleError(ErrorFromEtcd, err)
	}

	wg := &sync.WaitGroup{}
	setc := make(chan set)
	concurrent := c.Int("c")
	fmt.Printf("starting to import snapshot %s with %d clients\n", c.String("snap"), concurrent)
	for i := 0; i < concurrent; i++ {
		client := etcd.NewClient(endpoints)
		client.SetTransport(tr)

		if c.GlobalBool("debug") {
			go dumpCURL(client)
		}

		if ok := client.SyncCluster(); !ok {
			handleError(FailedToConnectToHost, errors.New("cannot sync with the cluster using endpoints "+strings.Join(endpoints, ", ")))
		}
		wg.Add(1)
		go runSet(client, setc, wg)
	}

	all, err := st.Get("/", true, true)
	if err != nil {
		handleError(ErrorFromEtcd, err)
	}
	n := copyKeys(all.Node, setc)

	hiddens := c.StringSlice("hidden")
	for _, h := range hiddens {
		allh, err := st.Get(h, true, true)
		if err != nil {
			handleError(ErrorFromEtcd, err)
		}
		n += copyKeys(allh.Node, setc)
	}
	close(setc)
	wg.Wait()
	fmt.Printf("finished importing %d keys\n", n)
}

func copyKeys(n *store.NodeExtern, setc chan set) int {
	num := 0
	if !n.Dir {
		setc <- set{n.Key, *n.Value, n.TTL}
		return 1
	}
	log.Println("entering dir:", n.Key)
	for _, nn := range n.Nodes {
		sub := copyKeys(nn, setc)
		num += sub
	}
	return num
}

func runSet(c *etcd.Client, setc chan set, wg *sync.WaitGroup) {
	for s := range setc {
		log.Println("copying key:", s.key)
		if s.ttl != 0 && s.ttl < 300 {
			log.Printf("extending key %s's ttl to 300 seconds", s.key)
			s.ttl = 5 * 60
		}
		_, err := c.Set(s.key, s.value, uint64(s.ttl))
		if err != nil {
			log.Fatalf("failed to copy key: %v\n", err)
		}
	}
	wg.Done()
}
<commit_msg>etcdctl: refactor message in import command<commit_after>package command

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"sync"

	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/coreos/go-etcd/etcd"
	"github.com/coreos/etcd/store"
)

type set struct {
	key   string
	value string
	ttl   int64
}

func NewImportSnapCommand() cli.Command {
	return cli.Command{
		Name:  "import",
		Usage: "import a snapshot to a cluster",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "snap", Value: "", Usage: "Path to the vaild etcd 0.4.x snapshot."},
			cli.StringSliceFlag{Name: "hidden", Value: new(cli.StringSlice), Usage: "Hidden key spaces to import from snapshot"},
			cli.IntFlag{Name: "c", Value: 10, Usage: "Number of concurrent clients to import the data"},
		},
		Action: handleImportSnap,
	}
}

func handleImportSnap(c *cli.Context) {
	d, err := ioutil.ReadFile(c.String("snap"))
	if err != nil {
		if c.String("snap") == "" {
			fmt.Printf("no snapshot file provided (use --snap)\n")
		} else {
			fmt.Printf("cannot read snapshot file %s\n", c.String("snap"))
		}
		os.Exit(1)
	}

	st := store.New()
	err = st.Recovery(d)
	if err != nil {
		fmt.Printf("cannot recover the snapshot file: %v\n", err)
		os.Exit(1)
	}

	endpoints, err := getEndpoints(c)
	if err != nil {
		handleError(ErrorFromEtcd, err)
	}
	tr, err := getTransport(c)
	if err != nil {
		handleError(ErrorFromEtcd, err)
	}

	wg := &sync.WaitGroup{}
	setc := make(chan set)
	concurrent := c.Int("c")
	fmt.Printf("starting to import snapshot %s with %d clients\n", c.String("snap"), concurrent)
	for i := 0; i < concurrent; i++ {
		client := etcd.NewClient(endpoints)
		client.SetTransport(tr)

		if c.GlobalBool("debug") {
			go dumpCURL(client)
		}

		if ok := client.SyncCluster(); !ok {
			handleError(FailedToConnectToHost, errors.New("cannot sync with the cluster using endpoints "+strings.Join(endpoints, ", ")))
		}
		wg.Add(1)
		go runSet(client, setc, wg)
	}

	all, err := st.Get("/", true, true)
	if err != nil {
		handleError(ErrorFromEtcd, err)
	}
	n := copyKeys(all.Node, setc)

	hiddens := c.StringSlice("hidden")
	for _, h := range hiddens {
		allh, err := st.Get(h, true, true)
		if err != nil {
			handleError(ErrorFromEtcd, err)
		}
		n += copyKeys(allh.Node, setc)
	}
	close(setc)
	wg.Wait()
	fmt.Printf("finished importing %d keys\n", n)
}

func copyKeys(n *store.NodeExtern, setc chan set) int {
	num := 0
	if !n.Dir {
		setc <- set{n.Key, *n.Value, n.TTL}
		return 1
	}
	log.Println("entering dir:", n.Key)
	for _, nn := range n.Nodes {
		sub := copyKeys(nn, setc)
		num += sub
	}
	return num
}

func runSet(c *etcd.Client, setc chan set, wg *sync.WaitGroup) {
	for s := range setc {
		log.Println("copying key:", s.key)
		if s.ttl != 0 && s.ttl < 300 {
			log.Printf("extending key %s's ttl to 300 seconds", s.key)
			s.ttl = 5 * 60
		}
		_, err := c.Set(s.key, s.value, uint64(s.ttl))
		if err != nil {
			log.Fatalf("failed to copy key: %v\n", err)
		}
	}
	wg.Done()
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/google/mtail/metrics"
)

// Unparser is for converting program syntax trees back to program text.
type Unparser struct {
	pos       int
	output    string
	line      string
	emitTypes bool
}

func (u *Unparser) indent() {
	u.pos += 2
}

func (u *Unparser) outdent() {
	u.pos -= 2
}

func (u *Unparser) prefix() (s string) {
	for i := 0; i < u.pos; i++ {
		s += " "
	}
	return
}

func (u *Unparser) emit(s string) {
	u.line += s
}

func (u *Unparser) newline() {
	u.output += u.prefix() + u.line + "\n"
	u.line = ""
}

func (u *Unparser) VisitBefore(n astNode) Visitor {
	if u.emitTypes {
		u.emit(fmt.Sprintf("<%s> ", n.Type()))
	}
	switch v := n.(type) {
	case *stmtlistNode:
		for _, child := range v.children {
			Walk(u, child)
			u.newline()
		}

	case *exprlistNode:
		if len(v.children) > 0 {
			Walk(u, v.children[0])
			for _, child := range v.children[1:] {
				u.emit(", ")
				Walk(u, child)
			}
		}

	case *condNode:
		if v.cond != nil {
			Walk(u, v.cond)
		}
		u.emit(" {")
		u.newline()
		u.indent()
		Walk(u, v.truthNode)
		if v.elseNode != nil {
			u.outdent()
			u.emit("} else {")
			u.indent()
			Walk(u, v.elseNode)
		}
		u.outdent()
		u.emit("}")

	case *regexNode:
		u.emit("/" + strings.Replace(v.pattern, "/", "\\/", -1) + "/")

	case *binaryExprNode:
		Walk(u, v.lhs)
		switch v.op {
		case LT:
			u.emit(" < ")
		case GT:
			u.emit(" > ")
		case LE:
			u.emit(" <= ")
		case GE:
			u.emit(" >= ")
		case EQ:
			u.emit(" == ")
		case NE:
			u.emit(" != ")
		case SHL:
			u.emit(" << ")
		case SHR:
			u.emit(" >> ")
		case AND:
			u.emit(" & ")
		case OR:
			u.emit(" | ")
		case XOR:
			u.emit(" ^ ")
		case NOT:
			u.emit(" ~ ")
		case PLUS:
			u.emit(" + ")
		case MINUS:
			u.emit(" - ")
		case MUL:
			u.emit(" * ")
		case DIV:
			u.emit(" / ")
		case POW:
			u.emit(" ** ")
		case ASSIGN:
			u.emit(" = ")
		case MOD:
			u.emit(" % ")
		}
		Walk(u, v.rhs)

	case *idNode:
		u.emit(v.name)

	case *caprefNode:
		u.emit("$" + v.name)

	case *builtinNode:
		u.emit(v.name + "(")
		if v.args != nil {
			Walk(u, v.args)
		}
		u.emit(")")

	case *indexedExprNode:
		Walk(u, v.lhs)
		u.emit("[")
		Walk(u, v.index)
		u.emit("]")

	case *declNode:
		switch v.kind {
		case metrics.Counter:
			u.emit("counter ")
		case metrics.Gauge:
			u.emit("gauge ")
		case metrics.Timer:
			u.emit("timer ")
		}
		u.emit(v.name)
		if len(v.keys) > 0 {
			u.emit(" by " + strings.Join(v.keys, ", "))
		}

	case *unaryExprNode:
		switch v.op {
		case INC:
			Walk(u, v.expr)
			u.emit("++")
		case NOT:
			u.emit(" ~")
			Walk(u, v.expr)
		}

	case *stringConstNode:
		u.emit("\"" + v.text + "\"")

	case *intConstNode:
		u.emit(strconv.FormatInt(v.i, 10))

	case *floatConstNode:
		u.emit(strconv.FormatFloat(v.f, 'g', -1, 64))

	case *defNode:
		u.emit(fmt.Sprintf("def %s {", v.name))
		u.newline()
		u.indent()
		Walk(u, v.block)
		u.outdent()
		u.emit("}")

	case *decoNode:
		u.emit(fmt.Sprintf("@%s {", v.name))
		u.newline()
		u.indent()
		Walk(u, v.block)
		u.outdent()
		u.emit("}")

	case *nextNode:
		u.emit("next")

	case *otherwiseNode:
		u.emit("otherwise")

	case *delNode:
		u.emit("del ")
		Walk(u, v.n)
		u.newline()

	default:
		panic(fmt.Sprintf("unparser found undefined type %T", n))
	}

	return nil
}

func (u *Unparser) VisitAfter(n astNode) {}

// Unparse begins the unparsing of the syntax tree, returning the program text as a single string.
func (u *Unparser) Unparse(n astNode) string {
	Walk(u, n)
	return u.output
}
<commit_msg>Emit type information with correct grouping.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/google/mtail/metrics"
)

// Unparser is for converting program syntax trees back to program text.
type Unparser struct {
	pos       int
	output    string
	line      string
	emitTypes bool
}

func (u *Unparser) indent() {
	u.pos += 2
}

func (u *Unparser) outdent() {
	u.pos -= 2
}

func (u *Unparser) prefix() (s string) {
	for i := 0; i < u.pos; i++ {
		s += " "
	}
	return
}

func (u *Unparser) emit(s string) {
	u.line += s
}

func (u *Unparser) newline() {
	u.output += u.prefix() + u.line + "\n"
	u.line = ""
}

func (u *Unparser) VisitBefore(n astNode) Visitor {
	if u.emitTypes {
		u.emit(fmt.Sprintf("<%s>(", n.Type()))
	}
	switch v := n.(type) {
	case *stmtlistNode:
		for _, child := range v.children {
			Walk(u, child)
			u.newline()
		}

	case *exprlistNode:
		if len(v.children) > 0 {
			Walk(u, v.children[0])
			for _, child := range v.children[1:] {
				u.emit(", ")
				Walk(u, child)
			}
		}

	case *condNode:
		if v.cond != nil {
			Walk(u, v.cond)
		}
		u.emit(" {")
		u.newline()
		u.indent()
		Walk(u, v.truthNode)
		if v.elseNode != nil {
			u.outdent()
			u.emit("} else {")
			u.indent()
			Walk(u, v.elseNode)
		}
		u.outdent()
		u.emit("}")

	case *regexNode:
		u.emit("/" + strings.Replace(v.pattern, "/", "\\/", -1) + "/")

	case *binaryExprNode:
		Walk(u, v.lhs)
		switch v.op {
		case LT:
			u.emit(" < ")
		case GT:
			u.emit(" > ")
		case LE:
			u.emit(" <= ")
		case GE:
			u.emit(" >= ")
		case EQ:
			u.emit(" == ")
		case NE:
			u.emit(" != ")
		case SHL:
			u.emit(" << ")
		case SHR:
			u.emit(" >> ")
		case AND:
			u.emit(" & ")
		case OR:
			u.emit(" | ")
		case XOR:
			u.emit(" ^ ")
		case NOT:
			u.emit(" ~ ")
		case PLUS:
			u.emit(" + ")
		case MINUS:
			u.emit(" - ")
		case MUL:
			u.emit(" * ")
		case DIV:
			u.emit(" / ")
		case POW:
			u.emit(" ** ")
		case ASSIGN:
			u.emit(" = ")
		case MOD:
			u.emit(" % ")
		}
		Walk(u, v.rhs)

	case *idNode:
		u.emit(v.name)

	case *caprefNode:
		u.emit("$" + v.name)

	case *builtinNode:
		u.emit(v.name + "(")
		if v.args != nil {
			Walk(u, v.args)
		}
		u.emit(")")

	case *indexedExprNode:
		Walk(u, v.lhs)
		u.emit("[")
		Walk(u, v.index)
		u.emit("]")

	case *declNode:
		switch v.kind {
		case metrics.Counter:
			u.emit("counter ")
		case metrics.Gauge:
			u.emit("gauge ")
		case metrics.Timer:
			u.emit("timer ")
		}
		u.emit(v.name)
		if len(v.keys) > 0 {
			u.emit(" by " + strings.Join(v.keys, ", "))
		}

	case *unaryExprNode:
		switch v.op {
		case INC:
			Walk(u, v.expr)
			u.emit("++")
		case NOT:
			u.emit(" ~")
			Walk(u, v.expr)
		}

	case *stringConstNode:
		u.emit("\"" + v.text + "\"")

	case *intConstNode:
		u.emit(strconv.FormatInt(v.i, 10))

	case *floatConstNode:
		u.emit(strconv.FormatFloat(v.f, 'g', -1, 64))

	case *defNode:
		u.emit(fmt.Sprintf("def %s {", v.name))
		u.newline()
		u.indent()
		Walk(u, v.block)
		u.outdent()
		u.emit("}")

	case *decoNode:
		u.emit(fmt.Sprintf("@%s {", v.name))
		u.newline()
		u.indent()
		Walk(u, v.block)
		u.outdent()
		u.emit("}")

	case *nextNode:
		u.emit("next")

	case *otherwiseNode:
		u.emit("otherwise")

	case *delNode:
		u.emit("del ")
		Walk(u, v.n)
		u.newline()

	default:
		panic(fmt.Sprintf("unparser found undefined type %T", n))
	}
	if u.emitTypes {
		u.emit(")")
	}
	return nil
}

func (u *Unparser) VisitAfter(n astNode) {
}

// Unparse begins the unparsing of the syntax tree, returning the program text as a single string.
func (u *Unparser) Unparse(n astNode) string {
	Walk(u, n)
	return u.output
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative 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

    https://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 e2e

import (
	"flag"
)

var _ = initializeFlags()

func initializeFlags() bool {
	// emitmetrics is a required flag for running periodic test jobs, add it here as a no-op to avoid the error
	emitMetrics := flag.Bool("emitmetrics", false,
		"Set this flag to true if you would like tests to emit metrics, e.g. latency of resources being realized in the system.")
	flag.Parse()
	return *emitMetrics
}
<commit_msg>Remove explicit flag parsing so that it works with go1.13 (#522)<commit_after>/*
Copyright 2019 The Knative 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

    https://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 e2e

import (
	"flag"
)

var _ = initializeFlags()

func initializeFlags() bool {
	// emitmetrics is a required flag for running periodic test jobs, add it here as a no-op to avoid the error
	emitMetrics := flag.Bool("emitmetrics", false,
		"Set this flag to true if you would like tests to emit metrics, e.g. latency of resources being realized in the system.")
	return *emitMetrics
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors 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 e2e

import (
	"bytes"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/clientset_generated/release_1_2"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/metrics"

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

const (
	maxKubectlExecRetries = 5
)

// Framework supports common operations used by e2e tests; it will keep a client & a namespace for you.
// Eventual goal is to merge this with integration test framework.
type Framework struct {
	BaseName string

	Client        *client.Client
	Clientset_1_2 *release_1_2.Clientset

	Namespace                *api.Namespace   // Every test has at least one namespace
	namespacesToDelete       []*api.Namespace // Some tests have more than one.
	NamespaceDeletionTimeout time.Duration

	gatherer containerResourceGatherer
	// Constraints that passed to a check which is exectued after data is gathered to
	// see if 99% of results are within acceptable bounds. It as to be injected in the test,
	// as expectations vary greatly. Constraints are groupped by the container names.
	addonResourceConstraints map[string]resourceConstraint

	logsSizeWaitGroup    sync.WaitGroup
	logsSizeCloseChannel chan bool
	logsSizeVerifier     *LogsSizeVerifier

	// To make sure that this framework cleans up after itself, no matter what,
	// we install a cleanup action before each test and clear it after.  If we
	// should abort, the AfterSuite hook should run all cleanup actions.
	cleanupHandle CleanupActionHandle
}

type TestDataSummary interface {
	PrintHumanReadable() string
	PrintJSON() string
}

// NewFramework makes a new framework and sets up a BeforeEach/AfterEach for
// you (you can write additional before/after each functions).
func NewFramework(baseName string) *Framework {
	f := &Framework{
		BaseName:                 baseName,
		addonResourceConstraints: make(map[string]resourceConstraint),
	}

	BeforeEach(f.beforeEach)
	AfterEach(f.afterEach)

	return f
}

// beforeEach gets a client and makes a namespace.
func (f *Framework) beforeEach() {
	// The fact that we need this feels like a bug in ginkgo.
	// https://github.com/onsi/ginkgo/issues/222
	f.cleanupHandle = AddCleanupAction(f.afterEach)

	By("Creating a kubernetes client")
	c, err := loadClient()
	Expect(err).NotTo(HaveOccurred())

	f.Client = c
	f.Clientset_1_2 = release_1_2.FromUnversionedClient(c)

	By("Building a namespace api object")
	namespace, err := f.CreateNamespace(f.BaseName, map[string]string{
		"e2e-framework": f.BaseName,
	})
	Expect(err).NotTo(HaveOccurred())

	f.Namespace = namespace

	if testContext.VerifyServiceAccount {
		By("Waiting for a default service account to be provisioned in namespace")
		err = waitForDefaultServiceAccountInNamespace(c, namespace.Name)
		Expect(err).NotTo(HaveOccurred())
	} else {
		Logf("Skipping waiting for service account")
	}

	if testContext.GatherKubeSystemResourceUsageData {
		f.gatherer.startGatheringData(c, resourceDataGatheringPeriodSeconds*time.Second)
	}

	if testContext.GatherLogsSizes {
		f.logsSizeWaitGroup = sync.WaitGroup{}
		f.logsSizeWaitGroup.Add(1)
		f.logsSizeCloseChannel = make(chan bool)
		f.logsSizeVerifier = NewLogsVerifier(c, f.logsSizeCloseChannel)
		go func() {
			f.logsSizeVerifier.Run()
			f.logsSizeWaitGroup.Done()
		}()
	}
}

// afterEach deletes the namespace, after reading its events.
func (f *Framework) afterEach() {
	RemoveCleanupAction(f.cleanupHandle)

	// DeleteNamespace at the very end in defer, to avoid any
	// expectation failures preventing deleting the namespace.
	defer func() {
		if testContext.DeleteNamespace {
			for _, ns := range f.namespacesToDelete {
				By(fmt.Sprintf("Destroying namespace %q for this suite.", ns.Name))

				timeout := 5 * time.Minute
				if f.NamespaceDeletionTimeout != 0 {
					timeout = f.NamespaceDeletionTimeout
				}
				if err := deleteNS(f.Client, ns.Name, timeout); err != nil {
					Failf("Couldn't delete ns %q: %s", ns.Name, err)
				}
			}
			f.namespacesToDelete = nil
		} else {
			Logf("Found DeleteNamespace=false, skipping namespace deletion!")
		}

		// Paranoia-- prevent reuse!
		f.Namespace = nil
		f.Client = nil
	}()

	// Print events if the test failed.
	if CurrentGinkgoTestDescription().Failed {
		By(fmt.Sprintf("Collecting events from namespace %q.", f.Namespace.Name))
		events, err := f.Client.Events(f.Namespace.Name).List(api.ListOptions{})
		Expect(err).NotTo(HaveOccurred())

		for _, e := range events.Items {
			Logf("event for %v: %v %v: %v", e.InvolvedObject.Name, e.Source, e.Reason, e.Message)
		}
		// Note that we don't wait for any cleanup to propagate, which means
		// that if you delete a bunch of pods right before ending your test,
		// you may or may not see the killing/deletion/cleanup events.

		dumpAllPodInfo(f.Client)

		dumpAllNodeInfo(f.Client)
	}

	summaries := make([]TestDataSummary, 0)
	if testContext.GatherKubeSystemResourceUsageData {
		summaries = append(summaries, f.gatherer.stopAndSummarize([]int{90, 99}, f.addonResourceConstraints))
	}

	if testContext.GatherLogsSizes {
		close(f.logsSizeCloseChannel)
		f.logsSizeWaitGroup.Wait()
		summaries = append(summaries, f.logsSizeVerifier.GetSummary())
	}

	if testContext.GatherMetricsAfterTest {
		// TODO: enable Scheduler and ControllerManager metrics grabbing when Master's Kubelet will be registered.
		grabber, err := metrics.NewMetricsGrabber(f.Client, true, false, false, true)
		if err != nil {
			Logf("Failed to create MetricsGrabber. Skipping metrics gathering.")
		} else {
			received, err := grabber.Grab(nil)
			if err != nil {
				Logf("MetricsGrabber failed grab metrics. Skipping metrics gathering.")
			} else {
				summaries = append(summaries, (*MetricsForE2E)(&received))
			}
		}
	}

	outputTypes := strings.Split(testContext.OutputPrintType, ",")
	for _, printType := range outputTypes {
		switch printType {
		case "hr":
			for i := range summaries {
				Logf(summaries[i].PrintHumanReadable())
			}
		case "json":
			for i := range summaries {
				typeName := reflect.TypeOf(summaries[i]).String()
				Logf("%v JSON\n%v", typeName[strings.LastIndex(typeName, ".")+1:len(typeName)], summaries[i].PrintJSON())
				Logf("Finished")
			}
		default:
			Logf("Unknown ouptut type: %v. Skipping.", printType)
		}
	}

	// Check whether all nodes are ready after the test.
	// This is explicitly done at the very end of the test, to avoid
	// e.g. not removing namespace in case of this failure.
	if err := allNodesReady(f.Client, time.Minute); err != nil {
		Failf("All nodes should be ready after test, %v", err)
	}
}

func (f *Framework) CreateNamespace(baseName string, labels map[string]string) (*api.Namespace, error) {
	ns, err := createTestingNS(baseName, f.Client, labels)
	if err == nil {
		f.namespacesToDelete = append(f.namespacesToDelete, ns)
	}
	return ns, err
}

// WaitForPodTerminated waits for the pod to be terminated with the given reason.
func (f *Framework) WaitForPodTerminated(podName, reason string) error {
	return waitForPodTerminatedInNamespace(f.Client, podName, reason, f.Namespace.Name)
}

// WaitForPodRunning waits for the pod to run in the namespace.
func (f *Framework) WaitForPodRunning(podName string) error {
	return waitForPodRunningInNamespace(f.Client, podName, f.Namespace.Name)
}

// WaitForPodRunningSlow waits for the pod to run in the namespace.
// It has a longer timeout then WaitForPodRunning (util.slowPodStartTimeout).
func (f *Framework) WaitForPodRunningSlow(podName string) error {
	return waitForPodRunningInNamespaceSlow(f.Client, podName, f.Namespace.Name)
}

// Runs the given pod and verifies that the output of exact container matches the desired output.
func (f *Framework) TestContainerOutput(scenarioName string, pod *api.Pod, containerIndex int, expectedOutput []string) {
	testContainerOutput(scenarioName, f.Client, pod, containerIndex, expectedOutput, f.Namespace.Name)
}

// Runs the given pod and verifies that the output of exact container matches the desired regexps.
func (f *Framework) TestContainerOutputRegexp(scenarioName string, pod *api.Pod, containerIndex int, expectedOutput []string) {
	testContainerOutputRegexp(scenarioName, f.Client, pod, containerIndex, expectedOutput, f.Namespace.Name)
}

// WaitForAnEndpoint waits for at least one endpoint to become available in the
// service's corresponding endpoints object.
func (f *Framework) WaitForAnEndpoint(serviceName string) error {
	for {
		// TODO: Endpoints client should take a field selector so we
		// don't have to list everything.
		list, err := f.Client.Endpoints(f.Namespace.Name).List(api.ListOptions{})
		if err != nil {
			return err
		}
		rv := list.ResourceVersion

		isOK := func(e *api.Endpoints) bool {
			return e.Name == serviceName && len(e.Subsets) > 0 && len(e.Subsets[0].Addresses) > 0
		}
		for i := range list.Items {
			if isOK(&list.Items[i]) {
				return nil
			}
		}

		options := api.ListOptions{
			FieldSelector:   fields.Set{"metadata.name": serviceName}.AsSelector(),
			ResourceVersion: rv,
		}
		w, err := f.Client.Endpoints(f.Namespace.Name).Watch(options)
		if err != nil {
			return err
		}
		defer w.Stop()

		for {
			val, ok := <-w.ResultChan()
			if !ok {
				// reget and re-watch
				break
			}
			if e, ok := val.Object.(*api.Endpoints); ok {
				if isOK(e) {
					return nil
				}
			}
		}
	}
}

// Write a file using kubectl exec echo <contents> > <path> via specified container
// Because of the primitive technique we're using here, we only allow ASCII alphanumeric characters
func (f *Framework) WriteFileViaContainer(podName, containerName string, path string, contents string) error {
	By("writing a file in the container")
	allowedCharacters := "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	for _, c := range contents {
		if !strings.ContainsRune(allowedCharacters, c) {
			return fmt.Errorf("Unsupported character in string to write: %v", c)
		}
	}
	command := fmt.Sprintf("echo '%s' > '%s'", contents, path)
	stdout, stderr, err := kubectlExecWithRetry(f.Namespace.Name, podName, containerName, "--", "/bin/sh", "-c", command)
	if err != nil {
		Logf("error running kubectl exec to write file: %v\nstdout=%v\nstderr=%v)", err, string(stdout), string(stderr))
	}
	return err
}

// Read a file using kubectl exec cat <path>
func (f *Framework) ReadFileViaContainer(podName, containerName string, path string) (string, error) {
	By("reading a file in the container")

	stdout, stderr, err := kubectlExecWithRetry(f.Namespace.Name, podName, containerName, "--", "cat", path)
	if err != nil {
		Logf("error running kubectl exec to read file: %v\nstdout=%v\nstderr=%v)", err, string(stdout), string(stderr))
	}
	return string(stdout), err
}

func kubectlExecWithRetry(namespace string, podName, containerName string, args ...string) ([]byte, []byte, error) {
	for numRetries := 0; numRetries < maxKubectlExecRetries; numRetries++ {
		if numRetries > 0 {
			Logf("Retrying kubectl exec (retry count=%v/%v)", numRetries+1, maxKubectlExecRetries)
		}

		stdOutBytes, stdErrBytes, err := kubectlExec(namespace, podName, containerName, args...)
		if err != nil {
			if strings.Contains(strings.ToLower(string(stdErrBytes)), "i/o timeout") {
				// Retry on "i/o timeout" errors
				Logf("Warning: kubectl exec encountered i/o timeout.\nerr=%v\nstdout=%v\nstderr=%v)", err, string(stdOutBytes), string(stdErrBytes))
				continue
			}
		}

		return stdOutBytes, stdErrBytes, err
	}
	err := fmt.Errorf("Failed: kubectl exec failed %d times with \"i/o timeout\". Giving up.", maxKubectlExecRetries)
	return nil, nil, err
}

func kubectlExec(namespace string, podName, containerName string, args ...string) ([]byte, []byte, error) {
	var stdout, stderr bytes.Buffer
	cmdArgs := []string{
		"exec",
		fmt.Sprintf("--namespace=%v", namespace),
		podName,
		fmt.Sprintf("-c=%v", containerName),
	}
	cmdArgs = append(cmdArgs, args...)

	cmd := kubectlCmd(cmdArgs...)
	cmd.Stdout, cmd.Stderr = &stdout, &stderr

	Logf("Running '%s %s'", cmd.Path, strings.Join(cmd.Args, " "))
	err := cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}
<commit_msg>Add logss to debug scalability failure<commit_after>/*
Copyright 2015 The Kubernetes Authors 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 e2e

import (
	"bytes"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/clientset_generated/release_1_2"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/metrics"

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

const (
	maxKubectlExecRetries = 5
)

// Framework supports common operations used by e2e tests; it will keep a client & a namespace for you.
// Eventual goal is to merge this with integration test framework.
type Framework struct {
	BaseName string

	Client        *client.Client
	Clientset_1_2 *release_1_2.Clientset

	Namespace                *api.Namespace   // Every test has at least one namespace
	namespacesToDelete       []*api.Namespace // Some tests have more than one.
	NamespaceDeletionTimeout time.Duration

	gatherer containerResourceGatherer
	// Constraints that passed to a check which is exectued after data is gathered to
	// see if 99% of results are within acceptable bounds. It as to be injected in the test,
	// as expectations vary greatly. Constraints are groupped by the container names.
	addonResourceConstraints map[string]resourceConstraint

	logsSizeWaitGroup    sync.WaitGroup
	logsSizeCloseChannel chan bool
	logsSizeVerifier     *LogsSizeVerifier

	// To make sure that this framework cleans up after itself, no matter what,
	// we install a cleanup action before each test and clear it after.  If we
	// should abort, the AfterSuite hook should run all cleanup actions.
	cleanupHandle CleanupActionHandle
}

type TestDataSummary interface {
	PrintHumanReadable() string
	PrintJSON() string
}

// NewFramework makes a new framework and sets up a BeforeEach/AfterEach for
// you (you can write additional before/after each functions).
func NewFramework(baseName string) *Framework {
	f := &Framework{
		BaseName:                 baseName,
		addonResourceConstraints: make(map[string]resourceConstraint),
	}

	BeforeEach(f.beforeEach)
	AfterEach(f.afterEach)

	return f
}

// beforeEach gets a client and makes a namespace.
func (f *Framework) beforeEach() {
	// The fact that we need this feels like a bug in ginkgo.
	// https://github.com/onsi/ginkgo/issues/222
	f.cleanupHandle = AddCleanupAction(f.afterEach)

	By("Creating a kubernetes client")
	c, err := loadClient()
	Expect(err).NotTo(HaveOccurred())

	f.Client = c
	f.Clientset_1_2 = release_1_2.FromUnversionedClient(c)

	By("Building a namespace api object")
	namespace, err := f.CreateNamespace(f.BaseName, map[string]string{
		"e2e-framework": f.BaseName,
	})
	Expect(err).NotTo(HaveOccurred())

	f.Namespace = namespace

	if testContext.VerifyServiceAccount {
		By("Waiting for a default service account to be provisioned in namespace")
		err = waitForDefaultServiceAccountInNamespace(c, namespace.Name)
		Expect(err).NotTo(HaveOccurred())
	} else {
		Logf("Skipping waiting for service account")
	}

	if testContext.GatherKubeSystemResourceUsageData {
		f.gatherer.startGatheringData(c, resourceDataGatheringPeriodSeconds*time.Second)
	}

	if testContext.GatherLogsSizes {
		f.logsSizeWaitGroup = sync.WaitGroup{}
		f.logsSizeWaitGroup.Add(1)
		f.logsSizeCloseChannel = make(chan bool)
		f.logsSizeVerifier = NewLogsVerifier(c, f.logsSizeCloseChannel)
		go func() {
			f.logsSizeVerifier.Run()
			f.logsSizeWaitGroup.Done()
		}()
	}
}

// afterEach deletes the namespace, after reading its events.
func (f *Framework) afterEach() {
	RemoveCleanupAction(f.cleanupHandle)

	// DeleteNamespace at the very end in defer, to avoid any
	// expectation failures preventing deleting the namespace.
	defer func() {
		if testContext.DeleteNamespace {
			for _, ns := range f.namespacesToDelete {
				By(fmt.Sprintf("Destroying namespace %q for this suite.", ns.Name))

				timeout := 5 * time.Minute
				if f.NamespaceDeletionTimeout != 0 {
					timeout = f.NamespaceDeletionTimeout
				}
				if err := deleteNS(f.Client, ns.Name, timeout); err != nil {
					Failf("Couldn't delete ns %q: %s", ns.Name, err)
				}
			}
			f.namespacesToDelete = nil
		} else {
			Logf("Found DeleteNamespace=false, skipping namespace deletion!")
		}

		// Paranoia-- prevent reuse!
		f.Namespace = nil
		f.Client = nil
	}()

	// Print events if the test failed.
	if CurrentGinkgoTestDescription().Failed {
		By(fmt.Sprintf("Collecting events from namespace %q.", f.Namespace.Name))
		events, err := f.Client.Events(f.Namespace.Name).List(api.ListOptions{})
		Expect(err).NotTo(HaveOccurred())

		for _, e := range events.Items {
			Logf("event for %v: %v %v: %v", e.InvolvedObject.Name, e.Source, e.Reason, e.Message)
		}
		// Note that we don't wait for any cleanup to propagate, which means
		// that if you delete a bunch of pods right before ending your test,
		// you may or may not see the killing/deletion/cleanup events.

		dumpAllPodInfo(f.Client)

		dumpAllNodeInfo(f.Client)
	}

	summaries := make([]TestDataSummary, 0)
	if testContext.GatherKubeSystemResourceUsageData {
		By("Collecting resource usage data")
		summaries = append(summaries, f.gatherer.stopAndSummarize([]int{90, 99}, f.addonResourceConstraints))
	}

	if testContext.GatherLogsSizes {
		By("Gathering log sizes data")
		close(f.logsSizeCloseChannel)
		f.logsSizeWaitGroup.Wait()
		summaries = append(summaries, f.logsSizeVerifier.GetSummary())
	}

	if testContext.GatherMetricsAfterTest {
		By("Gathering metrics")
		// TODO: enable Scheduler and ControllerManager metrics grabbing when Master's Kubelet will be registered.
		grabber, err := metrics.NewMetricsGrabber(f.Client, true, false, false, true)
		if err != nil {
			Logf("Failed to create MetricsGrabber. Skipping metrics gathering.")
		} else {
			received, err := grabber.Grab(nil)
			if err != nil {
				Logf("MetricsGrabber failed grab metrics. Skipping metrics gathering.")
			} else {
				summaries = append(summaries, (*MetricsForE2E)(&received))
			}
		}
	}

	outputTypes := strings.Split(testContext.OutputPrintType, ",")
	for _, printType := range outputTypes {
		switch printType {
		case "hr":
			for i := range summaries {
				Logf(summaries[i].PrintHumanReadable())
			}
		case "json":
			for i := range summaries {
				typeName := reflect.TypeOf(summaries[i]).String()
				Logf("%v JSON\n%v", typeName[strings.LastIndex(typeName, ".")+1:len(typeName)], summaries[i].PrintJSON())
				Logf("Finished")
			}
		default:
			Logf("Unknown ouptut type: %v. Skipping.", printType)
		}
	}

	// Check whether all nodes are ready after the test.
	// This is explicitly done at the very end of the test, to avoid
	// e.g. not removing namespace in case of this failure.
	if err := allNodesReady(f.Client, time.Minute); err != nil {
		Failf("All nodes should be ready after test, %v", err)
	}
}

func (f *Framework) CreateNamespace(baseName string, labels map[string]string) (*api.Namespace, error) {
	ns, err := createTestingNS(baseName, f.Client, labels)
	if err == nil {
		f.namespacesToDelete = append(f.namespacesToDelete, ns)
	}
	return ns, err
}

// WaitForPodTerminated waits for the pod to be terminated with the given reason.
func (f *Framework) WaitForPodTerminated(podName, reason string) error {
	return waitForPodTerminatedInNamespace(f.Client, podName, reason, f.Namespace.Name)
}

// WaitForPodRunning waits for the pod to run in the namespace.
func (f *Framework) WaitForPodRunning(podName string) error {
	return waitForPodRunningInNamespace(f.Client, podName, f.Namespace.Name)
}

// WaitForPodRunningSlow waits for the pod to run in the namespace.
// It has a longer timeout then WaitForPodRunning (util.slowPodStartTimeout).
func (f *Framework) WaitForPodRunningSlow(podName string) error {
	return waitForPodRunningInNamespaceSlow(f.Client, podName, f.Namespace.Name)
}

// Runs the given pod and verifies that the output of exact container matches the desired output.
func (f *Framework) TestContainerOutput(scenarioName string, pod *api.Pod, containerIndex int, expectedOutput []string) {
	testContainerOutput(scenarioName, f.Client, pod, containerIndex, expectedOutput, f.Namespace.Name)
}

// Runs the given pod and verifies that the output of exact container matches the desired regexps.
func (f *Framework) TestContainerOutputRegexp(scenarioName string, pod *api.Pod, containerIndex int, expectedOutput []string) {
	testContainerOutputRegexp(scenarioName, f.Client, pod, containerIndex, expectedOutput, f.Namespace.Name)
}

// WaitForAnEndpoint waits for at least one endpoint to become available in the
// service's corresponding endpoints object.
func (f *Framework) WaitForAnEndpoint(serviceName string) error {
	for {
		// TODO: Endpoints client should take a field selector so we
		// don't have to list everything.
		list, err := f.Client.Endpoints(f.Namespace.Name).List(api.ListOptions{})
		if err != nil {
			return err
		}
		rv := list.ResourceVersion

		isOK := func(e *api.Endpoints) bool {
			return e.Name == serviceName && len(e.Subsets) > 0 && len(e.Subsets[0].Addresses) > 0
		}
		for i := range list.Items {
			if isOK(&list.Items[i]) {
				return nil
			}
		}

		options := api.ListOptions{
			FieldSelector:   fields.Set{"metadata.name": serviceName}.AsSelector(),
			ResourceVersion: rv,
		}
		w, err := f.Client.Endpoints(f.Namespace.Name).Watch(options)
		if err != nil {
			return err
		}
		defer w.Stop()

		for {
			val, ok := <-w.ResultChan()
			if !ok {
				// reget and re-watch
				break
			}
			if e, ok := val.Object.(*api.Endpoints); ok {
				if isOK(e) {
					return nil
				}
			}
		}
	}
}

// Write a file using kubectl exec echo <contents> > <path> via specified container
// Because of the primitive technique we're using here, we only allow ASCII alphanumeric characters
func (f *Framework) WriteFileViaContainer(podName, containerName string, path string, contents string) error {
	By("writing a file in the container")
	allowedCharacters := "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	for _, c := range contents {
		if !strings.ContainsRune(allowedCharacters, c) {
			return fmt.Errorf("Unsupported character in string to write: %v", c)
		}
	}
	command := fmt.Sprintf("echo '%s' > '%s'", contents, path)
	stdout, stderr, err := kubectlExecWithRetry(f.Namespace.Name, podName, containerName, "--", "/bin/sh", "-c", command)
	if err != nil {
		Logf("error running kubectl exec to write file: %v\nstdout=%v\nstderr=%v)", err, string(stdout), string(stderr))
	}
	return err
}

// Read a file using kubectl exec cat <path>
func (f *Framework) ReadFileViaContainer(podName, containerName string, path string) (string, error) {
	By("reading a file in the container")

	stdout, stderr, err := kubectlExecWithRetry(f.Namespace.Name, podName, containerName, "--", "cat", path)
	if err != nil {
		Logf("error running kubectl exec to read file: %v\nstdout=%v\nstderr=%v)", err, string(stdout), string(stderr))
	}
	return string(stdout), err
}

func kubectlExecWithRetry(namespace string, podName, containerName string, args ...string) ([]byte, []byte, error) {
	for numRetries := 0; numRetries < maxKubectlExecRetries; numRetries++ {
		if numRetries > 0 {
			Logf("Retrying kubectl exec (retry count=%v/%v)", numRetries+1, maxKubectlExecRetries)
		}

		stdOutBytes, stdErrBytes, err := kubectlExec(namespace, podName, containerName, args...)
		if err != nil {
			if strings.Contains(strings.ToLower(string(stdErrBytes)), "i/o timeout") {
				// Retry on "i/o timeout" errors
				Logf("Warning: kubectl exec encountered i/o timeout.\nerr=%v\nstdout=%v\nstderr=%v)", err, string(stdOutBytes), string(stdErrBytes))
				continue
			}
		}

		return stdOutBytes, stdErrBytes, err
	}
	err := fmt.Errorf("Failed: kubectl exec failed %d times with \"i/o timeout\". Giving up.", maxKubectlExecRetries)
	return nil, nil, err
}

func kubectlExec(namespace string, podName, containerName string, args ...string) ([]byte, []byte, error) {
	var stdout, stderr bytes.Buffer
	cmdArgs := []string{
		"exec",
		fmt.Sprintf("--namespace=%v", namespace),
		podName,
		fmt.Sprintf("-c=%v", containerName),
	}
	cmdArgs = append(cmdArgs, args...)

	cmd := kubectlCmd(cmdArgs...)
	cmd.Stdout, cmd.Stderr = &stdout, &stderr

	Logf("Running '%s %s'", cmd.Path, strings.Join(cmd.Args, " "))
	err := cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}
<|endoftext|>


<commit_before>package vip_test

import (
	"testing"

	"github.com/luizbafilho/fusis/api/types"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/net"
	"github.com/luizbafilho/fusis/state/mocks"
	"github.com/luizbafilho/fusis/vip"
	"github.com/stretchr/testify/assert"
)

func TestVipsSync(t *testing.T) {
	s1 := types.Service{
		Name:     "test",
		Address:  "10.0.1.1",
		Port:     80,
		Mode:     "nat",
		Protocol: "tcp",
	}

	s2 := types.Service{
		Name:     "test2",
		Address:  "10.0.1.2",
		Port:     80,
		Protocol: "tcp",
		Mode:     "nat",
	}

	state := &mocks.State{}
	state.On("GetServices").Return([]types.Service{s1, s2})

	iface := "eth0"
	config := &config.BalancerConfig{
		Interfaces: config.Interfaces{
			Inbound:  iface,
			Outbound: iface,
		},
	}
	vipMngr, err := vip.New(config)
	assert.Nil(t, err)

	vipMngr.Sync(state)

	vips, err := net.GetFusisVipsIps(iface)
	assert.Nil(t, err)

	assert.Len(t, vips, 2)
	assert.Contains(t, vips, "10.0.1.1")
	assert.Contains(t, vips, "10.0.1.2")

	// Asserting remove
	state = &mocks.State{}
	state.On("GetServices").Return([]types.Service{s2})

	err = vipMngr.Sync(state)

	vips, err = net.GetFusisVipsIps(iface)
	assert.Nil(t, err)
	assert.Len(t, vips, 1)
	assert.Contains(t, vips, "10.0.1.2")

	net.DelVips(iface)
}
<commit_msg>Fix vip tests to use lo interface<commit_after>package vip_test

import (
	"testing"

	"github.com/luizbafilho/fusis/api/types"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/net"
	"github.com/luizbafilho/fusis/state/mocks"
	"github.com/luizbafilho/fusis/vip"
	"github.com/stretchr/testify/assert"
)

func TestVipsSync(t *testing.T) {
	s1 := types.Service{
		Name:     "test",
		Address:  "10.0.1.1",
		Port:     80,
		Mode:     "nat",
		Protocol: "tcp",
	}

	s2 := types.Service{
		Name:     "test2",
		Address:  "10.0.1.2",
		Port:     80,
		Protocol: "tcp",
		Mode:     "nat",
	}

	state := &mocks.State{}
	state.On("GetServices").Return([]types.Service{s1, s2})

	iface := "lo"
	config := &config.BalancerConfig{
		Interfaces: config.Interfaces{
			Inbound:  iface,
			Outbound: iface,
		},
	}
	vipMngr, err := vip.New(config)
	assert.Nil(t, err)

	vipMngr.Sync(state)

	vips, err := net.GetFusisVipsIps(iface)
	assert.Nil(t, err)

	assert.Len(t, vips, 2)
	assert.Contains(t, vips, "10.0.1.1")
	assert.Contains(t, vips, "10.0.1.2")

	// Asserting remove
	state = &mocks.State{}
	state.On("GetServices").Return([]types.Service{s2})

	err = vipMngr.Sync(state)

	vips, err = net.GetFusisVipsIps(iface)
	assert.Nil(t, err)
	assert.Len(t, vips, 1)
	assert.Contains(t, vips, "10.0.1.2")

	net.DelVips(iface)
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package qemu

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	. "github.com/google/syzkaller/log"
	"github.com/google/syzkaller/vm"
)

const (
	hostAddr = "10.0.2.10"
)

func init() {
	vm.Register("qemu", ctor)
}

type instance struct {
	cfg     *vm.Config
	port    int
	rpipe   io.ReadCloser
	wpipe   io.WriteCloser
	qemu    *exec.Cmd
	waiterC chan error
	merger  *vm.OutputMerger
}

func ctor(cfg *vm.Config) (vm.Instance, error) {
	for i := 0; ; i++ {
		inst, err := ctorImpl(cfg)
		if err == nil {
			return inst, nil
		}
		if i < 1000 && strings.Contains(err.Error(), "could not set up host forwarding rule") {
			continue
		}
		os.RemoveAll(cfg.Workdir)
		return nil, err
	}
}

func ctorImpl(cfg *vm.Config) (vm.Instance, error) {
	inst := &instance{cfg: cfg}
	closeInst := inst
	defer func() {
		if closeInst != nil {
			closeInst.close(false)
		}
	}()

	if err := validateConfig(cfg); err != nil {
		return nil, err
	}

	if cfg.Image == "9p" {
		inst.cfg.Sshkey = filepath.Join(inst.cfg.Workdir, "key")
		keygen := exec.Command("ssh-keygen", "-t", "rsa", "-b", "2048", "-N", "", "-C", "", "-f", inst.cfg.Sshkey)
		if out, err := keygen.CombinedOutput(); err != nil {
			return nil, fmt.Errorf("failed to execute ssh-keygen: %v\n%s", err, out)
		}
		initFile := filepath.Join(cfg.Workdir, "init.sh")
		if err := ioutil.WriteFile(initFile, []byte(strings.Replace(initScript, "{{KEY}}", inst.cfg.Sshkey, -1)), 0777); err != nil {
			return nil, fmt.Errorf("failed to create init file: %v", err)
		}
	}

	var err error
	inst.rpipe, inst.wpipe, err = vm.LongPipe()
	if err != nil {
		return nil, err
	}

	if err := inst.Boot(); err != nil {
		return nil, err
	}

	closeInst = nil
	return inst, nil
}

func validateConfig(cfg *vm.Config) error {
	if cfg.Bin == "" {
		cfg.Bin = "qemu-system-x86_64"
	}
	if cfg.Image == "9p" {
		if cfg.Kernel == "" {
			return fmt.Errorf("9p image requires kernel")
		}
	} else {
		if _, err := os.Stat(cfg.Image); err != nil {
			return fmt.Errorf("image file '%v' does not exist: %v", cfg.Image, err)
		}
		if _, err := os.Stat(cfg.Sshkey); err != nil {
			return fmt.Errorf("ssh key '%v' does not exist: %v", cfg.Sshkey, err)
		}
	}
	if cfg.Cpu <= 0 || cfg.Cpu > 1024 {
		return fmt.Errorf("bad qemu cpu: %v, want [1-1024]", cfg.Cpu)
	}
	if cfg.Mem < 128 || cfg.Mem > 1048576 {
		return fmt.Errorf("bad qemu mem: %v, want [128-1048576]", cfg.Mem)
	}
	return nil
}

func (inst *instance) Close() {
	inst.close(true)
}

func (inst *instance) close(removeWorkDir bool) {
	if inst.qemu != nil {
		inst.qemu.Process.Kill()
		err := <-inst.waiterC
		inst.waiterC <- err // repost it for waiting goroutines
	}
	if inst.merger != nil {
		inst.merger.Wait()
	}
	if inst.rpipe != nil {
		inst.rpipe.Close()
	}
	if inst.wpipe != nil {
		inst.wpipe.Close()
	}
	os.Remove(filepath.Join(inst.cfg.Workdir, "key"))
	if removeWorkDir {
		os.RemoveAll(inst.cfg.Workdir)
	}
}

func (inst *instance) Boot() error {
	for {
		// Find an unused TCP port.
		inst.port = rand.Intn(64<<10-1<<10) + 1<<10
		ln, err := net.Listen("tcp", fmt.Sprintf("localhost:%v", inst.port))
		if err == nil {
			ln.Close()
			break
		}
	}
	// TODO: ignores inst.cfg.Cpu
	args := []string{
		"-m", strconv.Itoa(inst.cfg.Mem),
		"-net", "nic",
		"-net", fmt.Sprintf("user,host=%v,hostfwd=tcp::%v-:22", hostAddr, inst.port),
		"-display", "none",
		"-serial", "stdio",
		"-no-reboot",
		"-numa", "node,nodeid=0,cpus=0-1", "-numa", "node,nodeid=1,cpus=2-3",
		"-smp", "sockets=2,cores=2,threads=1",
	}
	if inst.cfg.BinArgs == "" {
		// This is reasonable defaults for x86 kvm-enabled host.
		args = append(args,
			"-enable-kvm",
			"-usb", "-usbdevice", "mouse", "-usbdevice", "tablet",
			"-soundhw", "all",
		)
	} else {
		args = append(args, strings.Split(inst.cfg.BinArgs, " ")...)
	}
	if inst.cfg.Image == "9p" {
		args = append(args,
			"-fsdev", "local,id=fsdev0,path=/,security_model=none,readonly",
			"-device", "virtio-9p-pci,fsdev=fsdev0,mount_tag=/dev/root",
		)
	} else {
		args = append(args,
			"-hda", inst.cfg.Image,
			"-snapshot",
		)
	}
	if inst.cfg.Initrd != "" {
		args = append(args,
			"-initrd", inst.cfg.Initrd,
		)
	}
	if inst.cfg.Kernel != "" {
		cmdline := "console=ttyS0 vsyscall=native rodata=n oops=panic panic_on_warn=1 panic=-1" +
			" ftrace_dump_on_oops=orig_cpu earlyprintk=serial slub_debug=UZ "
		if inst.cfg.Image == "9p" {
			cmdline += "root=/dev/root rootfstype=9p rootflags=trans=virtio,version=9p2000.L,cache=loose "
			cmdline += "init=" + filepath.Join(inst.cfg.Workdir, "init.sh") + " "
		} else {
			cmdline += "root=/dev/sda "
		}
		args = append(args,
			"-kernel", inst.cfg.Kernel,
			"-append", cmdline+inst.cfg.Cmdline,
		)
	}
	if inst.cfg.Debug {
		Logf(0, "running command: %v %#v", inst.cfg.Bin, args)
	}
	qemu := exec.Command(inst.cfg.Bin, args...)
	qemu.Stdout = inst.wpipe
	qemu.Stderr = inst.wpipe
	if err := qemu.Start(); err != nil {
		return fmt.Errorf("failed to start %v %+v: %v", inst.cfg.Bin, args, err)
	}
	inst.wpipe.Close()
	inst.wpipe = nil
	inst.qemu = qemu
	// Qemu has started.

	// Start output merger.
	var tee io.Writer
	if inst.cfg.Debug {
		tee = os.Stdout
	}
	inst.merger = vm.NewOutputMerger(tee)
	inst.merger.Add(inst.rpipe)
	inst.rpipe = nil

	var bootOutput []byte
	bootOutputStop := make(chan bool)
	go func() {
		for {
			select {
			case out := <-inst.merger.Output:
				bootOutput = append(bootOutput, out...)
			case <-bootOutputStop:
				close(bootOutputStop)
				return
			}
		}
	}()

	// Wait for the qemu asynchronously.
	inst.waiterC = make(chan error, 1)
	go func() {
		err := qemu.Wait()
		inst.waiterC <- err
	}()

	// Wait for ssh server to come up.
	time.Sleep(10 * time.Second)
	start := time.Now()
	for {
		c, err := net.DialTimeout("tcp", fmt.Sprintf("localhost:%v", inst.port), 3*time.Second)
		if err == nil {
			c.SetDeadline(time.Now().Add(3 * time.Second))
			var tmp [1]byte
			n, err := c.Read(tmp[:])
			c.Close()
			if err == nil && n > 0 {
				break // ssh is up and responding
			}
			time.Sleep(3 * time.Second)
		}
		select {
		case err := <-inst.waiterC:
			inst.waiterC <- err     // repost it for Close
			time.Sleep(time.Second) // wait for any pending output
			bootOutputStop <- true
			<-bootOutputStop
			return fmt.Errorf("qemu stopped:\n%v\n", string(bootOutput))
		default:
		}
		if time.Since(start) > 10*time.Minute {
			bootOutputStop <- true
			<-bootOutputStop
			return fmt.Errorf("ssh server did not start:\n%v\n", string(bootOutput))
		}
	}
	bootOutputStop <- true
	return nil
}

func (inst *instance) Forward(port int) (string, error) {
	return fmt.Sprintf("%v:%v", hostAddr, port), nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	basePath := "/"
	if inst.cfg.Image == "9p" {
		basePath = "/tmp"
	}
	vmDst := filepath.Join(basePath, filepath.Base(hostSrc))
	args := append(inst.sshArgs("-P"), hostSrc, "root@localhost:"+vmDst)
	cmd := exec.Command("scp", args...)
	if inst.cfg.Debug {
		Logf(0, "running command: scp %#v", args)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stdout
	}
	if err := cmd.Start(); err != nil {
		return "", err
	}
	done := make(chan bool)
	go func() {
		select {
		case <-time.After(3 * time.Minute):
			cmd.Process.Kill()
		case <-done:
		}
	}()
	err := cmd.Wait()
	close(done)
	if err != nil {
		return "", err
	}
	return vmDst, nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (<-chan []byte, <-chan error, error) {
	rpipe, wpipe, err := vm.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	inst.merger.Add(rpipe)

	args := append(inst.sshArgs("-p"), "root@localhost", command)
	if inst.cfg.Debug {
		Logf(0, "running command: ssh %#v", args)
	}
	cmd := exec.Command("ssh", args...)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		return nil, nil, err
	}
	wpipe.Close()
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	done := make(chan bool)
	go func() {
		select {
		case <-time.After(timeout):
			signal(vm.TimeoutErr)
			cmd.Process.Kill()
		case <-stop:
			signal(vm.TimeoutErr)
			cmd.Process.Kill()
		case <-done:
		}
	}()
	go func() {
		err := cmd.Wait()
		close(done)
		signal(err)
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) sshArgs(portArg string) []string {
	args := []string{
		"-i", inst.cfg.Sshkey,
		portArg, strconv.Itoa(inst.port),
		"-F", "/dev/null",
		"-o", "ConnectionAttempts=10",
		"-o", "ConnectTimeout=10",
		"-o", "BatchMode=yes",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "IdentitiesOnly=yes",
		"-o", "StrictHostKeyChecking=no",
		"-o", "LogLevel=error",
	}
	if inst.cfg.Debug {
		args = append(args, "-v")
	}
	return args
}

const initScript = `#! /bin/bash
set -eux
mount -t proc none /proc
mount -t sysfs none /sys
mount -t debugfs nodev /sys/kernel/debug/
mount -t tmpfs none /tmp
mount -t tmpfs none /var
mount -t tmpfs none /etc
mount -t tmpfs none /root
touch /etc/fstab
echo "root::0:0:root:/root:/bin/bash" > /etc/passwd
mkdir -p /etc/ssh
cp {{KEY}}.pub /root/
chmod 0700 /root
chmod 0600 /root/key.pub
mkdir -p /var/run/sshd/
chmod 700 /var/run/sshd
cat > /etc/ssh/sshd_config <<EOF
          Port 22
          Protocol 2
          UsePrivilegeSeparation no
          HostKey {{KEY}}
          PermitRootLogin yes
          AuthenticationMethods publickey
          ChallengeResponseAuthentication no
          AuthorizedKeysFile /root/key.pub
          IgnoreUserKnownHosts yes
          AllowUsers root
          LogLevel INFO
          TCPKeepAlive yes
          RSAAuthentication yes
          PubkeyAuthentication yes
EOF
/sbin/dhclient eth0
/usr/sbin/sshd -e -D
/sbin/halt -f
`
<commit_msg>vm/qemu: prevent eth0 renaming<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package qemu

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	. "github.com/google/syzkaller/log"
	"github.com/google/syzkaller/vm"
)

const (
	hostAddr = "10.0.2.10"
)

func init() {
	vm.Register("qemu", ctor)
}

type instance struct {
	cfg     *vm.Config
	port    int
	rpipe   io.ReadCloser
	wpipe   io.WriteCloser
	qemu    *exec.Cmd
	waiterC chan error
	merger  *vm.OutputMerger
}

func ctor(cfg *vm.Config) (vm.Instance, error) {
	for i := 0; ; i++ {
		inst, err := ctorImpl(cfg)
		if err == nil {
			return inst, nil
		}
		if i < 1000 && strings.Contains(err.Error(), "could not set up host forwarding rule") {
			continue
		}
		os.RemoveAll(cfg.Workdir)
		return nil, err
	}
}

func ctorImpl(cfg *vm.Config) (vm.Instance, error) {
	inst := &instance{cfg: cfg}
	closeInst := inst
	defer func() {
		if closeInst != nil {
			closeInst.close(false)
		}
	}()

	if err := validateConfig(cfg); err != nil {
		return nil, err
	}

	if cfg.Image == "9p" {
		inst.cfg.Sshkey = filepath.Join(inst.cfg.Workdir, "key")
		keygen := exec.Command("ssh-keygen", "-t", "rsa", "-b", "2048", "-N", "", "-C", "", "-f", inst.cfg.Sshkey)
		if out, err := keygen.CombinedOutput(); err != nil {
			return nil, fmt.Errorf("failed to execute ssh-keygen: %v\n%s", err, out)
		}
		initFile := filepath.Join(cfg.Workdir, "init.sh")
		if err := ioutil.WriteFile(initFile, []byte(strings.Replace(initScript, "{{KEY}}", inst.cfg.Sshkey, -1)), 0777); err != nil {
			return nil, fmt.Errorf("failed to create init file: %v", err)
		}
	}

	var err error
	inst.rpipe, inst.wpipe, err = vm.LongPipe()
	if err != nil {
		return nil, err
	}

	if err := inst.Boot(); err != nil {
		return nil, err
	}

	closeInst = nil
	return inst, nil
}

func validateConfig(cfg *vm.Config) error {
	if cfg.Bin == "" {
		cfg.Bin = "qemu-system-x86_64"
	}
	if cfg.Image == "9p" {
		if cfg.Kernel == "" {
			return fmt.Errorf("9p image requires kernel")
		}
	} else {
		if _, err := os.Stat(cfg.Image); err != nil {
			return fmt.Errorf("image file '%v' does not exist: %v", cfg.Image, err)
		}
		if _, err := os.Stat(cfg.Sshkey); err != nil {
			return fmt.Errorf("ssh key '%v' does not exist: %v", cfg.Sshkey, err)
		}
	}
	if cfg.Cpu <= 0 || cfg.Cpu > 1024 {
		return fmt.Errorf("bad qemu cpu: %v, want [1-1024]", cfg.Cpu)
	}
	if cfg.Mem < 128 || cfg.Mem > 1048576 {
		return fmt.Errorf("bad qemu mem: %v, want [128-1048576]", cfg.Mem)
	}
	return nil
}

func (inst *instance) Close() {
	inst.close(true)
}

func (inst *instance) close(removeWorkDir bool) {
	if inst.qemu != nil {
		inst.qemu.Process.Kill()
		err := <-inst.waiterC
		inst.waiterC <- err // repost it for waiting goroutines
	}
	if inst.merger != nil {
		inst.merger.Wait()
	}
	if inst.rpipe != nil {
		inst.rpipe.Close()
	}
	if inst.wpipe != nil {
		inst.wpipe.Close()
	}
	os.Remove(filepath.Join(inst.cfg.Workdir, "key"))
	if removeWorkDir {
		os.RemoveAll(inst.cfg.Workdir)
	}
}

func (inst *instance) Boot() error {
	for {
		// Find an unused TCP port.
		inst.port = rand.Intn(64<<10-1<<10) + 1<<10
		ln, err := net.Listen("tcp", fmt.Sprintf("localhost:%v", inst.port))
		if err == nil {
			ln.Close()
			break
		}
	}
	// TODO: ignores inst.cfg.Cpu
	args := []string{
		"-m", strconv.Itoa(inst.cfg.Mem),
		"-net", "nic",
		"-net", fmt.Sprintf("user,host=%v,hostfwd=tcp::%v-:22", hostAddr, inst.port),
		"-display", "none",
		"-serial", "stdio",
		"-no-reboot",
		"-numa", "node,nodeid=0,cpus=0-1", "-numa", "node,nodeid=1,cpus=2-3",
		"-smp", "sockets=2,cores=2,threads=1",
	}
	if inst.cfg.BinArgs == "" {
		// This is reasonable defaults for x86 kvm-enabled host.
		args = append(args,
			"-enable-kvm",
			"-usb", "-usbdevice", "mouse", "-usbdevice", "tablet",
			"-soundhw", "all",
		)
	} else {
		args = append(args, strings.Split(inst.cfg.BinArgs, " ")...)
	}
	if inst.cfg.Image == "9p" {
		args = append(args,
			"-fsdev", "local,id=fsdev0,path=/,security_model=none,readonly",
			"-device", "virtio-9p-pci,fsdev=fsdev0,mount_tag=/dev/root",
		)
	} else {
		args = append(args,
			"-hda", inst.cfg.Image,
			"-snapshot",
		)
	}
	if inst.cfg.Initrd != "" {
		args = append(args,
			"-initrd", inst.cfg.Initrd,
		)
	}
	if inst.cfg.Kernel != "" {
		cmdline := "console=ttyS0 vsyscall=native rodata=n oops=panic panic_on_warn=1 panic=-1" +
			" ftrace_dump_on_oops=orig_cpu earlyprintk=serial slub_debug=UZ net.ifnames=0 biosdevname=0 "
		if inst.cfg.Image == "9p" {
			cmdline += "root=/dev/root rootfstype=9p rootflags=trans=virtio,version=9p2000.L,cache=loose "
			cmdline += "init=" + filepath.Join(inst.cfg.Workdir, "init.sh") + " "
		} else {
			cmdline += "root=/dev/sda "
		}
		args = append(args,
			"-kernel", inst.cfg.Kernel,
			"-append", cmdline+inst.cfg.Cmdline,
		)
	}
	if inst.cfg.Debug {
		Logf(0, "running command: %v %#v", inst.cfg.Bin, args)
	}
	qemu := exec.Command(inst.cfg.Bin, args...)
	qemu.Stdout = inst.wpipe
	qemu.Stderr = inst.wpipe
	if err := qemu.Start(); err != nil {
		return fmt.Errorf("failed to start %v %+v: %v", inst.cfg.Bin, args, err)
	}
	inst.wpipe.Close()
	inst.wpipe = nil
	inst.qemu = qemu
	// Qemu has started.

	// Start output merger.
	var tee io.Writer
	if inst.cfg.Debug {
		tee = os.Stdout
	}
	inst.merger = vm.NewOutputMerger(tee)
	inst.merger.Add(inst.rpipe)
	inst.rpipe = nil

	var bootOutput []byte
	bootOutputStop := make(chan bool)
	go func() {
		for {
			select {
			case out := <-inst.merger.Output:
				bootOutput = append(bootOutput, out...)
			case <-bootOutputStop:
				close(bootOutputStop)
				return
			}
		}
	}()

	// Wait for the qemu asynchronously.
	inst.waiterC = make(chan error, 1)
	go func() {
		err := qemu.Wait()
		inst.waiterC <- err
	}()

	// Wait for ssh server to come up.
	time.Sleep(10 * time.Second)
	start := time.Now()
	for {
		c, err := net.DialTimeout("tcp", fmt.Sprintf("localhost:%v", inst.port), 3*time.Second)
		if err == nil {
			c.SetDeadline(time.Now().Add(3 * time.Second))
			var tmp [1]byte
			n, err := c.Read(tmp[:])
			c.Close()
			if err == nil && n > 0 {
				break // ssh is up and responding
			}
			time.Sleep(3 * time.Second)
		}
		select {
		case err := <-inst.waiterC:
			inst.waiterC <- err     // repost it for Close
			time.Sleep(time.Second) // wait for any pending output
			bootOutputStop <- true
			<-bootOutputStop
			return fmt.Errorf("qemu stopped:\n%v\n", string(bootOutput))
		default:
		}
		if time.Since(start) > 10*time.Minute {
			bootOutputStop <- true
			<-bootOutputStop
			return fmt.Errorf("ssh server did not start:\n%v\n", string(bootOutput))
		}
	}
	bootOutputStop <- true
	return nil
}

func (inst *instance) Forward(port int) (string, error) {
	return fmt.Sprintf("%v:%v", hostAddr, port), nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	basePath := "/"
	if inst.cfg.Image == "9p" {
		basePath = "/tmp"
	}
	vmDst := filepath.Join(basePath, filepath.Base(hostSrc))
	args := append(inst.sshArgs("-P"), hostSrc, "root@localhost:"+vmDst)
	cmd := exec.Command("scp", args...)
	if inst.cfg.Debug {
		Logf(0, "running command: scp %#v", args)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stdout
	}
	if err := cmd.Start(); err != nil {
		return "", err
	}
	done := make(chan bool)
	go func() {
		select {
		case <-time.After(3 * time.Minute):
			cmd.Process.Kill()
		case <-done:
		}
	}()
	err := cmd.Wait()
	close(done)
	if err != nil {
		return "", err
	}
	return vmDst, nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (<-chan []byte, <-chan error, error) {
	rpipe, wpipe, err := vm.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	inst.merger.Add(rpipe)

	args := append(inst.sshArgs("-p"), "root@localhost", command)
	if inst.cfg.Debug {
		Logf(0, "running command: ssh %#v", args)
	}
	cmd := exec.Command("ssh", args...)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		return nil, nil, err
	}
	wpipe.Close()
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	done := make(chan bool)
	go func() {
		select {
		case <-time.After(timeout):
			signal(vm.TimeoutErr)
			cmd.Process.Kill()
		case <-stop:
			signal(vm.TimeoutErr)
			cmd.Process.Kill()
		case <-done:
		}
	}()
	go func() {
		err := cmd.Wait()
		close(done)
		signal(err)
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) sshArgs(portArg string) []string {
	args := []string{
		"-i", inst.cfg.Sshkey,
		portArg, strconv.Itoa(inst.port),
		"-F", "/dev/null",
		"-o", "ConnectionAttempts=10",
		"-o", "ConnectTimeout=10",
		"-o", "BatchMode=yes",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "IdentitiesOnly=yes",
		"-o", "StrictHostKeyChecking=no",
		"-o", "LogLevel=error",
	}
	if inst.cfg.Debug {
		args = append(args, "-v")
	}
	return args
}

const initScript = `#! /bin/bash
set -eux
mount -t proc none /proc
mount -t sysfs none /sys
mount -t debugfs nodev /sys/kernel/debug/
mount -t tmpfs none /tmp
mount -t tmpfs none /var
mount -t tmpfs none /etc
mount -t tmpfs none /root
touch /etc/fstab
echo "root::0:0:root:/root:/bin/bash" > /etc/passwd
mkdir -p /etc/ssh
cp {{KEY}}.pub /root/
chmod 0700 /root
chmod 0600 /root/key.pub
mkdir -p /var/run/sshd/
chmod 700 /var/run/sshd
cat > /etc/ssh/sshd_config <<EOF
          Port 22
          Protocol 2
          UsePrivilegeSeparation no
          HostKey {{KEY}}
          PermitRootLogin yes
          AuthenticationMethods publickey
          ChallengeResponseAuthentication no
          AuthorizedKeysFile /root/key.pub
          IgnoreUserKnownHosts yes
          AllowUsers root
          LogLevel INFO
          TCPKeepAlive yes
          RSAAuthentication yes
          PubkeyAuthentication yes
EOF
/sbin/dhclient eth0
/usr/sbin/sshd -e -D
/sbin/halt -f
`
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/motemen/ghq/logger"
)

type LocalRepository struct {
	FullPath  string
	RelPath   string
	PathParts []string
}

func LocalRepositoryFromFullPath(fullPath string) (*LocalRepository, error) {
	var relPath string

	for _, root := range localRepositoryRoots() {
		if !strings.HasPrefix(fullPath, root) {
			continue
		}

		var err error
		relPath, err = filepath.Rel(root, fullPath)
		if err == nil {
			break
		}
	}

	if relPath == "" {
		return nil, fmt.Errorf("no local repository found for: %s", fullPath)
	}

	pathParts := strings.Split(relPath, string(filepath.Separator))

	return &LocalRepository{fullPath, filepath.ToSlash(relPath), pathParts}, nil
}

func LocalRepositoryFromURL(remoteURL *url.URL) *LocalRepository {
	pathParts := append(
		[]string{remoteURL.Host}, strings.Split(remoteURL.Path, "/")...,
	)
	relPath := strings.TrimSuffix(path.Join(pathParts...), ".git")

	var localRepository *LocalRepository

	// Find existing local repository first
	walkLocalRepositories(func(repo *LocalRepository) {
		if repo.RelPath == relPath {
			localRepository = repo
		}
	})

	if localRepository != nil {
		return localRepository
	}

	// No local repository found, returning new one
	return &LocalRepository{
		path.Join(primaryLocalRepositoryRoot(), relPath),
		relPath,
		pathParts,
	}
}

// Subpaths returns lists of tail parts of relative path from the root directory (shortest first)
// for example, {"ghq", "motemen/ghq", "github.com/motemen/ghq"} for $root/github.com/motemen/ghq.
func (repo *LocalRepository) Subpaths() []string {
	tails := make([]string, len(repo.PathParts))

	for i := range repo.PathParts {
		tails[i] = strings.Join(repo.PathParts[len(repo.PathParts)-(i+1):], "/")
	}

	return tails
}

func (repo *LocalRepository) NonHostPath() string {
	return strings.Join(repo.PathParts[1:], "/")
}

func (repo *LocalRepository) IsUnderPrimaryRoot() bool {
	return strings.HasPrefix(repo.FullPath, primaryLocalRepositoryRoot())
}

// Matches checks if any subpath of the local repository equals the query.
func (repo *LocalRepository) Matches(pathQuery string) bool {
	for _, p := range repo.Subpaths() {
		if p == pathQuery {
			return true
		}
	}

	return false
}

// TODO return err
func (repo *LocalRepository) VCS() *VCSBackend {
	var (
		fi  os.FileInfo
		err error
	)

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".git/svn"))
	if err == nil && fi.IsDir() {
		return GitsvnBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".git"))
	if err == nil && fi.IsDir() {
		return GitBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".svn"))
	if err == nil && fi.IsDir() {
		return SubversionBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".hg"))
	if err == nil && fi.IsDir() {
		return MercurialBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, "_darcs"))
	if err == nil && fi.IsDir() {
		return DarcsBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".fslckout"))
	if err == nil && fi.IsDir() {
		return FossilBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, "_FOSSIL_"))
	if err == nil && fi.IsDir() {
		return FossilBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, "CVS"))
	if err == nil && fi.IsDir() {
		return cvsDummyBackend
	}
	return nil
}

var vcsDirs = []string{".git", ".svn", ".hg", "_darcs", ".fslckout", "_FOSSIL_", "CVS"}

func walkLocalRepositories(callback func(*LocalRepository)) {
	for _, root := range localRepositoryRoots() {
		filepath.Walk(root, func(path string, fileInfo os.FileInfo, err error) error {
			if err != nil || fileInfo == nil {
				return nil
			}

			if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
				realpath, err := filepath.EvalSymlinks(path)
				if err != nil {
					return nil
				}
				fileInfo, err = os.Stat(realpath)
				if err != nil {
					return nil
				}
			}
			if !fileInfo.IsDir() {
				return nil
			}

			vcsDirFound := false
			for _, d := range vcsDirs {
				_, err := os.Stat(filepath.Join(path, d))
				if err == nil {
					vcsDirFound = true
					break
				}
			}

			if !vcsDirFound {
				return nil
			}

			repo, err := LocalRepositoryFromFullPath(path)
			if err != nil {
				return nil
			}

			if repo == nil {
				return nil
			}
			callback(repo)
			return filepath.SkipDir
		})
	}
}

var _localRepositoryRoots []string

// localRepositoryRoots returns locally cloned repositories' root directories.
// The root dirs are determined as following:
//
//   - If GHQ_ROOT environment variable is nonempty, use it as the only root dir.
//   - Otherwise, use the result of `git config --get-all ghq.root` as the dirs.
//   - Otherwise, fallback to the default root, `~/.ghq`.
//
// TODO: More fancy default directory path?
func localRepositoryRoots() []string {
	if len(_localRepositoryRoots) != 0 {
		return _localRepositoryRoots
	}

	envRoot := os.Getenv("GHQ_ROOT")
	if envRoot != "" {
		_localRepositoryRoots = filepath.SplitList(envRoot)
	} else {
		var err error
		_localRepositoryRoots, err = GitConfigAll("ghq.root")
		logger.PanicIf(err)
	}

	if len(_localRepositoryRoots) == 0 {
		homeDir, err := os.UserHomeDir()
		logger.PanicIf(err)

		_localRepositoryRoots = []string{filepath.Join(homeDir, ".ghq")}
	}

	for i, v := range _localRepositoryRoots {
		path := filepath.Clean(v)
		if _, err := os.Stat(path); err == nil {
			_localRepositoryRoots[i], err = filepath.EvalSymlinks(path)
			logger.PanicIf(err)
		} else {
			_localRepositoryRoots[i] = path
		}
	}

	return _localRepositoryRoots
}

func primaryLocalRepositoryRoot() string {
	return localRepositoryRoots()[0]
}
<commit_msg>canonicalize root paths to absolute<commit_after>package main

import (
	"fmt"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/motemen/ghq/logger"
)

type LocalRepository struct {
	FullPath  string
	RelPath   string
	PathParts []string
}

func LocalRepositoryFromFullPath(fullPath string) (*LocalRepository, error) {
	var relPath string

	for _, root := range localRepositoryRoots() {
		if !strings.HasPrefix(fullPath, root) {
			continue
		}

		var err error
		relPath, err = filepath.Rel(root, fullPath)
		if err == nil {
			break
		}
	}

	if relPath == "" {
		return nil, fmt.Errorf("no local repository found for: %s", fullPath)
	}

	pathParts := strings.Split(relPath, string(filepath.Separator))

	return &LocalRepository{fullPath, filepath.ToSlash(relPath), pathParts}, nil
}

func LocalRepositoryFromURL(remoteURL *url.URL) *LocalRepository {
	pathParts := append(
		[]string{remoteURL.Host}, strings.Split(remoteURL.Path, "/")...,
	)
	relPath := strings.TrimSuffix(path.Join(pathParts...), ".git")

	var localRepository *LocalRepository

	// Find existing local repository first
	walkLocalRepositories(func(repo *LocalRepository) {
		if repo.RelPath == relPath {
			localRepository = repo
		}
	})

	if localRepository != nil {
		return localRepository
	}

	// No local repository found, returning new one
	return &LocalRepository{
		path.Join(primaryLocalRepositoryRoot(), relPath),
		relPath,
		pathParts,
	}
}

// Subpaths returns lists of tail parts of relative path from the root directory (shortest first)
// for example, {"ghq", "motemen/ghq", "github.com/motemen/ghq"} for $root/github.com/motemen/ghq.
func (repo *LocalRepository) Subpaths() []string {
	tails := make([]string, len(repo.PathParts))

	for i := range repo.PathParts {
		tails[i] = strings.Join(repo.PathParts[len(repo.PathParts)-(i+1):], "/")
	}

	return tails
}

func (repo *LocalRepository) NonHostPath() string {
	return strings.Join(repo.PathParts[1:], "/")
}

func (repo *LocalRepository) IsUnderPrimaryRoot() bool {
	return strings.HasPrefix(repo.FullPath, primaryLocalRepositoryRoot())
}

// Matches checks if any subpath of the local repository equals the query.
func (repo *LocalRepository) Matches(pathQuery string) bool {
	for _, p := range repo.Subpaths() {
		if p == pathQuery {
			return true
		}
	}

	return false
}

// TODO return err
func (repo *LocalRepository) VCS() *VCSBackend {
	var (
		fi  os.FileInfo
		err error
	)

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".git/svn"))
	if err == nil && fi.IsDir() {
		return GitsvnBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".git"))
	if err == nil && fi.IsDir() {
		return GitBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".svn"))
	if err == nil && fi.IsDir() {
		return SubversionBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".hg"))
	if err == nil && fi.IsDir() {
		return MercurialBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, "_darcs"))
	if err == nil && fi.IsDir() {
		return DarcsBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, ".fslckout"))
	if err == nil && fi.IsDir() {
		return FossilBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, "_FOSSIL_"))
	if err == nil && fi.IsDir() {
		return FossilBackend
	}

	fi, err = os.Stat(filepath.Join(repo.FullPath, "CVS"))
	if err == nil && fi.IsDir() {
		return cvsDummyBackend
	}
	return nil
}

var vcsDirs = []string{".git", ".svn", ".hg", "_darcs", ".fslckout", "_FOSSIL_", "CVS"}

func walkLocalRepositories(callback func(*LocalRepository)) {
	for _, root := range localRepositoryRoots() {
		filepath.Walk(root, func(path string, fileInfo os.FileInfo, err error) error {
			if err != nil || fileInfo == nil {
				return nil
			}

			if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
				realpath, err := filepath.EvalSymlinks(path)
				if err != nil {
					return nil
				}
				fileInfo, err = os.Stat(realpath)
				if err != nil {
					return nil
				}
			}
			if !fileInfo.IsDir() {
				return nil
			}

			vcsDirFound := false
			for _, d := range vcsDirs {
				_, err := os.Stat(filepath.Join(path, d))
				if err == nil {
					vcsDirFound = true
					break
				}
			}

			if !vcsDirFound {
				return nil
			}

			repo, err := LocalRepositoryFromFullPath(path)
			if err != nil {
				return nil
			}

			if repo == nil {
				return nil
			}
			callback(repo)
			return filepath.SkipDir
		})
	}
}

var _localRepositoryRoots []string

// localRepositoryRoots returns locally cloned repositories' root directories.
// The root dirs are determined as following:
//
//   - If GHQ_ROOT environment variable is nonempty, use it as the only root dir.
//   - Otherwise, use the result of `git config --get-all ghq.root` as the dirs.
//   - Otherwise, fallback to the default root, `~/.ghq`.
//
// TODO: More fancy default directory path?
func localRepositoryRoots() []string {
	if len(_localRepositoryRoots) != 0 {
		return _localRepositoryRoots
	}

	envRoot := os.Getenv("GHQ_ROOT")
	if envRoot != "" {
		_localRepositoryRoots = filepath.SplitList(envRoot)
	} else {
		var err error
		_localRepositoryRoots, err = GitConfigAll("ghq.root")
		logger.PanicIf(err)
	}

	if len(_localRepositoryRoots) == 0 {
		homeDir, err := os.UserHomeDir()
		logger.PanicIf(err)

		_localRepositoryRoots = []string{filepath.Join(homeDir, ".ghq")}
	}

	for i, v := range _localRepositoryRoots {
		path := filepath.Clean(v)
		if _, err := os.Stat(path); err == nil {
			path, err = filepath.EvalSymlinks(path)
			logger.PanicIf(err)
		}
		if !filepath.IsAbs(path) {
			var err error
			path, err = filepath.Abs(path)
			logger.PanicIf(err)
		}
		_localRepositoryRoots[i] = path
	}

	return _localRepositoryRoots
}

func primaryLocalRepositoryRoot() string {
	return localRepositoryRoots()[0]
}
<|endoftext|>