<commit_before>// Copyright 2020 The Cloudprober 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 grpc implements a gRPC probe.

This probes a cloudprober gRPC server and reports success rate, latency, and
validation failures.
*/
package grpc

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strconv"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/cloudprober/common/oauth"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	configpb "github.com/google/cloudprober/probes/grpc/proto"
	"github.com/google/cloudprober/probes/options"
	"github.com/google/cloudprober/probes/probeutils"
	"github.com/google/cloudprober/sysvars"
	"github.com/google/cloudprober/targets/endpoint"

	grpcprobepb "github.com/google/cloudprober/servers/grpc/proto"
	servicepb "github.com/google/cloudprober/servers/grpc/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/alts"
	grpcoauth "google.golang.org/grpc/credentials/oauth"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/resolver"

	// Import grpclb module so it can be used by name for DirectPath connections.
	_ "google.golang.org/grpc/balancer/grpclb"
)

const loadBalancingPolicy = `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`

// TargetsUpdateInterval controls frequency of target updates.
var (
	TargetsUpdateInterval = 1 * time.Minute
)

// Probe holds aggregate information about all probe runs, per-target.
type Probe struct {
	name     string
	src      string
	opts     *options.Options
	c        *configpb.ProbeConf
	l        *logger.Logger
	dialOpts []grpc.DialOption

	// Targets and cancellation function for each target.
	targets     []endpoint.Endpoint
	cancelFuncs map[string]context.CancelFunc
	targetsMu   sync.Mutex

	// Results by target.
	results map[string]*probeRunResult
}

// probeRunResult captures the metrics for a single target. Multiple threads
// can update metrics at the same time and the main thread periodically
// outputs the values in this struct.
type probeRunResult struct {
	sync.Mutex
	target        string
	total         metrics.Int
	success       metrics.Int
	latency       metrics.Value
	connectErrors metrics.Int
}

func (p *Probe) setupDialOpts() error {
	oauthCfg := p.c.GetOauthConfig()
	if oauthCfg != nil {
		oauthTS, err := oauth.TokenSourceFromConfig(oauthCfg, p.l)
		if err != nil {
			return err
		}
		p.dialOpts = append(p.dialOpts, grpc.WithPerRPCCredentials(grpcoauth.TokenSource{oauthTS}))
	}
	altsCfg := p.c.GetAltsConfig()
	if altsCfg != nil {
		altsOpts := &alts.ClientOptions{
			TargetServiceAccounts:    altsCfg.GetTargetServiceAccount(),
			HandshakerServiceAddress: altsCfg.GetHandshakerServiceAddress(),
		}
		p.dialOpts = append(p.dialOpts, grpc.WithTransportCredentials(alts.NewClientCreds(altsOpts)))
	}

	if oauthCfg == nil && altsCfg == nil {
		p.dialOpts = append(p.dialOpts, grpc.WithInsecure())
	}
	p.dialOpts = append(p.dialOpts, grpc.WithDefaultServiceConfig(loadBalancingPolicy))
	p.dialOpts = append(p.dialOpts, grpc.WithBlock())
	return nil
}

// Init initializes the probe with the given params.
func (p *Probe) Init(name string, opts *options.Options) error {
	c, ok := opts.ProbeConf.(*configpb.ProbeConf)
	if !ok {
		return errors.New("not a gRPC probe config")
	}
	p.c = c
	p.name = name
	p.opts = opts
	if p.l = opts.Logger; p.l == nil {
		p.l = &logger.Logger{}
	}
	p.targets = p.opts.Targets.ListEndpoints()
	p.cancelFuncs = make(map[string]context.CancelFunc)
	p.src = sysvars.Vars()["hostname"]
	if err := p.setupDialOpts(); err != nil {
		return err
	}
	resolver.SetDefaultScheme("dns")
	return nil
}

func (p *Probe) updateTargetsAndStartProbes(ctx context.Context) {
	newTargets := p.opts.Targets.ListEndpoints()
	numNewTargets := len(newTargets)

	p.targetsMu.Lock()
	defer p.targetsMu.Unlock()
	if numNewTargets == 0 || numNewTargets < (len(p.targets)/2) {
		p.l.Errorf("Too few new targets, retaining old targets. New targets: %v, old count: %d", newTargets, len(p.targets))
		return
	}

	updatedTargets := make(map[string]string)
	defer func() {
		if len(updatedTargets) > 0 {
			p.l.Infof("Probe(%s) targets updated: %v", p.name, updatedTargets)
		}
	}()

	activeTargets := make(map[string]bool)
	// Create results structure and start probe loop for new targets.
	for _, tgtEp := range newTargets {
		tgt := net.JoinHostPort(tgtEp.Name, strconv.Itoa(tgtEp.Port))
		activeTargets[tgt] = true
		if _, ok := p.results[tgt]; ok {
			continue
		}
		updatedTargets[tgt] = "ADD"
		p.results[tgt] = p.newResult(tgt)
		probeCtx, probeCancelFunc := context.WithCancel(ctx)
		for i := 0; i < int(p.c.GetNumConns()); i++ {
			go p.oneTargetLoop(probeCtx, tgt, i, p.results[tgt])
		}
		p.cancelFuncs[tgt] = probeCancelFunc
	}

	// Stop probing for deleted targets by invoking cancelFunc.
	for tgt := range p.results {
		if activeTargets[tgt] {
			continue
		}
		p.cancelFuncs[tgt]()
		updatedTargets[tgt] = "DELETE"
		delete(p.results, tgt)
		delete(p.cancelFuncs, tgt)
	}
	p.targets = newTargets
}

// connectWithRetry attempts to connect to a target. On failure, it retries in
// an infinite loop until successful, incrementing connectErrors for every
// connection error. On success, it returns a client immediately.
// Interval between connects is controlled by connect_timeout_msec, defaulting
// to probe timeout.
func (p *Probe) connectWithRetry(ctx context.Context, tgt, msgPattern string, result *probeRunResult) *grpc.ClientConn {
	connectTimeout := p.opts.Timeout
	if p.c.GetConnectTimeoutMsec() > 0 {
		connectTimeout = time.Duration(p.c.GetConnectTimeoutMsec()) * time.Millisecond
	}
	var conn *grpc.ClientConn
	var err error
	for {
		select {
		case <-ctx.Done():
			p.l.Warningf("ProbeId(%s): context cancelled in connect loop.", msgPattern)
			return nil
		default:
		}
		connCtx, cancelFunc := context.WithTimeout(ctx, connectTimeout)
		conn, err = grpc.DialContext(connCtx, tgt, p.dialOpts...)
		cancelFunc()
		if err != nil {
			p.l.Warningf("ProbeId(%v) connect error: %v", msgPattern, err)
		} else {
			p.l.Infof("ProbeId(%v) connection established.", msgPattern)
			break
		}
		result.Lock()
		result.total.Inc()
		result.connectErrors.Inc()
		result.Unlock()
	}
	return conn
}

// oneTargetLoop connects to and then continuously probes a single target.
func (p *Probe) oneTargetLoop(ctx context.Context, tgt string, index int, result *probeRunResult) {
	msgPattern := fmt.Sprintf("%s,%s,%03d", p.src, tgt, index)

	conn := p.connectWithRetry(ctx, tgt, msgPattern, result)
	if conn == nil {
		return
	}
	defer conn.Close()

	client := servicepb.NewProberClient(conn)
	timeout := p.opts.Timeout
	method := p.c.GetMethod()

	msgSize := p.c.GetBlobSize()
	msg := make([]byte, msgSize)
	probeutils.PatternPayload(msg, []byte(msgPattern))
	ticker := time.NewTicker(p.opts.Interval)
	for {
		select {
		case <-ctx.Done():
			p.l.Warningf("ProbeId(%s): context cancelled in request loop.", msgPattern)
			ticker.Stop()
			return
		case <-ticker.C:
		}

		reqCtx, cancelFunc := context.WithTimeout(ctx, timeout)
		var success int64
		var delta time.Duration
		start := time.Now()
		var err error
		var peer peer.Peer
		opts := []grpc.CallOption{
			grpc.WaitForReady(true),
			grpc.Peer(&peer),
		}
		switch method {
		case configpb.ProbeConf_ECHO:
			req := &grpcprobepb.EchoMessage{
				Blob: []byte(msg),
			}
			_, err = client.Echo(reqCtx, req, opts...)
		case configpb.ProbeConf_READ:
			req := &grpcprobepb.BlobReadRequest{
				Size: proto.Int32(msgSize),
			}
			_, err = client.BlobRead(reqCtx, req, opts...)
		case configpb.ProbeConf_WRITE:
			req := &grpcprobepb.BlobWriteRequest{
				Blob: []byte(msg),
			}
			_, err = client.BlobWrite(reqCtx, req, opts...)
		default:
			p.l.Criticalf("Method %v not implemented", method)
		}
		cancelFunc()
		if err != nil {
			peerAddr := "unknown"
			if peer.Addr != nil {
				peerAddr = peer.Addr.String()
			}
			p.l.Warningf("ProbeId(%s) request failed: %v. ConnState: %v. Peer: %v", msgPattern, err, conn.GetState(), peerAddr)
		} else {
			success = 1
			delta = time.Since(start)
		}
		// TODO(ls692): add validators for probe result.
		result.Lock()
		result.total.Inc()
		result.success.AddInt64(success)
		result.latency.AddFloat64(delta.Seconds() / p.opts.LatencyUnit.Seconds())
		result.Unlock()
	}
}

func (p *Probe) newResult(tgt string) *probeRunResult {
	var latencyValue metrics.Value
	if p.opts.LatencyDist != nil {
		latencyValue = p.opts.LatencyDist.Clone()
	} else {
		latencyValue = metrics.NewFloat(0)
	}
	return &probeRunResult{
		target:  tgt,
		latency: latencyValue,
	}
}

// Start starts and runs the probe indefinitely.
func (p *Probe) Start(ctx context.Context, dataChan chan *metrics.EventMetrics) {
	p.results = make(map[string]*probeRunResult)
	p.updateTargetsAndStartProbes(ctx)

	ticker := time.NewTicker(p.opts.StatsExportInterval)
	defer ticker.Stop()

	targetsUpdateTicker := time.NewTicker(TargetsUpdateInterval)
	defer targetsUpdateTicker.Stop()

	for ts := range ticker.C {
		// Stop further processing and exit if context is canceled.
		// Same context is used by probe loops.
		select {
		case <-ctx.Done():
			return
		default:
		}

		// Output results.
		for targetName, result := range p.results {
			result.Lock()
			em := metrics.NewEventMetrics(ts).
				AddMetric("total", result.total.Clone()).
				AddMetric("success", result.success.Clone()).
				AddMetric("latency", result.latency.Clone()).
				AddMetric("connecterrors", result.connectErrors.Clone()).
				AddLabel("ptype", "grpc").
				AddLabel("probe", p.name).
				AddLabel("dst", targetName)
			result.Unlock()
			em.LatencyUnit = p.opts.LatencyUnit
			for _, al := range p.opts.AdditionalLabels {
				em.AddLabel(al.KeyValueForTarget(targetName))
			}
			p.opts.LogMetrics(em)
			dataChan <- em
		}

		// Finally, update targets and start new probe loops if necessary.
		// Executing this as the last step in the loop also ensures that new
		// targets have at least one cycle of probes before next output cycle.
		select {
		case <-targetsUpdateTicker.C:
			p.updateTargetsAndStartProbes(ctx)
		default:
		}
	}
}
<commit_msg>Migrate gRPC insecure API call sites to use local credentials (#664)<commit_after>// Copyright 2020 The Cloudprober 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 grpc implements a gRPC probe.

This probes a cloudprober gRPC server and reports success rate, latency, and
validation failures.
*/
package grpc

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strconv"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/cloudprober/common/oauth"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	configpb "github.com/google/cloudprober/probes/grpc/proto"
	"github.com/google/cloudprober/probes/options"
	"github.com/google/cloudprober/probes/probeutils"
	"github.com/google/cloudprober/sysvars"
	"github.com/google/cloudprober/targets/endpoint"

	grpcprobepb "github.com/google/cloudprober/servers/grpc/proto"
	servicepb "github.com/google/cloudprober/servers/grpc/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/credentials/local"
	grpcoauth "google.golang.org/grpc/credentials/oauth"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/resolver"

	// Import grpclb module so it can be used by name for DirectPath connections.
	_ "google.golang.org/grpc/balancer/grpclb"
)

const loadBalancingPolicy = `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`

// TargetsUpdateInterval controls frequency of target updates.
var (
	TargetsUpdateInterval = 1 * time.Minute
)

// Probe holds aggregate information about all probe runs, per-target.
type Probe struct {
	name     string
	src      string
	opts     *options.Options
	c        *configpb.ProbeConf
	l        *logger.Logger
	dialOpts []grpc.DialOption

	// Targets and cancellation function for each target.
	targets     []endpoint.Endpoint
	cancelFuncs map[string]context.CancelFunc
	targetsMu   sync.Mutex

	// Results by target.
	results map[string]*probeRunResult
}

// probeRunResult captures the metrics for a single target. Multiple threads
// can update metrics at the same time and the main thread periodically
// outputs the values in this struct.
type probeRunResult struct {
	sync.Mutex
	target        string
	total         metrics.Int
	success       metrics.Int
	latency       metrics.Value
	connectErrors metrics.Int
}

func (p *Probe) setupDialOpts() error {
	oauthCfg := p.c.GetOauthConfig()
	if oauthCfg != nil {
		oauthTS, err := oauth.TokenSourceFromConfig(oauthCfg, p.l)
		if err != nil {
			return err
		}
		p.dialOpts = append(p.dialOpts, grpc.WithPerRPCCredentials(grpcoauth.TokenSource{oauthTS}))
	}
	altsCfg := p.c.GetAltsConfig()
	if altsCfg != nil {
		altsOpts := &alts.ClientOptions{
			TargetServiceAccounts:    altsCfg.GetTargetServiceAccount(),
			HandshakerServiceAddress: altsCfg.GetHandshakerServiceAddress(),
		}
		p.dialOpts = append(p.dialOpts, grpc.WithTransportCredentials(alts.NewClientCreds(altsOpts)))
	}

	if oauthCfg == nil && altsCfg == nil {
		p.dialOpts = append(p.dialOpts, grpc.WithTransportCredentials(local.NewCredentials()))
	}
	p.dialOpts = append(p.dialOpts, grpc.WithDefaultServiceConfig(loadBalancingPolicy))
	p.dialOpts = append(p.dialOpts, grpc.WithBlock())
	return nil
}

// Init initializes the probe with the given params.
func (p *Probe) Init(name string, opts *options.Options) error {
	c, ok := opts.ProbeConf.(*configpb.ProbeConf)
	if !ok {
		return errors.New("not a gRPC probe config")
	}
	p.c = c
	p.name = name
	p.opts = opts
	if p.l = opts.Logger; p.l == nil {
		p.l = &logger.Logger{}
	}
	p.targets = p.opts.Targets.ListEndpoints()
	p.cancelFuncs = make(map[string]context.CancelFunc)
	p.src = sysvars.Vars()["hostname"]
	if err := p.setupDialOpts(); err != nil {
		return err
	}
	resolver.SetDefaultScheme("dns")
	return nil
}

func (p *Probe) updateTargetsAndStartProbes(ctx context.Context) {
	newTargets := p.opts.Targets.ListEndpoints()
	numNewTargets := len(newTargets)

	p.targetsMu.Lock()
	defer p.targetsMu.Unlock()
	if numNewTargets == 0 || numNewTargets < (len(p.targets)/2) {
		p.l.Errorf("Too few new targets, retaining old targets. New targets: %v, old count: %d", newTargets, len(p.targets))
		return
	}

	updatedTargets := make(map[string]string)
	defer func() {
		if len(updatedTargets) > 0 {
			p.l.Infof("Probe(%s) targets updated: %v", p.name, updatedTargets)
		}
	}()

	activeTargets := make(map[string]bool)
	// Create results structure and start probe loop for new targets.
	for _, tgtEp := range newTargets {
		tgt := net.JoinHostPort(tgtEp.Name, strconv.Itoa(tgtEp.Port))
		activeTargets[tgt] = true
		if _, ok := p.results[tgt]; ok {
			continue
		}
		updatedTargets[tgt] = "ADD"
		p.results[tgt] = p.newResult(tgt)
		probeCtx, probeCancelFunc := context.WithCancel(ctx)
		for i := 0; i < int(p.c.GetNumConns()); i++ {
			go p.oneTargetLoop(probeCtx, tgt, i, p.results[tgt])
		}
		p.cancelFuncs[tgt] = probeCancelFunc
	}

	// Stop probing for deleted targets by invoking cancelFunc.
	for tgt := range p.results {
		if activeTargets[tgt] {
			continue
		}
		p.cancelFuncs[tgt]()
		updatedTargets[tgt] = "DELETE"
		delete(p.results, tgt)
		delete(p.cancelFuncs, tgt)
	}
	p.targets = newTargets
}

// connectWithRetry attempts to connect to a target. On failure, it retries in
// an infinite loop until successful, incrementing connectErrors for every
// connection error. On success, it returns a client immediately.
// Interval between connects is controlled by connect_timeout_msec, defaulting
// to probe timeout.
func (p *Probe) connectWithRetry(ctx context.Context, tgt, msgPattern string, result *probeRunResult) *grpc.ClientConn {
	connectTimeout := p.opts.Timeout
	if p.c.GetConnectTimeoutMsec() > 0 {
		connectTimeout = time.Duration(p.c.GetConnectTimeoutMsec()) * time.Millisecond
	}
	var conn *grpc.ClientConn
	var err error
	for {
		select {
		case <-ctx.Done():
			p.l.Warningf("ProbeId(%s): context cancelled in connect loop.", msgPattern)
			return nil
		default:
		}
		connCtx, cancelFunc := context.WithTimeout(ctx, connectTimeout)
		conn, err = grpc.DialContext(connCtx, tgt, p.dialOpts...)
		cancelFunc()
		if err != nil {
			p.l.Warningf("ProbeId(%v) connect error: %v", msgPattern, err)
		} else {
			p.l.Infof("ProbeId(%v) connection established.", msgPattern)
			break
		}
		result.Lock()
		result.total.Inc()
		result.connectErrors.Inc()
		result.Unlock()
	}
	return conn
}

// oneTargetLoop connects to and then continuously probes a single target.
func (p *Probe) oneTargetLoop(ctx context.Context, tgt string, index int, result *probeRunResult) {
	msgPattern := fmt.Sprintf("%s,%s,%03d", p.src, tgt, index)

	conn := p.connectWithRetry(ctx, tgt, msgPattern, result)
	if conn == nil {
		return
	}
	defer conn.Close()

	client := servicepb.NewProberClient(conn)
	timeout := p.opts.Timeout
	method := p.c.GetMethod()

	msgSize := p.c.GetBlobSize()
	msg := make([]byte, msgSize)
	probeutils.PatternPayload(msg, []byte(msgPattern))
	ticker := time.NewTicker(p.opts.Interval)
	for {
		select {
		case <-ctx.Done():
			p.l.Warningf("ProbeId(%s): context cancelled in request loop.", msgPattern)
			ticker.Stop()
			return
		case <-ticker.C:
		}

		reqCtx, cancelFunc := context.WithTimeout(ctx, timeout)
		var success int64
		var delta time.Duration
		start := time.Now()
		var err error
		var peer peer.Peer
		opts := []grpc.CallOption{
			grpc.WaitForReady(true),
			grpc.Peer(&peer),
		}
		switch method {
		case configpb.ProbeConf_ECHO:
			req := &grpcprobepb.EchoMessage{
				Blob: []byte(msg),
			}
			_, err = client.Echo(reqCtx, req, opts...)
		case configpb.ProbeConf_READ:
			req := &grpcprobepb.BlobReadRequest{
				Size: proto.Int32(msgSize),
			}
			_, err = client.BlobRead(reqCtx, req, opts...)
		case configpb.ProbeConf_WRITE:
			req := &grpcprobepb.BlobWriteRequest{
				Blob: []byte(msg),
			}
			_, err = client.BlobWrite(reqCtx, req, opts...)
		default:
			p.l.Criticalf("Method %v not implemented", method)
		}
		cancelFunc()
		if err != nil {
			peerAddr := "unknown"
			if peer.Addr != nil {
				peerAddr = peer.Addr.String()
			}
			p.l.Warningf("ProbeId(%s) request failed: %v. ConnState: %v. Peer: %v", msgPattern, err, conn.GetState(), peerAddr)
		} else {
			success = 1
			delta = time.Since(start)
		}
		// TODO(ls692): add validators for probe result.
		result.Lock()
		result.total.Inc()
		result.success.AddInt64(success)
		result.latency.AddFloat64(delta.Seconds() / p.opts.LatencyUnit.Seconds())
		result.Unlock()
	}
}

func (p *Probe) newResult(tgt string) *probeRunResult {
	var latencyValue metrics.Value
	if p.opts.LatencyDist != nil {
		latencyValue = p.opts.LatencyDist.Clone()
	} else {
		latencyValue = metrics.NewFloat(0)
	}
	return &probeRunResult{
		target:  tgt,
		latency: latencyValue,
	}
}

// Start starts and runs the probe indefinitely.
func (p *Probe) Start(ctx context.Context, dataChan chan *metrics.EventMetrics) {
	p.results = make(map[string]*probeRunResult)
	p.updateTargetsAndStartProbes(ctx)

	ticker := time.NewTicker(p.opts.StatsExportInterval)
	defer ticker.Stop()

	targetsUpdateTicker := time.NewTicker(TargetsUpdateInterval)
	defer targetsUpdateTicker.Stop()

	for ts := range ticker.C {
		// Stop further processing and exit if context is canceled.
		// Same context is used by probe loops.
		select {
		case <-ctx.Done():
			return
		default:
		}

		// Output results.
		for targetName, result := range p.results {
			result.Lock()
			em := metrics.NewEventMetrics(ts).
				AddMetric("total", result.total.Clone()).
				AddMetric("success", result.success.Clone()).
				AddMetric("latency", result.latency.Clone()).
				AddMetric("connecterrors", result.connectErrors.Clone()).
				AddLabel("ptype", "grpc").
				AddLabel("probe", p.name).
				AddLabel("dst", targetName)
			result.Unlock()
			em.LatencyUnit = p.opts.LatencyUnit
			for _, al := range p.opts.AdditionalLabels {
				em.AddLabel(al.KeyValueForTarget(targetName))
			}
			p.opts.LogMetrics(em)
			dataChan <- em
		}

		// Finally, update targets and start new probe loops if necessary.
		// Executing this as the last step in the loop also ensures that new
		// targets have at least one cycle of probes before next output cycle.
		select {
		case <-targetsUpdateTicker.C:
			p.updateTargetsAndStartProbes(ctx)
		default:
		}
	}
}
<|endoftext|>


<commit_before>package logging

import (
	"fmt"
	"io"
	stdlog "log"
	"os"

	kitlog "github.com/go-kit/kit/log"
	"gopkg.in/errgo.v1"
)

var closeChan chan<- os.Signal

func SetCloseChan(c chan<- os.Signal) {
	closeChan = c
}

// CheckFatal exits the process if err != nil
func CheckFatal(err error) {
	if err != nil {
		l := internal
		if l == nil {
			l = kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stderr))
			l = kitlog.NewContext(l).With("module", "logging", kitlog.DefaultCaller)
		}

		l.Log("check", "fatal", "err", errgo.Details(err))
		if closeChan != nil {
			l.Log("check", "notice", "msg", "Sending close message")
			closeChan <- os.Interrupt
		}
		os.Exit(1)
	}
}

var internal kitlog.Logger

// SetupLogging will initialize the logger backend and set the flags.
func SetupLogging(w io.Writer) {
	if w == nil {
		w = os.Stderr
	}
	logger := kitlog.NewLogfmtLogger(w)

	if lvl := os.Getenv("CRYPTIX_LOGLVL"); lvl != "" {
		logger.Log("module", "logging", "error", "CRYPTIX_LOGLVL is obsolete. levels are bad, mkay?")
	}
	// wrap logger to error-check the writes only once
	internal = kitlog.LoggerFunc(func(keyvals ...interface{}) error {
		if err := logger.Log(keyvals...); err != nil {
			fmt.Fprintf(os.Stderr, "warning: logger.Write() failed! %s", err)
			panic(err) // no other way to escalate this
		}
		return nil
	})
	stdlog.SetOutput(kitlog.NewStdlibAdapter(kitlog.NewContext(internal).With("module", "stdlib")))
	internal = kitlog.NewContext(internal).With("ts", kitlog.DefaultTimestamp, "caller", kitlog.DefaultCaller)
}

// Logger returns an Entry where the module field is set to name
func Logger(name string) *kitlog.Context {
	l := internal
	if l == nil {
		l = kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stderr))
		l = kitlog.NewContext(l).With("warning", "uninitizialized", kitlog.DefaultCaller)
	}
	if name == "" {
		l.Log("module", "logger", "error", "missing name parameter")
		name = "undefined"
	}
	return kitlog.NewContext(l).With("module", name)
}
<commit_msg>logging: update to new go-kit/log api<commit_after>package logging

import (
	"fmt"
	"io"
	stdlog "log"
	"os"

	kitlog "github.com/go-kit/kit/log"
	"gopkg.in/errgo.v1"
)

var closeChan chan<- os.Signal

func SetCloseChan(c chan<- os.Signal) {
	closeChan = c
}

// CheckFatal exits the process if err != nil
func CheckFatal(err error) {
	if err != nil {
		l := internal
		if l == nil {
			l = kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stderr))
			l = kitlog.With(l, "module", "logging", kitlog.DefaultCaller)
		}

		l.Log("check", "fatal", "err", errgo.Details(err))
		if closeChan != nil {
			l.Log("check", "notice", "msg", "Sending close message")
			closeChan <- os.Interrupt
		}
		os.Exit(1)
	}
}

var internal kitlog.Logger

// SetupLogging will initialize the logger backend and set the flags.
func SetupLogging(w io.Writer) {
	if w == nil {
		w = os.Stderr
	}
	logger := kitlog.NewLogfmtLogger(w)

	if lvl := os.Getenv("CRYPTIX_LOGLVL"); lvl != "" {
		logger.Log("module", "logging", "error", "CRYPTIX_LOGLVL is obsolete. levels are bad, mkay?")
	}
	// wrap logger to error-check the writes only once
	internal = kitlog.LoggerFunc(func(keyvals ...interface{}) error {
		if err := logger.Log(keyvals...); err != nil {
			fmt.Fprintf(os.Stderr, "warning: logger.Write() failed! %s", err)
			panic(err) // no other way to escalate this
		}
		return nil
	})
	stdlog.SetOutput(kitlog.NewStdlibAdapter(kitlog.With(internal, "module", "stdlib")))
	internal = kitlog.With(internal, "ts", kitlog.DefaultTimestamp, "caller", kitlog.DefaultCaller)
}

// Logger returns an Entry where the module field is set to name
func Logger(name string) kitlog.Logger {
	l := internal
	if l == nil {
		l = kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stderr))
		l = kitlog.With(l, "warning", "uninitizialized", kitlog.DefaultCaller)
	}
	if name == "" {
		l.Log("module", "logger", "error", "missing name parameter")
		name = "undefined"
	}
	return kitlog.With(l, "module", name)
}
<|endoftext|>


<commit_before>package processors

import (
	"fmt"
	"net/http"
	"strconv"
	"sync"
	"time"
)

type batchedRequest struct {
	Sequence int
	Request  *http.Request
}

type batchedResponse struct {
	Sequence int
	Response *http.Response
}

// ProcessBatch sends a batch of HTTP requests using http.Client.
// Each request is sent concurrently in a seperate goroutine.
// The HTTP responses are returned in the same sequence as their corresponding requests.
func ProcessBatch(requests []*http.Request, timeout time.Duration) ([]*http.Response, error) {
	z := len(requests)
	// Setup a buffered channel to queue up the requests for processing by individual HTTP Client goroutines
	batchedRequests := make(chan batchedRequest, z)
	for i := 0; i < z; i++ {
		batchedRequests <- batchedRequest{i, requests[i]}
	}
	// Close the channel - nothing else is sent to it
	close(batchedRequests)
	// Setup a second buffered channel for collecting the BatchedResponses from the individual HTTP Client goroutines
	batchedResponses := make(chan batchedResponse, z)
	// Setup a wait group so we know when all the BatchedRequests have been processed
	var wg sync.WaitGroup
	wg.Add(z)
	// Start our individual HTTP Client goroutines to process the BatchedRequests
	for i := 0; i < z; i++ {
		go func() {
			defer wg.Done()
			r := <-batchedRequests
			client := &http.Client{Timeout: timeout}
			response, err := client.Do(r.Request)
			if err != nil {
				// Create an error response for any HTTP Client errors - Status 400 (Bad Request)
				errorResponse := http.Response{}
				errorResponse.Proto = r.Request.Proto
				errorResponse.StatusCode = http.StatusBadRequest
				errorResponse.Status = strconv.Itoa(http.StatusBadRequest) + " " + err.Error()
				batchedResponses <- batchedResponse{r.Sequence, &errorResponse}
			} else {
				batchedResponses <- batchedResponse{r.Sequence, response}
			}
		}()
	}
	// Wait for all the requests to be processed
	wg.Wait()
	// Close the second buffered channel that we used to collect the BatchedResponses
	close(batchedResponses)
	// Check we have the correct number of BatchedResponses
	if len(batchedResponses) == z {
		// Return the BatchedResponses in their correct sequence
		result := make([]*http.Response, z)
		for i := 0; i < z; i++ {
			r := <-batchedResponses
			result[r.Sequence] = r.Response
		}
		return result, nil
	}
	err := fmt.Errorf("expected %d responses for this batch but only recieved %d", z, len(batchedResponses))
	return nil, err
}
<commit_msg>Continued work on initial release<commit_after>package processors

import (
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"
)

type batchedRequest struct {
	Sequence int
	Request  *http.Request
}

type batchedResponse struct {
	Sequence int
	Response *http.Response
}

// ProcessBatch sends a batch of HTTP requests using http.Transport.
// Each request is sent concurrently in a seperate goroutine.
// The HTTP responses are returned in the same sequence as their corresponding requests.
func ProcessBatch(requests []*http.Request, timeout time.Duration) ([]*http.Response, error) {
	z := len(requests)
	// Setup a buffered channel to queue up the requests for processing by individual HTTP Transport goroutines
	batchedRequests := make(chan batchedRequest, z)
	for i := 0; i < z; i++ {
		batchedRequests <- batchedRequest{i, requests[i]}
	}
	// Close the channel - nothing else is sent to it
	close(batchedRequests)
	// Setup a second buffered channel for collecting the BatchedResponses from the individual HTTP Transport goroutines
	batchedResponses := make(chan batchedResponse, z)
	// Setup a wait group so we know when all the BatchedRequests have been processed
	var wg sync.WaitGroup
	wg.Add(z)

	// Start our individual HTTP Transport goroutines to process the BatchedRequests
	for i := 0; i < z; i++ {
		go func() {
			defer wg.Done()
			r := <-batchedRequests
			transport := &http.Transport{ResponseHeaderTimeout: timeout}
			transport.DisableCompression = true
			response, err := transport.RoundTrip(r.Request)
			// TODO add support for all possible redirect status codes, see line 249 of https://golang.org/src/net/http/client.go
			if response.StatusCode == 302 {
				location := response.Header.Get("Location")
				if location != "" {
					redirectURL, err := url.Parse(location)
					if err == nil {
						if !redirectURL.IsAbs() { // handle relative URLs
							redirectURL, err = url.Parse(r.Request.URL.Scheme + "://" + r.Request.Host + "/" + location)
						}
						queryString := ""
						if len(redirectURL.Query()) > 0 {
							queryString = "?" + redirectURL.Query().Encode()
						}
						redirect, err := http.NewRequest("GET", redirectURL.Scheme+"://"+redirectURL.Host+redirectURL.Path+queryString, nil)
						if err == nil {
							response, err = transport.RoundTrip(redirect)
						}
					}
				}
			}

			if err != nil {
				// Create an error response for any HTTP Transport errors - Status 400 (Bad Request)
				errorResponse := http.Response{}
				errorResponse.Proto = r.Request.Proto
				errorResponse.StatusCode = http.StatusBadRequest
				errorResponse.Status = strconv.Itoa(http.StatusBadRequest) + " " + err.Error()
				batchedResponses <- batchedResponse{r.Sequence, &errorResponse}
			} else {
				batchedResponses <- batchedResponse{r.Sequence, response}
			}
		}()
	}

	// Wait for all the requests to be processed
	wg.Wait()
	// Close the second buffered channel that we used to collect the BatchedResponses
	close(batchedResponses)
	// Check we have the correct number of BatchedResponses
	if len(batchedResponses) == z {
		// Return the BatchedResponses in their correct sequence
		result := make([]*http.Response, z)
		for i := 0; i < z; i++ {
			r := <-batchedResponses
			result[r.Sequence] = r.Response
		}
		return result, nil
	}
	err := fmt.Errorf("expected %d responses for this batch but only recieved %d", z, len(batchedResponses))
	return nil, err
}
<|endoftext|>


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

package registry

import (
	"crypto/rand"
	"fmt"
	"io"
	"time"

	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/gcloud/gcs"
)

// Create a registry that stores data in the supplied GCS bucket, deriving a
// crypto key from the supplied password and ensuring that the bucket may not
// in the future be used with any other key and has not in the past, either.
// Return a crypter configured to use the key.
func NewGCSRegistry(
	bucket gcs.Bucket,
	cryptoPassword string,
	deriver crypto.KeyDeriver) (r Registry, crypter crypto.Crypter, err error) {
	return newGCSRegistry(
		bucket,
		cryptoPassword,
		deriver,
		crypto.NewCrypter,
		rand.Reader)
}

const (
	gcsJobKeyPrefix      = "jobs/"
	gcsMetadataKey_Name  = "job_name"
	gcsMetadataKey_Score = "hex_score"
)

// A registry that stores job records in a GCS bucket. Object names are of the
// form
//
//     <gcsJobKeyPrefix><time>
//
// where <time> is a time.Time with UTC location formatted according to
// time.RFC3339Nano. Additional information is stored as object metadata fields
// keyed by the constants above. Metadata fields are used in preference to
// object content so that they are accessible on a ListObjects request.
//
// The bucket additionally contains a "marker" object (named by the constant
// markerItemName) with metadata keys specifying a salt and a ciphertext for
// some random plaintext, generated ant written the first time the bucket is
// used. This marker allows us to verify that the user-provided crypto password
// is correct by deriving a key using the password and the salt and making sure
// that the ciphertext can be decrypted using that key.
type gcsRegistry struct {
	bucket gcs.Bucket
}

// Like NewGCSRegistry, but with more injected.
func newGCSRegistry(
	bucket gcs.Bucket,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
	createCrypter func(key []byte) (crypto.Crypter, error),
	cryptoRandSrc io.Reader) (r Registry, crypter crypto.Crypter, err error) {
}

func (r *gcsRegistry) RecordBackup(j CompletedJob) (err error) {
	err = fmt.Errorf("gcsRegistry.RecordBackup is not implemented.")
	return
}

func (r *gcsRegistry) ListRecentBackups() (jobs []CompletedJob, err error) {
	err = fmt.Errorf("gcsRegistry.ListRecentBackups is not implemented.")
	return
}

func (r *gcsRegistry) FindBackup(
	startTime time.Time) (job CompletedJob, err error) {
	err = fmt.Errorf("gcsRegistry.FindBackup is not implemented.")
	return
}
<commit_msg>Pasted code from before the GCS change.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package registry

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"time"

	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/gcloud/gcs"
)

// Create a registry that stores data in the supplied GCS bucket, deriving a
// crypto key from the supplied password and ensuring that the bucket may not
// in the future be used with any other key and has not in the past, either.
// Return a crypter configured to use the key.
func NewGCSRegistry(
	bucket gcs.Bucket,
	cryptoPassword string,
	deriver crypto.KeyDeriver) (r Registry, crypter crypto.Crypter, err error) {
	return newGCSRegistry(
		bucket,
		cryptoPassword,
		deriver,
		crypto.NewCrypter,
		rand.Reader)
}

const (
	gcsJobKeyPrefix      = "jobs/"
	gcsMetadataKey_Name  = "job_name"
	gcsMetadataKey_Score = "hex_score"
)

// A registry that stores job records in a GCS bucket. Object names are of the
// form
//
//     <gcsJobKeyPrefix><time>
//
// where <time> is a time.Time with UTC location formatted according to
// time.RFC3339Nano. Additional information is stored as object metadata fields
// keyed by the constants above. Metadata fields are used in preference to
// object content so that they are accessible on a ListObjects request.
//
// The bucket additionally contains a "marker" object (named by the constant
// markerItemName) with metadata keys specifying a salt and a ciphertext for
// some random plaintext, generated ant written the first time the bucket is
// used. This marker allows us to verify that the user-provided crypto password
// is correct by deriving a key using the password and the salt and making sure
// that the ciphertext can be decrypted using that key.
type gcsRegistry struct {
	bucket gcs.Bucket
}

// Like NewGCSRegistry, but with more injected.
func newGCSRegistry(
	bucket gcs.Bucket,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
	createCrypter func(key []byte) (crypto.Crypter, error),
	cryptoRandSrc io.Reader) (r Registry, crypter crypto.Crypter, err error) {
}

func (r *gcsRegistry) RecordBackup(j CompletedJob) (err error) {
	err = fmt.Errorf("gcsRegistry.RecordBackup is not implemented.")
	return
}

func (r *gcsRegistry) ListRecentBackups() (jobs []CompletedJob, err error) {
	err = fmt.Errorf("gcsRegistry.ListRecentBackups is not implemented.")
	return
}

func (r *gcsRegistry) FindBackup(
	startTime time.Time) (job CompletedJob, err error) {
	err = fmt.Errorf("gcsRegistry.FindBackup is not implemented.")
	return
}

func verifyCompatibleAndSetUpCrypter(
	markerAttrs []sdb.Attribute,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
	createCrypter func(key []byte) (crypto.Crypter, error),
) (crypter crypto.Crypter, err error) {
	// Look through the attributes for what we need.
	var ciphertext []byte
	var salt []byte

	for _, attr := range markerAttrs {
		var dest *[]byte
		switch attr.Name {
		case encryptedDataMarker:
			dest = &ciphertext
		case passwordSaltMarker:
			dest = &salt
		default:
			continue
		}

		// The data is base64-encoded.
		if *dest, err = base64.StdEncoding.DecodeString(attr.Value); err != nil {
			err = fmt.Errorf("Decoding %s (%s): %v", attr.Name, attr.Value, err)
			return
		}
	}

	// Did we get both ciphertext and salt?
	if ciphertext == nil {
		err = fmt.Errorf("Missing encrypted data marker.")
		return
	}

	if salt == nil {
		err = fmt.Errorf("Missing password salt marker.")
		return
	}

	// Derive a key and create a crypter.
	cryptoKey := deriver.DeriveKey(cryptoPassword, salt)
	if crypter, err = createCrypter(cryptoKey); err != nil {
		err = fmt.Errorf("createCrypter: %v", err)
		return
	}

	// Attempt to decrypt the ciphertext.
	if _, err = crypter.Decrypt(ciphertext); err != nil {
		// Special case: Did the crypter signal that the key was wrong?
		if _, ok := err.(*crypto.NotAuthenticError); ok {
			err = fmt.Errorf("The supplied password is incorrect.")
			return
		}

		// Generic error.
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	return
}
<|endoftext|>


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

package registry

import (
	"errors"
	"time"

	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/gcloud/gcs"
)

// Create a registry that stores data in the supplied GCS bucket, deriving a
// crypto key from the supplied password and ensuring that the bucket may not
// in the future be used with any other key and has not in the past, either.
// Return a crypter configured to use the key.
func NewGCSRegistry(
	bucket gcs.Bucket,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
) (r Registry, crypter crypto.Crypter, err error) {
	err = errors.New("NewGCSRegistry is not implemented.")
	return
}

const (
	gcsJobKeyPrefix      = "jobs/"
	gcsMetadataKey_Name  = "job_name"
	gcsMetadataKey_Score = "hex_score"
)

// A registry that stores job records in a GCS bucket. Object names are of the
// form
//
//     <gcsJobKeyPrefix><time>
//
// where <time> is a time.Time with UTC location formatted according to
// time.RFC3339Nano. Additional information is stored as object metadata fields
// keyed by the constants above. Metadata fields are used in preference to
// object content so that they are accessible on a ListObjects request.
type gcsRegistry struct {
	bucket gcs.Bucket
}

func (r *gcsRegistry) RecordBackup(j CompletedJob) (err error) {
	err = errors.New("gcsRegistry.RecordBackup is not implemented.")
	return
}

func (r *gcsRegistry) ListRecentBackups() (jobs []CompletedJob, err error) {
	err = errors.New("gcsRegistry.ListRecentBackups is not implemented.")
	return
}

func (r *gcsRegistry) FindBackup(
	startTime time.Time) (job CompletedJob, err error) {
	err = errors.New("gcsRegistry.FindBackup is not implemented.")
	return
}
<commit_msg>Switched to fmt.Errorf.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package registry

import (
	"fmt"
	"time"

	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/gcloud/gcs"
)

// Create a registry that stores data in the supplied GCS bucket, deriving a
// crypto key from the supplied password and ensuring that the bucket may not
// in the future be used with any other key and has not in the past, either.
// Return a crypter configured to use the key.
func NewGCSRegistry(
	bucket gcs.Bucket,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
) (r Registry, crypter crypto.Crypter, err error) {
	err = fmt.Errorf("NewGCSRegistry is not implemented.")
	return
}

const (
	gcsJobKeyPrefix      = "jobs/"
	gcsMetadataKey_Name  = "job_name"
	gcsMetadataKey_Score = "hex_score"
)

// A registry that stores job records in a GCS bucket. Object names are of the
// form
//
//     <gcsJobKeyPrefix><time>
//
// where <time> is a time.Time with UTC location formatted according to
// time.RFC3339Nano. Additional information is stored as object metadata fields
// keyed by the constants above. Metadata fields are used in preference to
// object content so that they are accessible on a ListObjects request.
type gcsRegistry struct {
	bucket gcs.Bucket
}

func (r *gcsRegistry) RecordBackup(j CompletedJob) (err error) {
	err = fmt.Errorf("gcsRegistry.RecordBackup is not implemented.")
	return
}

func (r *gcsRegistry) ListRecentBackups() (jobs []CompletedJob, err error) {
	err = fmt.Errorf("gcsRegistry.ListRecentBackups is not implemented.")
	return
}

func (r *gcsRegistry) FindBackup(
	startTime time.Time) (job CompletedJob, err error) {
	err = fmt.Errorf("gcsRegistry.FindBackup is not implemented.")
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 registry

import (
	"errors"
	"time"

	"strings"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
	"github.com/kubernetes-incubator/external-dns/provider"
	log "github.com/sirupsen/logrus"
)

// TXTRegistry implements registry interface with ownership implemented via associated TXT records
type TXTRegistry struct {
	provider provider.Provider
	ownerID  string //refers to the owner id of the current instance
	mapper   nameMapper

	// cache the records in memory and update on an interval instead.
	recordsCache            []*endpoint.Endpoint
	recordsCacheRefreshTime time.Time
	cacheInterval           time.Duration
}

// NewTXTRegistry returns new TXTRegistry object
func NewTXTRegistry(provider provider.Provider, txtPrefix, ownerID string, cacheInterval time.Duration) (*TXTRegistry, error) {
	if ownerID == "" {
		return nil, errors.New("owner id cannot be empty")
	}

	mapper := newPrefixNameMapper(txtPrefix)

	return &TXTRegistry{
		provider:      provider,
		ownerID:       ownerID,
		mapper:        mapper,
		cacheInterval: cacheInterval,
	}, nil
}

// Records returns the current records from the registry excluding TXT Records
// If TXT records was created previously to indicate ownership its corresponding value
// will be added to the endpoints Labels map
func (im *TXTRegistry) Records() ([]*endpoint.Endpoint, error) {
	// If we have the zones cached AND we have refreshed the cache since the
	// last given interval, then just use the cached results.
	if im.recordsCache != nil && time.Since(im.recordsCacheRefreshTime) < im.cacheInterval {
		log.Debug("Using cached records.")
		return im.recordsCache, nil
	}

	records, err := im.provider.Records()
	if err != nil {
		return nil, err
	}

	endpoints := []*endpoint.Endpoint{}

	labelMap := map[string]endpoint.Labels{}

	for _, record := range records {
		if record.RecordType != endpoint.RecordTypeTXT {
			endpoints = append(endpoints, record)
			continue
		}
		// We simply assume that TXT records for the registry will always have only one target.
		labels, err := endpoint.NewLabelsFromString(record.Targets[0])
		if err == endpoint.ErrInvalidHeritage {
			//if no heritage is found or it is invalid
			//case when value of txt record cannot be identified
			//record will not be removed as it will have empty owner
			endpoints = append(endpoints, record)
			continue
		}
		if err != nil {
			return nil, err
		}
		endpointDNSName := im.mapper.toEndpointName(record.DNSName)
		labelMap[endpointDNSName] = labels
	}

	for _, ep := range endpoints {
		if labels, ok := labelMap[ep.DNSName]; ok {
			ep.Labels = labels
		} else {
			//this indicates that owner could not be identified, as there is no corresponding TXT record
			ep.Labels = endpoint.NewLabels()
		}
	}

	// Update the cache.
	if im.cacheInterval > 0 {
		im.recordsCache = endpoints
		im.recordsCacheRefreshTime = time.Now()
	}

	return endpoints, nil
}

// ApplyChanges updates dns provider with the changes
// for each created/deleted record it will also take into account TXT records for creation/deletion
func (im *TXTRegistry) ApplyChanges(changes *plan.Changes) error {
	filteredChanges := &plan.Changes{
		Create:    changes.Create,
		UpdateNew: filterOwnedRecords(im.ownerID, changes.UpdateNew),
		UpdateOld: filterOwnedRecords(im.ownerID, changes.UpdateOld),
		Delete:    filterOwnedRecords(im.ownerID, changes.Delete),
	}
	for _, r := range filteredChanges.Create {
		r.Labels[endpoint.OwnerLabelKey] = im.ownerID
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))
		filteredChanges.Create = append(filteredChanges.Create, txt)

		if im.cacheInterval > 0 {
			im.addToCache(r)
		}
	}

	for _, r := range filteredChanges.Delete {
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))

		// when we delete TXT records for which value has changed (due to new label) this would still work because
		// !!! TXT record value is uniquely generated from the Labels of the endpoint. Hence old TXT record can be uniquely reconstructed
		filteredChanges.Delete = append(filteredChanges.Delete, txt)

		if im.cacheInterval > 0 {
			im.removeFromCache(r)
		}
	}

	// make sure TXT records are consistently updated as well
	for _, r := range filteredChanges.UpdateOld {
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))
		// when we updateOld TXT records for which value has changed (due to new label) this would still work because
		// !!! TXT record value is uniquely generated from the Labels of the endpoint. Hence old TXT record can be uniquely reconstructed
		filteredChanges.UpdateOld = append(filteredChanges.UpdateOld, txt)
		// remove old version of record from cache
		if im.cacheInterval > 0 {
			im.removeFromCache(r)
		}
	}

	// make sure TXT records are consistently updated as well
	for _, r := range filteredChanges.UpdateNew {
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))
		filteredChanges.UpdateNew = append(filteredChanges.UpdateNew, txt)
		// add new version of record to cache
		if im.cacheInterval > 0 {
			im.addToCache(r)
		}
	}

	return im.provider.ApplyChanges(filteredChanges)
}

/**
  TXT registry specific private methods
*/

/**
  nameMapper defines interface which maps the dns name defined for the source
  to the dns name which TXT record will be created with
*/

type nameMapper interface {
	toEndpointName(string) string
	toTXTName(string) string
}

type prefixNameMapper struct {
	prefix string
}

var _ nameMapper = prefixNameMapper{}

func newPrefixNameMapper(prefix string) prefixNameMapper {
	return prefixNameMapper{prefix: prefix}
}

func (pr prefixNameMapper) toEndpointName(txtDNSName string) string {
	if strings.HasPrefix(txtDNSName, pr.prefix) {
		return strings.TrimPrefix(txtDNSName, pr.prefix)
	}
	return ""
}

func (pr prefixNameMapper) toTXTName(endpointDNSName string) string {
	return pr.prefix + endpointDNSName
}

func (im *TXTRegistry) addToCache(ep *endpoint.Endpoint) {
	if im.recordsCache != nil {
		im.recordsCache = append(im.recordsCache, ep)
	}
}

func (im *TXTRegistry) removeFromCache(ep *endpoint.Endpoint) {
	if im.recordsCache == nil || ep == nil {
		// return early.
		return
	}

	for i, e := range im.recordsCache {
		if e.DNSName == ep.DNSName && e.RecordType == ep.RecordType && e.Targets.Same(ep.Targets) {
			// We found a match delete the endpoint from the cache.
			im.recordsCache = append(im.recordsCache[:i], im.recordsCache[i+1:]...)
			return
		}
	}
}
<commit_msg>Fix a possible nil map access<commit_after>/*
Copyright 2017 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 registry

import (
	"errors"
	"time"

	"strings"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
	"github.com/kubernetes-incubator/external-dns/provider"
	log "github.com/sirupsen/logrus"
)

// TXTRegistry implements registry interface with ownership implemented via associated TXT records
type TXTRegistry struct {
	provider provider.Provider
	ownerID  string //refers to the owner id of the current instance
	mapper   nameMapper

	// cache the records in memory and update on an interval instead.
	recordsCache            []*endpoint.Endpoint
	recordsCacheRefreshTime time.Time
	cacheInterval           time.Duration
}

// NewTXTRegistry returns new TXTRegistry object
func NewTXTRegistry(provider provider.Provider, txtPrefix, ownerID string, cacheInterval time.Duration) (*TXTRegistry, error) {
	if ownerID == "" {
		return nil, errors.New("owner id cannot be empty")
	}

	mapper := newPrefixNameMapper(txtPrefix)

	return &TXTRegistry{
		provider:      provider,
		ownerID:       ownerID,
		mapper:        mapper,
		cacheInterval: cacheInterval,
	}, nil
}

// Records returns the current records from the registry excluding TXT Records
// If TXT records was created previously to indicate ownership its corresponding value
// will be added to the endpoints Labels map
func (im *TXTRegistry) Records() ([]*endpoint.Endpoint, error) {
	// If we have the zones cached AND we have refreshed the cache since the
	// last given interval, then just use the cached results.
	if im.recordsCache != nil && time.Since(im.recordsCacheRefreshTime) < im.cacheInterval {
		log.Debug("Using cached records.")
		return im.recordsCache, nil
	}

	records, err := im.provider.Records()
	if err != nil {
		return nil, err
	}

	endpoints := []*endpoint.Endpoint{}

	labelMap := map[string]endpoint.Labels{}

	for _, record := range records {
		if record.RecordType != endpoint.RecordTypeTXT {
			endpoints = append(endpoints, record)
			continue
		}
		// We simply assume that TXT records for the registry will always have only one target.
		labels, err := endpoint.NewLabelsFromString(record.Targets[0])
		if err == endpoint.ErrInvalidHeritage {
			//if no heritage is found or it is invalid
			//case when value of txt record cannot be identified
			//record will not be removed as it will have empty owner
			endpoints = append(endpoints, record)
			continue
		}
		if err != nil {
			return nil, err
		}
		endpointDNSName := im.mapper.toEndpointName(record.DNSName)
		labelMap[endpointDNSName] = labels
	}

	for _, ep := range endpoints {
		if labels, ok := labelMap[ep.DNSName]; ok {
			ep.Labels = labels
		} else {
			//this indicates that owner could not be identified, as there is no corresponding TXT record
			ep.Labels = endpoint.NewLabels()
		}
	}

	// Update the cache.
	if im.cacheInterval > 0 {
		im.recordsCache = endpoints
		im.recordsCacheRefreshTime = time.Now()
	}

	return endpoints, nil
}

// ApplyChanges updates dns provider with the changes
// for each created/deleted record it will also take into account TXT records for creation/deletion
func (im *TXTRegistry) ApplyChanges(changes *plan.Changes) error {
	filteredChanges := &plan.Changes{
		Create:    changes.Create,
		UpdateNew: filterOwnedRecords(im.ownerID, changes.UpdateNew),
		UpdateOld: filterOwnedRecords(im.ownerID, changes.UpdateOld),
		Delete:    filterOwnedRecords(im.ownerID, changes.Delete),
	}
	for _, r := range filteredChanges.Create {
		if r.Labels == nil {
			r.Labels = make(map[string]string)
		}
		r.Labels[endpoint.OwnerLabelKey] = im.ownerID
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))
		filteredChanges.Create = append(filteredChanges.Create, txt)

		if im.cacheInterval > 0 {
			im.addToCache(r)
		}
	}

	for _, r := range filteredChanges.Delete {
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))

		// when we delete TXT records for which value has changed (due to new label) this would still work because
		// !!! TXT record value is uniquely generated from the Labels of the endpoint. Hence old TXT record can be uniquely reconstructed
		filteredChanges.Delete = append(filteredChanges.Delete, txt)

		if im.cacheInterval > 0 {
			im.removeFromCache(r)
		}
	}

	// make sure TXT records are consistently updated as well
	for _, r := range filteredChanges.UpdateOld {
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))
		// when we updateOld TXT records for which value has changed (due to new label) this would still work because
		// !!! TXT record value is uniquely generated from the Labels of the endpoint. Hence old TXT record can be uniquely reconstructed
		filteredChanges.UpdateOld = append(filteredChanges.UpdateOld, txt)
		// remove old version of record from cache
		if im.cacheInterval > 0 {
			im.removeFromCache(r)
		}
	}

	// make sure TXT records are consistently updated as well
	for _, r := range filteredChanges.UpdateNew {
		txt := endpoint.NewEndpoint(im.mapper.toTXTName(r.DNSName), endpoint.RecordTypeTXT, r.Labels.Serialize(true))
		filteredChanges.UpdateNew = append(filteredChanges.UpdateNew, txt)
		// add new version of record to cache
		if im.cacheInterval > 0 {
			im.addToCache(r)
		}
	}

	return im.provider.ApplyChanges(filteredChanges)
}

/**
  TXT registry specific private methods
*/

/**
  nameMapper defines interface which maps the dns name defined for the source
  to the dns name which TXT record will be created with
*/

type nameMapper interface {
	toEndpointName(string) string
	toTXTName(string) string
}

type prefixNameMapper struct {
	prefix string
}

var _ nameMapper = prefixNameMapper{}

func newPrefixNameMapper(prefix string) prefixNameMapper {
	return prefixNameMapper{prefix: prefix}
}

func (pr prefixNameMapper) toEndpointName(txtDNSName string) string {
	if strings.HasPrefix(txtDNSName, pr.prefix) {
		return strings.TrimPrefix(txtDNSName, pr.prefix)
	}
	return ""
}

func (pr prefixNameMapper) toTXTName(endpointDNSName string) string {
	return pr.prefix + endpointDNSName
}

func (im *TXTRegistry) addToCache(ep *endpoint.Endpoint) {
	if im.recordsCache != nil {
		im.recordsCache = append(im.recordsCache, ep)
	}
}

func (im *TXTRegistry) removeFromCache(ep *endpoint.Endpoint) {
	if im.recordsCache == nil || ep == nil {
		// return early.
		return
	}

	for i, e := range im.recordsCache {
		if e.DNSName == ep.DNSName && e.RecordType == ep.RecordType && e.Targets.Same(ep.Targets) {
			// We found a match delete the endpoint from the cache.
			im.recordsCache = append(im.recordsCache[:i], im.recordsCache[i+1:]...)
			return
		}
	}
}
<|endoftext|>


<commit_before>package travel

import (
	"fmt"
	"os"

	"github.com/akerl/voyager/cartogram"

	"github.com/akerl/speculate/creds"
	"github.com/akerl/speculate/executors"
)

type hop struct {
	Profile string
	Account string
	Role    string
	Mfa     bool
}

type voyage struct {
	pack    cartogram.Pack
	account cartogram.Account
	role    string
	hops    []hop
	creds   creds.Creds
}

// Itinerary describes a travel request
type Itinerary struct {
	Args        []string
	RoleName    string
	SessionName string
	Policy      string
	Lifetime    int64
	MfaCode     string
	MfaSerial   string
}

// TravelWithOptions loads creds from a full set of parameters
func Travel(i Itinerary) (creds.Creds, error) {
	var creds creds.Creds
	v := voyage{}

	if err := v.loadPack(); err != nil {
		return creds, err
	}
	if err := v.loadAccount(i.Args); err != nil {
		return creds, err
	}
	if err := v.loadRole(i.RoleName); err != nil {
		return creds, err
	}
	if err := v.loadHops(); err != nil {
		return creds, err
	}
	if err := v.loadCreds(i); err != nil {
		return creds, err
	}
	return v.creds, nil
}

func (v *voyage) loadPack() error {
	return v.pack.Load()
}

func (v *voyage) loadAccount(args []string) error {
	var err error
	v.account, err = v.pack.Find(args)
	return err
}

func (v *voyage) loadRole(roleName string) error {
	var err error
	v.role, err = v.account.PickRole(roleName)
	return err
}

func (v *voyage) loadHops() error {
	if err := parseHops(&v.hops, v.pack, v.account, v.role); err != nil {
		return err
	}
	for i, j := 0, len(v.hops)-1; i < j; i, j = i+1, j-1 {
		v.hops[i], v.hops[j] = v.hops[j], v.hops[i]
	}
	return nil
}

func (v *voyage) loadCreds(i Itinerary) error {
	var creds creds.Creds
	var err error

	profileHop, stack := v.hops[0], v.hops[1:]
	os.Setenv("AWS_PROFILE", profileHop.Profile)

	last := len(stack) - 1
	for index, thisHop := range stack {
		a := executors.Assumption{}
		if err := a.SetAccountID(thisHop.Account); err != nil {
			return err
		}
		if err := a.SetRoleName(thisHop.Role); err != nil {
			return err
		}
		if err := a.SetSessionName(i.SessionName); err != nil {
			return err
		}
		if err := a.SetLifetime(i.Lifetime); err != nil {
			return err
		}
		if index == last {
			if err := a.SetPolicy(i.Policy); err != nil {
				return err
			}
		}
		if thisHop.Mfa {
			if err := a.SetMfa(true); err != nil {
				return err
			}
			if err := a.SetMfaSerial(i.MfaSerial); err != nil {
				return err
			}
			if err := a.SetMfaCode(i.MfaCode); err != nil {
				return err
			}
		}
		creds, err = assumption.ExecuteWithCreds(creds)
		if err != nil {
			return err
		}
	}
	v.creds = creds
	return nil
}

func parseHops(stack *[]hop, cp cartogram.Pack, a cartogram.Account, r string) error {
	*stack = append(*stack, hop{Account: a.Account, Role: r, Mfa: a.Roles[r].Mfa})
	accountMatch := cartogram.AccountRegex.FindStringSubmatch(a.Source)
	if len(accountMatch) != 4 {
		*stack = append(*stack, hop{Profile: a.Source})
		return nil
	}
	sAccountID := accountMatch[1]
	sRole := accountMatch[3]
	found, sAccount := cp.Lookup(sAccountID)
	if !found {
		return fmt.Errorf("Failed to resolve hop for %s", sAccountID)
	}
	return parseHops(stack, cp, sAccount, sRole)
}
<commit_msg>fix typos<commit_after>package travel

import (
	"fmt"
	"os"

	"github.com/akerl/voyager/cartogram"

	"github.com/akerl/speculate/creds"
	"github.com/akerl/speculate/executors"
)

type hop struct {
	Profile string
	Account string
	Role    string
	Mfa     bool
}

type voyage struct {
	pack    cartogram.Pack
	account cartogram.Account
	role    string
	hops    []hop
	creds   creds.Creds
}

// Itinerary describes a travel request
type Itinerary struct {
	Args        []string
	RoleName    string
	SessionName string
	Policy      string
	Lifetime    int64
	MfaCode     string
	MfaSerial   string
}

// Travel loads creds from a full set of parameters
func Travel(i Itinerary) (creds.Creds, error) {
	var creds creds.Creds
	v := voyage{}

	if err := v.loadPack(); err != nil {
		return creds, err
	}
	if err := v.loadAccount(i.Args); err != nil {
		return creds, err
	}
	if err := v.loadRole(i.RoleName); err != nil {
		return creds, err
	}
	if err := v.loadHops(); err != nil {
		return creds, err
	}
	if err := v.loadCreds(i); err != nil {
		return creds, err
	}
	return v.creds, nil
}

func (v *voyage) loadPack() error {
	return v.pack.Load()
}

func (v *voyage) loadAccount(args []string) error {
	var err error
	v.account, err = v.pack.Find(args)
	return err
}

func (v *voyage) loadRole(roleName string) error {
	var err error
	v.role, err = v.account.PickRole(roleName)
	return err
}

func (v *voyage) loadHops() error {
	if err := parseHops(&v.hops, v.pack, v.account, v.role); err != nil {
		return err
	}
	for i, j := 0, len(v.hops)-1; i < j; i, j = i+1, j-1 {
		v.hops[i], v.hops[j] = v.hops[j], v.hops[i]
	}
	return nil
}

func (v *voyage) loadCreds(i Itinerary) error {
	var creds creds.Creds
	var err error

	profileHop, stack := v.hops[0], v.hops[1:]
	os.Setenv("AWS_PROFILE", profileHop.Profile)

	last := len(stack) - 1
	for index, thisHop := range stack {
		a := executors.Assumption{}
		if err := a.SetAccountID(thisHop.Account); err != nil {
			return err
		}
		if err := a.SetRoleName(thisHop.Role); err != nil {
			return err
		}
		if err := a.SetSessionName(i.SessionName); err != nil {
			return err
		}
		if err := a.SetLifetime(i.Lifetime); err != nil {
			return err
		}
		if index == last {
			if err := a.SetPolicy(i.Policy); err != nil {
				return err
			}
		}
		if thisHop.Mfa {
			if err := a.SetMfa(true); err != nil {
				return err
			}
			if err := a.SetMfaSerial(i.MfaSerial); err != nil {
				return err
			}
			if err := a.SetMfaCode(i.MfaCode); err != nil {
				return err
			}
		}
		creds, err = a.ExecuteWithCreds(creds)
		if err != nil {
			return err
		}
	}
	v.creds = creds
	return nil
}

func parseHops(stack *[]hop, cp cartogram.Pack, a cartogram.Account, r string) error {
	*stack = append(*stack, hop{Account: a.Account, Role: r, Mfa: a.Roles[r].Mfa})
	accountMatch := cartogram.AccountRegex.FindStringSubmatch(a.Source)
	if len(accountMatch) != 4 {
		*stack = append(*stack, hop{Profile: a.Source})
		return nil
	}
	sAccountID := accountMatch[1]
	sRole := accountMatch[3]
	found, sAccount := cp.Lookup(sAccountID)
	if !found {
		return fmt.Errorf("Failed to resolve hop for %s", sAccountID)
	}
	return parseHops(stack, cp, sAccount, sRole)
}
<|endoftext|>


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

package oracle

import (
	"fmt"
	"go/token"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/oracle/serial"
)

// Callers reports the possible callers of the function
// immediately enclosing the specified source location.
//
func callers(q *Query) error {
	lconf := loader.Config{Build: q.Build}

	if err := setPTAScope(&lconf, q.Scope); err != nil {
		return err
	}

	// Load/parse/type-check the program.
	lprog, err := lconf.Load()
	if err != nil {
		return err
	}
	q.Fset = lprog.Fset

	qpos, err := parseQueryPos(lprog, q.Pos, false)
	if err != nil {
		return err
	}

	prog := ssa.Create(lprog, 0)

	ptaConfig, err := setupPTA(prog, lprog, q.PTALog, q.Reflection)
	if err != nil {
		return err
	}

	pkg := prog.Package(qpos.info.Pkg)
	if pkg == nil {
		return fmt.Errorf("no SSA package")
	}
	if !ssa.HasEnclosingFunction(pkg, qpos.path) {
		return fmt.Errorf("this position is not inside a function")
	}

	// Defer SSA construction till after errors are reported.
	prog.BuildAll()

	target := ssa.EnclosingFunction(pkg, qpos.path)
	if target == nil {
		return fmt.Errorf("no SSA function built for this location (dead code?)")
	}

	// Run the pointer analysis, recording each
	// call found to originate from target.
	ptaConfig.BuildCallGraph = true
	cg := ptrAnalysis(ptaConfig).CallGraph
	cg.DeleteSyntheticNodes()
	edges := cg.CreateNode(target).In
	// TODO(adonovan): sort + dedup calls to ensure test determinism.

	q.result = &callersResult{
		target:    target,
		callgraph: cg,
		edges:     edges,
	}
	return nil
}

type callersResult struct {
	target    *ssa.Function
	callgraph *callgraph.Graph
	edges     []*callgraph.Edge
}

func (r *callersResult) display(printf printfFunc) {
	root := r.callgraph.Root
	if r.edges == nil {
		printf(r.target, "%s is not reachable in this program.", r.target)
	} else {
		printf(r.target, "%s is called from these %d sites:", r.target, len(r.edges))
		for _, edge := range r.edges {
			if edge.Caller == root {
				printf(r.target, "the root of the call graph")
			} else {
				printf(edge, "\t%s from %s", edge.Description(), edge.Caller.Func)
			}
		}
	}
}

func (r *callersResult) toSerial(res *serial.Result, fset *token.FileSet) {
	var callers []serial.Caller
	for _, edge := range r.edges {
		callers = append(callers, serial.Caller{
			Caller: edge.Caller.Func.String(),
			Pos:    fset.Position(edge.Pos()).String(),
			Desc:   edge.Description(),
		})
	}
	res.Callers = callers
}
<commit_msg>oracle: add TODO<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package oracle

import (
	"fmt"
	"go/token"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/oracle/serial"
)

// Callers reports the possible callers of the function
// immediately enclosing the specified source location.
//
func callers(q *Query) error {
	lconf := loader.Config{Build: q.Build}

	if err := setPTAScope(&lconf, q.Scope); err != nil {
		return err
	}

	// Load/parse/type-check the program.
	lprog, err := lconf.Load()
	if err != nil {
		return err
	}
	q.Fset = lprog.Fset

	qpos, err := parseQueryPos(lprog, q.Pos, false)
	if err != nil {
		return err
	}

	prog := ssa.Create(lprog, 0)

	ptaConfig, err := setupPTA(prog, lprog, q.PTALog, q.Reflection)
	if err != nil {
		return err
	}

	pkg := prog.Package(qpos.info.Pkg)
	if pkg == nil {
		return fmt.Errorf("no SSA package")
	}
	if !ssa.HasEnclosingFunction(pkg, qpos.path) {
		return fmt.Errorf("this position is not inside a function")
	}

	// Defer SSA construction till after errors are reported.
	prog.BuildAll()

	target := ssa.EnclosingFunction(pkg, qpos.path)
	if target == nil {
		return fmt.Errorf("no SSA function built for this location (dead code?)")
	}

	// TODO(adonovan): opt: if function is never address-taken, skip
	// the pointer analysis.  Just look for direct calls.  This can
	// be done in a single pass over the SSA.

	// Run the pointer analysis, recording each
	// call found to originate from target.
	ptaConfig.BuildCallGraph = true
	cg := ptrAnalysis(ptaConfig).CallGraph
	cg.DeleteSyntheticNodes()
	edges := cg.CreateNode(target).In
	// TODO(adonovan): sort + dedup calls to ensure test determinism.

	q.result = &callersResult{
		target:    target,
		callgraph: cg,
		edges:     edges,
	}
	return nil
}

type callersResult struct {
	target    *ssa.Function
	callgraph *callgraph.Graph
	edges     []*callgraph.Edge
}

func (r *callersResult) display(printf printfFunc) {
	root := r.callgraph.Root
	if r.edges == nil {
		printf(r.target, "%s is not reachable in this program.", r.target)
	} else {
		printf(r.target, "%s is called from these %d sites:", r.target, len(r.edges))
		for _, edge := range r.edges {
			if edge.Caller == root {
				printf(r.target, "the root of the call graph")
			} else {
				printf(edge, "\t%s from %s", edge.Description(), edge.Caller.Func)
			}
		}
	}
}

func (r *callersResult) toSerial(res *serial.Result, fset *token.FileSet) {
	var callers []serial.Caller
	for _, edge := range r.edges {
		callers = append(callers, serial.Caller{
			Caller: edge.Caller.Func.String(),
			Pos:    fset.Position(edge.Pos()).String(),
			Desc:   edge.Description(),
		})
	}
	res.Callers = callers
}
<|endoftext|>


<commit_before>package orgnetsim

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func IsFalse(t *testing.T, condition bool, msg string) {
	if condition {
		t.Error(msg)
	}
}

func IsTrue(t *testing.T, condition bool, msg string) {
	if !condition {
		t.Error(msg)
	}
}

func AreEqual(t *testing.T, expected interface{}, actual interface{}, msg string) {
	if expected != actual {
		t.Errorf("%s Expected = %v Actual = %v", msg, expected, actual)
	}
}

func NotEqual(t *testing.T, expected interface{}, actual interface{}, msg string) {
	if expected == actual {
		t.Errorf("%s Expected = %v Actual = %v", msg, expected, actual)
	}
}

func AssertSuccess(t *testing.T, err error) {
	if err != nil {
		t.Errorf(err.Error())
	}
}

//Convenience method to dump the colors and conversations arrays into a csv file
func WriteOutput(t *testing.T, filename string, s HierarchySpec, n RelationshipMgr, colors [][]int, conversations []int) {
	f, err := os.Create(filename)
	AssertSuccess(t, err)
	defer f.Close()

	var buffer bytes.Buffer

	for c := 0; c < n.MaxColors(); c++ {
		buffer.WriteString(fmt.Sprintf("%s,", Color(c).String()))
	}

	buffer.WriteString("Conversations,,Node,Influence,Susceptibility,Contrariness,Color,Change Count,,Link,Strength,,Levels,TeamSize,TeamLinkLevel,LinkTeamPeers,LinkTeams,InitColors,MaxColors,EvangelistAgents,LoneEvangelist,AgentsWithMemory\n")

	agents := n.Agents()
	links := n.Links()
	iterations := len(conversations)
	agentCount := len(agents)
	linkCount := len(links)
	totalLines := iterations
	if agentCount > totalLines {
		totalLines = agentCount
	}
	if linkCount > totalLines {
		totalLines = linkCount
	}

	for i := 0; i < totalLines; i++ {
		if i < iterations {
			for j := 0; j < n.MaxColors(); j++ {
				buffer.WriteString(fmt.Sprintf("%d,", colors[i][j]))
			}
			buffer.WriteString(fmt.Sprintf("%d", conversations[i]))
		} else {
			for j := 0; j < n.MaxColors(); j++ {
				buffer.WriteString(",")
			}
		}
		if i < agentCount {
			buffer.WriteString(fmt.Sprintf(",,%s,%f,%f,%f,%s,%d", agents[i].Identifier(), agents[i].State().Influence, agents[i].State().Susceptability, agents[i].State().Contrariness, agents[i].State().Color.String(), agents[i].State().ChangeCount))
		} else {
			buffer.WriteString(",,,,,,,")
		}
		if i < linkCount {
			buffer.WriteString(fmt.Sprintf(",,%s-%s,%d", links[i].Agent1ID, links[i].Agent2ID, links[i].Strength))
		} else {
			buffer.WriteString(",,,")
		}
		if i == 0 {
			var initColors string
			for x := 0; x < len(s.InitColors); x++ {
				initColors = initColors + Color(s.InitColors[x]).String()
			}
			buffer.WriteString(fmt.Sprintf(",,%d,%d,%d,%t,%t,%s,%d,%t,%t,%t\n", s.Levels, s.TeamSize, s.TeamLinkLevel, s.LinkTeamPeers, s.LinkTeams, initColors, s.MaxColors, s.EvangelistAgents, s.LoneEvangelist, s.AgentsWithMemory))
		} else {
			buffer.WriteString("\n")
		}
	}

	_, err = f.Write(buffer.Bytes())
	AssertSuccess(t, err)
}

func TestRunSim(t *testing.T) {
	s := GenerateNetwork(t)
	RunSimFromJSON(t, s)
}

func GenerateNetwork(t *testing.T) HierarchySpec {

	s := HierarchySpec{
		Levels:           4,
		TeamSize:         5,
		TeamLinkLevel:    3,
		LinkTeamPeers:    true,
		LinkTeams:        false,
		InitColors:       []Color{Grey, Red},
		MaxColors:        4,
		EvangelistAgents: false,
		LoneEvangelist:   false,
		AgentsWithMemory: true,
	}

	n, err := GenerateHierarchy(s)
	AssertSuccess(t, err)

	json := n.Serialise()

	f, err := os.Create("out.json")
	AssertSuccess(t, err)
	defer f.Close()

	_, err = f.Write([]byte(json))
	AssertSuccess(t, err)

	return s
}

func RunSimFromJSON(t *testing.T, s HierarchySpec) {
	infile := "out.json"
	json, err := ioutil.ReadFile(infile)
	AssertSuccess(t, err)

	n, err := NewNetwork(string(json))
	AssertSuccess(t, err)

	colors, conversations := RunSim(n, 2000)
	outfile := strings.Replace(infile, ".json", ".csv", 1)
	WriteOutput(t, outfile, s, n, colors, conversations)
}
<commit_msg>Add an example integration test that runs a sim from a json file.<commit_after>package orgnetsim

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func IsFalse(t *testing.T, condition bool, msg string) {
	if condition {
		t.Error(msg)
	}
}

func IsTrue(t *testing.T, condition bool, msg string) {
	if !condition {
		t.Error(msg)
	}
}

func AreEqual(t *testing.T, expected interface{}, actual interface{}, msg string) {
	if expected != actual {
		t.Errorf("%s Expected = %v Actual = %v", msg, expected, actual)
	}
}

func NotEqual(t *testing.T, expected interface{}, actual interface{}, msg string) {
	if expected == actual {
		t.Errorf("%s Expected = %v Actual = %v", msg, expected, actual)
	}
}

func AssertSuccess(t *testing.T, err error) {
	if err != nil {
		t.Errorf(err.Error())
	}
}

//Convenience method to dump the colors and conversations arrays into a csv file
func WriteOutput(t *testing.T, filename string, s HierarchySpec, n RelationshipMgr, colors [][]int, conversations []int) {
	f, err := os.Create(filename)
	AssertSuccess(t, err)
	defer f.Close()

	var buffer bytes.Buffer

	for c := 0; c < n.MaxColors(); c++ {
		buffer.WriteString(fmt.Sprintf("%s,", Color(c).String()))
	}

	buffer.WriteString("Conversations,,Node,Influence,Susceptibility,Contrariness,Color,Change Count,,Link,Strength,,Levels,TeamSize,TeamLinkLevel,LinkTeamPeers,LinkTeams,InitColors,MaxColors,EvangelistAgents,LoneEvangelist,AgentsWithMemory\n")

	agents := n.Agents()
	links := n.Links()
	iterations := len(conversations)
	agentCount := len(agents)
	linkCount := len(links)
	totalLines := iterations
	if agentCount > totalLines {
		totalLines = agentCount
	}
	if linkCount > totalLines {
		totalLines = linkCount
	}

	for i := 0; i < totalLines; i++ {
		if i < iterations {
			for j := 0; j < n.MaxColors(); j++ {
				buffer.WriteString(fmt.Sprintf("%d,", colors[i][j]))
			}
			buffer.WriteString(fmt.Sprintf("%d", conversations[i]))
		} else {
			for j := 0; j < n.MaxColors(); j++ {
				buffer.WriteString(",")
			}
		}
		if i < agentCount {
			buffer.WriteString(fmt.Sprintf(",,%s,%f,%f,%f,%s,%d", agents[i].Identifier(), agents[i].State().Influence, agents[i].State().Susceptability, agents[i].State().Contrariness, agents[i].State().Color.String(), agents[i].State().ChangeCount))
		} else {
			buffer.WriteString(",,,,,,,")
		}
		if i < linkCount {
			buffer.WriteString(fmt.Sprintf(",,%s-%s,%d", links[i].Agent1ID, links[i].Agent2ID, links[i].Strength))
		} else {
			buffer.WriteString(",,,")
		}
		if i == 0 {
			var initColors string
			for x := 0; x < len(s.InitColors); x++ {
				initColors = initColors + Color(s.InitColors[x]).String()
			}
			buffer.WriteString(fmt.Sprintf(",,%d,%d,%d,%t,%t,%s,%d,%t,%t,%t\n", s.Levels, s.TeamSize, s.TeamLinkLevel, s.LinkTeamPeers, s.LinkTeams, initColors, s.MaxColors, s.EvangelistAgents, s.LoneEvangelist, s.AgentsWithMemory))
		} else {
			buffer.WriteString("\n")
		}
	}

	_, err = f.Write(buffer.Bytes())
	AssertSuccess(t, err)
}

func GenerateNetwork(t *testing.T, filename string) HierarchySpec {

	s := HierarchySpec{
		Levels:           4,
		TeamSize:         5,
		TeamLinkLevel:    3,
		LinkTeamPeers:    true,
		LinkTeams:        false,
		InitColors:       []Color{Grey, Red},
		MaxColors:        4,
		EvangelistAgents: false,
		LoneEvangelist:   false,
		AgentsWithMemory: true,
	}

	n, err := GenerateHierarchy(s)
	AssertSuccess(t, err)

	json := n.Serialise()

	f, err := os.Create(filename)
	AssertSuccess(t, err)
	defer f.Close()

	_, err = f.Write([]byte(json))
	AssertSuccess(t, err)

	return s
}

func RunSimFromJSON(t *testing.T, filename string, s HierarchySpec) {
	json, err := ioutil.ReadFile(filename)
	AssertSuccess(t, err)

	n, err := NewNetwork(string(json))
	AssertSuccess(t, err)

	colors, conversations := RunSim(n, 2000)
	outfile := strings.Replace(filename, ".json", ".csv", 1)
	WriteOutput(t, outfile, s, n, colors, conversations)
}

func TestGenerateAndRunSim(t *testing.T) {
	filename := "tst.json"
	s := GenerateNetwork(t, filename)
	RunSimFromJSON(t, filename, s)
}

//Run a sim from a specific JSON file
func TestRunSim(t *testing.T) {
	t.SkipNow() //Comment this line to stop skipping the test
	filename := "out.json"
	s := HierarchySpec{}
	RunSimFromJSON(t, filename, s)
}
<|endoftext|>


<commit_before>package apollostats

import (
	"strings"
	"time"
)

type Ban struct {
	ID         int64     `gorm:"column:id;primary_key"`
	Timestamp  time.Time `gorm:"column:bantime"`
	CKey       string    `gorm:"column:ckey"`
	CID        string    `gorm:"column:computerid"`
	IP         string    `gorm:"column:ip"`
	Bantype    string    `gorm:"column:bantype"`
	Admin      string    `gorm:"column:a_ckey"`
	Reason     string    `gorm:"column:reason"`
	Duration   int       `gorm:"column:duration"`
	Expiration time.Time `gorm:"column:expiration_time"`
}

func (b *Ban) TableName() string {
	return "ban"
}

type AccountItem struct {
	ID        int64     `gorm:"column:id;primary_key"`
	Timestamp time.Time `gorm:"column:time"`
	CKey      string    `gorm:"column:ckey"`
	Item      string    `gorm:"column:item"`
}

func (a *AccountItem) TableName() string {
	return "acc_items"
}

type Death struct {
	ID        int64     `gorm:"column:id;primary_key"`
	Timestamp time.Time `gorm:"column:tod"`
	Name      string    `gorm:"column:name"`
	Job       string    `gorm:"column:job"`
	Room      string    `gorm:"column:pod"`
	Position  string    `gorm:"column:coord"`
	Brute     int       `gorm:"column:bruteloss"`
	Brain     int       `gorm:"column:brainloss"`
	Fire      int       `gorm:"column:fireloss"`
	Oxygen    int       `gorm:"column:oxyloss"`
}

func (d *Death) TableName() string {
	return "death"
}

func (d *Death) RoomName() string {
	// Cleanup the room name
	return strings.Trim(d.Room, "Ã¿")
}
<commit_msg>Adds func to pretty format a ban expire time.<commit_after>package apollostats

import (
	"strings"
	"time"
)

type Ban struct {
	ID         int64     `gorm:"column:id;primary_key"`
	Timestamp  time.Time `gorm:"column:bantime"`
	CKey       string    `gorm:"column:ckey"`
	CID        string    `gorm:"column:computerid"`
	IP         string    `gorm:"column:ip"`
	Bantype    string    `gorm:"column:bantype"`
	Admin      string    `gorm:"column:a_ckey"`
	Reason     string    `gorm:"column:reason"`
	Duration   int       `gorm:"column:duration"`
	Expiration time.Time `gorm:"column:expiration_time"`
}

func (b *Ban) TableName() string {
	return "ban"
}

func (b *Ban) Expires() string {
	if b.Duration < 0 {
		return "Permanent"
	} else {
		return b.Expiration.Format("2006-01-02 15:04 MST")
	}
}

type AccountItem struct {
	ID        int64     `gorm:"column:id;primary_key"`
	Timestamp time.Time `gorm:"column:time"`
	CKey      string    `gorm:"column:ckey"`
	Item      string    `gorm:"column:item"`
}

func (a *AccountItem) TableName() string {
	return "acc_items"
}

type Death struct {
	ID        int64     `gorm:"column:id;primary_key"`
	Timestamp time.Time `gorm:"column:tod"`
	Name      string    `gorm:"column:name"`
	Job       string    `gorm:"column:job"`
	Room      string    `gorm:"column:pod"`
	Position  string    `gorm:"column:coord"`
	Brute     int       `gorm:"column:bruteloss"`
	Brain     int       `gorm:"column:brainloss"`
	Fire      int       `gorm:"column:fireloss"`
	Oxygen    int       `gorm:"column:oxyloss"`
}

func (d *Death) TableName() string {
	return "death"
}

func (d *Death) RoomName() string {
	// Cleanup the room name
	return strings.Trim(d.Room, "Ã¿")
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package upgrades_test

import (
	gc "launchpad.net/gocheck"

	jujutesting "launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/upgrades"
)

type processDeprecatedAttributesSuite struct {
	jujutesting.JujuConnSuite
	ctx upgrades.Context
}

var _ = gc.Suite(&processDeprecatedAttributesSuite{})

func (s *processDeprecatedAttributesSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	apiState, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	s.ctx = &mockContext{
		agentConfig: &mockAgentConfig{dataDir: s.DataDir()},
		apiState:    apiState,
		state:       s.State,
	}
	cfg, err := s.State.EnvironConfig()
	c.Assert(err, gc.IsNil)
	// Add in old public bucket config.
	newCfg, err := cfg.Apply(map[string]interface{}{
		"public-bucket":         "foo",
		"public-bucket-region":  "bar",
		"public-bucket-url":     "shazbot",
		"default-instance-type": "vulch",
		"default-image-id":      "1234",
	})
	c.Assert(err, gc.IsNil)
	err = s.State.SetEnvironConfig(newCfg, cfg)
	c.Assert(err, gc.IsNil)
	cfg, err = s.State.EnvironConfig()
	c.Assert(err, gc.IsNil)
	allAttrs := cfg.AllAttrs()
	c.Assert(allAttrs["public-bucket"], gc.Equals, "foo")
	c.Assert(allAttrs["public-bucket-region"], gc.Equals, "bar")
	c.Assert(allAttrs["public-bucket-url"], gc.Equals, "shazbot")
	c.Assert(allAttrs["default-instance-type"], gc.Equals, "vulch")
	c.Assert(allAttrs["default-image-id"], gc.Equals, "1234")
}

func (s *processDeprecatedAttributesSuite) assertConfigProcessed(c *gc.C) {
	cfg, err := s.State.EnvironConfig()
	c.Assert(err, gc.IsNil)
	allAttrs := cfg.AllAttrs()
	for _, deprecated := range []string{
		"public-bucket", "public-bucket-region", "public-bucket-url",
		"default-image-id", "default-instance-type",
	} {
		_, ok := allAttrs[deprecated]
		c.Assert(ok, jc.IsFalse)
	}
}

func (s *processDeprecatedAttributesSuite) TestOldConfigRemoved(c *gc.C) {
	err := upgrades.ProcessDeprecatedAttributes(s.ctx)
	c.Assert(err, gc.IsNil)
	s.assertConfigProcessed(c)
}

func (s *processDeprecatedAttributesSuite) TestIdempotent(c *gc.C) {
	err := upgrades.ProcessDeprecatedAttributes(s.ctx)
	c.Assert(err, gc.IsNil)
	s.assertConfigProcessed(c)

	err = upgrades.ProcessDeprecatedAttributes(s.ctx)
	c.Assert(err, gc.IsNil)
	s.assertConfigProcessed(c)
}
<commit_msg>Move setup assets to explicit test<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package upgrades_test

import (
	gc "launchpad.net/gocheck"

	jujutesting "launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/upgrades"
)

type processDeprecatedAttributesSuite struct {
	jujutesting.JujuConnSuite
	ctx upgrades.Context
}

var _ = gc.Suite(&processDeprecatedAttributesSuite{})

func (s *processDeprecatedAttributesSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	apiState, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	s.ctx = &mockContext{
		agentConfig: &mockAgentConfig{dataDir: s.DataDir()},
		apiState:    apiState,
		state:       s.State,
	}
	cfg, err := s.State.EnvironConfig()
	c.Assert(err, gc.IsNil)
	// Add in old public bucket config.
	newCfg, err := cfg.Apply(map[string]interface{}{
		"public-bucket":         "foo",
		"public-bucket-region":  "bar",
		"public-bucket-url":     "shazbot",
		"default-instance-type": "vulch",
		"default-image-id":      "1234",
	})
	c.Assert(err, gc.IsNil)
	err = s.State.SetEnvironConfig(newCfg, cfg)
	c.Assert(err, gc.IsNil)
}

func (s *processDeprecatedAttributesSuite) TestAttributesSet(c *gc.C) {
	cfg, err := s.State.EnvironConfig()
	c.Assert(err, gc.IsNil)
	allAttrs := cfg.AllAttrs()
	c.Assert(allAttrs["public-bucket"], gc.Equals, "foo")
	c.Assert(allAttrs["public-bucket-region"], gc.Equals, "bar")
	c.Assert(allAttrs["public-bucket-url"], gc.Equals, "shazbot")
	c.Assert(allAttrs["default-instance-type"], gc.Equals, "vulch")
	c.Assert(allAttrs["default-image-id"], gc.Equals, "1234")
}

func (s *processDeprecatedAttributesSuite) assertConfigProcessed(c *gc.C) {
	cfg, err := s.State.EnvironConfig()
	c.Assert(err, gc.IsNil)
	allAttrs := cfg.AllAttrs()
	for _, deprecated := range []string{
		"public-bucket", "public-bucket-region", "public-bucket-url",
		"default-image-id", "default-instance-type",
	} {
		_, ok := allAttrs[deprecated]
		c.Assert(ok, jc.IsFalse)
	}
}

func (s *processDeprecatedAttributesSuite) TestOldConfigRemoved(c *gc.C) {
	err := upgrades.ProcessDeprecatedAttributes(s.ctx)
	c.Assert(err, gc.IsNil)
	s.assertConfigProcessed(c)
}

func (s *processDeprecatedAttributesSuite) TestIdempotent(c *gc.C) {
	err := upgrades.ProcessDeprecatedAttributes(s.ctx)
	c.Assert(err, gc.IsNil)
	s.assertConfigProcessed(c)

	err = upgrades.ProcessDeprecatedAttributes(s.ctx)
	c.Assert(err, gc.IsNil)
	s.assertConfigProcessed(c)
}
<|endoftext|>