<commit_before>package speed

import (
	"errors"
	"regexp"
	"sync"
)

// Registry defines a valid set of instance domains and metrics
type Registry interface {
	// checks if an instance domain of the passed name is already present or not
	HasInstanceDomain(name string) bool

	// checks if an metric of the passed name is already present or not
	HasMetric(name string) bool

	// returns the number of Metrics in the current registry
	MetricCount() int

	// returns the number of Instance Domains in the current registry
	InstanceDomainCount() int

	// returns the number of instances across all instance domains in the current registry
	InstanceCount() int

	// returns the number of non null strings initialized in the current registry
	StringCount() int

	// adds a InstanceDomain object to the writer
	AddInstanceDomain(InstanceDomain) error

	// adds a InstanceDomain object after constructing it using passed name and instances
	AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error)

	// adds a Metric object to the writer
	AddMetric(Metric) error

	// adds a Metric object after parsing the passed string for Instances and InstanceDomains
	AddSingletonMetricByString(name string, initialval interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error)
}

// PCPRegistry implements a registry for PCP as the client
type PCPRegistry struct {
	instanceDomains map[string]*PCPInstanceDomain // a cache for instanceDomains
	metrics         map[string]PCPMetric          // a cache for metrics

	// locks
	indomlock   sync.RWMutex
	metricslock sync.RWMutex

	// offsets
	instanceoffset int
	indomoffset    int
	metricsoffset  int
	valuesoffset   int
	stringsoffset  int

	// counts
	instanceCount int
	stringcount   int

	mapped bool
}

// NewPCPRegistry creates a new PCPRegistry object
func NewPCPRegistry() *PCPRegistry {
	return &PCPRegistry{
		instanceDomains: make(map[string]*PCPInstanceDomain),
		metrics:         make(map[string]PCPMetric),
	}
}

// InstanceCount returns the number of instances across all indoms in the registry
func (r *PCPRegistry) InstanceCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return r.instanceCount
}

// InstanceDomainCount returns the number of instance domains in the registry
func (r *PCPRegistry) InstanceDomainCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return len(r.instanceDomains)
}

// MetricCount returns the number of metrics in the registry
func (r *PCPRegistry) MetricCount() int {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	return len(r.metrics)
}

// StringCount returns the number of strings in the registry
func (r *PCPRegistry) StringCount() int { return r.stringcount }

// HasInstanceDomain returns true if the registry already has an indom of the specified name
func (r *PCPRegistry) HasInstanceDomain(name string) bool {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	_, present := r.instanceDomains[name]
	return present
}

// HasMetric returns true if the registry already has a metric of the specified name
func (r *PCPRegistry) HasMetric(name string) bool {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	_, present := r.metrics[name]
	return present
}

// AddInstanceDomain will add a new instance domain to the current registry
func (r *PCPRegistry) AddInstanceDomain(indom InstanceDomain) error {
	if r.HasInstanceDomain(indom.Name()) {
		return errors.New("InstanceDomain is already defined for the current registry")
	}

	r.indomlock.Lock()
	defer r.indomlock.Unlock()

	if r.mapped {
		return errors.New("Cannot add an indom when a mapping is active")
	}

	r.instanceDomains[indom.Name()] = indom.(*PCPInstanceDomain)
	r.instanceCount += indom.InstanceCount()

	if indom.(*PCPInstanceDomain).shortDescription.val != "" {
		r.stringcount++
	}

	if indom.(*PCPInstanceDomain).longDescription.val != "" {
		r.stringcount++
	}

	return nil
}

// AddMetric will add a new metric to the current registry
func (r *PCPRegistry) AddMetric(m Metric) error {
	if r.HasMetric(m.Name()) {
		return errors.New("Metric is already defined for the current registry")
	}

	r.metricslock.Lock()
	defer r.metricslock.Unlock()

	if r.mapped {
		return errors.New("Cannot add a metric when a mapping is active")
	}

	pcpm := m.(PCPMetric)

	r.metrics[m.Name()] = pcpm

	if pcpm.Indom() != nil {
		err := r.AddInstanceDomain(pcpm.Indom())
		if err != nil {
			return err
		}
	}

	if pcpm.ShortDescription().String() != "" {
		r.stringcount++
	}

	if pcpm.LongDescription().String() != "" {
		r.stringcount++
	}

	return nil
}

// AddInstanceDomainByName adds an instance domain using passed parameters
func (r *PCPRegistry) AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error) {
	if r.HasInstanceDomain(name) {
		return nil, errors.New("The InstanceDomain already exists for this registry")
	}

	indom, err := NewPCPInstanceDomain(name, "", "")
	if err != nil {
		return nil, err
	}

	for _, i := range instances {
		err = indom.AddInstance(i)
		if err != nil {
			return nil, err
		}
	}

	err = r.AddInstanceDomain(indom)
	if err != nil {
		return nil, err
	}

	return indom, nil
}

// AddSingletonMetricByString adds a new Singleton Metric
func (r *PCPRegistry) AddSingletonMetricByString(name string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error) {
	if r.HasMetric(name) {
		return nil, errors.New("The Metric already exists for this registry")
	}

	m, err := NewPCPSingletonMetric(val, name, t, s, u, "", "")
	if err != nil {
		return nil, err
	}

	err = r.AddMetric(m)
	if err != nil {
		return nil, err
	}

	return m, nil
}

// IdentifierPat contains the pattern for a valid name identifier
const identifierPat = "[\\p{L}\\p{N}]+"

const p = "\\A((" + identifierPat + ")(\\." + identifierPat + ")*?)(\\[(" + identifierPat + ")\\])?((\\." + identifierPat + ")*)\\z"

// identifierRegex gets the *regexp.Regexp object representing a valid metric identifier
var identifierRegex, _ = regexp.Compile(p)

func parseString(name string) (iname string, indomname string, mname string, err error) {
	if !identifierRegex.MatchString(name) {
		return "", "", "", errors.New("I don't know this")
	}

	matches := identifierRegex.FindStringSubmatch(name)
	iname, indomname, mname, err = matches[5], matches[1], matches[1]+matches[6], nil
	return
}
<commit_msg>registry: track value count separately<commit_after>package speed

import (
	"errors"
	"regexp"
	"sync"
)

// Registry defines a valid set of instance domains and metrics
type Registry interface {
	// checks if an instance domain of the passed name is already present or not
	HasInstanceDomain(name string) bool

	// checks if an metric of the passed name is already present or not
	HasMetric(name string) bool

	// returns the number of Metrics in the current registry
	MetricCount() int

	// returns the number of Values in the current registry
	ValuesCount() int

	// returns the number of Instance Domains in the current registry
	InstanceDomainCount() int

	// returns the number of instances across all instance domains in the current registry
	InstanceCount() int

	// returns the number of non null strings initialized in the current registry
	StringCount() int

	// adds a InstanceDomain object to the writer
	AddInstanceDomain(InstanceDomain) error

	// adds a InstanceDomain object after constructing it using passed name and instances
	AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error)

	// adds a Metric object to the writer
	AddMetric(Metric) error

	// adds a Metric object after parsing the passed string for Instances and InstanceDomains
	AddSingletonMetricByString(name string, initialval interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error)
}

// PCPRegistry implements a registry for PCP as the client
type PCPRegistry struct {
	instanceDomains map[string]*PCPInstanceDomain // a cache for instanceDomains
	metrics         map[string]PCPMetric          // a cache for metrics

	// locks
	indomlock   sync.RWMutex
	metricslock sync.RWMutex

	// offsets
	instanceoffset int
	indomoffset    int
	metricsoffset  int
	valuesoffset   int
	stringsoffset  int

	// counts
	instanceCount int
	valueCount    int
	stringcount   int

	mapped bool
}

// NewPCPRegistry creates a new PCPRegistry object
func NewPCPRegistry() *PCPRegistry {
	return &PCPRegistry{
		instanceDomains: make(map[string]*PCPInstanceDomain),
		metrics:         make(map[string]PCPMetric),
	}
}

// InstanceCount returns the number of instances across all indoms in the registry
func (r *PCPRegistry) InstanceCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return r.instanceCount
}

// InstanceDomainCount returns the number of instance domains in the registry
func (r *PCPRegistry) InstanceDomainCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return len(r.instanceDomains)
}

// MetricCount returns the number of metrics in the registry
func (r *PCPRegistry) MetricCount() int {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	return len(r.metrics)
}

// ValuesCount returns the number of values in the registry
func (r *PCPRegistry) ValuesCount() int { return r.valueCount }

// StringCount returns the number of strings in the registry
func (r *PCPRegistry) StringCount() int { return r.stringcount }

// HasInstanceDomain returns true if the registry already has an indom of the specified name
func (r *PCPRegistry) HasInstanceDomain(name string) bool {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	_, present := r.instanceDomains[name]
	return present
}

// HasMetric returns true if the registry already has a metric of the specified name
func (r *PCPRegistry) HasMetric(name string) bool {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	_, present := r.metrics[name]
	return present
}

// AddInstanceDomain will add a new instance domain to the current registry
func (r *PCPRegistry) AddInstanceDomain(indom InstanceDomain) error {
	if r.HasInstanceDomain(indom.Name()) {
		return errors.New("InstanceDomain is already defined for the current registry")
	}

	r.indomlock.Lock()
	defer r.indomlock.Unlock()

	if r.mapped {
		return errors.New("Cannot add an indom when a mapping is active")
	}

	r.instanceDomains[indom.Name()] = indom.(*PCPInstanceDomain)
	r.instanceCount += indom.InstanceCount()

	if indom.(*PCPInstanceDomain).shortDescription.val != "" {
		r.stringcount++
	}

	if indom.(*PCPInstanceDomain).longDescription.val != "" {
		r.stringcount++
	}

	return nil
}

// AddMetric will add a new metric to the current registry
func (r *PCPRegistry) AddMetric(m Metric) error {
	if r.HasMetric(m.Name()) {
		return errors.New("Metric is already defined for the current registry")
	}

	pcpm := m.(PCPMetric)

	// if it is an indom metric
	if pcpm.Indom() != nil && !r.HasInstanceDomain(pcpm.Indom().Name()) {
		return errors.New("Instance Domain is not defined for current registry")
	}

	if r.mapped {
		return errors.New("Cannot add a metric when a mapping is active")
	}

	r.metricslock.Lock()
	defer r.metricslock.Unlock()

	r.metrics[m.Name()] = pcpm

	if pcpm.Indom() != nil {
		r.valueCount += pcpm.Indom().InstanceCount()
	} else {
		r.valueCount++
	}

	if pcpm.ShortDescription().String() != "" {
		r.stringcount++
	}

	if pcpm.LongDescription().String() != "" {
		r.stringcount++
	}

	return nil
}

// AddInstanceDomainByName adds an instance domain using passed parameters
func (r *PCPRegistry) AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error) {
	if r.HasInstanceDomain(name) {
		return nil, errors.New("The InstanceDomain already exists for this registry")
	}

	indom, err := NewPCPInstanceDomain(name, instances, "", "")
	if err != nil {
		return nil, err
	}

	err = r.AddInstanceDomain(indom)
	if err != nil {
		return nil, err
	}

	return indom, nil
}

// AddSingletonMetricByString adds a new Singleton Metric
func (r *PCPRegistry) AddSingletonMetricByString(name string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error) {
	if r.HasMetric(name) {
		return nil, errors.New("The Metric already exists for this registry")
	}

	m, err := NewPCPSingletonMetric(val, name, t, s, u, "", "")
	if err != nil {
		return nil, err
	}

	err = r.AddMetric(m)
	if err != nil {
		return nil, err
	}

	return m, nil
}

// IdentifierPat contains the pattern for a valid name identifier
const identifierPat = "[\\p{L}\\p{N}]+"

const p = "\\A((" + identifierPat + ")(\\." + identifierPat + ")*?)(\\[(" + identifierPat + ")\\])?((\\." + identifierPat + ")*)\\z"

// identifierRegex gets the *regexp.Regexp object representing a valid metric identifier
var identifierRegex, _ = regexp.Compile(p)

func parseString(name string) (iname string, indomname string, mname string, err error) {
	if !identifierRegex.MatchString(name) {
		return "", "", "", errors.New("I don't know this")
	}

	matches := identifierRegex.FindStringSubmatch(name)
	iname, indomname, mname, err = matches[5], matches[1], matches[1]+matches[6], nil
	return
}
<|endoftext|>


<commit_before>package w32registry

import (
	"syscall"
	"unsafe"

	"github.com/hnakamur/w32syscall"
)

func CreateKey(key syscall.Handle, subkey string, class string, options uint32, desiredAccess uint32, securityAttributes *syscall.SecurityAttributes, disposition *uint32) (result syscall.Handle, err error) {
	var (
		subkeyp, classp *uint16
		reserved        uint32
	)
	if len(subkey) > 0 {
		subkeyp, err = syscall.UTF16PtrFromString(subkey)
		if err != nil {
			return
		}
	}
	if len(class) > 0 {
		classp, err = syscall.UTF16PtrFromString(class)
		if err != nil {
			return
		}
	}
	err = w32syscall.RegCreateKeyEx(key, subkeyp, reserved, classp, options, desiredAccess, securityAttributes, &result, disposition)
	return
}

func DeleteKeyValue(key syscall.Handle, subkey string, valname string) error {
	var (
		subkeyp, valnamep *uint16
		err               error
	)
	if len(subkey) > 0 {
		subkeyp, err = syscall.UTF16PtrFromString(subkey)
		if err != nil {
			return err
		}
	}
	if len(valname) > 0 {
		valnamep, err = syscall.UTF16PtrFromString(valname)
		if err != nil {
			return err
		}
	}
	return w32syscall.RegDeleteKeyValue(key, subkeyp, valnamep)
}

func DeleteTree(key syscall.Handle, subkey string) error {
	var (
		subkeyp *uint16
		err     error
	)
	if len(subkey) > 0 {
		subkeyp, err = syscall.UTF16PtrFromString(subkey)
		if err != nil {
			return err
		}
	}
	return w32syscall.RegDeleteTree(key, subkeyp)
}

func SetKeyValueString(key syscall.Handle, subkey string, valname string, value string) error {
	var buf []uint16
	buf, err := syscall.UTF16FromString(value)
	if err != nil {
		return err
	}
	bufLen := uint32(len(buf) + 2) // 2 for the terminating null character
	return setKeyValue(key, subkey, valname, syscall.REG_SZ, (*byte)(unsafe.Pointer(&buf[0])), bufLen)
}

func SetKeyValueUint32(key syscall.Handle, subkey string, valname string, value uint32) error {
	valLen := uint32(4) // uint32 size in bytes
	return setKeyValue(key, subkey, valname, syscall.REG_DWORD, (*byte)(unsafe.Pointer(&value)), valLen)
}

func setKeyValue(key syscall.Handle, subkey string, valname string, valtype uint32, buf *byte, buflen uint32) error {
	subkeyp, err := syscall.UTF16PtrFromString(subkey)
	if err != nil {
		return err
	}
	valnamep, err := syscall.UTF16PtrFromString(valname)
	if err != nil {
		return err
	}
	return w32syscall.RegSetKeyValue(key, subkeyp, valnamep, valtype, buf, buflen)
}
<commit_msg>Add GetValueString and GetValueUint32<commit_after>package w32registry

import (
	"syscall"
	"unsafe"

	"github.com/hnakamur/w32syscall"
)

const (
	dwordSize  = 4
	uint16Size = 2
)

func CreateKey(key syscall.Handle, subkey string, class string, options uint32, desiredAccess uint32, securityAttributes *syscall.SecurityAttributes, disposition *uint32) (result syscall.Handle, err error) {
	var (
		subkeyp, classp *uint16
		reserved        uint32
	)
	if len(subkey) > 0 {
		subkeyp, err = syscall.UTF16PtrFromString(subkey)
		if err != nil {
			return
		}
	}
	if len(class) > 0 {
		classp, err = syscall.UTF16PtrFromString(class)
		if err != nil {
			return
		}
	}
	err = w32syscall.RegCreateKeyEx(key, subkeyp, reserved, classp, options, desiredAccess, securityAttributes, &result, disposition)
	return
}

func DeleteKeyValue(key syscall.Handle, subkey string, valname string) error {
	var (
		subkeyp, valnamep *uint16
		err               error
	)
	if len(subkey) > 0 {
		subkeyp, err = syscall.UTF16PtrFromString(subkey)
		if err != nil {
			return err
		}
	}
	if len(valname) > 0 {
		valnamep, err = syscall.UTF16PtrFromString(valname)
		if err != nil {
			return err
		}
	}
	return w32syscall.RegDeleteKeyValue(key, subkeyp, valnamep)
}

func DeleteTree(key syscall.Handle, subkey string) error {
	var (
		subkeyp *uint16
		err     error
	)
	if len(subkey) > 0 {
		subkeyp, err = syscall.UTF16PtrFromString(subkey)
		if err != nil {
			return err
		}
	}
	return w32syscall.RegDeleteTree(key, subkeyp)
}

func GetValueString(key syscall.Handle, subkey string, valname string) (value string, err error) {
	var bufLen uint32
	subkeyp, err := syscall.UTF16PtrFromString(subkey)
	if err != nil {
		return
	}
	valnamep, err := syscall.UTF16PtrFromString(valname)
	if err != nil {
		return
	}
	var flags uint32 = w32syscall.RRF_RT_REG_SZ
	err = w32syscall.RegGetValue(key, subkeyp, valnamep, flags, nil, nil, &bufLen)
	if err != nil {
		return
	}

	buf := make([]uint16, bufLen)
	err = w32syscall.RegGetValue(key, subkeyp, valnamep, flags, nil, (*byte)(unsafe.Pointer(&buf[0])), &bufLen)
	if err != nil {
		return
	}

	value = syscall.UTF16ToString(buf)
	return
}

// GetValueUint32 returns the DWORD value for the specified key, subkey and valname. It sets err to syscall.ERROR_FILE_NOT_FOUND when key, subkey, or valname is not found.
func GetValueUint32(key syscall.Handle, subkey string, valname string) (value uint32, err error) {
	valLen := uint32(dwordSize)
	err = getValue(key, subkey, valname, w32syscall.RRF_RT_REG_DWORD, nil, (*byte)(unsafe.Pointer(&value)), &valLen)
	return
}

func getValue(key syscall.Handle, subkey string, valname string, flags uint32, valtype *uint32, buf *byte, buflen *uint32) error {
	subkeyp, err := syscall.UTF16PtrFromString(subkey)
	if err != nil {
		return err
	}
	valnamep, err := syscall.UTF16PtrFromString(valname)
	if err != nil {
		return err
	}
	return w32syscall.RegGetValue(key, subkeyp, valnamep, flags, valtype, buf, buflen)
}

// SetKeyValueString returns the string value for the specified key, subkey and valname. It sets err to syscall.ERROR_FILE_NOT_FOUND when key, subkey, or valname is not found.
func SetKeyValueString(key syscall.Handle, subkey string, valname string, value string) error {
	var buf []uint16
	buf, err := syscall.UTF16FromString(value)
	if err != nil {
		return err
	}
	bufLen := uint32(len(buf) * uint16Size)
	return setKeyValue(key, subkey, valname, syscall.REG_SZ, (*byte)(unsafe.Pointer(&buf[0])), bufLen)
}

func SetKeyValueUint32(key syscall.Handle, subkey string, valname string, value uint32) error {
	valLen := uint32(dwordSize)
	return setKeyValue(key, subkey, valname, syscall.REG_DWORD, (*byte)(unsafe.Pointer(&value)), valLen)
}

func setKeyValue(key syscall.Handle, subkey string, valname string, valtype uint32, buf *byte, buflen uint32) error {
	subkeyp, err := syscall.UTF16PtrFromString(subkey)
	if err != nil {
		return err
	}
	valnamep, err := syscall.UTF16PtrFromString(valname)
	if err != nil {
		return err
	}
	return w32syscall.RegSetKeyValue(key, subkeyp, valnamep, valtype, buf, buflen)
}
<|endoftext|>


<commit_before>// Package http implements a BitTorrent frontend via the HTTP protocol as
// described in BEP 3 and BEP 23.
package http

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"net/http"
	"time"

	"github.com/julienschmidt/httprouter"

	"github.com/chihaya/chihaya/bittorrent"
	"github.com/chihaya/chihaya/frontend"
	"github.com/chihaya/chihaya/pkg/log"
	"github.com/chihaya/chihaya/pkg/stop"
)

// Config represents all of the configurable options for an HTTP BitTorrent
// Frontend.
type Config struct {
	Addr                string        `yaml:"addr"`
	HTTPSAddr           string        `yaml:"https_addr"`
	ReadTimeout         time.Duration `yaml:"read_timeout"`
	WriteTimeout        time.Duration `yaml:"write_timeout"`
	IdleTimeout         time.Duration `yaml:"idle_timeout"`
	EnableKeepAlive     bool          `yaml:"enable_keepalive"`
	TLSCertPath         string        `yaml:"tls_cert_path"`
	TLSKeyPath          string        `yaml:"tls_key_path"`
	EnableLegacyPHPURLs bool          `yaml:"enable_legacy_php_urls"`
	EnableRequestTiming bool          `yaml:"enable_request_timing"`
	ParseOptions        `yaml:",inline"`
}

// LogFields renders the current config as a set of Logrus fields.
func (cfg Config) LogFields() log.Fields {
	return log.Fields{
		"addr":                cfg.Addr,
		"httpsAddr":           cfg.HTTPSAddr,
		"readTimeout":         cfg.ReadTimeout,
		"writeTimeout":        cfg.WriteTimeout,
		"idleTimeout":         cfg.IdleTimeout,
		"enableKeepAlive":     cfg.EnableKeepAlive,
		"tlsCertPath":         cfg.TLSCertPath,
		"tlsKeyPath":          cfg.TLSKeyPath,
		"enableLegacyPHPURLs": cfg.EnableLegacyPHPURLs,
		"enableRequestTiming": cfg.EnableRequestTiming,
		"allowIPSpoofing":     cfg.AllowIPSpoofing,
		"realIPHeader":        cfg.RealIPHeader,
		"maxNumWant":          cfg.MaxNumWant,
		"defaultNumWant":      cfg.DefaultNumWant,
		"maxScrapeInfoHashes": cfg.MaxScrapeInfoHashes,
	}
}

// Default config constants.
const (
	defaultReadTimeout  = 2 * time.Second
	defaultWriteTimeout = 2 * time.Second
	defaultIdleTimeout  = 30 * time.Second
)

// Validate sanity checks values set in a config and returns a new config with
// default values replacing anything that is invalid.
//
// This function warns to the logger when a value is changed.
func (cfg Config) Validate() Config {
	validcfg := cfg

	if cfg.ReadTimeout <= 0 {
		validcfg.ReadTimeout = defaultReadTimeout
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.ReadTimeout",
			"provided": cfg.ReadTimeout,
			"default":  validcfg.ReadTimeout,
		})
	}

	if cfg.WriteTimeout <= 0 {
		validcfg.WriteTimeout = defaultWriteTimeout
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.WriteTimeout",
			"provided": cfg.WriteTimeout,
			"default":  validcfg.WriteTimeout,
		})
	}

	if cfg.IdleTimeout <= 0 {
		validcfg.IdleTimeout = defaultIdleTimeout

		if cfg.EnableKeepAlive {
			// If keepalive is disabled, this configuration isn't used anyway.
			log.Warn("falling back to default configuration", log.Fields{
				"name":     "http.IdleTimeout",
				"provided": cfg.IdleTimeout,
				"default":  validcfg.IdleTimeout,
			})
		}
	}

	if cfg.MaxNumWant <= 0 {
		validcfg.MaxNumWant = defaultMaxNumWant
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.MaxNumWant",
			"provided": cfg.MaxNumWant,
			"default":  validcfg.MaxNumWant,
		})
	}

	if cfg.DefaultNumWant <= 0 {
		validcfg.DefaultNumWant = defaultDefaultNumWant
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.DefaultNumWant",
			"provided": cfg.DefaultNumWant,
			"default":  validcfg.DefaultNumWant,
		})
	}

	if cfg.MaxScrapeInfoHashes <= 0 {
		validcfg.MaxScrapeInfoHashes = defaultMaxScrapeInfoHashes
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.MaxScrapeInfoHashes",
			"provided": cfg.MaxScrapeInfoHashes,
			"default":  validcfg.MaxScrapeInfoHashes,
		})
	}

	return validcfg
}

// Frontend represents the state of an HTTP BitTorrent Frontend.
type Frontend struct {
	srv    *http.Server
	tlsSrv *http.Server
	tlsCfg *tls.Config

	logic frontend.TrackerLogic
	Config
}

// NewFrontend creates a new instance of an HTTP Frontend that asynchronously
// serves requests.
func NewFrontend(logic frontend.TrackerLogic, provided Config) (*Frontend, error) {
	cfg := provided.Validate()

	f := &Frontend{
		logic:  logic,
		Config: cfg,
	}

	if cfg.Addr == "" && cfg.HTTPSAddr == "" {
		return nil, errors.New("must specify addr or https_addr or both")
	}

	// If TLS is enabled, create a key pair.
	if cfg.TLSCertPath != "" && cfg.TLSKeyPath != "" {
		var err error
		f.tlsCfg = &tls.Config{
			Certificates: make([]tls.Certificate, 1),
		}
		f.tlsCfg.Certificates[0], err = tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath)
		if err != nil {
			return nil, err
		}
	}

	if cfg.HTTPSAddr != "" && f.tlsCfg == nil {
		return nil, errors.New("must specify tls_cert_path and tls_key_path when using https_addr")
	}
	if cfg.HTTPSAddr == "" && f.tlsCfg != nil {
		return nil, errors.New("must specify https_addr when using tls_cert_path and tls_key_path")
	}

	var listenerHTTP, listenerHTTPS net.Listener
	var err error
	if cfg.Addr != "" {
		listenerHTTP, err = net.Listen("tcp", f.Addr)
		if err != nil {
			return nil, err
		}
	}
	if cfg.HTTPSAddr != "" {
		listenerHTTPS, err = net.Listen("tcp", f.HTTPSAddr)
		if err != nil {
			if listenerHTTP != nil {
				listenerHTTP.Close()
			}
			return nil, err
		}
	}

	if cfg.Addr != "" {
		go func() {
			if err := f.serveHTTP(listenerHTTP); err != nil {
				log.Fatal("failed while serving http", log.Err(err))
			}
		}()
	}

	if cfg.HTTPSAddr != "" {
		go func() {
			if err := f.serveHTTPS(listenerHTTPS); err != nil {
				log.Fatal("failed while serving https", log.Err(err))
			}
		}()
	}

	return f, nil
}

// Stop provides a thread-safe way to shutdown a currently running Frontend.
func (f *Frontend) Stop() stop.Result {
	stopGroup := stop.NewGroup()

	if f.srv != nil {
		stopGroup.AddFunc(f.makeStopFunc(f.srv))
	}
	if f.tlsSrv != nil {
		stopGroup.AddFunc(f.makeStopFunc(f.tlsSrv))
	}

	return stopGroup.Stop()
}

func (f *Frontend) makeStopFunc(stopSrv *http.Server) stop.Func {
	return func() stop.Result {
		c := make(stop.Channel)
		go func() {
			c.Done(stopSrv.Shutdown(context.Background()))
		}()
		return c.Result()
	}
}

func (f *Frontend) handler() http.Handler {
	router := httprouter.New()
	router.GET("/announce", f.announceRoute)
	router.GET("/scrape", f.scrapeRoute)

	if f.EnableLegacyPHPURLs {
		log.Info("http: enabling legacy PHP URLs")
		router.GET("/announce.php", f.announceRoute)
		router.GET("/scrape.php", f.scrapeRoute)
	}

	return router
}

// serveHTTP blocks while listening and serving non-TLS HTTP BitTorrent
// requests until Stop() is called or an error is returned.
func (f *Frontend) serveHTTP(l net.Listener) error {
	f.srv = &http.Server{
		Addr:         f.Addr,
		Handler:      f.handler(),
		ReadTimeout:  f.ReadTimeout,
		WriteTimeout: f.WriteTimeout,
		IdleTimeout:  f.IdleTimeout,
	}

	f.srv.SetKeepAlivesEnabled(f.EnableKeepAlive)

	// Start the HTTP server.
	if err := f.srv.Serve(l); err != http.ErrServerClosed {
		return err
	}
	return nil
}

// serveHTTPS blocks while listening and serving TLS HTTP BitTorrent
// requests until Stop() is called or an error is returned.
func (f *Frontend) serveHTTPS(l net.Listener) error {
	f.tlsSrv = &http.Server{
		Addr:         f.HTTPSAddr,
		TLSConfig:    f.tlsCfg,
		Handler:      f.handler(),
		ReadTimeout:  f.ReadTimeout,
		WriteTimeout: f.WriteTimeout,
	}

	f.tlsSrv.SetKeepAlivesEnabled(f.EnableKeepAlive)

	// Start the HTTP server.
	if err := f.tlsSrv.ServeTLS(l, "", ""); err != http.ErrServerClosed {
		return err
	}
	return nil
}

// announceRoute parses and responds to an Announce.
func (f *Frontend) announceRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	var err error
	var start time.Time
	if f.EnableRequestTiming {
		start = time.Now()
	}
	var af *bittorrent.AddressFamily
	defer func() {
		if f.EnableRequestTiming {
			recordResponseDuration("announce", af, err, time.Since(start))
		} else {
			recordResponseDuration("announce", af, err, time.Duration(0))
		}
	}()

	req, err := ParseAnnounce(r, f.ParseOptions)
	if err != nil {
		WriteError(w, err)
		return
	}
	af = new(bittorrent.AddressFamily)
	*af = req.IP.AddressFamily

	ctx, resp, err := f.logic.HandleAnnounce(context.Background(), req)
	if err != nil {
		WriteError(w, err)
		return
	}

	err = WriteAnnounceResponse(w, resp)
	if err != nil {
		WriteError(w, err)
		return
	}

	go f.logic.AfterAnnounce(ctx, req, resp)
}

// scrapeRoute parses and responds to a Scrape.
func (f *Frontend) scrapeRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	var err error
	var start time.Time
	if f.EnableRequestTiming {
		start = time.Now()
	}
	var af *bittorrent.AddressFamily
	defer func() {
		if f.EnableRequestTiming {
			recordResponseDuration("scrape", af, err, time.Since(start))
		} else {
			recordResponseDuration("scrape", af, err, time.Duration(0))
		}
	}()

	req, err := ParseScrape(r, f.ParseOptions)
	if err != nil {
		WriteError(w, err)
		return
	}

	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		log.Error("http: unable to determine remote address for scrape", log.Err(err))
		WriteError(w, err)
		return
	}

	reqIP := net.ParseIP(host)
	if reqIP.To4() != nil {
		req.AddressFamily = bittorrent.IPv4
	} else if len(reqIP) == net.IPv6len { // implies reqIP.To4() == nil
		req.AddressFamily = bittorrent.IPv6
	} else {
		log.Error("http: invalid IP: neither v4 nor v6", log.Fields{"RemoteAddr": r.RemoteAddr})
		WriteError(w, bittorrent.ErrInvalidIP)
		return
	}
	af = new(bittorrent.AddressFamily)
	*af = req.AddressFamily

	ctx, resp, err := f.logic.HandleScrape(context.Background(), req)
	if err != nil {
		WriteError(w, err)
		return
	}

	err = WriteScrapeResponse(w, resp)
	if err != nil {
		WriteError(w, err)
		return
	}

	go f.logic.AfterScrape(ctx, req, resp)
}
<commit_msg>http: explicitly set Content-Type header<commit_after>// Package http implements a BitTorrent frontend via the HTTP protocol as
// described in BEP 3 and BEP 23.
package http

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"net/http"
	"time"

	"github.com/julienschmidt/httprouter"

	"github.com/chihaya/chihaya/bittorrent"
	"github.com/chihaya/chihaya/frontend"
	"github.com/chihaya/chihaya/pkg/log"
	"github.com/chihaya/chihaya/pkg/stop"
)

// Config represents all of the configurable options for an HTTP BitTorrent
// Frontend.
type Config struct {
	Addr                string        `yaml:"addr"`
	HTTPSAddr           string        `yaml:"https_addr"`
	ReadTimeout         time.Duration `yaml:"read_timeout"`
	WriteTimeout        time.Duration `yaml:"write_timeout"`
	IdleTimeout         time.Duration `yaml:"idle_timeout"`
	EnableKeepAlive     bool          `yaml:"enable_keepalive"`
	TLSCertPath         string        `yaml:"tls_cert_path"`
	TLSKeyPath          string        `yaml:"tls_key_path"`
	EnableLegacyPHPURLs bool          `yaml:"enable_legacy_php_urls"`
	EnableRequestTiming bool          `yaml:"enable_request_timing"`
	ParseOptions        `yaml:",inline"`
}

// LogFields renders the current config as a set of Logrus fields.
func (cfg Config) LogFields() log.Fields {
	return log.Fields{
		"addr":                cfg.Addr,
		"httpsAddr":           cfg.HTTPSAddr,
		"readTimeout":         cfg.ReadTimeout,
		"writeTimeout":        cfg.WriteTimeout,
		"idleTimeout":         cfg.IdleTimeout,
		"enableKeepAlive":     cfg.EnableKeepAlive,
		"tlsCertPath":         cfg.TLSCertPath,
		"tlsKeyPath":          cfg.TLSKeyPath,
		"enableLegacyPHPURLs": cfg.EnableLegacyPHPURLs,
		"enableRequestTiming": cfg.EnableRequestTiming,
		"allowIPSpoofing":     cfg.AllowIPSpoofing,
		"realIPHeader":        cfg.RealIPHeader,
		"maxNumWant":          cfg.MaxNumWant,
		"defaultNumWant":      cfg.DefaultNumWant,
		"maxScrapeInfoHashes": cfg.MaxScrapeInfoHashes,
	}
}

// Default config constants.
const (
	defaultReadTimeout  = 2 * time.Second
	defaultWriteTimeout = 2 * time.Second
	defaultIdleTimeout  = 30 * time.Second
)

// Validate sanity checks values set in a config and returns a new config with
// default values replacing anything that is invalid.
//
// This function warns to the logger when a value is changed.
func (cfg Config) Validate() Config {
	validcfg := cfg

	if cfg.ReadTimeout <= 0 {
		validcfg.ReadTimeout = defaultReadTimeout
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.ReadTimeout",
			"provided": cfg.ReadTimeout,
			"default":  validcfg.ReadTimeout,
		})
	}

	if cfg.WriteTimeout <= 0 {
		validcfg.WriteTimeout = defaultWriteTimeout
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.WriteTimeout",
			"provided": cfg.WriteTimeout,
			"default":  validcfg.WriteTimeout,
		})
	}

	if cfg.IdleTimeout <= 0 {
		validcfg.IdleTimeout = defaultIdleTimeout

		if cfg.EnableKeepAlive {
			// If keepalive is disabled, this configuration isn't used anyway.
			log.Warn("falling back to default configuration", log.Fields{
				"name":     "http.IdleTimeout",
				"provided": cfg.IdleTimeout,
				"default":  validcfg.IdleTimeout,
			})
		}
	}

	if cfg.MaxNumWant <= 0 {
		validcfg.MaxNumWant = defaultMaxNumWant
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.MaxNumWant",
			"provided": cfg.MaxNumWant,
			"default":  validcfg.MaxNumWant,
		})
	}

	if cfg.DefaultNumWant <= 0 {
		validcfg.DefaultNumWant = defaultDefaultNumWant
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.DefaultNumWant",
			"provided": cfg.DefaultNumWant,
			"default":  validcfg.DefaultNumWant,
		})
	}

	if cfg.MaxScrapeInfoHashes <= 0 {
		validcfg.MaxScrapeInfoHashes = defaultMaxScrapeInfoHashes
		log.Warn("falling back to default configuration", log.Fields{
			"name":     "http.MaxScrapeInfoHashes",
			"provided": cfg.MaxScrapeInfoHashes,
			"default":  validcfg.MaxScrapeInfoHashes,
		})
	}

	return validcfg
}

// Frontend represents the state of an HTTP BitTorrent Frontend.
type Frontend struct {
	srv    *http.Server
	tlsSrv *http.Server
	tlsCfg *tls.Config

	logic frontend.TrackerLogic
	Config
}

// NewFrontend creates a new instance of an HTTP Frontend that asynchronously
// serves requests.
func NewFrontend(logic frontend.TrackerLogic, provided Config) (*Frontend, error) {
	cfg := provided.Validate()

	f := &Frontend{
		logic:  logic,
		Config: cfg,
	}

	if cfg.Addr == "" && cfg.HTTPSAddr == "" {
		return nil, errors.New("must specify addr or https_addr or both")
	}

	// If TLS is enabled, create a key pair.
	if cfg.TLSCertPath != "" && cfg.TLSKeyPath != "" {
		var err error
		f.tlsCfg = &tls.Config{
			Certificates: make([]tls.Certificate, 1),
		}
		f.tlsCfg.Certificates[0], err = tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath)
		if err != nil {
			return nil, err
		}
	}

	if cfg.HTTPSAddr != "" && f.tlsCfg == nil {
		return nil, errors.New("must specify tls_cert_path and tls_key_path when using https_addr")
	}
	if cfg.HTTPSAddr == "" && f.tlsCfg != nil {
		return nil, errors.New("must specify https_addr when using tls_cert_path and tls_key_path")
	}

	var listenerHTTP, listenerHTTPS net.Listener
	var err error
	if cfg.Addr != "" {
		listenerHTTP, err = net.Listen("tcp", f.Addr)
		if err != nil {
			return nil, err
		}
	}
	if cfg.HTTPSAddr != "" {
		listenerHTTPS, err = net.Listen("tcp", f.HTTPSAddr)
		if err != nil {
			if listenerHTTP != nil {
				listenerHTTP.Close()
			}
			return nil, err
		}
	}

	if cfg.Addr != "" {
		go func() {
			if err := f.serveHTTP(listenerHTTP); err != nil {
				log.Fatal("failed while serving http", log.Err(err))
			}
		}()
	}

	if cfg.HTTPSAddr != "" {
		go func() {
			if err := f.serveHTTPS(listenerHTTPS); err != nil {
				log.Fatal("failed while serving https", log.Err(err))
			}
		}()
	}

	return f, nil
}

// Stop provides a thread-safe way to shutdown a currently running Frontend.
func (f *Frontend) Stop() stop.Result {
	stopGroup := stop.NewGroup()

	if f.srv != nil {
		stopGroup.AddFunc(f.makeStopFunc(f.srv))
	}
	if f.tlsSrv != nil {
		stopGroup.AddFunc(f.makeStopFunc(f.tlsSrv))
	}

	return stopGroup.Stop()
}

func (f *Frontend) makeStopFunc(stopSrv *http.Server) stop.Func {
	return func() stop.Result {
		c := make(stop.Channel)
		go func() {
			c.Done(stopSrv.Shutdown(context.Background()))
		}()
		return c.Result()
	}
}

func (f *Frontend) handler() http.Handler {
	router := httprouter.New()
	router.GET("/announce", f.announceRoute)
	router.GET("/scrape", f.scrapeRoute)

	if f.EnableLegacyPHPURLs {
		log.Info("http: enabling legacy PHP URLs")
		router.GET("/announce.php", f.announceRoute)
		router.GET("/scrape.php", f.scrapeRoute)
	}

	return router
}

// serveHTTP blocks while listening and serving non-TLS HTTP BitTorrent
// requests until Stop() is called or an error is returned.
func (f *Frontend) serveHTTP(l net.Listener) error {
	f.srv = &http.Server{
		Addr:         f.Addr,
		Handler:      f.handler(),
		ReadTimeout:  f.ReadTimeout,
		WriteTimeout: f.WriteTimeout,
		IdleTimeout:  f.IdleTimeout,
	}

	f.srv.SetKeepAlivesEnabled(f.EnableKeepAlive)

	// Start the HTTP server.
	if err := f.srv.Serve(l); err != http.ErrServerClosed {
		return err
	}
	return nil
}

// serveHTTPS blocks while listening and serving TLS HTTP BitTorrent
// requests until Stop() is called or an error is returned.
func (f *Frontend) serveHTTPS(l net.Listener) error {
	f.tlsSrv = &http.Server{
		Addr:         f.HTTPSAddr,
		TLSConfig:    f.tlsCfg,
		Handler:      f.handler(),
		ReadTimeout:  f.ReadTimeout,
		WriteTimeout: f.WriteTimeout,
	}

	f.tlsSrv.SetKeepAlivesEnabled(f.EnableKeepAlive)

	// Start the HTTP server.
	if err := f.tlsSrv.ServeTLS(l, "", ""); err != http.ErrServerClosed {
		return err
	}
	return nil
}

// announceRoute parses and responds to an Announce.
func (f *Frontend) announceRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	var err error
	var start time.Time
	if f.EnableRequestTiming {
		start = time.Now()
	}
	var af *bittorrent.AddressFamily
	defer func() {
		if f.EnableRequestTiming {
			recordResponseDuration("announce", af, err, time.Since(start))
		} else {
			recordResponseDuration("announce", af, err, time.Duration(0))
		}
	}()

	req, err := ParseAnnounce(r, f.ParseOptions)
	if err != nil {
		WriteError(w, err)
		return
	}
	af = new(bittorrent.AddressFamily)
	*af = req.IP.AddressFamily

	ctx, resp, err := f.logic.HandleAnnounce(context.Background(), req)
	if err != nil {
		WriteError(w, err)
		return
	}

	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	err = WriteAnnounceResponse(w, resp)
	if err != nil {
		WriteError(w, err)
		return
	}

	go f.logic.AfterAnnounce(ctx, req, resp)
}

// scrapeRoute parses and responds to a Scrape.
func (f *Frontend) scrapeRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	var err error
	var start time.Time
	if f.EnableRequestTiming {
		start = time.Now()
	}
	var af *bittorrent.AddressFamily
	defer func() {
		if f.EnableRequestTiming {
			recordResponseDuration("scrape", af, err, time.Since(start))
		} else {
			recordResponseDuration("scrape", af, err, time.Duration(0))
		}
	}()

	req, err := ParseScrape(r, f.ParseOptions)
	if err != nil {
		WriteError(w, err)
		return
	}

	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		log.Error("http: unable to determine remote address for scrape", log.Err(err))
		WriteError(w, err)
		return
	}

	reqIP := net.ParseIP(host)
	if reqIP.To4() != nil {
		req.AddressFamily = bittorrent.IPv4
	} else if len(reqIP) == net.IPv6len { // implies reqIP.To4() == nil
		req.AddressFamily = bittorrent.IPv6
	} else {
		log.Error("http: invalid IP: neither v4 nor v6", log.Fields{"RemoteAddr": r.RemoteAddr})
		WriteError(w, bittorrent.ErrInvalidIP)
		return
	}
	af = new(bittorrent.AddressFamily)
	*af = req.AddressFamily

	ctx, resp, err := f.logic.HandleScrape(context.Background(), req)
	if err != nil {
		WriteError(w, err)
		return
	}

	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	err = WriteScrapeResponse(w, resp)
	if err != nil {
		WriteError(w, err)
		return
	}

	go f.logic.AfterScrape(ctx, req, resp)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"reflect"
	"testing"
)

var runTests = []struct {
	args     []string
	in       string
	expected string
}{
	{
		args:     []string{"echo", "Hello,", "world!"},
		in:       ``,
		expected: `echo Hello, world\!`,
	},
	{
		args: []string{"echo", "{{foo}},", "{{bar}}"},
		in: `Hello test
world test!
`,
		expected: `echo 'Hello test,' 'world test!'`,
	},
	{
		args: []string{"echo", "{{foo}},", "{{bar}},", "{{foo}}-{{bar}}-{{baz}}"},
		in: `Hello
wonderful
world!
`,
		expected: `echo Hello, wonderful, Hello-wonderful-world\!`,
	},
}

func TestRun(t *testing.T) {
	path := "./.test/run.json"
	for _, test := range runTests {
		in := bufio.NewReader(bytes.NewBufferString(test.in))
		cmd, err := Run(path, test.args, in)
		if err != nil {
			t.Errorf("error occurred unexpectedly: %+v", err)
		}
		if !reflect.DeepEqual(cmd, test.expected) {
			t.Errorf("command not correct (expected: %+v, got: %+v)", test.expected, cmd)
		}
	}
}
<commit_msg>add test for concurrent execution to check the atomic write<commit_after>package main

import (
	"bufio"
	"bytes"
	"reflect"
	"sync"
	"testing"
)

var runTests = []struct {
	args     []string
	in       string
	expected string
}{
	{
		args:     []string{"echo", "Hello,", "world!"},
		in:       ``,
		expected: `echo Hello, world\!`,
	},
	{
		args: []string{"echo", "{{foo}},", "{{bar}}"},
		in: `Hello test
world test!
`,
		expected: `echo 'Hello test,' 'world test!'`,
	},
	{
		args: []string{"echo", "{{foo}},", "{{bar}},", "{{foo}}-{{bar}}-{{baz}}"},
		in: `Hello
wonderful
world!
`,
		expected: `echo Hello, wonderful, Hello-wonderful-world\!`,
	},
}

func TestRun(t *testing.T) {
	path := "./.test/run.json"
	for _, test := range runTests {
		in := bufio.NewReader(bytes.NewBufferString(test.in))
		cmd, err := Run(path, test.args, in)
		if err != nil {
			t.Errorf("error occurred unexpectedly: %+v", err)
		}
		if !reflect.DeepEqual(cmd, test.expected) {
			t.Errorf("command not correct (expected: %+v, got: %+v)", test.expected, cmd)
		}
	}
}

func TestRun_concurrently(t *testing.T) {
	path := "./.test/concurrently.json"
	test := runTests[1]
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		go func() {
			wg.Add(1)
			defer wg.Done()
			in := bufio.NewReader(bytes.NewBufferString(test.in))
			cmd, err := Run(path, test.args, in)
			if err != nil {
				t.Errorf("error occurred unexpectedly: %+v", err)
			}
			if !reflect.DeepEqual(cmd, test.expected) {
				t.Errorf("command not correct (expected: %+v, got: %+v)", test.expected, cmd)
			}
		}()
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package config

import (
	"net/http"

	"github.com/Cloud-Foundations/golib/pkg/log"
)

type BaseConfig struct {
	Gen_Cert_URLS string `yaml:"gen_cert_urls"`
	Username      string `yaml:"username"`
	FilePrefix    string `yaml:"file_prefix"`
	AddGroups     bool   `yaml:"add_groups"`
}

// AppConfigFile represents a keymaster client configuration file
type AppConfigFile struct {
	Base BaseConfig
}

// LoadVerifyConfigFile reads, verifies, and returns the contents of
// a keymaster configuration file. LoadVerifyConfigFile returns an error if the
// configuration file is invalid.
func LoadVerifyConfigFile(configFilename string) (AppConfigFile, error) {
	return loadVerifyConfigFile(configFilename)
}

// GetConfigFromHost grabs a default config file from a given host and stores
// it in the local file system.
func GetConfigFromHost(
	configFilename string,
	hostname string,
	client *http.Client,
	logger log.Logger) error {
	return getConfigFromHost(configFilename, hostname, client, logger)
}
<commit_msg>Add WebauthBrowser to lib/client/config.BaseConfig.<commit_after>package config

import (
	"net/http"

	"github.com/Cloud-Foundations/golib/pkg/log"
)

type BaseConfig struct {
	Gen_Cert_URLS  string `yaml:"gen_cert_urls"`
	Username       string `yaml:"username"`
	FilePrefix     string `yaml:"file_prefix"`
	AddGroups      bool   `yaml:"add_groups"`
	WebauthBrowser string `yaml:"webauth_browser"`
}

// AppConfigFile represents a keymaster client configuration file
type AppConfigFile struct {
	Base BaseConfig
}

// LoadVerifyConfigFile reads, verifies, and returns the contents of
// a keymaster configuration file. LoadVerifyConfigFile returns an error if the
// configuration file is invalid.
func LoadVerifyConfigFile(configFilename string) (AppConfigFile, error) {
	return loadVerifyConfigFile(configFilename)
}

// GetConfigFromHost grabs a default config file from a given host and stores
// it in the local file system.
func GetConfigFromHost(
	configFilename string,
	hostname string,
	client *http.Client,
	logger log.Logger) error {
	return getConfigFromHost(configFilename, hostname, client, logger)
}
<|endoftext|>


<commit_before>// Copyright 2014 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package gin

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"runtime"
	"strings"
	"time"
)

var (
	dunno     = []byte("???")
	centerDot = []byte("·")
	dot       = []byte(".")
	slash     = []byte("/")
)

// RecoveryFunc defines the function passable to CustomRecovery.
type RecoveryFunc func(c *Context, err interface{})

// Recovery returns a middleware that recovers from any panics and writes a 500 if there was one.
func Recovery() HandlerFunc {
	return RecoveryWithWriter(DefaultErrorWriter)
}

// CustomRecovery returns a middleware that recovers from any panics and calls the provided handle func to handle it.
func CustomRecovery(handle RecoveryFunc) HandlerFunc {
	return RecoveryWithWriter(DefaultErrorWriter, handle)
}

// RecoveryWithWriter returns a middleware for a given writer that recovers from any panics and writes a 500 if there was one.
func RecoveryWithWriter(out io.Writer, recovery ...RecoveryFunc) HandlerFunc {
	if len(recovery) > 0 {
		return CustomRecoveryWithWriter(out, recovery[0])
	}
	return CustomRecoveryWithWriter(out, defaultHandleRecovery)
}

// CustomRecoveryWithWriter returns a middleware for a given writer that recovers from any panics and calls the provided handle func to handle it.
func CustomRecoveryWithWriter(out io.Writer, handle RecoveryFunc) HandlerFunc {
	var logger *log.Logger
	if out != nil {
		logger = log.New(out, "\n\n\x1b[31m", log.LstdFlags)
	}
	return func(c *Context) {
		defer func() {
			if err := recover(); err != nil {
				// Check for a broken connection, as it is not really a
				// condition that warrants a panic stack trace.
				var brokenPipe bool
				if ne, ok := err.(*net.OpError); ok {
					if se, ok := ne.Err.(*os.SyscallError); ok {
						if strings.Contains(strings.ToLower(se.Error()), "broken pipe") || strings.Contains(strings.ToLower(se.Error()), "connection reset by peer") {
							brokenPipe = true
						}
					}
				}
				if logger != nil {
					stack := stack(3)
					httpRequest, _ := httputil.DumpRequest(c.Request, false)
					headers := strings.Split(string(httpRequest), "\r\n")
					for idx, header := range headers {
						current := strings.Split(header, ":")
						if current[0] == "Authorization" {
							headers[idx] = current[0] + ": *"
						}
					}
					headersToStr := strings.Join(headers, "\r\n")
					if brokenPipe {
						logger.Printf("%s\n%s%s", err, headersToStr, reset)
					} else if IsDebugging() {
						logger.Printf("[Recovery] %s panic recovered:\n%s\n%s\n%s%s",
							timeFormat(time.Now()), headersToStr, err, stack, reset)
					} else {
						logger.Printf("[Recovery] %s panic recovered:\n%s\n%s%s",
							timeFormat(time.Now()), err, stack, reset)
					}
				}
				if brokenPipe {
					// If the connection is dead, we can't write a status to it.
					c.Error(err.(error)) // nolint: errcheck
					c.Abort()
				} else {
					handle(c, err)
				}
			}
		}()
		c.Next()
	}
}

func defaultHandleRecovery(c *Context, err interface{}) {
	c.AbortWithStatus(http.StatusInternalServerError)
}

// stack returns a nicely formatted stack frame, skipping skip frames.
func stack(skip int) []byte {
	buf := new(bytes.Buffer) // the returned data
	// As we loop, we open files and read them. These variables record the currently
	// loaded file.
	var lines [][]byte
	var lastFile string
	for i := skip; ; i++ { // Skip the expected number of frames
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		// Print this much at least.  If we can't find the source, it won't show.
		fmt.Fprintf(buf, "%s:%d (0x%x)\n", file, line, pc)
		if file != lastFile {
			data, err := ioutil.ReadFile(file)
			if err != nil {
				continue
			}
			lines = bytes.Split(data, []byte{'\n'})
			lastFile = file
		}
		fmt.Fprintf(buf, "\t%s: %s\n", function(pc), source(lines, line))
	}
	return buf.Bytes()
}

// source returns a space-trimmed slice of the n'th line.
func source(lines [][]byte, n int) []byte {
	n-- // in stack trace, lines are 1-indexed but our array is 0-indexed
	if n < 0 || n >= len(lines) {
		return dunno
	}
	return bytes.TrimSpace(lines[n])
}

// function returns, if possible, the name of the function containing the PC.
func function(pc uintptr) []byte {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return dunno
	}
	name := []byte(fn.Name())
	// The name includes the path name to the package, which is unnecessary
	// since the file name is already included.  Plus, it has center dots.
	// That is, we see
	//	runtime/debug.*T·ptrmethod
	// and want
	//	*T.ptrmethod
	// Also the package path might contains dot (e.g. code.google.com/...),
	// so first eliminate the path prefix
	if lastSlash := bytes.LastIndex(name, slash); lastSlash >= 0 {
		name = name[lastSlash+1:]
	}
	if period := bytes.Index(name, dot); period >= 0 {
		name = name[period+1:]
	}
	name = bytes.Replace(name, centerDot, dot, -1)
	return name
}

// timeFormat returns a customized time string for logger.
func timeFormat(t time.Time) string {
	timeString := t.Format("2006/01/02 - 15:04:05")
	return timeString
}
<commit_msg>skip unnecessary variable assignment in timeFormat (#2761)<commit_after>// Copyright 2014 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package gin

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"runtime"
	"strings"
	"time"
)

var (
	dunno     = []byte("???")
	centerDot = []byte("·")
	dot       = []byte(".")
	slash     = []byte("/")
)

// RecoveryFunc defines the function passable to CustomRecovery.
type RecoveryFunc func(c *Context, err interface{})

// Recovery returns a middleware that recovers from any panics and writes a 500 if there was one.
func Recovery() HandlerFunc {
	return RecoveryWithWriter(DefaultErrorWriter)
}

// CustomRecovery returns a middleware that recovers from any panics and calls the provided handle func to handle it.
func CustomRecovery(handle RecoveryFunc) HandlerFunc {
	return RecoveryWithWriter(DefaultErrorWriter, handle)
}

// RecoveryWithWriter returns a middleware for a given writer that recovers from any panics and writes a 500 if there was one.
func RecoveryWithWriter(out io.Writer, recovery ...RecoveryFunc) HandlerFunc {
	if len(recovery) > 0 {
		return CustomRecoveryWithWriter(out, recovery[0])
	}
	return CustomRecoveryWithWriter(out, defaultHandleRecovery)
}

// CustomRecoveryWithWriter returns a middleware for a given writer that recovers from any panics and calls the provided handle func to handle it.
func CustomRecoveryWithWriter(out io.Writer, handle RecoveryFunc) HandlerFunc {
	var logger *log.Logger
	if out != nil {
		logger = log.New(out, "\n\n\x1b[31m", log.LstdFlags)
	}
	return func(c *Context) {
		defer func() {
			if err := recover(); err != nil {
				// Check for a broken connection, as it is not really a
				// condition that warrants a panic stack trace.
				var brokenPipe bool
				if ne, ok := err.(*net.OpError); ok {
					if se, ok := ne.Err.(*os.SyscallError); ok {
						if strings.Contains(strings.ToLower(se.Error()), "broken pipe") || strings.Contains(strings.ToLower(se.Error()), "connection reset by peer") {
							brokenPipe = true
						}
					}
				}
				if logger != nil {
					stack := stack(3)
					httpRequest, _ := httputil.DumpRequest(c.Request, false)
					headers := strings.Split(string(httpRequest), "\r\n")
					for idx, header := range headers {
						current := strings.Split(header, ":")
						if current[0] == "Authorization" {
							headers[idx] = current[0] + ": *"
						}
					}
					headersToStr := strings.Join(headers, "\r\n")
					if brokenPipe {
						logger.Printf("%s\n%s%s", err, headersToStr, reset)
					} else if IsDebugging() {
						logger.Printf("[Recovery] %s panic recovered:\n%s\n%s\n%s%s",
							timeFormat(time.Now()), headersToStr, err, stack, reset)
					} else {
						logger.Printf("[Recovery] %s panic recovered:\n%s\n%s%s",
							timeFormat(time.Now()), err, stack, reset)
					}
				}
				if brokenPipe {
					// If the connection is dead, we can't write a status to it.
					c.Error(err.(error)) // nolint: errcheck
					c.Abort()
				} else {
					handle(c, err)
				}
			}
		}()
		c.Next()
	}
}

func defaultHandleRecovery(c *Context, err interface{}) {
	c.AbortWithStatus(http.StatusInternalServerError)
}

// stack returns a nicely formatted stack frame, skipping skip frames.
func stack(skip int) []byte {
	buf := new(bytes.Buffer) // the returned data
	// As we loop, we open files and read them. These variables record the currently
	// loaded file.
	var lines [][]byte
	var lastFile string
	for i := skip; ; i++ { // Skip the expected number of frames
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		// Print this much at least.  If we can't find the source, it won't show.
		fmt.Fprintf(buf, "%s:%d (0x%x)\n", file, line, pc)
		if file != lastFile {
			data, err := ioutil.ReadFile(file)
			if err != nil {
				continue
			}
			lines = bytes.Split(data, []byte{'\n'})
			lastFile = file
		}
		fmt.Fprintf(buf, "\t%s: %s\n", function(pc), source(lines, line))
	}
	return buf.Bytes()
}

// source returns a space-trimmed slice of the n'th line.
func source(lines [][]byte, n int) []byte {
	n-- // in stack trace, lines are 1-indexed but our array is 0-indexed
	if n < 0 || n >= len(lines) {
		return dunno
	}
	return bytes.TrimSpace(lines[n])
}

// function returns, if possible, the name of the function containing the PC.
func function(pc uintptr) []byte {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return dunno
	}
	name := []byte(fn.Name())
	// The name includes the path name to the package, which is unnecessary
	// since the file name is already included.  Plus, it has center dots.
	// That is, we see
	//	runtime/debug.*T·ptrmethod
	// and want
	//	*T.ptrmethod
	// Also the package path might contains dot (e.g. code.google.com/...),
	// so first eliminate the path prefix
	if lastSlash := bytes.LastIndex(name, slash); lastSlash >= 0 {
		name = name[lastSlash+1:]
	}
	if period := bytes.Index(name, dot); period >= 0 {
		name = name[period+1:]
	}
	name = bytes.Replace(name, centerDot, dot, -1)
	return name
}

// timeFormat returns a customized time string for logger.
func timeFormat(t time.Time) string {
	return t.Format("2006/01/02 - 15:04:05")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/getgauge/common"
	"github.com/getgauge/gauge/config"
	"github.com/getgauge/gauge/gauge_messages"
	"github.com/golang/protobuf/proto"
	"os"
	"strings"
)

type rephraseRefactorer struct {
	oldStep   *step
	newStep   *step
	isConcept bool
}

type refactoringResult struct {
	success            bool
	specsChanged       []string
	conceptsChanged    []string
	runnerFilesChanged []string
	errors             []string
	warnings           []string
}

func performRephraseRefactoring(oldStep, newStep string) *refactoringResult {
	if newStep == oldStep {
		return rephraseFailure("Same old step name and new step name.")
	}
	agent, err := getRefactorAgent(oldStep, newStep)

	if err != nil {
		return rephraseFailure(err.Error())
	}

	projectRoot, err := common.GetProjectRoot()
	if err != nil {
		return rephraseFailure(err.Error())
	}

	result := &refactoringResult{success: true, errors: make([]string, 0), warnings: make([]string, 0)}
	specs, specParseResults := findSpecs(projectRoot, &conceptDictionary{})
	addErrorsAndWarningsToRefactoringResult(result, specParseResults...)
	if !result.success {
		return result
	}
	conceptDictionary, parseResult := createConceptsDictionary(false)

	addErrorsAndWarningsToRefactoringResult(result, parseResult)
	if !result.success {
		return result
	}

	refactorResult := agent.performRefactoringOn(specs, conceptDictionary)
	refactorResult.warnings = append(refactorResult.warnings, result.warnings...)
	return refactorResult
}

func rephraseFailure(errors ...string) *refactoringResult {
	return &refactoringResult{success: false, errors: errors}
}

func addErrorsAndWarningsToRefactoringResult(refactorResult *refactoringResult, parseResults ...*parseResult) {
	for _, parseResult := range parseResults {
		if !parseResult.ok {
			refactorResult.success = false
			refactorResult.errors = append(refactorResult.errors, parseResult.error.Error())
		}
		refactorResult.appendWarnings(parseResult.warnings)
	}
}

func (agent *rephraseRefactorer) performRefactoringOn(specs []*specification, conceptDictionary *conceptDictionary) *refactoringResult {
	runner := agent.startRunner()
	defer runner.kill()
	err, stepName, isStepPresent := agent.getStepNameFromRunner(runner)
	if err != nil {
		return rephraseFailure(fmt.Sprintf("Failed to perform refactoring: %s", err))
	}
	specsRefactored, conceptFilesRefactored := agent.rephraseInSpecsAndConcepts(&specs, conceptDictionary)
	specFiles, conceptFiles := writeToConceptAndSpecFiles(specs, conceptDictionary, specsRefactored, conceptFilesRefactored)
	refactoringResult := &refactoringResult{specsChanged: specFiles, success:true, conceptsChanged: conceptFiles, errors: make([]string, 0)}
	if isStepPresent {
		filesChanged, err := agent.requestRunnerForRefactoring(runner, stepName)
		refactoringResult.runnerFilesChanged = filesChanged
		if err != nil {
			refactoringResult.errors = append(refactoringResult.errors, fmt.Sprintf("Only spec files and concepts refactored: %s", err))
			refactoringResult.success = false
		}
	}
	return refactoringResult
}

func (agent *rephraseRefactorer) rephraseInSpecsAndConcepts(specs *[]*specification, conceptDictionary *conceptDictionary) (map[*specification]bool, map[string]bool) {
	specsRefactored := make(map[*specification]bool, 0)
	conceptFilesRefactored := make(map[string]bool, 0)
	orderMap := agent.createOrderOfArgs()
	for _, spec := range *specs {
		specsRefactored[spec] = spec.renameSteps(*agent.oldStep, *agent.newStep, orderMap)
	}
	isConcept := false
	for _, concept := range conceptDictionary.conceptsMap {
		_, ok := conceptFilesRefactored[concept.fileName]
		conceptFilesRefactored[concept.fileName] = !ok && false || conceptFilesRefactored[concept.fileName]
		for _, item := range concept.conceptStep.items {
			isRefactored := conceptFilesRefactored[concept.fileName]
			conceptFilesRefactored[concept.fileName] = item.kind() == stepKind &&
				item.(*step).rename(*agent.oldStep, *agent.newStep, isRefactored, orderMap, &isConcept) ||
				isRefactored
		}
	}
	agent.isConcept = isConcept
	return specsRefactored, conceptFilesRefactored
}

func (agent *rephraseRefactorer) createOrderOfArgs() map[int]int {
	orderMap := make(map[int]int, len(agent.newStep.args))
	for i, arg := range agent.newStep.args {
		orderMap[i] = SliceIndex(len(agent.oldStep.args), func(i int) bool { return agent.oldStep.args[i].String() == arg.String() })
	}
	return orderMap
}

func SliceIndex(limit int, predicate func(i int) bool) int {
	for i := 0; i < limit; i++ {
		if predicate(i) {
			return i
		}
	}
	return -1
}

func getRefactorAgent(oldStepText, newStepText string) (*rephraseRefactorer, error) {
	parser := new(specParser)
	stepTokens, err := parser.generateTokens("* " + oldStepText + "\n" + "*" + newStepText)
	if err != nil {
		return nil, err
	}
	spec := &specification{}
	steps := make([]*step, 0)
	for _, stepToken := range stepTokens {
		step, err := spec.createStepUsingLookup(stepToken, nil)
		if err != nil {
			return nil, err
		}
		steps = append(steps, step)
	}
	return &rephraseRefactorer{oldStep: steps[0], newStep: steps[1]}, nil
}

func (agent *rephraseRefactorer) requestRunnerForRefactoring(testRunner *testRunner, stepName string) ([]string, error) {
	refactorRequest, err := agent.createRefactorRequest(testRunner, stepName)
	if err != nil {
		return nil, err
	}
	refactorResponse := agent.sendRefactorRequest(testRunner, refactorRequest)
	var runnerError error
	if !refactorResponse.GetSuccess() {
		runnerError = errors.New(refactorResponse.GetError())
	}
	return refactorResponse.GetFilesChanged(), runnerError
}

func (agent *rephraseRefactorer) startRunner() *testRunner {
	loadGaugeEnvironment()
	startAPIService(0)
	testRunner, err := startRunnerAndMakeConnection(getProjectManifest())
	if err != nil {
		fmt.Printf("Failed to connect to test runner: %s", err)
		os.Exit(1)
	}
	return testRunner
}

func (agent *rephraseRefactorer) sendRefactorRequest(testRunner *testRunner, refactorRequest *gauge_messages.Message) *gauge_messages.RefactorResponse {
	response, err := getResponseForMessageWithTimeout(refactorRequest, testRunner.connection, config.RefactorTimeout())
	if err != nil {
		return &gauge_messages.RefactorResponse{Success: proto.Bool(false), Error: proto.String(err.Error())}
	}
	return response.GetRefactorResponse()
}

//Todo: Check for inline tables
func (agent *rephraseRefactorer) createRefactorRequest(runner *testRunner, stepName string) (*gauge_messages.Message, error) {
	oldStepValue, err := agent.getStepValueFor(agent.oldStep, stepName)
	if err != nil {
		return nil, err
	}
	orderMap := agent.createOrderOfArgs()
	newStepName := agent.generateNewStepName(oldStepValue.args, orderMap)
	newStepValue, err := extractStepValueAndParams(newStepName, false)
	if err != nil {
		return nil, err
	}
	oldProtoStepValue := convertToProtoStepValue(oldStepValue)
	newProtoStepValue := convertToProtoStepValue(newStepValue)
	return &gauge_messages.Message{MessageType: gauge_messages.Message_RefactorRequest.Enum(), RefactorRequest: &gauge_messages.RefactorRequest{OldStepValue: oldProtoStepValue, NewStepValue: newProtoStepValue, ParamPositions: agent.createParameterPositions(orderMap)}}, nil
}

func (agent *rephraseRefactorer) generateNewStepName(args []string, orderMap map[int]int) string {
	agent.newStep.populateFragments()
	paramIndex := 0
	for _, fragment := range agent.newStep.fragments {
		if fragment.GetFragmentType() == gauge_messages.Fragment_Parameter {
			if orderMap[paramIndex] != -1 {
				fragment.GetParameter().Value = proto.String(args[orderMap[paramIndex]])
			}
			paramIndex++
		}
	}
	return convertToStepText(agent.newStep.fragments)
}

func (agent *rephraseRefactorer) getStepNameFromRunner(runner *testRunner) (error, string, bool) {
	stepNameMessage := &gauge_messages.Message{MessageType: gauge_messages.Message_StepNameRequest.Enum(), StepNameRequest: &gauge_messages.StepNameRequest{StepValue: proto.String(agent.oldStep.value)}}
	responseMessage, err := getResponseForMessageWithTimeout(stepNameMessage, runner.connection, config.RunnerAPIRequestTimeout())
	if err != nil {
		return err, "", false
	}
	if !(responseMessage.GetStepNameResponse().GetIsStepPresent()) {
		fmt.Println("Step implementation not found: " + agent.oldStep.lineText)
		return nil, "", false
	}
	if responseMessage.GetStepNameResponse().GetHasAlias() {
		return errors.New(fmt.Sprintf("steps with aliases : '%s' cannot be refactored.", strings.Join(responseMessage.GetStepNameResponse().GetStepName(), "', '"))), "", false
	}

	return nil, responseMessage.GetStepNameResponse().GetStepName()[0], true
}

func (agent *rephraseRefactorer) createParameterPositions(orderMap map[int]int) []*gauge_messages.ParameterPosition {
	paramPositions := make([]*gauge_messages.ParameterPosition, 0)
	for k, v := range orderMap {
		paramPositions = append(paramPositions, &gauge_messages.ParameterPosition{NewPosition: proto.Int(k), OldPosition: proto.Int(v)})
	}
	return paramPositions
}

func (agent *rephraseRefactorer) getStepValueFor(step *step, stepName string) (*stepValue, error) {
	return extractStepValueAndParams(stepName, false)
}

func writeToConceptAndSpecFiles(specs []*specification, conceptDictionary *conceptDictionary, specsRefactored map[*specification]bool, conceptFilesRefactored map[string]bool) ([]string, []string) {
	specFiles := make([]string, 0)
	conceptFiles := make([]string, 0)
	for _, spec := range specs {
		if specsRefactored[spec] {
			specFiles = append(specFiles, spec.fileName)
			formatted := formatSpecification(spec)
			saveFile(spec.fileName, formatted, true)
		}
	}
	conceptMap := formatConcepts(conceptDictionary)
	for fileName, concept := range conceptMap {
		if conceptFilesRefactored[fileName] {
			conceptFiles = append(conceptFiles, fileName)
			saveFile(fileName, concept, true)
		}
	}
	return specFiles, conceptFiles
}

func (refactoringResult *refactoringResult) appendWarnings(warnings []*warning) {
	if refactoringResult.warnings == nil {
		refactoringResult.warnings = make([]string, 0)
	}
	for _, warning := range warnings {
		refactoringResult.warnings = append(refactoringResult.warnings, warning.message)
	}
}

func (refactoringResult *refactoringResult) allFilesChanges() []string {
	filesChanged := make([]string, 0)
	filesChanged = append(filesChanged, refactoringResult.specsChanged...)
	filesChanged = append(filesChanged, refactoringResult.conceptsChanged...)
	filesChanged = append(filesChanged, refactoringResult.runnerFilesChanged...)
	return filesChanged

}
<commit_msg>Fixing error messaging on perfroming refactoring<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/getgauge/common"
	"github.com/getgauge/gauge/config"
	"github.com/getgauge/gauge/gauge_messages"
	"github.com/golang/protobuf/proto"
	"os"
	"strings"
)

type rephraseRefactorer struct {
	oldStep   *step
	newStep   *step
	isConcept bool
}

type refactoringResult struct {
	success            bool
	specsChanged       []string
	conceptsChanged    []string
	runnerFilesChanged []string
	errors             []string
	warnings           []string
}

func performRephraseRefactoring(oldStep, newStep string) *refactoringResult {
	if newStep == oldStep {
		return rephraseFailure("Same old step name and new step name.")
	}
	agent, err := getRefactorAgent(oldStep, newStep)

	if err != nil {
		return rephraseFailure(err.Error())
	}

	projectRoot, err := common.GetProjectRoot()
	if err != nil {
		return rephraseFailure(err.Error())
	}

	result := &refactoringResult{success: true, errors: make([]string, 0), warnings: make([]string, 0)}
	specs, specParseResults := findSpecs(projectRoot, &conceptDictionary{})
	addErrorsAndWarningsToRefactoringResult(result, specParseResults...)
	if !result.success {
		return result
	}
	conceptDictionary, parseResult := createConceptsDictionary(false)

	addErrorsAndWarningsToRefactoringResult(result, parseResult)
	if !result.success {
		return result
	}

	refactorResult := agent.performRefactoringOn(specs, conceptDictionary)
	refactorResult.warnings = append(refactorResult.warnings, result.warnings...)
	return refactorResult
}

func rephraseFailure(errors ...string) *refactoringResult {
	return &refactoringResult{success: false, errors: errors}
}

func addErrorsAndWarningsToRefactoringResult(refactorResult *refactoringResult, parseResults ...*parseResult) {
	for _, parseResult := range parseResults {
		if !parseResult.ok {
			refactorResult.success = false
			refactorResult.errors = append(refactorResult.errors, parseResult.error.Error())
		}
		refactorResult.appendWarnings(parseResult.warnings)
	}
}

func (agent *rephraseRefactorer) performRefactoringOn(specs []*specification, conceptDictionary *conceptDictionary) *refactoringResult {
	specsRefactored, conceptFilesRefactored := agent.rephraseInSpecsAndConcepts(&specs, conceptDictionary)
	specFiles, conceptFiles := writeToConceptAndSpecFiles(specs, conceptDictionary, specsRefactored, conceptFilesRefactored)
	refactoringResult := &refactoringResult{specsChanged: specFiles, success: false, conceptsChanged: conceptFiles, errors: make([]string, 0)}

	runner := agent.startRunner()
	defer runner.kill()
	stepName, err := agent.getStepNameFromRunner(runner)
	if err != nil {
		refactoringResult.errors = append(refactoringResult.errors, err.Error())
		return refactoringResult
	}
	runnerFilesChanged, err := agent.requestRunnerForRefactoring(runner, stepName)
	if err != nil {
		refactoringResult.errors = append(refactoringResult.errors, fmt.Sprintf("Only spec files and concepts refactored: %s", err))
		return refactoringResult
	}
	refactoringResult.success = true
	refactoringResult.runnerFilesChanged = runnerFilesChanged
	return refactoringResult
}

func (agent *rephraseRefactorer) rephraseInSpecsAndConcepts(specs *[]*specification, conceptDictionary *conceptDictionary) (map[*specification]bool, map[string]bool) {
	specsRefactored := make(map[*specification]bool, 0)
	conceptFilesRefactored := make(map[string]bool, 0)
	orderMap := agent.createOrderOfArgs()
	for _, spec := range *specs {
		specsRefactored[spec] = spec.renameSteps(*agent.oldStep, *agent.newStep, orderMap)
	}
	isConcept := false
	for _, concept := range conceptDictionary.conceptsMap {
		_, ok := conceptFilesRefactored[concept.fileName]
		conceptFilesRefactored[concept.fileName] = !ok && false || conceptFilesRefactored[concept.fileName]
		for _, item := range concept.conceptStep.items {
			isRefactored := conceptFilesRefactored[concept.fileName]
			conceptFilesRefactored[concept.fileName] = item.kind() == stepKind &&
				item.(*step).rename(*agent.oldStep, *agent.newStep, isRefactored, orderMap, &isConcept) ||
				isRefactored
		}
	}
	agent.isConcept = isConcept
	return specsRefactored, conceptFilesRefactored
}

func (agent *rephraseRefactorer) createOrderOfArgs() map[int]int {
	orderMap := make(map[int]int, len(agent.newStep.args))
	for i, arg := range agent.newStep.args {
		orderMap[i] = SliceIndex(len(agent.oldStep.args), func(i int) bool { return agent.oldStep.args[i].String() == arg.String() })
	}
	return orderMap
}

func SliceIndex(limit int, predicate func(i int) bool) int {
	for i := 0; i < limit; i++ {
		if predicate(i) {
			return i
		}
	}
	return -1
}

func getRefactorAgent(oldStepText, newStepText string) (*rephraseRefactorer, error) {
	parser := new(specParser)
	stepTokens, err := parser.generateTokens("* " + oldStepText + "\n" + "*" + newStepText)
	if err != nil {
		return nil, err
	}
	spec := &specification{}
	steps := make([]*step, 0)
	for _, stepToken := range stepTokens {
		step, err := spec.createStepUsingLookup(stepToken, nil)
		if err != nil {
			return nil, err
		}
		steps = append(steps, step)
	}
	return &rephraseRefactorer{oldStep: steps[0], newStep: steps[1]}, nil
}

func (agent *rephraseRefactorer) requestRunnerForRefactoring(testRunner *testRunner, stepName string) ([]string, error) {
	refactorRequest, err := agent.createRefactorRequest(testRunner, stepName)
	if err != nil {
		return nil, err
	}
	refactorResponse := agent.sendRefactorRequest(testRunner, refactorRequest)
	var runnerError error
	if !refactorResponse.GetSuccess() {
		runnerError = errors.New(refactorResponse.GetError())
	}
	return refactorResponse.GetFilesChanged(), runnerError
}

func (agent *rephraseRefactorer) startRunner() *testRunner {
	loadGaugeEnvironment()
	startAPIService(0)
	testRunner, err := startRunnerAndMakeConnection(getProjectManifest())
	if err != nil {
		fmt.Printf("Failed to connect to test runner: %s", err)
		os.Exit(1)
	}
	return testRunner
}

func (agent *rephraseRefactorer) sendRefactorRequest(testRunner *testRunner, refactorRequest *gauge_messages.Message) *gauge_messages.RefactorResponse {
	response, err := getResponseForMessageWithTimeout(refactorRequest, testRunner.connection, config.RefactorTimeout())
	if err != nil {
		return &gauge_messages.RefactorResponse{Success: proto.Bool(false), Error: proto.String(err.Error())}
	}
	return response.GetRefactorResponse()
}

//Todo: Check for inline tables
func (agent *rephraseRefactorer) createRefactorRequest(runner *testRunner, stepName string) (*gauge_messages.Message, error) {
	oldStepValue, err := agent.getStepValueFor(agent.oldStep, stepName)
	if err != nil {
		return nil, err
	}
	orderMap := agent.createOrderOfArgs()
	newStepName := agent.generateNewStepName(oldStepValue.args, orderMap)
	newStepValue, err := extractStepValueAndParams(newStepName, false)
	if err != nil {
		return nil, err
	}
	oldProtoStepValue := convertToProtoStepValue(oldStepValue)
	newProtoStepValue := convertToProtoStepValue(newStepValue)
	return &gauge_messages.Message{MessageType: gauge_messages.Message_RefactorRequest.Enum(), RefactorRequest: &gauge_messages.RefactorRequest{OldStepValue: oldProtoStepValue, NewStepValue: newProtoStepValue, ParamPositions: agent.createParameterPositions(orderMap)}}, nil
}

func (agent *rephraseRefactorer) generateNewStepName(args []string, orderMap map[int]int) string {
	agent.newStep.populateFragments()
	paramIndex := 0
	for _, fragment := range agent.newStep.fragments {
		if fragment.GetFragmentType() == gauge_messages.Fragment_Parameter {
			if orderMap[paramIndex] != -1 {
				fragment.GetParameter().Value = proto.String(args[orderMap[paramIndex]])
			}
			paramIndex++
		}
	}
	return convertToStepText(agent.newStep.fragments)
}

func (agent *rephraseRefactorer) getStepNameFromRunner(runner *testRunner) (string, error) {
	stepNameMessage := &gauge_messages.Message{MessageType: gauge_messages.Message_StepNameRequest.Enum(), StepNameRequest: &gauge_messages.StepNameRequest{StepValue: proto.String(agent.oldStep.value)}}
	responseMessage, err := getResponseForMessageWithTimeout(stepNameMessage, runner.connection, config.RunnerAPIRequestTimeout())
	if err != nil {
		return "", err
	}
	if !(responseMessage.GetStepNameResponse().GetIsStepPresent()) {
		return "", errors.New(fmt.Sprintf("Step implementation not found: %s", agent.oldStep.lineText))
	}
	if responseMessage.GetStepNameResponse().GetHasAlias() {
		return "", errors.New(fmt.Sprintf("steps with aliases : '%s' cannot be refactored.", strings.Join(responseMessage.GetStepNameResponse().GetStepName(), "', '")))
	}
	return responseMessage.GetStepNameResponse().GetStepName()[0], nil
}

func (agent *rephraseRefactorer) createParameterPositions(orderMap map[int]int) []*gauge_messages.ParameterPosition {
	paramPositions := make([]*gauge_messages.ParameterPosition, 0)
	for k, v := range orderMap {
		paramPositions = append(paramPositions, &gauge_messages.ParameterPosition{NewPosition: proto.Int(k), OldPosition: proto.Int(v)})
	}
	return paramPositions
}

func (agent *rephraseRefactorer) getStepValueFor(step *step, stepName string) (*stepValue, error) {
	return extractStepValueAndParams(stepName, false)
}

func writeToConceptAndSpecFiles(specs []*specification, conceptDictionary *conceptDictionary, specsRefactored map[*specification]bool, conceptFilesRefactored map[string]bool) ([]string, []string) {
	specFiles := make([]string, 0)
	conceptFiles := make([]string, 0)
	for _, spec := range specs {
		if specsRefactored[spec] {
			specFiles = append(specFiles, spec.fileName)
			formatted := formatSpecification(spec)
			saveFile(spec.fileName, formatted, true)
		}
	}
	conceptMap := formatConcepts(conceptDictionary)
	for fileName, concept := range conceptMap {
		if conceptFilesRefactored[fileName] {
			conceptFiles = append(conceptFiles, fileName)
			saveFile(fileName, concept, true)
		}
	}
	return specFiles, conceptFiles
}

func (refactoringResult *refactoringResult) appendWarnings(warnings []*warning) {
	if refactoringResult.warnings == nil {
		refactoringResult.warnings = make([]string, 0)
	}
	for _, warning := range warnings {
		refactoringResult.warnings = append(refactoringResult.warnings, warning.message)
	}
}

func (refactoringResult *refactoringResult) allFilesChanges() []string {
	filesChanged := make([]string, 0)
	filesChanged = append(filesChanged, refactoringResult.specsChanged...)
	filesChanged = append(filesChanged, refactoringResult.conceptsChanged...)
	filesChanged = append(filesChanged, refactoringResult.runnerFilesChanged...)
	return filesChanged

}
<|endoftext|>


<commit_before>package main

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

	"github.com/codegangsta/cli"
	"github.com/coreos/go-etcd/etcd"
	"github.com/dotcloud/docker/api"
	docker "github.com/dotcloud/docker/api/client"
)

const (
	dockerIP = "172.17.42.1"
	ttl      = 60
)

type registerError struct {
	message string
}

func (e *registerError) Error() string {
	return e.message
}

type dockerInspectWriter struct {
	lastCommand []dockerInspectPortMapping
}

type dockerInspectPortMapping struct {
	HostConfig struct {
		PortBindings map[string][]struct {
			HostIp   string
			HostPort string
		}
	}
	State struct {
		Running bool
	}
}

type dockerPortMapping struct {
	ContainerPort string
	Port          string
	Host          string
}

func (dpr dockerInspectPortMapping) portMappingsList() []*dockerPortMapping {
	dockerPortMappings := make([]*dockerPortMapping, 0, len(dpr.HostConfig.PortBindings))

	for ContainerPort, Binding := range dpr.HostConfig.PortBindings {
		currentDockerPortMapping := dockerPortMapping{}

		if len(Binding) == 0 {
			continue
		}

		if pos := strings.Index(ContainerPort, "/"); pos >= 0 {
			currentDockerPortMapping.ContainerPort = ContainerPort[:pos]
		} else {
			currentDockerPortMapping.ContainerPort = ContainerPort
		}

		if pos := strings.Index(Binding[0].HostPort, "/"); pos >= 0 {
			currentDockerPortMapping.Port = Binding[0].HostPort[:pos]
		} else {
			currentDockerPortMapping.Port = Binding[0].HostPort
		}

		if Binding[0].HostIp == "0.0.0.0" {
			currentDockerPortMapping.Host = dockerIP
		} else {
			currentDockerPortMapping.Host = Binding[0].HostIp
		}

		dockerPortMappings = append(dockerPortMappings, &currentDockerPortMapping)
	}

	return dockerPortMappings
}

func (diw *dockerInspectWriter) Write(p []byte) (n int, err error) {
	json.Unmarshal(p, &diw.lastCommand)
	return len(p), nil
}

func parseCliOptions(c *cli.Context) {
	if !c.IsSet("container") {
		fmt.Println("--container argument is required")
		return
	}

	if err := register(c.GlobalString("container")); err != nil {
		fmt.Println(err)
	}
}

func getContainerInfo(container string) (*dockerInspectPortMapping, error) {
	dockerWriter := dockerInspectWriter{}
	dockerClient := docker.NewDockerCli(nil, &dockerWriter, os.Stderr, "unix", api.DEFAULTUNIXSOCKET, nil)
	dockerClient.CmdInspect(container)

	if len(dockerWriter.lastCommand) == 0 {
		return nil, &registerError{message: "Container does not exist"}
	}

	if dockerWriter.lastCommand[0].State.Running == false {
		return nil, &registerError{message: "Container is not running"}
	}

	return &dockerWriter.lastCommand[0], nil
}

func register(container string) error {
	etcdClient := etcd.NewClient([]string{fmt.Sprintf("http://%s:4001", dockerIP)})
	containerInfo, err := getContainerInfo(container)

	if err != nil {
		return err
	}

	if _, err := etcdClient.UpdateDir(fmt.Sprint("containers/", container), ttl); err != nil {
		// If update dir fails is because the directory doesn't exist, so, let's create it
		if _, err := etcdClient.SetDir(fmt.Sprint("containers/", container), ttl); err != nil {
			return err
		}
	}

	for _, dockerPortMapping := range containerInfo.portMappingsList() {
		if _, err := etcdClient.Set(fmt.Sprintf("containers/%s/ports/%s/host/", container, dockerPortMapping.ContainerPort), dockerPortMapping.Host, ttl); err != nil {
			return err
		}

		if _, err := etcdClient.Set(fmt.Sprintf("containers/%s/ports/%s/port/", container, dockerPortMapping.ContainerPort), dockerPortMapping.Port, ttl); err != nil {
			return err
		}
	}

	return nil
}

func deregister(container string) error {
	etcdClient := etcd.NewClient([]string{fmt.Sprintf("http://%s:4001", dockerIP)})

	if _, err := etcdClient.Delete(fmt.Sprint("containers/", container), true); err != nil {
		return err
	}

	return nil
}

func init() {
}

func main() {
	app := cli.NewApp()
	app.Name = "register"
	app.Usage = "Register the ports of a specfied Docker container with Etcd"
	app.Action = parseCliOptions
	app.Version = "0.0.0"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "container", Usage: "The container name or id"},
	}
	app.Run(os.Args)
}
<commit_msg>Changed name of the function which execute the script functionality<commit_after>package main

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

	"github.com/codegangsta/cli"
	"github.com/coreos/go-etcd/etcd"
	"github.com/dotcloud/docker/api"
	docker "github.com/dotcloud/docker/api/client"
)

const (
	dockerIP = "172.17.42.1"
	ttl      = 60
)

type registerError struct {
	message string
}

func (e *registerError) Error() string {
	return e.message
}

type dockerInspectWriter struct {
	lastCommand []dockerInspectPortMapping
}

type dockerInspectPortMapping struct {
	HostConfig struct {
		PortBindings map[string][]struct {
			HostIp   string
			HostPort string
		}
	}
	State struct {
		Running bool
	}
}

type dockerPortMapping struct {
	ContainerPort string
	Port          string
	Host          string
}

func (dpr dockerInspectPortMapping) portMappingsList() []*dockerPortMapping {
	dockerPortMappings := make([]*dockerPortMapping, 0, len(dpr.HostConfig.PortBindings))

	for ContainerPort, Binding := range dpr.HostConfig.PortBindings {
		currentDockerPortMapping := dockerPortMapping{}

		if len(Binding) == 0 {
			continue
		}

		if pos := strings.Index(ContainerPort, "/"); pos >= 0 {
			currentDockerPortMapping.ContainerPort = ContainerPort[:pos]
		} else {
			currentDockerPortMapping.ContainerPort = ContainerPort
		}

		if pos := strings.Index(Binding[0].HostPort, "/"); pos >= 0 {
			currentDockerPortMapping.Port = Binding[0].HostPort[:pos]
		} else {
			currentDockerPortMapping.Port = Binding[0].HostPort
		}

		if Binding[0].HostIp == "0.0.0.0" {
			currentDockerPortMapping.Host = dockerIP
		} else {
			currentDockerPortMapping.Host = Binding[0].HostIp
		}

		dockerPortMappings = append(dockerPortMappings, &currentDockerPortMapping)
	}

	return dockerPortMappings
}

func (diw *dockerInspectWriter) Write(p []byte) (n int, err error) {
	json.Unmarshal(p, &diw.lastCommand)
	return len(p), nil
}

func startRegistration(c *cli.Context) {
	if !c.IsSet("container") {
		fmt.Println("--container argument is required")
		return
	}

	if err := register(c.GlobalString("container")); err != nil {
		fmt.Println(err)
	}
}

func getContainerInfo(container string) (*dockerInspectPortMapping, error) {
	dockerWriter := dockerInspectWriter{}
	dockerClient := docker.NewDockerCli(nil, &dockerWriter, os.Stderr, "unix", api.DEFAULTUNIXSOCKET, nil)
	dockerClient.CmdInspect(container)

	if len(dockerWriter.lastCommand) == 0 {
		return nil, &registerError{message: "Container does not exist"}
	}

	if dockerWriter.lastCommand[0].State.Running == false {
		return nil, &registerError{message: "Container is not running"}
	}

	return &dockerWriter.lastCommand[0], nil
}

func register(container string) error {
	etcdClient := etcd.NewClient([]string{fmt.Sprintf("http://%s:4001", dockerIP)})
	containerInfo, err := getContainerInfo(container)

	if err != nil {
		return err
	}

	if _, err := etcdClient.UpdateDir(fmt.Sprint("containers/", container), ttl); err != nil {
		// If update dir fails is because the directory doesn't exist, so, let's create it
		if _, err := etcdClient.SetDir(fmt.Sprint("containers/", container), ttl); err != nil {
			return err
		}
	}

	for _, dockerPortMapping := range containerInfo.portMappingsList() {
		if _, err := etcdClient.Set(fmt.Sprintf("containers/%s/ports/%s/host/", container, dockerPortMapping.ContainerPort), dockerPortMapping.Host, ttl); err != nil {
			return err
		}

		if _, err := etcdClient.Set(fmt.Sprintf("containers/%s/ports/%s/port/", container, dockerPortMapping.ContainerPort), dockerPortMapping.Port, ttl); err != nil {
			return err
		}
	}

	return nil
}

func deregister(container string) error {
	etcdClient := etcd.NewClient([]string{fmt.Sprintf("http://%s:4001", dockerIP)})

	if _, err := etcdClient.Delete(fmt.Sprint("containers/", container), true); err != nil {
		return err
	}

	return nil
}

func init() {
}

func main() {
	app := cli.NewApp()
	app.Name = "register"
	app.Usage = "Register the ports of a specfied Docker container with Etcd"
	app.Action = startRegistration
	app.Version = "0.0.0"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "container", Usage: "The container name or id"},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package TF2RconWrapper

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/james4k/rcon"
)

// TF2RconConnection represents a rcon connection to a TF2 server
type TF2RconConnection struct {
	rcLock *sync.RWMutex
	rc     *rcon.RemoteConsole

	host     string
	password string
}

var (
	ErrUnknownCommand = errors.New("Unknown Command")
	CVarValueRegex    = regexp.MustCompile(`^"(?:.*?)" = "(.*?)"`)
	//# userid name                uniqueid            connected ping loss state  adr
	rePlayerInfo = regexp.MustCompile(`^#\s+(\d+)\s+"(.+)"\s+(\[U:1:\d+\])\s+\d+:\d+\s+\d+\s+\d+\s+\w+\s+(\d+\.+\d+\.\d+\.\d+:\d+)`)
)

func (c *TF2RconConnection) QueryNoResp(req string) error {
	_, err := c.rc.Write(req)
	return err
}

// Query executes a query and returns the server responses
func (c *TF2RconConnection) Query(req string) (string, error) {
	if c.rc == nil {
		return "", errors.New("Not connected to RCON host")
	}

	c.rcLock.RLock()
	defer c.rcLock.Unlock()

	reqID, reqErr := c.rc.Write(req)
	if reqErr != nil {
		// log.Println(reqErr)
		return "", reqErr
	}

	resp, respID, respErr := c.rc.Read()
	if respErr != nil {
		// log.Println(respErr)
		return "", respErr
	}

	counter := 10
	// retry 10 times
	for {
		if reqID == respID {
			break
		} else if counter < 0 {
			return "", errors.New("Couldn't get a response.")
		} else {
			counter--
			resp, respID, respErr = c.rc.Read()
			if respErr != nil {
				// log.Println(respErr)
				return "", reqErr
			}
		}
	}

	if strings.HasPrefix(resp, "Unknown command") {
		return resp, ErrUnknownCommand
	}

	return resp, nil
}

func (c *TF2RconConnection) GetConVar(cvar string) (string, error) {
	raw, err := c.Query(cvar)

	if err != nil {
		return "", err
	}

	// Querying just a variable's name sends back a message like the
	// following:
	//
	// "cvar_name" = "current value" ( def. "default value" )
	//  var flags like notify replicated
	//  - short description of cvar

	firstLine := strings.Split(raw, "\n")[0]
	matches := CVarValueRegex.FindStringSubmatch(firstLine)
	if len(matches) != 2 {
		return "", errors.New("Unknown cvar.")
	}

	return matches[1], nil
}

func (c *TF2RconConnection) SetConVar(cvar string, val string) (string, error) {
	return c.Query(fmt.Sprintf("%s \"%s\"", cvar, val))
}

// GetPlayers returns a list of players in the server. Includes bots.
func (c *TF2RconConnection) GetPlayers() ([]Player, error) {
	statusString, err := c.Query("status")
	if err != nil {
		return nil, err
	}

	index := strings.Index(statusString, "#")
	i := 0
	for index == -1 {
		statusString, _ = c.Query("status")
		index = strings.Index(statusString, "#")
		i++
		if i == 5 {
			return nil, errors.New("Couldn't get output of status")
		}
	}

	users := strings.Split(statusString[index:], "\n")
	var list []Player
	for _, userString := range users {
		if !rePlayerInfo.MatchString(userString) {
			continue
		}
		matches := rePlayerInfo.FindStringSubmatch(userString)
		player := Player{
			UserID:   matches[1],
			Username: matches[2],
			SteamID:  matches[3],
			Ip:       matches[4],
		}
		list = append(list, player)
	}

	return list, nil
}

// KickPlayer kicks a player
func (c *TF2RconConnection) KickPlayer(p Player, message string) error {
	return c.KickPlayerID(p.UserID, message)
}

// Kicks a player with the given player ID
func (c *TF2RconConnection) KickPlayerID(userID string, message string) error {
	query := fmt.Sprintf("kickid %s %s", userID, message)
	_, err := c.Query(query)
	return err
}

// BanPlayer bans a player
func (c *TF2RconConnection) BanPlayer(minutes int, p Player, message string) error {
	query := "banid " + fmt.Sprintf("%v", minutes) + " " + p.UserID
	if message != "" {
		query += " \"" + message + "\""
	}
	_, err := c.Query(query)
	return err
}

// UnbanPlayer unbans a player
func (c *TF2RconConnection) UnbanPlayer(p Player) error {
	query := "unbanid " + p.UserID
	_, err := c.Query(query)
	return err
}

// Say sends a message to the TF2 server chat
func (c *TF2RconConnection) Say(message string) error {
	query := "say " + message
	_, err := c.Query(query)
	return err
}

func (c *TF2RconConnection) Sayf(format string, a ...interface{}) error {
	err := c.Say(fmt.Sprintf(format, a...))
	return err
}

// ChangeRconPassword changes the rcon password and updates the current connection
// to use the new password
func (c *TF2RconConnection) ChangeRconPassword(password string) error {
	_, err := c.SetConVar("rcon_password", password)

	if err == nil {
		c.rc.Close()
		newConnection, _ := rcon.Dial(c.host, password)
		c.rc = newConnection
	}

	return err
}

// ChangeMap changes the map
func (c *TF2RconConnection) ChangeMap(mapname string) error {
	query := "changelevel \"" + mapname + "\""
	res, err := c.Query(query)
	if res != "" {
		return errors.New("Map not found.")
	}
	return err
}

// ChangeServerPassword changes the server password
func (c *TF2RconConnection) ChangeServerPassword(password string) error {
	_, err := c.SetConVar("sv_password", password)
	return err
}

// GetServerPassword returns the server password
func (c *TF2RconConnection) GetServerPassword() (string, error) {
	return c.GetConVar("sv_password")
}

func (c *TF2RconConnection) AddTag(newTag string) error {
	tags, err := c.GetConVar("sv_tags")
	if err != nil {
		return err
	}

	// Source servers don't auto-remove duplicate tags, and noone
	tagExists := false
	for _, tag := range strings.Split(tags, ",") {
		if tag == newTag {
			tagExists = true
			break
		}
	}

	if !tagExists {
		newTags := strings.Join([]string{tags, newTag}, ",")
		_, err := c.SetConVar("sv_tags", newTags)
		return err
	}

	return nil
}

func (c *TF2RconConnection) RemoveTag(tagName string) error {
	tags, err := c.GetConVar("sv_tags")
	if err != nil {
		return err
	}

	if strings.Contains(tags, tagName) {
		// Replace all instances of the given tagName. This may leave
		// duplicated or trailing commas in the sv_tags string; however
		// Source servers clean up the value of sv_tags to remove those
		// anyways
		_, err := c.SetConVar("sv_tags", strings.Replace(tags, tagName, "", -1))
		return err
	}

	return nil
}

// RedirectLogs send the logaddress_add command
func (c *TF2RconConnection) RedirectLogs(addr string) error {
	query := "logaddress_add " + addr
	_, err := c.Query(query)
	return err
}

func (c *TF2RconConnection) StopLogRedirection(addr string) {
	query := fmt.Sprintf("logaddress_del %s", addr)
	c.QueryNoResp(query)
}

// Close closes the connection
func (c *TF2RconConnection) Close() {
	c.rc.Close()
}

// ExecConfig accepts a string and executes its lines one by one. Assumes
// UNiX line endings
func (c *TF2RconConnection) ExecConfig(config string) error {
	lines := strings.Split(config, "\n")
	for _, line := range lines {
		_, err := c.Query(line)
		if err != nil {
			return err
		}
	}
	return nil
}

// NewTF2RconConnection builds a new TF2RconConnection to a server at address ("ip:port") using
// a rcon_password password
func NewTF2RconConnection(address, password string) (*TF2RconConnection, error) {
	rc, err := rcon.Dial(address, password)
	if err != nil {
		return nil, err
	}

	return &TF2RconConnection{new(sync.RWMutex), rc, address, password}, nil
}

func (c *TF2RconConnection) Reconnect(duration time.Duration) error {
	var err error
	var cur time.Duration

	c.rcLock.Lock()
	defer c.rcLock.Unlock()

	c.Close()
	for cur += time.Second; cur <= duration; {
		c.rc, err = rcon.Dial(c.host, c.password)
		if err == nil {
			return nil
		}
	}

	return err
}
<commit_msg>Remove nil check<commit_after>package TF2RconWrapper

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/james4k/rcon"
)

// TF2RconConnection represents a rcon connection to a TF2 server
type TF2RconConnection struct {
	rcLock *sync.RWMutex
	rc     *rcon.RemoteConsole

	host     string
	password string
}

var (
	ErrUnknownCommand = errors.New("Unknown Command")
	CVarValueRegex    = regexp.MustCompile(`^"(?:.*?)" = "(.*?)"`)
	//# userid name                uniqueid            connected ping loss state  adr
	rePlayerInfo = regexp.MustCompile(`^#\s+(\d+)\s+"(.+)"\s+(\[U:1:\d+\])\s+\d+:\d+\s+\d+\s+\d+\s+\w+\s+(\d+\.+\d+\.\d+\.\d+:\d+)`)
)

func (c *TF2RconConnection) QueryNoResp(req string) error {
	_, err := c.rc.Write(req)
	return err
}

// Query executes a query and returns the server responses
func (c *TF2RconConnection) Query(req string) (string, error) {
	c.rcLock.RLock()
	defer c.rcLock.Unlock()

	reqID, reqErr := c.rc.Write(req)
	if reqErr != nil {
		// log.Println(reqErr)
		return "", reqErr
	}

	resp, respID, respErr := c.rc.Read()
	if respErr != nil {
		// log.Println(respErr)
		return "", respErr
	}

	counter := 10
	// retry 10 times
	for {
		if reqID == respID {
			break
		} else if counter < 0 {
			return "", errors.New("Couldn't get a response.")
		} else {
			counter--
			resp, respID, respErr = c.rc.Read()
			if respErr != nil {
				// log.Println(respErr)
				return "", reqErr
			}
		}
	}

	if strings.HasPrefix(resp, "Unknown command") {
		return resp, ErrUnknownCommand
	}

	return resp, nil
}

func (c *TF2RconConnection) GetConVar(cvar string) (string, error) {
	raw, err := c.Query(cvar)

	if err != nil {
		return "", err
	}

	// Querying just a variable's name sends back a message like the
	// following:
	//
	// "cvar_name" = "current value" ( def. "default value" )
	//  var flags like notify replicated
	//  - short description of cvar

	firstLine := strings.Split(raw, "\n")[0]
	matches := CVarValueRegex.FindStringSubmatch(firstLine)
	if len(matches) != 2 {
		return "", errors.New("Unknown cvar.")
	}

	return matches[1], nil
}

func (c *TF2RconConnection) SetConVar(cvar string, val string) (string, error) {
	return c.Query(fmt.Sprintf("%s \"%s\"", cvar, val))
}

// GetPlayers returns a list of players in the server. Includes bots.
func (c *TF2RconConnection) GetPlayers() ([]Player, error) {
	statusString, err := c.Query("status")
	if err != nil {
		return nil, err
	}

	index := strings.Index(statusString, "#")
	i := 0
	for index == -1 {
		statusString, _ = c.Query("status")
		index = strings.Index(statusString, "#")
		i++
		if i == 5 {
			return nil, errors.New("Couldn't get output of status")
		}
	}

	users := strings.Split(statusString[index:], "\n")
	var list []Player
	for _, userString := range users {
		if !rePlayerInfo.MatchString(userString) {
			continue
		}
		matches := rePlayerInfo.FindStringSubmatch(userString)
		player := Player{
			UserID:   matches[1],
			Username: matches[2],
			SteamID:  matches[3],
			Ip:       matches[4],
		}
		list = append(list, player)
	}

	return list, nil
}

// KickPlayer kicks a player
func (c *TF2RconConnection) KickPlayer(p Player, message string) error {
	return c.KickPlayerID(p.UserID, message)
}

// Kicks a player with the given player ID
func (c *TF2RconConnection) KickPlayerID(userID string, message string) error {
	query := fmt.Sprintf("kickid %s %s", userID, message)
	_, err := c.Query(query)
	return err
}

// BanPlayer bans a player
func (c *TF2RconConnection) BanPlayer(minutes int, p Player, message string) error {
	query := "banid " + fmt.Sprintf("%v", minutes) + " " + p.UserID
	if message != "" {
		query += " \"" + message + "\""
	}
	_, err := c.Query(query)
	return err
}

// UnbanPlayer unbans a player
func (c *TF2RconConnection) UnbanPlayer(p Player) error {
	query := "unbanid " + p.UserID
	_, err := c.Query(query)
	return err
}

// Say sends a message to the TF2 server chat
func (c *TF2RconConnection) Say(message string) error {
	query := "say " + message
	_, err := c.Query(query)
	return err
}

func (c *TF2RconConnection) Sayf(format string, a ...interface{}) error {
	err := c.Say(fmt.Sprintf(format, a...))
	return err
}

// ChangeRconPassword changes the rcon password and updates the current connection
// to use the new password
func (c *TF2RconConnection) ChangeRconPassword(password string) error {
	_, err := c.SetConVar("rcon_password", password)

	if err == nil {
		c.rc.Close()
		newConnection, _ := rcon.Dial(c.host, password)
		c.rc = newConnection
	}

	return err
}

// ChangeMap changes the map
func (c *TF2RconConnection) ChangeMap(mapname string) error {
	query := "changelevel \"" + mapname + "\""
	res, err := c.Query(query)
	if res != "" {
		return errors.New("Map not found.")
	}
	return err
}

// ChangeServerPassword changes the server password
func (c *TF2RconConnection) ChangeServerPassword(password string) error {
	_, err := c.SetConVar("sv_password", password)
	return err
}

// GetServerPassword returns the server password
func (c *TF2RconConnection) GetServerPassword() (string, error) {
	return c.GetConVar("sv_password")
}

func (c *TF2RconConnection) AddTag(newTag string) error {
	tags, err := c.GetConVar("sv_tags")
	if err != nil {
		return err
	}

	// Source servers don't auto-remove duplicate tags, and noone
	tagExists := false
	for _, tag := range strings.Split(tags, ",") {
		if tag == newTag {
			tagExists = true
			break
		}
	}

	if !tagExists {
		newTags := strings.Join([]string{tags, newTag}, ",")
		_, err := c.SetConVar("sv_tags", newTags)
		return err
	}

	return nil
}

func (c *TF2RconConnection) RemoveTag(tagName string) error {
	tags, err := c.GetConVar("sv_tags")
	if err != nil {
		return err
	}

	if strings.Contains(tags, tagName) {
		// Replace all instances of the given tagName. This may leave
		// duplicated or trailing commas in the sv_tags string; however
		// Source servers clean up the value of sv_tags to remove those
		// anyways
		_, err := c.SetConVar("sv_tags", strings.Replace(tags, tagName, "", -1))
		return err
	}

	return nil
}

// RedirectLogs send the logaddress_add command
func (c *TF2RconConnection) RedirectLogs(addr string) error {
	query := "logaddress_add " + addr
	_, err := c.Query(query)
	return err
}

func (c *TF2RconConnection) StopLogRedirection(addr string) {
	query := fmt.Sprintf("logaddress_del %s", addr)
	c.QueryNoResp(query)
}

// Close closes the connection
func (c *TF2RconConnection) Close() {
	c.rc.Close()
}

// ExecConfig accepts a string and executes its lines one by one. Assumes
// UNiX line endings
func (c *TF2RconConnection) ExecConfig(config string) error {
	lines := strings.Split(config, "\n")
	for _, line := range lines {
		_, err := c.Query(line)
		if err != nil {
			return err
		}
	}
	return nil
}

// NewTF2RconConnection builds a new TF2RconConnection to a server at address ("ip:port") using
// a rcon_password password
func NewTF2RconConnection(address, password string) (*TF2RconConnection, error) {
	rc, err := rcon.Dial(address, password)
	if err != nil {
		return nil, err
	}

	return &TF2RconConnection{new(sync.RWMutex), rc, address, password}, nil
}

func (c *TF2RconConnection) Reconnect(duration time.Duration) error {
	var err error
	var cur time.Duration

	c.rcLock.Lock()
	defer c.rcLock.Unlock()

	c.Close()
	for cur += time.Second; cur <= duration; {
		c.rc, err = rcon.Dial(c.host, c.password)
		if err == nil {
			return nil
		}
	}

	return err
}
<|endoftext|>