<commit_before>package common

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"os"
	"time"

	"fmt"
	"github.com/BurntSushi/toml"
	log "github.com/Sirupsen/logrus"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/docker"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/ssh"
	"path/filepath"
)

type DockerConfig struct {
	docker_helpers.DockerCredentials
	Hostname               *string        `toml:"hostname" json:"hostname" long:"hostname" env:"DOCKER_HOSTNAME" description:"Custom container hostname"`
	Image                  string         `toml:"image" json:"image" long:"image" env:"DOCKER_IMAGE" description:"Docker image to be used"`
	Privileged             bool           `toml:"privileged" json:"privileged" long:"privileged" env:"DOCKER_PRIVILEGED" description:"Give extended privileges to container"`
	DisableCache           *bool          `toml:"disable_cache" json:"disable_cache" long:"disable-cache" env:"DOCKER_DISABLE_CACHE" description:"Disable all container caching"`
	Volumes                []string       `toml:"volumes" json:"volumes" long:"volumes" env:"DOCKER_VOLUMES" description:"Bind mount a volumes"`
	CacheDir               *string        `toml:"cache_dir" json:"cache_dir" long:"cache-dir" env:"DOCKER_CACHE_DIR" description:"Directory where to store caches"`
	ExtraHosts             []string       `toml:"extra_hosts" json:"extra_hosts" long:"extra-hosts" env:"DOCKER_EXTRA_HOSTS" description:"Add a custom host-to-IP mapping"`
	Links                  []string       `toml:"links" json:"links" long:"links" env:"DOCKER_LINKS" description:"Add link to another container"`
	Services               []string       `toml:"services" json:"services" long:"services" env:"DOCKER_SERVICES" description:"Add service that is started with container"`
	WaitForServicesTimeout *int           `toml:"wait_for_services_timeout" json:"wait_for_services_timeout" long:"wait-for-services-timeout" env:"DOCKER_WAIT_FOR_SERVICES_TIMEOUT" description:"How long to wait for service startup"`
	AllowedImages          []string       `toml:"allowed_images" json:"allowed_images" long:"allowed-images" env:"DOCKER_ALLOWED_IMAGES" description:"Whitelist allowed images"`
	AllowedServices        []string       `toml:"allowed_services" json:"allowed_services" long:"allowed-services" env:"DOCKER_ALLOWED_SERVICES" description:"Whitelist allowed services"`
	ImageTTL               *time.Duration `toml:"image_ttl" json:"image_ttl" long:"image-ttl" env:"DOCKER_IMAGE_TTL"`
}

type ParallelsConfig struct {
	BaseName         string  `toml:"base_name" json:"base_name" long:"base-name" env:"PARALLELS_BASE_NAME" description:"VM name to be used"`
	TemplateName     *string `toml:"template_name" json:"template_name" long:"template-name" env:"PARALLELS_TEMPLATE_NAME" description:"VM template to be created"`
	DisableSnapshots *bool   `toml:"disable_snapshots" json:"disable_snapshots" long:"disable-snapshots" env:"PARALLELS_DISABLE_SNAPSHOTS" description:"Disable snapshoting to speedup VM creation"`
}

type RunnerCredentials struct {
	URL       string `toml:"url" json:"url" short:"u" long:"url" env:"CI_SERVER_URL" required:"true" description:"Runner URL"`
	Token     string `toml:"token" json:"token" short:"t" long:"token" env:"CI_SERVER_TOKEN" required:"true" description:"Runner token"`
	TLSCAFile string `toml:"tls-ca-file" json:"tls-ca-file" long:"tls-ca-file" env:"CI_SERVER_TLS_CA_FILE" description:"File containing the certificates to verify the peer when using HTTPS"`
}

type RunnerSettings struct {
	Executor  string  `toml:"executor" json:"executor" long:"executor" env:"RUNNER_EXECUTOR" required:"true" description:"Select executor, eg. shell, docker, etc."`
	BuildsDir *string `toml:"builds_dir" json:"builds_dir" long:"builds-dir" env:"RUNNER_BUILDS_DIR" description:"Directory where builds are stored"`
	CacheDir  *string `toml:"cache_dir" json:"cache_dir" long:"cache-dir" env:"RUNNER_CACHE_DIR" description:"Directory where build cache is stored"`

	Environment []string `toml:"environment" json:"environment" long:"env" env:"RUNNER_ENV" description:"Custom environment variables injected to build environment"`

	Shell *string `toml:"shell" json:"shell" long:"shell" env:"RUNNER_SHELL" description:"Select bash, cmd or powershell"`

	SSH       *ssh.Config      `toml:"ssh" json:"ssh" group:"ssh executor" namespace:"ssh"`
	Docker    *DockerConfig    `toml:"docker" json:"docker" group:"docker executor" namespace:"docker"`
	Parallels *ParallelsConfig `toml:"parallels" json:"parallels" group:"parallels executor" namespace:"parallels"`
}

type RunnerConfig struct {
	Name           string `toml:"name" json:"name" long:"name" env:"RUNNER_NAME" description:"Runner name"`
	Limit          *int   `toml:"limit" json:"limit" long:"limit" env:"RUNNER_LIMIT" description:"Maximum number of builds processed by this runner"`
	DisableVerbose *bool  `toml:"disable_verbose" json:"disable_verbose"`
	OutputLimit    *int   `toml:"output_limit" long:"ouput-limit" env:"RUNNER_OUTPUT_LIMIT" description:"Maximum build trace size"`

	RunnerCredentials
	RunnerSettings
}

type BaseConfig struct {
	Concurrent int             `toml:"concurrent" json:"concurrent"`
	User       *string         `toml:"user" json:"user"`
	Runners    []*RunnerConfig `toml:"runners" json:"runners"`
}

type Config struct {
	BaseConfig
	ModTime time.Time `json:"-"`
	Loaded  bool      `json:"-"`
}

func (c *RunnerCredentials) ShortDescription() string {
	return helpers.ShortenToken(c.Token)
}

func (c *RunnerCredentials) UniqueID() string {
	return c.URL + c.Token
}

func (c *RunnerCredentials) Log() *log.Entry {
	return log.WithField("runner", c.ShortDescription())
}

func (c *RunnerConfig) String() string {
	return fmt.Sprintf("%v url=%v token=%v executor=%v", c.Name, c.URL, c.Token, c.Executor)
}

func (c *RunnerConfig) GetVariables() BuildVariables {
	var variables BuildVariables

	for _, environment := range c.Environment {
		if variable, err := ParseVariable(environment); err == nil {
			variable.Internal = true
			variables = append(variables, variable)
		}
	}

	return variables
}

func NewConfig() *Config {
	return &Config{
		BaseConfig: BaseConfig{
			Concurrent: 1,
		},
	}
}

func (c *Config) StatConfig(configFile string) error {
	_, err := os.Stat(configFile)
	if err != nil {
		return err
	}
	return nil
}

func (c *Config) LoadConfig(configFile string) error {
	info, err := os.Stat(configFile)

	// permission denied is soft error
	if os.IsNotExist(err) {
		return nil
	} else if err != nil {
		return err
	}

	if _, err = toml.DecodeFile(configFile, &c.BaseConfig); err != nil {
		return err
	}

	c.ModTime = info.ModTime()
	c.Loaded = true
	return nil
}

func (c *Config) SaveConfig(configFile string) error {
	var newConfig bytes.Buffer
	newBuffer := bufio.NewWriter(&newConfig)

	if err := toml.NewEncoder(newBuffer).Encode(&c.BaseConfig); err != nil {
		log.Fatalf("Error encoding TOML: %s", err)
		return err
	}

	if err := newBuffer.Flush(); err != nil {
		return err
	}

	// create directory to store configuration
	os.MkdirAll(filepath.Dir(configFile), 0700)

	// write config file
	if err := ioutil.WriteFile(configFile, newConfig.Bytes(), 0600); err != nil {
		return err
	}

	c.Loaded = true
	return nil
}
<commit_msg>Bring back --description to `gitlab-runner register`<commit_after>package common

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"os"
	"time"

	"fmt"
	"github.com/BurntSushi/toml"
	log "github.com/Sirupsen/logrus"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/docker"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/ssh"
	"path/filepath"
)

type DockerConfig struct {
	docker_helpers.DockerCredentials
	Hostname               *string        `toml:"hostname" json:"hostname" long:"hostname" env:"DOCKER_HOSTNAME" description:"Custom container hostname"`
	Image                  string         `toml:"image" json:"image" long:"image" env:"DOCKER_IMAGE" description:"Docker image to be used"`
	Privileged             bool           `toml:"privileged" json:"privileged" long:"privileged" env:"DOCKER_PRIVILEGED" description:"Give extended privileges to container"`
	DisableCache           *bool          `toml:"disable_cache" json:"disable_cache" long:"disable-cache" env:"DOCKER_DISABLE_CACHE" description:"Disable all container caching"`
	Volumes                []string       `toml:"volumes" json:"volumes" long:"volumes" env:"DOCKER_VOLUMES" description:"Bind mount a volumes"`
	CacheDir               *string        `toml:"cache_dir" json:"cache_dir" long:"cache-dir" env:"DOCKER_CACHE_DIR" description:"Directory where to store caches"`
	ExtraHosts             []string       `toml:"extra_hosts" json:"extra_hosts" long:"extra-hosts" env:"DOCKER_EXTRA_HOSTS" description:"Add a custom host-to-IP mapping"`
	Links                  []string       `toml:"links" json:"links" long:"links" env:"DOCKER_LINKS" description:"Add link to another container"`
	Services               []string       `toml:"services" json:"services" long:"services" env:"DOCKER_SERVICES" description:"Add service that is started with container"`
	WaitForServicesTimeout *int           `toml:"wait_for_services_timeout" json:"wait_for_services_timeout" long:"wait-for-services-timeout" env:"DOCKER_WAIT_FOR_SERVICES_TIMEOUT" description:"How long to wait for service startup"`
	AllowedImages          []string       `toml:"allowed_images" json:"allowed_images" long:"allowed-images" env:"DOCKER_ALLOWED_IMAGES" description:"Whitelist allowed images"`
	AllowedServices        []string       `toml:"allowed_services" json:"allowed_services" long:"allowed-services" env:"DOCKER_ALLOWED_SERVICES" description:"Whitelist allowed services"`
	ImageTTL               *time.Duration `toml:"image_ttl" json:"image_ttl" long:"image-ttl" env:"DOCKER_IMAGE_TTL"`
}

type ParallelsConfig struct {
	BaseName         string  `toml:"base_name" json:"base_name" long:"base-name" env:"PARALLELS_BASE_NAME" description:"VM name to be used"`
	TemplateName     *string `toml:"template_name" json:"template_name" long:"template-name" env:"PARALLELS_TEMPLATE_NAME" description:"VM template to be created"`
	DisableSnapshots *bool   `toml:"disable_snapshots" json:"disable_snapshots" long:"disable-snapshots" env:"PARALLELS_DISABLE_SNAPSHOTS" description:"Disable snapshoting to speedup VM creation"`
}

type RunnerCredentials struct {
	URL       string `toml:"url" json:"url" short:"u" long:"url" env:"CI_SERVER_URL" required:"true" description:"Runner URL"`
	Token     string `toml:"token" json:"token" short:"t" long:"token" env:"CI_SERVER_TOKEN" required:"true" description:"Runner token"`
	TLSCAFile string `toml:"tls-ca-file" json:"tls-ca-file" long:"tls-ca-file" env:"CI_SERVER_TLS_CA_FILE" description:"File containing the certificates to verify the peer when using HTTPS"`
}

type RunnerSettings struct {
	Executor  string  `toml:"executor" json:"executor" long:"executor" env:"RUNNER_EXECUTOR" required:"true" description:"Select executor, eg. shell, docker, etc."`
	BuildsDir *string `toml:"builds_dir" json:"builds_dir" long:"builds-dir" env:"RUNNER_BUILDS_DIR" description:"Directory where builds are stored"`
	CacheDir  *string `toml:"cache_dir" json:"cache_dir" long:"cache-dir" env:"RUNNER_CACHE_DIR" description:"Directory where build cache is stored"`

	Environment []string `toml:"environment" json:"environment" long:"env" env:"RUNNER_ENV" description:"Custom environment variables injected to build environment"`

	Shell *string `toml:"shell" json:"shell" long:"shell" env:"RUNNER_SHELL" description:"Select bash, cmd or powershell"`

	SSH       *ssh.Config      `toml:"ssh" json:"ssh" group:"ssh executor" namespace:"ssh"`
	Docker    *DockerConfig    `toml:"docker" json:"docker" group:"docker executor" namespace:"docker"`
	Parallels *ParallelsConfig `toml:"parallels" json:"parallels" group:"parallels executor" namespace:"parallels"`
}

type RunnerConfig struct {
	Name           string `toml:"name" json:"name" short:"name" long:"description" env:"RUNNER_NAME" description:"Runner name"`
	Limit          *int   `toml:"limit" json:"limit" long:"limit" env:"RUNNER_LIMIT" description:"Maximum number of builds processed by this runner"`
	DisableVerbose *bool  `toml:"disable_verbose" json:"disable_verbose"`
	OutputLimit    *int   `toml:"output_limit" long:"ouput-limit" env:"RUNNER_OUTPUT_LIMIT" description:"Maximum build trace size"`

	RunnerCredentials
	RunnerSettings
}

type BaseConfig struct {
	Concurrent int             `toml:"concurrent" json:"concurrent"`
	User       *string         `toml:"user" json:"user"`
	Runners    []*RunnerConfig `toml:"runners" json:"runners"`
}

type Config struct {
	BaseConfig
	ModTime time.Time `json:"-"`
	Loaded  bool      `json:"-"`
}

func (c *RunnerCredentials) ShortDescription() string {
	return helpers.ShortenToken(c.Token)
}

func (c *RunnerCredentials) UniqueID() string {
	return c.URL + c.Token
}

func (c *RunnerCredentials) Log() *log.Entry {
	return log.WithField("runner", c.ShortDescription())
}

func (c *RunnerConfig) String() string {
	return fmt.Sprintf("%v url=%v token=%v executor=%v", c.Name, c.URL, c.Token, c.Executor)
}

func (c *RunnerConfig) GetVariables() BuildVariables {
	var variables BuildVariables

	for _, environment := range c.Environment {
		if variable, err := ParseVariable(environment); err == nil {
			variable.Internal = true
			variables = append(variables, variable)
		}
	}

	return variables
}

func NewConfig() *Config {
	return &Config{
		BaseConfig: BaseConfig{
			Concurrent: 1,
		},
	}
}

func (c *Config) StatConfig(configFile string) error {
	_, err := os.Stat(configFile)
	if err != nil {
		return err
	}
	return nil
}

func (c *Config) LoadConfig(configFile string) error {
	info, err := os.Stat(configFile)

	// permission denied is soft error
	if os.IsNotExist(err) {
		return nil
	} else if err != nil {
		return err
	}

	if _, err = toml.DecodeFile(configFile, &c.BaseConfig); err != nil {
		return err
	}

	c.ModTime = info.ModTime()
	c.Loaded = true
	return nil
}

func (c *Config) SaveConfig(configFile string) error {
	var newConfig bytes.Buffer
	newBuffer := bufio.NewWriter(&newConfig)

	if err := toml.NewEncoder(newBuffer).Encode(&c.BaseConfig); err != nil {
		log.Fatalf("Error encoding TOML: %s", err)
		return err
	}

	if err := newBuffer.Flush(); err != nil {
		return err
	}

	// create directory to store configuration
	os.MkdirAll(filepath.Dir(configFile), 0700)

	// write config file
	if err := ioutil.WriteFile(configFile, newConfig.Bytes(), 0600); err != nil {
		return err
	}

	c.Loaded = true
	return nil
}
<|endoftext|>


<commit_before>package common

import (
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"
)

// PackerKeyEnv is used to specify the key interval (delay) between keystrokes
// sent to the VM, typically in boot commands. This is to prevent host CPU
// utilization from causing key presses to be skipped or repeated incorrectly.
const PackerKeyEnv = "PACKER_KEY_INTERVAL"

// PackerKeyDefault 100ms is appropriate for shared build infrastructure while a
// shorter delay (e.g. 10ms) can be used on a workstation. See PackerKeyEnv.
const PackerKeyDefault = 100 * time.Millisecond

// ScrubConfig is a helper that returns a string representation of
// any struct with the given values stripped out.
func ScrubConfig(target interface{}, values ...string) string {
	conf := fmt.Sprintf("Config: %+v", target)
	for _, value := range values {
		if value == "" {
			continue
		}
		conf = strings.Replace(conf, value, "<Filtered>", -1)
	}
	return conf
}

// ChooseString returns the first non-empty value.
func ChooseString(vals ...string) string {
	for _, el := range vals {
		if el != "" {
			return el
		}
	}

	return ""
}

// DownloadableURL processes a URL that may also be a file path and returns
// a completely valid URL. For example, the original URL might be "local/file.iso"
// which isn't a valid URL. DownloadableURL will return "file:///local/file.iso"
func DownloadableURL(original string) (string, error) {
	if runtime.GOOS == "windows" {
		// If the distance to the first ":" is just one character, assume
		// we're dealing with a drive letter and thus a file path.
		idx := strings.Index(original, ":")
		if idx == 1 {
			original = "file:///" + original
		}
	}

	url, err := url.Parse(original)
	if err != nil {
		return "", err
	}

	if url.Scheme == "" {
		url.Scheme = "file"
	}

	if url.Scheme == "file" {
		// Windows file handling is all sorts of tricky...
		if runtime.GOOS == "windows" {
			// If the path is using Windows-style slashes, URL parses
			// it into the host field.
			if url.Path == "" && strings.Contains(url.Host, `\`) {
				url.Path = url.Host
				url.Host = ""
			}

			// For Windows absolute file paths, remove leading / prior to processing
			// since net/url turns "C:/" into "/C:/"
			if len(url.Path) > 0 && url.Path[0] == '/' {
				url.Path = url.Path[1:len(url.Path)]
			}
		}

		// Only do the filepath transformations if the file appears
		// to actually exist.
		if _, err := os.Stat(url.Path); err == nil {
			url.Path, err = filepath.Abs(url.Path)
			if err != nil {
				return "", err
			}

			url.Path, err = filepath.EvalSymlinks(url.Path)
			if err != nil {
				return "", err
			}

			url.Path = filepath.Clean(url.Path)
		}

		if runtime.GOOS == "windows" {
			// Also replace all backslashes with forwardslashes since Windows
			// users are likely to do this but the URL should actually only
			// contain forward slashes.
			url.Path = strings.Replace(url.Path, `\`, `/`, -1)
		}
	}

	// Make sure it is lowercased
	url.Scheme = strings.ToLower(url.Scheme)

	// Verify that the scheme is something we support in our common downloader.
	supported := []string{"file", "http", "https"}
	found := false
	for _, s := range supported {
		if url.Scheme == s {
			found = true
			break
		}
	}

	if !found {
		return "", fmt.Errorf("Unsupported URL scheme: %s", url.Scheme)
	}

	return url.String(), nil
}
<commit_msg>fix regression :(<commit_after>package common

import (
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"
)

// PackerKeyEnv is used to specify the key interval (delay) between keystrokes
// sent to the VM, typically in boot commands. This is to prevent host CPU
// utilization from causing key presses to be skipped or repeated incorrectly.
const PackerKeyEnv = "PACKER_KEY_INTERVAL"

// PackerKeyDefault 100ms is appropriate for shared build infrastructure while a
// shorter delay (e.g. 10ms) can be used on a workstation. See PackerKeyEnv.
const PackerKeyDefault = 100 * time.Millisecond

// ScrubConfig is a helper that returns a string representation of
// any struct with the given values stripped out.
func ScrubConfig(target interface{}, values ...string) string {
	conf := fmt.Sprintf("Config: %+v", target)
	for _, value := range values {
		if value == "" {
			continue
		}
		conf = strings.Replace(conf, value, "<Filtered>", -1)
	}
	return conf
}

// ChooseString returns the first non-empty value.
func ChooseString(vals ...string) string {
	for _, el := range vals {
		if el != "" {
			return el
		}
	}

	return ""
}

// DownloadableURL processes a URL that may also be a file path and returns
// a completely valid URL. For example, the original URL might be "local/file.iso"
// which isn't a valid URL. DownloadableURL will return "file:///local/file.iso"
func DownloadableURL(original string) (string, error) {
	if runtime.GOOS == "windows" {
		// If the distance to the first ":" is just one character, assume
		// we're dealing with a drive letter and thus a file path.
		idx := strings.Index(original, ":")
		if idx == 1 {
			original = "file:///" + original
		}
	}

	url, err := url.Parse(original)
	if err != nil {
		return "", err
	}

	if url.Scheme == "" {
		url.Scheme = "file"
	}

	if url.Scheme == "file" {
		// Windows file handling is all sorts of tricky...
		if runtime.GOOS == "windows" {
			// If the path is using Windows-style slashes, URL parses
			// it into the host field.
			if url.Path == "" && strings.Contains(url.Host, `\`) {
				url.Path = url.Host
				url.Host = ""
			}

			// For Windows absolute file paths, remove leading / prior to processing
			// since net/url turns "C:/" into "/C:/"
			if len(url.Path) > 0 && url.Path[0] == '/' {
				url.Path = url.Path[1:len(url.Path)]
			}
		}

		// Only do the filepath transformations if the file appears
		// to actually exist.
		if _, err := os.Stat(url.Path); err == nil {
			url.Path, err = filepath.Abs(url.Path)
			if err != nil {
				return "", err
			}

			url.Path, err = filepath.EvalSymlinks(url.Path)
			if err != nil {
				return "", err
			}

			url.Path = filepath.Clean(url.Path)
		}

		if runtime.GOOS == "windows" {
			// Also replace all backslashes with forwardslashes since Windows
			// users are likely to do this but the URL should actually only
			// contain forward slashes.
			url.Path = strings.Replace(url.Path, `\`, `/`, -1)
		}
	}

	// Make sure it is lowercased
	url.Scheme = strings.ToLower(url.Scheme)

	// This is to work around issue #5927. This can safely be removed once
	// we distribute with a version of Go that fixes that bug.
	//
	// See: https://code.google.com/p/go/issues/detail?id=5927
	if url.Path != "" && url.Path[0] != '/' {
		url.Path = "/" + url.Path
	}

	// Verify that the scheme is something we support in our common downloader.
	supported := []string{"file", "http", "https"}
	found := false
	for _, s := range supported {
		if url.Scheme == s {
			found = true
			break
		}
	}

	if !found {
		return "", fmt.Errorf("Unsupported URL scheme: %s", url.Scheme)
	}

	return url.String(), nil
}
<|endoftext|>


<commit_before>package common

import (
	crand "crypto/rand"
	mrand "math/rand"
	"sync"
	"time"
)

const (
	strChars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" // 62 characters
)

// pseudo random number generator.
// seeded with OS randomness (crand)

type Rand struct {
	sync.Mutex
	rand *mrand.Rand
}

var grand *Rand

func init() {
	grand = NewRand()
	grand.init()
}

func NewRand() *Rand {
	rand := &Rand{}
	rand.init()
	return rand
}

func (r *Rand) init() {
	bz := cRandBytes(8)
	var seed uint64
	for i := 0; i < 8; i++ {
		seed |= uint64(bz[i])
		seed <<= 8
	}
	r.reset(int64(seed))
}

func (r *Rand) reset(seed int64) {
	r.rand = mrand.New(mrand.NewSource(seed))
}

//----------------------------------------
// Global functions

func Seed(seed int64) {
	grand.Seed(seed)
}

func RandStr(length int) string {
	return grand.Str(length)
}

func RandUint16() uint16 {
	return grand.Uint16()
}

func RandUint32() uint32 {
	return grand.Uint32()
}

func RandUint64() uint64 {
	return grand.Uint64()
}

func RandUint() uint {
	return grand.Uint()
}

func RandInt16() int16 {
	return grand.Int16()
}

func RandInt32() int32 {
	return grand.Int32()
}

func RandInt64() int64 {
	return grand.Int64()
}

func RandInt() int {
	return grand.Int()
}

func RandInt31() int32 {
	return grand.Int31()
}

func RandInt63() int64 {
	return grand.Int63()
}

func RandUint16Exp() uint16 {
	return grand.Uint16Exp()
}

func RandUint32Exp() uint32 {
	return grand.Uint32Exp()
}

func RandUint64Exp() uint64 {
	return grand.Uint64Exp()
}

func RandFloat32() float32 {
	return grand.Float32()
}

func RandTime() time.Time {
	return grand.Time()
}

func RandBytes(n int) []byte {
	return grand.Bytes(n)
}

func RandIntn(n int) int {
	return grand.Intn(n)
}

func RandPerm(n int) []int {
	return grand.Perm(n)
}

//----------------------------------------
// Rand methods

func (r *Rand) Seed(seed int64) {
	r.Lock()
	r.reset(seed)
	r.Unlock()
}

// Constructs an alphanumeric string of given length.
// It is not safe for cryptographic usage.
func (r *Rand) Str(length int) string {
	chars := []byte{}
MAIN_LOOP:
	for {
		val := r.Int63()
		for i := 0; i < 10; i++ {
			v := int(val & 0x3f) // rightmost 6 bits
			if v >= 62 {         // only 62 characters in strChars
				val >>= 6
				continue
			} else {
				chars = append(chars, strChars[v])
				if len(chars) == length {
					break MAIN_LOOP
				}
				val >>= 6
			}
		}
	}

	return string(chars)
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint16() uint16 {
	return uint16(r.rand.Uint32() & (1<<16 - 1))
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint32() uint32 {
	r.Lock()
	u32 := r.rand.Uint32()
	r.Unlock()
	return u32
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint64() uint64 {
	return uint64(r.rand.Uint32())<<32 + uint64(r.rand.Uint32())
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint() uint {
	r.Lock()
	i := r.rand.Int()
	r.Unlock()
	return uint(i)
}

// It is not safe for cryptographic usage.
func (r *Rand) Int16() int16 {
	return int16(r.rand.Uint32() & (1<<16 - 1))
}

// It is not safe for cryptographic usage.
func (r *Rand) Int32() int32 {
	return int32(r.rand.Uint32())
}

// It is not safe for cryptographic usage.
func (r *Rand) Int64() int64 {
	return int64(r.rand.Uint64())
}

// It is not safe for cryptographic usage.
func (r *Rand) Int() int {
	r.Lock()
	i := r.rand.Int()
	r.Unlock()
	return i
}

// It is not safe for cryptographic usage.
func (r *Rand) Int31() int32 {
	r.Lock()
	i31 := r.rand.Int31()
	r.Unlock()
	return i31
}

// It is not safe for cryptographic usage.
func (r *Rand) Int63() int64 {
	r.Lock()
	i63 := r.rand.Int63()
	r.Unlock()
	return i63
}

// Distributed pseudo-exponentially to test for various cases
// It is not safe for cryptographic usage.
func (r *Rand) Uint16Exp() uint16 {
	bits := r.rand.Uint32() % 16
	if bits == 0 {
		return 0
	}
	n := uint16(1 << (bits - 1))
	n += uint16(r.rand.Int31()) & ((1 << (bits - 1)) - 1)
	return n
}

// Distributed pseudo-exponentially to test for various cases
// It is not safe for cryptographic usage.
func (r *Rand) Uint32Exp() uint32 {
	bits := r.rand.Uint32() % 32
	if bits == 0 {
		return 0
	}
	n := uint32(1 << (bits - 1))
	n += uint32(r.rand.Int31()) & ((1 << (bits - 1)) - 1)
	return n
}

// Distributed pseudo-exponentially to test for various cases
// It is not safe for cryptographic usage.
func (r *Rand) Uint64Exp() uint64 {
	bits := r.rand.Uint32() % 64
	if bits == 0 {
		return 0
	}
	n := uint64(1 << (bits - 1))
	n += uint64(r.rand.Int63()) & ((1 << (bits - 1)) - 1)
	return n
}

// It is not safe for cryptographic usage.
func (r *Rand) Float32() float32 {
	r.Lock()
	f32 := r.rand.Float32()
	r.Unlock()
	return f32
}

// It is not safe for cryptographic usage.
func (r *Rand) Time() time.Time {
	return time.Unix(int64(r.Uint64Exp()), 0)
}

// RandBytes returns n random bytes from the OS's source of entropy ie. via crypto/rand.
// It is not safe for cryptographic usage.
func (r *Rand) Bytes(n int) []byte {
	// cRandBytes isn't guaranteed to be fast so instead
	// use random bytes generated from the internal PRNG
	bs := make([]byte, n)
	for i := 0; i < len(bs); i++ {
		bs[i] = byte(r.rand.Int() & 0xFF)
	}
	return bs
}

// RandIntn returns, as an int, a non-negative pseudo-random number in [0, n).
// It panics if n <= 0.
// It is not safe for cryptographic usage.
func (r *Rand) Intn(n int) int {
	r.Lock()
	i := r.rand.Intn(n)
	r.Unlock()
	return i
}

// RandPerm returns a pseudo-random permutation of n integers in [0, n).
// It is not safe for cryptographic usage.
func (r *Rand) Perm(n int) []int {
	r.Lock()
	perm := r.rand.Perm(n)
	r.Unlock()
	return perm
}

// NOTE: This relies on the os's random number generator.
// For real security, we should salt that with some seed.
// See github.com/tendermint/go-crypto for a more secure reader.
func cRandBytes(numBytes int) []byte {
	b := make([]byte, numBytes)
	_, err := crand.Read(b)
	if err != nil {
		PanicCrisis(err)
	}
	return b
}
<commit_msg>Fix race condition in random.go<commit_after>package common

import (
	crand "crypto/rand"
	mrand "math/rand"
	"sync"
	"time"
)

const (
	strChars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" // 62 characters
)

// pseudo random number generator.
// seeded with OS randomness (crand)

type Rand struct {
	sync.Mutex
	rand *mrand.Rand
}

var grand *Rand

func init() {
	grand = NewRand()
	grand.init()
}

func NewRand() *Rand {
	rand := &Rand{}
	rand.init()
	return rand
}

func (r *Rand) init() {
	bz := cRandBytes(8)
	var seed uint64
	for i := 0; i < 8; i++ {
		seed |= uint64(bz[i])
		seed <<= 8
	}
	r.reset(int64(seed))
}

func (r *Rand) reset(seed int64) {
	r.rand = mrand.New(mrand.NewSource(seed))
}

//----------------------------------------
// Global functions

func Seed(seed int64) {
	grand.Seed(seed)
}

func RandStr(length int) string {
	return grand.Str(length)
}

func RandUint16() uint16 {
	return grand.Uint16()
}

func RandUint32() uint32 {
	return grand.Uint32()
}

func RandUint64() uint64 {
	return grand.Uint64()
}

func RandUint() uint {
	return grand.Uint()
}

func RandInt16() int16 {
	return grand.Int16()
}

func RandInt32() int32 {
	return grand.Int32()
}

func RandInt64() int64 {
	return grand.Int64()
}

func RandInt() int {
	return grand.Int()
}

func RandInt31() int32 {
	return grand.Int31()
}

func RandInt63() int64 {
	return grand.Int63()
}

func RandUint16Exp() uint16 {
	return grand.Uint16Exp()
}

func RandUint32Exp() uint32 {
	return grand.Uint32Exp()
}

func RandUint64Exp() uint64 {
	return grand.Uint64Exp()
}

func RandFloat32() float32 {
	return grand.Float32()
}

func RandTime() time.Time {
	return grand.Time()
}

func RandBytes(n int) []byte {
	return grand.Bytes(n)
}

func RandIntn(n int) int {
	return grand.Intn(n)
}

func RandPerm(n int) []int {
	return grand.Perm(n)
}

//----------------------------------------
// Rand methods

func (r *Rand) Seed(seed int64) {
	r.Lock()
	r.reset(seed)
	r.Unlock()
}

// Constructs an alphanumeric string of given length.
// It is not safe for cryptographic usage.
func (r *Rand) Str(length int) string {
	chars := []byte{}
MAIN_LOOP:
	for {
		val := r.Int63()
		for i := 0; i < 10; i++ {
			v := int(val & 0x3f) // rightmost 6 bits
			if v >= 62 {         // only 62 characters in strChars
				val >>= 6
				continue
			} else {
				chars = append(chars, strChars[v])
				if len(chars) == length {
					break MAIN_LOOP
				}
				val >>= 6
			}
		}
	}

	return string(chars)
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint16() uint16 {
	return uint16(r.Uint32() & (1<<16 - 1))
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint32() uint32 {
	r.Lock()
	u32 := r.rand.Uint32()
	r.Unlock()
	return u32
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint64() uint64 {
	return uint64(r.Uint32())<<32 + uint64(r.Uint32())
}

// It is not safe for cryptographic usage.
func (r *Rand) Uint() uint {
	r.Lock()
	i := r.rand.Int()
	r.Unlock()
	return uint(i)
}

// It is not safe for cryptographic usage.
func (r *Rand) Int16() int16 {
	return int16(r.Uint32() & (1<<16 - 1))
}

// It is not safe for cryptographic usage.
func (r *Rand) Int32() int32 {
	return int32(r.Uint32())
}

// It is not safe for cryptographic usage.
func (r *Rand) Int64() int64 {
	return int64(r.Uint64())
}

// It is not safe for cryptographic usage.
func (r *Rand) Int() int {
	r.Lock()
	i := r.rand.Int()
	r.Unlock()
	return i
}

// It is not safe for cryptographic usage.
func (r *Rand) Int31() int32 {
	r.Lock()
	i31 := r.rand.Int31()
	r.Unlock()
	return i31
}

// It is not safe for cryptographic usage.
func (r *Rand) Int63() int64 {
	r.Lock()
	i63 := r.rand.Int63()
	r.Unlock()
	return i63
}

// Distributed pseudo-exponentially to test for various cases
// It is not safe for cryptographic usage.
func (r *Rand) Uint16Exp() uint16 {
	bits := r.Uint32() % 16
	if bits == 0 {
		return 0
	}
	n := uint16(1 << (bits - 1))
	n += uint16(r.Int31()) & ((1 << (bits - 1)) - 1)
	return n
}

// Distributed pseudo-exponentially to test for various cases
// It is not safe for cryptographic usage.
func (r *Rand) Uint32Exp() uint32 {
	bits := r.Uint32() % 32
	if bits == 0 {
		return 0
	}
	n := uint32(1 << (bits - 1))
	n += uint32(r.Int31()) & ((1 << (bits - 1)) - 1)
	return n
}

// Distributed pseudo-exponentially to test for various cases
// It is not safe for cryptographic usage.
func (r *Rand) Uint64Exp() uint64 {
	bits := r.Uint32() % 64
	if bits == 0 {
		return 0
	}
	n := uint64(1 << (bits - 1))
	n += uint64(r.Int63()) & ((1 << (bits - 1)) - 1)
	return n
}

// It is not safe for cryptographic usage.
func (r *Rand) Float32() float32 {
	r.Lock()
	f32 := r.rand.Float32()
	r.Unlock()
	return f32
}

// It is not safe for cryptographic usage.
func (r *Rand) Time() time.Time {
	return time.Unix(int64(r.Uint64Exp()), 0)
}

// RandBytes returns n random bytes from the OS's source of entropy ie. via crypto/rand.
// It is not safe for cryptographic usage.
func (r *Rand) Bytes(n int) []byte {
	// cRandBytes isn't guaranteed to be fast so instead
	// use random bytes generated from the internal PRNG
	bs := make([]byte, n)
	for i := 0; i < len(bs); i++ {
		bs[i] = byte(r.Int() & 0xFF)
	}
	return bs
}

// RandIntn returns, as an int, a non-negative pseudo-random number in [0, n).
// It panics if n <= 0.
// It is not safe for cryptographic usage.
func (r *Rand) Intn(n int) int {
	r.Lock()
	i := r.rand.Intn(n)
	r.Unlock()
	return i
}

// RandPerm returns a pseudo-random permutation of n integers in [0, n).
// It is not safe for cryptographic usage.
func (r *Rand) Perm(n int) []int {
	r.Lock()
	perm := r.rand.Perm(n)
	r.Unlock()
	return perm
}

// NOTE: This relies on the os's random number generator.
// For real security, we should salt that with some seed.
// See github.com/tendermint/go-crypto for a more secure reader.
func cRandBytes(numBytes int) []byte {
	b := make([]byte, numBytes)
	_, err := crand.Read(b)
	if err != nil {
		PanicCrisis(err)
	}
	return b
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"math"
	"os"
	"sync"
	"time"
)

// TrafficGenerator represents the traffic generation object
type TrafficGenerator struct {
	stats       Stats
	trafficFunc func(string) Request
	wg          sync.WaitGroup
}

// Worker represents a client making the requests
type Worker struct {
	id         int
	trafficGen *TrafficGenerator
}

var trafficMap = map[string]func(string) Request{
	"http": getURL,
	"dns":  lookupURL,
}

var statsMap = map[string]func() Stats{
	"http": newHTTPStats,
	"dns":  newDNSStats,
}

// NewTrafficGenerator will return a new TrafficGenerator object
func NewTrafficGenerator(trafficType string) (*TrafficGenerator, error) {
	tFunc, ok := trafficMap[trafficType]
	if !ok {
		return nil, ErrInvalidTrafficType
	}
	stats, err := newStats(trafficType)
	if err != nil {
		return nil, err
	}
	return &TrafficGenerator{
		trafficFunc: tFunc,
		stats:       stats,
	}, nil
}

// Generate generates traffic
func (trafficGen *TrafficGenerator) Generate() {
	for i := 1; i <= nbOfClients; i++ {
		trafficGen.wg.Add(1)
		go trafficGen.NewWorker(i).work()
	}
	trafficGen.wg.Wait()
}

// NewWorker creates a new worker for traffic generation
func (trafficGen *TrafficGenerator) NewWorker(i int) *Worker {
	return &Worker{
		id:         i,
		trafficGen: trafficGen,
	}
}

// DisplayStats renders the statistics of the traffic generation
func (trafficGen *TrafficGenerator) DisplayStats() {
	trafficGen.stats.Render()
}

func (w *Worker) work() {
	defer w.trafficGen.wg.Done()
	start := time.Now()

	workerFmt := fmt.Sprintf("worker#%%0%dd", getPadding(nbOfClients))
	counterFmt := fmt.Sprintf(" - %%0%dd/%%d ", getPadding(nbOfRequests))

	prefix := fmt.Sprintf(workerFmt, w.id)
	logger := log.New(os.Stdout, prefix, 0)
	// Repeat nbOfRequests requests
	for i := 1; i <= nbOfRequests; i++ {
		logger.SetPrefix(prefix + fmt.Sprintf(counterFmt, i, nbOfRequests))
		// Find an URL
		url := findRandomURL()
		// Make the request
		r := w.trafficGen.trafficFunc(url)
		// Add the request to the stats
		w.trafficGen.stats.AddRequest(r)
		// Print the request
		logger.Print(r.String())

		time.Sleep(time.Duration(avgMillisecondsToWait) * time.Millisecond)
	}
	w.trafficGen.stats.SetDuration(time.Since(start))
}

// getPadding returns the padding size of the int given
func getPadding(nb int) int {
	// Get the padding size : floor(log10(nb)) + 1
	return int(math.Log10(float64(nb))) + 1
}

func getAvgDuration(total time.Duration, number int) string {
	if number == 0 {
		return "NaN"
	}
	return (total / time.Duration(number)).String()
}
<commit_msg>Allow to send SIGINT / SIGKILL to stop the workers<commit_after>package main

import (
	"fmt"
	"log"
	"math"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// TrafficGenerator represents the traffic generation object
type TrafficGenerator struct {
	stats       Stats
	trafficFunc func(string) Request
	wg          sync.WaitGroup
}

// Worker represents a client making the requests
type Worker struct {
	id         int
	trafficGen *TrafficGenerator
}

var trafficMap = map[string]func(string) Request{
	"http": getURL,
	"dns":  lookupURL,
}

var statsMap = map[string]func() Stats{
	"http": newHTTPStats,
	"dns":  newDNSStats,
}

var exitChan = make(chan struct{})

// NewTrafficGenerator will return a new TrafficGenerator object
func NewTrafficGenerator(trafficType string) (*TrafficGenerator, error) {
	tFunc, ok := trafficMap[trafficType]
	if !ok {
		return nil, ErrInvalidTrafficType
	}
	stats, err := newStats(trafficType)
	if err != nil {
		return nil, err
	}
	return &TrafficGenerator{
		trafficFunc: tFunc,
		stats:       stats,
	}, nil
}

// Generate generates traffic
func (trafficGen *TrafficGenerator) Generate() {
	// Create a channel that will listen to SIGINT / SIGTERM
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT)
	signal.Notify(c, syscall.SIGTERM)

	for i := 1; i <= nbOfClients; i++ {
		trafficGen.wg.Add(1)
		// Launch the workers in a go routine
		go trafficGen.NewWorker(i).work()
	}

	// Done channel to stop the loop when all the workers are done
	var done = make(chan struct{})
	go func() {
		// Wait for the workerss
		trafficGen.wg.Wait()
		// All the workers are done
		done <- struct{}{}
	}()

	// Wait for the workers to end
	// or a signal in the loop
	var forceShutdown bool
	for {
		select {
		case <-done:
			// All the workers are done, we quit
			return
		case sig := <-c:
			// We listen for signals
			switch sig {
			case syscall.SIGINT, syscall.SIGKILL:
				// If it's the second time we get a signal, quit
				if forceShutdown {
					os.Exit(1)
				}

				// Notify all the workers that they need to stop
				for i := 1; i <= nbOfClients; i++ {
					exitChan <- struct{}{}
				}

				// Next time we get a signal, need to quit directly
				forceShutdown = true
			}
		}
	}
}

// NewWorker creates a new worker for traffic generation
func (trafficGen *TrafficGenerator) NewWorker(i int) *Worker {
	return &Worker{
		id:         i,
		trafficGen: trafficGen,
	}
}

// DisplayStats renders the statistics of the traffic generation
func (trafficGen *TrafficGenerator) DisplayStats() {
	trafficGen.stats.Render()
}

func (w *Worker) work() {
	var exit bool
	defer w.trafficGen.wg.Done()
	start := time.Now()

	var done = make(chan struct{})
	// When the work is done, notify the watching go routine
	defer func() { done <- struct{}{} }()

	// Create the watching go routine that will watch if we need to quit early
	go func() {
		for {
			select {
			// If we need to exit
			case <-exitChan:
				exit = true
			// If everything is done
			case <-done:
				return
			}
		}
	}()

	workerFmt := fmt.Sprintf("worker#%%0%dd", getPadding(nbOfClients))
	counterFmt := fmt.Sprintf(" - %%0%dd/%%d ", getPadding(nbOfRequests))

	prefix := fmt.Sprintf(workerFmt, w.id)
	logger := log.New(os.Stdout, prefix, 0)

	// Repeat nbOfRequests requests
	for i := 1; i <= nbOfRequests; i++ {
		// If we got an exit signal, quit
		if exit {
			return
		}
		logger.SetPrefix(prefix + fmt.Sprintf(counterFmt, i, nbOfRequests))
		// Find an URL
		url := findRandomURL()
		// Make the request
		r := w.trafficGen.trafficFunc(url)
		// Add the request to the stats
		w.trafficGen.stats.AddRequest(r)
		// Print the request
		logger.Print(r.String())

		time.Sleep(time.Duration(avgMillisecondsToWait) * time.Millisecond)
	}
	w.trafficGen.stats.SetDuration(time.Since(start))
}

// getPadding returns the padding size of the int given
func getPadding(nb int) int {
	// Get the padding size : floor(log10(nb)) + 1
	return int(math.Log10(float64(nb))) + 1
}

func getAvgDuration(total time.Duration, number int) string {
	if number == 0 {
		return "NaN"
	}
	return (total / time.Duration(number)).String()
}
<|endoftext|>


<commit_before>// This example demonstrates a priority queue built using the heap interface.
package main

import (
  "container/heap"
  "fmt"
)

// An Item is something we manage in a priority queue.
type Item struct {
  x         int
  y         int
  angle     float32
  index     int
}

// A Frontier implements heap.Interface and holds Items.
type Frontier []*Item

func (pq Frontier) Len() int { return len(pq) }

func (pq Frontier) Less(i, j int) bool {
  // We want Pop to give us the highest, not lowest, angle so we use greater than here.
  return pq[i].angle < pq[j].angle
}

func (pq Frontier) Swap(i, j int) {
  pq[i], pq[j] = pq[j], pq[i]
  pq[i].index = i
  pq[j].index = j
}

func (pq *Frontier) Push(x interface{}) {
  n := len(*pq)
  item := x.(*Item)
  item.index = n
  *pq = append(*pq, item)
}

func (pq *Frontier) Pop() interface{} {
  old := *pq
  n := len(old)
  item := old[n-1]
  item.index = -1 // for safety
  *pq = old[0 : n-1]
  return item
}

// update modifies the angle and value of an Item in the queue.
func (pq *Frontier) update(item *Item, angle float32, angle float32) {
  item.value = value
  item.angle = angle
  heap.Fix(pq, item.index)
}

// This example creates a Frontier with some items, adds and manipulates an item,
// and then removes the items in angle order.
func main() {
  // Create a angle queue, put the items in it, and
  // establish the angle queue (heap) invariants.
  pq := make(Frontier, 4)
  pq[0] = Item{1,0,0.1,0}
  pq[1] = Item{0,1,0.2,1}
  pq[2] = Item{-1,0,1,2}
  pq[3] = Item{0,-1,3,3}
  heap.Init(&pq)

  // Take the items out; they arrive in decreasing angle order.
  for pq.Len() > 0 {
    item := heap.Pop(&pq).(*Item)
    fmt.Printf("%2.f:(%d,%d) ", item.angle, item.x, item.y)
  }
}
<commit_msg>Sigh, not ideal<commit_after>// This example demonstrates a priority queue built using the heap interface.
package main

import (
  "container/heap"
  "fmt"
)

// An Item is something we manage in a priority queue.
type Item struct {
  x         int
  y         int
  angle     float32
  index     int
}

// A Frontier implements heap.Interface and holds Items.
type Frontier []*Item

func (pq Frontier) Len() int { return len(pq) }

func (pq Frontier) Less(i, j int) bool {
  // We want Pop to give us the highest, not lowest, angle so we use greater than here.
  return pq[i].angle < pq[j].angle
}

func (pq Frontier) Swap(i, j int) {
  pq[i], pq[j] = pq[j], pq[i]
  pq[i].index = i
  pq[j].index = j
}

func (pq *Frontier) Push(x interface{}) {
  n := len(*pq)
  item := x.(*Item)
  item.index = n
  *pq = append(*pq, item)
}

func (pq *Frontier) Pop() interface{} {
  old := *pq
  n := len(old)
  item := old[n-1]
  item.index = -1 // for safety
  *pq = old[0 : n-1]
  return item
}

// update modifies the angle and value of an Item in the queue.
func (pq *Frontier) update(item *Item, x int, y int, angle float32) {
  item.x = x
  item.y = y
  item.angle = angle
  heap.Fix(pq, item.index)
}

// This example creates a Frontier with some items, adds and manipulates an item,
// and then removes the items in angle order.
func main() {
  // Create a angle queue, put the items in it, and
  // establish the angle queue (heap) invariants.
  pq := make(Frontier, 4)
  pq[0] = Item{1,0,0.1,0}
  pq[1] = Item{0,1,0.2,1}
  pq[2] = Item{-1,0,1,2}
  pq[3] = Item{0,-1,3,3}
  heap.Init(&pq)

  // Take the items out; they arrive in decreasing angle order.
  for pq.Len() > 0 {
    item := heap.Pop(&pq).(*Item)
    fmt.Printf("%2.f:(%d,%d) ", item.angle, item.x, item.y)
  }
}
<|endoftext|>


<commit_before>package state

import (
	"bufio"
	"fmt"
	"io"
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/trivial"
	"launchpad.net/tomb"
	"net"
	"os/exec"
	"strings"
	"time"
)

// These variables would be constants but they are varied
// for testing purposes.
var (
	sshRemotePort    = 22
	sshRetryInterval = time.Second
	sshUser          = "ubuntu@"
)

// sshDial dials the MongoDB instance at addr through
// an SSH proxy.
func sshDial(addr, keyFile string) (*sshForwarder, *mgo.Session, error) {
	fwd, err := newSSHForwarder(addr, keyFile)
	if err != nil {
		return nil, nil, err
	}
	defer trivial.ErrorContextf(&err, "cannot dial MongoDB via SSH at address %s", addr)
	conn, err := mgo.Dial(fwd.localAddr)
	if err != nil {
		fwd.stop()
		return nil, nil, err
	}
	return fwd, conn, nil
}

type sshForwarder struct {
	tomb.Tomb
	localAddr  string
	remoteHost string
	remotePort string
	keyFile    string
}

// newSSHForwarder starts an ssh proxy connecting to the
// remote TCP address. The localAddr field holds the
// name of the local proxy address. If keyFile is non-empty,
// it should name a file containing a private identity key.
func newSSHForwarder(remoteAddr, keyFile string) (fwd *sshForwarder, err error) {
	defer trivial.ErrorContextf(&err, "cannot start SSH proxy for address %s", remoteAddr)
	remoteHost, remotePort, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		return nil, err
	}
	localPort, err := choosePort()
	if err != nil {
		return nil, fmt.Errorf("cannot choose local port: %v", err)
	}
	fwd = &sshForwarder{
		localAddr:  fmt.Sprintf("localhost:%d", localPort),
		remoteHost: remoteHost,
		remotePort: remotePort,
		keyFile:    keyFile,
	}
	proc, err := fwd.start()
	if err != nil {
		return nil, err
	}
	go fwd.run(proc)
	return fwd, nil
}

func (fwd *sshForwarder) stop() error {
	fwd.Kill(nil)
	return fwd.Wait()
}

// run is called with the ssh process already active.
// It loops, restarting ssh until it gets an unrecoverable error.
func (fwd *sshForwarder) run(proc *sshProc) {
	defer fwd.Done()
	restartCount := 0
	startTime := time.Now()

	// waitErrTimeout and waitErr become valid when the ssh client has exited.
	var (
		waitErrTimeout <-chan time.Time
		waitErr        error
	)
	for {
		select {
		case <-fwd.Dying():
			proc.stop()
			return

		case sshErr := <-proc.error:
			log.Printf("state: ssh error (will retry: %v): %v", !sshErr.fatal, sshErr)
			if sshErr.fatal {
				proc.stop()
				fwd.Kill(sshErr)
				return
			}
			// If the ssh process dies repeatedly after running
			// for a very short time and we don't recognise the
			// error, we assume that something unknown is
			// going wrong and quit.
			if sshErr.unknown && time.Now().Sub(startTime) < 200*time.Millisecond {
				if restartCount++; restartCount > 10 {
					proc.stop()
					log.Printf("state: too many ssh errors")
					fwd.Kill(fmt.Errorf("too many errors: %v", sshErr))
					return
				}
			} else {
				restartCount = 0
			}

		case waitErr = <-proc.wait:
			// If ssh has exited, we'll wait a little while
			// in case we've got the exit status before we've
			// received the error.
			waitErrTimeout = time.After(100 * time.Millisecond)
			continue

		case <-waitErrTimeout:
			log.Printf("state: ssh client exited silently: %v", waitErr)
			// We only get here if ssh exits when no fatal
			// errors have been printed by ssh.  In that
			// case, it's probably best to treat it as fatal
			// rather than restarting blindly.
			proc.stop()
			if waitErr == nil {
				waitErr = fmt.Errorf("ssh exited silently")
			} else {
				waitErr = fmt.Errorf("ssh exited silently: %v", waitErr)
			}
			fwd.Kill(waitErr)
			return
		}
		proc.stop()
		var err error
		time.Sleep(sshRetryInterval)
		startTime = time.Now()
		proc, err = fwd.start()
		if err != nil {
			fwd.Kill(err)
			return
		}
	}
}

// start starts an ssh client to forward connections
// from a local port to the remote port.
func (fwd *sshForwarder) start() (p *sshProc, err error) {
	defer trivial.ErrorContextf(&err, "cannot start SSH client")
	args := []string{
		"-T",
		"-N",
		"-o", "StrictHostKeyChecking no",
		"-o", "PasswordAuthentication no",
		"-L", fmt.Sprintf(fmt.Sprintf("%s:localhost:%s", fwd.localAddr, fwd.remotePort)),
		"-p", fmt.Sprint(sshRemotePort),
	}
	if fwd.keyFile != "" {
		args = append(args, "-i", fwd.keyFile)
	}
	args = append(args, sshUser+fwd.remoteHost)

	c := exec.Command("ssh", args...)
	log.Printf("state: starting ssh client: %q", c.Args)
	output, err := c.StdoutPipe()
	if err != nil {
		return nil, err
	}
	c.Stderr = c.Stdout

	err = c.Start()
	if err != nil {
		return nil, err
	}

	wait := make(chan error, 1)
	go func() {
		wait <- c.Wait()
	}()

	errorc := make(chan *sshError, 1)
	go func() {
		r := bufio.NewReader(output)
		for {
			line, err := r.ReadString('\n')
			if err != nil {
				// Discard any error, because it's likely to be
				// from the output being closed, and we can't do
				// much about it anyway - we're more interested
				// in the ssh exit status.
				return
			}
			if err := parseSSHError(strings.TrimRight(line, "\r\n")); err != nil {
				errorc <- err
				return
			}
		}
	}()

	return &sshProc{
		error:  errorc,
		wait:   wait,
		output: output,
		cmd:    c,
	}, nil
}

// sshProc represents a running ssh process.
type sshProc struct {
	error  <-chan *sshError // Error printed by ssh.
	wait   <-chan error     // SSH exit status.
	output io.Closer        // The output pipe.
	cmd    *exec.Cmd        // The running command.
}

func (p *sshProc) stop() {
	p.output.Close()
	p.cmd.Process.Kill()
}

// sshError represents an error printed by ssh.
type sshError struct {
	fatal   bool // If true, there's no point in retrying.
	unknown bool // Whether we've failed to recognise the error message.
	msg     string
}

func (e *sshError) Error() string {
	return "ssh: " + e.msg
}

// parseSSHError parses an error as printed by ssh.
// If it's not actually an error, it returns nil.
func parseSSHError(s string) *sshError {
	if strings.HasPrefix(s, "ssh: ") {
		s = s[len("ssh: "):]
	}
	log.Printf("state: ssh: %s", s)
	err := &sshError{msg: s}
	switch {
	case strings.HasPrefix(s, "Warning: Permanently added"):
		// Even with a null host file, and ignoring strict host checking
		// we'll end up with a "Permanently added" warning.
		// suppress it as it's effectively normal for our usage...
		return nil

	case strings.HasPrefix(s, "Could not resolve hostname"):
		err.fatal = true

	case strings.Contains(s, "Permission denied"):
		err.fatal = true

	default:
		err.unknown = true
	}
	return err
}

func choosePort() (int, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}
<commit_msg>state: increase mgo dial timeout when using ssh<commit_after>package state

import (
	"bufio"
	"fmt"
	"io"
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/trivial"
	"launchpad.net/tomb"
	"net"
	"os/exec"
	"strings"
	"time"
)

// These variables would be constants but they are varied
// for testing purposes.
var (
	sshRemotePort    = 22
	sshRetryInterval = time.Second
	sshUser          = "ubuntu@"
)

// sshDial dials the MongoDB instance at addr through
// an SSH proxy.
func sshDial(addr, keyFile string) (*sshForwarder, *mgo.Session, error) {
	fwd, err := newSSHForwarder(addr, keyFile)
	if err != nil {
		return nil, nil, err
	}
	defer trivial.ErrorContextf(&err, "cannot dial MongoDB via SSH at address %s", addr)
	conn, err := mgo.DialWithTimeout(fwd.localAddr, 10*time.Minute)
	if err != nil {
		fwd.stop()
		return nil, nil, err
	}
	return fwd, conn, nil
}

type sshForwarder struct {
	tomb.Tomb
	localAddr  string
	remoteHost string
	remotePort string
	keyFile    string
}

// newSSHForwarder starts an ssh proxy connecting to the
// remote TCP address. The localAddr field holds the
// name of the local proxy address. If keyFile is non-empty,
// it should name a file containing a private identity key.
func newSSHForwarder(remoteAddr, keyFile string) (fwd *sshForwarder, err error) {
	defer trivial.ErrorContextf(&err, "cannot start SSH proxy for address %s", remoteAddr)
	remoteHost, remotePort, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		return nil, err
	}
	localPort, err := choosePort()
	if err != nil {
		return nil, fmt.Errorf("cannot choose local port: %v", err)
	}
	fwd = &sshForwarder{
		localAddr:  fmt.Sprintf("localhost:%d", localPort),
		remoteHost: remoteHost,
		remotePort: remotePort,
		keyFile:    keyFile,
	}
	proc, err := fwd.start()
	if err != nil {
		return nil, err
	}
	go fwd.run(proc)
	return fwd, nil
}

func (fwd *sshForwarder) stop() error {
	fwd.Kill(nil)
	return fwd.Wait()
}

// run is called with the ssh process already active.
// It loops, restarting ssh until it gets an unrecoverable error.
func (fwd *sshForwarder) run(proc *sshProc) {
	defer fwd.Done()
	restartCount := 0
	startTime := time.Now()

	// waitErrTimeout and waitErr become valid when the ssh client has exited.
	var (
		waitErrTimeout <-chan time.Time
		waitErr        error
	)
	for {
		select {
		case <-fwd.Dying():
			proc.stop()
			return

		case sshErr := <-proc.error:
			log.Printf("state: ssh error (will retry: %v): %v", !sshErr.fatal, sshErr)
			if sshErr.fatal {
				proc.stop()
				fwd.Kill(sshErr)
				return
			}
			// If the ssh process dies repeatedly after running
			// for a very short time and we don't recognise the
			// error, we assume that something unknown is
			// going wrong and quit.
			if sshErr.unknown && time.Now().Sub(startTime) < 200*time.Millisecond {
				if restartCount++; restartCount > 10 {
					proc.stop()
					log.Printf("state: too many ssh errors")
					fwd.Kill(fmt.Errorf("too many errors: %v", sshErr))
					return
				}
			} else {
				restartCount = 0
			}

		case waitErr = <-proc.wait:
			// If ssh has exited, we'll wait a little while
			// in case we've got the exit status before we've
			// received the error.
			waitErrTimeout = time.After(100 * time.Millisecond)
			continue

		case <-waitErrTimeout:
			log.Printf("state: ssh client exited silently: %v", waitErr)
			// We only get here if ssh exits when no fatal
			// errors have been printed by ssh.  In that
			// case, it's probably best to treat it as fatal
			// rather than restarting blindly.
			proc.stop()
			if waitErr == nil {
				waitErr = fmt.Errorf("ssh exited silently")
			} else {
				waitErr = fmt.Errorf("ssh exited silently: %v", waitErr)
			}
			fwd.Kill(waitErr)
			return
		}
		proc.stop()
		var err error
		time.Sleep(sshRetryInterval)
		startTime = time.Now()
		proc, err = fwd.start()
		if err != nil {
			fwd.Kill(err)
			return
		}
	}
}

// start starts an ssh client to forward connections
// from a local port to the remote port.
func (fwd *sshForwarder) start() (p *sshProc, err error) {
	defer trivial.ErrorContextf(&err, "cannot start SSH client")
	args := []string{
		"-T",
		"-N",
		"-o", "StrictHostKeyChecking no",
		"-o", "PasswordAuthentication no",
		"-L", fmt.Sprintf(fmt.Sprintf("%s:localhost:%s", fwd.localAddr, fwd.remotePort)),
		"-p", fmt.Sprint(sshRemotePort),
	}
	if fwd.keyFile != "" {
		args = append(args, "-i", fwd.keyFile)
	}
	args = append(args, sshUser+fwd.remoteHost)

	c := exec.Command("ssh", args...)
	log.Printf("state: starting ssh client: %q", c.Args)
	output, err := c.StdoutPipe()
	if err != nil {
		return nil, err
	}
	c.Stderr = c.Stdout

	err = c.Start()
	if err != nil {
		return nil, err
	}

	wait := make(chan error, 1)
	go func() {
		wait <- c.Wait()
	}()

	errorc := make(chan *sshError, 1)
	go func() {
		r := bufio.NewReader(output)
		for {
			line, err := r.ReadString('\n')
			if err != nil {
				// Discard any error, because it's likely to be
				// from the output being closed, and we can't do
				// much about it anyway - we're more interested
				// in the ssh exit status.
				return
			}
			if err := parseSSHError(strings.TrimRight(line, "\r\n")); err != nil {
				errorc <- err
				return
			}
		}
	}()

	return &sshProc{
		error:  errorc,
		wait:   wait,
		output: output,
		cmd:    c,
	}, nil
}

// sshProc represents a running ssh process.
type sshProc struct {
	error  <-chan *sshError // Error printed by ssh.
	wait   <-chan error     // SSH exit status.
	output io.Closer        // The output pipe.
	cmd    *exec.Cmd        // The running command.
}

func (p *sshProc) stop() {
	p.output.Close()
	p.cmd.Process.Kill()
}

// sshError represents an error printed by ssh.
type sshError struct {
	fatal   bool // If true, there's no point in retrying.
	unknown bool // Whether we've failed to recognise the error message.
	msg     string
}

func (e *sshError) Error() string {
	return "ssh: " + e.msg
}

// parseSSHError parses an error as printed by ssh.
// If it's not actually an error, it returns nil.
func parseSSHError(s string) *sshError {
	if strings.HasPrefix(s, "ssh: ") {
		s = s[len("ssh: "):]
	}
	log.Printf("state: ssh: %s", s)
	err := &sshError{msg: s}
	switch {
	case strings.HasPrefix(s, "Warning: Permanently added"):
		// Even with a null host file, and ignoring strict host checking
		// we'll end up with a "Permanently added" warning.
		// suppress it as it's effectively normal for our usage...
		return nil

	case strings.HasPrefix(s, "Could not resolve hostname"):
		err.fatal = true

	case strings.Contains(s, "Permission denied"):
		err.fatal = true

	default:
		err.unknown = true
	}
	return err
}

func choosePort() (int, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}
<|endoftext|>


<commit_before>package mastodon

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"
)

// UpdateEvent is struct for passing status event to app.
type UpdateEvent struct {
	Status *Status `json:"status"`
}

func (e *UpdateEvent) event() {}

// NotificationEvent is struct for passing notification event to app.
type NotificationEvent struct{}

func (e *NotificationEvent) event() {}

// DeleteEvent is struct for passing deletion event to app.
type DeleteEvent struct{ ID int64 }

func (e *DeleteEvent) event() {}

// ErrorEvent is struct for passing errors to app.
type ErrorEvent struct{ err error }

func (e *ErrorEvent) event()        {}
func (e *ErrorEvent) Error() string { return e.err.Error() }

// Event is interface passing events to app.
type Event interface {
	event()
}

func handleReader(ctx context.Context, q chan Event, r io.Reader) error {
	name := ""
	s := bufio.NewScanner(r)
	for s.Scan() {
		line := s.Text()
		token := strings.SplitN(line, ":", 2)
		if len(token) != 2 {
			continue
		}
		switch strings.TrimSpace(token[0]) {
		case "event":
			name = strings.TrimSpace(token[1])
		case "data":
			switch name {
			case "update":
				var status Status
				err := json.Unmarshal([]byte(token[1]), &status)
				if err == nil {
					q <- &UpdateEvent{&status}
				}
			case "notification":
			case "delete":
			}
		default:
		}
	}
	return ctx.Err()
}

func (c *Client) streaming(ctx context.Context, p string, tag string) (chan Event, error) {
	u, err := url.Parse(c.config.Server)
	if err != nil {
		return nil, err
	}
	u.Path = path.Join(u.Path, "/api/v1/streaming/"+p)

	params := url.Values{}
	params.Set("tag", tag)
	var resp *http.Response

	q := make(chan Event, 10)
	go func() {
		defer ctx.Done()

		for {
			var in io.Reader
			if tag != "" {
				in = strings.NewReader(params.Encode())
			}
			req, err := http.NewRequest(http.MethodGet, u.String(), in)
			if err == nil {
				req.Header.Set("Authorization", "Bearer "+c.config.AccessToken)
				resp, err = c.Do(req)
				if resp.StatusCode != 200 {
					err = fmt.Errorf("bad request: %v", resp.Status)
				}
			}
			if err == nil {
				err = handleReader(ctx, q, resp.Body)
				if err == nil {
					break
				}
			} else {
				q <- &ErrorEvent{err}
			}
			resp.Body.Close()
			time.Sleep(3 * time.Second)
		}
	}()
	go func() {
		<-ctx.Done()
		if resp != nil && resp.Body != nil {
			resp.Body.Close()
		}
	}()
	return q, nil

}

// StreamingPublic return channel to read events on public.
func (c *Client) StreamingPublic(ctx context.Context) (chan Event, error) {
	return c.streaming(ctx, "public", "")
}

// StreamingHome return channel to read events on home.
func (c *Client) StreamingHome(ctx context.Context) (chan Event, error) {
	return c.streaming(ctx, "home", "")
}

// StreamingHashtag return channel to read events on tagged timeline.
func (c *Client) StreamingHashtag(ctx context.Context, tag string) (chan Event, error) {
	return c.streaming(ctx, "hashtag", tag)
}
<commit_msg>check response<commit_after>package mastodon

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"
)

// UpdateEvent is struct for passing status event to app.
type UpdateEvent struct {
	Status *Status `json:"status"`
}

func (e *UpdateEvent) event() {}

// NotificationEvent is struct for passing notification event to app.
type NotificationEvent struct{}

func (e *NotificationEvent) event() {}

// DeleteEvent is struct for passing deletion event to app.
type DeleteEvent struct{ ID int64 }

func (e *DeleteEvent) event() {}

// ErrorEvent is struct for passing errors to app.
type ErrorEvent struct{ err error }

func (e *ErrorEvent) event()        {}
func (e *ErrorEvent) Error() string { return e.err.Error() }

// Event is interface passing events to app.
type Event interface {
	event()
}

func handleReader(ctx context.Context, q chan Event, r io.Reader) error {
	name := ""
	s := bufio.NewScanner(r)
	for s.Scan() {
		line := s.Text()
		token := strings.SplitN(line, ":", 2)
		if len(token) != 2 {
			continue
		}
		switch strings.TrimSpace(token[0]) {
		case "event":
			name = strings.TrimSpace(token[1])
		case "data":
			switch name {
			case "update":
				var status Status
				err := json.Unmarshal([]byte(token[1]), &status)
				if err == nil {
					q <- &UpdateEvent{&status}
				}
			case "notification":
			case "delete":
			}
		default:
		}
	}
	return ctx.Err()
}

func (c *Client) streaming(ctx context.Context, p string, tag string) (chan Event, error) {
	u, err := url.Parse(c.config.Server)
	if err != nil {
		return nil, err
	}
	u.Path = path.Join(u.Path, "/api/v1/streaming/"+p)

	params := url.Values{}
	params.Set("tag", tag)
	var resp *http.Response

	q := make(chan Event, 10)
	go func() {
		defer ctx.Done()

		for {
			var in io.Reader
			if tag != "" {
				in = strings.NewReader(params.Encode())
			}
			req, err := http.NewRequest(http.MethodGet, u.String(), in)
			if err == nil {
				req.Header.Set("Authorization", "Bearer "+c.config.AccessToken)
				resp, err = c.Do(req)
				if resp != nil && resp.StatusCode != 200 {
					err = fmt.Errorf("bad request: %v", resp.Status)
				}
			}
			if err == nil {
				err = handleReader(ctx, q, resp.Body)
				if err == nil {
					break
				}
			} else {
				q <- &ErrorEvent{err}
			}
			resp.Body.Close()
			time.Sleep(3 * time.Second)
		}
	}()
	go func() {
		<-ctx.Done()
		if resp != nil && resp.Body != nil {
			resp.Body.Close()
		}
	}()
	return q, nil

}

// StreamingPublic return channel to read events on public.
func (c *Client) StreamingPublic(ctx context.Context) (chan Event, error) {
	return c.streaming(ctx, "public", "")
}

// StreamingHome return channel to read events on home.
func (c *Client) StreamingHome(ctx context.Context) (chan Event, error) {
	return c.streaming(ctx, "home", "")
}

// StreamingHashtag return channel to read events on tagged timeline.
func (c *Client) StreamingHashtag(ctx context.Context, tag string) (chan Event, error) {
	return c.streaming(ctx, "hashtag", tag)
}
<|endoftext|>


<commit_before>package adapt

import (
	"fmt"
)

// Surrogate is an interpolant for a function.
type Surrogate struct {
	Inputs    uint      // Number of inputs
	Outputs   uint      // Number of outputs
	Level     uint      // Level of interpolation
	Nodes     uint      // Number of nodes
	Indices   []uint64  // Indices of the nodes
	Surpluses []float64 // Hierarchical surpluses
	Accept    []uint    // Number of nodes accepted at each iteration
	Reject    []uint    // Number of nodes rejected at each iteration
}

func newSurrogate(ni, no uint) *Surrogate {
	return &Surrogate{
		Inputs:    ni,
		Outputs:   no,
		Indices:   make([]uint64, 0, ni),
		Surpluses: make([]float64, 0, no),
		Accept:    make([]uint, 0),
		Reject:    make([]uint, 0),
	}
}

func (s *Surrogate) push(indices []uint64, surpluses []float64) {
	s.Indices = append(s.Indices, indices...)
	s.Surpluses = append(s.Surpluses, surpluses...)
}

func (s *Surrogate) step(level, accepted, rejected uint) {
	s.Level = level
	s.Nodes += accepted
	s.Accept = append(s.Accept, accepted)
	s.Reject = append(s.Reject, rejected)
}

// String returns human-friendly information about the interpolant.
func (s *Surrogate) String() string {
	return fmt.Sprintf("Surrogate{inputs: %d, outputs: %d, level: %d, nodes: %d}",
		s.Inputs, s.Outputs, s.Level, s.Nodes)
}
<commit_msg>Adjust the description of Surrogate<commit_after>package adapt

import (
	"fmt"
)

// Surrogate is an interpolant for a function.
type Surrogate struct {
	Inputs    uint      // Number of inputs
	Outputs   uint      // Number of outputs
	Level     uint      // Interpolation level
	Nodes     uint      // Number of nodes
	Indices   []uint64  // Indices of the nodes
	Surpluses []float64 // Hierarchical surpluses of the nodes
	Accept    []uint    // Number of nodes accepted at each iteration
	Reject    []uint    // Number of nodes rejected at each iteration
}

func newSurrogate(ni, no uint) *Surrogate {
	return &Surrogate{
		Inputs:    ni,
		Outputs:   no,
		Indices:   make([]uint64, 0, ni),
		Surpluses: make([]float64, 0, no),
		Accept:    make([]uint, 0),
		Reject:    make([]uint, 0),
	}
}

func (s *Surrogate) push(indices []uint64, surpluses []float64) {
	s.Indices = append(s.Indices, indices...)
	s.Surpluses = append(s.Surpluses, surpluses...)
}

func (s *Surrogate) step(level, accepted, rejected uint) {
	s.Level = level
	s.Nodes += accepted
	s.Accept = append(s.Accept, accepted)
	s.Reject = append(s.Reject, rejected)
}

// String returns human-friendly information about the interpolant.
func (s *Surrogate) String() string {
	return fmt.Sprintf("Surrogate{inputs: %d, outputs: %d, level: %d, nodes: %d}",
		s.Inputs, s.Outputs, s.Level, s.Nodes)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
)

var (
	source = flag.String("s", path.Join(".", "templates"), "Location of templates")
	output = flag.String("o", "", "Output file")
)

func main() {
	flag.Parse()

	buf := new(bytes.Buffer)
	fmt.Fprint(buf, `package templates

	import "text/template"

  	var templates = map[string]string{`)

	if err := filepath.Walk(*source, func(path string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Ignore non-templates files
		if filepath.Ext(path) != ".tmpl" {
			return nil
		}

		b, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}

		fmt.Fprintf(buf, "\"%s\": `%s`,\n", filepath.Base(path), b)

		return nil
	}); err != nil {
		log.Fatal(err)
	}

	fmt.Fprint(buf, `}

	// Parse parses declared templates.
	func Parse(t *template.Template) (*template.Template, error) {
  		for name, s := range templates {
  			var tmpl *template.Template
  			if t == nil {
  				t = template.New(name)
  			}
  			if name == t.Name() {
  				tmpl = t
  			} else {
  				tmpl = t.New(name)
  			}
	  		if _, err := tmpl.Parse(s); err != nil {
  				return nil, err
  			}
  		}
  		return t, nil
  	}`)

	clean, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}

	file := os.Stdout
	if *output != "" {
		file, err = os.Create(*output)
		if err != nil {
			log.Fatal(err)
		}
	}
	fmt.Fprintln(file, string(clean))
}
<commit_msg>Add HTML template support<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
)

var (
	templateType = flag.String("t", "text", "Type of template (text/html)")
	source       = flag.String("s", path.Join(".", "templates"), "Location of templates")
	output       = flag.String("o", "", "Output file")
)

func main() {
	flag.Parse()

	buf := new(bytes.Buffer)
	fmt.Fprint(buf, fmt.Sprintf(`package templates

	import "%s/template"

  	var templates = map[string]string{`, *templateType))

	if err := filepath.Walk(*source, func(path string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if *templateType == "text" {

			// Ignore non-templates files
			if filepath.Ext(path) != ".tmpl" {
				return nil
			}
		} else {
			// Ignore non-templates files
			if filepath.Ext(path) != ".html" {
				return nil
			}
		}

		b, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}

		fmt.Fprintf(buf, "\"%s\": `%s`,\n", filepath.Base(path), b)

		return nil
	}); err != nil {
		log.Fatal(err)
	}

	fmt.Fprint(buf, `}

	// Parse parses declared templates.
	func Parse(t *template.Template) (*template.Template, error) {
  		for name, s := range templates {
  			var tmpl *template.Template
  			if t == nil {
  				t = template.New(name)
  			}
  			if name == t.Name() {
  				tmpl = t
  			} else {
  				tmpl = t.New(name)
  			}
	  		if _, err := tmpl.Parse(s); err != nil {
  				return nil, err
  			}
  		}
  		return t, nil
  	}`)

	clean, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}

	file := os.Stdout
	if *output != "" {
		file, err = os.Create(*output)
		if err != nil {
			log.Fatal(err)
		}
	}
	fmt.Fprintln(file, string(clean))
}
<|endoftext|>


<commit_before>package jira

var all_templates = map[string]string{
	"debug":          default_debug_template,
	"fields":         default_debug_template,
	"editmeta":       default_debug_template,
	"transmeta":      default_debug_template,
	"createmeta":     default_debug_template,
	"issuelinktypes": default_debug_template,
	"list":           default_list_template,
	"table":          default_table_template,
	"view":           default_view_template,
	"edit":           default_edit_template,
	"transitions":    default_transitions_template,
	"issuetypes":     default_issuetypes_template,
	"create":         default_create_template,
	"comment":        default_comment_template,
	"transition":     default_transition_template,
	"request":        default_debug_template,
}

const default_debug_template = "{{ . | toJson}}\n"

const default_list_template = "{{ range .issues }}{{ .key | append \":\" | printf \"%-12s\"}} {{ .fields.summary }}\n{{ end }}"

const default_table_template = `+{{ "-" | rep 16 }}+{{ "-" | rep 57 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+{{ "-" | rep 12 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+
| {{ "Issue" | printf "%-14s" }} | {{ "Summary" | printf "%-55s" }} | {{ "Priority" | printf "%-12s" }} | {{ "Status" | printf "%-12s" }} | {{ "Age" | printf "%-10s" }} | {{ "Reporter" | printf "%-12s" }} | {{ "Assignee" | printf "%-12s" }} |
+{{ "-" | rep 16 }}+{{ "-" | rep 57 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+{{ "-" | rep 12 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+
{{ range .issues }}| {{ .key | printf "%-14s"}} | {{ .fields.summary | abbrev 55 | printf "%-55s" }} | {{.fields.priority.name | printf "%-12s" }} | {{.fields.status.name | printf "%-12s" }} | {{.fields.created | age | printf "%-10s" }} | {{if .fields.reporter}}{{ .fields.reporter.name | printf "%-12s"}}{{else}}<unassigned>{{end}} | {{if .fields.assignee }}{{.fields.assignee.name | printf "%-12s" }}{{else}}<unassigned>{{end}} |
{{ end }}+{{ "-" | rep 16 }}+{{ "-" | rep 57 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+{{ "-" | rep 12 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+
`

const default_view_template = `issue: {{ .key }}
created: {{ .fields.created }}
status: {{ .fields.status.name }}
summary: {{ .fields.summary }}
project: {{ .fields.project.key }}
components: {{ range .fields.components }}{{ .name }} {{end}}
issuetype: {{ .fields.issuetype.name }}
assignee: {{ if .fields.assignee }}{{ .fields.assignee.name }}{{end}}
reporter: {{ if .fields.reporter }}{{ .fields.reporter.name }}{{end}}
watchers: {{ range .fields.customfield_10110 }}{{ .name }} {{end}}
blockers: {{ range .fields.issuelinks }}{{if .outwardIssue}}{{ .outwardIssue.key }}[{{.outwardIssue.fields.status.name}}]{{end}}{{end}}
depends: {{ range .fields.issuelinks }}{{if .inwardIssue}}{{ .inwardIssue.key }}[{{.inwardIssue.fields.status.name}}]{{end}}{{end}}
priority: {{ .fields.priority.name }}
description: |
  {{ or .fields.description "" | indent 2 }}

comments:
{{ range .fields.comment.comments }}  - | # {{.author.name}} at {{.created}}
    {{ or .body "" | indent 4}}
{{end}}
`
const default_edit_template = `# issue: {{ .key }}
update:
  comment:
    - add:
        body: |~
          {{ or .overrides.comment "" | indent 10 }}
fields:
  summary: {{ or .overrides.summary .fields.summary }}
  components: # Values: {{ range .meta.fields.components.allowedValues }}{{.name}}, {{end}}{{if .overrides.components }}{{ range (split "," .overrides.components)}}
    - name: {{.}}{{end}}{{else}}{{ range .fields.components }}
    - name: {{ .name }}{{end}}{{end}}
  assignee:
    name: {{ if .overrides.assignee }}{{.overrides.assignee}}{{else}}{{if .fields.assignee }}{{ .fields.assignee.name }}{{end}}{{end}}
  reporter:
    name: {{ if .overrides.reporter }}{{ .overrides.reporter }}{{else if .fields.reporter}}{{ .fields.reporter.name }}{{end}}
  # watchers
  customfield_10110: {{ range .fields.customfield_10110 }}
    - name: {{ .name }}{{end}}{{if .overrides.watcher}}
    - name: {{ .overrides.watcher}}{{end}}
  priority: # Values: {{ range .meta.fields.priority.allowedValues }}{{.name}}, {{end}}
    name: {{ or .overrides.priority .fields.priority.name }}
  description: |~
    {{ or .overrides.description (or .fields.description "") | indent 4 }}
# comments:
# {{ range .fields.comment.comments }}  - | # {{.author.name}} at {{.created}}
#     {{ or .body "" | indent 4 | comment}}
# {{end}}
`
const default_transitions_template = `{{ range .transitions }}{{.id }}: {{.name}}
{{end}}`

const default_issuetypes_template = `{{ range .projects }}{{ range .issuetypes }}{{color "+bh"}}{{.name | append ":" | printf "%-13s" }}{{color "reset"}} {{.description}}
{{end}}{{end}}`

const default_create_template = `fields:
  project:
    key: {{ .overrides.project }}
  issuetype:
    name: {{ .overrides.issuetype }}
  summary: {{ or .overrides.summary "" }}
  priority: # Values: {{ range .meta.fields.priority.allowedValues }}{{.name}}, {{end}}
    name: {{ or .overrides.priority "unassigned" }}
  components: # Values: {{ range .meta.fields.components.allowedValues }}{{.name}}, {{end}}{{ range split "," (or .overrides.components "")}}
    - name: {{ . }}{{end}}
  description: |~
    {{ or .overrides.description "" | indent 4 }}
  assignee:
    name: {{ or .overrides.assignee "" }}
  reporter:
    name: {{ or .overrides.reporter .overrides.user }}
  # watchers
  customfield_10110: {{ range split "," (or .overrides.watchers "")}}
    - name: {{.}}{{end}}
    - name:
`

const default_comment_template = `body: |~
  {{ or .overrides.comment "" | indent 2 }}
`

const default_transition_template = `update:
  comment:
    - add:
        body: |~
          {{ or .overrides.comment "" | indent 10 }}
fields:{{if .meta.fields.assignee}}
  assignee:
    name: {{if .overrides.assignee}}{{.overrides.assignee}}{{else}}{{if .fields.assignee}}{{.fields.assignee.name}}{{end}}{{end}}{{end}}{{if .meta.fields.components}}
  components: # Values: {{ range .meta.fields.components.allowedValues }}{{.name}}, {{end}}{{if .overrides.components }}{{ range (split "," .overrides.components)}}
    - name: {{.}}{{end}}{{else}}{{ range .fields.components }}
    - name: {{ .name }}{{end}}{{end}}{{end}}{{if .meta.fields.description}}
  description: {{or .overrides.description .fields.description }}{{end}}{{if .meta.fields.fixVersions}}{{if .meta.fields.fixVersions.allowedValues}}
  fixVersions: # Values: {{ range .meta.fields.fixVersions.allowedValues }}{{.name}}, {{end}}{{if .overrides.fixVersions}}{{ range (split "," .overrides.fixVersions)}}
    - name: {{.}}{{end}}{{else}}{{range .fields.fixVersions}}
    - name: {{.}}{{end}}{{end}}{{end}}{{end}}{{if .meta.fields.issuetype}}
  issuetype: # Values: {{ range .meta.fields.issuetype.allowedValues }}{{.name}}, {{end}}
    name: {{if .overrides.issuetype}}{{.overrides.issuetype}}{{else}}{{if .fields.issuetype}}{{.fields.issuetype.name}}{{end}}{{end}}{{end}}{{if .meta.fields.labels}}
  labels: {{range .fields.labels}}
    - {{.}}{{end}}{{if .overrides.labels}}{{range (split "," .overrides.labels)}}
    - {{.}}{{end}}{{end}}{{end}}{{if .meta.fields.priority}}
  priority: # Values: {{ range .meta.fields.priority.allowedValues }}{{.name}}, {{end}}
    name: {{ or .overrides.priority "unassigned" }}{{end}}{{if .meta.fields.reporter}}
  reporter:
    name: {{if .overrides.reporter}}{{.overrides.reporter}}{{else}}{{if .fields.reporter}}{{.fields.reporter.name}}{{end}}{{end}}{{end}}{{if .meta.fields.resolution}}
  resolution: # Values: {{ range .meta.fields.resolution.allowedValues }}{{.name}}, {{end}}
    name: {{if .overrides.resolution}}{{.overrides.resolution}}{{else if .fields.resolution}}{{.fields.resolution.name}}{{else}}Fixed{{end}}{{end}}{{if .meta.fields.summary}}
  summary: {{or .overrides.summary .fields.summary}}{{end}}{{if .meta.fields.versions.allowedValues}}
  versions: # Values: {{ range .meta.fields.versions.allowedValues }}{{.name}}, {{end}}{{if .overrides.versions}}{{ range (split "," .overrides.versions)}}
    - name: {{.}}{{end}}{{else}}{{range .fields.versions}}
    - name: {{.}}{{end}}{{end}}{{end}}
transition:
  id: {{ .transition.id }}
  name: {{ .transition.name }}
`
<commit_msg>Reporter is not generally editable.<commit_after>package jira

var all_templates = map[string]string{
	"debug":          default_debug_template,
	"fields":         default_debug_template,
	"editmeta":       default_debug_template,
	"transmeta":      default_debug_template,
	"createmeta":     default_debug_template,
	"issuelinktypes": default_debug_template,
	"list":           default_list_template,
	"table":          default_table_template,
	"view":           default_view_template,
	"edit":           default_edit_template,
	"transitions":    default_transitions_template,
	"issuetypes":     default_issuetypes_template,
	"create":         default_create_template,
	"comment":        default_comment_template,
	"transition":     default_transition_template,
	"request":        default_debug_template,
}

const default_debug_template = "{{ . | toJson}}\n"

const default_list_template = "{{ range .issues }}{{ .key | append \":\" | printf \"%-12s\"}} {{ .fields.summary }}\n{{ end }}"

const default_table_template = `+{{ "-" | rep 16 }}+{{ "-" | rep 57 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+{{ "-" | rep 12 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+
| {{ "Issue" | printf "%-14s" }} | {{ "Summary" | printf "%-55s" }} | {{ "Priority" | printf "%-12s" }} | {{ "Status" | printf "%-12s" }} | {{ "Age" | printf "%-10s" }} | {{ "Reporter" | printf "%-12s" }} | {{ "Assignee" | printf "%-12s" }} |
+{{ "-" | rep 16 }}+{{ "-" | rep 57 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+{{ "-" | rep 12 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+
{{ range .issues }}| {{ .key | printf "%-14s"}} | {{ .fields.summary | abbrev 55 | printf "%-55s" }} | {{.fields.priority.name | printf "%-12s" }} | {{.fields.status.name | printf "%-12s" }} | {{.fields.created | age | printf "%-10s" }} | {{if .fields.reporter}}{{ .fields.reporter.name | printf "%-12s"}}{{else}}<unassigned>{{end}} | {{if .fields.assignee }}{{.fields.assignee.name | printf "%-12s" }}{{else}}<unassigned>{{end}} |
{{ end }}+{{ "-" | rep 16 }}+{{ "-" | rep 57 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+{{ "-" | rep 12 }}+{{ "-" | rep 14 }}+{{ "-" | rep 14 }}+
`

const default_view_template = `issue: {{ .key }}
created: {{ .fields.created }}
status: {{ .fields.status.name }}
summary: {{ .fields.summary }}
project: {{ .fields.project.key }}
components: {{ range .fields.components }}{{ .name }} {{end}}
issuetype: {{ .fields.issuetype.name }}
assignee: {{ if .fields.assignee }}{{ .fields.assignee.name }}{{end}}
reporter: {{ if .fields.reporter }}{{ .fields.reporter.name }}{{end}}
watchers: {{ range .fields.customfield_10110 }}{{ .name }} {{end}}
blockers: {{ range .fields.issuelinks }}{{if .outwardIssue}}{{ .outwardIssue.key }}[{{.outwardIssue.fields.status.name}}]{{end}}{{end}}
depends: {{ range .fields.issuelinks }}{{if .inwardIssue}}{{ .inwardIssue.key }}[{{.inwardIssue.fields.status.name}}]{{end}}{{end}}
priority: {{ .fields.priority.name }}
description: |
  {{ or .fields.description "" | indent 2 }}

comments:
{{ range .fields.comment.comments }}  - | # {{.author.name}} at {{.created}}
    {{ or .body "" | indent 4}}
{{end}}
`
const default_edit_template = `# issue: {{ .key }}
update:
  comment:
    - add:
        body: |~
          {{ or .overrides.comment "" | indent 10 }}
fields:
  summary: {{ or .overrides.summary .fields.summary }}
  components: # Values: {{ range .meta.fields.components.allowedValues }}{{.name}}, {{end}}{{if .overrides.components }}{{ range (split "," .overrides.components)}}
    - name: {{.}}{{end}}{{else}}{{ range .fields.components }}
    - name: {{ .name }}{{end}}{{end}}
  assignee:
    name: {{ if .overrides.assignee }}{{.overrides.assignee}}{{else}}{{if .fields.assignee }}{{ .fields.assignee.name }}{{end}}{{end}}
#	reporter:
#	  name: {{ if .overrides.reporter }}{{ .overrides.reporter }}{{else if .fields.reporter}}{{ .fields.reporter.name }}{{end}}
  # watchers
  customfield_10110: {{ range .fields.customfield_10110 }}
    - name: {{ .name }}{{end}}{{if .overrides.watcher}}
    - name: {{ .overrides.watcher}}{{end}}
  priority: # Values: {{ range .meta.fields.priority.allowedValues }}{{.name}}, {{end}}
    name: {{ or .overrides.priority .fields.priority.name }}
  description: |~
    {{ or .overrides.description (or .fields.description "") | indent 4 }}
# comments:
# {{ range .fields.comment.comments }}  - | # {{.author.name}} at {{.created}}
#     {{ or .body "" | indent 4 | comment}}
# {{end}}
`
const default_transitions_template = `{{ range .transitions }}{{.id }}: {{.name}}
{{end}}`

const default_issuetypes_template = `{{ range .projects }}{{ range .issuetypes }}{{color "+bh"}}{{.name | append ":" | printf "%-13s" }}{{color "reset"}} {{.description}}
{{end}}{{end}}`

const default_create_template = `fields:
  project:
    key: {{ .overrides.project }}
  issuetype:
    name: {{ .overrides.issuetype }}
  summary: {{ or .overrides.summary "" }}
  priority: # Values: {{ range .meta.fields.priority.allowedValues }}{{.name}}, {{end}}
    name: {{ or .overrides.priority "unassigned" }}
  components: # Values: {{ range .meta.fields.components.allowedValues }}{{.name}}, {{end}}{{ range split "," (or .overrides.components "")}}
    - name: {{ . }}{{end}}
  description: |~
    {{ or .overrides.description "" | indent 4 }}
  assignee:
    name: {{ or .overrides.assignee "" }}
  reporter:
    name: {{ or .overrides.reporter .overrides.user }}
  # watchers
  customfield_10110: {{ range split "," (or .overrides.watchers "")}}
    - name: {{.}}{{end}}
    - name:
`

const default_comment_template = `body: |~
  {{ or .overrides.comment "" | indent 2 }}
`

const default_transition_template = `update:
  comment:
    - add:
        body: |~
          {{ or .overrides.comment "" | indent 10 }}
fields:{{if .meta.fields.assignee}}
  assignee:
    name: {{if .overrides.assignee}}{{.overrides.assignee}}{{else}}{{if .fields.assignee}}{{.fields.assignee.name}}{{end}}{{end}}{{end}}{{if .meta.fields.components}}
  components: # Values: {{ range .meta.fields.components.allowedValues }}{{.name}}, {{end}}{{if .overrides.components }}{{ range (split "," .overrides.components)}}
    - name: {{.}}{{end}}{{else}}{{ range .fields.components }}
    - name: {{ .name }}{{end}}{{end}}{{end}}{{if .meta.fields.description}}
  description: {{or .overrides.description .fields.description }}{{end}}{{if .meta.fields.fixVersions}}{{if .meta.fields.fixVersions.allowedValues}}
  fixVersions: # Values: {{ range .meta.fields.fixVersions.allowedValues }}{{.name}}, {{end}}{{if .overrides.fixVersions}}{{ range (split "," .overrides.fixVersions)}}
    - name: {{.}}{{end}}{{else}}{{range .fields.fixVersions}}
    - name: {{.}}{{end}}{{end}}{{end}}{{end}}{{if .meta.fields.issuetype}}
  issuetype: # Values: {{ range .meta.fields.issuetype.allowedValues }}{{.name}}, {{end}}
    name: {{if .overrides.issuetype}}{{.overrides.issuetype}}{{else}}{{if .fields.issuetype}}{{.fields.issuetype.name}}{{end}}{{end}}{{end}}{{if .meta.fields.labels}}
  labels: {{range .fields.labels}}
    - {{.}}{{end}}{{if .overrides.labels}}{{range (split "," .overrides.labels)}}
    - {{.}}{{end}}{{end}}{{end}}{{if .meta.fields.priority}}
  priority: # Values: {{ range .meta.fields.priority.allowedValues }}{{.name}}, {{end}}
    name: {{ or .overrides.priority "unassigned" }}{{end}}{{if .meta.fields.reporter}}
  reporter:
    name: {{if .overrides.reporter}}{{.overrides.reporter}}{{else}}{{if .fields.reporter}}{{.fields.reporter.name}}{{end}}{{end}}{{end}}{{if .meta.fields.resolution}}
  resolution: # Values: {{ range .meta.fields.resolution.allowedValues }}{{.name}}, {{end}}
    name: {{if .overrides.resolution}}{{.overrides.resolution}}{{else if .fields.resolution}}{{.fields.resolution.name}}{{else}}Fixed{{end}}{{end}}{{if .meta.fields.summary}}
  summary: {{or .overrides.summary .fields.summary}}{{end}}{{if .meta.fields.versions.allowedValues}}
  versions: # Values: {{ range .meta.fields.versions.allowedValues }}{{.name}}, {{end}}{{if .overrides.versions}}{{ range (split "," .overrides.versions)}}
    - name: {{.}}{{end}}{{else}}{{range .fields.versions}}
    - name: {{.}}{{end}}{{end}}{{end}}
transition:
  id: {{ .transition.id }}
  name: {{ .transition.name }}
`
<|endoftext|>