<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package leveldb

import (
	"bytes"

	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/errors"
	"github.com/syndtr/goleveldb/leveldb/filter"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/util"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/utils"
	"github.com/ava-labs/avalanchego/utils/logging"
)

const (
	// minBlockCacheSize is the minimum number of bytes to use for block caching
	// in leveldb.
	minBlockCacheSize = 12 * opt.MiB

	// minWriteBufferSize is the minimum number of bytes to use for buffers in
	// leveldb.
	minWriteBufferSize = 12 * opt.MiB

	// minHandleCap is the minimum number of files descriptors to cap levelDB to
	// use
	minHandleCap = 64

	// levelDBByteOverhead is the number of bytes of constant overhead that
	// should be added to a batch size per operation.
	levelDBByteOverhead = 8
)

// Database is a persistent key-value store. Apart from basic data storage
// functionality it also supports batch writes and iterating over the keyspace
// in binary-alphabetical order.
type Database struct {
	*leveldb.DB
	log logging.Logger

	// True if there was previously an error other than "not found" or "closed"
	// while performing a db operation. If [errored] == true, Has, Get, Put,
	// Delete and batch writes fail with ErrAvoidCorruption.
	// The node should shut down.
	errored bool
}

// New returns a wrapped LevelDB object.
func New(file string, log logging.Logger, blockCacheSize, writeBufferSize, handleCap int) (*Database, error) {
	// Enforce minimums
	if blockCacheSize < minBlockCacheSize {
		blockCacheSize = minBlockCacheSize
	}
	if writeBufferSize < minWriteBufferSize {
		writeBufferSize = minWriteBufferSize
	}
	if handleCap < minHandleCap {
		handleCap = minHandleCap
	}

	// Open the db and recover any potential corruptions
	db, err := leveldb.OpenFile(file, &opt.Options{
		OpenFilesCacheCapacity: handleCap,
		BlockCacheCapacity:     blockCacheSize,
		// There are two buffers of size WriteBuffer used.
		WriteBuffer: writeBufferSize / 2,
		Filter:      filter.NewBloomFilter(10),
	})
	if _, corrupted := err.(*errors.ErrCorrupted); corrupted {
		db, err = leveldb.RecoverFile(file, nil)
	}
	if err != nil {
		return nil, err
	}
	return &Database{
		DB:  db,
		log: log,
	}, nil
}

// Has returns if the key is set in the database
func (db *Database) Has(key []byte) (bool, error) {
	if db.errored {
		return false, database.ErrAvoidCorruption
	}
	has, err := db.DB.Has(key, nil)
	return has, db.handleError(err)
}

// Get returns the value the key maps to in the database
func (db *Database) Get(key []byte) ([]byte, error) {
	if db.errored {
		return nil, database.ErrAvoidCorruption
	}
	value, err := db.DB.Get(key, nil)
	return value, db.handleError(err)
}

// Put sets the value of the provided key to the provided value
func (db *Database) Put(key []byte, value []byte) error {
	if db.errored {
		return database.ErrAvoidCorruption
	}
	return db.handleError(db.DB.Put(key, value, nil))
}

// Delete removes the key from the database
func (db *Database) Delete(key []byte) error {
	if db.errored {
		return database.ErrAvoidCorruption
	}
	return db.handleError(db.DB.Delete(key, nil))
}

// NewBatch creates a write/delete-only buffer that is atomically committed to
// the database when write is called
func (db *Database) NewBatch() database.Batch { return &batch{db: db} }

// NewIterator creates a lexicographically ordered iterator over the database
func (db *Database) NewIterator() database.Iterator {
	return &iter{db.DB.NewIterator(new(util.Range), nil)}
}

// NewIteratorWithStart creates a lexicographically ordered iterator over the
// database starting at the provided key
func (db *Database) NewIteratorWithStart(start []byte) database.Iterator {
	return &iter{db.DB.NewIterator(&util.Range{Start: start}, nil)}
}

// NewIteratorWithPrefix creates a lexicographically ordered iterator over the
// database ignoring keys that do not start with the provided prefix
func (db *Database) NewIteratorWithPrefix(prefix []byte) database.Iterator {
	return &iter{db.DB.NewIterator(util.BytesPrefix(prefix), nil)}
}

// NewIteratorWithStartAndPrefix creates a lexicographically ordered iterator
// over the database starting at start and ignoring keys that do not start with
// the provided prefix
func (db *Database) NewIteratorWithStartAndPrefix(start, prefix []byte) database.Iterator {
	iterRange := util.BytesPrefix(prefix)
	if bytes.Compare(start, prefix) == 1 {
		iterRange.Start = start
	}
	return &iter{db.DB.NewIterator(iterRange, nil)}
}

// Stat returns a particular internal stat of the database.
func (db *Database) Stat(property string) (string, error) {
	stat, err := db.DB.GetProperty(property)
	return stat, db.handleError(err)
}

// This comment is basically copy pasted from the underlying levelDB library:

// Compact the underlying DB for the given key range.
// Specifically, deleted and overwritten versions are discarded,
// and the data is rearranged to reduce the cost of operations
// needed to access the data. This operation should typically only
// be invoked by users who understand the underlying implementation.
//
// A nil start is treated as a key before all keys in the DB.
// And a nil limit is treated as a key after all keys in the DB.
// Therefore if both are nil then it will compact entire DB.
func (db *Database) Compact(start []byte, limit []byte) error {
	return db.handleError(db.DB.CompactRange(util.Range{Start: start, Limit: limit}))
}

// Close implements the Database interface
func (db *Database) Close() error { return db.handleError(db.DB.Close()) }

func (db *Database) handleError(err error) error {
	err = updateError(err)
	// If we get an error other than "not found" or "closed", disallow future
	// database operations to avoid possible corruption
	if err != nil && err != database.ErrNotFound && err != database.ErrClosed {
		db.log.Fatal("leveldb error: %w", err)
		db.errored = true
	}
	return err
}

// batch is a wrapper around a levelDB batch to contain sizes.
type batch struct {
	leveldb.Batch
	db   *Database
	size int
}

// Put the value into the batch for later writing
func (b *batch) Put(key, value []byte) error {
	b.Batch.Put(key, value)
	b.size += len(key) + len(value) + levelDBByteOverhead
	return nil
}

// Delete the key during writing
func (b *batch) Delete(key []byte) error {
	b.Batch.Delete(key)
	b.size += len(key) + levelDBByteOverhead
	return nil
}

// Size retrieves the amount of data queued up for writing.
func (b *batch) Size() int { return b.size }

// Write flushes any accumulated data to disk.
func (b *batch) Write() error {
	if b.db.errored {
		return database.ErrAvoidCorruption
	}
	return b.db.handleError(b.db.DB.Write(&b.Batch, nil))
}

// Reset resets the batch for reuse.
func (b *batch) Reset() {
	b.Batch.Reset()
	b.size = 0
}

// Replay the batch contents.
func (b *batch) Replay(w database.KeyValueWriter) error {
	replay := &replayer{writer: w}
	if err := b.Batch.Replay(replay); err != nil {
		// Never actually returns an error, because Replay just returns nil
		return b.db.handleError(err)
	}
	return b.db.handleError(replay.err)
}

// Inner returns itself
func (b *batch) Inner() database.Batch { return b }

type replayer struct {
	writer database.KeyValueWriter
	err    error
}

func (r *replayer) Put(key, value []byte) {
	if r.err != nil {
		return
	}
	r.err = r.writer.Put(key, value)
}

func (r *replayer) Delete(key []byte) {
	if r.err != nil {
		return
	}
	r.err = r.writer.Delete(key)
}

type iter struct{ iterator.Iterator }

// Error implements the Iterator interface
func (it *iter) Error() error { return updateError(it.Iterator.Error()) }

// Key implements the Iterator interface
func (it *iter) Key() []byte { return utils.CopyBytes(it.Iterator.Key()) }

// Value implements the Iterator interface
func (it *iter) Value() []byte { return utils.CopyBytes(it.Iterator.Value()) }

func updateError(err error) error {
	switch err {
	case leveldb.ErrClosed:
		return database.ErrClosed
	case leveldb.ErrNotFound:
		return database.ErrNotFound
	default:
		return err
	}
}
<commit_msg>fixed potential leveldb shutdown race<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package leveldb

import (
	"bytes"
	"sync/atomic"

	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/errors"
	"github.com/syndtr/goleveldb/leveldb/filter"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/util"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/utils"
	"github.com/ava-labs/avalanchego/utils/logging"
)

const (
	// minBlockCacheSize is the minimum number of bytes to use for block caching
	// in leveldb.
	minBlockCacheSize = 12 * opt.MiB

	// minWriteBufferSize is the minimum number of bytes to use for buffers in
	// leveldb.
	minWriteBufferSize = 12 * opt.MiB

	// minHandleCap is the minimum number of files descriptors to cap levelDB to
	// use
	minHandleCap = 64

	// levelDBByteOverhead is the number of bytes of constant overhead that
	// should be added to a batch size per operation.
	levelDBByteOverhead = 8
)

// Database is a persistent key-value store. Apart from basic data storage
// functionality it also supports batch writes and iterating over the keyspace
// in binary-alphabetical order.
type Database struct {
	*leveldb.DB
	log logging.Logger

	// 1 if there was previously an error other than "not found" or "closed"
	// while performing a db operation. If [errored] == 1, Has, Get, Put,
	// Delete and batch writes fail with ErrAvoidCorruption.
	errored uint64
}

// New returns a wrapped LevelDB object.
func New(file string, log logging.Logger, blockCacheSize, writeBufferSize, handleCap int) (*Database, error) {
	// Enforce minimums
	if blockCacheSize < minBlockCacheSize {
		blockCacheSize = minBlockCacheSize
	}
	if writeBufferSize < minWriteBufferSize {
		writeBufferSize = minWriteBufferSize
	}
	if handleCap < minHandleCap {
		handleCap = minHandleCap
	}

	// Open the db and recover any potential corruptions
	db, err := leveldb.OpenFile(file, &opt.Options{
		OpenFilesCacheCapacity: handleCap,
		BlockCacheCapacity:     blockCacheSize,
		// There are two buffers of size WriteBuffer used.
		WriteBuffer: writeBufferSize / 2,
		Filter:      filter.NewBloomFilter(10),
	})
	if _, corrupted := err.(*errors.ErrCorrupted); corrupted {
		db, err = leveldb.RecoverFile(file, nil)
	}
	if err != nil {
		return nil, err
	}
	return &Database{
		DB:  db,
		log: log,
	}, nil
}

// Has returns if the key is set in the database
func (db *Database) Has(key []byte) (bool, error) {
	if db.corrupted() {
		return false, database.ErrAvoidCorruption
	}
	has, err := db.DB.Has(key, nil)
	return has, db.handleError(err)
}

// Get returns the value the key maps to in the database
func (db *Database) Get(key []byte) ([]byte, error) {
	if db.corrupted() {
		return nil, database.ErrAvoidCorruption
	}
	value, err := db.DB.Get(key, nil)
	return value, db.handleError(err)
}

// Put sets the value of the provided key to the provided value
func (db *Database) Put(key []byte, value []byte) error {
	if db.corrupted() {
		return database.ErrAvoidCorruption
	}
	return db.handleError(db.DB.Put(key, value, nil))
}

// Delete removes the key from the database
func (db *Database) Delete(key []byte) error {
	if db.corrupted() {
		return database.ErrAvoidCorruption
	}
	return db.handleError(db.DB.Delete(key, nil))
}

// NewBatch creates a write/delete-only buffer that is atomically committed to
// the database when write is called
func (db *Database) NewBatch() database.Batch { return &batch{db: db} }

// NewIterator creates a lexicographically ordered iterator over the database
func (db *Database) NewIterator() database.Iterator {
	return &iter{db.DB.NewIterator(new(util.Range), nil)}
}

// NewIteratorWithStart creates a lexicographically ordered iterator over the
// database starting at the provided key
func (db *Database) NewIteratorWithStart(start []byte) database.Iterator {
	return &iter{db.DB.NewIterator(&util.Range{Start: start}, nil)}
}

// NewIteratorWithPrefix creates a lexicographically ordered iterator over the
// database ignoring keys that do not start with the provided prefix
func (db *Database) NewIteratorWithPrefix(prefix []byte) database.Iterator {
	return &iter{db.DB.NewIterator(util.BytesPrefix(prefix), nil)}
}

// NewIteratorWithStartAndPrefix creates a lexicographically ordered iterator
// over the database starting at start and ignoring keys that do not start with
// the provided prefix
func (db *Database) NewIteratorWithStartAndPrefix(start, prefix []byte) database.Iterator {
	iterRange := util.BytesPrefix(prefix)
	if bytes.Compare(start, prefix) == 1 {
		iterRange.Start = start
	}
	return &iter{db.DB.NewIterator(iterRange, nil)}
}

// Stat returns a particular internal stat of the database.
func (db *Database) Stat(property string) (string, error) {
	stat, err := db.DB.GetProperty(property)
	return stat, db.handleError(err)
}

// This comment is basically copy pasted from the underlying levelDB library:

// Compact the underlying DB for the given key range.
// Specifically, deleted and overwritten versions are discarded,
// and the data is rearranged to reduce the cost of operations
// needed to access the data. This operation should typically only
// be invoked by users who understand the underlying implementation.
//
// A nil start is treated as a key before all keys in the DB.
// And a nil limit is treated as a key after all keys in the DB.
// Therefore if both are nil then it will compact entire DB.
func (db *Database) Compact(start []byte, limit []byte) error {
	return db.handleError(db.DB.CompactRange(util.Range{Start: start, Limit: limit}))
}

// Close implements the Database interface
func (db *Database) Close() error { return db.handleError(db.DB.Close()) }

func (db *Database) corrupted() bool {
	return atomic.LoadUint64(&db.errored) == 1
}

func (db *Database) handleError(err error) error {
	err = updateError(err)
	switch err {
	case nil, database.ErrNotFound, database.ErrClosed:
	// If we get an error other than "not found" or "closed", disallow future
	// database operations to avoid possible corruption
	default:
		db.log.Fatal("leveldb error: %w", err)
		atomic.StoreUint64(&db.errored, 1)
	}
	return err
}

// batch is a wrapper around a levelDB batch to contain sizes.
type batch struct {
	leveldb.Batch
	db   *Database
	size int
}

// Put the value into the batch for later writing
func (b *batch) Put(key, value []byte) error {
	b.Batch.Put(key, value)
	b.size += len(key) + len(value) + levelDBByteOverhead
	return nil
}

// Delete the key during writing
func (b *batch) Delete(key []byte) error {
	b.Batch.Delete(key)
	b.size += len(key) + levelDBByteOverhead
	return nil
}

// Size retrieves the amount of data queued up for writing.
func (b *batch) Size() int { return b.size }

// Write flushes any accumulated data to disk.
func (b *batch) Write() error {
	if b.db.corrupted() {
		return database.ErrAvoidCorruption
	}
	return b.db.handleError(b.db.DB.Write(&b.Batch, nil))
}

// Reset resets the batch for reuse.
func (b *batch) Reset() {
	b.Batch.Reset()
	b.size = 0
}

// Replay the batch contents.
func (b *batch) Replay(w database.KeyValueWriter) error {
	replay := &replayer{writer: w}
	if err := b.Batch.Replay(replay); err != nil {
		// Never actually returns an error, because Replay just returns nil
		return b.db.handleError(err)
	}
	return b.db.handleError(replay.err)
}

// Inner returns itself
func (b *batch) Inner() database.Batch { return b }

type replayer struct {
	writer database.KeyValueWriter
	err    error
}

func (r *replayer) Put(key, value []byte) {
	if r.err != nil {
		return
	}
	r.err = r.writer.Put(key, value)
}

func (r *replayer) Delete(key []byte) {
	if r.err != nil {
		return
	}
	r.err = r.writer.Delete(key)
}

type iter struct{ iterator.Iterator }

// Error implements the Iterator interface
func (it *iter) Error() error { return updateError(it.Iterator.Error()) }

// Key implements the Iterator interface
func (it *iter) Key() []byte { return utils.CopyBytes(it.Iterator.Key()) }

// Value implements the Iterator interface
func (it *iter) Value() []byte { return utils.CopyBytes(it.Iterator.Value()) }

func updateError(err error) error {
	switch err {
	case leveldb.ErrClosed:
		return database.ErrClosed
	case leveldb.ErrNotFound:
		return database.ErrNotFound
	default:
		return err
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Bret Jordan, All rights reserved.
//
// Use of this source code is governed by an Apache 2.0 license
// that can be found in the LICENSE file in the root of the source
// tree.

package datastore

import (
	"github.com/freetaxii/libstix2/objects"
	"github.com/freetaxii/libstix2/resources"
)

// Datastorer - This type enables access to the TAXII datastore.
// GetObject returns a STIX object
type Datastorer interface {
	Close() error
	GetEnabledCollections() resources.CollectionsType
	GetObject(stixid, version string) (interface{}, error)
	GetObjectsFromCollection(query QueryType) (*objects.BundleType, *QueryReturnDataType, error)
	GetListOfObjectsFromCollection(query QueryType) (*[]CollectionRawDataType, *QueryReturnDataType, error)
}

// QueryType - This struct will hold all of the variables that a user can
// query a collection on.
type QueryType struct {
	CollectionID string
	STIXID       []string
	STIXType     []string
	STIXVersion  []string
	AddedAfter   []string
	RangeBegin   int
	RangeEnd     int
	RangeMax     int
}

// QueryReturnDataType - This struct contains the various bits of meta data that
// are returned from a query against a TAXII server. This is done so that the
// method signatures do not need to change as time goes on and we add more meta
// data that needs to be returned.
type QueryReturnDataType struct {
	Size           int
	DateAddedFirst string
	DateAddedLast  string
}

// CollectionRawDataType -
type CollectionRawDataType struct {
	STIXID      string
	DateAdded   string
	STIXVersion string
	SpecVersion string
}
<commit_msg>added manifest<commit_after>// Copyright 2017 Bret Jordan, All rights reserved.
//
// Use of this source code is governed by an Apache 2.0 license
// that can be found in the LICENSE file in the root of the source
// tree.

package datastore

import (
	"github.com/freetaxii/libstix2/objects"
	"github.com/freetaxii/libstix2/resources"
)

// Datastorer - This type enables access to the TAXII datastore.
// GetObject returns a STIX object
type Datastorer interface {
	Close() error
	GetEnabledCollections() resources.CollectionsType
	GetObject(stixid, version string) (interface{}, error)
	GetObjectsFromCollection(query QueryType) (*objects.BundleType, *QueryReturnDataType, error)
	GetListOfObjectsFromCollection(query QueryType) (*[]CollectionRawDataType, *QueryReturnDataType, error)
	GetManifestFromCollection(query QueryType) (*resources.ManifestType, *QueryReturnDataType, error)
}

// QueryType - This struct will hold all of the variables that a user can
// query a collection on.
type QueryType struct {
	CollectionID string
	STIXID       []string
	STIXType     []string
	STIXVersion  []string
	AddedAfter   []string
	RangeBegin   int
	RangeEnd     int
	RangeMax     int
}

// QueryReturnDataType - This struct contains the various bits of meta data that
// are returned from a query against a TAXII server. This is done so that the
// method signatures do not need to change as time goes on and we add more meta
// data that needs to be returned.
type QueryReturnDataType struct {
	Size           int
	DateAddedFirst string
	DateAddedLast  string
}

// CollectionRawDataType -
type CollectionRawDataType struct {
	STIXID      string
	DateAdded   string
	STIXVersion string
	SpecVersion string
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Christian Saide <Supernomad>
// Licensed under the MPL-2.0, for details see https://github.com/Supernomad/quantum/blob/master/LICENSE

package datastore

import (
	"errors"
	"sync"
	"time"

	"github.com/Supernomad/quantum/common"
)

// Type represents the datastore backend to use for synchronizing mapping objects over the quantum network.
type Type int

const (
	// ETCDDatastore will tell quantum to use etcd as the backend datastore.
	ETCDDatastore Type = iota

	// MOCKDatastore will tell quantum to use a moked out backend datastore for testing.
	MOCKDatastore

	lockTTL = 10 * time.Second
)

// Datastore interface for quantum to use for retrieving mapping data from the backend datastore.
type Datastore interface {
	Init() error
	Mapping(ip uint32) (*common.Mapping, bool)
	Start(wg *sync.WaitGroup)
	Stop()
}

// New generates a datastore object based on the passed in Type and user configuration.
func New(datastoreType Type, log *common.Logger, cfg *common.Config) (Datastore, error) {
	switch datastoreType {
	case ETCDDatastore:
		return newEtcd(log, cfg)
	case MOCKDatastore:
		return newMock(log, cfg)
	default:
		return nil, errors.New("specified backend doesn't exist")
	}
}
<commit_msg>Another round of comments<commit_after>// Copyright (c) 2016 Christian Saide <Supernomad>
// Licensed under the MPL-2.0, for details see https://github.com/Supernomad/quantum/blob/master/LICENSE

package datastore

import (
	"errors"
	"sync"
	"time"

	"github.com/Supernomad/quantum/common"
)

// Type represents the datastore backend to use for synchronizing mapping objects over the quantum network.
type Type int

const (
	// ETCDDatastore will tell quantum to use etcd as the backend datastore.
	ETCDDatastore Type = iota

	// MOCKDatastore will tell quantum to use a moked out backend datastore for testing.
	MOCKDatastore

	lockTTL = 10 * time.Second
)

// Datastore interface for quantum to use for retrieving mapping data from the backend datastore.
type Datastore interface {
	// Init should handle setting up the datastore connections, and initializing the mappings/local mapping.
	Init() error

	// Mapping should return the mapping and true if it exists, if not the mapping should be nil and false should be returned along with it.
	Mapping(ip uint32) (*common.Mapping, bool)

	// Start should kick off any routines that need to run in the background to groom the mappings and manage the datastore state.
	Start(wg *sync.WaitGroup)

	// Stop should fully shutdown all operation and ensure that all connections are terminated gracefully.
	Stop()
}

// New generates a datastore object based on the passed in Type and user configuration.
func New(datastoreType Type, log *common.Logger, cfg *common.Config) (Datastore, error) {
	switch datastoreType {
	case ETCDDatastore:
		return newEtcd(log, cfg)
	case MOCKDatastore:
		return newMock(log, cfg)
	default:
		return nil, errors.New("specified backend doesn't exist")
	}
}
<|endoftext|>


<commit_before>package identify

import (
	"context"
	"sync"
	"time"

	net "github.com/libp2p/go-libp2p-net"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
)

const ActivationThresh = 4

var GCInterval = 10 * time.Minute

type observation struct {
	seenTime      time.Time
	connDirection net.Direction
}

// ObservedAddr is an entry for an address reported by our peers.
// We only use addresses that:
// - have been observed at least 4 times in last 1h. (counter symmetric nats)
// - have been observed at least once recently (1h), because our position in the
//   network, or network port mapppings, may have changed.
type ObservedAddr struct {
	Addr     ma.Multiaddr
	SeenBy   map[string]observation // peer(observer) address -> observation info
	LastSeen time.Time
}

func (oa *ObservedAddr) activated(ttl time.Duration) bool {
	// cleanup SeenBy set
	now := time.Now()

	for k, ob := range oa.SeenBy {
		if now.Sub(ob.seenTime) > ttl*ActivationThresh {
			delete(oa.SeenBy, k)
		}
	}

	// We only activate if in the TTL other peers observed the same address
	// of ours at least 4 times.
	return len(oa.SeenBy) >= ActivationThresh
}

type newObservation struct {
	observed, local, observer ma.Multiaddr
	direction                 net.Direction
}

// ObservedAddrSet keeps track of a set of ObservedAddrs
// the zero-value is ready to be used.
type ObservedAddrSet struct {
	sync.Mutex // guards whole datastruct.

	// local(internal) address -> list of observed(external) addresses
	addrs map[string][]*ObservedAddr
	ttl   time.Duration

	// this is the worker channel
	wch chan newObservation
}

func NewObservedAddrSet(ctx context.Context) *ObservedAddrSet {
	oas := &ObservedAddrSet{
		addrs: make(map[string][]*ObservedAddr),
		wch:   make(chan newObservation, 1),
	}
	go oas.worker(ctx)
	return oas
}

// AddrsFor return all activated observed addresses associated with the given
// (resolved) listen address.
func (oas *ObservedAddrSet) AddrsFor(addr ma.Multiaddr) (addrs []ma.Multiaddr) {
	oas.Lock()
	defer oas.Unlock()

	// for zero-value.
	if len(oas.addrs) == 0 {
		return nil
	}

	key := string(addr.Bytes())
	observedAddrs, ok := oas.addrs[key]
	if !ok {
		return
	}

	now := time.Now()
	for _, a := range observedAddrs {
		if now.Sub(a.LastSeen) <= oas.ttl && a.activated(oas.ttl) {
			addrs = append(addrs, a.Addr)
		}
	}

	return addrs
}

// Addrs return all activated observed addresses
func (oas *ObservedAddrSet) Addrs() (addrs []ma.Multiaddr) {
	oas.Lock()
	defer oas.Unlock()

	// for zero-value.
	if len(oas.addrs) == 0 {
		return nil
	}

	now := time.Now()
	for _, observedAddrs := range oas.addrs {
		for _, a := range observedAddrs {
			if now.Sub(a.LastSeen) <= oas.ttl && a.activated(oas.ttl) {
				addrs = append(addrs, a.Addr)
			}
		}
	}
	return addrs
}

func (oas *ObservedAddrSet) Add(observed, local, observer ma.Multiaddr,
	direction net.Direction) {
	select {
	case oas.wch <- newObservation{observed: observed, local: local, observer: observer, direction: direction}:
	default:
		log.Debugf("dropping address observation of %s; buffer full", observed)
	}
}

func (oas *ObservedAddrSet) worker(ctx context.Context) {
	ticker := time.NewTicker(GCInterval)
	defer ticker.Stop()

	for {
		select {
		case obs := <-oas.wch:
			oas.doAdd(obs.observed, obs.local, obs.observer, obs.direction)

		case <-ticker.C:
			oas.gc()

		case <-ctx.Done():
			return
		}
	}
}

func (oas *ObservedAddrSet) gc() {
	oas.Lock()
	defer oas.Unlock()

	now := time.Now()
	for local, observedAddrs := range oas.addrs {
		// TODO we can do this without allocating by compacting the array in place
		filteredAddrs := make([]*ObservedAddr, 0, len(observedAddrs))
		for _, a := range observedAddrs {
			// leave only alive observed addresses
			if now.Sub(a.LastSeen) <= oas.ttl {
				filteredAddrs = append(filteredAddrs, a)
			}
		}
		oas.addrs[local] = filteredAddrs
	}
}

func (oas *ObservedAddrSet) doAdd(observed, local, observer ma.Multiaddr,
	direction net.Direction) {

	now := time.Now()
	observerString := observerGroup(observer)
	localString := string(local.Bytes())
	ob := observation{
		seenTime:      now,
		connDirection: direction,
	}

	oas.Lock()
	defer oas.Unlock()

	// for zero-value.
	if oas.addrs == nil {
		oas.addrs = make(map[string][]*ObservedAddr)
		oas.ttl = pstore.OwnObservedAddrTTL
	}

	observedAddrs := oas.addrs[localString]
	// check if observed address seen yet, if so, update it
	for i, previousObserved := range observedAddrs {
		if previousObserved.Addr.Equal(observed) {
			observedAddrs[i].SeenBy[observerString] = ob
			observedAddrs[i].LastSeen = now
			return
		}
	}
	// observed address not seen yet, append it
	oas.addrs[localString] = append(oas.addrs[localString], &ObservedAddr{
		Addr: observed,
		SeenBy: map[string]observation{
			observerString: ob,
		},
		LastSeen: now,
	})
}

// observerGroup is a function that determines what part of
// a multiaddr counts as a different observer. for example,
// two ipfs nodes at the same IP/TCP transport would get
// the exact same NAT mapping; they would count as the
// same observer. This may protect against NATs who assign
// different ports to addresses at different IP hosts, but
// not TCP ports.
//
// Here, we use the root multiaddr address. This is mostly
// IP addresses. In practice, this is what we want.
func observerGroup(m ma.Multiaddr) string {
	//TODO: If IPv6 rolls out we should mark /64 routing zones as one group
	first, _ := ma.SplitFirst(m)
	return string(first.Bytes())
}

func (oas *ObservedAddrSet) SetTTL(ttl time.Duration) {
	oas.Lock()
	defer oas.Unlock()
	oas.ttl = ttl
}

func (oas *ObservedAddrSet) TTL() time.Duration {
	oas.Lock()
	defer oas.Unlock()
	// for zero-value.
	if oas.addrs == nil {
		oas.ttl = pstore.OwnObservedAddrTTL
	}
	return oas.ttl
}
<commit_msg>initialize ttl in ObservedAddrSet<commit_after>package identify

import (
	"context"
	"sync"
	"time"

	net "github.com/libp2p/go-libp2p-net"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
)

const ActivationThresh = 4

var GCInterval = 10 * time.Minute

type observation struct {
	seenTime      time.Time
	connDirection net.Direction
}

// ObservedAddr is an entry for an address reported by our peers.
// We only use addresses that:
// - have been observed at least 4 times in last 1h. (counter symmetric nats)
// - have been observed at least once recently (1h), because our position in the
//   network, or network port mapppings, may have changed.
type ObservedAddr struct {
	Addr     ma.Multiaddr
	SeenBy   map[string]observation // peer(observer) address -> observation info
	LastSeen time.Time
}

func (oa *ObservedAddr) activated(ttl time.Duration) bool {
	// cleanup SeenBy set
	now := time.Now()

	for k, ob := range oa.SeenBy {
		if now.Sub(ob.seenTime) > ttl*ActivationThresh {
			delete(oa.SeenBy, k)
		}
	}

	// We only activate if in the TTL other peers observed the same address
	// of ours at least 4 times.
	return len(oa.SeenBy) >= ActivationThresh
}

type newObservation struct {
	observed, local, observer ma.Multiaddr
	direction                 net.Direction
}

// ObservedAddrSet keeps track of a set of ObservedAddrs
// the zero-value is ready to be used.
type ObservedAddrSet struct {
	sync.Mutex // guards whole datastruct.

	// local(internal) address -> list of observed(external) addresses
	addrs map[string][]*ObservedAddr
	ttl   time.Duration

	// this is the worker channel
	wch chan newObservation
}

func NewObservedAddrSet(ctx context.Context) *ObservedAddrSet {
	oas := &ObservedAddrSet{
		addrs: make(map[string][]*ObservedAddr),
		ttl:   pstore.OwnObservedAddrTTL,
		wch:   make(chan newObservation, 1),
	}
	go oas.worker(ctx)
	return oas
}

// AddrsFor return all activated observed addresses associated with the given
// (resolved) listen address.
func (oas *ObservedAddrSet) AddrsFor(addr ma.Multiaddr) (addrs []ma.Multiaddr) {
	oas.Lock()
	defer oas.Unlock()

	if len(oas.addrs) == 0 {
		return nil
	}

	key := string(addr.Bytes())
	observedAddrs, ok := oas.addrs[key]
	if !ok {
		return
	}

	now := time.Now()
	for _, a := range observedAddrs {
		if now.Sub(a.LastSeen) <= oas.ttl && a.activated(oas.ttl) {
			addrs = append(addrs, a.Addr)
		}
	}

	return addrs
}

// Addrs return all activated observed addresses
func (oas *ObservedAddrSet) Addrs() (addrs []ma.Multiaddr) {
	oas.Lock()
	defer oas.Unlock()

	if len(oas.addrs) == 0 {
		return nil
	}

	now := time.Now()
	for _, observedAddrs := range oas.addrs {
		for _, a := range observedAddrs {
			if now.Sub(a.LastSeen) <= oas.ttl && a.activated(oas.ttl) {
				addrs = append(addrs, a.Addr)
			}
		}
	}
	return addrs
}

func (oas *ObservedAddrSet) Add(observed, local, observer ma.Multiaddr,
	direction net.Direction) {
	select {
	case oas.wch <- newObservation{observed: observed, local: local, observer: observer, direction: direction}:
	default:
		log.Debugf("dropping address observation of %s; buffer full", observed)
	}
}

func (oas *ObservedAddrSet) worker(ctx context.Context) {
	ticker := time.NewTicker(GCInterval)
	defer ticker.Stop()

	for {
		select {
		case obs := <-oas.wch:
			oas.doAdd(obs.observed, obs.local, obs.observer, obs.direction)

		case <-ticker.C:
			oas.gc()

		case <-ctx.Done():
			return
		}
	}
}

func (oas *ObservedAddrSet) gc() {
	oas.Lock()
	defer oas.Unlock()

	now := time.Now()
	for local, observedAddrs := range oas.addrs {
		// TODO we can do this without allocating by compacting the array in place
		filteredAddrs := make([]*ObservedAddr, 0, len(observedAddrs))
		for _, a := range observedAddrs {
			// leave only alive observed addresses
			if now.Sub(a.LastSeen) <= oas.ttl {
				filteredAddrs = append(filteredAddrs, a)
			}
		}
		oas.addrs[local] = filteredAddrs
	}
}

func (oas *ObservedAddrSet) doAdd(observed, local, observer ma.Multiaddr,
	direction net.Direction) {

	now := time.Now()
	observerString := observerGroup(observer)
	localString := string(local.Bytes())
	ob := observation{
		seenTime:      now,
		connDirection: direction,
	}

	oas.Lock()
	defer oas.Unlock()

	observedAddrs := oas.addrs[localString]
	// check if observed address seen yet, if so, update it
	for i, previousObserved := range observedAddrs {
		if previousObserved.Addr.Equal(observed) {
			observedAddrs[i].SeenBy[observerString] = ob
			observedAddrs[i].LastSeen = now
			return
		}
	}
	// observed address not seen yet, append it
	oas.addrs[localString] = append(oas.addrs[localString], &ObservedAddr{
		Addr: observed,
		SeenBy: map[string]observation{
			observerString: ob,
		},
		LastSeen: now,
	})
}

// observerGroup is a function that determines what part of
// a multiaddr counts as a different observer. for example,
// two ipfs nodes at the same IP/TCP transport would get
// the exact same NAT mapping; they would count as the
// same observer. This may protect against NATs who assign
// different ports to addresses at different IP hosts, but
// not TCP ports.
//
// Here, we use the root multiaddr address. This is mostly
// IP addresses. In practice, this is what we want.
func observerGroup(m ma.Multiaddr) string {
	//TODO: If IPv6 rolls out we should mark /64 routing zones as one group
	first, _ := ma.SplitFirst(m)
	return string(first.Bytes())
}

func (oas *ObservedAddrSet) SetTTL(ttl time.Duration) {
	oas.Lock()
	defer oas.Unlock()
	oas.ttl = ttl
}

func (oas *ObservedAddrSet) TTL() time.Duration {
	oas.Lock()
	defer oas.Unlock()
	// for zero-value.
	if oas.addrs == nil {
		oas.ttl = pstore.OwnObservedAddrTTL
	}
	return oas.ttl
}
<|endoftext|>


<commit_before>/**
 * (C) Copyright 2014, Deft Labs
 *
 * 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 dlshared

import (
	"time"
	"testing"
)

func TestMetrics(t *testing.T) {

	relayFuncCallCount := 0

	relayFunc := func(sourceName string, metrics []Metric) {
		relayFuncCallCount++
	}

	metrics := NewMetrics("TestSource", []MetricsRelayFunction { relayFunc }, 1, 100)

	if err := metrics.Start(); err != nil {
		t.Errorf("TestMetrics Start is broken: %v", err)
	}

	startTime := CurrentTimeInMillis()

	for idx := 0; idx < 100000; idx++ {
		metrics.Gauge("test.guage", 100)
		metrics.Count("test.count.by.one")
		metrics.CountWithValue("test.count.with.value", 10000000)
	}

	execTime := CurrentTimeInMillis() - startTime

	if execTime > 500 {
		t.Errorf("TestMetrics is too slow - time in ms: %d", execTime)
	}

	time.Sleep(1 * time.Second)

	if err := metrics.Stop(); err != nil {
		t.Errorf("TestMetrics Stop is broken: %v", err)
	}

	if relayFuncCallCount == 0 {
		t.Errorf("TestMetrics relay function call is broken - not called")
	}
}

<commit_msg>changed threshold.<commit_after>/**
 * (C) Copyright 2014, Deft Labs
 *
 * 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 dlshared

import (
	"time"
	"testing"
)

func TestMetrics(t *testing.T) {

	relayFuncCallCount := 0

	relayFunc := func(sourceName string, metrics []Metric) {
		relayFuncCallCount++
	}

	metrics := NewMetrics("TestSource", []MetricsRelayFunction { relayFunc }, 1, 100)

	if err := metrics.Start(); err != nil {
		t.Errorf("TestMetrics Start is broken: %v", err)
	}

	startTime := CurrentTimeInMillis()

	for idx := 0; idx < 100000; idx++ {
		metrics.Gauge("test.guage", 100)
		metrics.Count("test.count.by.one")
		metrics.CountWithValue("test.count.with.value", 10000000)
	}

	execTime := CurrentTimeInMillis() - startTime

	if execTime > 1000 {
		t.Errorf("TestMetrics is too slow - time in ms: %d", execTime)
	}

	time.Sleep(1 * time.Second)

	if err := metrics.Stop(); err != nil {
		t.Errorf("TestMetrics Stop is broken: %v", err)
	}

	if relayFuncCallCount == 0 {
		t.Errorf("TestMetrics relay function call is broken - not called")
	}
}

<|endoftext|>


<commit_before>//Package node provides the node side logic of the node.
package node

import (
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/satori/go.uuid"
	"github.com/tywkeene/autobd/client"
	"github.com/tywkeene/autobd/index"
	"github.com/tywkeene/autobd/options"
	"github.com/tywkeene/autobd/utils"
	"github.com/tywkeene/autobd/version"
	"io/ioutil"
	"os"
	"strings"
	"time"
)

type Node struct {
	Servers map[string]*client.Client
	UUID    string
	Synced  bool
	Config  options.NodeConf
}

var localNode *Node
var nodeUseragent string = "Autobd-node/" + version.GetNodeVersion()

func newNode(config options.NodeConf) *Node {
	servers := make(map[string]*client.Client, 0)
	for _, url := range config.Servers {
		servers[url] = client.NewClient(url)
	}
	return &Node{servers, "", false, config}
}

func InitNode(config options.NodeConf) *Node {
	node := newNode(config)
	//Check to see if we already have a UUID stored in a file, if not, generate one and
	//write it to node.Config.UUIDPath
	if _, err := os.Stat(config.UUIDPath); os.IsNotExist(err) {
		node.UUID = uuid.NewV4().String()
		node.WriteNodeUUID()
		log.Infof("Generated and wrote node UUID (%s) to (%s) ", node.UUID, node.Config.UUIDPath)
	} else {
		node.ReadNodeUUID()
		log.Infof("Read node UUID (%s) from (%s) ", node.UUID, node.Config.UUIDPath)
	}
	return node
}

func (node *Node) WriteNodeUUID() error {
	outfile, err := os.Create(node.Config.UUIDPath)
	if err != nil {
		return err
	}
	defer outfile.Close()
	serial, err := json.MarshalIndent(&node.UUID, " ", " ")
	if err != nil {
		return err
	}
	_, err = outfile.WriteString(string(serial))
	return err
}

func (node *Node) ReadNodeUUID() error {
	if _, err := os.Stat(node.Config.UUIDPath); err != nil {
		return err
	}
	serial, err := ioutil.ReadFile(node.Config.UUIDPath)
	if err != nil {
		return err
	}
	return json.Unmarshal(serial, &node.UUID)
}

func (node *Node) validateServerVersion(remote *version.VersionInfo) error {
	if version.GetAPIVersion() != remote.APIVersion {
		return fmt.Errorf("Mismatched version with server. Server: %s Local: %s",
			remote.APIVersion, version.GetAPIVersion())
	}
	remoteMajor := strings.Split(remote.APIVersion, ".")[0]
	if version.GetMajor() != remoteMajor {
		return fmt.Errorf("Mismatched API version with server. Server: %s Local: %s",
			remoteMajor, version.GetMajor())
	}
	return nil
}

func (node *Node) StartHeart() {
	go func(config options.NodeConf) {
		interval, _ := time.ParseDuration(config.HeartbeatInterval)
		log.Info("Started heartbeat, updating every ", interval)
		for {
			time.Sleep(interval)
			for _, server := range node.Servers {
				if server.Online == false {
					continue
				}
				_, err := server.SendHeartbeat(node.UUID, node.Synced, nodeUseragent)
				if utils.HandleError("node/StartHeart()", err, utils.ErrorActionErr) == true {
					server.MissedBeats++
					if server.MissedBeats == node.Config.MaxMissedBeats {
						server.Online = false
						log.Error(server.Address + " has missed max heartbeats, ignoring")
					}
				}
			}
		}
	}(node.Config)
}

func (node *Node) CountOnlineServers() int {
	var count int = 0
	for _, server := range node.Servers {
		if server.Online == true {
			count++
		}
	}
	return count
}

func (node *Node) ValidateAndIdentifyWithServers() error {
	for _, server := range node.Servers {
		serial, err := server.RequestVersion()
		if err != nil {
			return err
		}
		var remoteVer *version.VersionInfo
		if err := json.Unmarshal(serial, &remoteVer); err != nil {
			return err
		}

		if options.Config.NodeConfig.IgnoreVersionMismatch == false {
			if err := node.validateServerVersion(remoteVer); err != nil {
				return err
			}
		}
		_, err = server.IdentifyWithServer(node.UUID, nodeUseragent)
		if utils.HandleError("node/ValidateAndIdentifyWithServers", err, utils.ErrorActionErr) == true {
			continue
		}
	}
	node.StartHeart()
	return nil
}

func CompareDirs(local map[string]*index.Index, remote map[string]*index.Index) []*index.Index {
	need := make([]*index.Index, 0)
	for objName, object := range remote {
		_, existsLocally := local[object.Name] //Does it exist on the node?

		//If it doesn't exist on the node at all, we obviously need it
		if existsLocally == false {
			need = append(need, remote[objName])
			continue
		}

		// If it does, and it's a directory, and it has children
		if existsLocally == true && object.IsDir == true && object.Files != nil {
			dirNeed := CompareDirs(local[objName].Files, object.Files) //Scan the children
			need = append(need, dirNeed...)
			continue
		}

		//If it isn't a directory and doesn't exist
		if existsLocally == false && object.IsDir == false {
			need = append(need, remote[objName])
			continue
		}

		//If it is a file and does exist, compare checksums
		if existsLocally == true && object.IsDir == false {
			if local[objName].Checksum != remote[objName].Checksum {
				log.Info("Checksum mismatch:", objName)
				need = append(need, remote[objName])
				continue
			}
		}
	}
	return need
}

//Compare a local and remote index, return a slice of needed indexes (or nil)
func (node *Node) CompareIndex(target string, uuid string, userAgent string, client *client.Client) ([]*index.Index, error) {
	serial, err := client.RequestIndex(target, uuid, userAgent)
	if err != nil {
		return nil, err
	}
	var remoteIndex map[string]*index.Index
	if err := json.Unmarshal(serial, &remoteIndex); err != nil {
		return nil, err
	}
	localIndex, err := index.GetIndex("/")
	if err != nil {
		return nil, err
	}
	need := CompareDirs(localIndex, remoteIndex)
	return need, nil
}

func (node *Node) SyncUp(need []*index.Index, s *client.Client) {
	for _, object := range need {
		log.Printf("Need %s from %s", object.Name, s.Address)
		if object.IsDir == true {
			err := s.RequestSyncDir(object.Name, node.UUID, nodeUseragent)
			if utils.HandleError("node/SyncUp()", err, utils.ErrorActionInfo) == true {
				continue
			}
		} else if object.IsDir == false {
			err := s.RequestSyncFile(object.Name, node.UUID, nodeUseragent)
			if err != nil {
				//EOF just means the sync is finished, don't log an error
				utils.HandleError("node/SyncUp()", err, utils.ErrorActionInfo)
				continue
			}
		}
	}
}

func (node *Node) UpdateLoop() error {
	err := node.ValidateAndIdentifyWithServers()
	utils.HandlePanic("node/UpdateLoop", err)
	log.Printf("Running as a node. Updating every %s with %s",
		node.Config.UpdateInterval, node.Config.Servers)

	updateInterval, err := time.ParseDuration(node.Config.UpdateInterval)
	utils.HandlePanic("node/UpdateLoop()", err)
	for {
		time.Sleep(updateInterval)
		if node.CountOnlineServers() == 0 {
			utils.HandlePanic("node/UpdateLoop()", fmt.Errorf("No servers online, dying"))
		}
		for _, server := range node.Servers {
			if server.Online == false {
				log.Info("Skipping offline server: ", server.Address)
				continue
			}
			need, err := node.CompareIndex(node.Config.TargetDirectory, node.UUID, nodeUseragent, server)
			if utils.HandleError("node/UpdateLoop()", err, utils.ErrorActionWarn) == true {
				continue
			}

			if len(need) == 0 {
				node.Synced = true
				continue
			}
			node.SyncUp(need, server)
		}
	}
}
<commit_msg>Refactored node/node.go:UpdateLoop() and node/node.go:SyncUp(). Added node/node.go:SetSync()<commit_after>//Package node provides the node side logic of the node.
package node

import (
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/satori/go.uuid"
	"github.com/tywkeene/autobd/client"
	"github.com/tywkeene/autobd/index"
	"github.com/tywkeene/autobd/options"
	"github.com/tywkeene/autobd/utils"
	"github.com/tywkeene/autobd/version"
	"io/ioutil"
	"os"
	"strings"
	"time"
)

type Node struct {
	Servers map[string]*client.Client
	UUID    string
	Synced  bool
	Config  options.NodeConf
}

var localNode *Node
var nodeUseragent string = "Autobd-node/" + version.GetNodeVersion()

func newNode(config options.NodeConf) *Node {
	servers := make(map[string]*client.Client, 0)
	for _, url := range config.Servers {
		servers[url] = client.NewClient(url)
	}
	return &Node{servers, "", false, config}
}

func InitNode(config options.NodeConf) *Node {
	node := newNode(config)
	//Check to see if we already have a UUID stored in a file, if not, generate one and
	//write it to node.Config.UUIDPath
	if _, err := os.Stat(config.UUIDPath); os.IsNotExist(err) {
		node.UUID = uuid.NewV4().String()
		node.WriteNodeUUID()
		log.Infof("Generated and wrote node UUID (%s) to (%s) ", node.UUID, node.Config.UUIDPath)
	} else {
		node.ReadNodeUUID()
		log.Infof("Read node UUID (%s) from (%s) ", node.UUID, node.Config.UUIDPath)
	}
	return node
}

func (node *Node) WriteNodeUUID() error {
	outfile, err := os.Create(node.Config.UUIDPath)
	if err != nil {
		return err
	}
	defer outfile.Close()
	serial, err := json.MarshalIndent(&node.UUID, " ", " ")
	if err != nil {
		return err
	}
	_, err = outfile.WriteString(string(serial))
	return err
}

func (node *Node) ReadNodeUUID() error {
	if _, err := os.Stat(node.Config.UUIDPath); err != nil {
		return err
	}
	serial, err := ioutil.ReadFile(node.Config.UUIDPath)
	if err != nil {
		return err
	}
	return json.Unmarshal(serial, &node.UUID)
}

func (node *Node) validateServerVersion(remote *version.VersionInfo) error {
	if version.GetAPIVersion() != remote.APIVersion {
		return fmt.Errorf("Mismatched version with server. Server: %s Local: %s",
			remote.APIVersion, version.GetAPIVersion())
	}
	remoteMajor := strings.Split(remote.APIVersion, ".")[0]
	if version.GetMajor() != remoteMajor {
		return fmt.Errorf("Mismatched API version with server. Server: %s Local: %s",
			remoteMajor, version.GetMajor())
	}
	return nil
}

func (node *Node) StartHeart() {
	go func(config options.NodeConf) {
		interval, _ := time.ParseDuration(config.HeartbeatInterval)
		log.Info("Started heartbeat, updating every ", interval)
		for {
			time.Sleep(interval)
			for _, server := range node.Servers {
				if server.Online == false {
					continue
				}
				_, err := server.SendHeartbeat(node.UUID, node.Synced, nodeUseragent)
				if utils.HandleError("node/StartHeart()", err, utils.ErrorActionErr) == true {
					server.MissedBeats++
					if server.MissedBeats == node.Config.MaxMissedBeats {
						server.Online = false
						log.Error(server.Address + " has missed max heartbeats, ignoring")
					}
				}
			}
		}
	}(node.Config)
}

func (node *Node) CountOnlineServers() int {
	var count int = 0
	for _, server := range node.Servers {
		if server.Online == true {
			count++
		}
	}
	return count
}

func (node *Node) ValidateAndIdentifyWithServers() error {
	for _, server := range node.Servers {
		serial, err := server.RequestVersion()
		if err != nil {
			return err
		}
		var remoteVer *version.VersionInfo
		if err := json.Unmarshal(serial, &remoteVer); err != nil {
			return err
		}

		if options.Config.NodeConfig.IgnoreVersionMismatch == false {
			if err := node.validateServerVersion(remoteVer); err != nil {
				return err
			}
		}
		_, err = server.IdentifyWithServer(node.UUID, nodeUseragent)
		if utils.HandleError("node/ValidateAndIdentifyWithServers", err, utils.ErrorActionErr) == true {
			continue
		}
	}
	node.StartHeart()
	return nil
}

func CompareDirs(local map[string]*index.Index, remote map[string]*index.Index) []*index.Index {
	need := make([]*index.Index, 0)
	for objName, remoteObject := range remote {
		_, existsLocally := local[remoteObject.Name] //Does it exist on the node?
		//If it doesn't exist on the node at all, we obviously need it
		if existsLocally == false {
			need = append(need, remoteObject)
			continue
		}
		// If it does, and it's a directory, and it has children
		if existsLocally == true && remoteObject.IsDir == true && remoteObject.Files != nil {
			dirNeed := CompareDirs(local[objName].Files, remoteObject.Files) //Scan the children
			need = append(need, dirNeed...)
			continue
		}
		//If it is a file and does exist, compare checksums
		if existsLocally == true && remoteObject.IsDir == false {
			if local[objName].Checksum != remoteObject.Checksum {
				log.Info("Checksum mismatch:", objName)
				need = append(need, remoteObject)
				continue
			}
		}
	}
	return need
}

//Compare a local and remote index, return a slice of needed indexes (or nil)
func (node *Node) CompareIndex(target string, uuid string, userAgent string, client *client.Client) ([]*index.Index, error) {
	serial, err := client.RequestIndex(target, uuid, userAgent)
	if err != nil {
		return nil, err
	}
	var remoteIndex map[string]*index.Index
	if err := json.Unmarshal(serial, &remoteIndex); err != nil {
		return nil, err
	}
	if _, err := os.Stat(target); os.IsNotExist(err) {
		os.Mkdir(target, 0755)
	}
	localIndex, err := index.GetIndex(target)
	if err != nil {
		return nil, err
	}
	need := CompareDirs(localIndex, remoteIndex)
	return need, nil
}

func (node *Node) SetSync(value bool) {
	if value != node.Synced {
		log.Infof("node.Synced changed from %v to %v", node.Synced, value)
		node.Synced = value
	}
}

func (node *Node) SyncUp(server *client.Client) error {
	need, err := node.CompareIndex(node.Config.TargetDirectory, node.UUID, nodeUseragent, server)
	if err != nil {
		return err
	}
	if len(need) > 0 {
		node.SetSync(false)
		for _, object := range need {
			log.Printf("Need %s from %s", object.Name, server.Address)
			if object.IsDir == true {
				err := server.RequestSyncDir(object.Name, node.UUID, nodeUseragent)
				if utils.HandleError("node/SyncUp()", err, utils.ErrorActionInfo) == true {
					continue
				}
			} else if object.IsDir == false {
				err := server.RequestSyncFile(object.Name, node.UUID, nodeUseragent)
				if err != nil {
					//EOF just means the sync is finished, don't log an error
					utils.HandleError("node/SyncUp()", err, utils.ErrorActionInfo)
					continue
				}
			}
		}
	} else {
		node.SetSync(true)
	}
	return nil
}

func (node *Node) UpdateLoop() error {
	err := node.ValidateAndIdentifyWithServers()
	utils.HandlePanic("node/UpdateLoop", err)
	log.Printf("Running as a node. Updating every %s with %s",
		node.Config.UpdateInterval, node.Config.Servers)

	updateInterval, err := time.ParseDuration(node.Config.UpdateInterval)
	utils.HandlePanic("node/UpdateLoop()", err)
	for {
		time.Sleep(updateInterval)
		if node.CountOnlineServers() == 0 {
			utils.HandlePanic("node/UpdateLoop()", fmt.Errorf("No servers online, dying"))
		}
		for _, server := range node.Servers {
			if server.Online == false {
				log.Info("Skipping offline server: ", server.Address)
				continue
			}
			err := node.SyncUp(server)
			if utils.HandleError("node/UpdateLoop()", err, utils.ErrorActionWarn) == true {
				break
			}
		}
	}
}
<|endoftext|>


<commit_before>// +build darwin dragonfly freebsd linux netbsd openbsd solaris

package allocdir

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"

	"golang.org/x/sys/unix"
)

var (
	// SharedAllocContainerPath is the path inside container for mounted
	// directory shared across tasks in a task group.
	SharedAllocContainerPath = filepath.Join("/", SharedAllocName)

	// TaskLocalContainer is the path inside a container for mounted directory
	// for local storage.
	TaskLocalContainerPath = filepath.Join("/", TaskLocal)

	// TaskSecretsContainerPath is the path inside a container for mounted
	// secrets directory
	TaskSecretsContainerPath = filepath.Join("/", TaskSecrets)
)

// dropDirPermissions gives full access to a directory to all users and sets
// the owner to nobody.
func dropDirPermissions(path string, desired os.FileMode) error {
	if err := os.Chmod(path, desired|0777); err != nil {
		return fmt.Errorf("Chmod(%v) failed: %v", path, err)
	}

	// Can't change owner if not root.
	if unix.Geteuid() != 0 {
		return nil
	}

	u, err := user.Lookup("nobody")
	if err != nil {
		return err
	}

	uid, err := getUid(u)
	if err != nil {
		return err
	}

	gid, err := getGid(u)
	if err != nil {
		return err
	}

	if err := os.Chown(path, uid, gid); err != nil {
		return fmt.Errorf("Couldn't change owner/group of %v to (uid: %v, gid: %v): %v", path, uid, gid, err)
	}

	return nil
}

// getUid for a user
func getUid(u *user.User) (int, error) {
	uid, err := strconv.Atoi(u.Uid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Uid to an int: %v", err)
	}

	return uid, nil
}

// getGid for a user
func getGid(u *user.User) (int, error) {
	gid, err := strconv.Atoi(u.Gid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Gid to an int: %v", err)
	}

	return gid, nil
}

// linkOrCopy attempts to hardlink dst to src and fallsback to copying if the
// hardlink fails.
func linkOrCopy(src, dst string, uid, gid int, perm os.FileMode) error {
	// Avoid link/copy if the file already exists in the chroot
	// TODO 0.6 clean this up. This was needed because chroot creation fails
	// when a process restarts.
	if fileInfo, _ := os.Stat(dst); fileInfo != nil {
		return nil
	}
	// Attempt to hardlink.
	//if err := os.Link(src, dst); err == nil {
	//	return nil
	//}

	return fileCopy(src, dst, uid, gid, perm)
}
<commit_msg>Don't disable hardlinking!<commit_after>// +build darwin dragonfly freebsd linux netbsd openbsd solaris

package allocdir

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"

	"golang.org/x/sys/unix"
)

var (
	// SharedAllocContainerPath is the path inside container for mounted
	// directory shared across tasks in a task group.
	SharedAllocContainerPath = filepath.Join("/", SharedAllocName)

	// TaskLocalContainer is the path inside a container for mounted directory
	// for local storage.
	TaskLocalContainerPath = filepath.Join("/", TaskLocal)

	// TaskSecretsContainerPath is the path inside a container for mounted
	// secrets directory
	TaskSecretsContainerPath = filepath.Join("/", TaskSecrets)
)

// dropDirPermissions gives full access to a directory to all users and sets
// the owner to nobody.
func dropDirPermissions(path string, desired os.FileMode) error {
	if err := os.Chmod(path, desired|0777); err != nil {
		return fmt.Errorf("Chmod(%v) failed: %v", path, err)
	}

	// Can't change owner if not root.
	if unix.Geteuid() != 0 {
		return nil
	}

	u, err := user.Lookup("nobody")
	if err != nil {
		return err
	}

	uid, err := getUid(u)
	if err != nil {
		return err
	}

	gid, err := getGid(u)
	if err != nil {
		return err
	}

	if err := os.Chown(path, uid, gid); err != nil {
		return fmt.Errorf("Couldn't change owner/group of %v to (uid: %v, gid: %v): %v", path, uid, gid, err)
	}

	return nil
}

// getUid for a user
func getUid(u *user.User) (int, error) {
	uid, err := strconv.Atoi(u.Uid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Uid to an int: %v", err)
	}

	return uid, nil
}

// getGid for a user
func getGid(u *user.User) (int, error) {
	gid, err := strconv.Atoi(u.Gid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Gid to an int: %v", err)
	}

	return gid, nil
}

// linkOrCopy attempts to hardlink dst to src and fallsback to copying if the
// hardlink fails.
func linkOrCopy(src, dst string, uid, gid int, perm os.FileMode) error {
	// Avoid link/copy if the file already exists in the chroot
	// TODO 0.6 clean this up. This was needed because chroot creation fails
	// when a process restarts.
	if fileInfo, _ := os.Stat(dst); fileInfo != nil {
		return nil
	}
	// Attempt to hardlink.
	if err := os.Link(src, dst); err == nil {
		return nil
	}

	return fileCopy(src, dst, uid, gid, perm)
}
<|endoftext|>


<commit_before>package yamlpatch

import "fmt"

// Op is a type alias
type Op string

// Ops
const (
	opAdd     Op = "add"
	opRemove  Op = "remove"
	opReplace Op = "replace"
	opMove    Op = "move"
	opCopy    Op = "copy"
)

// Operation is an RFC6902 'Operation'
// https://tools.ietf.org/html/rfc6902#section-4
type Operation struct {
	Op       Op           `yaml:"op,omitempty"`
	Path     string       `yaml:"path,omitempty"`
	From     string       `yaml:"from,omitempty"`
	RawValue *interface{} `yaml:"value,omitempty"`
}

// Value returns the operation's value as a *Node
func (o *Operation) Value() *Node {
	return NewNode(o.RawValue)
}

// Perform executes the operation on the given container
func (o *Operation) Perform(c Container) error {
	var err error

	switch o.Op {
	case opAdd:
		err = tryAdd(c, o)
	case opRemove:
		err = tryRemove(c, o)
	case opReplace:
		err = tryReplace(c, o)
	case opMove:
		err = tryMove(c, o)
	case opCopy:
		err = tryCopy(c, o)
	default:
		err = fmt.Errorf("Unexpected op: %s", o.Op)
	}

	return err
}

func tryAdd(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.Path)

	if con == nil {
		return fmt.Errorf("yamlpatch add operation does not apply: doc is missing path: %s", op.Path)
	}

	return con.Add(key, op.Value())
}

func tryRemove(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.Path)

	if con == nil {
		return fmt.Errorf("yamlpatch remove operation does not apply: doc is missing path: %s", op.Path)
	}

	return con.Remove(key)
}

func tryReplace(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.Path)

	if con == nil {
		return fmt.Errorf("yamlpatch replace operation does not apply: doc is missing path: %s", op.Path)
	}

	val, err := con.Get(key)
	if val == nil || err != nil {
		return fmt.Errorf("yamlpatch replace operation does not apply: doc is missing key: %s", op.Path)
	}

	return con.Set(key, op.Value())
}

func tryMove(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.From)
	if con == nil {
		return fmt.Errorf("yamlpatch move operation does not apply: doc is missing from path: %s", op.From)
	}

	val, err := con.Get(key)
	if err != nil {
		return err
	}

	err = con.Remove(key)
	if err != nil {
		return err
	}

	con, key = findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("yamlpatch move operation does not apply: doc is missing destination path: %s", op.Path)
	}

	return con.Set(key, val)
}

func tryCopy(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.From)
	if con == nil {
		return fmt.Errorf("copy operation does not apply: doc is missing from path: %s", op.From)
	}

	val, err := con.Get(key)
	if err != nil {
		return err
	}

	con, key = findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("copy operation does not apply: doc is missing destination path: %s", op.Path)
	}

	return con.Set(key, val)
}
<commit_msg>Whitespace<commit_after>package yamlpatch

import "fmt"

// Op is a type alias
type Op string

// Ops
const (
	opAdd     Op = "add"
	opRemove  Op = "remove"
	opReplace Op = "replace"
	opMove    Op = "move"
	opCopy    Op = "copy"
)

// Operation is an RFC6902 'Operation'
// https://tools.ietf.org/html/rfc6902#section-4
type Operation struct {
	Op       Op           `yaml:"op,omitempty"`
	Path     string       `yaml:"path,omitempty"`
	From     string       `yaml:"from,omitempty"`
	RawValue *interface{} `yaml:"value,omitempty"`
}

// Value returns the operation's value as a *Node
func (o *Operation) Value() *Node {
	return NewNode(o.RawValue)
}

// Perform executes the operation on the given container
func (o *Operation) Perform(c Container) error {
	var err error

	switch o.Op {
	case opAdd:
		err = tryAdd(c, o)
	case opRemove:
		err = tryRemove(c, o)
	case opReplace:
		err = tryReplace(c, o)
	case opMove:
		err = tryMove(c, o)
	case opCopy:
		err = tryCopy(c, o)
	default:
		err = fmt.Errorf("Unexpected op: %s", o.Op)
	}

	return err
}

func tryAdd(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("yamlpatch add operation does not apply: doc is missing path: %s", op.Path)
	}

	return con.Add(key, op.Value())
}

func tryRemove(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("yamlpatch remove operation does not apply: doc is missing path: %s", op.Path)
	}

	return con.Remove(key)
}

func tryReplace(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("yamlpatch replace operation does not apply: doc is missing path: %s", op.Path)
	}

	val, err := con.Get(key)
	if val == nil || err != nil {
		return fmt.Errorf("yamlpatch replace operation does not apply: doc is missing key: %s", op.Path)
	}

	return con.Set(key, op.Value())
}

func tryMove(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.From)
	if con == nil {
		return fmt.Errorf("yamlpatch move operation does not apply: doc is missing from path: %s", op.From)
	}

	val, err := con.Get(key)
	if err != nil {
		return err
	}

	err = con.Remove(key)
	if err != nil {
		return err
	}

	con, key = findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("yamlpatch move operation does not apply: doc is missing destination path: %s", op.Path)
	}

	return con.Set(key, val)
}

func tryCopy(doc Container, op *Operation) error {
	con, key := findContainer(doc, op.From)
	if con == nil {
		return fmt.Errorf("copy operation does not apply: doc is missing from path: %s", op.From)
	}

	val, err := con.Get(key)
	if err != nil {
		return err
	}

	con, key = findContainer(doc, op.Path)
	if con == nil {
		return fmt.Errorf("copy operation does not apply: doc is missing destination path: %s", op.Path)
	}

	return con.Set(key, val)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/jessevdk/go-operators/types"
	"fmt"
	"github.com/jessevdk/go-flags"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path"
	"strings"
)

type TypeCheck struct {
	Path string
	ImportPath string
	Package *types.Package
	FileSet *token.FileSet
	ParseFiles   []string
	ProcessFiles []string
	FilesToAst map[string]*ast.File
	Ast []*ast.File
}

var typechecks = make(map[string]*TypeCheck)

func packageLocations(args []string) []string {
	if len(args) == 0 {
		dirname, err := os.Getwd()

		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to obtain current working directory\n")
			os.Exit(1)
		}

		return []string{dirname}
	}

	return args
}

func buildFiles(packageDir string) (parse []string, process []string, pkgname string) {
	ctx := build.Default

	ctx.BuildTags = []string{"operators"}

	p, err := ctx.ImportDir(packageDir, 0)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error while importing build: %s\n", err)
		os.Exit(1)
	}

	for _, f := range p.GoFiles {
		parse = append(parse, path.Join(packageDir, f))

		if strings.HasSuffix(f, ".op.go") {
			process = append(process, f)
		}
	}

	return parse, process, p.Name
}

func parseAST(files []string) (fs *token.FileSet, afs []*ast.File, afsmap map[string]*ast.File) {
	fs = token.NewFileSet()
	afs = make([]*ast.File, 0, len(files))
	afsmap = make(map[string]*ast.File)

	for _, f := range files {
		af, err := parser.ParseFile(fs, f, nil, 0)

		if err != nil {
			fmt.Fprintf(os.Stderr, "Error while parsing AST: %s\n", err)
			os.Exit(1)
		}

		afsmap[f] = af
		afs = append(afs, af)
	}

	return fs, afs, afsmap
}

func checkTypes(pkgpath string, importpath string) *TypeCheck {
	if ret, ok := typechecks[importpath]; ok {
		return ret
	}

	parse, process, pkgname := buildFiles(pkgpath)
	fs, afs, afsmap := parseAST(parse)

	var conf types.Config

	conf.Import = importSources

	pp, err := conf.Check(pkgname, fs, afs, nil)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error while type checking: %s\n", err)
		os.Exit(1)
	}

	ret := &TypeCheck{
		Path: pkgpath,
		Package: pp,
		ImportPath: importpath,
		ParseFiles: parse,
		ProcessFiles: process,
		FileSet: fs,
		Ast: afs,
		FilesToAst: afsmap,
	}

	typechecks[importpath] = ret
	return ret
}

func resolvePackage(importpath string, bool tryLocal) *TypeCheck {
	if ret, ok := typechecks[importpath]; ok {
		return ret
	}

	// Try local first
	if tryLocal {
		if _, err := os.Stat(importpath); err == nil {
			return checkTypes(importpath, importpath)
		}
	}

	paths := strings.Split(os.Getenv("GOPATH"), ":")

	for _, p := range paths {
		src := path.Join(p, "src", importpath)

		if _, err := os.Stat(src); err != nil {
			continue
		}

		return checkTypes(src, importpath)
	}

	return nil
}

func importSource(imports map[string]*types.Package, path string) (pkg *types.Package, err error) {
	// Try from source
	ct := resolvePackage(path, false)

	if ct == nil {
		return nil, fmt.Errorf("Could not locate import path %s", path)
	}

	imports[ct.ImportPath] = ct.Package
	return ct.Package, nil

}

func importSources(imports map[string]*types.Package, path string) (pkg *types.Package, err error) {
	if operatorPackages[path] {
		return importSource(imports, path)
	}

	pkg, err = types.GcImport(imports, path)

	if err != nil {
		if path == "C" {
			return nil, fmt.Errorf("go-operators does not have support for packages that use cgo at the moment")
		}

		return importSource(imports, path)
	}

	return pkg, err
}

func replacer(overloads map[ast.Expr]types.OverloadInfo, node ast.Node) ast.Node {
	expr, ok := node.(ast.Expr)

	if !ok {
		return node
	}

	info, ok := overloads[expr]

	if !ok {
		return node
	}

	sel := &ast.SelectorExpr{
		X:   info.Recv,
		Sel: ast.NewIdent(info.Func.Name()),
	}

	args := []ast.Expr{}

	if info.Oper != nil {
		args = append(args, info.Oper)
	}

	// Create function call expression
	call := &ast.CallExpr{
		Fun:  sel,
		Args: args,
	}

	return call
}

func replaceOperators(ct *TypeCheck) {
	overloads := ct.Package.Overloads()

	for _, f := range ct.ProcessFiles {
		af := ct.FilesToAst[path.Join(ct.Path, f)]

		af = replace(func(node ast.Node) ast.Node {
			return replacer(overloads, node)
		}, af).(*ast.File)

		suffix := ".op.go"
		outname := f[:len(f)-len(suffix)] + ".go"

		fn := path.Join(ct.Path, outname)

		of, err := os.Create(fn)

		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create output file: %s\n", err)
			os.Exit(1)
		}

		defer of.Close()

		if opts.Verbose {
			fmt.Println(fn)
		}

		// Write build constraint
		fmt.Fprintln(of, "// +build !operators\n")

		if err := format.Node(of, ct.FileSet, af); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to write code: %s\n", err)
			os.Exit(1)
		}
	}
}

var opts struct {
	Verbose bool `short:"v" long:"verbose" description:"Enable verbose mode"`
}

var operatorPackages = make(map[string]bool)

func main() {
	fp := flags.NewParser(&opts, flags.Default)

	args, err := fp.Parse()

	if err != nil {
		os.Exit(1)
	}

	packageDirs := packageLocations(args)

	for _, p := range packageDirs {
		operatorPackages[p] = true
	}

	for _, p := range packageDirs {
		ct := resolvePackage(p, true)
		replaceOperators(ct)
	}
}
<commit_msg>Error out when package cannot be found<commit_after>package main

import (
	"github.com/jessevdk/go-operators/types"
	"fmt"
	"github.com/jessevdk/go-flags"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path"
	"strings"
)

type TypeCheck struct {
	Path string
	ImportPath string
	Package *types.Package
	FileSet *token.FileSet
	ParseFiles   []string
	ProcessFiles []string
	FilesToAst map[string]*ast.File
	Ast []*ast.File
}

var typechecks = make(map[string]*TypeCheck)

func packageLocations(args []string) []string {
	if len(args) == 0 {
		dirname, err := os.Getwd()

		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to obtain current working directory\n")
			os.Exit(1)
		}

		return []string{dirname}
	}

	return args
}

func buildFiles(packageDir string) (parse []string, process []string, pkgname string) {
	ctx := build.Default

	ctx.BuildTags = []string{"operators"}

	p, err := ctx.ImportDir(packageDir, 0)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error while importing build: %s\n", err)
		os.Exit(1)
	}

	for _, f := range p.GoFiles {
		parse = append(parse, path.Join(packageDir, f))

		if strings.HasSuffix(f, ".op.go") {
			process = append(process, f)
		}
	}

	return parse, process, p.Name
}

func parseAST(files []string) (fs *token.FileSet, afs []*ast.File, afsmap map[string]*ast.File) {
	fs = token.NewFileSet()
	afs = make([]*ast.File, 0, len(files))
	afsmap = make(map[string]*ast.File)

	for _, f := range files {
		af, err := parser.ParseFile(fs, f, nil, 0)

		if err != nil {
			fmt.Fprintf(os.Stderr, "Error while parsing AST: %s\n", err)
			os.Exit(1)
		}

		afsmap[f] = af
		afs = append(afs, af)
	}

	return fs, afs, afsmap
}

func checkTypes(pkgpath string, importpath string) *TypeCheck {
	if ret, ok := typechecks[importpath]; ok {
		return ret
	}

	parse, process, pkgname := buildFiles(pkgpath)
	fs, afs, afsmap := parseAST(parse)

	var conf types.Config

	conf.Import = importSources

	pp, err := conf.Check(pkgname, fs, afs, nil)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error while type checking: %s\n", err)
		os.Exit(1)
	}

	ret := &TypeCheck{
		Path: pkgpath,
		Package: pp,
		ImportPath: importpath,
		ParseFiles: parse,
		ProcessFiles: process,
		FileSet: fs,
		Ast: afs,
		FilesToAst: afsmap,
	}

	typechecks[importpath] = ret
	return ret
}

func resolvePackage(importpath string, bool tryLocal) *TypeCheck {
	if ret, ok := typechecks[importpath]; ok {
		return ret
	}

	// Try local first
	if tryLocal {
		if _, err := os.Stat(importpath); err == nil {
			return checkTypes(importpath, importpath)
		}
	}

	paths := strings.Split(os.Getenv("GOPATH"), ":")

	for _, p := range paths {
		src := path.Join(p, "src", importpath)

		if _, err := os.Stat(src); err != nil {
			continue
		}

		return checkTypes(src, importpath)
	}

	fmt.Fprintf(os.Stderr, "Could not find package %s.\n", importpath)
	os.Exit(1)

	return nil
}

func importSource(imports map[string]*types.Package, path string) (pkg *types.Package, err error) {
	// Try from source
	ct := resolvePackage(path, false)

	if ct == nil {
		return nil, fmt.Errorf("Could not locate import path %s", path)
	}

	imports[ct.ImportPath] = ct.Package
	return ct.Package, nil

}

func importSources(imports map[string]*types.Package, path string) (pkg *types.Package, err error) {
	if operatorPackages[path] {
		return importSource(imports, path)
	}

	pkg, err = types.GcImport(imports, path)

	if err != nil {
		if path == "C" {
			return nil, fmt.Errorf("go-operators does not have support for packages that use cgo at the moment")
		}

		return importSource(imports, path)
	}

	return pkg, err
}

func replacer(overloads map[ast.Expr]types.OverloadInfo, node ast.Node) ast.Node {
	expr, ok := node.(ast.Expr)

	if !ok {
		return node
	}

	info, ok := overloads[expr]

	if !ok {
		return node
	}

	sel := &ast.SelectorExpr{
		X:   info.Recv,
		Sel: ast.NewIdent(info.Func.Name()),
	}

	args := []ast.Expr{}

	if info.Oper != nil {
		args = append(args, info.Oper)
	}

	// Create function call expression
	call := &ast.CallExpr{
		Fun:  sel,
		Args: args,
	}

	return call
}

func replaceOperators(ct *TypeCheck) {
	overloads := ct.Package.Overloads()

	for _, f := range ct.ProcessFiles {
		af := ct.FilesToAst[path.Join(ct.Path, f)]

		af = replace(func(node ast.Node) ast.Node {
			return replacer(overloads, node)
		}, af).(*ast.File)

		suffix := ".op.go"
		outname := f[:len(f)-len(suffix)] + ".go"

		fn := path.Join(ct.Path, outname)

		of, err := os.Create(fn)

		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create output file: %s\n", err)
			os.Exit(1)
		}

		defer of.Close()

		if opts.Verbose {
			fmt.Println(fn)
		}

		// Write build constraint
		fmt.Fprintln(of, "// +build !operators\n")

		if err := format.Node(of, ct.FileSet, af); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to write code: %s\n", err)
			os.Exit(1)
		}
	}
}

var opts struct {
	Verbose bool `short:"v" long:"verbose" description:"Enable verbose mode"`
}

var operatorPackages = make(map[string]bool)

func main() {
	fp := flags.NewParser(&opts, flags.Default)

	args, err := fp.Parse()

	if err != nil {
		os.Exit(1)
	}

	packageDirs := packageLocations(args)

	for _, p := range packageDirs {
		operatorPackages[p] = true
	}

	for _, p := range packageDirs {
		ct := resolvePackage(p, true)
		replaceOperators(ct)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"math"
	"os"
	"text/tabwriter"
)

// Money lets is print money numbers pretty
type Money float64

func (m Money) String() string {
	return fmt.Sprintf("$%.2f", m)
}

var (
	amount   = flag.Float64("amount", 0.0, "loan amount")
	interest = flag.Float64("rate", 0.0, "interest rate")
	months   = flag.Int("months", 360, "length of the loan (360 = 30 years)")
	table    = flag.Bool("table", false, "print amortization table")
	payment  = flag.Float64("payment", 0.0, "override the payment amount")
)

func main() {
	flag.Parse()
	if *amount == 0 || *interest == 0 {
		flag.Usage()
		return
	}

	if *table {
		printAmortizationTable()
	} else {
		periodic := (*interest / 100) / 12
		discount := calculateDiscount(*interest/100, *months)
		payment := *amount / discount
		interestPayment := *amount * periodic

		fmt.Println("Monthly Payment:", Money(payment))
		fmt.Println("First Interest Payment:", Money(interestPayment))
		fmt.Println("First Principal Payment:", Money(payment-interestPayment))
	}
}

func calculateDiscount(interest float64, months int) float64 {
	periodic := interest / 12
	daily := math.Pow(periodic+1, float64(months))
	return (daily - 1) / (periodic * daily)
}

func printAmortizationTable() {
	periodic := (*interest / 100) / 12
	discount := calculateDiscount(*interest/100, *months)
	balance := *amount
	writer := tabwriter.NewWriter(os.Stdout, 0, 8, 2, '\t', 0)
	monthlyPayment := *payment
	if monthlyPayment == 0 {
		monthlyPayment = balance / discount
	}

	headers := "Opening\tPayment\tInterest\tPrincipal\tEnding\n"
	writer.Write([]byte(headers))

	for period := 1; period <= *months && balance > 0; period++ {
		interestPayment := balance * periodic
		principalPayment := monthlyPayment - interestPayment

		if principalPayment > balance {
			principalPayment = balance
			monthlyPayment = interestPayment + principalPayment
		}

		line := fmt.Sprintf("%s\t%s\t%s\t%s\t%s\n",
			Money(balance), Money(monthlyPayment), Money(interestPayment),
			Money(principalPayment), Money(balance-principalPayment))

		writer.Write([]byte(line))
		balance -= principalPayment
	}

	writer.Flush()
}
<commit_msg>print payment period on amortization table<commit_after>package main

import (
	"flag"
	"fmt"
	"math"
	"os"
	"text/tabwriter"
)

// Money lets is print money numbers pretty
type Money float64

func (m Money) String() string {
	return fmt.Sprintf("$%.2f", m)
}

var (
	amount   = flag.Float64("amount", 0.0, "loan amount")
	interest = flag.Float64("rate", 0.0, "interest rate")
	months   = flag.Int("months", 360, "length of the loan (360 = 30 years)")
	table    = flag.Bool("table", false, "print amortization table")
	payment  = flag.Float64("payment", 0.0, "override the payment amount")
)

func main() {
	flag.Parse()
	if *amount == 0 || *interest == 0 {
		flag.Usage()
		return
	}

	if *table {
		printAmortizationTable()
	} else {
		periodic := (*interest / 100) / 12
		discount := calculateDiscount(*interest/100, *months)
		payment := *amount / discount
		interestPayment := *amount * periodic

		fmt.Println("Monthly Payment:", Money(payment))
		fmt.Println("First Interest Payment:", Money(interestPayment))
		fmt.Println("First Principal Payment:", Money(payment-interestPayment))
	}
}

func calculateDiscount(interest float64, months int) float64 {
	periodic := interest / 12
	daily := math.Pow(periodic+1, float64(months))
	return (daily - 1) / (periodic * daily)
}

func printAmortizationTable() {
	periodic := (*interest / 100) / 12
	discount := calculateDiscount(*interest/100, *months)
	balance := *amount
	writer := tabwriter.NewWriter(os.Stdout, 0, 8, 2, '\t', 0)
	monthlyPayment := *payment
	if monthlyPayment == 0 {
		monthlyPayment = balance / discount
	}

	headers := "Period\tOpening\tPayment\tInterest\tPrincipal\tEnding\n"
	writer.Write([]byte(headers))

	for period := 1; period <= *months && balance > 0; period++ {
		interestPayment := balance * periodic
		principalPayment := monthlyPayment - interestPayment

		if principalPayment > balance {
			principalPayment = balance
			monthlyPayment = interestPayment + principalPayment
		}

		line := fmt.Sprintf("%d\t%s\t%s\t%s\t%s\t%s\n", period,
			Money(balance), Money(monthlyPayment), Money(interestPayment),
			Money(principalPayment), Money(balance-principalPayment))

		writer.Write([]byte(line))
		balance -= principalPayment
	}

	writer.Flush()
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]

// The MIG loader is a simple bootstrapping tool for MIG. It can be scheduled
// to run on a host system and download the newest available version of the
// agent. If the loader identifies a newer version of the agent available, it
// will download the required files from the API, replace the existing files,
// and notify any existing agent it should terminate.
package main

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"github.com/jvehent/cljs"
	"io"
	"io/ioutil"
	"mig.ninja/mig"
	"mig.ninja/mig/mig-agent/agentcontext"
	"mig.ninja/mig/pgp"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
)

var ctx Context
var haveChanges bool
var apiManifest *mig.ManifestResponse
var wg sync.WaitGroup

func initializeHaveBundle() (ret []mig.BundleDictionaryEntry, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("initializeHaveBundle() -> %v", e)
		}
	}()

	ret, err = mig.GetHostBundle()
	if err != nil {
		panic(err)
	}
	ret, err = mig.HashBundle(ret)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{Desc: "initialized local bundle information"}
	for _, x := range ret {
		hv := x.SHA256
		if hv == "" {
			hv = "not found"
		}
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("%v %v -> %v", x.Name, x.Path, hv)}
	}
	return
}

func requestManifest() (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("requestManifest() -> %v", e)
		}
	}()

	murl := APIURL + "manifest/agent/"
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("requesting manifest from %v", murl)}

	mparam := mig.ManifestParameters{}
	mparam.AgentIdentifier = ctx.AgentIdentifier
	buf, err := json.Marshal(mparam)
	if err != nil {
		panic(err)
	}
	mstring := string(buf)
	data := url.Values{"parameters": {mstring}}
	r, err := http.NewRequest("POST", murl, strings.NewReader(data.Encode()))
	if err != nil {
		panic(err)
	}
	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	r.Header.Add("X-LOADERKEY", ctx.LoaderKey)
	client := http.Client{}
	resp, err := client.Do(r)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}

	var resource *cljs.Resource
	err = json.Unmarshal(body, &resource)
	if err != nil {
		panic(err)
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("HTTP %v, API call failed with error '%v' (code %s)", resp.StatusCode,
			resource.Collection.Error.Message, resource.Collection.Error.Code)
		panic(err)
	}

	// Extract our manifest from the response.
	manifest, err := valueToManifest(resource.Collection.Items[0].Data[0].Value)
	if err != nil {
		panic(err)
	}
	apiManifest = &manifest
	return checkManifestSignature(apiManifest)
}

func valueToManifest(v interface{}) (m mig.ManifestResponse, err error) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	err = json.Unmarshal(b, &m)
	return
}

func valueToFetchResponse(v interface{}) (m mig.ManifestFetchResponse, err error) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	err = json.Unmarshal(b, &m)
	return
}

func fetchFile(n string) (ret []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("fetchFile() -> %v", e)
		}
	}()

	murl := APIURL + "manifest/fetch/"

	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("fetching file from %v", murl)}
	mparam := mig.ManifestParameters{}
	mparam.AgentIdentifier = ctx.AgentIdentifier
	mparam.Object = n
	buf, err := json.Marshal(mparam)
	if err != nil {
		panic(err)
	}
	mstring := string(buf)
	data := url.Values{"parameters": {mstring}}
	r, err := http.NewRequest("POST", murl, strings.NewReader(data.Encode()))
	if err != nil {
		panic(err)
	}
	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	r.Header.Add("X-LOADERKEY", ctx.LoaderKey)
	client := http.Client{}
	resp, err := client.Do(r)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}

	var resource *cljs.Resource
	err = json.Unmarshal(body, &resource)
	if err != nil {
		panic(err)
	}

	// Extract fetch response.
	fetchresp, err := valueToFetchResponse(resource.Collection.Items[0].Data[0].Value)
	if err != nil {
		panic(err)
	}

	// Decompress the returned file and return it as a byte slice.
	b := bytes.NewBuffer(fetchresp.Data)
	gz, err := gzip.NewReader(b)
	if err != nil {
		panic(err)
	}
	ret, err = ioutil.ReadAll(gz)
	if err != nil {
		panic(err)
	}
	return
}

func fetchAndReplace(entry mig.BundleDictionaryEntry, sig string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("fetchAndReplace() -> %v", e)
		}
	}()

	// Grab the new file from the API.
	filebuf, err := fetchFile(entry.Name)
	if err != nil {
		panic(err)
	}

	// Stage the new file. Write the file recieved from the API to the
	// file system and validate the signature of the new file to make
	// sure it matches the signature from the manifest.
	//
	// Append .loader to the file name to use as the staged file path.
	reppath := entry.Path + ".loader"
	fd, err := os.OpenFile(reppath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0700)
	if err != nil {
		panic(err)
	}
	_, err = fd.Write(filebuf)
	if err != nil {
		panic(err)
	}
	fd.Close()

	// Validate the signature on the new file.
	ctx.Channels.Log <- mig.Log{Desc: "validating staged file signature"}
	h := sha256.New()
	fd, err = os.Open(reppath)
	if err != nil {
		panic(err)
	}
	buf := make([]byte, 4096)
	for {
		n, err := fd.Read(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			fd.Close()
			panic(err)
		}
		if n > 0 {
			h.Write(buf[:n])
		}
	}
	fd.Close()
	if sig != fmt.Sprintf("%x", h.Sum(nil)) {
		panic("staged file signature mismatch")
	}

	// Got this far, OK to proceed with the replacement.
	ctx.Channels.Log <- mig.Log{Desc: "installing staged file"}
	err = os.Rename(reppath, entry.Path)
	if err != nil {
		panic(err)
	}
	return
}

func checkEntry(entry mig.BundleDictionaryEntry) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("checkEntry() -> %v", e)
		}
	}()

	var compare mig.ManifestEntry
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("comparing %v %v", entry.Name, entry.Path)}
	found := false
	for _, x := range apiManifest.Entries {
		if x.Name == entry.Name {
			compare = x
			found = true
			break
		}
	}
	if !found {
		ctx.Channels.Log <- mig.Log{Desc: "entry not in API manifest, ignoring"}
		return
	}
	hv := entry.SHA256
	if hv == "" {
		hv = "not found"
	}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("we have %v", hv)}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("they have %v", compare.SHA256)}
	if entry.SHA256 == compare.SHA256 {
		ctx.Channels.Log <- mig.Log{Desc: "nothing to do here"}
		return
	}
	haveChanges = true
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("refreshing %v", entry.Name)}
	err = fetchAndReplace(entry, compare.SHA256)
	if err != nil {
		panic(err)
	}
	return
}

// Compare the manifest that the API sent with our knowledge of what is
// currently installed. For each case there is a difference, we will
// request the new file and replace the existing entry.
func compareManifest(have []mig.BundleDictionaryEntry) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("compareManifest() -> %v", e)
		}
	}()

	for _, x := range have {
		err := checkEntry(x)
		if err != nil {
			panic(err)
		}
	}
	return
}

func checkManifestSignature(mr *mig.ManifestResponse) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("checkManifestSignature() -> %v", e)
		}
	}()

	var keys [][]byte
	for _, pk := range MANIFESTPGPKEYS {
		keys = append(keys, []byte(pk))
	}
	keyring, _, err := pgp.ArmoredKeysToKeyring(keys)
	if err != nil {
		panic(err)
	}
	cnt, err := mr.VerifySignatures(keyring)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("%v valid signatures in manifest", cnt)}
	if cnt < REQUIREDSIGNATURES {
		err = fmt.Errorf("Not enough valid signatures (got %v, need %v), rejecting",
			cnt, REQUIREDSIGNATURES)
		panic(err)
	}

	return
}

func initContext() (ctx Context, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("initContext() -> %v", e)
		}
	}()

	ctx.Channels.Log = make(chan mig.Log, 37)
	ctx.Logging = LOGGINGCONF
	ctx.Logging, err = mig.InitLogger(ctx.Logging, "mig-loader")
	if err != nil {
		panic(err)
	}
	ctx.LoaderKey = LOADERKEY

	hints := agentcontext.AgentContextHints{
		DiscoverPublicIP: DISCOVERPUBLICIP,
		DiscoverAWSMeta:  DISCOVERAWSMETA,
		APIUrl:           APIURL,
		Proxies:          PROXIES[:],
	}
	actx, err := agentcontext.NewAgentContext(ctx.Channels.Log, hints)
	if err != nil {
		panic(err)
	}
	ctx.AgentIdentifier = actx.ToAgent()

	return
}

func doExit(v int) {
	close(ctx.Channels.Log)
	wg.Wait()
	os.Exit(v)
}

func main() {
	var err error
	runtime.GOMAXPROCS(1)

	ctx, err = initContext()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	wg.Add(1)
	go func() {
		var stop bool
		for event := range ctx.Channels.Log {
			stop, err = mig.ProcessLog(ctx.Logging, event)
			if err != nil {
				panic("unable to process log")
			}
			if stop {
				break
			}
		}
		wg.Done()
	}()
	ctx.Channels.Log <- mig.Log{Desc: "logging routine started"}

	// Get our current status from the file system.
	have, err := initializeHaveBundle()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	// Retrieve our manifest from the API.
	err = requestManifest()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	err = compareManifest(have)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	if haveChanges {
		err = runTriggers()
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			doExit(1)
		}
	}
	doExit(0)
}
<commit_msg>[minor] use loader key on file system if present<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]

// The MIG loader is a simple bootstrapping tool for MIG. It can be scheduled
// to run on a host system and download the newest available version of the
// agent. If the loader identifies a newer version of the agent available, it
// will download the required files from the API, replace the existing files,
// and notify any existing agent it should terminate.
package main

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"github.com/jvehent/cljs"
	"io"
	"io/ioutil"
	"mig.ninja/mig"
	"mig.ninja/mig/mig-agent/agentcontext"
	"mig.ninja/mig/pgp"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
)

var ctx Context
var haveChanges bool
var apiManifest *mig.ManifestResponse
var wg sync.WaitGroup

func initializeHaveBundle() (ret []mig.BundleDictionaryEntry, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("initializeHaveBundle() -> %v", e)
		}
	}()

	ret, err = mig.GetHostBundle()
	if err != nil {
		panic(err)
	}
	ret, err = mig.HashBundle(ret)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{Desc: "initialized local bundle information"}
	for _, x := range ret {
		hv := x.SHA256
		if hv == "" {
			hv = "not found"
		}
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("%v %v -> %v", x.Name, x.Path, hv)}
	}
	return
}

func requestManifest() (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("requestManifest() -> %v", e)
		}
	}()

	murl := APIURL + "manifest/agent/"
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("requesting manifest from %v", murl)}

	mparam := mig.ManifestParameters{}
	mparam.AgentIdentifier = ctx.AgentIdentifier
	buf, err := json.Marshal(mparam)
	if err != nil {
		panic(err)
	}
	mstring := string(buf)
	data := url.Values{"parameters": {mstring}}
	r, err := http.NewRequest("POST", murl, strings.NewReader(data.Encode()))
	if err != nil {
		panic(err)
	}
	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	r.Header.Add("X-LOADERKEY", ctx.LoaderKey)
	client := http.Client{}
	resp, err := client.Do(r)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}

	var resource *cljs.Resource
	err = json.Unmarshal(body, &resource)
	if err != nil {
		panic(err)
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("HTTP %v, API call failed with error '%v' (code %s)", resp.StatusCode,
			resource.Collection.Error.Message, resource.Collection.Error.Code)
		panic(err)
	}

	// Extract our manifest from the response.
	manifest, err := valueToManifest(resource.Collection.Items[0].Data[0].Value)
	if err != nil {
		panic(err)
	}
	apiManifest = &manifest
	return checkManifestSignature(apiManifest)
}

func valueToManifest(v interface{}) (m mig.ManifestResponse, err error) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	err = json.Unmarshal(b, &m)
	return
}

func valueToFetchResponse(v interface{}) (m mig.ManifestFetchResponse, err error) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	err = json.Unmarshal(b, &m)
	return
}

func fetchFile(n string) (ret []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("fetchFile() -> %v", e)
		}
	}()

	murl := APIURL + "manifest/fetch/"

	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("fetching file from %v", murl)}
	mparam := mig.ManifestParameters{}
	mparam.AgentIdentifier = ctx.AgentIdentifier
	mparam.Object = n
	buf, err := json.Marshal(mparam)
	if err != nil {
		panic(err)
	}
	mstring := string(buf)
	data := url.Values{"parameters": {mstring}}
	r, err := http.NewRequest("POST", murl, strings.NewReader(data.Encode()))
	if err != nil {
		panic(err)
	}
	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	r.Header.Add("X-LOADERKEY", ctx.LoaderKey)
	client := http.Client{}
	resp, err := client.Do(r)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}

	var resource *cljs.Resource
	err = json.Unmarshal(body, &resource)
	if err != nil {
		panic(err)
	}

	// Extract fetch response.
	fetchresp, err := valueToFetchResponse(resource.Collection.Items[0].Data[0].Value)
	if err != nil {
		panic(err)
	}

	// Decompress the returned file and return it as a byte slice.
	b := bytes.NewBuffer(fetchresp.Data)
	gz, err := gzip.NewReader(b)
	if err != nil {
		panic(err)
	}
	ret, err = ioutil.ReadAll(gz)
	if err != nil {
		panic(err)
	}
	return
}

func fetchAndReplace(entry mig.BundleDictionaryEntry, sig string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("fetchAndReplace() -> %v", e)
		}
	}()

	// Grab the new file from the API.
	filebuf, err := fetchFile(entry.Name)
	if err != nil {
		panic(err)
	}

	// Stage the new file. Write the file recieved from the API to the
	// file system and validate the signature of the new file to make
	// sure it matches the signature from the manifest.
	//
	// Append .loader to the file name to use as the staged file path.
	reppath := entry.Path + ".loader"
	fd, err := os.OpenFile(reppath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0700)
	if err != nil {
		panic(err)
	}
	_, err = fd.Write(filebuf)
	if err != nil {
		panic(err)
	}
	fd.Close()

	// Validate the signature on the new file.
	ctx.Channels.Log <- mig.Log{Desc: "validating staged file signature"}
	h := sha256.New()
	fd, err = os.Open(reppath)
	if err != nil {
		panic(err)
	}
	buf := make([]byte, 4096)
	for {
		n, err := fd.Read(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			fd.Close()
			panic(err)
		}
		if n > 0 {
			h.Write(buf[:n])
		}
	}
	fd.Close()
	if sig != fmt.Sprintf("%x", h.Sum(nil)) {
		panic("staged file signature mismatch")
	}

	// Got this far, OK to proceed with the replacement.
	ctx.Channels.Log <- mig.Log{Desc: "installing staged file"}
	err = os.Rename(reppath, entry.Path)
	if err != nil {
		panic(err)
	}
	return
}

func checkEntry(entry mig.BundleDictionaryEntry) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("checkEntry() -> %v", e)
		}
	}()

	var compare mig.ManifestEntry
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("comparing %v %v", entry.Name, entry.Path)}
	found := false
	for _, x := range apiManifest.Entries {
		if x.Name == entry.Name {
			compare = x
			found = true
			break
		}
	}
	if !found {
		ctx.Channels.Log <- mig.Log{Desc: "entry not in API manifest, ignoring"}
		return
	}
	hv := entry.SHA256
	if hv == "" {
		hv = "not found"
	}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("we have %v", hv)}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("they have %v", compare.SHA256)}
	if entry.SHA256 == compare.SHA256 {
		ctx.Channels.Log <- mig.Log{Desc: "nothing to do here"}
		return
	}
	haveChanges = true
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("refreshing %v", entry.Name)}
	err = fetchAndReplace(entry, compare.SHA256)
	if err != nil {
		panic(err)
	}
	return
}

// Compare the manifest that the API sent with our knowledge of what is
// currently installed. For each case there is a difference, we will
// request the new file and replace the existing entry.
func compareManifest(have []mig.BundleDictionaryEntry) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("compareManifest() -> %v", e)
		}
	}()

	for _, x := range have {
		err := checkEntry(x)
		if err != nil {
			panic(err)
		}
	}
	return
}

func checkManifestSignature(mr *mig.ManifestResponse) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("checkManifestSignature() -> %v", e)
		}
	}()

	var keys [][]byte
	for _, pk := range MANIFESTPGPKEYS {
		keys = append(keys, []byte(pk))
	}
	keyring, _, err := pgp.ArmoredKeysToKeyring(keys)
	if err != nil {
		panic(err)
	}
	cnt, err := mr.VerifySignatures(keyring)
	if err != nil {
		panic(err)
	}
	ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("%v valid signatures in manifest", cnt)}
	if cnt < REQUIREDSIGNATURES {
		err = fmt.Errorf("Not enough valid signatures (got %v, need %v), rejecting",
			cnt, REQUIREDSIGNATURES)
		panic(err)
	}

	return
}

func initContext() (ctx Context, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("initContext() -> %v", e)
		}
	}()

	ctx.Channels.Log = make(chan mig.Log, 37)
	ctx.Logging = LOGGINGCONF
	ctx.Logging, err = mig.InitLogger(ctx.Logging, "mig-loader")
	if err != nil {
		panic(err)
	}
	ctx.LoaderKey = LOADERKEY

	hints := agentcontext.AgentContextHints{
		DiscoverPublicIP: DISCOVERPUBLICIP,
		DiscoverAWSMeta:  DISCOVERAWSMETA,
		APIUrl:           APIURL,
		Proxies:          PROXIES[:],
	}
	actx, err := agentcontext.NewAgentContext(ctx.Channels.Log, hints)
	if err != nil {
		panic(err)
	}
	ctx.AgentIdentifier = actx.ToAgent()

	return
}

func doExit(v int) {
	close(ctx.Channels.Log)
	wg.Wait()
	os.Exit(v)
}

// Return the path to the expected loader key file location
func getLoaderKeyfile() string {
	switch runtime.GOOS {
	case "linux":
		return "/etc/mig/mig-loader.key"
	case "darwin":
		return "/etc/mig/mig-loader.key"
	}
	panic("loader does not support this operating system")
	return ""
}

// Attempt to obtain the loader key from the file system and override the
// built-in secret
func loadLoaderKey() error {
	fd, err := os.Open(getLoaderKeyfile())
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	defer fd.Close()
	buf, _, err := bufio.NewReader(fd).ReadLine()
	if err != nil {
		// Nothing in the loader key file
		if err == io.EOF {
			return nil
		}
		return err
	}
	LOADERKEY = string(buf)
	return nil
}

func main() {
	var err error
	runtime.GOMAXPROCS(1)

	err = loadLoaderKey()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	ctx, err = initContext()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	wg.Add(1)
	go func() {
		var stop bool
		for event := range ctx.Channels.Log {
			stop, err = mig.ProcessLog(ctx.Logging, event)
			if err != nil {
				panic("unable to process log")
			}
			if stop {
				break
			}
		}
		wg.Done()
	}()
	ctx.Channels.Log <- mig.Log{Desc: "logging routine started"}

	// Get our current status from the file system.
	have, err := initializeHaveBundle()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	// Retrieve our manifest from the API.
	err = requestManifest()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	err = compareManifest(have)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		doExit(1)
	}

	if haveChanges {
		err = runTriggers()
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			doExit(1)
		}
	}
	doExit(0)
}
<|endoftext|>