<commit_before>package qln

import (
	"fmt"
	"path/filepath"
	"sync"

	"github.com/boltdb/bolt"
	"github.com/mit-dci/lit/btcutil"
	"github.com/mit-dci/lit/btcutil/hdkeychain"
	"github.com/mit-dci/lit/coinparam"
	"github.com/mit-dci/lit/db/lnbolt" // TODO Abstract this more.
	"github.com/mit-dci/lit/dlc"
	"github.com/mit-dci/lit/eventbus"
	"github.com/mit-dci/lit/lncore"
	"github.com/mit-dci/lit/lnp2p"
	"github.com/mit-dci/lit/logging"
	"github.com/mit-dci/lit/portxo"
	"github.com/mit-dci/lit/wallit"
	"github.com/mit-dci/lit/watchtower"
)

// NewLitNode starts up a lit node.  Needs priv key, and a path.
// Does not activate a subwallet; do that after init.
func NewLitNode(privKey *[32]byte, path string, trackerURL string, proxyURL string, nat string) (*LitNode, error) {

	var err error

	// Maybe make a new parameter set for "LN".. meh
	// TODO change this to a non-coin
	rootPrivKey, err := hdkeychain.NewMaster(privKey[:], &coinparam.TestNet3Params)
	if err != nil {
		return nil, err
	}

	nd := new(LitNode)
	nd.LitFolder = path

	litdbpath := filepath.Join(nd.LitFolder, "ln.db")
	err = nd.OpenDB(litdbpath)
	if err != nil {
		return nil, err
	}

	db2 := &lnbolt.LitBoltDB{}
	var dbx lncore.LitStorage
	dbx = db2
	err = dbx.Open(filepath.Join(nd.LitFolder, "db2"))
	if err != nil {
		return nil, err
	}
	nd.NewLitDB = dbx
	err = nd.NewLitDB.Check()
	if err != nil {
		return nil, err
	}

	// Event system setup.
	ebus := eventbus.NewEventBus()
	nd.Events = &ebus

	// Peer manager
	nd.PeerMan, err = lnp2p.NewPeerManager(rootPrivKey, nd.NewLitDB.GetPeerDB(), trackerURL, &ebus, nil) // TODO proxy/nat stuff
	if err != nil {
		return nil, err
	}

	// Actually start the thread to send messages.
	nd.PeerMan.StartSending()

	// Register adapter event handlers.  These are for hooking in the new peer management with the old one.
	h1 := makeTmpNewPeerHandler(nd)
	nd.Events.RegisterHandler("lnp2p.peer.new", h1)
	h2 := makeTmpDisconnectPeerHandler(nd)
	nd.Events.RegisterHandler("lnp2p.peer.disconnect", h2)
	h3 := makeTmpSigProofHandler(nd)
	nd.Events.RegisterHandler("qln.chanupdate.sigproof", h3)

	// Sets up handlers for all the messages we need to handle.
	nd.registerHandlers()

	var kg portxo.KeyGen
	kg.Depth = 5
	kg.Step[0] = 44 | 1<<31
	kg.Step[1] = 513 | 1<<31
	kg.Step[2] = 9 | 1<<31
	kg.Step[3] = 0 | 1<<31
	kg.Step[4] = 0 | 1<<31
	nd.IdentityKey, err = kg.DerivePrivateKey(rootPrivKey)
	if err != nil {
		return nil, err
	}

	kg.Step[3] = 1 | 1<<31
	rcPriv, err := kg.DerivePrivateKey(rootPrivKey)
	nd.DefaultRemoteControlKey = rcPriv.PubKey()
	rcPriv = nil

	nd.TrackerURL = trackerURL

	nd.ProxyURL = proxyURL

	nd.Nat = nat

	nd.InitRouting()

	// optional tower activation

	nd.Tower = new(watchtower.WatchTower)

	// Create a new manager for the discreet log contracts
	nd.DlcManager, err = dlc.NewManager(filepath.Join(nd.LitFolder, "dlc.db"))
	if err != nil {
		return nil, err
	}

	// make maps and channels
	nd.UserMessageBox = make(chan string, 32)

	nd.InProg = new(InFlightFund)
	nd.InProg.done = make(chan uint32, 1)

	nd.InProgDual = new(InFlightDualFund)
	nd.InProgDual.done = make(chan *DualFundingResult, 1)

	nd.InProgMultihop, err = nd.GetAllMultihopPayments()
	if err != nil {
		return nil, err
	}

	nd.RemoteMtx.Lock()
	nd.RemoteCons = make(map[uint32]*RemotePeer)
	nd.RemoteMtx.Unlock()

	nd.SubWallet = make(map[uint32]UWallet)

	// REFACTORING STUFF
	nd.PeerMap = map[*lnp2p.Peer]*RemotePeer{}
	nd.PeerMapMtx = &sync.Mutex{}

	return nd, nil
}

// LinkBaseWallet activates a wallet and hooks it into the litnode.
func (nd *LitNode) LinkBaseWallet(
	privKey *[32]byte, birthHeight int32, resync bool, tower bool,
	host string, proxy string, param *coinparam.Params) error {

	rootpriv, err := hdkeychain.NewMaster(privKey[:], param)
	if err != nil {
		return err
	}

	WallitIdx := param.HDCoinType

	// see if we've already attached a wallet for this coin type
	if nd.SubWallet[WallitIdx] != nil {
		return fmt.Errorf("coin type %d already linked", WallitIdx)
	}
	// see if startheight is below allowed with coinparam
	if birthHeight < param.StartHeight {
		return fmt.Errorf("%s birth height give as %d, but parameters start at %d",
			param.Name, birthHeight, param.StartHeight)
	}

	// see if there are other wallets already linked
	if len(nd.SubWallet) != 0 {
		// there are; assert multiwallet (may already be asserted)
		nd.MultiWallet = true
	}

	// if there aren't, Multiwallet will still be false; set new wallit to
	// be the first & default
	var cointype int
	nd.SubWallet[WallitIdx], cointype, err = wallit.NewWallit(
		rootpriv, birthHeight, resync, host, nd.LitFolder, proxy, param)

	if err != nil {
		logging.Error(err)
		return err
	}

	if nd.ConnectedCoinTypes == nil {
		nd.ConnectedCoinTypes = make(map[uint32]bool)
		nd.ConnectedCoinTypes[uint32(cointype)] = true
	}
	nd.ConnectedCoinTypes[uint32(cointype)] = true
	// re-register channel addresses
	qChans, err := nd.GetAllQchans()
	if err != nil {
		return err
	}

	for _, qChan := range qChans {
		var pkh [20]byte
		pkhSlice := btcutil.Hash160(qChan.MyRefundPub[:])
		copy(pkh[:], pkhSlice)
		nd.SubWallet[WallitIdx].ExportHook().RegisterAddress(pkh)

		logging.Infof("Registering outpoint %v", qChan.PorTxo.Op)

		nd.SubWallet[WallitIdx].WatchThis(qChan.PorTxo.Op)
	}

	go nd.OPEventHandler(nd.SubWallet[WallitIdx].LetMeKnow())
	go nd.HeightEventHandler(nd.SubWallet[WallitIdx].LetMeKnowHeight())

	if !nd.MultiWallet {
		nd.DefaultCoin = param.HDCoinType
	}

	// if this node is running a watchtower, link the watchtower to the
	// new wallet block events

	if tower {
		err = nd.Tower.HookLink(
			nd.LitFolder, param, nd.SubWallet[WallitIdx].ExportHook())
		if err != nil {
			return err
		}
	}

	return nil
}

// Opens the DB file for the LnNode
func (nd *LitNode) OpenDB(filename string) error {
	var err error

	nd.LitDB, err = bolt.Open(filename, 0644, nil)
	if err != nil {
		return err
	}
	// create buckets if they're not already there
	err = nd.LitDB.Update(func(btx *bolt.Tx) error {
		_, err := btx.CreateBucketIfNotExists(BKTChannelData)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTPeers)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTChanMap)
		if err != nil {
			return err
		}
		_, err = btx.CreateBucketIfNotExists(BKTPeerMap)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTWatch)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTHTLCOPs)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTPayments)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTRCAuth)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Add println to know where you are<commit_after>package qln

import (
	"fmt"
	"log"
	"path/filepath"
	"sync"

	"github.com/boltdb/bolt"
	"github.com/mit-dci/lit/btcutil"
	"github.com/mit-dci/lit/btcutil/hdkeychain"
	"github.com/mit-dci/lit/coinparam"
	"github.com/mit-dci/lit/db/lnbolt" // TODO Abstract this more.
	"github.com/mit-dci/lit/dlc"
	"github.com/mit-dci/lit/eventbus"
	"github.com/mit-dci/lit/lncore"
	"github.com/mit-dci/lit/lnp2p"
	"github.com/mit-dci/lit/logging"
	"github.com/mit-dci/lit/portxo"
	"github.com/mit-dci/lit/wallit"
	"github.com/mit-dci/lit/watchtower"
)

// NewLitNode starts up a lit node.  Needs priv key, and a path.
// Does not activate a subwallet; do that after init.
func NewLitNode(privKey *[32]byte, path string, trackerURL string, proxyURL string, nat string) (*LitNode, error) {

	var err error

	log.Printf("====== MODIFIED LIT NODE ======\n")

	// Maybe make a new parameter set for "LN".. meh
	// TODO change this to a non-coin
	rootPrivKey, err := hdkeychain.NewMaster(privKey[:], &coinparam.TestNet3Params)
	if err != nil {
		return nil, err
	}

	nd := new(LitNode)
	nd.LitFolder = path

	litdbpath := filepath.Join(nd.LitFolder, "ln.db")
	err = nd.OpenDB(litdbpath)
	if err != nil {
		return nil, err
	}

	db2 := &lnbolt.LitBoltDB{}
	var dbx lncore.LitStorage
	dbx = db2
	err = dbx.Open(filepath.Join(nd.LitFolder, "db2"))
	if err != nil {
		return nil, err
	}
	nd.NewLitDB = dbx
	err = nd.NewLitDB.Check()
	if err != nil {
		return nil, err
	}

	// Event system setup.
	ebus := eventbus.NewEventBus()
	nd.Events = &ebus

	// Peer manager
	nd.PeerMan, err = lnp2p.NewPeerManager(rootPrivKey, nd.NewLitDB.GetPeerDB(), trackerURL, &ebus, nil) // TODO proxy/nat stuff
	if err != nil {
		return nil, err
	}

	// Actually start the thread to send messages.
	nd.PeerMan.StartSending()

	// Register adapter event handlers.  These are for hooking in the new peer management with the old one.
	h1 := makeTmpNewPeerHandler(nd)
	nd.Events.RegisterHandler("lnp2p.peer.new", h1)
	h2 := makeTmpDisconnectPeerHandler(nd)
	nd.Events.RegisterHandler("lnp2p.peer.disconnect", h2)
	h3 := makeTmpSigProofHandler(nd)
	nd.Events.RegisterHandler("qln.chanupdate.sigproof", h3)

	// Sets up handlers for all the messages we need to handle.
	nd.registerHandlers()

	var kg portxo.KeyGen
	kg.Depth = 5
	kg.Step[0] = 44 | 1<<31
	kg.Step[1] = 513 | 1<<31
	kg.Step[2] = 9 | 1<<31
	kg.Step[3] = 0 | 1<<31
	kg.Step[4] = 0 | 1<<31
	nd.IdentityKey, err = kg.DerivePrivateKey(rootPrivKey)
	if err != nil {
		return nil, err
	}

	kg.Step[3] = 1 | 1<<31
	rcPriv, err := kg.DerivePrivateKey(rootPrivKey)
	nd.DefaultRemoteControlKey = rcPriv.PubKey()
	rcPriv = nil

	nd.TrackerURL = trackerURL

	nd.ProxyURL = proxyURL

	nd.Nat = nat

	nd.InitRouting()

	// optional tower activation

	nd.Tower = new(watchtower.WatchTower)

	// Create a new manager for the discreet log contracts
	nd.DlcManager, err = dlc.NewManager(filepath.Join(nd.LitFolder, "dlc.db"))
	if err != nil {
		return nil, err
	}

	// make maps and channels
	nd.UserMessageBox = make(chan string, 32)

	nd.InProg = new(InFlightFund)
	nd.InProg.done = make(chan uint32, 1)

	nd.InProgDual = new(InFlightDualFund)
	nd.InProgDual.done = make(chan *DualFundingResult, 1)

	nd.InProgMultihop, err = nd.GetAllMultihopPayments()
	if err != nil {
		return nil, err
	}

	nd.RemoteMtx.Lock()
	nd.RemoteCons = make(map[uint32]*RemotePeer)
	nd.RemoteMtx.Unlock()

	nd.SubWallet = make(map[uint32]UWallet)

	// REFACTORING STUFF
	nd.PeerMap = map[*lnp2p.Peer]*RemotePeer{}
	nd.PeerMapMtx = &sync.Mutex{}

	return nd, nil
}

// LinkBaseWallet activates a wallet and hooks it into the litnode.
func (nd *LitNode) LinkBaseWallet(
	privKey *[32]byte, birthHeight int32, resync bool, tower bool,
	host string, proxy string, param *coinparam.Params) error {

	rootpriv, err := hdkeychain.NewMaster(privKey[:], param)
	if err != nil {
		return err
	}

	WallitIdx := param.HDCoinType

	// see if we've already attached a wallet for this coin type
	if nd.SubWallet[WallitIdx] != nil {
		return fmt.Errorf("coin type %d already linked", WallitIdx)
	}
	// see if startheight is below allowed with coinparam
	if birthHeight < param.StartHeight {
		return fmt.Errorf("%s birth height give as %d, but parameters start at %d",
			param.Name, birthHeight, param.StartHeight)
	}

	// see if there are other wallets already linked
	if len(nd.SubWallet) != 0 {
		// there are; assert multiwallet (may already be asserted)
		nd.MultiWallet = true
	}

	// if there aren't, Multiwallet will still be false; set new wallit to
	// be the first & default
	var cointype int
	nd.SubWallet[WallitIdx], cointype, err = wallit.NewWallit(
		rootpriv, birthHeight, resync, host, nd.LitFolder, proxy, param)

	if err != nil {
		logging.Error(err)
		return err
	}

	if nd.ConnectedCoinTypes == nil {
		nd.ConnectedCoinTypes = make(map[uint32]bool)
		nd.ConnectedCoinTypes[uint32(cointype)] = true
	}
	nd.ConnectedCoinTypes[uint32(cointype)] = true
	// re-register channel addresses
	qChans, err := nd.GetAllQchans()
	if err != nil {
		return err
	}

	for _, qChan := range qChans {
		var pkh [20]byte
		pkhSlice := btcutil.Hash160(qChan.MyRefundPub[:])
		copy(pkh[:], pkhSlice)
		nd.SubWallet[WallitIdx].ExportHook().RegisterAddress(pkh)

		logging.Infof("Registering outpoint %v", qChan.PorTxo.Op)

		nd.SubWallet[WallitIdx].WatchThis(qChan.PorTxo.Op)
	}

	go nd.OPEventHandler(nd.SubWallet[WallitIdx].LetMeKnow())
	go nd.HeightEventHandler(nd.SubWallet[WallitIdx].LetMeKnowHeight())

	if !nd.MultiWallet {
		nd.DefaultCoin = param.HDCoinType
	}

	// if this node is running a watchtower, link the watchtower to the
	// new wallet block events

	if tower {
		err = nd.Tower.HookLink(
			nd.LitFolder, param, nd.SubWallet[WallitIdx].ExportHook())
		if err != nil {
			return err
		}
	}

	return nil
}

// Opens the DB file for the LnNode
func (nd *LitNode) OpenDB(filename string) error {
	var err error

	nd.LitDB, err = bolt.Open(filename, 0644, nil)
	if err != nil {
		return err
	}
	// create buckets if they're not already there
	err = nd.LitDB.Update(func(btx *bolt.Tx) error {
		_, err := btx.CreateBucketIfNotExists(BKTChannelData)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTPeers)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTChanMap)
		if err != nil {
			return err
		}
		_, err = btx.CreateBucketIfNotExists(BKTPeerMap)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTWatch)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTHTLCOPs)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTPayments)
		if err != nil {
			return err
		}

		_, err = btx.CreateBucketIfNotExists(BKTRCAuth)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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

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

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

package mungers

import (
	"fmt"
	"time"

	"k8s.io/contrib/mungegithub/github"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	stalePendingCIHours = 24
)

// StalePendingCI will ask the k8s-bot to test any PR with a LGTM that has
// been pending for more than 24 hours. This can happen when the jenkins VM
// is restarted.
//
// The real fix would be for the jenkins VM restart to not move every single
// PR to pending without actually testing...
//
// But this is our world and so we should really do this for all PRs which
// aren't likely to get another push (everything that is mergeable). Since that
// can be a lot of PRs, I'm just doing it for the LGTM PRs automatically...
//
// With minor modification this can be run easily by hand. Remove the LGTM check
// godep go build
// ./mungegithub --token-file=/PATH/TO/YOUR/TOKEN --pr-mungers=rerun-old-pending-ci --once (--dry-run)
type StalePendingCI struct{}

func init() {
	RegisterMungerOrDie(StalePendingCI{})
}

// Name is the name usable in --pr-mungers
func (StalePendingCI) Name() string { return "stale-pending-ci" }

// Initialize will initialize the munger
func (StalePendingCI) Initialize(config *github.Config) error { return nil }

// EachLoop is called at the start of every munge loop
func (StalePendingCI) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (StalePendingCI) AddFlags(cmd *cobra.Command, config *github.Config) {}

// Munge is the workhorse the will actually make updates to the PR
func (StalePendingCI) Munge(obj *github.MungeObject) {
	requiredContexts := []string{jenkinsUnitContext, jenkinsE2EContext}

	if !obj.IsPR() {
		return
	}

	if !obj.HasLabel("lgtm") {
		return
	}

	if mergeable, err := obj.IsMergeable(); !mergeable || err != nil {
		return
	}

	status := obj.GetStatusState(requiredContexts)
	if status != "pending" {
		return
	}

	for _, context := range requiredContexts {
		statusTime := obj.GetStatusTime(context)
		if statusTime == nil {
			glog.Errorf("%d: unable to determine time %q context was set", *obj.Issue.Number, context)
			return
		}
		if time.Since(*statusTime) > stalePendingCIHours*time.Hour {
			msgFormat := `@k8s-bot test this issue: #IGNORE

Tests have been pending for %d hours`
			msg := fmt.Sprintf(msgFormat, stalePendingCIHours)
			obj.WriteComment(msg)
			return
		}
	}
}
<commit_msg>Minor comment fix in stale-pending-ci<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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

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

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

package mungers

import (
	"fmt"
	"time"

	"k8s.io/contrib/mungegithub/github"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	stalePendingCIHours = 24
)

// StalePendingCI will ask the k8s-bot to test any PR with a LGTM that has
// been pending for more than 24 hours. This can happen when the jenkins VM
// is restarted.
//
// The real fix would be for the jenkins VM restart to not move every single
// PR to pending without actually testing...
//
// But this is our world and so we should really do this for all PRs which
// aren't likely to get another push (everything that is mergeable). Since that
// can be a lot of PRs, I'm just doing it for the LGTM PRs automatically...
//
// With minor modification this can be run easily by hand. Remove the LGTM check
// godep go build
// ./mungegithub --token-file=/PATH/TO/YOUR/TOKEN --pr-mungers=stale-pending-ci --once (--dry-run)
type StalePendingCI struct{}

func init() {
	RegisterMungerOrDie(StalePendingCI{})
}

// Name is the name usable in --pr-mungers
func (StalePendingCI) Name() string { return "stale-pending-ci" }

// Initialize will initialize the munger
func (StalePendingCI) Initialize(config *github.Config) error { return nil }

// EachLoop is called at the start of every munge loop
func (StalePendingCI) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (StalePendingCI) AddFlags(cmd *cobra.Command, config *github.Config) {}

// Munge is the workhorse the will actually make updates to the PR
func (StalePendingCI) Munge(obj *github.MungeObject) {
	requiredContexts := []string{jenkinsUnitContext, jenkinsE2EContext}

	if !obj.IsPR() {
		return
	}

	if !obj.HasLabel("lgtm") {
		return
	}

	if mergeable, err := obj.IsMergeable(); !mergeable || err != nil {
		return
	}

	status := obj.GetStatusState(requiredContexts)
	if status != "pending" {
		return
	}

	for _, context := range requiredContexts {
		statusTime := obj.GetStatusTime(context)
		if statusTime == nil {
			glog.Errorf("%d: unable to determine time %q context was set", *obj.Issue.Number, context)
			return
		}
		if time.Since(*statusTime) > stalePendingCIHours*time.Hour {
			msgFormat := `@k8s-bot test this issue: #IGNORE

Tests have been pending for %d hours`
			msg := fmt.Sprintf(msgFormat, stalePendingCIHours)
			obj.WriteComment(msg)
			return
		}
	}
}
<|endoftext|>


<commit_before>package redis

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"strings"
	"sync"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type Redis struct {
	Servers []string
}

var sampleConfig = `
[[inputs.redis]]
  ## specify servers via a url matching:
  ##  [protocol://][:password]@address[:port]
  ##  e.g.
  ##    tcp://localhost:6379
  ##    tcp://:password@192.168.99.100
  ##
  ## If no servers are specified, then localhost is used as the host.
  ## If no port is specified, 6379 is used
  servers = ["tcp://localhost:6379"]
`

func (r *Redis) SampleConfig() string {
	return sampleConfig
}

func (r *Redis) Description() string {
	return "Read metrics from one or many redis servers"
}

var Tracking = map[string]string{
	"uptime_in_seconds":           "uptime",
	"connected_clients":           "clients",
	"used_memory":                 "used_memory",
	"used_memory_rss":             "used_memory_rss",
	"used_memory_peak":            "used_memory_peak",
	"used_memory_lua":             "used_memory_lua",
	"rdb_changes_since_last_save": "rdb_changes_since_last_save",
	"total_connections_received":  "total_connections_received",
	"total_commands_processed":    "total_commands_processed",
	"instantaneous_ops_per_sec":   "instantaneous_ops_per_sec",
	"instantaneous_input_kbps":    "instantaneous_input_kbps",
	"instantaneous_output_kbps":   "instantaneous_output_kbps",
	"sync_full":                   "sync_full",
	"sync_partial_ok":             "sync_partial_ok",
	"sync_partial_err":            "sync_partial_err",
	"expired_keys":                "expired_keys",
	"evicted_keys":                "evicted_keys",
	"keyspace_hits":               "keyspace_hits",
	"keyspace_misses":             "keyspace_misses",
	"pubsub_channels":             "pubsub_channels",
	"pubsub_patterns":             "pubsub_patterns",
	"latest_fork_usec":            "latest_fork_usec",
	"connected_slaves":            "connected_slaves",
	"master_repl_offset":          "master_repl_offset",
	"repl_backlog_active":         "repl_backlog_active",
	"repl_backlog_size":           "repl_backlog_size",
	"repl_backlog_histlen":        "repl_backlog_histlen",
	"mem_fragmentation_ratio":     "mem_fragmentation_ratio",
	"used_cpu_sys":                "used_cpu_sys",
	"used_cpu_user":               "used_cpu_user",
	"used_cpu_sys_children":       "used_cpu_sys_children",
	"used_cpu_user_children":      "used_cpu_user_children",
}

var ErrProtocolError = errors.New("redis protocol error")

// Reads stats from all configured servers accumulates stats.
// Returns one of the errors encountered while gather stats (if any).
func (r *Redis) Gather(acc telegraf.Accumulator) error {
	if len(r.Servers) == 0 {
		url := &url.URL{
			Host: ":6379",
		}
		r.gatherServer(url, acc)
		return nil
	}

	var wg sync.WaitGroup

	var outerr error

	for _, serv := range r.Servers {
		u, err := url.Parse(serv)
		if err != nil {
			return fmt.Errorf("Unable to parse to address '%s': %s", serv, err)
		} else if u.Scheme == "" {
			// fallback to simple string based address (i.e. "10.0.0.1:10000")
			u.Scheme = "tcp"
			u.Host = serv
			u.Path = ""
		}
		wg.Add(1)
		go func(serv string) {
			defer wg.Done()
			outerr = r.gatherServer(u, acc)
		}(serv)
	}

	wg.Wait()

	return outerr
}

const defaultPort = "6379"

func (r *Redis) gatherServer(addr *url.URL, acc telegraf.Accumulator) error {
	_, _, err := net.SplitHostPort(addr.Host)
	if err != nil {
		addr.Host = addr.Host + ":" + defaultPort
	}

	c, err := net.Dial("tcp", addr.Host)
	if err != nil {
		return fmt.Errorf("Unable to connect to redis server '%s': %s", addr.Host, err)
	}
	defer c.Close()

	if addr.User != nil {
		pwd, set := addr.User.Password()
		if set && pwd != "" {
			c.Write([]byte(fmt.Sprintf("AUTH %s\r\n", pwd)))

			rdr := bufio.NewReader(c)

			line, err := rdr.ReadString('\n')
			if err != nil {
				return err
			}
			if line[0] != '+' {
				return fmt.Errorf("%s", strings.TrimSpace(line)[1:])
			}
		}
	}

	c.Write([]byte("INFO\r\n"))
	c.Write([]byte("EOF\r\n"))
	rdr := bufio.NewReader(c)

	// Setup tags for all redis metrics
	host, port := "unknown", "unknown"
	// If there's an error, ignore and use 'unknown' tags
	host, port, _ = net.SplitHostPort(addr.Host)
	tags := map[string]string{"server": host, "port": port}

	return gatherInfoOutput(rdr, acc, tags)
}

// gatherInfoOutput gathers
func gatherInfoOutput(
	rdr *bufio.Reader,
	acc telegraf.Accumulator,
	tags map[string]string,
) error {
	var keyspace_hits, keyspace_misses uint64 = 0, 0

	scanner := bufio.NewScanner(rdr)
	fields := make(map[string]interface{})
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "ERR") {
			break
		}

		if len(line) == 0 || line[0] == '#' {
			continue
		}

		parts := strings.SplitN(line, ":", 2)
		if len(parts) < 2 {
			continue
		}

		name := string(parts[0])
		metric, ok := Tracking[name]
		if !ok {
			kline := strings.TrimSpace(string(parts[1]))
			gatherKeyspaceLine(name, kline, acc, tags)
			continue
		}

		val := strings.TrimSpace(parts[1])
		ival, err := strconv.ParseUint(val, 10, 64)

		if name == "keyspace_hits" {
			keyspace_hits = ival
		}

		if name == "keyspace_misses" {
			keyspace_misses = ival
		}

		if err == nil {
			fields[metric] = ival
			continue
		}

		fval, err := strconv.ParseFloat(val, 64)
		if err != nil {
			return err
		}

		fields[metric] = fval
	}
	var keyspace_hitrate float64 = 0.0
	if keyspace_hits != 0 || keyspace_misses != 0 {
		keyspace_hitrate = float64(keyspace_hits) / float64(keyspace_hits+keyspace_misses)
	}
	fields["keyspace_hitrate"] = keyspace_hitrate
	acc.AddFields("redis", fields, tags)
	return nil
}

// Parse the special Keyspace line at end of redis stats
// This is a special line that looks something like:
//     db0:keys=2,expires=0,avg_ttl=0
// And there is one for each db on the redis instance
func gatherKeyspaceLine(
	name string,
	line string,
	acc telegraf.Accumulator,
	tags map[string]string,
) {
	if strings.Contains(line, "keys=") {
		fields := make(map[string]interface{})
		tags["database"] = name
		dbparts := strings.Split(line, ",")
		for _, dbp := range dbparts {
			kv := strings.Split(dbp, "=")
			ival, err := strconv.ParseUint(kv[1], 10, 64)
			if err == nil {
				fields[kv[0]] = ival
			}
		}
		acc.AddFields("redis_keyspace", fields, tags)
	}
}

func init() {
	inputs.Add("redis", func() telegraf.Input {
		return &Redis{}
	})
}
<commit_msg>Update redis.go<commit_after>package redis

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"strings"
	"sync"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type Redis struct {
	Servers []string
}

var sampleConfig = `
  ## specify servers via a url matching:
  ##  [protocol://][:password]@address[:port]
  ##  e.g.
  ##    tcp://localhost:6379
  ##    tcp://:password@192.168.99.100
  ##
  ## If no servers are specified, then localhost is used as the host.
  ## If no port is specified, 6379 is used
  servers = ["tcp://localhost:6379"]
`

func (r *Redis) SampleConfig() string {
	return sampleConfig
}

func (r *Redis) Description() string {
	return "Read metrics from one or many redis servers"
}

var Tracking = map[string]string{
	"uptime_in_seconds":           "uptime",
	"connected_clients":           "clients",
	"used_memory":                 "used_memory",
	"used_memory_rss":             "used_memory_rss",
	"used_memory_peak":            "used_memory_peak",
	"used_memory_lua":             "used_memory_lua",
	"rdb_changes_since_last_save": "rdb_changes_since_last_save",
	"total_connections_received":  "total_connections_received",
	"total_commands_processed":    "total_commands_processed",
	"instantaneous_ops_per_sec":   "instantaneous_ops_per_sec",
	"instantaneous_input_kbps":    "instantaneous_input_kbps",
	"instantaneous_output_kbps":   "instantaneous_output_kbps",
	"sync_full":                   "sync_full",
	"sync_partial_ok":             "sync_partial_ok",
	"sync_partial_err":            "sync_partial_err",
	"expired_keys":                "expired_keys",
	"evicted_keys":                "evicted_keys",
	"keyspace_hits":               "keyspace_hits",
	"keyspace_misses":             "keyspace_misses",
	"pubsub_channels":             "pubsub_channels",
	"pubsub_patterns":             "pubsub_patterns",
	"latest_fork_usec":            "latest_fork_usec",
	"connected_slaves":            "connected_slaves",
	"master_repl_offset":          "master_repl_offset",
	"repl_backlog_active":         "repl_backlog_active",
	"repl_backlog_size":           "repl_backlog_size",
	"repl_backlog_histlen":        "repl_backlog_histlen",
	"mem_fragmentation_ratio":     "mem_fragmentation_ratio",
	"used_cpu_sys":                "used_cpu_sys",
	"used_cpu_user":               "used_cpu_user",
	"used_cpu_sys_children":       "used_cpu_sys_children",
	"used_cpu_user_children":      "used_cpu_user_children",
}

var ErrProtocolError = errors.New("redis protocol error")

// Reads stats from all configured servers accumulates stats.
// Returns one of the errors encountered while gather stats (if any).
func (r *Redis) Gather(acc telegraf.Accumulator) error {
	if len(r.Servers) == 0 {
		url := &url.URL{
			Host: ":6379",
		}
		r.gatherServer(url, acc)
		return nil
	}

	var wg sync.WaitGroup

	var outerr error

	for _, serv := range r.Servers {
		u, err := url.Parse(serv)
		if err != nil {
			return fmt.Errorf("Unable to parse to address '%s': %s", serv, err)
		} else if u.Scheme == "" {
			// fallback to simple string based address (i.e. "10.0.0.1:10000")
			u.Scheme = "tcp"
			u.Host = serv
			u.Path = ""
		}
		wg.Add(1)
		go func(serv string) {
			defer wg.Done()
			outerr = r.gatherServer(u, acc)
		}(serv)
	}

	wg.Wait()

	return outerr
}

const defaultPort = "6379"

func (r *Redis) gatherServer(addr *url.URL, acc telegraf.Accumulator) error {
	_, _, err := net.SplitHostPort(addr.Host)
	if err != nil {
		addr.Host = addr.Host + ":" + defaultPort
	}

	c, err := net.Dial("tcp", addr.Host)
	if err != nil {
		return fmt.Errorf("Unable to connect to redis server '%s': %s", addr.Host, err)
	}
	defer c.Close()

	if addr.User != nil {
		pwd, set := addr.User.Password()
		if set && pwd != "" {
			c.Write([]byte(fmt.Sprintf("AUTH %s\r\n", pwd)))

			rdr := bufio.NewReader(c)

			line, err := rdr.ReadString('\n')
			if err != nil {
				return err
			}
			if line[0] != '+' {
				return fmt.Errorf("%s", strings.TrimSpace(line)[1:])
			}
		}
	}

	c.Write([]byte("INFO\r\n"))
	c.Write([]byte("EOF\r\n"))
	rdr := bufio.NewReader(c)

	// Setup tags for all redis metrics
	host, port := "unknown", "unknown"
	// If there's an error, ignore and use 'unknown' tags
	host, port, _ = net.SplitHostPort(addr.Host)
	tags := map[string]string{"server": host, "port": port}

	return gatherInfoOutput(rdr, acc, tags)
}

// gatherInfoOutput gathers
func gatherInfoOutput(
	rdr *bufio.Reader,
	acc telegraf.Accumulator,
	tags map[string]string,
) error {
	var keyspace_hits, keyspace_misses uint64 = 0, 0

	scanner := bufio.NewScanner(rdr)
	fields := make(map[string]interface{})
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "ERR") {
			break
		}

		if len(line) == 0 || line[0] == '#' {
			continue
		}

		parts := strings.SplitN(line, ":", 2)
		if len(parts) < 2 {
			continue
		}

		name := string(parts[0])
		metric, ok := Tracking[name]
		if !ok {
			kline := strings.TrimSpace(string(parts[1]))
			gatherKeyspaceLine(name, kline, acc, tags)
			continue
		}

		val := strings.TrimSpace(parts[1])
		ival, err := strconv.ParseUint(val, 10, 64)

		if name == "keyspace_hits" {
			keyspace_hits = ival
		}

		if name == "keyspace_misses" {
			keyspace_misses = ival
		}

		if err == nil {
			fields[metric] = ival
			continue
		}

		fval, err := strconv.ParseFloat(val, 64)
		if err != nil {
			return err
		}

		fields[metric] = fval
	}
	var keyspace_hitrate float64 = 0.0
	if keyspace_hits != 0 || keyspace_misses != 0 {
		keyspace_hitrate = float64(keyspace_hits) / float64(keyspace_hits+keyspace_misses)
	}
	fields["keyspace_hitrate"] = keyspace_hitrate
	acc.AddFields("redis", fields, tags)
	return nil
}

// Parse the special Keyspace line at end of redis stats
// This is a special line that looks something like:
//     db0:keys=2,expires=0,avg_ttl=0
// And there is one for each db on the redis instance
func gatherKeyspaceLine(
	name string,
	line string,
	acc telegraf.Accumulator,
	tags map[string]string,
) {
	if strings.Contains(line, "keys=") {
		fields := make(map[string]interface{})
		tags["database"] = name
		dbparts := strings.Split(line, ",")
		for _, dbp := range dbparts {
			kv := strings.Split(dbp, "=")
			ival, err := strconv.ParseUint(kv[1], 10, 64)
			if err == nil {
				fields[kv[0]] = ival
			}
		}
		acc.AddFields("redis_keyspace", fields, tags)
	}
}

func init() {
	inputs.Add("redis", func() telegraf.Input {
		return &Redis{}
	})
}
<|endoftext|>


<commit_before>// Package quickfix provides functions for fixing Go ASTs
// that are well typed but "go build" refuses to build.
package quickfix

import (
	"fmt"
	"regexp"
	"strings"

	"go/ast"
	"go/token"

	"golang.org/x/tools/go/ast/astutil"
	_ "golang.org/x/tools/go/gcimporter"
	"golang.org/x/tools/go/types"
)

var (
	declaredNotUsed        = regexp.MustCompile(`^([a-zA-Z0-9_]+) declared but not used$`)
	importedNotUsed        = regexp.MustCompile(`^(".+") imported but not used$`)
	noNewVariablesOnDefine = "no new variables on left side of :="
)

type Config struct {
	Fset     *token.FileSet
	Files    []*ast.File
	TypeInfo *types.Info
	MaxTries int
}

func QuickFix(fset *token.FileSet, files []*ast.File) (err error) {
	config := Config{
		Fset:     fset,
		Files:    files,
		MaxTries: 10,
	}
	return config.QuickFix()
}

// QuickFix rewrites AST files of same package so that they pass go build.
// For example:
//   v declared but not used             -> append `_ = v`
//   "p" imported but not used           -> rewrite to `import _ "p"`
//   no new variables on left side of := -> rewrite `:=` to `=`
//
// TODO implement hardMode, which removes errorneous code rather than adding
func (c Config) QuickFix() (err error) {
	maxTries := 10
	if c.MaxTries > 0 {
		maxTries = c.MaxTries
	}
	for i := 0; i < maxTries; i++ {
		var foundError bool
		foundError, err = c.QuickFixOnce()
		if !foundError {
			return nil
		}
	}

	return
}

type tracedVisitor struct {
	path  []ast.Node
	visit func(ast.Node, []ast.Node) bool
}

func (v tracedVisitor) Visit(node ast.Node) ast.Visitor {
	if v.visit(node, v.path) {
		return tracedVisitor{
			path:  append([]ast.Node{node}, v.path...),
			visit: v.visit,
		}
	}

	return nil
}

func traverseAST(node ast.Node, visit func(ast.Node, []ast.Node) bool) {
	v := tracedVisitor{
		visit: visit,
	}
	ast.Walk(v, node)
}

// pkgsWithSideEffect are set of packages which are known to provide APIs by
// blank identifier import (import _ "p").
var pkgsWithSideEffect = map[string]bool{}

func init() {
	for _, path := range []string{
		"expvar",
		"image/gif",
		"image/jpeg",
		"image/png",
		"net/http/pprof",
		"unsafe",
		"golang.org/x/image/bmp",
		"golang.org/x/image/tiff",
		"golang.org/x/image/vp8",
		"golang.org/x/image/vp81",
		"golang.org/x/image/webp",
		"golang.org/x/tools/go/gcimporter",
	} {
		pkgsWithSideEffect[`"`+path+`"`] = true
	}
}

func RevertQuickFix(fset *token.FileSet, files []*ast.File) error {
	config := Config{
		Fset:     fset,
		Files:    files,
		MaxTries: 10,
	}
	return config.RevertQuickFix()
}

// RevertQuickFix reverts possible quickfixes introduced by QuickFix.
// This may result to non-buildable source, and cannot reproduce the original
// code before prior QuickFix.
// For example:
//   `_ = v`        -> removed
//   `import _ "p"` -> rewritten to `import "p"`
func (c Config) RevertQuickFix() (err error) {
	fset := c.Fset
	files := c.Files

	nodeToRemove := map[ast.Node]bool{}

	for _, f := range files {
		ast.Inspect(f, func(node ast.Node) bool {
			if assign, ok := node.(*ast.AssignStmt); ok {
				if len(assign.Lhs) == 1 && isBlankIdent(assign.Lhs[0]) &&
					len(assign.Rhs) == 1 && isIdent(assign.Rhs[0]) {
					// The statement is `_ = v`
					nodeToRemove[node] = true
				}

				return false
			} else if imp, ok := node.(*ast.ImportSpec); ok {
				if isBlankIdent(imp.Name) && !pkgsWithSideEffect[imp.Path.Value] {
					// The spec is `import _ "p"` and p is not a package that
					// provides "side effects"
					imp.Name = nil
				}

				return false
			}

			return true
		})

		for len(nodeToRemove) > 0 {
			traverseAST(f, func(node ast.Node, nodepath []ast.Node) bool {
				if nodeToRemove[node] {
					parent := nodepath[0]
					if removeChildNode(node, parent) == false {
						err = fmt.Errorf(
							"BUG: could not remove node: %s (in: %s)",
							fset.Position(node.Pos()),
							fset.Position(parent.Pos()),
						)
					}
					delete(nodeToRemove, node)
					return false
				}

				return true
			})
		}
	}

	return
}

func (c Config) QuickFixOnce() (bool, error) {
	fset := c.Fset
	files := c.Files

	errs := []error{}
	config := &types.Config{
		Error: func(err error) {
			errs = append(errs, err)
		},
	}

	_, err := config.Check("_quickfix", fset, files, c.TypeInfo)
	if err == nil {
		return false, nil
	}

	// apply fixes on AST later so that we won't break funcs that inspect AST by positions
	fixes := map[error]func() bool{}
	unhandled := ErrorList{}

	foundError := len(errs) > 0

	for _, err := range errs {
		err, ok := err.(types.Error)
		if !ok {
			unhandled = append(unhandled, err)
			continue
		}

		f := findFile(c.Files, err.Pos)
		if f == nil {
			e := ErrCouldNotLocate{
				Err:  err,
				Fset: fset,
			}
			unhandled = append(unhandled, e)
			continue
		}

		nodepath, _ := astutil.PathEnclosingInterval(f, err.Pos, err.Pos)

		var fix func() bool

		// - "%s declared but not used"
		// - "%q imported but not used"
		// - "label %s declared but not used" TODO
		// - "no new variables on left side of :="
		if m := declaredNotUsed.FindStringSubmatch(err.Msg); m != nil {
			identName := m[1]
			fix = func() bool {
				return fixDeclaredNotUsed(nodepath, identName)
			}
		} else if m := importedNotUsed.FindStringSubmatch(err.Msg); m != nil {
			pkgPath := m[1] // quoted string, but it's okay because this will be compared to ast.BasicLit.Value.
			fix = func() bool {
				return fixImportedNotUsed(nodepath, pkgPath)
			}
		} else if err.Msg == noNewVariablesOnDefine {
			fix = func() bool {
				return fixNoNewVariables(nodepath)
			}
		} else {
			unhandled = append(unhandled, err)
		}

		if fix != nil {
			fixes[err] = fix
		}
	}

	for err, fix := range fixes {
		if fix() == false {
			unhandled = append(unhandled, err)
		}
	}

	return foundError, unhandled.any()
}

func fixDeclaredNotUsed(nodepath []ast.Node, identName string) bool {
	// insert "_ = x" to supress "declared but not used" error
	stmt := &ast.AssignStmt{
		Lhs: []ast.Expr{ast.NewIdent("_")},
		Tok: token.ASSIGN,
		Rhs: []ast.Expr{ast.NewIdent(identName)},
	}
	return appendStmt(nodepath, stmt)
}

func fixImportedNotUsed(nodepath []ast.Node, pkgPath string) bool {
	for _, node := range nodepath {
		if f, ok := node.(*ast.File); ok {
			for _, imp := range f.Imports {
				if imp.Path.Value == pkgPath {
					// make this import spec anonymous one
					imp.Name = ast.NewIdent("_")
					return true
				}
			}
		}
	}
	return false
}

func fixNoNewVariables(nodepath []ast.Node) bool {
	for _, node := range nodepath {
		switch node := node.(type) {
		case *ast.AssignStmt:
			if node.Tok == token.DEFINE {
				node.Tok = token.ASSIGN
				return true
			}

		case *ast.RangeStmt:
			if node.Tok == token.DEFINE {
				node.Tok = token.ASSIGN
				return true
			}
		}
	}
	return false
}

type ErrorList []error

func (errs ErrorList) any() error {
	if len(errs) == 0 {
		return nil
	}

	return errs
}

func (errs ErrorList) Error() string {
	s := []string{fmt.Sprintf("%d error(s):", len(errs))}
	for _, e := range errs {
		s = append(s, fmt.Sprintf("- %s", e))
	}
	return strings.Join(s, "\n")
}

func appendStmt(nodepath []ast.Node, stmt ast.Stmt) bool {
	for _, node := range nodepath {
		switch node := node.(type) {
		case *ast.BlockStmt:
			if node.List == nil {
				node.List = []ast.Stmt{}
			}
			node.List = append(node.List, stmt)

		case *ast.CaseClause:
			if node.Body == nil {
				node.Body = []ast.Stmt{}
			}
			node.Body = append(node.Body, stmt)

		case *ast.CommClause:
			if node.Body == nil {
				node.Body = []ast.Stmt{}
			}
			node.Body = append(node.Body, stmt)

		case *ast.RangeStmt:
			if node.Body == nil {
				node.Body = &ast.BlockStmt{}
			}
			if node.Body.List == nil {
				node.Body.List = []ast.Stmt{}
			}
			node.Body.List = append(node.Body.List, stmt)

		default:
			continue
		}

		return true
	}

	return false
}

func removeChildNode(child, parent ast.Node) bool {
	switch parent := parent.(type) {
	case *ast.BlockStmt:
		removeFromStmtList(child, parent.List)
		return true
	case *ast.CaseClause:
		removeFromStmtList(child, parent.Body)
		return true
	case *ast.CommClause:
		removeFromStmtList(child, parent.Body)
		return true
	case *ast.RangeStmt:
		removeFromStmtList(child, parent.Body.List)
		return true
	}

	return false
}

// removeFromStmtList remove node from slice of statements list. This function
// modifies list in-place and pads rest of the slice with ast.EmptyStmt.
func removeFromStmtList(node ast.Node, list []ast.Stmt) bool {
	for i, s := range list {
		if s == node {
			for ; i < len(list)-1; i++ {
				list[i] = list[i+1]
			}
			list[len(list)-1] = &ast.EmptyStmt{}
			return true
		}
	}

	return false
}

func findFile(files []*ast.File, pos token.Pos) *ast.File {
	for _, f := range files {
		if f.Pos() <= pos && pos < f.End() {
			return f
		}
	}

	return nil
}

func isIdent(node ast.Node) bool {
	if node == nil {
		return false
	}

	_, ok := node.(*ast.Ident)
	return ok
}

func isBlankIdent(node ast.Node) bool {
	if node == nil {
		return false
	}

	ident, ok := node.(*ast.Ident)
	return ok && ident != nil && ident.Name == "_"
}

type ErrCouldNotLocate struct {
	Err  types.Error
	Fset *token.FileSet
}

func (e ErrCouldNotLocate) Error() string {
	return fmt.Sprintf("cannot find file for error %q: %s (%d)", e.Err.Error(), e.Fset.Position(e.Err.Pos), e.Err.Pos)
}
<commit_msg>support 'imported but not used as'<commit_after>// Package quickfix provides functions for fixing Go ASTs
// that are well typed but "go build" refuses to build.
package quickfix

import (
	"fmt"
	"regexp"
	"strings"

	"go/ast"
	"go/token"

	"golang.org/x/tools/go/ast/astutil"
	_ "golang.org/x/tools/go/gcimporter"
	"golang.org/x/tools/go/types"
)

var (
	declaredNotUsed        = regexp.MustCompile(`^([a-zA-Z0-9_]+) declared but not used$`)
	importedNotUsed        = regexp.MustCompile(`^(".+") imported but not used`)
	noNewVariablesOnDefine = "no new variables on left side of :="
)

type Config struct {
	Fset     *token.FileSet
	Files    []*ast.File
	TypeInfo *types.Info
	MaxTries int
}

func QuickFix(fset *token.FileSet, files []*ast.File) (err error) {
	config := Config{
		Fset:     fset,
		Files:    files,
		MaxTries: 10,
	}
	return config.QuickFix()
}

// QuickFix rewrites AST files of same package so that they pass go build.
// For example:
//   v declared but not used             -> append `_ = v`
//   "p" imported but not used           -> rewrite to `import _ "p"`
//   no new variables on left side of := -> rewrite `:=` to `=`
//
// TODO implement hardMode, which removes errorneous code rather than adding
func (c Config) QuickFix() (err error) {
	maxTries := 10
	if c.MaxTries > 0 {
		maxTries = c.MaxTries
	}
	for i := 0; i < maxTries; i++ {
		var foundError bool
		foundError, err = c.QuickFixOnce()
		if !foundError {
			return nil
		}
	}

	return
}

type tracedVisitor struct {
	path  []ast.Node
	visit func(ast.Node, []ast.Node) bool
}

func (v tracedVisitor) Visit(node ast.Node) ast.Visitor {
	if v.visit(node, v.path) {
		return tracedVisitor{
			path:  append([]ast.Node{node}, v.path...),
			visit: v.visit,
		}
	}

	return nil
}

func traverseAST(node ast.Node, visit func(ast.Node, []ast.Node) bool) {
	v := tracedVisitor{
		visit: visit,
	}
	ast.Walk(v, node)
}

// pkgsWithSideEffect are set of packages which are known to provide APIs by
// blank identifier import (import _ "p").
var pkgsWithSideEffect = map[string]bool{}

func init() {
	for _, path := range []string{
		"expvar",
		"image/gif",
		"image/jpeg",
		"image/png",
		"net/http/pprof",
		"unsafe",
		"golang.org/x/image/bmp",
		"golang.org/x/image/tiff",
		"golang.org/x/image/vp8",
		"golang.org/x/image/vp81",
		"golang.org/x/image/webp",
		"golang.org/x/tools/go/gcimporter",
	} {
		pkgsWithSideEffect[`"`+path+`"`] = true
	}
}

func RevertQuickFix(fset *token.FileSet, files []*ast.File) error {
	config := Config{
		Fset:     fset,
		Files:    files,
		MaxTries: 10,
	}
	return config.RevertQuickFix()
}

// RevertQuickFix reverts possible quickfixes introduced by QuickFix.
// This may result to non-buildable source, and cannot reproduce the original
// code before prior QuickFix.
// For example:
//   `_ = v`        -> removed
//   `import _ "p"` -> rewritten to `import "p"`
func (c Config) RevertQuickFix() (err error) {
	fset := c.Fset
	files := c.Files

	nodeToRemove := map[ast.Node]bool{}

	for _, f := range files {
		ast.Inspect(f, func(node ast.Node) bool {
			if assign, ok := node.(*ast.AssignStmt); ok {
				if len(assign.Lhs) == 1 && isBlankIdent(assign.Lhs[0]) &&
					len(assign.Rhs) == 1 && isIdent(assign.Rhs[0]) {
					// The statement is `_ = v`
					nodeToRemove[node] = true
				}

				return false
			} else if imp, ok := node.(*ast.ImportSpec); ok {
				if isBlankIdent(imp.Name) && !pkgsWithSideEffect[imp.Path.Value] {
					// The spec is `import _ "p"` and p is not a package that
					// provides "side effects"
					imp.Name = nil
				}

				return false
			}

			return true
		})

		for len(nodeToRemove) > 0 {
			traverseAST(f, func(node ast.Node, nodepath []ast.Node) bool {
				if nodeToRemove[node] {
					parent := nodepath[0]
					if removeChildNode(node, parent) == false {
						err = fmt.Errorf(
							"BUG: could not remove node: %s (in: %s)",
							fset.Position(node.Pos()),
							fset.Position(parent.Pos()),
						)
					}
					delete(nodeToRemove, node)
					return false
				}

				return true
			})
		}
	}

	return
}

func (c Config) QuickFixOnce() (bool, error) {
	fset := c.Fset
	files := c.Files

	errs := []error{}
	config := &types.Config{
		Error: func(err error) {
			errs = append(errs, err)
		},
	}

	_, err := config.Check("_quickfix", fset, files, c.TypeInfo)
	if err == nil {
		return false, nil
	}

	// apply fixes on AST later so that we won't break funcs that inspect AST by positions
	fixes := map[error]func() bool{}
	unhandled := ErrorList{}

	foundError := len(errs) > 0

	for _, err := range errs {
		err, ok := err.(types.Error)
		if !ok {
			unhandled = append(unhandled, err)
			continue
		}

		f := findFile(c.Files, err.Pos)
		if f == nil {
			e := ErrCouldNotLocate{
				Err:  err,
				Fset: fset,
			}
			unhandled = append(unhandled, e)
			continue
		}

		nodepath, _ := astutil.PathEnclosingInterval(f, err.Pos, err.Pos)

		var fix func() bool

		// - "%s declared but not used"
		// - "%q imported but not used" (+ " as %s")
		// - "label %s declared but not used" TODO
		// - "no new variables on left side of :="
		if m := declaredNotUsed.FindStringSubmatch(err.Msg); m != nil {
			identName := m[1]
			fix = func() bool {
				return fixDeclaredNotUsed(nodepath, identName)
			}
		} else if m := importedNotUsed.FindStringSubmatch(err.Msg); m != nil {
			pkgPath := m[1] // quoted string, but it's okay because this will be compared to ast.BasicLit.Value.
			fix = func() bool {
				return fixImportedNotUsed(nodepath, pkgPath)
			}
		} else if err.Msg == noNewVariablesOnDefine {
			fix = func() bool {
				return fixNoNewVariables(nodepath)
			}
		} else {
			unhandled = append(unhandled, err)
		}

		if fix != nil {
			fixes[err] = fix
		}
	}

	for err, fix := range fixes {
		if fix() == false {
			unhandled = append(unhandled, err)
		}
	}

	return foundError, unhandled.any()
}

func fixDeclaredNotUsed(nodepath []ast.Node, identName string) bool {
	// insert "_ = x" to supress "declared but not used" error
	stmt := &ast.AssignStmt{
		Lhs: []ast.Expr{ast.NewIdent("_")},
		Tok: token.ASSIGN,
		Rhs: []ast.Expr{ast.NewIdent(identName)},
	}
	return appendStmt(nodepath, stmt)
}

func fixImportedNotUsed(nodepath []ast.Node, pkgPath string) bool {
	for _, node := range nodepath {
		if f, ok := node.(*ast.File); ok {
			for _, imp := range f.Imports {
				if imp.Path.Value == pkgPath {
					// make this import spec anonymous one
					imp.Name = ast.NewIdent("_")
					return true
				}
			}
		}
	}
	return false
}

func fixNoNewVariables(nodepath []ast.Node) bool {
	for _, node := range nodepath {
		switch node := node.(type) {
		case *ast.AssignStmt:
			if node.Tok == token.DEFINE {
				node.Tok = token.ASSIGN
				return true
			}

		case *ast.RangeStmt:
			if node.Tok == token.DEFINE {
				node.Tok = token.ASSIGN
				return true
			}
		}
	}
	return false
}

type ErrorList []error

func (errs ErrorList) any() error {
	if len(errs) == 0 {
		return nil
	}

	return errs
}

func (errs ErrorList) Error() string {
	s := []string{fmt.Sprintf("%d error(s):", len(errs))}
	for _, e := range errs {
		s = append(s, fmt.Sprintf("- %s", e))
	}
	return strings.Join(s, "\n")
}

func appendStmt(nodepath []ast.Node, stmt ast.Stmt) bool {
	for _, node := range nodepath {
		switch node := node.(type) {
		case *ast.BlockStmt:
			if node.List == nil {
				node.List = []ast.Stmt{}
			}
			node.List = append(node.List, stmt)

		case *ast.CaseClause:
			if node.Body == nil {
				node.Body = []ast.Stmt{}
			}
			node.Body = append(node.Body, stmt)

		case *ast.CommClause:
			if node.Body == nil {
				node.Body = []ast.Stmt{}
			}
			node.Body = append(node.Body, stmt)

		case *ast.RangeStmt:
			if node.Body == nil {
				node.Body = &ast.BlockStmt{}
			}
			if node.Body.List == nil {
				node.Body.List = []ast.Stmt{}
			}
			node.Body.List = append(node.Body.List, stmt)

		default:
			continue
		}

		return true
	}

	return false
}

func removeChildNode(child, parent ast.Node) bool {
	switch parent := parent.(type) {
	case *ast.BlockStmt:
		removeFromStmtList(child, parent.List)
		return true
	case *ast.CaseClause:
		removeFromStmtList(child, parent.Body)
		return true
	case *ast.CommClause:
		removeFromStmtList(child, parent.Body)
		return true
	case *ast.RangeStmt:
		removeFromStmtList(child, parent.Body.List)
		return true
	}

	return false
}

// removeFromStmtList remove node from slice of statements list. This function
// modifies list in-place and pads rest of the slice with ast.EmptyStmt.
func removeFromStmtList(node ast.Node, list []ast.Stmt) bool {
	for i, s := range list {
		if s == node {
			for ; i < len(list)-1; i++ {
				list[i] = list[i+1]
			}
			list[len(list)-1] = &ast.EmptyStmt{}
			return true
		}
	}

	return false
}

func findFile(files []*ast.File, pos token.Pos) *ast.File {
	for _, f := range files {
		if f.Pos() <= pos && pos < f.End() {
			return f
		}
	}

	return nil
}

func isIdent(node ast.Node) bool {
	if node == nil {
		return false
	}

	_, ok := node.(*ast.Ident)
	return ok
}

func isBlankIdent(node ast.Node) bool {
	if node == nil {
		return false
	}

	ident, ok := node.(*ast.Ident)
	return ok && ident != nil && ident.Name == "_"
}

type ErrCouldNotLocate struct {
	Err  types.Error
	Fset *token.FileSet
}

func (e ErrCouldNotLocate) Error() string {
	return fmt.Sprintf("cannot find file for error %q: %s (%d)", e.Err.Error(), e.Fset.Position(e.Err.Pos), e.Err.Pos)
}
<|endoftext|>


<commit_before>package limit

import (
	"context"
	"fmt"
	"net/http"
	"psychic-rat/log"
	"sync"
	"time"
)

type (
	Limiter struct {
		limits      map[string]*bucket
		mu          sync.Mutex
		max         int
		increment   int
		interval    int
		idGenerator IdGeneratorFunc
	}

	IdGeneratorFunc func(*http.Request) string

	bucket struct {
		tokens int
		mu     sync.Mutex
	}
)

func New(max, increment, interval int, idGen IdGeneratorFunc) *Limiter {
	return &Limiter{limits: map[string]*bucket{}, max: max, increment: increment, interval: interval, idGenerator: idGen}
}

func (l *Limiter) CheckLimit(r *http.Request) error {
	id := l.idGenerator(r)
	bh := l.getBucketFor(id)
	if !bh.consumeToken() {
		return fmt.Errorf("tokens exhausted for %s", id)
	}
	return nil
}

func (l *Limiter) getBucketFor(s string) *bucket {
	l.mu.Lock()
	defer l.mu.Unlock()

	b, ok := l.limits[s]
	if !ok {
		b = &bucket{tokens: l.max}
		l.limits[s] = b
		go b.tokenFiller(l.interval, l.increment, l.max, func() { l.deleteBucket(s) })
		log.Logf(context.Background(), "bucket filler started for %s", s)
	}
	return b
}

func (l *Limiter) deleteBucket(s string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.limits, s)
}

func (b *bucket) consumeToken() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.tokens == 0 {
		return false
	}
	b.tokens--
	return true
}

func (b *bucket) tokenFiller(interval, increment, max int, done func()) {
	for {
		time.Sleep(time.Duration(interval) * time.Second)
		b.tokens += increment
		if b.tokens >= max {
			done()
			return
		}
	}
}
<commit_msg>changed tokenFiller<commit_after>package limit

import (
	"context"
	"fmt"
	"net/http"
	"psychic-rat/log"
	"sync"
	"time"
)

type (
	Limiter struct {
		limits      map[string]*bucket
		mu          sync.Mutex
		max         int
		increment   int
		interval    int
		idGenerator IdGeneratorFunc
	}

	IdGeneratorFunc func(*http.Request) string

	bucket struct {
		tokens int
		mu     sync.Mutex
	}
)

func New(max, increment, interval int, idGen IdGeneratorFunc) *Limiter {
	return &Limiter{limits: map[string]*bucket{}, max: max, increment: increment, interval: interval, idGenerator: idGen}
}

func (l *Limiter) CheckLimit(r *http.Request) error {
	id := l.idGenerator(r)
	bh := l.getBucketFor(id)
	if !bh.consumeToken() {
		return fmt.Errorf("tokens exhausted for %s", id)
	}
	return nil
}

func (l *Limiter) getBucketFor(s string) *bucket {
	l.mu.Lock()
	defer l.mu.Unlock()

	b, ok := l.limits[s]
	if !ok {
		b = &bucket{tokens: l.max}
		l.limits[s] = b
		go func() {
			b.tokenFiller(l.interval, l.increment, l.max)
			l.deleteBucket(s)
		}()
		log.Logf(context.Background(), "bucket filler started for %s", s)
	}
	return b
}

func (l *Limiter) deleteBucket(s string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.limits, s)
}

func (b *bucket) consumeToken() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.tokens == 0 {
		return false
	}
	b.tokens--
	return true
}

func (b *bucket) tokenFiller(interval, increment, max int) {
	for {
		time.Sleep(time.Duration(interval) * time.Second)
		b.tokens += increment
		if b.tokens >= max {
			return
		}
	}
}
<|endoftext|>


<commit_before>package dfs

import (
	"github.com/zenoss/glog"
	"github.com/zenoss/serviced/dao"
	"github.com/zenoss/serviced/volume"

	docker "github.com/fsouza/go-dockerclient"

	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os/exec"
	"os/user"
	"path"
	"strings"
	"sync"
	"time"
)

const (
	USER_ROOT        string = "root"
	DOCKER_ENDPOINT  string = "unix:///var/run/docker.sock"
	DOCKER_LATEST    string = "latest"
	DOCKER_IMAGEJSON string = "images.json"

	ERR_STATENOTFOUND string = "Service state not found for serviceId: %s"
)

var (
	unused interface{}
	Lock   *sync.Mutex = new(sync.Mutex)
	// stubs
	getCurrentUser = user.Current
)

var runServiceCommand = func(state *dao.ServiceState, command string) ([]byte, error) {
	lxcAttach, err := exec.LookPath("lxc-attach")
	if err != nil {
		return []byte{}, err
	}
	cmd := exec.Command(lxcAttach, "-n", state.DockerId, "-e", "--", "bin/bash", "-c", command)

	output, err := cmd.CombinedOutput()
	if err != nil {
		glog.Errorf("Error running command: `%s` for serviceId: %s out: %s err: %s", command, state.ServiceId, output, err)
		return output, err
	}
	glog.V(0).Infof("Successfully ran command: `%s` for serviceId: %s out: %s", command, state.ServiceId, output)
	return output, nil
}

type DistributedFileSystem struct {
	client       dao.ControlPlane
	dockerClient *docker.Client
}

func NewDistributedFileSystem(client dao.ControlPlane) (*DistributedFileSystem, error) {
	dockerClient, err := docker.NewClient(DOCKER_ENDPOINT)
	if err != nil {
		glog.V(2).Infof("snapshot.NewDockerClient client=%+v err=%s", client, err)
		return nil, err
	}

	return &DistributedFileSystem{
		client:       client,
		dockerClient: dockerClient,
	}, nil
}

func (d *DistributedFileSystem) getServiceState(serviceId string, state *dao.ServiceState) error {
	var states []*dao.ServiceState
	if err := d.client.GetServiceStates(serviceId, &states); err != nil {
		return err
	}
	for i, s := range states {
		glog.V(3).Infof("DEBUG states[%d]: service:%+v state:%+v", i, serviceId, s)
		if s.DockerId != "" {
			*state = *s
			return nil
		}
	}
	return errors.New(fmt.Sprintf(ERR_STATENOTFOUND, serviceId))
}

func (d *DistributedFileSystem) Pause(service *dao.Service, state *dao.ServiceState) error {
	if output, err := runServiceCommand(state, service.Snapshot.Pause); err != nil {
		errmsg := fmt.Sprintf("output: %s, err: %s", output, err)
		glog.V(2).Infof("DistributedFileSystem.Pause service=%+v err=%s", service, err)
		return errors.New(errmsg)
	}
	return nil
}

func (d *DistributedFileSystem) Resume(service *dao.Service, state *dao.ServiceState) error {
	if output, err := runServiceCommand(state, service.Snapshot.Resume); err != nil {
		errmsg := fmt.Sprintf("output: %s, err: %s", output, err)
		glog.V(2).Infof("DistributedFileSystem.Resume service=%+v err=%s", service, err)
		return errors.New(errmsg)
	}
	return nil
}

func (d *DistributedFileSystem) Snapshot(serviceId string, label *string) error {
	var tenantId string
	if err := d.client.GetTenantId(serviceId, &tenantId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	}

	var service dao.Service
	if err := d.client.GetService(tenantId, &service); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	}

	// simplest case - do everything here

	// call quiesce for services with 'DistributedFileSystem.Pause' and
	// 'DistributedFileSystem.Resume' definition.  Only root can run
	// lxc-attach
	if whoami, err := getCurrentUser(); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	} else if USER_ROOT != whoami.Username {
		glog.Warningf("Unable to pause/resume service - User is not %s - whoami:%+v", USER_ROOT, whoami)
	} else {
		var servicesList []*dao.Service
		if err := d.client.GetServices(unused, &servicesList); err != nil {
			glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
			return err
		}
		for _, service := range servicesList {
			if service.Snapshot.Pause == "" || service.Snapshot.Resume == "" {
				continue
			}

			var state dao.ServiceState
			if err := d.getServiceState(service.Id, &state); err != nil {
				glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
				return err
			}

			err := d.Pause(service, &state)
			defer d.Resume(service, &state)
			if err != nil {
				glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
				return err
			}
		}
	}

	// create a snapshot
	var (
		theVolume volume.Volume
	)
	if err := d.client.GetVolume(tenantId, &theVolume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	} else {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v volume=%+v", serviceId, theVolume)
		*label = getSnapshotLabel(&theVolume)
		if err := theVolume.Snapshot(*label); err != nil {
			glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
			return err
		}
	}

	glog.V(2).Infof("Successfully created snapshot for service Id:%s Name:%s Label:%s", service.Id, service.Name, *label)
	return nil
}

func (d *DistributedFileSystem) Commit(dockerId string, label *string) error {
	Lock.Lock()
	defer Lock.Unlock()

	// Get the container
	container, err := d.dockerClient.InspectContainer(dockerId)
	if err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	} else if container.State.Running {
		err := errors.New("cannot commit a running container")
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}

	// Get the service id
	name := strings.Split(container.Name, "_")
	serviceId := name[0]

	// Get tag & repo information
	var (
		images []docker.APIImages
		image  *docker.APIImages
	)

	if err := d.getLatestImages(&images); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}
	for _, i := range images {
		if i.ID == container.Image {
			image = &i
		}
	}
	if image == nil {
		err := errors.New("cannot commit a stale container")
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}

	// Snapshot the DFS
	if err := d.Snapshot(serviceId, label); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}

	// Commit the container to the image
	newImage, err := d.dockerClient.CommitContainer(docker.CommitContainerOptions{
		Container:  container.ID,
		Repository: image.Repository,
		Tag:        image.Tag,
	})
	if err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	// Copy images to the DFS
	*image = docker.APIImages{
		ID:          newImage.ID,
		RepoTags:    image.RepoTags,
		Created:     newImage.Created.Unix(),
		Size:        newImage.Size,
		VirtualSize: image.VirtualSize,
		ParentId:    newImage.Parent,
		Repository:  image.Repository,
		Tag:         image.Tag,
	}

	// Get the path to the volume and write the images
	var (
		tenantId string
		volume   volume.Volume
	)

	if err := d.client.GetTenantId(serviceId, &tenantId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	if err := d.client.GetVolume(tenantId, &volume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	config := path.Join(volume.Path(), *label, DOCKER_IMAGEJSON)
	if data, err := json.Marshal(images); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	} else if err := ioutil.WriteFile(config, data, 0644); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	return nil
}

func (d *DistributedFileSystem) getLatestImages(images *[]docker.APIImages) error {
	// Get all of the images from docker and find the ones tagged as the latest
	if allImages, err := d.dockerClient.ListImages(true); err != nil {
		return err
	} else {
		for _, image := range allImages {
			if image.Tag == DOCKER_LATEST {
				*images = append(*images, image)
			}
		}
	}

	return nil
}

// Gets the images from the snapshot; returns error if image does not exist in docker
func (d *DistributedFileSystem) getSnapshotImages(snapshotId string, volume *volume.Volume, images *[]docker.APIImages) error {
	config := path.Join(volume.Path(), snapshotId, DOCKER_IMAGEJSON)

	if data, err := ioutil.ReadFile(config); err != nil {
		return err
	} else if err := json.Unmarshal(data, images); err != nil {
		return err
	}

	// Check if the images still exist
	for _, image := range *images {
		if _, err := d.dockerClient.InspectImage(image.ID); err != nil {
			return err
		}
	}

	return nil
}

// Retags containers with the given snapshot
func (d *DistributedFileSystem) retag(images *[]docker.APIImages, volume *volume.Volume) error {

	// Set the tag of the new image
	dockerBin, err := exec.LookPath("docker")
	if err != nil {
		return err
	}

	for _, image := range *images {
		// docker tag %{image.ID} %{image.Repo}:%{image.Tag}
		cmd := exec.Command(dockerBin, "tag", image.ID, fmt.Sprintf("%s:%s", image.Repository, image.Tag))
		if err := cmd.Run(); err != nil {
			return err
		}
	}

	return nil
}

func (d *DistributedFileSystem) Rollback(snapshotId string) error {
	Lock.Lock()
	defer Lock.Unlock()

	var (
		tenantId string
		volume   volume.Volume
	)

	parts := strings.Split(snapshotId, "_")
	if len(parts) != 2 {
		glog.V(2).Infof("DistributedFileSystem.Rollback malformed snapshot Id: %s", snapshotId)
		return errors.New("malformed snapshotId")
	}
	serviceId := parts[0]
	label := parts[1]
	if err := d.client.GetTenantId(serviceId, &tenantId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	// Validate existance of images for this snapshot
	var service dao.Service
	err := d.client.GetService(tenantId, &service)
	glog.V(2).Infof("Getting service instance: %s", tenantId)
	if err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	if err := d.client.GetVolume(tenantId, &volume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	var latestImages []docker.APIImages
	if err := d.getLatestImages(&latestImages); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	var snapshotImages []docker.APIImages
	if err := d.getSnapshotImages(snapshotId, &volume, &snapshotImages); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	d.client.StopService(tenantId, nil)
	// TODO: Wait for real event that confirms shutdown
	time.Sleep(time.Second * 5) // wait for shutdown

	if err := d.retag(&snapshotImages, &volume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		d.retag(&latestImages, &volume)
		return err
	}

	glog.V(2).Infof("performing rollback on %s to %s", tenantId, label)
	if err := volume.Rollback(snapshotId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		d.retag(&latestImages, &volume)
		return err
	}
	var unusedStr string = ""

	return d.client.StartService(tenantId, &unusedStr)
}

func getSnapshotLabel(v *volume.Volume) string {
	format := "20060102-150405"
	loc := time.Now()
	utc := loc.UTC()
	return fmt.Sprintf("%s_%s", v.Name(), utc.Format(format))
}<commit_msg>break when image is found<commit_after>package dfs

import (
	"github.com/zenoss/glog"
	"github.com/zenoss/serviced/dao"
	"github.com/zenoss/serviced/volume"

	docker "github.com/fsouza/go-dockerclient"

	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os/exec"
	"os/user"
	"path"
	"strings"
	"sync"
	"time"
)

const (
	USER_ROOT        string = "root"
	DOCKER_ENDPOINT  string = "unix:///var/run/docker.sock"
	DOCKER_LATEST    string = "latest"
	DOCKER_IMAGEJSON string = "images.json"

	ERR_STATENOTFOUND string = "Service state not found for serviceId: %s"
)

var (
	unused interface{}
	Lock   *sync.Mutex = new(sync.Mutex)
	// stubs
	getCurrentUser = user.Current
)

var runServiceCommand = func(state *dao.ServiceState, command string) ([]byte, error) {
	lxcAttach, err := exec.LookPath("lxc-attach")
	if err != nil {
		return []byte{}, err
	}
	cmd := exec.Command(lxcAttach, "-n", state.DockerId, "-e", "--", "bin/bash", "-c", command)

	output, err := cmd.CombinedOutput()
	if err != nil {
		glog.Errorf("Error running command: `%s` for serviceId: %s out: %s err: %s", command, state.ServiceId, output, err)
		return output, err
	}
	glog.V(0).Infof("Successfully ran command: `%s` for serviceId: %s out: %s", command, state.ServiceId, output)
	return output, nil
}

type DistributedFileSystem struct {
	client       dao.ControlPlane
	dockerClient *docker.Client
}

func NewDistributedFileSystem(client dao.ControlPlane) (*DistributedFileSystem, error) {
	dockerClient, err := docker.NewClient(DOCKER_ENDPOINT)
	if err != nil {
		glog.V(2).Infof("snapshot.NewDockerClient client=%+v err=%s", client, err)
		return nil, err
	}

	return &DistributedFileSystem{
		client:       client,
		dockerClient: dockerClient,
	}, nil
}

func (d *DistributedFileSystem) getServiceState(serviceId string, state *dao.ServiceState) error {
	var states []*dao.ServiceState
	if err := d.client.GetServiceStates(serviceId, &states); err != nil {
		return err
	}
	for i, s := range states {
		glog.V(3).Infof("DEBUG states[%d]: service:%+v state:%+v", i, serviceId, s)
		if s.DockerId != "" {
			*state = *s
			return nil
		}
	}
	return errors.New(fmt.Sprintf(ERR_STATENOTFOUND, serviceId))
}

func (d *DistributedFileSystem) Pause(service *dao.Service, state *dao.ServiceState) error {
	if output, err := runServiceCommand(state, service.Snapshot.Pause); err != nil {
		errmsg := fmt.Sprintf("output: %s, err: %s", output, err)
		glog.V(2).Infof("DistributedFileSystem.Pause service=%+v err=%s", service, err)
		return errors.New(errmsg)
	}
	return nil
}

func (d *DistributedFileSystem) Resume(service *dao.Service, state *dao.ServiceState) error {
	if output, err := runServiceCommand(state, service.Snapshot.Resume); err != nil {
		errmsg := fmt.Sprintf("output: %s, err: %s", output, err)
		glog.V(2).Infof("DistributedFileSystem.Resume service=%+v err=%s", service, err)
		return errors.New(errmsg)
	}
	return nil
}

func (d *DistributedFileSystem) Snapshot(serviceId string, label *string) error {
	var tenantId string
	if err := d.client.GetTenantId(serviceId, &tenantId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	}

	var service dao.Service
	if err := d.client.GetService(tenantId, &service); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	}

	// simplest case - do everything here

	// call quiesce for services with 'DistributedFileSystem.Pause' and
	// 'DistributedFileSystem.Resume' definition.  Only root can run
	// lxc-attach
	if whoami, err := getCurrentUser(); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	} else if USER_ROOT != whoami.Username {
		glog.Warningf("Unable to pause/resume service - User is not %s - whoami:%+v", USER_ROOT, whoami)
	} else {
		var servicesList []*dao.Service
		if err := d.client.GetServices(unused, &servicesList); err != nil {
			glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
			return err
		}
		for _, service := range servicesList {
			if service.Snapshot.Pause == "" || service.Snapshot.Resume == "" {
				continue
			}

			var state dao.ServiceState
			if err := d.getServiceState(service.Id, &state); err != nil {
				glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
				return err
			}

			err := d.Pause(service, &state)
			defer d.Resume(service, &state)
			if err != nil {
				glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
				return err
			}
		}
	}

	// create a snapshot
	var (
		theVolume volume.Volume
	)
	if err := d.client.GetVolume(tenantId, &theVolume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
		return err
	} else {
		glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v volume=%+v", serviceId, theVolume)
		*label = getSnapshotLabel(&theVolume)
		if err := theVolume.Snapshot(*label); err != nil {
			glog.V(2).Infof("DistributedFileSystem.Snapshot service=%+v err=%s", serviceId, err)
			return err
		}
	}

	glog.V(2).Infof("Successfully created snapshot for service Id:%s Name:%s Label:%s", service.Id, service.Name, *label)
	return nil
}

func (d *DistributedFileSystem) Commit(dockerId string, label *string) error {
	Lock.Lock()
	defer Lock.Unlock()

	// Get the container
	container, err := d.dockerClient.InspectContainer(dockerId)
	if err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	} else if container.State.Running {
		err := errors.New("cannot commit a running container")
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}

	// Get the service id
	name := strings.Split(container.Name, "_")
	serviceId := name[0]

	// Get tag & repo information
	var (
		images []docker.APIImages
		image  *docker.APIImages
	)

	if err := d.getLatestImages(&images); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}
	for _, i := range images {
		if i.ID == container.Image {
			image = &i
			break
		}
	}
	if image == nil {
		err := errors.New("cannot commit a stale container")
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}

	// Snapshot the DFS
	if err := d.Snapshot(serviceId, label); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit dockerId=%+v err=%s", dockerId, err)
		return err
	}

	// Commit the container to the image
	newImage, err := d.dockerClient.CommitContainer(docker.CommitContainerOptions{
		Container:  container.ID,
		Repository: image.Repository,
		Tag:        image.Tag,
	})
	if err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	// Copy images to the DFS
	*image = docker.APIImages{
		ID:          newImage.ID,
		RepoTags:    image.RepoTags,
		Created:     newImage.Created.Unix(),
		Size:        newImage.Size,
		VirtualSize: image.VirtualSize,
		ParentId:    newImage.Parent,
		Repository:  image.Repository,
		Tag:         image.Tag,
	}

	// Get the path to the volume and write the images
	var (
		tenantId string
		volume   volume.Volume
	)

	if err := d.client.GetTenantId(serviceId, &tenantId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	if err := d.client.GetVolume(tenantId, &volume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	config := path.Join(volume.Path(), *label, DOCKER_IMAGEJSON)
	if data, err := json.Marshal(images); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	} else if err := ioutil.WriteFile(config, data, 0644); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Commit container=%+v err=%s", dockerId, err)
		return err
	}

	return nil
}

func (d *DistributedFileSystem) getLatestImages(images *[]docker.APIImages) error {
	// Get all of the images from docker and find the ones tagged as the latest
	if allImages, err := d.dockerClient.ListImages(true); err != nil {
		return err
	} else {
		for _, image := range allImages {
			if image.Tag == DOCKER_LATEST {
				*images = append(*images, image)
			}
		}
	}

	return nil
}

// Gets the images from the snapshot; returns error if image does not exist in docker
func (d *DistributedFileSystem) getSnapshotImages(snapshotId string, volume *volume.Volume, images *[]docker.APIImages) error {
	config := path.Join(volume.Path(), snapshotId, DOCKER_IMAGEJSON)

	if data, err := ioutil.ReadFile(config); err != nil {
		return err
	} else if err := json.Unmarshal(data, images); err != nil {
		return err
	}

	// Check if the images still exist
	for _, image := range *images {
		if _, err := d.dockerClient.InspectImage(image.ID); err != nil {
			return err
		}
	}

	return nil
}

// Retags containers with the given snapshot
func (d *DistributedFileSystem) retag(images *[]docker.APIImages, volume *volume.Volume) error {

	// Set the tag of the new image
	dockerBin, err := exec.LookPath("docker")
	if err != nil {
		return err
	}

	for _, image := range *images {
		// docker tag %{image.ID} %{image.Repo}:%{image.Tag}
		cmd := exec.Command(dockerBin, "tag", image.ID, fmt.Sprintf("%s:%s", image.Repository, image.Tag))
		if err := cmd.Run(); err != nil {
			return err
		}
	}

	return nil
}

func (d *DistributedFileSystem) Rollback(snapshotId string) error {
	Lock.Lock()
	defer Lock.Unlock()

	var (
		tenantId string
		volume   volume.Volume
	)

	parts := strings.Split(snapshotId, "_")
	if len(parts) != 2 {
		glog.V(2).Infof("DistributedFileSystem.Rollback malformed snapshot Id: %s", snapshotId)
		return errors.New("malformed snapshotId")
	}
	serviceId := parts[0]
	label := parts[1]
	if err := d.client.GetTenantId(serviceId, &tenantId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	// Validate existance of images for this snapshot
	var service dao.Service
	err := d.client.GetService(tenantId, &service)
	glog.V(2).Infof("Getting service instance: %s", tenantId)
	if err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	if err := d.client.GetVolume(tenantId, &volume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	var latestImages []docker.APIImages
	if err := d.getLatestImages(&latestImages); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	var snapshotImages []docker.APIImages
	if err := d.getSnapshotImages(snapshotId, &volume, &snapshotImages); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		return err
	}

	d.client.StopService(tenantId, nil)
	// TODO: Wait for real event that confirms shutdown
	time.Sleep(time.Second * 5) // wait for shutdown

	if err := d.retag(&snapshotImages, &volume); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		d.retag(&latestImages, &volume)
		return err
	}

	glog.V(2).Infof("performing rollback on %s to %s", tenantId, label)
	if err := volume.Rollback(snapshotId); err != nil {
		glog.V(2).Infof("DistributedFileSystem.Rollback service=%+v err=%s", serviceId, err)
		d.retag(&latestImages, &volume)
		return err
	}
	var unusedStr string = ""

	return d.client.StartService(tenantId, &unusedStr)
}

func getSnapshotLabel(v *volume.Volume) string {
	format := "20060102-150405"
	loc := time.Now()
	utc := loc.UTC()
	return fmt.Sprintf("%s_%s", v.Name(), utc.Format(format))
}<|endoftext|>


<commit_before>// 13 december 2015

package ui

// #cgo LDFLAGS: -L${SRCDIR} -lui -framework CoreFoundation -lpthread
// #include <CoreFoundation/CoreFoundation.h>
// #include <pthread.h>
// extern void _CFRunLoopSetCurrent(CFRunLoopRef);
// extern pthread_t _CFMainPThread;
import "C"

// OS X cares very deeply if we don't run on the very first thread the OS creates
// why? who knows. it's stupid and completely indefensible. let's use undocumented APIs to get around it.
// apple uses them too: http://www.opensource.apple.com/source/kext_tools/kext_tools-19.2/kextd_main.c?txt
// apple HAS SUGGESTED them too: http://lists.apple.com/archives/darwin-development/2002/Sep/msg00250.html
// gstreamer uses them too: http://cgit.freedesktop.org/gstreamer/gst-plugins-good/tree/sys/osxvideo/osxvideosink.m
func ensureMainThread() {
	// TODO set to nil like the apple code?
	C._CFRunLoopSetCurrent(C.CFRunLoopGetMain())
	// TODO is this part necessary?
	C._CFMainPThread = C.pthread_self()
}
<commit_msg>Set up rpath properly on OS X.<commit_after>// 13 december 2015

package ui

// #cgo LDFLAGS: -L${SRCDIR} -lui -framework CoreFoundation -lpthread -rpath @executable_path
// /* (thanks to http://jorgen.tjer.no/post/2014/05/20/dt-rpath-ld-and-at-rpath-dyld/ for the @executable_path clarifiaction */
// #include <CoreFoundation/CoreFoundation.h>
// #include <pthread.h>
// extern void _CFRunLoopSetCurrent(CFRunLoopRef);
// extern pthread_t _CFMainPThread;
import "C"

// OS X cares very deeply if we don't run on the very first thread the OS creates
// why? who knows. it's stupid and completely indefensible. let's use undocumented APIs to get around it.
// apple uses them too: http://www.opensource.apple.com/source/kext_tools/kext_tools-19.2/kextd_main.c?txt
// apple HAS SUGGESTED them too: http://lists.apple.com/archives/darwin-development/2002/Sep/msg00250.html
// gstreamer uses them too: http://cgit.freedesktop.org/gstreamer/gst-plugins-good/tree/sys/osxvideo/osxvideosink.m
func ensureMainThread() {
	// TODO set to nil like the apple code?
	C._CFRunLoopSetCurrent(C.CFRunLoopGetMain())
	// TODO is this part necessary?
	C._CFMainPThread = C.pthread_self()
}
<|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 datasrcs

import (
	"context"
	"fmt"
	"regexp"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/http"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/systems"
	"github.com/caffix/eventbus"
	"github.com/caffix/service"
	"github.com/caffix/stringset"
)

// ViewDNS is the Service that handles access to the ViewDNS data source.
type ViewDNS struct {
	service.BaseService

	SourceType string
}

// NewViewDNS returns he object initialized, but not yet started.
func NewViewDNS(sys systems.System) *ViewDNS {
	v := &ViewDNS{SourceType: requests.SCRAPE}

	v.BaseService = *service.NewBaseService(v, "ViewDNS")
	return v
}

// Description implements the Service interface.
func (v *ViewDNS) Description() string {
	return v.SourceType
}

// OnStart implements the Service interface.
func (v *ViewDNS) OnStart() error {
	v.SetRateLimit(1)
	return nil
}

// OnRequest implements the Service interface.
func (v *ViewDNS) OnRequest(ctx context.Context, args service.Args) {
	if req, ok := args.(*requests.DNSRequest); ok {
		v.dnsRequest(ctx, req)
	}
}

func (v *ViewDNS) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
	cfg, bus, err := ContextConfigBus(ctx)
	if err != nil {
		return
	}

	if !cfg.IsDomainInScope(req.Domain) {
		return
	}

	numRateLimitChecks(v, 9)
	bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
		fmt.Sprintf("Querying %s for %s subdomains", v.String(), req.Domain))

	var unique []string
	u := v.getIPHistoryURL(req.Domain)
	// The ViewDNS IP History lookup sometimes reveals interesting results
	page, err := http.RequestWebPage(ctx, u, nil, nil, nil)
	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh, fmt.Sprintf("%s: %s: %v", v.String(), u, err))
		return
	}

	// Look for IP addresses in the web page returned
	re := regexp.MustCompile(net.IPv4RE)
	for _, sd := range re.FindAllString(page, -1) {
		addr := NewUniqueElements(unique, sd)

		if len(addr) > 0 {
			bus.Publish(requests.NewAddrTopic, eventbus.PriorityHigh, &requests.AddrRequest{
				Address: addr[0],
				Domain:  req.Domain,
				Tag:     v.SourceType,
				Source:  v.String(),
			})
		}
	}
}

// OnWhoisRequest implements the Service interface.
func (v *ViewDNS) OnWhoisRequest(ctx context.Context, req *requests.WhoisRequest) {
	cfg, bus, err := ContextConfigBus(ctx)
	if err != nil {
		return
	}

	if !cfg.IsDomainInScope(req.Domain) {
		return
	}

	numRateLimitChecks(v, 9)
	u := v.getReverseWhoisURL(req.Domain)
	page, err := http.RequestWebPage(ctx, u, nil, nil, nil)
	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh, fmt.Sprintf("%s: %s: %v", v.String(), u, err))
		return
	}
	// Pull the table we need from the page content
	table := getViewDNSTable(page)
	if table == "" {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: %s: Failed to discover the table of results", v.String(), u),
		)
		return
	}
	// Get the list of domain names discovered through the reverse DNS service
	re := regexp.MustCompile("
([a-zA-Z0-9]{1}[a-zA-Z0-9-]{0,61}[a-zA-Z0-9]{1}[.]{1}[a-zA-Z0-9-]+)	")
	subs := re.FindAllStringSubmatch(table, -1)

	matches := stringset.New()
	for _, match := range subs {
		sub := match[1]
		if sub != "" {
			matches.Insert(strings.TrimSpace(sub))
		}
	}

	if len(matches) > 0 {
		bus.Publish(requests.NewWhoisTopic, eventbus.PriorityHigh, &requests.WhoisRequest{
			Domain:     req.Domain,
			NewDomains: matches.Slice(),
			Tag:        v.SourceType,
			Source:     v.String(),
		})
	}
}

func getViewDNSTable(page string) string {
	var begin, end int

	s := page
	for i := 0; i < 4; i++ {
		b := strings.Index(s, "<table")
		if b == -1 {
			return ""
		}
		begin += b + 6

		if e := strings.Index(s[b:], "
"); e != -1 {
			end = begin + e
		} else {
			return ""
		}
		s = page[end+8:]
	}
	i := strings.Index(page[begin:end], "<table")
	i = strings.Index(page[begin+i+6:end], "<table")
	return page[begin+i : end]
}

func (v *ViewDNS) getReverseWhoisURL(domain string) string {
	format := "https://viewdns.info/reversewhois/?q=%s"
	return fmt.Sprintf(format, domain)
}

func (v *ViewDNS) getIPHistoryURL(domain string) string {
	format := "https://viewdns.info/iphistory/?domain=%s"
	return fmt.Sprintf(format, domain)
}

// NewUniqueElements removes elements that have duplicates in the original or new elements.
func NewUniqueElements(orig []string, add ...string) []string {
	var n []string

	for _, av := range add {
		found := false
		s := strings.ToLower(av)

		// Check the original slice for duplicates
		for _, ov := range orig {
			if s == strings.ToLower(ov) {
				found = true
				break
			}
		}
		// Check that we didn't already add it in
		if !found {
			for _, nv := range n {
				if s == nv {
					found = true
					break
				}
			}
		}
		// If no duplicates were found, add the entry in
		if !found {
			n = append(n, s)
		}
	}
	return n
}
<commit_msg>Unreliable data source<commit_after><|endoftext|>


<commit_before>package duplicatedetector

import (
	"bytes"
	"fmt"

	"github.com/bradfitz/gomemcache/memcache"
)

// checker wraps some utility methods to work with memcached as duplicate detector
type checker struct {
	cache      *memcache.Client
	prefix     string
	expiration int32
	value      []byte
}

// NewChecker returns a new instance of a Duplicate Detector
// It can be configured with a special prefix (useful for namespacing different apps)
// and a TTL (in seconds, either a relative time from now, up to 1 month, or an absolute
// Unix epoch time. Zero means the items have no expiration time)
func NewChecker(mc *memcache.Client, prefix string, ttl int32) *checker {
	return &checker{cache: mc, prefix: prefix, expiration: ttl, value: []byte("x")}
}

// getKeyFor prepends the prefix to the item key
func (c *checker) getKeyFor(id string) string {
	return c.prefix + "_" + id
}

// getItemFor returns an Item object ready to be stored in memcache
func (c *checker) getItemFor(id string) *memcache.Item {
	return &memcache.Item{Key: c.getKeyFor(id), Value: c.value, Expiration: c.expiration}
}

// Set will unconditionally add the current item to the cache, even if it's already there
func (c *checker) Set(id string) error {
	return c.cache.Set(c.getItemFor(id))
}

// Has will check if the item has been previously seen already
// The function could return an error in case Memcache is not reachable or
// the retrieved value is not what was stored by the duplicate detector
func (c *checker) Has(id string) (bool, error) {
	k := c.getKeyFor(id)
	v, err := c.cache.Get(k)
	if err != nil {
		return false, err
	}
	if bytes.Equal(v.Value, c.value) {
		return true, nil
	}
	return false, fmt.Errorf("key '%s' in cache, but unrecognised value %v", k, v.Value)
}

// Delete will remove the item from the cache, allowing a new Item with the same key in
func (c *checker) Delete(id string) error {
	err := c.cache.Delete(c.getKeyFor(id))
	if err == nil || err == memcache.ErrCacheMiss {
		return nil
	}
	return err
}

// IsDuplicate checks if the ID has been seen before (true) or if it's the first time (false).
// This counts as a touch: the first time an ID is checked, it is added to the cache;
// the second time the same ID is checked, it is considered as a duplicate
// The function could return an error in case Memcache is not reachable
func (c *checker) IsDuplicate(id string) (bool, error) {
	err := c.cache.Add(c.getItemFor(id))
	if memcache.ErrNotStored == err {
		return true, nil
	}
	return false, err
}
<commit_msg>make struct public<commit_after>package duplicatedetector

import (
	"bytes"
	"fmt"

	"github.com/bradfitz/gomemcache/memcache"
)

// checker wraps some utility methods to work with memcached as duplicate detector
type Checker struct {
	cache      *memcache.Client
	prefix     string
	expiration int32
	value      []byte
}

// NewChecker returns a new instance of a Duplicate Detector
// It can be configured with a special prefix (useful for namespacing different apps)
// and a TTL (in seconds, either a relative time from now, up to 1 month, or an absolute
// Unix epoch time. Zero means the items have no expiration time)
func NewChecker(mc *memcache.Client, prefix string, ttl int32) *Checker {
	return &checker{cache: mc, prefix: prefix, expiration: ttl, value: []byte("x")}
}

// getKeyFor prepends the prefix to the item key
func (c *Checker) getKeyFor(id string) string {
	return c.prefix + "_" + id
}

// getItemFor returns an Item object ready to be stored in memcache
func (c *Checker) getItemFor(id string) *memcache.Item {
	return &memcache.Item{Key: c.getKeyFor(id), Value: c.value, Expiration: c.expiration}
}

// Set will unconditionally add the current item to the cache, even if it's already there
func (c *Checker) Set(id string) error {
	return c.cache.Set(c.getItemFor(id))
}

// Has will check if the item has been previously seen already
// The function could return an error in case Memcache is not reachable or
// the retrieved value is not what was stored by the duplicate detector
func (c *Checker) Has(id string) (bool, error) {
	k := c.getKeyFor(id)
	v, err := c.cache.Get(k)
	if err != nil {
		return false, err
	}
	if bytes.Equal(v.Value, c.value) {
		return true, nil
	}
	return false, fmt.Errorf("key '%s' in cache, but unrecognised value %v", k, v.Value)
}

// Delete will remove the item from the cache, allowing a new Item with the same key in
func (c *Checker) Delete(id string) error {
	err := c.cache.Delete(c.getKeyFor(id))
	if err == nil || err == memcache.ErrCacheMiss {
		return nil
	}
	return err
}

// IsDuplicate checks if the ID has been seen before (true) or if it's the first time (false).
// This counts as a touch: the first time an ID is checked, it is added to the cache;
// the second time the same ID is checked, it is considered as a duplicate
// The function could return an error in case Memcache is not reachable
func (c *Checker) IsDuplicate(id string) (bool, error) {
	err := c.cache.Add(c.getItemFor(id))
	if memcache.ErrNotStored == err {
		return true, nil
	}
	return false, err
}
<|endoftext|>