<commit_before>package halo

import (
	"sort"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/table"
)

// halos allows for arrays of halo properties to be sorted simultaneously.
type halos struct {
	rids []int
	xs, ys, zs, ms, rs []float64
}

func (hs *halos) Len() int { return len(hs.rs) }
func (hs *halos) Less(i, j int) bool { return hs.rs[i] < hs.rs[j] }
func (hs *halos) Swap(i, j int) {
	hs.rs[i], hs.rs[j] = hs.rs[j], hs.rs[i]
	hs.ms[i], hs.ms[j] = hs.ms[j], hs.ms[i]
	hs.xs[i], hs.xs[j] = hs.xs[j], hs.xs[i]
	hs.ys[i], hs.ys[j] = hs.ys[j], hs.ys[i]
	hs.zs[i], hs.zs[j] = hs.zs[j], hs.zs[i]
	hs.rids[i], hs.rids[j] = hs.rids[j], hs.rids[i]
}

// ReadRockstar reads halo information from the given Rockstar catalog, sorted
// from largest to smallest.
func ReadRockstar(
	file string, rType Radius, cosmo *io.CosmologyHeader,
) (rids []int, xs, ys, zs, ms, rs []float64, err error) {
	rCol := rType.RockstarColumn()
	idCol, xCol, yCol, zCol := 1, 17, 18, 19

	colIdxs := []int{ idCol, xCol, yCol, zCol, rCol }
	cols, err := table.ReadTable(file, colIdxs, nil)
	if err != nil { return nil, nil, nil, nil, nil, nil, err }

	ids := cols[0]
	xs, ys, zs = cols[1], cols[2], cols[3]
	if rType.RockstarMass() {
		ms = cols[4]
		rs = make([]float64, len(ms))
		rType.Radius(cosmo, ms, rs)
	} else {
		rs = cols[4]
		ms = make([]float64, len(rs))
		for i := range rs { rs[i] /= 1000 } // kpc -> Mpc
		rType.Mass(cosmo, rs, ms)
	}

	rids = make([]int, len(ids))
	for i := range rids { rids[i] = int(ids[i]) }

	sort.Sort(sort.Reverse(&halos{ rids, xs, ys, zs, ms, rs }))
	return rids, xs, ys, zs, ms, rs, nil
}

type Val int
const (
	Scale Val = iota
	ID
	DescScale
	DescID
	NumProg
	PID
	UPID
	DescPID
	Phantom
	SAMMVir
	MVir
	RVir
	Rs
	Vrms
	MMP
	ScaleOfLastMMP
	VMax
	X
	Y
	Z
	Vx
	Vy
	Vz
	Jx
	Jy
	Jz
	Spin
	BreadthFirstID
	DepthFirstID
	TreeRootID
	OrigHaloID
	SnapNum
	NextCoprogenitorDepthFirstID
	LastProgenitorDepthFirstID
	RsKylpin
	MVirAll
	M200b
	M200c
	M500c
	M2500c
	XOff
	Voff
	SpinBullock
	BToA
	CToA
	Ax
	Ay
	Az
	BToA500c
	CToA500c
	Ax500c
	Ay500c
	Az500c
	TU
	MAcc
	MPeak
	VAcc
	VPeak
	HalfmassScale
	AccRateInst
	AccRate100Myr
	AccRateTdyn
	valNum
	RadVir
	Rad200b
	Rad200c
	Rad500c
	Rad2500c
)

func ReadRockstarVals(
	file string, cosmo *io.CosmologyHeader, valFlags ...Val,
) (ids []int, vals[][]float64, err error) {
	colIdxs := []int{ int(ID) }
	for _, val := range valFlags {
		switch val {
		case RadVir: colIdxs = append(colIdxs, int(MVir))
		case Rad200b: colIdxs = append(colIdxs, int(M200b))
		case Rad200c: colIdxs = append(colIdxs, int(M200c))
		case Rad500c: colIdxs = append(colIdxs, int(M500c))
		case Rad2500c: colIdxs = append(colIdxs, int(M2500c))
		default:
			colIdxs = append(colIdxs, int(val))
		}
	}
	vals, err = table.ReadTable(file, colIdxs, nil)
	if err != nil { return nil, nil, err }

	ids = make([]int, len(vals[0]))
	for i := range vals[0] {
		ids[i] = int(vals[0][i])
	}

	for i, val := range valFlags {
		switch val {
		case RadVir: RVirial.Radius(cosmo, vals[i], vals[i])
		case Rad200b: R200m.Radius(cosmo, vals[i], vals[i])
		case Rad200c: R200c.Radius(cosmo, vals[i], vals[i])
		case Rad500c: R500c.Radius(cosmo, vals[i], vals[i])
		case Rad2500c: R2500c.Radius(cosmo, vals[i], vals[i])
		}
	}

	if len(vals) == 1 {
		return ids, [][]float64{}, nil
	} else {
		return ids, vals[1:], nil
	}
}

func init() {
	if valNum != 62 { panic("Internal gotetra setup error.") }
}
<commit_msg>Buf fix in ReadRockstarVals() on radii.<commit_after>package halo

import (
	"sort"

	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/table"
)

// halos allows for arrays of halo properties to be sorted simultaneously.
type halos struct {
	rids []int
	xs, ys, zs, ms, rs []float64
}

func (hs *halos) Len() int { return len(hs.rs) }
func (hs *halos) Less(i, j int) bool { return hs.rs[i] < hs.rs[j] }
func (hs *halos) Swap(i, j int) {
	hs.rs[i], hs.rs[j] = hs.rs[j], hs.rs[i]
	hs.ms[i], hs.ms[j] = hs.ms[j], hs.ms[i]
	hs.xs[i], hs.xs[j] = hs.xs[j], hs.xs[i]
	hs.ys[i], hs.ys[j] = hs.ys[j], hs.ys[i]
	hs.zs[i], hs.zs[j] = hs.zs[j], hs.zs[i]
	hs.rids[i], hs.rids[j] = hs.rids[j], hs.rids[i]
}

// ReadRockstar reads halo information from the given Rockstar catalog, sorted
// from largest to smallest.
func ReadRockstar(
	file string, rType Radius, cosmo *io.CosmologyHeader,
) (rids []int, xs, ys, zs, ms, rs []float64, err error) {
	rCol := rType.RockstarColumn()
	idCol, xCol, yCol, zCol := 1, 17, 18, 19

	colIdxs := []int{ idCol, xCol, yCol, zCol, rCol }
	cols, err := table.ReadTable(file, colIdxs, nil)
	if err != nil { return nil, nil, nil, nil, nil, nil, err }

	ids := cols[0]
	xs, ys, zs = cols[1], cols[2], cols[3]
	if rType.RockstarMass() {
		ms = cols[4]
		rs = make([]float64, len(ms))
		rType.Radius(cosmo, ms, rs)
	} else {
		rs = cols[4]
		ms = make([]float64, len(rs))
		for i := range rs { rs[i] /= 1000 } // kpc -> Mpc
		rType.Mass(cosmo, rs, ms)
	}

	rids = make([]int, len(ids))
	for i := range rids { rids[i] = int(ids[i]) }

	sort.Sort(sort.Reverse(&halos{ rids, xs, ys, zs, ms, rs }))
	return rids, xs, ys, zs, ms, rs, nil
}

type Val int
const (
	Scale Val = iota
	ID
	DescScale
	DescID
	NumProg
	PID
	UPID
	DescPID
	Phantom
	SAMMVir
	MVir
	RVir
	Rs
	Vrms
	MMP
	ScaleOfLastMMP
	VMax
	X
	Y
	Z
	Vx
	Vy
	Vz
	Jx
	Jy
	Jz
	Spin
	BreadthFirstID
	DepthFirstID
	TreeRootID
	OrigHaloID
	SnapNum
	NextCoprogenitorDepthFirstID
	LastProgenitorDepthFirstID
	RsKylpin
	MVirAll
	M200b
	M200c
	M500c
	M2500c
	XOff
	Voff
	SpinBullock
	BToA
	CToA
	Ax
	Ay
	Az
	BToA500c
	CToA500c
	Ax500c
	Ay500c
	Az500c
	TU
	MAcc
	MPeak
	VAcc
	VPeak
	HalfmassScale
	AccRateInst
	AccRate100Myr
	AccRateTdyn
	valNum
	RadVir
	Rad200b
	Rad200c
	Rad500c
	Rad2500c
)

func ReadRockstarVals(
	file string, cosmo *io.CosmologyHeader, valFlags ...Val,
) (ids []int, vals[][]float64, err error) {
	colIdxs := []int{ int(ID) }
	for _, val := range valFlags {
		switch val {
		case RadVir: colIdxs = append(colIdxs, int(MVir))
		case Rad200b: colIdxs = append(colIdxs, int(M200b))
		case Rad200c: colIdxs = append(colIdxs, int(M200c))
		case Rad500c: colIdxs = append(colIdxs, int(M500c))
		case Rad2500c: colIdxs = append(colIdxs, int(M2500c))
		default:
			colIdxs = append(colIdxs, int(val))
		}
	}
	vals, err = table.ReadTable(file, colIdxs, nil)
	if err != nil { return nil, nil, err }

	ids = make([]int, len(vals[0]))
	for i := range vals[0] {
		ids[i] = int(vals[0][i])
	}

	for i, val := range valFlags {
		switch val {
		case RadVir: RVirial.Radius(cosmo, vals[i+1], vals[i+1])
		case Rad200b: R200m.Radius(cosmo, vals[i+1], vals[i+1])
		case Rad200c: R200c.Radius(cosmo, vals[i+1], vals[i+1])
		case Rad500c: R500c.Radius(cosmo, vals[i+1], vals[i+1])
		case Rad2500c: R2500c.Radius(cosmo, vals[i+1], vals[i+1])
		}
	}

	if len(vals) == 1 {
		return ids, [][]float64{}, nil
	} else {
		return ids, vals[1:], nil
	}
}

func init() {
	if valNum != 62 { panic("Internal gotetra setup error.") }
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package resourcelock

import (
	"context"
	"fmt"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	coordinationv1 "k8s.io/client-go/kubernetes/typed/coordination/v1"
	corev1 "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	LeaderElectionRecordAnnotationKey = "control-plane.alpha.kubernetes.io/leader"
	EndpointsResourceLock             = "endpoints"
	ConfigMapsResourceLock            = "configmaps"
	LeasesResourceLock                = "leases"
	EndpointsLeasesResourceLock       = "endpointsleases"
	ConfigMapsLeasesResourceLock      = "configmapsleases"
)

// LeaderElectionRecord is the record that is stored in the leader election annotation.
// This information should be used for observational purposes only and could be replaced
// with a random string (e.g. UUID) with only slight modification of this code.
// TODO(mikedanese): this should potentially be versioned
type LeaderElectionRecord struct {
	// HolderIdentity is the ID that owns the lease. If empty, no one owns this lease and
	// all callers may acquire. Versions of this library prior to Kubernetes 1.14 will not
	// attempt to acquire leases with empty identities and will wait for the full lease
	// interval to expire before attempting to reacquire. This value is set to empty when
	// a client voluntarily steps down.
	HolderIdentity       string      `json:"holderIdentity"`
	LeaseDurationSeconds int         `json:"leaseDurationSeconds"`
	AcquireTime          metav1.Time `json:"acquireTime"`
	RenewTime            metav1.Time `json:"renewTime"`
	LeaderTransitions    int         `json:"leaderTransitions"`
}

// EventRecorder records a change in the ResourceLock.
type EventRecorder interface {
	Eventf(obj runtime.Object, eventType, reason, message string, args ...interface{})
}

// ResourceLockConfig common data that exists across different
// resource locks
type ResourceLockConfig struct {
	// Identity is the unique string identifying a lease holder across
	// all participants in an election.
	Identity string
	// EventRecorder is optional.
	EventRecorder EventRecorder
}

// Interface offers a common interface for locking on arbitrary
// resources used in leader election.  The Interface is used
// to hide the details on specific implementations in order to allow
// them to change over time.  This interface is strictly for use
// by the leaderelection code.
type Interface interface {
	// Get returns the LeaderElectionRecord
	Get(ctx context.Context) (*LeaderElectionRecord, []byte, error)

	// Create attempts to create a LeaderElectionRecord
	Create(ctx context.Context, ler LeaderElectionRecord) error

	// Update will update and existing LeaderElectionRecord
	Update(ctx context.Context, ler LeaderElectionRecord) error

	// RecordEvent is used to record events
	RecordEvent(string)

	// Identity will return the locks Identity
	Identity() string

	// Describe is used to convert details on current resource lock
	// into a string
	Describe() string
}

// Manufacture will create a lock of a given type according to the input parameters
func New(lockType string, ns string, name string, coreClient corev1.CoreV1Interface, coordinationClient coordinationv1.CoordinationV1Interface, rlc ResourceLockConfig) (Interface, error) {
	endpointsLock := &EndpointsLock{
		EndpointsMeta: metav1.ObjectMeta{
			Namespace: ns,
			Name:      name,
		},
		Client:     coreClient,
		LockConfig: rlc,
	}
	configmapLock := &ConfigMapLock{
		ConfigMapMeta: metav1.ObjectMeta{
			Namespace: ns,
			Name:      name,
		},
		Client:     coreClient,
		LockConfig: rlc,
	}
	leaseLock := &LeaseLock{
		LeaseMeta: metav1.ObjectMeta{
			Namespace: ns,
			Name:      name,
		},
		Client:     coordinationClient,
		LockConfig: rlc,
	}
	switch lockType {
	case EndpointsResourceLock:
		return endpointsLock, nil
	case ConfigMapsResourceLock:
		return configmapLock, nil
	case LeasesResourceLock:
		return leaseLock, nil
	case EndpointsLeasesResourceLock:
		return &MultiLock{
			Primary:   endpointsLock,
			Secondary: leaseLock,
		}, nil
	case ConfigMapsLeasesResourceLock:
		return &MultiLock{
			Primary:   configmapLock,
			Secondary: leaseLock,
		}, nil
	default:
		return nil, fmt.Errorf("Invalid lock-type %s", lockType)
	}
}

// NewFromKubeconfig will create a lock of a given type according to the input parameters.
func NewFromKubeconfig(lockType string, ns string, name string, rlc ResourceLockConfig, kubeconfig *restclient.Config, renewDeadline time.Duration) (Interface, error) {
	// shallow copy, do not modify the kubeconfig
	config := *kubeconfig
	timeout := ((renewDeadline / time.Millisecond) / 2) * time.Millisecond
	if timeout < time.Second {
		timeout = time.Second
	}
	config.Timeout = timeout
	leaderElectionClient := clientset.NewForConfigOrDie(restclient.AddUserAgent(&config, "leader-election"))
	return New(lockType, ns, name, leaderElectionClient.CoreV1(), leaderElectionClient.CoordinationV1(), rlc)
}
<commit_msg>[Leader election] Add documentation to function<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package resourcelock

import (
	"context"
	"fmt"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	coordinationv1 "k8s.io/client-go/kubernetes/typed/coordination/v1"
	corev1 "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	LeaderElectionRecordAnnotationKey = "control-plane.alpha.kubernetes.io/leader"
	EndpointsResourceLock             = "endpoints"
	ConfigMapsResourceLock            = "configmaps"
	LeasesResourceLock                = "leases"
	EndpointsLeasesResourceLock       = "endpointsleases"
	ConfigMapsLeasesResourceLock      = "configmapsleases"
)

// LeaderElectionRecord is the record that is stored in the leader election annotation.
// This information should be used for observational purposes only and could be replaced
// with a random string (e.g. UUID) with only slight modification of this code.
// TODO(mikedanese): this should potentially be versioned
type LeaderElectionRecord struct {
	// HolderIdentity is the ID that owns the lease. If empty, no one owns this lease and
	// all callers may acquire. Versions of this library prior to Kubernetes 1.14 will not
	// attempt to acquire leases with empty identities and will wait for the full lease
	// interval to expire before attempting to reacquire. This value is set to empty when
	// a client voluntarily steps down.
	HolderIdentity       string      `json:"holderIdentity"`
	LeaseDurationSeconds int         `json:"leaseDurationSeconds"`
	AcquireTime          metav1.Time `json:"acquireTime"`
	RenewTime            metav1.Time `json:"renewTime"`
	LeaderTransitions    int         `json:"leaderTransitions"`
}

// EventRecorder records a change in the ResourceLock.
type EventRecorder interface {
	Eventf(obj runtime.Object, eventType, reason, message string, args ...interface{})
}

// ResourceLockConfig common data that exists across different
// resource locks
type ResourceLockConfig struct {
	// Identity is the unique string identifying a lease holder across
	// all participants in an election.
	Identity string
	// EventRecorder is optional.
	EventRecorder EventRecorder
}

// Interface offers a common interface for locking on arbitrary
// resources used in leader election.  The Interface is used
// to hide the details on specific implementations in order to allow
// them to change over time.  This interface is strictly for use
// by the leaderelection code.
type Interface interface {
	// Get returns the LeaderElectionRecord
	Get(ctx context.Context) (*LeaderElectionRecord, []byte, error)

	// Create attempts to create a LeaderElectionRecord
	Create(ctx context.Context, ler LeaderElectionRecord) error

	// Update will update and existing LeaderElectionRecord
	Update(ctx context.Context, ler LeaderElectionRecord) error

	// RecordEvent is used to record events
	RecordEvent(string)

	// Identity will return the locks Identity
	Identity() string

	// Describe is used to convert details on current resource lock
	// into a string
	Describe() string
}

// Manufacture will create a lock of a given type according to the input parameters
func New(lockType string, ns string, name string, coreClient corev1.CoreV1Interface, coordinationClient coordinationv1.CoordinationV1Interface, rlc ResourceLockConfig) (Interface, error) {
	endpointsLock := &EndpointsLock{
		EndpointsMeta: metav1.ObjectMeta{
			Namespace: ns,
			Name:      name,
		},
		Client:     coreClient,
		LockConfig: rlc,
	}
	configmapLock := &ConfigMapLock{
		ConfigMapMeta: metav1.ObjectMeta{
			Namespace: ns,
			Name:      name,
		},
		Client:     coreClient,
		LockConfig: rlc,
	}
	leaseLock := &LeaseLock{
		LeaseMeta: metav1.ObjectMeta{
			Namespace: ns,
			Name:      name,
		},
		Client:     coordinationClient,
		LockConfig: rlc,
	}
	switch lockType {
	case EndpointsResourceLock:
		return endpointsLock, nil
	case ConfigMapsResourceLock:
		return configmapLock, nil
	case LeasesResourceLock:
		return leaseLock, nil
	case EndpointsLeasesResourceLock:
		return &MultiLock{
			Primary:   endpointsLock,
			Secondary: leaseLock,
		}, nil
	case ConfigMapsLeasesResourceLock:
		return &MultiLock{
			Primary:   configmapLock,
			Secondary: leaseLock,
		}, nil
	default:
		return nil, fmt.Errorf("Invalid lock-type %s", lockType)
	}
}

// NewFromKubeconfig will create a lock of a given type according to the input parameters.
// Timeout set for a client used to contact to Kubernetes should be lower than
// RenewDeadline to keep a single hung request from forcing a leader loss.
// Setting it to max(time.Second, RenewDeadline/2) as a reasonable heuristic.
func NewFromKubeconfig(lockType string, ns string, name string, rlc ResourceLockConfig, kubeconfig *restclient.Config, renewDeadline time.Duration) (Interface, error) {
	// shallow copy, do not modify the kubeconfig
	config := *kubeconfig
	timeout := ((renewDeadline / time.Millisecond) / 2) * time.Millisecond
	if timeout < time.Second {
		timeout = time.Second
	}
	config.Timeout = timeout
	leaderElectionClient := clientset.NewForConfigOrDie(restclient.AddUserAgent(&config, "leader-election"))
	return New(lockType, ns, name, leaderElectionClient.CoreV1(), leaderElectionClient.CoordinationV1(), rlc)
}
<|endoftext|>


<commit_before>package sftp

import (
	"context"
	"encoding"
	"io"
	"path"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"

	"github.com/pkg/errors"
)

var maxTxPacket uint32 = 1 << 15

// Handlers contains the 4 SFTP server request handlers.
type Handlers struct {
	FileGet  FileReader
	FilePut  FileWriter
	FileCmd  FileCmder
	FileList FileLister
}

// RequestServer abstracts the sftp protocol with an http request-like protocol
type RequestServer struct {
	*serverConn
	Handlers        Handlers
	pktMgr          *packetManager
	openRequests    map[string]*Request
	openRequestLock sync.RWMutex
	handleCount     int
}

// NewRequestServer creates/allocates/returns new RequestServer.
// Normally there there will be one server per user-session.
func NewRequestServer(rwc io.ReadWriteCloser, h Handlers) *RequestServer {
	svrConn := &serverConn{
		conn: conn{
			Reader:      rwc,
			WriteCloser: rwc,
		},
	}
	return &RequestServer{
		serverConn:   svrConn,
		Handlers:     h,
		pktMgr:       newPktMgr(svrConn),
		openRequests: make(map[string]*Request),
	}
}

// New Open packet/Request
func (rs *RequestServer) nextRequest(r *Request) string {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	rs.handleCount++
	handle := strconv.Itoa(rs.handleCount)
	rs.openRequests[handle] = r
	return handle
}

// Returns Request from openRequests, bool is false if it is missing
// If the method is different, save/return a new Request w/ that Method.
//
// The Requests in openRequests work essentially as open file descriptors that
// you can do different things with. What you are doing with it are denoted by
// the first packet of that type (read/write/etc). We create a new Request when
// it changes to set the request.Method attribute in a thread safe way.
func (rs *RequestServer) getRequest(handle, method string) (*Request, bool) {
	rs.openRequestLock.RLock()
	r, ok := rs.openRequests[handle]
	rs.openRequestLock.RUnlock()
	if !ok || r.Method == method {
		return r, ok
	}
	// if we make it here we need to replace the request
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	r, ok = rs.openRequests[handle]
	if !ok || r.Method == method { // re-check needed b/c lock race
		return r, ok
	}
	r = r.copy()
	r.Method = method
	rs.openRequests[handle] = r
	return r, ok
}

func (rs *RequestServer) closeRequest(handle string) error {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	if r, ok := rs.openRequests[handle]; ok {
		delete(rs.openRequests, handle)
		return r.close()
	}
	return syscall.EBADF
}

// Close the read/write/closer to trigger exiting the main server loop
func (rs *RequestServer) Close() error { return rs.conn.Close() }

// Serve requests for user session
func (rs *RequestServer) Serve() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var wg sync.WaitGroup
	runWorker := func(ch requestChan) {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := rs.packetWorker(ctx, ch); err != nil {
				rs.conn.Close() // shuts down recvPacket
			}
		}()
	}
	pktChan := rs.pktMgr.workerChan(runWorker)

	var err error
	var pkt requestPacket
	var pktType uint8
	var pktBytes []byte
	for {
		pktType, pktBytes, err = rs.recvPacket()
		if err != nil {
			break
		}

		pkt, err = makePacket(rxPacket{fxp(pktType), pktBytes})
		if err != nil {
			switch errors.Cause(err) {
			case errUnknownExtendedPacket:
				rs.serverConn.sendError(pkt, ErrSshFxOpUnsupported)
			default:
				debug("makePacket err: %v", err)
				rs.conn.Close() // shuts down recvPacket
				break
			}
		}

		pktChan <- pkt
	}

	close(pktChan) // shuts down sftpServerWorkers
	wg.Wait()      // wait for all workers to exit

	// make sure all open requests are properly closed
	// (eg. possible on dropped connections, client crashes, etc.)
	for handle, req := range rs.openRequests {
		delete(rs.openRequests, handle)
		req.close()
	}

	return err
}

func (rs *RequestServer) packetWorker(
	ctx context.Context, pktChan chan requestPacket,
) error {
	for pkt := range pktChan {
		var rpkt responsePacket
		switch pkt := pkt.(type) {
		case *sshFxInitPacket:
			rpkt = sshFxVersionPacket{sftpProtocolVersion, nil}
		case *sshFxpClosePacket:
			handle := pkt.getHandle()
			rpkt = statusFromError(pkt, rs.closeRequest(handle))
		case *sshFxpRealpathPacket:
			rpkt = cleanPacketPath(pkt)
		case *sshFxpOpendirPacket:
			request := requestFromPacket(ctx, pkt)
			handle := rs.nextRequest(request)
			rpkt = sshFxpHandlePacket{pkt.id(), handle}
		case *sshFxpOpenPacket:
			request := requestFromPacket(ctx, pkt)
			handle := rs.nextRequest(request)
			rpkt = sshFxpHandlePacket{pkt.id(), handle}
			if pkt.hasPflags(ssh_FXF_CREAT) {
				if p := request.call(rs.Handlers, pkt); !statusOk(p) {
					rpkt = p // if error in write, return it
				}
			}
		case hasHandle:
			handle := pkt.getHandle()
			request, ok := rs.getRequest(handle, requestMethod(pkt))
			if !ok {
				rpkt = statusFromError(pkt, syscall.EBADF)
			} else {
				rpkt = request.call(rs.Handlers, pkt)
			}
		case hasPath:
			request := requestFromPacket(ctx, pkt)
			rpkt = request.call(rs.Handlers, pkt)
			request.close()
		default:
			return errors.Errorf("unexpected packet type %T", pkt)
		}

		err := rs.sendPacket(rpkt)
		if err != nil {
			return err
		}
	}
	return nil
}

// True is responsePacket is an OK status packet
func statusOk(rpkt responsePacket) bool {
	p, ok := rpkt.(sshFxpStatusPacket)
	return ok && p.StatusError.Code == ssh_FX_OK
}

// clean and return name packet for file
func cleanPacketPath(pkt *sshFxpRealpathPacket) responsePacket {
	path := cleanPath(pkt.getPath())
	return &sshFxpNamePacket{
		ID: pkt.id(),
		NameAttrs: []sshFxpNameAttr{{
			Name:     path,
			LongName: path,
			Attrs:    emptyFileStat,
		}},
	}
}

// Makes sure we have a clean POSIX (/) absolute path to work with
func cleanPath(p string) string {
	p = filepath.ToSlash(p)
	if !filepath.IsAbs(p) {
		p = "/" + p
	}
	return path.Clean(p)
}

// Wrap underlying connection methods to use packetManager
func (rs *RequestServer) sendPacket(m encoding.BinaryMarshaler) error {
	if pkt, ok := m.(responsePacket); ok {
		rs.pktMgr.readyPacket(pkt)
	} else {
		return errors.Errorf("unexpected packet type %T", m)
	}
	return nil
}
<commit_msg>Handle RequestServer the same<commit_after>package sftp

import (
	"context"
	"encoding"
	"io"
	"path"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"

	"github.com/pkg/errors"
)

var maxTxPacket uint32 = 1 << 15

// Handlers contains the 4 SFTP server request handlers.
type Handlers struct {
	FileGet  FileReader
	FilePut  FileWriter
	FileCmd  FileCmder
	FileList FileLister
}

// RequestServer abstracts the sftp protocol with an http request-like protocol
type RequestServer struct {
	*serverConn
	Handlers        Handlers
	pktMgr          *packetManager
	openRequests    map[string]*Request
	openRequestLock sync.RWMutex
	handleCount     int
}

// NewRequestServer creates/allocates/returns new RequestServer.
// Normally there there will be one server per user-session.
func NewRequestServer(rwc io.ReadWriteCloser, h Handlers) *RequestServer {
	svrConn := &serverConn{
		conn: conn{
			Reader:      rwc,
			WriteCloser: rwc,
		},
	}
	return &RequestServer{
		serverConn:   svrConn,
		Handlers:     h,
		pktMgr:       newPktMgr(svrConn),
		openRequests: make(map[string]*Request),
	}
}

// New Open packet/Request
func (rs *RequestServer) nextRequest(r *Request) string {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	rs.handleCount++
	handle := strconv.Itoa(rs.handleCount)
	rs.openRequests[handle] = r
	return handle
}

// Returns Request from openRequests, bool is false if it is missing
// If the method is different, save/return a new Request w/ that Method.
//
// The Requests in openRequests work essentially as open file descriptors that
// you can do different things with. What you are doing with it are denoted by
// the first packet of that type (read/write/etc). We create a new Request when
// it changes to set the request.Method attribute in a thread safe way.
func (rs *RequestServer) getRequest(handle, method string) (*Request, bool) {
	rs.openRequestLock.RLock()
	r, ok := rs.openRequests[handle]
	rs.openRequestLock.RUnlock()
	if !ok || r.Method == method {
		return r, ok
	}
	// if we make it here we need to replace the request
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	r, ok = rs.openRequests[handle]
	if !ok || r.Method == method { // re-check needed b/c lock race
		return r, ok
	}
	r = r.copy()
	r.Method = method
	rs.openRequests[handle] = r
	return r, ok
}

func (rs *RequestServer) closeRequest(handle string) error {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	if r, ok := rs.openRequests[handle]; ok {
		delete(rs.openRequests, handle)
		return r.close()
	}
	return syscall.EBADF
}

// Close the read/write/closer to trigger exiting the main server loop
func (rs *RequestServer) Close() error { return rs.conn.Close() }

// Serve requests for user session
func (rs *RequestServer) Serve() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var wg sync.WaitGroup
	runWorker := func(ch requestChan) {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := rs.packetWorker(ctx, ch); err != nil {
				rs.conn.Close() // shuts down recvPacket
			}
		}()
	}
	pktChan := rs.pktMgr.workerChan(runWorker)

	var err error
	var pkt requestPacket
	var pktType uint8
	var pktBytes []byte
	for {
		pktType, pktBytes, err = rs.recvPacket()
		if err != nil {
			break
		}

		pkt, err = makePacket(rxPacket{fxp(pktType), pktBytes})
		if err != nil {
			switch errors.Cause(err) {
			case errUnknownExtendedPacket:
				if err := rs.serverConn.sendError(pkt, ErrSshFxOpUnsupported); err != nil {
					debug("failed to send err packet: %v", err)
					rs.conn.Close() // shuts down recvPacket
					break
				}
			default:
				debug("makePacket err: %v", err)
				rs.conn.Close() // shuts down recvPacket
				break
			}
		}

		pktChan <- pkt
	}

	close(pktChan) // shuts down sftpServerWorkers
	wg.Wait()      // wait for all workers to exit

	// make sure all open requests are properly closed
	// (eg. possible on dropped connections, client crashes, etc.)
	for handle, req := range rs.openRequests {
		delete(rs.openRequests, handle)
		req.close()
	}

	return err
}

func (rs *RequestServer) packetWorker(
	ctx context.Context, pktChan chan requestPacket,
) error {
	for pkt := range pktChan {
		var rpkt responsePacket
		switch pkt := pkt.(type) {
		case *sshFxInitPacket:
			rpkt = sshFxVersionPacket{sftpProtocolVersion, nil}
		case *sshFxpClosePacket:
			handle := pkt.getHandle()
			rpkt = statusFromError(pkt, rs.closeRequest(handle))
		case *sshFxpRealpathPacket:
			rpkt = cleanPacketPath(pkt)
		case *sshFxpOpendirPacket:
			request := requestFromPacket(ctx, pkt)
			handle := rs.nextRequest(request)
			rpkt = sshFxpHandlePacket{pkt.id(), handle}
		case *sshFxpOpenPacket:
			request := requestFromPacket(ctx, pkt)
			handle := rs.nextRequest(request)
			rpkt = sshFxpHandlePacket{pkt.id(), handle}
			if pkt.hasPflags(ssh_FXF_CREAT) {
				if p := request.call(rs.Handlers, pkt); !statusOk(p) {
					rpkt = p // if error in write, return it
				}
			}
		case hasHandle:
			handle := pkt.getHandle()
			request, ok := rs.getRequest(handle, requestMethod(pkt))
			if !ok {
				rpkt = statusFromError(pkt, syscall.EBADF)
			} else {
				rpkt = request.call(rs.Handlers, pkt)
			}
		case hasPath:
			request := requestFromPacket(ctx, pkt)
			rpkt = request.call(rs.Handlers, pkt)
			request.close()
		default:
			return errors.Errorf("unexpected packet type %T", pkt)
		}

		err := rs.sendPacket(rpkt)
		if err != nil {
			return err
		}
	}
	return nil
}

// True is responsePacket is an OK status packet
func statusOk(rpkt responsePacket) bool {
	p, ok := rpkt.(sshFxpStatusPacket)
	return ok && p.StatusError.Code == ssh_FX_OK
}

// clean and return name packet for file
func cleanPacketPath(pkt *sshFxpRealpathPacket) responsePacket {
	path := cleanPath(pkt.getPath())
	return &sshFxpNamePacket{
		ID: pkt.id(),
		NameAttrs: []sshFxpNameAttr{{
			Name:     path,
			LongName: path,
			Attrs:    emptyFileStat,
		}},
	}
}

// Makes sure we have a clean POSIX (/) absolute path to work with
func cleanPath(p string) string {
	p = filepath.ToSlash(p)
	if !filepath.IsAbs(p) {
		p = "/" + p
	}
	return path.Clean(p)
}

// Wrap underlying connection methods to use packetManager
func (rs *RequestServer) sendPacket(m encoding.BinaryMarshaler) error {
	if pkt, ok := m.(responsePacket); ok {
		rs.pktMgr.readyPacket(pkt)
	} else {
		return errors.Errorf("unexpected packet type %T", m)
	}
	return nil
}
<|endoftext|>


<commit_before>package consul

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/fatih/color"
	"github.com/hashicorp/consul/api"

	"github.com/servehub/utils"
)

var upstreamNameRegex = regexp.MustCompile("[^\\w]+")
var spacesRegex = regexp.MustCompile("\\s+")

func NginxTemplateContextCommand() cli.Command {
	return cli.Command{
		Name:  "nginx-template-context",
		Usage: "Collect and return data for consul-template",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "filter"},
		},
		Action: func(c *cli.Context) error {
			consul, _ := api.NewClient(api.DefaultConfig())

			filters := make(map[string]string)
			for _, filter := range strings.Split(c.String("filter"), ",") {
				if filter != "" {
					fvs := strings.SplitN(filter, "=", 2)

					if len(fvs) > 1 {
						filters[fvs[0]] = strings.TrimSpace(fvs[1])
					} else {
						filters[fvs[0]] = "true"
					}
				}
			}

			upstreams := make(map[string]map[string]map[string]interface{})
			services := make(map[string]map[string][]map[string]interface{})
			duplicates := make(map[string]string)

			allServicesRoutes, _, err := consul.KV().List("services/routes/", nil)
			if err != nil {
				return fmt.Errorf("Error on load routes: %s", err)
			}

			for _, kv := range allServicesRoutes {
				name := strings.TrimPrefix(kv.Key, "services/routes/")

				instances, _, err := consul.Health().Service(name, "", true, nil)
				if err != nil {
					return fmt.Errorf("Error on get service `%s` health: %s", name, err)
				}

				if len(instances) == 0 {
					continue
				}

				routes := consulRoutes{}
				if err := json.Unmarshal(kv.Value, &routes); err != nil {
					return fmt.Errorf("Error on parse route json: %s. Serive `%s`, json: %s", err, name, string(kv.Value))
				}

				upstream := upstreamNameRegex.ReplaceAllString("serve_"+name, "_")
				if instances[0].Service.Port == 0 {
					upstream += "_static"
				}

				for _, route := range routes.Routes {
					for _, host := range spacesRegex.Split(route.Host, -1) {

						skipedByFilters := false

						for fk, fv := range filters {
							if vval, ok := route.Vars[fk]; !ok || vval != fv {
								skipedByFilters = true
								break
							}

							delete(route.Vars, fk)
						}

						if skipedByFilters {
							break
						}

						delete(route.Vars, "public") // todo: remove hardcoded filter

						location := route.Location
						if location == "" {
							location = "/"
						}

						for _, inst := range instances {
							putUpstream(upstream, inst, upstreams)
						}

						if _, ok := services[host]; !ok {
							services[host] = make(map[string][]map[string]interface{}, 0)
						}

						if _, ok := services[host][location]; !ok {
							services[host][location] = make([]map[string]interface{}, 0)
						}

						routeKeys := "-"
						routeValues := "-"
						for k, v := range route.Vars {
							routeKeys += "${" + k + "}-"
							routeValues += v + "-"
						}

						if exists, ok := duplicates[host+location+routeKeys+routeValues]; !ok {
							duplicates[host+location+routeKeys+routeValues] = name
						} else {
							fmt.Fprintln(os.Stderr, color.RedString("Service with the same routes already exists! exists: %s, skipped: %s", exists, name))
							continue
						}

						ssl := route.Ssl
						if _, ok := duplicates[host+":ssl"]; !ok {
							duplicates[host+":ssl"] = "exist"
						} else {
							ssl = nil
						}

						if route.Params == nil {
							route.Params = make(map[string]interface{})
						}

						services[host][location] = append(services[host][location], map[string]interface{}{
							"upstream":    upstream,
							"routeKeys":   routeKeys,
							"routeValues": routeValues,
							"sortIndex":   strconv.Itoa(len(route.Vars)),
							"cache":       route.Cache,
							"ssl":         ssl,
							"extra":       route.Extra,
							"params":      route.Params,
						})
					}
				}
			}

			// sort routes by sort index
			for _, hh := range services {
				for _, ll := range hh {
					sort.Sort(utils.BySortIndex(ll))
				}
			}

			out, _ := json.MarshalIndent(map[string]interface{}{
				"upstreams": upstreams,
				"services":  services,
			}, "", "  ")

			fmt.Fprintln(os.Stdout, string(out))
			return nil
		},
	}
}

func putUpstream(upstream string, inst *api.ServiceEntry, upstreams map[string]map[string]map[string]interface{}) {
	port := inst.Service.Port

	if _, ok := upstreams[upstream]; !ok {
		upstreams[upstream] = make(map[string]map[string]interface{}, 0)
	}

	address := inst.Node.Address
	if inst.Service.Address != "" {
		address = inst.Service.Address
	}

	upstreams[upstream][fmt.Sprintf("%s:%d", address, port)] = map[string]interface{}{
		"address": address,
		"port":    port,
	}
}

type consulRoutes struct {
	Routes []consulRoute `json:"routes"`
}

type consulRoute struct {
	Host     string                 `json:"host"`
	Location string                 `json:"location,omitempty"`
	Vars     map[string]string      `json:"vars,omitempty"`
	Cache    map[string]string      `json:"cache,omitempty"`
	Ssl      map[string]string      `json:"ssl,omitempty"`
	Extra    string                 `json:"extra,omitempty"`
	Params   map[string]interface{} `json:"params,omitempty"`
}
<commit_msg>serve nginx-template-context: add hostsParams var<commit_after>package consul

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/fatih/color"
	"github.com/hashicorp/consul/api"

	"github.com/servehub/utils"
	"github.com/servehub/utils/mergemap"
)

var upstreamNameRegex = regexp.MustCompile("[^\\w]+")
var spacesRegex = regexp.MustCompile("\\s+")

func NginxTemplateContextCommand() cli.Command {
	return cli.Command{
		Name:  "nginx-template-context",
		Usage: "Collect and return data for consul-template",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "filter"},
		},
		Action: func(c *cli.Context) error {
			consul, _ := api.NewClient(api.DefaultConfig())

			filters := make(map[string]string)
			for _, filter := range strings.Split(c.String("filter"), ",") {
				if filter != "" {
					fvs := strings.SplitN(filter, "=", 2)

					if len(fvs) > 1 {
						filters[fvs[0]] = strings.TrimSpace(fvs[1])
					} else {
						filters[fvs[0]] = "true"
					}
				}
			}

			upstreams := make(map[string]map[string]map[string]interface{})
			services := make(map[string]map[string][]map[string]interface{})
			hostsParams := make(map[string]map[string]interface{})
			duplicates := make(map[string]string)

			allServicesRoutes, _, err := consul.KV().List("services/routes/", nil)
			if err != nil {
				return fmt.Errorf("Error on load routes: %s", err)
			}

			for _, kv := range allServicesRoutes {
				name := strings.TrimPrefix(kv.Key, "services/routes/")

				instances, _, err := consul.Health().Service(name, "", true, nil)
				if err != nil {
					return fmt.Errorf("Error on get service `%s` health: %s", name, err)
				}

				if len(instances) == 0 {
					continue
				}

				routes := consulRoutes{}
				if err := json.Unmarshal(kv.Value, &routes); err != nil {
					return fmt.Errorf("Error on parse route json: %s. Serive `%s`, json: %s", err, name, string(kv.Value))
				}

				upstream := upstreamNameRegex.ReplaceAllString("serve_"+name, "_")
				if instances[0].Service.Port == 0 {
					upstream += "_static"
				}

				for _, route := range routes.Routes {
					for _, host := range spacesRegex.Split(route.Host, -1) {

						skipedByFilters := false

						for fk, fv := range filters {
							if vval, ok := route.Vars[fk]; !ok || vval != fv {
								skipedByFilters = true
								break
							}

							delete(route.Vars, fk)
						}

						if skipedByFilters {
							break
						}

						delete(route.Vars, "public") // todo: remove hardcoded filter

						location := route.Location
						if location == "" {
							location = "/"
						}

						for _, inst := range instances {
							putUpstream(upstream, inst, upstreams)
						}

						if _, ok := services[host]; !ok {
							services[host] = make(map[string][]map[string]interface{}, 0)
						}

						if _, ok := services[host][location]; !ok {
							services[host][location] = make([]map[string]interface{}, 0)
						}

						routeKeys := "-"
						routeValues := "-"
						for k, v := range route.Vars {
							routeKeys += "${" + k + "}-"
							routeValues += v + "-"
						}

						if exists, ok := duplicates[host+location+routeKeys+routeValues]; !ok {
							duplicates[host+location+routeKeys+routeValues] = name
						} else {
							fmt.Fprintln(os.Stderr, color.RedString("Service with the same routes already exists! exists: %s, skipped: %s", exists, name))
							continue
						}

						ssl := route.Ssl
						if _, ok := duplicates[host+":ssl"]; !ok {
							duplicates[host+":ssl"] = "exist"
						} else {
							ssl = nil
						}

						if route.Params == nil {
							route.Params = make(map[string]interface{})
						}

						if _, ok := hostsParams[host]; !ok {
							hostsParams[host] = make(map[string]interface{}, 0)
						}

						if res, err := mergemap.Merge(hostsParams[host], route.Params); err == nil {
							hostsParams[host] = res
						}

						services[host][location] = append(services[host][location], map[string]interface{}{
							"upstream":    upstream,
							"routeKeys":   routeKeys,
							"routeValues": routeValues,
							"sortIndex":   strconv.Itoa(len(route.Vars)),
							"cache":       route.Cache,
							"ssl":         ssl,
							"extra":       route.Extra,
							"params":      route.Params,
						})
					}
				}
			}

			// sort routes by sort index
			for _, hh := range services {
				for _, ll := range hh {
					sort.Sort(utils.BySortIndex(ll))
				}
			}

			out, _ := json.MarshalIndent(map[string]interface{}{
				"upstreams": upstreams,
				"services":  services,
				"hostsParams": hostsParams,
			}, "", "  ")

			fmt.Fprintln(os.Stdout, string(out))
			return nil
		},
	}
}

func putUpstream(upstream string, inst *api.ServiceEntry, upstreams map[string]map[string]map[string]interface{}) {
	port := inst.Service.Port

	if _, ok := upstreams[upstream]; !ok {
		upstreams[upstream] = make(map[string]map[string]interface{}, 0)
	}

	address := inst.Node.Address
	if inst.Service.Address != "" {
		address = inst.Service.Address
	}

	upstreams[upstream][fmt.Sprintf("%s:%d", address, port)] = map[string]interface{}{
		"address": address,
		"port":    port,
	}
}

type consulRoutes struct {
	Routes []consulRoute `json:"routes"`
}

type consulRoute struct {
	Host     string                 `json:"host"`
	Location string                 `json:"location,omitempty"`
	Vars     map[string]string      `json:"vars,omitempty"`
	Cache    map[string]string      `json:"cache,omitempty"`
	Ssl      map[string]string      `json:"ssl,omitempty"`
	Extra    string                 `json:"extra,omitempty"`
	Params   map[string]interface{} `json:"params,omitempty"`
}
<|endoftext|>


<commit_before>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package kbfsedits

import (
	"fmt"
	"sort"
	"sync"

	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/kbfs/tlf"
)

type tlfKey struct {
	tlfName tlf.CanonicalName
	tlfType tlf.Type
}

type tlfWithHistory struct {
	key     tlfKey
	history writersByRevision
}

type tlfByTime struct {
	histories []tlfWithHistory
	indices   map[tlfKey]int
}

func (tbm tlfByTime) Len() int {
	return len(tbm.histories)
}

func (tbm tlfByTime) Less(i, j int) bool {
	iHistory := tbm.histories[i].history
	jHistory := tbm.histories[j].history

	if len(iHistory) == 0 || len(iHistory[0].notifications) == 0 {
		return false
	} else if len(jHistory) == 0 || len(jHistory[0].notifications) == 0 {
		return true
	}

	iTime := iHistory[0].notifications[0].Time
	jTime := jHistory[0].notifications[0].Time

	return iTime.After(jTime)
}

func (tbm tlfByTime) Swap(i, j int) {
	tbm.histories[i], tbm.histories[j] = tbm.histories[j], tbm.histories[i]
	tbm.indices[tbm.histories[i].key] = i
	tbm.indices[tbm.histories[j].key] = j
}

// UserHistory keeps a sorted list of the top known TLF edit
// histories, and can convert those histories into keybase1 protocol
// structs.  TLF histories must be updated by an external caller
// whenever they change.
type UserHistory struct {
	lock sync.RWMutex
	tlfs tlfByTime
}

// NewUserHistory constructs a UserHistory instance.
func NewUserHistory() *UserHistory {
	return &UserHistory{
		tlfs: tlfByTime{
			indices: make(map[tlfKey]int),
		},
	}
}

// UpdateHistory should be called whenever the edit history for a
// given TLF gets new information.
func (uh *UserHistory) UpdateHistory(
	tlfName tlf.CanonicalName, tlfType tlf.Type, tlfHistory *TlfHistory) {
	history := tlfHistory.getHistory()
	key := tlfKey{tlfName, tlfType}

	uh.lock.Lock()
	defer uh.lock.Unlock()
	if currIndex, ok := uh.tlfs.indices[key]; ok {
		uh.tlfs.histories[currIndex].history = history
	} else {
		uh.tlfs.indices[key] = len(uh.tlfs.indices)
		uh.tlfs.histories = append(
			uh.tlfs.histories, tlfWithHistory{key, history})
	}
	sort.Sort(uh.tlfs)
}

func (uh *UserHistory) getTlfHistoryLocked(
	tlfName tlf.CanonicalName, tlfType tlf.Type) (
	history keybase1.FSFolderEditHistory) {
	key := tlfKey{tlfName, tlfType}
	currIndex, ok := uh.tlfs.indices[key]
	if !ok {
		return keybase1.FSFolderEditHistory{}
	}
	tlfHistory := uh.tlfs.histories[currIndex].history

	folder := keybase1.Folder{
		Name:       string(tlfName),
		FolderType: tlfType.FolderType(),
		Private:    tlfType == tlf.Private,
	}
	history.Folder = folder
	if len(tlfHistory) == 0 {
		return history
	}
	if len(tlfHistory[0].notifications) > 0 {
		history.ServerTime = keybase1.ToTime(
			tlfHistory[0].notifications[0].Time)
	}
	history.History = make(
		[]keybase1.FSFolderWriterEditHistory, len(tlfHistory))
	for i, wn := range tlfHistory {
		history.History[i].WriterName = wn.writerName
		history.History[i].Edits = make(
			[]keybase1.FSFolderWriterEdit, len(wn.notifications))
		for j, n := range wn.notifications {
			history.History[i].Edits[j].Filename = n.Filename
			history.History[i].Edits[j].ServerTime = keybase1.ToTime(n.Time)
			switch n.Type {
			case NotificationCreate:
				history.History[i].Edits[j].NotificationType =
					keybase1.FSNotificationType_FILE_CREATED
			case NotificationModify:
				history.History[i].Edits[j].NotificationType =
					keybase1.FSNotificationType_FILE_MODIFIED
			default:
				panic(fmt.Sprintf("Unknown notification type %s", n.Type))
			}
		}
	}
	return history
}

// GetTlfHistory returns the edit history of a given TLF, converted to
// keybase1 protocol structs.
func (uh *UserHistory) GetTlfHistory(
	tlfName tlf.CanonicalName, tlfType tlf.Type) (
	history keybase1.FSFolderEditHistory) {
	uh.lock.RLock()
	defer uh.lock.RUnlock()
	return uh.getTlfHistoryLocked(tlfName, tlfType)
}

// Clear erases all saved histories; TLFs must be re-added.
func (uh *UserHistory) Clear() {
	uh.lock.Lock()
	defer uh.lock.Unlock()
	uh.tlfs = tlfByTime{
		indices: make(map[tlfKey]int),
	}
}
<commit_msg>kbfsedits: add a Get() method to UserHistory<commit_after>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package kbfsedits

import (
	"fmt"
	"sort"
	"sync"

	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/kbfs/tlf"
)

const (
	// The max number of TLFs to return in a user history
	maxTlfs = 10
)

type tlfKey struct {
	tlfName tlf.CanonicalName
	tlfType tlf.Type
}

type tlfWithHistory struct {
	key     tlfKey
	history writersByRevision
}

type tlfByTime struct {
	histories []tlfWithHistory
	indices   map[tlfKey]int
}

func (tbm tlfByTime) Len() int {
	return len(tbm.histories)
}

func (tbm tlfByTime) Less(i, j int) bool {
	iHistory := tbm.histories[i].history
	jHistory := tbm.histories[j].history

	if len(iHistory) == 0 || len(iHistory[0].notifications) == 0 {
		return false
	} else if len(jHistory) == 0 || len(jHistory[0].notifications) == 0 {
		return true
	}

	iTime := iHistory[0].notifications[0].Time
	jTime := jHistory[0].notifications[0].Time

	return iTime.After(jTime)
}

func (tbm tlfByTime) Swap(i, j int) {
	tbm.histories[i], tbm.histories[j] = tbm.histories[j], tbm.histories[i]
	tbm.indices[tbm.histories[i].key] = i
	tbm.indices[tbm.histories[j].key] = j
}

// UserHistory keeps a sorted list of the top known TLF edit
// histories, and can convert those histories into keybase1 protocol
// structs.  TLF histories must be updated by an external caller
// whenever they change.
type UserHistory struct {
	lock sync.RWMutex
	tlfs tlfByTime
}

// NewUserHistory constructs a UserHistory instance.
func NewUserHistory() *UserHistory {
	return &UserHistory{
		tlfs: tlfByTime{
			indices: make(map[tlfKey]int),
		},
	}
}

// UpdateHistory should be called whenever the edit history for a
// given TLF gets new information.
func (uh *UserHistory) UpdateHistory(
	tlfName tlf.CanonicalName, tlfType tlf.Type, tlfHistory *TlfHistory) {
	history := tlfHistory.getHistory()
	key := tlfKey{tlfName, tlfType}

	uh.lock.Lock()
	defer uh.lock.Unlock()
	if currIndex, ok := uh.tlfs.indices[key]; ok {
		uh.tlfs.histories[currIndex].history = history
	} else {
		uh.tlfs.indices[key] = len(uh.tlfs.indices)
		uh.tlfs.histories = append(
			uh.tlfs.histories, tlfWithHistory{key, history})
	}
	sort.Sort(uh.tlfs)
}

func (uh *UserHistory) getTlfHistoryLocked(
	tlfName tlf.CanonicalName, tlfType tlf.Type) (
	history keybase1.FSFolderEditHistory) {
	key := tlfKey{tlfName, tlfType}
	currIndex, ok := uh.tlfs.indices[key]
	if !ok {
		return keybase1.FSFolderEditHistory{}
	}
	tlfHistory := uh.tlfs.histories[currIndex].history

	folder := keybase1.Folder{
		Name:       string(tlfName),
		FolderType: tlfType.FolderType(),
		Private:    tlfType == tlf.Private,
	}
	history.Folder = folder
	if len(tlfHistory) == 0 {
		return history
	}
	if len(tlfHistory[0].notifications) > 0 {
		history.ServerTime = keybase1.ToTime(
			tlfHistory[0].notifications[0].Time)
	}
	history.History = make(
		[]keybase1.FSFolderWriterEditHistory, len(tlfHistory))
	for i, wn := range tlfHistory {
		history.History[i].WriterName = wn.writerName
		history.History[i].Edits = make(
			[]keybase1.FSFolderWriterEdit, len(wn.notifications))
		for j, n := range wn.notifications {
			history.History[i].Edits[j].Filename = n.Filename
			history.History[i].Edits[j].ServerTime = keybase1.ToTime(n.Time)
			switch n.Type {
			case NotificationCreate:
				history.History[i].Edits[j].NotificationType =
					keybase1.FSNotificationType_FILE_CREATED
			case NotificationModify:
				history.History[i].Edits[j].NotificationType =
					keybase1.FSNotificationType_FILE_MODIFIED
			default:
				panic(fmt.Sprintf("Unknown notification type %s", n.Type))
			}
		}
	}
	return history
}

// GetTlfHistory returns the edit history of a given TLF, converted to
// keybase1 protocol structs.
func (uh *UserHistory) GetTlfHistory(
	tlfName tlf.CanonicalName, tlfType tlf.Type) (
	history keybase1.FSFolderEditHistory) {
	uh.lock.RLock()
	defer uh.lock.RUnlock()
	return uh.getTlfHistoryLocked(tlfName, tlfType)
}

// Get returns the full edit history for the user, converted to
// keybase1 protocol structs.
func (uh *UserHistory) Get() (history []keybase1.FSFolderEditHistory) {
	uh.lock.RLock()
	defer uh.lock.RUnlock()
	for _, h := range uh.tlfs.histories {
		history = append(history, uh.getTlfHistoryLocked(
			h.key.tlfName, h.key.tlfType))
		if len(history) == maxTlfs {
			break
		}
	}
	return history
}

// Clear erases all saved histories; TLFs must be re-added.
func (uh *UserHistory) Clear() {
	uh.lock.Lock()
	defer uh.lock.Unlock()
	uh.tlfs = tlfByTime{
		indices: make(map[tlfKey]int),
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package resolvers

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/limits"
	amassnet "github.com/OWASP/Amass/v3/net"
	amassdns "github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/miekg/dns"
)

var (
	retryCodes = []int{
		dns.RcodeRefused,
		dns.RcodeServerFailure,
		dns.RcodeNotImplemented,
	}

	maxRetries = 3
)

// ResolverPool manages many DNS resolvers for high-performance use, such as brute forcing attacks.
type ResolverPool struct {
	Resolvers []Resolver
	Done      chan struct{}
	// Logger for error messages
	Log          *log.Logger
	wildcardLock sync.Mutex
	wildcards    map[string]*wildcard
	// Domains discovered by the SubdomainToDomain function
	domainLock     sync.Mutex
	domainCache    map[string]struct{}
	hasBeenStopped bool
}

// SetupResolverPool initializes a ResolverPool with the type of resolvers indicated by the parameters.
func SetupResolverPool(addrs []string, scoring, ratemon bool, log *log.Logger) *ResolverPool {
	if len(addrs) <= 0 {
		return nil
	}

	// Do not allow the number of resolvers to exceed the ulimit
	temp := addrs
	addrs = []string{}
	max := int(float64(limits.GetFileLimit())*0.7) / 2
	for i, r := range temp {
		if i > max {
			break
		}
		addrs = append(addrs, r)
	}

	finished := make(chan Resolver, 100)
	for _, addr := range addrs {
		go func(ip string, ch chan Resolver) {
			if n := NewBaseResolver(ip); n != nil {
				ch <- n
				return
			}
			ch <- nil
		}(addr, finished)
	}

	l := len(addrs)
	var resolvers []Resolver
	t := time.NewTimer(5 * time.Second)
	defer t.Stop()
loop:
	for i := 0; i < l; i++ {
		select {
		case <-t.C:
			break loop
		case r := <-finished:
			if r == nil {
				continue loop
			}
			if scoring {
				if r = NewScoredResolver(r); r == nil {
					continue loop
				}
			}
			if ratemon {
				if r = NewRateMonitoredResolver(r); r == nil {
					continue loop
				}
			}
			resolvers = append(resolvers, r)
		}
	}

	if len(resolvers) == 0 {
		return nil
	}

	return NewResolverPool(resolvers, log)
}

// NewResolverPool initializes a ResolverPool that uses the provided Resolvers.
func NewResolverPool(res []Resolver, logger *log.Logger) *ResolverPool {
	rp := &ResolverPool{
		Resolvers:   res,
		Done:        make(chan struct{}, 2),
		Log:         logger,
		wildcards:   make(map[string]*wildcard),
		domainCache: make(map[string]struct{}),
	}

	// Assign a null logger when one is not provided
	if rp.Log == nil {
		rp.Log = log.New(ioutil.Discard, "", 0)
	}

	rp.SanityChecks()
	return rp
}

// Stop calls the Stop method for each Resolver object in the pool.
func (rp *ResolverPool) Stop() error {
	rp.hasBeenStopped = true

	for _, r := range rp.Resolvers {
		r.Stop()
	}

	rp.Resolvers = []Resolver{}
	return nil
}

// IsStopped implements the Resolver interface.
func (rp *ResolverPool) IsStopped() bool {
	return rp.hasBeenStopped
}

// Address implements the Resolver interface.
func (rp *ResolverPool) Address() string {
	return "N/A"
}

// Port implements the Resolver interface.
func (rp *ResolverPool) Port() int {
	return 0
}

// Available returns true if the Resolver can handle another DNS request.
func (rp *ResolverPool) Available() (bool, error) {
	return true, nil
}

// Stats returns performance counters.
func (rp *ResolverPool) Stats() map[int]int64 {
	stats := make(map[int]int64)

	for _, r := range rp.Resolvers {
		for k, v := range r.Stats() {
			if cur, found := stats[k]; found {
				stats[k] = cur + v
			} else {
				stats[k] = v
			}
		}
	}

	return stats
}

// WipeStats clears the performance counters.
func (rp *ResolverPool) WipeStats() {
	return
}

// ReportError implements the Resolver interface.
func (rp *ResolverPool) ReportError() {
	return
}

// SubdomainToDomain returns the first subdomain name of the provided
// parameter that responds to a DNS query for the NS record type.
func (rp *ResolverPool) SubdomainToDomain(name string) string {
	rp.domainLock.Lock()
	defer rp.domainLock.Unlock()

	var domain string
	// Obtain all parts of the subdomain name
	labels := strings.Split(strings.TrimSpace(name), ".")
	// Check the cache for all parts of the name
	for i := len(labels); i >= 0; i-- {
		sub := strings.Join(labels[i:], ".")

		if _, ok := rp.domainCache[sub]; ok {
			domain = sub
			break
		}
	}
	if domain != "" {
		return domain
	}
	// Check the DNS for all parts of the name
	for i := 0; i < len(labels)-1; i++ {
		sub := strings.Join(labels[i:], ".")

		if ns, _, err := rp.Resolve(context.TODO(), sub, "NS", PriorityHigh); err == nil {
			pieces := strings.Split(ns[0].Data, ",")
			rp.domainCache[pieces[0]] = struct{}{}
			domain = pieces[0]
			break
		}
	}
	return domain
}

// NextResolver returns a randomly selected Resolver from the pool that has availability.
func (rp *ResolverPool) NextResolver() Resolver {
	var attempts int
	max := len(rp.Resolvers)

	if max == 0 || rp.numUsableResolvers() == 0 {
		return nil
	}

	for {
		r := rp.Resolvers[rand.Int()%max]

		if stopped := r.IsStopped(); !stopped {
			return r
		}

		attempts++
		if attempts > max {
			// Check every resolver sequentially
			for _, r := range rp.Resolvers {
				if stopped := r.IsStopped(); !stopped {
					return r
				}
			}
			break
		}
	}

	return nil
}

// Reverse is performs reverse DNS queries using available Resolvers in the pool.
func (rp *ResolverPool) Reverse(ctx context.Context, addr string, priority int) (string, string, error) {
	var name, ptr string

	if ip := net.ParseIP(addr); amassnet.IsIPv4(ip) {
		ptr = amassdns.ReverseIP(addr) + ".in-addr.arpa"
	} else if amassnet.IsIPv6(ip) {
		ptr = amassdns.IPv6NibbleFormat(ip.String()) + ".ip6.arpa"
	} else {
		return ptr, "", &ResolveError{
			Err:   fmt.Sprintf("Invalid IP address parameter: %s", addr),
			Rcode: 100,
		}
	}

	answers, _, err := rp.Resolve(ctx, ptr, "PTR", priority)
	if err != nil {
		return ptr, name, err
	}

	for _, a := range answers {
		if a.Type == 12 {
			name = RemoveLastDot(a.Data)
			break
		}
	}

	if name == "" {
		err = &ResolveError{
			Err:   fmt.Sprintf("PTR record not found for IP address: %s", addr),
			Rcode: 100,
		}
	} else if strings.HasSuffix(name, ".in-addr.arpa") || strings.HasSuffix(name, ".ip6.arpa") {
		err = &ResolveError{
			Err:   fmt.Sprintf("Invalid target in PTR record answer: %s", name),
			Rcode: 100,
		}
	}

	return ptr, name, err
}

// Resolve performs a DNS request using available Resolvers in the pool.
func (rp *ResolverPool) Resolve(ctx context.Context, name, qtype string, priority int) ([]requests.DNSAnswer, bool, error) {
	var attempts int
	switch priority {
	case PriorityCritical:
		attempts = 1000
	case PriorityHigh:
		attempts = 100
	case PriorityLow:
		attempts = 25
	}

	// This loop ensures the correct number of attempts of the DNS query
	for count := 0; count < attempts; count++ {
		r := rp.NextResolver()
		if r == nil {
			// Give the system a chance to breathe before trying again
			time.Sleep(time.Duration(randomInt(1000, 1500)) * time.Millisecond)
			continue
		}

		ans, again, err := r.Resolve(ctx, name, qtype, priority)
		if again {
			continue
		}

		return ans, again, err
	}

	return []requests.DNSAnswer{}, false, &ResolveError{
		Err: fmt.Sprintf("Resolver: %d attempts for %s type %s returned 0 results", attempts, name, qtype),
	}
}

func (rp *ResolverPool) numUsableResolvers() int {
	var num int

	for _, r := range rp.Resolvers {
		if stopped := r.IsStopped(); !stopped {
			num++
		}
	}
	return num
}

func randomInt(min, max int) int {
	return min + rand.Intn((max-min)+1)
}
<commit_msg>fixed #349 DNS resolver issue<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package resolvers

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/limits"
	amassnet "github.com/OWASP/Amass/v3/net"
	amassdns "github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/miekg/dns"
)

var (
	retryCodes = []int{
		dns.RcodeRefused,
		dns.RcodeServerFailure,
		dns.RcodeNotImplemented,
	}

	maxRetries = 3
)

// ResolverPool manages many DNS resolvers for high-performance use, such as brute forcing attacks.
type ResolverPool struct {
	Resolvers []Resolver
	Done      chan struct{}
	// Logger for error messages
	Log          *log.Logger
	wildcardLock sync.Mutex
	wildcards    map[string]*wildcard
	// Domains discovered by the SubdomainToDomain function
	domainLock     sync.Mutex
	domainCache    map[string]struct{}
	hasBeenStopped bool
}

// SetupResolverPool initializes a ResolverPool with the type of resolvers indicated by the parameters.
func SetupResolverPool(addrs []string, scoring, ratemon bool, log *log.Logger) *ResolverPool {
	if len(addrs) <= 0 {
		return nil
	}

	// Do not allow the number of resolvers to exceed the ulimit
	temp := addrs
	addrs = []string{}
	max := int(float64(limits.GetFileLimit())*0.7) / 2
	for i, r := range temp {
		if i > max {
			break
		}
		addrs = append(addrs, r)
	}

	finished := make(chan Resolver, 100)
	for _, addr := range addrs {
		go func(ip string, ch chan Resolver) {
			if n := NewBaseResolver(ip); n != nil {
				ch <- n
				return
			}
			ch <- nil
		}(addr, finished)
	}

	l := len(addrs)
	var resolvers []Resolver
	t := time.NewTimer(5 * time.Second)
	defer t.Stop()
loop:
	for i := 0; i < l; i++ {
		select {
		case <-t.C:
			break loop
		case r := <-finished:
			if r == nil {
				continue loop
			}
			if scoring {
				if r = NewScoredResolver(r); r == nil {
					continue loop
				}
			}
			if ratemon {
				if r = NewRateMonitoredResolver(r); r == nil {
					continue loop
				}
			}
			resolvers = append(resolvers, r)
		}
	}

	if len(resolvers) == 0 {
		return nil
	}

	return NewResolverPool(resolvers, log)
}

// NewResolverPool initializes a ResolverPool that uses the provided Resolvers.
func NewResolverPool(res []Resolver, logger *log.Logger) *ResolverPool {
	rp := &ResolverPool{
		Resolvers:   res,
		Done:        make(chan struct{}, 2),
		Log:         logger,
		wildcards:   make(map[string]*wildcard),
		domainCache: make(map[string]struct{}),
	}

	// Assign a null logger when one is not provided
	if rp.Log == nil {
		rp.Log = log.New(ioutil.Discard, "", 0)
	}

	rp.SanityChecks()
	return rp
}

// Stop calls the Stop method for each Resolver object in the pool.
func (rp *ResolverPool) Stop() error {
	rp.hasBeenStopped = true

	for _, r := range rp.Resolvers {
		r.Stop()
	}

	rp.Resolvers = []Resolver{}
	return nil
}

// IsStopped implements the Resolver interface.
func (rp *ResolverPool) IsStopped() bool {
	return rp.hasBeenStopped
}

// Address implements the Resolver interface.
func (rp *ResolverPool) Address() string {
	return "N/A"
}

// Port implements the Resolver interface.
func (rp *ResolverPool) Port() int {
	return 0
}

// Available returns true if the Resolver can handle another DNS request.
func (rp *ResolverPool) Available() (bool, error) {
	return true, nil
}

// Stats returns performance counters.
func (rp *ResolverPool) Stats() map[int]int64 {
	stats := make(map[int]int64)

	for _, r := range rp.Resolvers {
		for k, v := range r.Stats() {
			if cur, found := stats[k]; found {
				stats[k] = cur + v
			} else {
				stats[k] = v
			}
		}
	}

	return stats
}

// WipeStats clears the performance counters.
func (rp *ResolverPool) WipeStats() {
	return
}

// ReportError implements the Resolver interface.
func (rp *ResolverPool) ReportError() {
	return
}

// SubdomainToDomain returns the first subdomain name of the provided
// parameter that responds to a DNS query for the NS record type.
func (rp *ResolverPool) SubdomainToDomain(name string) string {
	rp.domainLock.Lock()
	defer rp.domainLock.Unlock()

	var domain string
	// Obtain all parts of the subdomain name
	labels := strings.Split(strings.TrimSpace(name), ".")
	// Check the cache for all parts of the name
	for i := len(labels); i >= 0; i-- {
		sub := strings.Join(labels[i:], ".")

		if _, ok := rp.domainCache[sub]; ok {
			domain = sub
			break
		}
	}
	if domain != "" {
		return domain
	}
	// Check the DNS for all parts of the name
	for i := 0; i < len(labels)-1; i++ {
		sub := strings.Join(labels[i:], ".")

		if ns, _, err := rp.Resolve(context.TODO(), sub, "NS", PriorityHigh); err == nil {
			pieces := strings.Split(ns[0].Data, ",")
			rp.domainCache[pieces[0]] = struct{}{}
			domain = pieces[0]
			break
		}
	}
	return domain
}

// NextResolver returns a randomly selected Resolver from the pool that has availability.
func (rp *ResolverPool) NextResolver() Resolver {
	var attempts int
	max := len(rp.Resolvers)

	if max == 0 || rp.numUsableResolvers() == 0 {
		return nil
	}

	for {
		r := rp.Resolvers[rand.Int()%max]

		if stopped := r.IsStopped(); !stopped {
			return r
		}

		attempts++
		if attempts > max {
			// Check every resolver sequentially
			for _, r := range rp.Resolvers {
				if stopped := r.IsStopped(); !stopped {
					return r
				}
			}
			break
		}
	}

	return nil
}

// Reverse is performs reverse DNS queries using available Resolvers in the pool.
func (rp *ResolverPool) Reverse(ctx context.Context, addr string, priority int) (string, string, error) {
	var name, ptr string

	if ip := net.ParseIP(addr); amassnet.IsIPv4(ip) {
		ptr = amassdns.ReverseIP(addr) + ".in-addr.arpa"
	} else if amassnet.IsIPv6(ip) {
		ptr = amassdns.IPv6NibbleFormat(ip.String()) + ".ip6.arpa"
	} else {
		return ptr, "", &ResolveError{
			Err:   fmt.Sprintf("Invalid IP address parameter: %s", addr),
			Rcode: 100,
		}
	}

	answers, _, err := rp.Resolve(ctx, ptr, "PTR", priority)
	if err != nil {
		return ptr, name, err
	}

	for _, a := range answers {
		if a.Type == 12 {
			name = RemoveLastDot(a.Data)
			break
		}
	}

	if name == "" {
		err = &ResolveError{
			Err:   fmt.Sprintf("PTR record not found for IP address: %s", addr),
			Rcode: 100,
		}
	} else if strings.HasSuffix(name, ".in-addr.arpa") || strings.HasSuffix(name, ".ip6.arpa") {
		err = &ResolveError{
			Err:   fmt.Sprintf("Invalid target in PTR record answer: %s", name),
			Rcode: 100,
		}
	}

	return ptr, name, err
}

// Resolve performs a DNS request using available Resolvers in the pool.
func (rp *ResolverPool) Resolve(ctx context.Context, name, qtype string, priority int) ([]requests.DNSAnswer, bool, error) {
	var attempts int
	switch priority {
	case PriorityCritical:
		attempts = 1000
	case PriorityHigh:
		attempts = 100
	case PriorityLow:
		attempts = 25
	}

	// This loop ensures the correct number of attempts of the DNS query
	for count := 0; count < attempts; count++ {
		r := rp.NextResolver()
		if r == nil {
			// Give the system a chance to breathe before trying again
			time.Sleep(time.Duration(randomInt(1000, 1500)) * time.Millisecond)
			continue
		}

		success := true
		ans, again, err := r.Resolve(ctx, name, qtype, priority)
		if again {
			success = false
		} else if err != nil {
			if rc := (err.(*ResolveError)).Rcode; rc == NotAvailableRcode ||
				rc == dns.RcodeServerFailure || rc == dns.RcodeRefused || rc == dns.RcodeNotImplemented {
				success = false
			}
		}

		if success {
			return ans, again, err
		}
	}

	return []requests.DNSAnswer{}, false, &ResolveError{
		Err: fmt.Sprintf("Resolver: %d attempts for %s type %s returned 0 results", attempts, name, qtype),
	}
}

func (rp *ResolverPool) numUsableResolvers() int {
	var num int

	for _, r := range rp.Resolvers {
		if stopped := r.IsStopped(); !stopped {
			num++
		}
	}
	return num
}

func randomInt(min, max int) int {
	return min + rand.Intn((max-min)+1)
}
<|endoftext|>


<commit_before>package manta

import (
	"sort"
	"testing"

	"github.com/dotabuff/manta/dota"
	"github.com/stretchr/testify/assert"
)

func TestParseOneMatch(t *testing.T) {
	assert := assert.New(t)

	buf := mustGetReplayData("1560315800", "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560315800.dem")
	parser, err := NewParser(buf)
	assert.Nil(err)
	err = parser.Start()
	assert.Nil(err)
}

// Simply tests that we can read the outer structure of a real match
func TestParseRealMatches(t *testing.T) {
	assert := assert.New(t)

	// Important: make sure to include the most recent test last. These generate fixtures to easily
	// detect diffs in various data structures upon commit, and the latest replay should always be
	// last to provide a most up-to-date baseline.
	scenarios := []struct {
		matchId                string
		replayUrl              string
		expectCombatLogDamage  int32
		expectCombatLogHealing int32
		expectCombatLogDeaths  int32
		expectCombatLogEvents  int32
		expectUnitOrderEvents  int32
	}{
		{
			matchId:                "1560289528",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560289528.dem",
			expectCombatLogDamage:  1180993,
			expectCombatLogHealing: 57511,
			expectCombatLogDeaths:  1449,
			expectCombatLogEvents:  49146,
			expectUnitOrderEvents:  63387,
		},
		{
			matchId:                "1560294294",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560294294.dem",
			expectCombatLogDamage:  768154,
			expectCombatLogHealing: 11565,
			expectCombatLogDeaths:  954,
			expectCombatLogEvents:  24535,
			expectUnitOrderEvents:  30657,
		},
		{
			matchId:                "1560315800",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560315800.dem",
			expectCombatLogDamage:  1332418,
			expectCombatLogHealing: 57874,
			expectCombatLogDeaths:  1645,
			expectCombatLogEvents:  51288,
			expectUnitOrderEvents:  63992,
		},
		{
			matchId:                "1582611189",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1582611189.dem",
			expectCombatLogDamage:  599388,
			expectCombatLogHealing: 28576,
			expectCombatLogDeaths:  930,
			expectCombatLogEvents:  23800,
			expectUnitOrderEvents:  40237,
		},
		{
			matchId:                "1648457986",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1648457986.dem",
			expectCombatLogDamage:  224773,
			expectCombatLogHealing: 5914,
			expectCombatLogDeaths:  466,
			expectCombatLogEvents:  10170,
			expectUnitOrderEvents:  17822,
		},
		{
			matchId:                "1712853372",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1712853372.dem",
			expectCombatLogDamage:  671297,
			expectCombatLogHealing: 23467,
			expectCombatLogDeaths:  1099,
			expectCombatLogEvents:  30381,
			expectUnitOrderEvents:  48107,
		},
	}

	for _, s := range scenarios {
		buf := mustGetReplayData(s.matchId, s.replayUrl)
		parser, err := NewParser(buf)
		if err != nil {
			t.Errorf("unable to instantiate parser: %s", err)
			continue
		}

		gotCombatLogDamage := int32(0)
		gotCombatLogHealing := int32(0)
		gotCombatLogDeaths := int32(0)
		gotCombatLogEvents := int32(0)
		gotUnitOrderEvents := int32(0)

		parser.Callbacks.OnCDOTAUserMsg_SpectatorPlayerUnitOrders(func(m *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) error {
			gotUnitOrderEvents += 1
			return nil
		})

		parser.OnGameEvent("dota_combatlog", func(m *GameEvent) error {
			gotCombatLogEvents += 1

			t, err := m.GetInt32("type")
			assert.Nil(err)

			switch dota.DOTA_COMBATLOG_TYPES(t) {
			case dota.DOTA_COMBATLOG_TYPES_DOTA_COMBATLOG_DAMAGE:
				v, err := m.GetInt32("value")
				assert.Nil(err)
				gotCombatLogDamage += v
			case dota.DOTA_COMBATLOG_TYPES_DOTA_COMBATLOG_DEATH:
				gotCombatLogDeaths += 1
			case dota.DOTA_COMBATLOG_TYPES_DOTA_COMBATLOG_HEAL:
				v, err := m.GetInt32("value")
				assert.Nil(err)
				gotCombatLogHealing += v
			}

			return nil
		})

		if fixturesMode {
			// Writes out the source_1_legacy_game_events_list.json fixture so that we can identify changes to schema.
			parser.Callbacks.OnCMsgSource1LegacyGameEventList(func(m *dota.CMsgSource1LegacyGameEventList) error {
				_dump_fixture("legacy_game_events/list_latest.json", _json_marshal(m))
				return nil
			})
		}

		err = parser.Start()
		assert.Nil(err, s.matchId)

		if fixturesMode {
			// Use this to write out instancebaseline fixtures
			t, _ := parser.StringTables.GetTableByName("instancebaseline")
			for _, i := range t.Items {
				classId, _ := atoi32(i.Key)
				className := parser.ClassInfo[classId]
				_dump_fixture(_sprintf("instancebaseline/%s_%s.rawbuf", className, s.matchId), i.Value)
			}

			// Use this to dump layout of sendtables
			out := _sprintf("")
			classIds := make([]int, 0)
			for classId, _ := range parser.ClassInfo {
				classIds = append(classIds, int(classId))
			}
			sort.Ints(classIds)
			for _, classId := range classIds {
				className := parser.ClassInfo[int32(classId)]
				st := parser.SendTables.Tables[className]
				out += _sprintf("class id=%d name=%s\n", classId, className)
				out += _sprintf("table id=%d name=%s version=%d\n", st.Index, st.Name, st.Version)
				for i, p := range st.Props {
					out += _sprintf(" -> prop %d: %s\n", i, p.Describe())
				}
				out += "\n"
			}
		}

		assert.Equal(s.expectCombatLogDamage, gotCombatLogDamage, s.matchId)
		assert.Equal(s.expectCombatLogHealing, gotCombatLogHealing, s.matchId)
		assert.Equal(s.expectCombatLogDeaths, gotCombatLogDeaths, s.matchId)
		assert.Equal(s.expectCombatLogEvents, gotCombatLogEvents, s.matchId)
		assert.Equal(s.expectUnitOrderEvents, gotUnitOrderEvents, s.matchId)
	}
}

func BenchmarkParseMatch(b *testing.B) {
	assert := assert.New(b)

	buf := mustGetReplayData("1560315800", "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560315800.dem")

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		parser, err := NewParser(buf)
		assert.Nil(err)
		err = parser.Start()
		assert.Nil(err)
	}

	b.ReportAllocs()
}
<commit_msg>Add new replay<commit_after>package manta

import (
	"sort"
	"testing"

	"github.com/dotabuff/manta/dota"
	"github.com/stretchr/testify/assert"
)

func TestParseOneMatch(t *testing.T) {
	assert := assert.New(t)

	buf := mustGetReplayData("1716444111", "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1716444111.dem")
	parser, err := NewParser(buf)
	assert.Nil(err)
	err = parser.Start()
	assert.Nil(err)
}

// Simply tests that we can read the outer structure of a real match
func TestParseRealMatches(t *testing.T) {
	assert := assert.New(t)

	// Important: make sure to include the most recent test last. These generate fixtures to easily
	// detect diffs in various data structures upon commit, and the latest replay should always be
	// last to provide a most up-to-date baseline.
	scenarios := []struct {
		matchId                string
		replayUrl              string
		expectCombatLogDamage  int32
		expectCombatLogHealing int32
		expectCombatLogDeaths  int32
		expectCombatLogEvents  int32
		expectUnitOrderEvents  int32
	}{
		{
			matchId:                "1560289528",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560289528.dem",
			expectCombatLogDamage:  1180993,
			expectCombatLogHealing: 57511,
			expectCombatLogDeaths:  1449,
			expectCombatLogEvents:  49146,
			expectUnitOrderEvents:  63387,
		},
		{
			matchId:                "1560294294",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560294294.dem",
			expectCombatLogDamage:  768154,
			expectCombatLogHealing: 11565,
			expectCombatLogDeaths:  954,
			expectCombatLogEvents:  24535,
			expectUnitOrderEvents:  30657,
		},
		{
			matchId:                "1560315800",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560315800.dem",
			expectCombatLogDamage:  1332418,
			expectCombatLogHealing: 57874,
			expectCombatLogDeaths:  1645,
			expectCombatLogEvents:  51288,
			expectUnitOrderEvents:  63992,
		},
		{
			matchId:                "1582611189",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1582611189.dem",
			expectCombatLogDamage:  599388,
			expectCombatLogHealing: 28576,
			expectCombatLogDeaths:  930,
			expectCombatLogEvents:  23800,
			expectUnitOrderEvents:  40237,
		},
		{
			matchId:                "1648457986",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1648457986.dem",
			expectCombatLogDamage:  224773,
			expectCombatLogHealing: 5914,
			expectCombatLogDeaths:  466,
			expectCombatLogEvents:  10170,
			expectUnitOrderEvents:  17822,
		},
		{
			matchId:                "1712853372",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1712853372.dem",
			expectCombatLogDamage:  671297,
			expectCombatLogHealing: 23467,
			expectCombatLogDeaths:  1099,
			expectCombatLogEvents:  30381,
			expectUnitOrderEvents:  48107,
		},
		{
			matchId:                "1716444111",
			replayUrl:              "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1716444111.dem",
			expectCombatLogDamage:  1398735,
			expectCombatLogHealing: 49659,
			expectCombatLogDeaths:  2169,
			expectCombatLogEvents:  76921,
			expectUnitOrderEvents:  48822,
		},
	}

	for _, s := range scenarios {
		buf := mustGetReplayData(s.matchId, s.replayUrl)
		parser, err := NewParser(buf)
		if err != nil {
			t.Errorf("unable to instantiate parser: %s", err)
			continue
		}

		gotCombatLogDamage := int32(0)
		gotCombatLogHealing := int32(0)
		gotCombatLogDeaths := int32(0)
		gotCombatLogEvents := int32(0)
		gotUnitOrderEvents := int32(0)

		parser.Callbacks.OnCDOTAUserMsg_SpectatorPlayerUnitOrders(func(m *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) error {
			gotUnitOrderEvents += 1
			return nil
		})

		parser.OnGameEvent("dota_combatlog", func(m *GameEvent) error {
			gotCombatLogEvents += 1

			t, err := m.GetInt32("type")
			assert.Nil(err)

			switch dota.DOTA_COMBATLOG_TYPES(t) {
			case dota.DOTA_COMBATLOG_TYPES_DOTA_COMBATLOG_DAMAGE:
				v, err := m.GetInt32("value")
				assert.Nil(err)
				gotCombatLogDamage += v
			case dota.DOTA_COMBATLOG_TYPES_DOTA_COMBATLOG_DEATH:
				gotCombatLogDeaths += 1
			case dota.DOTA_COMBATLOG_TYPES_DOTA_COMBATLOG_HEAL:
				v, err := m.GetInt32("value")
				assert.Nil(err)
				gotCombatLogHealing += v
			}

			return nil
		})

		if fixturesMode {
			// Writes out the source_1_legacy_game_events_list.json fixture so that we can identify changes to schema.
			parser.Callbacks.OnCMsgSource1LegacyGameEventList(func(m *dota.CMsgSource1LegacyGameEventList) error {
				_dump_fixture("legacy_game_events/list_latest.json", _json_marshal(m))
				return nil
			})
		}

		err = parser.Start()
		assert.Nil(err, s.matchId)

		if fixturesMode {
			// Use this to write out instancebaseline fixtures
			t, _ := parser.StringTables.GetTableByName("instancebaseline")
			for _, i := range t.Items {
				classId, _ := atoi32(i.Key)
				className := parser.ClassInfo[classId]
				_dump_fixture(_sprintf("instancebaseline/%s_%s.rawbuf", className, s.matchId), i.Value)
			}

			// Use this to dump layout of sendtables
			out := _sprintf("")
			classIds := make([]int, 0)
			for classId, _ := range parser.ClassInfo {
				classIds = append(classIds, int(classId))
			}
			sort.Ints(classIds)
			for _, classId := range classIds {
				className := parser.ClassInfo[int32(classId)]
				st := parser.SendTables.Tables[className]
				out += _sprintf("class id=%d name=%s\n", classId, className)
				out += _sprintf("table id=%d name=%s version=%d\n", st.Index, st.Name, st.Version)
				for i, p := range st.Props {
					out += _sprintf(" -> prop %d: %s\n", i, p.Describe())
				}
				out += "\n"
			}
		}

		assert.Equal(s.expectCombatLogDamage, gotCombatLogDamage, s.matchId)
		assert.Equal(s.expectCombatLogHealing, gotCombatLogHealing, s.matchId)
		assert.Equal(s.expectCombatLogDeaths, gotCombatLogDeaths, s.matchId)
		assert.Equal(s.expectCombatLogEvents, gotCombatLogEvents, s.matchId)
		assert.Equal(s.expectUnitOrderEvents, gotUnitOrderEvents, s.matchId)
	}
}

func BenchmarkParseMatch(b *testing.B) {
	assert := assert.New(b)

	buf := mustGetReplayData("1560315800", "https://s3-us-west-2.amazonaws.com/manta.dotabuff/1560315800.dem")

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		parser, err := NewParser(buf)
		assert.Nil(err)
		err = parser.Start()
		assert.Nil(err)
	}

	b.ReportAllocs()
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/service/directconnect"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
)

func resourceAwsDxConnection() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsDxConnectionCreate,
		Read:   resourceAwsDxConnectionRead,
		Update: resourceAwsDxConnectionUpdate,
		Delete: resourceAwsDxConnectionDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"bandwidth": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateDxConnectionBandWidth(),
			},
			"location": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"jumbo_frame_capable": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"tags": tagsSchema(),
			"has_logical_redundancy": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"aws_device": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsDxConnectionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	req := &directconnect.CreateConnectionInput{
		Bandwidth:      aws.String(d.Get("bandwidth").(string)),
		ConnectionName: aws.String(d.Get("name").(string)),
		Location:       aws.String(d.Get("location").(string)),
	}

	log.Printf("[DEBUG] Creating Direct Connect connection: %#v", req)
	resp, err := conn.CreateConnection(req)
	if err != nil {
		return err
	}

	d.SetId(aws.StringValue(resp.ConnectionId))
	return resourceAwsDxConnectionUpdate(d, meta)
}

func resourceAwsDxConnectionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	resp, err := conn.DescribeConnections(&directconnect.DescribeConnectionsInput{
		ConnectionId: aws.String(d.Id()),
	})
	if err != nil {
		if isNoSuchDxConnectionErr(err) {
			log.Printf("[WARN] Direct Connect connection (%s) not found, removing from state", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	if len(resp.Connections) < 1 {
		log.Printf("[WARN] Direct Connect connection (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}
	if len(resp.Connections) != 1 {
		return fmt.Errorf("Number of Direct Connect connections (%s) isn't one, got %d", d.Id(), len(resp.Connections))
	}
	connection := resp.Connections[0]
	if d.Id() != aws.StringValue(connection.ConnectionId) {
		return fmt.Errorf("Direct Connect connection (%s) not found", d.Id())
	}
	if aws.StringValue(connection.ConnectionState) == directconnect.ConnectionStateDeleted {
		log.Printf("[WARN] Direct Connect connection (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	arn := arn.ARN{
		Partition: meta.(*AWSClient).partition,
		Region:    meta.(*AWSClient).region,
		Service:   "directconnect",
		AccountID: meta.(*AWSClient).accountid,
		Resource:  fmt.Sprintf("dxcon/%s", d.Id()),
	}.String()
	d.Set("arn", arn)
	d.Set("name", connection.ConnectionName)
	d.Set("bandwidth", connection.Bandwidth)
	d.Set("location", connection.Location)
	d.Set("jumbo_frame_capable", connection.JumboFrameCapable)
	d.Set("has_logical_redundancy", connection.HasLogicalRedundancy)
	d.Set("aws_device", connection.AwsDeviceV2)

	err1 := getTagsDX(conn, d, arn)
	return err1
}

func resourceAwsDxConnectionUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	arn := arn.ARN{
		Partition: meta.(*AWSClient).partition,
		Region:    meta.(*AWSClient).region,
		Service:   "directconnect",
		AccountID: meta.(*AWSClient).accountid,
		Resource:  fmt.Sprintf("dxcon/%s", d.Id()),
	}.String()
	if err := setTagsDX(conn, d, arn); err != nil {
		return err
	}

	return resourceAwsDxConnectionRead(d, meta)
}

func resourceAwsDxConnectionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	log.Printf("[DEBUG] Deleting Direct Connect connection: %s", d.Id())
	_, err := conn.DeleteConnection(&directconnect.DeleteConnectionInput{
		ConnectionId: aws.String(d.Id()),
	})
	if err != nil {
		if isNoSuchDxConnectionErr(err) {
			return nil
		}
		return err
	}

	deleteStateConf := &resource.StateChangeConf{
		Pending:    []string{directconnect.ConnectionStatePending, directconnect.ConnectionStateOrdering, directconnect.ConnectionStateAvailable, directconnect.ConnectionStateRequested, directconnect.ConnectionStateDeleting},
		Target:     []string{directconnect.ConnectionStateDeleted},
		Refresh:    dxConnectionRefreshStateFunc(conn, d.Id()),
		Timeout:    10 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}
	_, err = deleteStateConf.WaitForState()
	if err != nil {
		return fmt.Errorf("Error waiting for Direct Connect connection (%s) to be deleted: %s", d.Id(), err)
	}

	return nil
}

func dxConnectionRefreshStateFunc(conn *directconnect.DirectConnect, connId string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		input := &directconnect.DescribeConnectionsInput{
			ConnectionId: aws.String(connId),
		}
		resp, err := conn.DescribeConnections(input)
		if err != nil {
			return nil, "failed", err
		}
		if len(resp.Connections) < 1 {
			return resp, directconnect.ConnectionStateDeleted, nil
		}
		return resp, *resp.Connections[0].ConnectionState, nil
	}
}

func isNoSuchDxConnectionErr(err error) bool {
	return isAWSErr(err, "DirectConnectClientException", "Could not find Connection with ID")
}
<commit_msg>r/aws_dx_connection: Refactor tagging logic to keyvaluetags package.<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/service/directconnect"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/terraform-providers/terraform-provider-aws/aws/internal/keyvaluetags"
)

func resourceAwsDxConnection() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsDxConnectionCreate,
		Read:   resourceAwsDxConnectionRead,
		Update: resourceAwsDxConnectionUpdate,
		Delete: resourceAwsDxConnectionDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"bandwidth": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateDxConnectionBandWidth(),
			},
			"location": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"jumbo_frame_capable": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"tags": tagsSchema(),
			"has_logical_redundancy": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"aws_device": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsDxConnectionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	req := &directconnect.CreateConnectionInput{
		Bandwidth:      aws.String(d.Get("bandwidth").(string)),
		ConnectionName: aws.String(d.Get("name").(string)),
		Location:       aws.String(d.Get("location").(string)),
	}

	if v := d.Get("tags").(map[string]interface{}); len(v) > 0 {
		req.Tags = keyvaluetags.New(v).IgnoreAws().DirectconnectTags()
	}

	log.Printf("[DEBUG] Creating Direct Connect connection: %#v", req)
	resp, err := conn.CreateConnection(req)
	if err != nil {
		return err
	}

	d.SetId(aws.StringValue(resp.ConnectionId))

	return resourceAwsDxConnectionRead(d, meta)
}

func resourceAwsDxConnectionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	resp, err := conn.DescribeConnections(&directconnect.DescribeConnectionsInput{
		ConnectionId: aws.String(d.Id()),
	})
	if err != nil {
		if isNoSuchDxConnectionErr(err) {
			log.Printf("[WARN] Direct Connect connection (%s) not found, removing from state", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	if len(resp.Connections) < 1 {
		log.Printf("[WARN] Direct Connect connection (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}
	if len(resp.Connections) != 1 {
		return fmt.Errorf("Number of Direct Connect connections (%s) isn't one, got %d", d.Id(), len(resp.Connections))
	}
	connection := resp.Connections[0]
	if d.Id() != aws.StringValue(connection.ConnectionId) {
		return fmt.Errorf("Direct Connect connection (%s) not found", d.Id())
	}
	if aws.StringValue(connection.ConnectionState) == directconnect.ConnectionStateDeleted {
		log.Printf("[WARN] Direct Connect connection (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	arn := arn.ARN{
		Partition: meta.(*AWSClient).partition,
		Region:    meta.(*AWSClient).region,
		Service:   "directconnect",
		AccountID: meta.(*AWSClient).accountid,
		Resource:  fmt.Sprintf("dxcon/%s", d.Id()),
	}.String()
	d.Set("arn", arn)
	d.Set("name", connection.ConnectionName)
	d.Set("bandwidth", connection.Bandwidth)
	d.Set("location", connection.Location)
	d.Set("jumbo_frame_capable", connection.JumboFrameCapable)
	d.Set("has_logical_redundancy", connection.HasLogicalRedundancy)
	d.Set("aws_device", connection.AwsDeviceV2)

	tags, err := keyvaluetags.DirectconnectListTags(conn, arn)

	if err != nil {
		return fmt.Errorf("error listing tags for Direct Connect connection (%s): %s", arn, err)
	}

	if err := d.Set("tags", tags.IgnoreAws().Map()); err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}

func resourceAwsDxConnectionUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	arn := d.Get("arn").(string)
	if d.HasChange("tags") {
		o, n := d.GetChange("tags")

		if err := keyvaluetags.DirectconnectUpdateTags(conn, arn, o, n); err != nil {
			return fmt.Errorf("error updating Direct Connect connection (%s) tags: %s", arn, err)
		}
	}

	return resourceAwsDxConnectionRead(d, meta)
}

func resourceAwsDxConnectionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).dxconn

	log.Printf("[DEBUG] Deleting Direct Connect connection: %s", d.Id())
	_, err := conn.DeleteConnection(&directconnect.DeleteConnectionInput{
		ConnectionId: aws.String(d.Id()),
	})
	if err != nil {
		if isNoSuchDxConnectionErr(err) {
			return nil
		}
		return err
	}

	deleteStateConf := &resource.StateChangeConf{
		Pending:    []string{directconnect.ConnectionStatePending, directconnect.ConnectionStateOrdering, directconnect.ConnectionStateAvailable, directconnect.ConnectionStateRequested, directconnect.ConnectionStateDeleting},
		Target:     []string{directconnect.ConnectionStateDeleted},
		Refresh:    dxConnectionRefreshStateFunc(conn, d.Id()),
		Timeout:    10 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}
	_, err = deleteStateConf.WaitForState()
	if err != nil {
		return fmt.Errorf("Error waiting for Direct Connect connection (%s) to be deleted: %s", d.Id(), err)
	}

	return nil
}

func dxConnectionRefreshStateFunc(conn *directconnect.DirectConnect, connId string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		input := &directconnect.DescribeConnectionsInput{
			ConnectionId: aws.String(connId),
		}
		resp, err := conn.DescribeConnections(input)
		if err != nil {
			return nil, "failed", err
		}
		if len(resp.Connections) < 1 {
			return resp, directconnect.ConnectionStateDeleted, nil
		}
		return resp, *resp.Connections[0].ConnectionState, nil
	}
}

func isNoSuchDxConnectionErr(err error) bool {
	return isAWSErr(err, "DirectConnectClientException", "Could not find Connection with ID")
}
<|endoftext|>


<commit_before>package envoy

import (
	"context"
	"errors"
	"os"
	"testing"
	"time"

	"github.com/cilium/cilium/pkg/completion"
	"github.com/cilium/cilium/pkg/policy"
	"github.com/cilium/cilium/pkg/policy/api"

	"github.com/sirupsen/logrus"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

type EnvoySuite struct {
	waitGroup *completion.WaitGroup
}

var _ = Suite(&EnvoySuite{})

type testRedirect struct {
	name string
}

func (t *testRedirect) Log(pblog *HttpLogEntry) {
	log.Infof("%s/%s: Access log message: %s", t.name, pblog.CiliumResourceName, pblog.String())
}

func (s *EnvoySuite) waitForProxyCompletion() error {
	start := time.Now()
	log.Debug("Waiting for proxy updates to complete...")
	err := s.waitGroup.Wait()
	if err != nil {
		return errors.New("proxy state changes failed")
	}
	log.Debug("Wait time for proxy updates: ", time.Since(start))
	return nil
}

func (s *EnvoySuite) TestEnvoy(c *C) {
	log.SetLevel(logrus.DebugLevel)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	s.waitGroup = completion.NewWaitGroup(ctx)

	if os.Getenv("CILIUM_ENABLE_ENVOY_UNIT_TEST") == "" {
		c.Skip("skipping envoy unit test; CILIUM_ENABLE_ENVOY_UNIT_TEST not set")
	}

	stateLogDir := c.MkDir()

	// launch debug variant of the Envoy proxy
	Envoy := StartEnvoy(9942, stateLogDir, stateLogDir, 42)
	c.Assert(Envoy, NotNil)

	sel := api.NewWildcardEndpointSelector()

	// TODO: Test for success once we get feedback from Envoy.
	Envoy.AddListener("listener1", 8081, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Path: "foo"},
			{Method: "POST"},
			{Host: "cilium"},
			{Headers: []string{"via"}}}}},
		true, &testRedirect{name: "listener1"}, s.waitGroup)
	Envoy.AddListener("listener2", 8082, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Headers: []string{"via", "x-foo: bar"}}}}},
		true, &testRedirect{name: "listener2"}, s.waitGroup)
	Envoy.AddListener("listener3", 8083, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Method: "GET", Path: ".*public"}}}},
		false, &testRedirect{name: "listener3"}, s.waitGroup)

	err := s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Update listener2
	Envoy.UpdateListener("listener2", policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Headers: []string{"via: home", "x-foo: bar"}}}}}, s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Update listener1
	Envoy.UpdateListener("listener1", policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Headers: []string{"via"}}}}}, s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Remove listener3
	Envoy.RemoveListener("listener3", s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Add listener3 again
	Envoy.AddListener("listener3", 8083, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Method: "GET", Path: ".*public"}}}},
		false, &testRedirect{name: "listener3"}, s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Remove listener3 again, and wait for timeout after stopping Envoy.
	Envoy.RemoveListener("listener3", s.waitGroup)
	err = Envoy.StopEnvoy()
	c.Assert(err, IsNil)
	err = s.waitForProxyCompletion()
	c.Assert(err, NotNil)
	log.Debug("Proxy updates failed: ", err)
}
<commit_msg>envoy: Add logging into envoy_test.go<commit_after>package envoy

import (
	"context"
	"errors"
	"os"
	"testing"
	"time"

	"github.com/cilium/cilium/pkg/completion"
	"github.com/cilium/cilium/pkg/policy"
	"github.com/cilium/cilium/pkg/policy/api"

	"github.com/sirupsen/logrus"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

type EnvoySuite struct {
	waitGroup *completion.WaitGroup
}

var _ = Suite(&EnvoySuite{})

type testRedirect struct {
	name string
}

func (t *testRedirect) Log(pblog *HttpLogEntry) {
	log.Infof("%s/%s: Access log message: %s", t.name, pblog.CiliumResourceName, pblog.String())
}

func (s *EnvoySuite) waitForProxyCompletion() error {
	start := time.Now()
	log.Debug("Waiting for proxy updates to complete...")
	err := s.waitGroup.Wait()
	if err != nil {
		return errors.New("proxy state changes failed")
	}
	log.Debug("Wait time for proxy updates: ", time.Since(start))
	return nil
}

func (s *EnvoySuite) TestEnvoy(c *C) {
	log.SetLevel(logrus.DebugLevel)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	s.waitGroup = completion.NewWaitGroup(ctx)

	if os.Getenv("CILIUM_ENABLE_ENVOY_UNIT_TEST") == "" {
		c.Skip("skipping envoy unit test; CILIUM_ENABLE_ENVOY_UNIT_TEST not set")
	}

	stateLogDir := c.MkDir()
	log.Debugf("state log directory: %s", stateLogDir)

	// launch debug variant of the Envoy proxy
	Envoy := StartEnvoy(9942, stateLogDir, stateLogDir, 42)
	c.Assert(Envoy, NotNil)
	log.Debug("started Envoy")

	sel := api.NewWildcardEndpointSelector()

	// TODO: Test for success once we get feedback from Envoy.
	log.Debug("adding listener1")
	Envoy.AddListener("listener1", 8081, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Path: "foo"},
			{Method: "POST"},
			{Host: "cilium"},
			{Headers: []string{"via"}}}}},
		true, &testRedirect{name: "listener1"}, s.waitGroup)

	log.Debug("adding listener2")
	Envoy.AddListener("listener2", 8082, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Headers: []string{"via", "x-foo: bar"}}}}},
		true, &testRedirect{name: "listener2"}, s.waitGroup)

	log.Debug("adding listener3")
	Envoy.AddListener("listener3", 8083, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Method: "GET", Path: ".*public"}}}},
		false, &testRedirect{name: "listener3"}, s.waitGroup)

	err := s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	log.Debug("completed adding listener1, listener2, listener3")
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Update listener2
	log.Debug("updating listener 2")
	Envoy.UpdateListener("listener2", policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Headers: []string{"via: home", "x-foo: bar"}}}}}, s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	log.Debug("completed updating listener 2")
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Update listener1
	log.Debug("updating listener 1")
	Envoy.UpdateListener("listener1", policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Headers: []string{"via"}}}}}, s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	log.Debug("completed updating listener 1")
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Remove listener3
	log.Debug("removing listener 3")
	Envoy.RemoveListener("listener3", s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	log.Debug("completed removing listener 3")
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Add listener3 again
	log.Debug("adding listener 3")
	Envoy.AddListener("listener3", 8083, policy.L7DataMap{
		sel: api.L7Rules{HTTP: []api.PortRuleHTTP{
			{Method: "GET", Path: ".*public"}}}},
		false, &testRedirect{name: "listener3"}, s.waitGroup)

	err = s.waitForProxyCompletion()
	c.Assert(err, IsNil)
	log.Debug("completed adding listener 3")
	s.waitGroup = completion.NewWaitGroup(ctx)

	// Remove listener3 again, and wait for timeout after stopping Envoy.
	log.Debug("removing listener 3")
	Envoy.RemoveListener("listener3", s.waitGroup)
	err = Envoy.StopEnvoy()
	c.Assert(err, IsNil)
	err = s.waitForProxyCompletion()
	c.Assert(err, NotNil)
	log.Debugf("failed to remove listener 3: %s", err)
}
<|endoftext|>