<commit_before>// Copyright (c) 2016 Pani Networks
// 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 (
	"bytes"
	"crypto/rand"
	"flag"
	"fmt"
	_ "github.com/cgilmour/maxopen"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"
)

var (
	port         = flag.Uint("port", 8080, "Port to listen for HTTP requests")
	endpoint     = flag.String("endpoint", "/endpoint", "Name of the endpoint that HTTP requests can be sent to")
	responseSize = flag.Uint("response-size", 512, "Number of bytes of random data to send in HTTP responses")
)

func main() {
	flag.Parse()

	randomData := make([]byte, *responseSize)
	if _, err := rand.Read(randomData); err != nil {
		fmt.Fprintf(os.Stderr, "Initialization error: Failed to read %d bytes of random data: %s", *responseSize, err)
		return
	}

	http.HandleFunc(*endpoint, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", fmt.Sprintf("%d", *responseSize))
		w.WriteHeader(http.StatusOK)
		io.Copy(w, bytes.NewReader(randomData))
	})

	server := http.Server{Addr: fmt.Sprintf(":%d", *port)}
	server.SetKeepAlivesEnabled(false)
	errCh := make(chan struct{})
	bc := &byteCounter{}
	cc := &connCounter{}
	go func() {
		defer close(errCh)
		ln, err := net.Listen("tcp", server.Addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error from net.Listen: %s\n", err)
			return
		}
		err = server.Serve(&timedListener{Listener: ln, bc: bc, cc: cc})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error from server.Serve: %s\n", err)
			return
		}
	}()
	go func() {
		tickCh := time.Tick(100 * time.Millisecond)
		idle := false
		idleTime := time.Time{}
		for range tickCh {
			rx, tx := bc.Sample()
			n := cc.Sample()
			if rx == 0 && tx == 0 && n == 0 {
				if !idle {
					idleTime = time.Now()
				}
				idle = true
				continue
			}
			if idle {
				fmt.Println("Idle for", time.Now().Sub(idleTime))
				idle = false
			}
			fmt.Println("connections:", n, "bytes sent:", tx, "bytes received:", rx)
		}
	}()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, os.Kill)

	select {
	case <-sigCh:
	case <-errCh:
	}
}

type timedListener struct {
	net.Listener
	bc *byteCounter
	cc *connCounter
}

func (tl *timedListener) Accept() (net.Conn, error) {
	conn, err := tl.Listener.Accept()
	tl.cc.m.Lock()
	tl.cc.n++
	tl.cc.m.Unlock()
	if err != nil {
		fmt.Println("Accept() error was", err)
	}
	return &timedConn{Conn: conn, openTS: time.Now(), bc: tl.bc}, err
}

type connCounter struct {
	m sync.Mutex
	n int
}

func (c *connCounter) Sample() int {
	c.m.Lock()
	n := c.n
	c.n = 0
	c.m.Unlock()
	return n
}

type timedConn struct {
	net.Conn
	openTS  time.Time
	closeTS time.Time
	bc      *byteCounter
}

func (tc *timedConn) Read(b []byte) (int, error) {
	n, err := tc.Conn.Read(b)
	tc.bc.m.Lock()
	tc.bc.rx += n
	tc.bc.m.Unlock()
	return n, err
}

func (tc *timedConn) Write(b []byte) (int, error) {
	n, err := tc.Conn.Write(b)
	tc.bc.m.Lock()
	tc.bc.tx += n
	tc.bc.m.Unlock()
	return n, err
}

func (tc timedConn) Close() error {
	tc.closeTS = time.Now()
	// fmt.Printf("Opened: %s, Closed: %s, Delta: %s\n", tc.openTS, tc.closeTS, tc.closeTS.Sub(tc.openTS))
	return tc.Conn.Close()
}

type byteCounter struct {
	m  sync.Mutex
	rx int
	tx int
}

func (bc *byteCounter) Sample() (int, int) {
	bc.m.Lock()
	rx, tx := bc.rx, bc.tx
	bc.rx, bc.tx = 0, 0
	bc.m.Unlock()
	return rx, tx
}
<commit_msg>Additional comments in test-http-server.<commit_after>// Copyright (c) 2016 Pani Networks
// 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 (
	"bytes"
	"crypto/rand"
	"flag"
	"fmt"
	_ "github.com/cgilmour/maxopen"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"
)

var (
	port         = flag.Uint("port", 8080, "Port to listen for HTTP requests")
	endpoint     = flag.String("endpoint", "/endpoint", "Name of the endpoint that HTTP requests can be sent to")
	responseSize = flag.Uint("response-size", 512, "Number of bytes of random data to send in HTTP responses")
)

func main() {
	flag.Parse()

	// Generate 'responseSize' bytes of random data.
	// This can be slow for large response sizes, delaying startup.
	randomData := make([]byte, *responseSize)
	if _, err := rand.Read(randomData); err != nil {
		fmt.Fprintf(os.Stderr, "Initialization error: Failed to read %d bytes of random data: %s", *responseSize, err)
		return
	}

	// a simple HandleFunc that responds to a request with a Content-Length (to prevent chunk encoding),
	// an OK status and the contents of randomData.
	http.HandleFunc(*endpoint, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", fmt.Sprintf("%d", *responseSize))
		w.WriteHeader(http.StatusOK)
		io.Copy(w, bytes.NewReader(randomData))
	})

	// Initialize a server
	server := http.Server{Addr: fmt.Sprintf(":%d", *port)}
	server.SetKeepAlivesEnabled(false)
	errCh := make(chan struct{})

	// Counters for reporting server activity at regular intervals.
	bc := &byteCounter{}
	cc := &connCounter{}

	// try to start the server.
	go func() {
		defer close(errCh)
		ln, err := net.Listen("tcp", server.Addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error from net.Listen: %s\n", err)
			return
		}
		err = server.Serve(&timedListener{Listener: ln, bc: bc, cc: cc})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error from server.Serve: %s\n", err)
			return
		}
	}()

	// check every 100ms if there's something to report.
	go func() {
		tickCh := time.Tick(100 * time.Millisecond)
		idle := false
		idleTime := time.Time{}
		for range tickCh {
			rx, tx := bc.Sample()
			n := cc.Sample()
			if rx == 0 && tx == 0 && n == 0 {
				// transition to idle state
				if !idle {
					idleTime = time.Now()
				}
				idle = true
				continue
			}
			if idle {
				// transition from idle state
				fmt.Println("Idle for", time.Now().Sub(idleTime))
				idle = false
			}
			fmt.Println("connections:", n, "bytes sent:", tx, "bytes received:", rx)
		}
	}()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, os.Kill)

	// stop on a ^C or an error while initializing or running the server.
	select {
	case <-sigCh:
	case <-errCh:
	}
}

// this is a net.Listener that captures information for connections that arrive on the server.
type timedListener struct {
	net.Listener
	bc *byteCounter
	cc *connCounter
}

func (tl *timedListener) Accept() (net.Conn, error) {
	conn, err := tl.Listener.Accept()
	// increase connection count for this sample
	tl.cc.m.Lock()
	tl.cc.n++
	tl.cc.m.Unlock()
	if err != nil {
		fmt.Println("Accept() error was", err)
	}
	// wrap the connection, so it can capture additional byteCounter information
	return &timedConn{Conn: conn, openTS: time.Now(), bc: tl.bc}, err
}

type connCounter struct {
	m sync.Mutex
	n int
}

// Sample() reports and resets the number of connections since it was last called.
func (c *connCounter) Sample() int {
	c.m.Lock()
	n := c.n
	c.n = 0
	c.m.Unlock()
	return n
}

// a timedConnection captures when the connection was opened/closed, and increases counters during
// Read() and Write() calls.
type timedConn struct {
	net.Conn
	openTS  time.Time
	closeTS time.Time
	bc      *byteCounter
}

func (tc *timedConn) Read(b []byte) (int, error) {
	n, err := tc.Conn.Read(b)
	tc.bc.m.Lock()
	tc.bc.rx += n
	tc.bc.m.Unlock()
	return n, err
}

func (tc *timedConn) Write(b []byte) (int, error) {
	n, err := tc.Conn.Write(b)
	tc.bc.m.Lock()
	tc.bc.tx += n
	tc.bc.m.Unlock()
	return n, err
}

func (tc timedConn) Close() error {
	tc.closeTS = time.Now()
	// fmt.Printf("Opened: %s, Closed: %s, Delta: %s\n", tc.openTS, tc.closeTS, tc.closeTS.Sub(tc.openTS))
	return tc.Conn.Close()
}

// The byteCounter captures bytes read and written for the timedConn it is related to.
type byteCounter struct {
	m  sync.Mutex
	rx int
	tx int
}

// Sample() reports and resets the number of bytes read and written since it was last called
func (bc *byteCounter) Sample() (int, int) {
	bc.m.Lock()
	rx, tx := bc.rx, bc.tx
	bc.rx, bc.tx = 0, 0
	bc.m.Unlock()
	return rx, tx
}
<|endoftext|>


<commit_before>package watchdog

import (
	"errors"
	"sync"
	"time"

	schematypes "github.com/taskcluster/go-schematypes"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/plugins"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
)

// default timeout, setting to 45 minutes is safe, notice that downloading
// images could take 10-15 minutes in some extreme cases.
const defaultTimeout = 45 // minutes

type provider struct {
	plugins.PluginProviderBase
}

type plugin struct {
	plugins.PluginBase
	Timeout     time.Duration
	Environment runtime.Environment
	Monitor     runtime.Monitor
	m           sync.Mutex
	timer       *time.Timer
	done        atomics.Once
	running     int
}

type taskPlugin struct {
	plugins.TaskPluginBase
	Plugin  *plugin
	started atomics.Once
	stopped atomics.Once
}

func init() {
	plugins.Register("watchdog", provider{})
}

func (provider) ConfigSchema() schematypes.Schema {
	return configSchema
}

func (provider) NewPlugin(options plugins.PluginOptions) (plugins.Plugin, error) {
	var c config
	schematypes.MustValidateAndMap(configSchema, options.Config, &c)

	// Ensure that a default value is set
	if c.Timeout == 0 {
		debug("watchdog configured with default timeout")
		c.Timeout = defaultTimeout * time.Minute
	}

	p := &plugin{
		Timeout:     c.Timeout,
		Environment: *options.Environment,
		Monitor:     options.Monitor,
		timer:       time.NewTimer(c.Timeout),
	}
	go p.waitForTimeout()

	return p, nil
}

func (p *plugin) waitForTimeout() {
	select {
	case <-p.done.Done():
	case <-p.timer.C:
		var reported atomics.Once
		go reported.Do(func() {
			p.Monitor.ReportError(errors.New("watchdog: worker timeout exceeded"))
		})

		// Call stopNow within 30s, even if the ReportError blocks and fails...
		// this is just extra hardening...
		select {
		case <-reported.Done():
		case <-time.After(30 * time.Second):
		}
		p.Environment.Worker.StopNow()
	}
}

func (p *plugin) Touch() {
	p.Monitor.Debug("watchdog touched")

	p.m.Lock()
	defer p.m.Unlock()

	p.timer.Stop()
	if p.running == 0 {
		p.timer.Reset(p.Timeout)
	} else {
		p.Monitor.Debugf("watchdog timer paused, currently running: %d", p.running)
	}
}

func (p *plugin) Increment() {
	p.m.Lock()
	p.running++
	p.Monitor.Debugf("watchdog running task incremented to: %d", p.running)
	p.m.Unlock()

	p.Touch()
}

func (p *plugin) Decrement() {
	p.m.Lock()
	p.running--
	p.Monitor.Debugf("watchdog running task decremented to: %d", p.running)
	p.m.Unlock()

	p.Touch()
}

func (p *plugin) ReportIdle(time.Duration) {
	p.Touch()
}

func (p *plugin) NewTaskPlugin(plugins.TaskPluginOptions) (plugins.TaskPlugin, error) {
	p.Touch()
	return &taskPlugin{
		Plugin: p,
	}, nil
}

func (p *plugin) Dispose() error {
	p.m.Lock()
	defer p.m.Unlock()

	// Stop waitForTimeout
	p.done.Do(nil)
	p.timer.Stop()

	return nil
}

func (p *taskPlugin) BuildSandbox(sandboxBuilder engines.SandboxBuilder) error {
	p.Plugin.Touch()
	return nil
}

func (p *taskPlugin) Started(sandbox engines.Sandbox) error {
	// Increment running, and remember this
	// The use of atomics.Once is just to ensure robustness... worker should never
	// call started twice. But we have to remember if started was called in
	// exception and dispose, otherwise, we don't know if we should decrement
	p.started.Do(p.Plugin.Increment)
	return nil
}

func (p *taskPlugin) Stopped(result engines.ResultSet) (bool, error) {
	if p.started.IsDone() {
		p.stopped.Do(p.Plugin.Decrement)
	} else {
		// if started wasn't called then that's contract violation
		panic("TaskPlugin.Stopped() called before TaskPlugin.Started()")
	}
	return true, nil
}

func (p *taskPlugin) Finished(success bool) error {
	p.Plugin.Touch()
	return nil
}

func (p *taskPlugin) Exception(reason runtime.ExceptionReason) error {
	// decrement only if started was called
	if p.started.IsDone() {
		p.stopped.Do(p.Plugin.Decrement)
	}
	p.Plugin.Touch()
	return nil
}

func (p *taskPlugin) Dispose() error {
	// decrement only if started was called
	if p.started.IsDone() {
		p.stopped.Do(p.Plugin.Decrement)
	}
	p.Plugin.Touch()
	return nil
}
<commit_msg>Added stacktrace to watchdog, fixing #327<commit_after>package watchdog

import (
	"errors"
	goruntime "runtime"
	"strings"
	"sync"
	"time"

	schematypes "github.com/taskcluster/go-schematypes"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/plugins"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
)

// default timeout, setting to 45 minutes is safe, notice that downloading
// images could take 10-15 minutes in some extreme cases.
const defaultTimeout = 45 // minutes

type provider struct {
	plugins.PluginProviderBase
}

type plugin struct {
	plugins.PluginBase
	Timeout     time.Duration
	Environment runtime.Environment
	Monitor     runtime.Monitor
	m           sync.Mutex
	timer       *time.Timer
	done        atomics.Once
	running     int
}

type taskPlugin struct {
	plugins.TaskPluginBase
	Plugin  *plugin
	started atomics.Once
	stopped atomics.Once
}

func init() {
	plugins.Register("watchdog", provider{})
}

func (provider) ConfigSchema() schematypes.Schema {
	return configSchema
}

func (provider) NewPlugin(options plugins.PluginOptions) (plugins.Plugin, error) {
	var c config
	schematypes.MustValidateAndMap(configSchema, options.Config, &c)

	// Ensure that a default value is set
	if c.Timeout == 0 {
		debug("watchdog configured with default timeout")
		c.Timeout = defaultTimeout * time.Minute
	}

	p := &plugin{
		Timeout:     c.Timeout,
		Environment: *options.Environment,
		Monitor:     options.Monitor,
		timer:       time.NewTimer(c.Timeout),
	}
	go p.waitForTimeout()

	return p, nil
}

// allStacks returns a dump of stacks for all goroutines
func allStacks() string {
	b := make([]byte, 4096)
	N := goruntime.Stack(b, true)
	for N >= len(b) {
		b = make([]byte, 2*len(b))
		N = goruntime.Stack(b, true)
	}
	return string(b[:N])
}

func (p *plugin) waitForTimeout() {
	select {
	case <-p.done.Done():
	case <-p.timer.C:
		var reported atomics.Once
		go reported.Do(func() {
			// Create a dump of all stacks, and prefix with watch-stackdump, so that
			// it's obvious that this is an intentional stack dump and not the result
			// of some panic that was caught and returned as an error.
			prefix := "watchdog-stackdump: "
			stacks := prefix + strings.Join(strings.Split(allStacks(), "\n"), "\n"+prefix)
			p.Monitor.ReportError(errors.New("watchdog: worker timeout exceeded"), stacks)
		})

		// Call stopNow within 30s, even if the ReportError blocks and fails...
		// this is just extra hardening...
		select {
		case <-reported.Done():
		case <-time.After(30 * time.Second):
		}
		p.Environment.Worker.StopNow()
	}
}

func (p *plugin) Touch() {
	p.Monitor.Debug("watchdog touched")

	p.m.Lock()
	defer p.m.Unlock()

	p.timer.Stop()
	if p.running == 0 {
		p.timer.Reset(p.Timeout)
	} else {
		p.Monitor.Debugf("watchdog timer paused, currently running: %d", p.running)
	}
}

func (p *plugin) Increment() {
	p.m.Lock()
	p.running++
	p.Monitor.Debugf("watchdog running task incremented to: %d", p.running)
	p.m.Unlock()

	p.Touch()
}

func (p *plugin) Decrement() {
	p.m.Lock()
	p.running--
	p.Monitor.Debugf("watchdog running task decremented to: %d", p.running)
	p.m.Unlock()

	p.Touch()
}

func (p *plugin) ReportIdle(time.Duration) {
	p.Touch()
}

func (p *plugin) NewTaskPlugin(plugins.TaskPluginOptions) (plugins.TaskPlugin, error) {
	p.Touch()
	return &taskPlugin{
		Plugin: p,
	}, nil
}

func (p *plugin) Dispose() error {
	p.m.Lock()
	defer p.m.Unlock()

	// Stop waitForTimeout
	p.done.Do(nil)
	p.timer.Stop()

	return nil
}

func (p *taskPlugin) BuildSandbox(sandboxBuilder engines.SandboxBuilder) error {
	p.Plugin.Touch()
	return nil
}

func (p *taskPlugin) Started(sandbox engines.Sandbox) error {
	// Increment running, and remember this
	// The use of atomics.Once is just to ensure robustness... worker should never
	// call started twice. But we have to remember if started was called in
	// exception and dispose, otherwise, we don't know if we should decrement
	p.started.Do(p.Plugin.Increment)
	return nil
}

func (p *taskPlugin) Stopped(result engines.ResultSet) (bool, error) {
	if p.started.IsDone() {
		p.stopped.Do(p.Plugin.Decrement)
	} else {
		// if started wasn't called then that's contract violation
		panic("TaskPlugin.Stopped() called before TaskPlugin.Started()")
	}
	return true, nil
}

func (p *taskPlugin) Finished(success bool) error {
	p.Plugin.Touch()
	return nil
}

func (p *taskPlugin) Exception(reason runtime.ExceptionReason) error {
	// decrement only if started was called
	if p.started.IsDone() {
		p.stopped.Do(p.Plugin.Decrement)
	}
	p.Plugin.Touch()
	return nil
}

func (p *taskPlugin) Dispose() error {
	// decrement only if started was called
	if p.started.IsDone() {
		p.stopped.Do(p.Plugin.Decrement)
	}
	p.Plugin.Touch()
	return nil
}
<|endoftext|>


<commit_before>// Package server provides a sample HTTP/Websocket server which registers
// itself in consul using one or more url prefixes to demonstrate and
// test the automatic fabio routing table update.
//
// During startup the server performs the following steps:
//
// * Add a handler for each prefix which provides a unique
//   response for that instance and endpoint
// * Add a `/health` handler for the consul health check
// * Register the service in consul with the listen address,
//   a health check under the given name and with one `urlprefix-`
//   tag per prefix
// * Install a signal handler to deregister the service on exit
//
// If the protocol is set to "ws" the registered endpoints function
// as websocket echo servers.
//
// Example:
//
//   # http server
//   ./server -addr 127.0.0.1:5000 -name svc-a -prefix /foo,/bar
//   ./server -addr 127.0.0.1:5001 -name svc-b -prefix /baz,/bar
//
//   # https server
//   ./server -addr 127.0.0.1:5000 -name svc-a -proto https -certFile ... -keyFile ... -prefix /foo
//   ./server -addr 127.0.0.1:5000 -name svc-a -proto https -certFile ... -keyFile ... -prefix "/foo tlsskipverify=true"
//
//   # websocket server
//   ./server -addr 127.0.0.1:6000 -name ws-a -proto ws -prefix /echo1,/echo2
//
//   # tcp server
//   ./server -addr 127.0.0.1:7000 -name tcp-a -proto tcp -prefix :1234
//
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/net/websocket"

	"github.com/fabiolb/fabio/proxy/tcp"
	"github.com/hashicorp/consul/api"
)

type Args struct {
	addr     string
	consul   string
	name     string
	prefix   string
	proto    string
	rawtags  string
	token    string
	certFile string
	keyFile  string
	status   int
	tags     []string
}

func main() {
	var args Args
	flag.StringVar(&args.addr, "addr", "127.0.0.1:5000", "host:port of the service")
	flag.StringVar(&args.consul, "consul", "127.0.0.1:8500", "host:port of the consul agent")
	flag.StringVar(&args.name, "name", filepath.Base(os.Args[0]), "name of the service")
	flag.StringVar(&args.prefix, "prefix", "", "comma-sep list of 'host/path' or ':port' prefixes to register")
	flag.StringVar(&args.proto, "proto", "http", "protocol for endpoints: http, ws or tcp")
	flag.StringVar(&args.rawtags, "tags", "", "additional tags to register in consul")
	flag.StringVar(&args.token, "token", "", "consul ACL token")
	flag.StringVar(&args.certFile, "cert", "", "path to cert file")
	flag.StringVar(&args.keyFile, "key", "", "path to key file")
	flag.IntVar(&args.status, "status", http.StatusOK, "http status code")
	flag.Parse()

	if args.prefix == "" {
		flag.Usage()
		os.Exit(1)
	}

	if (args.proto == "https" || args.proto == "wss") && args.certFile == "" {
		log.Fatalf("Proto %s requires a certificate. Please provide -cert/-key", args.proto)
	}

	if args.rawtags != "" {
		args.tags = strings.Split(args.rawtags, ",")
	}

	type server interface {
		ListenAndServe() error
		ListenAndServeTLS(certFile, keyFile string) error
	}

	var srv server
	var tags []string
	var check *api.AgentServiceCheck
	switch args.proto {
	case "http", "https", "ws", "wss":
		srv, tags, check = newHTTPServer(args)
	case "tcp":
		srv, tags, check = newTCPServer(args)
	default:
		log.Fatal("Invalid protocol ", args.proto)
	}

	// start server
	go func() {
		var err error
		if args.certFile != "" {
			err = srv.ListenAndServeTLS(args.certFile, args.keyFile)
		} else {
			err = srv.ListenAndServe()
		}
		if err != nil {
			log.Fatal(err)
		}
	}()

	// get host and port as string/int
	host, portstr, err := net.SplitHostPort(args.addr)
	if err != nil {
		log.Fatal(err)
	}
	port, err := strconv.Atoi(portstr)
	if err != nil {
		log.Fatal(err)
	}

	// register service with health check
	serviceID := args.name + "-" + args.addr
	service := &api.AgentServiceRegistration{
		ID:      serviceID,
		Name:    args.name,
		Port:    port,
		Address: host,
		Tags:    tags,
		Check:   check,
	}

	config := &api.Config{Address: args.consul, Scheme: "http", Token: args.token}
	client, err := api.NewClient(config)
	if err != nil {
		log.Fatal(err)
	}

	if err := client.Agent().ServiceRegister(service); err != nil {
		log.Fatal(err)
	}
	log.Printf("Registered %s service %q in consul with tags %q", args.proto, args.name, strings.Join(tags, ","))

	// run until we get a signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt, os.Kill)
	<-quit

	// deregister service
	if err := client.Agent().ServiceDeregister(serviceID); err != nil {
		log.Fatal(err)
	}
	log.Printf("Deregistered service %q in consul", args.name)
}

func newHTTPServer(args Args) (*http.Server, []string, *api.AgentServiceCheck) {
	addr, proto, tags := args.addr, args.proto, args.tags

	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "OK")
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "not found", 404)
		log.Printf("%s -> 404", r.URL)
	})

	prefixes := strings.Split(args.prefix, ",")
	for _, p := range prefixes {
		uri := strings.Fields(p)[0]
		switch proto {
		case "http", "https":
			mux.HandleFunc(uri, func(w http.ResponseWriter, r *http.Request) {
				scheme := "http"
				if r.TLS != nil {
					scheme = "https"
				}
				w.WriteHeader(args.status)
				fmt.Fprintf(w, "Serving %s via %s from %s on %s\n", r.RequestURI, scheme, args.name, addr)
			})
		case "ws", "wss":
			mux.Handle(uri, websocket.Handler(WSEchoServer))
		}

		tag := "urlprefix-" + p
		if proto == "https" || proto == "wss" {
			tag += " proto=https"
		}
		tags = append(tags, tag)
	}

	checkScheme := "http"
	if args.certFile != "" {
		checkScheme = "https"
	}
	check := &api.AgentServiceCheck{
		HTTP:          checkScheme + "://" + addr + "/health",
		Interval:      "1s",
		Timeout:       "1s",
		TLSSkipVerify: true,
	}
	return &http.Server{Addr: addr, Handler: mux}, tags, check
}

func WSEchoServer(ws *websocket.Conn) {
	addr := ws.LocalAddr().String()
	pfx := []byte("[" + addr + "] ")

	log.Printf("ws connect on %s", addr)

	// the following could be done with io.Copy(ws, ws)
	// but I want to add some meta data
	var msg = make([]byte, 1024)
	for {
		n, err := ws.Read(msg)
		if err != nil && err != io.EOF {
			log.Printf("ws error on %s. %s", addr, err)
			break
		}
		_, err = ws.Write(append(pfx, msg[:n]...))
		if err != nil && err != io.EOF {
			log.Printf("ws error on %s. %s", addr, err)
			break
		}
	}
	log.Printf("ws disconnect on %s", addr)
}

func newTCPServer(args Args) (*tcp.Server, []string, *api.AgentServiceCheck) {
	tags := args.tags
	for _, p := range strings.Split(args.prefix, ",") {
		tags = append(tags, "urlprefix-"+p+" proto=tcp")
	}
	check := &api.AgentServiceCheck{TCP: args.addr, Interval: "2s", Timeout: "1s"}
	return &tcp.Server{Addr: args.addr, Handler: tcp.HandlerFunc(TCPEchoHandler)}, tags, check
}

func TCPEchoHandler(c net.Conn) error {
	defer c.Close()

	addr := c.LocalAddr().String()
	_, err := fmt.Fprintf(c, "[%s] Welcome\n", addr)
	if err != nil {
		return err
	}

	for {
		line, _, err := bufio.NewReader(c).ReadLine()
		if err != nil {
			return err
		}
		_, err = fmt.Fprintf(c, "[%s] %s\n", addr, string(line))
		if err != nil {
			return err
		}
	}
}
<commit_msg>fix demo server to handle redirect opts<commit_after>// Package server provides a sample HTTP/Websocket server which registers
// itself in consul using one or more url prefixes to demonstrate and
// test the automatic fabio routing table update.
//
// During startup the server performs the following steps:
//
// * Add a handler for each prefix which provides a unique
//   response for that instance and endpoint
// * Add a `/health` handler for the consul health check
// * Register the service in consul with the listen address,
//   a health check under the given name and with one `urlprefix-`
//   tag per prefix
// * Install a signal handler to deregister the service on exit
//
// If the protocol is set to "ws" the registered endpoints function
// as websocket echo servers.
//
// Example:
//
//   # http server
//   ./server -addr 127.0.0.1:5000 -name svc-a -prefix /foo -prefix /bar
//   ./server -addr 127.0.0.1:5001 -name svc-b -prefix /baz -prefix /bar
//   ./server -addr 127.0.0.1:5002 -name svc-c -prefix "/gogl redirect=301,https://www.google.de/"
//
//   # https server
//   ./server -addr 127.0.0.1:5000 -name svc-a -proto https -certFile ... -keyFile ... -prefix /foo
//   ./server -addr 127.0.0.1:5000 -name svc-a -proto https -certFile ... -keyFile ... -prefix "/foo tlsskipverify=true"
//
//   # websocket server
//   ./server -addr 127.0.0.1:6000 -name ws-a -proto ws -prefix /echo1 -prefix /echo2
//
//   # tcp server
//   ./server -addr 127.0.0.1:7000 -name tcp-a -proto tcp -prefix :1234
//
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/fabiolb/fabio/proxy/tcp"

	"github.com/hashicorp/consul/api"
	"golang.org/x/net/websocket"
)

type Args struct {
	addr     string
	consul   string
	name     string
	proto    string
	token    string
	certFile string
	keyFile  string
	status   int
	prefixes []string
	tags     []string
}

type stringsVar []string

func (s *stringsVar) Set(value string) error {
	*s = append(*s, value)
	return nil
}

func (s stringsVar) String() string {
	return strings.Join(s, " ")
}

func main() {
	var args Args

	flag.StringVar(&args.addr, "addr", "127.0.0.1:5000", "host:port of the service")
	flag.StringVar(&args.consul, "consul", "127.0.0.1:8500", "host:port of the consul agent")
	flag.StringVar(&args.name, "name", filepath.Base(os.Args[0]), "name of the service")
	flag.StringVar(&args.proto, "proto", "http", "protocol for endpoints: http, ws or tcp")
	flag.StringVar(&args.token, "token", "", "consul ACL token")
	flag.StringVar(&args.certFile, "cert", "", "path to cert file")
	flag.StringVar(&args.keyFile, "key", "", "path to key file")
	flag.IntVar(&args.status, "status", http.StatusOK, "http status code")
	flag.Var((*stringsVar)(&args.prefixes), "prefix", "'host/path' or ':port' prefix to register. Can be specified multiple times")
	flag.Var((*stringsVar)(&args.tags), "tags", "additional tags to register in consul. Can be specified multiple times")
	flag.Parse()

	if len(args.prefixes) == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if (args.proto == "https" || args.proto == "wss") && args.certFile == "" {
		log.Fatalf("Proto %s requires a certificate. Please provide -cert/-key", args.proto)
	}

	type server interface {
		ListenAndServe() error
		ListenAndServeTLS(certFile, keyFile string) error
	}

	var srv server
	var tags []string
	var check *api.AgentServiceCheck
	switch args.proto {
	case "http", "https", "ws", "wss":
		srv, tags, check = newHTTPServer(args)
	case "tcp":
		srv, tags, check = newTCPServer(args)
	default:
		log.Fatal("Invalid protocol ", args.proto)
	}

	// start server
	go func() {
		var err error
		if args.certFile != "" {
			err = srv.ListenAndServeTLS(args.certFile, args.keyFile)
		} else {
			err = srv.ListenAndServe()
		}
		if err != nil {
			log.Fatal(err)
		}
	}()

	// get host and port as string/int
	host, portstr, err := net.SplitHostPort(args.addr)
	if err != nil {
		log.Fatal(err)
	}
	port, err := strconv.Atoi(portstr)
	if err != nil {
		log.Fatal(err)
	}

	// register service with health check
	serviceID := args.name + "-" + args.addr
	service := &api.AgentServiceRegistration{
		ID:      serviceID,
		Name:    args.name,
		Port:    port,
		Address: host,
		Tags:    tags,
		Check:   check,
	}

	config := &api.Config{Address: args.consul, Scheme: "http", Token: args.token}
	client, err := api.NewClient(config)
	if err != nil {
		log.Fatal(err)
	}

	if err := client.Agent().ServiceRegister(service); err != nil {
		log.Fatal(err)
	}
	log.Printf("Registered %s service %q in consul with tags %q", args.proto, args.name, strings.Join(tags, ","))

	// run until we get a signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt, os.Kill)
	<-quit

	// deregister service
	if err := client.Agent().ServiceDeregister(serviceID); err != nil {
		log.Fatal(err)
	}
	log.Printf("Deregistered service %q in consul", args.name)
}

func newHTTPServer(args Args) (*http.Server, []string, *api.AgentServiceCheck) {
	addr, proto, tags := args.addr, args.proto, args.tags

	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "OK")
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "not found", 404)
		log.Printf("%s -> 404", r.URL)
	})

	for _, p := range args.prefixes {
		uri := strings.Fields(p)[0]
		switch proto {
		case "http", "https":
			mux.HandleFunc(uri, func(w http.ResponseWriter, r *http.Request) {
				scheme := "http"
				if r.TLS != nil {
					scheme = "https"
				}
				w.WriteHeader(args.status)
				fmt.Fprintf(w, "Serving %s via %s from %s on %s\n", r.RequestURI, scheme, args.name, addr)
			})
		case "ws", "wss":
			mux.Handle(uri, websocket.Handler(WSEchoServer))
		}

		tag := "urlprefix-" + p
		if proto == "https" || proto == "wss" {
			tag += " proto=https"
		}
		tags = append(tags, tag)
	}

	checkScheme := "http"
	if args.certFile != "" {
		checkScheme = "https"
	}
	check := &api.AgentServiceCheck{
		HTTP:          checkScheme + "://" + addr + "/health",
		Interval:      "1s",
		Timeout:       "1s",
		TLSSkipVerify: true,
	}
	return &http.Server{Addr: addr, Handler: mux}, tags, check
}

func WSEchoServer(ws *websocket.Conn) {
	addr := ws.LocalAddr().String()
	pfx := []byte("[" + addr + "] ")

	log.Printf("ws connect on %s", addr)

	// the following could be done with io.Copy(ws, ws)
	// but I want to add some meta data
	var msg = make([]byte, 1024)
	for {
		n, err := ws.Read(msg)
		if err != nil && err != io.EOF {
			log.Printf("ws error on %s. %s", addr, err)
			break
		}
		_, err = ws.Write(append(pfx, msg[:n]...))
		if err != nil && err != io.EOF {
			log.Printf("ws error on %s. %s", addr, err)
			break
		}
	}
	log.Printf("ws disconnect on %s", addr)
}

func newTCPServer(args Args) (*tcp.Server, []string, *api.AgentServiceCheck) {
	tags := args.tags
	for _, p := range args.prefixes {
		tags = append(tags, "urlprefix-"+p+" proto=tcp")
	}
	check := &api.AgentServiceCheck{TCP: args.addr, Interval: "2s", Timeout: "1s"}
	return &tcp.Server{Addr: args.addr, Handler: tcp.HandlerFunc(TCPEchoHandler)}, tags, check
}

func TCPEchoHandler(c net.Conn) error {
	defer c.Close()

	addr := c.LocalAddr().String()
	_, err := fmt.Fprintf(c, "[%s] Welcome\n", addr)
	if err != nil {
		return err
	}

	for {
		line, _, err := bufio.NewReader(c).ReadLine()
		if err != nil {
			return err
		}
		_, err = fmt.Fprintf(c, "[%s] %s\n", addr, string(line))
		if err != nil {
			return err
		}
	}
}
<|endoftext|>


<commit_before>package now

import (
    "bufio"
    "container/list"
    "fmt"
    "os"
    "strings"
)

type Thing struct {
    Name string
}

var (
    theList *list.List = list.New()
)

const filename = ".next.things"

func init() { readList() }

func openList(fresh bool) (*os.File, bool) {
    home := os.Getenv("HOME")
    if home == "" {
        return nil, false
    }

    filePath := home + "/" + filename
    flags := os.O_CREAT|os.O_RDWR
    if fresh {
        flags |= os.O_TRUNC
    } else {
        flags |= os.O_APPEND
    }
    file, err := os.Open(filePath, flags, 0600)

    if err != nil {
        return file, false
    }

    return file, true
}

func readList() {
    listFile, ok := openList(false)
    defer listFile.Close()

    if !ok { return }

    reader := bufio.NewReader(listFile)

    theList.Init()
    nextThing, err := reader.ReadString('\n')
    for ; err == nil; nextThing, err = reader.ReadString('\n') {
        theList.PushBack(strings.Split(nextThing, "\n", -1)[0])
    }
}

func writeList() {
    listFile, ok := openList(true)
    defer listFile.Close()

    if !ok { return }

    things := theList.Iter()
    for thing := range things {
        listFile.WriteString(fmt.Sprintf("%s\n", thing))
    }
}

func AddThing(addition string) {
    if addition != "" {
        readList()
        theList.PushBack(addition)
        writeList()
    }
}

func GetNext() string {
    readList()
    if front := theList.Front(); front != nil {
        return theList.Front().Value.(string)
    }
    return ""
}

func Pop() {
    readList()
    if front := theList.Front(); front != nil {
        theList.Remove(front)
    }
    writeList()
}

func PushNext(addition string) {
    readList()
    theList.PushFront(addition)
    writeList()
}
<commit_msg>type Thing is completely unneeded.<commit_after>package now

import (
    "bufio"
    "container/list"
    "fmt"
    "os"
    "strings"
)

var (
    theList *list.List = list.New()
)

const filename = ".next.things"

func init() { readList() }

func openList(fresh bool) (*os.File, bool) {
    home := os.Getenv("HOME")
    if home == "" {
        return nil, false
    }

    filePath := home + "/" + filename
    flags := os.O_CREAT|os.O_RDWR
    if fresh {
        flags |= os.O_TRUNC
    } else {
        flags |= os.O_APPEND
    }
    file, err := os.Open(filePath, flags, 0600)

    if err != nil {
        return file, false
    }

    return file, true
}

func readList() {
    listFile, ok := openList(false)
    defer listFile.Close()

    if !ok { return }

    reader := bufio.NewReader(listFile)

    theList.Init()
    nextThing, err := reader.ReadString('\n')
    for ; err == nil; nextThing, err = reader.ReadString('\n') {
        theList.PushBack(strings.Split(nextThing, "\n", -1)[0])
    }
}

func writeList() {
    listFile, ok := openList(true)
    defer listFile.Close()

    if !ok { return }

    things := theList.Iter()
    for thing := range things {
        listFile.WriteString(fmt.Sprintf("%s\n", thing))
    }
}

func AddThing(addition string) {
    if addition != "" {
        readList()
        theList.PushBack(addition)
        writeList()
    }
}

func GetNext() string {
    readList()
    if front := theList.Front(); front != nil {
        return theList.Front().Value.(string)
    }
    return ""
}

func Pop() {
    readList()
    if front := theList.Front(); front != nil {
        theList.Remove(front)
    }
    writeList()
}

func PushNext(addition string) {
    readList()
    theList.PushFront(addition)
    writeList()
}
<|endoftext|>


<commit_before>package master

type pqSlice struct {
	Slice          []interface{}
	LessComparator func(i, j interface{}) bool
}

func (s pqSlice) Len() int {
	return len(s.Slice)
}

func (s pqSlice) Less(left, right int) bool {
	return s.LessComparator(s.Slice[left], s.Slice[right])
}

func (s *pqSlice) Swap(i, j int) {
	s.Slice[i], s.Slice[j] = s.Slice[j], s.Slice[i]
}

func (s pqSlice) Push(i interface{}) {
	s.Slice = append(s.Slice, i)
}

func (s pqSlice) Pop() interface{} {
	ret := s.Slice[len(s.Slice)-1]
	s.Slice = s.Slice[0 : len(s.Slice)-1]
	return ret
}
<commit_msg>Use call-by-reference methods for Push and Pop<commit_after>package master

type pqSlice struct {
	Slice          []interface{}
	LessComparator func(i, j interface{}) bool
}

func (s pqSlice) Len() int {
	return len(s.Slice)
}

func (s pqSlice) Less(left, right int) bool {
	return s.LessComparator(s.Slice[left], s.Slice[right])
}

func (s *pqSlice) Swap(i, j int) {
	s.Slice[i], s.Slice[j] = s.Slice[j], s.Slice[i]
}

func (s *pqSlice) Push(i interface{}) {
	s.Slice = append(s.Slice, i)
}

func (s *pqSlice) Pop() interface{} {
	ret := s.Slice[len(s.Slice)-1]
	s.Slice = s.Slice[0 : len(s.Slice)-1]
	return ret
}
<|endoftext|>


<commit_before>package gootp

import (
	"bytes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"math"
	"strconv"
	"strings"
	"time"
)

func hmacMsg(c, key []byte, h func() hash.Hash) ([]byte, error) {
	//Hash the message derived from the time with the key from the TOTP secret
	mac := hmac.New(h, key)
	_, err := mac.Write(c)
	if err != nil {
		return make([]byte, 0), err
	}
	return mac.Sum(nil), nil
}

func generateHashKey(s string) ([]byte, error) {
	//Decode and convert the base32 encoded secret string to a byte array
	key, err := base32.StdEncoding.DecodeString(strings.ToUpper(s))
	if err != nil {
		return make([]byte, 0), err
	}
	return key, nil
}

func generateHashCountFromTime(t time.Time) ([]byte, error) {
	//The message to be hashed is the count of the number of time intervals
	c := t.Unix() / 30
	return generateHashCount(c)
}

func generateHashCount(c int64) ([]byte, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, c)
	if err != nil {
		return make([]byte, 0), err
	}
	return buf.Bytes(), nil
}

func extractFromHash(hash []byte) uint32 {
	//Get the last 4 bits of the hash (a value ranging from 0-15)
	// This will be the index into the 20-byte value
	iBits := hash[len(hash)-1] & 0xf
	iInt := int64(iBits)

	//Extract the next 4 bytes starting at the index and convert to uint32
	r := bytes.NewReader(hash)
	var b [4]byte
	r.ReadAt(b[:], iInt)
	b[0] = b[0] & 0x7f
	return binary.BigEndian.Uint32(b[:])
}

/*
 Get the HMAC-based One Time Password (RFC 4226). Providing the following inputs:
 	- Secret string at least 16 bytes / 128 bits in length
 	- Counter value, the moving factor (see RFC 4226 section 5.2).  This counter MUST be synchronized between the HOTP generator (client) and the HOTP validator (server).
 	- A hash function to use, eg SHA1, SHA256, SHA512
 	- The number of digits to be returned in the OTP. Must be a minimum of 6.

 Note that the returned OTP is a string as a leading zero is valid so an integer type is not appropriate
*/
func GetHOTP(secret string, count int64, mode func() hash.Hash, digits int) (otp string, err error) {
	if digits < 6 {
		err = errors.New("The number of digits of the OTP generated must be at least 6")
		return
	}
	if len(secret) < 16 {
		err = errors.New("The secret string used to generate the OTP must be at least 128 bits")
		return
	}
	key, err := generateHashKey(secret)
	if err != nil {
		return
	}
	msg, err := generateHashCount(count)
	if err != nil {
		return
	}
	if mode == nil {
		mode = sha1.New
	}
	hash, err := hmacMsg(msg, key, mode)
	if err != nil {
		return
	}
	otpInt := int(math.Mod(float64(extractFromHash(hash)), math.Pow(10, float64(digits))))
	otp = fmt.Sprintf("%0"+strconv.Itoa(digits)+"d", otpInt)
	return
}

/*
 Get the Time-based One Time Password (RFC 6238) for the current time. Providing the following inputs:
 	- Secret string at least 16 bytes / 128 bits in length.
 	- A hash function to use, eg SHA1, SHA256, SHA512.
 	- The number of digits to be returned in the OTP. Must be a minimum of 6.

 Note that the returned OTP is a string as a leading zero is valid so an integer type is not appropriate.
 The number of seconds the OTP is valid for is also returned.
*/
func GetTOTPNow(secret string, mode func() hash.Hash, digits int) (otp string, timeRemaining int, err error) {
	otp, timeRemaining, err = GetTOTPAt(secret, time.Now().UTC(), mode, digits)
	return
}

/*
 Get the Time-based One Time Password (RFC 6238) for a specific time. Providing the following inputs:
 	- Secret string at least 16 bytes / 128 bits in length.
 	- The UTC time for which the TOTP should be generated.
 	- A hash function to use, eg SHA1, SHA256, SHA512.
 	- The number of digits to be returned in the OTP. Must be a minimum of 6.

 Note that the returned OTP is a string as a leading zero is valid so an integer type is not appropriate.
 The number of seconds the OTP is valid for is also returned.
*/
func GetTOTPAt(secret string, t time.Time, h func() hash.Hash, digits int) (otp string, timeRemaining int, err error) {
	key, err := generateHashKey(secret)
	if err != nil {
		return
	}
	msg, err := generateHashCountFromTime(t)
	if err != nil {
		return
	}
	hash, err := hmacMsg(msg, key, h)
	if err != nil {
		return
	}
	timeRemaining = int(30 - math.Mod(float64(t.Unix()), 30))
	otpInt := int(math.Mod(float64(extractFromHash(hash)), math.Pow(10, float64(digits))))
	otp = fmt.Sprintf("%0"+strconv.Itoa(digits)+"d", otpInt)
	return
}

// Get a Time-based One Time Password history (RFC 6238). Providing the following inputs:
//
// - Secret string at least 16 bytes / 128 bits in length.
//
// - A hash function to use, eg SHA1, SHA256, SHA512.
//
// - The number of digits to be returned in the OTP. Must be a minimum of 6.
//
//
// Note that the returned OTP is an array of strings as a leading zero is valid so an integer type is not appropriate. The first element in the array is the current OTP.
//
// The number of seconds the current OTP is valid for is also returned.
func GetTOTPHistory(secret string, h func() hash.Hash, digits int, history int) (otps []string, timeRemaining int, err error) {
	key, err := generateHashKey(secret)
	if err != nil {
		return
	}
	c := time.Now().UTC().Unix() / 30
	timeRemaining = int(30 - math.Mod(float64(time.Now().UTC().Unix()), 30))
	for i := 0; i < history; i++ {
		c = c - int64(i)
		msg, ierr := generateHashCount(c)
		if ierr != nil {
			err = ierr
			return
		}
		hash, ierr := hmacMsg(msg, key, h)
		if ierr != nil {
			err = ierr
			return
		}
		otpInt := int(math.Mod(float64(extractFromHash(hash)), math.Pow(10, float64(digits))))
		otp := fmt.Sprintf("%0"+strconv.Itoa(digits)+"d", otpInt)
		otps = append(otps, otp)
	}
	return
}

// Generate a base32 encoded secret string to be shared between the client and the server for OTPs
// Specify the length of the secret to generate in bytes. Note this needs to be at least 16 bytes / 128 bits.
func GenerateOTPSecret(s int) (string, error) {
	if s < 16 {
		return "", errors.New("The secret size needs to be at least 16 bytes / 128 bits")
	}
	b, err := generateRandomBytes(s)
	return base32.StdEncoding.EncodeToString(b), err
}

func generateRandomBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	_, err := rand.Read(b)
	if err != nil {
		return nil, err
	}

	return b, nil
}
<commit_msg>update<commit_after>package gootp

import (
	"bytes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"math"
	"strconv"
	"strings"
	"time"
)

func hmacMsg(c, key []byte, h func() hash.Hash) ([]byte, error) {
	//Hash the message derived from the time with the key from the TOTP secret
	mac := hmac.New(h, key)
	_, err := mac.Write(c)
	if err != nil {
		return make([]byte, 0), err
	}
	return mac.Sum(nil), nil
}

func generateHashKey(s string) ([]byte, error) {
	//Decode and convert the base32 encoded secret string to a byte array
	key, err := base32.StdEncoding.DecodeString(strings.ToUpper(s))
	if err != nil {
		return make([]byte, 0), err
	}
	return key, nil
}

func generateHashCountFromTime(t time.Time) ([]byte, error) {
	//The message to be hashed is the count of the number of time intervals
	c := t.Unix() / 30
	return generateHashCount(c)
}

func generateHashCount(c int64) ([]byte, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, c)
	if err != nil {
		return make([]byte, 0), err
	}
	return buf.Bytes(), nil
}

func extractFromHash(hash []byte) uint32 {
	//Get the last 4 bits of the hash (a value ranging from 0-15)
	// This will be the index into the 20-byte value
	iBits := hash[len(hash)-1] & 0xf
	iInt := int64(iBits)

	//Extract the next 4 bytes starting at the index and convert to uint32
	r := bytes.NewReader(hash)
	var b [4]byte
	r.ReadAt(b[:], iInt)
	b[0] = b[0] & 0x7f
	return binary.BigEndian.Uint32(b[:])
}

// Get the HMAC-based One Time Password (RFC 4226). Providing the following inputs:
//    - Secret string at least 16 bytes / 128 bits in length
//
//    - Counter value, the moving factor (see RFC 4226 section 5.2).  This counter MUST be synchronized between the HOTP generator (client) and the HOTP validator (server).
//
//    - A hash function to use, eg SHA1, SHA256, SHA512
//
//    - The number of digits to be returned in the OTP. Must be a minimum of 6.
//
// Note that the returned OTP is a string as a leading zero is valid so an integer type is not appropriate
func GetHOTP(secret string, count int64, mode func() hash.Hash, digits int) (otp string, err error) {
	if digits < 6 {
		err = errors.New("The number of digits of the OTP generated must be at least 6")
		return
	}
	if len(secret) < 16 {
		err = errors.New("The secret string used to generate the OTP must be at least 128 bits")
		return
	}
	key, err := generateHashKey(secret)
	if err != nil {
		return
	}
	msg, err := generateHashCount(count)
	if err != nil {
		return
	}
	if mode == nil {
		mode = sha1.New
	}
	hash, err := hmacMsg(msg, key, mode)
	if err != nil {
		return
	}
	otpInt := int(math.Mod(float64(extractFromHash(hash)), math.Pow(10, float64(digits))))
	otp = fmt.Sprintf("%0"+strconv.Itoa(digits)+"d", otpInt)
	return
}

// Get the Time-based One Time Password (RFC 6238) for the current time. Providing the following inputs:
//    - Secret string at least 16 bytes / 128 bits in length.
//    - A hash function to use, eg SHA1, SHA256, SHA512.
//    - The number of digits to be returned in the OTP. Must be a minimum of 6.
//
//
// Note that the returned OTP is a string as a leading zero is valid so an integer type is not appropriate.
//
// The number of seconds the OTP is valid for is also returned.
func GetTOTPNow(secret string, mode func() hash.Hash, digits int) (otp string, timeRemaining int, err error) {
	otp, timeRemaining, err = GetTOTPAt(secret, time.Now().UTC(), mode, digits)
	return
}

// Get the Time-based One Time Password (RFC 6238) for a specific time. Providing the following inputs:
//    - Secret string at least 16 bytes / 128 bits in length.
//    - The UTC time for which the TOTP should be generated.
//    - A hash function to use, eg SHA1, SHA256, SHA512.
//    - The number of digits to be returned in the OTP. Must be a minimum of 6.
//
//
// Note that the returned OTP is a string as a leading zero is valid so an integer type is not appropriate.
//
// The number of seconds the OTP is valid for is also returned.
func GetTOTPAt(secret string, t time.Time, h func() hash.Hash, digits int) (otp string, timeRemaining int, err error) {
	key, err := generateHashKey(secret)
	if err != nil {
		return
	}
	msg, err := generateHashCountFromTime(t)
	if err != nil {
		return
	}
	hash, err := hmacMsg(msg, key, h)
	if err != nil {
		return
	}
	timeRemaining = int(30 - math.Mod(float64(t.Unix()), 30))
	otpInt := int(math.Mod(float64(extractFromHash(hash)), math.Pow(10, float64(digits))))
	otp = fmt.Sprintf("%0"+strconv.Itoa(digits)+"d", otpInt)
	return
}

// Get a Time-based One Time Password history (RFC 6238). Providing the following inputs:
//
//    - Secret string at least 16 bytes / 128 bits in length.
//
//    - A hash function to use, eg SHA1, SHA256, SHA512.
//
//    - The number of digits to be returned in the OTP. Must be a minimum of 6.
//
//
// Note that the returned OTP is an array of strings as a leading zero is valid so an integer type is not appropriate. The first element in the array is the current OTP.
//
// The number of seconds the current OTP is valid for is also returned.
func GetTOTPHistory(secret string, h func() hash.Hash, digits int, history int) (otps []string, timeRemaining int, err error) {
	key, err := generateHashKey(secret)
	if err != nil {
		return
	}
	c := time.Now().UTC().Unix() / 30
	timeRemaining = int(30 - math.Mod(float64(time.Now().UTC().Unix()), 30))
	for i := 0; i < history; i++ {
		c = c - int64(i)
		msg, ierr := generateHashCount(c)
		if ierr != nil {
			err = ierr
			return
		}
		hash, ierr := hmacMsg(msg, key, h)
		if ierr != nil {
			err = ierr
			return
		}
		otpInt := int(math.Mod(float64(extractFromHash(hash)), math.Pow(10, float64(digits))))
		otp := fmt.Sprintf("%0"+strconv.Itoa(digits)+"d", otpInt)
		otps = append(otps, otp)
	}
	return
}

// Generate a base32 encoded secret string to be shared between the client and the server for OTPs
// Specify the length of the secret to generate in bytes. Note this needs to be at least 16 bytes / 128 bits.
func GenerateOTPSecret(s int) (string, error) {
	if s < 16 {
		return "", errors.New("The secret size needs to be at least 16 bytes / 128 bits")
	}
	b, err := generateRandomBytes(s)
	return base32.StdEncoding.EncodeToString(b), err
}

func generateRandomBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	_, err := rand.Read(b)
	if err != nil {
		return nil, err
	}

	return b, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"fmt"
	"log"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/elb"
)

func resourceAwsElb() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsElbCreate,
		Read:   resourceAwsElbRead,
		Update: resourceAwsElbUpdate,
		Delete: resourceAwsElbDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"internal": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
				Computed: true,
			},

			"cross_zone_load_balancing": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
			},

			"availability_zones": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				ForceNew: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"instances": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			// TODO: could be not ForceNew
			"security_groups": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				ForceNew: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"subnets": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			// TODO: could be not ForceNew
			"listener": &schema.Schema{
				Type:     schema.TypeSet,
				Required: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"instance_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"instance_protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"lb_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"lb_protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"ssl_certificate_id": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
				Set: resourceAwsElbListenerHash,
			},

			// TODO: could be not ForceNew
			"health_check": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"healthy_threshold": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"unhealthy_threshold": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"target": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"interval": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"timeout": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
					},
				},
				Set: resourceAwsElbHealthCheckHash,
			},

			"dns_name": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsElbCreate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	// Expand the "listener" set to goamz compat []elb.Listener
	listeners, err := expandListeners(d.Get("listener").(*schema.Set).List())
	if err != nil {
		return err
	}

	// Provision the elb
	elbOpts := &elb.CreateLoadBalancer{
		LoadBalancerName: d.Get("name").(string),
		Listeners:        listeners,
		Internal:         d.Get("internal").(bool),
	}

	if v, ok := d.GetOk("availability_zones"); ok {
		elbOpts.AvailZone = expandStringList(v.(*schema.Set).List())
	}

	if v, ok := d.GetOk("security_groups"); ok {
		elbOpts.SecurityGroups = expandStringList(v.(*schema.Set).List())
	}

	if v, ok := d.GetOk("subnets"); ok {
		elbOpts.Subnets = expandStringList(v.(*schema.Set).List())
	}

	log.Printf("[DEBUG] ELB create configuration: %#v", elbOpts)
	if _, err := elbconn.CreateLoadBalancer(elbOpts); err != nil {
		return fmt.Errorf("Error creating ELB: %s", err)
	}

	// Assign the elb's unique identifier for use later
	d.SetId(d.Get("name").(string))
	log.Printf("[INFO] ELB ID: %s", d.Id())

	// Enable partial mode and record what we set
	d.Partial(true)
	d.SetPartial("name")
	d.SetPartial("internal")
	d.SetPartial("availability_zones")
	d.SetPartial("listener")
	d.SetPartial("security_groups")
	d.SetPartial("subnets")

	if d.HasChange("health_check") {
		vs := d.Get("health_check").(*schema.Set).List()
		if len(vs) > 0 {
			check := vs[0].(map[string]interface{})

			configureHealthCheckOpts := elb.ConfigureHealthCheck{
				LoadBalancerName: d.Id(),
				Check: elb.HealthCheck{
					HealthyThreshold:   int64(check["healthy_threshold"].(int)),
					UnhealthyThreshold: int64(check["unhealthy_threshold"].(int)),
					Interval:           int64(check["interval"].(int)),
					Target:             check["target"].(string),
					Timeout:            int64(check["timeout"].(int)),
				},
			}

			_, err = elbconn.ConfigureHealthCheck(&configureHealthCheckOpts)
			if err != nil {
				return fmt.Errorf("Failure configuring health check: %s", err)
			}
		}
	}

	return resourceAwsElbUpdate(d, meta)
}

func resourceAwsElbRead(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	// Retrieve the ELB properties for updating the state
	describeElbOpts := &elb.DescribeLoadBalancer{
		Names: []string{d.Id()},
	}

	describeResp, err := elbconn.DescribeLoadBalancers(describeElbOpts)
	if err != nil {
		if ec2err, ok := err.(*elb.Error); ok && ec2err.Code == "LoadBalancerNotFound" {
			// The ELB is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}

		return fmt.Errorf("Error retrieving ELB: %s", err)
	}
	if len(describeResp.LoadBalancers) != 1 {
		return fmt.Errorf("Unable to find ELB: %#v", describeResp.LoadBalancers)
	}

	lb := describeResp.LoadBalancers[0]

	d.Set("name", lb.LoadBalancerName)
	d.Set("dns_name", lb.DNSName)
	d.Set("internal", lb.Scheme == "internal")
	d.Set("availability_zones", lb.AvailabilityZones)
	d.Set("instances", flattenInstances(lb.Instances))
	d.Set("listener", flattenListeners(lb.Listeners))
	d.Set("security_groups", lb.SecurityGroups)
	d.Set("subnets", lb.Subnets)

	// There's only one health check, so save that to state as we
	// currently can
	if lb.HealthCheck.Target != "" {
		d.Set("health_check", flattenHealthCheck(lb.HealthCheck))
	}

	return nil
}

func resourceAwsElbUpdate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	d.Partial(true)

	// If we currently have instances, or did have instances,
	// we want to figure out what to add and remove from the load
	// balancer
	if d.HasChange("instances") {
		o, n := d.GetChange("instances")
		os := o.(*schema.Set)
		ns := n.(*schema.Set)
		remove := expandStringList(os.Difference(ns).List())
		add := expandStringList(ns.Difference(os).List())

		if len(add) > 0 {
			registerInstancesOpts := elb.RegisterInstancesWithLoadBalancer{
				LoadBalancerName: d.Id(),
				Instances:        add,
			}

			_, err := elbconn.RegisterInstancesWithLoadBalancer(&registerInstancesOpts)
			if err != nil {
				return fmt.Errorf("Failure registering instances: %s", err)
			}
		}
		if len(remove) > 0 {
			deRegisterInstancesOpts := elb.DeregisterInstancesFromLoadBalancer{
				LoadBalancerName: d.Id(),
				Instances:        remove,
			}

			_, err := elbconn.DeregisterInstancesFromLoadBalancer(&deRegisterInstancesOpts)
			if err != nil {
				return fmt.Errorf("Failure deregistering instances: %s", err)
			}
		}

		d.SetPartial("instances")
	}

	log.Println("[INFO] outside modify attributes")
	if d.HasChange("cross_zone_load_balancing") {
		log.Println("[INFO] inside modify attributes")
		attrs := elb.ModifyLoadBalancerAttributes{
			LoadBalancerName: d.Get("name").(string),
			LoadBalancerAttributes: elb.LoadBalancerAttributes{
				CrossZoneLoadBalancingEnabled: d.Get("cross_zone_load_balancing").(bool),
			},
		}
		_, err := elbconn.ModifyLoadBalancerAttributes(&attrs)
		if err != nil {
			return fmt.Errorf("Failure configuring cross zone balancing: %s", err)
		}
		d.SetPartial("cross_zone_load_balancing")
	}

	if d.HasChange("health_check") {
		vs := d.Get("health_check").(*schema.Set).List()
		if len(vs) > 0 {
			check := vs[0].(map[string]interface{})
			configureHealthCheckOpts := elb.ConfigureHealthCheck{
				LoadBalancerName: d.Id(),
				Check: elb.HealthCheck{
					HealthyThreshold:   int64(check["healthy_threshold"].(int)),
					UnhealthyThreshold: int64(check["unhealthy_threshold"].(int)),
					Interval:           int64(check["interval"].(int)),
					Target:             check["target"].(string),
					Timeout:            int64(check["timeout"].(int)),
				},
			}
			_, err := elbconn.ConfigureHealthCheck(&configureHealthCheckOpts)
			if err != nil {
				return fmt.Errorf("Failure configuring health check: %s", err)
			}
			d.SetPartial("health_check")
		}
	}

	d.Partial(false)
	return resourceAwsElbRead(d, meta)
}

func resourceAwsElbDelete(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	log.Printf("[INFO] Deleting ELB: %s", d.Id())

	// Destroy the load balancer
	deleteElbOpts := elb.DeleteLoadBalancer{
		LoadBalancerName: d.Id(),
	}
	if _, err := elbconn.DeleteLoadBalancer(&deleteElbOpts); err != nil {
		return fmt.Errorf("Error deleting ELB: %s", err)
	}

	return nil
}

func resourceAwsElbHealthCheckHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["healthy_threshold"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["unhealthy_threshold"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["target"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["interval"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["timeout"].(int)))

	return hashcode.String(buf.String())
}

func resourceAwsElbListenerHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["instance_port"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["instance_protocol"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["lb_port"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["lb_protocol"].(string)))

	if v, ok := m["ssl_certificate_id"]; ok {
		buf.WriteString(fmt.Sprintf("%s-", v.(string)))
	}

	return hashcode.String(buf.String())
}
<commit_msg>provider/aws: fix for #915, disabling ForceNew while updating elb healthchecks<commit_after>package aws

import (
	"bytes"
	"fmt"
	"log"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/elb"
)

func resourceAwsElb() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsElbCreate,
		Read:   resourceAwsElbRead,
		Update: resourceAwsElbUpdate,
		Delete: resourceAwsElbDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"internal": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
				Computed: true,
			},

			"cross_zone_load_balancing": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
			},

			"availability_zones": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				ForceNew: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"instances": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			// TODO: could be not ForceNew
			"security_groups": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				ForceNew: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"subnets": &schema.Schema{
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Optional: true,
				Computed: true,
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			// TODO: could be not ForceNew
			"listener": &schema.Schema{
				Type:     schema.TypeSet,
				Required: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"instance_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"instance_protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"lb_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"lb_protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"ssl_certificate_id": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
				Set: resourceAwsElbListenerHash,
			},

			// TODO: could be not ForceNew
			"health_check": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"healthy_threshold": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"unhealthy_threshold": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"target": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"interval": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"timeout": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
					},
				},
				Set: resourceAwsElbHealthCheckHash,
			},

			"dns_name": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsElbCreate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	// Expand the "listener" set to goamz compat []elb.Listener
	listeners, err := expandListeners(d.Get("listener").(*schema.Set).List())
	if err != nil {
		return err
	}

	// Provision the elb
	elbOpts := &elb.CreateLoadBalancer{
		LoadBalancerName: d.Get("name").(string),
		Listeners:        listeners,
		Internal:         d.Get("internal").(bool),
	}

	if v, ok := d.GetOk("availability_zones"); ok {
		elbOpts.AvailZone = expandStringList(v.(*schema.Set).List())
	}

	if v, ok := d.GetOk("security_groups"); ok {
		elbOpts.SecurityGroups = expandStringList(v.(*schema.Set).List())
	}

	if v, ok := d.GetOk("subnets"); ok {
		elbOpts.Subnets = expandStringList(v.(*schema.Set).List())
	}

	log.Printf("[DEBUG] ELB create configuration: %#v", elbOpts)
	if _, err := elbconn.CreateLoadBalancer(elbOpts); err != nil {
		return fmt.Errorf("Error creating ELB: %s", err)
	}

	// Assign the elb's unique identifier for use later
	d.SetId(d.Get("name").(string))
	log.Printf("[INFO] ELB ID: %s", d.Id())

	// Enable partial mode and record what we set
	d.Partial(true)
	d.SetPartial("name")
	d.SetPartial("internal")
	d.SetPartial("availability_zones")
	d.SetPartial("listener")
	d.SetPartial("security_groups")
	d.SetPartial("subnets")

	if d.HasChange("health_check") {
		vs := d.Get("health_check").(*schema.Set).List()
		if len(vs) > 0 {
			check := vs[0].(map[string]interface{})

			configureHealthCheckOpts := elb.ConfigureHealthCheck{
				LoadBalancerName: d.Id(),
				Check: elb.HealthCheck{
					HealthyThreshold:   int64(check["healthy_threshold"].(int)),
					UnhealthyThreshold: int64(check["unhealthy_threshold"].(int)),
					Interval:           int64(check["interval"].(int)),
					Target:             check["target"].(string),
					Timeout:            int64(check["timeout"].(int)),
				},
			}

			_, err = elbconn.ConfigureHealthCheck(&configureHealthCheckOpts)
			if err != nil {
				return fmt.Errorf("Failure configuring health check: %s", err)
			}
		}
	}

	return resourceAwsElbUpdate(d, meta)
}

func resourceAwsElbRead(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	// Retrieve the ELB properties for updating the state
	describeElbOpts := &elb.DescribeLoadBalancer{
		Names: []string{d.Id()},
	}

	describeResp, err := elbconn.DescribeLoadBalancers(describeElbOpts)
	if err != nil {
		if ec2err, ok := err.(*elb.Error); ok && ec2err.Code == "LoadBalancerNotFound" {
			// The ELB is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}

		return fmt.Errorf("Error retrieving ELB: %s", err)
	}
	if len(describeResp.LoadBalancers) != 1 {
		return fmt.Errorf("Unable to find ELB: %#v", describeResp.LoadBalancers)
	}

	lb := describeResp.LoadBalancers[0]

	d.Set("name", lb.LoadBalancerName)
	d.Set("dns_name", lb.DNSName)
	d.Set("internal", lb.Scheme == "internal")
	d.Set("availability_zones", lb.AvailabilityZones)
	d.Set("instances", flattenInstances(lb.Instances))
	d.Set("listener", flattenListeners(lb.Listeners))
	d.Set("security_groups", lb.SecurityGroups)
	d.Set("subnets", lb.Subnets)

	// There's only one health check, so save that to state as we
	// currently can
	if lb.HealthCheck.Target != "" {
		d.Set("health_check", flattenHealthCheck(lb.HealthCheck))
	}

	return nil
}

func resourceAwsElbUpdate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	d.Partial(true)

	// If we currently have instances, or did have instances,
	// we want to figure out what to add and remove from the load
	// balancer
	if d.HasChange("instances") {
		o, n := d.GetChange("instances")
		os := o.(*schema.Set)
		ns := n.(*schema.Set)
		remove := expandStringList(os.Difference(ns).List())
		add := expandStringList(ns.Difference(os).List())

		if len(add) > 0 {
			registerInstancesOpts := elb.RegisterInstancesWithLoadBalancer{
				LoadBalancerName: d.Id(),
				Instances:        add,
			}

			_, err := elbconn.RegisterInstancesWithLoadBalancer(&registerInstancesOpts)
			if err != nil {
				return fmt.Errorf("Failure registering instances: %s", err)
			}
		}
		if len(remove) > 0 {
			deRegisterInstancesOpts := elb.DeregisterInstancesFromLoadBalancer{
				LoadBalancerName: d.Id(),
				Instances:        remove,
			}

			_, err := elbconn.DeregisterInstancesFromLoadBalancer(&deRegisterInstancesOpts)
			if err != nil {
				return fmt.Errorf("Failure deregistering instances: %s", err)
			}
		}

		d.SetPartial("instances")
	}

	log.Println("[INFO] outside modify attributes")
	if d.HasChange("cross_zone_load_balancing") {
		log.Println("[INFO] inside modify attributes")
		attrs := elb.ModifyLoadBalancerAttributes{
			LoadBalancerName: d.Get("name").(string),
			LoadBalancerAttributes: elb.LoadBalancerAttributes{
				CrossZoneLoadBalancingEnabled: d.Get("cross_zone_load_balancing").(bool),
			},
		}
		_, err := elbconn.ModifyLoadBalancerAttributes(&attrs)
		if err != nil {
			return fmt.Errorf("Failure configuring cross zone balancing: %s", err)
		}
		d.SetPartial("cross_zone_load_balancing")
	}

	if d.HasChange("health_check") {
		vs := d.Get("health_check").(*schema.Set).List()
		if len(vs) > 0 {
			check := vs[0].(map[string]interface{})
			configureHealthCheckOpts := elb.ConfigureHealthCheck{
				LoadBalancerName: d.Id(),
				Check: elb.HealthCheck{
					HealthyThreshold:   int64(check["healthy_threshold"].(int)),
					UnhealthyThreshold: int64(check["unhealthy_threshold"].(int)),
					Interval:           int64(check["interval"].(int)),
					Target:             check["target"].(string),
					Timeout:            int64(check["timeout"].(int)),
				},
			}
			_, err := elbconn.ConfigureHealthCheck(&configureHealthCheckOpts)
			if err != nil {
				return fmt.Errorf("Failure configuring health check: %s", err)
			}
			d.SetPartial("health_check")
		}
	}

	d.Partial(false)
	return resourceAwsElbRead(d, meta)
}

func resourceAwsElbDelete(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn

	log.Printf("[INFO] Deleting ELB: %s", d.Id())

	// Destroy the load balancer
	deleteElbOpts := elb.DeleteLoadBalancer{
		LoadBalancerName: d.Id(),
	}
	if _, err := elbconn.DeleteLoadBalancer(&deleteElbOpts); err != nil {
		return fmt.Errorf("Error deleting ELB: %s", err)
	}

	return nil
}

func resourceAwsElbHealthCheckHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["healthy_threshold"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["unhealthy_threshold"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["target"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["interval"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["timeout"].(int)))

	return hashcode.String(buf.String())
}

func resourceAwsElbListenerHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["instance_port"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["instance_protocol"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["lb_port"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["lb_protocol"].(string)))

	if v, ok := m["ssl_certificate_id"]; ok {
		buf.WriteString(fmt.Sprintf("%s-", v.(string)))
	}

	return hashcode.String(buf.String())
}
<|endoftext|>


<commit_before>// Package marathon scheduler integration
package marathon

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

	"github.com/op/go-logging"

	"github.com/microscaling/microscaling/demand"
	"github.com/microscaling/microscaling/scheduler"
)

var log = logging.MustGetLogger("mssscheduler")

// MarathonScheduler holds Marathon API URL and the running services.
type MarathonScheduler struct {
	baseMarathonURL string
	services        map[string][]string
}

// AppsMessage from the Marathon Apps API.
type AppsMessage struct {
	Apps []App `json:"apps"`
}

// App from the Marathon Apps API.
type App struct {
	ID        string `json:"id"`
	Instances int    `json:"instances"`
}

var (
	httpClient = &http.Client{
		// TODO Make timeout configurable.
		Timeout: 10 * time.Second,
	}
)

// NewScheduler returns a pointer to the scheduler.
func NewScheduler(marathonAPI string) *MarathonScheduler {
	return &MarathonScheduler{
		baseMarathonURL: getBaseMarathonURL(marathonAPI),
		services:        make(map[string][]string),
	}
}

type startStopPayload struct {
	Instances int `json:"instances"`
}

// compile-time assert that we implement the right interface
var _ scheduler.Scheduler = (*MarathonScheduler)(nil)

// InitScheduler initializes the scheduler.
func (m *MarathonScheduler) InitScheduler(task *demand.Task) (err error) {
	return err
}

// StopStartTasks by calling the Marathon scaling API.
func (m *MarathonScheduler) StopStartTasks(tasks *demand.Tasks) error {
	// Create tasks if there aren't enough of them, and stop them if there are too many
	var tooMany []*demand.Task
	var tooFew []*demand.Task
	var diff int
	var err error = nil

	// TODO: Consider checking the number running before we start & stop
	for _, task := range tasks.Tasks {
		if task.Demand > task.Requested {
			// There aren't enough of these containers yet
			tooFew = append(tooFew, task)
		}
		if task.Demand < task.Requested {
			// there aren't enough of these containers yet
			tooMany = append(tooMany, task)
		}
	}

	// Scale down first to free up resources
	for _, task := range tooMany {
		diff = task.Requested - task.Demand
		log.Infof("Stop %d of task %s", diff, task.Name)
		err = m.stopStartTask(task)
		if err != nil {
			log.Errorf("Couldn't stop %s: %v ", task.Name, err)
		}
		log.Infof("now have %d", task.Requested)
		// task.Requested = task.Demand
		//tasks[name] = task
	}

	// Now we can scale up
	for _, task := range tooFew {
		diff = task.Demand - task.Requested
		log.Infof("Start %d of task %s", diff, task.Name)
		err = m.stopStartTask(task)
		if err != nil {
			log.Errorf("Couldn't start %s: %v ", task.Name, err)
		}
		log.Infof("now have %d", task.Requested)
		//tasks[name] = task
	}

	log.Infof("%v", tasks)
	return err
}

// Marathon API to get the current running tasks.
func getJSONGet(url string) (body []byte, err error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Errorf("Failed to build API GET request err %v", err)
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	resp, err := httpClient.Do(req)
	if err != nil {
		log.Errorf("Failed to GET err %v", err)
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		log.Errorf("Http error %d: %s", resp.StatusCode, resp.Status)
	}

	body, err = ioutil.ReadAll(resp.Body)

	return body, err
}

// CountAllTasks tells us how many instances of each task are currently running.
func (m *MarathonScheduler) CountAllTasks(running *demand.Tasks) error {
	var (
		err         error
		appsMessage AppsMessage
	)

	url := m.baseMarathonURL + "apps/"

	body, err := getJSONGet(url)
	if err != nil {
		log.Errorf("Error getting Marathon Apps %v", err)
	}

	err = json.Unmarshal(body, &appsMessage)
	if err != nil {
		log.Errorf("Error %v unmarshalling from %s", err, string(body[:]))
	}

	appCounts := make(map[string]int)

	// Remove leading slash from App IDs and set the instance counts.
	for _, app := range appsMessage.Apps {
		appCounts[strings.Replace(app.ID, "/", "", 1)] = app.Instances
	}

	// Set running counts. Defaults to 0 if the App does not exist.
	tasks := running.Tasks
	for _, t := range tasks {
		t.Running = appCounts[t.Name]
	}

	return err
}

// stopStartTask updates the number of running tasks using the Marathon API.
func (m *MarathonScheduler) stopStartTask(task *demand.Task) error {
	// Submit a post request to Marathon to match the requested number of the requested app
	// format looks like:
	// PUT http://marathon:8080/v2/apps/<app>
	//  Request:
	//  {
	//    "instances": 8
	//  }
	url := m.baseMarathonURL + "apps/" + task.Name
	log.Infof("Start/stop PUT: %s", url)

	payload := startStopPayload{
		Instances: task.Demand,
	}
	w := &bytes.Buffer{}
	encoder := json.NewEncoder(w)
	err := encoder.Encode(&payload)
	if err != nil {
		return fmt.Errorf("Failed to encode json. %v", err)
	}

	req, err := http.NewRequest("PUT", url, w)
	if err != nil {
		return fmt.Errorf("Failed to build PUT request err %v", err)
	}

	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	defer resp.Body.Close()

	if err != nil {
		return fmt.Errorf("start/stop err %v", err)
	}

	if resp.StatusCode > 299 {
		return fmt.Errorf("error response from marathon. %s", resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("start/stop read err %v", err)
	}

	// We do nothing with this body
	s := string(body)
	log.Infof("start/stop json: %s", s)

	// Now we've asked for this many, update the currentcount
	task.Requested = task.Demand

	return nil
}

// getBaseMarathonURL returns the base API path.
func getBaseMarathonURL(marathonAPI string) string {
	return marathonAPI + "/v2/"
}
<commit_msg>Remove map of running services as its not required.<commit_after>// Package marathon scheduler integration
package marathon

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

	"github.com/op/go-logging"

	"github.com/microscaling/microscaling/demand"
	"github.com/microscaling/microscaling/scheduler"
)

var log = logging.MustGetLogger("mssscheduler")

// MarathonScheduler holds Marathon API URL.
type MarathonScheduler struct {
	baseMarathonURL string
}

// AppsMessage from the Marathon Apps API.
type AppsMessage struct {
	Apps []App `json:"apps"`
}

// App from the Marathon Apps API.
type App struct {
	ID        string `json:"id"`
	Instances int    `json:"instances"`
}

var (
	httpClient = &http.Client{
		// TODO Make timeout configurable.
		Timeout: 10 * time.Second,
	}
)

// NewScheduler returns a pointer to the scheduler.
func NewScheduler(marathonAPI string) *MarathonScheduler {
	return &MarathonScheduler{
		baseMarathonURL: getBaseMarathonURL(marathonAPI),
	}
}

type startStopPayload struct {
	Instances int `json:"instances"`
}

// compile-time assert that we implement the right interface
var _ scheduler.Scheduler = (*MarathonScheduler)(nil)

// InitScheduler initializes the scheduler.
func (m *MarathonScheduler) InitScheduler(task *demand.Task) (err error) {
	return err
}

// StopStartTasks by calling the Marathon scaling API.
func (m *MarathonScheduler) StopStartTasks(tasks *demand.Tasks) error {
	// Create tasks if there aren't enough of them, and stop them if there are too many
	var tooMany []*demand.Task
	var tooFew []*demand.Task
	var diff int
	var err error = nil

	// TODO: Consider checking the number running before we start & stop
	for _, task := range tasks.Tasks {
		if task.Demand > task.Requested {
			// There aren't enough of these containers yet
			tooFew = append(tooFew, task)
		}
		if task.Demand < task.Requested {
			// there aren't enough of these containers yet
			tooMany = append(tooMany, task)
		}
	}

	// Scale down first to free up resources
	for _, task := range tooMany {
		diff = task.Requested - task.Demand
		log.Infof("Stop %d of task %s", diff, task.Name)
		err = m.stopStartTask(task)
		if err != nil {
			log.Errorf("Couldn't stop %s: %v ", task.Name, err)
		}
		log.Infof("now have %d", task.Requested)
		// task.Requested = task.Demand
		//tasks[name] = task
	}

	// Now we can scale up
	for _, task := range tooFew {
		diff = task.Demand - task.Requested
		log.Infof("Start %d of task %s", diff, task.Name)
		err = m.stopStartTask(task)
		if err != nil {
			log.Errorf("Couldn't start %s: %v ", task.Name, err)
		}
		log.Infof("now have %d", task.Requested)
		//tasks[name] = task
	}

	log.Infof("%v", tasks)
	return err
}

// Marathon API to get the current running tasks.
func getJSONGet(url string) (body []byte, err error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Errorf("Failed to build API GET request err %v", err)
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	resp, err := httpClient.Do(req)
	if err != nil {
		log.Errorf("Failed to GET err %v", err)
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		log.Errorf("Http error %d: %s", resp.StatusCode, resp.Status)
	}

	body, err = ioutil.ReadAll(resp.Body)

	return body, err
}

// CountAllTasks tells us how many instances of each task are currently running.
func (m *MarathonScheduler) CountAllTasks(running *demand.Tasks) error {
	var (
		err         error
		appsMessage AppsMessage
	)

	url := m.baseMarathonURL + "apps/"

	body, err := getJSONGet(url)
	if err != nil {
		log.Errorf("Error getting Marathon Apps %v", err)
	}

	err = json.Unmarshal(body, &appsMessage)
	if err != nil {
		log.Errorf("Error %v unmarshalling from %s", err, string(body[:]))
	}

	appCounts := make(map[string]int)

	// Remove leading slash from App IDs and set the instance counts.
	for _, app := range appsMessage.Apps {
		appCounts[strings.Replace(app.ID, "/", "", 1)] = app.Instances
	}

	// Set running counts. Defaults to 0 if the App does not exist.
	tasks := running.Tasks
	for _, t := range tasks {
		t.Running = appCounts[t.Name]
	}

	return err
}

// stopStartTask updates the number of running tasks using the Marathon API.
func (m *MarathonScheduler) stopStartTask(task *demand.Task) error {
	// Submit a post request to Marathon to match the requested number of the requested app
	// format looks like:
	// PUT http://marathon:8080/v2/apps/<app>
	//  Request:
	//  {
	//    "instances": 8
	//  }
	url := m.baseMarathonURL + "apps/" + task.Name
	log.Infof("Start/stop PUT: %s", url)

	payload := startStopPayload{
		Instances: task.Demand,
	}
	w := &bytes.Buffer{}
	encoder := json.NewEncoder(w)
	err := encoder.Encode(&payload)
	if err != nil {
		return fmt.Errorf("Failed to encode json. %v", err)
	}

	req, err := http.NewRequest("PUT", url, w)
	if err != nil {
		return fmt.Errorf("Failed to build PUT request err %v", err)
	}

	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	defer resp.Body.Close()

	if err != nil {
		return fmt.Errorf("start/stop err %v", err)
	}

	if resp.StatusCode > 299 {
		return fmt.Errorf("error response from marathon. %s", resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("start/stop read err %v", err)
	}

	// We do nothing with this body
	s := string(body)
	log.Infof("start/stop json: %s", s)

	// Now we've asked for this many, update the currentcount
	task.Requested = task.Demand

	return nil
}

// getBaseMarathonURL returns the base API path.
func getBaseMarathonURL(marathonAPI string) string {
	return marathonAPI + "/v2/"
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
)

func NewPkg(dir, importPathRoot, importPath, rev string) *Pkg {
	return &Pkg{
		Dir:            dir,
		ImportPathRoot: importPathRoot,
		ImportPath:     importPath,
		Rev:            rev,
		goFiles:        newSet(),
	}
}

type Pkg struct {
	Dir            string
	ImportPathRoot string
	ImportPath     string
	Rev            string
	goFiles        set
}

func (p Pkg) String() string {
	return fmt.Sprintf("Dir=%s, ImportPathRoot=%s, ImportPath=%s, Rev=%s, Files=%v", p.Dir, p.ImportPathRoot, p.ImportPath, p.Rev, p.GoFiles())
}

func (p *Pkg) addGoFiles(files []string) {
	for _, file := range files {
		p.goFiles.Add(file)
	}
}

func (p *Pkg) GoFiles() []string {
	return p.goFiles.ToSliceString()
}

type PkgLoader struct {
	GoPath string
}

func (pl *PkgLoader) Load(importPaths ...string) ([]*Pkg, error) {
	err := goGet(pl.GoPath, importPaths...)
	if err != nil {
		return nil, err
	}

	ps, err := listPkgs(importPaths...)
	if err != nil {
		return nil, err
	}

	pkgs, err := pl.loadPkgs(ps)
	if err != nil {
		return nil, err
	}

	return pkgs, nil
}

func (pl *PkgLoader) loadPkgs(a []*pkg) ([]*Pkg, error) {
	pkgMap := make(map[string]*Pkg)
	seen := newSet()

	err := pl.recursiveLoadPkgs(a, pkgMap, seen)
	if err != nil {
		return nil, err
	}

	var pkgs []*Pkg
	for _, pkg := range pkgMap {
		pkgs = append(pkgs, pkg)
	}

	return pkgs, nil
}

func (pl *PkgLoader) getUnloadPkgs(ps []*pkg, seen set) []*pkg {
	unloadPkgsMap := make(map[string]*pkg)
	for _, p := range ps {
		if p.Standard {
			continue
		}

		if seen.Contains(p.ImportPath) {
			continue
		}

		if _, ok := unloadPkgsMap[p.ImportPath]; !ok {
			unloadPkgsMap[p.ImportPath] = p
		}
	}

	var unloadPkgs []*pkg
	for _, p := range unloadPkgsMap {
		unloadPkgs = append(unloadPkgs, p)
	}

	return unloadPkgs
}

func (pl *PkgLoader) getDepPkgs(ps []*pkg) ([]*pkg, error) {
	s := newSet()
	for _, p := range ps {
		// for its test dependencies
		imports := append(p.TestImports, p.XTestImports...)
		// for its dependencies
		imports = append(imports, p.Deps...)

		for _, i := range imports {
			s.Add(i)
		}
	}

	return listPkgs(s.ToSliceString()...)
}

func (pl *PkgLoader) recursiveLoadPkgs(ps []*pkg, pkgMap map[string]*Pkg, seen set) error {
	unloadPkgs := pl.getUnloadPkgs(ps, seen)
	if len(unloadPkgs) == 0 {
		return nil
	}

	err := pl.doLoadPkgs(unloadPkgs, pkgMap, seen)
	if err != nil {
		return err
	}

	depPkgs, err := pl.getDepPkgs(ps)
	if err != nil {
		return err
	}

	return pl.recursiveLoadPkgs(depPkgs, pkgMap, seen)
}

func (pl *PkgLoader) getImportPaths(ps []*pkg) []string {
	var importPaths []string
	for _, p := range ps {
		importPaths = append(importPaths, p.ImportPath)
	}

	return importPaths
}

func (pl *PkgLoader) doLoadPkgs(ps []*pkg, pkgMap map[string]*Pkg, seen set) error {
	importPaths := pl.getImportPaths(ps)

	// make sure dependencies are downloaded
	err := goGet(pl.GoPath, importPaths...)
	if err != nil {
		return err
	}

	return pl.cachePkgs(ps, pkgMap, seen)
}

func (pl *PkgLoader) cachePkgs(ps []*pkg, pkgMap map[string]*Pkg, seen set) error {
	for _, p := range ps {
		vcs, importPathRoot, err := VCSFromDir(p.Dir, filepath.Join(p.Root, "src"))
		if err != nil {
			return err
		}

		pkg, ok := pkgMap[p.ImportPath]
		if !ok {
			rev, err := vcs.Identify(p.Dir)
			if err != nil {
				return err
			}

			pkg = NewPkg(p.Dir, importPathRoot, p.ImportPath, rev)
			pkgMap[p.ImportPath] = pkg
		}

		pkg.addGoFiles(p.AllGoFiles())

		seen.Add(p.ImportPath)
	}

	return nil
}

type pkg struct {
	Dir        string
	Root       string
	ImportPath string
	Deps       []string
	Standard   bool

	GoFiles        []string
	CgoFiles       []string
	IgnoredGoFiles []string

	TestGoFiles  []string
	TestImports  []string
	XTestGoFiles []string
	XTestImports []string

	Error struct {
		Err string
	}
}

func (p *pkg) AllGoFiles() (a []string) {
	a = append(a, pathOf(p.Dir, p.GoFiles)...)
	a = append(a, pathOf(p.Dir, p.CgoFiles)...)
	a = append(a, pathOf(p.Dir, p.TestGoFiles)...)
	a = append(a, pathOf(p.Dir, p.XTestGoFiles)...)
	a = append(a, pathOf(p.Dir, p.IgnoredGoFiles)...)

	return a
}

func pathOf(dir string, files []string) []string {
	paths := make([]string, 0)
	for _, f := range files {
		paths = append(paths, filepath.Join(dir, f))
	}

	return paths
}

func listPkgs(name ...string) ([]*pkg, error) {
	l := pkgLister{
		Env: goCmdEnv(),
	}
	return l.List(name...)
}

type pkgLister struct {
	Env []string
}

func (l *pkgLister) List(name ...string) (a []*pkg, err error) {
	if len(name) == 0 {
		return nil, nil
	}

	args := []string{"list", "-e", "-json"}
	cmd := exec.Command("go", append(args, name...)...)
	cmd.Env = l.Env

	r, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	cmd.Stderr = os.Stderr
	err = cmd.Start()
	if err != nil {
		return nil, err
	}
	d := json.NewDecoder(r)
	for {
		info := new(pkg)
		err = d.Decode(info)
		if err == io.EOF {
			break
		}
		if err != nil {
			info.Error.Err = err.Error()
		}
		a = append(a, info)
	}

	err = cmd.Wait()
	if err != nil {
		return nil, err
	}

	return a, nil
}
<commit_msg>Modified importPaths to add subdirectories<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
)

func NewPkg(dir, importPathRoot, importPath, rev string) *Pkg {
	return &Pkg{
		Dir:            dir,
		ImportPathRoot: importPathRoot,
		ImportPath:     importPath,
		Rev:            rev,
		goFiles:        newSet(),
	}
}

type Pkg struct {
	Dir            string
	ImportPathRoot string
	ImportPath     string
	Rev            string
	goFiles        set
}

func (p Pkg) String() string {
	return fmt.Sprintf("Dir=%s, ImportPathRoot=%s, ImportPath=%s, Rev=%s, Files=%v", p.Dir, p.ImportPathRoot, p.ImportPath, p.Rev, p.GoFiles())
}

func (p *Pkg) addGoFiles(files []string) {
	for _, file := range files {
		p.goFiles.Add(file)
	}
}

func (p *Pkg) GoFiles() []string {
	return p.goFiles.ToSliceString()
}

type PkgLoader struct {
	GoPath string
}

func (pl *PkgLoader) Load(importPaths ...string) ([]*Pkg, error) {
	err := goGet(pl.GoPath, importPaths...)
	if err != nil {
		return nil, err
	}

	ps, err := listPkgs(importPaths...)
	if err != nil {
		return nil, err
	}

	pkgs, err := pl.loadPkgs(ps)
	if err != nil {
		return nil, err
	}

	return pkgs, nil
}

func (pl *PkgLoader) loadPkgs(a []*pkg) ([]*Pkg, error) {
	pkgMap := make(map[string]*Pkg)
	seen := newSet()

	err := pl.recursiveLoadPkgs(a, pkgMap, seen)
	if err != nil {
		return nil, err
	}

	var pkgs []*Pkg
	for _, pkg := range pkgMap {
		pkgs = append(pkgs, pkg)
	}

	return pkgs, nil
}

func (pl *PkgLoader) getUnloadPkgs(ps []*pkg, seen set) []*pkg {
	unloadPkgsMap := make(map[string]*pkg)
	for _, p := range ps {
		if p.Standard {
			continue
		}

		if seen.Contains(p.ImportPath) {
			continue
		}

		if _, ok := unloadPkgsMap[p.ImportPath]; !ok {
			unloadPkgsMap[p.ImportPath] = p
		}
	}

	var unloadPkgs []*pkg
	for _, p := range unloadPkgsMap {
		unloadPkgs = append(unloadPkgs, p)
	}

	return unloadPkgs
}

func (pl *PkgLoader) getDepPkgs(ps []*pkg) ([]*pkg, error) {
	s := newSet()
	for _, p := range ps {
		// for its test dependencies
		imports := append(p.TestImports, p.XTestImports...)
		// for its dependencies
		imports = append(imports, p.Deps...)

		for _, i := range imports {
			s.Add(i)
		}
	}

	return listPkgs(s.ToSliceString()...)
}

func (pl *PkgLoader) recursiveLoadPkgs(ps []*pkg, pkgMap map[string]*Pkg, seen set) error {
	unloadPkgs := pl.getUnloadPkgs(ps, seen)
	if len(unloadPkgs) == 0 {
		return nil
	}

	err := pl.doLoadPkgs(unloadPkgs, pkgMap, seen)
	if err != nil {
		return err
	}

	depPkgs, err := pl.getDepPkgs(ps)
	if err != nil {
		return err
	}

	return pl.recursiveLoadPkgs(depPkgs, pkgMap, seen)
}

func (pl *PkgLoader) getImportPaths(ps []*pkg) []string {
	var importPaths []string
	for _, p := range ps {
		importPaths = append(importPaths, p.ImportPath)
	}

	return importPaths
}

func (pl *PkgLoader) doLoadPkgs(ps []*pkg, pkgMap map[string]*Pkg, seen set) error {
	importPaths := pl.getImportPaths(ps)

	// make sure dependencies are downloaded
	err := goGet(pl.GoPath, importPaths...)
	if err != nil {
		return err
	}

	return pl.cachePkgs(ps, pkgMap, seen)
}

func (pl *PkgLoader) cachePkgs(ps []*pkg, pkgMap map[string]*Pkg, seen set) error {
	for _, p := range ps {
		vcs, importPathRoot, err := VCSFromDir(p.Dir, filepath.Join(p.Root, "src"))
		if err != nil {
			return err
		}

		pkg, ok := pkgMap[p.ImportPath]
		if !ok {
			rev, err := vcs.Identify(p.Dir)
			if err != nil {
				return err
			}

			pkg = NewPkg(p.Dir, importPathRoot, p.ImportPath, rev)
			pkgMap[p.ImportPath] = pkg
		}

		pkg.addGoFiles(p.AllGoFiles())

		seen.Add(p.ImportPath)
	}

	return nil
}

type pkg struct {
	Dir        string
	Root       string
	ImportPath string
	Deps       []string
	Standard   bool

	GoFiles        []string
	CgoFiles       []string
	IgnoredGoFiles []string

	TestGoFiles  []string
	TestImports  []string
	XTestGoFiles []string
	XTestImports []string

	Error struct {
		Err string
	}
}

func (p *pkg) AllGoFiles() (a []string) {
	a = append(a, pathOf(p.Dir, p.GoFiles)...)
	a = append(a, pathOf(p.Dir, p.CgoFiles)...)
	a = append(a, pathOf(p.Dir, p.TestGoFiles)...)
	a = append(a, pathOf(p.Dir, p.XTestGoFiles)...)
	a = append(a, pathOf(p.Dir, p.IgnoredGoFiles)...)

	return a
}

func pathOf(dir string, files []string) []string {
	paths := make([]string, 0)
	for _, f := range files {
		paths = append(paths, filepath.Join(dir, f))
	}

	return paths
}

func listPkgs(name ...string) ([]*pkg, error) {
	l := pkgLister{
		Env: goCmdEnv(),
	}
	name = addSubPaths(name)
	return l.List(name...)
}

type pkgLister struct {
	Env []string
}

func (l *pkgLister) List(name ...string) (a []*pkg, err error) {
	if len(name) == 0 {
		return nil, nil
	}

	args := []string{"list", "-e", "-json"}
	cmd := exec.Command("go", append(args, name...)...)
	cmd.Env = l.Env

	r, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	cmd.Stderr = os.Stderr
	err = cmd.Start()
	if err != nil {
		return nil, err
	}
	d := json.NewDecoder(r)
	for {
		info := new(pkg)
		err = d.Decode(info)
		if err == io.EOF {
			break
		}
		if err != nil {
			info.Error.Err = err.Error()
		}
		a = append(a, info)
	}

	err = cmd.Wait()
	if err != nil {
		return nil, err
	}

	return a, nil
}

func addSubPaths(name []string) []string {
	var paths []string
	for _, path := range name {
		paths = append(paths, addSubPath(path))
	}
	return paths
}

func addSubPath(path string) string {
	return path + "/..."
}
<|endoftext|>