<commit_before>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/routing"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus

	cachedAddrs       []ma.Multiaddr
	cachedAddrsExpiry time.Time
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			ar.mx.Lock()
			ar.cachedAddrs = nil
			ar.mx.Unlock()
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	if ar.numRelays() >= DesiredRelays {
		return false
	}

	update := false
	for retry := 0; retry < 5; retry++ {
		if retry > 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
			case <-ctx.Done():
				return update
			}
		}

		update = ar.findRelaysOnce(ctx) || update
		if ar.numRelays() > 0 {
			return update
		}
	}
	return update
}

func (ar *AutoRelay) findRelaysOnce(ctx context.Context) bool {
	pis, err := ar.discoverRelays(ctx)
	if err != nil {
		log.Debugf("error discovering relays: %s", err)
		return false
	}
	log.Debugf("discovered %d relays", len(pis))
	pis = ar.selectRelays(ctx, pis)
	log.Debugf("selected %d relays", len(pis))

	update := false
	for _, pi := range pis {
		update = ar.tryRelay(ctx, pi) || update
		if ar.numRelays() >= DesiredRelays {
			break
		}
	}
	return update
}

func (ar *AutoRelay) numRelays() int {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	return len(ar.relays)
}

// usingRelay returns if we're currently using the given relay.
func (ar *AutoRelay) usingRelay(p peer.ID) bool {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	_, ok := ar.relays[p]
	return ok
}

// addRelay adds the given relay to our set of relays.
// returns true when we add a new relay
func (ar *AutoRelay) tryRelay(ctx context.Context, pi peer.AddrInfo) bool {
	if ar.usingRelay(pi.ID) {
		return false
	}

	if !ar.connect(ctx, pi) {
		return false
	}

	ar.mx.Lock()
	defer ar.mx.Unlock()

	// make sure we're still connected.
	if ar.host.Network().Connectedness(pi.ID) != network.Connected {
		return false
	}
	ar.relays[pi.ID] = struct{}{}

	return true
}

func (ar *AutoRelay) connect(ctx context.Context, pi peer.AddrInfo) bool {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	if len(pi.Addrs) == 0 {
		var err error
		pi, err = ar.router.FindPeer(ctx, pi.ID)
		if err != nil {
			log.Debugf("error finding relay peer %s: %s", pi.ID, err.Error())
			return false
		}
	}

	err := ar.host.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
		return false
	}

	// tag the connection as very important
	ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)
	return true
}

func (ar *AutoRelay) discoverRelays(ctx context.Context) ([]peer.AddrInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return discovery.FindPeers(ctx, ar.discover, RelayRendezvous, discovery.Limit(1000))
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []peer.AddrInfo) []peer.AddrInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	shuffleRelays(pis)
	return pis
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.status != autonat.NATStatusPrivate {
		return addrs
	}

	if ar.cachedAddrs != nil && time.Now().Before(ar.cachedAddrsExpiry) {
		return ar.cachedAddrs
	}

	raddrs := make([]ma.Multiaddr, 0, 4*len(ar.relays)+4)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for p := range ar.relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	ar.cachedAddrs = raddrs
	ar.cachedAddrsExpiry = time.Now().Add(30 * time.Second)

	return raddrs
}

func shuffleRelays(pis []peer.AddrInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(network.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(network.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(network.Network, network.Conn)   {}

func (ar *AutoRelay) Disconnected(net network.Network, c network.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == network.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(network.Network, network.Stream) {}
func (ar *AutoRelay) ClosedStream(network.Network, network.Stream) {}
<commit_msg>autorelay: ensure candidate relays can hop<commit_after>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/routing"

	autonat "github.com/libp2p/go-libp2p-autonat"
	circuit "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus

	cachedAddrs       []ma.Multiaddr
	cachedAddrsExpiry time.Time
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			ar.mx.Lock()
			ar.cachedAddrs = nil
			ar.mx.Unlock()
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	if ar.numRelays() >= DesiredRelays {
		return false
	}

	update := false
	for retry := 0; retry < 5; retry++ {
		if retry > 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
			case <-ctx.Done():
				return update
			}
		}

		update = ar.findRelaysOnce(ctx) || update
		if ar.numRelays() > 0 {
			return update
		}
	}
	return update
}

func (ar *AutoRelay) findRelaysOnce(ctx context.Context) bool {
	pis, err := ar.discoverRelays(ctx)
	if err != nil {
		log.Debugf("error discovering relays: %s", err)
		return false
	}
	log.Debugf("discovered %d relays", len(pis))
	pis = ar.selectRelays(ctx, pis)
	log.Debugf("selected %d relays", len(pis))

	update := false
	for _, pi := range pis {
		update = ar.tryRelay(ctx, pi) || update
		if ar.numRelays() >= DesiredRelays {
			break
		}
	}
	return update
}

func (ar *AutoRelay) numRelays() int {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	return len(ar.relays)
}

// usingRelay returns if we're currently using the given relay.
func (ar *AutoRelay) usingRelay(p peer.ID) bool {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	_, ok := ar.relays[p]
	return ok
}

// addRelay adds the given relay to our set of relays.
// returns true when we add a new relay
func (ar *AutoRelay) tryRelay(ctx context.Context, pi peer.AddrInfo) bool {
	if ar.usingRelay(pi.ID) {
		return false
	}

	if !ar.connect(ctx, pi) {
		return false
	}

	ok, err := circuit.CanHop(ctx, ar.host, pi.ID)
	if err != nil {
		log.Debugf("error querying relay: %s", err.Error())
		return false
	}

	if !ok {
		// not a hop relay
		return false
	}

	ar.mx.Lock()
	defer ar.mx.Unlock()

	// make sure we're still connected.
	if ar.host.Network().Connectedness(pi.ID) != network.Connected {
		return false
	}
	ar.relays[pi.ID] = struct{}{}

	return true
}

func (ar *AutoRelay) connect(ctx context.Context, pi peer.AddrInfo) bool {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	if len(pi.Addrs) == 0 {
		var err error
		pi, err = ar.router.FindPeer(ctx, pi.ID)
		if err != nil {
			log.Debugf("error finding relay peer %s: %s", pi.ID, err.Error())
			return false
		}
	}

	err := ar.host.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
		return false
	}

	// tag the connection as very important
	ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)
	return true
}

func (ar *AutoRelay) discoverRelays(ctx context.Context) ([]peer.AddrInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return discovery.FindPeers(ctx, ar.discover, RelayRendezvous, discovery.Limit(1000))
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []peer.AddrInfo) []peer.AddrInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	shuffleRelays(pis)
	return pis
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.status != autonat.NATStatusPrivate {
		return addrs
	}

	if ar.cachedAddrs != nil && time.Now().Before(ar.cachedAddrsExpiry) {
		return ar.cachedAddrs
	}

	raddrs := make([]ma.Multiaddr, 0, 4*len(ar.relays)+4)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for p := range ar.relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	ar.cachedAddrs = raddrs
	ar.cachedAddrsExpiry = time.Now().Add(30 * time.Second)

	return raddrs
}

func shuffleRelays(pis []peer.AddrInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(network.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(network.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(network.Network, network.Conn)   {}

func (ar *AutoRelay) Disconnected(net network.Network, c network.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == network.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(network.Network, network.Stream) {}
func (ar *AutoRelay) ClosedStream(network.Network, network.Stream) {}
<|endoftext|>


<commit_before>package keystore

import (
	"errors"
	"fmt"
	"github.com/cossacklabs/acra/utils"
	"github.com/cossacklabs/acra/zone"
	"github.com/cossacklabs/themis/gothemis/keys"
	log "github.com/sirupsen/logrus"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"crypto/rand"
	"sync"
)

type FilesystemKeyStore struct {
	keys                map[string][]byte
	privateKeyDirectory string
	publicKeyDirectory  string
	directory           string
	lock                *sync.RWMutex
	encryptor           KeyEncryptor
}

func NewFilesystemKeyStore(directory string, encryptor KeyEncryptor) (*FilesystemKeyStore, error) {
	return NewFilesystemKeyStoreTwoPath(directory, directory, encryptor)
}

func NewFilesystemKeyStoreTwoPath(privateKeyFolder, publicKeyFolder string, encryptor KeyEncryptor) (*FilesystemKeyStore, error) {
	// check folder for private key
	directory, err := utils.AbsPath(privateKeyFolder)
	if err != nil {
		return nil, err
	}
	fi, err := os.Stat(directory)
	if nil == err && runtime.GOOS == "linux" && fi.Mode().Perm().String() != "-rwx------" {
		log.Errorln(" key store folder has an incorrect permissions")
		return nil, errors.New("key store folder has an incorrect permissions")
	}
	if privateKeyFolder != publicKeyFolder {
		// check folder for public key
		directory, err = utils.AbsPath(privateKeyFolder)
		if err != nil {
			return nil, err
		}
		fi, err = os.Stat(directory)
		if nil != err && !os.IsNotExist(err) {
			return nil, err
		}
	}
	return &FilesystemKeyStore{privateKeyDirectory: privateKeyFolder, publicKeyDirectory: publicKeyFolder,
		keys: make(map[string][]byte), lock: &sync.RWMutex{}, encryptor: encryptor}, nil
}

func (store *FilesystemKeyStore) generateKeyPair(filename string, clientId []byte) (*keys.Keypair, error) {
	keypair, err := keys.New(keys.KEYTYPE_EC)
	if err != nil {
		return nil, err
	}
	dirpath := filepath.Dir(store.getPrivateKeyFilePath(filename))
	err = os.MkdirAll(dirpath, 0700)
	if err != nil {
		return nil, err
	}

	encryptedPrivate, err := store.encryptor.Encrypt(keypair.Private.Value, clientId)
	if err != nil {
		return nil, err
	}
	err = ioutil.WriteFile(store.getPrivateKeyFilePath(filename), encryptedPrivate, 0600)
	if err != nil {
		return nil, err
	}
	err = ioutil.WriteFile(store.getPublicKeyFilePath(fmt.Sprintf("%s.pub", filename)), keypair.Public.Value, 0644)
	if err != nil {
		return nil, err
	}
	return keypair, nil
}

func (store *FilesystemKeyStore) generateKey(filename string, length uint8) ([]byte, error) {
	randomBytes := make([]byte, length)
	_, err := rand.Read(randomBytes)
	// Note that err == nil only if we read len(b) bytes.
	if err != nil {
		log.Error(err)
		return nil, err
	}
	dirpath := filepath.Dir(store.getPrivateKeyFilePath(filename))
	err = os.MkdirAll(dirpath, 0700)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	err = ioutil.WriteFile(store.getPrivateKeyFilePath(filename), randomBytes, 0600)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	return randomBytes, nil
}

func (store *FilesystemKeyStore) GenerateZoneKey() ([]byte, []byte, error) {
	/* save private key in fs, return id and public key*/
	var id []byte
	for {
		// generate until key not exists
		id = zone.GenerateZoneId()
		if !store.HasZonePrivateKey(id) {
			break
		}
	}

	keypair, err := store.generateKeyPair(getZoneKeyFilename(id), id)
	if err != nil {
		return []byte{}, []byte{}, err
	}
	store.lock.Lock()
	defer store.lock.Unlock()
	// cache key
	store.keys[getZoneKeyFilename(id)] = keypair.Private.Value
	return id, keypair.Public.Value, nil
}

func (store *FilesystemKeyStore) getPrivateKeyFilePath(filename string) string {
	return fmt.Sprintf("%s%s%s", store.privateKeyDirectory, string(os.PathSeparator), filename)
}

func (store *FilesystemKeyStore) getPublicKeyFilePath(filename string) string {
	return fmt.Sprintf("%s%s%s", store.publicKeyDirectory, string(os.PathSeparator), filename)
}

func (store *FilesystemKeyStore) GetZonePrivateKey(id []byte) (*keys.PrivateKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getZoneKeyFilename(id)
	store.lock.Lock()
	defer store.lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PrivateKey{Value: key}, nil
	}
	privateKey, err := utils.LoadPrivateKey(store.getPrivateKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	if privateKey.Value, err = store.encryptor.Decrypt(privateKey.Value, id); err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = privateKey.Value
	return privateKey, nil
}

func (store *FilesystemKeyStore) HasZonePrivateKey(id []byte) bool {
	if !ValidateId(id) {
		return false
	}
	// add caching false answers. now if key doesn't exists than always checks on fs
	// it's system call and slow.
	if len(id) == 0 {
		return false
	}
	fname := getZoneKeyFilename(id)
	store.lock.RLock()
	defer store.lock.RUnlock()
	_, ok := store.keys[fname]
	if ok {
		return true
	}
	exists, _ := utils.FileExists(store.getPrivateKeyFilePath(fname))
	return exists
}

func (store *FilesystemKeyStore) GetPeerPublicKey(id []byte) (*keys.PublicKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getPublicKeyFilename(id)
	store.lock.Lock()
	defer store.lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PublicKey{Value: key}, nil
	}
	publicKey, err := utils.LoadPublicKey(store.getPublicKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = publicKey.Value
	return publicKey, nil
}

func (store *FilesystemKeyStore) GetPrivateKey(id []byte) (*keys.PrivateKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getServerKeyFilename(id)
	lock.Lock()
	defer lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PrivateKey{Value: key}, nil
	}
	privateKey, err := utils.LoadPrivateKey(store.getPrivateKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	if privateKey.Value, err = store.encryptor.Decrypt(privateKey.Value, id); err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = privateKey.Value
	return privateKey, nil
}

func (store *FilesystemKeyStore) GetServerDecryptionPrivateKey(id []byte) (*keys.PrivateKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getServerDecryptionKeyFilename(id)
	store.lock.Lock()
	defer store.lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PrivateKey{Value: key}, nil
	}
	privateKey, err := utils.LoadPrivateKey(store.getPrivateKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	if privateKey.Value, err = store.encryptor.Decrypt(privateKey.Value, id); err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = privateKey.Value
	return privateKey, nil
}

func (store *FilesystemKeyStore) GenerateProxyKeys(id []byte) error {
	if !ValidateId(id) {
		return ErrInvalidClientId
	}
	filename := getProxyKeyFilename(id)

	_, err := store.generateKeyPair(filename, id)
	if err != nil {
		return err
	}
	return nil
}
func (store *FilesystemKeyStore) GenerateServerKeys(id []byte) error {
	if !ValidateId(id) {
		return ErrInvalidClientId
	}
	filename := getServerKeyFilename(id)
	_, err := store.generateKeyPair(filename, id)
	if err != nil {
		return err
	}
	return nil
}

// generate key pair for data encryption/decryption
func (store *FilesystemKeyStore) GenerateDataEncryptionKeys(id []byte) error {
	if !ValidateId(id) {
		return ErrInvalidClientId
	}
	_, err := store.generateKeyPair(getServerDecryptionKeyFilename(id), id)
	if err != nil {
		return err
	}
	return nil
}

// clear all cached keys
func (store *FilesystemKeyStore) Reset() {
	store.keys = make(map[string][]byte)
}

func (store *FilesystemKeyStore) GetPoisonKeyPair() (*keys.Keypair, error) {
	privatePath := store.getPrivateKeyFilePath(POISON_KEY_FILENAME)
	publicPath := store.getPublicKeyFilePath(fmt.Sprintf("%s.pub", POISON_KEY_FILENAME))
	privateExists, err := utils.FileExists(privatePath)
	if err != nil {
		return nil, err
	}
	publicExists, err := utils.FileExists(publicPath)
	if err != nil {
		return nil, err
	}
	if privateExists && publicExists {
		private, err := utils.LoadPrivateKey(privatePath)
		if err != nil {
			return nil, err
		}
		if private.Value, err = store.encryptor.Decrypt(private.Value, []byte(POISON_KEY_FILENAME)); err != nil {
			return nil, err
		}
		public, err := utils.LoadPublicKey(publicPath)
		if err != nil {
			return nil, err
		}
		return &keys.Keypair{Public: public, Private: private}, nil
	}
	log.Infoln("Generate poison key pair")
	return store.generateKeyPair(POISON_KEY_FILENAME, []byte(POISON_KEY_FILENAME))
}

func (store *FilesystemKeyStore) GetAuthKey(remove bool) ([]byte, error) {
	keyPath := store.getPrivateKeyFilePath(BASIC_AUTH_KEY_FILENAME)
	keyExists, err := utils.FileExists(keyPath)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	if keyExists && !remove {
		key, err := utils.ReadFile(keyPath)
		if err != nil {
			log.Error(err)
			return nil, err
		}
		return key, nil
	}
	log.Infof("Generate basic auth key for AcraConfigUI to %v", keyPath)
	return store.generateKey(BASIC_AUTH_KEY_FILENAME, BASIC_AUTH_KEY_LENGTH)
}
<commit_msg>create dir if not exists for public keys (#148)<commit_after>package keystore

import (
	"crypto/rand"
	"errors"
	"fmt"
	"github.com/cossacklabs/acra/utils"
	"github.com/cossacklabs/acra/zone"
	"github.com/cossacklabs/themis/gothemis/keys"
	log "github.com/sirupsen/logrus"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sync"
)

type FilesystemKeyStore struct {
	keys                map[string][]byte
	privateKeyDirectory string
	publicKeyDirectory  string
	directory           string
	lock                *sync.RWMutex
	encryptor           KeyEncryptor
}

func NewFilesystemKeyStore(directory string, encryptor KeyEncryptor) (*FilesystemKeyStore, error) {
	return NewFilesystemKeyStoreTwoPath(directory, directory, encryptor)
}

func NewFilesystemKeyStoreTwoPath(privateKeyFolder, publicKeyFolder string, encryptor KeyEncryptor) (*FilesystemKeyStore, error) {
	// check folder for private key
	directory, err := utils.AbsPath(privateKeyFolder)
	if err != nil {
		return nil, err
	}
	fi, err := os.Stat(directory)
	if nil == err && runtime.GOOS == "linux" && fi.Mode().Perm().String() != "-rwx------" {
		log.Errorln(" key store folder has an incorrect permissions")
		return nil, errors.New("key store folder has an incorrect permissions")
	}
	if privateKeyFolder != publicKeyFolder {
		// check folder for public key
		directory, err = utils.AbsPath(privateKeyFolder)
		if err != nil {
			return nil, err
		}
		fi, err = os.Stat(directory)
		if nil != err && !os.IsNotExist(err) {
			return nil, err
		}
	}
	return &FilesystemKeyStore{privateKeyDirectory: privateKeyFolder, publicKeyDirectory: publicKeyFolder,
		keys: make(map[string][]byte), lock: &sync.RWMutex{}, encryptor: encryptor}, nil
}

func (store *FilesystemKeyStore) generateKeyPair(filename string, clientId []byte) (*keys.Keypair, error) {
	keypair, err := keys.New(keys.KEYTYPE_EC)
	if err != nil {
		return nil, err
	}
	privateKeysFolder := filepath.Dir(store.getPrivateKeyFilePath(filename))
	err = os.MkdirAll(privateKeysFolder, 0700)
	if err != nil {
		return nil, err
	}

	publicKeysFolder := filepath.Dir(store.getPublicKeyFilePath(filename))
	err = os.MkdirAll(publicKeysFolder, 0700)
	if err != nil {
		return nil, err
	}

	encryptedPrivate, err := store.encryptor.Encrypt(keypair.Private.Value, clientId)
	if err != nil {
		return nil, err
	}
	err = ioutil.WriteFile(store.getPrivateKeyFilePath(filename), encryptedPrivate, 0600)
	if err != nil {
		return nil, err
	}
	err = ioutil.WriteFile(store.getPublicKeyFilePath(fmt.Sprintf("%s.pub", filename)), keypair.Public.Value, 0644)
	if err != nil {
		return nil, err
	}
	return keypair, nil
}

func (store *FilesystemKeyStore) generateKey(filename string, length uint8) ([]byte, error) {
	randomBytes := make([]byte, length)
	_, err := rand.Read(randomBytes)
	// Note that err == nil only if we read len(b) bytes.
	if err != nil {
		log.Error(err)
		return nil, err
	}
	dirpath := filepath.Dir(store.getPrivateKeyFilePath(filename))
	err = os.MkdirAll(dirpath, 0700)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	err = ioutil.WriteFile(store.getPrivateKeyFilePath(filename), randomBytes, 0600)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	return randomBytes, nil
}

func (store *FilesystemKeyStore) GenerateZoneKey() ([]byte, []byte, error) {
	/* save private key in fs, return id and public key*/
	var id []byte
	for {
		// generate until key not exists
		id = zone.GenerateZoneId()
		if !store.HasZonePrivateKey(id) {
			break
		}
	}

	keypair, err := store.generateKeyPair(getZoneKeyFilename(id), id)
	if err != nil {
		return []byte{}, []byte{}, err
	}
	store.lock.Lock()
	defer store.lock.Unlock()
	// cache key
	store.keys[getZoneKeyFilename(id)] = keypair.Private.Value
	return id, keypair.Public.Value, nil
}

func (store *FilesystemKeyStore) getPrivateKeyFilePath(filename string) string {
	return fmt.Sprintf("%s%s%s", store.privateKeyDirectory, string(os.PathSeparator), filename)
}

func (store *FilesystemKeyStore) getPublicKeyFilePath(filename string) string {
	return fmt.Sprintf("%s%s%s", store.publicKeyDirectory, string(os.PathSeparator), filename)
}

func (store *FilesystemKeyStore) GetZonePrivateKey(id []byte) (*keys.PrivateKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getZoneKeyFilename(id)
	store.lock.Lock()
	defer store.lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PrivateKey{Value: key}, nil
	}
	privateKey, err := utils.LoadPrivateKey(store.getPrivateKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	if privateKey.Value, err = store.encryptor.Decrypt(privateKey.Value, id); err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = privateKey.Value
	return privateKey, nil
}

func (store *FilesystemKeyStore) HasZonePrivateKey(id []byte) bool {
	if !ValidateId(id) {
		return false
	}
	// add caching false answers. now if key doesn't exists than always checks on fs
	// it's system call and slow.
	if len(id) == 0 {
		return false
	}
	fname := getZoneKeyFilename(id)
	store.lock.RLock()
	defer store.lock.RUnlock()
	_, ok := store.keys[fname]
	if ok {
		return true
	}
	exists, _ := utils.FileExists(store.getPrivateKeyFilePath(fname))
	return exists
}

func (store *FilesystemKeyStore) GetPeerPublicKey(id []byte) (*keys.PublicKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getPublicKeyFilename(id)
	store.lock.Lock()
	defer store.lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PublicKey{Value: key}, nil
	}
	publicKey, err := utils.LoadPublicKey(store.getPublicKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = publicKey.Value
	return publicKey, nil
}

func (store *FilesystemKeyStore) GetPrivateKey(id []byte) (*keys.PrivateKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getServerKeyFilename(id)
	lock.Lock()
	defer lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PrivateKey{Value: key}, nil
	}
	privateKey, err := utils.LoadPrivateKey(store.getPrivateKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	if privateKey.Value, err = store.encryptor.Decrypt(privateKey.Value, id); err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = privateKey.Value
	return privateKey, nil
}

func (store *FilesystemKeyStore) GetServerDecryptionPrivateKey(id []byte) (*keys.PrivateKey, error) {
	if !ValidateId(id) {
		return nil, ErrInvalidClientId
	}
	fname := getServerDecryptionKeyFilename(id)
	store.lock.Lock()
	defer store.lock.Unlock()
	key, ok := store.keys[fname]
	if ok {
		log.Debugf("load cached key: %s", fname)
		return &keys.PrivateKey{Value: key}, nil
	}
	privateKey, err := utils.LoadPrivateKey(store.getPrivateKeyFilePath(fname))
	if err != nil {
		return nil, err
	}
	if privateKey.Value, err = store.encryptor.Decrypt(privateKey.Value, id); err != nil {
		return nil, err
	}
	log.Debugf("load key from fs: %s", fname)
	store.keys[fname] = privateKey.Value
	return privateKey, nil
}

func (store *FilesystemKeyStore) GenerateProxyKeys(id []byte) error {
	if !ValidateId(id) {
		return ErrInvalidClientId
	}
	filename := getProxyKeyFilename(id)

	_, err := store.generateKeyPair(filename, id)
	if err != nil {
		return err
	}
	return nil
}
func (store *FilesystemKeyStore) GenerateServerKeys(id []byte) error {
	if !ValidateId(id) {
		return ErrInvalidClientId
	}
	filename := getServerKeyFilename(id)
	_, err := store.generateKeyPair(filename, id)
	if err != nil {
		return err
	}
	return nil
}

// generate key pair for data encryption/decryption
func (store *FilesystemKeyStore) GenerateDataEncryptionKeys(id []byte) error {
	if !ValidateId(id) {
		return ErrInvalidClientId
	}
	_, err := store.generateKeyPair(getServerDecryptionKeyFilename(id), id)
	if err != nil {
		return err
	}
	return nil
}

// clear all cached keys
func (store *FilesystemKeyStore) Reset() {
	store.keys = make(map[string][]byte)
}

func (store *FilesystemKeyStore) GetPoisonKeyPair() (*keys.Keypair, error) {
	privatePath := store.getPrivateKeyFilePath(POISON_KEY_FILENAME)
	publicPath := store.getPublicKeyFilePath(fmt.Sprintf("%s.pub", POISON_KEY_FILENAME))
	privateExists, err := utils.FileExists(privatePath)
	if err != nil {
		return nil, err
	}
	publicExists, err := utils.FileExists(publicPath)
	if err != nil {
		return nil, err
	}
	if privateExists && publicExists {
		private, err := utils.LoadPrivateKey(privatePath)
		if err != nil {
			return nil, err
		}
		if private.Value, err = store.encryptor.Decrypt(private.Value, []byte(POISON_KEY_FILENAME)); err != nil {
			return nil, err
		}
		public, err := utils.LoadPublicKey(publicPath)
		if err != nil {
			return nil, err
		}
		return &keys.Keypair{Public: public, Private: private}, nil
	}
	log.Infoln("Generate poison key pair")
	return store.generateKeyPair(POISON_KEY_FILENAME, []byte(POISON_KEY_FILENAME))
}

func (store *FilesystemKeyStore) GetAuthKey(remove bool) ([]byte, error) {
	keyPath := store.getPrivateKeyFilePath(BASIC_AUTH_KEY_FILENAME)
	keyExists, err := utils.FileExists(keyPath)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	if keyExists && !remove {
		key, err := utils.ReadFile(keyPath)
		if err != nil {
			log.Error(err)
			return nil, err
		}
		return key, nil
	}
	log.Infof("Generate basic auth key for AcraConfigUI to %v", keyPath)
	return store.generateKey(BASIC_AUTH_KEY_FILENAME, BASIC_AUTH_KEY_LENGTH)
}
<|endoftext|>


<commit_before>package indicators

import (
	"errors"
	"github.com/thetruetrade/gotrade"
)

// A Relative Strength Indicator (Rsi), no storage, for use in other indicators
type RsiWithoutStorage struct {
	*baseIndicator
	*baseFloatBounds

	// private variables
	valueAvailableAction ValueAvailableActionFloat
	periodCounter        int
	previousClose        float64
	previousGain         float64
	previousLoss         float64
	timePeriod           int
}

// NewRsiWithoutStorage creates a Relative Strength Indicator (Rsi) without storage
func NewRsiWithoutStorage(timePeriod int, valueAvailableAction ValueAvailableActionFloat) (indicator *RsiWithoutStorage, err error) {

	// an indicator without storage MUST have a value available action
	if valueAvailableAction == nil {
		return nil, ErrValueAvailableActionIsNil
	}

	// the minimum timeperiod for this indicator is 2
	if timePeriod < 2 {
		return nil, errors.New("timePeriod is less than the minimum (2)")
	}

	// check the maximum timeperiod
	if timePeriod > MaximumLookbackPeriod {
		return nil, errors.New("timePeriod is greater than the maximum (100000)")
	}

	lookback := timePeriod
	ind := RsiWithoutStorage{
		baseIndicator:        newBaseIndicator(lookback),
		baseFloatBounds:      newBaseFloatBounds(),
		periodCounter:        (timePeriod * -1) - 1,
		previousClose:        0.0,
		previousGain:         0.0,
		previousLoss:         0.0,
		valueAvailableAction: valueAvailableAction,
		timePeriod:           timePeriod,
	}

	return &ind, err
}

// A Relative Strength Indicator (Rsi)
type Rsi struct {
	*RsiWithoutStorage
	selectData gotrade.DataSelectionFunc

	// public variables
	Data []float64
}

// NewRsi creates a Relative Strength Indicator (Rsi) for online usage
func NewRsi(timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind := Rsi{selectData: selectData}
	ind.RsiWithoutStorage, err = NewRsiWithoutStorage(timePeriod,
		func(dataItem float64, streamBarIndex int) {
			ind.Data = append(ind.Data, dataItem)
		})

	return &ind, err
}

// NewDefaultRsi creates a Relative Strength Indicator (Rsi) for online usage with default parameters
//	- timePeriod: 14
func NewDefaultRsi() (indicator *Rsi, err error) {
	timePeriod := 14
	return NewRsi(timePeriod, gotrade.UseClosePrice)
}

// NewRsiWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage
func NewRsiWithSrcLen(sourceLength uint, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind, err := NewRsi(timePeriod, selectData)

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewDefaultRsiWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage with default parameters
func NewDefaultRsiWithSrcLen(sourceLength uint) (indicator *Rsi, err error) {
	ind, err := NewDefaultRsi()

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewRsiForStream creates a Relative Strength Indicator (Rsi) for online usage with a source data stream
func NewRsiForStream(priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind, err := NewRsi(timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultRsiForStream creates a Relative Strength Indicator (Rsi) for online usage with a source data stream
func NewDefaultRsiForStream(priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Rsi, err error) {
	ind, err := NewDefaultRsi()
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewRsiForStreamWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage with a source data stream
func NewRsiForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind, err := NewRsiWithSrcLen(sourceLength, timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultRsiForStreamWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage with a source data stream
func NewDefaultRsiForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Rsi, err error) {
	ind, err := NewDefaultRsiWithSrcLen(sourceLength)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// ReceiveDOHLCVTick consumes a source data DOHLCV price tick
func (ind *Rsi) ReceiveDOHLCVTick(tickData gotrade.DOHLCV, streamBarIndex int) {
	var selectedData = ind.selectData(tickData)
	ind.ReceiveTick(selectedData, streamBarIndex)
}

func (ind *RsiWithoutStorage) ReceiveTick(tickData float64, streamBarIndex int) {
	ind.periodCounter += 1

	if ind.periodCounter > ind.timePeriod*-1 {

		if ind.periodCounter <= 0 {

			if tickData > ind.previousClose {
				ind.previousGain += (tickData - ind.previousClose)
			} else {
				ind.previousLoss -= (tickData - ind.previousClose)
			}
		}

		if ind.periodCounter == 0 {
			ind.previousGain /= float64(ind.timePeriod)
			ind.previousLoss /= float64(ind.timePeriod)

			// increment the number of results this indicator can be expected to return
			ind.dataLength += 1
			if ind.validFromBar == -1 {
				// set the streamBarIndex from which this indicator returns valid results
				ind.validFromBar = streamBarIndex
			}

			var result float64
			//    Rsi = 100 * (prevGain/(prevGain+prevLoss))
			if ind.previousGain+ind.previousLoss == 0.0 {
				result = 0.0
			} else {
				result = 100.0 * (ind.previousGain / (ind.previousGain + ind.previousLoss))
			}

			// update the maximum result value
			if result > ind.maxValue {
				ind.maxValue = result
			}

			// update the minimum result value
			if result < ind.minValue {
				ind.minValue = result
			}

			// notify of a new result value though the value available action
			ind.valueAvailableAction(result, streamBarIndex)
		}

		if ind.periodCounter > 0 {
			ind.previousGain *= float64(ind.timePeriod - 1)
			ind.previousLoss *= float64(ind.timePeriod - 1)

			if tickData > ind.previousClose {
				ind.previousGain += (tickData - ind.previousClose)
			} else {
				ind.previousLoss -= (tickData - ind.previousClose)
			}

			ind.previousGain /= float64(ind.timePeriod)
			ind.previousLoss /= float64(ind.timePeriod)

			// increment the number of results this indicator can be expected to return
			ind.dataLength += 1

			var result float64
			//    Rsi = 100 * (prevGain/(prevGain+prevLoss))
			if ind.previousGain+ind.previousLoss == 0.0 {
				result = 0.0
			} else {
				result = 100.0 * (ind.previousGain / (ind.previousGain + ind.previousLoss))
			}

			// update the maximum result value
			if result > ind.maxValue {
				ind.maxValue = result
			}

			// update the minimum result value
			if result < ind.minValue {
				ind.minValue = result
			}

			// notify of a new result value though the value available action
			ind.valueAvailableAction(result, streamBarIndex)
		}
	}
	ind.previousClose = tickData
}
<commit_msg>#76 Remove duplication - rsi<commit_after>package indicators

import (
	"errors"
	"github.com/thetruetrade/gotrade"
)

// A Relative Strength Indicator (Rsi), no storage, for use in other indicators
type RsiWithoutStorage struct {
	*baseIndicatorWithFloatBounds

	// private variables
	periodCounter int
	previousClose float64
	previousGain  float64
	previousLoss  float64
	timePeriod    int
}

// NewRsiWithoutStorage creates a Relative Strength Indicator (Rsi) without storage
func NewRsiWithoutStorage(timePeriod int, valueAvailableAction ValueAvailableActionFloat) (indicator *RsiWithoutStorage, err error) {

	// an indicator without storage MUST have a value available action
	if valueAvailableAction == nil {
		return nil, ErrValueAvailableActionIsNil
	}

	// the minimum timeperiod for this indicator is 2
	if timePeriod < 2 {
		return nil, errors.New("timePeriod is less than the minimum (2)")
	}

	// check the maximum timeperiod
	if timePeriod > MaximumLookbackPeriod {
		return nil, errors.New("timePeriod is greater than the maximum (100000)")
	}

	lookback := timePeriod
	ind := RsiWithoutStorage{
		baseIndicatorWithFloatBounds: newBaseIndicatorWithFloatBounds(lookback, valueAvailableAction),
		periodCounter:                (timePeriod * -1) - 1,
		previousClose:                0.0,
		previousGain:                 0.0,
		previousLoss:                 0.0,
		timePeriod:                   timePeriod,
	}

	return &ind, err
}

// A Relative Strength Indicator (Rsi)
type Rsi struct {
	*RsiWithoutStorage
	selectData gotrade.DataSelectionFunc

	// public variables
	Data []float64
}

// NewRsi creates a Relative Strength Indicator (Rsi) for online usage
func NewRsi(timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind := Rsi{selectData: selectData}
	ind.RsiWithoutStorage, err = NewRsiWithoutStorage(timePeriod,
		func(dataItem float64, streamBarIndex int) {
			ind.Data = append(ind.Data, dataItem)
		})

	return &ind, err
}

// NewDefaultRsi creates a Relative Strength Indicator (Rsi) for online usage with default parameters
//	- timePeriod: 14
func NewDefaultRsi() (indicator *Rsi, err error) {
	timePeriod := 14
	return NewRsi(timePeriod, gotrade.UseClosePrice)
}

// NewRsiWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage
func NewRsiWithSrcLen(sourceLength uint, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind, err := NewRsi(timePeriod, selectData)

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewDefaultRsiWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage with default parameters
func NewDefaultRsiWithSrcLen(sourceLength uint) (indicator *Rsi, err error) {
	ind, err := NewDefaultRsi()

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewRsiForStream creates a Relative Strength Indicator (Rsi) for online usage with a source data stream
func NewRsiForStream(priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind, err := NewRsi(timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultRsiForStream creates a Relative Strength Indicator (Rsi) for online usage with a source data stream
func NewDefaultRsiForStream(priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Rsi, err error) {
	ind, err := NewDefaultRsi()
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewRsiForStreamWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage with a source data stream
func NewRsiForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Rsi, err error) {
	ind, err := NewRsiWithSrcLen(sourceLength, timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultRsiForStreamWithSrcLen creates a Relative Strength Indicator (Rsi) for offline usage with a source data stream
func NewDefaultRsiForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Rsi, err error) {
	ind, err := NewDefaultRsiWithSrcLen(sourceLength)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// ReceiveDOHLCVTick consumes a source data DOHLCV price tick
func (ind *Rsi) ReceiveDOHLCVTick(tickData gotrade.DOHLCV, streamBarIndex int) {
	var selectedData = ind.selectData(tickData)
	ind.ReceiveTick(selectedData, streamBarIndex)
}

func (ind *RsiWithoutStorage) ReceiveTick(tickData float64, streamBarIndex int) {
	ind.periodCounter += 1

	if ind.periodCounter > ind.timePeriod*-1 {

		if ind.periodCounter <= 0 {

			if tickData > ind.previousClose {
				ind.previousGain += (tickData - ind.previousClose)
			} else {
				ind.previousLoss -= (tickData - ind.previousClose)
			}
		}

		if ind.periodCounter == 0 {
			ind.previousGain /= float64(ind.timePeriod)
			ind.previousLoss /= float64(ind.timePeriod)

			var result float64
			//    Rsi = 100 * (prevGain/(prevGain+prevLoss))
			if ind.previousGain+ind.previousLoss == 0.0 {
				result = 0.0
			} else {
				result = 100.0 * (ind.previousGain / (ind.previousGain + ind.previousLoss))
			}

			ind.UpdateIndicatorWithNewValue(result, streamBarIndex)
		}

		if ind.periodCounter > 0 {
			ind.previousGain *= float64(ind.timePeriod - 1)
			ind.previousLoss *= float64(ind.timePeriod - 1)

			if tickData > ind.previousClose {
				ind.previousGain += (tickData - ind.previousClose)
			} else {
				ind.previousLoss -= (tickData - ind.previousClose)
			}

			ind.previousGain /= float64(ind.timePeriod)
			ind.previousLoss /= float64(ind.timePeriod)

			var result float64
			//    Rsi = 100 * (prevGain/(prevGain+prevLoss))
			if ind.previousGain+ind.previousLoss == 0.0 {
				result = 0.0
			} else {
				result = 100.0 * (ind.previousGain / (ind.previousGain + ind.previousLoss))
			}

			ind.UpdateIndicatorWithNewValue(result, streamBarIndex)
		}
	}
	ind.previousClose = tickData
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
)

// Message to send to other nodes
type BinaryTransportMessage struct {
	Version uint32                     // Message format (e.g. 1)
	Type    BinaryTransportMessageType // Sort message
	Len     uint32                     // Length of data
	Data    []byte                     // Actual data
}

// This version
const BINARY_TRANSPORT_MESSAGE_VERSION uint32 = 1

// Message type
type BinaryTransportMessageType uint32

// Message types
const (
	EmptyBinaryTransportMessageType       BinaryTransportMessageType = iota // 0 = not set
	ShardIdxBinaryTransportMessageType                                      // 1 = shard index
	FileBinaryTransportMessageType                                          // 2 = file
	CreateShardBinaryTransportMessageType                                   // 3 = create shard (replication)
)

// To bytes
func (this *BinaryTransportMessage) Bytes() []byte {
	if this.Type < 1 || this.Len < 1 {
		panic("Can not convert empty message to bytes")
	}
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, this.Version) // meta version
	binary.Write(buf, binary.BigEndian, this.Type)    // type
	binary.Write(buf, binary.BigEndian, this.Len)     // data length
	if this.Data != nil {
		buf.Write(this.Data)
	}
	return buf.Bytes()
}

// From bytes
func (this *BinaryTransportMessage) FromBytes(b []byte) {
	buf := bytes.NewReader(b)
	var err error
	err = binary.Read(buf, binary.BigEndian, &this.Version) // meta version
	panicErr(err)
	err = binary.Read(buf, binary.BigEndian, &this.Type) // type
	panicErr(err)
	err = binary.Read(buf, binary.BigEndian, &this.Len) // length
	panicErr(err)

	// Read data
	this.Data = make([]byte, this.Len)
	buf.Read(this.Data)
}

// New message
func newBinaryTransportMessage(t BinaryTransportMessageType, data []byte) *BinaryTransportMessage {
	var l uint32 = 0
	if data != nil {
		l = uint32(len(data))
	}
	return &BinaryTransportMessage{
		Version: BINARY_TRANSPORT_MESSAGE_VERSION,
		Type:    t,
		Len:     l,
		Data:    data,
	}
}
<commit_msg>Update comment<commit_after>package main

import (
	"bytes"
	"encoding/binary"
)

// Message to send to other nodes
type BinaryTransportMessage struct {
	Version uint32                     // Message format (e.g. 1)
	Type    BinaryTransportMessageType // Sort message
	Len     uint32                     // Length of data
	Data    []byte                     // Actual data
}

// This version
const BINARY_TRANSPORT_MESSAGE_VERSION uint32 = 1

// Message type
type BinaryTransportMessageType uint32

// Message types
const (
	EmptyBinaryTransportMessageType       BinaryTransportMessageType = iota // 0 = not set
	ShardIdxBinaryTransportMessageType                                      // 1 = shard index
	FileBinaryTransportMessageType                                          // 2 = file binary transport
	CreateShardBinaryTransportMessageType                                   // 3 = create shard (replication)
)

// To bytes
func (this *BinaryTransportMessage) Bytes() []byte {
	if this.Type < 1 || this.Len < 1 {
		panic("Can not convert empty message to bytes")
	}
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, this.Version) // meta version
	binary.Write(buf, binary.BigEndian, this.Type)    // type
	binary.Write(buf, binary.BigEndian, this.Len)     // data length
	if this.Data != nil {
		buf.Write(this.Data)
	}
	return buf.Bytes()
}

// From bytes
func (this *BinaryTransportMessage) FromBytes(b []byte) {
	buf := bytes.NewReader(b)
	var err error
	err = binary.Read(buf, binary.BigEndian, &this.Version) // meta version
	panicErr(err)
	err = binary.Read(buf, binary.BigEndian, &this.Type) // type
	panicErr(err)
	err = binary.Read(buf, binary.BigEndian, &this.Len) // length
	panicErr(err)

	// Read data
	this.Data = make([]byte, this.Len)
	buf.Read(this.Data)
}

// New message
func newBinaryTransportMessage(t BinaryTransportMessageType, data []byte) *BinaryTransportMessage {
	var l uint32 = 0
	if data != nil {
		l = uint32(len(data))
	}
	return &BinaryTransportMessage{
		Version: BINARY_TRANSPORT_MESSAGE_VERSION,
		Type:    t,
		Len:     l,
		Data:    data,
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsSpotInstanceRequest() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSpotInstanceRequestCreate,
		Read:   resourceAwsSpotInstanceRequestRead,
		Delete: resourceAwsSpotInstanceRequestDelete,
		Update: resourceAwsSpotInstanceRequestUpdate,

		Schema: func() map[string]*schema.Schema {
			// The Spot Instance Request Schema is based on the AWS Instance schema.
			s := resourceAwsInstance().Schema

			// Everything on a spot instance is ForceNew except tags
			for k, v := range s {
				if k == "tags" {
					continue
				}
				v.ForceNew = true
			}

			s["spot_price"] = &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			}
			s["spot_type"] = &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "persistent",
			}
			s["wait_for_fulfillment"] = &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			}
			s["spot_bid_status"] = &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			}
			s["spot_request_state"] = &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			}
			s["spot_instance_id"] = &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			}

			return s
		}(),
	}
}

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

	instanceOpts, err := buildAwsInstanceOpts(d, meta)
	if err != nil {
		return err
	}

	spotOpts := &ec2.RequestSpotInstancesInput{
		SpotPrice: aws.String(d.Get("spot_price").(string)),
		Type:      aws.String(d.Get("spot_type").(string)),

		// Though the AWS API supports creating spot instance requests for multiple
		// instances, for TF purposes we fix this to one instance per request.
		// Users can get equivalent behavior out of TF's "count" meta-parameter.
		InstanceCount: aws.Int64(1),

		LaunchSpecification: &ec2.RequestSpotLaunchSpecification{
			BlockDeviceMappings: instanceOpts.BlockDeviceMappings,
			EbsOptimized:        instanceOpts.EBSOptimized,
			Monitoring:          instanceOpts.Monitoring,
			IamInstanceProfile:  instanceOpts.IAMInstanceProfile,
			ImageId:             instanceOpts.ImageID,
			InstanceType:        instanceOpts.InstanceType,
			KeyName:             instanceOpts.KeyName,
			Placement:           instanceOpts.SpotPlacement,
			SecurityGroupIds:    instanceOpts.SecurityGroupIDs,
			SecurityGroups:      instanceOpts.SecurityGroups,
			SubnetId:            instanceOpts.SubnetID,
			UserData:            instanceOpts.UserData64,
		},
	}

	// If the instance is configured with a Network Interface (a subnet, has
	// public IP, etc), then the instanceOpts.SecurityGroupIds and SubnetId will
	// be nil
	if len(instanceOpts.NetworkInterfaces) > 0 {
		spotOpts.LaunchSpecification.SecurityGroupIds = instanceOpts.NetworkInterfaces[0].Groups
		spotOpts.LaunchSpecification.SubnetId = instanceOpts.NetworkInterfaces[0].SubnetId
	}

	// Make the spot instance request
	log.Printf("[DEBUG] Requesting spot bid opts: %s", spotOpts)
	resp, err := conn.RequestSpotInstances(spotOpts)
	if err != nil {
		return fmt.Errorf("Error requesting spot instances: %s", err)
	}
	if len(resp.SpotInstanceRequests) != 1 {
		return fmt.Errorf(
			"Expected response with length 1, got: %s", resp)
	}

	sir := *resp.SpotInstanceRequests[0]
	d.SetId(*sir.SpotInstanceRequestId)

	if d.Get("wait_for_fulfillment").(bool) {
		spotStateConf := &resource.StateChangeConf{
			// http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/spot-bid-status.html
			Pending:    []string{"start", "pending-evaluation", "pending-fulfillment"},
			Target:     "fulfilled",
			Refresh:    SpotInstanceStateRefreshFunc(conn, sir),
			Timeout:    10 * time.Minute,
			Delay:      10 * time.Second,
			MinTimeout: 3 * time.Second,
		}

		log.Printf("[DEBUG] waiting for spot bid to resolve... this may take several minutes.")
		_, err = spotStateConf.WaitForState()

		if err != nil {
			return fmt.Errorf("Error while waiting for spot request (%s) to resolve: %s", sir, err)
		}
	}

	return resourceAwsSpotInstanceRequestUpdate(d, meta)
}

// Update spot state, etc
func resourceAwsSpotInstanceRequestRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	req := &ec2.DescribeSpotInstanceRequestsInput{
		SpotInstanceRequestIds: []*string{aws.String(d.Id())},
	}
	resp, err := conn.DescribeSpotInstanceRequests(req)

	if err != nil {
		// If the spot request was not found, return nil so that we can show
		// that it is gone.
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSpotInstanceRequestID.NotFound" {
			d.SetId("")
			return nil
		}

		// Some other error, report it
		return err
	}

	// If nothing was found, then return no state
	if len(resp.SpotInstanceRequests) == 0 {
		d.SetId("")
		return nil
	}

	request := resp.SpotInstanceRequests[0]

	// if the request is cancelled, then it is gone
	if *request.State == "cancelled" {
		d.SetId("")
		return nil
	}

	d.Set("spot_bid_status", *request.Status.Code)
	// Instance ID is not set if the request is still pending
	if request.InstanceId != nil {
		d.Set("spot_instance_id", *request.InstanceId)
		if err := readInstance(d, meta); err != nil {
			return fmt.Errorf("[ERR] Error reading Spot Instance Data: %s", err)
		}
	}
	d.Set("spot_request_state", *request.State)
	d.Set("tags", tagsToMap(request.Tags))

	// return nil
	// let's read the instance data...
	return readInstance(d, meta)
}

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

	resp, err := conn.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(d.Get("spot_instance_id").(string))},
	})
	if err != nil {
		// If the instance was not found, return nil so that we can show
		// that the instance is gone.
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidInstanceID.NotFound" {
			return fmt.Errorf("no instance found")
		}

		// Some other error, report it
		return err
	}

	// If nothing was found, then return no state
	if len(resp.Reservations) == 0 {
		return fmt.Errorf("no instances found")
	}

	instance := resp.Reservations[0].Instances[0]

	// Set these fields for connection information
	if instance != nil {
		d.Set("public_dns", instance.PublicDnsName)
		d.Set("public_ip", instance.PublicIpAddress)
		d.Set("private_dns", instance.PrivateDnsName)
		d.Set("private_ip", instance.PrivateIpAddress)
	}

	// set connection information
	if instance.PublicIpAddress != nil {
		d.SetConnInfo(map[string]string{
			"type": "ssh",
			"host": *instance.PublicIpAddress,
		})
	} else if instance.PrivateIpAddress != nil {
		d.SetConnInfo(map[string]string{
			"type": "ssh",
			"host": *instance.PrivateIpAddress,
		})
	}

	return nil
}

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

	d.Partial(true)
	if err := setTags(conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	d.Partial(false)

	return resourceAwsSpotInstanceRequestRead(d, meta)
}

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

	log.Printf("[INFO] Cancelling spot request: %s", d.Id())
	_, err := conn.CancelSpotInstanceRequests(&ec2.CancelSpotInstanceRequestsInput{
		SpotInstanceRequestIds: []*string{aws.String(d.Id())},
	})

	if err != nil {
		return fmt.Errorf("Error cancelling spot request (%s): %s", d.Id(), err)
	}

	if instanceId := d.Get("spot_instance_id").(string); instanceId != "" {
		log.Printf("[INFO] Terminating instance: %s", instanceId)
		if err := awsTerminateInstance(conn, instanceId); err != nil {
			return fmt.Errorf("Error terminating spot instance: %s", err)
		}
	}

	return nil
}

// SpotInstanceStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// an EC2 spot instance request
func SpotInstanceStateRefreshFunc(
	conn *ec2.EC2, sir ec2.SpotInstanceRequest) resource.StateRefreshFunc {

	return func() (interface{}, string, error) {
		resp, err := conn.DescribeSpotInstanceRequests(&ec2.DescribeSpotInstanceRequestsInput{
			SpotInstanceRequestIds: []*string{sir.SpotInstanceRequestId},
		})

		if err != nil {
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSpotInstanceRequestID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else {
				log.Printf("Error on StateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.SpotInstanceRequests) == 0 {
			// Sometimes AWS just has consistency issues and doesn't see
			// our request yet. Return an empty state.
			return nil, "", nil
		}

		req := resp.SpotInstanceRequests[0]
		return req, *req.Status.Code, nil
	}
}
<commit_msg>minor tweaks to connection info setup<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsSpotInstanceRequest() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSpotInstanceRequestCreate,
		Read:   resourceAwsSpotInstanceRequestRead,
		Delete: resourceAwsSpotInstanceRequestDelete,
		Update: resourceAwsSpotInstanceRequestUpdate,

		Schema: func() map[string]*schema.Schema {
			// The Spot Instance Request Schema is based on the AWS Instance schema.
			s := resourceAwsInstance().Schema

			// Everything on a spot instance is ForceNew except tags
			for k, v := range s {
				if k == "tags" {
					continue
				}
				v.ForceNew = true
			}

			s["spot_price"] = &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			}
			s["spot_type"] = &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "persistent",
			}
			s["wait_for_fulfillment"] = &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			}
			s["spot_bid_status"] = &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			}
			s["spot_request_state"] = &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			}
			s["spot_instance_id"] = &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			}

			return s
		}(),
	}
}

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

	instanceOpts, err := buildAwsInstanceOpts(d, meta)
	if err != nil {
		return err
	}

	spotOpts := &ec2.RequestSpotInstancesInput{
		SpotPrice: aws.String(d.Get("spot_price").(string)),
		Type:      aws.String(d.Get("spot_type").(string)),

		// Though the AWS API supports creating spot instance requests for multiple
		// instances, for TF purposes we fix this to one instance per request.
		// Users can get equivalent behavior out of TF's "count" meta-parameter.
		InstanceCount: aws.Int64(1),

		LaunchSpecification: &ec2.RequestSpotLaunchSpecification{
			BlockDeviceMappings: instanceOpts.BlockDeviceMappings,
			EbsOptimized:        instanceOpts.EBSOptimized,
			Monitoring:          instanceOpts.Monitoring,
			IamInstanceProfile:  instanceOpts.IAMInstanceProfile,
			ImageId:             instanceOpts.ImageID,
			InstanceType:        instanceOpts.InstanceType,
			KeyName:             instanceOpts.KeyName,
			Placement:           instanceOpts.SpotPlacement,
			SecurityGroupIds:    instanceOpts.SecurityGroupIDs,
			SecurityGroups:      instanceOpts.SecurityGroups,
			SubnetId:            instanceOpts.SubnetID,
			UserData:            instanceOpts.UserData64,
		},
	}

	// If the instance is configured with a Network Interface (a subnet, has
	// public IP, etc), then the instanceOpts.SecurityGroupIds and SubnetId will
	// be nil
	if len(instanceOpts.NetworkInterfaces) > 0 {
		spotOpts.LaunchSpecification.SecurityGroupIds = instanceOpts.NetworkInterfaces[0].Groups
		spotOpts.LaunchSpecification.SubnetId = instanceOpts.NetworkInterfaces[0].SubnetId
	}

	// Make the spot instance request
	log.Printf("[DEBUG] Requesting spot bid opts: %s", spotOpts)
	resp, err := conn.RequestSpotInstances(spotOpts)
	if err != nil {
		return fmt.Errorf("Error requesting spot instances: %s", err)
	}
	if len(resp.SpotInstanceRequests) != 1 {
		return fmt.Errorf(
			"Expected response with length 1, got: %s", resp)
	}

	sir := *resp.SpotInstanceRequests[0]
	d.SetId(*sir.SpotInstanceRequestId)

	if d.Get("wait_for_fulfillment").(bool) {
		spotStateConf := &resource.StateChangeConf{
			// http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/spot-bid-status.html
			Pending:    []string{"start", "pending-evaluation", "pending-fulfillment"},
			Target:     "fulfilled",
			Refresh:    SpotInstanceStateRefreshFunc(conn, sir),
			Timeout:    10 * time.Minute,
			Delay:      10 * time.Second,
			MinTimeout: 3 * time.Second,
		}

		log.Printf("[DEBUG] waiting for spot bid to resolve... this may take several minutes.")
		_, err = spotStateConf.WaitForState()

		if err != nil {
			return fmt.Errorf("Error while waiting for spot request (%s) to resolve: %s", sir, err)
		}
	}

	return resourceAwsSpotInstanceRequestUpdate(d, meta)
}

// Update spot state, etc
func resourceAwsSpotInstanceRequestRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	req := &ec2.DescribeSpotInstanceRequestsInput{
		SpotInstanceRequestIds: []*string{aws.String(d.Id())},
	}
	resp, err := conn.DescribeSpotInstanceRequests(req)

	if err != nil {
		// If the spot request was not found, return nil so that we can show
		// that it is gone.
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSpotInstanceRequestID.NotFound" {
			d.SetId("")
			return nil
		}

		// Some other error, report it
		return err
	}

	// If nothing was found, then return no state
	if len(resp.SpotInstanceRequests) == 0 {
		d.SetId("")
		return nil
	}

	request := resp.SpotInstanceRequests[0]

	// if the request is cancelled, then it is gone
	if *request.State == "cancelled" {
		d.SetId("")
		return nil
	}

	d.Set("spot_bid_status", *request.Status.Code)
	// Instance ID is not set if the request is still pending
	if request.InstanceId != nil {
		d.Set("spot_instance_id", *request.InstanceId)
		if err := readInstance(d, meta); err != nil {
			return fmt.Errorf("[ERR] Error reading Spot Instance Data: %s", err)
		}
	}
	d.Set("spot_request_state", *request.State)
	d.Set("tags", tagsToMap(request.Tags))

	// Read the instance data, setting up connection information
	return readInstance(d, meta)
}

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

	resp, err := conn.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(d.Get("spot_instance_id").(string))},
	})
	if err != nil {
		// If the instance was not found, return nil so that we can show
		// that the instance is gone.
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidInstanceID.NotFound" {
			return fmt.Errorf("no instance found")
		}

		// Some other error, report it
		return err
	}

	// If nothing was found, then return no state
	if len(resp.Reservations) == 0 {
		return fmt.Errorf("no instances found")
	}

	instance := resp.Reservations[0].Instances[0]

	// Set these fields for connection information
	if instance != nil {
		d.Set("public_dns", instance.PublicDnsName)
		d.Set("public_ip", instance.PublicIpAddress)
		d.Set("private_dns", instance.PrivateDnsName)
		d.Set("private_ip", instance.PrivateIpAddress)

		// set connection information
		if instance.PublicIpAddress != nil {
			d.SetConnInfo(map[string]string{
				"type": "ssh",
				"host": *instance.PublicIpAddress,
			})
		} else if instance.PrivateIpAddress != nil {
			d.SetConnInfo(map[string]string{
				"type": "ssh",
				"host": *instance.PrivateIpAddress,
			})
		}
	}

	return nil
}

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

	d.Partial(true)
	if err := setTags(conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	d.Partial(false)

	return resourceAwsSpotInstanceRequestRead(d, meta)
}

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

	log.Printf("[INFO] Cancelling spot request: %s", d.Id())
	_, err := conn.CancelSpotInstanceRequests(&ec2.CancelSpotInstanceRequestsInput{
		SpotInstanceRequestIds: []*string{aws.String(d.Id())},
	})

	if err != nil {
		return fmt.Errorf("Error cancelling spot request (%s): %s", d.Id(), err)
	}

	if instanceId := d.Get("spot_instance_id").(string); instanceId != "" {
		log.Printf("[INFO] Terminating instance: %s", instanceId)
		if err := awsTerminateInstance(conn, instanceId); err != nil {
			return fmt.Errorf("Error terminating spot instance: %s", err)
		}
	}

	return nil
}

// SpotInstanceStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// an EC2 spot instance request
func SpotInstanceStateRefreshFunc(
	conn *ec2.EC2, sir ec2.SpotInstanceRequest) resource.StateRefreshFunc {

	return func() (interface{}, string, error) {
		resp, err := conn.DescribeSpotInstanceRequests(&ec2.DescribeSpotInstanceRequestsInput{
			SpotInstanceRequestIds: []*string{sir.SpotInstanceRequestId},
		})

		if err != nil {
			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidSpotInstanceRequestID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else {
				log.Printf("Error on StateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.SpotInstanceRequests) == 0 {
			// Sometimes AWS just has consistency issues and doesn't see
			// our request yet. Return an empty state.
			return nil, "", nil
		}

		req := resp.SpotInstanceRequests[0]
		return req, *req.Status.Code, nil
	}
}
<|endoftext|>


<commit_before>package repair

import (
	"fmt"
	"io"
	"koding/klient/remote/req"
	"time"

	"github.com/koding/logging"
)

type KontrolRepair struct {
	Log logging.Logger

	// The klient we will be communicating with.
	Klient interface {
		RemoteStatus(req.Status) error
	}

	// The options that this repairer will use.
	RetryOptions RetryOptions

	// The way
	Stdout io.Writer

	// A struct that runs the given command.
	Exec interface {
		Run(string, ...string) error
	}
}

func (r *KontrolRepair) String() string {
	return "kontrolrepair"
}

func (r *KontrolRepair) Status() error {
	var (
		needNewline bool
		err         error
	)

	for i := uint(0); i <= r.RetryOptions.StatusRetries; i++ {
		err = r.status()
		if err == nil {
			break
		}

		switch i {
		case 0:
			needNewline = true
			fmt.Fprint(r.Stdout, "Kontrol not connected. Waiting for reconnect.")
		default:
			fmt.Fprint(r.Stdout, " .")
		}

		time.Sleep(r.RetryOptions.StatusDelay)
	}

	if needNewline {
		fmt.Fprint(r.Stdout, "\n")
	}

	return err
}

func (r *KontrolRepair) status() error {
	return r.Klient.RemoteStatus(req.Status{
		Item: req.KontrolStatus,
	})
}

func (r *KontrolRepair) Repair() error {
	fmt.Fprintln(r.Stdout, "Kontrol has not reconnected in expected time. Restarting.")

	if err := r.Exec.Run("sudo", "kd", "restart"); err != nil {
		return fmt.Errorf(
			"Subprocess kd failed to start. err:%s", err,
		)
	}

	// Run status again, to confirm it's running as best we can. If not, we've
	// tried and failed.
	if err := r.status(); err != nil {
		fmt.Fprintln(r.Stdout, "Unable to reconnect to kontrol.")
	}

	return nil
}
<commit_msg>klientctl: Added a retry on the final kontrol status check<commit_after>package repair

import (
	"fmt"
	"io"
	"koding/klient/remote/req"
	"time"

	"github.com/koding/logging"
)

type KontrolRepair struct {
	Log logging.Logger

	// The klient we will be communicating with.
	Klient interface {
		RemoteStatus(req.Status) error
	}

	// The options that this repairer will use.
	RetryOptions RetryOptions

	// The way
	Stdout io.Writer

	// A struct that runs the given command.
	Exec interface {
		Run(string, ...string) error
	}
}

func (r *KontrolRepair) String() string {
	return "kontrolrepair"
}

func (r *KontrolRepair) Status() error {
	if err := r.remoteStatus(); err == nil {
		return nil
	}

	fmt.Fprint(r.Stdout, "Kontrol not connected. Waiting for reconnect.")

	return r.statusLoop()
}

// statusLoop runs the status loop, optionally printing a dot to indicate progress.
func (r *KontrolRepair) statusLoop() error {
	var err error

	for i := uint(0); i <= r.RetryOptions.StatusRetries; i++ {
		if err = r.remoteStatus(); err == nil {
			break
		}

		fmt.Fprint(r.Stdout, ".")

		time.Sleep(r.RetryOptions.StatusDelay)
	}

	return err
}

func (r *KontrolRepair) remoteStatus() error {
	return r.Klient.RemoteStatus(req.Status{
		Item: req.KontrolStatus,
	})
}

func (r *KontrolRepair) Repair() error {
	fmt.Fprintln(r.Stdout, "Kontrol has not reconnected in expected time. Restarting.")

	if err := r.Exec.Run("sudo", "kd", "restart"); err != nil {
		return fmt.Errorf(
			"Subprocess kd failed to start. err:%s", err,
		)
	}

	// Run status again, to confirm it's running as best we can. If not, we've
	// tried and failed.
	fmt.Fprint(r.Stdout, "Waiting for reconnect.")

	if err := r.statusLoop(); err != nil {
		fmt.Fprintln(r.Stdout, "Unable to reconnect to kontrol.")
		return err
	}

	fmt.Fprint(r.Stdout, "\nReconnected to kontrol.")
	return nil
}
<|endoftext|>


<commit_before>package checks

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/armon/circbuf"
	"github.com/docker/go-connections/sockets"
)

// DockerClient is a simplified client for the Docker Engine API
// to execute the health checks and avoid significant dependencies.
// It also consumes all data returned from the Docker API through
// a ring buffer with a fixed limit to avoid excessive resource
// consumption.
type DockerClient struct {
	host     string
	scheme   string
	proto    string
	addr     string
	basepath string
	maxbuf   int64
	client   *http.Client
}

func NewDockerClient(host string, maxbuf int64) (*DockerClient, error) {
	if host == "" {
		host = DefaultDockerHost
	}

	proto, addr, basepath, err := ParseHost(host)
	if err != nil {
		return nil, err
	}

	transport := new(http.Transport)
	sockets.ConfigureTransport(transport, proto, addr)
	client := &http.Client{Transport: transport}

	return &DockerClient{
		host:     host,
		scheme:   "http",
		proto:    proto,
		addr:     addr,
		basepath: basepath,
		maxbuf:   maxbuf,
		client:   client,
	}, nil
}

func (c *DockerClient) Host() string {
	return c.host
}

// ParseHost verifies that the given host strings is valid.
// copied from github.com/docker/docker/client.go
func ParseHost(host string) (string, string, string, error) {
	protoAddrParts := strings.SplitN(host, "://", 2)
	if len(protoAddrParts) == 1 {
		return "", "", "", fmt.Errorf("unable to parse docker host `%s`", host)
	}

	var basePath string
	proto, addr := protoAddrParts[0], protoAddrParts[1]
	if proto == "tcp" {
		parsed, err := url.Parse("tcp://" + addr)
		if err != nil {
			return "", "", "", err
		}
		addr = parsed.Host
		basePath = parsed.Path
	}
	return proto, addr, basePath, nil
}

func (c *DockerClient) call(method, uri string, v interface{}) (*circbuf.Buffer, int, error) {
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, 0, err
	}

	if c.proto == "unix" || c.proto == "npipe" {
		// For local communications, it doesn't matter what the host is. We just
		// need a valid and meaningful host name. (See #189)
		req.Host = "docker"
	}

	req.URL.Host = c.addr
	req.URL.Scheme = c.scheme

	if v != nil {
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(v); err != nil {
			return nil, 0, err
		}
		req.Body = ioutil.NopCloser(&b)
		req.Header.Set("Content-Type", "application/json")
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()

	b, err := circbuf.NewBuffer(c.maxbuf)
	if err != nil {
		return nil, 0, err
	}
	_, err = io.Copy(b, resp.Body)
	return b, resp.StatusCode, err
}

func (c *DockerClient) CreateExec(containerID string, cmd []string) (string, error) {
	data := struct {
		AttachStdin  bool
		AttachStdout bool
		AttachStderr bool
		Tty          bool
		Cmd          []string
	}{
		AttachStderr: true,
		AttachStdout: true,
		Cmd:          cmd,
	}

	uri := fmt.Sprintf("/containers/%s/exec", url.QueryEscape(containerID))
	b, code, err := c.call("POST", uri, data)
	switch {
	case err != nil:
		return "", fmt.Errorf("create exec failed for container %s: %s", containerID, err)
	case code == 201:
		var resp struct{ Id string }
		if err = json.NewDecoder(bytes.NewReader(b.Bytes())).Decode(&resp); err != nil {
			return "", fmt.Errorf("create exec response for container %s cannot be parsed: %s", containerID, err)
		}
		return resp.Id, nil
	case code == 404:
		return "", fmt.Errorf("create exec failed for unknown container %s", containerID)
	case code == 409:
		return "", fmt.Errorf("create exec failed since container %s is paused or stopped", containerID)
	default:
		return "", fmt.Errorf("create exec failed for container %s with status %d: %s", containerID, code, b)
	}
}

func (c *DockerClient) StartExec(containerID, execID string) (*circbuf.Buffer, error) {
	data := struct{ Detach, Tty bool }{Detach: false, Tty: true}
	uri := fmt.Sprintf("/exec/%s/start", execID)
	b, code, err := c.call("POST", uri, data)
	switch {
	case err != nil && !strings.Contains(err.Error(), "connection reset by peer"):
		return nil, fmt.Errorf("start exec failed for container %s: %s", containerID, err)
	case code == 200:
		return b, nil
	case code == 404:
		return nil, fmt.Errorf("start exec failed for container %s: invalid exec id %s", containerID, execID)
	case code == 409:
		return nil, fmt.Errorf("start exec failed since container %s is paused or stopped", containerID)
	default:
		return nil, fmt.Errorf("start exec failed for container %s with status %d: body: %s err: %s", containerID, code, b, err)
	}
}

func (c *DockerClient) InspectExec(containerID, execID string) (int, error) {
	uri := fmt.Sprintf("/exec/%s/json", execID)
	b, code, err := c.call("GET", uri, nil)
	switch {
	case err != nil:
		return 0, fmt.Errorf("inspect exec failed for container %s: %s", containerID, err)
	case code == 200:
		var resp struct{ ExitCode int }
		if err := json.NewDecoder(bytes.NewReader(b.Bytes())).Decode(&resp); err != nil {
			return 0, fmt.Errorf("inspect exec response for container %s cannot be parsed: %s", containerID, err)
		}
		return resp.ExitCode, nil
	case code == 404:
		return 0, fmt.Errorf("inspect exec failed for container %s: invalid exec id %s", containerID, execID)
	default:
		return 0, fmt.Errorf("inspect exec failed for container %s with status %d: %s", containerID, code, b)
	}
}
<commit_msg>docker: do not alloc a tty since this is not interactive<commit_after>package checks

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/armon/circbuf"
	"github.com/docker/go-connections/sockets"
)

// DockerClient is a simplified client for the Docker Engine API
// to execute the health checks and avoid significant dependencies.
// It also consumes all data returned from the Docker API through
// a ring buffer with a fixed limit to avoid excessive resource
// consumption.
type DockerClient struct {
	host     string
	scheme   string
	proto    string
	addr     string
	basepath string
	maxbuf   int64
	client   *http.Client
}

func NewDockerClient(host string, maxbuf int64) (*DockerClient, error) {
	if host == "" {
		host = DefaultDockerHost
	}

	proto, addr, basepath, err := ParseHost(host)
	if err != nil {
		return nil, err
	}

	transport := new(http.Transport)
	sockets.ConfigureTransport(transport, proto, addr)
	client := &http.Client{Transport: transport}

	return &DockerClient{
		host:     host,
		scheme:   "http",
		proto:    proto,
		addr:     addr,
		basepath: basepath,
		maxbuf:   maxbuf,
		client:   client,
	}, nil
}

func (c *DockerClient) Host() string {
	return c.host
}

// ParseHost verifies that the given host strings is valid.
// copied from github.com/docker/docker/client.go
func ParseHost(host string) (string, string, string, error) {
	protoAddrParts := strings.SplitN(host, "://", 2)
	if len(protoAddrParts) == 1 {
		return "", "", "", fmt.Errorf("unable to parse docker host `%s`", host)
	}

	var basePath string
	proto, addr := protoAddrParts[0], protoAddrParts[1]
	if proto == "tcp" {
		parsed, err := url.Parse("tcp://" + addr)
		if err != nil {
			return "", "", "", err
		}
		addr = parsed.Host
		basePath = parsed.Path
	}
	return proto, addr, basePath, nil
}

func (c *DockerClient) call(method, uri string, v interface{}) (*circbuf.Buffer, int, error) {
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, 0, err
	}

	if c.proto == "unix" || c.proto == "npipe" {
		// For local communications, it doesn't matter what the host is. We just
		// need a valid and meaningful host name. (See #189)
		req.Host = "docker"
	}

	req.URL.Host = c.addr
	req.URL.Scheme = c.scheme

	if v != nil {
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(v); err != nil {
			return nil, 0, err
		}
		req.Body = ioutil.NopCloser(&b)
		req.Header.Set("Content-Type", "application/json")
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()

	b, err := circbuf.NewBuffer(c.maxbuf)
	if err != nil {
		return nil, 0, err
	}
	_, err = io.Copy(b, resp.Body)
	return b, resp.StatusCode, err
}

func (c *DockerClient) CreateExec(containerID string, cmd []string) (string, error) {
	data := struct {
		AttachStdin  bool
		AttachStdout bool
		AttachStderr bool
		Tty          bool
		Cmd          []string
	}{
		AttachStderr: true,
		AttachStdout: true,
		Cmd:          cmd,
	}

	uri := fmt.Sprintf("/containers/%s/exec", url.QueryEscape(containerID))
	b, code, err := c.call("POST", uri, data)
	switch {
	case err != nil:
		return "", fmt.Errorf("create exec failed for container %s: %s", containerID, err)
	case code == 201:
		var resp struct{ Id string }
		if err = json.NewDecoder(bytes.NewReader(b.Bytes())).Decode(&resp); err != nil {
			return "", fmt.Errorf("create exec response for container %s cannot be parsed: %s", containerID, err)
		}
		return resp.Id, nil
	case code == 404:
		return "", fmt.Errorf("create exec failed for unknown container %s", containerID)
	case code == 409:
		return "", fmt.Errorf("create exec failed since container %s is paused or stopped", containerID)
	default:
		return "", fmt.Errorf("create exec failed for container %s with status %d: %s", containerID, code, b)
	}
}

func (c *DockerClient) StartExec(containerID, execID string) (*circbuf.Buffer, error) {
	data := struct{ Detach, Tty bool }{Detach: false, Tty: false}
	uri := fmt.Sprintf("/exec/%s/start", execID)
	b, code, err := c.call("POST", uri, data)
	switch {
	case err != nil && !strings.Contains(err.Error(), "connection reset by peer"):
		return nil, fmt.Errorf("start exec failed for container %s: %s", containerID, err)
	case code == 200:
		return b, nil
	case code == 404:
		return nil, fmt.Errorf("start exec failed for container %s: invalid exec id %s", containerID, execID)
	case code == 409:
		return nil, fmt.Errorf("start exec failed since container %s is paused or stopped", containerID)
	default:
		return nil, fmt.Errorf("start exec failed for container %s with status %d: body: %s err: %s", containerID, code, b, err)
	}
}

func (c *DockerClient) InspectExec(containerID, execID string) (int, error) {
	uri := fmt.Sprintf("/exec/%s/json", execID)
	b, code, err := c.call("GET", uri, nil)
	switch {
	case err != nil:
		return 0, fmt.Errorf("inspect exec failed for container %s: %s", containerID, err)
	case code == 200:
		var resp struct{ ExitCode int }
		if err := json.NewDecoder(bytes.NewReader(b.Bytes())).Decode(&resp); err != nil {
			return 0, fmt.Errorf("inspect exec response for container %s cannot be parsed: %s", containerID, err)
		}
		return resp.ExitCode, nil
	case code == 404:
		return 0, fmt.Errorf("inspect exec failed for container %s: invalid exec id %s", containerID, execID)
	default:
		return 0, fmt.Errorf("inspect exec failed for container %s with status %d: %s", containerID, code, b)
	}
}
<|endoftext|>


<commit_before>package hwaflib

func (ctx *Context) Version() string {
	version := "20131204"
	return version
}

func (ctx *Context) Revision() string {
	revision := "6dc9a4b"
	return revision
}

// EOF

<commit_msg>version: 20131206<commit_after>package hwaflib

func (ctx *Context) Version() string {
	version := "20131206"
	return version
}

func (ctx *Context) Revision() string {
	revision := "76b771a"
	return revision
}

// EOF

<|endoftext|>


<commit_before>// Tinkering with go, zmq, and Jupyter
package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"

	zmq "github.com/pebbe/zmq4"
)

// MessageHeader is a Jupyter message header
// See: http://jupyter-client.readthedocs.org/en/latest/messaging.html
type MessageHeader struct {
	MessageID   string `json:"msg_id"`
	Username    string `json:"username"`
	Session     string `json:"session"`
	MessageType string `json:"msg_type"`
	Version     string `json:"version"`
}

// Message is a generic Jupyter message (not a wire message)
type Message struct {
	Header       MessageHeader          `json:"header"`
	ParentHeader MessageHeader          `json:"parent_header"`
	Metadata     map[string]interface{} `json:"metadata"`
	Content      map[string]interface{} `json:"content"`
}

// ConnectionInfo represents the runtime connection data used by Jupyter kernels
type ConnectionInfo struct {
	IOPubPort       int    `json:"iopub_port"`
	StdinPort       int    `json:"stdin_port"`
	IP              string `json:"ip"`
	Transport       string `json:"transport"`
	HBPort          int    `json:"hb_port"`
	SignatureScheme string `json:"signature_scheme"`
	ShellPort       int    `json:"shell_port"`
	ControlPort     int    `json:"control_port"`
	Key             string `json:"key"`
}

// DELIMITER denotes the Jupyter multipart message
var DELIMITER = "<IDS|MSG>"

// ParseWireProtocol fills a Message with all the juicy Jupyter bits
func (m *Message) ParseWireProtocol(wireMessage [][]byte, key []byte) (err error) {
	var i int
	var el []byte

	// Wire protocol
	/**
		[
	  		b'u-u-i-d',         # zmq identity(ies)
	  		b'<IDS|MSG>',       # delimiter
	  		b'baddad42',        # HMAC signature
			b'{header}',        # serialized header dict
			b'{parent_header}', # serialized parent header dict
			b'{metadata}',      # serialized metadata dict
			b'{content}',       # serialized content dict
			b'blob',            # extra raw data buffer(s)
	  		...
		]
	*/
	// Determine where the delimiter is
	for i, el = range wireMessage {
		if string(el) == DELIMITER {
			break // Found our delimiter
		}
	}

	if string(wireMessage[i]) != DELIMITER {
		return errors.New("Couldn't find delimeter")
	}

	// Extract the zmq identiti(es)
	//identities := wireMessage[:delimiterLocation]

	// If the message was signed
	if len(key) != 0 {
		mac := hmac.New(sha256.New, key)
		for _, msgpart := range wireMessage[i+2 : i+6] {
			mac.Write(msgpart)
		}
		signature := make([]byte, hex.DecodedLen(len(wireMessage[i+1])))
		hex.Decode(signature, wireMessage[i+1])
		if !hmac.Equal(mac.Sum(nil), signature) {
			return errors.New("Invalid signature")
		}
	}

	json.Unmarshal(wireMessage[i+2], &m.Header)
	json.Unmarshal(wireMessage[i+3], &m.ParentHeader)
	json.Unmarshal(wireMessage[i+4], &m.Metadata)
	json.Unmarshal(wireMessage[i+5], &m.Content)

	return nil
}

func main() {
	flag.Parse()
	if flag.NArg() < 1 {
		log.Fatalln("Need a command line argument for the connection file.")
	}

	connFile, err := os.Open(flag.Arg(0))
	if err != nil {
		fmt.Errorf("Couldn't open connection file: %v", err)
		os.Exit(1)
	}

	jsonParser := json.NewDecoder(connFile)

	var connInfo ConnectionInfo

	if err = jsonParser.Decode(&connInfo); err != nil {
		fmt.Errorf("Couldn't parse connection file: %v", err)
		os.Exit(2)
	}

	iopubSocket, err := zmq.NewSocket(zmq.SUB)
	if err != nil {
		fmt.Errorf("Couldn't start the iopub socket: %v", err)
	}

	defer iopubSocket.Close()

	connectionString := fmt.Sprintf("%s://%s:%d",
		connInfo.Transport,
		connInfo.IP,
		connInfo.IOPubPort)

	iopubSocket.Connect(connectionString)
	iopubSocket.SetSubscribe("")

	fmt.Println("Connected to")
	fmt.Println(connInfo)

	for {
		wireMessage, err := iopubSocket.RecvMessageBytes(0)
		if err != nil {
			fmt.Printf("Error on receive: %v\n", err)
			continue
		}

		var message Message
		message.ParseWireProtocol(wireMessage, []byte(connInfo.Key))

		fmt.Println(message.Content)

	}

}
<commit_msg>Flesh out the mimebundle<commit_after>// Tinkering with go, zmq, and Jupyter
package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"

	zmq "github.com/pebbe/zmq4"
)

// MessageHeader is a Jupyter message header
// See: http://jupyter-client.readthedocs.org/en/latest/messaging.html
type MessageHeader struct {
	MessageID   string `json:"msg_id"`
	Username    string `json:"username"`
	Session     string `json:"session"`
	MessageType string `json:"msg_type"`
	Version     string `json:"version"`
}

// Message is a generic Jupyter message (not a wire message)
type Message struct {
	Header       MessageHeader          `json:"header"`
	ParentHeader MessageHeader          `json:"parent_header"`
	Metadata     map[string]interface{} `json:"metadata"`
	Content      map[string]interface{} `json:"content"`
}

// MimeBundle is a collection of mimetypes -> data
// Example:
//     'text/html' -> '
Hey!
'
//     'image/png' -> 'R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7'
type MimeBundle map[string]string

// ConnectionInfo represents the runtime connection data used by Jupyter kernels
type ConnectionInfo struct {
	IOPubPort       int    `json:"iopub_port"`
	StdinPort       int    `json:"stdin_port"`
	IP              string `json:"ip"`
	Transport       string `json:"transport"`
	HBPort          int    `json:"hb_port"`
	SignatureScheme string `json:"signature_scheme"`
	ShellPort       int    `json:"shell_port"`
	ControlPort     int    `json:"control_port"`
	Key             string `json:"key"`
}

// DELIMITER denotes the Jupyter multipart message
var DELIMITER = "<IDS|MSG>"

// ParseWireProtocol fills a Message with all the juicy Jupyter bits
func (m *Message) ParseWireProtocol(wireMessage [][]byte, key []byte) (err error) {
	var i int
	var el []byte

	// Wire protocol
	/**
		[
	  		b'u-u-i-d',         # zmq identity(ies)
	  		b'<IDS|MSG>',       # delimiter
	  		b'baddad42',        # HMAC signature
			b'{header}',        # serialized header dict
			b'{parent_header}', # serialized parent header dict
			b'{metadata}',      # serialized metadata dict
			b'{content}',       # serialized content dict
			b'blob',            # extra raw data buffer(s)
	  		...
		]
	*/
	// Determine where the delimiter is
	for i, el = range wireMessage {
		if string(el) == DELIMITER {
			break // Found our delimiter
		}
	}

	if string(wireMessage[i]) != DELIMITER {
		return errors.New("Couldn't find delimeter")
	}

	// Extract the zmq identiti(es)
	//identities := wireMessage[:delimiterLocation]

	// If the message was signed
	if len(key) != 0 {
		// TODO: Programmatic selection of scheme
		mac := hmac.New(sha256.New, key)
		for _, msgpart := range wireMessage[i+2 : i+6] {
			mac.Write(msgpart)
		}
		signature := make([]byte, hex.DecodedLen(len(wireMessage[i+1])))
		hex.Decode(signature, wireMessage[i+1])
		if !hmac.Equal(mac.Sum(nil), signature) {
			return errors.New("Invalid signature")
		}
	}

	json.Unmarshal(wireMessage[i+2], &m.Header)
	json.Unmarshal(wireMessage[i+3], &m.ParentHeader)
	json.Unmarshal(wireMessage[i+4], &m.Metadata)
	json.Unmarshal(wireMessage[i+5], &m.Content)

	return nil
}

func main() {
	flag.Parse()
	if flag.NArg() < 1 {
		log.Fatalln("Need a command line argument for the connection file.")
	}

	connFile, err := os.Open(flag.Arg(0))
	if err != nil {
		fmt.Errorf("Couldn't open connection file: %v", err)
		os.Exit(1)
	}

	jsonParser := json.NewDecoder(connFile)

	var connInfo ConnectionInfo

	if err = jsonParser.Decode(&connInfo); err != nil {
		fmt.Errorf("Couldn't parse connection file: %v", err)
		os.Exit(2)
	}

	iopubSocket, err := zmq.NewSocket(zmq.SUB)
	if err != nil {
		fmt.Errorf("Couldn't start the iopub socket: %v", err)
	}

	defer iopubSocket.Close()

	connectionString := fmt.Sprintf("%s://%s:%d",
		connInfo.Transport,
		connInfo.IP,
		connInfo.IOPubPort)

	iopubSocket.Connect(connectionString)
	iopubSocket.SetSubscribe("")

	fmt.Printf("Connected to %v\n", connectionString)

	for {
		wireMessage, err := iopubSocket.RecvMessageBytes(0)
		if err != nil {
			fmt.Printf("Error on receive: %v\n", err)
			continue
		}

		var message Message
		message.ParseWireProtocol(wireMessage, []byte(connInfo.Key))

		_, ok := message.Content["data"]

		if !ok {
			continue
		}

		mimeBundle := (message.Content["data"]).(map[string]interface{})

		fmt.Println(mimeBundle)

	}

}
<|endoftext|>