<commit_before>//
// command/run.go
//
// Copyright (c) 2016-2017 Junpei Kawamoto
//
// This software is released under the MIT License.
//
// http://opensource.org/licenses/mit-license.php
//

package command

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	colorable "github.com/mattn/go-colorable"
	gitconfig "github.com/tcnksm/go-gitconfig"
	"github.com/ttacon/chalk"
	"github.com/urfave/cli"
)

// SourceArchive defines a name of source archive file.
const SourceArchive = "source.tar.gz"

// RunOpt defines a option parameter for run function.
type RunOpt struct {
	// Same options as DockerfileOpt.
	*DockerfileOpt
	// Travis configuration file.
	Filename string
	// Container name.
	Name string
	// Runtime version to which only versions matching will be run.
	Version string
	// Image tag.
	Tag string
	// Max processors to be used.
	Processors int
	// If true, print Dockerfile and entrypoint.sh.
	Verbose bool
	// If true, not using cache during buidling a docker image.
	NoCache bool
	// If true, omit printing color codes.
	NoColor bool
}

// Run implements the action of this command.
func Run(c *cli.Context) error {

	opt := RunOpt{
		DockerfileOpt: &DockerfileOpt{
			BaseImage:  c.String("base"),
			AptProxy:   c.String("apt-proxy"),
			PypiProxy:  c.String("pypi-proxy"),
			HTTPProxy:  c.String("http-proxy"),
			HTTPSProxy: c.String("https-proxy"),
			NoProxy:    c.String("no-proxy"),
		},
		Filename:   c.Args().First(),
		Name:       c.String("name"),
		Version:    c.String("select"),
		Tag:        c.String("tag"),
		Processors: c.Int("max-processors"),
		Verbose:    c.Bool("verbose"),
		NoCache:    c.Bool("no-cache"),
		NoColor:    c.Bool("no-color"),
	}
	if err := run(&opt); err != nil {
		return cli.NewExitError(err.Error(), 1)
	}
	return nil
}

func run(opt *RunOpt) (err error) {

	var stdout io.Writer
	if opt.NoColor {
		stdout = colorable.NewNonColorable(os.Stdout)
	} else {
		stdout = colorable.NewColorableStdout()
	}

	var dstout io.Writer
	if opt.Verbose {
		dstout = stdout
	} else {
		dstout = ioutil.Discard
	}

	// Prepare to be canceled.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill, syscall.SIGQUIT)
	go func() {
		<-sig
		cancel()
	}()

	// Load a Travis's script file.
	if opt.Filename == "" {
		opt.Filename = ".travis.yml"
	}
	travis, err := NewTravisFromFile(opt.Filename)
	if err != nil {
		return
	}

	// Get repository information.
	origin, err := gitconfig.OriginURL()
	if err != nil {
		return
	}
	opt.Repository = getRepository(origin)

	// Set up the tag name of the container image.
	if opt.Tag == "" {
		opt.Tag = fmt.Sprintf("loci/%s", strings.ToLower(path.Base(opt.Repository)))
	}
	tempDir := filepath.Join(os.TempDir(), opt.Tag)
	if err = os.MkdirAll(tempDir, 0777); err != nil {
		return
	}
	defer os.RemoveAll(tempDir)

	// Archive source files.
	pwd, err := os.Getwd()
	if err != nil {
		return
	}

	fmt.Fprintln(dstout, chalk.Yellow.Color("Creating an archive of source codes."))
	if err = Archive(ctx, pwd, filepath.Join(tempDir, SourceArchive), dstout, os.Stderr); err != nil {
		return
	}

	// Create Dockerfile.
	fmt.Fprintln(dstout, chalk.Yellow.Color("Creating Dockerfile"))
	docker, err := Dockerfile(travis, opt.DockerfileOpt, SourceArchive)
	if err != nil {
		return
	}
	if err = ioutil.WriteFile(filepath.Join(tempDir, "Dockerfile"), docker, 0644); err != nil {
		return
	}
	fmt.Fprintln(dstout, string(docker))

	// Create entrypoint.sh.
	fmt.Fprintln(dstout, chalk.Yellow.Color("Creating entrypoint.sh"))
	entry, err := Entrypoint(travis)
	if err != nil {
		return
	}
	if err = ioutil.WriteFile(filepath.Join(tempDir, "entrypoint.sh"), entry, 0644); err != nil {
		return
	}
	fmt.Fprintln(dstout, string(entry))

	argset, err := travis.ArgumentSet()
	if err != nil {
		return
	}

	// Start testing with goroutines.
	fmt.Fprintln(dstout, chalk.Yellow.Color("Start testing."))
	display, err := NewDisplay(cancel)
	if err != nil {
		return
	}

	var i int
	var wg sync.WaitGroup
	semaphore := make(chan struct{}, opt.Processors)
	errs := NewErrorSet()
	for version, set := range argset {

		if opt.Version != "" && version != opt.Version {
			continue
		}

		wg.Add(1)
		go func(version string, set [][]string) (err error) {
			semaphore <- struct{}{}
			defer func() {
				<-semaphore
				wg.Done()
			}()

			// Build a container image.
			fp, err := os.OpenFile(fmt.Sprintf("loci-build-v%v.log", version), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
			if err != nil {
				errs.Add(version, err)
				return
			}
			defer fp.Close()

			sec := display.AddSection(fmt.Sprintf("Building a image for v%v", version))
			defer display.DeleteSection(sec)

			output := sec.Writer()
			defer output.Close()

			tag := fmt.Sprintf("%v/%v", opt.Tag, version)
			err = Build(ctx, tempDir, tag, version, opt.NoCache, io.MultiWriter(output, colorable.NewColorable(fp)))
			if err != nil {
				errs.Add(version, err)
				return
			}

			for _, envs := range set {

				wg.Add(1)
				go func(envs []string) {
					semaphore <- struct{}{}
					defer func() {
						<-semaphore
						wg.Done()
					}()

					// Run tests in a sandbox.
					fp, err := os.OpenFile(
						fmt.Sprintf("loci-v%v.log", strings.Join(append([]string{version}, envs...), "-")),
						os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
					if err != nil {
						errs.Add(fmt.Sprintf("%v:%v", version, envs), err)
						return
					}
					defer fp.Close()

					sec := display.AddSection(fmt.Sprintf("Running tests (v%v: %v)", version, envs))
					defer display.DeleteSection(sec)

					output := sec.Writer()
					defer output.Close()

					name := opt.Name
					if name != "" {
						i++
						name = fmt.Sprintf("%s-%d", name, i)
					}

					err = Start(ctx, tag, name, envs, io.MultiWriter(output, colorable.NewColorable(fp)))
					if err == context.Canceled {
						errs.Add(fmt.Sprintf("%v:%v", version, envs), fmt.Errorf(sec.String()))
					} else if err != nil {
						errs.Add(fmt.Sprintf("%v:%v", version, envs), fmt.Errorf("%s\n%s", chalk.Red.Color(err.Error()), sec.String()))
					}
					return

				}(envs)

			}

			return

		}(version, set)

	}

	wg.Wait()
	display.Close()

	if errs.Size() == 0 {
		fmt.Fprintln(stdout, chalk.Green.Color("All tests have been passed."))
	} else {
		err = cli.NewMultiError(errs.GetList()...)
	}
	return

}

// getRepository returns the repository path from a given remote URL of
// origin repository. The repository path consists of a URL without
// sheme, user name, password, and .git suffix.
func getRepository(origin string) (res string) {

	switch {
	case strings.Contains(origin, "@"):
		res = strings.Replace(strings.Split(origin, "@")[1], ":", "/", 1)
	case strings.HasPrefix(origin, "http://"):
		res = origin[len("http://"):]
	case strings.HasPrefix(origin, "https://"):
		res = origin[len("https://"):]
	default:
		res = strings.Replace(origin, ":", "/", 1)
	}
	if strings.HasSuffix(res, ".git") {
		res = res[:len(res)-len(".git")]
	}

	return

}
<commit_msg>Fixed to support no-color mode.<commit_after>//
// command/run.go
//
// Copyright (c) 2016-2017 Junpei Kawamoto
//
// This software is released under the MIT License.
//
// http://opensource.org/licenses/mit-license.php
//

package command

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	colorable "github.com/mattn/go-colorable"
	gitconfig "github.com/tcnksm/go-gitconfig"
	"github.com/ttacon/chalk"
	"github.com/urfave/cli"
)

// SourceArchive defines a name of source archive file.
const SourceArchive = "source.tar.gz"

// RunOpt defines a option parameter for run function.
type RunOpt struct {
	// Same options as DockerfileOpt.
	*DockerfileOpt
	// Travis configuration file.
	Filename string
	// Container name.
	Name string
	// Runtime version to which only versions matching will be run.
	Version string
	// Image tag.
	Tag string
	// Max processors to be used.
	Processors int
	// If true, print Dockerfile and entrypoint.sh.
	Verbose bool
	// If true, not using cache during buidling a docker image.
	NoCache bool
	// If true, omit printing color codes.
	NoColor bool
}

// Run implements the action of this command.
func Run(c *cli.Context) error {

	opt := RunOpt{
		DockerfileOpt: &DockerfileOpt{
			BaseImage:  c.String("base"),
			AptProxy:   c.String("apt-proxy"),
			PypiProxy:  c.String("pypi-proxy"),
			HTTPProxy:  c.String("http-proxy"),
			HTTPSProxy: c.String("https-proxy"),
			NoProxy:    c.String("no-proxy"),
		},
		Filename:   c.Args().First(),
		Name:       c.String("name"),
		Version:    c.String("select"),
		Tag:        c.String("tag"),
		Processors: c.Int("max-processors"),
		Verbose:    c.Bool("verbose"),
		NoCache:    c.Bool("no-cache"),
		NoColor:    c.Bool("no-color"),
	}
	if err := run(&opt); err != nil {
		return cli.NewExitError(err.Error(), 1)
	}
	return nil
}

func run(opt *RunOpt) (err error) {

	var stdout io.Writer
	if opt.NoColor {
		stdout = colorable.NewNonColorable(os.Stdout)
		cli.ErrWriter = colorable.NewNonColorable(cli.ErrWriter)
	} else {
		stdout = colorable.NewColorableStdout()
	}

	var dstout io.Writer
	if opt.Verbose {
		dstout = stdout
	} else {
		dstout = ioutil.Discard
	}

	// Prepare to be canceled.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill, syscall.SIGQUIT)
	go func() {
		<-sig
		cancel()
	}()

	// Load a Travis's script file.
	if opt.Filename == "" {
		opt.Filename = ".travis.yml"
	}
	travis, err := NewTravisFromFile(opt.Filename)
	if err != nil {
		return
	}

	// Get repository information.
	origin, err := gitconfig.OriginURL()
	if err != nil {
		return
	}
	opt.Repository = getRepository(origin)

	// Set up the tag name of the container image.
	if opt.Tag == "" {
		opt.Tag = fmt.Sprintf("loci/%s", strings.ToLower(path.Base(opt.Repository)))
	}
	tempDir := filepath.Join(os.TempDir(), opt.Tag)
	if err = os.MkdirAll(tempDir, 0777); err != nil {
		return
	}
	defer os.RemoveAll(tempDir)

	// Archive source files.
	pwd, err := os.Getwd()
	if err != nil {
		return
	}

	fmt.Fprintln(dstout, chalk.Yellow.Color("Creating an archive of source codes."))
	if err = Archive(ctx, pwd, filepath.Join(tempDir, SourceArchive), dstout, os.Stderr); err != nil {
		return
	}

	// Create Dockerfile.
	fmt.Fprintln(dstout, chalk.Yellow.Color("Creating Dockerfile"))
	docker, err := Dockerfile(travis, opt.DockerfileOpt, SourceArchive)
	if err != nil {
		return
	}
	if err = ioutil.WriteFile(filepath.Join(tempDir, "Dockerfile"), docker, 0644); err != nil {
		return
	}
	fmt.Fprintln(dstout, string(docker))

	// Create entrypoint.sh.
	fmt.Fprintln(dstout, chalk.Yellow.Color("Creating entrypoint.sh"))
	entry, err := Entrypoint(travis)
	if err != nil {
		return
	}
	if err = ioutil.WriteFile(filepath.Join(tempDir, "entrypoint.sh"), entry, 0644); err != nil {
		return
	}
	fmt.Fprintln(dstout, string(entry))

	argset, err := travis.ArgumentSet()
	if err != nil {
		return
	}

	// Start testing with goroutines.
	fmt.Fprintln(dstout, chalk.Yellow.Color("Start testing."))
	display, err := NewDisplay(cancel)
	if err != nil {
		return
	}

	var i int
	var wg sync.WaitGroup
	semaphore := make(chan struct{}, opt.Processors)
	errs := NewErrorSet()
	for version, set := range argset {

		if opt.Version != "" && version != opt.Version {
			continue
		}

		wg.Add(1)
		go func(version string, set [][]string) (err error) {
			semaphore <- struct{}{}
			defer func() {
				<-semaphore
				wg.Done()
			}()

			// Build a container image.
			fp, err := os.OpenFile(fmt.Sprintf("loci-build-v%v.log", version), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
			if err != nil {
				errs.Add(version, err)
				return
			}
			defer fp.Close()

			sec := display.AddSection(fmt.Sprintf("Building a image for v%v", version))
			defer display.DeleteSection(sec)

			var output io.Writer
			writer := sec.Writer()
			defer writer.Close()
			output = writer
			if opt.NoColor {
				output = colorable.NewNonColorable(output)
			}

			tag := fmt.Sprintf("%v/%v", opt.Tag, version)
			err = Build(ctx, tempDir, tag, version, opt.NoCache, io.MultiWriter(output, colorable.NewColorable(fp)))
			if err != nil {
				errs.Add(version, err)
				return
			}

			for _, envs := range set {

				wg.Add(1)
				go func(envs []string) {
					semaphore <- struct{}{}
					defer func() {
						<-semaphore
						wg.Done()
					}()

					// Run tests in a sandbox.
					fp, err := os.OpenFile(
						fmt.Sprintf("loci-v%v.log", strings.Join(append([]string{version}, envs...), "-")),
						os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
					if err != nil {
						errs.Add(fmt.Sprintf("%v:%v", version, envs), err)
						return
					}
					defer fp.Close()

					sec := display.AddSection(fmt.Sprintf("Running tests (v%v: %v)", version, envs))
					defer display.DeleteSection(sec)

					var output io.Writer
					writer := sec.Writer()
					defer writer.Close()
					output = writer
					if opt.NoColor {
						output = colorable.NewNonColorable(output)
					}

					name := opt.Name
					if name != "" {
						i++
						name = fmt.Sprintf("%s-%d", name, i)
					}

					err = Start(ctx, tag, name, envs, io.MultiWriter(output, colorable.NewColorable(fp)))
					if err == context.Canceled {
						errs.Add(fmt.Sprintf("%v:%v", version, envs), fmt.Errorf(sec.String()))
					} else if err != nil {
						errs.Add(fmt.Sprintf("%v:%v", version, envs), fmt.Errorf("%s\n%s", chalk.Red.Color(err.Error()), sec.String()))
					}
					return

				}(envs)

			}

			return

		}(version, set)

	}

	wg.Wait()
	display.Close()

	if errs.Size() == 0 {
		fmt.Fprintln(stdout, chalk.Green.Color("All tests have been passed."))
	} else {
		err = cli.NewMultiError(errs.GetList()...)
	}
	return

}

// getRepository returns the repository path from a given remote URL of
// origin repository. The repository path consists of a URL without
// sheme, user name, password, and .git suffix.
func getRepository(origin string) (res string) {

	switch {
	case strings.Contains(origin, "@"):
		res = strings.Replace(strings.Split(origin, "@")[1], ":", "/", 1)
	case strings.HasPrefix(origin, "http://"):
		res = origin[len("http://"):]
	case strings.HasPrefix(origin, "https://"):
		res = origin[len("https://"):]
	default:
		res = strings.Replace(origin, ":", "/", 1)
	}
	if strings.HasSuffix(res, ".git") {
		res = res[:len(res)-len(".git")]
	}

	return

}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd Authors
//
// 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 mvcc

import (
	"sync/atomic"
	"testing"

	"github.com/coreos/etcd/lease"
	"github.com/coreos/etcd/mvcc/backend"
)

type fakeConsistentIndex uint64

func (i *fakeConsistentIndex) ConsistentIndex() uint64 {
	return atomic.LoadUint64((*uint64)(i))
}

func BenchmarkStorePut(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// arbitrary number of bytes
	bytesN := 64
	keys := createBytesSlice(bytesN, b.N)
	vals := createBytesSlice(bytesN, b.N)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Put(keys[i], vals[i], lease.NoLease)
	}
}

func BenchmarkStoreRangeOneKey(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// 64 byte key/val
	key, val := createBytesSlice(64, 1), createBytesSlice(64, 1)
	s.Put(key[0], val[0], lease.NoLease)
	// Force into boltdb tx instead of backend read tx.
	s.Commit()

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Range(key[0], nil, RangeOptions{})
	}
}

func BenchmarkConsistentIndex(b *testing.B) {
	fci := fakeConsistentIndex(10)
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &fci)
	defer cleanup(s, be, tmpPath)

	tx := s.b.BatchTx()
	tx.Lock()
	s.saveIndex(tx)
	tx.Unlock()

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.ConsistentIndex()
	}
}

// BenchmarkStoreTxnPutUpdate is same as above, but instead updates single key
func BenchmarkStorePutUpdate(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// arbitrary number of bytes
	keys := createBytesSlice(64, 1)
	vals := createBytesSlice(1024, 1)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Put(keys[0], vals[0], lease.NoLease)
	}
}

// BenchmarkStoreTxnPut benchmarks the Put operation
// with transaction begin and end, where transaction involves
// some synchronization operations, such as mutex locking.
func BenchmarkStoreTxnPut(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// arbitrary number of bytes
	bytesN := 64
	keys := createBytesSlice(bytesN, b.N)
	vals := createBytesSlice(bytesN, b.N)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		txn := s.Write()
		txn.Put(keys[i], vals[i], lease.NoLease)
		txn.End()
	}
}

// benchmarkStoreRestore benchmarks the restore operation
func benchmarkStoreRestore(revsPerKey int, b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	// use closure to capture 's' to pick up the reassignment
	defer func() { cleanup(s, be, tmpPath) }()

	// arbitrary number of bytes
	bytesN := 64
	keys := createBytesSlice(bytesN, b.N)
	vals := createBytesSlice(bytesN, b.N)

	for i := 0; i < b.N; i++ {
		for j := 0; j < revsPerKey; j++ {
			txn := s.Write()
			txn.Put(keys[i], vals[i], lease.NoLease)
			txn.End()
		}
	}
	s.Close()

	b.ReportAllocs()
	b.ResetTimer()
	s = NewStore(be, &lease.FakeLessor{}, &i)
}

func BenchmarkStoreRestoreRevs1(b *testing.B) {
	benchmarkStoreRestore(1, b)
}

func BenchmarkStoreRestoreRevs10(b *testing.B) {
	benchmarkStoreRestore(10, b)
}

func BenchmarkStoreRestoreRevs20(b *testing.B) {
	benchmarkStoreRestore(20, b)
}
<commit_msg>mvcc: add range benchmark for fetching 100 keys<commit_after>// Copyright 2015 The etcd Authors
//
// 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 mvcc

import (
	"sync/atomic"
	"testing"

	"github.com/coreos/etcd/lease"
	"github.com/coreos/etcd/mvcc/backend"
)

type fakeConsistentIndex uint64

func (i *fakeConsistentIndex) ConsistentIndex() uint64 {
	return atomic.LoadUint64((*uint64)(i))
}

func BenchmarkStorePut(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// arbitrary number of bytes
	bytesN := 64
	keys := createBytesSlice(bytesN, b.N)
	vals := createBytesSlice(bytesN, b.N)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Put(keys[i], vals[i], lease.NoLease)
	}
}

func BenchmarkStoreRangeKey1(b *testing.B)   { benchmarkStoreRange(b, 1) }
func BenchmarkStoreRangeKey100(b *testing.B) { benchmarkStoreRange(b, 100) }

func benchmarkStoreRange(b *testing.B, n int) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// 64 byte key/val
	keys, val := createBytesSlice(64, n), createBytesSlice(64, 1)
	for i := range keys {
		s.Put(keys[i], val[0], lease.NoLease)
	}
	// Force into boltdb tx instead of backend read tx.
	s.Commit()

	var begin, end []byte
	if n == 1 {
		begin, end = keys[0], nil
	} else {
		begin, end = []byte{}, []byte{}
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Range(begin, end, RangeOptions{})
	}
}

func BenchmarkConsistentIndex(b *testing.B) {
	fci := fakeConsistentIndex(10)
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &fci)
	defer cleanup(s, be, tmpPath)

	tx := s.b.BatchTx()
	tx.Lock()
	s.saveIndex(tx)
	tx.Unlock()

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.ConsistentIndex()
	}
}

// BenchmarkStoreTxnPutUpdate is same as above, but instead updates single key
func BenchmarkStorePutUpdate(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// arbitrary number of bytes
	keys := createBytesSlice(64, 1)
	vals := createBytesSlice(1024, 1)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		s.Put(keys[0], vals[0], lease.NoLease)
	}
}

// BenchmarkStoreTxnPut benchmarks the Put operation
// with transaction begin and end, where transaction involves
// some synchronization operations, such as mutex locking.
func BenchmarkStoreTxnPut(b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	defer cleanup(s, be, tmpPath)

	// arbitrary number of bytes
	bytesN := 64
	keys := createBytesSlice(bytesN, b.N)
	vals := createBytesSlice(bytesN, b.N)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		txn := s.Write()
		txn.Put(keys[i], vals[i], lease.NoLease)
		txn.End()
	}
}

// benchmarkStoreRestore benchmarks the restore operation
func benchmarkStoreRestore(revsPerKey int, b *testing.B) {
	var i fakeConsistentIndex
	be, tmpPath := backend.NewDefaultTmpBackend()
	s := NewStore(be, &lease.FakeLessor{}, &i)
	// use closure to capture 's' to pick up the reassignment
	defer func() { cleanup(s, be, tmpPath) }()

	// arbitrary number of bytes
	bytesN := 64
	keys := createBytesSlice(bytesN, b.N)
	vals := createBytesSlice(bytesN, b.N)

	for i := 0; i < b.N; i++ {
		for j := 0; j < revsPerKey; j++ {
			txn := s.Write()
			txn.Put(keys[i], vals[i], lease.NoLease)
			txn.End()
		}
	}
	s.Close()

	b.ReportAllocs()
	b.ResetTimer()
	s = NewStore(be, &lease.FakeLessor{}, &i)
}

func BenchmarkStoreRestoreRevs1(b *testing.B) {
	benchmarkStoreRestore(1, b)
}

func BenchmarkStoreRestoreRevs10(b *testing.B) {
	benchmarkStoreRestore(10, b)
}

func BenchmarkStoreRestoreRevs20(b *testing.B) {
	benchmarkStoreRestore(20, b)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"time"

	"github.com/getlantern/golog"
	"github.com/getlantern/zenodb"
	"github.com/getlantern/zenodb/rpc"
)

var (
	log = golog.LoggerFor("zeno")

	dbdir             = flag.String("dbdir", "zenodb", "The directory in which to store the database files, defaults to ./zenodb")
	schema            = flag.String("schema", "schema.yaml", "Location of schema file, defaults to ./schema.yaml")
	ispdb             = flag.String("ispdb", "", "In order to enable ISP functions, point this to an IP2Location Lite ISP database file like the one here - https://lite.ip2location.com/database/ip-asn")
	fresh             = flag.Bool("fresh", false, "Set this flag to include data not yet flushed from memstore in query results")
	vtime             = flag.Bool("vtime", false, "Set this flag to use virtual instead of real time. When using virtual time, the advancement of time will be governed by the timestamps received via insterts.")
	walSync           = flag.Duration("walsync", 5*time.Second, "How frequently to sync the WAL to disk. Set to 0 to sync after every write. Defaults to 5 seconds.")
	maxWALAge         = flag.Duration("maxwalage", 336*time.Hour, "Maximum age for WAL files. Files older than this will be deleted. Defaults to 336 hours (2 weeks).")
	walCompressionAge = flag.Duration("walcompressage", 24*time.Hour, "Age at which to start compressing WAL files with gzip. Defaults to 24 hours.")
	addr              = flag.String("addr", "localhost:17712", "The address at which to listen for gRPC connections, defaults to localhost:17712")
	httpAddr          = flag.String("http-addr", "localhost:17713", "The address at which to listen for JSON over HTTP connections, defaults to localhost:17713")
	pprofAddr         = flag.String("pprofaddr", "localhost:4000", "if specified, will listen for pprof connections at the specified tcp address")
)

func main() {
	flag.Parse()

	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatalf("Unable to listen for gRPC connections at %v: %v", *addr, err)
	}

	hl, err := net.Listen("tcp", *httpAddr)
	if err != nil {
		log.Fatalf("Unable to listen for HTTP connections at %v: %v", *httpAddr, err)
	}

	db, err := zenodb.NewDB(&zenodb.DBOpts{
		Dir:                    *dbdir,
		SchemaFile:             *schema,
		ISPDatabase:            *ispdb,
		IncludeMemStoreInQuery: *fresh,
		VirtualTime:            *vtime,
		WALSyncInterval:        *walSync,
		MaxWALAge:              *maxWALAge,
		WALCompressionAge:      *walCompressionAge,
	})

	if err != nil {
		log.Fatalf("Unable to open database at %v: %v", *dbdir, err)
	}
	fmt.Printf("Opened database at %v\n", *dbdir)

	fmt.Printf("Listening for gRPC connections at %v\n", l.Addr())
	fmt.Printf("Listening for HTTP connections at %v\n", hl.Addr())

	go serveHTTP(db, hl)
	serveRPC(db, l)
}

func serveRPC(db *zenodb.DB, l net.Listener) {
	err := rpc.Serve(db, l)
	if err != nil {
		log.Fatalf("Error serving gRPC: %v", err)
	}
}
<commit_msg>Compressing WAL at 1 hour<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"time"

	"github.com/getlantern/golog"
	"github.com/getlantern/zenodb"
	"github.com/getlantern/zenodb/rpc"
)

var (
	log = golog.LoggerFor("zeno")

	dbdir             = flag.String("dbdir", "zenodb", "The directory in which to store the database files, defaults to ./zenodb")
	schema            = flag.String("schema", "schema.yaml", "Location of schema file, defaults to ./schema.yaml")
	ispdb             = flag.String("ispdb", "", "In order to enable ISP functions, point this to an IP2Location Lite ISP database file like the one here - https://lite.ip2location.com/database/ip-asn")
	fresh             = flag.Bool("fresh", false, "Set this flag to include data not yet flushed from memstore in query results")
	vtime             = flag.Bool("vtime", false, "Set this flag to use virtual instead of real time. When using virtual time, the advancement of time will be governed by the timestamps received via insterts.")
	walSync           = flag.Duration("walsync", 5*time.Second, "How frequently to sync the WAL to disk. Set to 0 to sync after every write. Defaults to 5 seconds.")
	maxWALAge         = flag.Duration("maxwalage", 336*time.Hour, "Maximum age for WAL files. Files older than this will be deleted. Defaults to 336 hours (2 weeks).")
	walCompressionAge = flag.Duration("walcompressage", 1*time.Hour, "Age at which to start compressing WAL files with gzip. Defaults to 1 hour.")
	addr              = flag.String("addr", "localhost:17712", "The address at which to listen for gRPC connections, defaults to localhost:17712")
	httpAddr          = flag.String("http-addr", "localhost:17713", "The address at which to listen for JSON over HTTP connections, defaults to localhost:17713")
	pprofAddr         = flag.String("pprofaddr", "localhost:4000", "if specified, will listen for pprof connections at the specified tcp address")
)

func main() {
	flag.Parse()

	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatalf("Unable to listen for gRPC connections at %v: %v", *addr, err)
	}

	hl, err := net.Listen("tcp", *httpAddr)
	if err != nil {
		log.Fatalf("Unable to listen for HTTP connections at %v: %v", *httpAddr, err)
	}

	db, err := zenodb.NewDB(&zenodb.DBOpts{
		Dir:                    *dbdir,
		SchemaFile:             *schema,
		ISPDatabase:            *ispdb,
		IncludeMemStoreInQuery: *fresh,
		VirtualTime:            *vtime,
		WALSyncInterval:        *walSync,
		MaxWALAge:              *maxWALAge,
		WALCompressionAge:      *walCompressionAge,
	})

	if err != nil {
		log.Fatalf("Unable to open database at %v: %v", *dbdir, err)
	}
	fmt.Printf("Opened database at %v\n", *dbdir)

	fmt.Printf("Listening for gRPC connections at %v\n", l.Addr())
	fmt.Printf("Listening for HTTP connections at %v\n", hl.Addr())

	go serveHTTP(db, hl)
	serveRPC(db, l)
}

func serveRPC(db *zenodb.DB, l net.Listener) {
	err := rpc.Serve(db, l)
	if err != nil {
		log.Fatalf("Error serving gRPC: %v", err)
	}
}
<|endoftext|>


<commit_before>/*
https://godoc.org/code.google.com/p/go.net/websocket
http://blog.golang.org/spotlight-on-external-go-libraries
https://gist.github.com/jweir/4528042
https://github.com/golang-samples/websocket/blob/master/simple/main.go
http://blog.jupo.org/2013/02/23/a-tale-of-two-queues/
*/

package main

import (
	"code.google.com/p/go.net/websocket"
	"gopkg.in/redis.v1"
	"net/http"
	"fmt"
	_ "reflect"
)

var (
    foo = "bar"
)

func haltOnErr(err error){
	if err != nil { panic(err) }
}

func echoHandler(ws *websocket.Conn) {

	fmt.Println("connecting!")

	client := redis.NewTCPClient(&redis.Options{
	    Addr: "127.0.0.1:6379",
	})

	defer client.Close()

	pubsub := client.PubSub()
	defer pubsub.Close()

	err := pubsub.Subscribe("mychannel")
	haltOnErr(err)

	if err != nil{
		return
	}

	for {
		i, _ := pubsub.Receive()
		msg, _ := i.(*redis.Message)

		if msg != nil {
		   // fmt.Println(msg.Payload)
		   websocket.JSON.Send(ws, msg.Payload)
		}
	}
}

func main() {

	/* http://stackoverflow.com/questions/19708330/serving-a-websocket-in-go */

	http.HandleFunc("/", func (w http.ResponseWriter, req *http.Request){
        	s := websocket.Server{Handler: websocket.Handler(echoHandler)}
        	s.ServeHTTP(w, req)
    	});

	http_err := http.ListenAndServe("127.0.0.1:8080", nil)

	if http_err != nil {
		panic("ListenAndServe: " + http_err.Error())
	}

	fmt.Println("Listening on 127.0.0.1:8080")
}
<commit_msg>go is weird<commit_after>/*
https://godoc.org/code.google.com/p/go.net/websocket
http://blog.golang.org/spotlight-on-external-go-libraries
https://gist.github.com/jweir/4528042
https://github.com/golang-samples/websocket/blob/master/simple/main.go
http://blog.jupo.org/2013/02/23/a-tale-of-two-queues/
*/

package main

import (
	"code.google.com/p/go.net/websocket"
	"gopkg.in/redis.v1"
	"net/http"
	"fmt"
	_ "reflect"
)

var (
    redis_host = "127.0.0.1"
    redis_port = "6379"
    redis_channel = "pubsocketd"
    /* see below - I have no idea what I am doing */
    /* pubsub *redis.PubSub */
)

/*
func init() {

	addr := redis_host + ":" + redis_port

	client := redis.NewTCPClient(&redis.Options{
	    Addr: addr,
	})

	defer client.Close()

	pubsub := client.PubSub()
	defer pubsub.Close()

	err := pubsub.Subscribe(redis_channel)
}
*/

func echoHandler(ws *websocket.Conn) {

	fmt.Println("connecting!")
	// fmt.Println(ws.RemoteAddr())

	addr := redis_host + ":" + redis_port

	client := redis.NewTCPClient(&redis.Options{
	    Addr: addr,
	})

	defer client.Close()

	pubsub := client.PubSub()
	defer pubsub.Close()

	err := pubsub.Subscribe(redis_channel)

	if err != nil {
	   websocket.JSON.Send(ws, "FUBAR")
	   ws.Close()
	}

	for {
		i, _ := pubsub.Receive()
		msg, _ := i.(*redis.Message)

		if msg != nil {
		   // fmt.Println(msg.Payload)
		   websocket.JSON.Send(ws, msg.Payload)
		}
	}
}

func main() {

	/* http://stackoverflow.com/questions/19708330/serving-a-websocket-in-go */

	http.HandleFunc("/", func (w http.ResponseWriter, req *http.Request){
        	s := websocket.Server{Handler: websocket.Handler(echoHandler)}
        	s.ServeHTTP(w, req)
    	});

	fmt.Println("Listening on " + redis_host + ":" + redis_port + " and relaying messages from '" + redis_channel + "'")

	http_err := http.ListenAndServe("127.0.0.1:8080", nil)

	if http_err != nil {
		panic("ListenAndServe: " + http_err.Error())
	}
}
<|endoftext|>


<commit_before>package webrtc

import (
	"math"
)

// RTCSctpTransport provides details about the SCTP transport.
type RTCSctpTransport struct {
	// Transport represents the transport over which all SCTP packets for data
	// channels will be sent and received.
	Transport RTCDtlsTransport

	// State represents the current state of the SCTP transport.
	State RTCSctpTransportState

	// MaxMessageSize represents the maximum size of data that can be passed to
	// RTCDataChannel's send() method.
	MaxMessageSize float64

	// MaxChannels represents the maximum amount of RTCDataChannel's that can
	// be used simultaneously.
	MaxChannels *uint16

	// OnStateChange  func()
}

func newRTCSctpTransport() *RTCSctpTransport {
	res := &RTCSctpTransport{
		State: RTCSctpTransportStateConnecting,
	}

	res.updateMessageSize()
	res.updateMaxChannels()

	return res
}

func (r *RTCSctpTransport) updateMessageSize() {
	var remoteMaxMessageSize float64 = 65536 // TODO: get from SDP
	var canSendSize float64 = 65536          // TODO: Get from SCTP implementation

	r.MaxMessageSize = r.calcMessageSize(remoteMaxMessageSize, canSendSize)
}

func (r *RTCSctpTransport) calcMessageSize(remoteMaxMessageSize, canSendSize float64) float64 {
	switch {
	case remoteMaxMessageSize == 0 &&
		canSendSize == 0:
		return math.Inf(1)

	case remoteMaxMessageSize == 0:
		return canSendSize

	case canSendSize == 0:
		return remoteMaxMessageSize

	case canSendSize > remoteMaxMessageSize:
		return remoteMaxMessageSize

	default:
		return canSendSize
	}
}

func (r *RTCSctpTransport) updateMaxChannels() {
	val := uint16(65535)
	r.MaxChannels = &val // TODO: Get from implementation
}
<commit_msg>Add RTCSctpTransport.Transport in preparation for spec compliance<commit_after>package webrtc

import (
	"math"
)

// RTCSctpTransport provides details about the SCTP transport.
type RTCSctpTransport struct {
	// Transport represents the transport over which all SCTP packets for data
	// channels will be sent and received.
	Transport *RTCDtlsTransport

	// State represents the current state of the SCTP transport.
	State RTCSctpTransportState

	// MaxMessageSize represents the maximum size of data that can be passed to
	// RTCDataChannel's send() method.
	MaxMessageSize float64

	// MaxChannels represents the maximum amount of RTCDataChannel's that can
	// be used simultaneously.
	MaxChannels *uint16

	// OnStateChange  func()

	// dataChannels
	// dataChannels map[uint16]*RTCDataChannel
}

func newRTCSctpTransport() *RTCSctpTransport {
	res := &RTCSctpTransport{
		State: RTCSctpTransportStateConnecting,
	}

	res.updateMessageSize()
	res.updateMaxChannels()

	return res
}

func (r *RTCSctpTransport) updateMessageSize() {
	var remoteMaxMessageSize float64 = 65536 // TODO: get from SDP
	var canSendSize float64 = 65536          // TODO: Get from SCTP implementation

	r.MaxMessageSize = r.calcMessageSize(remoteMaxMessageSize, canSendSize)
}

func (r *RTCSctpTransport) calcMessageSize(remoteMaxMessageSize, canSendSize float64) float64 {
	switch {
	case remoteMaxMessageSize == 0 &&
		canSendSize == 0:
		return math.Inf(1)

	case remoteMaxMessageSize == 0:
		return canSendSize

	case canSendSize == 0:
		return remoteMaxMessageSize

	case canSendSize > remoteMaxMessageSize:
		return remoteMaxMessageSize

	default:
		return canSendSize
	}
}

func (r *RTCSctpTransport) updateMaxChannels() {
	val := uint16(65535)
	r.MaxChannels = &val // TODO: Get from implementation
}
<|endoftext|>


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

// TODO(adg): floating nav bar + legend
// TODO(adg): set-specific legend

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"go/build"
	"html/template"
	"io"
	"io/ioutil"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strconv"
)

// htmlOutput reads the profile data from profile and generates an HTML
// coverage report, writing it to outfile. If outfile is empty,
// it writes the report to a temporary file and opens it in a web browser.
func htmlOutput(profile, outfile string) error {
	pf, err := os.Open(profile)
	if err != nil {
		return err
	}
	defer pf.Close()

	profiles, err := ParseProfiles(pf)
	if err != nil {
		return err
	}

	var files []*templateFile

	for fn, profile := range profiles {
		dir, file := filepath.Split(fn)
		pkg, err := build.Import(dir, ".", build.FindOnly)
		if err != nil {
			return fmt.Errorf("can't find %q: %v", fn, err)
		}
		src, err := ioutil.ReadFile(filepath.Join(pkg.Dir, file))
		if err != nil {
			return fmt.Errorf("can't read %q: %v", fn, err)
		}
		var buf bytes.Buffer
		err = htmlGen(&buf, src, profile.Tokens(src))
		if err != nil {
			return err
		}
		files = append(files, &templateFile{
			Name: fn,
			Body: template.HTML(buf.String()),
		})
	}

	var out *os.File
	if outfile == "" {
		var dir string
		dir, err = ioutil.TempDir("", "cover")
		if err != nil {
			return err
		}
		out, err = os.Create(filepath.Join(dir, "coverage.html"))
	} else {
		out, err = os.Create(outfile)
	}
	err = htmlTemplate.Execute(out, templateData{Files: files})
	if err == nil {
		err = out.Close()
	}
	if err != nil {
		return err
	}

	if outfile == "" {
		if !startBrowser("file://" + out.Name()) {
			fmt.Fprintf(os.Stderr, "HTML output written to %s\n", out.Name())
		}
	}

	return nil
}

// Profile represents the profiling data for a specific file.
type Profile struct {
	Blocks []ProfileBlock
}

// ProfileBlock represents a single block of profiling data.
type ProfileBlock struct {
	StartLine, StartCol int
	EndLine, EndCol     int
	NumStmt, Count      int
}

// ParseProfiles parses profile data from the given Reader and returns a
// Profile for each file.
func ParseProfiles(r io.Reader) (map[string]*Profile, error) {
	files := make(map[string]*Profile)
	buf := bufio.NewReader(r)
	// First line is mode.
	mode, err := buf.ReadString('\n')
	if err != nil {
		return nil, err
	}
	_ = mode // TODO: Use the mode to affect the display.
	// Rest of file is in the format
	//	encoding/base64/base64.go:34.44,37.40 3 1
	// where the fields are: name.go:line.column,line.column numberOfStatements count
	s := bufio.NewScanner(buf)
	for s.Scan() {
		line := s.Text()
		m := lineRe.FindStringSubmatch(line)
		if m == nil {
			return nil, fmt.Errorf("line %q doesn't match expected format: %v", m, lineRe)
		}
		fn := m[1]
		p := files[fn]
		if p == nil {
			p = new(Profile)
			files[fn] = p
		}
		p.Blocks = append(p.Blocks, ProfileBlock{
			StartLine: toInt(m[2]),
			StartCol:  toInt(m[3]),
			EndLine:   toInt(m[4]),
			EndCol:    toInt(m[5]),
			NumStmt:   toInt(m[6]),
			Count:     toInt(m[7]),
		})
	}
	if err := s.Err(); err != nil {
		return nil, err
	}
	for _, p := range files {
		sort.Sort(blocksByStart(p.Blocks))
	}
	return files, nil
}

type blocksByStart []ProfileBlock

func (b blocksByStart) Len() int      { return len(b) }
func (b blocksByStart) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
func (b blocksByStart) Less(i, j int) bool {
	return b[i].StartLine < b[j].StartLine || b[i].StartLine == b[j].StartLine && b[i].StartCol < b[j].StartCol
}

var lineRe = regexp.MustCompile(`^(.+):([0-9]+).([0-9]+),([0-9]+).([0-9]+) ([0-9]+) ([0-9]+)$`)

func toInt(s string) int {
	i, err := strconv.ParseInt(s, 10, 64)
	if err != nil {
		panic(err)
	}
	return int(i)
}

// Token represents the position in a source file of an opening or closing
//  tag. These are used to colorize the source.
type Token struct {
	Pos   int
	Start bool
	Count int
	Norm  float64 // count normalized to 0-1
}

// Tokens returns a Profile as a set of Tokens within the provided src.
func (p *Profile) Tokens(src []byte) (tokens []Token) {
	// Find maximum counts.
	max := 0
	for _, b := range p.Blocks {
		if b.Count > max {
			max = b.Count
		}
	}
	// Divisor for normalization.
	divisor := math.Log(float64(max))

	// tok returns a Token, populating the Norm field with a normalized Count.
	tok := func(pos int, start bool, count int) Token {
		t := Token{Pos: pos, Start: start, Count: count}
		if !start || count == 0 {
			return t
		}
		if max <= 1 {
			t.Norm = 0.4 // "set" mode; use pale color
		} else if count > 0 {
			t.Norm = math.Log(float64(count)) / divisor
		}
		return t
	}

	line, col := 1, 2
	for si, bi := 0, 0; si < len(src) && bi < len(p.Blocks); {
		b := p.Blocks[bi]
		if b.StartLine == line && b.StartCol == col {
			tokens = append(tokens, tok(si, true, b.Count))
		}
		if b.EndLine == line && b.EndCol == col {
			tokens = append(tokens, tok(si, false, 0))
			bi++
			continue // Don't advance through src; maybe the next block starts here.
		}
		if src[si] == '\n' {
			line++
			col = 0
		}
		col++
		si++
	}
	sort.Sort(tokensByPos(tokens))
	return
}

type tokensByPos []Token

func (t tokensByPos) Len() int      { return len(t) }
func (t tokensByPos) Swap(i, j int) { t[i], t[j] = t[j], t[i] }
func (t tokensByPos) Less(i, j int) bool {
	if t[i].Pos == t[j].Pos {
		return !t[i].Start && t[j].Start
	}
	return t[i].Pos < t[j].Pos
}

// htmlGen generates an HTML coverage report with the provided filename,
// source code, and tokens, and writes it to the given Writer.
func htmlGen(w io.Writer, src []byte, tokens []Token) error {
	dst := bufio.NewWriter(w)
	for i := range src {
		for len(tokens) > 0 && tokens[0].Pos == i {
			t := tokens[0]
			if t.Start {
				n := 0
				if t.Count > 0 {
					n = int(math.Floor(t.Norm*9)) + 1
				}
				fmt.Fprintf(dst, `<span class="cov%v" title="%v">`, n, t.Count)
			} else {
				dst.WriteString("")
			}
			tokens = tokens[1:]
		}
		switch b := src[i]; b {
		case '>':
			dst.WriteString(">")
		case '<':
			dst.WriteString("<")
		case '&':
			dst.WriteString("&")
		case '\t':
			dst.WriteString("        ")
		default:
			dst.WriteByte(b)
		}
	}
	return dst.Flush()
}

// startBrowser tries to open the URL in a browser
// and returns whether it succeed.
func startBrowser(url string) bool {
	// try to start the browser
	var args []string
	switch runtime.GOOS {
	case "darwin":
		args = []string{"open"}
	case "windows":
		args = []string{"cmd", "/c", "start"}
	default:
		args = []string{"xdg-open"}
	}
	cmd := exec.Command(args[0], append(args[1:], url)...)
	return cmd.Start() == nil
}

// rgb returns an rgb value for the specified coverage value
// between 0 (no coverage) and 10 (max coverage).
func rgb(n int) string {
	if n == 0 {
		return "rgb(192, 0, 0)" // Red
	}
	// Gradient from gray to green.
	r := 128 - 12*(n-1)
	g := 128 + 12*(n-1)
	b := 128 + 3*(n-1)
	return fmt.Sprintf("rgb(%v, %v, %v)", r, g, b)
}

// colors generates the CSS rules for coverage colors.
func colors() template.CSS {
	var buf bytes.Buffer
	for i := 0; i < 11; i++ {
		fmt.Fprintf(&buf, ".cov%v { color: %v }\n", i, rgb(i))
	}
	return template.CSS(buf.String())
}

var htmlTemplate = template.Must(template.New("html").Funcs(template.FuncMap{
	"colors": colors,
}).Parse(tmplHTML))

type templateData struct {
	Files []*templateFile
}

type templateFile struct {
	Name string
	Body template.HTML
}

const tmplHTML = `
<!DOCTYPE html>

			<select id="files">
			{{range $i, $f := .Files}}
			<option value="file{{$i}}">{{$f.Name}}</option>
			{{end}}
			</select>

				not tracked
				no coverage
				low coverage
				*
				*
				*
				*
				*
				*
				*
				*
				high coverage

		{{range $i, $f := .Files}}

{{$f.Body}}

		{{end}}

`
<commit_msg>go.tools/cmd/cover: add floating bar with file selector and legend<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(adg): set-specific legend

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"go/build"
	"html/template"
	"io"
	"io/ioutil"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strconv"
)

// htmlOutput reads the profile data from profile and generates an HTML
// coverage report, writing it to outfile. If outfile is empty,
// it writes the report to a temporary file and opens it in a web browser.
func htmlOutput(profile, outfile string) error {
	pf, err := os.Open(profile)
	if err != nil {
		return err
	}
	defer pf.Close()

	profiles, err := ParseProfiles(pf)
	if err != nil {
		return err
	}

	var files []*templateFile

	for fn, profile := range profiles {
		dir, file := filepath.Split(fn)
		pkg, err := build.Import(dir, ".", build.FindOnly)
		if err != nil {
			return fmt.Errorf("can't find %q: %v", fn, err)
		}
		src, err := ioutil.ReadFile(filepath.Join(pkg.Dir, file))
		if err != nil {
			return fmt.Errorf("can't read %q: %v", fn, err)
		}
		var buf bytes.Buffer
		err = htmlGen(&buf, src, profile.Tokens(src))
		if err != nil {
			return err
		}
		files = append(files, &templateFile{
			Name: fn,
			Body: template.HTML(buf.String()),
		})
	}

	var out *os.File
	if outfile == "" {
		var dir string
		dir, err = ioutil.TempDir("", "cover")
		if err != nil {
			return err
		}
		out, err = os.Create(filepath.Join(dir, "coverage.html"))
	} else {
		out, err = os.Create(outfile)
	}
	err = htmlTemplate.Execute(out, templateData{Files: files})
	if err == nil {
		err = out.Close()
	}
	if err != nil {
		return err
	}

	if outfile == "" {
		if !startBrowser("file://" + out.Name()) {
			fmt.Fprintf(os.Stderr, "HTML output written to %s\n", out.Name())
		}
	}

	return nil
}

// Profile represents the profiling data for a specific file.
type Profile struct {
	Blocks []ProfileBlock
}

// ProfileBlock represents a single block of profiling data.
type ProfileBlock struct {
	StartLine, StartCol int
	EndLine, EndCol     int
	NumStmt, Count      int
}

// ParseProfiles parses profile data from the given Reader and returns a
// Profile for each file.
func ParseProfiles(r io.Reader) (map[string]*Profile, error) {
	files := make(map[string]*Profile)
	buf := bufio.NewReader(r)
	// First line is mode.
	mode, err := buf.ReadString('\n')
	if err != nil {
		return nil, err
	}
	_ = mode // TODO: Use the mode to affect the display.
	// Rest of file is in the format
	//	encoding/base64/base64.go:34.44,37.40 3 1
	// where the fields are: name.go:line.column,line.column numberOfStatements count
	s := bufio.NewScanner(buf)
	for s.Scan() {
		line := s.Text()
		m := lineRe.FindStringSubmatch(line)
		if m == nil {
			return nil, fmt.Errorf("line %q doesn't match expected format: %v", m, lineRe)
		}
		fn := m[1]
		p := files[fn]
		if p == nil {
			p = new(Profile)
			files[fn] = p
		}
		p.Blocks = append(p.Blocks, ProfileBlock{
			StartLine: toInt(m[2]),
			StartCol:  toInt(m[3]),
			EndLine:   toInt(m[4]),
			EndCol:    toInt(m[5]),
			NumStmt:   toInt(m[6]),
			Count:     toInt(m[7]),
		})
	}
	if err := s.Err(); err != nil {
		return nil, err
	}
	for _, p := range files {
		sort.Sort(blocksByStart(p.Blocks))
	}
	return files, nil
}

type blocksByStart []ProfileBlock

func (b blocksByStart) Len() int      { return len(b) }
func (b blocksByStart) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
func (b blocksByStart) Less(i, j int) bool {
	return b[i].StartLine < b[j].StartLine || b[i].StartLine == b[j].StartLine && b[i].StartCol < b[j].StartCol
}

var lineRe = regexp.MustCompile(`^(.+):([0-9]+).([0-9]+),([0-9]+).([0-9]+) ([0-9]+) ([0-9]+)$`)

func toInt(s string) int {
	i, err := strconv.ParseInt(s, 10, 64)
	if err != nil {
		panic(err)
	}
	return int(i)
}

// Token represents the position in a source file of an opening or closing
//  tag. These are used to colorize the source.
type Token struct {
	Pos   int
	Start bool
	Count int
	Norm  float64 // count normalized to 0-1
}

// Tokens returns a Profile as a set of Tokens within the provided src.
func (p *Profile) Tokens(src []byte) (tokens []Token) {
	// Find maximum counts.
	max := 0
	for _, b := range p.Blocks {
		if b.Count > max {
			max = b.Count
		}
	}
	// Divisor for normalization.
	divisor := math.Log(float64(max))

	// tok returns a Token, populating the Norm field with a normalized Count.
	tok := func(pos int, start bool, count int) Token {
		t := Token{Pos: pos, Start: start, Count: count}
		if !start || count == 0 {
			return t
		}
		if max <= 1 {
			t.Norm = 0.4 // "set" mode; use pale color
		} else if count > 0 {
			t.Norm = math.Log(float64(count)) / divisor
		}
		return t
	}

	line, col := 1, 2
	for si, bi := 0, 0; si < len(src) && bi < len(p.Blocks); {
		b := p.Blocks[bi]
		if b.StartLine == line && b.StartCol == col {
			tokens = append(tokens, tok(si, true, b.Count))
		}
		if b.EndLine == line && b.EndCol == col {
			tokens = append(tokens, tok(si, false, 0))
			bi++
			continue // Don't advance through src; maybe the next block starts here.
		}
		if src[si] == '\n' {
			line++
			col = 0
		}
		col++
		si++
	}
	sort.Sort(tokensByPos(tokens))
	return
}

type tokensByPos []Token

func (t tokensByPos) Len() int      { return len(t) }
func (t tokensByPos) Swap(i, j int) { t[i], t[j] = t[j], t[i] }
func (t tokensByPos) Less(i, j int) bool {
	if t[i].Pos == t[j].Pos {
		return !t[i].Start && t[j].Start
	}
	return t[i].Pos < t[j].Pos
}

// htmlGen generates an HTML coverage report with the provided filename,
// source code, and tokens, and writes it to the given Writer.
func htmlGen(w io.Writer, src []byte, tokens []Token) error {
	dst := bufio.NewWriter(w)
	for i := range src {
		for len(tokens) > 0 && tokens[0].Pos == i {
			t := tokens[0]
			if t.Start {
				n := 0
				if t.Count > 0 {
					n = int(math.Floor(t.Norm*9)) + 1
				}
				fmt.Fprintf(dst, `<span class="cov%v" title="%v">`, n, t.Count)
			} else {
				dst.WriteString("")
			}
			tokens = tokens[1:]
		}
		switch b := src[i]; b {
		case '>':
			dst.WriteString(">")
		case '<':
			dst.WriteString("<")
		case '&':
			dst.WriteString("&")
		case '\t':
			dst.WriteString("        ")
		default:
			dst.WriteByte(b)
		}
	}
	return dst.Flush()
}

// startBrowser tries to open the URL in a browser
// and returns whether it succeed.
func startBrowser(url string) bool {
	// try to start the browser
	var args []string
	switch runtime.GOOS {
	case "darwin":
		args = []string{"open"}
	case "windows":
		args = []string{"cmd", "/c", "start"}
	default:
		args = []string{"xdg-open"}
	}
	cmd := exec.Command(args[0], append(args[1:], url)...)
	return cmd.Start() == nil
}

// rgb returns an rgb value for the specified coverage value
// between 0 (no coverage) and 10 (max coverage).
func rgb(n int) string {
	if n == 0 {
		return "rgb(192, 0, 0)" // Red
	}
	// Gradient from gray to green.
	r := 128 - 12*(n-1)
	g := 128 + 12*(n-1)
	b := 128 + 3*(n-1)
	return fmt.Sprintf("rgb(%v, %v, %v)", r, g, b)
}

// colors generates the CSS rules for coverage colors.
func colors() template.CSS {
	var buf bytes.Buffer
	for i := 0; i < 11; i++ {
		fmt.Fprintf(&buf, ".cov%v { color: %v }\n", i, rgb(i))
	}
	return template.CSS(buf.String())
}

var htmlTemplate = template.Must(template.New("html").Funcs(template.FuncMap{
	"colors": colors,
}).Parse(tmplHTML))

type templateData struct {
	Files []*templateFile
}

type templateFile struct {
	Name string
	Body template.HTML
}

const tmplHTML = `
<!DOCTYPE html>

				<select id="files">
				{{range $i, $f := .Files}}
				<option value="file{{$i}}">{{$f.Name}}</option>
				{{end}}
				</select>

				not tracked
				no coverage
				low coverage
				*
				*
				*
				*
				*
				*
				*
				*
				high coverage

		{{range $i, $f := .Files}}

{{$f.Body}}

		{{end}}

`
<|endoftext|>


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

package main

import (
	"github.com/globocom/tsuru/cmd"
	"os"
)

const (
	version = "0.4.7"
	header  = "Supported-Crane"
)

func buildManager(name string) *cmd.Manager {
	m := cmd.BuildBaseManager(name, version, header)
	m.Register(&ServiceCreate{})
	m.Register(&ServiceRemove{})
	m.Register(&ServiceList{})
	m.Register(&ServiceUpdate{})
	m.Register(&ServiceDocGet{})
	m.Register(&ServiceDocAdd{})
	m.Register(&ServiceTemplate{})
	return m
}

func main() {
	name := cmd.ExtractProgramName(os.Args[0])
	manager := buildManager(name)
	args := os.Args[1:]
	manager.Run(args)
}
<commit_msg>cmd/crane: version 0.5<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"github.com/globocom/tsuru/cmd"
	"os"
)

const (
	version = "0.5"
	header  = "Supported-Crane"
)

func buildManager(name string) *cmd.Manager {
	m := cmd.BuildBaseManager(name, version, header)
	m.Register(&ServiceCreate{})
	m.Register(&ServiceRemove{})
	m.Register(&ServiceList{})
	m.Register(&ServiceUpdate{})
	m.Register(&ServiceDocGet{})
	m.Register(&ServiceDocAdd{})
	m.Register(&ServiceTemplate{})
	return m
}

func main() {
	name := cmd.ExtractProgramName(os.Args[0])
	manager := buildManager(name)
	args := os.Args[1:]
	manager.Run(args)
}
<|endoftext|>


<commit_before>package main

import (
	contextpkg "context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/docker/containerd/log"
	"github.com/docker/containerd/remotes"
	digest "github.com/opencontainers/go-digest"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
	"golang.org/x/net/context/ctxhttp"
)

// TODO(stevvooe): Create "multi-fetch" mode that just takes a remote
// then receives object/hint lines on stdin, returning content as
// needed.

var fetchCommand = cli.Command{
	Name:        "fetch",
	Usage:       "retrieve objects from a remote",
	ArgsUsage:   "[flags] <remote> <object> [<hint>, ...]",
	Description: `Fetch objects by identifier from a remote.`,
	Flags: []cli.Flag{
		cli.DurationFlag{
			Name:   "timeout",
			Usage:  "total timeout for fetch",
			EnvVar: "CONTAINERD_FETCH_TIMEOUT",
		},
	},
	Action: func(context *cli.Context) error {
		var (
			ctx     = contextpkg.Background()
			timeout = context.Duration("timeout")
			locator = context.Args().First()
			args    = context.Args().Tail()
		)

		if timeout > 0 {
			var cancel func()
			ctx, cancel = contextpkg.WithTimeout(ctx, timeout)
			defer cancel()
		}

		if locator == "" {
			return errors.New("containerd: remote required")
		}

		if len(args) < 1 {
			return errors.New("containerd: object required")
		}

		object := args[0]
		hints := args[1:]

		resolver, err := getResolver(ctx)
		if err != nil {
			return err
		}

		remote, err := resolver.Resolve(ctx, locator)
		if err != nil {
			return err
		}

		ctx = log.WithLogger(ctx, log.G(ctx).WithFields(
			logrus.Fields{
				"remote": locator,
				"object": object,
			}))

		log.G(ctx).Infof("fetching")
		rc, err := remote.Fetch(ctx, object, hints...)
		if err != nil {
			return err
		}
		defer rc.Close()

		if _, err := io.Copy(os.Stdout, rc); err != nil {
			return err
		}

		return nil
	},
}

// NOTE(stevvooe): Most of the code below this point is prototype code to
// demonstrate a very simplified docker.io fetcher. We have a lot of hard coded
// values but we leave many of the details down to the fetcher, creating a lot
// of room for ways to fetch content.

// getResolver prepares the resolver from the environment and options.
func getResolver(ctx contextpkg.Context) (remotes.Resolver, error) {
	return remotes.ResolverFunc(func(ctx contextpkg.Context, locator string) (remotes.Fetcher, error) {
		if !strings.HasPrefix(locator, "docker.io") && !strings.HasPrefix(locator, "localhost:5000") {
			return nil, errors.Errorf("unsupported locator: %q", locator)
		}

		var (
			base = url.URL{
				Scheme: "https",
				Host:   "registry-1.docker.io",
			}
			prefix = strings.TrimPrefix(locator, "docker.io/")
		)

		if strings.HasPrefix(locator, "localhost:5000") {
			base.Scheme = "http"
			base.Host = "localhost:5000"
			prefix = strings.TrimPrefix(locator, "localhost:5000/")
		}

		token, err := getToken(ctx, "repository:"+prefix+":pull")
		if err != nil {
			return nil, err
		}

		return remotes.FetcherFunc(func(ctx contextpkg.Context, object string, hints ...string) (io.ReadCloser, error) {
			ctx = log.WithLogger(ctx, log.G(ctx).WithFields(
				logrus.Fields{
					"prefix": prefix, // or repo?
					"base":   base.String(),
					"hints":  hints,
				},
			))

			paths, err := getV2URLPaths(prefix, object, hints...)
			if err != nil {
				return nil, err
			}

			for _, path := range paths {
				base.Path = path
				url := base.String()
				log.G(ctx).WithField("url", url).Debug("fetch content")

				req, err := http.NewRequest(http.MethodGet, url, nil)
				if err != nil {
					return nil, err
				}

				req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
				for _, mediatype := range remotes.HintValues("mediatype", hints...) {
					req.Header.Set("Accept", mediatype)
				}

				resp, err := ctxhttp.Do(ctx, http.DefaultClient, req)
				if err != nil {
					return nil, err
				}

				if resp.StatusCode > 299 {
					if resp.StatusCode == http.StatusNotFound {
						continue // try one of the other urls.
					}
					resp.Body.Close()
					return nil, errors.Errorf("unexpected status code %v: %v", url, resp.Status)
				}

				return resp.Body, nil
			}

			return nil, errors.New("not found")
		}), nil
	}), nil
}

func getToken(ctx contextpkg.Context, scopes ...string) (string, error) {
	var (
		u = url.URL{
			Scheme: "https",
			Host:   "auth.docker.io",
			Path:   "/token",
		}

		q = url.Values{
			"scope":   scopes,
			"service": []string{"registry.docker.io"}, // usually comes from auth challenge
		}
	)

	u.RawQuery = q.Encode()

	log.G(ctx).WithField("token.url", u.String()).Debug("requesting token")
	resp, err := ctxhttp.Get(ctx, http.DefaultClient, u.String())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode > 299 {
		return "", errors.Errorf("unexpected status code: %v %v", resp.StatusCode, resp.Status)
	}

	p, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	var tokenResponse struct {
		Token string `json:"token"`
	}

	if err := json.Unmarshal(p, &tokenResponse); err != nil {
		return "", err
	}

	return tokenResponse.Token, nil
}

// getV2URLPaths generates the candidate urls paths for the object based on the
// set of hints and the provided object id. URLs are returned in the order of
// most to least likely succeed.
func getV2URLPaths(prefix, object string, hints ...string) ([]string, error) {
	var urls []string

	// TODO(stevvooe): We can probably define a higher-level "type" hint to
	// avoid having to do extra round trips to resolve content, as well as
	// avoid the tedium of providing media types.

	if remotes.HintExists("mediatype", "application/vnd.docker.distribution.manifest.v2+json", hints...) { // TODO(stevvooe): make this handle oci types, as well.
		// fast path out if we know we are getting a manifest. Arguably, we
		// should fallback to blobs, just in case.
		urls = append(urls, path.Join("/v2", prefix, "manifests", object))
	}

	_, err := digest.Parse(object)
	if err == nil {
		// we have a digest, use blob or manifest path, depending on hints, may
		// need to try both.
		urls = append(urls, path.Join("/v2", prefix, "blobs", object))
	}

	// probably a take, so we go through the manifests endpoint
	urls = append(urls, path.Join("/v2", prefix, "manifests", object))

	var (
		noduplicates []string
		seen         = map[string]struct{}{}
	)
	for _, u := range urls {
		if _, ok := seen[u]; !ok {
			seen[u] = struct{}{}
			noduplicates = append(noduplicates, u)
		}
	}

	return noduplicates, nil
}
<commit_msg>cmd/dist/fetch: address subtle concurrency bug<commit_after>package main

import (
	contextpkg "context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/docker/containerd/log"
	"github.com/docker/containerd/remotes"
	digest "github.com/opencontainers/go-digest"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
	"golang.org/x/net/context/ctxhttp"
)

// TODO(stevvooe): Create "multi-fetch" mode that just takes a remote
// then receives object/hint lines on stdin, returning content as
// needed.

var fetchCommand = cli.Command{
	Name:        "fetch",
	Usage:       "retrieve objects from a remote",
	ArgsUsage:   "[flags] <remote> <object> [<hint>, ...]",
	Description: `Fetch objects by identifier from a remote.`,
	Flags: []cli.Flag{
		cli.DurationFlag{
			Name:   "timeout",
			Usage:  "total timeout for fetch",
			EnvVar: "CONTAINERD_FETCH_TIMEOUT",
		},
	},
	Action: func(context *cli.Context) error {
		var (
			ctx     = contextpkg.Background()
			timeout = context.Duration("timeout")
			locator = context.Args().First()
			args    = context.Args().Tail()
		)

		if timeout > 0 {
			var cancel func()
			ctx, cancel = contextpkg.WithTimeout(ctx, timeout)
			defer cancel()
		}

		if locator == "" {
			return errors.New("containerd: remote required")
		}

		if len(args) < 1 {
			return errors.New("containerd: object required")
		}

		object := args[0]
		hints := args[1:]

		resolver, err := getResolver(ctx)
		if err != nil {
			return err
		}

		remote, err := resolver.Resolve(ctx, locator)
		if err != nil {
			return err
		}

		ctx = log.WithLogger(ctx, log.G(ctx).WithFields(
			logrus.Fields{
				"remote": locator,
				"object": object,
			}))

		log.G(ctx).Infof("fetching")
		rc, err := remote.Fetch(ctx, object, hints...)
		if err != nil {
			return err
		}
		defer rc.Close()

		if _, err := io.Copy(os.Stdout, rc); err != nil {
			return err
		}

		return nil
	},
}

// NOTE(stevvooe): Most of the code below this point is prototype code to
// demonstrate a very simplified docker.io fetcher. We have a lot of hard coded
// values but we leave many of the details down to the fetcher, creating a lot
// of room for ways to fetch content.

// getResolver prepares the resolver from the environment and options.
func getResolver(ctx contextpkg.Context) (remotes.Resolver, error) {
	return remotes.ResolverFunc(func(ctx contextpkg.Context, locator string) (remotes.Fetcher, error) {
		if !strings.HasPrefix(locator, "docker.io") && !strings.HasPrefix(locator, "localhost:5000") {
			return nil, errors.Errorf("unsupported locator: %q", locator)
		}

		var (
			base = url.URL{
				Scheme: "https",
				Host:   "registry-1.docker.io",
			}
			prefix = strings.TrimPrefix(locator, "docker.io/")
		)

		if strings.HasPrefix(locator, "localhost:5000") {
			base.Scheme = "http"
			base.Host = "localhost:5000"
			prefix = strings.TrimPrefix(locator, "localhost:5000/")
		}

		token, err := getToken(ctx, "repository:"+prefix+":pull")
		if err != nil {
			return nil, err
		}

		return remotes.FetcherFunc(func(ctx contextpkg.Context, object string, hints ...string) (io.ReadCloser, error) {
			ctx = log.WithLogger(ctx, log.G(ctx).WithFields(
				logrus.Fields{
					"prefix": prefix, // or repo?
					"base":   base.String(),
					"hints":  hints,
				},
			))

			paths, err := getV2URLPaths(prefix, object, hints...)
			if err != nil {
				return nil, err
			}

			for _, path := range paths {
				url := base
				url.Path = path
				log.G(ctx).WithField("url", url).Debug("fetch content")

				req, err := http.NewRequest(http.MethodGet, url.String(), nil)
				if err != nil {
					return nil, err
				}

				req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
				for _, mediatype := range remotes.HintValues("mediatype", hints...) {
					req.Header.Set("Accept", mediatype)
				}

				resp, err := ctxhttp.Do(ctx, http.DefaultClient, req)
				if err != nil {
					return nil, err
				}

				if resp.StatusCode > 299 {
					if resp.StatusCode == http.StatusNotFound {
						continue // try one of the other urls.
					}
					resp.Body.Close()
					return nil, errors.Errorf("unexpected status code %v: %v", url, resp.Status)
				}

				return resp.Body, nil
			}

			return nil, errors.New("not found")
		}), nil
	}), nil
}

func getToken(ctx contextpkg.Context, scopes ...string) (string, error) {
	var (
		u = url.URL{
			Scheme: "https",
			Host:   "auth.docker.io",
			Path:   "/token",
		}

		q = url.Values{
			"scope":   scopes,
			"service": []string{"registry.docker.io"}, // usually comes from auth challenge
		}
	)

	u.RawQuery = q.Encode()

	log.G(ctx).WithField("token.url", u.String()).Debug("requesting token")
	resp, err := ctxhttp.Get(ctx, http.DefaultClient, u.String())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode > 299 {
		return "", errors.Errorf("unexpected status code: %v %v", resp.StatusCode, resp.Status)
	}

	p, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	var tokenResponse struct {
		Token string `json:"token"`
	}

	if err := json.Unmarshal(p, &tokenResponse); err != nil {
		return "", err
	}

	return tokenResponse.Token, nil
}

// getV2URLPaths generates the candidate urls paths for the object based on the
// set of hints and the provided object id. URLs are returned in the order of
// most to least likely succeed.
func getV2URLPaths(prefix, object string, hints ...string) ([]string, error) {
	var urls []string

	// TODO(stevvooe): We can probably define a higher-level "type" hint to
	// avoid having to do extra round trips to resolve content, as well as
	// avoid the tedium of providing media types.

	if remotes.HintExists("mediatype", "application/vnd.docker.distribution.manifest.v2+json", hints...) { // TODO(stevvooe): make this handle oci types, as well.
		// fast path out if we know we are getting a manifest. Arguably, we
		// should fallback to blobs, just in case.
		urls = append(urls, path.Join("/v2", prefix, "manifests", object))
	}

	_, err := digest.Parse(object)
	if err == nil {
		// we have a digest, use blob or manifest path, depending on hints, may
		// need to try both.
		urls = append(urls, path.Join("/v2", prefix, "blobs", object))
	}

	// probably a take, so we go through the manifests endpoint
	urls = append(urls, path.Join("/v2", prefix, "manifests", object))

	var (
		noduplicates []string
		seen         = map[string]struct{}{}
	)
	for _, u := range urls {
		if _, ok := seen[u]; !ok {
			seen[u] = struct{}{}
			noduplicates = append(noduplicates, u)
		}
	}

	return noduplicates, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"

	"gopkg.in/alecthomas/kingpin.v2"
)

func main() {
	os.Exit(Main(os.Args[1:], os.Stdin, os.Stdout, os.Stderr))
}

var (
	app = kingpin.New("hitch", "Repeatr release control")

	releaseCmd                = app.Command("release", "Create new releases.")
	releaseStartCmd           = releaseCmd.Command("start", "Start staging a new release.  (Use more release commands to add data, then commit when done.)")
	releaseAddStepCmd         = releaseCmd.Command("add-step", "Add a step to the replay instructions in the release currently being staged.")
	releaseAddStep_NameArg    = releaseAddStepCmd.Arg("name", "The name to assign this new step.").Required().String()
	releaseAddStep_FormulaArg = releaseAddStepCmd.Arg("formula", "The path to the formula file that runs this step.").Required().String()
	releaseAddStep_ImportsArg = releaseAddStepCmd.Arg("imports", "The path to an imports file which explains the catalogs and release names for wares in the step formula.").String()

	showCmd = app.Command("show", "Show release info objects, or specific content hashes.")
)

func Main(args []string, stdin io.Reader, stdout, stderr io.Writer) (exitCode int) {
	app.HelpFlag.Short('h')

	// Rigging kingpin to use our in/out/err/code.
	var kingpinTerminate bool
	app.Terminate(func(status int) {
		kingpinTerminate = true
		exitCode = status
	})
	app.UsageWriter(stderr)
	app.ErrorWriter(stderr)

	// Parse flags; raise any immediate parse errors.
	cmd, err := app.Parse(args)
	if kingpinTerminate {
		return
	}
	if err != nil {
		fmt.Fprintf(stderr, "hitch: invalid command: %s\n", err)
		return 1
	}

	// Switch for command to invoke.
	switch cmd {
	case releaseStartCmd.FullCommand():
		return 0
	case releaseAddStepCmd.FullCommand():
		fmt.Fprintf(stdout, "whee!\n%q\n%q\n%q\n", *releaseAddStep_NameArg, *releaseAddStep_FormulaArg, *releaseAddStep_ImportsArg)
		return 0
	default:
		panic(fmt.Errorf("hitch: unhandled command %q", cmd))
		return 1
	}
}
<commit_msg>cli: outlines of some command funcs.<commit_after>package main

import (
	"fmt"
	"io"
	"os"

	"gopkg.in/alecthomas/kingpin.v2"
)

func main() {
	os.Exit(Main(os.Args[1:], os.Stdin, os.Stdout, os.Stderr))
}

var (
	app = kingpin.New("hitch", "Repeatr release control")

	initCmd = app.Command("init", "Initialize a new hitch DB.")

	releaseCmd                = app.Command("release", "Create new releases.")
	releaseStartCmd           = releaseCmd.Command("start", "Start staging a new release.  (Use more release commands to add data, then commit when done.)")
	releaseStart_CatalogArg   = releaseStartCmd.Arg("catalogName", "The name to assign this new step.").Required().String()
	releaseStart_ReleaseArg   = releaseStartCmd.Arg("releaseName", "The path to the formula file that runs this step.").Required().String()
	releaseAddStepCmd         = releaseCmd.Command("add-step", "Add a step to the replay instructions in the release currently being staged.")
	releaseAddStep_NameArg    = releaseAddStepCmd.Arg("name", "The name to assign this new step.").Required().String()
	releaseAddStep_FormulaArg = releaseAddStepCmd.Arg("formula", "The path to the formula file that runs this step.").Required().String()
	releaseAddStep_ImportsArg = releaseAddStepCmd.Arg("imports", "The path to an imports file which explains the catalogs and release names for wares in the step formula.").String()

	showCmd = app.Command("show", "Show release info objects, or specific content hashes.")
)

func Main(args []string, stdin io.Reader, stdout, stderr io.Writer) (exitCode int) {
	app.HelpFlag.Short('h')

	// Rigging kingpin to use our in/out/err/code.
	var kingpinTerminate bool
	app.Terminate(func(status int) {
		kingpinTerminate = true
		exitCode = status
	})
	app.UsageWriter(stderr)
	app.ErrorWriter(stderr)

	// Parse flags; raise any immediate parse errors.
	cmd, err := app.Parse(args)
	if kingpinTerminate {
		return
	}
	if err != nil {
		fmt.Fprintf(stderr, "hitch: invalid command: %s\n", err)
		return 1
	}

	// Switch for command to invoke.
	switch cmd {
	case releaseStartCmd.FullCommand():
		return 0
	case releaseAddStepCmd.FullCommand():
		fmt.Fprintf(stdout, "whee!\n%q\n%q\n%q\n", *releaseAddStep_NameArg, *releaseAddStep_FormulaArg, *releaseAddStep_ImportsArg)
		return 0
	default:
		panic(fmt.Errorf("hitch: unhandled command %q", cmd))
		return 1
	}
}

func init() {
	// Check for existing hitch.db root anywhere above this.  Reject if exists.  Nested repos would be silly.

	// Make hitch.db sigil file in cwd.
}

func releaseStart(catalogName, releaseName string) {
	// Find hitch.db root.

	// Check for staging file.  Reject command if staging file already exists.

	// Check for catalog already existing.  Reject if not (this is fat-finger avoidance).
	// REVIEW : is this helpful?  skipping for now.

	// Check for catalog+release already existing.  Reject if released before.

	// If catalog has signing keys set up, check that we have those keys.

	// All checks passed.
	// Make staging file.  It's just a very skeletal catalog.
}

func releaseAddStep(stepName, formulaFilename, importsFilename string) {
	// TODO programming style: do we need a bag for all the IO and code stuff?  seems yes.
	// Because you might still want to flip the "-" args to a single read of stdin here.

	// Find hitch.db root.

	// Check for staging file.  Must exist; reject command if no release has been started.

	// Check that this step name is not already in use.  Reject if is.

	// Check both files exist and parse.
	// If the args are '-', read stdin. (If both filenames are dash, only do so once; use the same buffer twice.)

	// Conditionally: if per-op validation is enabled:
	//  Look over all the imports; if any are "wire", check those are already available.

	// Append the replay object.  Write back out.
}
<|endoftext|>


<commit_before>// Copyright (C) 2017  Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package main

import (
	"net"

	"github.com/aristanetworks/glog"
	kcp "github.com/xtaci/kcp-go"
)

type udpClient struct {
	addr string
	conn net.Conn
}

func newUDPClient(addr string) OpenTSDBConn {
	return &udpClient{
		addr: addr,
	}
}

func (c *udpClient) Put(d *DataPoint) error {
	var err error
	if c.conn == nil {
		c.conn, err = kcp.DialWithOptions(c.addr, nil, 10, 3)
		if err != nil {
			return err
		}
	}
	_, err = c.conn.Write([]byte(d.String()))
	if err != nil {
		c.conn.Close()
		c.conn = nil
	}
	return err
}

type udpServer struct {
	lis    *kcp.Listener
	telnet *telnetClient
}

func newUDPServer(udpAddr, tsdbAddr string) (*udpServer, error) {
	lis, err := kcp.ListenWithOptions(udpAddr, nil, 10, 3)
	if err != nil {
		return nil, err
	}
	return &udpServer{
		lis:    lis,
		telnet: newTelnetClient(tsdbAddr).(*telnetClient),
	}, nil
}

func (c *udpServer) Run() error {
	for {
		conn, err := c.lis.AcceptKCP()
		if err != nil {
			return err
		}

		go func() {
			defer conn.Close()
			var buf [4096]byte
			for {
				n, err := conn.Read(buf[:])
				if err != nil {
					if n != 0 { // Not EOF
						glog.Error(err)
					}
					return
				}
				err = c.telnet.PutBytes(buf[:n])
				if err != nil {
					glog.Error(err)
					return
				}
			}
		}()
	}
}
<commit_msg>cmd/octsdb: Add some logging at a higher verbosity<commit_after>// Copyright (C) 2017  Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package main

import (
	"net"

	"github.com/aristanetworks/glog"
	kcp "github.com/xtaci/kcp-go"
)

type udpClient struct {
	addr string
	conn net.Conn
}

func newUDPClient(addr string) OpenTSDBConn {
	return &udpClient{
		addr: addr,
	}
}

func (c *udpClient) Put(d *DataPoint) error {
	var err error
	if c.conn == nil {
		c.conn, err = kcp.DialWithOptions(c.addr, nil, 10, 3)
		if err != nil {
			return err
		}
	}
	if glog.V(3) {
		glog.Info(d.String())
	}
	_, err = c.conn.Write([]byte(d.String()))
	if err != nil {
		c.conn.Close()
		c.conn = nil
	}
	return err
}

type udpServer struct {
	lis    *kcp.Listener
	telnet *telnetClient
}

func newUDPServer(udpAddr, tsdbAddr string) (*udpServer, error) {
	lis, err := kcp.ListenWithOptions(udpAddr, nil, 10, 3)
	if err != nil {
		return nil, err
	}
	return &udpServer{
		lis:    lis,
		telnet: newTelnetClient(tsdbAddr).(*telnetClient),
	}, nil
}

func (c *udpServer) Run() error {
	for {
		conn, err := c.lis.AcceptKCP()
		if err != nil {
			return err
		}

		go func() {
			defer conn.Close()
			var buf [4096]byte
			for {
				n, err := conn.Read(buf[:])
				if err != nil {
					if n != 0 { // Not EOF
						glog.Error(err)
					}
					return
				}
				if glog.V(3) {
					glog.Info(string(buf[:n]))
				}
				err = c.telnet.PutBytes(buf[:n])
				if err != nil {
					glog.Error(err)
					return
				}
			}
		}()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net/url"
	"os"

	"github.com/BurntSushi/toml"
	"github.com/ensonmj/expect"
)

// Host config
type Host struct {
	HostName string
	User     string
	Pass     string
	Opt      string
}

// Conf for relay.toml
type Conf struct {
	Relay     string
	RelayUser string
	RelayPass string
	Hosts     map[string]Host
}

var (
	fVerbose bool
)

func init() {
	flag.BoolVar(&fVerbose, "v", false, "open debug log")
}

func main() {
	var conf Conf
	if _, err := toml.DecodeFile("relay.toml", &conf); err != nil {
		fmt.Println(err)
		return
	}

	flag.Parse()

	var host Host

	abbrOrURL := flag.Arg(0)
	if abbrOrURL != "" {
		var ok bool
		host, ok = conf.Hosts[abbrOrURL]
		if !ok {
			u, err := url.Parse(abbrOrURL)
			if err != nil {
				errExit(err)
			}
			hostName := u.Hostname()
			user := u.User.Username()
			pass, _ := u.User.Password()
			var opt string
			if u.Scheme != "" {
				opt = "--" + u.Scheme
			} else {
				q, err := url.ParseQuery(u.RawQuery)
				if err != nil {
					errExit(err)
				}
				opt = q["opt"][0]
			}
			host = Host{
				HostName: hostName,
				User:     user,
				Pass:     pass,
				Opt:      opt,
			}
		}
	}
	remoteCmd := flag.Arg(1)

	cmd := fmt.Sprintf("zssh -t %s@%s", conf.RelayUser, conf.Relay)
	child, err := expect.Spawn(cmd)
	if err != nil {
		errExit(err)
	}
	if fVerbose {
		child.Debug(true)
	}
	if abbrOrURL == "" {
		child.Interact()
		return
	}

	pairs := []expect.ExpectPair{
		{"password", func(_ []byte) error {
			err := expect.WithoutEcho(func() error {
				pass := conf.RelayPass
				if pass == "" {
					err := child.SendLineUser("\nPlease enter your password:")
					if err != nil {
						return err
					}
					pass, err = child.ReadLineUser()
					if err != nil {
						return err
					}
				}
				return child.SendLine(pass)
			})
			if err != nil {
				return err
			}
			err = child.SendLineUser("Checking password...")
			if err != nil {
				return err
			}

			err = child.Expect("bash-baidu-ssl$")
			if err != nil {
				return err
			}
			return loginOrRun(child, host, remoteCmd)
		}},
		{"bash-baidu-ssl$", func(_ []byte) error {
			return loginOrRun(child, host, remoteCmd)
		}},
	}
	err = child.ExpectMulti(pairs)
	if err != nil {
		errExit(err)
	}
}

func loginOrRun(child *expect.ExpectSubproc, host Host, remoteCmd string) error {
	var cmd string
	if host.User != "" {
		cmd = fmt.Sprintf("ssh %s %s@%s", host.Opt, host.User, host.HostName)
	} else {
		cmd = fmt.Sprintf("ssh %s %s", host.Opt, host.HostName)
	}
	if remoteCmd != "" {
		cmd = fmt.Sprintf("%s %s", cmd, remoteCmd)
	}
	fmt.Println(cmd)
	err := child.SendLine(cmd)
	if err != nil {
		return err
	}

	if host.Pass != "" {
		err := child.Expect("password")
		if err != nil {
			return err
		}
		err = child.SendLine(host.Pass)
		if err != nil {
			return err
		}
	}

	if remoteCmd == "" {
		child.Interact()
	}

	return nil
}

func errExit(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}
<commit_msg>replace zssh with ssh<commit_after>package main

import (
	"flag"
	"fmt"
	"net/url"
	"os"

	"github.com/BurntSushi/toml"
	"github.com/ensonmj/expect"
)

// Host config
type Host struct {
	HostName string
	User     string
	Pass     string
	Opt      string
}

// Conf for relay.toml
type Conf struct {
	Relay     string
	RelayUser string
	RelayPass string
	Hosts     map[string]Host
}

var (
	fVerbose bool
)

func init() {
	flag.BoolVar(&fVerbose, "v", false, "open debug log")
}

func main() {
	var conf Conf
	if _, err := toml.DecodeFile("relay.toml", &conf); err != nil {
		fmt.Println(err)
		return
	}

	flag.Parse()

	var host Host

	abbrOrURL := flag.Arg(0)
	if abbrOrURL != "" {
		var ok bool
		host, ok = conf.Hosts[abbrOrURL]
		if !ok {
			u, err := url.Parse(abbrOrURL)
			if err != nil {
				errExit(err)
			}
			hostName := u.Hostname()
			user := u.User.Username()
			pass, _ := u.User.Password()
			var opt string
			if u.Scheme != "" {
				opt = "--" + u.Scheme
			} else {
				q, err := url.ParseQuery(u.RawQuery)
				if err != nil {
					errExit(err)
				}
				opt = q["opt"][0]
			}
			host = Host{
				HostName: hostName,
				User:     user,
				Pass:     pass,
				Opt:      opt,
			}
		}
	}
	remoteCmd := flag.Arg(1)

	cmd := fmt.Sprintf("ssh -t %s@%s", conf.RelayUser, conf.Relay)
	child, err := expect.Spawn(cmd)
	if err != nil {
		errExit(err)
	}
	if fVerbose {
		child.Debug(true)
	}
	if abbrOrURL == "" {
		child.Interact()
		return
	}

	pairs := []expect.ExpectPair{
		{"password:", func(_ []byte) error {
			err := expect.WithoutEcho(func() error {
				pass := conf.RelayPass
				if pass == "" {
					err := child.SendLineUser("\nPlease enter your password:")
					if err != nil {
						return err
					}
					pass, err = child.ReadLineUser()
					if err != nil {
						return err
					}
				}
				return child.SendLine(pass)
			})
			if err != nil {
				return err
			}
			err = child.SendLineUser("Checking password...")
			if err != nil {
				return err
			}

			err = child.Expect("bash-baidu-ssl$")
			if err != nil {
				return err
			}
			return loginOrRun(child, host, remoteCmd)
		}},
		{"bash-baidu-ssl$", func(_ []byte) error {
			return loginOrRun(child, host, remoteCmd)
		}},
	}
	err = child.ExpectMulti(pairs)
	if err != nil {
		errExit(err)
	}
}

func loginOrRun(child *expect.ExpectSubproc, host Host, remoteCmd string) error {
	var cmd string
	if host.User != "" {
		cmd = fmt.Sprintf("ssh %s %s@%s", host.Opt, host.User, host.HostName)
	} else {
		cmd = fmt.Sprintf("ssh %s %s", host.Opt, host.HostName)
	}
	if remoteCmd != "" {
		cmd = fmt.Sprintf("%s %s", cmd, remoteCmd)
	}
	fmt.Println(cmd)
	err := child.SendLine(cmd)
	if err != nil {
		return err
	}

	if host.Pass != "" {
		err := child.Expect("password")
		if err != nil {
			return err
		}
		err = child.SendLine(host.Pass)
		if err != nil {
			return err
		}
	}

	if remoteCmd == "" {
		child.Interact()
	}

	return nil
}

func errExit(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"time"

	gclient "github.com/cloudfoundry-incubator/garden/client"
	gconn "github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/atc"
	"github.com/concourse/tsa"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/rata"
	"golang.org/x/crypto/ssh"
)

type registrarSSHServer struct {
	logger            lager.Logger
	atcEndpoint       *rata.RequestGenerator
	heartbeatInterval time.Duration
	forwardHost       string
	config            *ssh.ServerConfig
	httpClient        *http.Client
}

func (server *registrarSSHServer) Serve(listener net.Listener) {
	for {
		c, err := listener.Accept()
		if err != nil {
			server.logger.Error("failed-to-accept", err)
			return
		}

		logger := server.logger.Session("connection")

		conn, chans, reqs, err := ssh.NewServerConn(c, server.config)
		if err != nil {
			logger.Info("handshake-failed", lager.Data{"error": err.Error()})
			continue
		}

		go server.handleConn(logger, conn, chans, reqs)
	}
}

type forwardedTCPIP struct {
	process   ifrit.Process
	boundPort uint32
}

func (server *registrarSSHServer) handleConn(logger lager.Logger, conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) {
	defer conn.Close()

	forwardedTCPIPs := make(chan forwardedTCPIP, 1)
	go server.handleForwardRequests(logger, conn, reqs, forwardedTCPIPs)

	var processes []ifrit.Process

	// ensure processes get cleaned up
	defer func() {
		cleanupLog := logger.Session("cleanup")

		for _, p := range processes {
			cleanupLog.Debug("interrupting")

			p.Signal(os.Interrupt)
		}

		for _, p := range processes {
			err := <-p.Wait()
			if err != nil {
				cleanupLog.Error("process-exited-with-failure", err)
			} else {
				cleanupLog.Debug("process-exited-successfully")
			}
		}
	}()

	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			logger.Info("rejecting-unknown-channel-type", lager.Data{
				"type": newChannel.ChannelType(),
			})

			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		channel, requests, err := newChannel.Accept()
		if err != nil {
			logger.Error("failed-to-accept-channel", err)
			return
		}

		defer channel.Close()

		for req := range requests {
			logger.Info("channel-request", lager.Data{
				"type": req.Type,
			})

			if req.Type != "exec" {
				logger.Info("rejecting")

				req.Reply(false, nil)
				continue
			}

			var request execRequest
			err = ssh.Unmarshal(req.Payload, &request)
			if err != nil {
				logger.Error("malformed-exec-request", err)

				req.Reply(false, nil)

				return
			}

			switch request.Command {
			case "register-worker":
				logger := logger.Session("register-worker")

				req.Reply(true, nil)

				process, err := server.continuouslyRegisterWorkerDirectly(logger, channel)
				if err != nil {
					logger.Error("failed-to-register", err)
					return
				}

				processes = append(processes, process)

				err = conn.Wait()
				logger.Error("connection-closed", err)

			case "forward-worker":
				logger := logger.Session("forward-worker")

				var forwarded forwardedTCPIP

				select {
				case forwarded = <-forwardedTCPIPs:
					logger.Info("forwarded-tcpip", lager.Data{
						"bound-port": forwarded.boundPort,
					})

					processes = append(processes, forwarded.process)

					process, err := server.continuouslyRegisterForwardedWorker(logger, channel, forwarded.boundPort)
					if err != nil {
						logger.Error("failed-to-register", err)
						return
					}

					processes = append(processes, process)

					err = conn.Wait()
					logger.Error("connection-closed", err)

				case <-time.After(10 * time.Second): // todo better?
					logger.Info("never-forwarded-tcpip")
				}

			default:
				logger.Info("invalid-command", lager.Data{
					"command": request.Command,
				})

				req.Reply(false, nil)
			}
		}
	}
}

func (server *registrarSSHServer) continuouslyRegisterWorkerDirectly(
	logger lager.Logger,
	channel ssh.Channel,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	return server.heartbeatWorker(logger, worker), nil
}

func (server *registrarSSHServer) continuouslyRegisterForwardedWorker(
	logger lager.Logger,
	channel ssh.Channel,
	boundPort uint32,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	worker.Addr = fmt.Sprintf("%s:%d", server.forwardHost, boundPort)

	return server.heartbeatWorker(logger, worker), nil
}

func (server *registrarSSHServer) heartbeatWorker(logger lager.Logger, worker atc.Worker) ifrit.Process {
	return ifrit.Invoke(tsa.NewHeartbeater(
		logger,
		server.heartbeatInterval,
		gclient.New(gconn.New("tcp", worker.Addr)),
		server.atcEndpoint,
		worker,
	))
}

func (server *registrarSSHServer) handleForwardRequests(
	logger lager.Logger,
	conn *ssh.ServerConn,
	reqs <-chan *ssh.Request,
	forwardedTCPIPs chan<- forwardedTCPIP,
) {
	var forwardedThings int

	for r := range reqs {
		switch r.Type {
		case "tcpip-forward":
			logger := logger.Session("tcpip-forward")

			forwardedThings++

			if forwardedThings > 1 {
				logger.Info("rejecting-extra-forward-request")
				r.Reply(false, nil)
				continue
			}

			var req tcpipForwardRequest
			err := ssh.Unmarshal(r.Payload, &req)
			if err != nil {
				logger.Error("malformed-tcpip-request", err)
				r.Reply(false, nil)
				continue
			}

			bindAddr := net.JoinHostPort(req.BindIP, fmt.Sprintf("%d", req.BindPort))

			logger.Info("forwarding-tcpip", lager.Data{
				"bind-addr": bindAddr,
			})

			listener, err := net.Listen("tcp", bindAddr)
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			defer listener.Close()

			_, port, err := net.SplitHostPort(listener.Addr().String())
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			var res tcpipForwardResponse
			_, err = fmt.Sscanf(port, "%d", &res.BoundPort)
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			process := server.forwardTCPIP(logger, conn, listener, req.BindIP, res.BoundPort)

			forwardedTCPIPs <- forwardedTCPIP{
				boundPort: res.BoundPort,
				process:   process,
			}

			r.Reply(true, ssh.Marshal(res))

		default:
			logger.Info("ignoring-request", lager.Data{"type": r.Type})
			r.Reply(false, nil)
		}
	}
}

func (server *registrarSSHServer) forwardTCPIP(
	logger lager.Logger,
	conn *ssh.ServerConn,
	listener net.Listener,
	forwardIP string,
	forwardPort uint32,
) ifrit.Process {
	return ifrit.Invoke(ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		go func() {
			<-signals

			listener.Close()
		}()

		close(ready)

		for {
			localConn, err := listener.Accept()
			if err != nil {
				logger.Error("failed-to-accept", err)
				break
			}

			var req forwardTCPIPChannelRequest
			req.ForwardIP = forwardIP
			req.ForwardPort = forwardPort

			host, port, err := net.SplitHostPort(localConn.RemoteAddr().String())
			if err != nil {
				logger.Error("failed-to-split-host-port", err)
				continue
			}

			req.OriginIP = host
			_, err = fmt.Sscanf(port, "%d", &req.OriginPort)
			if err != nil {
				logger.Error("failed-to-parse-port", err)
				continue
			}

			channel, reqs, err := conn.OpenChannel("forwarded-tcpip", ssh.Marshal(req))
			if err != nil {
				logger.Error("failed-to-open-channel", err)
				continue
			}

			go func() {
				for r := range reqs {
					logger.Info("ignoring-request", lager.Data{
						"type": r.Type,
					})

					r.Reply(false, nil)
				}
			}()

			go io.Copy(channel, localConn)
			go io.Copy(localConn, channel)
		}

		return nil
	}))
}
<commit_msg>prevent blocking on heartbeater being ready<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"time"

	gclient "github.com/cloudfoundry-incubator/garden/client"
	gconn "github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/atc"
	"github.com/concourse/tsa"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/rata"
	"golang.org/x/crypto/ssh"
)

type registrarSSHServer struct {
	logger            lager.Logger
	atcEndpoint       *rata.RequestGenerator
	heartbeatInterval time.Duration
	forwardHost       string
	config            *ssh.ServerConfig
	httpClient        *http.Client
}

func (server *registrarSSHServer) Serve(listener net.Listener) {
	for {
		c, err := listener.Accept()
		if err != nil {
			server.logger.Error("failed-to-accept", err)
			return
		}

		logger := server.logger.Session("connection")

		conn, chans, reqs, err := ssh.NewServerConn(c, server.config)
		if err != nil {
			logger.Info("handshake-failed", lager.Data{"error": err.Error()})
			continue
		}

		go server.handleConn(logger, conn, chans, reqs)
	}
}

type forwardedTCPIP struct {
	process   ifrit.Process
	boundPort uint32
}

func (server *registrarSSHServer) handleConn(logger lager.Logger, conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) {
	defer conn.Close()

	forwardedTCPIPs := make(chan forwardedTCPIP, 1)
	go server.handleForwardRequests(logger, conn, reqs, forwardedTCPIPs)

	var processes []ifrit.Process

	// ensure processes get cleaned up
	defer func() {
		cleanupLog := logger.Session("cleanup")

		for _, p := range processes {
			cleanupLog.Debug("interrupting")

			p.Signal(os.Interrupt)
		}

		for _, p := range processes {
			err := <-p.Wait()
			if err != nil {
				cleanupLog.Error("process-exited-with-failure", err)
			} else {
				cleanupLog.Debug("process-exited-successfully")
			}
		}
	}()

	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			logger.Info("rejecting-unknown-channel-type", lager.Data{
				"type": newChannel.ChannelType(),
			})

			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		channel, requests, err := newChannel.Accept()
		if err != nil {
			logger.Error("failed-to-accept-channel", err)
			return
		}

		defer channel.Close()

		for req := range requests {
			logger.Info("channel-request", lager.Data{
				"type": req.Type,
			})

			if req.Type != "exec" {
				logger.Info("rejecting")

				req.Reply(false, nil)
				continue
			}

			var request execRequest
			err = ssh.Unmarshal(req.Payload, &request)
			if err != nil {
				logger.Error("malformed-exec-request", err)

				req.Reply(false, nil)

				return
			}

			switch request.Command {
			case "register-worker":
				logger := logger.Session("register-worker")

				req.Reply(true, nil)

				process, err := server.continuouslyRegisterWorkerDirectly(logger, channel)
				if err != nil {
					logger.Error("failed-to-register", err)
					return
				}

				processes = append(processes, process)

				err = conn.Wait()
				logger.Error("connection-closed", err)

			case "forward-worker":
				logger := logger.Session("forward-worker")

				var forwarded forwardedTCPIP

				select {
				case forwarded = <-forwardedTCPIPs:
					logger.Info("forwarded-tcpip", lager.Data{
						"bound-port": forwarded.boundPort,
					})

					processes = append(processes, forwarded.process)

					process, err := server.continuouslyRegisterForwardedWorker(logger, channel, forwarded.boundPort)
					if err != nil {
						logger.Error("failed-to-register", err)
						return
					}

					processes = append(processes, process)

					err = conn.Wait()
					logger.Error("connection-closed", err)

				case <-time.After(10 * time.Second): // todo better?
					logger.Info("never-forwarded-tcpip")
				}

			default:
				logger.Info("invalid-command", lager.Data{
					"command": request.Command,
				})

				req.Reply(false, nil)
			}
		}
	}
}

func (server *registrarSSHServer) continuouslyRegisterWorkerDirectly(
	logger lager.Logger,
	channel ssh.Channel,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	return server.heartbeatWorker(logger, worker), nil
}

func (server *registrarSSHServer) continuouslyRegisterForwardedWorker(
	logger lager.Logger,
	channel ssh.Channel,
	boundPort uint32,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	worker.Addr = fmt.Sprintf("%s:%d", server.forwardHost, boundPort)

	return server.heartbeatWorker(logger, worker), nil
}

func (server *registrarSSHServer) heartbeatWorker(logger lager.Logger, worker atc.Worker) ifrit.Process {
	return ifrit.Background(tsa.NewHeartbeater(
		logger,
		server.heartbeatInterval,
		gclient.New(gconn.New("tcp", worker.Addr)),
		server.atcEndpoint,
		worker,
	))
}

func (server *registrarSSHServer) handleForwardRequests(
	logger lager.Logger,
	conn *ssh.ServerConn,
	reqs <-chan *ssh.Request,
	forwardedTCPIPs chan<- forwardedTCPIP,
) {
	var forwardedThings int

	for r := range reqs {
		switch r.Type {
		case "tcpip-forward":
			logger := logger.Session("tcpip-forward")

			forwardedThings++

			if forwardedThings > 1 {
				logger.Info("rejecting-extra-forward-request")
				r.Reply(false, nil)
				continue
			}

			var req tcpipForwardRequest
			err := ssh.Unmarshal(r.Payload, &req)
			if err != nil {
				logger.Error("malformed-tcpip-request", err)
				r.Reply(false, nil)
				continue
			}

			bindAddr := net.JoinHostPort(req.BindIP, fmt.Sprintf("%d", req.BindPort))

			logger.Info("forwarding-tcpip", lager.Data{
				"bind-addr": bindAddr,
			})

			listener, err := net.Listen("tcp", bindAddr)
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			defer listener.Close()

			_, port, err := net.SplitHostPort(listener.Addr().String())
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			var res tcpipForwardResponse
			_, err = fmt.Sscanf(port, "%d", &res.BoundPort)
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			process := server.forwardTCPIP(logger, conn, listener, req.BindIP, res.BoundPort)

			forwardedTCPIPs <- forwardedTCPIP{
				boundPort: res.BoundPort,
				process:   process,
			}

			r.Reply(true, ssh.Marshal(res))

		default:
			logger.Info("ignoring-request", lager.Data{"type": r.Type})
			r.Reply(false, nil)
		}
	}
}

func (server *registrarSSHServer) forwardTCPIP(
	logger lager.Logger,
	conn *ssh.ServerConn,
	listener net.Listener,
	forwardIP string,
	forwardPort uint32,
) ifrit.Process {
	return ifrit.Background(ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		go func() {
			<-signals

			listener.Close()
		}()

		close(ready)

		for {
			localConn, err := listener.Accept()
			if err != nil {
				logger.Error("failed-to-accept", err)
				break
			}

			var req forwardTCPIPChannelRequest
			req.ForwardIP = forwardIP
			req.ForwardPort = forwardPort

			host, port, err := net.SplitHostPort(localConn.RemoteAddr().String())
			if err != nil {
				logger.Error("failed-to-split-host-port", err)
				continue
			}

			req.OriginIP = host
			_, err = fmt.Sscanf(port, "%d", &req.OriginPort)
			if err != nil {
				logger.Error("failed-to-parse-port", err)
				continue
			}

			channel, reqs, err := conn.OpenChannel("forwarded-tcpip", ssh.Marshal(req))
			if err != nil {
				logger.Error("failed-to-open-channel", err)
				continue
			}

			go func() {
				for r := range reqs {
					logger.Info("ignoring-request", lager.Data{
						"type": r.Type,
					})

					r.Reply(false, nil)
				}
			}()

			go io.Copy(channel, localConn)
			go io.Copy(localConn, channel)
		}

		return nil
	}))
}
<|endoftext|>