<commit_before>package pubnub

import (
	"errors"
	"sync"

	"github.com/pubnub/go/messaging"
)

var (
	ErrChannelNotSet    = errors.New("channel name is not set")
	ErrTimeout          = errors.New("timeout reached")
	ErrConnectionAbort  = errors.New("connection aborted")
	ErrConnectionClosed = errors.New("connection closed")
)

type PubNubClient struct {
	pub      *messaging.Pubnub
	channels map[string]*Channel
	mu       sync.RWMutex
	closed   bool
}

func NewPubNubClient(cs *ClientSettings) *PubNubClient {
	pub := messaging.NewPubnub(cs.PublishKey, cs.SubscribeKey, cs.SecretKey, cs.Cipher, cs.SSL, cs.ID)
	return &PubNubClient{
		pub:      pub,
		channels: make(map[string]*Channel),
	}
}

////////////// ClientSettings ////////////////////

type ClientSettings struct {
	PublishKey   string
	SubscribeKey string
	SecretKey    string
	Cipher       string
	SSL          bool
	ID           string
}

// Push sends a message to the channel with channelName. If Access Manager is enabled
// access must be granted first.
func (p *PubNubClient) Push(channelName string, body interface{}) error {
	if p.closed {
		return ErrConnectionClosed
	}

	pr := NewPubNubRequest(channelName, nil, nil)
	defer pr.Close()

	go pr.handleResponse()

	go p.pub.Publish(channelName, body, pr.successCh, pr.errorCh)

	return pr.Do()
}

// Grant read/write access to the given token for TTL period. If token
func (p *PubNubClient) Grant(a *AuthSettings) error {
	if p.closed {
		return ErrConnectionClosed
	}

	if a.ChannelName == "" {
		return ErrChannelNotSet
	}

	pr := NewPubNubRequest(a.ChannelName, nil, nil)
	defer pr.Close()

	p.pub.SetAuthenticationKey(a.Token)

	go pr.handleResponse()

	// channel name, read access, write access, TTL, success channel, error channel
	go p.pub.GrantSubscribe(a.ChannelName, a.CanRead, a.CanWrite, a.TTL, pr.successCh, pr.errorCh)

	return pr.Do()
}

// Subscribe to given channel.
// Returns a message listener channel
// TODO add support for multiple channel subscription
func (p *PubNubClient) Subscribe(channelName string) (*Channel, error) {
	if p.closed {
		return nil, ErrConnectionClosed
	}

	if channelName == "" {
		return nil, ErrChannelNotSet
	}

	return p.fetchOrCreateChannel(channelName)
}

func (p *PubNubClient) Close() {
	p.pub.CloseExistingConnection()

	for _, channel := range p.channels {
		channel.Close()
	}

	p.closed = true
}

func (p *PubNubClient) SetAuthToken(token string) {
	p.pub.SetAuthenticationKey(token)
}

func (p *PubNubClient) fetchOrCreateChannel(channelName string) (*Channel, error) {
	p.mu.RLock()
	channel, ok := p.channels[channelName]
	p.mu.RUnlock()
	if ok {
		return channel, nil
	}

	channel, err := p.NewChannel(channelName)
	if err != nil {
		return nil, err
	}

	p.mu.Lock()

	p.channels[channelName] = channel

	p.mu.Unlock()

	return channel, nil
}

//////////////////// PubnubAuthSettings /////////////////////

type AuthSettings struct {

	// PubNub channel name
	ChannelName string

	// Grant access for the token. When token is an
	// empty string it provides public access for the channel.
	Token string

	// Grant read access
	CanRead bool

	// Grant write Access
	CanWrite bool

	// Time to live value in minutes.
	// Access is revoked after TTL period
	// Min-max values can be consecutively: 1 and 525600
	// 0 value will grant access indefinitely
	// -1 causes default value (1440) to be set
	TTL int
}
<commit_msg>fix package according to new pubnub/go package<commit_after>package pubnub

import (
	"errors"
	"sync"

	"github.com/pubnub/go/messaging"
)

var (
	ErrChannelNotSet    = errors.New("channel name is not set")
	ErrTimeout          = errors.New("timeout reached")
	ErrConnectionAbort  = errors.New("connection aborted")
	ErrConnectionClosed = errors.New("connection closed")
)

type PubNubClient struct {
	pub      *messaging.Pubnub
	channels map[string]*Channel
	mu       sync.RWMutex
	closed   bool
}

func NewPubNubClient(cs *ClientSettings) *PubNubClient {
	pub := messaging.NewPubnub(cs.PublishKey, cs.SubscribeKey, cs.SecretKey, cs.Cipher, cs.SSL, cs.ID, nil)
	return &PubNubClient{
		pub:      pub,
		channels: make(map[string]*Channel),
	}
}

////////////// ClientSettings ////////////////////

type ClientSettings struct {
	PublishKey   string
	SubscribeKey string
	SecretKey    string
	Cipher       string
	SSL          bool
	ID           string
}

// Push sends a message to the channel with channelName. If Access Manager is enabled
// access must be granted first.
func (p *PubNubClient) Push(channelName string, body interface{}) error {
	if p.closed {
		return ErrConnectionClosed
	}

	pr := NewPubNubRequest(channelName, nil, nil)
	defer pr.Close()

	go pr.handleResponse()

	go p.pub.Publish(channelName, body, pr.successCh, pr.errorCh)

	return pr.Do()
}

// Grant read/write access to the given token for TTL period. If token
func (p *PubNubClient) Grant(a *AuthSettings) error {
	if p.closed {
		return ErrConnectionClosed
	}

	if a.ChannelName == "" {
		return ErrChannelNotSet
	}

	pr := NewPubNubRequest(a.ChannelName, nil, nil)
	defer pr.Close()

	p.pub.SetAuthenticationKey(a.Token)

	go pr.handleResponse()

	// channel name, read access, write access, TTL, success channel, error channel
	go p.pub.GrantSubscribe(a.ChannelName, a.CanRead, a.CanWrite, a.TTL, a.Token, pr.successCh, pr.errorCh)

	return pr.Do()
}

// Subscribe to given channel.
// Returns a message listener channel
// TODO add support for multiple channel subscription
func (p *PubNubClient) Subscribe(channelName string) (*Channel, error) {
	if p.closed {
		return nil, ErrConnectionClosed
	}

	if channelName == "" {
		return nil, ErrChannelNotSet
	}

	return p.fetchOrCreateChannel(channelName)
}

func (p *PubNubClient) Close() {
	p.pub.CloseExistingConnection()

	for _, channel := range p.channels {
		channel.Close()
	}

	p.closed = true
}

func (p *PubNubClient) SetAuthToken(token string) {
	p.pub.SetAuthenticationKey(token)
}

func (p *PubNubClient) fetchOrCreateChannel(channelName string) (*Channel, error) {
	p.mu.RLock()
	channel, ok := p.channels[channelName]
	p.mu.RUnlock()
	if ok {
		return channel, nil
	}

	channel, err := p.NewChannel(channelName)
	if err != nil {
		return nil, err
	}

	p.mu.Lock()

	p.channels[channelName] = channel

	p.mu.Unlock()

	return channel, nil
}

//////////////////// PubnubAuthSettings /////////////////////

type AuthSettings struct {

	// PubNub channel name
	ChannelName string

	// Grant access for the token. When token is an
	// empty string it provides public access for the channel.
	Token string

	// Grant read access
	CanRead bool

	// Grant write Access
	CanWrite bool

	// Time to live value in minutes.
	// Access is revoked after TTL period
	// Min-max values can be consecutively: 1 and 525600
	// 0 value will grant access indefinitely
	// -1 causes default value (1440) to be set
	TTL int
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess 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 tabletmanager

import (
	"testing"
	"time"

	"context"

	"github.com/stretchr/testify/assert"

	"vitess.io/vitess/go/vt/mysqlctl"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

func TestReplManagerSetTabletType(t *testing.T) {
	defer func(saved bool) { *mysqlctl.DisableActiveReparents = saved }(*mysqlctl.DisableActiveReparents)
	*mysqlctl.DisableActiveReparents = true

	tm := &TabletManager{}
	tm.replManager = newReplManager(context.Background(), tm, 100*time.Millisecond)

	// DisableActiveReparents == true should result in no-op
	*mysqlctl.DisableActiveReparents = true
	tm.replManager.SetTabletType(topodatapb.TabletType_REPLICA)
	assert.False(t, tm.replManager.ticks.Running())

	// primary should stop the manager
	*mysqlctl.DisableActiveReparents = false
	tm.replManager.ticks.Start(nil)
	tm.replManager.SetTabletType(topodatapb.TabletType_PRIMARY)
	assert.False(t, tm.replManager.ticks.Running())

	// If replcation is stopped, the manager should not start
	tm.replManager.setReplicationStopped(true)
	tm.replManager.ticks.Start(nil)
	tm.replManager.SetTabletType(topodatapb.TabletType_REPLICA)
	assert.False(t, tm.replManager.ticks.Running())

	// If the manager is already running, rm.check should not be called
	tm.replManager.setReplicationStopped(false)
	tm.replManager.ticks.Start(nil)
	tm.replManager.SetTabletType(topodatapb.TabletType_REPLICA)
	assert.True(t, tm.replManager.ticks.Running())
	tm.replManager.ticks.Stop()
}

func TestReplManagerSetReplicationStopped(t *testing.T) {
	defer func(saved bool) { *mysqlctl.DisableActiveReparents = saved }(*mysqlctl.DisableActiveReparents)
	*mysqlctl.DisableActiveReparents = true

	tm := &TabletManager{}
	tm.replManager = newReplManager(context.Background(), tm, 100*time.Millisecond)

	// DisableActiveReparents == true should result in no-op
	*mysqlctl.DisableActiveReparents = true
	tm.replManager.setReplicationStopped(true)
	assert.False(t, tm.replManager.ticks.Running())
	tm.replManager.setReplicationStopped(false)
	assert.False(t, tm.replManager.ticks.Running())

	*mysqlctl.DisableActiveReparents = false
	tm.replManager.setReplicationStopped(false)
	assert.True(t, tm.replManager.ticks.Running())
	tm.replManager.setReplicationStopped(true)
	assert.False(t, tm.replManager.ticks.Running())
}

func TestReplManagerReset(t *testing.T) {
	defer func(saved bool) { *mysqlctl.DisableActiveReparents = saved }(*mysqlctl.DisableActiveReparents)

	tm := &TabletManager{}
	tm.replManager = newReplManager(context.Background(), tm, 100*time.Millisecond)

	*mysqlctl.DisableActiveReparents = false
	tm.replManager.setReplicationStopped(false)
	assert.True(t, tm.replManager.ticks.Running())
	// reset should not affect the ticks, but the replication stopped should be false
	tm.replManager.reset()
	assert.True(t, tm.replManager.ticks.Running())
	assert.False(t, tm.replManager.replicationStopped())

	tm.replManager.setReplicationStopped(true)
	assert.False(t, tm.replManager.ticks.Running())
	tm.replManager.reset()
	// reset should not affect the ticks, but the replication stopped should be false
	assert.False(t, tm.replManager.ticks.Running())
	assert.False(t, tm.replManager.replicationStopped())
}
<commit_msg>test: added a test to check reset is a no-op when disableActiveReparents has been set<commit_after>/*
Copyright 2020 The Vitess 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 tabletmanager

import (
	"testing"
	"time"

	"context"

	"github.com/stretchr/testify/assert"

	"vitess.io/vitess/go/vt/mysqlctl"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

func TestReplManagerSetTabletType(t *testing.T) {
	defer func(saved bool) { *mysqlctl.DisableActiveReparents = saved }(*mysqlctl.DisableActiveReparents)
	*mysqlctl.DisableActiveReparents = true

	tm := &TabletManager{}
	tm.replManager = newReplManager(context.Background(), tm, 100*time.Millisecond)

	// DisableActiveReparents == true should result in no-op
	*mysqlctl.DisableActiveReparents = true
	tm.replManager.SetTabletType(topodatapb.TabletType_REPLICA)
	assert.False(t, tm.replManager.ticks.Running())

	// primary should stop the manager
	*mysqlctl.DisableActiveReparents = false
	tm.replManager.ticks.Start(nil)
	tm.replManager.SetTabletType(topodatapb.TabletType_PRIMARY)
	assert.False(t, tm.replManager.ticks.Running())

	// If replcation is stopped, the manager should not start
	tm.replManager.setReplicationStopped(true)
	tm.replManager.ticks.Start(nil)
	tm.replManager.SetTabletType(topodatapb.TabletType_REPLICA)
	assert.False(t, tm.replManager.ticks.Running())

	// If the manager is already running, rm.check should not be called
	tm.replManager.setReplicationStopped(false)
	tm.replManager.ticks.Start(nil)
	tm.replManager.SetTabletType(topodatapb.TabletType_REPLICA)
	assert.True(t, tm.replManager.ticks.Running())
	tm.replManager.ticks.Stop()
}

func TestReplManagerSetReplicationStopped(t *testing.T) {
	defer func(saved bool) { *mysqlctl.DisableActiveReparents = saved }(*mysqlctl.DisableActiveReparents)
	*mysqlctl.DisableActiveReparents = true

	tm := &TabletManager{}
	tm.replManager = newReplManager(context.Background(), tm, 100*time.Millisecond)

	// DisableActiveReparents == true should result in no-op
	*mysqlctl.DisableActiveReparents = true
	tm.replManager.setReplicationStopped(true)
	assert.False(t, tm.replManager.ticks.Running())
	tm.replManager.setReplicationStopped(false)
	assert.False(t, tm.replManager.ticks.Running())

	*mysqlctl.DisableActiveReparents = false
	tm.replManager.setReplicationStopped(false)
	assert.True(t, tm.replManager.ticks.Running())
	tm.replManager.setReplicationStopped(true)
	assert.False(t, tm.replManager.ticks.Running())
}

func TestReplManagerReset(t *testing.T) {
	defer func(saved bool) { *mysqlctl.DisableActiveReparents = saved }(*mysqlctl.DisableActiveReparents)

	tm := &TabletManager{}
	tm.replManager = newReplManager(context.Background(), tm, 100*time.Millisecond)

	*mysqlctl.DisableActiveReparents = false
	tm.replManager.setReplicationStopped(false)
	assert.True(t, tm.replManager.ticks.Running())
	// reset should not affect the ticks, but the replication stopped should be false
	tm.replManager.reset()
	assert.True(t, tm.replManager.ticks.Running())
	assert.False(t, tm.replManager.replicationStopped())

	tm.replManager.setReplicationStopped(true)
	assert.False(t, tm.replManager.ticks.Running())
	tm.replManager.reset()
	// reset should not affect the ticks, but the replication stopped should be false
	assert.False(t, tm.replManager.ticks.Running())
	assert.False(t, tm.replManager.replicationStopped())

	tm.replManager.setReplicationStopped(true)
	assert.True(t, tm.replManager.replicationStopped())
	// DisableActiveReparents == true should result in no-op
	*mysqlctl.DisableActiveReparents = true
	tm.replManager.reset()
	assert.True(t, tm.replManager.replicationStopped())
}
<|endoftext|>


<commit_before>package txdb

import (
	"context"
	"reflect"
	"testing"

	"chain/database/pg"
	"chain/database/pg/pgtest"
	"chain/database/sql"
	"chain/errors"
	"chain/protocol/bc"
	"chain/testutil"
)

func mustParseHash(s string) bc.Hash {
	h, err := bc.ParseHash(s)
	if err != nil {
		panic(err)
	}
	return h
}

func TestPoolTxs(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()

	_, err := dbtx.Exec(ctx, `
		INSERT INTO pool_txs (tx_hash, data)
		VALUES (
			'70dd2c70c5c1e859c32bcc4415d15c70a95c977e0c35c115541f692f00ffbe9c',
			decode('07010200000000000568656c6c6f', 'hex')
		);
	`)
	if err != nil {
		testutil.FatalErr(t, err)
	}

	got, err := dumpPoolTxs(ctx, dbtx)
	if err != nil {
		t.Fatalf("err got = %v want nil", err)
	}

	wantTx := bc.NewTx(bc.TxData{
		Version:       1,
		ReferenceData: []byte("hello"),
	})
	want := []*bc.Tx{wantTx}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("txs do not match")
		for _, tx := range got {
			t.Logf("\tgot %v", tx)
		}
		for _, tx := range want {
			t.Logf("\twant %v", tx)
		}
	}
}

func TestGetTxs(t *testing.T) {
	_, db := pgtest.NewDB(t, pgtest.SchemaPath)
	dbctx := pg.NewContext(context.Background(), db)
	pool := NewPool(pg.FromContext(dbctx).(*sql.DB))
	ctx := context.Background()

	tx := bc.NewTx(bc.TxData{ReferenceData: []byte("tx")})
	err := pool.Insert(ctx, tx)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}
	poolTxs, err := pool.GetTxs(ctx, tx.Hash)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}

	if !reflect.DeepEqual(poolTxs[tx.Hash], tx) {
		t.Errorf("got:\n\t%+v\nwant:\n\t%+v", poolTxs[tx.Hash], tx)
	}

	nonexistentHash := mustParseHash("beefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeef")
	_, gotErr := pool.GetTxs(ctx, tx.Hash, nonexistentHash)
	if gotErr != nil {
		t.Errorf("got err=%v want nil", gotErr)
	}
}

func TestInsertPoolTx(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()
	tx := bc.NewTx(bc.TxData{ReferenceData: []byte("tx")})
	err := (&Pool{dbtx}).Insert(ctx, tx)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}
}

func TestGetBlock(t *testing.T) {
	_, db := pgtest.NewDB(t, pgtest.SchemaPath)
	dbctx := pg.NewContext(context.Background(), db)
	pgtest.Exec(dbctx, t, `
		INSERT INTO blocks (block_hash, height, data, header)
		VALUES
		('0000000000000000000000000000000000000000000000000000000000000000', 0, '', ''),
		(
			'1f20d89dd393f452b4396589ed5d6f90465cb032aa3f9fe42a99d47c7089b0a3',
			1,
			decode('03010131323300000000000000000000000000000000000000000000000000000000006453414243000000000000000000000000000000000000000000000000000000000058595a000000000000000000000000000000000000000000000000000000000012746573742d6f75747075742d73637269707411010f746573742d7369672d73637269707401070102000000000007746573742d7478', 'hex'),
			''
		);
	`)
	store := NewStore(pg.FromContext(dbctx).(*sql.DB))
	ctx := context.Background()
	got, err := store.GetBlock(ctx, 1)
	if err != nil {
		t.Fatalf("err got = %v want nil", err)
	}
	want := &bc.Block{
		BlockHeader: bc.BlockHeader{
			Version:           1,
			Height:            1,
			PreviousBlockHash: [32]byte{'1', '2', '3'},
			TransactionsMerkleRoot: bc.Hash{
				'A', 'B', 'C', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
				0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
			},
			AssetsMerkleRoot: bc.Hash{
				'X', 'Y', 'Z', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
				0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
			},
			TimestampMS:      100,
			Witness:          [][]byte{[]byte("test-sig-script")},
			ConsensusProgram: []byte("test-output-script"),
		},
		Transactions: []*bc.Tx{
			bc.NewTx(bc.TxData{Version: 1, ReferenceData: []byte("test-tx")}),
		},
	}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("latest block:\ngot:  %+v\nwant: %+v", got, want)
	}
}

func getBlockByHash(ctx context.Context, db pg.DB, hash string) (*bc.Block, error) {
	const q = `SELECT data FROM blocks WHERE block_hash=$1`
	block := new(bc.Block)
	err := db.QueryRow(ctx, q, hash).Scan(block)
	if err == sql.ErrNoRows {
		err = pg.ErrUserInputNotFound
	}
	return block, errors.WithDetailf(err, "block hash=%v", hash)
}

func TestInsertBlock(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()
	blk := &bc.Block{
		BlockHeader: bc.BlockHeader{
			Version: 1,
			Height:  1,
		},
		Transactions: []*bc.Tx{
			bc.NewTx(bc.TxData{
				ReferenceData: []byte("a"),
			}),
			bc.NewTx(bc.TxData{
				ReferenceData: []byte("b"),
			}),
		},
	}
	err := (&Store{dbtx}).SaveBlock(ctx, blk)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}

	// block in database
	_, err = getBlockByHash(ctx, dbtx, blk.Hash().String())
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}
}

func TestListBlocks(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()
	blks := []*bc.Block{
		{
			BlockHeader: bc.BlockHeader{
				Version: 1,
				Height:  2,
			},
		},
		{
			BlockHeader: bc.BlockHeader{
				Version: 1,
				Height:  1,
			},
		},
	}
	for _, blk := range blks {
		err := (&Store{dbtx}).SaveBlock(ctx, blk)
		if err != nil {
			t.Log(errors.Stack(err))
			t.Fatal(err)
		}
	}
	cases := []struct {
		prev  string
		limit int
		want  []*bc.Block
	}{{
		prev:  "",
		limit: 50,
		want:  blks,
	}, {
		prev:  "2",
		limit: 50,
		want:  []*bc.Block{blks[1]},
	}, {
		prev:  "",
		limit: 1,
		want:  []*bc.Block{blks[0]},
	}, {
		prev:  "1",
		limit: 50,
		want:  nil,
	}}

	for _, c := range cases {
		got, err := listBlocks(ctx, dbtx, c.prev, c.limit)
		if err != nil {
			t.Log(errors.Stack(err))
			t.Errorf("ListBlocks(%q, %d) error = %q", c.prev, c.limit, err)
			continue
		}

		if !reflect.DeepEqual(got, c.want) {
			t.Errorf("got ListBlocks(%q, %d):\n\t%+v\nwant:\n\t%+v", c.prev, c.limit, got, c.want)
		}
	}
}
<commit_msg>core/txdb: make tests faster<commit_after>package txdb

import (
	"context"
	"reflect"
	"testing"

	"chain/database/pg"
	"chain/database/pg/pgtest"
	"chain/database/sql"
	"chain/errors"
	"chain/protocol/bc"
	"chain/testutil"
)

func mustParseHash(s string) bc.Hash {
	h, err := bc.ParseHash(s)
	if err != nil {
		panic(err)
	}
	return h
}

func TestPoolTxs(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()

	_, err := dbtx.Exec(ctx, `
		INSERT INTO pool_txs (tx_hash, data)
		VALUES (
			'70dd2c70c5c1e859c32bcc4415d15c70a95c977e0c35c115541f692f00ffbe9c',
			decode('07010200000000000568656c6c6f', 'hex')
		);
	`)
	if err != nil {
		testutil.FatalErr(t, err)
	}

	got, err := dumpPoolTxs(ctx, dbtx)
	if err != nil {
		t.Fatalf("err got = %v want nil", err)
	}

	wantTx := bc.NewTx(bc.TxData{
		Version:       1,
		ReferenceData: []byte("hello"),
	})
	want := []*bc.Tx{wantTx}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("txs do not match")
		for _, tx := range got {
			t.Logf("\tgot %v", tx)
		}
		for _, tx := range want {
			t.Logf("\twant %v", tx)
		}
	}
}

func TestGetTxs(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	pool := NewPool(dbtx)
	ctx := context.Background()

	tx := bc.NewTx(bc.TxData{ReferenceData: []byte("tx")})
	err := pool.Insert(ctx, tx)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}
	poolTxs, err := pool.GetTxs(ctx, tx.Hash)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}

	if !reflect.DeepEqual(poolTxs[tx.Hash], tx) {
		t.Errorf("got:\n\t%+v\nwant:\n\t%+v", poolTxs[tx.Hash], tx)
	}

	nonexistentHash := mustParseHash("beefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeefbeef")
	_, gotErr := pool.GetTxs(ctx, tx.Hash, nonexistentHash)
	if gotErr != nil {
		t.Errorf("got err=%v want nil", gotErr)
	}
}

func TestInsertPoolTx(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()
	tx := bc.NewTx(bc.TxData{ReferenceData: []byte("tx")})
	err := (&Pool{dbtx}).Insert(ctx, tx)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}
}

func TestGetBlock(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	pgtest.Exec(pg.NewContext(context.Background(), dbtx), t, `
		INSERT INTO blocks (block_hash, height, data, header)
		VALUES
		('0000000000000000000000000000000000000000000000000000000000000000', 0, '', ''),
		(
			'1f20d89dd393f452b4396589ed5d6f90465cb032aa3f9fe42a99d47c7089b0a3',
			1,
			decode('03010131323300000000000000000000000000000000000000000000000000000000006453414243000000000000000000000000000000000000000000000000000000000058595a000000000000000000000000000000000000000000000000000000000012746573742d6f75747075742d73637269707411010f746573742d7369672d73637269707401070102000000000007746573742d7478', 'hex'),
			''
		);
	`)
	store := NewStore(dbtx)
	ctx := context.Background()
	got, err := store.GetBlock(ctx, 1)
	if err != nil {
		t.Fatalf("err got = %v want nil", err)
	}
	want := &bc.Block{
		BlockHeader: bc.BlockHeader{
			Version:           1,
			Height:            1,
			PreviousBlockHash: [32]byte{'1', '2', '3'},
			TransactionsMerkleRoot: bc.Hash{
				'A', 'B', 'C', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
				0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
			},
			AssetsMerkleRoot: bc.Hash{
				'X', 'Y', 'Z', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
				0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
			},
			TimestampMS:      100,
			Witness:          [][]byte{[]byte("test-sig-script")},
			ConsensusProgram: []byte("test-output-script"),
		},
		Transactions: []*bc.Tx{
			bc.NewTx(bc.TxData{Version: 1, ReferenceData: []byte("test-tx")}),
		},
	}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("latest block:\ngot:  %+v\nwant: %+v", got, want)
	}
}

func getBlockByHash(ctx context.Context, db pg.DB, hash string) (*bc.Block, error) {
	const q = `SELECT data FROM blocks WHERE block_hash=$1`
	block := new(bc.Block)
	err := db.QueryRow(ctx, q, hash).Scan(block)
	if err == sql.ErrNoRows {
		err = pg.ErrUserInputNotFound
	}
	return block, errors.WithDetailf(err, "block hash=%v", hash)
}

func TestInsertBlock(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()
	blk := &bc.Block{
		BlockHeader: bc.BlockHeader{
			Version: 1,
			Height:  1,
		},
		Transactions: []*bc.Tx{
			bc.NewTx(bc.TxData{
				ReferenceData: []byte("a"),
			}),
			bc.NewTx(bc.TxData{
				ReferenceData: []byte("b"),
			}),
		},
	}
	err := (&Store{dbtx}).SaveBlock(ctx, blk)
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}

	// block in database
	_, err = getBlockByHash(ctx, dbtx, blk.Hash().String())
	if err != nil {
		t.Log(errors.Stack(err))
		t.Fatal(err)
	}
}

func TestListBlocks(t *testing.T) {
	dbtx := pgtest.NewTx(t)
	ctx := context.Background()
	blks := []*bc.Block{
		{
			BlockHeader: bc.BlockHeader{
				Version: 1,
				Height:  2,
			},
		},
		{
			BlockHeader: bc.BlockHeader{
				Version: 1,
				Height:  1,
			},
		},
	}
	for _, blk := range blks {
		err := (&Store{dbtx}).SaveBlock(ctx, blk)
		if err != nil {
			t.Log(errors.Stack(err))
			t.Fatal(err)
		}
	}
	cases := []struct {
		prev  string
		limit int
		want  []*bc.Block
	}{{
		prev:  "",
		limit: 50,
		want:  blks,
	}, {
		prev:  "2",
		limit: 50,
		want:  []*bc.Block{blks[1]},
	}, {
		prev:  "",
		limit: 1,
		want:  []*bc.Block{blks[0]},
	}, {
		prev:  "1",
		limit: 50,
		want:  nil,
	}}

	for _, c := range cases {
		got, err := listBlocks(ctx, dbtx, c.prev, c.limit)
		if err != nil {
			t.Log(errors.Stack(err))
			t.Errorf("ListBlocks(%q, %d) error = %q", c.prev, c.limit, err)
			continue
		}

		if !reflect.DeepEqual(got, c.want) {
			t.Errorf("got ListBlocks(%q, %d):\n\t%+v\nwant:\n\t%+v", c.prev, c.limit, got, c.want)
		}
	}
}
<|endoftext|>


<commit_before>// Package main is used for testing of generated 'views' listing.
// There is no way to include a new import dynamically, thus
// we are running this test from generate_test.go
// as a new command using exec package.
package main

import (
	"../assets/views"

	"github.com/anonx/sunplate/log"
)

func main() {
	if l := len(views.Context); l != 2 {
		log.Error.Fatalf("Length of views.Context expected to be equal to 2, it is %d instead.", l)
	}

	//
	// Make sure templates are presented in the format we expect.
	//
	for k, v := range expectedValues {
		if views.Context[k] != v {
			log.Error.Fatalf("'%s' wasn't found in %#v.", k, views.Context)
		}
	}
}

var expectedValues = map[string]string{
	"testdata/views/test1.template": "testdata/views/test1.template",
	"testdata/views/test2.template": "testdata/views/test2.template",
}
<commit_msg>Update expected results of generated liting<commit_after>// Package main is used for testing of generated 'views' listing.
// There is no way to include a new import dynamically, thus
// we are running this test from generate_test.go
// as a new command using exec package.
package main

import (
	"../assets/views"

	"github.com/anonx/sunplate/log"
)

func main() {
	if l := len(views.Context); l != 2 {
		log.Error.Fatalf("Length of views.Context expected to be equal to 2, it is %d instead.", l)
	}

	//
	// Make sure templates are presented in the format we expect.
	//
	for k, v := range expectedValues {
		if views.Context[k] != v {
			log.Error.Fatalf("'%s' wasn't found in %#v.", k, views.Context)
		}
	}
}

var expectedValues = map[string]string{
	"test1.template": "testdata/views/test1.template",
	"test2.template": "testdata/views/test2.template",
}
<|endoftext|>


<commit_before>// Package boxer is a streaming encryption implementation, based on Adam
// Langley's article: https://www.imperialviolet.org/2014/06/27/streamingencryption.html
//
// In short, nacl/secretbox is used to seal a file in chunks, with each chunk
// being prefixed with its length. The nonce is incrementally marked so
// chunks are guaranteed to be in order. The encrypted blob is prepended with
// a header containing a version ID, the maximum chunk size, and flags. The
// flags are currently unused, but may be used in future versions.
package boxer

import (
	"encoding/binary"
	"errors"
	"io"
	"math"

	"golang.org/x/crypto/nacl/secretbox"
)

var (
	ErrAlreadyClosed = errors.New("encryptor: already closed")
	ErrInvalidData   = errors.New("decryptor: encrypted message is invalid")
	ErrChunkSize     = errors.New("boxer: invalid chunk size")
)

const (
	// DefaultChunkSize is the default maximum chunk size for reading and
	// writing.
	DefaultChunkSize = 65536

	// Overhead is the number of bytes of overhead when boxing a message.
	Overhead = secretbox.Overhead

	// offset is the number of bytes used to advise the length of the
	// chunk. It should be large enough to advise the entirety of
	// the chunk.
	offset = 4

	ver1 = 1
)

func nonceKey(nonce *[16]byte, key *[32]byte) (*[24]byte, *[32]byte) {
	var n [24]byte
	copy(n[:], nonce[:])
	var k [32]byte
	copy(k[:], key[:])
	return &n, &k
}

type chunk uint32

// Encryptor is an io.WriteCloser. Writes to an Encryptor are encrypted
// and written to w.
type Encryptor struct {
	w     io.Writer // underlying writer
	nonce *[24]byte // nacl nonce, increments per chunk
	key   *[32]byte // encryption key
	in    []byte    // input buffer
	out   []byte    // encryption buffer
	size  int       // chunk size
	n     int       // end of buffer
	err   error     // last error
}

// NewEncryptor returns a new Encryptor. Writes to the returned Encryptor
// are encrypted and written to w. The size parameter dictates the maximum
// chunk size. It should be a positive integer in the range [0, 1 << 32 - 1].
// Writes will always be chunk size +
//
// All writes will not be flushed until Close is called. Not closing an
// Encryptor will rsult in an invalid stream.
//
// Neither nonce or key are modified.
func NewEncryptorSize(w io.Writer, nonce *[16]byte, key *[32]byte, size int) (*Encryptor, error) {
	if size > math.MaxInt32 {
		return nil, ErrChunkSize
	}
	e := Encryptor{w: w, size: size}
	err := e.writeHeaders()
	if err != nil {
		return nil, err
	}
	// Save the allocations until after we've determined everything is kosher.
	e.in = make([]byte, e.size)
	e.out = make([]byte, offset+Overhead+e.size)
	e.nonce, e.key = nonceKey(nonce, key)
	return &e, nil
}

// NewEncryptor creates an Encryptor with the default chunk size.
func NewEncryptor(w io.Writer, nonce *[16]byte, key *[32]byte) *Encryptor {
	enc, _ := NewEncryptorSize(w, nonce, key, DefaultChunkSize)
	return enc
}

func (e *Encryptor) writeHeaders() error {
	_, err := e.w.Write([]byte{ver1 /* version */, 0 /* flags */})
	if err != nil {
		return err
	}
	return binary.Write(e.w, binary.LittleEndian, uint32(e.size))
}

// Writer writes an encrypted form of p to the underlying io.Writer. The
// compressed bytes are not necessarily flushed until the Encryptor is closed.
func (e *Encryptor) Write(p []byte) (n int, err error) {
	if e.err != nil {
		return 0, e.err
	}
	var m int
	for n < len(p) {
		m = copy(e.in[e.n:], p[n:])
		e.n += m
		n += m
		if e.n == e.size {
			err := e.flush()
			if err != nil {
				return n, err
			}
		}
	}
	return n, e.err
}

func (e *Encryptor) flush() error {
	if e.err != nil {
		return e.err
	}
	enc := secretbox.Seal(e.out[offset:offset], e.in[:e.n], e.nonce, e.key)
	binary.LittleEndian.PutUint32(e.out[0:], uint32(len(enc)))
	_, e.err = e.w.Write(e.out[:offset+len(enc)])
	e.n = 0
	incrCounter(e.nonce)
	return e.err
}

// Close closes the Encryptor, flushing any unwritten data to the underlying
// io.Writer but does not close the underlying io.Writer.
func (e *Encryptor) Close() (err error) {
	if e.err == ErrAlreadyClosed {
		return ErrAlreadyClosed
	}
	// Write out any pending data, mark the nonce, then write our EOF byte.
	e.flush()
	e.nonce[23] |= 0x80
	_, err = e.Write([]byte{1})
	e.flush()

	for i := range e.in {
		e.in[i] = 0
	}
	for i := range e.out {
		e.out[i] = 0
	}
	for i := range e.key {
		e.key[i] = 0
	}
	for i := range e.nonce {
		e.nonce[i] = 0
	}
	e.err = ErrAlreadyClosed
	return err
}

func incrCounter(nonce *[24]byte) {
	for i := 16; i < 24; i++ {
		nonce[i]++
		if nonce[i] != 0 {
			break
		}
	}
}

// Decryptor is an io.ReadCloser that reads encrypted data written by an
// Encryptor.
type Decryptor struct {
	r     io.Reader
	nonce *[24]byte
	key   *[32]byte
	rp    int // read position
	eb    int // end of chunk, meaning depends on part of code
	in    []byte
	out   []byte
	size  chunk // chunk size
	err   error
	next  chunk
	last  bool
}

// NewDecryptor returns a new Decryptor. Nonce and key should be identical to
// the values originally passed to NewEncryptor.
//
// Neither nonce or key are modified.
func NewDecryptor(r io.Reader, nonce *[16]byte, key *[32]byte) (*Decryptor, error) {
	d := Decryptor{r: r}
	err := d.readHeaders()
	if err != nil {
		return nil, err
	}
	d.out = make([]byte, d.size)
	d.in = make([]byte, offset+Overhead+d.size)
	d.nonce, d.key = nonceKey(nonce, key)
	return &d, nil
}

func (d *Decryptor) readHeaders() error {
	var buf [1 /* ver */ + 1 /* flags */ + 4 /* chunk */ + 4 /* next */ + 0]byte
	_, err := io.ReadFull(d.r, buf[:])
	if err != nil {
		return err
	}
	if buf[0] != ver1 {
		return errors.New("boxer: invalid version ID")
	}
	_ = buf[1] // Future: flags.
	d.size = chunk(binary.LittleEndian.Uint32(buf[2:]))
	if d.size >= math.MaxInt32 {
		return ErrChunkSize
	}
	d.next = chunk(binary.LittleEndian.Uint32(buf[6:]))
	return nil
}

// Read implements io.Reader.
func (d *Decryptor) Read(p []byte) (n int, err error) {
	if d.err != nil || len(p) == 0 {
		return 0, d.err
	}
	var m int
	for n < len(p) {
		if d.rp >= d.eb {
			d.err = d.fill()
			if d.err != nil {
				return n, d.err
			}
		}
		m = copy(p[n:], d.out[d.rp:d.eb])
		d.rp += m
		n += m
	}
	return n, nil
}

func (d *Decryptor) fill() (err error) {
	if d.err != nil {
		return d.err
	}

	d.eb, err = d.r.Read(d.in[:d.next+offset])
	if err != nil {
		return err
	}

	d.rp = 0
	d.next = chunk(binary.LittleEndian.Uint32(d.in[d.eb-offset:]))

	// The minimum read should be 18 bytes. The only time we'll
	// have less is the very end where our buffer looks like:
	// [ x x x x x x x x x x x x x x x x 1 ]
	//   |                             | |
	//   |_____________________________| |_ EOF byte
	//                  |
	//        16 bytes of authenticator
	if d.eb < Overhead+offset {
		d.last = true
		d.nonce[23] |= 0x80
	} else {
		d.eb -= offset
	}

	// If we're reading the last chunk it's okay to have an invalid next chunk.
	// It might be left over data from the previous read.
	if !d.last && (d.next <= 0 || d.next > d.size+Overhead) {
		return ErrInvalidData
	}

	m, ok := secretbox.Open(d.out[:0], d.in[:d.eb], d.nonce, d.key)
	if !ok {
		return ErrInvalidData
	}
	d.eb = len(m)
	if d.last {
		if d.out[0] != 1 {
			return ErrInvalidData
		}
		return io.EOF
	}
	incrCounter(d.nonce)
	return nil
}

// Close closes the Decryptor but does not close the underlying io.Reader.
func (d *Decryptor) Close() error {
	if d.err == ErrAlreadyClosed {
		return ErrAlreadyClosed
	}
	for i := range d.in {
		d.in[i] = 0
	}
	for i := range d.out {
		d.out[i] = 0
	}
	for i := range d.key {
		d.key[i] = 0
	}
	for i := range d.nonce {
		d.nonce[i] = 0
	}
	d.err = ErrAlreadyClosed
	return nil
}
<commit_msg>finish comment<commit_after>// Package boxer is a streaming encryption implementation, based on Adam
// Langley's article: https://www.imperialviolet.org/2014/06/27/streamingencryption.html
//
// In short, nacl/secretbox is used to seal a file in chunks, with each chunk
// being prefixed with its length. The nonce is incrementally marked so
// chunks are guaranteed to be in order. The encrypted blob is prepended with
// a header containing a version ID, the maximum chunk size, and flags. The
// flags are currently unused, but may be used in future versions.
package boxer

import (
	"encoding/binary"
	"errors"
	"io"
	"math"

	"golang.org/x/crypto/nacl/secretbox"
)

var (
	ErrAlreadyClosed = errors.New("encryptor: already closed")
	ErrInvalidData   = errors.New("decryptor: encrypted message is invalid")
	ErrChunkSize     = errors.New("boxer: invalid chunk size")
)

const (
	// DefaultChunkSize is the default maximum chunk size for reading and
	// writing.
	DefaultChunkSize = 65536

	// Overhead is the number of bytes of overhead when boxing a message.
	Overhead = secretbox.Overhead

	// offset is the number of bytes used to advise the length of the
	// chunk. It should be large enough to advise the entirety of
	// the chunk.
	offset = 4

	ver1 = 1
)

func nonceKey(nonce *[16]byte, key *[32]byte) (*[24]byte, *[32]byte) {
	var n [24]byte
	copy(n[:], nonce[:])
	var k [32]byte
	copy(k[:], key[:])
	return &n, &k
}

type chunk uint32

// Encryptor is an io.WriteCloser. Writes to an Encryptor are encrypted
// and written to w.
type Encryptor struct {
	w     io.Writer // underlying writer
	nonce *[24]byte // nacl nonce, increments per chunk
	key   *[32]byte // encryption key
	in    []byte    // input buffer
	out   []byte    // encryption buffer
	size  int       // chunk size
	n     int       // end of buffer
	err   error     // last error
}

// NewEncryptor returns a new Encryptor. Writes to the returned Encryptor
// are encrypted and written to w. The size parameter dictates the maximum
// chunk size. It should be a positive integer in the range [0, 1 << 32 - 1].
// Writes will always be chunk size + Overhead.
//
// All writes will not be flushed until Close is called. Not closing an
// Encryptor will rsult in an invalid stream.
//
// Neither nonce or key are modified.
func NewEncryptorSize(w io.Writer, nonce *[16]byte, key *[32]byte, size int) (*Encryptor, error) {
	if size > math.MaxInt32 {
		return nil, ErrChunkSize
	}
	e := Encryptor{w: w, size: size}
	err := e.writeHeaders()
	if err != nil {
		return nil, err
	}
	// Save the allocations until after we've determined everything is kosher.
	e.in = make([]byte, e.size)
	e.out = make([]byte, offset+Overhead+e.size)
	e.nonce, e.key = nonceKey(nonce, key)
	return &e, nil
}

// NewEncryptor creates an Encryptor with the default chunk size.
func NewEncryptor(w io.Writer, nonce *[16]byte, key *[32]byte) *Encryptor {
	enc, _ := NewEncryptorSize(w, nonce, key, DefaultChunkSize)
	return enc
}

func (e *Encryptor) writeHeaders() error {
	_, err := e.w.Write([]byte{ver1 /* version */, 0 /* flags */})
	if err != nil {
		return err
	}
	return binary.Write(e.w, binary.LittleEndian, uint32(e.size))
}

// Writer writes an encrypted form of p to the underlying io.Writer. The
// compressed bytes are not necessarily flushed until the Encryptor is closed.
func (e *Encryptor) Write(p []byte) (n int, err error) {
	if e.err != nil {
		return 0, e.err
	}
	var m int
	for n < len(p) {
		m = copy(e.in[e.n:], p[n:])
		e.n += m
		n += m
		if e.n == e.size {
			err := e.flush()
			if err != nil {
				return n, err
			}
		}
	}
	return n, e.err
}

func (e *Encryptor) flush() error {
	if e.err != nil {
		return e.err
	}
	enc := secretbox.Seal(e.out[offset:offset], e.in[:e.n], e.nonce, e.key)
	binary.LittleEndian.PutUint32(e.out[0:], uint32(len(enc)))
	_, e.err = e.w.Write(e.out[:offset+len(enc)])
	e.n = 0
	incrCounter(e.nonce)
	return e.err
}

// Close closes the Encryptor, flushing any unwritten data to the underlying
// io.Writer but does not close the underlying io.Writer.
func (e *Encryptor) Close() (err error) {
	if e.err == ErrAlreadyClosed {
		return ErrAlreadyClosed
	}
	// Write out any pending data, mark the nonce, then write our EOF byte.
	e.flush()
	e.nonce[23] |= 0x80
	_, err = e.Write([]byte{1})
	e.flush()

	for i := range e.in {
		e.in[i] = 0
	}
	for i := range e.out {
		e.out[i] = 0
	}
	for i := range e.key {
		e.key[i] = 0
	}
	for i := range e.nonce {
		e.nonce[i] = 0
	}
	e.err = ErrAlreadyClosed
	return err
}

func incrCounter(nonce *[24]byte) {
	for i := 16; i < 24; i++ {
		nonce[i]++
		if nonce[i] != 0 {
			break
		}
	}
}

// Decryptor is an io.ReadCloser that reads encrypted data written by an
// Encryptor.
type Decryptor struct {
	r     io.Reader
	nonce *[24]byte
	key   *[32]byte
	rp    int // read position
	eb    int // end of chunk, meaning depends on part of code
	in    []byte
	out   []byte
	size  chunk // chunk size
	err   error
	next  chunk
	last  bool
}

// NewDecryptor returns a new Decryptor. Nonce and key should be identical to
// the values originally passed to NewEncryptor.
//
// Neither nonce or key are modified.
func NewDecryptor(r io.Reader, nonce *[16]byte, key *[32]byte) (*Decryptor, error) {
	d := Decryptor{r: r}
	err := d.readHeaders()
	if err != nil {
		return nil, err
	}
	d.out = make([]byte, d.size)
	d.in = make([]byte, offset+Overhead+d.size)
	d.nonce, d.key = nonceKey(nonce, key)
	return &d, nil
}

func (d *Decryptor) readHeaders() error {
	var buf [1 /* ver */ + 1 /* flags */ + 4 /* chunk */ + 4 /* next */ + 0]byte
	_, err := io.ReadFull(d.r, buf[:])
	if err != nil {
		return err
	}
	if buf[0] != ver1 {
		return errors.New("boxer: invalid version ID")
	}
	_ = buf[1] // Future: flags.
	d.size = chunk(binary.LittleEndian.Uint32(buf[2:]))
	if d.size >= math.MaxInt32 {
		return ErrChunkSize
	}
	d.next = chunk(binary.LittleEndian.Uint32(buf[6:]))
	return nil
}

// Read implements io.Reader.
func (d *Decryptor) Read(p []byte) (n int, err error) {
	if d.err != nil || len(p) == 0 {
		return 0, d.err
	}
	var m int
	for n < len(p) {
		if d.rp >= d.eb {
			d.err = d.fill()
			if d.err != nil {
				return n, d.err
			}
		}
		m = copy(p[n:], d.out[d.rp:d.eb])
		d.rp += m
		n += m
	}
	return n, nil
}

func (d *Decryptor) fill() (err error) {
	if d.err != nil {
		return d.err
	}

	d.eb, err = d.r.Read(d.in[:d.next+offset])
	if err != nil {
		return err
	}

	d.rp = 0
	d.next = chunk(binary.LittleEndian.Uint32(d.in[d.eb-offset:]))

	// The minimum read should be 18 bytes. The only time we'll
	// have less is the very end where our buffer looks like:
	// [ x x x x x x x x x x x x x x x x 1 ]
	//   |                             | |
	//   |_____________________________| |_ EOF byte
	//                  |
	//        16 bytes of authenticator
	if d.eb < Overhead+offset {
		d.last = true
		d.nonce[23] |= 0x80
	} else {
		d.eb -= offset
	}

	// If we're reading the last chunk it's okay to have an invalid next chunk.
	// It might be left over data from the previous read.
	if !d.last && (d.next <= 0 || d.next > d.size+Overhead) {
		return ErrInvalidData
	}

	m, ok := secretbox.Open(d.out[:0], d.in[:d.eb], d.nonce, d.key)
	if !ok {
		return ErrInvalidData
	}
	d.eb = len(m)
	if d.last {
		if d.out[0] != 1 {
			return ErrInvalidData
		}
		return io.EOF
	}
	incrCounter(d.nonce)
	return nil
}

// Close closes the Decryptor but does not close the underlying io.Reader.
func (d *Decryptor) Close() error {
	if d.err == ErrAlreadyClosed {
		return ErrAlreadyClosed
	}
	for i := range d.in {
		d.in[i] = 0
	}
	for i := range d.out {
		d.out[i] = 0
	}
	for i := range d.key {
		d.key[i] = 0
	}
	for i := range d.nonce {
		d.nonce[i] = 0
	}
	d.err = ErrAlreadyClosed
	return nil
}
<|endoftext|>


<commit_before>package cryptanalysis

import(
    "strings"
)

func BreakSingleByteXor(data []byte) (float64, byte, string) {
    low := 1000.0
    msg := ""
    key := byte(0)

    // Bruteforce the key by XORing each possible key, analyzing the decrypted
    // message, and scoring it. Lowest score wins.
    for i:=0; i<256; i++ {
        k := []byte{byte(i)}
        dec := XOR(data, k)
        str := strings.ToLower(string(dec))
        score := ScoreEnglish(str)

        if score < low {
            low = score
            msg = string(dec)
            key = k[0]
        }
    }

    return low, key, msg
}
<commit_msg>Use appropriate xor function.<commit_after>package cryptanalysis

import(
    "strings"
)

func BreakSingleByteXor(data []byte) (float64, byte, string) {
    low := 1000.0
    msg := ""
    key := byte(0)

    // Bruteforce the key by XORing each possible key, analyzing the decrypted
    // message, and scoring it. Lowest score wins.
    for i:=0; i<256; i++ {
        k := byte(i)
        dec := XorArrayByte(data, k)
        str := strings.ToLower(string(dec))
        score := ScoreEnglish(str)

        if score < low {
            low = score
            msg = string(dec)
            key = k[0]
        }
    }

    return low, key, msg
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
)

var cmdBuild = &Command {
  UsageLine: "build",
  Short: "Preprocess Squirel",
  Long:"",
}

func init() {
  cmdBuild.Run = RunMe
}

func RunMe(cmd *Command, args []string) {
  fmt.Printf("BUILDING!!\n")
}<commit_msg>More work<commit_after>package main

import (
  "fmt"
  "github.com/gyuho/goraph"
)

var cmdBuild = &Command {
  UsageLine: "build",
  Short: "Preprocess Squirel",
  Long:"",
}

func init() {
  cmdBuild.Run = RunMe
}

func ProcessLineMacro()

/*
 Example build file
 {
    agent_file:"path/to/file"
    device_file:"path/to/file"
    library_dirs:[
      "libs"
    ]
 }
 */

func RunMe(cmd *Command, args []string) {
  fmt.Printf("BUILDING!!\n")
  // Load the build file

  // Start with the agent file, and device files

  // Get a list of all the #import statements
    // Then open and scan all imported files for a #library name
    //   If a file is invalid throw error
    // Make a graph of library and file imports

  // Check for acyclical graph with a topological sort
  // If loops fail
  // Else keep going

  // Now that we have a list of includes
  // Preprocess them for __FILE__, and __LINE__

}<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"koding/tools/utils"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var actions = map[string]func(){
	"shutdown": func() {
		withAll("lxc-shutdown")
	},
	"stop": func() {
		withAll("lxc-stop")
	},
	"unprepare": func() {
		for _, vm := range selectVMs(os.Args[2]) {
			err := vm.Unprepare()
			fmt.Printf("%v: %v\n", vm, err)
		}
	},
	"create-test-vms": func() {
		startIP := net.IPv4(172, 16, 0, 2)
		if len(os.Args) >= 4 {
			startIP = net.ParseIP(os.Args[3])
		}
		ipPoolFetch, _ := utils.NewIntPool(utils.IPToInt(startIP), nil)
		count, _ := strconv.Atoi(os.Args[2])
		for i := 0; i < count; i++ {
			go func(i int) {
				vm := virt.VM{
					Id: bson.NewObjectId(),
					IP: utils.IntToIP(<-ipPoolFetch),
				}
				vm.Prepare(nil)
				vm.StartCommand().Run()
				fmt.Println(i)
			}(i)
		}
	},
}

func main() {
	virt.LoadTemplates("templates")
	action := actions[os.Args[1]]
	action()
}

func selectVMs(selector string) []*virt.VM {
	if selector == "all" {
		dirs, err := ioutil.ReadDir("/var/lib/lxc")
		if err != nil {
			panic(err)
		}
		vms := make([]*virt.VM, 0)
		for _, dir := range dirs {
			if strings.HasPrefix(dir.Name(), "vm-") {
				vms = append(vms, &virt.VM{Id: bson.ObjectIdHex(dir.Name()[3:])})
			}
		}
		return vms
	}
	fmt.Println("Invalid selector: " + selector)
	os.Exit(1)
	return nil
}

func withAll(action string) {
	for _, vm := range selectVMs(os.Args[2]) {
		cmd := exec.Command(action, "-n", vm.String())
		out, err := cmd.CombinedOutput()
		fmt.Println(strings.Join(cmd.Args, " ") + ":")
		fmt.Println(err)
		fmt.Println(string(out))
	}
}
<commit_msg>virt: vmtool sync<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"koding/tools/utils"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var actions = map[string]func(){
	"shutdown": func() {
		withAll("lxc-shutdown")
	},
	"stop": func() {
		withAll("lxc-stop")
	},
	"unprepare": func() {
		for _, vm := range selectVMs(os.Args[2]) {
			err := vm.Unprepare()
			fmt.Printf("%v: %v\n", vm, err)
		}
	},
	"create-test-vms": func() {
		startIP := net.IPv4(172, 16, 0, 2)
		if len(os.Args) >= 4 {
			startIP = net.ParseIP(os.Args[3])
		}
		ipPoolFetch, _ := utils.NewIntPool(utils.IPToInt(startIP), nil)
		count, _ := strconv.Atoi(os.Args[2])
		done := make(chan int)
		for i := 0; i < count; i++ {
			go func(i int) {
				vm := virt.VM{
					Id: bson.NewObjectId(),
					IP: utils.IntToIP(<-ipPoolFetch),
				}
				vm.Prepare(nil)
				vm.StartCommand().Run()
				done <- i
			}(i)
		}
		for i := 0; i < count; i++ {
			fmt.Println(<-done)
		}
	},
}

func main() {
	virt.LoadTemplates("templates")
	action := actions[os.Args[1]]
	action()
}

func selectVMs(selector string) []*virt.VM {
	if selector == "all" {
		dirs, err := ioutil.ReadDir("/var/lib/lxc")
		if err != nil {
			panic(err)
		}
		vms := make([]*virt.VM, 0)
		for _, dir := range dirs {
			if strings.HasPrefix(dir.Name(), "vm-") {
				vms = append(vms, &virt.VM{Id: bson.ObjectIdHex(dir.Name()[3:])})
			}
		}
		return vms
	}
	fmt.Println("Invalid selector: " + selector)
	os.Exit(1)
	return nil
}

func withAll(action string) {
	for _, vm := range selectVMs(os.Args[2]) {
		cmd := exec.Command(action, "-n", vm.String())
		out, err := cmd.CombinedOutput()
		fmt.Println(strings.Join(cmd.Args, " ") + ":")
		fmt.Println(err)
		fmt.Println(string(out))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"koding/tools/utils"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var actions = map[string]func(){
	"shutdown": func() {
		withAll("lxc-shutdown")
	},
	"stop": func() {
		withAll("lxc-stop")
	},
	"unprepare": func() {
		for _, vm := range selectVMs(os.Args[2]) {
			err := vm.Unprepare()
			fmt.Printf("%v: %v\n", vm, err)
		}
	},
	"create-test-vms": func() {
		startIP := net.IPv4(172, 16, 0, 2)
		if len(os.Args) >= 4 {
			startIP = net.ParseIP(os.Args[3])
		}
		ipPoolFetch, _ := utils.NewIntPool(utils.IPToInt(startIP), nil)
		count, _ := strconv.Atoi(os.Args[2])
		done := make(chan int)
		for i := 0; i < count; i++ {
			go func(i int) {
				vm := virt.VM{
					Id: bson.NewObjectId(),
					IP: utils.IntToIP(<-ipPoolFetch),
				}
				vm.Prepare(nil)
				vm.StartCommand().Run()
				done <- i
			}(i)
		}
		for i := 0; i < count; i++ {
			fmt.Println(<-done)
		}
	},
}

func main() {
	virt.LoadTemplates("templates")
	action := actions[os.Args[1]]
	action()
}

func selectVMs(selector string) []*virt.VM {
	if selector == "all" {
		dirs, err := ioutil.ReadDir("/var/lib/lxc")
		if err != nil {
			panic(err)
		}
		vms := make([]*virt.VM, 0)
		for _, dir := range dirs {
			if strings.HasPrefix(dir.Name(), "vm-") {
				vms = append(vms, &virt.VM{Id: bson.ObjectIdHex(dir.Name()[3:])})
			}
		}
		return vms
	}
	fmt.Println("Invalid selector: " + selector)
	os.Exit(1)
	return nil
}

func withAll(action string) {
	for _, vm := range selectVMs(os.Args[2]) {
		cmd := exec.Command(action, "-n", vm.String())
		out, err := cmd.CombinedOutput()
		fmt.Println(strings.Join(cmd.Args, " ") + ":")
		fmt.Println(err)
		fmt.Println(string(out))
	}
}
<commit_msg>virt: vmtool with multiple threads.<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"koding/tools/utils"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
)

var actions = map[string]func(){
	"shutdown": func() {
		withAll("lxc-shutdown")
	},
	"stop": func() {
		withAll("lxc-stop")
	},
	"unprepare": func() {
		for _, vm := range selectVMs(os.Args[2]) {
			err := vm.Unprepare()
			fmt.Printf("%v: %v\n", vm, err)
		}
	},
	"create-test-vms": func() {
		startIP := net.IPv4(172, 16, 0, 2)
		if len(os.Args) >= 4 {
			startIP = net.ParseIP(os.Args[3])
		}
		ipPoolFetch, _ := utils.NewIntPool(utils.IPToInt(startIP), nil)
		count, _ := strconv.Atoi(os.Args[2])
		done := make(chan int)
		for i := 0; i < count; i++ {
			go func(i int) {
				vm := virt.VM{
					Id: bson.NewObjectId(),
					IP: utils.IntToIP(<-ipPoolFetch),
				}
				vm.Prepare(nil)
				vm.StartCommand().Run()
				done <- i
			}(i)
		}
		for i := 0; i < count; i++ {
			fmt.Println(<-done)
		}
	},
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	virt.LoadTemplates("templates")

	action := actions[os.Args[1]]
	action()
}

func selectVMs(selector string) []*virt.VM {
	if selector == "all" {
		dirs, err := ioutil.ReadDir("/var/lib/lxc")
		if err != nil {
			panic(err)
		}
		vms := make([]*virt.VM, 0)
		for _, dir := range dirs {
			if strings.HasPrefix(dir.Name(), "vm-") {
				vms = append(vms, &virt.VM{Id: bson.ObjectIdHex(dir.Name()[3:])})
			}
		}
		return vms
	}
	fmt.Println("Invalid selector: " + selector)
	os.Exit(1)
	return nil
}

func withAll(action string) {
	for _, vm := range selectVMs(os.Args[2]) {
		cmd := exec.Command(action, "-n", vm.String())
		out, err := cmd.CombinedOutput()
		fmt.Println(strings.Join(cmd.Args, " ") + ":")
		fmt.Println(err)
		fmt.Println(string(out))
	}
}
<|endoftext|>


<commit_before>// Code generated by "go run gen_schema.go"; DO NOT EDIT

package taskname

var SCHEMA_FROM_GIT = map[string]*Schema{
	"Build":       {Keys: []string{"os", "compiler", "target_arch", "configuration"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Calmbench":   {Keys: []string{"os", "compiler", "model", "cpu_or_gpu", "cpu_or_gpu_value", "arch", "configuration", "test_filter"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Housekeeper": {Keys: []string{"frequency"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Infra":       {Keys: []string{"frequency"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Perf":        {Keys: []string{"os", "compiler", "model", "cpu_or_gpu", "cpu_or_gpu_value", "arch", "configuration", "test_filter"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Test":        {Keys: []string{"os", "compiler", "model", "cpu_or_gpu", "cpu_or_gpu_value", "arch", "configuration", "test_filter"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Upload":      {Keys: []string(nil), OptionalKeys: []string(nil), RecurseRoles: []string{"Build", "Calmbench", "Perf", "Test"}},
}

var SEPARATOR_FROM_GIT = "-"
<commit_msg>Regenerate taskname schema<commit_after>// Code generated by "go run gen_schema.go"; DO NOT EDIT

package taskname

var SCHEMA_FROM_GIT = map[string]*Schema{
	"Build":       {Keys: []string{"os", "compiler", "target_arch", "configuration"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"BuildStats":  {Keys: []string{"os", "compiler", "target_arch", "configuration"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Calmbench":   {Keys: []string{"os", "compiler", "model", "cpu_or_gpu", "cpu_or_gpu_value", "arch", "configuration", "test_filter"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Housekeeper": {Keys: []string{"frequency"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Infra":       {Keys: []string{"frequency"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Perf":        {Keys: []string{"os", "compiler", "model", "cpu_or_gpu", "cpu_or_gpu_value", "arch", "configuration", "test_filter"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Test":        {Keys: []string{"os", "compiler", "model", "cpu_or_gpu", "cpu_or_gpu_value", "arch", "configuration", "test_filter"}, OptionalKeys: []string{"extra_config"}, RecurseRoles: []string(nil)},
	"Upload":      {Keys: []string(nil), OptionalKeys: []string(nil), RecurseRoles: []string{"Build", "BuildStats", "Calmbench", "Perf", "Test"}},
}

var SEPARATOR_FROM_GIT = "-"
<|endoftext|>


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

package tabletserver

import (
	"reflect"
	"testing"

	"github.com/youtube/vitess/go/vt/tabletserver/proto"
)

func TestComments(t *testing.T) {
	query := proto.Query{
		Sql:           "foo /*** bar */",
		BindVariables: make(map[string]interface{}),
	}
	stripTrailing(&query)
	want := proto.Query{
		Sql:           "foo",
		BindVariables: map[string]interface{}{"_trailingComment": " /*** bar */"},
	}
	if !reflect.DeepEqual(query, want) {
		t.Errorf("got \n%+v, want \n%+v", query, want)
	}
}
<commit_msg>fix typos in comments test.<commit_after>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"reflect"
	"testing"

	"github.com/youtube/vitess/go/vt/tabletserver/proto"
)

func TestComments(t *testing.T) {
	query := proto.Query{
		Sql:           "foo /*** bar */",
		BindVariables: make(map[string]interface{}),
	}
	stripTrailing(&query)
	want := proto.Query{
		Sql: "foo",
		BindVariables: map[string]interface{}{
			"_trailingComment": " /*** bar */",
		},
	}
	if !reflect.DeepEqual(query, want) {
		t.Errorf("got\n%+v, want\n%+v", query, want)
	}
}
<|endoftext|>


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

package sed

import (
	"bytes"
	"fmt"
	"os"
)

type c_cmd struct {
	addr	*address
	text	[]byte
}

func (c *c_cmd) match(line []byte, lineNumber int) bool {
	return c.addr.match(line, lineNumber)
}

func (c *c_cmd) String() string {
	if c != nil {
		if c.addr != nil {
			return fmt.Sprintf("{Append Cmd addr:%s text:%s}", c.addr.String(), c.text)
		}
		return fmt.Sprintf("{Append Cmd text:%s}", c.text)
	}
	return fmt.Sprintf("{Append Cmd}")
}

func (c *c_cmd) printText(s *Sed) {
	// we are going to get the newline from the
	fmt.Fprint(s.outputFile, string(c.text))
}

func (c *c_cmd) processLine(s *Sed) (bool, os.Error) {
  s.patternSpace = s.patternSpace[0:0]
	if c.addr != nil {
	  switch c.addr.address_type {
	    case ADDRESS_RANGE:
  	    if s.lineNumber+1 == c.addr.rangeEnd {
					c.printText(s);
					return true, nil
  	    }
	    case ADDRESS_LINE, ADDRESS_REGEX, ADDRESS_LAST_LINE:
				c.printText(s);
				return true, nil
	    case ADDRESS_TO_END_OF_FILE:
	      // FIX need to output at end of file
    }
	} else {
		c.printText(s);
		return true, nil
	}
	return false, nil
}

func NewCCmd(s *Sed, line []byte, addr *address) (*c_cmd, os.Error) {
	cmd := new(c_cmd)
	cmd.addr = addr
	cmd.text = line[1:]
	for bytes.HasSuffix(cmd.text, []byte{'\\'}) {
		cmd.text = cmd.text[0 : len(cmd.text)-1]
		line, err := s.getNextScriptLine()
		if err != nil {
			break
		}
		cmd.text = bytes.AddByte(cmd.text, '\n')
		cmd.text = bytes.Add(cmd.text, line)
	}
	return cmd, nil
}
<commit_msg>Make c work a bit more like sed<commit_after>//
//  c_cmd.go
//  sed
//
// Copyright (c) 2009 Geoffrey Clements
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

package sed

import (
	"bytes"
	"fmt"
	"os"
)

type c_cmd struct {
	addr	*address
	text	[]byte
}

func (c *c_cmd) match(line []byte, lineNumber int) bool {
	return c.addr.match(line, lineNumber)
}

func (c *c_cmd) String() string {
	if c != nil {
		if c.addr != nil {
			return fmt.Sprintf("{c command addr:%s text:%s}", c.addr.String(), c.text)
		}
		return fmt.Sprintf("{c command text:%s}", c.text)
	}
	return fmt.Sprintf("{c command}")
}

func (c *c_cmd) printText(s *Sed) {
	// we are going to get the newline from the
	s.outputFile.Write(c.text)
}

func (c *c_cmd) processLine(s *Sed) (bool, os.Error) {
	s.patternSpace = s.patternSpace[0:0]
	if c.addr != nil {
		switch c.addr.address_type {
		case ADDRESS_RANGE:
			if s.lineNumber+1 == c.addr.rangeEnd {
				c.printText(s)
				return true, nil
			}
		case ADDRESS_LINE, ADDRESS_REGEX, ADDRESS_LAST_LINE:
			c.printText(s)
			return true, nil
		case ADDRESS_TO_END_OF_FILE:
			// FIX need to output at end of file
		}
	} else {
		c.printText(s)
		return true, nil
	}
	return false, nil
}

func NewCCmd(s *Sed, line []byte, addr *address) (*c_cmd, os.Error) {
	cmd := new(c_cmd)
	cmd.addr = addr
	cmd.text = line[1:]
	for bytes.HasSuffix(cmd.text, []byte{'\\'}) {
		cmd.text = cmd.text[0 : len(cmd.text)-1]
		line, err := s.getNextScriptLine()
		if err != nil {
			break
		}
		cmd.text = bytes.AddByte(cmd.text, '\n')
		cmd.text = bytes.Add(cmd.text, line)
	}
	return cmd, nil
}
<|endoftext|>


<commit_before>package pilosa

import (
	"bytes"
	"fmt"
	"io"
	"sort"
	"sync"
	"time"

	"github.com/golang/groupcache/lru"
	"github.com/pilosa/pilosa/internal"
)

// Cache represents a cache for bitmap counts.
type Cache interface {
	Add(bitmapID uint64, n uint64)
	BulkAdd(bitmapID uint64, n uint64)
	Get(bitmapID uint64) uint64
	Len() int

	// Returns a list of all bitmap IDs.
	BitmapIDs() []uint64

	// Updates the cache, if necessary.
	Invalidate()

	// Returns an ordered list of the top ranked bitmaps.
	Top() []BitmapPair
}

// LRUCache represents a least recently used Cache implemenation.
type LRUCache struct {
	cache  *lru.Cache
	counts map[uint64]uint64
}

// NewLRUCache returns a new instance of LRUCache.
func NewLRUCache(maxEntries int) *LRUCache {
	c := &LRUCache{
		cache:  lru.New(maxEntries),
		counts: make(map[uint64]uint64),
	}
	c.cache.OnEvicted = c.onEvicted
	return c
}

func (c *LRUCache) BulkAdd(bitmapID, n uint64) {
	c.Add(bitmapID, n)
}

// Add adds a bitmap to the cache.
func (c *LRUCache) Add(bitmapID, n uint64) {
	c.cache.Add(bitmapID, n)
	c.counts[bitmapID] = n
}

// Get returns a bitmap with a given id.
func (c *LRUCache) Get(bitmapID uint64) uint64 {
	n, _ := c.cache.Get(bitmapID)
	nn, _ := n.(uint64)
	return nn
}

// Len returns the number of items in the cache.
func (c *LRUCache) Len() int { return c.cache.Len() }

// Invalidate is a no-op.
func (c *LRUCache) Invalidate() {}

// BitmapIDs returns a list of all bitmap IDs in the cache.
func (c *LRUCache) BitmapIDs() []uint64 {
	a := make([]uint64, 0, len(c.counts))
	for id := range c.counts {
		a = append(a, id)
	}
	sort.Sort(uint64Slice(a))
	return a
}

// Top returns all counts in the cache.
func (c *LRUCache) Top() []BitmapPair {
	a := make([]BitmapPair, 0, len(c.counts))
	for id, n := range c.counts {
		a = append(a, BitmapPair{
			ID:    id,
			Count: uint64(n),
		})
	}
	sort.Sort(BitmapPairs(a))
	return a
}

func (c *LRUCache) onEvicted(key lru.Key, _ interface{}) { delete(c.counts, key.(uint64)) }

// Ensure LRUCache implements Cache.
var _ Cache = &LRUCache{}

// RankCache represents a cache with sorted entries.
type RankCache struct {
	mu       sync.Mutex
	entries  map[uint64]uint64
	rankings []BitmapPair // cached, ordered list

	updateN    int
	updateTime time.Time

	ThresholdLength int
	ThresholdIndex  int
	ThresholdValue  uint64
}

// NewRankCache returns a new instance of RankCache.
func NewRankCache() *RankCache {
	return &RankCache{
		entries: make(map[uint64]uint64),
	}
}

// Add adds a bitmap to the cache.
func (c *RankCache) Add(bitmapID uint64, n uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	// Ignore if the bit count on the bitmap is below the threshold.
	if n < c.ThresholdValue {
		return
	}

	c.entries[bitmapID] = n

	c.invalidate()
}

// BulkAdd adds a bitmap to the cache unsorted. You should Invalidate after completion.
func (c *RankCache) BulkAdd(bitmapID uint64, n uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if n < c.ThresholdValue {
		return
	}

	c.entries[bitmapID] = n
}

// Get returns a bitmap with a given id.
func (c *RankCache) Get(bitmapID uint64) uint64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.entries[bitmapID]
}

// Len returns the number of items in the cache.
func (c *RankCache) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.entries)
}

// BitmapIDs returns a list of all bitmap IDs in the cache.
func (c *RankCache) BitmapIDs() []uint64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	a := make([]uint64, 0, len(c.entries))
	for id := range c.entries {
		a = append(a, id)
	}
	sort.Sort(uint64Slice(a))
	return a
}

// update reorders the entries by rank.
func (c *RankCache) Invalidate() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.invalidate()

}
func (c *RankCache) invalidate() {
	// Don't invalidate more than once every X seconds.
	// TODO: consider making this configurable.
	if time.Now().Sub(c.updateTime).Seconds() < 10 {
		return
	}
	// Convert cache to a sorted list.
	rankings := make([]BitmapPair, 0, len(c.entries))
	for id, cnt := range c.entries {
		rankings = append(rankings, BitmapPair{
			ID:    id,
			Count: cnt,
		})
	}
	sort.Sort(BitmapPairs(rankings))

	// Store the count of the item at the threshold index.
	c.rankings = rankings
	if len(c.rankings) > c.ThresholdIndex {
		c.ThresholdValue = rankings[c.ThresholdIndex].Count
	} else {
		c.ThresholdValue = 1
	}

	// Reset counters.
	c.updateTime, c.updateN = time.Now(), 0

	// If size is larger than the threshold then trim it.
	if len(c.entries) > c.ThresholdLength {
		for id, cnt := range c.entries {
			if cnt <= c.ThresholdValue {
				delete(c.entries, id)
			}
		}
	}
}

// Top returns an ordered list of bitmaps.
func (c *RankCache) Top() []BitmapPair { return c.rankings }

// WriteTo writes the cache to w.
func (c *RankCache) WriteTo(w io.Writer) (n int64, err error) {
	panic("FIXME: TODO")
}

// ReadFrom read from r into the cache.
func (c *RankCache) ReadFrom(r io.Reader) (n int64, err error) {
	panic("FIXME: TODO")
}

// Ensure RankCache implements Cache.
var _ Cache = &RankCache{}

// BitmapPair represents a bitmap with an associated identifier.
type BitmapPair struct {
	ID    uint64
	Count uint64
}

// BitmapPairs is a sortable list of BitmapPair objects.
type BitmapPairs []BitmapPair

func (p BitmapPairs) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p BitmapPairs) Len() int           { return len(p) }
func (p BitmapPairs) Less(i, j int) bool { return p[i].Count > p[j].Count }

type Pair struct {
	Key   uint64 `json:"key"`
	Count uint64 `json:"count"`
}

func encodePair(p Pair) *internal.Pair {
	return &internal.Pair{
		Key:   p.Key,
		Count: p.Count,
	}
}

func decodePair(pb *internal.Pair) Pair {
	return Pair{
		Key:   pb.Key,
		Count: pb.Count,
	}
}

type Pairs []Pair

func (p Pairs) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p Pairs) Len() int           { return len(p) }
func (p Pairs) Less(i, j int) bool { return p[i].Count > p[j].Count }

type PairHeap struct {
	Pairs
}

func (p PairHeap) Less(i, j int) bool { return p.Pairs[i].Count < p.Pairs[j].Count }

func (h *Pairs) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(Pair))
}

func (h *Pairs) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// Add merges other into p and returns a new slice.
func (p Pairs) Add(other []Pair) []Pair {
	// Create lookup of key/counts.
	m := make(map[uint64]uint64, len(p))
	for _, pair := range p {
		m[pair.Key] = pair.Count
	}

	// Add/merge from other.
	for _, pair := range other {
		m[pair.Key] += pair.Count
	}

	// Convert back to slice.
	a := make([]Pair, 0, len(m))
	for k, v := range m {
		a = append(a, Pair{Key: k, Count: v})
	}
	return a
}

// Keys returns a slice of all keys in p.
func (p Pairs) Keys() []uint64 {
	a := make([]uint64, len(p))
	for i := range p {
		a[i] = p[i].Key
	}
	return a
}

func (p Pairs) String() string {
	var buf bytes.Buffer
	buf.WriteString("Pairs(")
	for i := range p {
		fmt.Fprintf(&buf, "%d/%d", p[i].Key, p[i].Count)
		if i < len(p)-1 {
			buf.WriteString(", ")
		}
	}
	buf.WriteString(")")
	return buf.String()
}

func encodePairs(a Pairs) []*internal.Pair {
	other := make([]*internal.Pair, len(a))
	for i := range a {
		other[i] = encodePair(a[i])
	}
	return other
}

func decodePairs(a []*internal.Pair) []Pair {
	other := make([]Pair, len(a))
	for i := range a {
		other[i] = decodePair(a[i])
	}
	return other
}

// uint64Slice represents a sortable slice of uint64 numbers.
type uint64Slice []uint64

func (p uint64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p uint64Slice) Len() int           { return len(p) }
func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }

// merge combines p and other to a unique sorted set of values.
// p and other must both have unique sets and be sorted.
func (p uint64Slice) merge(other []uint64) []uint64 {
	ret := make([]uint64, 0, len(p))

	i, j := 0, 0
	for i < len(p) && j < len(other) {
		a, b := p[i], other[j]
		if a == b {
			ret = append(ret, a)
			i, j = i+1, j+1
		} else if a < b {
			ret = append(ret, a)
			i++
		} else {
			ret = append(ret, b)
			j++
		}
	}

	if i < len(p) {
		ret = append(ret, p[i:]...)
	} else if j < len(other) {
		ret = append(ret, other[j:]...)
	}

	return ret
}

// BitmapCache provides an interface for caching full bitmaps.
type BitmapCache interface {
	Fetch(id uint64) (*Bitmap, bool)
	Add(id uint64, b *Bitmap)
}

// SimpleCache implements BitmapCache
// it is meant to be a short-lived cache for cases where writes are continuing to access
// the same bit within a short time frame (i.e. good for write-heavy loads)
// A read-heavy use case would cause the cache to get bigger, potentially causing the
// node to run out of memory.
type SimpleCache struct {
	cache map[uint64]*Bitmap
}

func (s *SimpleCache) Fetch(id uint64) (*Bitmap, bool) {
	return s.cache[id]
}

func (s *SimpleCache) Add(id uint64, b *Bitmap) {
	s.cache[id] = b
}
<commit_msg>bug fix in SimpleCache.Fetch (really just reverting my change)<commit_after>package pilosa

import (
	"bytes"
	"fmt"
	"io"
	"sort"
	"sync"
	"time"

	"github.com/golang/groupcache/lru"
	"github.com/pilosa/pilosa/internal"
)

// Cache represents a cache for bitmap counts.
type Cache interface {
	Add(bitmapID uint64, n uint64)
	BulkAdd(bitmapID uint64, n uint64)
	Get(bitmapID uint64) uint64
	Len() int

	// Returns a list of all bitmap IDs.
	BitmapIDs() []uint64

	// Updates the cache, if necessary.
	Invalidate()

	// Returns an ordered list of the top ranked bitmaps.
	Top() []BitmapPair
}

// LRUCache represents a least recently used Cache implemenation.
type LRUCache struct {
	cache  *lru.Cache
	counts map[uint64]uint64
}

// NewLRUCache returns a new instance of LRUCache.
func NewLRUCache(maxEntries int) *LRUCache {
	c := &LRUCache{
		cache:  lru.New(maxEntries),
		counts: make(map[uint64]uint64),
	}
	c.cache.OnEvicted = c.onEvicted
	return c
}

func (c *LRUCache) BulkAdd(bitmapID, n uint64) {
	c.Add(bitmapID, n)
}

// Add adds a bitmap to the cache.
func (c *LRUCache) Add(bitmapID, n uint64) {
	c.cache.Add(bitmapID, n)
	c.counts[bitmapID] = n
}

// Get returns a bitmap with a given id.
func (c *LRUCache) Get(bitmapID uint64) uint64 {
	n, _ := c.cache.Get(bitmapID)
	nn, _ := n.(uint64)
	return nn
}

// Len returns the number of items in the cache.
func (c *LRUCache) Len() int { return c.cache.Len() }

// Invalidate is a no-op.
func (c *LRUCache) Invalidate() {}

// BitmapIDs returns a list of all bitmap IDs in the cache.
func (c *LRUCache) BitmapIDs() []uint64 {
	a := make([]uint64, 0, len(c.counts))
	for id := range c.counts {
		a = append(a, id)
	}
	sort.Sort(uint64Slice(a))
	return a
}

// Top returns all counts in the cache.
func (c *LRUCache) Top() []BitmapPair {
	a := make([]BitmapPair, 0, len(c.counts))
	for id, n := range c.counts {
		a = append(a, BitmapPair{
			ID:    id,
			Count: uint64(n),
		})
	}
	sort.Sort(BitmapPairs(a))
	return a
}

func (c *LRUCache) onEvicted(key lru.Key, _ interface{}) { delete(c.counts, key.(uint64)) }

// Ensure LRUCache implements Cache.
var _ Cache = &LRUCache{}

// RankCache represents a cache with sorted entries.
type RankCache struct {
	mu       sync.Mutex
	entries  map[uint64]uint64
	rankings []BitmapPair // cached, ordered list

	updateN    int
	updateTime time.Time

	ThresholdLength int
	ThresholdIndex  int
	ThresholdValue  uint64
}

// NewRankCache returns a new instance of RankCache.
func NewRankCache() *RankCache {
	return &RankCache{
		entries: make(map[uint64]uint64),
	}
}

// Add adds a bitmap to the cache.
func (c *RankCache) Add(bitmapID uint64, n uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	// Ignore if the bit count on the bitmap is below the threshold.
	if n < c.ThresholdValue {
		return
	}

	c.entries[bitmapID] = n

	c.invalidate()
}

// BulkAdd adds a bitmap to the cache unsorted. You should Invalidate after completion.
func (c *RankCache) BulkAdd(bitmapID uint64, n uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if n < c.ThresholdValue {
		return
	}

	c.entries[bitmapID] = n
}

// Get returns a bitmap with a given id.
func (c *RankCache) Get(bitmapID uint64) uint64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.entries[bitmapID]
}

// Len returns the number of items in the cache.
func (c *RankCache) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.entries)
}

// BitmapIDs returns a list of all bitmap IDs in the cache.
func (c *RankCache) BitmapIDs() []uint64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	a := make([]uint64, 0, len(c.entries))
	for id := range c.entries {
		a = append(a, id)
	}
	sort.Sort(uint64Slice(a))
	return a
}

// update reorders the entries by rank.
func (c *RankCache) Invalidate() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.invalidate()

}
func (c *RankCache) invalidate() {
	// Don't invalidate more than once every X seconds.
	// TODO: consider making this configurable.
	if time.Now().Sub(c.updateTime).Seconds() < 10 {
		return
	}
	// Convert cache to a sorted list.
	rankings := make([]BitmapPair, 0, len(c.entries))
	for id, cnt := range c.entries {
		rankings = append(rankings, BitmapPair{
			ID:    id,
			Count: cnt,
		})
	}
	sort.Sort(BitmapPairs(rankings))

	// Store the count of the item at the threshold index.
	c.rankings = rankings
	if len(c.rankings) > c.ThresholdIndex {
		c.ThresholdValue = rankings[c.ThresholdIndex].Count
	} else {
		c.ThresholdValue = 1
	}

	// Reset counters.
	c.updateTime, c.updateN = time.Now(), 0

	// If size is larger than the threshold then trim it.
	if len(c.entries) > c.ThresholdLength {
		for id, cnt := range c.entries {
			if cnt <= c.ThresholdValue {
				delete(c.entries, id)
			}
		}
	}
}

// Top returns an ordered list of bitmaps.
func (c *RankCache) Top() []BitmapPair { return c.rankings }

// WriteTo writes the cache to w.
func (c *RankCache) WriteTo(w io.Writer) (n int64, err error) {
	panic("FIXME: TODO")
}

// ReadFrom read from r into the cache.
func (c *RankCache) ReadFrom(r io.Reader) (n int64, err error) {
	panic("FIXME: TODO")
}

// Ensure RankCache implements Cache.
var _ Cache = &RankCache{}

// BitmapPair represents a bitmap with an associated identifier.
type BitmapPair struct {
	ID    uint64
	Count uint64
}

// BitmapPairs is a sortable list of BitmapPair objects.
type BitmapPairs []BitmapPair

func (p BitmapPairs) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p BitmapPairs) Len() int           { return len(p) }
func (p BitmapPairs) Less(i, j int) bool { return p[i].Count > p[j].Count }

type Pair struct {
	Key   uint64 `json:"key"`
	Count uint64 `json:"count"`
}

func encodePair(p Pair) *internal.Pair {
	return &internal.Pair{
		Key:   p.Key,
		Count: p.Count,
	}
}

func decodePair(pb *internal.Pair) Pair {
	return Pair{
		Key:   pb.Key,
		Count: pb.Count,
	}
}

type Pairs []Pair

func (p Pairs) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p Pairs) Len() int           { return len(p) }
func (p Pairs) Less(i, j int) bool { return p[i].Count > p[j].Count }

type PairHeap struct {
	Pairs
}

func (p PairHeap) Less(i, j int) bool { return p.Pairs[i].Count < p.Pairs[j].Count }

func (h *Pairs) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(Pair))
}

func (h *Pairs) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// Add merges other into p and returns a new slice.
func (p Pairs) Add(other []Pair) []Pair {
	// Create lookup of key/counts.
	m := make(map[uint64]uint64, len(p))
	for _, pair := range p {
		m[pair.Key] = pair.Count
	}

	// Add/merge from other.
	for _, pair := range other {
		m[pair.Key] += pair.Count
	}

	// Convert back to slice.
	a := make([]Pair, 0, len(m))
	for k, v := range m {
		a = append(a, Pair{Key: k, Count: v})
	}
	return a
}

// Keys returns a slice of all keys in p.
func (p Pairs) Keys() []uint64 {
	a := make([]uint64, len(p))
	for i := range p {
		a[i] = p[i].Key
	}
	return a
}

func (p Pairs) String() string {
	var buf bytes.Buffer
	buf.WriteString("Pairs(")
	for i := range p {
		fmt.Fprintf(&buf, "%d/%d", p[i].Key, p[i].Count)
		if i < len(p)-1 {
			buf.WriteString(", ")
		}
	}
	buf.WriteString(")")
	return buf.String()
}

func encodePairs(a Pairs) []*internal.Pair {
	other := make([]*internal.Pair, len(a))
	for i := range a {
		other[i] = encodePair(a[i])
	}
	return other
}

func decodePairs(a []*internal.Pair) []Pair {
	other := make([]Pair, len(a))
	for i := range a {
		other[i] = decodePair(a[i])
	}
	return other
}

// uint64Slice represents a sortable slice of uint64 numbers.
type uint64Slice []uint64

func (p uint64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p uint64Slice) Len() int           { return len(p) }
func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }

// merge combines p and other to a unique sorted set of values.
// p and other must both have unique sets and be sorted.
func (p uint64Slice) merge(other []uint64) []uint64 {
	ret := make([]uint64, 0, len(p))

	i, j := 0, 0
	for i < len(p) && j < len(other) {
		a, b := p[i], other[j]
		if a == b {
			ret = append(ret, a)
			i, j = i+1, j+1
		} else if a < b {
			ret = append(ret, a)
			i++
		} else {
			ret = append(ret, b)
			j++
		}
	}

	if i < len(p) {
		ret = append(ret, p[i:]...)
	} else if j < len(other) {
		ret = append(ret, other[j:]...)
	}

	return ret
}

// BitmapCache provides an interface for caching full bitmaps.
type BitmapCache interface {
	Fetch(id uint64) (*Bitmap, bool)
	Add(id uint64, b *Bitmap)
}

// SimpleCache implements BitmapCache
// it is meant to be a short-lived cache for cases where writes are continuing to access
// the same bit within a short time frame (i.e. good for write-heavy loads)
// A read-heavy use case would cause the cache to get bigger, potentially causing the
// node to run out of memory.
type SimpleCache struct {
	cache map[uint64]*Bitmap
}

func (s *SimpleCache) Fetch(id uint64) (*Bitmap, bool) {
	m, ok := s.cache[id]
	return m, ok
}

func (s *SimpleCache) Add(id uint64, b *Bitmap) {
	s.cache[id] = b
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus 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 main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"

	"github.com/prometheus/node_exporter/collector"
)

const subsystem = "exporter"

var (
	// Version of node_exporter. Set at build time.
	Version = "0.0.0.dev"

	memProfile        = flag.String("debug.memprofile-file", "", "Write memory profile to this file upon receipt of SIGUSR1.")
	listenAddress     = flag.String("web.listen-address", ":9100", "Address on which to expose metrics and web interface.")
	metricsPath       = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
	enabledCollectors = flag.String("collectors.enabled", "diskstats,filefd,filesystem,loadavg,mdadm,meminfo,netdev,netstat,sockstat,stat,textfile,time,uname", "Comma-separated list of collectors to use.")
	printCollectors   = flag.Bool("collectors.print", false, "If true, print available collectors and exit.")

	collectorLabelNames = []string{"collector", "result"}

	scrapeDurations = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace: collector.Namespace,
			Subsystem: subsystem,
			Name:      "scrape_duration_seconds",
			Help:      "node_exporter: Duration of a scrape job.",
		},
		collectorLabelNames,
	)
)

// NodeCollector implements the prometheus.Collector interface.
type NodeCollector struct {
	collectors map[string]collector.Collector
}

// Describe implements the prometheus.Collector interface.
func (n NodeCollector) Describe(ch chan<- *prometheus.Desc) {
	scrapeDurations.Describe(ch)
}

// Collect implements the prometheus.Collector interface.
func (n NodeCollector) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(n.collectors))
	for name, c := range n.collectors {
		go func(name string, c collector.Collector) {
			execute(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
	scrapeDurations.Collect(ch)
}

func execute(name string, c collector.Collector, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Update(ch)
	duration := time.Since(begin)
	var result string

	if err != nil {
		log.Errorf("ERROR: %s collector failed after %fs: %s", name, duration.Seconds(), err)
		result = "error"
	} else {
		log.Debugf("OK: %s collector succeeded after %fs.", name, duration.Seconds())
		result = "success"
	}
	scrapeDurations.WithLabelValues(name, result).Observe(duration.Seconds())
}

func loadCollectors() (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	for _, name := range strings.Split(*enabledCollectors, ",") {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func main() {
	flag.Parse()

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}
	collectors, err := loadCollectors()
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors:")
	for n := range collectors {
		log.Infof(" - %s", n)
	}

	nodeCollector := NodeCollector{collectors: collectors}
	prometheus.MustRegister(nodeCollector)

	sigUsr1 := make(chan os.Signal)
	signal.Notify(sigUsr1, syscall.SIGUSR1)

	handler := prometheus.Handler()

	http.Handle(*metricsPath, handler)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Node Exporter</title>

Node Exporter

<a href="` + *metricsPath + `">Metrics

			`))
	})

	log.Infof("Starting node_exporter v%s at %s", Version, *listenAddress)
	err = http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Filter list of collectors enabled by default<commit_after>// Copyright 2015 The Prometheus 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 main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"

	"github.com/prometheus/node_exporter/collector"
)

const subsystem = "exporter"

var (
	// Version of node_exporter. Set at build time.
	Version = "0.0.0.dev"

	memProfile        = flag.String("debug.memprofile-file", "", "Write memory profile to this file upon receipt of SIGUSR1.")
	listenAddress     = flag.String("web.listen-address", ":9100", "Address on which to expose metrics and web interface.")
	metricsPath       = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
	enabledCollectors = flag.String("collectors.enabled",
		filterAvailableCollectors("diskstats,filefd,filesystem,loadavg,mdadm,meminfo,netdev,netstat,sockstat,stat,textfile,time,uname"),
		"Comma-separated list of collectors to use.")
	printCollectors = flag.Bool("collectors.print", false, "If true, print available collectors and exit.")

	collectorLabelNames = []string{"collector", "result"}

	scrapeDurations = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace: collector.Namespace,
			Subsystem: subsystem,
			Name:      "scrape_duration_seconds",
			Help:      "node_exporter: Duration of a scrape job.",
		},
		collectorLabelNames,
	)
)

// NodeCollector implements the prometheus.Collector interface.
type NodeCollector struct {
	collectors map[string]collector.Collector
}

// Describe implements the prometheus.Collector interface.
func (n NodeCollector) Describe(ch chan<- *prometheus.Desc) {
	scrapeDurations.Describe(ch)
}

// Collect implements the prometheus.Collector interface.
func (n NodeCollector) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(n.collectors))
	for name, c := range n.collectors {
		go func(name string, c collector.Collector) {
			execute(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
	scrapeDurations.Collect(ch)
}

func filterAvailableCollectors(collectors string) string {
	availableCollectors := make([]string, 0)
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Update(ch)
	duration := time.Since(begin)
	var result string

	if err != nil {
		log.Errorf("ERROR: %s collector failed after %fs: %s", name, duration.Seconds(), err)
		result = "error"
	} else {
		log.Debugf("OK: %s collector succeeded after %fs.", name, duration.Seconds())
		result = "success"
	}
	scrapeDurations.WithLabelValues(name, result).Observe(duration.Seconds())
}

func loadCollectors() (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	for _, name := range strings.Split(*enabledCollectors, ",") {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func main() {
	flag.Parse()

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}
	collectors, err := loadCollectors()
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors:")
	for n := range collectors {
		log.Infof(" - %s", n)
	}

	nodeCollector := NodeCollector{collectors: collectors}
	prometheus.MustRegister(nodeCollector)

	sigUsr1 := make(chan os.Signal)
	signal.Notify(sigUsr1, syscall.SIGUSR1)

	handler := prometheus.Handler()

	http.Handle(*metricsPath, handler)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Node Exporter</title>

Node Exporter

<a href="` + *metricsPath + `">Metrics

			`))
	})

	log.Infof("Starting node_exporter v%s at %s", Version, *listenAddress)
	err = http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>