<commit_before>// Copyright 2016 IBM Corporation
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

package config

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"reflect"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/urfave/cli"
	"gopkg.in/yaml.v2"
)

// Command TODO
type Command struct {
	Cmd     []string `yaml:"cmd"`
	Env     []string `yaml:"env"`
	Primary bool     `yaml:"primary"`
}

// Service configuration
type Service struct {
	Name string   `yaml:"name"`
	Tags []string `yaml:"tags"`
}

// Endpoint configuration
type Endpoint struct {
	Host string `yaml:"host"`
	Port int    `yaml:"port"`
	Type string `yaml:"type"`
}

// Registry configuration
type Registry struct {
	URL   string        `yaml:"url"`
	Token string        `yaml:"token"`
	Poll  time.Duration `yaml:"poll"`
}

// Controller configuration
type Controller struct {
	URL   string        `yaml:"url"`
	Token string        `yaml:"token"`
	Poll  time.Duration `yaml:"poll"`
}

// HealthCheck configuration.
type HealthCheck struct {
	Type     string        `yaml:"type"`
	Value    string        `yaml:"value"`
	Interval time.Duration `yaml:"interval"`
	Timeout  time.Duration `yaml:"timeout"`
	Method   string        `yaml:"method"`
	Code     int           `yaml:"code"`
}

// Config stores the various configuration options for the sidecar
type Config struct {
	Register bool `yaml:"register"`
	Proxy    bool `yaml:"proxy"`

	Service  Service  `yaml:"service"`
	Endpoint Endpoint `yaml:"endpoint"`

	Registry   Registry   `yaml:"registry"`
	Controller Controller `yaml:"controller"`

	Supervise bool `yaml:"supervise"`

	HealthChecks []HealthCheck `yaml:"healthchecks"`

	LogLevel string `yaml:"log_level"`

	Log            bool   `yaml:"log"`
	LogstashServer string `yaml:"logstash_server"`

	Commands []Command `yaml:"commands"`

	Debug string
}

// New creates a new Config object from the given commandline flags, environment variables, and configuration file context.
func New(context *cli.Context) (*Config, error) {

	// Initialize configuration with default values
	config := *&DefaultConfig

	// Load configuration from file, if specified
	if context.IsSet(configFlag) {
		err := config.loadFromFile(context.String(configFlag))
		if err != nil {
			return nil, err
		}
	}

	// Load configuration from context (commandline flags and environment variables)
	err := config.loadFromContext(context)
	if err != nil {
		return nil, err
	}

	if config.Endpoint.Host == "" {
		logrus.Infof("No hostname is configured. Using local IP instead...")
		config.Endpoint.Host = waitForLocalIP()
		logrus.Infof("Obtained local IP %s", config.Endpoint.Host)
	}

	return &config, nil
}

func (c *Config) loadFromFile(configFile string) error {
	bytes, err := ioutil.ReadFile(configFile)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(bytes, c)
	if err != nil {
		return nil
	}

	return nil
}

func (c *Config) loadFromContext(context *cli.Context) error {
	loadFromContextIfSet := func(ptr interface{}, flagName string) {
		if !context.IsSet(flagName) {
			return
		}

		configValue := reflect.ValueOf(ptr).Elem()
		var flagValue interface{}
		switch configValue.Kind() {
		case reflect.Bool:
			flagValue = context.Bool(flagName)
		case reflect.String:
			flagValue = context.String(flagName)
		case reflect.Int:
			flagValue = context.Int(flagName)
		case reflect.Int64:
			flagValue = context.Duration(flagName)
		case reflect.Float64:
			flagValue = context.Float64(flagName)
		case reflect.Slice:
			switch configValue.Type().Elem().Kind() {
			case reflect.String:
				flagValue = context.StringSlice(flagName)
			case reflect.Int:
				flagValue = context.IntSlice(flagName)
			default:
				logrus.Errorf("unsupported configuration type '%v' for '%v'", configValue.Kind(), flagName)
			}
		default:
			logrus.Errorf("unsupported configuration type '%v' for '%v'", configValue.Kind(), flagName)
		}

		configValue.Set(reflect.ValueOf(flagValue))
	}

	loadFromContextIfSet(&c.Register, registerFlag)
	loadFromContextIfSet(&c.Proxy, proxyFlag)
	loadFromContextIfSet(&c.Endpoint.Host, endpointHostFlag)
	loadFromContextIfSet(&c.Endpoint.Port, endpointPortFlag)
	loadFromContextIfSet(&c.Endpoint.Type, endpointTypeFlag)
	loadFromContextIfSet(&c.Registry.URL, registryURLFlag)
	loadFromContextIfSet(&c.Registry.Token, registryTokenFlag)
	loadFromContextIfSet(&c.Registry.Poll, registryPollFlag)
	loadFromContextIfSet(&c.Controller.URL, controllerURLFlag)
	loadFromContextIfSet(&c.Controller.Token, controllerTokenFlag)
	loadFromContextIfSet(&c.Controller.Poll, controllerPollFlag)
	loadFromContextIfSet(&c.Supervise, superviseFlag)
	loadFromContextIfSet(&c.Log, logFlag)
	loadFromContextIfSet(&c.LogLevel, logLevelFlag)
	loadFromContextIfSet(&c.LogstashServer, logstashServerFlag)
	loadFromContextIfSet(&c.LogLevel, logLevelFlag)
	loadFromContextIfSet(&c.Debug, debugFlag)

	if context.IsSet(serviceFlag) {
		name, tags := parseServiceNameAndTags(context.String(serviceFlag))
		c.Service.Name = name
		c.Service.Tags = tags
	}

	// For healthchecks flags, we take the raw flag value as the healthcheck value,
	// and let the 'register.BuildHealthChecks' do the hard work and figure out what
	// kind of healthcheck it is.
	if context.IsSet(healthchecksFlag) {
		hcValues := context.StringSlice(healthchecksFlag)
		for _, hcValue := range hcValues {
			hc := HealthCheck{
				Value: hcValue,
			}
			c.HealthChecks = append(c.HealthChecks, hc)
		}
	}

	if context.Args().Present() {
		cmd := Command{
			Cmd:     context.Args(),
			Primary: true,
		}
		c.Commands = []Command{cmd}
	}

	return nil
}

// Validate the configuration
func (c *Config) Validate() error {

	if !c.Register && !c.Proxy {
		return errors.New("Sidecar serves no purpose. Please enable either proxy or registry or both")
	}

	// Create list of validation checks
	validators := []ValidatorFunc{}

	if c.Supervise {
		validators = append(validators,
			func() error {
				if len(c.Commands) == 0 {
					return fmt.Errorf("Supervision mode requires application launch arguments")
				}
				return nil
			},
		)
	}

	if c.Log {
		validators = append(validators,
			IsNotEmpty("Logstash Host", c.LogstashServer),
		)
	}

	// Registry URL is needed for both proxying and registering.  Registry token is not required in all auth cases
	validators = append(validators, IsValidURL("Registry URL", c.Registry.URL))

	if c.Register {
		validators = append(validators,
			IsNotEmpty("Service Name", c.Service.Name),
			IsInRange("Service Endpoint Port", c.Endpoint.Port, 1, 65535),
			IsInSet("Service Endpoint Type", c.Endpoint.Type, []string{"http", "https", "tcp", "udp", "user"}),
		)
	}

	if c.Proxy {
		validators = append(validators,
			IsValidURL("Controller URL", c.Controller.URL),
			IsInRangeDuration("Controller polling interval", c.Controller.Poll, 5*time.Second, 1*time.Hour),
		)

	}

	return Validate(validators)
}

// waitForLocalIP waits until a local IP is available
func waitForLocalIP() string {
	ip := ""
	for {
		ip = localIP()
		if ip != "" {
			break
		}
		logrus.Warn("Could not obtain local IP")
		time.Sleep(time.Second * 10)
	}
	return ip
}

// localIP retrieves the IP address of the system
func localIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}

	for _, address := range addrs {
		// check the address type and if it is not a loopback return it
		if ipNet, ok := address.(*net.IPNet); ok && !ipNet.IP.IsLoopback() {
			if ipNet.IP.To4() != nil {
				return ipNet.IP.String()
			}
		}
	}

	return ""
}

func parseServiceNameAndTags(service string) (name string, tags []string) {
	i := strings.Index(service, ":")
	if i == -1 {
		name = service
		tags = []string{}
	} else {
		name = service[:i]
		tags = strings.Split(service[i+1:], ",")
	}
	return
}
<commit_msg>removed duplicate line<commit_after>// Copyright 2016 IBM Corporation
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

package config

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"reflect"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/urfave/cli"
	"gopkg.in/yaml.v2"
)

// Command TODO
type Command struct {
	Cmd     []string `yaml:"cmd"`
	Env     []string `yaml:"env"`
	Primary bool     `yaml:"primary"`
}

// Service configuration
type Service struct {
	Name string   `yaml:"name"`
	Tags []string `yaml:"tags"`
}

// Endpoint configuration
type Endpoint struct {
	Host string `yaml:"host"`
	Port int    `yaml:"port"`
	Type string `yaml:"type"`
}

// Registry configuration
type Registry struct {
	URL   string        `yaml:"url"`
	Token string        `yaml:"token"`
	Poll  time.Duration `yaml:"poll"`
}

// Controller configuration
type Controller struct {
	URL   string        `yaml:"url"`
	Token string        `yaml:"token"`
	Poll  time.Duration `yaml:"poll"`
}

// HealthCheck configuration.
type HealthCheck struct {
	Type     string        `yaml:"type"`
	Value    string        `yaml:"value"`
	Interval time.Duration `yaml:"interval"`
	Timeout  time.Duration `yaml:"timeout"`
	Method   string        `yaml:"method"`
	Code     int           `yaml:"code"`
}

// Config stores the various configuration options for the sidecar
type Config struct {
	Register bool `yaml:"register"`
	Proxy    bool `yaml:"proxy"`

	Service  Service  `yaml:"service"`
	Endpoint Endpoint `yaml:"endpoint"`

	Registry   Registry   `yaml:"registry"`
	Controller Controller `yaml:"controller"`

	Supervise bool `yaml:"supervise"`

	HealthChecks []HealthCheck `yaml:"healthchecks"`

	LogLevel string `yaml:"log_level"`

	Log            bool   `yaml:"log"`
	LogstashServer string `yaml:"logstash_server"`

	Commands []Command `yaml:"commands"`

	Debug string
}

// New creates a new Config object from the given commandline flags, environment variables, and configuration file context.
func New(context *cli.Context) (*Config, error) {

	// Initialize configuration with default values
	config := *&DefaultConfig

	// Load configuration from file, if specified
	if context.IsSet(configFlag) {
		err := config.loadFromFile(context.String(configFlag))
		if err != nil {
			return nil, err
		}
	}

	// Load configuration from context (commandline flags and environment variables)
	err := config.loadFromContext(context)
	if err != nil {
		return nil, err
	}

	if config.Endpoint.Host == "" {
		logrus.Infof("No hostname is configured. Using local IP instead...")
		config.Endpoint.Host = waitForLocalIP()
		logrus.Infof("Obtained local IP %s", config.Endpoint.Host)
	}

	return &config, nil
}

func (c *Config) loadFromFile(configFile string) error {
	bytes, err := ioutil.ReadFile(configFile)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(bytes, c)
	if err != nil {
		return nil
	}

	return nil
}

func (c *Config) loadFromContext(context *cli.Context) error {
	loadFromContextIfSet := func(ptr interface{}, flagName string) {
		if !context.IsSet(flagName) {
			return
		}

		configValue := reflect.ValueOf(ptr).Elem()
		var flagValue interface{}
		switch configValue.Kind() {
		case reflect.Bool:
			flagValue = context.Bool(flagName)
		case reflect.String:
			flagValue = context.String(flagName)
		case reflect.Int:
			flagValue = context.Int(flagName)
		case reflect.Int64:
			flagValue = context.Duration(flagName)
		case reflect.Float64:
			flagValue = context.Float64(flagName)
		case reflect.Slice:
			switch configValue.Type().Elem().Kind() {
			case reflect.String:
				flagValue = context.StringSlice(flagName)
			case reflect.Int:
				flagValue = context.IntSlice(flagName)
			default:
				logrus.Errorf("unsupported configuration type '%v' for '%v'", configValue.Kind(), flagName)
			}
		default:
			logrus.Errorf("unsupported configuration type '%v' for '%v'", configValue.Kind(), flagName)
		}

		configValue.Set(reflect.ValueOf(flagValue))
	}

	loadFromContextIfSet(&c.Register, registerFlag)
	loadFromContextIfSet(&c.Proxy, proxyFlag)
	loadFromContextIfSet(&c.Endpoint.Host, endpointHostFlag)
	loadFromContextIfSet(&c.Endpoint.Port, endpointPortFlag)
	loadFromContextIfSet(&c.Endpoint.Type, endpointTypeFlag)
	loadFromContextIfSet(&c.Registry.URL, registryURLFlag)
	loadFromContextIfSet(&c.Registry.Token, registryTokenFlag)
	loadFromContextIfSet(&c.Registry.Poll, registryPollFlag)
	loadFromContextIfSet(&c.Controller.URL, controllerURLFlag)
	loadFromContextIfSet(&c.Controller.Token, controllerTokenFlag)
	loadFromContextIfSet(&c.Controller.Poll, controllerPollFlag)
	loadFromContextIfSet(&c.Supervise, superviseFlag)
	loadFromContextIfSet(&c.Log, logFlag)
	loadFromContextIfSet(&c.LogstashServer, logstashServerFlag)
	loadFromContextIfSet(&c.LogLevel, logLevelFlag)
	loadFromContextIfSet(&c.Debug, debugFlag)

	if context.IsSet(serviceFlag) {
		name, tags := parseServiceNameAndTags(context.String(serviceFlag))
		c.Service.Name = name
		c.Service.Tags = tags
	}

	// For healthchecks flags, we take the raw flag value as the healthcheck value,
	// and let the 'register.BuildHealthChecks' do the hard work and figure out what
	// kind of healthcheck it is.
	if context.IsSet(healthchecksFlag) {
		hcValues := context.StringSlice(healthchecksFlag)
		for _, hcValue := range hcValues {
			hc := HealthCheck{
				Value: hcValue,
			}
			c.HealthChecks = append(c.HealthChecks, hc)
		}
	}

	if context.Args().Present() {
		cmd := Command{
			Cmd:     context.Args(),
			Primary: true,
		}
		c.Commands = []Command{cmd}
	}

	return nil
}

// Validate the configuration
func (c *Config) Validate() error {

	if !c.Register && !c.Proxy {
		return errors.New("Sidecar serves no purpose. Please enable either proxy or registry or both")
	}

	// Create list of validation checks
	validators := []ValidatorFunc{}

	if c.Supervise {
		validators = append(validators,
			func() error {
				if len(c.Commands) == 0 {
					return fmt.Errorf("Supervision mode requires application launch arguments")
				}
				return nil
			},
		)
	}

	if c.Log {
		validators = append(validators,
			IsNotEmpty("Logstash Host", c.LogstashServer),
		)
	}

	// Registry URL is needed for both proxying and registering.  Registry token is not required in all auth cases
	validators = append(validators, IsValidURL("Registry URL", c.Registry.URL))

	if c.Register {
		validators = append(validators,
			IsNotEmpty("Service Name", c.Service.Name),
			IsInRange("Service Endpoint Port", c.Endpoint.Port, 1, 65535),
			IsInSet("Service Endpoint Type", c.Endpoint.Type, []string{"http", "https", "tcp", "udp", "user"}),
		)
	}

	if c.Proxy {
		validators = append(validators,
			IsValidURL("Controller URL", c.Controller.URL),
			IsInRangeDuration("Controller polling interval", c.Controller.Poll, 5*time.Second, 1*time.Hour),
		)

	}

	return Validate(validators)
}

// waitForLocalIP waits until a local IP is available
func waitForLocalIP() string {
	ip := ""
	for {
		ip = localIP()
		if ip != "" {
			break
		}
		logrus.Warn("Could not obtain local IP")
		time.Sleep(time.Second * 10)
	}
	return ip
}

// localIP retrieves the IP address of the system
func localIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}

	for _, address := range addrs {
		// check the address type and if it is not a loopback return it
		if ipNet, ok := address.(*net.IPNet); ok && !ipNet.IP.IsLoopback() {
			if ipNet.IP.To4() != nil {
				return ipNet.IP.String()
			}
		}
	}

	return ""
}

func parseServiceNameAndTags(service string) (name string, tags []string) {
	i := strings.Index(service, ":")
	if i == -1 {
		name = service
		tags = []string{}
	} else {
		name = service[:i]
		tags = strings.Split(service[i+1:], ",")
	}
	return
}
<|endoftext|>


<commit_before>package exporter

import (
	"testing"
)

func TestAllRegexpsCompile(t *testing.T) {
	patterns := loadPatternDir(t)
	for pattern := range *patterns {
		_, err := Compile("%{"+pattern+"}", patterns)
		if err != nil {
			t.Errorf("%v", err.Error())
		}
	}
}
<commit_msg>some more tests in grok_test.go<commit_after>package exporter

import (
	"testing"
	"strings"
)

func TestAllRegexpsCompile(t *testing.T) {
	patterns := loadPatternDir(t)
	for pattern := range *patterns {
		_, err := Compile("%{"+pattern+"}", patterns)
		if err != nil {
			t.Errorf("%v", err.Error())
		}
	}
}

func TestUnknownGrokPattern(t *testing.T) {
	patterns := loadPatternDir(t)
	_, err := Compile("%{USER} [a-z] %{SOME_UNKNOWN_PATTERN}.*", patterns)
	if err == nil || ! strings.Contains(err.Error(), "SOME_UNKNOWN_PATTERN") {
		t.Error("expected error message saying which pattern is undefined.")
	}
}

func TestInvalidRegexp(t *testing.T) {
	patterns := loadPatternDir(t)
	_, err := Compile("%{USER} [a-z] \\", patterns) // wrong because regex cannot end with backslash
	if err == nil || ! strings.Contains(err.Error(), "%{USER} [a-z] \\") {
		t.Error("expected error message saying which pattern is invalid.")
	}
}

func TestNamedCaptureGroup(t *testing.T) {
	patterns := loadPatternDir(t)
	regex, err := Compile("User %{USER:user} has logged in.", patterns)
	if err != nil {
		t.Error(err)
	}
	found := regex.Gsub("User fabian has logged in.", "\\k<user>")
	if found != "fabian" {
		t.Errorf("Expected to capture 'fabian', but captured '%v'.", found)
	}
}
<|endoftext|>


<commit_before>package config

import (
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/user"
	"path"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/golog"
	"github.com/getlantern/nattywad"
	"github.com/getlantern/yaml"
	"github.com/getlantern/yamlconf"

	"github.com/getlantern/flashlight/client"
	"github.com/getlantern/flashlight/globals"
	"github.com/getlantern/flashlight/server"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/util"
)

const (
	CloudConfigPollInterval = 1 * time.Minute

	cloudflare  = "cloudflare"
	etag        = "ETag"
	ifNoneMatch = "If-None-Match"
)

var (
	log                 = golog.LoggerFor("flashlight.config")
	m                   *yamlconf.Manager
	lastCloudConfigETag = ""
)

type Config struct {
	Version       int
	CloudConfig   string
	CloudConfigCA string
	Addr          string
	Role          string
	InstanceId    string
	Country       string
	StatsAddr     string
	CpuProfile    string
	MemProfile    string
	WaddellCert   string
	Stats         *statreporter.Config
	Server        *server.ServerConfig
	Client        *client.ClientConfig
	TrustedCAs    []*CA
}

// CA represents a certificate authority
type CA struct {
	CommonName string
	Cert       string // PEM-encoded
}

// Start starts the configuration system.
func Start(updateHandler func(updated *Config)) (*Config, error) {
	m = &yamlconf.Manager{
		FilePath:         InConfigDir("lantern.yaml"),
		FilePollInterval: 1 * time.Second,
		ConfigServerAddr: *configaddr,
		EmptyConfig: func() yamlconf.Config {
			return &Config{}
		},
		OneTimeSetup: func(ycfg yamlconf.Config) error {
			cfg := ycfg.(*Config)
			return cfg.applyFlags()
		},
		CustomPoll: func(currentCfg yamlconf.Config) (mutate func(yamlconf.Config) error, waitTime time.Duration, err error) {
			cfg := currentCfg.(*Config)
			waitTime = cfg.cloudPollSleepTime()
			if cfg.CloudConfig == "" {
				// Config doesn't have a CloudConfig, just ignore
				mutate = func(ycfg yamlconf.Config) error {
					// do nothing
					return nil
				}
				return
			}

			var bytes []byte
			bytes, err = cfg.fetchCloudConfig()
			if err == nil {
				mutate = func(ycfg yamlconf.Config) error {
					log.Debugf("Merging cloud configuration")
					cfg := ycfg.(*Config)
					return cfg.updateFrom(bytes)
				}
			}
			return
		},
	}
	initial, err := m.Start()
	var cfg *Config
	if err == nil {
		cfg = initial.(*Config)
		updateGlobals(cfg)
		go func() {
			// Read updates
			for {
				next := m.Next()
				nextCfg := next.(*Config)
				updateGlobals(nextCfg)
				updateHandler(nextCfg)
			}
		}()
	}
	return cfg, err
}

func updateGlobals(cfg *Config) {
	globals.InstanceId = cfg.InstanceId
	globals.Country = cfg.Country
	if cfg.WaddellCert != "" {
		globals.WaddellCert = cfg.WaddellCert
	}
	err := globals.SetTrustedCAs(cfg.TrustedCACerts())
	if err != nil {
		log.Fatalf("Unable to configure trusted CAs: %s", err)
	}
}

// Update updates the configuration using the given mutator function.
func Update(mutate func(cfg *Config) error) error {
	return m.Update(func(ycfg yamlconf.Config) error {
		return mutate(ycfg.(*Config))
	})
}

// InConfigDir returns the path to the given filename inside of the configdir.
func InConfigDir(filename string) string {
	cdir := *configdir
	if cdir == "" {
		cdir = platformSpecificConfigDir()
	}
	log.Debugf("Placing configuration in %v", cdir)
	if _, err := os.Stat(cdir); err != nil {
		if os.IsNotExist(err) {
			// Create config dir
			if err := os.MkdirAll(cdir, 0755); err != nil {
				log.Fatalf("Unable to create configdir at %s: %s", cdir, err)
			}
		}
	}
	return fmt.Sprintf("%s%c%s", cdir, os.PathSeparator, filename)
}

// TrustedCACerts returns a slice of PEM-encoded certs for the trusted CAs
func (cfg *Config) TrustedCACerts() []string {
	certs := make([]string, 0, len(cfg.TrustedCAs))
	for _, ca := range cfg.TrustedCAs {
		certs = append(certs, ca.Cert)
	}
	return certs
}

// GetVersion implements the method from interface yamlconf.Config
func (cfg *Config) GetVersion() int {
	return cfg.Version
}

// SetVersion implements the method from interface yamlconf.Config
func (cfg *Config) SetVersion(version int) {
	cfg.Version = version
}

// ApplyDefaults implements the method from interface yamlconf.Config
//
// ApplyDefaults populates default values on a Config to make sure that we have
// a minimum viable config for running.  As new settings are added to
// flashlight, this function should be updated to provide sensible defaults for
// those settings.
func (cfg *Config) ApplyDefaults() {
	if cfg.Role == "" {
		cfg.Role = "client"
	}

	if cfg.Addr == "" {
		cfg.Addr = "localhost:8787"
	}

	// Default country
	if cfg.Country == "" {
		cfg.Country = *country
	}

	// Make sure we always have a stats config
	if cfg.Stats == nil {
		cfg.Stats = &statreporter.Config{}
	}

	if cfg.Stats.StatshubAddr == "" {
		cfg.Stats.StatshubAddr = *statshubAddr
	}

	if cfg.Client != nil && cfg.Role == "client" {
		cfg.applyClientDefaults()
	}

	if cfg.TrustedCAs == nil || len(cfg.TrustedCAs) == 0 {
		cfg.TrustedCAs = defaultTrustedCAs
	}
}

func (cfg *Config) applyClientDefaults() {
	// Make sure we always have at least one masquerade set
	if cfg.Client.MasqueradeSets == nil {
		cfg.Client.MasqueradeSets = make(map[string][]*fronted.Masquerade)
	}
	if len(cfg.Client.MasqueradeSets) == 0 {
		cfg.Client.MasqueradeSets[cloudflare] = cloudflareMasquerades
	}

	// Make sure we always have at least one server
	if cfg.Client.FrontedServers == nil {
		cfg.Client.FrontedServers = make([]*client.FrontedServerInfo, 0)
	}
	if len(cfg.Client.FrontedServers) == 0 && len(cfg.Client.ChainedServers) == 0 {
		cfg.Client.FrontedServers = append(cfg.Client.FrontedServers, &client.FrontedServerInfo{
			Host:           "fallbacks.getiantem.org",
			Port:           443,
			PoolSize:       30,
			MasqueradeSet:  cloudflare,
			MaxMasquerades: 20,
			QOS:            10,
			Weight:         4000,
		})
	}

	// Make sure all servers have a QOS and Weight configured
	for _, server := range cfg.Client.FrontedServers {
		if server.QOS == 0 {
			server.QOS = 5
		}
		if server.Weight == 0 {
			server.Weight = 100
		}
		if server.RedialAttempts == 0 {
			server.RedialAttempts = 2
		}
	}

	// Always make sure we have a map of ChainedServers
	if cfg.Client.ChainedServers == nil {
		cfg.Client.ChainedServers = make(map[string]*client.ChainedServerInfo)
	}

	// Always make sure that we have a map of Peers
	if cfg.Client.Peers == nil {
		cfg.Client.Peers = make(map[string]*nattywad.ServerPeer)
	}

	// Sort servers so that they're always in a predictable order
	cfg.Client.SortServers()
}

func (cfg *Config) IsDownstream() bool {
	return cfg.Role == "client"
}

func (cfg *Config) IsUpstream() bool {
	return !cfg.IsDownstream()
}

func (cfg Config) cloudPollSleepTime() time.Duration {
	return time.Duration((CloudConfigPollInterval.Nanoseconds() / 2) + rand.Int63n(CloudConfigPollInterval.Nanoseconds()))
}

func (cfg Config) fetchCloudConfig() ([]byte, error) {
	log.Debugf("Fetching cloud config from: %s", cfg.CloudConfig)
	// Try it unproxied first
	bytes, err := cfg.doFetchCloudConfig("")
	if err != nil && cfg.IsDownstream() {
		// If that failed, try it proxied
		bytes, err = cfg.doFetchCloudConfig(cfg.Addr)
	}
	if err != nil {
		return nil, fmt.Errorf("Unable to read yaml from %s: %s", cfg.CloudConfig, err)
	}
	return bytes, err
}

func (cfg Config) doFetchCloudConfig(proxyAddr string) ([]byte, error) {
	client, err := util.HTTPClient(cfg.CloudConfigCA, proxyAddr)
	if err != nil {
		return nil, fmt.Errorf("Unable to initialize HTTP client: %s", err)
	}
	log.Debugf("Checking for cloud configuration at: %s", cfg.CloudConfig)
	req, err := http.NewRequest("GET", cfg.CloudConfig, nil)
	if err != nil {
		return nil, fmt.Errorf("Unable to construct request for cloud config at %s: %s", cfg.CloudConfig, err)
	}
	if lastCloudConfigETag != "" {
		// Don't bother fetching if unchanged
		req.Header.Set(ifNoneMatch, lastCloudConfigETag)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Unable to fetch cloud config at %s: %s", cfg.CloudConfig, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode == 304 {
		log.Debugf("Config unchanged in cloud")
		return nil, nil
	} else if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Unexpected response status: %d", resp.StatusCode)
	}
	lastCloudConfigETag = resp.Header.Get(etag)
	gzReader, err := gzip.NewReader(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("Unable to open gzip reader: %s", err)
	}
	return ioutil.ReadAll(gzReader)
}

// updateFrom creates a new Config by merging the given yaml into this Config.
// Any servers in the updated yaml replace ones in the original Config and any
// masquerade sets in the updated yaml replace ones in the original Config.
func (updated *Config) updateFrom(updateBytes []byte) error {
	err := yaml.Unmarshal(updateBytes, updated)
	if err != nil {
		return fmt.Errorf("Unable to unmarshal YAML for update: %s", err)
	}
	// Need to de-duplicate servers, since yaml appends them
	servers := make(map[string]*client.FrontedServerInfo)
	for _, server := range updated.Client.FrontedServers {
		servers[server.Host] = server
	}
	updated.Client.FrontedServers = make([]*client.FrontedServerInfo, 0, len(servers))
	for _, server := range servers {
		updated.Client.FrontedServers = append(updated.Client.FrontedServers, server)
	}
	return nil
}

func inHomeDir(filename string) string {
	log.Tracef("Determining user's home directory")
	usr, err := user.Current()
	if err != nil {
		panic(fmt.Errorf("Unable to determine user's home directory: %s", err))
	}
	return path.Join(usr.HomeDir, filename)
}
<commit_msg>Made path formatting consistent<commit_after>package config

import (
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/user"
	"path/filepath"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/golog"
	"github.com/getlantern/nattywad"
	"github.com/getlantern/yaml"
	"github.com/getlantern/yamlconf"

	"github.com/getlantern/flashlight/client"
	"github.com/getlantern/flashlight/globals"
	"github.com/getlantern/flashlight/server"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/util"
)

const (
	CloudConfigPollInterval = 1 * time.Minute

	cloudflare  = "cloudflare"
	etag        = "ETag"
	ifNoneMatch = "If-None-Match"
)

var (
	log                 = golog.LoggerFor("flashlight.config")
	m                   *yamlconf.Manager
	lastCloudConfigETag = ""
)

type Config struct {
	Version       int
	CloudConfig   string
	CloudConfigCA string
	Addr          string
	Role          string
	InstanceId    string
	Country       string
	StatsAddr     string
	CpuProfile    string
	MemProfile    string
	WaddellCert   string
	Stats         *statreporter.Config
	Server        *server.ServerConfig
	Client        *client.ClientConfig
	TrustedCAs    []*CA
}

// CA represents a certificate authority
type CA struct {
	CommonName string
	Cert       string // PEM-encoded
}

// Start starts the configuration system.
func Start(updateHandler func(updated *Config)) (*Config, error) {
	m = &yamlconf.Manager{
		FilePath:         InConfigDir("lantern.yaml"),
		FilePollInterval: 1 * time.Second,
		ConfigServerAddr: *configaddr,
		EmptyConfig: func() yamlconf.Config {
			return &Config{}
		},
		OneTimeSetup: func(ycfg yamlconf.Config) error {
			cfg := ycfg.(*Config)
			return cfg.applyFlags()
		},
		CustomPoll: func(currentCfg yamlconf.Config) (mutate func(yamlconf.Config) error, waitTime time.Duration, err error) {
			cfg := currentCfg.(*Config)
			waitTime = cfg.cloudPollSleepTime()
			if cfg.CloudConfig == "" {
				// Config doesn't have a CloudConfig, just ignore
				mutate = func(ycfg yamlconf.Config) error {
					// do nothing
					return nil
				}
				return
			}

			var bytes []byte
			bytes, err = cfg.fetchCloudConfig()
			if err == nil {
				mutate = func(ycfg yamlconf.Config) error {
					log.Debugf("Merging cloud configuration")
					cfg := ycfg.(*Config)
					return cfg.updateFrom(bytes)
				}
			}
			return
		},
	}
	initial, err := m.Start()
	var cfg *Config
	if err == nil {
		cfg = initial.(*Config)
		updateGlobals(cfg)
		go func() {
			// Read updates
			for {
				next := m.Next()
				nextCfg := next.(*Config)
				updateGlobals(nextCfg)
				updateHandler(nextCfg)
			}
		}()
	}
	return cfg, err
}

func updateGlobals(cfg *Config) {
	globals.InstanceId = cfg.InstanceId
	globals.Country = cfg.Country
	if cfg.WaddellCert != "" {
		globals.WaddellCert = cfg.WaddellCert
	}
	err := globals.SetTrustedCAs(cfg.TrustedCACerts())
	if err != nil {
		log.Fatalf("Unable to configure trusted CAs: %s", err)
	}
}

// Update updates the configuration using the given mutator function.
func Update(mutate func(cfg *Config) error) error {
	return m.Update(func(ycfg yamlconf.Config) error {
		return mutate(ycfg.(*Config))
	})
}

// InConfigDir returns the path to the given filename inside of the configdir.
func InConfigDir(filename string) string {
	cdir := *configdir
	if cdir == "" {
		cdir = platformSpecificConfigDir()
	}
	log.Debugf("Placing configuration in %v", cdir)
	if _, err := os.Stat(cdir); err != nil {
		if os.IsNotExist(err) {
			// Create config dir
			if err := os.MkdirAll(cdir, 0755); err != nil {
				log.Fatalf("Unable to create configdir at %s: %s", cdir, err)
			}
		}
	}
	return filepath.Join(cdir, filename)
}

// TrustedCACerts returns a slice of PEM-encoded certs for the trusted CAs
func (cfg *Config) TrustedCACerts() []string {
	certs := make([]string, 0, len(cfg.TrustedCAs))
	for _, ca := range cfg.TrustedCAs {
		certs = append(certs, ca.Cert)
	}
	return certs
}

// GetVersion implements the method from interface yamlconf.Config
func (cfg *Config) GetVersion() int {
	return cfg.Version
}

// SetVersion implements the method from interface yamlconf.Config
func (cfg *Config) SetVersion(version int) {
	cfg.Version = version
}

// ApplyDefaults implements the method from interface yamlconf.Config
//
// ApplyDefaults populates default values on a Config to make sure that we have
// a minimum viable config for running.  As new settings are added to
// flashlight, this function should be updated to provide sensible defaults for
// those settings.
func (cfg *Config) ApplyDefaults() {
	if cfg.Role == "" {
		cfg.Role = "client"
	}

	if cfg.Addr == "" {
		cfg.Addr = "localhost:8787"
	}

	// Default country
	if cfg.Country == "" {
		cfg.Country = *country
	}

	// Make sure we always have a stats config
	if cfg.Stats == nil {
		cfg.Stats = &statreporter.Config{}
	}

	if cfg.Stats.StatshubAddr == "" {
		cfg.Stats.StatshubAddr = *statshubAddr
	}

	if cfg.Client != nil && cfg.Role == "client" {
		cfg.applyClientDefaults()
	}

	if cfg.TrustedCAs == nil || len(cfg.TrustedCAs) == 0 {
		cfg.TrustedCAs = defaultTrustedCAs
	}
}

func (cfg *Config) applyClientDefaults() {
	// Make sure we always have at least one masquerade set
	if cfg.Client.MasqueradeSets == nil {
		cfg.Client.MasqueradeSets = make(map[string][]*fronted.Masquerade)
	}
	if len(cfg.Client.MasqueradeSets) == 0 {
		cfg.Client.MasqueradeSets[cloudflare] = cloudflareMasquerades
	}

	// Make sure we always have at least one server
	if cfg.Client.FrontedServers == nil {
		cfg.Client.FrontedServers = make([]*client.FrontedServerInfo, 0)
	}
	if len(cfg.Client.FrontedServers) == 0 && len(cfg.Client.ChainedServers) == 0 {
		cfg.Client.FrontedServers = append(cfg.Client.FrontedServers, &client.FrontedServerInfo{
			Host:           "fallbacks.getiantem.org",
			Port:           443,
			PoolSize:       30,
			MasqueradeSet:  cloudflare,
			MaxMasquerades: 20,
			QOS:            10,
			Weight:         4000,
		})
	}

	// Make sure all servers have a QOS and Weight configured
	for _, server := range cfg.Client.FrontedServers {
		if server.QOS == 0 {
			server.QOS = 5
		}
		if server.Weight == 0 {
			server.Weight = 100
		}
		if server.RedialAttempts == 0 {
			server.RedialAttempts = 2
		}
	}

	// Always make sure we have a map of ChainedServers
	if cfg.Client.ChainedServers == nil {
		cfg.Client.ChainedServers = make(map[string]*client.ChainedServerInfo)
	}

	// Always make sure that we have a map of Peers
	if cfg.Client.Peers == nil {
		cfg.Client.Peers = make(map[string]*nattywad.ServerPeer)
	}

	// Sort servers so that they're always in a predictable order
	cfg.Client.SortServers()
}

func (cfg *Config) IsDownstream() bool {
	return cfg.Role == "client"
}

func (cfg *Config) IsUpstream() bool {
	return !cfg.IsDownstream()
}

func (cfg Config) cloudPollSleepTime() time.Duration {
	return time.Duration((CloudConfigPollInterval.Nanoseconds() / 2) + rand.Int63n(CloudConfigPollInterval.Nanoseconds()))
}

func (cfg Config) fetchCloudConfig() ([]byte, error) {
	log.Debugf("Fetching cloud config from: %s", cfg.CloudConfig)
	// Try it unproxied first
	bytes, err := cfg.doFetchCloudConfig("")
	if err != nil && cfg.IsDownstream() {
		// If that failed, try it proxied
		bytes, err = cfg.doFetchCloudConfig(cfg.Addr)
	}
	if err != nil {
		return nil, fmt.Errorf("Unable to read yaml from %s: %s", cfg.CloudConfig, err)
	}
	return bytes, err
}

func (cfg Config) doFetchCloudConfig(proxyAddr string) ([]byte, error) {
	client, err := util.HTTPClient(cfg.CloudConfigCA, proxyAddr)
	if err != nil {
		return nil, fmt.Errorf("Unable to initialize HTTP client: %s", err)
	}
	log.Debugf("Checking for cloud configuration at: %s", cfg.CloudConfig)
	req, err := http.NewRequest("GET", cfg.CloudConfig, nil)
	if err != nil {
		return nil, fmt.Errorf("Unable to construct request for cloud config at %s: %s", cfg.CloudConfig, err)
	}
	if lastCloudConfigETag != "" {
		// Don't bother fetching if unchanged
		req.Header.Set(ifNoneMatch, lastCloudConfigETag)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Unable to fetch cloud config at %s: %s", cfg.CloudConfig, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode == 304 {
		log.Debugf("Config unchanged in cloud")
		return nil, nil
	} else if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Unexpected response status: %d", resp.StatusCode)
	}
	lastCloudConfigETag = resp.Header.Get(etag)
	gzReader, err := gzip.NewReader(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("Unable to open gzip reader: %s", err)
	}
	return ioutil.ReadAll(gzReader)
}

// updateFrom creates a new Config by merging the given yaml into this Config.
// Any servers in the updated yaml replace ones in the original Config and any
// masquerade sets in the updated yaml replace ones in the original Config.
func (updated *Config) updateFrom(updateBytes []byte) error {
	err := yaml.Unmarshal(updateBytes, updated)
	if err != nil {
		return fmt.Errorf("Unable to unmarshal YAML for update: %s", err)
	}
	// Need to de-duplicate servers, since yaml appends them
	servers := make(map[string]*client.FrontedServerInfo)
	for _, server := range updated.Client.FrontedServers {
		servers[server.Host] = server
	}
	updated.Client.FrontedServers = make([]*client.FrontedServerInfo, 0, len(servers))
	for _, server := range servers {
		updated.Client.FrontedServers = append(updated.Client.FrontedServers, server)
	}
	return nil
}

func inHomeDir(filename string) string {
	log.Tracef("Determining user's home directory")
	usr, err := user.Current()
	if err != nil {
		panic(fmt.Errorf("Unable to determine user's home directory: %s", err))
	}
	return filepath.Join(usr.HomeDir, filename)
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/testing"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/cmd/envcmd"
)

var _ = gc.Suite(&commandsSuite{})

type commandsSuite struct {
	stub    *testing.Stub
	command *stubCommand
}

func (s *commandsSuite) SetUpTest(c *gc.C) {
	s.stub = &testing.Stub{}
	s.command = &stubCommand{stub: s.stub}
}

func (s *commandsSuite) TearDownTest(c *gc.C) {
	registeredCommands = nil
}

func (s *commandsSuite) TestRegisterCommand(c *gc.C) {
	RegisterCommand(func() cmd.Command {
		return s.command
	})

	// We can't compare functions directly, so...
	c.Check(registeredEnvCommands, gc.HasLen, 0)
	c.Assert(registeredCommands, gc.HasLen, 1)
	command := registeredCommands[0]()
	c.Check(command, gc.Equals, s.command)
}

func (s *commandsSuite) TestRegisterEnvCommand(c *gc.C) {
	RegisterEnvCommand(func() envcmd.EnvironCommand {
		return s.command
	})

	// We can't compare functions directly, so...
	c.Assert(registeredCommands, gc.HasLen, 0)
	c.Assert(registeredEnvCommands, gc.HasLen, 1)
	command := registeredEnvCommands[0]()
	c.Check(command, gc.Equals, s.command)
}

type stubCommand struct {
	cmd.CommandBase
	stub *testing.Stub
	info *cmd.Info
}

func (c *stubCommand) Info() *cmd.Info {
	c.stub.AddCall("Info")
	c.stub.NextErr() // pop one off

	if c.info == nil {
		return &cmd.Info{
			Name: "some-command",
		}
	}
	return c.info
}

func (c *stubCommand) Run(ctx *cmd.Context) error {
	c.stub.AddCall("Run", ctx)
	if err := c.stub.NextErr(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (c *stubCommand) SetEnvName(name string) {
	c.stub.AddCall("SetEnvName", name)
	c.stub.NextErr() // pop one off

	// Do nothing.
}

type stubRegistry struct {
	stub *testing.Stub

	names []string
}

func (r *stubRegistry) Register(subcmd cmd.Command) {
	r.stub.AddCall("Register", subcmd)
	r.stub.NextErr() // pop one off

	r.names = append(r.names, subcmd.Info().Name)
	for _, name := range subcmd.Info().Aliases {
		r.names = append(r.names, name)
	}
}

func (r *stubRegistry) RegisterSuperAlias(name, super, forName string, check cmd.DeprecationCheck) {
	r.stub.AddCall("RegisterSuperAlias", name, super, forName)
	r.stub.NextErr() // pop one off

	r.names = append(r.names, name)
}

func (r *stubRegistry) RegisterDeprecated(subcmd cmd.Command, check cmd.DeprecationCheck) {
	r.stub.AddCall("RegisterDeprecated", subcmd, check)
	r.stub.NextErr() // pop one off

	r.names = append(r.names, subcmd.Info().Name)
	for _, name := range subcmd.Info().Aliases {
		r.names = append(r.names, name)
	}
}
<commit_msg>Make sure the test commands get properly cleared.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/testing"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/cmd/envcmd"
)

var _ = gc.Suite(&commandsSuite{})

type commandsSuite struct {
	stub    *testing.Stub
	command *stubCommand
}

func (s *commandsSuite) SetUpTest(c *gc.C) {
	s.stub = &testing.Stub{}
	s.command = &stubCommand{stub: s.stub}
}

func (s *commandsSuite) TearDownTest(c *gc.C) {
	registeredCommands = nil
	registeredEnvCommands = nil
}

func (s *commandsSuite) TestRegisterCommand(c *gc.C) {
	RegisterCommand(func() cmd.Command {
		return s.command
	})

	// We can't compare functions directly, so...
	c.Check(registeredEnvCommands, gc.HasLen, 0)
	c.Assert(registeredCommands, gc.HasLen, 1)
	command := registeredCommands[0]()
	c.Check(command, gc.Equals, s.command)
}

func (s *commandsSuite) TestRegisterEnvCommand(c *gc.C) {
	RegisterEnvCommand(func() envcmd.EnvironCommand {
		return s.command
	})

	// We can't compare functions directly, so...
	c.Assert(registeredCommands, gc.HasLen, 0)
	c.Assert(registeredEnvCommands, gc.HasLen, 1)
	command := registeredEnvCommands[0]()
	c.Check(command, gc.Equals, s.command)
}

type stubCommand struct {
	cmd.CommandBase
	stub *testing.Stub
	info *cmd.Info
}

func (c *stubCommand) Info() *cmd.Info {
	c.stub.AddCall("Info")
	c.stub.NextErr() // pop one off

	if c.info == nil {
		return &cmd.Info{
			Name: "some-command",
		}
	}
	return c.info
}

func (c *stubCommand) Run(ctx *cmd.Context) error {
	c.stub.AddCall("Run", ctx)
	if err := c.stub.NextErr(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (c *stubCommand) SetEnvName(name string) {
	c.stub.AddCall("SetEnvName", name)
	c.stub.NextErr() // pop one off

	// Do nothing.
}

type stubRegistry struct {
	stub *testing.Stub

	names []string
}

func (r *stubRegistry) Register(subcmd cmd.Command) {
	r.stub.AddCall("Register", subcmd)
	r.stub.NextErr() // pop one off

	r.names = append(r.names, subcmd.Info().Name)
	for _, name := range subcmd.Info().Aliases {
		r.names = append(r.names, name)
	}
}

func (r *stubRegistry) RegisterSuperAlias(name, super, forName string, check cmd.DeprecationCheck) {
	r.stub.AddCall("RegisterSuperAlias", name, super, forName)
	r.stub.NextErr() // pop one off

	r.names = append(r.names, name)
}

func (r *stubRegistry) RegisterDeprecated(subcmd cmd.Command, check cmd.DeprecationCheck) {
	r.stub.AddCall("RegisterDeprecated", subcmd, check)
	r.stub.NextErr() // pop one off

	r.names = append(r.names, subcmd.Info().Name)
	for _, name := range subcmd.Info().Aliases {
		r.names = append(r.names, name)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/go-igdman/igdman"
	"github.com/getlantern/golog"
	"github.com/getlantern/nattywad"
	"github.com/getlantern/waddell"

	"github.com/getlantern/flashlight/globals"
	"github.com/getlantern/flashlight/nattest"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/statserver"
)

const (
	PortmapFailure = 50
)

var (
	log = golog.LoggerFor("flashlight.server")
)

type Server struct {
	// Addr: listen address in form of host:port
	Addr string

	// Host: FQDN that is guaranteed to hit this server
	Host string

	// ReadTimeout: (optional) timeout for read ops
	ReadTimeout time.Duration

	// WriteTimeout: (optional) timeout for write ops
	WriteTimeout time.Duration

	CertContext                *fronted.CertContext // context for certificate management
	AllowNonGlobalDestinations bool                 // if true, requests to LAN, Loopback, etc. will be allowed

	waddellClient  *waddell.Client
	nattywadServer *nattywad.Server
	cfg            *ServerConfig
	cfgMutex       sync.Mutex
}

func (server *Server) Configure(newCfg *ServerConfig) {
	server.cfgMutex.Lock()
	defer server.cfgMutex.Unlock()

	oldCfg := server.cfg

	log.Debug("Server.Configure() called")
	if oldCfg != nil && reflect.DeepEqual(oldCfg, newCfg) {
		log.Debugf("Server configuration unchanged")
		return
	}

	if oldCfg == nil || newCfg.Portmap != oldCfg.Portmap {
		// Portmap changed
		if oldCfg != nil && oldCfg.Portmap > 0 {
			log.Debugf("Attempting to unmap old external port %d", oldCfg.Portmap)
			err := unmapPort(oldCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to unmap old external port: %s", err)
			}
			log.Debugf("Unmapped old external port %d", oldCfg.Portmap)
		}

		if newCfg.Portmap > 0 {
			log.Debugf("Attempting to map new external port %d", newCfg.Portmap)
			err := mapPort(server.Addr, newCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to map new external port: %s", err)
				os.Exit(PortmapFailure)
			}
			log.Debugf("Mapped new external port %d", newCfg.Portmap)
		}
	}

	nattywadIsEnabled := newCfg.WaddellAddr != ""
	nattywadWasEnabled := server.nattywadServer != nil
	waddellAddrChanged := oldCfg == nil && newCfg.WaddellAddr != "" || oldCfg != nil && oldCfg.WaddellAddr != newCfg.WaddellAddr

	if waddellAddrChanged {
		if nattywadWasEnabled {
			server.stopNattywad()
		}
		if nattywadIsEnabled {
			server.startNattywad(newCfg.WaddellAddr)
		}
	}

	server.cfg = newCfg
}

func (server *Server) ListenAndServe() error {
	if server.Host != "" {
		log.Debugf("Running as host %s", server.Host)
	}

	fs := &fronted.Server{
		Addr:                       server.Addr,
		Host:                       server.Host,
		ReadTimeout:                server.ReadTimeout,
		WriteTimeout:               server.WriteTimeout,
		CertContext:                server.CertContext,
		AllowNonGlobalDestinations: server.AllowNonGlobalDestinations,
	}

	if server.cfg.Unencrypted {
		log.Debug("Running in unencrypted")
		fs.CertContext = nil
	}

	// Add callbacks to track bytes given
	fs.OnBytesReceived = func(ip string, bytes int64) {
		onBytesGiven(bytes)
		statserver.OnBytesReceived(ip, bytes)
	}
	fs.OnBytesSent = func(ip string, bytes int64) {
		onBytesGiven(bytes)
		statserver.OnBytesSent(ip, bytes)
	}

	l, err := fs.Listen()
	if err != nil {
		return fmt.Errorf("Unable to listen at %s: %s", server.Addr, err)
	}
	return fs.Serve(l)
}

func (server *Server) startNattywad(waddellAddr string) {
	log.Debugf("Connecting to waddell at: %s", waddellAddr)
	var err error
	server.waddellClient, err = waddell.NewClient(&waddell.ClientConfig{
		Dial: func() (net.Conn, error) {
			return net.Dial("tcp", waddellAddr)
		},
		ServerCert:        globals.WaddellCert,
		ReconnectAttempts: 10,
		OnId: func(id waddell.PeerId) {
			log.Debugf("Connected to Waddell!! Id is: %s", id)
		},
	})
	if err != nil {
		log.Errorf("Unable to connect to waddell: %s", err)
		server.waddellClient = nil
		return
	}
	server.nattywadServer = &nattywad.Server{
		Client: server.waddellClient,
		OnSuccess: func(local *net.UDPAddr, remote *net.UDPAddr) bool {
			err := nattest.Serve(local)
			if err != nil {
				log.Error(err.Error())
				return false
			}
			return true
		},
	}
	server.nattywadServer.Start()
}

func (server *Server) stopNattywad() {
	log.Debug("Stopping nattywad server")
	server.nattywadServer.Stop()
	server.nattywadServer = nil
	log.Debug("Stopping waddell client")
	server.waddellClient.Close()
	server.waddellClient = nil
}

func mapPort(addr string, port int) error {
	parts := strings.Split(addr, ":")

	internalPort, err := strconv.Atoi(parts[1])
	if err != nil {
		return fmt.Errorf("Unable to parse local port: ")
	}

	internalIP := parts[0]
	if internalIP == "" {
		internalIP, err = determineInternalIP()
		if err != nil {
			return fmt.Errorf("Unable to determine internal IP: %s", err)
		}
	}

	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	err = igd.AddPortMapping(igdman.TCP, internalIP, internalPort, port, 0)
	if err != nil {
		return fmt.Errorf("Unable to map port with igdman %d: %s", port, err)
	}

	return nil
}

func unmapPort(port int) error {
	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	if err != nil {
		return fmt.Errorf("Unable to unmap port with igdman %d: %s", port, err)
	}

	return nil
}

// determineInternalIP determines the internal IP to use for mapping ports. It
// does this by dialing a website on the public Internet and then finding out
// the LocalAddr for the corresponding connection. This gives us an interface
// that we know has Internet access, which makes it suitable for port mapping.
func determineInternalIP() (string, error) {
	conn, err := net.DialTimeout("tcp", "s3.amazonaws.com:443", 20*time.Second)
	if err != nil {
		return "", fmt.Errorf("Unable to determine local IP: %s", err)
	}
	defer conn.Close()
	return strings.Split(conn.LocalAddr().String(), ":")[0], nil
}

func onBytesGiven(bytes int64) {
	dims := statreporter.CountryDim()
	dims.Increment("bytesGiven").Add(bytes)
	dims.Increment("bytesGivenByFlashlight").Add(bytes)
}
<commit_msg>Log update<commit_after>package server

import (
	"fmt"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/go-igdman/igdman"
	"github.com/getlantern/golog"
	"github.com/getlantern/nattywad"
	"github.com/getlantern/waddell"

	"github.com/getlantern/flashlight/globals"
	"github.com/getlantern/flashlight/nattest"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/statserver"
)

const (
	PortmapFailure = 50
)

var (
	log = golog.LoggerFor("flashlight.server")
)

type Server struct {
	// Addr: listen address in form of host:port
	Addr string

	// Host: FQDN that is guaranteed to hit this server
	Host string

	// ReadTimeout: (optional) timeout for read ops
	ReadTimeout time.Duration

	// WriteTimeout: (optional) timeout for write ops
	WriteTimeout time.Duration

	CertContext                *fronted.CertContext // context for certificate management
	AllowNonGlobalDestinations bool                 // if true, requests to LAN, Loopback, etc. will be allowed

	waddellClient  *waddell.Client
	nattywadServer *nattywad.Server
	cfg            *ServerConfig
	cfgMutex       sync.Mutex
}

func (server *Server) Configure(newCfg *ServerConfig) {
	server.cfgMutex.Lock()
	defer server.cfgMutex.Unlock()

	oldCfg := server.cfg

	log.Debug("Server.Configure() called")
	if oldCfg != nil && reflect.DeepEqual(oldCfg, newCfg) {
		log.Debugf("Server configuration unchanged")
		return
	}

	if oldCfg == nil || newCfg.Portmap != oldCfg.Portmap {
		// Portmap changed
		if oldCfg != nil && oldCfg.Portmap > 0 {
			log.Debugf("Attempting to unmap old external port %d", oldCfg.Portmap)
			err := unmapPort(oldCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to unmap old external port: %s", err)
			}
			log.Debugf("Unmapped old external port %d", oldCfg.Portmap)
		}

		if newCfg.Portmap > 0 {
			log.Debugf("Attempting to map new external port %d", newCfg.Portmap)
			err := mapPort(server.Addr, newCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to map new external port: %s", err)
				os.Exit(PortmapFailure)
			}
			log.Debugf("Mapped new external port %d", newCfg.Portmap)
		}
	}

	nattywadIsEnabled := newCfg.WaddellAddr != ""
	nattywadWasEnabled := server.nattywadServer != nil
	waddellAddrChanged := oldCfg == nil && newCfg.WaddellAddr != "" || oldCfg != nil && oldCfg.WaddellAddr != newCfg.WaddellAddr

	if waddellAddrChanged {
		if nattywadWasEnabled {
			server.stopNattywad()
		}
		if nattywadIsEnabled {
			server.startNattywad(newCfg.WaddellAddr)
		}
	}

	server.cfg = newCfg
}

func (server *Server) ListenAndServe() error {
	if server.Host != "" {
		log.Debugf("Running as host %s", server.Host)
	}

	fs := &fronted.Server{
		Addr:                       server.Addr,
		Host:                       server.Host,
		ReadTimeout:                server.ReadTimeout,
		WriteTimeout:               server.WriteTimeout,
		CertContext:                server.CertContext,
		AllowNonGlobalDestinations: server.AllowNonGlobalDestinations,
	}

	if server.cfg.Unencrypted {
		log.Debug("Running in unencrypted mode")
		fs.CertContext = nil
	}

	// Add callbacks to track bytes given
	fs.OnBytesReceived = func(ip string, bytes int64) {
		onBytesGiven(bytes)
		statserver.OnBytesReceived(ip, bytes)
	}
	fs.OnBytesSent = func(ip string, bytes int64) {
		onBytesGiven(bytes)
		statserver.OnBytesSent(ip, bytes)
	}

	l, err := fs.Listen()
	if err != nil {
		return fmt.Errorf("Unable to listen at %s: %s", server.Addr, err)
	}
	return fs.Serve(l)
}

func (server *Server) startNattywad(waddellAddr string) {
	log.Debugf("Connecting to waddell at: %s", waddellAddr)
	var err error
	server.waddellClient, err = waddell.NewClient(&waddell.ClientConfig{
		Dial: func() (net.Conn, error) {
			return net.Dial("tcp", waddellAddr)
		},
		ServerCert:        globals.WaddellCert,
		ReconnectAttempts: 10,
		OnId: func(id waddell.PeerId) {
			log.Debugf("Connected to Waddell!! Id is: %s", id)
		},
	})
	if err != nil {
		log.Errorf("Unable to connect to waddell: %s", err)
		server.waddellClient = nil
		return
	}
	server.nattywadServer = &nattywad.Server{
		Client: server.waddellClient,
		OnSuccess: func(local *net.UDPAddr, remote *net.UDPAddr) bool {
			err := nattest.Serve(local)
			if err != nil {
				log.Error(err.Error())
				return false
			}
			return true
		},
	}
	server.nattywadServer.Start()
}

func (server *Server) stopNattywad() {
	log.Debug("Stopping nattywad server")
	server.nattywadServer.Stop()
	server.nattywadServer = nil
	log.Debug("Stopping waddell client")
	server.waddellClient.Close()
	server.waddellClient = nil
}

func mapPort(addr string, port int) error {
	parts := strings.Split(addr, ":")

	internalPort, err := strconv.Atoi(parts[1])
	if err != nil {
		return fmt.Errorf("Unable to parse local port: ")
	}

	internalIP := parts[0]
	if internalIP == "" {
		internalIP, err = determineInternalIP()
		if err != nil {
			return fmt.Errorf("Unable to determine internal IP: %s", err)
		}
	}

	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	err = igd.AddPortMapping(igdman.TCP, internalIP, internalPort, port, 0)
	if err != nil {
		return fmt.Errorf("Unable to map port with igdman %d: %s", port, err)
	}

	return nil
}

func unmapPort(port int) error {
	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	if err != nil {
		return fmt.Errorf("Unable to unmap port with igdman %d: %s", port, err)
	}

	return nil
}

// determineInternalIP determines the internal IP to use for mapping ports. It
// does this by dialing a website on the public Internet and then finding out
// the LocalAddr for the corresponding connection. This gives us an interface
// that we know has Internet access, which makes it suitable for port mapping.
func determineInternalIP() (string, error) {
	conn, err := net.DialTimeout("tcp", "s3.amazonaws.com:443", 20*time.Second)
	if err != nil {
		return "", fmt.Errorf("Unable to determine local IP: %s", err)
	}
	defer conn.Close()
	return strings.Split(conn.LocalAddr().String(), ":")[0], nil
}

func onBytesGiven(bytes int64) {
	dims := statreporter.CountryDim()
	dims.Increment("bytesGiven").Add(bytes)
	dims.Increment("bytesGivenByFlashlight").Add(bytes)
}
<|endoftext|>


<commit_before>package multicast

import (
	"encoding/json"
	stdLog "log"
	"sync"
	"syscall"
	"testing"
	"time"
)

const verbose = false

func TestMulticast(t *testing.T) {
	mc1 := JoinMulticast(nil, nil)
	if mc1 == nil {
		t.Fatal("Unable to join multicast group")
	} else if verbose {
		stdLog.Println("Joined and listening to multicast IP", mc1.addr.IP, "on port", mc1.addr.Port)
	}

	// Enable Multicast looping for testing
	f, err := mc1.conn.File()
	err = syscall.SetsockoptInt(int(f.Fd()), syscall.IPPROTO_IP, syscall.IP_MULTICAST_LOOP, 1)
	if err != nil {
		t.Fatal("Unable to set socket for multicast looping")
	}

	var wg sync.WaitGroup

	wg.Add(1)
	// Sender node
	go func(mc *Multicast) {
		defer wg.Done()

		// Give some time to the other goroutines to build up
		time.Sleep(time.Millisecond * 200)

		msg := "Multicast Hello World!"
		n, e := mc.write(([]byte)(msg))
		if e != nil {
			t.Fatal("Unable to multicast message")
		}

		if verbose {
			stdLog.Printf("--> Sent %d bytes: %s\n", n, msg)
		}

		mc.LeaveMulticast()
		if verbose {
			stdLog.Println("Leaving multicast IP", mc.addr.IP, "on port", mc.addr.Port)
		}
	}(mc1)

	nNodes := 9
	for i := 0; i < nNodes; i++ {
		wg.Add(1)
		go receiverNode(t, &wg, i+1)
	}
	wg.Wait()
}

func receiverNode(t *testing.T, wg *sync.WaitGroup, id int) {
	defer wg.Done()

	mc := JoinMulticast(nil, nil)
	if mc == nil {
		t.Fatal("Unable to join multicast group")
	} else if verbose {
		stdLog.Println("Joined and listening to multicast IP", mc.addr.IP, "on port", mc.addr.Port)
	}

	b := make([]byte, 1000)
	n, _, e := mc.read(b)
	if e != nil {
		t.Fatal("Unable to multicast message")
	}
	if n <= 0 {
		t.Fatal("No data received in multicast messages")
	}
	if verbose {
		stdLog.Println("Node", id, "<-- Received", n, "bytes:", string(b))
	}

	mc.LeaveMulticast()
	if verbose {
		stdLog.Println("Node", id, "leaving multicast IP", mc.addr.IP, "on port", mc.addr.Port)
	}
}

func TestMulticastMessages(t *testing.T) {
	mc1 := JoinMulticast(nil, nil)
	if mc1 == nil {
		t.Fatal("Unable to join multicast group")
	}

	// Enable Multicast looping for testing
	f, err := mc1.conn.File()
	err = syscall.SetsockoptInt(int(f.Fd()), syscall.IPPROTO_IP, syscall.IP_MULTICAST_LOOP, 1)
	if err != nil {
		t.Fatal("Unable to set socket for multicast looping")
	}

	mc1.SetPayload("testHello")
	mc1.SetPeriod(1)
	mc1.StartMulticast()
	mc1.ListenPeers()

	mc2 := JoinMulticast(nil, nil)
	if mc2 == nil {
		t.Fatal("Unable to join multicast group")
	}

	b := make([]byte, messageMaxSize)
	n, _, e := mc2.read(b)
	if e != nil {
		t.Fatal("Unable to multicast message")
	}
	if n > 0 {
		var msg multicastMessage
		if e = json.Unmarshal(b[:n], &msg); e != nil || msg.Type != typeHello || msg.Payload != "testHello" {
			stdLog.Println(string(b[:n]))
			stdLog.Println(msg)
			t.Fatal("Multicast Hello message is incorrectly formatted")
		}
	} else {
		stdLog.Println("Received 0 bytes")
	}

	mc1.LeaveMulticast()
	mc2.LeaveMulticast()
}

func TestMulticastAnnouncing(t *testing.T) {
	mc1 := JoinMulticast(nil, nil)
	if mc1 == nil {
		t.Fatal("Unable to join multicast group")
	}

	// Enable Multicast looping for testing
	f, err := mc1.conn.File()
	err = syscall.SetsockoptInt(int(f.Fd()), syscall.IPPROTO_IP, syscall.IP_MULTICAST_LOOP, 1)
	if err != nil {
		t.Fatal("Unable to set socket for multicast looping")
	}

	mc1.SetPeriod(1)
	go func() {
		if e := mc1.sendHellos(); e != nil {
			log.Fatal("Error sending hellos")
		}
	}()

	mc2 := JoinMulticast(
		func(string, []PeerInfo) {
			if verbose {
				stdLog.Println("Adding Peer")
			}
		},
		func(string, []PeerInfo) {
			if verbose {
				stdLog.Println("Removing Peer")
			}
		})
	if mc2 == nil {
		t.Fatal("Unable to join multicast group")
	}

	mc2.StartMulticast()
	mc2.ListenPeers()

	time.Sleep(time.Millisecond * 1100) // Just enough to let the multicast run

	// Should be zero because we don't add ourselves to the peers map
	if len(mc2.peers) != 0 {
		stdLog.Println("Peers in MC1", mc1.peers)
		stdLog.Println("Peers in MC2", mc2.peers)
		t.Fatal("Wrong count of peers")
	}
}
<commit_msg>Possible fix for issue #2810<commit_after>package multicast

import (
	"encoding/json"
	stdLog "log"
	"sync"
	"syscall"
	"testing"
	"time"
)

const verbose = false

func TestMulticast(t *testing.T) {
	mc1 := JoinMulticast(nil, nil)
	if mc1 == nil {
		t.Fatal("Unable to join multicast group")
	} else if verbose {
		stdLog.Println("Joined and listening to multicast IP", mc1.addr.IP, "on port", mc1.addr.Port)
	}

	// Enable Multicast looping for testing
	f, err := mc1.conn.File()
	err = syscall.SetsockoptInt(int(f.Fd()), syscall.IPPROTO_IP, syscall.IP_MULTICAST_LOOP, 1)
	if err != nil {
		t.Fatal("Unable to set socket for multicast looping")
	}

	var wg sync.WaitGroup

	wg.Add(1)
	// Sender node
	go func(mc *Multicast) {
		defer wg.Done()

		// Give some time to the other goroutines to build up
		time.Sleep(time.Millisecond * 200)

		msg := "Multicast Hello World!"
		n, e := mc.write(([]byte)(msg))
		if e != nil {
			t.Fatal("Unable to multicast message")
		}

		if verbose {
			stdLog.Printf("--> Sent %d bytes: %s\n", n, msg)
		}

		mc.LeaveMulticast()
		if verbose {
			stdLog.Println("Leaving multicast IP", mc.addr.IP, "on port", mc.addr.Port)
		}
	}(mc1)

	nNodes := 9
	for i := 0; i < nNodes; i++ {
		wg.Add(1)
		go receiverNode(t, &wg, i+1)
	}
	wg.Wait()
}

func receiverNode(t *testing.T, wg *sync.WaitGroup, id int) {
	defer wg.Done()

	mc := JoinMulticast(nil, nil)
	if mc == nil {
		t.Fatal("Unable to join multicast group")
	} else if verbose {
		stdLog.Println("Joined and listening to multicast IP", mc.addr.IP, "on port", mc.addr.Port)
	}

	b := make([]byte, 1000)
	n, _, e := mc.read(b)
	if e != nil {
		t.Fatal("Unable to multicast message")
	}
	if n <= 0 {
		t.Fatal("No data received in multicast messages")
	}
	if verbose {
		stdLog.Println("Node", id, "<-- Received", n, "bytes:", string(b))
	}

	mc.LeaveMulticast()
	if verbose {
		stdLog.Println("Node", id, "leaving multicast IP", mc.addr.IP, "on port", mc.addr.Port)
	}
}

func TestMulticastMessages(t *testing.T) {
	mc1 := JoinMulticast(nil, nil)
	if mc1 == nil {
		t.Fatal("Unable to join multicast group")
	}

	// Enable Multicast looping for testing
	f, err := mc1.conn.File()
	err = syscall.SetsockoptInt(int(f.Fd()), syscall.IPPROTO_IP, syscall.IP_MULTICAST_LOOP, 1)
	if err != nil {
		t.Fatal("Unable to set socket for multicast looping")
	}

	mc1.SetPayload("testHello")
	mc1.SetPeriod(1)
	mc1.StartMulticast()
	mc1.ListenPeers()

	mc2 := JoinMulticast(nil, nil)
	if mc2 == nil {
		t.Fatal("Unable to join multicast group")
	}

	b := make([]byte, messageMaxSize)
	var msg multicastMessage
Out:
	for {
		if n, _, err := mc2.read(b); err != nil {
			t.Fatal("Error reading multicast message")
		} else {
			if err := json.Unmarshal(b[:n], &msg); err != nil {
				t.Fatal("Error unmarshalling multicast message")
			} else {
				switch msg.Type {
				case typeHello:
					if msg.Payload != "testHello" {
						stdLog.Println(string(b[:n]))
						stdLog.Println(msg)
						t.Fatal("Multicast Hello message is incorrectly formatted")
					}
					break Out
				case typeBye:
					continue
				default:
					t.Fatal("Unknown multicast message type")
				}
			}
		}
	}

	mc1.LeaveMulticast()
	mc2.LeaveMulticast()
}

func TestMulticastAnnouncing(t *testing.T) {
	mc1 := JoinMulticast(nil, nil)
	if mc1 == nil {
		t.Fatal("Unable to join multicast group")
	}

	// Enable Multicast looping for testing
	f, err := mc1.conn.File()
	err = syscall.SetsockoptInt(int(f.Fd()), syscall.IPPROTO_IP, syscall.IP_MULTICAST_LOOP, 1)
	if err != nil {
		t.Fatal("Unable to set socket for multicast looping")
	}

	mc1.SetPeriod(1)
	go func() {
		if e := mc1.sendHellos(); e != nil {
			log.Fatal("Error sending hellos")
		}
	}()

	mc2 := JoinMulticast(
		func(string, []PeerInfo) {
			if verbose {
				stdLog.Println("Adding Peer")
			}
		},
		func(string, []PeerInfo) {
			if verbose {
				stdLog.Println("Removing Peer")
			}
		})
	if mc2 == nil {
		t.Fatal("Unable to join multicast group")
	}

	mc2.StartMulticast()
	mc2.ListenPeers()

	time.Sleep(time.Millisecond * 1100) // Just enough to let the multicast run

	// Should be zero because we don't add ourselves to the peers map
	if len(mc2.peers) != 0 {
		stdLog.Println("Peers in MC1", mc1.peers)
		stdLog.Println("Peers in MC2", mc2.peers)
		t.Fatal("Wrong count of peers")
	}
}
<|endoftext|>


<commit_before>package disk

import (
	"os"
	"testing"

	"github.com/funkygao/assert"
)

func TestCursor(t *testing.T) {
	q := newQueue("cluster", "topic", "dir", -1)
	defer os.RemoveAll("dir")

	err := q.Open() // will open cursor internally
	assert.Equal(t, nil, err)
	defer q.Close()

	q.cursor.pos.Offset = 90
	q.cursor.pos.SegmentId = 5
	err = q.cursor.dump()
	assert.Equal(t, nil, err)

	q.cursor.pos = position{} // reset
	err = q.cursor.open()
	assert.Equal(t, nil, err)
	assert.Equal(t, int64(90), q.cursor.pos.Offset)
	assert.Equal(t, uint64(5), q.cursor.pos.SegmentId)
}
<commit_msg>fix unit test compile err<commit_after>package disk

import (
	"os"
	"testing"

	"github.com/funkygao/assert"
)

func TestCursor(t *testing.T) {
	ct := clusterTopic{cluster: "cluster", topic: "topic"}
	q := newQueue(ct, "dir", -1)
	defer os.RemoveAll("dir")

	err := q.Open() // will open cursor internally
	assert.Equal(t, nil, err)
	defer q.Close()

	q.cursor.pos.Offset = 90
	q.cursor.pos.SegmentId = 5
	err = q.cursor.dump()
	assert.Equal(t, nil, err)

	q.cursor.pos = position{} // reset
	err = q.cursor.open()
	assert.Equal(t, nil, err)
	assert.Equal(t, int64(90), q.cursor.pos.Offset)
	assert.Equal(t, uint64(5), q.cursor.pos.SegmentId)
}
<|endoftext|>


<commit_before>package worker

import (
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/getsentry/raven-go"
)

var (
	severityMap = map[logrus.Level]raven.Severity{
		logrus.DebugLevel: raven.DEBUG,
		logrus.InfoLevel:  raven.INFO,
		logrus.WarnLevel:  raven.WARNING,
		logrus.ErrorLevel: raven.ERROR,
		logrus.FatalLevel: raven.FATAL,
		logrus.PanicLevel: raven.FATAL,
	}
)

// SentryHook delivers logs to a sentry server
type SentryHook struct {
	Timeout time.Duration
	client  *raven.Client
	levels  []logrus.Level
}

// NewSentryHook creates a hook to be added to an instance of logger and
// initializes the raven client. This method sets the timeout to 100
// milliseconds.
func NewSentryHook(DSN string, levels []logrus.Level) (*SentryHook, error) {
	client, err := raven.New(DSN)
	if err != nil {
		return nil, err
	}
	return &SentryHook{100 * time.Millisecond, client, levels}, nil
}

// Fire is called when an event should be sent to sentry
func (hook *SentryHook) Fire(entry *logrus.Entry) error {
	defer func() {
		if r := recover(); r != nil {
			entry.Logger.WithField("panic", r).Error("paniced when trying to send log to sentry")
		}
	}()

	packet := &raven.Packet{
		Message:   entry.Message,
		Timestamp: raven.Timestamp(entry.Time),
		Level:     severityMap[entry.Level],
		Platform:  "go",
	}

	if serverName, ok := entry.Data["server_name"]; ok {
		packet.ServerName = serverName.(string)
		delete(entry.Data, "server_name")
	}
	packet.Extra = map[string]interface{}(entry.Data)
	packet.Interfaces = append(packet.Interfaces, raven.NewStacktrace(4, 3, []string{"github.com/travis-ci/worker"}))

	_, errCh := hook.client.Capture(packet, nil)
	if hook.Timeout != 0 {
		timeoutCh := time.After(hook.Timeout)
		select {
		case err := <-errCh:
			return err
		case <-timeoutCh:
			return fmt.Errorf("no response from sentry server in %s", hook.Timeout)
		}
	}
	return nil
}

// Levels returns the available logging levels.
func (hook *SentryHook) Levels() []logrus.Level {
	return hook.levels
}
<commit_msg>sentry: force err field to be a string if it's an error<commit_after>package worker

import (
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/getsentry/raven-go"
)

var (
	severityMap = map[logrus.Level]raven.Severity{
		logrus.DebugLevel: raven.DEBUG,
		logrus.InfoLevel:  raven.INFO,
		logrus.WarnLevel:  raven.WARNING,
		logrus.ErrorLevel: raven.ERROR,
		logrus.FatalLevel: raven.FATAL,
		logrus.PanicLevel: raven.FATAL,
	}
)

// SentryHook delivers logs to a sentry server
type SentryHook struct {
	Timeout time.Duration
	client  *raven.Client
	levels  []logrus.Level
}

// NewSentryHook creates a hook to be added to an instance of logger and
// initializes the raven client. This method sets the timeout to 100
// milliseconds.
func NewSentryHook(DSN string, levels []logrus.Level) (*SentryHook, error) {
	client, err := raven.New(DSN)
	if err != nil {
		return nil, err
	}
	return &SentryHook{100 * time.Millisecond, client, levels}, nil
}

// Fire is called when an event should be sent to sentry
func (hook *SentryHook) Fire(entry *logrus.Entry) error {
	defer func() {
		if r := recover(); r != nil {
			entry.Logger.WithField("panic", r).Error("paniced when trying to send log to sentry")
		}
	}()

	packet := &raven.Packet{
		Message:   entry.Message,
		Timestamp: raven.Timestamp(entry.Time),
		Level:     severityMap[entry.Level],
		Platform:  "go",
	}

	if serverName, ok := entry.Data["server_name"]; ok {
		packet.ServerName = serverName.(string)
		delete(entry.Data, "server_name")
	}
	packet.Extra = map[string]interface{}(entry.Data)

	if errMaybe, ok := packet.Extra["err"]; ok {
		if err, ok := errMaybe.(error); ok {
			packet.Extra["err"] = err.Error()
		}
	}

	packet.Interfaces = append(packet.Interfaces, raven.NewStacktrace(4, 3, []string{"github.com/travis-ci/worker"}))

	_, errCh := hook.client.Capture(packet, nil)
	if hook.Timeout != 0 {
		timeoutCh := time.After(hook.Timeout)
		select {
		case err := <-errCh:
			return err
		case <-timeoutCh:
			return fmt.Errorf("no response from sentry server in %s", hook.Timeout)
		}
	}
	return nil
}

// Levels returns the available logging levels.
func (hook *SentryHook) Levels() []logrus.Level {
	return hook.levels
}
<|endoftext|>


<commit_before>package tsdb

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/influxdb/models"
)

func NewStore(path string) *Store {
	opts := NewEngineOptions()
	opts.Config = NewConfig()

	return &Store{
		path:          path,
		EngineOptions: opts,
		Logger:        log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

var (
	ErrShardNotFound = fmt.Errorf("shard not found")
)

const (
	MaintenanceCheckInterval = time.Minute
)

type Store struct {
	mu   sync.RWMutex
	path string

	databaseIndexes map[string]*DatabaseIndex
	shards          map[uint64]*Shard

	EngineOptions EngineOptions
	Logger        *log.Logger

	closing chan struct{}
	wg      sync.WaitGroup
}

// Path returns the store's root path.
func (s *Store) Path() string { return s.path }

// DatabaseIndexN returns the number of databases indicies in the store.
func (s *Store) DatabaseIndexN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.databaseIndexes)
}

// Shard returns a shard by id.
func (s *Store) Shard(id uint64) *Shard {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.shards[id]
}

// ShardN returns the number of shard in the store.
func (s *Store) ShardN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.shards)
}

func (s *Store) CreateShard(database, retentionPolicy string, shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.closing:
		return fmt.Errorf("closing")
	default:
	}

	// shard already exists
	if _, ok := s.shards[shardID]; ok {
		return nil
	}

	// created the db and retention policy dirs if they don't exist
	if err := os.MkdirAll(filepath.Join(s.path, database, retentionPolicy), 0700); err != nil {
		return err
	}

	// create the WAL directory
	walPath := filepath.Join(s.EngineOptions.Config.WALDir, database, retentionPolicy, fmt.Sprintf("%d", shardID))
	if err := os.MkdirAll(walPath, 0700); err != nil {
		return err
	}

	// create the database index if it does not exist
	db, ok := s.databaseIndexes[database]
	if !ok {
		db = NewDatabaseIndex()
		s.databaseIndexes[database] = db
	}

	shardPath := filepath.Join(s.path, database, retentionPolicy, strconv.FormatUint(shardID, 10))
	shard := NewShard(shardID, db, shardPath, walPath, s.EngineOptions)
	if err := shard.Open(); err != nil {
		return err
	}

	s.shards[shardID] = shard

	return nil
}

// DeleteShard removes a shard from disk.
func (s *Store) DeleteShard(shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// ensure shard exists
	sh, ok := s.shards[shardID]
	if !ok {
		return nil
	}

	if err := sh.Close(); err != nil {
		return err
	}

	if err := os.Remove(sh.path); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.walPath); err != nil {
		return err
	}

	delete(s.shards, shardID)

	return nil
}

// DeleteDatabase will close all shards associated with a database and remove the directory and files from disk.
func (s *Store) DeleteDatabase(name string, shardIDs []uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, id := range shardIDs {
		shard := s.shards[id]
		if shard != nil {
			shard.Close()
		}
	}
	if err := os.RemoveAll(filepath.Join(s.path, name)); err != nil {
		return err
	}
	if err := os.RemoveAll(filepath.Join(s.EngineOptions.Config.WALDir, name)); err != nil {
		return err
	}
	delete(s.databaseIndexes, name)
	return nil
}

// ShardIDs returns a slice of all ShardIDs under management.
func (s *Store) ShardIDs() []uint64 {
	ids := make([]uint64, 0, len(s.shards))
	for i, _ := range s.shards {
		ids = append(ids, i)
	}
	return ids
}

func (s *Store) ValidateAggregateFieldsInStatement(shardID uint64, measurementName string, stmt *influxql.SelectStatement) error {
	s.mu.RLock()
	shard := s.shards[shardID]
	s.mu.RUnlock()
	if shard == nil {
		return ErrShardNotFound
	}
	return shard.ValidateAggregateFieldsInStatement(measurementName, stmt)
}

func (s *Store) DatabaseIndex(name string) *DatabaseIndex {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.databaseIndexes[name]
}

// Databases returns all the databases in the indexes
func (s *Store) Databases() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	databases := []string{}
	for db := range s.databaseIndexes {
		databases = append(databases, db)
	}
	return databases
}

func (s *Store) Measurement(database, name string) *Measurement {
	s.mu.RLock()
	db := s.databaseIndexes[database]
	s.mu.RUnlock()
	if db == nil {
		return nil
	}
	return db.Measurement(name)
}

// DiskSize returns the size of all the shard files in bytes.  This size does not include the WAL size.
func (s *Store) DiskSize() (int64, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	var size int64
	for _, shardID := range s.ShardIDs() {
		shard := s.Shard(shardID)
		sz, err := shard.DiskSize()
		if err != nil {
			return 0, err
		}
		size += sz
	}
	return size, nil
}

// deleteSeries loops through the local shards and deletes the series data and metadata for the passed in series keys
func (s *Store) deleteSeries(keys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteSeries(keys); err != nil {
			return err
		}
	}
	return nil
}

// deleteMeasurement loops through the local shards and removes the measurement field encodings from each shard
func (s *Store) deleteMeasurement(name string, seriesKeys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteMeasurement(name, seriesKeys); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) loadIndexes() error {
	dbs, err := ioutil.ReadDir(s.path)
	if err != nil {
		return err
	}
	for _, db := range dbs {
		if !db.IsDir() {
			s.Logger.Printf("Skipping database dir: %s. Not a directory", db.Name())
			continue
		}
		s.databaseIndexes[db.Name()] = NewDatabaseIndex()
	}
	return nil
}

func (s *Store) loadShards() error {
	// loop through the current database indexes
	for db := range s.databaseIndexes {
		rps, err := ioutil.ReadDir(filepath.Join(s.path, db))
		if err != nil {
			return err
		}

		for _, rp := range rps {
			// retention policies should be directories.  Skip anything that is not a dir.
			if !rp.IsDir() {
				s.Logger.Printf("Skipping retention policy dir: %s. Not a directory", rp.Name())
				continue
			}

			shards, err := ioutil.ReadDir(filepath.Join(s.path, db, rp.Name()))
			if err != nil {
				return err
			}
			for _, sh := range shards {
				path := filepath.Join(s.path, db, rp.Name(), sh.Name())
				walPath := filepath.Join(s.EngineOptions.Config.WALDir, db, rp.Name(), sh.Name())

				// Shard file names are numeric shardIDs
				shardID, err := strconv.ParseUint(sh.Name(), 10, 64)
				if err != nil {
					s.Logger.Printf("Skipping shard: %s. Not a valid path", rp.Name())
					continue
				}

				shard := NewShard(shardID, s.databaseIndexes[db], path, walPath, s.EngineOptions)
				err = shard.Open()
				if err != nil {
					return fmt.Errorf("failed to open shard %d: %s", shardID, err)
				}
				s.shards[shardID] = shard
			}
		}
	}
	return nil

}

// periodicMaintenance is the method called in a goroutine on the opening of the store
// to perform periodic maintenance of the shards.
func (s *Store) periodicMaintenance() {
	t := time.NewTicker(MaintenanceCheckInterval)
	for {
		select {
		case <-t.C:
			s.performMaintenance()
		case <-s.closing:
			t.Stop()
			return
		}
	}
}

// performMaintenance will loop through the shars and tell them
// to perform any maintenance tasks. Those tasks should kick off
// their own goroutines if it's anything that could take time.
func (s *Store) performMaintenance() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, sh := range s.shards {
		s.performMaintenanceOnShard(sh)
	}
}

func (s *Store) performMaintenanceOnShard(shard *Shard) {
	defer func() {
		if r := recover(); r != nil {
			s.Logger.Printf("recovered eror in maintenance on shard %d", shard.id)
		}
	}()
	shard.PerformMaintenance()
}

func (s *Store) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.closing = make(chan struct{})

	s.shards = map[uint64]*Shard{}
	s.databaseIndexes = map[string]*DatabaseIndex{}

	s.Logger.Printf("Using data dir: %v", s.Path())

	// Create directory.
	if err := os.MkdirAll(s.path, 0777); err != nil {
		return err
	}

	// TODO: Start AE for Node
	if err := s.loadIndexes(); err != nil {
		return err
	}

	if err := s.loadShards(); err != nil {
		return err
	}

	go s.periodicMaintenance()

	return nil
}

func (s *Store) WriteToShard(shardID uint64, points []models.Point) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sh, ok := s.shards[shardID]
	if !ok {
		return ErrShardNotFound
	}

	return sh.WritePoints(points)
}

func (s *Store) CreateMapper(shardID uint64, stmt influxql.Statement, chunkSize int) (Mapper, error) {
	shard := s.Shard(shardID)

	switch stmt := stmt.(type) {
	case *influxql.SelectStatement:
		if (stmt.IsRawQuery && !stmt.HasDistinct()) || stmt.IsSimpleDerivative() {
			m := NewRawMapper(shard, stmt)
			m.ChunkSize = chunkSize
			return m, nil
		}
		return NewAggregateMapper(shard, stmt), nil

	case *influxql.ShowMeasurementsStatement:
		m := NewShowMeasurementsMapper(shard, stmt)
		m.ChunkSize = chunkSize
		return m, nil
	case *influxql.ShowTagKeysStatement:
		return NewShowTagKeysMapper(shard, stmt, chunkSize), nil
	default:
		return nil, fmt.Errorf("can't create mapper for statement type: %T", stmt)
	}
}

func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closing != nil {
		close(s.closing)
		s.closing = nil
	}
	s.wg.Wait()

	for _, sh := range s.shards {
		if err := sh.Close(); err != nil {
			return err
		}
	}
	if s.closing != nil {
		close(s.closing)
	}
	s.closing = nil
	s.shards = nil
	s.databaseIndexes = nil

	return nil
}

// IsRetryable returns true if this error is temporary and could be retried
func IsRetryable(err error) bool {
	if err == nil {
		return true
	}

	if strings.Contains(err.Error(), "field type conflict") {
		return false
	}
	return true
}
<commit_msg>Fix TestStoreOpenShardCreateDelete<commit_after>package tsdb

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/influxdb/models"
)

func NewStore(path string) *Store {
	opts := NewEngineOptions()
	opts.Config = NewConfig()

	return &Store{
		path:          path,
		EngineOptions: opts,
		Logger:        log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

var (
	ErrShardNotFound = fmt.Errorf("shard not found")
)

const (
	MaintenanceCheckInterval = time.Minute
)

type Store struct {
	mu   sync.RWMutex
	path string

	databaseIndexes map[string]*DatabaseIndex
	shards          map[uint64]*Shard

	EngineOptions EngineOptions
	Logger        *log.Logger

	closing chan struct{}
	wg      sync.WaitGroup
}

// Path returns the store's root path.
func (s *Store) Path() string { return s.path }

// DatabaseIndexN returns the number of databases indicies in the store.
func (s *Store) DatabaseIndexN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.databaseIndexes)
}

// Shard returns a shard by id.
func (s *Store) Shard(id uint64) *Shard {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.shards[id]
}

// ShardN returns the number of shard in the store.
func (s *Store) ShardN() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.shards)
}

func (s *Store) CreateShard(database, retentionPolicy string, shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.closing:
		return fmt.Errorf("closing")
	default:
	}

	// shard already exists
	if _, ok := s.shards[shardID]; ok {
		return nil
	}

	// created the db and retention policy dirs if they don't exist
	if err := os.MkdirAll(filepath.Join(s.path, database, retentionPolicy), 0700); err != nil {
		return err
	}

	// create the WAL directory
	walPath := filepath.Join(s.EngineOptions.Config.WALDir, database, retentionPolicy, fmt.Sprintf("%d", shardID))
	if err := os.MkdirAll(walPath, 0700); err != nil {
		return err
	}

	// create the database index if it does not exist
	db, ok := s.databaseIndexes[database]
	if !ok {
		db = NewDatabaseIndex()
		s.databaseIndexes[database] = db
	}

	shardPath := filepath.Join(s.path, database, retentionPolicy, strconv.FormatUint(shardID, 10))
	shard := NewShard(shardID, db, shardPath, walPath, s.EngineOptions)
	if err := shard.Open(); err != nil {
		return err
	}

	s.shards[shardID] = shard

	return nil
}

// DeleteShard removes a shard from disk.
func (s *Store) DeleteShard(shardID uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// ensure shard exists
	sh, ok := s.shards[shardID]
	if !ok {
		return nil
	}

	if err := sh.Close(); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.path); err != nil {
		return err
	}

	if err := os.RemoveAll(sh.walPath); err != nil {
		return err
	}

	delete(s.shards, shardID)

	return nil
}

// DeleteDatabase will close all shards associated with a database and remove the directory and files from disk.
func (s *Store) DeleteDatabase(name string, shardIDs []uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, id := range shardIDs {
		shard := s.shards[id]
		if shard != nil {
			shard.Close()
		}
	}
	if err := os.RemoveAll(filepath.Join(s.path, name)); err != nil {
		return err
	}
	if err := os.RemoveAll(filepath.Join(s.EngineOptions.Config.WALDir, name)); err != nil {
		return err
	}
	delete(s.databaseIndexes, name)
	return nil
}

// ShardIDs returns a slice of all ShardIDs under management.
func (s *Store) ShardIDs() []uint64 {
	ids := make([]uint64, 0, len(s.shards))
	for i, _ := range s.shards {
		ids = append(ids, i)
	}
	return ids
}

func (s *Store) ValidateAggregateFieldsInStatement(shardID uint64, measurementName string, stmt *influxql.SelectStatement) error {
	s.mu.RLock()
	shard := s.shards[shardID]
	s.mu.RUnlock()
	if shard == nil {
		return ErrShardNotFound
	}
	return shard.ValidateAggregateFieldsInStatement(measurementName, stmt)
}

func (s *Store) DatabaseIndex(name string) *DatabaseIndex {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.databaseIndexes[name]
}

// Databases returns all the databases in the indexes
func (s *Store) Databases() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	databases := []string{}
	for db := range s.databaseIndexes {
		databases = append(databases, db)
	}
	return databases
}

func (s *Store) Measurement(database, name string) *Measurement {
	s.mu.RLock()
	db := s.databaseIndexes[database]
	s.mu.RUnlock()
	if db == nil {
		return nil
	}
	return db.Measurement(name)
}

// DiskSize returns the size of all the shard files in bytes.  This size does not include the WAL size.
func (s *Store) DiskSize() (int64, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	var size int64
	for _, shardID := range s.ShardIDs() {
		shard := s.Shard(shardID)
		sz, err := shard.DiskSize()
		if err != nil {
			return 0, err
		}
		size += sz
	}
	return size, nil
}

// deleteSeries loops through the local shards and deletes the series data and metadata for the passed in series keys
func (s *Store) deleteSeries(keys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteSeries(keys); err != nil {
			return err
		}
	}
	return nil
}

// deleteMeasurement loops through the local shards and removes the measurement field encodings from each shard
func (s *Store) deleteMeasurement(name string, seriesKeys []string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sh := range s.shards {
		if err := sh.DeleteMeasurement(name, seriesKeys); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) loadIndexes() error {
	dbs, err := ioutil.ReadDir(s.path)
	if err != nil {
		return err
	}
	for _, db := range dbs {
		if !db.IsDir() {
			s.Logger.Printf("Skipping database dir: %s. Not a directory", db.Name())
			continue
		}
		s.databaseIndexes[db.Name()] = NewDatabaseIndex()
	}
	return nil
}

func (s *Store) loadShards() error {
	// loop through the current database indexes
	for db := range s.databaseIndexes {
		rps, err := ioutil.ReadDir(filepath.Join(s.path, db))
		if err != nil {
			return err
		}

		for _, rp := range rps {
			// retention policies should be directories.  Skip anything that is not a dir.
			if !rp.IsDir() {
				s.Logger.Printf("Skipping retention policy dir: %s. Not a directory", rp.Name())
				continue
			}

			shards, err := ioutil.ReadDir(filepath.Join(s.path, db, rp.Name()))
			if err != nil {
				return err
			}
			for _, sh := range shards {
				path := filepath.Join(s.path, db, rp.Name(), sh.Name())
				walPath := filepath.Join(s.EngineOptions.Config.WALDir, db, rp.Name(), sh.Name())

				// Shard file names are numeric shardIDs
				shardID, err := strconv.ParseUint(sh.Name(), 10, 64)
				if err != nil {
					s.Logger.Printf("Skipping shard: %s. Not a valid path", rp.Name())
					continue
				}

				shard := NewShard(shardID, s.databaseIndexes[db], path, walPath, s.EngineOptions)
				err = shard.Open()
				if err != nil {
					return fmt.Errorf("failed to open shard %d: %s", shardID, err)
				}
				s.shards[shardID] = shard
			}
		}
	}
	return nil

}

// periodicMaintenance is the method called in a goroutine on the opening of the store
// to perform periodic maintenance of the shards.
func (s *Store) periodicMaintenance() {
	t := time.NewTicker(MaintenanceCheckInterval)
	for {
		select {
		case <-t.C:
			s.performMaintenance()
		case <-s.closing:
			t.Stop()
			return
		}
	}
}

// performMaintenance will loop through the shars and tell them
// to perform any maintenance tasks. Those tasks should kick off
// their own goroutines if it's anything that could take time.
func (s *Store) performMaintenance() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, sh := range s.shards {
		s.performMaintenanceOnShard(sh)
	}
}

func (s *Store) performMaintenanceOnShard(shard *Shard) {
	defer func() {
		if r := recover(); r != nil {
			s.Logger.Printf("recovered eror in maintenance on shard %d", shard.id)
		}
	}()
	shard.PerformMaintenance()
}

func (s *Store) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.closing = make(chan struct{})

	s.shards = map[uint64]*Shard{}
	s.databaseIndexes = map[string]*DatabaseIndex{}

	s.Logger.Printf("Using data dir: %v", s.Path())

	// Create directory.
	if err := os.MkdirAll(s.path, 0777); err != nil {
		return err
	}

	// TODO: Start AE for Node
	if err := s.loadIndexes(); err != nil {
		return err
	}

	if err := s.loadShards(); err != nil {
		return err
	}

	go s.periodicMaintenance()

	return nil
}

func (s *Store) WriteToShard(shardID uint64, points []models.Point) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sh, ok := s.shards[shardID]
	if !ok {
		return ErrShardNotFound
	}

	return sh.WritePoints(points)
}

func (s *Store) CreateMapper(shardID uint64, stmt influxql.Statement, chunkSize int) (Mapper, error) {
	shard := s.Shard(shardID)

	switch stmt := stmt.(type) {
	case *influxql.SelectStatement:
		if (stmt.IsRawQuery && !stmt.HasDistinct()) || stmt.IsSimpleDerivative() {
			m := NewRawMapper(shard, stmt)
			m.ChunkSize = chunkSize
			return m, nil
		}
		return NewAggregateMapper(shard, stmt), nil

	case *influxql.ShowMeasurementsStatement:
		m := NewShowMeasurementsMapper(shard, stmt)
		m.ChunkSize = chunkSize
		return m, nil
	case *influxql.ShowTagKeysStatement:
		return NewShowTagKeysMapper(shard, stmt, chunkSize), nil
	default:
		return nil, fmt.Errorf("can't create mapper for statement type: %T", stmt)
	}
}

func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closing != nil {
		close(s.closing)
		s.closing = nil
	}
	s.wg.Wait()

	for _, sh := range s.shards {
		if err := sh.Close(); err != nil {
			return err
		}
	}
	if s.closing != nil {
		close(s.closing)
	}
	s.closing = nil
	s.shards = nil
	s.databaseIndexes = nil

	return nil
}

// IsRetryable returns true if this error is temporary and could be retried
func IsRetryable(err error) bool {
	if err == nil {
		return true
	}

	if strings.Contains(err.Error(), "field type conflict") {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package needle

import (
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"path"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func parseMultipart(r *http.Request, sizeLimit int64) (
	fileName string, data []byte, mimeType string, isGzipped bool, originalDataSize int, isChunkedFile bool, e error) {
	defer func() {
		if e != nil && r.Body != nil {
			io.Copy(ioutil.Discard, r.Body)
			r.Body.Close()
		}
	}()
	form, fe := r.MultipartReader()
	if fe != nil {
		glog.V(0).Infoln("MultipartReader [ERROR]", fe)
		e = fe
		return
	}

	//first multi-part item
	part, fe := form.NextPart()
	if fe != nil {
		glog.V(0).Infoln("Reading Multi part [ERROR]", fe)
		e = fe
		return
	}

	fileName = part.FileName()
	if fileName != "" {
		fileName = path.Base(fileName)
	}

	println("reading part", sizeLimit)

	data, e = ioutil.ReadAll(io.LimitReader(part, sizeLimit+1))
	if e != nil {
		glog.V(0).Infoln("Reading Content [ERROR]", e)
		return
	}
	if len(data) == int(sizeLimit)+1 {
		e = fmt.Errorf("file over the limited %d bytes", sizeLimit)
		return
	}

	//if the filename is empty string, do a search on the other multi-part items
	for fileName == "" {
		part2, fe := form.NextPart()
		if fe != nil {
			break // no more or on error, just safely break
		}

		fName := part2.FileName()

		//found the first <file type> multi-part has filename
		if fName != "" {
			data2, fe2 := ioutil.ReadAll(io.LimitReader(part2, sizeLimit+1))
			if fe2 != nil {
				glog.V(0).Infoln("Reading Content [ERROR]", fe2)
				e = fe2
				return
			}
			if len(data) == int(sizeLimit)+1 {
				e = fmt.Errorf("file over the limited %d bytes", sizeLimit)
				return
			}

			//update
			data = data2
			fileName = path.Base(fName)
			break
		}
	}

	originalDataSize = len(data)

	isChunkedFile, _ = strconv.ParseBool(r.FormValue("cm"))

	if !isChunkedFile {

		dotIndex := strings.LastIndex(fileName, ".")
		ext, mtype := "", ""
		if dotIndex > 0 {
			ext = strings.ToLower(fileName[dotIndex:])
			mtype = mime.TypeByExtension(ext)
		}
		contentType := part.Header.Get("Content-Type")
		if contentType != "" && mtype != contentType {
			mimeType = contentType //only return mime type if not deductable
			mtype = contentType
		}

		if part.Header.Get("Content-Encoding") == "gzip" {
			if unzipped, e := util.UnGzipData(data); e == nil {
				originalDataSize = len(unzipped)
			}
			isGzipped = true
		} else if util.IsGzippable(ext, mtype, data) {
			if compressedData, err := util.GzipData(data); err == nil {
				if len(data) > len(compressedData) {
					data = compressedData
					isGzipped = true
				}
			}
		}
	}

	return
}
<commit_msg>remove println<commit_after>package needle

import (
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"path"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func parseMultipart(r *http.Request, sizeLimit int64) (
	fileName string, data []byte, mimeType string, isGzipped bool, originalDataSize int, isChunkedFile bool, e error) {
	defer func() {
		if e != nil && r.Body != nil {
			io.Copy(ioutil.Discard, r.Body)
			r.Body.Close()
		}
	}()
	form, fe := r.MultipartReader()
	if fe != nil {
		glog.V(0).Infoln("MultipartReader [ERROR]", fe)
		e = fe
		return
	}

	//first multi-part item
	part, fe := form.NextPart()
	if fe != nil {
		glog.V(0).Infoln("Reading Multi part [ERROR]", fe)
		e = fe
		return
	}

	fileName = part.FileName()
	if fileName != "" {
		fileName = path.Base(fileName)
	}

	data, e = ioutil.ReadAll(io.LimitReader(part, sizeLimit+1))
	if e != nil {
		glog.V(0).Infoln("Reading Content [ERROR]", e)
		return
	}
	if len(data) == int(sizeLimit)+1 {
		e = fmt.Errorf("file over the limited %d bytes", sizeLimit)
		return
	}

	//if the filename is empty string, do a search on the other multi-part items
	for fileName == "" {
		part2, fe := form.NextPart()
		if fe != nil {
			break // no more or on error, just safely break
		}

		fName := part2.FileName()

		//found the first <file type> multi-part has filename
		if fName != "" {
			data2, fe2 := ioutil.ReadAll(io.LimitReader(part2, sizeLimit+1))
			if fe2 != nil {
				glog.V(0).Infoln("Reading Content [ERROR]", fe2)
				e = fe2
				return
			}
			if len(data) == int(sizeLimit)+1 {
				e = fmt.Errorf("file over the limited %d bytes", sizeLimit)
				return
			}

			//update
			data = data2
			fileName = path.Base(fName)
			break
		}
	}

	originalDataSize = len(data)

	isChunkedFile, _ = strconv.ParseBool(r.FormValue("cm"))

	if !isChunkedFile {

		dotIndex := strings.LastIndex(fileName, ".")
		ext, mtype := "", ""
		if dotIndex > 0 {
			ext = strings.ToLower(fileName[dotIndex:])
			mtype = mime.TypeByExtension(ext)
		}
		contentType := part.Header.Get("Content-Type")
		if contentType != "" && mtype != contentType {
			mimeType = contentType //only return mime type if not deductable
			mtype = contentType
		}

		if part.Header.Get("Content-Encoding") == "gzip" {
			if unzipped, e := util.UnGzipData(data); e == nil {
				originalDataSize = len(unzipped)
			}
			isGzipped = true
		} else if util.IsGzippable(ext, mtype, data) {
			if compressedData, err := util.GzipData(data); err == nil {
				if len(data) > len(compressedData) {
					data = compressedData
					isGzipped = true
				}
			}
		}
	}

	return
}
<|endoftext|>