<commit_before>/*
Copyright 2014 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 proxy

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/golang/glog"
)

// Proxier is a simple proxy for tcp connections between a localhost:lport and services that provide
// the actual implementations.
type Proxier struct {
	loadBalancer LoadBalancer
	serviceMap   map[string]int
}

func NewProxier(loadBalancer LoadBalancer) *Proxier {
	return &Proxier{loadBalancer: loadBalancer, serviceMap: make(map[string]int)}
}

func CopyBytes(in, out *net.TCPConn) {
	glog.Infof("Copying from %v <-> %v <-> %v <-> %v",
		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
	_, err := io.Copy(in, out)
	if err != nil && err != io.EOF {
		glog.Errorf("I/O error: %v", err)
	}

	in.CloseRead()
	out.CloseWrite()
}

// Create a bidirectional byte shuffler. Copies bytes to/from each connection.
func ProxyConnection(in, out *net.TCPConn) {
	glog.Infof("Creating proxy between %v <-> %v <-> %v <-> %v",
		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
	go CopyBytes(in, out)
	go CopyBytes(out, in)
}

func (proxier Proxier) AcceptHandler(service string, listener net.Listener) {
	for {
		inConn, err := listener.Accept()
		if err != nil {
			glog.Errorf("Accept failed: %v", err)
			continue
		}
		glog.Infof("Accepted connection from: %v to %v", inConn.RemoteAddr(), inConn.LocalAddr())

		// Figure out where this request should go.
		endpoint, err := proxier.loadBalancer.LoadBalance(service, inConn.RemoteAddr())
		if err != nil {
			glog.Errorf("Couldn't find an endpoint for %s %v", service, err)
			inConn.Close()
			continue
		}

		glog.Infof("Mapped service %s to endpoint %s", service, endpoint)
		outConn, err := net.DialTimeout("tcp", endpoint, time.Duration(5)*time.Second)
		// We basically need to take everything from inConn and send to outConn
		// and anything coming from outConn needs to be sent to inConn.
		if err != nil {
			glog.Errorf("Dial failed: %v", err)
			inConn.Close()
			continue
		}
		go ProxyConnection(inConn.(*net.TCPConn), outConn.(*net.TCPConn))
	}
}

// AddService starts listening for a new service on a given port.
func (proxier Proxier) AddService(service string, port int) error {
	// Make sure we can start listening on the port before saying all's well.
	l, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return err
	}
	proxier.addServiceCommon(service, l)
	return nil
}

// addService starts listening for a new service, returning the port it's using.
// For testing on a system with unknown ports used.
func (proxier Proxier) addServiceOnUnusedPort(service string) (string, error) {
	// Make sure we can start listening on the port before saying all's well.
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		return "", err
	}
	proxier.addServiceCommon(service, l)
	_, port, err := net.SplitHostPort(l.Addr().String())
	return port, nil
}

func (proxier Proxier) addServiceCommon(service string, l net.Listener) {
	glog.Infof("Listening for %s on %s", service, l.Addr().String())
	// If that succeeds, start the accepting loop.
	go proxier.AcceptHandler(service, l)
}

func (proxier Proxier) OnUpdate(services []api.Service) {
	glog.Infof("Received update notice: %+v", services)
	for _, service := range services {
		port, exists := proxier.serviceMap[service.ID]
		if !exists || port != service.Port {
			glog.Infof("Adding a new service %s on port %d", service.ID, service.Port)
			err := proxier.AddService(service.ID, service.Port)
			if err == nil {
				proxier.serviceMap[service.ID] = service.Port
			} else {
				glog.Infof("Failed to start listening for %s on %d", service.ID, service.Port)
			}
		}
	}
}
<commit_msg>pkg/proxy: remove unnecessary io.EOF checking<commit_after>/*
Copyright 2014 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 proxy

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/golang/glog"
)

// Proxier is a simple proxy for tcp connections between a localhost:lport and services that provide
// the actual implementations.
type Proxier struct {
	loadBalancer LoadBalancer
	serviceMap   map[string]int
}

func NewProxier(loadBalancer LoadBalancer) *Proxier {
	return &Proxier{loadBalancer: loadBalancer, serviceMap: make(map[string]int)}
}

func CopyBytes(in, out *net.TCPConn) {
	glog.Infof("Copying from %v <-> %v <-> %v <-> %v",
		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
	_, err := io.Copy(in, out)
	if err != nil {
		glog.Errorf("I/O error: %v", err)
	}

	in.CloseRead()
	out.CloseWrite()
}

// Create a bidirectional byte shuffler. Copies bytes to/from each connection.
func ProxyConnection(in, out *net.TCPConn) {
	glog.Infof("Creating proxy between %v <-> %v <-> %v <-> %v",
		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
	go CopyBytes(in, out)
	go CopyBytes(out, in)
}

func (proxier Proxier) AcceptHandler(service string, listener net.Listener) {
	for {
		inConn, err := listener.Accept()
		if err != nil {
			glog.Errorf("Accept failed: %v", err)
			continue
		}
		glog.Infof("Accepted connection from: %v to %v", inConn.RemoteAddr(), inConn.LocalAddr())

		// Figure out where this request should go.
		endpoint, err := proxier.loadBalancer.LoadBalance(service, inConn.RemoteAddr())
		if err != nil {
			glog.Errorf("Couldn't find an endpoint for %s %v", service, err)
			inConn.Close()
			continue
		}

		glog.Infof("Mapped service %s to endpoint %s", service, endpoint)
		outConn, err := net.DialTimeout("tcp", endpoint, time.Duration(5)*time.Second)
		// We basically need to take everything from inConn and send to outConn
		// and anything coming from outConn needs to be sent to inConn.
		if err != nil {
			glog.Errorf("Dial failed: %v", err)
			inConn.Close()
			continue
		}
		go ProxyConnection(inConn.(*net.TCPConn), outConn.(*net.TCPConn))
	}
}

// AddService starts listening for a new service on a given port.
func (proxier Proxier) AddService(service string, port int) error {
	// Make sure we can start listening on the port before saying all's well.
	l, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return err
	}
	proxier.addServiceCommon(service, l)
	return nil
}

// addService starts listening for a new service, returning the port it's using.
// For testing on a system with unknown ports used.
func (proxier Proxier) addServiceOnUnusedPort(service string) (string, error) {
	// Make sure we can start listening on the port before saying all's well.
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		return "", err
	}
	proxier.addServiceCommon(service, l)
	_, port, err := net.SplitHostPort(l.Addr().String())
	return port, nil
}

func (proxier Proxier) addServiceCommon(service string, l net.Listener) {
	glog.Infof("Listening for %s on %s", service, l.Addr().String())
	// If that succeeds, start the accepting loop.
	go proxier.AcceptHandler(service, l)
}

func (proxier Proxier) OnUpdate(services []api.Service) {
	glog.Infof("Received update notice: %+v", services)
	for _, service := range services {
		port, exists := proxier.serviceMap[service.ID]
		if !exists || port != service.Port {
			glog.Infof("Adding a new service %s on port %d", service.ID, service.Port)
			err := proxier.AddService(service.ID, service.Port)
			if err == nil {
				proxier.serviceMap[service.ID] = service.Port
			} else {
				glog.Infof("Failed to start listening for %s on %d", service.ID, service.Port)
			}
		}
	}
}
<|endoftext|>


<commit_before>package update

import (
	"io/ioutil"

	"crypto"

	"fmt"

	"os"

	"bufio"

	"strings"

	"encoding/base64"

	"github.com/containerum/chkit/cmd/util"
	"github.com/containerum/chkit/pkg/chkitErrors"
	"github.com/inconshreveable/go-update"
	"golang.org/x/crypto/ssh/terminal"
	"gopkg.in/urfave/cli.v2"
)

var PublicKeyB64 = "cHVibGljIGtleQo="

const (
	ErrUpdateApply = chkitErrors.Err("update apply failed")
)

func verifiedUpdate(upd *Package) error {
	checksum, err := ioutil.ReadAll(upd.Hash)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}

	signature, err := ioutil.ReadAll(upd.Signature)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}

	opts := update.Options{
		Checksum:  checksum,
		Signature: signature,
		Verifier:  update.NewECDSAVerifier(),
		Hash:      crypto.SHA256,
	}
	publicKey, err := base64.StdEncoding.DecodeString(PublicKeyB64)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}
	err = opts.SetPublicKeyPEM(publicKey)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}
	err = update.Apply(upd.Binary, opts)

	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}

	return nil
}

func Update(ctx *cli.Context, downloader LatestCheckerDownloader, restartAfter bool) error {
	latestVersion, err := downloader.LatestVersion()
	if err != nil {
		return err
	}

	if latestVersion.LE(util.GetVersion(ctx)) {
		return nil
	}

	// check if we have terminal supports escape sequences
	var colorStart, colorEnd string
	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		colorStart = "\x1b[31;1m"
		colorEnd = "\x1b[0m"
	}
	fmt.Printf("%sYou are using version %s, however version %s is available%s\n",
		colorStart, util.GetVersion(ctx), latestVersion, colorEnd)

	scanner := bufio.NewScanner(os.Stdin)
	scanner.Split(bufio.ScanWords)
askLoop:
	for {
		fmt.Println("Do you want to update [Y/n]: ")
		for !scanner.Scan() {
			break
		}
		if scanner.Err() != nil {
			util.GetLog(ctx).WithError(err).Error("scan failed")
			continue
		}
		switch strings.ToLower(scanner.Text()) {
		case "", "y":
			break askLoop
		case "n":
			return nil
		default:
			continue
		}
	}

	archive, err := downloader.LatestDownload()
	if err != nil {
		return err
	}
	defer archive.Close()

	pkg, err := unpack(archive)
	if err != nil {
		return err
	}
	defer pkg.Close()

	err = verifiedUpdate(pkg)
	if err != nil {
		return err
	}

	if restartAfter {
		gracefulRestart(ctx)
	}

	return nil
}
<commit_msg>Extract ask for update logic to function<commit_after>package update

import (
	"io/ioutil"

	"crypto"

	"fmt"

	"os"

	"bufio"

	"strings"

	"encoding/base64"

	"github.com/blang/semver"
	"github.com/containerum/chkit/cmd/util"
	"github.com/containerum/chkit/pkg/chkitErrors"
	"github.com/inconshreveable/go-update"
	"golang.org/x/crypto/ssh/terminal"
	"gopkg.in/urfave/cli.v2"
)

var PublicKeyB64 = "cHVibGljIGtleQo="

const (
	ErrUpdateApply = chkitErrors.Err("update apply failed")
)

func verifiedUpdate(upd *Package) error {
	checksum, err := ioutil.ReadAll(upd.Hash)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}

	signature, err := ioutil.ReadAll(upd.Signature)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}

	opts := update.Options{
		Checksum:  checksum,
		Signature: signature,
		Verifier:  update.NewECDSAVerifier(),
		Hash:      crypto.SHA256,
	}
	publicKey, err := base64.StdEncoding.DecodeString(PublicKeyB64)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}
	err = opts.SetPublicKeyPEM(publicKey)
	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}
	err = update.Apply(upd.Binary, opts)

	if err != nil {
		return chkitErrors.Wrap(ErrUpdateApply, err)
	}

	return nil
}

func AskForUpdate(ctx *cli.Context, latestVersion semver.Version) (bool, error) {
	// check if we have terminal supports escape sequences
	var colorStart, colorEnd string
	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		colorStart = "\x1b[31;1m"
		colorEnd = "\x1b[0m"
	}
	fmt.Printf("%sYou are using version %s, however version %s is available%s\n",
		colorStart, util.GetVersion(ctx), latestVersion, colorEnd)

	scanner := bufio.NewScanner(os.Stdin)
	scanner.Split(bufio.ScanWords)
	for {
		fmt.Println("Do you want to update [Y/n]: ")
		for !scanner.Scan() {
			break
		}
		if err := scanner.Err(); err != nil {
			util.GetLog(ctx).WithError(err).Error("scan failed")
			return false, err
		}
		switch strings.ToLower(scanner.Text()) {
		case "", "y":
			return true, nil
		case "n":
			return false, nil
		default:
			continue
		}
	}
}

func Update(ctx *cli.Context, downloader LatestCheckerDownloader, restartAfter bool) error {
	latestVersion, err := downloader.LatestVersion()
	if err != nil {
		return err
	}

	if latestVersion.LE(util.GetVersion(ctx)) {
		return nil
	}

	doUpdate, err := AskForUpdate(ctx, latestVersion)
	if err != nil {
		return err
	}
	if !doUpdate {
		return nil
	}

	archive, err := downloader.LatestDownload()
	if err != nil {
		return err
	}
	defer archive.Close()

	pkg, err := unpack(archive)
	if err != nil {
		return err
	}
	defer pkg.Close()

	err = verifiedUpdate(pkg)
	if err != nil {
		return err
	}

	if restartAfter {
		gracefulRestart(ctx)
	}

	return nil
}
<|endoftext|>


<commit_before>package plugin

import (
	"database/sql"
	"time"

	"github.com/vektra/cypress"
)

const cEnableHstore = `
CREATE EXTENSION hstore
`

const cCreateTable = `
CREATE TABLE cypress_messages (
	timestamp TIMESTAMP,
	version INTEGER,
	type INTEGER,
	session_id TEXT,
	attributes HSTORE,
	tags HSTORE
)`

const cAddRow = `
INSERT INTO cypress_messages (
	timestamp,
	version,
	type,
	session_id,
	attributes,
	tags
) VALUES ($1, $2, $3, $4, $5, $6)`

type DBInterface interface {
	Ping() error
	Exec(query string, args ...interface{}) (sql.Result, error)
}

type ResultInterface interface {
	LastInsertId() (int64, error)
	RowsAffected() (int64, error)
}

type PostgreSQL struct {
	DB DBInterface
}

func (p *PostgreSQL) Init(db DBInterface) {
	p.DB = db
}

func (p *PostgreSQL) SetupDB() error {
	err := p.DB.Ping()
	if err != nil {
		return err
	}

	// TODO: first check if already enabled
	_, err = p.DB.Exec(cEnableHstore)
	if err != nil {
		// return err
	}

	// TODO: first check if already created
	_, err = p.DB.Exec(cCreateTable)
	if err != nil {
		// return err
	}

	// TODO: alter table if schema doesnt match

	return nil
}

func (p *PostgreSQL) Receive(m *cypress.Message) error {
	_, err := p.DB.Exec(cAddRow,
		m.GetTimestamp().Time().Format(time.RFC3339Nano),
		m.Version,
		m.Type,
		m.SessionId,
		m.HstoreAttributes(),
		m.HstoreTags(),
	)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Create db table and extension only if they don't exist<commit_after>package plugin

import (
	"database/sql"
	"time"

	"github.com/vektra/cypress"
)

const cEnableHstore = `
CREATE EXTENSION IF NOT EXISTS hstore
`

const cCreateTable = `
CREATE TABLE IF NOT EXISTS cypress_messages (
	timestamp TIMESTAMP,
	version INTEGER,
	type INTEGER,
	session_id TEXT,
	attributes HSTORE,
	tags HSTORE
)`

const cAddRow = `
INSERT INTO cypress_messages (
	timestamp,
	version,
	type,
	session_id,
	attributes,
	tags
) VALUES ($1, $2, $3, $4, $5, $6)`

type DBInterface interface {
	Ping() error
	Exec(query string, args ...interface{}) (sql.Result, error)
}

type ResultInterface interface {
	LastInsertId() (int64, error)
	RowsAffected() (int64, error)
}

type PostgreSQL struct {
	DB DBInterface
}

func (p *PostgreSQL) Init(db DBInterface) {
	p.DB = db
}

func (p *PostgreSQL) SetupDB() error {
	err := p.DB.Ping()
	if err != nil {
		return err
	}

	_, err = p.DB.Exec(cEnableHstore)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec(cCreateTable)
	if err != nil {
		return err
	}

	// TODO: alter table if schema doesnt match ?

	return nil
}

func (p *PostgreSQL) Receive(m *cypress.Message) error {
	_, err := p.DB.Exec(cAddRow,
		m.GetTimestamp().Time().Format(time.RFC3339Nano),
		m.Version,
		m.Type,
		m.SessionId,
		m.HstoreAttributes(),
		m.HstoreTags(),
	)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/maniksurtani/quotaservice/master/LICENSE

package config

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"sync"
	"time"

	"github.com/maniksurtani/quotaservice/logging"
	"github.com/samuel/go-zookeeper/zk"
)

const (
	sessionTimeout = 3 * time.Second
	createRetries  = 3
)

type eventListener func() (<-chan zk.Event, error)

type ZkWatch struct {
	listener eventListener
	channel  <-chan zk.Event
	stopper  chan struct{}
}

type ZkConfigPersister struct {
	// Current configuration hash
	config string

	// Historical map of configurations
	// hash -> marshalled config
	configs map[string][]byte

	// Base Zookeeper path
	path string

	watcher chan struct{}

	conn  *zk.Conn
	watch *ZkWatch

	wg sync.WaitGroup
	sync.RWMutex
}

type zkLogger struct{}

var _ zk.Logger = &zkLogger{}

func (_ *zkLogger) Printf(format string, args ...interface{}) {
	logging.Printf(format, args...)
}

func NewZkConfigPersister(path string, servers []string) (ConfigPersister, error) {
	conn, _, err := zk.Connect(servers, sessionTimeout, func(c *zk.Conn) {
		c.SetLogger(&zkLogger{})
	})

	if err != nil {
		return nil, err
	}

	err = createPath(conn, path)

	if err != nil {
		conn.Close()
		return nil, err
	}

	persister := &ZkConfigPersister{
		conn:    conn,
		path:    path,
		watcher: make(chan struct{}, 1),
		configs: make(map[string][]byte)}

	watch, err := persister.createWatch(persister.currentConfigEventListener)

	if err != nil {
		conn.Close()
		return nil, err
	}

	persister.watch = watch

	return persister, nil
}

func (z *ZkConfigPersister) createWatch(listener eventListener) (*ZkWatch, error) {
	watch := &ZkWatch{
		listener: listener,
		stopper:  make(chan struct{}, 1)}

	channel, err := listener()

	if err != nil {
		return nil, err
	}

	watch.channel = channel

	z.wg.Add(1)
	go z.waitForEvents(watch)

	return watch, nil
}

func (z *ZkConfigPersister) waitForEvents(watch *ZkWatch) {
	defer z.wg.Done()

	for {
		select {
		case event := <-watch.channel:
			if event.Err != nil {
				logging.Print("Received error from zookeeper", event)
			}
		case <-watch.stopper:
			return
		}

		channel, err := watch.listener()

		if err != nil {
			logging.Printf("Received error from zookeeper executing listener: %+v", err)
			continue
		}

		watch.channel = channel
	}
}

// Tries to create the configuration path, if it doesn't exist
// It tries multiple times in case there's a race with another quotaservice node coming up
func createPath(conn *zk.Conn, path string) (err error) {
	for i := 0; i < createRetries; i++ {
		exists, _, err := conn.Exists(path)

		if exists && err == nil {
			return nil
		}

		_, err = conn.Create(path, []byte{}, 0, zk.WorldACL(zk.PermAll))

		if err == nil {
			return nil
		}

		logging.Printf("Could not create zk path, sleeping for 100ms")
		time.Sleep(100 * time.Millisecond)
	}

	if err == nil {
		err = errors.New("could not create and get path " + path)
	}

	return err
}

// PersistAndNotify persists a marshalled configuration passed in.
func (z *ZkConfigPersister) PersistAndNotify(marshalledConfig io.Reader) error {
	z.Lock()
	defer z.Unlock()

	b, e := ioutil.ReadAll(marshalledConfig)

	if e != nil {
		return e
	}

	key := HashConfig(b)

	if key == z.config {
		return nil
	}

	if err := z.archiveConfig(key, b); err != nil {
		return err
	}

	_, err := z.conn.Set(z.path, []byte(key), -1)

	// There is no notification, that happens when zookeeper alerts the watcher

	return err
}

// ReadPersistedConfig provides a reader to a marshalled config previously persisted.
func (z *ZkConfigPersister) ReadPersistedConfig() (io.Reader, error) {
	z.RLock()
	defer z.RUnlock()

	return bytes.NewReader(z.configs[z.config]), nil
}

func (z *ZkConfigPersister) currentConfigEventListener() (<-chan zk.Event, error) {
	z.Lock()
	defer z.Unlock()

	children, _, err := z.conn.Children(z.path)

	if err != nil {
		logging.Printf("Received error from zookeeper when fetching children of %s: %+v", z.path, err)
		return nil, err
	}

	z.configs = make(map[string][]byte)

	for _, child := range children {
		path := fmt.Sprintf("%s/%s", z.path, child)
		data, _, err := z.conn.Get(path)

		if err != nil {
			logging.Printf("Received error from zookeeper when fetching %s: %+v", path, err)
		} else {
			z.configs[child] = data
		}
	}

	config, _, ch, err := z.conn.GetW(z.path)

	if err != nil {
		logging.Printf("Received error from zookeeper when fetching %s: %+v", z.path, err)
		return nil, err
	}

	z.config = string(config)

	select {
	case z.watcher <- struct{}{}:
		// Notified
	default:
		// Doesn't matter; another notification is pending.
	}

	return ch, nil
}

func (z *ZkConfigPersister) archiveConfig(key string, config []byte) error {
	path := fmt.Sprintf("%s/%s", z.path, key)
	_, err := z.conn.Create(path, config, 0, zk.WorldACL(zk.PermAll))
	return err
}

// ConfigChangedWatcher returns a channel that is notified whenever configuration changes are
// detected. Changes are coalesced so that a single notification may be emitted for multiple
// changes.
func (z *ZkConfigPersister) ConfigChangedWatcher() <-chan struct{} {
	return z.watcher
}

// ReadHistoricalConfigs returns an array of previously persisted configs
func (z *ZkConfigPersister) ReadHistoricalConfigs() ([]io.Reader, error) {
	z.RLock()
	defer z.RUnlock()

	readers := make([]io.Reader, 0)

	for _, config := range z.configs {
		readers = append(readers, bytes.NewReader(config))
	}

	return readers, nil
}

// Closes makes sure all event listeners are done
// and then closes the connection
func (z *ZkConfigPersister) Close() {
	z.watch.stopper <- struct{}{}
	z.wg.Wait()

	close(z.watch.stopper)
	close(z.watcher)

	z.conn.Close()
}
<commit_msg>zookeeper: allow passing options to underlying connection (#61)<commit_after>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/maniksurtani/quotaservice/master/LICENSE

package config

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"sync"
	"time"

	"github.com/maniksurtani/quotaservice/logging"
	"github.com/samuel/go-zookeeper/zk"
)

const (
	sessionTimeout = 3 * time.Second
	createRetries  = 3
)

type eventListener func() (<-chan zk.Event, error)

type ZkWatch struct {
	listener eventListener
	channel  <-chan zk.Event
	stopper  chan struct{}
}

type ZkConfigPersister struct {
	// Current configuration hash
	config string

	// Historical map of configurations
	// hash -> marshalled config
	configs map[string][]byte

	// Base Zookeeper path
	path string

	watcher chan struct{}

	conn  *zk.Conn
	watch *ZkWatch

	wg sync.WaitGroup
	sync.RWMutex
}

// Mirrors go-zookeeper's connOption
type connOption func(c *zk.Conn)

func NewZkConfigPersister(path string, servers []string, options ...connOption) (ConfigPersister, error) {
	conn, _, err := zk.Connect(servers, sessionTimeout, func(c *zk.Conn) {
		c.SetLogger(logging.CurrentLogger())

		// Allows overriding logger/etc. go-zookeeper options
		for _, option := range options {
			option(c)
		}
	})

	if err != nil {
		return nil, err
	}

	err = createPath(conn, path)

	if err != nil {
		conn.Close()
		return nil, err
	}

	persister := &ZkConfigPersister{
		conn:    conn,
		path:    path,
		watcher: make(chan struct{}, 1),
		configs: make(map[string][]byte)}

	watch, err := persister.createWatch(persister.currentConfigEventListener)

	if err != nil {
		conn.Close()
		return nil, err
	}

	persister.watch = watch

	return persister, nil
}

func (z *ZkConfigPersister) createWatch(listener eventListener) (*ZkWatch, error) {
	watch := &ZkWatch{
		listener: listener,
		stopper:  make(chan struct{}, 1)}

	channel, err := listener()

	if err != nil {
		return nil, err
	}

	watch.channel = channel

	z.wg.Add(1)
	go z.waitForEvents(watch)

	return watch, nil
}

func (z *ZkConfigPersister) waitForEvents(watch *ZkWatch) {
	defer z.wg.Done()

	for {
		select {
		case event := <-watch.channel:
			if event.Err != nil {
				logging.Print("Received error from zookeeper", event)
			}
		case <-watch.stopper:
			return
		}

		channel, err := watch.listener()

		if err != nil {
			logging.Printf("Received error from zookeeper executing listener: %+v", err)
			continue
		}

		watch.channel = channel
	}
}

// Tries to create the configuration path, if it doesn't exist
// It tries multiple times in case there's a race with another quotaservice node coming up
func createPath(conn *zk.Conn, path string) (err error) {
	for i := 0; i < createRetries; i++ {
		exists, _, err := conn.Exists(path)

		if exists && err == nil {
			return nil
		}

		_, err = conn.Create(path, []byte{}, 0, zk.WorldACL(zk.PermAll))

		if err == nil {
			return nil
		}

		logging.Printf("Could not create zk path, sleeping for 100ms")
		time.Sleep(100 * time.Millisecond)
	}

	if err == nil {
		err = errors.New("could not create and get path " + path)
	}

	return err
}

// PersistAndNotify persists a marshalled configuration passed in.
func (z *ZkConfigPersister) PersistAndNotify(marshalledConfig io.Reader) error {
	z.Lock()
	defer z.Unlock()

	b, e := ioutil.ReadAll(marshalledConfig)

	if e != nil {
		return e
	}

	key := HashConfig(b)

	if key == z.config {
		return nil
	}

	if err := z.archiveConfig(key, b); err != nil {
		return err
	}

	_, err := z.conn.Set(z.path, []byte(key), -1)

	// There is no notification, that happens when zookeeper alerts the watcher

	return err
}

// ReadPersistedConfig provides a reader to a marshalled config previously persisted.
func (z *ZkConfigPersister) ReadPersistedConfig() (io.Reader, error) {
	z.RLock()
	defer z.RUnlock()

	return bytes.NewReader(z.configs[z.config]), nil
}

func (z *ZkConfigPersister) currentConfigEventListener() (<-chan zk.Event, error) {
	z.Lock()
	defer z.Unlock()

	children, _, err := z.conn.Children(z.path)

	if err != nil {
		logging.Printf("Received error from zookeeper when fetching children of %s: %+v", z.path, err)
		return nil, err
	}

	z.configs = make(map[string][]byte)

	for _, child := range children {
		path := fmt.Sprintf("%s/%s", z.path, child)
		data, _, err := z.conn.Get(path)

		if err != nil {
			logging.Printf("Received error from zookeeper when fetching %s: %+v", path, err)
		} else {
			z.configs[child] = data
		}
	}

	config, _, ch, err := z.conn.GetW(z.path)

	if err != nil {
		logging.Printf("Received error from zookeeper when fetching %s: %+v", z.path, err)
		return nil, err
	}

	z.config = string(config)

	select {
	case z.watcher <- struct{}{}:
		// Notified
	default:
		// Doesn't matter; another notification is pending.
	}

	return ch, nil
}

func (z *ZkConfigPersister) archiveConfig(key string, config []byte) error {
	path := fmt.Sprintf("%s/%s", z.path, key)
	_, err := z.conn.Create(path, config, 0, zk.WorldACL(zk.PermAll))
	return err
}

// ConfigChangedWatcher returns a channel that is notified whenever configuration changes are
// detected. Changes are coalesced so that a single notification may be emitted for multiple
// changes.
func (z *ZkConfigPersister) ConfigChangedWatcher() <-chan struct{} {
	return z.watcher
}

// ReadHistoricalConfigs returns an array of previously persisted configs
func (z *ZkConfigPersister) ReadHistoricalConfigs() ([]io.Reader, error) {
	z.RLock()
	defer z.RUnlock()

	readers := make([]io.Reader, 0)

	for _, config := range z.configs {
		readers = append(readers, bytes.NewReader(config))
	}

	return readers, nil
}

// Closes makes sure all event listeners are done
// and then closes the connection
func (z *ZkConfigPersister) Close() {
	z.watch.stopper <- struct{}{}
	z.wg.Wait()

	close(z.watch.stopper)
	close(z.watcher)

	z.conn.Close()
}
<|endoftext|>


<commit_before>package consulrunner

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"time"

	"code.cloudfoundry.org/cfhttp"
	"code.cloudfoundry.org/consuladapter"
	"github.com/hashicorp/consul/api"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	. "github.com/onsi/gomega"
)

type ClusterRunner struct {
	startingPort    int
	numNodes        int
	consulProcesses []ifrit.Process
	running         bool
	dataDir         string
	configDir       string
	scheme          string
	sessionTTL      time.Duration

	mutex *sync.RWMutex
}

const defaultDataDirPrefix = "consul_data"
const defaultConfigDirPrefix = "consul_config"

func NewClusterRunner(startingPort int, numNodes int, scheme string) *ClusterRunner {
	Expect(startingPort).To(BeNumerically(">", 0))
	Expect(startingPort).To(BeNumerically("<", 1<<16))
	Expect(numNodes).To(BeNumerically(">", 0))

	return &ClusterRunner{
		startingPort: startingPort,
		numNodes:     numNodes,
		scheme:       scheme,
		sessionTTL:   5 * time.Second,

		mutex: &sync.RWMutex{},
	}
}

func (cr *ClusterRunner) SessionTTL() time.Duration {
	return cr.sessionTTL
}

func (cr *ClusterRunner) Start() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if cr.running {
		return
	}

	tmpDir, err := ioutil.TempDir("", defaultDataDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.dataDir = tmpDir

	tmpDir, err = ioutil.TempDir("", defaultConfigDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.configDir = tmpDir

	cr.consulProcesses = make([]ifrit.Process, cr.numNodes)

	for i := 0; i < cr.numNodes; i++ {
		iStr := fmt.Sprintf("%d", i)
		nodeDataDir := path.Join(cr.dataDir, iStr)
		os.MkdirAll(nodeDataDir, 0700)

		configFilePath := writeConfigFile(
			cr.configDir,
			nodeDataDir,
			iStr,
			cr.startingPort,
			i,
			cr.numNodes,
			cr.sessionTTL,
		)

		process := ginkgomon.Invoke(ginkgomon.New(ginkgomon.Config{
			Name:              fmt.Sprintf("consul_cluster[%d]", i),
			AnsiColorCode:     "35m",
			StartCheck:        "agent: Join completed.",
			StartCheckTimeout: 10 * time.Second,
			Command: exec.Command(
				"consul",
				"agent",
				"--log-level", "trace",
				"--config-file", configFilePath,
			),
		}))
		cr.consulProcesses[i] = process

		ready := process.Ready()
		Eventually(ready, 10, 0.05).Should(BeClosed(), "Expected consul to be up and running")
	}

	cr.running = true
}

func (cr *ClusterRunner) NewClient() consuladapter.Client {
	client, err := api.NewClient(&api.Config{
		Address:    cr.Address(),
		Scheme:     cr.scheme,
		HttpClient: cfhttp.NewStreamingClient(),
	})
	Expect(err).NotTo(HaveOccurred())

	return consuladapter.NewConsulClient(client)
}

func (cr *ClusterRunner) WaitUntilReady() {
	client := cr.NewClient()
	catalog := client.Catalog()

	Eventually(func() error {
		_, qm, err := catalog.Nodes(nil)
		if err != nil {
			return err
		}
		if qm.KnownLeader && qm.LastIndex > 0 {
			return nil
		}
		return errors.New("not ready")
	}, 10, 100*time.Millisecond).Should(BeNil())
}

func (cr *ClusterRunner) Stop() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if !cr.running {
		return
	}

	for i := 0; i < cr.numNodes; i++ {
		stopSignal(cr.consulProcesses[i], 5*time.Second)
	}

	os.RemoveAll(cr.dataDir)
	os.RemoveAll(cr.configDir)
	cr.consulProcesses = nil
	cr.running = false
}

func (cr *ClusterRunner) ConsulCluster() string {
	urls := make([]string, cr.numNodes)
	for i := 0; i < cr.numNodes; i++ {
		urls[i] = fmt.Sprintf("%s://127.0.0.1:%d", cr.scheme, cr.startingPort+i*PortOffsetLength+PortOffsetHTTP)
	}

	return strings.Join(urls, ",")
}

func (cr *ClusterRunner) Address() string {
	return fmt.Sprintf("127.0.0.1:%d", cr.startingPort+PortOffsetHTTP)
}

func (cr *ClusterRunner) URL() string {
	return fmt.Sprintf("%s://%s", cr.scheme, cr.Address())
}

func (cr *ClusterRunner) Reset() error {
	client := cr.NewClient()

	sessions, _, err := client.Session().List(nil)
	if err == nil {
		for _, session := range sessions {
			_, err1 := client.Session().Destroy(session.ID, nil)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	services, err := client.Agent().Services()
	if err == nil {
		for _, service := range services {
			err1 := client.Agent().ServiceDeregister(service.ID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	checks, err := client.Agent().Checks()
	if err == nil {
		for _, check := range checks {
			err1 := client.Agent().CheckDeregister(check.CheckID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	_, err1 := client.KV().DeleteTree("", nil)

	return err1
}
<commit_msg>Skip over consul when resetting Consul services<commit_after>package consulrunner

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"time"

	"code.cloudfoundry.org/cfhttp"
	"code.cloudfoundry.org/consuladapter"
	"github.com/hashicorp/consul/api"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	. "github.com/onsi/gomega"
)

type ClusterRunner struct {
	startingPort    int
	numNodes        int
	consulProcesses []ifrit.Process
	running         bool
	dataDir         string
	configDir       string
	scheme          string
	sessionTTL      time.Duration

	mutex *sync.RWMutex
}

const defaultDataDirPrefix = "consul_data"
const defaultConfigDirPrefix = "consul_config"

func NewClusterRunner(startingPort int, numNodes int, scheme string) *ClusterRunner {
	Expect(startingPort).To(BeNumerically(">", 0))
	Expect(startingPort).To(BeNumerically("<", 1<<16))
	Expect(numNodes).To(BeNumerically(">", 0))

	return &ClusterRunner{
		startingPort: startingPort,
		numNodes:     numNodes,
		scheme:       scheme,
		sessionTTL:   5 * time.Second,

		mutex: &sync.RWMutex{},
	}
}

func (cr *ClusterRunner) SessionTTL() time.Duration {
	return cr.sessionTTL
}

func (cr *ClusterRunner) Start() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if cr.running {
		return
	}

	tmpDir, err := ioutil.TempDir("", defaultDataDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.dataDir = tmpDir

	tmpDir, err = ioutil.TempDir("", defaultConfigDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.configDir = tmpDir

	cr.consulProcesses = make([]ifrit.Process, cr.numNodes)

	for i := 0; i < cr.numNodes; i++ {
		iStr := fmt.Sprintf("%d", i)
		nodeDataDir := path.Join(cr.dataDir, iStr)
		os.MkdirAll(nodeDataDir, 0700)

		configFilePath := writeConfigFile(
			cr.configDir,
			nodeDataDir,
			iStr,
			cr.startingPort,
			i,
			cr.numNodes,
			cr.sessionTTL,
		)

		process := ginkgomon.Invoke(ginkgomon.New(ginkgomon.Config{
			Name:              fmt.Sprintf("consul_cluster[%d]", i),
			AnsiColorCode:     "35m",
			StartCheck:        "agent: Join completed.",
			StartCheckTimeout: 10 * time.Second,
			Command: exec.Command(
				"consul",
				"agent",
				"--log-level", "trace",
				"--config-file", configFilePath,
			),
		}))
		cr.consulProcesses[i] = process

		ready := process.Ready()
		Eventually(ready, 10, 0.05).Should(BeClosed(), "Expected consul to be up and running")
	}

	cr.running = true
}

func (cr *ClusterRunner) NewClient() consuladapter.Client {
	client, err := api.NewClient(&api.Config{
		Address:    cr.Address(),
		Scheme:     cr.scheme,
		HttpClient: cfhttp.NewStreamingClient(),
	})
	Expect(err).NotTo(HaveOccurred())

	return consuladapter.NewConsulClient(client)
}

func (cr *ClusterRunner) WaitUntilReady() {
	client := cr.NewClient()
	catalog := client.Catalog()

	Eventually(func() error {
		_, qm, err := catalog.Nodes(nil)
		if err != nil {
			return err
		}
		if qm.KnownLeader && qm.LastIndex > 0 {
			return nil
		}
		return errors.New("not ready")
	}, 10, 100*time.Millisecond).Should(BeNil())
}

func (cr *ClusterRunner) Stop() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if !cr.running {
		return
	}

	for i := 0; i < cr.numNodes; i++ {
		stopSignal(cr.consulProcesses[i], 5*time.Second)
	}

	os.RemoveAll(cr.dataDir)
	os.RemoveAll(cr.configDir)
	cr.consulProcesses = nil
	cr.running = false
}

func (cr *ClusterRunner) ConsulCluster() string {
	urls := make([]string, cr.numNodes)
	for i := 0; i < cr.numNodes; i++ {
		urls[i] = fmt.Sprintf("%s://127.0.0.1:%d", cr.scheme, cr.startingPort+i*PortOffsetLength+PortOffsetHTTP)
	}

	return strings.Join(urls, ",")
}

func (cr *ClusterRunner) Address() string {
	return fmt.Sprintf("127.0.0.1:%d", cr.startingPort+PortOffsetHTTP)
}

func (cr *ClusterRunner) URL() string {
	return fmt.Sprintf("%s://%s", cr.scheme, cr.Address())
}

func (cr *ClusterRunner) Reset() error {
	client := cr.NewClient()

	sessions, _, err := client.Session().List(nil)
	if err == nil {
		for _, session := range sessions {
			_, err1 := client.Session().Destroy(session.ID, nil)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	services, err := client.Agent().Services()
	if err == nil {
		for _, service := range services {
			if service.Service == "consul" {
				continue
			}
			err1 := client.Agent().ServiceDeregister(service.ID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	checks, err := client.Agent().Checks()
	if err == nil {
		for _, check := range checks {
			err1 := client.Agent().CheckDeregister(check.CheckID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	_, err1 := client.KV().DeleteTree("", nil)

	return err1
}
<|endoftext|>


<commit_before>package consulrunner

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"time"

	"code.cloudfoundry.org/cfhttp"
	"code.cloudfoundry.org/consuladapter"
	"github.com/hashicorp/consul/api"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

type ClusterRunner struct {
	startingPort    int
	numNodes        int
	consulProcesses []ifrit.Process
	running         bool
	dataDir         string
	configDir       string
	scheme          string
	sessionTTL      time.Duration

	mutex *sync.RWMutex
}

const defaultDataDirPrefix = "consul_data"
const defaultConfigDirPrefix = "consul_config"

func NewClusterRunner(startingPort int, numNodes int, scheme string) *ClusterRunner {
	Expect(startingPort).To(BeNumerically(">", 0))
	Expect(startingPort).To(BeNumerically("<", 1<<16))
	Expect(numNodes).To(BeNumerically(">", 0))

	return &ClusterRunner{
		startingPort: startingPort,
		numNodes:     numNodes,
		scheme:       scheme,
		sessionTTL:   5 * time.Second,

		mutex: &sync.RWMutex{},
	}
}

func (cr *ClusterRunner) SessionTTL() time.Duration {
	return cr.sessionTTL
}

func (cr *ClusterRunner) ConsulVersion() string {
	cmd := exec.Command("consul", "-v")
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
	Expect(session.Out).To(gbytes.Say("Consul v"))
	lines := strings.Split(string(session.Out.Contents()), "\n")
	versionLine := lines[0]
	return strings.TrimPrefix(versionLine, "Consul v")
}

func (cr *ClusterRunner) HasPerformanceFlag() bool {
	return !strings.HasPrefix(cr.ConsulVersion(), "0.6")
}

func (cr *ClusterRunner) Start() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if cr.running {
		return
	}

	tmpDir, err := ioutil.TempDir("", defaultDataDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.dataDir = tmpDir

	tmpDir, err = ioutil.TempDir("", defaultConfigDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.configDir = tmpDir

	cr.consulProcesses = make([]ifrit.Process, cr.numNodes)

	for i := 0; i < cr.numNodes; i++ {
		iStr := fmt.Sprintf("%d", i)
		nodeDataDir := path.Join(cr.dataDir, iStr)
		os.MkdirAll(nodeDataDir, 0700)

		configFilePath := writeConfigFile(
			cr.HasPerformanceFlag(),
			cr.configDir,
			nodeDataDir,
			iStr,
			cr.startingPort,
			i,
			cr.numNodes,
			cr.sessionTTL,
		)

		process := ginkgomon.Invoke(ginkgomon.New(ginkgomon.Config{
			Name:              fmt.Sprintf("consul_cluster[%d]", i),
			AnsiColorCode:     "35m",
			StartCheck:        "agent: Join completed.",
			StartCheckTimeout: 10 * time.Second,
			Command: exec.Command(
				"consul",
				"agent",
				"--log-level", "trace",
				"--config-file", configFilePath,
			),
		}))
		cr.consulProcesses[i] = process

		ready := process.Ready()
		Eventually(ready, 10, 0.05).Should(BeClosed(), "Expected consul to be up and running")
	}

	cr.running = true
}

func (cr *ClusterRunner) NewClient() consuladapter.Client {
	client, err := api.NewClient(&api.Config{
		Address:    cr.Address(),
		Scheme:     cr.scheme,
		HttpClient: cfhttp.NewStreamingClient(),
	})
	Expect(err).NotTo(HaveOccurred())

	return consuladapter.NewConsulClient(client)
}

func (cr *ClusterRunner) WaitUntilReady() {
	client := cr.NewClient()
	catalog := client.Catalog()

	Eventually(func() error {
		_, qm, err := catalog.Nodes(nil)
		if err != nil {
			return err
		}
		if qm.KnownLeader && qm.LastIndex > 0 {
			return nil
		}
		return errors.New("not ready")
	}, 10, 100*time.Millisecond).Should(BeNil())
}

func (cr *ClusterRunner) Stop() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if !cr.running {
		return
	}

	for i := 0; i < cr.numNodes; i++ {
		stopSignal(cr.consulProcesses[i], 5*time.Second)
	}

	os.RemoveAll(cr.dataDir)
	os.RemoveAll(cr.configDir)
	cr.consulProcesses = nil
	cr.running = false
}

func (cr *ClusterRunner) ConsulCluster() string {
	urls := make([]string, cr.numNodes)
	for i := 0; i < cr.numNodes; i++ {
		urls[i] = fmt.Sprintf("%s://127.0.0.1:%d", cr.scheme, cr.startingPort+i*PortOffsetLength+PortOffsetHTTP)
	}

	return strings.Join(urls, ",")
}

func (cr *ClusterRunner) Address() string {
	return fmt.Sprintf("127.0.0.1:%d", cr.startingPort+PortOffsetHTTP)
}

func (cr *ClusterRunner) URL() string {
	return fmt.Sprintf("%s://%s", cr.scheme, cr.Address())
}

func (cr *ClusterRunner) Reset() error {
	client := cr.NewClient()

	sessions, _, err := client.Session().List(nil)
	if err == nil {
		for _, session := range sessions {
			_, err1 := client.Session().Destroy(session.ID, nil)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	services, err := client.Agent().Services()
	if err == nil {
		for _, service := range services {
			if service.Service == "consul" {
				continue
			}
			err1 := client.Agent().ServiceDeregister(service.ID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	checks, err := client.Agent().Checks()
	if err == nil {
		for _, check := range checks {
			err1 := client.Agent().CheckDeregister(check.CheckID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	_, err1 := client.KV().DeleteTree("", nil)

	return err1
}
<commit_msg>Remove prefix 'v' when retrieving version<commit_after>package consulrunner

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"time"

	"code.cloudfoundry.org/cfhttp"
	"code.cloudfoundry.org/consuladapter"
	"github.com/hashicorp/consul/api"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

type ClusterRunner struct {
	startingPort    int
	numNodes        int
	consulProcesses []ifrit.Process
	running         bool
	dataDir         string
	configDir       string
	scheme          string
	sessionTTL      time.Duration

	mutex *sync.RWMutex
}

const defaultDataDirPrefix = "consul_data"
const defaultConfigDirPrefix = "consul_config"

func NewClusterRunner(startingPort int, numNodes int, scheme string) *ClusterRunner {
	Expect(startingPort).To(BeNumerically(">", 0))
	Expect(startingPort).To(BeNumerically("<", 1<<16))
	Expect(numNodes).To(BeNumerically(">", 0))

	return &ClusterRunner{
		startingPort: startingPort,
		numNodes:     numNodes,
		scheme:       scheme,
		sessionTTL:   5 * time.Second,

		mutex: &sync.RWMutex{},
	}
}

func (cr *ClusterRunner) SessionTTL() time.Duration {
	return cr.sessionTTL
}

func (cr *ClusterRunner) ConsulVersion() string {
	cmd := exec.Command("consul", "-v")
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
	Expect(session.Out).To(gbytes.Say("Consul "))
	lines := strings.Split(string(session.Out.Contents()), "\n")
	versionLine := lines[0]
	//Consul in 'dev' mode does not contain the prefix 'v', only 'Consul 0.7.1-dev'
	return strings.TrimPrefix(strings.TrimPrefix(versionLine, "Consul "), "v")
}

func (cr *ClusterRunner) HasPerformanceFlag() bool {
	return !strings.HasPrefix(cr.ConsulVersion(), "0.6")
}

func (cr *ClusterRunner) Start() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if cr.running {
		return
	}

	tmpDir, err := ioutil.TempDir("", defaultDataDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.dataDir = tmpDir

	tmpDir, err = ioutil.TempDir("", defaultConfigDirPrefix)
	Expect(err).NotTo(HaveOccurred())
	cr.configDir = tmpDir

	cr.consulProcesses = make([]ifrit.Process, cr.numNodes)

	for i := 0; i < cr.numNodes; i++ {
		iStr := fmt.Sprintf("%d", i)
		nodeDataDir := path.Join(cr.dataDir, iStr)
		os.MkdirAll(nodeDataDir, 0700)

		configFilePath := writeConfigFile(
			cr.HasPerformanceFlag(),
			cr.configDir,
			nodeDataDir,
			iStr,
			cr.startingPort,
			i,
			cr.numNodes,
			cr.sessionTTL,
		)

		process := ginkgomon.Invoke(ginkgomon.New(ginkgomon.Config{
			Name:              fmt.Sprintf("consul_cluster[%d]", i),
			AnsiColorCode:     "35m",
			StartCheck:        "agent: Join completed.",
			StartCheckTimeout: 10 * time.Second,
			Command: exec.Command(
				"consul",
				"agent",
				"--log-level", "trace",
				"--config-file", configFilePath,
			),
		}))
		cr.consulProcesses[i] = process

		ready := process.Ready()
		Eventually(ready, 10, 0.05).Should(BeClosed(), "Expected consul to be up and running")
	}

	cr.running = true
}

func (cr *ClusterRunner) NewClient() consuladapter.Client {
	client, err := api.NewClient(&api.Config{
		Address:    cr.Address(),
		Scheme:     cr.scheme,
		HttpClient: cfhttp.NewStreamingClient(),
	})
	Expect(err).NotTo(HaveOccurred())

	return consuladapter.NewConsulClient(client)
}

func (cr *ClusterRunner) WaitUntilReady() {
	client := cr.NewClient()
	catalog := client.Catalog()

	Eventually(func() error {
		_, qm, err := catalog.Nodes(nil)
		if err != nil {
			return err
		}
		if qm.KnownLeader && qm.LastIndex > 0 {
			return nil
		}
		return errors.New("not ready")
	}, 10, 100*time.Millisecond).Should(BeNil())
}

func (cr *ClusterRunner) Stop() {
	cr.mutex.Lock()
	defer cr.mutex.Unlock()

	if !cr.running {
		return
	}

	for i := 0; i < cr.numNodes; i++ {
		stopSignal(cr.consulProcesses[i], 5*time.Second)
	}

	os.RemoveAll(cr.dataDir)
	os.RemoveAll(cr.configDir)
	cr.consulProcesses = nil
	cr.running = false
}

func (cr *ClusterRunner) ConsulCluster() string {
	urls := make([]string, cr.numNodes)
	for i := 0; i < cr.numNodes; i++ {
		urls[i] = fmt.Sprintf("%s://127.0.0.1:%d", cr.scheme, cr.startingPort+i*PortOffsetLength+PortOffsetHTTP)
	}

	return strings.Join(urls, ",")
}

func (cr *ClusterRunner) Address() string {
	return fmt.Sprintf("127.0.0.1:%d", cr.startingPort+PortOffsetHTTP)
}

func (cr *ClusterRunner) URL() string {
	return fmt.Sprintf("%s://%s", cr.scheme, cr.Address())
}

func (cr *ClusterRunner) Reset() error {
	client := cr.NewClient()

	sessions, _, err := client.Session().List(nil)
	if err == nil {
		for _, session := range sessions {
			_, err1 := client.Session().Destroy(session.ID, nil)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	services, err := client.Agent().Services()
	if err == nil {
		for _, service := range services {
			if service.Service == "consul" {
				continue
			}
			err1 := client.Agent().ServiceDeregister(service.ID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	checks, err := client.Agent().Checks()
	if err == nil {
		for _, check := range checks {
			err1 := client.Agent().CheckDeregister(check.CheckID)
			if err1 != nil {
				err = err1
			}
		}
	}

	if err != nil {
		return err
	}

	_, err1 := client.KV().DeleteTree("", nil)

	return err1
}
<|endoftext|>


<commit_before>package workflowhelpers

import (
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/commandstarter"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers/internal"
)

type remoteResource interface {
	Create()
	Destroy()
	ShouldRemain() bool
}

type ReproducibleTestSuiteSetup struct {
	config internal.TestSuiteConfig

	shortTimeout time.Duration
	longTimeout  time.Duration

	organizationName string
	spaceName        string

	TestUser  remoteResource
	TestSpace remoteResource

	regularUserContext UserContext
	adminUserContext   UserContext

	SkipSSLValidation bool

	isPersistent bool

	originalCfHomeDir string
	currentCfHomeDir  string
}

const RUNAWAY_QUOTA_MEM_LIMIT = "99999G"

func NewTestSuiteSetup(config internal.TestSuiteConfig) *ReproducibleTestSuiteSetup {
	testSpace := internal.NewRegularTestSpace(config, "10G")
	testUser := internal.NewTestUser(config, commandstarter.NewCommandStarter())
	adminUser := internal.NewAdminUser(config, commandstarter.NewCommandStarter())

	shortTimeout := config.GetScaledTimeout(1 * time.Minute)
	regularUserContext := NewUserContext(config.GetApiEndpoint(), testUser, testSpace, config.GetSkipSSLValidation(), shortTimeout)
	adminUserContext := NewUserContext(config.GetApiEndpoint(), adminUser, nil, config.GetSkipSSLValidation(), shortTimeout)

	return NewBaseTestSuiteSetup(config, testSpace, testUser, regularUserContext, adminUserContext)
}

func NewPersistentAppTestSuiteSetup(config internal.TestSuiteConfig) *ReproducibleTestSuiteSetup {
	testSpace := internal.NewPersistentAppTestSpace(config)
	testUser := internal.NewTestUser(config, commandstarter.NewCommandStarter())
	adminUser := internal.NewAdminUser(config, commandstarter.NewCommandStarter())

	shortTimeout := config.GetScaledTimeout(1 * time.Minute)
	regularUserContext := NewUserContext(config.GetApiEndpoint(), testUser, testSpace, config.GetSkipSSLValidation(), shortTimeout)
	adminUserContext := NewUserContext(config.GetApiEndpoint(), adminUser, nil, config.GetSkipSSLValidation(), shortTimeout)

	testSuiteSetup := NewBaseTestSuiteSetup(config, testSpace, testUser, regularUserContext, adminUserContext)
	testSuiteSetup.isPersistent = true

	return testSuiteSetup
}

func NewRunawayAppTestSuiteSetup(config internal.TestSuiteConfig) *ReproducibleTestSuiteSetup {
	testSpace := internal.NewRegularTestSpace(config, RUNAWAY_QUOTA_MEM_LIMIT)
	testUser := internal.NewTestUser(config, commandstarter.NewCommandStarter())
	adminUser := internal.NewAdminUser(config, commandstarter.NewCommandStarter())

	shortTimeout := config.GetScaledTimeout(1 * time.Minute)
	regularUserContext := NewUserContext(config.GetApiEndpoint(), testUser, testSpace, config.GetSkipSSLValidation(), shortTimeout)
	adminUserContext := NewUserContext(config.GetApiEndpoint(), adminUser, nil, config.GetSkipSSLValidation(), shortTimeout)

	return NewBaseTestSuiteSetup(config, testSpace, testUser, regularUserContext, adminUserContext)
}

func NewBaseTestSuiteSetup(config internal.TestSuiteConfig, testSpace, testUser remoteResource, regularUserContext, adminUserContext UserContext) *ReproducibleTestSuiteSetup {
	shortTimeout := config.GetScaledTimeout(1 * time.Minute)

	return &ReproducibleTestSuiteSetup{
		config: config,

		shortTimeout: shortTimeout,
		longTimeout:  config.GetScaledTimeout(5 * time.Minute),

		organizationName: generator.PrefixedRandomName(config.GetNamePrefix(), "ORG"),
		spaceName:        generator.PrefixedRandomName(config.GetNamePrefix(), "SPACE"),

		regularUserContext: regularUserContext,
		adminUserContext:   adminUserContext,

		isPersistent: false,
		TestSpace:    testSpace,
		TestUser:     testUser,
	}
}

func (testSetup ReproducibleTestSuiteSetup) ShortTimeout() time.Duration {
	return testSetup.shortTimeout
}

func (testSetup ReproducibleTestSuiteSetup) LongTimeout() time.Duration {
	return testSetup.longTimeout
}

func (testSetup *ReproducibleTestSuiteSetup) Setup() {
	AsUser(testSetup.AdminUserContext(), testSetup.shortTimeout, func() {
		testSetup.TestSpace.Create()
		testSetup.TestUser.Create()
		testSetup.regularUserContext.AddUserToSpace()
	})

	testSetup.originalCfHomeDir, testSetup.currentCfHomeDir = testSetup.regularUserContext.SetCfHomeDir()
	testSetup.regularUserContext.Login()
	testSetup.regularUserContext.TargetSpace()
}

func (testSetup *ReproducibleTestSuiteSetup) Teardown() {
	testSetup.regularUserContext.Logout()
	testSetup.regularUserContext.UnsetCfHomeDir(testSetup.originalCfHomeDir, testSetup.currentCfHomeDir)
	AsUser(testSetup.AdminUserContext(), testSetup.shortTimeout, func() {
		if !testSetup.TestUser.ShouldRemain() {
			testSetup.TestUser.Destroy()
		}

		if !testSetup.TestSpace.ShouldRemain() {
			testSetup.TestSpace.Destroy()
		}
	})
}

func (testSetup *ReproducibleTestSuiteSetup) AdminUserContext() UserContext {
	return testSetup.adminUserContext
}

func (testSetup *ReproducibleTestSuiteSetup) RegularUserContext() UserContext {
	return testSetup.regularUserContext
}
<commit_msg>Remove unused config from ReproducibleTestSetup struct<commit_after>package workflowhelpers

import (
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/commandstarter"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers/internal"
)

type remoteResource interface {
	Create()
	Destroy()
	ShouldRemain() bool
}

type ReproducibleTestSuiteSetup struct {
	shortTimeout time.Duration
	longTimeout  time.Duration

	organizationName string
	spaceName        string

	TestUser  remoteResource
	TestSpace remoteResource

	regularUserContext UserContext
	adminUserContext   UserContext

	SkipSSLValidation bool

	isPersistent bool

	originalCfHomeDir string
	currentCfHomeDir  string
}

const RUNAWAY_QUOTA_MEM_LIMIT = "99999G"

func NewTestSuiteSetup(config internal.TestSuiteConfig) *ReproducibleTestSuiteSetup {
	testSpace := internal.NewRegularTestSpace(config, "10G")
	testUser := internal.NewTestUser(config, commandstarter.NewCommandStarter())
	adminUser := internal.NewAdminUser(config, commandstarter.NewCommandStarter())

	shortTimeout := config.GetScaledTimeout(1 * time.Minute)
	regularUserContext := NewUserContext(config.GetApiEndpoint(), testUser, testSpace, config.GetSkipSSLValidation(), shortTimeout)
	adminUserContext := NewUserContext(config.GetApiEndpoint(), adminUser, nil, config.GetSkipSSLValidation(), shortTimeout)

	return NewBaseTestSuiteSetup(config, testSpace, testUser, regularUserContext, adminUserContext)
}

func NewPersistentAppTestSuiteSetup(config internal.TestSuiteConfig) *ReproducibleTestSuiteSetup {
	testSpace := internal.NewPersistentAppTestSpace(config)
	testUser := internal.NewTestUser(config, commandstarter.NewCommandStarter())
	adminUser := internal.NewAdminUser(config, commandstarter.NewCommandStarter())

	shortTimeout := config.GetScaledTimeout(1 * time.Minute)
	regularUserContext := NewUserContext(config.GetApiEndpoint(), testUser, testSpace, config.GetSkipSSLValidation(), shortTimeout)
	adminUserContext := NewUserContext(config.GetApiEndpoint(), adminUser, nil, config.GetSkipSSLValidation(), shortTimeout)

	testSuiteSetup := NewBaseTestSuiteSetup(config, testSpace, testUser, regularUserContext, adminUserContext)
	testSuiteSetup.isPersistent = true

	return testSuiteSetup
}

func NewRunawayAppTestSuiteSetup(config internal.TestSuiteConfig) *ReproducibleTestSuiteSetup {
	testSpace := internal.NewRegularTestSpace(config, RUNAWAY_QUOTA_MEM_LIMIT)
	testUser := internal.NewTestUser(config, commandstarter.NewCommandStarter())
	adminUser := internal.NewAdminUser(config, commandstarter.NewCommandStarter())

	shortTimeout := config.GetScaledTimeout(1 * time.Minute)
	regularUserContext := NewUserContext(config.GetApiEndpoint(), testUser, testSpace, config.GetSkipSSLValidation(), shortTimeout)
	adminUserContext := NewUserContext(config.GetApiEndpoint(), adminUser, nil, config.GetSkipSSLValidation(), shortTimeout)

	return NewBaseTestSuiteSetup(config, testSpace, testUser, regularUserContext, adminUserContext)
}

func NewBaseTestSuiteSetup(config internal.TestSuiteConfig, testSpace, testUser remoteResource, regularUserContext, adminUserContext UserContext) *ReproducibleTestSuiteSetup {
	shortTimeout := config.GetScaledTimeout(1 * time.Minute)

	return &ReproducibleTestSuiteSetup{
		shortTimeout: shortTimeout,
		longTimeout:  config.GetScaledTimeout(5 * time.Minute),

		organizationName: generator.PrefixedRandomName(config.GetNamePrefix(), "ORG"),
		spaceName:        generator.PrefixedRandomName(config.GetNamePrefix(), "SPACE"),

		regularUserContext: regularUserContext,
		adminUserContext:   adminUserContext,

		isPersistent: false,
		TestSpace:    testSpace,
		TestUser:     testUser,
	}
}

func (testSetup ReproducibleTestSuiteSetup) ShortTimeout() time.Duration {
	return testSetup.shortTimeout
}

func (testSetup ReproducibleTestSuiteSetup) LongTimeout() time.Duration {
	return testSetup.longTimeout
}

func (testSetup *ReproducibleTestSuiteSetup) Setup() {
	AsUser(testSetup.AdminUserContext(), testSetup.shortTimeout, func() {
		testSetup.TestSpace.Create()
		testSetup.TestUser.Create()
		testSetup.regularUserContext.AddUserToSpace()
	})

	testSetup.originalCfHomeDir, testSetup.currentCfHomeDir = testSetup.regularUserContext.SetCfHomeDir()
	testSetup.regularUserContext.Login()
	testSetup.regularUserContext.TargetSpace()
}

func (testSetup *ReproducibleTestSuiteSetup) Teardown() {
	testSetup.regularUserContext.Logout()
	testSetup.regularUserContext.UnsetCfHomeDir(testSetup.originalCfHomeDir, testSetup.currentCfHomeDir)
	AsUser(testSetup.AdminUserContext(), testSetup.shortTimeout, func() {
		if !testSetup.TestUser.ShouldRemain() {
			testSetup.TestUser.Destroy()
		}

		if !testSetup.TestSpace.ShouldRemain() {
			testSetup.TestSpace.Destroy()
		}
	})
}

func (testSetup *ReproducibleTestSuiteSetup) AdminUserContext() UserContext {
	return testSetup.adminUserContext
}

func (testSetup *ReproducibleTestSuiteSetup) RegularUserContext() UserContext {
	return testSetup.regularUserContext
}
<|endoftext|>


<commit_before>// Copyright 2016 CoreOS, Inc.
//
// 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 util

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"syscall"

	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/systemd"

	"github.com/vincent-petithory/dataurl"
)

var (
	ErrSchemeUnsupported = errors.New("unsupported source scheme")
	ErrPathNotAbsolute   = errors.New("path is not absolute")
	ErrNotFound          = errors.New("resource not found")
	ErrFailed            = errors.New("failed to fetch resource")
)

const (
	oemDevicePath = "/dev/disk/by-label/OEM" // Device link where oem partition is found.
	oemDirPath    = "/usr/share/oem"         // OEM dir within root fs to consider for pxe scenarios.
	oemMountPath  = "/mnt/oem"               // Mountpoint where oem partition is mounted when present.
)

// FetchResource fetches a resource given a URL. The supported schemes are http, data, and oem.
func FetchResource(l *log.Logger, u url.URL) ([]byte, error) {
	switch u.Scheme {
	case "http":
		client := NewHttpClient(l)
		data, status, err := client.Get(u.String())
		if err != nil {
			return nil, err
		}

		l.Debug("GET result: %s", http.StatusText(status))
		switch status {
		case http.StatusOK, http.StatusNoContent:
			return data, nil
		case http.StatusNotFound:
			return nil, ErrNotFound
		default:
			return nil, ErrFailed
		}

	case "data":
		url, err := dataurl.DecodeString(u.String())
		if err != nil {
			return nil, err
		}

		return url.Data, nil

	case "oem":
		path := filepath.Clean(u.Path)
		if !filepath.IsAbs(path) {
			l.Err("oem path is not absolute: %q", u.Path)
			return nil, ErrPathNotAbsolute
		}

		// check if present under oemDirPath, if so use it.
		absPath := filepath.Join(oemDirPath, path)
		data, err := ioutil.ReadFile(absPath)
		if os.IsNotExist(err) {
			l.Info("oem config not found in %q, trying %q",
				oemDirPath, oemMountPath)

			// try oemMountPath, requires mounting it.
			err := mountOEM(l)
			if err != nil {
				l.Err("failed to mount oem partition: %v", err)
				return nil, ErrFailed
			}

			absPath := filepath.Join(oemMountPath, path)
			data, err = ioutil.ReadFile(absPath)
			umountOEM(l)
		} else if err != nil {
			l.Err("failed to read oem config: %v", err)
			return nil, ErrFailed
		}

		return data, nil

	default:
		return nil, ErrSchemeUnsupported
	}
}

// mountOEM waits for the presence of and mounts the oem partition at oemMountPath.
func mountOEM(l *log.Logger) error {
	dev := []string{oemDevicePath}
	if err := systemd.WaitOnDevices(dev, "oem-cmdline"); err != nil {
		l.Err("failed to wait for oem device: %v", err)
		return err
	}

	if err := os.MkdirAll(oemMountPath, 0700); err != nil {
		l.Err("failed to create oem mount point: %v", err)
		return err
	}

	if err := l.LogOp(
		func() error {
			return syscall.Mount(dev[0], oemMountPath, "ext4", 0, "")
		},
		"mounting %q at %q", oemDevicePath, oemMountPath,
	); err != nil {
		return fmt.Errorf("failed to mount device %q at %q: %v",
			oemDevicePath, oemMountPath, err)
	}

	return nil
}

// umountOEM unmounts the oem partition at oemMountPath.
func umountOEM(l *log.Logger) {
	l.LogOp(
		func() error { return syscall.Unmount(oemMountPath, 0) },
		"unmounting %q", oemMountPath,
	)
}
<commit_msg>util/url: allow https urls<commit_after>// Copyright 2016 CoreOS, Inc.
//
// 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 util

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"syscall"

	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/systemd"

	"github.com/vincent-petithory/dataurl"
)

var (
	ErrSchemeUnsupported = errors.New("unsupported source scheme")
	ErrPathNotAbsolute   = errors.New("path is not absolute")
	ErrNotFound          = errors.New("resource not found")
	ErrFailed            = errors.New("failed to fetch resource")
)

const (
	oemDevicePath = "/dev/disk/by-label/OEM" // Device link where oem partition is found.
	oemDirPath    = "/usr/share/oem"         // OEM dir within root fs to consider for pxe scenarios.
	oemMountPath  = "/mnt/oem"               // Mountpoint where oem partition is mounted when present.
)

// FetchResource fetches a resource given a URL. The supported schemes are http, data, and oem.
func FetchResource(l *log.Logger, u url.URL) ([]byte, error) {
	switch u.Scheme {
	case "http", "https":
		client := NewHttpClient(l)
		data, status, err := client.Get(u.String())
		if err != nil {
			return nil, err
		}

		l.Debug("GET result: %s", http.StatusText(status))
		switch status {
		case http.StatusOK, http.StatusNoContent:
			return data, nil
		case http.StatusNotFound:
			return nil, ErrNotFound
		default:
			return nil, ErrFailed
		}

	case "data":
		url, err := dataurl.DecodeString(u.String())
		if err != nil {
			return nil, err
		}

		return url.Data, nil

	case "oem":
		path := filepath.Clean(u.Path)
		if !filepath.IsAbs(path) {
			l.Err("oem path is not absolute: %q", u.Path)
			return nil, ErrPathNotAbsolute
		}

		// check if present under oemDirPath, if so use it.
		absPath := filepath.Join(oemDirPath, path)
		data, err := ioutil.ReadFile(absPath)
		if os.IsNotExist(err) {
			l.Info("oem config not found in %q, trying %q",
				oemDirPath, oemMountPath)

			// try oemMountPath, requires mounting it.
			err := mountOEM(l)
			if err != nil {
				l.Err("failed to mount oem partition: %v", err)
				return nil, ErrFailed
			}

			absPath := filepath.Join(oemMountPath, path)
			data, err = ioutil.ReadFile(absPath)
			umountOEM(l)
		} else if err != nil {
			l.Err("failed to read oem config: %v", err)
			return nil, ErrFailed
		}

		return data, nil

	default:
		return nil, ErrSchemeUnsupported
	}
}

// mountOEM waits for the presence of and mounts the oem partition at oemMountPath.
func mountOEM(l *log.Logger) error {
	dev := []string{oemDevicePath}
	if err := systemd.WaitOnDevices(dev, "oem-cmdline"); err != nil {
		l.Err("failed to wait for oem device: %v", err)
		return err
	}

	if err := os.MkdirAll(oemMountPath, 0700); err != nil {
		l.Err("failed to create oem mount point: %v", err)
		return err
	}

	if err := l.LogOp(
		func() error {
			return syscall.Mount(dev[0], oemMountPath, "ext4", 0, "")
		},
		"mounting %q at %q", oemDevicePath, oemMountPath,
	); err != nil {
		return fmt.Errorf("failed to mount device %q at %q: %v",
			oemDevicePath, oemMountPath, err)
	}

	return nil
}

// umountOEM unmounts the oem partition at oemMountPath.
func umountOEM(l *log.Logger) {
	l.LogOp(
		func() error { return syscall.Unmount(oemMountPath, 0) },
		"unmounting %q", oemMountPath,
	)
}
<|endoftext|>


<commit_before>package aws

import (
	"github.com/hashicorp/terraform/terraform"
)

type ResourceProvider struct {
}

func (p *ResourceProvider) Configure(map[string]interface{}) error {
	return nil
}

func (p *ResourceProvider) Diff(
	s *terraform.ResourceState,
	c map[string]interface{}) (*terraform.ResourceDiff, error) {
	return &terraform.ResourceDiff{
		Attributes: map[string]*terraform.ResourceAttrDiff{
			"id": &terraform.ResourceAttrDiff{
				Old:         "",
				NewComputed: true,
				RequiresNew: true,
			},
			"created": &terraform.ResourceAttrDiff{
				Old: "false",
				New: "true",
			},
		},
	}, nil
}

func (p *ResourceProvider) Resources() []terraform.ResourceType {
	return []terraform.ResourceType{
		terraform.ResourceType{
			Name: "aws_instance",
		},
	}
}
<commit_msg>providers/aws: get tests passing<commit_after>package aws

import (
	"github.com/hashicorp/terraform/terraform"
)

type ResourceProvider struct {
}

func (p *ResourceProvider) Configure(*terraform.ResourceConfig) error {
	return nil
}

func (p *ResourceProvider) Diff(
	s *terraform.ResourceState,
	c *terraform.ResourceConfig) (*terraform.ResourceDiff, error) {
	return &terraform.ResourceDiff{
		Attributes: map[string]*terraform.ResourceAttrDiff{
			"id": &terraform.ResourceAttrDiff{
				Old:         "",
				NewComputed: true,
				RequiresNew: true,
			},
			"created": &terraform.ResourceAttrDiff{
				Old: "false",
				New: "true",
			},
		},
	}, nil
}

func (p *ResourceProvider) Resources() []terraform.ResourceType {
	return []terraform.ResourceType{
		terraform.ResourceType{
			Name: "aws_instance",
		},
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Primitive HTTP client. See RFC 2616.

package http

import (
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"strings"
	"url"
)

// A Client is an HTTP client. Its zero value (DefaultClient) is a usable client
// that uses DefaultTransport.
//
// The Client's Transport typically has internal state (cached
// TCP connections), so Clients should be reused instead of created as
// needed. Clients are safe for concurrent use by multiple goroutines.
//
// Client is not yet very configurable.
type Client struct {
	Transport RoundTripper // if nil, DefaultTransport is used

	// If CheckRedirect is not nil, the client calls it before
	// following an HTTP redirect. The arguments req and via
	// are the upcoming request and the requests made already,
	// oldest first. If CheckRedirect returns an error, the client
	// returns that error instead of issue the Request req.
	//
	// If CheckRedirect is nil, the Client uses its default policy,
	// which is to stop after 10 consecutive requests.
	CheckRedirect func(req *Request, via []*Request) os.Error
}

// DefaultClient is the default Client and is used by Get, Head, and Post.
var DefaultClient = &Client{}

// RoundTripper is an interface representing the ability to execute a
// single HTTP transaction, obtaining the Response for a given Request.
//
// A RoundTripper must be safe for concurrent use by multiple
// goroutines.
type RoundTripper interface {
	// RoundTrip executes a single HTTP transaction, returning
	// the Response for the request req.  RoundTrip should not
	// attempt to interpret the response.  In particular,
	// RoundTrip must return err == nil if it obtained a response,
	// regardless of the response's HTTP status code.  A non-nil
	// err should be reserved for failure to obtain a response.
	// Similarly, RoundTrip should not attempt to handle
	// higher-level protocol details such as redirects,
	// authentication, or cookies.
	//
	// RoundTrip should not modify the request, except for
	// consuming the Body.  The request's URL and Header fields
	// are guaranteed to be initialized.
	RoundTrip(*Request) (*Response, os.Error)
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Used in Send to implement io.ReadCloser by bundling together the
// bufio.Reader through which we read the response, and the underlying
// network connection.
type readClose struct {
	io.Reader
	io.Closer
}

// Do sends an HTTP request and returns an HTTP response, following
// policy (e.g. redirects, cookies, auth) as configured on the client.
//
// A non-nil response always contains a non-nil resp.Body.
//
// Callers should close resp.Body when done reading from it. If
// resp.Body is not closed, the Client's underlying RoundTripper
// (typically Transport) may not be able to re-use a persistent TCP
// connection to the server for a subsequent "keep-alive" request.
//
// Generally Get, Post, or PostForm will be used instead of Do.
func (c *Client) Do(req *Request) (resp *Response, err os.Error) {
	if req.Method == "GET" || req.Method == "HEAD" {
		return c.doFollowingRedirects(req)
	}
	return send(req, c.Transport)
}

// send issues an HTTP request.  Caller should close resp.Body when done reading from it.
func send(req *Request, t RoundTripper) (resp *Response, err os.Error) {
	if t == nil {
		t = DefaultTransport
		if t == nil {
			err = os.NewError("http: no Client.Transport or DefaultTransport")
			return
		}
	}

	if req.URL == nil {
		return nil, os.NewError("http: nil Request.URL")
	}

	// Most the callers of send (Get, Post, et al) don't need
	// Headers, leaving it uninitialized.  We guarantee to the
	// Transport that this has been initialized, though.
	if req.Header == nil {
		req.Header = make(Header)
	}

	info := req.URL.RawUserinfo
	if len(info) > 0 {
		if req.Header == nil {
			req.Header = make(Header)
		}
		req.Header.Set("Authorization", "Basic "+base64.URLEncoding.EncodeToString([]byte(info)))
	}
	return t.RoundTrip(req)
}

// True if the specified HTTP status code is one for which the Get utility should
// automatically redirect.
func shouldRedirect(statusCode int) bool {
	switch statusCode {
	case StatusMovedPermanently, StatusFound, StatusSeeOther, StatusTemporaryRedirect:
		return true
	}
	return false
}

// Get issues a GET to the specified URL.  If the response is one of the following
// redirect codes, Get follows the redirect, up to a maximum of 10 redirects:
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// Caller should close r.Body when done reading from it.
//
// Get is a convenience wrapper around DefaultClient.Get.
func Get(url string) (r *Response, err os.Error) {
	return DefaultClient.Get(url)
}

// Get issues a GET to the specified URL.  If the response is one of the
// following redirect codes, Get follows the redirect after calling the
// Client's CheckRedirect function.
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// Caller should close r.Body when done reading from it.
func (c *Client) Get(url string) (r *Response, err os.Error) {
	req, err := NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	return c.doFollowingRedirects(req)
}

func (c *Client) doFollowingRedirects(ireq *Request) (r *Response, err os.Error) {
	// TODO: if/when we add cookie support, the redirected request shouldn't
	// necessarily supply the same cookies as the original.
	var base *url.URL
	redirectChecker := c.CheckRedirect
	if redirectChecker == nil {
		redirectChecker = defaultCheckRedirect
	}
	var via []*Request

	req := ireq
	urlStr := "" // next relative or absolute URL to fetch (after first request)
	for redirect := 0; ; redirect++ {
		if redirect != 0 {
			req = new(Request)
			req.Method = ireq.Method
			req.Header = make(Header)
			req.URL, err = base.Parse(urlStr)
			if err != nil {
				break
			}
			if len(via) > 0 {
				// Add the Referer header.
				lastReq := via[len(via)-1]
				if lastReq.URL.Scheme != "https" {
					req.Header.Set("Referer", lastReq.URL.String())
				}

				err = redirectChecker(req, via)
				if err != nil {
					break
				}
			}
		}

		urlStr = req.URL.String()
		if r, err = send(req, c.Transport); err != nil {
			break
		}
		if shouldRedirect(r.StatusCode) {
			r.Body.Close()
			if urlStr = r.Header.Get("Location"); urlStr == "" {
				err = os.NewError(fmt.Sprintf("%d response missing Location header", r.StatusCode))
				break
			}
			base = req.URL
			via = append(via, req)
			continue
		}
		return
	}

	method := ireq.Method
	err = &url.Error{method[0:1] + strings.ToLower(method[1:]), urlStr, err}
	return
}

func defaultCheckRedirect(req *Request, via []*Request) os.Error {
	if len(via) >= 10 {
		return os.NewError("stopped after 10 redirects")
	}
	return nil
}

// Post issues a POST to the specified URL.
//
// Caller should close r.Body when done reading from it.
//
// Post is a wrapper around DefaultClient.Post
func Post(url string, bodyType string, body io.Reader) (r *Response, err os.Error) {
	return DefaultClient.Post(url, bodyType, body)
}

// Post issues a POST to the specified URL.
//
// Caller should close r.Body when done reading from it.
func (c *Client) Post(url string, bodyType string, body io.Reader) (r *Response, err os.Error) {
	req, err := NewRequest("POST", url, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", bodyType)
	return send(req, c.Transport)
}

// PostForm issues a POST to the specified URL,
// with data's keys and values urlencoded as the request body.
//
// Caller should close r.Body when done reading from it.
//
// PostForm is a wrapper around DefaultClient.PostForm
func PostForm(url string, data url.Values) (r *Response, err os.Error) {
	return DefaultClient.PostForm(url, data)
}

// PostForm issues a POST to the specified URL,
// with data's keys and values urlencoded as the request body.
//
// Caller should close r.Body when done reading from it.
func (c *Client) PostForm(url string, data url.Values) (r *Response, err os.Error) {
	return c.Post(url, "application/x-www-form-urlencoded", strings.NewReader(data.Encode()))
}

// Head issues a HEAD to the specified URL.  If the response is one of the
// following redirect codes, Head follows the redirect after calling the
// Client's CheckRedirect function.
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// Head is a wrapper around DefaultClient.Head
func Head(url string) (r *Response, err os.Error) {
	return DefaultClient.Head(url)
}

// Head issues a HEAD to the specified URL.  If the response is one of the
// following redirect codes, Head follows the redirect after calling the
// Client's CheckRedirect function.
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
func (c *Client) Head(url string) (r *Response, err os.Error) {
	req, err := NewRequest("HEAD", url, nil)
	if err != nil {
		return nil, err
	}
	return c.doFollowingRedirects(req)
}
<commit_msg>http: avoid panic caused by nil URL<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Primitive HTTP client. See RFC 2616.

package http

import (
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"strings"
	"url"
)

// A Client is an HTTP client. Its zero value (DefaultClient) is a usable client
// that uses DefaultTransport.
//
// The Client's Transport typically has internal state (cached
// TCP connections), so Clients should be reused instead of created as
// needed. Clients are safe for concurrent use by multiple goroutines.
//
// Client is not yet very configurable.
type Client struct {
	Transport RoundTripper // if nil, DefaultTransport is used

	// If CheckRedirect is not nil, the client calls it before
	// following an HTTP redirect. The arguments req and via
	// are the upcoming request and the requests made already,
	// oldest first. If CheckRedirect returns an error, the client
	// returns that error instead of issue the Request req.
	//
	// If CheckRedirect is nil, the Client uses its default policy,
	// which is to stop after 10 consecutive requests.
	CheckRedirect func(req *Request, via []*Request) os.Error
}

// DefaultClient is the default Client and is used by Get, Head, and Post.
var DefaultClient = &Client{}

// RoundTripper is an interface representing the ability to execute a
// single HTTP transaction, obtaining the Response for a given Request.
//
// A RoundTripper must be safe for concurrent use by multiple
// goroutines.
type RoundTripper interface {
	// RoundTrip executes a single HTTP transaction, returning
	// the Response for the request req.  RoundTrip should not
	// attempt to interpret the response.  In particular,
	// RoundTrip must return err == nil if it obtained a response,
	// regardless of the response's HTTP status code.  A non-nil
	// err should be reserved for failure to obtain a response.
	// Similarly, RoundTrip should not attempt to handle
	// higher-level protocol details such as redirects,
	// authentication, or cookies.
	//
	// RoundTrip should not modify the request, except for
	// consuming the Body.  The request's URL and Header fields
	// are guaranteed to be initialized.
	RoundTrip(*Request) (*Response, os.Error)
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Used in Send to implement io.ReadCloser by bundling together the
// bufio.Reader through which we read the response, and the underlying
// network connection.
type readClose struct {
	io.Reader
	io.Closer
}

// Do sends an HTTP request and returns an HTTP response, following
// policy (e.g. redirects, cookies, auth) as configured on the client.
//
// A non-nil response always contains a non-nil resp.Body.
//
// Callers should close resp.Body when done reading from it. If
// resp.Body is not closed, the Client's underlying RoundTripper
// (typically Transport) may not be able to re-use a persistent TCP
// connection to the server for a subsequent "keep-alive" request.
//
// Generally Get, Post, or PostForm will be used instead of Do.
func (c *Client) Do(req *Request) (resp *Response, err os.Error) {
	if req.Method == "GET" || req.Method == "HEAD" {
		return c.doFollowingRedirects(req)
	}
	return send(req, c.Transport)
}

// send issues an HTTP request.  Caller should close resp.Body when done reading from it.
func send(req *Request, t RoundTripper) (resp *Response, err os.Error) {
	if t == nil {
		t = DefaultTransport
		if t == nil {
			err = os.NewError("http: no Client.Transport or DefaultTransport")
			return
		}
	}

	if req.URL == nil {
		return nil, os.NewError("http: nil Request.URL")
	}

	// Most the callers of send (Get, Post, et al) don't need
	// Headers, leaving it uninitialized.  We guarantee to the
	// Transport that this has been initialized, though.
	if req.Header == nil {
		req.Header = make(Header)
	}

	info := req.URL.RawUserinfo
	if len(info) > 0 {
		req.Header.Set("Authorization", "Basic "+base64.URLEncoding.EncodeToString([]byte(info)))
	}
	return t.RoundTrip(req)
}

// True if the specified HTTP status code is one for which the Get utility should
// automatically redirect.
func shouldRedirect(statusCode int) bool {
	switch statusCode {
	case StatusMovedPermanently, StatusFound, StatusSeeOther, StatusTemporaryRedirect:
		return true
	}
	return false
}

// Get issues a GET to the specified URL.  If the response is one of the following
// redirect codes, Get follows the redirect, up to a maximum of 10 redirects:
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// Caller should close r.Body when done reading from it.
//
// Get is a convenience wrapper around DefaultClient.Get.
func Get(url string) (r *Response, err os.Error) {
	return DefaultClient.Get(url)
}

// Get issues a GET to the specified URL.  If the response is one of the
// following redirect codes, Get follows the redirect after calling the
// Client's CheckRedirect function.
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// Caller should close r.Body when done reading from it.
func (c *Client) Get(url string) (r *Response, err os.Error) {
	req, err := NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	return c.doFollowingRedirects(req)
}

func (c *Client) doFollowingRedirects(ireq *Request) (r *Response, err os.Error) {
	// TODO: if/when we add cookie support, the redirected request shouldn't
	// necessarily supply the same cookies as the original.
	var base *url.URL
	redirectChecker := c.CheckRedirect
	if redirectChecker == nil {
		redirectChecker = defaultCheckRedirect
	}
	var via []*Request

	if ireq.URL == nil {
		return nil, os.NewError("http: nil Request.URL")
	}

	req := ireq
	urlStr := "" // next relative or absolute URL to fetch (after first request)
	for redirect := 0; ; redirect++ {
		if redirect != 0 {
			req = new(Request)
			req.Method = ireq.Method
			req.Header = make(Header)
			req.URL, err = base.Parse(urlStr)
			if err != nil {
				break
			}
			if len(via) > 0 {
				// Add the Referer header.
				lastReq := via[len(via)-1]
				if lastReq.URL.Scheme != "https" {
					req.Header.Set("Referer", lastReq.URL.String())
				}

				err = redirectChecker(req, via)
				if err != nil {
					break
				}
			}
		}

		urlStr = req.URL.String()
		if r, err = send(req, c.Transport); err != nil {
			break
		}
		if shouldRedirect(r.StatusCode) {
			r.Body.Close()
			if urlStr = r.Header.Get("Location"); urlStr == "" {
				err = os.NewError(fmt.Sprintf("%d response missing Location header", r.StatusCode))
				break
			}
			base = req.URL
			via = append(via, req)
			continue
		}
		return
	}

	method := ireq.Method
	err = &url.Error{method[0:1] + strings.ToLower(method[1:]), urlStr, err}
	return
}

func defaultCheckRedirect(req *Request, via []*Request) os.Error {
	if len(via) >= 10 {
		return os.NewError("stopped after 10 redirects")
	}
	return nil
}

// Post issues a POST to the specified URL.
//
// Caller should close r.Body when done reading from it.
//
// Post is a wrapper around DefaultClient.Post
func Post(url string, bodyType string, body io.Reader) (r *Response, err os.Error) {
	return DefaultClient.Post(url, bodyType, body)
}

// Post issues a POST to the specified URL.
//
// Caller should close r.Body when done reading from it.
func (c *Client) Post(url string, bodyType string, body io.Reader) (r *Response, err os.Error) {
	req, err := NewRequest("POST", url, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", bodyType)
	return send(req, c.Transport)
}

// PostForm issues a POST to the specified URL,
// with data's keys and values urlencoded as the request body.
//
// Caller should close r.Body when done reading from it.
//
// PostForm is a wrapper around DefaultClient.PostForm
func PostForm(url string, data url.Values) (r *Response, err os.Error) {
	return DefaultClient.PostForm(url, data)
}

// PostForm issues a POST to the specified URL,
// with data's keys and values urlencoded as the request body.
//
// Caller should close r.Body when done reading from it.
func (c *Client) PostForm(url string, data url.Values) (r *Response, err os.Error) {
	return c.Post(url, "application/x-www-form-urlencoded", strings.NewReader(data.Encode()))
}

// Head issues a HEAD to the specified URL.  If the response is one of the
// following redirect codes, Head follows the redirect after calling the
// Client's CheckRedirect function.
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// Head is a wrapper around DefaultClient.Head
func Head(url string) (r *Response, err os.Error) {
	return DefaultClient.Head(url)
}

// Head issues a HEAD to the specified URL.  If the response is one of the
// following redirect codes, Head follows the redirect after calling the
// Client's CheckRedirect function.
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
func (c *Client) Head(url string) (r *Response, err os.Error) {
	req, err := NewRequest("HEAD", url, nil)
	if err != nil {
		return nil, err
	}
	return c.doFollowingRedirects(req)
}
<|endoftext|>