<commit_before>package agent

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	mqtt "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/juju/loggo"
)

var AlreadyConfigured = errors.New("Already configured")
var AlreadyUnConfigured = errors.New("Already unconfigured")

//
// Acts as a bridge between local and cloud brokers, this includes reconnecting
// and emitting status changes.
//
// Once configured and started this will attempt to connect to connect to local
// and cloud brokers, if something dies it will reconnect based on the configured
// reconnect backoff.
//
type Bridge struct {
	conf   *Config
	local  *mqtt.MqttClient
	remote *mqtt.MqttClient
	log    loggo.Logger

	localTopics []replaceTopic
	cloudTopics []replaceTopic

	cloudUrl *url.URL
	token    string

	timer       *time.Timer
	reconnectCh chan bool
	shutdownCh  chan bool

	Configured bool
	Connected  bool
	Counter    int64

	LastError error

	bridgeLock sync.Mutex
}

type replaceTopic struct {
	on      string
	replace string
	with    string
}

func (r *replaceTopic) updated(originalTopic string) string {
	return strings.Replace(originalTopic, r.replace, r.with, 1)
}

var localTopics = []replaceTopic{
	// location related topics (TODO: move to cloud userspace RPC)
	{on: "$location/calibration", replace: "$location", with: "$cloud/location"},
	{on: "$location/delete", replace: "$location", with: "$cloud/location"},
	{on: "$device/+/+/rssi", replace: "$device", with: "$cloud/device"},

	// module health statistics
	{on: "$node/+/module/status", replace: "$node", with: "$cloud/node"},

	// cloud userspace RPC requests
	{on: "$ninja/services/rpc/+/+", replace: "$ninja", with: "$cloud/ninja"},
	{on: "$ninja/services/+", replace: "$ninja", with: "$cloud/ninja"},

	// temporary alternate topic to distinguish remote device replies from local-destined ones
	// used by the phone app for remote actuations
	// the alternate remote_ topic is to prevent a loopback with the below rule in the other direction
	// TODO: use a tag like $mesh-source to prevent loops (never re-proxy msgs with your source)
	{on: "$device/+/channel/+/reply", replace: "$device", with: "$cloud/remote_device"},

	// push up all local RPC methods in case the cloud is responding,
	// this is currently used for the push notification channel
	{on: "$device/+/channel/+", replace: "$device", with: "$cloud/device"},

	// push up state changes to the cloud
	{on: "$device/+/channel/+/event/state", replace: "$device", with: "$cloud/device"},
}

var cloudTopics = []replaceTopic{
	// location related topics
	{on: "$cloud/location/calibration/progress", replace: "$cloud/location", with: "$location"},
	{on: "$cloud/device/+/+/location", replace: "$cloud/device", with: "$device"},

	// cloud userspace RPC replies
	{on: "$cloud/ninja/services/rpc/+/+/reply", replace: "$cloud/ninja", with: "$ninja"},

	// see comment for $device/+/channel/+/reply above
	{on: "$cloud/remote_device/+/channel/+", replace: "$cloud/remote_device", with: "$device"},

	// allow cloud to announce devices and channels (used for phone on 3G and notification subscription channel)
	{on: "$cloud/device/+/event/announce", replace: "$cloud/device", with: "$device"},
	{on: "$cloud/device/+/channel/+/event/announce", replace: "$cloud/device", with: "$device"},

	// retrieve RPC responses from the cloud,
	// this is currently used for the push notification channel
	{on: "$cloud/device/+/channel/+/reply", replace: "$cloud/device", with: "$device"},
}

func createBridge(conf *Config) *Bridge {
	return &Bridge{conf: conf, localTopics: localTopics, cloudTopics: cloudTopics, log: loggo.GetLogger("bridge")}
}

func (b *Bridge) start(cloudUrl string, token string) (err error) {

	if b.Configured {
		b.log.Warningf("Already configured.")
		return AlreadyConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Connecting the bridge")

	b.Configured = true

	url, err := url.Parse(cloudUrl)

	if err != nil {
		return err
	}

	b.cloudUrl = url
	b.token = token

	b.reconnectCh = make(chan bool, 1)
	b.shutdownCh = make(chan bool, 1)

	if err = b.connect(); err != nil {
		b.log.Errorf("Connect failed %s", err)
		b.scheduleReconnect(err)
	}

	go b.mainBridgeLoop()

	return err
}

func (b *Bridge) stop() error {

	if !b.Configured {
		b.log.Warningf("Already unconfigured.")
		return AlreadyUnConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Disconnecting bridge")

	if b.Configured {
		// tell the worker to shutdown
		b.shutdownCh <- true

		b.Configured = false
	}

	b.resetTimer()

	b.disconnectAll()

	return nil
}

func (b *Bridge) connect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true

	return nil
}

func (b *Bridge) reconnect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true
	b.LastError = nil

	return nil
}

func (b *Bridge) subscriptions() (err error) {

	if err = b.subscribe(b.local, b.remote, b.localTopics, "local"); err != nil {
		return err
	}

	if err = b.subscribe(b.remote, b.local, b.cloudTopics, "cloud"); err != nil {
		return err
	}
	return nil

}

func (b *Bridge) disconnectAll() {
	b.log.Infof("disconnectAll")
	// we are now disconnected
	b.Connected = false
	if b.local != nil && b.local.IsConnected() {
		b.local.Disconnect(100)
	}
	if b.remote != nil && b.remote.IsConnected() {
		b.remote.Disconnect(100)
	}
}

func (b *Bridge) mainBridgeLoop() {

	for {
		select {
		case <-b.reconnectCh:
			b.log.Infof("reconnecting")
			if err := b.reconnect(); err != nil {
				b.log.Errorf("Reconnect failed %s", err)
				b.scheduleReconnect(err)
			}
		case <-b.shutdownCh:
			b.log.Infof("shutting down bridge")
			return
		}

	}

}

func (b *Bridge) buildClient(server string, token string) (*mqtt.MqttClient, error) {

	b.log.Infof("building client for %s", server)

	opts := mqtt.NewClientOptions().AddBroker(server).SetTlsConfig(&tls.Config{InsecureSkipVerify: true})

	if token != "" {
		opts.SetUsername(token)
	}

	opts.SetClientId(fmt.Sprintf("%d", time.Now().Unix()))

	opts.SetKeepAlive(15) // set a 15 second ping time for ELB

	// pretty much log the reason and quit
	opts.SetOnConnectionLost(b.onConnectionLoss)

	client := mqtt.NewClient(opts)
	_, err := client.Start()
	return client, err
}

func (b *Bridge) subscribe(src *mqtt.MqttClient, dst *mqtt.MqttClient, topics []replaceTopic, tag string) (err error) {
	for _, topic := range topics {

		topicFilter, _ := mqtt.NewTopicFilter(topic.on, 0)
		b.log.Infof("(%s) subscribed to %s", tag, topic.on)

		if receipt, err := src.StartSubscription(b.buildHandler(topic, tag, dst), topicFilter); err != nil {
			return err
		} else {
			<-receipt
			b.log.Infof("(%s) subscribed to %+v", tag, topicFilter)
		}
	}

	return nil
}

func (b *Bridge) unsubscribe(client *mqtt.MqttClient, topics []replaceTopic, tag string) {
	topicNames := []string{}
	for _, topic := range topics {
		topicNames = append(topicNames, topic.on)
	}
	b.log.Infof("(%s) unsubscribed to %s", tag, topicNames)
	client.EndSubscription(topicNames...)
}

func (b *Bridge) buildHandler(topic replaceTopic, tag string, dst *mqtt.MqttClient) mqtt.MessageHandler {
	return func(src *mqtt.MqttClient, msg mqtt.Message) {
		if b.log.IsDebugEnabled() {
			b.log.Debugf("(%s) topic: %s updated: %s len: %d", tag, msg.Topic(), topic.updated(msg.Topic()), len(msg.Payload()))
		}
		b.Counter++
		payload := b.updateSource(msg.Payload(), tag)
		dst.PublishMessage(topic.updated(msg.Topic()), mqtt.NewMessage(payload))
	}
}

func (b *Bridge) scheduleReconnect(reason error) {
	b.LastError = reason
	b.disconnectAll()
	b.resetTimer()

	switch reason {
	case mqtt.ErrBadCredentials:
		b.log.Warningf("Reconnect failed trying again in 30s")

		b.timer = time.AfterFunc(30*time.Second, func() {
			b.reconnectCh <- true
		})

	default:
		b.log.Warningf("Reconnect failed trying again in 5s")
		// TODO add exponential backoff
		b.timer = time.AfterFunc(5*time.Second, func() {
			b.reconnectCh <- true
		})
	}

}

func (b *Bridge) resetTimer() {
	if b.timer != nil {
		b.timer.Stop()
	}
}

func (b *Bridge) onConnectionLoss(client *mqtt.MqttClient, reason error) {
	b.log.Errorf("Connection failed %s", reason)

	// we are now disconnected
	b.Connected = false

	b.scheduleReconnect(reason)

}

func (b *Bridge) IsConnected() bool {
	if b.remote == nil || b.local == nil {
		return false
	}
	return (b.remote.IsConnected() && b.local.IsConnected())
}

func (b *Bridge) updateSource(payload []byte, tag string) []byte {

	if !bytes.Contains(payload, []byte("$mesh-source")) {
		switch tag {
		case "local":
			payload = addMeshSource(b.conf.SerialNo, payload)
		case "cloud":
			payload = addMeshSource("cloud-"+strings.Replace(b.cloudUrl.Host, ".", "_", -1), payload) // encoded to look less wierd
		}
	}

	b.log.Debugf("msg %s", string(payload))

	return payload
}

func addMeshSource(source string, payload []byte) []byte {
	return bytes.Replace(payload, []byte("{"), []byte(`{"$mesh-source":"`+source+`", `), 1)
}
<commit_msg>Doing less per message and bit of tidy.<commit_after>package agent

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	mqtt "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/juju/loggo"
)

var AlreadyConfigured = errors.New("Already configured")
var AlreadyUnConfigured = errors.New("Already unconfigured")

//
// Acts as a bridge between local and cloud brokers, this includes reconnecting
// and emitting status changes.
//
// Once configured and started this will attempt to connect to connect to local
// and cloud brokers, if something dies it will reconnect based on the configured
// reconnect backoff.
//
type Bridge struct {
	conf   *Config
	local  *mqtt.MqttClient
	remote *mqtt.MqttClient
	log    loggo.Logger

	localTopics []replaceTopic
	cloudTopics []replaceTopic

	cloudUrl *url.URL
	token    string

	timer       *time.Timer
	reconnectCh chan bool
	shutdownCh  chan bool

	Configured bool
	Connected  bool
	Counter    int64

	LastError error

	bridgeLock sync.Mutex
}

type replaceTopic struct {
	on      string
	replace string
	with    string
}

func (r *replaceTopic) updated(originalTopic string) string {
	return strings.Replace(originalTopic, r.replace, r.with, 1)
}

var localTopics = []replaceTopic{
	// location related topics (TODO: move to cloud userspace RPC)
	{on: "$location/calibration", replace: "$location", with: "$cloud/location"},
	{on: "$location/delete", replace: "$location", with: "$cloud/location"},
	{on: "$device/+/+/rssi", replace: "$device", with: "$cloud/device"},

	// module health statistics
	{on: "$node/+/module/status", replace: "$node", with: "$cloud/node"},

	// cloud userspace RPC requests
	{on: "$ninja/services/rpc/+/+", replace: "$ninja", with: "$cloud/ninja"},
	{on: "$ninja/services/+", replace: "$ninja", with: "$cloud/ninja"},

	// temporary alternate topic to distinguish remote device replies from local-destined ones
	// used by the phone app for remote actuations
	// the alternate remote_ topic is to prevent a loopback with the below rule in the other direction
	// TODO: use a tag like $mesh-source to prevent loops (never re-proxy msgs with your source)
	{on: "$device/+/channel/+/reply", replace: "$device", with: "$cloud/remote_device"},

	// push up all local RPC methods in case the cloud is responding,
	// this is currently used for the push notification channel
	{on: "$device/+/channel/+", replace: "$device", with: "$cloud/device"},

	// push up state changes to the cloud
	{on: "$device/+/channel/+/event/state", replace: "$device", with: "$cloud/device"},
}

var cloudTopics = []replaceTopic{
	// location related topics
	{on: "$cloud/location/calibration/progress", replace: "$cloud/location", with: "$location"},
	{on: "$cloud/device/+/+/location", replace: "$cloud/device", with: "$device"},

	// cloud userspace RPC replies
	{on: "$cloud/ninja/services/rpc/+/+/reply", replace: "$cloud/ninja", with: "$ninja"},

	// see comment for $device/+/channel/+/reply above
	{on: "$cloud/remote_device/+/channel/+", replace: "$cloud/remote_device", with: "$device"},

	// allow cloud to announce devices and channels (used for phone on 3G and notification subscription channel)
	{on: "$cloud/device/+/event/announce", replace: "$cloud/device", with: "$device"},
	{on: "$cloud/device/+/channel/+/event/announce", replace: "$cloud/device", with: "$device"},

	// retrieve RPC responses from the cloud,
	// this is currently used for the push notification channel
	{on: "$cloud/device/+/channel/+/reply", replace: "$cloud/device", with: "$device"},
}

func createBridge(conf *Config) *Bridge {
	return &Bridge{conf: conf, localTopics: localTopics, cloudTopics: cloudTopics, log: loggo.GetLogger("bridge")}
}

func (b *Bridge) start(cloudUrl string, token string) (err error) {

	if b.Configured {
		b.log.Warningf("Already configured.")
		return AlreadyConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Connecting the bridge")

	b.Configured = true

	url, err := url.Parse(cloudUrl)

	if err != nil {
		return err
	}

	b.cloudUrl = url
	b.token = token

	b.reconnectCh = make(chan bool, 1)
	b.shutdownCh = make(chan bool, 1)

	if err = b.connect(); err != nil {
		b.log.Errorf("Connect failed %s", err)
		b.scheduleReconnect(err)
	}

	go b.mainBridgeLoop()

	return err
}

func (b *Bridge) stop() error {

	if !b.Configured {
		b.log.Warningf("Already unconfigured.")
		return AlreadyUnConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Disconnecting bridge")

	if b.Configured {
		// tell the worker to shutdown
		b.shutdownCh <- true

		b.Configured = false
	}

	b.resetTimer()

	b.disconnectAll()

	return nil
}

func (b *Bridge) connect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true

	return nil
}

func (b *Bridge) reconnect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true
	b.LastError = nil

	return nil
}

func (b *Bridge) subscriptions() (err error) {

	if err = b.subscribe(b.local, b.remote, b.localTopics, "local"); err != nil {
		return err
	}

	if err = b.subscribe(b.remote, b.local, b.cloudTopics, "cloud"); err != nil {
		return err
	}
	return nil

}

func (b *Bridge) disconnectAll() {
	b.log.Infof("disconnectAll")
	// we are now disconnected
	b.Connected = false
	if b.local != nil && b.local.IsConnected() {
		b.local.Disconnect(100)
	}
	if b.remote != nil && b.remote.IsConnected() {
		b.remote.Disconnect(100)
	}
}

func (b *Bridge) mainBridgeLoop() {

	for {
		select {
		case <-b.reconnectCh:
			b.log.Infof("reconnecting")
			if err := b.reconnect(); err != nil {
				b.log.Errorf("Reconnect failed %s", err)
				b.scheduleReconnect(err)
			}
		case <-b.shutdownCh:
			b.log.Infof("shutting down bridge")
			return
		}

	}

}

func (b *Bridge) buildClient(server string, token string) (*mqtt.MqttClient, error) {

	b.log.Infof("building client for %s", server)

	opts := mqtt.NewClientOptions().AddBroker(server).SetTlsConfig(&tls.Config{InsecureSkipVerify: true})

	if token != "" {
		opts.SetUsername(token)
	}

	opts.SetClientId(fmt.Sprintf("%d", time.Now().Unix()))

	opts.SetKeepAlive(15) // set a 15 second ping time for ELB

	// pretty much log the reason and quit
	opts.SetOnConnectionLost(b.onConnectionLoss)

	client := mqtt.NewClient(opts)
	_, err := client.Start()

	return client, err
}

func (b *Bridge) subscribe(src *mqtt.MqttClient, dst *mqtt.MqttClient, topics []replaceTopic, tag string) (err error) {

	var (
		handler mqtt.MessageHandler
	)

	for _, topic := range topics {

		topicFilter, _ := mqtt.NewTopicFilter(topic.on, 0)
		b.log.Infof("(%s) subscribed to %s", tag, topic.on)

		handler = b.buildHandler(topic, b.buildSource(tag), b.remote)

		if _, err = src.StartSubscription(handler, topicFilter); err != nil {
			return err
		}

		b.log.Infof("(%s) subscribed to %+v", tag, topicFilter)
	}

	return nil
}

func (b *Bridge) unsubscribe(client *mqtt.MqttClient, topics []replaceTopic, tag string) {

	topicNames := []string{}

	for _, topic := range topics {
		topicNames = append(topicNames, topic.on)
	}

	b.log.Infof("(%s) unsubscribed to %s", tag, topicNames)
	client.EndSubscription(topicNames...)
}

func (b *Bridge) buildHandler(topic replaceTopic, source string, dst *mqtt.MqttClient) mqtt.MessageHandler {
	return func(src *mqtt.MqttClient, msg mqtt.Message) {
		if b.log.IsDebugEnabled() {
			b.log.Debugf("(%s) topic: %s updated: %s len: %d", source, msg.Topic(), topic.updated(msg.Topic()), len(msg.Payload()))
		}
		b.Counter++
		payload := b.updateSource(msg.Payload(), source)
		dst.PublishMessage(topic.updated(msg.Topic()), mqtt.NewMessage(payload))
	}
}

func (b *Bridge) scheduleReconnect(reason error) {
	b.LastError = reason
	b.disconnectAll()
	b.resetTimer()

	switch reason {
	case mqtt.ErrBadCredentials:
		b.log.Warningf("Reconnect failed trying again in 30s")

		b.timer = time.AfterFunc(30*time.Second, func() {
			b.reconnectCh <- true
		})

	default:
		b.log.Warningf("Reconnect failed trying again in 5s")
		// TODO add exponential backoff
		b.timer = time.AfterFunc(5*time.Second, func() {
			b.reconnectCh <- true
		})
	}

}

func (b *Bridge) resetTimer() {
	if b.timer != nil {
		b.timer.Stop()
	}
}

func (b *Bridge) onConnectionLoss(client *mqtt.MqttClient, reason error) {
	b.log.Errorf("Connection failed %s", reason)

	// we are now disconnected
	b.Connected = false

	b.scheduleReconnect(reason)

}

func (b *Bridge) IsConnected() bool {
	if b.remote == nil || b.local == nil {
		return false
	}
	return (b.remote.IsConnected() && b.local.IsConnected())
}

func (b *Bridge) buildSource(tag string) string {

	switch tag {
	case "local":
		return b.conf.SerialNo
	case "cloud":
		return "cloud-" + strings.Replace(b.cloudUrl.Host, ".", "_", -1) // encoded to look less wierd
	}

	return ""
}

func (b *Bridge) updateSource(payload []byte, source string) []byte {

	if !bytes.Contains(payload, []byte("$mesh-source")) {
		payload = bytes.Replace(payload, []byte("{"), []byte(`{"$mesh-source":"`+source+`", `), 1)
	}

	b.log.Debugf("msg %s", string(payload))

	return payload
}
<|endoftext|>


<commit_before>package backends

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"github.com/squaremo/ambergreen/common/data"

	etcd_errors "github.com/coreos/etcd/error"
	"github.com/coreos/go-etcd/etcd"
)

const servicePath = "/weave/service/"

type Backend struct {
	client *etcd.Client
}

func NewBackendFromEnv() *Backend {
	etcd_address := os.Getenv("ETCD_PORT")
	if etcd_address == "" {
		etcd_address = os.Getenv("ETCD_ADDRESS")
	}
	if strings.HasPrefix(etcd_address, "tcp:") {
		etcd_address = "http:" + etcd_address[4:]
	}
	if etcd_address == "" {
		etcd_address = "http://127.0.0.1:4001"
	}

	return NewBackend(etcd_address)
}

func NewBackend(addr string) *Backend {
	return &Backend{client: etcd.NewClient([]string{addr})}
}

// Check if we can talk to etcd
func (b *Backend) Ping() error {
	rr := etcd.NewRawRequest("GET", "version", nil, nil)
	_, err := b.client.SendRequest(rr)
	return err
}

func (b *Backend) CheckRegisteredService(serviceName string) error {
	_, err := b.client.Get(servicePath+serviceName, false, false)
	return err
}

func (b *Backend) AddService(serviceName string, details data.Service) error {
	json, err := json.Marshal(&details)
	if err != nil {
		return fmt.Errorf("Failed to encode: %s", err)
	}
	_, err = b.client.Set(servicePath+serviceName+"/details", string(json), 0)
	return err
}

func (b *Backend) RemoveService(serviceName string) error {
	_, err := b.client.Delete(servicePath+serviceName, true)
	return err
}

func (b *Backend) RemoveAllServices() error {
	_, err := b.client.Delete(servicePath, true)
	return err
}

func (b *Backend) GetServiceDetails(serviceName string) (data.Service, error) {
	var service data.Service
	details, err := b.client.Get(servicePath+serviceName+"/details", false, false)
	if err != nil {
		return service, err
	}
	if err := json.Unmarshal([]byte(details.Node.Value), &service); err != nil {
		return service, err
	}
	return service, nil
}

func (b *Backend) ForeachServiceInstance(fs ServiceFunc, fi ServiceInstanceFunc) error {
	r, err := b.client.Get(servicePath, true, fi != nil)
	if err != nil {
		if etcderr, ok := err.(*etcd.EtcdError); ok && etcderr.ErrorCode == etcd_errors.EcodeKeyNotFound {
			return nil
		}
		return err
	}
	for _, node := range r.Node.Nodes {
		serviceName := strings.TrimPrefix(node.Key, servicePath)
		serviceData, err := b.GetServiceDetails(serviceName)
		if err != nil {
			return err
		}
		if fs != nil {
			fs(serviceName, serviceData)
		}
		if fi != nil {
			for _, instance := range node.Nodes {
				if strings.HasSuffix(instance.Key, "/details") {
					continue
				}
				var instanceData data.Instance
				if err := json.Unmarshal([]byte(instance.Value), &instanceData); err != nil {
					return err
				}
				fi(serviceName, strings.TrimPrefix(instance.Key, node.Key+"/"), instanceData)
			}
		}
	}
	return nil
}

func (b *Backend) ForeachInstance(serviceName string, fi InstanceFunc) error {
	serviceKey := servicePath + serviceName + "/"
	r, err := b.client.Get(serviceKey, true, false)
	if err != nil {
		if etcderr, ok := err.(*etcd.EtcdError); ok && etcderr.ErrorCode == etcd_errors.EcodeKeyNotFound {
			return nil
		}
		return err
	}
	for _, instance := range r.Node.Nodes {
		if strings.HasSuffix(instance.Key, "/details") {
			continue
		}
		var instanceData data.Instance
		if err := json.Unmarshal([]byte(instance.Value), &instanceData); err != nil {
			return err
		}
		fi(strings.TrimPrefix(instance.Key, serviceKey), instanceData)
	}
	return nil
}

func (b *Backend) AddInstance(serviceName string, instanceName string, details data.Instance) error {
	json, err := json.Marshal(details)
	if err != nil {
		return fmt.Errorf("Failed to encode: %s", err)
	}
	if _, err := b.client.Set(servicePath+serviceName+"/"+instanceName, string(json), 0); err != nil {
		return fmt.Errorf("Unable to write: %s", err)
	}
	return nil
}

func (b *Backend) RemoveInstance(serviceName, instanceName string) error {
	_, err := b.client.Delete(servicePath+serviceName+"/"+instanceName, true)
	return err
}

func (b *Backend) WatchServices(resCh chan<- data.ServiceChange, stopCh <-chan struct{}, withInstanceChanges bool) {
	// XXX error handling

	etcdCh := make(chan *etcd.Response, 1)
	watchStopCh := make(chan bool, 1)
	go b.client.Watch(servicePath, 0, true, etcdCh, nil)

	svcs := make(map[string]struct{})

	handleResponse := func(r *etcd.Response) {
		path := r.Node.Key
		switch r.Action {
		case "delete":
			if len(path) <= len(servicePath) {
				// All services deleted
				for name := range svcs {
					resCh <- data.ServiceChange{name, true}
				}
				svcs = make(map[string]struct{})
				return
			}

			p := strings.Split(path[len(servicePath):], "/")
			if len(p) == 1 {
				// Service deleted
				delete(svcs, p[0])
				resCh <- data.ServiceChange{p[0], true}
				return
			}

			// Instance deletion
			if withInstanceChanges {
				resCh <- data.ServiceChange{p[0], false}
			}

		case "set":
			if len(path) <= len(servicePath) {
				return
			}

			p := strings.Split(path[len(servicePath):], "/")
			svcs[p[0]] = struct{}{}
			if withInstanceChanges ||
				len(p) == 2 && p[1] == "details" {
				resCh <- data.ServiceChange{p[0], false}
			}
		}
	}

	b.ForeachServiceInstance(func(name string, svc data.Service) {
		svcs[name] = struct{}{}
	}, nil)

	go func() {
		for {
			select {
			case <-stopCh:
				watchStopCh <- true
				return

			case r := <-etcdCh:
				handleResponse(r)
			}
		}
	}()
}
<commit_msg>Rearrange the etcd services structure<commit_after>package backends

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"github.com/squaremo/ambergreen/common/data"

	etcd_errors "github.com/coreos/etcd/error"
	"github.com/coreos/go-etcd/etcd"
)

type Backend struct {
	client *etcd.Client
}

func NewBackendFromEnv() *Backend {
	etcd_address := os.Getenv("ETCD_PORT")
	if etcd_address == "" {
		etcd_address = os.Getenv("ETCD_ADDRESS")
	}
	if strings.HasPrefix(etcd_address, "tcp:") {
		etcd_address = "http:" + etcd_address[4:]
	}
	if etcd_address == "" {
		etcd_address = "http://127.0.0.1:4001"
	}

	return NewBackend(etcd_address)
}

func NewBackend(addr string) *Backend {
	return &Backend{client: etcd.NewClient([]string{addr})}
}

// Check if we can talk to etcd
func (b *Backend) Ping() error {
	rr := etcd.NewRawRequest("GET", "version", nil, nil)
	_, err := b.client.SendRequest(rr)
	return err
}

const ROOT = "/weave/service/"

func serviceRootKey(serviceName string) string {
	return ROOT + serviceName
}

func serviceKey(serviceName string) string {
	return fmt.Sprintf("%s%s/details", ROOT, serviceName)
}

func instanceKey(serviceName string, instanceName string) string {
	return fmt.Sprintf("%s%s/instance/%s", ROOT, serviceName, instanceName)
}

type parsedRootKey struct {
}

type parsedServiceRootKey struct {
	serviceName string
}

type parsedServiceKey struct {
	serviceName string
}

type parsedInstanceKey struct {
	serviceName  string
	instanceName string
}

// Parse a path to find its type

func parseKey(key string) interface{} {
	if len(key) <= len(ROOT) {
		return parsedRootKey{}
	}

	p := strings.Split(key[len(ROOT):], "/")
	if len(p) == 1 {
		return parsedServiceRootKey{p[0]}
	}

	switch p[1] {
	case "details":
		return parsedServiceKey{p[0]}

	case "instance":
		if len(p) == 3 {
			return parsedInstanceKey{p[0], p[2]}
		}
	}

	return nil
}

func (b *Backend) CheckRegisteredService(serviceName string) error {
	_, err := b.client.Get(serviceRootKey(serviceName), false, false)
	return err
}

func (b *Backend) AddService(serviceName string, details data.Service) error {
	json, err := json.Marshal(&details)
	if err != nil {
		return fmt.Errorf("Failed to encode: %s", err)
	}
	_, err = b.client.Set(serviceKey(serviceName), string(json), 0)
	return err
}

func (b *Backend) RemoveService(serviceName string) error {
	_, err := b.client.Delete(serviceRootKey(serviceName), true)
	return err
}

func (b *Backend) RemoveAllServices() error {
	_, err := b.client.Delete(ROOT, true)
	return err
}

func (b *Backend) GetServiceDetails(serviceName string) (data.Service, error) {
	r, err := b.client.Get(serviceKey(serviceName), false, false)
	if err != nil {
		return data.Service{}, err
	}

	return unmarshalService(r.Node)
}

func unmarshalService(node *etcd.Node) (data.Service, error) {
	var service data.Service
	err := json.Unmarshal([]byte(node.Value), &service)
	return service, err
}

func unmarshalInstance(node *etcd.Node) (data.Instance, error) {
	var instance data.Instance
	err := json.Unmarshal([]byte(node.Value), &instance)
	return instance, err
}

func traverse(node *etcd.Node, f func(*etcd.Node) error) error {
	for _, child := range node.Nodes {
		var err error
		if child.Dir {
			err = traverse(child, f)
		} else {
			err = f(child)
		}

		if err != nil {
			return err
		}
	}

	return nil
}

func (b *Backend) traverse(key string, f func(*etcd.Node) error) error {
	r, err := b.client.Get(key, false, true)
	if err != nil {
		if etcderr, ok := err.(*etcd.EtcdError); ok && etcderr.ErrorCode == etcd_errors.EcodeKeyNotFound {
			return nil
		}
		return err
	}

	return traverse(r.Node, f)
}

func (b *Backend) ForeachServiceInstance(fs ServiceFunc, fi ServiceInstanceFunc) error {
	return b.traverse(ROOT, func(node *etcd.Node) error {
		switch key := parseKey(node.Key).(type) {
		case parsedServiceKey:
			if fs != nil {
				svc, err := unmarshalService(node)
				if err != nil {
					return err
				}

				fs(key.serviceName, svc)
			}

		case parsedInstanceKey:
			if fi != nil {
				inst, err := unmarshalInstance(node)
				if err != nil {
					return err
				}

				fi(key.serviceName, key.instanceName, inst)
			}
		}

		return nil
	})
}

func (b *Backend) AddInstance(serviceName string, instanceName string, details data.Instance) error {
	json, err := json.Marshal(details)
	if err != nil {
		return fmt.Errorf("Failed to encode: %s", err)
	}
	if _, err := b.client.Set(instanceKey(serviceName, instanceName), string(json), 0); err != nil {
		return fmt.Errorf("Unable to write: %s", err)
	}
	return nil
}

func (b *Backend) RemoveInstance(serviceName, instanceName string) error {
	_, err := b.client.Delete(instanceKey(serviceName, instanceName), true)
	return err
}

func (b *Backend) ForeachInstance(serviceName string, fi InstanceFunc) error {
	return b.traverse(serviceRootKey(serviceName), func(node *etcd.Node) error {
		switch key := parseKey(node.Key).(type) {
		case parsedInstanceKey:
			if fi != nil {
				inst, err := unmarshalInstance(node)
				if err != nil {
					return err
				}

				fi(key.instanceName, inst)
			}
		}

		return nil
	})
}

func (b *Backend) WatchServices(resCh chan<- data.ServiceChange, stopCh <-chan struct{}, withInstanceChanges bool) {
	// XXX error handling

	etcdCh := make(chan *etcd.Response, 1)
	watchStopCh := make(chan bool, 1)
	go b.client.Watch(ROOT, 0, true, etcdCh, nil)

	svcs := make(map[string]struct{})

	handleResponse := func(r *etcd.Response) {
		switch r.Action {
		case "delete":
			switch key := parseKey(r.Node.Key).(type) {
			case parsedRootKey:
				for name := range svcs {
					resCh <- data.ServiceChange{name, true}
				}
				svcs = make(map[string]struct{})

			case parsedServiceRootKey:
				delete(svcs, key.serviceName)
				resCh <- data.ServiceChange{key.serviceName, true}

			case parsedInstanceKey:
				if withInstanceChanges {
					resCh <- data.ServiceChange{key.serviceName, false}
				}
			}

		case "set":
			switch key := parseKey(r.Node.Key).(type) {
			case parsedServiceKey:
				svcs[key.serviceName] = struct{}{}
				resCh <- data.ServiceChange{key.serviceName, false}

			case parsedInstanceKey:
				if withInstanceChanges {
					resCh <- data.ServiceChange{key.serviceName, false}
				}
			}
		}
	}

	b.ForeachServiceInstance(func(name string, svc data.Service) {
		svcs[name] = struct{}{}
	}, nil)

	go func() {
		for {
			select {
			case <-stopCh:
				watchStopCh <- true
				return

			case r := <-etcdCh:
				handleResponse(r)
			}
		}
	}()
}
<|endoftext|>


<commit_before>package debug

import (
	"log"
	"sync/atomic"
	"time"
)

// LogDebugger is the simplest debugger which prints log messages to the STDERR
type LogDebugger struct {
	counter int32
	start   time.Time
}

// Init initializes the LogDebugger
func (l *LogDebugger) Init() error {
	l.counter = 0
	l.start = time.Now()
	return nil
}

// Event receives Collector events and prints them to STDERR
func (l *LogDebugger) Event(e *Event) {
	i := atomic.AddInt32(&l.counter, 1)
	log.Printf("[%06d] %d [%6d - %s] %q (%s)\n", i, e.CollectorId, e.RequestId, e.Type, e.Values, time.Since(l.start))
}
<commit_msg>[enh] add ability to set custom output and flags to logdebugger<commit_after>package debug

import (
	"io"
	"log"
	"os"
	"sync/atomic"
	"time"
)

// LogDebugger is the simplest debugger which prints log messages to the STDERR
type LogDebugger struct {
	// Output is the log destination, anything can be used which implements them
	// io.Writer interface. Leave it blank to use STDERR
	Output io.Writer
	// Prefix appears at the beginning of each generated log line
	Prefix string
	// Flag defines the logging properties.
	Flag    int
	logger  *log.Logger
	counter int32
	start   time.Time
}

// Init initializes the LogDebugger
func (l *LogDebugger) Init() error {
	l.counter = 0
	l.start = time.Now()
	if l.Output == nil {
		l.Output = os.Stderr
	}
	l.logger = log.New(l.Output, l.Prefix, l.Flag)
	return nil
}

// Event receives Collector events and prints them to STDERR
func (l *LogDebugger) Event(e *Event) {
	i := atomic.AddInt32(&l.counter, 1)
	l.logger.Printf("[%06d] %d [%6d - %s] %q (%s)\n", i, e.CollectorId, e.RequestId, e.Type, e.Values, time.Since(l.start))
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"fmt"
	"hash/crc64"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/googlecodelabs/tools/claat/parser"
	"github.com/googlecodelabs/tools/claat/types"
)

const (
	// supported codelab source types must be registered parsers
	// TODO: define these in claat/parser/..., e.g. in parser/gdoc
	srcInvalid   srcType = ""
	srcGoogleDoc srcType = "gdoc" // Google Docs doc
	srcMarkdown  srcType = "md"   // Markdown text

	// driveAPI is a base URL for Drive API
	driveAPI = "https://www.googleapis.com/drive/v3"
)

// srcType is codelab source type
type srcType string

// resource is a codelab resource, loaded from local file
// or fetched from remote location.
type resource struct {
	typ  srcType       // source type
	body io.ReadCloser // resource body
	mod  time.Time     // last update of content
}

// codelab wraps types.Codelab, while adding source type
// and modified timestamp fields.
type codelab struct {
	*types.Codelab
	typ srcType   //  source type
	mod time.Time // last modified timestamp
}

// slurpCodelab retrieves and parses codelab source.
// It returns parsed codelab and its source type.
//
// The function will also fetch and parse fragments included
// with types.ImportNode.
func slurpCodelab(src string) (*codelab, error) {
	res, err := fetch(src)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()
	clab, err := parser.Parse(string(res.typ), res.body)
	if err != nil {
		return nil, err
	}

	// fetch imports and parse them as fragments
	var imports []*types.ImportNode
	for _, st := range clab.Steps {
		imports = append(imports, importNodes(st.Content.Nodes)...)
	}
	ch := make(chan error, len(imports))
	defer close(ch)
	for _, imp := range imports {
		go func(n *types.ImportNode) {
			frag, err := slurpFragment(n.URL)
			if err != nil {
				ch <- fmt.Errorf("%s: %v", n.URL, err)
				return
			}
			n.Content.Nodes = frag
			ch <- nil
		}(imp)
	}
	for _ = range imports {
		if err := <-ch; err != nil {
			return nil, err
		}
	}

	v := &codelab{
		Codelab: clab,
		typ:     res.typ,
		mod:     res.mod,
	}
	return v, nil
}

func slurpFragment(url string) ([]types.Node, error) {
	res, err := fetchRemote(url, true)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()
	return parser.ParseFragment(string(res.typ), res.body)
}

// fetch retrieves codelab doc either from local disk
// or a remote location.
// The caller is responsible for closing returned stream.
func fetch(name string) (*resource, error) {
	fi, err := os.Stat(name)
	if os.IsNotExist(err) {
		return fetchRemote(name, false)
	}
	r, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	return &resource{
		body: r,
		typ:  srcMarkdown,
		mod:  fi.ModTime(),
	}, nil
}

// fetchRemote retrieves resource r from the network.
//
// If urlStr is not a URL, i.e. does not have the host part, it is considered to be
// a Google Doc ID and fetched accordingly. Otherwise, a simple GET request
// is used to retrieve the contents.
//
// The caller is responsible for closing returned stream.
// If nometa is true, resource.mod may have zero value.
func fetchRemote(urlStr string, nometa bool) (*resource, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}
	if u.Host == "" || u.Host == "docs.google.com" {
		return fetchDriveFile(urlStr, nometa)
	}
	return fetchRemoteFile(urlStr)
}

// fetchRemoteFile retrieves codelab resource from url.
// It is a special case of fetchRemote function.
func fetchRemoteFile(url string) (*resource, error) {
	res, err := retryGet(nil, url, 3)
	if err != nil {
		return nil, err
	}
	t, err := http.ParseTime(res.Header.Get("last-modified"))
	if err != nil {
		t = time.Now()
	}
	return &resource{
		body: res.Body,
		mod:  t,
		typ:  srcMarkdown,
	}, nil
}

// fetchDriveFile uses Drive API to retrieve HTML representation of a Google Doc.
// See https://developers.google.com/drive/web/manage-downloads#downloading_google_documents
// for more details.
//
// If nometa is true, resource.mod will have zero value.
func fetchDriveFile(id string, nometa bool) (*resource, error) {
	id = gdocID(id)
	exportURL := gdocExportURL(id)
	client, err := driveClient()
	if err != nil {
		return nil, err
	}

	if nometa {
		res, err := retryGet(client, exportURL, 7)
		if err != nil {
			return nil, err
		}
		return &resource{body: res.Body, typ: srcGoogleDoc}, nil
	}

	u := fmt.Sprintf("%s/files/%s?fields=id,mimeType,modifiedTime", driveAPI, id)
	res, err := retryGet(client, u, 7)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	meta := &struct {
		ID       string    `json:"id"`
		MimeType string    `json:"mimeType"`
		Modified time.Time `json:"modifiedTime"`
	}{}
	if err := json.NewDecoder(res.Body).Decode(meta); err != nil {
		return nil, err
	}
	if meta.MimeType != "application/vnd.google-apps.document" {
		return nil, fmt.Errorf("%s: invalid mime type: %s", id, meta.MimeType)
	}

	if res, err = retryGet(client, exportURL, 7); err != nil {
		return nil, err
	}
	return &resource{
		body: res.Body,
		mod:  meta.Modified,
		typ:  srcGoogleDoc,
	}, nil
}

var crcTable = crc64.MakeTable(crc64.ECMA)

func slurpBytes(client *http.Client, codelabSrc, dir, imgURL string, n int) (string, error) {
	// images can be local in Markdown cases or remote.
	// Only proceed a simple copy on local reference.
	var b []byte
	var ext string
	u, err := url.Parse(imgURL)
	if err != nil {
		return "", err
	}
	if u.Host == "" {
		if imgURL, err = restrictPathToParent(imgURL, filepath.Dir(codelabSrc)); err != nil {
			return "", err
		}
		b, err = ioutil.ReadFile(imgURL)
		ext = filepath.Ext(imgURL)
	} else {
		b, err = slurpRemoteBytes(client, dir, imgURL, 5)
		ext = ".png"
	}
	if err != nil {
		return "", err
	}

	crc := crc64.Checksum(b, crcTable)
	file := fmt.Sprintf("%x%s", crc, ext)
	dst := filepath.Join(dir, file)
	return file, ioutil.WriteFile(dst, b, 0644)
}

func slurpRemoteBytes(client *http.Client, dir, url string, n int) ([]byte, error) {
	res, err := retryGet(client, url, n)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// retryGet tries to GET specified url up to n times.
// Default client will be used if not provided.
func retryGet(client *http.Client, url string, n int) (*http.Response, error) {
	if client == nil {
		client = http.DefaultClient
	}
	for i := 0; i <= n; i++ {
		if i > 0 {
			t := time.Duration((math.Pow(2, float64(i)) + rand.Float64()) * float64(time.Second))
			time.Sleep(t)
		}
		res, err := client.Get(url)
		// return early with a good response
		// the rest is error handling
		if err == nil && res.StatusCode == http.StatusOK {
			return res, nil
		}

		// sometimes Drive API wouldn't even start a response,
		// we get net/http: TLS handshake timeout instead:
		// consider this a temporary failure and retry again
		if err != nil {
			continue
		}
		// otherwise, decode error response and check for "rate limit"
		defer res.Body.Close()
		var erres struct {
			Error struct {
				Errors []struct{ Reason string }
			}
		}
		b, _ := ioutil.ReadAll(res.Body)
		json.Unmarshal(b, &erres)
		var rateLimit bool
		for _, e := range erres.Error.Errors {
			if e.Reason == "rateLimitExceeded" || e.Reason == "userRateLimitExceeded" {
				rateLimit = true
				break
			}
		}
		// this is neither a rate limit error, nor a server error:
		// retrying is useless
		if !rateLimit && res.StatusCode < http.StatusInternalServerError {
			return nil, fmt.Errorf("fetch %s: %s; %s", url, res.Status, b)
		}
	}
	return nil, fmt.Errorf("%s: failed after %d retries", url, n)
}

func gdocID(url string) string {
	const s = "/document/d/"
	if i := strings.Index(url, s); i >= 0 {
		url = url[i+len(s):]
	}
	if i := strings.IndexRune(url, '/'); i > 0 {
		url = url[:i]
	}
	return url
}

func gdocExportURL(id string) string {
	return fmt.Sprintf("%s/files/%s/export?mimeType=text/html", driveAPI, id)
}

// restrictPathToParent will ensure that assetPath is in parent.
// It will thus return an absolute path to the asset.
func restrictPathToParent(assetPath, parent string) (string, error) {
	parent, err := filepath.Abs(parent)
	if err != nil {
		return "", err
	}
	if !filepath.IsAbs(assetPath) {
		assetPath = filepath.Join(parent, assetPath)
	}
	if !strings.HasPrefix(assetPath, parent) {
		return "", fmt.Errorf("%s isn't a subdirectory of %s", assetPath, parent)
	}
	return assetPath, nil
}
<commit_msg>Added team drive support from cli<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"fmt"
	"hash/crc64"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/googlecodelabs/tools/claat/parser"
	"github.com/googlecodelabs/tools/claat/types"
)

const (
	// supported codelab source types must be registered parsers
	// TODO: define these in claat/parser/..., e.g. in parser/gdoc
	srcInvalid   srcType = ""
	srcGoogleDoc srcType = "gdoc" // Google Docs doc
	srcMarkdown  srcType = "md"   // Markdown text

	// driveAPI is a base URL for Drive API
	driveAPI = "https://www.googleapis.com/drive/v3"
)

// srcType is codelab source type
type srcType string

// resource is a codelab resource, loaded from local file
// or fetched from remote location.
type resource struct {
	typ  srcType       // source type
	body io.ReadCloser // resource body
	mod  time.Time     // last update of content
}

// codelab wraps types.Codelab, while adding source type
// and modified timestamp fields.
type codelab struct {
	*types.Codelab
	typ srcType   //  source type
	mod time.Time // last modified timestamp
}

// slurpCodelab retrieves and parses codelab source.
// It returns parsed codelab and its source type.
//
// The function will also fetch and parse fragments included
// with types.ImportNode.
func slurpCodelab(src string) (*codelab, error) {
	res, err := fetch(src)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()
	clab, err := parser.Parse(string(res.typ), res.body)
	if err != nil {
		return nil, err
	}

	// fetch imports and parse them as fragments
	var imports []*types.ImportNode
	for _, st := range clab.Steps {
		imports = append(imports, importNodes(st.Content.Nodes)...)
	}
	ch := make(chan error, len(imports))
	defer close(ch)
	for _, imp := range imports {
		go func(n *types.ImportNode) {
			frag, err := slurpFragment(n.URL)
			if err != nil {
				ch <- fmt.Errorf("%s: %v", n.URL, err)
				return
			}
			n.Content.Nodes = frag
			ch <- nil
		}(imp)
	}
	for _ = range imports {
		if err := <-ch; err != nil {
			return nil, err
		}
	}

	v := &codelab{
		Codelab: clab,
		typ:     res.typ,
		mod:     res.mod,
	}
	return v, nil
}

func slurpFragment(url string) ([]types.Node, error) {
	res, err := fetchRemote(url, true)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()
	return parser.ParseFragment(string(res.typ), res.body)
}

// fetch retrieves codelab doc either from local disk
// or a remote location.
// The caller is responsible for closing returned stream.
func fetch(name string) (*resource, error) {
	fi, err := os.Stat(name)
	if os.IsNotExist(err) {
		return fetchRemote(name, false)
	}
	r, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	return &resource{
		body: r,
		typ:  srcMarkdown,
		mod:  fi.ModTime(),
	}, nil
}

// fetchRemote retrieves resource r from the network.
//
// If urlStr is not a URL, i.e. does not have the host part, it is considered to be
// a Google Doc ID and fetched accordingly. Otherwise, a simple GET request
// is used to retrieve the contents.
//
// The caller is responsible for closing returned stream.
// If nometa is true, resource.mod may have zero value.
func fetchRemote(urlStr string, nometa bool) (*resource, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}
	if u.Host == "" || u.Host == "docs.google.com" {
		return fetchDriveFile(urlStr, nometa)
	}
	return fetchRemoteFile(urlStr)
}

// fetchRemoteFile retrieves codelab resource from url.
// It is a special case of fetchRemote function.
func fetchRemoteFile(url string) (*resource, error) {
	res, err := retryGet(nil, url, 3)
	if err != nil {
		return nil, err
	}
	t, err := http.ParseTime(res.Header.Get("last-modified"))
	if err != nil {
		t = time.Now()
	}
	return &resource{
		body: res.Body,
		mod:  t,
		typ:  srcMarkdown,
	}, nil
}

// fetchDriveFile uses Drive API to retrieve HTML representation of a Google Doc.
// See https://developers.google.com/drive/web/manage-downloads#downloading_google_documents
// for more details.
//
// If nometa is true, resource.mod will have zero value.
func fetchDriveFile(id string, nometa bool) (*resource, error) {
	id = gdocID(id)
	exportURL := gdocExportURL(id)
	client, err := driveClient()
	if err != nil {
		return nil, err
	}

	if nometa {
		res, err := retryGet(client, exportURL, 7)
		if err != nil {
			return nil, err
		}
		return &resource{body: res.Body, typ: srcGoogleDoc}, nil
	}

	u := fmt.Sprintf("%s/files/%s?fields=id,mimeType,modifiedTime&supportsTeamDrives=true", driveAPI, id)
	res, err := retryGet(client, u, 7)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	meta := &struct {
		ID       string    `json:"id"`
		MimeType string    `json:"mimeType"`
		Modified time.Time `json:"modifiedTime"`
	}{}
	if err := json.NewDecoder(res.Body).Decode(meta); err != nil {
		return nil, err
	}
	if meta.MimeType != "application/vnd.google-apps.document" {
		return nil, fmt.Errorf("%s: invalid mime type: %s", id, meta.MimeType)
	}

	if res, err = retryGet(client, exportURL, 7); err != nil {
		return nil, err
	}
	return &resource{
		body: res.Body,
		mod:  meta.Modified,
		typ:  srcGoogleDoc,
	}, nil
}

var crcTable = crc64.MakeTable(crc64.ECMA)

func slurpBytes(client *http.Client, codelabSrc, dir, imgURL string, n int) (string, error) {
	// images can be local in Markdown cases or remote.
	// Only proceed a simple copy on local reference.
	var b []byte
	var ext string
	u, err := url.Parse(imgURL)
	if err != nil {
		return "", err
	}
	if u.Host == "" {
		if imgURL, err = restrictPathToParent(imgURL, filepath.Dir(codelabSrc)); err != nil {
			return "", err
		}
		b, err = ioutil.ReadFile(imgURL)
		ext = filepath.Ext(imgURL)
	} else {
		b, err = slurpRemoteBytes(client, dir, imgURL, 5)
		ext = ".png"
	}
	if err != nil {
		return "", err
	}

	crc := crc64.Checksum(b, crcTable)
	file := fmt.Sprintf("%x%s", crc, ext)
	dst := filepath.Join(dir, file)
	return file, ioutil.WriteFile(dst, b, 0644)
}

func slurpRemoteBytes(client *http.Client, dir, url string, n int) ([]byte, error) {
	res, err := retryGet(client, url, n)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// retryGet tries to GET specified url up to n times.
// Default client will be used if not provided.
func retryGet(client *http.Client, url string, n int) (*http.Response, error) {
	if client == nil {
		client = http.DefaultClient
	}
	for i := 0; i <= n; i++ {
		if i > 0 {
			t := time.Duration((math.Pow(2, float64(i)) + rand.Float64()) * float64(time.Second))
			time.Sleep(t)
		}
		res, err := client.Get(url)
		// return early with a good response
		// the rest is error handling
		if err == nil && res.StatusCode == http.StatusOK {
			return res, nil
		}

		// sometimes Drive API wouldn't even start a response,
		// we get net/http: TLS handshake timeout instead:
		// consider this a temporary failure and retry again
		if err != nil {
			continue
		}
		// otherwise, decode error response and check for "rate limit"
		defer res.Body.Close()
		var erres struct {
			Error struct {
				Errors []struct{ Reason string }
			}
		}
		b, _ := ioutil.ReadAll(res.Body)
		json.Unmarshal(b, &erres)
		var rateLimit bool
		for _, e := range erres.Error.Errors {
			if e.Reason == "rateLimitExceeded" || e.Reason == "userRateLimitExceeded" {
				rateLimit = true
				break
			}
		}
		// this is neither a rate limit error, nor a server error:
		// retrying is useless
		if !rateLimit && res.StatusCode < http.StatusInternalServerError {
			return nil, fmt.Errorf("fetch %s: %s; %s", url, res.Status, b)
		}
	}
	return nil, fmt.Errorf("%s: failed after %d retries", url, n)
}

func gdocID(url string) string {
	const s = "/document/d/"
	if i := strings.Index(url, s); i >= 0 {
		url = url[i+len(s):]
	}
	if i := strings.IndexRune(url, '/'); i > 0 {
		url = url[:i]
	}
	return url
}

func gdocExportURL(id string) string {
	return fmt.Sprintf("%s/files/%s/export?mimeType=text/html", driveAPI, id)
}

// restrictPathToParent will ensure that assetPath is in parent.
// It will thus return an absolute path to the asset.
func restrictPathToParent(assetPath, parent string) (string, error) {
	parent, err := filepath.Abs(parent)
	if err != nil {
		return "", err
	}
	if !filepath.IsAbs(assetPath) {
		assetPath = filepath.Join(parent, assetPath)
	}
	if !strings.HasPrefix(assetPath, parent) {
		return "", fmt.Errorf("%s isn't a subdirectory of %s", assetPath, parent)
	}
	return assetPath, nil
}
<|endoftext|>


<commit_before>package clair

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/optiopay/klar/docker"
)

const EMPTY_LAYER_BLOB_SUM = "sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4"

// Clair is representation of Clair server
type Clair struct {
	url string
	api API
}

type API interface {
	Analyze(image *docker.Image) ([]*Vulnerability, error)
	Push(image *docker.Image) error
}

type layer struct {
	Name       string
	Path       string
	ParentName string
	Format     string
	Features   []feature
	Headers    headers
}

type headers struct {
	Authorization string
}

type feature struct {
	Name            string          `json:"Name,omitempty"`
	NamespaceName   string          `json:"NamespaceName,omitempty"`
	Version         string          `json:"Version,omitempty"`
	Vulnerabilities []Vulnerability `json:"Vulnerabilities"`
	AddedBy         string          `json:"AddedBy,omitempty"`
}

// Vulnerability represents vulnerability entity returned by Clair
type Vulnerability struct {
	Name           string                 `json:"Name,omitempty"`
	NamespaceName  string                 `json:"NamespaceName,omitempty"`
	Description    string                 `json:"Description,omitempty"`
	Link           string                 `json:"Link,omitempty"`
	Severity       string                 `json:"Severity,omitempty"`
	Metadata       map[string]interface{} `json:"Metadata,omitempty"`
	FixedBy        string                 `json:"FixedBy,omitempty"`
	FixedIn        []feature              `json:"FixedIn,omitempty"`
	FeatureName    string                 `json:"featureName",omitempty`
	FeatureVersion string                 `json:"featureName",omitempty`
}

type layerError struct {
	Message string
}

type clairError struct {
	Message string `json:"Layer"`
}

type layerEnvelope struct {
	Layer *layer      `json:"Layer,omitempty"`
	Error *clairError `json:"Error,omitempty"`
}

// NewClair construct Clair entity using potentially incomplete server URL
// If protocol is missing HTTP will be used. If port is missing 6060 will be used
func NewClair(url string, version int, timeout time.Duration) Clair {
	api, err := newAPI(url, version, timeout)
	if err != nil {
		panic(fmt.Sprintf("cant't create API client version %d %s: %s", version, url, err))
	}

	return Clair{url, api}
}

func newLayer(image *docker.Image, index int) *layer {
	var parentName string
	if index != 0 {
		parentName = image.LayerName(index - 1)
	}

	return &layer{
		Name:       image.LayerName(index),
		Path:       strings.Join([]string{image.Registry, image.Name, "blobs", image.FsLayers[index].BlobSum}, "/"),
		ParentName: parentName,
		Format:     "Docker",
		Headers:    headers{image.Token},
	}
}

func filterEmptyLayers(fsLayers []docker.FsLayer) (filteredLayers []docker.FsLayer) {
	for _, layer := range fsLayers {
		if layer.BlobSum != EMPTY_LAYER_BLOB_SUM {
			filteredLayers = append(filteredLayers, layer)
		}
	}
	return
}

// Analyse sent each layer from Docker image to Clair and returns
// a list of found vulnerabilities
func (c *Clair) Analyse(image *docker.Image) ([]*Vulnerability, error) {
	// Filter the empty layers in image
	image.FsLayers = filterEmptyLayers(image.FsLayers)
	layerLength := len(image.FsLayers)
	if layerLength == 0 {
		fmt.Fprintf(os.Stderr, "no need to analyse image %s/%s:%s as there is no non-emtpy layer\n",
			image.Registry, image.Name, image.Tag)
		return nil, nil
	}

	if err := c.api.Push(image); err != nil {
		return nil, fmt.Errorf("push image %s/%s:%s to Clair failed: %s\n", image.Registry, image.Name, image.Tag, err.Error())
	}
	vs, err := c.api.Analyze(image)
	if err != nil {
		return nil, fmt.Errorf("analyse image %s/%s:%s failed: %s\n", image.Registry, image.Name, image.Tag, err.Error())
	}

	return vs, nil
}
<commit_msg>Fix feature version json tag<commit_after>package clair

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/optiopay/klar/docker"
)

const EMPTY_LAYER_BLOB_SUM = "sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4"

// Clair is representation of Clair server
type Clair struct {
	url string
	api API
}

type API interface {
	Analyze(image *docker.Image) ([]*Vulnerability, error)
	Push(image *docker.Image) error
}

type layer struct {
	Name       string
	Path       string
	ParentName string
	Format     string
	Features   []feature
	Headers    headers
}

type headers struct {
	Authorization string
}

type feature struct {
	Name            string          `json:"Name,omitempty"`
	NamespaceName   string          `json:"NamespaceName,omitempty"`
	Version         string          `json:"Version,omitempty"`
	Vulnerabilities []Vulnerability `json:"Vulnerabilities"`
	AddedBy         string          `json:"AddedBy,omitempty"`
}

// Vulnerability represents vulnerability entity returned by Clair
type Vulnerability struct {
	Name           string                 `json:"Name,omitempty"`
	NamespaceName  string                 `json:"NamespaceName,omitempty"`
	Description    string                 `json:"Description,omitempty"`
	Link           string                 `json:"Link,omitempty"`
	Severity       string                 `json:"Severity,omitempty"`
	Metadata       map[string]interface{} `json:"Metadata,omitempty"`
	FixedBy        string                 `json:"FixedBy,omitempty"`
	FixedIn        []feature              `json:"FixedIn,omitempty"`
	FeatureName    string                 `json:"FeatureName",omitempty`
	FeatureVersion string                 `json:"FeatureVersion",omitempty`
}

type layerError struct {
	Message string
}

type clairError struct {
	Message string `json:"Layer"`
}

type layerEnvelope struct {
	Layer *layer      `json:"Layer,omitempty"`
	Error *clairError `json:"Error,omitempty"`
}

// NewClair construct Clair entity using potentially incomplete server URL
// If protocol is missing HTTP will be used. If port is missing 6060 will be used
func NewClair(url string, version int, timeout time.Duration) Clair {
	api, err := newAPI(url, version, timeout)
	if err != nil {
		panic(fmt.Sprintf("cant't create API client version %d %s: %s", version, url, err))
	}

	return Clair{url, api}
}

func newLayer(image *docker.Image, index int) *layer {
	var parentName string
	if index != 0 {
		parentName = image.LayerName(index - 1)
	}

	return &layer{
		Name:       image.LayerName(index),
		Path:       strings.Join([]string{image.Registry, image.Name, "blobs", image.FsLayers[index].BlobSum}, "/"),
		ParentName: parentName,
		Format:     "Docker",
		Headers:    headers{image.Token},
	}
}

func filterEmptyLayers(fsLayers []docker.FsLayer) (filteredLayers []docker.FsLayer) {
	for _, layer := range fsLayers {
		if layer.BlobSum != EMPTY_LAYER_BLOB_SUM {
			filteredLayers = append(filteredLayers, layer)
		}
	}
	return
}

// Analyse sent each layer from Docker image to Clair and returns
// a list of found vulnerabilities
func (c *Clair) Analyse(image *docker.Image) ([]*Vulnerability, error) {
	// Filter the empty layers in image
	image.FsLayers = filterEmptyLayers(image.FsLayers)
	layerLength := len(image.FsLayers)
	if layerLength == 0 {
		fmt.Fprintf(os.Stderr, "no need to analyse image %s/%s:%s as there is no non-emtpy layer\n",
			image.Registry, image.Name, image.Tag)
		return nil, nil
	}

	if err := c.api.Push(image); err != nil {
		return nil, fmt.Errorf("push image %s/%s:%s to Clair failed: %s\n", image.Registry, image.Name, image.Tag, err.Error())
	}
	vs, err := c.api.Analyze(image)
	if err != nil {
		return nil, fmt.Errorf("analyse image %s/%s:%s failed: %s\n", image.Registry, image.Name, image.Tag, err.Error())
	}

	return vs, nil
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"fmt"
	"bytes"
	"io"
	"io/ioutil"
	"bufio"
	"net/http"
	//"net/http/httputil"
	"encoding/json"
	"strings"
	"path/filepath"

	"crypto/aes"
	"crypto/rand"
	"crypto/cipher"
	"encoding/hex"

	"github.com/waucka/secretshare/commonlib"
	"github.com/codegangsta/cli"
)

type clientConfig struct {
	EndpointBaseURL string `json:"endpointBaseUrl"`
	Bucket string `json:"bucket"`
}

var config clientConfig

func loadConfig(configPath string) error {
	configFile, err := os.Open(configPath)
	if os.IsNotExist(err) {
		// No file; use defaults.
		config.EndpointBaseURL = commonlib.EndpointBaseURL
		config.Bucket = commonlib.Bucket
		return nil
	}
	if err != nil {
		return err
	}
	configData, err := ioutil.ReadAll(configFile)
	if err != nil {
		return err
	}
	err = json.Unmarshal(configData, &config)
	if err != nil {
		return err
	}
	return nil
}

func uploadEncrypted(stream io.Reader, messageSize int64, putURL string, headers http.Header, key []byte) {
	encrypter, err := commonlib.NewEncrypter(stream, messageSize, key)
	if err != nil {
		fmt.Printf("Can't encrypt: %s\n", err.Error())
		os.Exit(1)
	}

	commonlib.DEBUGPrintf("Starting upload to %s\n", putURL)
	uploadClient := &http.Client{}
	req, err := http.NewRequest("PUT", putURL, bufio.NewReaderSize(encrypter, 4096))
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	headerStrings := make([]string, 0)
	for k, v := range headers {
		canonicalKey := http.CanonicalHeaderKey(k)
		if len(v) == 1 {
			commonlib.DEBUGPrintf("Adding header %s (%s): %s\n", canonicalKey, k, v)
			req.Header.Set(canonicalKey, v[0])
			headerStrings = append(headerStrings, fmt.Sprintf(`-H "%s: %s"`, canonicalKey, v[0]))
		} else {
			items, ok := req.Header[canonicalKey]
			if ok {
				for _, item := range v {
					commonlib.DEBUGPrintf("Appending %s to header %s (%s)\n", item, canonicalKey, k)
					items = append(items, item)
				}
				req.Header[canonicalKey] = items
			} else {
				commonlib.DEBUGPrintf("Adding header %s (%s): %s\n", canonicalKey, k, v[0])
				req.Header[canonicalKey] = v
			}
		}
	}

	commonlib.DEBUGPrintln("All custom headers set!")

	// Set Content-Length header to avoid HTTP 501 from S3.
	// Don't bother setting it in headerStrings; curl does this on its own.
	req.ContentLength = encrypter.TotalSize
	commonlib.DEBUGPrintln("Content-Length set!")

	/*dump, err := httputil.DumpRequestOut(req, false)
	if err == nil {
		fmt.Println("Request:")
		fmt.Printf("%q", dump)
		fmt.Println()
	} else {
		fmt.Println("Error dumping request!")
		fmt.Println(err.Error())
		os.Exit(1)
	}*/

	commonlib.DEBUGPrintf("Uploading %d bytes...\n", req.ContentLength)
	resp, err := uploadClient.Do(req)
	if err != nil {
		fmt.Println("Error uploading file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	if resp.StatusCode != http.StatusOK {
		commonlib.DEBUGPrintf("Failed to upload file!  Status code: %d\n", resp.StatusCode)
		commonlib.DEBUGPrintf(`curl -XPUT -d @$FILENAME %s '%s'`, strings.Join(headerStrings, " "), putURL)
		fmt.Println()
		os.Exit(1)
	}
}

func sendSecret(c *cli.Context) {
	config.EndpointBaseURL = c.Parent().String("endpoint")
	config.Bucket = c.Parent().String("bucket")
	filename := c.Args()[0]
	stats, err := os.Stat(filename)
	if err != nil {
		fmt.Printf("Can't read file %s: %s\n", filename, err.Error())
		os.Exit(1)
	}
	fileSize := stats.Size()
	basename := filepath.Base(filename)
	requestBytes, err := json.Marshal(&commonlib.UploadRequest{
		TTL: c.Int("ttl"),
	})
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	buf := bytes.NewBuffer(requestBytes)

	key := make([]byte, 32)
	num_key_bytes, err := rand.Read(key)
	if err != nil {
		fmt.Println("Failed to generate key!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	if num_key_bytes < 32 {
		fmt.Println("Failed to generate key!")
		fmt.Println(commonlib.NotEnoughKeyRandomnessError.Error())
		os.Exit(1)
	}
	keystr := hex.EncodeToString(key)

	resp, err := http.Post(config.EndpointBaseURL + "/upload", "application/json", buf)
	if err != nil {
		fmt.Println("Failed to connect to server!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	if resp.Body == nil {
		fmt.Println("No data received from server!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer resp.Body.Close()
	bodyBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("Malformed response received from server!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	var responseData commonlib.UploadResponse
	err = json.Unmarshal(bodyBytes, &responseData)
	if err != nil {
		fmt.Println("Malformed response received from server!")
		fmt.Println(err.Error())
		fmt.Println(string(bodyBytes))
		os.Exit(1)
	}

	f, err := os.Open(filename)
	if err != nil {
		fmt.Printf("Can't read file %s: %s\n", filename, err.Error())
		os.Exit(1)
	}
	defer f.Close()
	stream := bufio.NewReader(f)
	uploadEncrypted(stream, fileSize, responseData.PutURL, responseData.Headers, key)

	filemeta := commonlib.FileMetadata{
		Filename: basename,
		Filesize: fileSize,
	}
	metabytes, err := json.Marshal(filemeta)
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	metabuf := bytes.NewBuffer(metabytes)
	uploadEncrypted(metabuf, int64(len(metabytes)), responseData.MetaPutURL, responseData.MetaHeaders, key)

	fmt.Println("File uploaded!")
	fmt.Printf("Key: %s\nID: %s\nURL: https://s3-us-west-1.amazonaws.com/exosite-secretshare/%s\n",
		keystr, responseData.Id, responseData.Id)
	fmt.Println("To receive this secret:")
	fmt.Printf("secretshare receive %s %s\n", responseData.Id, keystr)
}

func decrypt(ciphertext, key []byte) []byte {
	paddingLen := ciphertext[0]
	commonlib.DEBUGPrintf("decrypt: paddingLen = %d\n", paddingLen)
	commonlib.DEBUGPrintf("decrypt: len(ciphertext) = %d\n", len(ciphertext))
	iv := ciphertext[1:aes.BlockSize + 1]
	raw := ciphertext[1 + aes.BlockSize:len(ciphertext)]
	commonlib.DEBUGPrintf("decrypt: len(raw) = %d\n", len(raw))

	if len(raw) % aes.BlockSize != 0 {
		fmt.Println("Data is malformed!")
		os.Exit(1)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	decrypter := cipher.NewCBCDecrypter(block, iv)
	decrypter.CryptBlocks(raw, raw)
	// Discard padding
	return raw[:len(raw) - int(paddingLen)]
}

func recvSecret(c *cli.Context) {
	config.EndpointBaseURL = c.Parent().String("endpoint")
	config.Bucket = c.Parent().String("bucket")
	id := c.Args()[0]
	keystr := c.Args()[1]
	key, err := hex.DecodeString(keystr)
	if err != nil {
		fmt.Println("Malformed key!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	// Download metadata
	resp, err := http.Get(fmt.Sprintf("https://s3-us-west-1.amazonaws.com/exosite-secretshare/meta/%s", id))
	if err != nil {
		fmt.Println("Failed to download file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer resp.Body.Close()
	metabytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("Failed to download metadata!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	realMeta := decrypt(metabytes, key)

	var filemeta commonlib.FileMetadata
	err = json.Unmarshal(realMeta, &filemeta)
	if err != nil {
		fmt.Println("Malformed metadata!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	// This is how you check if a file exists in Go.  Yep.
	if _, err := os.Stat(filemeta.Filename); err == nil {
		inreader := bufio.NewReader(os.Stdin)
		fmt.Printf("File %s already exists!  Overwrite (y/n)? ", filemeta.Filename)
		answer, err := inreader.ReadString('\n')
		if err != nil {
			os.Exit(1)
		}

		if answer == "y\n" || answer == "Y\n" {
			os.Remove(filemeta.Filename)
		} else {
			fmt.Printf("Download cancelled.\n")
			os.Exit(1)
		}
	}
	outf, err := os.Create(filemeta.Filename)
	if err != nil {
		fmt.Printf("Can't create file %s: %s\n", filemeta.Filename, err.Error())
		os.Exit(1)
	}
	defer outf.Close()

	// Download data
	resp, err = http.Get(fmt.Sprintf("https://s3-us-west-1.amazonaws.com/exosite-secretshare/%s", id))
	if err != nil {
		fmt.Println("Failed to download file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer resp.Body.Close()
	decrypter, err := commonlib.NewDecrypter(resp.Body, filemeta.Filesize, key)
	if err != nil {
		fmt.Println("Failed to set up decryption!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	bytesWritten, err := io.Copy(outf, decrypter)
	commonlib.DEBUGPrintf("Wrote %d bytes\n", bytesWritten)
	if err != nil {
		fmt.Println("Failed to save file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	fmt.Printf("File downloaded as %s\n", filemeta.Filename)
}

func main() {
	err := loadConfig("~/.secretsharerc")
	if err != nil {
		fmt.Println("Failed to load configuration")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	app := cli.NewApp()
	app.Name = "secretshare"
	app.Usage = "Securely share secrets"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name: "endpoint",
			Value: config.EndpointBaseURL,
			Usage: "API endpoint to connect to when requesting IDs",
		},
		cli.StringFlag{
			Name: "bucket",
			Value: config.Bucket,
			Usage: "S3 bucket to store files in",
		},
	}
	app.Commands = []cli.Command{
		{
			Name: "send",
			Usage: "Send a secret file",
			Action: sendSecret,
			Flags: []cli.Flag{
				cli.IntFlag{
					Name: "ttl",
					Value: 4 * 60,
					Usage: "Time in minutes that the file should be available (doesn't work yet)",
				},
			},
		},
		{
			Name: "receive",
			Usage: "Receive a secret file",
			Action: recvSecret,
		},
	}
	app.Run(os.Args)
}
<commit_msg>Clean up endpoint base URLs<commit_after>package main

import (
	"os"
	"fmt"
	"bytes"
	"io"
	"io/ioutil"
	"bufio"
	"net/http"
	//"net/http/httputil"
	"encoding/json"
	"strings"
	"path/filepath"

	"crypto/aes"
	"crypto/rand"
	"crypto/cipher"
	"encoding/hex"

	"github.com/waucka/secretshare/commonlib"
	"github.com/codegangsta/cli"
)

type clientConfig struct {
	EndpointBaseURL string `json:"endpointBaseUrl"`
	Bucket string `json:"bucket"`
}

var config clientConfig

func loadConfig(configPath string) error {
	configFile, err := os.Open(configPath)
	if os.IsNotExist(err) {
		// No file; use defaults.
		config.EndpointBaseURL = commonlib.EndpointBaseURL
		config.Bucket = commonlib.Bucket
		return nil
	}
	if err != nil {
		return err
	}
	configData, err := ioutil.ReadAll(configFile)
	if err != nil {
		return err
	}
	err = json.Unmarshal(configData, &config)
	if err != nil {
		return err
	}
	return nil
}

func cleanUrl(url string) string {
	if strings.HasSuffix(url, "/") {
		return url[:len(url) - 1]
	}
	return url
}

func uploadEncrypted(stream io.Reader, messageSize int64, putURL string, headers http.Header, key []byte) {
	encrypter, err := commonlib.NewEncrypter(stream, messageSize, key)
	if err != nil {
		fmt.Printf("Can't encrypt: %s\n", err.Error())
		os.Exit(1)
	}

	commonlib.DEBUGPrintf("Starting upload to %s\n", putURL)
	uploadClient := &http.Client{}
	req, err := http.NewRequest("PUT", putURL, bufio.NewReaderSize(encrypter, 4096))
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	headerStrings := make([]string, 0)
	for k, v := range headers {
		canonicalKey := http.CanonicalHeaderKey(k)
		if len(v) == 1 {
			commonlib.DEBUGPrintf("Adding header %s (%s): %s\n", canonicalKey, k, v)
			req.Header.Set(canonicalKey, v[0])
			headerStrings = append(headerStrings, fmt.Sprintf(`-H "%s: %s"`, canonicalKey, v[0]))
		} else {
			items, ok := req.Header[canonicalKey]
			if ok {
				for _, item := range v {
					commonlib.DEBUGPrintf("Appending %s to header %s (%s)\n", item, canonicalKey, k)
					items = append(items, item)
				}
				req.Header[canonicalKey] = items
			} else {
				commonlib.DEBUGPrintf("Adding header %s (%s): %s\n", canonicalKey, k, v[0])
				req.Header[canonicalKey] = v
			}
		}
	}

	commonlib.DEBUGPrintln("All custom headers set!")

	// Set Content-Length header to avoid HTTP 501 from S3.
	// Don't bother setting it in headerStrings; curl does this on its own.
	req.ContentLength = encrypter.TotalSize
	commonlib.DEBUGPrintln("Content-Length set!")

	/*dump, err := httputil.DumpRequestOut(req, false)
	if err == nil {
		fmt.Println("Request:")
		fmt.Printf("%q", dump)
		fmt.Println()
	} else {
		fmt.Println("Error dumping request!")
		fmt.Println(err.Error())
		os.Exit(1)
	}*/

	commonlib.DEBUGPrintf("Uploading %d bytes...\n", req.ContentLength)
	resp, err := uploadClient.Do(req)
	if err != nil {
		fmt.Println("Error uploading file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	if resp.StatusCode != http.StatusOK {
		commonlib.DEBUGPrintf("Failed to upload file!  Status code: %d\n", resp.StatusCode)
		commonlib.DEBUGPrintf(`curl -XPUT -d @$FILENAME %s '%s'`, strings.Join(headerStrings, " "), putURL)
		fmt.Println()
		os.Exit(1)
	}
}

func sendSecret(c *cli.Context) {
	config.EndpointBaseURL = cleanUrl(c.Parent().String("endpoint"))
	config.Bucket = c.Parent().String("bucket")
	filename := c.Args()[0]
	stats, err := os.Stat(filename)
	if err != nil {
		fmt.Printf("Can't read file %s: %s\n", filename, err.Error())
		os.Exit(1)
	}
	fileSize := stats.Size()
	basename := filepath.Base(filename)
	requestBytes, err := json.Marshal(&commonlib.UploadRequest{
		TTL: c.Int("ttl"),
	})
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	buf := bytes.NewBuffer(requestBytes)

	key := make([]byte, 32)
	num_key_bytes, err := rand.Read(key)
	if err != nil {
		fmt.Println("Failed to generate key!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	if num_key_bytes < 32 {
		fmt.Println("Failed to generate key!")
		fmt.Println(commonlib.NotEnoughKeyRandomnessError.Error())
		os.Exit(1)
	}
	keystr := hex.EncodeToString(key)

	resp, err := http.Post(config.EndpointBaseURL + "/upload", "application/json", buf)
	if err != nil {
		fmt.Println("Failed to connect to server!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	if resp.Body == nil {
		fmt.Println("No data received from server!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer resp.Body.Close()
	bodyBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("Malformed response received from server!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	var responseData commonlib.UploadResponse
	err = json.Unmarshal(bodyBytes, &responseData)
	if err != nil {
		fmt.Println("Malformed response received from server!")
		fmt.Println(err.Error())
		fmt.Println(string(bodyBytes))
		os.Exit(1)
	}

	f, err := os.Open(filename)
	if err != nil {
		fmt.Printf("Can't read file %s: %s\n", filename, err.Error())
		os.Exit(1)
	}
	defer f.Close()
	stream := bufio.NewReader(f)
	uploadEncrypted(stream, fileSize, responseData.PutURL, responseData.Headers, key)

	filemeta := commonlib.FileMetadata{
		Filename: basename,
		Filesize: fileSize,
	}
	metabytes, err := json.Marshal(filemeta)
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	metabuf := bytes.NewBuffer(metabytes)
	uploadEncrypted(metabuf, int64(len(metabytes)), responseData.MetaPutURL, responseData.MetaHeaders, key)

	fmt.Println("File uploaded!")
	fmt.Printf("Key: %s\nID: %s\nURL: https://s3-us-west-1.amazonaws.com/exosite-secretshare/%s\n",
		keystr, responseData.Id, responseData.Id)
	fmt.Println("To receive this secret:")
	fmt.Printf("secretshare receive %s %s\n", responseData.Id, keystr)
}

func decrypt(ciphertext, key []byte) []byte {
	paddingLen := ciphertext[0]
	commonlib.DEBUGPrintf("decrypt: paddingLen = %d\n", paddingLen)
	commonlib.DEBUGPrintf("decrypt: len(ciphertext) = %d\n", len(ciphertext))
	iv := ciphertext[1:aes.BlockSize + 1]
	raw := ciphertext[1 + aes.BlockSize:len(ciphertext)]
	commonlib.DEBUGPrintf("decrypt: len(raw) = %d\n", len(raw))

	if len(raw) % aes.BlockSize != 0 {
		fmt.Println("Data is malformed!")
		os.Exit(1)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println("Internal error!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	decrypter := cipher.NewCBCDecrypter(block, iv)
	decrypter.CryptBlocks(raw, raw)
	// Discard padding
	return raw[:len(raw) - int(paddingLen)]
}

func recvSecret(c *cli.Context) {
	config.EndpointBaseURL = cleanUrl(c.Parent().String("endpoint"))
	config.Bucket = c.Parent().String("bucket")
	id := c.Args()[0]
	keystr := c.Args()[1]
	key, err := hex.DecodeString(keystr)
	if err != nil {
		fmt.Println("Malformed key!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	// Download metadata
	resp, err := http.Get(fmt.Sprintf("https://s3-us-west-1.amazonaws.com/exosite-secretshare/meta/%s", id))
	if err != nil {
		fmt.Println("Failed to download file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer resp.Body.Close()
	metabytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("Failed to download metadata!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	realMeta := decrypt(metabytes, key)

	var filemeta commonlib.FileMetadata
	err = json.Unmarshal(realMeta, &filemeta)
	if err != nil {
		fmt.Println("Malformed metadata!")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	// This is how you check if a file exists in Go.  Yep.
	if _, err := os.Stat(filemeta.Filename); err == nil {
		inreader := bufio.NewReader(os.Stdin)
		fmt.Printf("File %s already exists!  Overwrite (y/n)? ", filemeta.Filename)
		answer, err := inreader.ReadString('\n')
		if err != nil {
			os.Exit(1)
		}

		if answer == "y\n" || answer == "Y\n" {
			os.Remove(filemeta.Filename)
		} else {
			fmt.Printf("Download cancelled.\n")
			os.Exit(1)
		}
	}
	outf, err := os.Create(filemeta.Filename)
	if err != nil {
		fmt.Printf("Can't create file %s: %s\n", filemeta.Filename, err.Error())
		os.Exit(1)
	}
	defer outf.Close()

	// Download data
	resp, err = http.Get(fmt.Sprintf("https://s3-us-west-1.amazonaws.com/exosite-secretshare/%s", id))
	if err != nil {
		fmt.Println("Failed to download file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer resp.Body.Close()
	decrypter, err := commonlib.NewDecrypter(resp.Body, filemeta.Filesize, key)
	if err != nil {
		fmt.Println("Failed to set up decryption!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	bytesWritten, err := io.Copy(outf, decrypter)
	commonlib.DEBUGPrintf("Wrote %d bytes\n", bytesWritten)
	if err != nil {
		fmt.Println("Failed to save file!")
		fmt.Println(err.Error())
		os.Exit(1)
	}
	fmt.Printf("File downloaded as %s\n", filemeta.Filename)
}

func main() {
	err := loadConfig("~/.secretsharerc")
	if err != nil {
		fmt.Println("Failed to load configuration")
		fmt.Println(err.Error())
		os.Exit(1)
	}

	app := cli.NewApp()
	app.Name = "secretshare"
	app.Usage = "Securely share secrets"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name: "endpoint",
			Value: config.EndpointBaseURL,
			Usage: "API endpoint to connect to when requesting IDs",
		},
		cli.StringFlag{
			Name: "bucket",
			Value: config.Bucket,
			Usage: "S3 bucket to store files in",
		},
	}
	app.Commands = []cli.Command{
		{
			Name: "send",
			Usage: "Send a secret file",
			Action: sendSecret,
			Flags: []cli.Flag{
				cli.IntFlag{
					Name: "ttl",
					Value: 4 * 60,
					Usage: "Time in minutes that the file should be available (doesn't work yet)",
				},
			},
		},
		{
			Name: "receive",
			Usage: "Receive a secret file",
			Action: recvSecret,
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/richo/roving/types"
)

type Server struct {
	hostport string
}

type Fuzzer struct {
	cmd *exec.Cmd
	Id  string
}

func newFuzzer() Fuzzer {
	name, err := os.Hostname()
	if err != nil {
		log.Fatal("Couldn't get hostname", err)
	}

	number := types.RandInt() & 0xffff

	return Fuzzer{
		Id: fmt.Sprintf("%s-%x", name, number),
	}
}

type WatchDog struct {
	Interval time.Duration
	Fuzzer   *Fuzzer
	Server   *Server
}

func (f *Fuzzer) run() {
	f.cmd = exec.Command(f.path(),
		"-o", "output",
		"-i", "input",
		"-S", f.Id,
		"./target",
	)
	stdout, err := f.cmd.StdoutPipe()
	if err != nil {
		log.Fatalf("Couldn't get stdout handle", err)
	}

	stderr, err := f.cmd.StderrPipe()
	if err != nil {
		log.Fatalf("Couldn't get stderr handle", err)
	}

	if err := f.cmd.Start(); err != nil {
		log.Fatalf("Couldn't start fuzzer", err)
	}

	go func() {
		io.Copy(os.Stdout, stdout)
	}()

	go func() {
		io.Copy(os.Stderr, stderr)
	}()

	log.Printf("Waiting for fuzzer to exit")

	err = f.cmd.Wait()
	if err != nil {
		log.Fatal(err)
	}

}

func (f *Fuzzer) stop() {
	log.Printf("Stopping the fuzzer")
	f.cmd.Process.Signal(syscall.SIGSTOP)
}

func (f *Fuzzer) start() {
	log.Printf("Starting the fuzzer")
	f.cmd.Process.Signal(syscall.SIGCONT)
}

func (f *Fuzzer) State() types.State {
	state := types.State{
		Id:    f.Id,
		Queue: types.ReadCorpus("output/queue"),
	}

	return state
}

func (f *Fuzzer) path() string {
	root := os.Getenv("AFL")
	if root == "" { // Not found, hopefully it's in PATH
		return "afl-fuzz"
	}
	return fmt.Sprintf("%s/afl-fuzz", root)
}

func (w *WatchDog) run() {
	log.Printf("Starting watchdog goroutine")
	ticker := time.NewTicker(w.Interval)
	for {
		select {
		case <-ticker.C:
			w.Fuzzer.stop()

			log.Printf("Uploading our corpus")
			state := w.Fuzzer.State()
			w.Server.UploadState(state)

			log.Printf("Downloading their corpus")
			other := w.Server.FetchState(w.Fuzzer.Id)

			log.Printf("Unpacking their state")
			w.Fuzzer.UnpackStates(other)

			w.Fuzzer.start()
		}
	}
}

func (s *Server) fetchToFile(resource, file string) {
	target := s.getPath(resource)
	defer target.Body.Close()

	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		log.Panicf("Couldn't open %s for writing", file, err)
	}

	io.Copy(f, target.Body)
}

func (s *Server) getPath(path string) *http.Response {
	resource := fmt.Sprintf("http://%s/%s", s.hostport, path)
	log.Printf("Fetching %s", resource)

	resp, err := http.Get(resource)
	if err != nil {
		log.Panicf("Couldn't fetch target", err)
	}

	return resp
}

func (s *Server) FetchTarget() {
	s.fetchToFile("target", "target")
}

func base64ToPath(content, path string) {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		log.Panicf("Couldn't open %s for writing", path, err)
	}

	body, err := base64.StdEncoding.DecodeString(content)
	if err != nil {
		log.Fatalf("Couldn't decode body for %s", path, err)
	}
	f.Write([]byte(body))
}

func (s *Server) FetchInputs() {
	os.Mkdir("input", 0755)

	inputs := s.getPath("inputs")
	defer inputs.Body.Close()

	inps := &types.InputCorpus{}

	encoder := json.NewDecoder(inputs.Body)
	encoder.Decode(&inps)

	for _, inp := range inps.Inputs {
		path := fmt.Sprintf("input/%s", inp.Name)
		base64ToPath(inp.Body, path)
	}
}

func (s *Fuzzer) UnpackStates(other []types.State) {
	for _, state := range other {
		log.Printf("Unpacking state from %s", state.Id)
		for _, input := range state.Queue.Inputs {
			base64ToPath(input.Body, fmt.Sprintf("output/queue/%s", input.Name))
		}
	}
}

func (s *Server) FetchState(id string) []types.State {
	var state []types.State
	resp := s.getPath(fmt.Sprintf("state/%s", id))

	encoder := json.NewDecoder(resp.Body)
	encoder.Decode(&state)

	return state
}

func (s *Server) UploadState(state types.State) {
	data, err := json.Marshal(state)
	if err != nil {
		log.Fatal("Couldn't marshall state", err)
	}

	buffer := bytes.NewReader(data)

	resource := fmt.Sprintf("http://%s/%s", s.hostport, "state")
	_, err = http.Post(resource, "application/json", buffer)

	if err != nil {
		log.Panicf("Couldn't upload state", err)
	}
}

func setupWorkDir() {
	var err error
	// TODO(richo) Ephemeral workdirs for concurrency
	if err = os.Mkdir("work", 0755); err != nil {
		log.Panicf("Couldn't make workdir", err)
	}
	if err = os.Chdir("work"); err != nil {
		log.Panicf("Couldn't change to workdir", err)
	}
}

func fetchCorpus(server string) {
}

func main() {
	args := os.Args
	if len(args) != 2 {
		log.Printf("Usage: ./client server:port")
		return
	}

	setupWorkDir()

	server := Server{args[1]}
	server.FetchTarget()
	server.FetchInputs()

	fuzzer := newFuzzer()

	log.Printf("Brought up a fuzzer with id %s", fuzzer.Id)

	watchdog := WatchDog{
		Interval: 5 * time.Minute,
		Fuzzer:   &fuzzer,
		Server:   &server,
	}
	go watchdog.run()

	fuzzer.run()

	// TODO(richo) Fetch the input corpus
}
<commit_msg>Send the state from the slave location<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/richo/roving/types"
)

type Server struct {
	hostport string
}

type Fuzzer struct {
	cmd *exec.Cmd
	Id  string
}

func newFuzzer() Fuzzer {
	name, err := os.Hostname()
	if err != nil {
		log.Fatal("Couldn't get hostname", err)
	}

	number := types.RandInt() & 0xffff

	return Fuzzer{
		Id: fmt.Sprintf("%s-%x", name, number),
	}
}

type WatchDog struct {
	Interval time.Duration
	Fuzzer   *Fuzzer
	Server   *Server
}

func (f *Fuzzer) run() {
	f.cmd = exec.Command(f.path(),
		"-o", "output",
		"-i", "input",
		"-S", f.Id,
		"./target",
	)
	stdout, err := f.cmd.StdoutPipe()
	if err != nil {
		log.Fatalf("Couldn't get stdout handle", err)
	}

	stderr, err := f.cmd.StderrPipe()
	if err != nil {
		log.Fatalf("Couldn't get stderr handle", err)
	}

	if err := f.cmd.Start(); err != nil {
		log.Fatalf("Couldn't start fuzzer", err)
	}

	go func() {
		io.Copy(os.Stdout, stdout)
	}()

	go func() {
		io.Copy(os.Stderr, stderr)
	}()

	log.Printf("Waiting for fuzzer to exit")

	err = f.cmd.Wait()
	if err != nil {
		log.Fatal(err)
	}

}

func (f *Fuzzer) stop() {
	log.Printf("Stopping the fuzzer")
	f.cmd.Process.Signal(syscall.SIGSTOP)
}

func (f *Fuzzer) start() {
	log.Printf("Starting the fuzzer")
	f.cmd.Process.Signal(syscall.SIGCONT)
}

func (f *Fuzzer) State() types.State {
	state := types.State{
		Id:    f.Id,
		Queue: types.ReadCorpus(fmt.Sprintf("output/%s/queue", f.Id)),
	}

	return state
}

func (f *Fuzzer) path() string {
	root := os.Getenv("AFL")
	if root == "" { // Not found, hopefully it's in PATH
		return "afl-fuzz"
	}
	return fmt.Sprintf("%s/afl-fuzz", root)
}

func (w *WatchDog) run() {
	log.Printf("Starting watchdog goroutine")
	ticker := time.NewTicker(w.Interval)
	for {
		select {
		case <-ticker.C:
			w.Fuzzer.stop()

			log.Printf("Uploading our corpus")
			state := w.Fuzzer.State()
			w.Server.UploadState(state)

			log.Printf("Downloading their corpus")
			other := w.Server.FetchState(w.Fuzzer.Id)

			log.Printf("Unpacking their state")
			w.Fuzzer.UnpackStates(other)

			w.Fuzzer.start()
		}
	}
}

func (s *Server) fetchToFile(resource, file string) {
	target := s.getPath(resource)
	defer target.Body.Close()

	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		log.Panicf("Couldn't open %s for writing", file, err)
	}

	io.Copy(f, target.Body)
}

func (s *Server) getPath(path string) *http.Response {
	resource := fmt.Sprintf("http://%s/%s", s.hostport, path)
	log.Printf("Fetching %s", resource)

	resp, err := http.Get(resource)
	if err != nil {
		log.Panicf("Couldn't fetch target", err)
	}

	return resp
}

func (s *Server) FetchTarget() {
	s.fetchToFile("target", "target")
}

func base64ToPath(content, path string) {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		log.Panicf("Couldn't open %s for writing", path, err)
	}

	body, err := base64.StdEncoding.DecodeString(content)
	if err != nil {
		log.Fatalf("Couldn't decode body for %s", path, err)
	}
	f.Write([]byte(body))
}

func (s *Server) FetchInputs() {
	os.Mkdir("input", 0755)

	inputs := s.getPath("inputs")
	defer inputs.Body.Close()

	inps := &types.InputCorpus{}

	encoder := json.NewDecoder(inputs.Body)
	encoder.Decode(&inps)

	for _, inp := range inps.Inputs {
		path := fmt.Sprintf("input/%s", inp.Name)
		base64ToPath(inp.Body, path)
	}
}

func (s *Fuzzer) UnpackStates(other []types.State) {
	for _, state := range other {
		log.Printf("Unpacking state from %s", state.Id)
		for _, input := range state.Queue.Inputs {
			base64ToPath(input.Body, fmt.Sprintf("output/%s/queue/%s", state.Id, input.Name))
		}
	}
}

func (s *Server) FetchState(id string) []types.State {
	var state []types.State
	resp := s.getPath(fmt.Sprintf("state/%s", id))

	encoder := json.NewDecoder(resp.Body)
	encoder.Decode(&state)

	return state
}

func (s *Server) UploadState(state types.State) {
	data, err := json.Marshal(state)
	if err != nil {
		log.Fatal("Couldn't marshall state", err)
	}

	buffer := bytes.NewReader(data)

	resource := fmt.Sprintf("http://%s/%s", s.hostport, "state")
	_, err = http.Post(resource, "application/json", buffer)

	if err != nil {
		log.Panicf("Couldn't upload state", err)
	}
}

func setupWorkDir() {
	var err error
	// TODO(richo) Ephemeral workdirs for concurrency
	if err = os.Mkdir("work", 0755); err != nil {
		log.Panicf("Couldn't make workdir", err)
	}
	if err = os.Chdir("work"); err != nil {
		log.Panicf("Couldn't change to workdir", err)
	}
}

func fetchCorpus(server string) {
}

func main() {
	args := os.Args
	if len(args) != 2 {
		log.Printf("Usage: ./client server:port")
		return
	}

	setupWorkDir()

	server := Server{args[1]}
	server.FetchTarget()
	server.FetchInputs()

	fuzzer := newFuzzer()

	log.Printf("Brought up a fuzzer with id %s", fuzzer.Id)

	watchdog := WatchDog{
		Interval: 3 * time.Second,
		Fuzzer:   &fuzzer,
		Server:   &server,
	}
	go watchdog.run()

	fuzzer.run()

	// TODO(richo) Fetch the input corpus
}
<|endoftext|>


<commit_before>package commands

import (
	"os"

	"github.com/spf13/cobra"

	"github.com/tendermint/tendermint/types"
)

var resetAllCmd = &cobra.Command{
	Use:   "unsafe_reset_all",
	Short: "(unsafe) Remove all the data and WAL, reset this node's validator",
	Run:   resetAll,
}

var resetPrivValidatorCmd = &cobra.Command{
	Use:   "unsafe_reset_priv_validator",
	Short: "(unsafe) Reset this node's validator",
	Run:   resetPrivValidator,
}

func init() {
	RootCmd.AddCommand(resetAllCmd)
	RootCmd.AddCommand(resetPrivValidatorCmd)
}

// XXX: this is totally unsafe.
// it's only suitable for testnets.
func resetAll(cmd *cobra.Command, args []string) {
	resetPrivValidator(cmd, args)
	os.RemoveAll(config.GetString("db_dir"))
	os.Remove(config.GetString("cs_wal_file"))
}

// XXX: this is totally unsafe.
// it's only suitable for testnets.
func resetPrivValidator(cmd *cobra.Command, args []string) {
	// Get PrivValidator
	var privValidator *types.PrivValidator
	privValidatorFile := config.GetString("priv_validator_file")
	if _, err := os.Stat(privValidatorFile); err == nil {
		privValidator = types.LoadPrivValidator(privValidatorFile)
		privValidator.Reset()
		log.Notice("Reset PrivValidator", "file", privValidatorFile)
	} else {
		privValidator = types.GenPrivValidator()
		privValidator.SetFile(privValidatorFile)
		privValidator.Save()
		log.Notice("Generated PrivValidator", "file", privValidatorFile)
	}
}
<commit_msg>export ResetAll cmd<commit_after>package commands

import (
	"os"

	"github.com/spf13/cobra"

	"github.com/tendermint/tendermint/types"
)

var resetAllCmd = &cobra.Command{
	Use:   "unsafe_reset_all",
	Short: "(unsafe) Remove all the data and WAL, reset this node's validator",
	Run:   ResetAll,
}

var resetPrivValidatorCmd = &cobra.Command{
	Use:   "unsafe_reset_priv_validator",
	Short: "(unsafe) Reset this node's validator",
	Run:   resetPrivValidator,
}

func init() {
	RootCmd.AddCommand(resetAllCmd)
	RootCmd.AddCommand(resetPrivValidatorCmd)
}

// XXX: this is totally unsafe.
// it's only suitable for testnets.
func ResetAll(cmd *cobra.Command, args []string) {
	resetPrivValidator(cmd, args)
	os.RemoveAll(config.GetString("db_dir"))
	os.Remove(config.GetString("cs_wal_file"))
}

// XXX: this is totally unsafe.
// it's only suitable for testnets.
func resetPrivValidator(cmd *cobra.Command, args []string) {
	// Get PrivValidator
	var privValidator *types.PrivValidator
	privValidatorFile := config.GetString("priv_validator_file")
	if _, err := os.Stat(privValidatorFile); err == nil {
		privValidator = types.LoadPrivValidator(privValidatorFile)
		privValidator.Reset()
		log.Notice("Reset PrivValidator", "file", privValidatorFile)
	} else {
		privValidator = types.GenPrivValidator()
		privValidator.SetFile(privValidatorFile)
		privValidator.Save()
		log.Notice("Generated PrivValidator", "file", privValidatorFile)
	}
}
<|endoftext|>


<commit_before>package common

import (
	"io/ioutil"
	"os"
	"strings"
	"testing"

	. "github.com/onsi/gomega"
)

var (
	testAppName  = "test-app-1"
	testAppDir   = strings.Join([]string{"/home/dokku/", testAppName}, "")
	testEnvFile  = strings.Join([]string{testAppDir, "/ENV"}, "")
	testEnvLine  = "export testKey=TESTING"
	testAppName2 = "01-test-app-1"
	testAppDir2  = strings.Join([]string{"/home/dokku/", testAppName2}, "")
	testEnvFile2 = strings.Join([]string{testAppDir2, "/ENV"}, "")
	testEnvLine2 = "export testKey=TESTING"
)

func setupTestApp() (err error) {
	Expect(os.MkdirAll(testAppDir, 0644)).To(Succeed())
	b := []byte(testEnvLine + "\n")
	if err = ioutil.WriteFile(testEnvFile, b, 0644); err != nil {
		return
	}
	return
}

func setupTestApp2() (err error) {
	Expect(os.MkdirAll(testAppDir2, 0644)).To(Succeed())
	b := []byte(testEnvLine2 + "\n")
	if err = ioutil.WriteFile(testEnvFile2, b, 0644); err != nil {
		return
	}
	return
}

func teardownTestApp() {
	os.RemoveAll(testAppDir)
}

func teardownTestApp2() {
	os.RemoveAll(testAppDir2)
}

func TestCommonGetEnv(t *testing.T) {
	RegisterTestingT(t)
	Expect(MustGetEnv("DOKKU_ROOT")).To(Equal("/home/dokku"))
}

func TestCommonGetAppImageRepo(t *testing.T) {
	RegisterTestingT(t)
	Expect(GetAppImageRepo("testapp")).To(Equal("dokku/testapp"))
}

func TestCommonVerifyImageInvalid(t *testing.T) {
	RegisterTestingT(t)
	Expect(VerifyImage("testapp")).To(Equal(false))
}

func TestCommonVerifyAppNameInvalid(t *testing.T) {
	RegisterTestingT(t)
	err := VerifyAppName("1994testApp")
	Expect(err).To(HaveOccurred())
}

func TestCommonVerifyAppName(t *testing.T) {
	RegisterTestingT(t)
	Expect(setupTestApp()).To(Succeed())
	Expect(VerifyAppName(testAppName)).To(Succeed())
	teardownTestApp()

	RegisterTestingT(t)
	Expect(setupTestApp2()).To(Succeed())
	Expect(VerifyAppName(testAppName2)).To(Succeed())
	teardownTestApp2()
}

func TestCommonDokkuAppsError(t *testing.T) {
	RegisterTestingT(t)
	_, err := DokkuApps()
	Expect(err).To(HaveOccurred())
}

func TestCommonDokkuApps(t *testing.T) {
	RegisterTestingT(t)
	Expect(setupTestApp()).To(Succeed())
	apps, err := DokkuApps()
	Expect(err).NotTo(HaveOccurred())
	Expect(apps).To(HaveLen(1))
	Expect(apps[0]).To(Equal(testAppName))
	teardownTestApp()
}

func TestCommonStripInlineComments(t *testing.T) {
	RegisterTestingT(t)
	text := StripInlineComments(strings.Join([]string{testEnvLine, "# testing comment"}, " "))
	Expect(text).To(Equal(testEnvLine))
}
<commit_msg>test: set a dummy path for PLUGIN_ENABLED_PATH<commit_after>package common

import (
	"io/ioutil"
	"os"
	"strings"
	"testing"

	. "github.com/onsi/gomega"
)

var (
	testAppName  = "test-app-1"
	testAppDir   = strings.Join([]string{"/home/dokku/", testAppName}, "")
	testEnvFile  = strings.Join([]string{testAppDir, "/ENV"}, "")
	testEnvLine  = "export testKey=TESTING"
	testAppName2 = "01-test-app-1"
	testAppDir2  = strings.Join([]string{"/home/dokku/", testAppName2}, "")
	testEnvFile2 = strings.Join([]string{testAppDir2, "/ENV"}, "")
	testEnvLine2 = "export testKey=TESTING"
)

func setupTestApp() (err error) {
	Expect(os.MkdirAll(testAppDir, 0644)).To(Succeed())
	b := []byte(testEnvLine + "\n")
	if err = ioutil.WriteFile(testEnvFile, b, 0644); err != nil {
		return
	}
	return
}

func setupTestApp2() (err error) {
	Expect(os.MkdirAll(testAppDir2, 0644)).To(Succeed())
	b := []byte(testEnvLine2 + "\n")
	if err = ioutil.WriteFile(testEnvFile2, b, 0644); err != nil {
		return
	}
	return
}

func teardownTestApp() {
	os.RemoveAll(testAppDir)
}

func teardownTestApp2() {
	os.RemoveAll(testAppDir2)
}

func TestCommonGetEnv(t *testing.T) {
	RegisterTestingT(t)
	Expect(MustGetEnv("DOKKU_ROOT")).To(Equal("/home/dokku"))
}

func TestCommonGetAppImageRepo(t *testing.T) {
	RegisterTestingT(t)
	Expect(GetAppImageRepo("testapp")).To(Equal("dokku/testapp"))
}

func TestCommonVerifyImageInvalid(t *testing.T) {
	RegisterTestingT(t)
	Expect(VerifyImage("testapp")).To(Equal(false))
}

func TestCommonVerifyAppNameInvalid(t *testing.T) {
	RegisterTestingT(t)
	err := VerifyAppName("1994testApp")
	Expect(err).To(HaveOccurred())
}

func TestCommonVerifyAppName(t *testing.T) {
	RegisterTestingT(t)
	Expect(setupTestApp()).To(Succeed())
	Expect(VerifyAppName(testAppName)).To(Succeed())
	teardownTestApp()

	RegisterTestingT(t)
	Expect(setupTestApp2()).To(Succeed())
	Expect(VerifyAppName(testAppName2)).To(Succeed())
	teardownTestApp2()
}

func TestCommonDokkuAppsError(t *testing.T) {
	RegisterTestingT(t)
	_, err := DokkuApps()
	Expect(err).To(HaveOccurred())
}

func TestCommonDokkuApps(t *testing.T) {
	RegisterTestingT(t)
	os.Setenv("PLUGIN_ENABLED_PATH", "/var/lib/dokku/plugins/enabled")
	Expect(setupTestApp()).To(Succeed())
	apps, err := DokkuApps()
	Expect(err).NotTo(HaveOccurred())
	Expect(apps).To(HaveLen(1))
	Expect(apps[0]).To(Equal(testAppName))
	teardownTestApp()
}

func TestCommonStripInlineComments(t *testing.T) {
	RegisterTestingT(t)
	text := StripInlineComments(strings.Join([]string{testEnvLine, "# testing comment"}, " "))
	Expect(text).To(Equal(testEnvLine))
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

var nopHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

func TestFwdProtoHandler(t *testing.T) {
	prevForwardedProto := "fakeproto"
	prevForwardedPort := "8080"

	rec := httptest.NewRecorder()
	request, _ := http.NewRequest("GET", "http://test.com", nil)
	request.RemoteAddr = "1.2.3.4:5678"
	h := fwdProtoHandler{Handler: nopHandler, Proto: "https", Port: "443"}
	h.ServeHTTP(rec, request)
	if v := request.Header.Get("X-Forwarded-For"); v != "1.2.3.4" {
		t.Errorf("want X-Forwarded-For %s, got %s", "1.2.3.4", v)
	}
	if v := request.Header.Get("X-Forwarded-Proto"); v != "https" {
		t.Errorf("want X-Forwarded-Proto %s, got %s", "https", v)
	}
	if v := request.Header.Get("X-Forwarded-Port"); v != "443" {
		t.Errorf("want X-Forwarded-Port %s, got %s", "443", v)
	}

	// test with headers already set, make sure we append correctly
	rec = httptest.NewRecorder()
	request, _ = http.NewRequest("GET", "http://test.com", nil)
	request.Header.Set("X-Forwarded-Proto", prevForwardedProto)
	request.Header.Set("X-Forwarded-Port", prevForwardedPort)

	h.ServeHTTP(rec, request)
	want := prevForwardedProto + ", https"
	got := request.Header.Get("X-Forwarded-Proto")
	if want != got {
		t.Errorf("want X-Forwarded-Proto %q, got %q", want, got)
	}
	want = prevForwardedPort + ", 443"
	got = request.Header.Get("X-Forwarded-Port")
	if want != got {
		t.Errorf("want X-Forwarded-Port %q, got %q", want, got)
	}
}
<commit_msg>router: switch fwdProtoHandler test to use gocheck<commit_after>package main

import (
	"net/http"
	"net/http/httptest"

	. "github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-check"
)

var nopHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

func (s *S) TestFwdProtoHandler(c *C) {
	prevForwardedProto := "fakeproto"
	prevForwardedPort := "8080"

	rec := httptest.NewRecorder()
	request, _ := http.NewRequest("GET", "http://test.com", nil)
	request.RemoteAddr = "1.2.3.4:5678"
	h := fwdProtoHandler{Handler: nopHandler, Proto: "https", Port: "443"}
	h.ServeHTTP(rec, request)
	c.Assert(request.Header.Get("X-Forwarded-For"), Equals, "1.2.3.4")
	c.Assert(request.Header.Get("X-Forwarded-Proto"), Equals, "https")
	c.Assert(request.Header.Get("X-Forwarded-Port"), Equals, "443")

	// test with headers already set, make sure we append correctly
	rec = httptest.NewRecorder()
	request, _ = http.NewRequest("GET", "http://test.com", nil)
	request.Header.Set("X-Forwarded-Proto", prevForwardedProto)
	request.Header.Set("X-Forwarded-Port", prevForwardedPort)

	h.ServeHTTP(rec, request)
	c.Assert(request.Header.Get("X-Forwarded-Proto"), Equals, prevForwardedProto+", https")
	c.Assert(request.Header.Get("X-Forwarded-Port"), Equals, prevForwardedPort+", 443")
}
<|endoftext|>