<commit_before>package config

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

	"github.com/getlantern/appdir"
	"github.com/getlantern/fronted"
	"github.com/getlantern/geolookup"
	"github.com/getlantern/golog"
	"github.com/getlantern/proxiedsites"
	"github.com/getlantern/waitforserver"
	"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
	CpuProfile    string
	MemProfile    string
	UIAddr        string // UI HTTP server address
	AutoReport    *bool  // Report anonymous usage to GA
	Stats         *statreporter.Config
	Server        *server.ServerConfig
	Client        *client.ClientConfig
	ProxiedSites  *proxiedsites.Config // List of proxied site domains that get routed through Lantern rather than accessed directly
	TrustedCAs    []*CA
}

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

// Init initializes the configuration system.
func Init() (*Config, error) {
	configPath, err := InConfigDir("lantern.yaml")
	if err != nil {
		return nil, err
	}
	m = &yamlconf.Manager{
		FilePath:         configPath,
		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) {
			// By default, do nothing
			mutate = func(ycfg yamlconf.Config) error {
				// do nothing
				return nil
			}
			cfg := currentCfg.(*Config)
			waitTime = cfg.cloudPollSleepTime()
			if cfg.CloudConfig == "" {
				// Config doesn't have a CloudConfig, just ignore
				return
			}

			var bytes []byte
			bytes, err = cfg.fetchCloudConfig()
			if err == nil && bytes != 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)
		err = updateGlobals(cfg)
		if err != nil {
			return nil, err
		}
	}
	return cfg, err
}

// Run runs the configuration system.
func Run(updateHandler func(updated *Config)) error {
	for {
		next := m.Next()
		nextCfg := next.(*Config)
		err := updateGlobals(nextCfg)
		if err != nil {
			return err
		}
		updateHandler(nextCfg)
	}
}

func updateGlobals(cfg *Config) error {
	globals.InstanceId = cfg.InstanceId
	loc := &geolookup.City{}
	loc.Country.IsoCode = cfg.Country
	globals.SetLocation(loc)
	err := globals.SetTrustedCAs(cfg.TrustedCACerts())
	if err != nil {
		return fmt.Errorf("Unable to configure trusted CAs: %s", err)
	}
	return nil
}

// 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, error) {
	cdir := *configdir
	if cdir == "" {
		cdir = appdir.General("Lantern")
	}
	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 {
				return "", fmt.Errorf("Unable to create configdir at %s: %s", cdir, err)
			}
		}
	}
	return filepath.Join(cdir, filename), nil
}

// 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"
	}

	if cfg.UIAddr == "" {
		cfg.UIAddr = "localhost:16823"
	}

	if cfg.CloudConfig == "" {
		cfg.CloudConfig = "https://s3.amazonaws.com/lantern_config/cloud.2.0.0-nl.yaml.gz"
	}

	// 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.ProxiedSites == nil {
		log.Debugf("Adding empty proxiedsites")
		cfg.ProxiedSites = &proxiedsites.Config{
			Delta: &proxiedsites.Delta{
				Additions: []string{},
				Deletions: []string{},
			},
			Cloud: []string{},
		}
	}

	if cfg.ProxiedSites.Cloud == nil || len(cfg.ProxiedSites.Cloud) == 0 {
		log.Debugf("Loading default cloud proxiedsites")
		cfg.ProxiedSites.Cloud = defaultProxiedSites
	}

	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 = []*client.FrontedServerInfo{
			&client.FrontedServerInfo{
				Host:           "nl.fallbacks.getiantem.org",
				Port:           443,
				PoolSize:       30,
				MasqueradeSet:  cloudflare,
				MaxMasquerades: 20,
				QOS:            10,
				Weight:         4000,
			},
		}

		cfg.Client.ChainedServers = make(map[string]*client.ChainedServerInfo, len(fallbacks))
		for key, fb := range fallbacks {
			cfg.Client.ChainedServers[key] = fb
		}
	}

	if cfg.AutoReport == nil {
		cfg.AutoReport = new(bool)
		*cfg.AutoReport = true
	}

	// 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)
	}

	// 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() (bytes []byte, err error) {
	log.Debugf("Fetching cloud config from: %s", cfg.CloudConfig)

	if cfg.IsDownstream() {
		// Clients must always proxy the request
		if cfg.Addr == "" {
			err = fmt.Errorf("No proxyAddr")
		} else {
			bytes, err = cfg.doFetchCloudConfig(cfg.Addr)
		}
	} else {
		bytes, err = cfg.doFetchCloudConfig("")
	}
	if err != nil {
		bytes = nil
		err = fmt.Errorf("Unable to read yaml from %s: %s", cfg.CloudConfig, err)
	}
	return
}

func (cfg Config) doFetchCloudConfig(proxyAddr string) ([]byte, error) {
	log.Tracef("doFetchCloudConfig via '%s'", proxyAddr)

	if proxyAddr != "" {
		// Wait for proxy to become available
		err := waitforserver.WaitForServer("tcp", proxyAddr, 30*time.Second)
		if err != nil {
			return nil, fmt.Errorf("Proxy never came up at %v: %v", proxyAddr, err)
		}
	}

	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)
	}

	// Same with global proxiedsites
	if len(updated.ProxiedSites.Cloud) > 0 {
		wlDomains := make(map[string]bool)
		for _, domain := range updated.ProxiedSites.Cloud {
			wlDomains[domain] = true
		}
		updated.ProxiedSites.Cloud = make([]string, 0, len(wlDomains))
		for domain, _ := range wlDomains {
			updated.ProxiedSites.Cloud = append(updated.ProxiedSites.Cloud, domain)
		}
		sort.Strings(updated.ProxiedSites.Cloud)
	}
	return nil
}
<commit_msg>add capability for country-specific cloud configs<commit_after>package config

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

	"github.com/getlantern/appdir"
	"github.com/getlantern/fronted"
	"github.com/getlantern/geolookup"
	"github.com/getlantern/golog"
	"github.com/getlantern/proxiedsites"
	"github.com/getlantern/waitforserver"
	"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"
	countryPlaceholder = "${COUNTRY}"
)

var (
	log                 = golog.LoggerFor("flashlight.config")
	m                   *yamlconf.Manager
	lastCloudConfigETag = map[string]string{}
)

type Config struct {
	Version       int
	CloudConfig   string
	CloudConfigCA string
	Addr          string
	Role          string
	InstanceId    string
	Country       string
	CpuProfile    string
	MemProfile    string
	UIAddr        string // UI HTTP server address
	AutoReport    *bool  // Report anonymous usage to GA
	Stats         *statreporter.Config
	Server        *server.ServerConfig
	Client        *client.ClientConfig
	ProxiedSites  *proxiedsites.Config // List of proxied site domains that get routed through Lantern rather than accessed directly
	TrustedCAs    []*CA
}

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

// Init initializes the configuration system.
func Init() (*Config, error) {
	configPath, err := InConfigDir("lantern.yaml")
	if err != nil {
		return nil, err
	}
	m = &yamlconf.Manager{
		FilePath:         configPath,
		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) {
			// By default, do nothing
			mutate = func(ycfg yamlconf.Config) error {
				// do nothing
				return nil
			}
			cfg := currentCfg.(*Config)
			waitTime = cfg.cloudPollSleepTime()
			if cfg.CloudConfig == "" {
				// Config doesn't have a CloudConfig, just ignore
				return
			}

			var bytes []byte
			bytes, err = cfg.fetchCloudConfig()
			if err == nil && bytes != 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)
		err = updateGlobals(cfg)
		if err != nil {
			return nil, err
		}
	}
	return cfg, err
}

// Run runs the configuration system.
func Run(updateHandler func(updated *Config)) error {
	for {
		next := m.Next()
		nextCfg := next.(*Config)
		err := updateGlobals(nextCfg)
		if err != nil {
			return err
		}
		updateHandler(nextCfg)
	}
}

func updateGlobals(cfg *Config) error {
	globals.InstanceId = cfg.InstanceId
	loc := &geolookup.City{}
	loc.Country.IsoCode = cfg.Country
	globals.SetLocation(loc)
	err := globals.SetTrustedCAs(cfg.TrustedCACerts())
	if err != nil {
		return fmt.Errorf("Unable to configure trusted CAs: %s", err)
	}
	return nil
}

// 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, error) {
	cdir := *configdir
	if cdir == "" {
		cdir = appdir.General("Lantern")
	}
	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 {
				return "", fmt.Errorf("Unable to create configdir at %s: %s", cdir, err)
			}
		}
	}
	return filepath.Join(cdir, filename), nil
}

// 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"
	}

	if cfg.UIAddr == "" {
		cfg.UIAddr = "localhost:16823"
	}

	if cfg.CloudConfig == "" {
		cfg.CloudConfig = fmt.Sprintf("https://s3.amazonaws.com/lantern_config/cloud.%v.yaml.gz", countryPlaceholder)
	}

	// 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.ProxiedSites == nil {
		log.Debugf("Adding empty proxiedsites")
		cfg.ProxiedSites = &proxiedsites.Config{
			Delta: &proxiedsites.Delta{
				Additions: []string{},
				Deletions: []string{},
			},
			Cloud: []string{},
		}
	}

	if cfg.ProxiedSites.Cloud == nil || len(cfg.ProxiedSites.Cloud) == 0 {
		log.Debugf("Loading default cloud proxiedsites")
		cfg.ProxiedSites.Cloud = defaultProxiedSites
	}

	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 = []*client.FrontedServerInfo{
			&client.FrontedServerInfo{
				Host:           "nl.fallbacks.getiantem.org",
				Port:           443,
				PoolSize:       30,
				MasqueradeSet:  cloudflare,
				MaxMasquerades: 20,
				QOS:            10,
				Weight:         4000,
			},
		}

		cfg.Client.ChainedServers = make(map[string]*client.ChainedServerInfo, len(fallbacks))
		for key, fb := range fallbacks {
			cfg.Client.ChainedServers[key] = fb
		}
	}

	if cfg.AutoReport == nil {
		cfg.AutoReport = new(bool)
		*cfg.AutoReport = true
	}

	// 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)
	}

	// 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() (bytes []byte, err error) {
	log.Debugf("Fetching cloud config...")

	if cfg.IsDownstream() {
		// Clients must always proxy the request
		if cfg.Addr == "" {
			err = fmt.Errorf("No proxyAddr")
		} else {
			bytes, err = cfg.fetchCloudConfigForAddr(cfg.Addr)
		}
	} else {
		bytes, err = cfg.fetchCloudConfigForAddr("")
	}
	if err != nil {
		bytes = nil
		err = fmt.Errorf("Unable to read yaml from cloud config: %s", err)
	}
	return
}

func (cfg Config) fetchCloudConfigForAddr(proxyAddr string) ([]byte, error) {
	log.Tracef("fetchCloudConfigForAddr '%s'", proxyAddr)

	if proxyAddr != "" {
		// Wait for proxy to become available
		err := waitforserver.WaitForServer("tcp", proxyAddr, 30*time.Second)
		if err != nil {
			return nil, fmt.Errorf("Proxy never came up at %v: %v", proxyAddr, err)
		}
	}

	client, err := util.HTTPClient(cfg.CloudConfigCA, proxyAddr)
	if err != nil {
		return nil, fmt.Errorf("Unable to initialize HTTP client: %s", err)
	}
	country := strings.ToLower(globals.GetCountry())
	if country != "" {
		ret, err := cfg.fetchCloudConfigForCountry(client, country)
		// We could check specifically for a 404, but S3 actually returns a 403 when
		// a resource is not available.  I thought we'd better lean on the side of
		// robustness by avoiding hardcoding an S3-ism here, than to try and save an
		// extra request every now and then.
		if err == nil {
			return ret, err
		} else {
			log.Debugf("Couldn't fetch cloud config for country '%s'; trying the default one", country)
		}
	}
	return cfg.fetchCloudConfigForCountry(client, "default")
}

func (cfg Config) fetchCloudConfigForCountry(client *http.Client, country string) ([]byte, error) {
	url := strings.Replace(cfg.CloudConfig, countryPlaceholder, country, 1)
	log.Debugf("Checking for cloud configuration at: %s", url)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, fmt.Errorf("Unable to construct request for cloud config at %s: %s", url, err)
	}
	if lastCloudConfigETag[url] != "" {
		// Don't bother fetching if unchanged
		req.Header.Set(ifNoneMatch, lastCloudConfigETag[url])
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Unable to fetch cloud config at %s: %s", url, 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[url] = 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)
	}

	// Same with global proxiedsites
	if len(updated.ProxiedSites.Cloud) > 0 {
		wlDomains := make(map[string]bool)
		for _, domain := range updated.ProxiedSites.Cloud {
			wlDomains[domain] = true
		}
		updated.ProxiedSites.Cloud = make([]string, 0, len(wlDomains))
		for domain, _ := range wlDomains {
			updated.ProxiedSites.Cloud = append(updated.ProxiedSites.Cloud, domain)
		}
		sort.Strings(updated.ProxiedSites.Cloud)
	}
	return nil
}
<|endoftext|>


<commit_before>package kallax

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"encoding/hex"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/oklog/ulid"
	uuid "github.com/satori/go.uuid"
)

// Model contains all the basic fields that make something a model, that is,
// the ID and some internal data used by kallax.
// To make a struct a model, it only needs to have Model embedded.
//
//	type MyModel struct {
//		kallax.Model
//		Foo string
//	}
//
// Custom name for the table can be specified using the struct tag `table` when
// embedding the Model.
//
//	type MyModel struct {
//		kallax.Model `table:"custom_name"`
//	}
//
// Otherwise, the default name of the table is the name of the model converted
// to lower snake case. E.g: MyModel => my_model.
// No pluralization is done right now, but might be done in the future, so
// please, set the name of the tables yourself.
type Model struct {
	virtualColumns map[string]Identifier
	persisted      bool
	writable       bool
}

// NewModel creates a new Model that is writable and not persisted.
func NewModel() Model {
	return Model{
		persisted:      false,
		writable:       true,
		virtualColumns: make(map[string]Identifier),
	}
}

// IsPersisted returns whether the Model has already been persisted to the
// database or not.
func (m *Model) IsPersisted() bool {
	return m.persisted
}

func (m *Model) setPersisted() {
	m.persisted = true
}

// IsWritable returns whether this Model can be saved into the database.
// For example, a model with partially retrieved data is not writable, so
// it is not saved by accident and the data is corrupted. For example, if
// you select only 2 columns out of all the ones the table has, it will not
// be writable.
func (m *Model) IsWritable() bool {
	return m.writable
}

func (m *Model) setWritable(w bool) {
	m.writable = w
}

// ClearVirtualColumns clears all the previous virtual columns.
// This method is only intended for internal use. It is only exposed for
// technical reasons.
func (m *Model) ClearVirtualColumns() {
	m.virtualColumns = make(map[string]Identifier)
}

// AddVirtualColumn adds a new virtual column with the given name and value.
// This method is only intended for internal use. It is only exposed for
// technical reasons.
func (m *Model) AddVirtualColumn(name string, v Identifier) {
	if m.virtualColumns == nil {
		m.ClearVirtualColumns()
	}
	m.virtualColumns[name] = v
}

// VirtualColumn returns the value of the virtual column with the given column name.
// This method is only intended for internal use. It is only exposed for
// technical reasons.
func (m *Model) VirtualColumn(name string) Identifier {
	if m.virtualColumns == nil {
		m.ClearVirtualColumns()
	}
	return m.virtualColumns[name]
}

// Identifier is a type used to identify a model.
type Identifier interface {
	sql.Scanner
	driver.Valuer
	// Equals reports whether the identifier and the given one are equal.
	Equals(Identifier) bool
	// IsEmpty returns whether the ID is empty or not.
	IsEmpty() bool
	// Raw returns the internal value of the identifier.
	Raw() interface{}
}

// Identifiable must be implemented by those values that can be identified by an ID.
type Identifiable interface {
	// GetID returns the ID.
	GetID() Identifier
}

// Persistable must be implemented by those values that can be persisted.
type Persistable interface {
	// IsPersisted returns whether this Model is new in the store or not.
	IsPersisted() bool
	setPersisted()
}

// Writable must be implemented by those values that defines internally
// if they can be sent back to the database to be stored with its changes.
type Writable interface {
	// IsWritable returns whether this Model can be saved into the database.
	IsWritable() bool
	setWritable(bool)
}

// ColumnAddresser provides the pointer addresses of columns.
type ColumnAddresser interface {
	// ColumnAddress returns the pointer to the column value of the given
	// column name, or an error if it does not exist in the model.
	ColumnAddress(string) (interface{}, error)
}

// Relationable can perform operations related to relationships of a record.
type Relationable interface {
	// NewRelationshipRecord returns a new Record for the relationship at the
	// given field.
	NewRelationshipRecord(string) (Record, error)
	// SetRelationship sets the relationship value at the given field.
	SetRelationship(string, interface{}) error
}

// Valuer provides the values for columns.
type Valuer interface {
	// Value returns the value of the given column, or an error if it does not
	// exist in the model.
	Value(string) (interface{}, error)
}

// VirtualColumnContainer contains a collection of virtual columns and
// manages them.
type VirtualColumnContainer interface {
	// ClearVirtualColumns removes all virtual columns.
	ClearVirtualColumns()
	// AddVirtualColumn adds a new virtual column with the given name and value
	AddVirtualColumn(string, Identifier)
	// VirtualColumn returns the virtual column with the given column name.
	VirtualColumn(string) Identifier
}

// RecordValues returns the values of a record at the given columns in the same
// order as the columns.
func RecordValues(record Valuer, columns ...string) ([]interface{}, error) {
	var values = make([]interface{}, len(columns))
	for i, col := range columns {
		v, err := record.Value(col)
		if err != nil {
			return nil, err
		}
		values[i] = v
	}
	return values, nil
}

// Record is something that can be stored as a row in the database.
type Record interface {
	Identifiable
	Persistable
	Writable
	Relationable
	ColumnAddresser
	Valuer
	VirtualColumnContainer
}

var randPool = &sync.Pool{
	New: func() interface{} {
		return rand.NewSource(time.Now().UnixNano())
	},
}

// ULID is an ID type provided by kallax that is a lexically sortable UUID.
// The internal representation is an ULID (https://github.com/oklog/ulid).
// It already implements sql.Scanner and driver.Valuer, so it's perfectly
// safe for database usage.
type ULID uuid.UUID

// NewULID returns a new ULID, which is a lexically sortable UUID.
func NewULID() ULID {
	entropy := randPool.Get().(rand.Source)
	id := ULID(ulid.MustNew(ulid.Timestamp(time.Now()), rand.New(entropy)))
	randPool.Put(entropy)

	return id
}

// Scan implements the Scanner interface.
func (id *ULID) Scan(src interface{}) error {
	switch src := src.(type) {
	case []byte:
		if len(src) != 16 {
			return id.UnmarshalText(src)
		}

		var ulid ulid.ULID
		if err := ulid.UnmarshalBinary(src); err != nil {
			return err
		}
		*id = ULID(ulid)
		return nil
	case string:
		return id.Scan([]byte(src))
	default:
		return fmt.Errorf("kallax: cannot scan %T into ULID", src)
	}
}

var (
	urnPrefix  = []byte("urn:uuid:")
	byteGroups = []int{8, 4, 4, 4, 12}
)

// UnmarshalText implements the encoding.TextUnmarshaler interface.
// Following formats are supported:
// "6ba7b810-9dad-11d1-80b4-00c04fd430c8",
// "{6ba7b810-9dad-11d1-80b4-00c04fd430c8}",
// "urn:uuid:6ba7b810-9dad-11d1-80b4-00c04fd430c8"
// Implements the exact same code as the UUID UnmarshalText removing the
// version check.
func (u *ULID) UnmarshalText(text []byte) (err error) {
	if len(text) < 32 {
		err = fmt.Errorf("uuid: UUID string too short: %s", text)
		return
	}

	t := text[:]
	braced := false

	if bytes.Equal(t[:9], urnPrefix) {
		t = t[9:]
	} else if t[0] == '{' {
		braced = true
		t = t[1:]
	}

	b := u[:]

	for i, byteGroup := range byteGroups {
		if i > 0 && t[0] == '-' {
			t = t[1:]
		} else if i > 0 && t[0] != '-' {
			err = fmt.Errorf("kallax: invalid ulid string format")
			return
		}

		if len(t) < byteGroup {
			err = fmt.Errorf("kallax: ulid string too short: %s", text)
			return
		}

		if i == 4 && len(t) > byteGroup &&
			((braced && t[byteGroup] != '}') || len(t[byteGroup:]) > 1 || !braced) {
			err = fmt.Errorf("kallax: ulid string too long: %s", t)
			return
		}

		_, err = hex.Decode(b[:byteGroup/2], t[:byteGroup])

		if err != nil {
			return
		}

		t = t[byteGroup:]
		b = b[byteGroup/2:]
	}

	return
}

// Value implements the Valuer interface.
func (id ULID) Value() (driver.Value, error) {
	return uuid.UUID(id).Value()
}

// IsEmpty returns whether the ID is empty or not. An empty ID means it has not
// been set yet.
func (id ULID) IsEmpty() bool {
	return uuid.Equal(uuid.UUID(id), uuid.Nil)
}

// String returns the string representation of the ID.
func (id ULID) String() string {
	return uuid.UUID(id).String()
}

// Equals reports whether the ID and the given one are equals.
func (id ULID) Equals(other Identifier) bool {
	v, ok := other.(*ULID)
	if !ok {
		return false
	}

	return uuid.Equal(uuid.UUID(id), uuid.UUID(*v))
}

// Raw returns the underlying raw value.
func (id ULID) Raw() interface{} {
	return id
}

// NumericID is a wrapper for int64 that implements the Identifier interface.
// You don't need to actually use this as a type in your model. They will be
// automatically converted to and from in the generated code.
type NumericID int64

// Scan implements the Scanner interface.
func (id *NumericID) Scan(src interface{}) error {
	switch src := src.(type) {
	case int64:
		*(*int64)(id) = src
	default:
		return fmt.Errorf("kallax: cannot scan value of type %T into a numeric ID", src)
	}

	return nil
}

// Value implements the Valuer interface.
func (id NumericID) Value() (driver.Value, error) {
	return int64(id), nil
}

// IsEmpty returns whether the ID is empty or not. An empty ID means it has not
// been set yet.
func (id NumericID) IsEmpty() bool {
	return int64(id) == 0
}

// String returns the string representation of the ID.
func (id NumericID) String() string {
	return fmt.Sprint(int64(id))
}

// Equals reports whether the ID and the given one are equals.
func (id NumericID) Equals(other Identifier) bool {
	v, ok := other.(*NumericID)
	if !ok {
		return false
	}

	return int64(id) == int64(*v)
}

// Raw returns the underlying raw value.
func (id NumericID) Raw() interface{} {
	return id
}

// UUID is a wrapper type for uuid.UUID that implements the Identifier
// interface.
// You don't need to actually use this as a type in your model. They will be
// automatically converted to and from in the generated code.
type UUID uuid.UUID

// Scan implements the Scanner interface.
func (id *UUID) Scan(src interface{}) error {
	return (*uuid.UUID)(id).Scan(src)
}

// Value implements the Valuer interface.
func (id UUID) Value() (driver.Value, error) {
	return uuid.UUID(id).Value()
}

// IsEmpty returns whether the ID is empty or not. An empty ID means it has not
// been set yet.
func (id UUID) IsEmpty() bool {
	return uuid.Equal(uuid.UUID(id), uuid.Nil)
}

// String returns the string representation of the ID.
func (id UUID) String() string {
	return uuid.UUID(id).String()
}

// Equals reports whether the ID and the given one are equals.
func (id UUID) Equals(other Identifier) bool {
	v, ok := other.(*UUID)
	if !ok {
		return false
	}

	return uuid.Equal(uuid.UUID(id), uuid.UUID(*v))
}

// Raw returns the underlying raw value.
func (id UUID) Raw() interface{} {
	return id
}

type virtualColumn struct {
	r   Record
	col string
	id  Identifier
}

// VirtualColumn returns a sql.Scanner that will scan the given column as a
// virtual column in the given record.
func VirtualColumn(col string, r Record, id Identifier) sql.Scanner {
	return &virtualColumn{r, col, id}
}

// Scan implements the scanner interface.
func (c *virtualColumn) Scan(src interface{}) error {
	if err := c.id.Scan(src); err != nil {
		return err
	}

	c.r.AddVirtualColumn(c.col, c.id)
	return nil
}
<commit_msg>implement TextMarshaler for ULID<commit_after>package kallax

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"encoding/hex"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/oklog/ulid"
	uuid "github.com/satori/go.uuid"
)

// Model contains all the basic fields that make something a model, that is,
// the ID and some internal data used by kallax.
// To make a struct a model, it only needs to have Model embedded.
//
//	type MyModel struct {
//		kallax.Model
//		Foo string
//	}
//
// Custom name for the table can be specified using the struct tag `table` when
// embedding the Model.
//
//	type MyModel struct {
//		kallax.Model `table:"custom_name"`
//	}
//
// Otherwise, the default name of the table is the name of the model converted
// to lower snake case. E.g: MyModel => my_model.
// No pluralization is done right now, but might be done in the future, so
// please, set the name of the tables yourself.
type Model struct {
	virtualColumns map[string]Identifier
	persisted      bool
	writable       bool
}

// NewModel creates a new Model that is writable and not persisted.
func NewModel() Model {
	return Model{
		persisted:      false,
		writable:       true,
		virtualColumns: make(map[string]Identifier),
	}
}

// IsPersisted returns whether the Model has already been persisted to the
// database or not.
func (m *Model) IsPersisted() bool {
	return m.persisted
}

func (m *Model) setPersisted() {
	m.persisted = true
}

// IsWritable returns whether this Model can be saved into the database.
// For example, a model with partially retrieved data is not writable, so
// it is not saved by accident and the data is corrupted. For example, if
// you select only 2 columns out of all the ones the table has, it will not
// be writable.
func (m *Model) IsWritable() bool {
	return m.writable
}

func (m *Model) setWritable(w bool) {
	m.writable = w
}

// ClearVirtualColumns clears all the previous virtual columns.
// This method is only intended for internal use. It is only exposed for
// technical reasons.
func (m *Model) ClearVirtualColumns() {
	m.virtualColumns = make(map[string]Identifier)
}

// AddVirtualColumn adds a new virtual column with the given name and value.
// This method is only intended for internal use. It is only exposed for
// technical reasons.
func (m *Model) AddVirtualColumn(name string, v Identifier) {
	if m.virtualColumns == nil {
		m.ClearVirtualColumns()
	}
	m.virtualColumns[name] = v
}

// VirtualColumn returns the value of the virtual column with the given column name.
// This method is only intended for internal use. It is only exposed for
// technical reasons.
func (m *Model) VirtualColumn(name string) Identifier {
	if m.virtualColumns == nil {
		m.ClearVirtualColumns()
	}
	return m.virtualColumns[name]
}

// Identifier is a type used to identify a model.
type Identifier interface {
	sql.Scanner
	driver.Valuer
	// Equals reports whether the identifier and the given one are equal.
	Equals(Identifier) bool
	// IsEmpty returns whether the ID is empty or not.
	IsEmpty() bool
	// Raw returns the internal value of the identifier.
	Raw() interface{}
}

// Identifiable must be implemented by those values that can be identified by an ID.
type Identifiable interface {
	// GetID returns the ID.
	GetID() Identifier
}

// Persistable must be implemented by those values that can be persisted.
type Persistable interface {
	// IsPersisted returns whether this Model is new in the store or not.
	IsPersisted() bool
	setPersisted()
}

// Writable must be implemented by those values that defines internally
// if they can be sent back to the database to be stored with its changes.
type Writable interface {
	// IsWritable returns whether this Model can be saved into the database.
	IsWritable() bool
	setWritable(bool)
}

// ColumnAddresser provides the pointer addresses of columns.
type ColumnAddresser interface {
	// ColumnAddress returns the pointer to the column value of the given
	// column name, or an error if it does not exist in the model.
	ColumnAddress(string) (interface{}, error)
}

// Relationable can perform operations related to relationships of a record.
type Relationable interface {
	// NewRelationshipRecord returns a new Record for the relationship at the
	// given field.
	NewRelationshipRecord(string) (Record, error)
	// SetRelationship sets the relationship value at the given field.
	SetRelationship(string, interface{}) error
}

// Valuer provides the values for columns.
type Valuer interface {
	// Value returns the value of the given column, or an error if it does not
	// exist in the model.
	Value(string) (interface{}, error)
}

// VirtualColumnContainer contains a collection of virtual columns and
// manages them.
type VirtualColumnContainer interface {
	// ClearVirtualColumns removes all virtual columns.
	ClearVirtualColumns()
	// AddVirtualColumn adds a new virtual column with the given name and value
	AddVirtualColumn(string, Identifier)
	// VirtualColumn returns the virtual column with the given column name.
	VirtualColumn(string) Identifier
}

// RecordValues returns the values of a record at the given columns in the same
// order as the columns.
func RecordValues(record Valuer, columns ...string) ([]interface{}, error) {
	var values = make([]interface{}, len(columns))
	for i, col := range columns {
		v, err := record.Value(col)
		if err != nil {
			return nil, err
		}
		values[i] = v
	}
	return values, nil
}

// Record is something that can be stored as a row in the database.
type Record interface {
	Identifiable
	Persistable
	Writable
	Relationable
	ColumnAddresser
	Valuer
	VirtualColumnContainer
}

var randPool = &sync.Pool{
	New: func() interface{} {
		return rand.NewSource(time.Now().UnixNano())
	},
}

// ULID is an ID type provided by kallax that is a lexically sortable UUID.
// The internal representation is an ULID (https://github.com/oklog/ulid).
// It already implements sql.Scanner and driver.Valuer, so it's perfectly
// safe for database usage.
type ULID uuid.UUID

// NewULID returns a new ULID, which is a lexically sortable UUID.
func NewULID() ULID {
	entropy := randPool.Get().(rand.Source)
	id := ULID(ulid.MustNew(ulid.Timestamp(time.Now()), rand.New(entropy)))
	randPool.Put(entropy)

	return id
}

// Scan implements the Scanner interface.
func (id *ULID) Scan(src interface{}) error {
	switch src := src.(type) {
	case []byte:
		if len(src) != 16 {
			return id.UnmarshalText(src)
		}

		var ulid ulid.ULID
		if err := ulid.UnmarshalBinary(src); err != nil {
			return err
		}
		*id = ULID(ulid)
		return nil
	case string:
		return id.Scan([]byte(src))
	default:
		return fmt.Errorf("kallax: cannot scan %T into ULID", src)
	}
}

var (
	urnPrefix  = []byte("urn:uuid:")
	byteGroups = []int{8, 4, 4, 4, 12}
)

// UnmarshalText implements the encoding.TextUnmarshaler interface.
// Following formats are supported:
// "6ba7b810-9dad-11d1-80b4-00c04fd430c8",
// "{6ba7b810-9dad-11d1-80b4-00c04fd430c8}",
// "urn:uuid:6ba7b810-9dad-11d1-80b4-00c04fd430c8"
// Implements the exact same code as the UUID UnmarshalText removing the
// version check.
func (u *ULID) UnmarshalText(text []byte) (err error) {
	if len(text) < 32 {
		err = fmt.Errorf("uuid: UUID string too short: %s", text)
		return
	}

	t := text[:]
	braced := false

	if bytes.Equal(t[:9], urnPrefix) {
		t = t[9:]
	} else if t[0] == '{' {
		braced = true
		t = t[1:]
	}

	b := u[:]

	for i, byteGroup := range byteGroups {
		if i > 0 && t[0] == '-' {
			t = t[1:]
		} else if i > 0 && t[0] != '-' {
			err = fmt.Errorf("kallax: invalid ulid string format")
			return
		}

		if len(t) < byteGroup {
			err = fmt.Errorf("kallax: ulid string too short: %s", text)
			return
		}

		if i == 4 && len(t) > byteGroup &&
			((braced && t[byteGroup] != '}') || len(t[byteGroup:]) > 1 || !braced) {
			err = fmt.Errorf("kallax: ulid string too long: %s", t)
			return
		}

		_, err = hex.Decode(b[:byteGroup/2], t[:byteGroup])

		if err != nil {
			return
		}

		t = t[byteGroup:]
		b = b[byteGroup/2:]
	}

	return
}

func (id ULID) MarshalText() ([]byte, error) {
	return []byte(id.String()), nil
}

// Value implements the Valuer interface.
func (id ULID) Value() (driver.Value, error) {
	return uuid.UUID(id).Value()
}

// IsEmpty returns whether the ID is empty or not. An empty ID means it has not
// been set yet.
func (id ULID) IsEmpty() bool {
	return uuid.Equal(uuid.UUID(id), uuid.Nil)
}

// String returns the string representation of the ID.
func (id ULID) String() string {
	return uuid.UUID(id).String()
}

// Equals reports whether the ID and the given one are equals.
func (id ULID) Equals(other Identifier) bool {
	v, ok := other.(*ULID)
	if !ok {
		return false
	}

	return uuid.Equal(uuid.UUID(id), uuid.UUID(*v))
}

// Raw returns the underlying raw value.
func (id ULID) Raw() interface{} {
	return id
}

// NumericID is a wrapper for int64 that implements the Identifier interface.
// You don't need to actually use this as a type in your model. They will be
// automatically converted to and from in the generated code.
type NumericID int64

// Scan implements the Scanner interface.
func (id *NumericID) Scan(src interface{}) error {
	switch src := src.(type) {
	case int64:
		*(*int64)(id) = src
	default:
		return fmt.Errorf("kallax: cannot scan value of type %T into a numeric ID", src)
	}

	return nil
}

// Value implements the Valuer interface.
func (id NumericID) Value() (driver.Value, error) {
	return int64(id), nil
}

// IsEmpty returns whether the ID is empty or not. An empty ID means it has not
// been set yet.
func (id NumericID) IsEmpty() bool {
	return int64(id) == 0
}

// String returns the string representation of the ID.
func (id NumericID) String() string {
	return fmt.Sprint(int64(id))
}

// Equals reports whether the ID and the given one are equals.
func (id NumericID) Equals(other Identifier) bool {
	v, ok := other.(*NumericID)
	if !ok {
		return false
	}

	return int64(id) == int64(*v)
}

// Raw returns the underlying raw value.
func (id NumericID) Raw() interface{} {
	return id
}

// UUID is a wrapper type for uuid.UUID that implements the Identifier
// interface.
// You don't need to actually use this as a type in your model. They will be
// automatically converted to and from in the generated code.
type UUID uuid.UUID

// Scan implements the Scanner interface.
func (id *UUID) Scan(src interface{}) error {
	return (*uuid.UUID)(id).Scan(src)
}

// Value implements the Valuer interface.
func (id UUID) Value() (driver.Value, error) {
	return uuid.UUID(id).Value()
}

// IsEmpty returns whether the ID is empty or not. An empty ID means it has not
// been set yet.
func (id UUID) IsEmpty() bool {
	return uuid.Equal(uuid.UUID(id), uuid.Nil)
}

// String returns the string representation of the ID.
func (id UUID) String() string {
	return uuid.UUID(id).String()
}

// Equals reports whether the ID and the given one are equals.
func (id UUID) Equals(other Identifier) bool {
	v, ok := other.(*UUID)
	if !ok {
		return false
	}

	return uuid.Equal(uuid.UUID(id), uuid.UUID(*v))
}

// Raw returns the underlying raw value.
func (id UUID) Raw() interface{} {
	return id
}

type virtualColumn struct {
	r   Record
	col string
	id  Identifier
}

// VirtualColumn returns a sql.Scanner that will scan the given column as a
// virtual column in the given record.
func VirtualColumn(col string, r Record, id Identifier) sql.Scanner {
	return &virtualColumn{r, col, id}
}

// Scan implements the scanner interface.
func (c *virtualColumn) Scan(src interface{}) error {
	if err := c.id.Scan(src); err != nil {
		return err
	}

	c.r.AddVirtualColumn(c.col, c.id)
	return nil
}
<|endoftext|>


<commit_before>package entities

import (
	"github.com/shopspring/decimal"
)

type Consignation struct {
	Type int
	Quantity int
	FutureId int
	Price decimal.Decimal
	OpenQuantity int
	Direction int
	FirmId int
}
<commit_msg>Implement BinaryMarshaler inteface<commit_after>package entities

import (
	"github.com/shopspring/decimal"
	"strconv"
)

type Consignation struct {
	Type int
	Quantity int
	FutureId int
	Price decimal.Decimal
	OpenQuantity int
	Direction int
	FirmId int
}

func (c Consignation) MarshalBinary() (data []byte, err error) {
	return []byte(strconv.Itoa(c.Type) + "," +
		strconv.Itoa(c.Quantity) + "," +
		strconv.Itoa(c.FutureId) + "," +
		strconv.Itoa(c.FutureId) + "," +
		strconv.Itoa(c.OpenQuantity) + "," +
		strconv.Itoa(c.Direction) + "," +
		strconv.Itoa(c.FirmId) + "," +
		c.Price.String() + ","),nil
}

func (c Consignation) UnmarshalBinary(data []byte) error  {
	return nil
}
<|endoftext|>


<commit_before>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/network"
	"github.com/APTrust/exchange/stats"
	"net/url"
	"time"
)

const UNKNOWN_TOPIC = "unknown_topic"

type APTQueue struct {
	Context      *context.Context
	NSQClient    *network.NSQClient
	topic        string
	stats        *stats.APTQueueStats
	dryRun       bool
	statsEnabled bool
}

// NewAPTQueue creates a new queue worker to push WorkItems from
// Pharos into NSQ, and marked them as queued. If param topic is
// specified, this will queue items destined for the specified topic;
// otherwise, it will queue items for all topics. If param enableStats
// is true, it will dump stats about what was queued to a JSON file.
// If param dryRun is true, it will log all the items it would have
// queued, without actually pushing anything to NSQ.
func NewAPTQueue(_context *context.Context, topic string, enableStats, dryRun bool) *APTQueue {
	_context.MessageLog.Info("NSQ address: %s", _context.Config.NsqdHttpAddress)
	nsqClient := network.NewNSQClient(_context.Config.NsqdHttpAddress)
	aptQueue := &APTQueue{
		Context:      _context,
		NSQClient:    nsqClient,
		topic:        topic,
		statsEnabled: enableStats,
		dryRun:       dryRun,
	}
	if enableStats {
		aptQueue.stats = stats.NewAPTQueueStats()
	}
	return aptQueue
}

// Run retrieves all unqueued work items from Pharos and pushes
// them into the appropriate NSQ topic.
func (aptQueue *APTQueue) Run() {
	aptQueue.printLogHeader()
	params := url.Values{}
	params.Set("queued", "false")
	params.Set("status", constants.StatusPending)
	params.Set("retry", "true")
	params.Set("node_empty", "true")
	params.Set("page", "1")
	params.Set("per_page", "100")
	for {
		resp := aptQueue.Context.PharosClient.WorkItemList(params)
		aptQueue.Context.MessageLog.Info("GET %s", resp.Request.URL)
		if resp.Error != nil {
			aptQueue.recordError(
				"Error getting WorkItem list from Pharos: %s",
				resp.Error)
		}
		for _, item := range resp.WorkItems() {
			if aptQueue.addToNSQ(item) {
				aptQueue.markAsQueued(item)
			}
		}
		if resp.HasNextPage() == false {
			break
		}
		params = resp.ParamsForNextPage()
	}
}

func (aptQueue *APTQueue) addToNSQ(workItem *models.WorkItem) bool {
	identifier := workItem.Name
	if workItem.ObjectIdentifier != "" {
		identifier = workItem.ObjectIdentifier
	}
	if workItem.GenericFileIdentifier != "" {
		identifier = workItem.GenericFileIdentifier
	}
	topic := aptQueue.getNSQTopic(workItem)
	if aptQueue.topic != "" && topic != aptQueue.topic {
		aptQueue.Context.MessageLog.Info(
			"Skipping WorkItem id %d - %s (%s/%s/%s) because topic would be %s",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status, topic)
		return false
	}
	if topic == UNKNOWN_TOPIC {
		aptQueue.recordError(
			"Unknown topic for WorkItem %d - %s (%s/%s/%s)",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status)
		return false
	}
	if aptQueue.dryRun {
		aptQueue.Context.MessageLog.Info(
			"[DRY RUN] Would add WorkItem id %d - %s (%s/%s/%s) - to %s",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status, topic)
		return false
	}
	err := aptQueue.NSQClient.Enqueue(topic, workItem.Id)
	if err != nil {
		aptQueue.recordError("Error sending WorkItem %d %s (%s/%s/%s) - to %s: %v",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status, topic, err)
		return false
	}
	aptQueue.Context.MessageLog.Info("Added WorkItem id %d - %s (%s/%s/%s) - to %s",
		workItem.Id, identifier, workItem.Action, workItem.Stage, workItem.Status, topic)
	if aptQueue.stats != nil {
		aptQueue.stats.AddWorkItem(topic, workItem)
	}
	return true
}

func (aptQueue *APTQueue) markAsQueued(workItem *models.WorkItem) *models.WorkItem {
	utcNow := time.Now().UTC()
	workItem.Date = utcNow
	workItem.QueuedAt = &utcNow
	resp := aptQueue.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		aptQueue.recordError("Error setting QueuedAt for WorkItem with id %d: %v",
			workItem.Id, resp.Error)
		return nil
	}
	if resp.Response.StatusCode != 200 {
		aptQueue.processPharosError(resp)
		return nil
	}
	aptQueue.Context.MessageLog.Info("Marked WorkItem id %d (%s/%s/%s) as queued in Pharos",
		workItem.Id, workItem.Action, workItem.Stage, workItem.Status)
	if aptQueue.stats != nil {
		aptQueue.stats.AddItemMarkedAsQueued(workItem)
	}
	return resp.WorkItem()
}

func (aptQueue *APTQueue) processPharosError(resp *network.PharosResponse) {
	respBody := ""
	bytesRead, aptQueuer := resp.RawResponseData()
	if aptQueuer == nil {
		respBody = string(bytesRead)
	} else {
		respBody = fmt.Sprintf("[Could not read response body: %v]", aptQueuer)
	}
	aptQueue.recordError("%s %s returned status code %d. Response body: %s",
		resp.Request.Method, resp.Request.URL, resp.Response.StatusCode, respBody)
}

func (aptQueue *APTQueue) recordError(format string, a ...interface{}) {
	msg := fmt.Sprintf(format, a...)
	if aptQueue.stats != nil {
		aptQueue.stats.AddError(msg)
	}
	aptQueue.Context.MessageLog.Error(msg)
}

func (aptQueue *APTQueue) getNSQTopic(workItem *models.WorkItem) string {
	config := aptQueue.Context.Config
	topic := UNKNOWN_TOPIC
	if workItem.Action == constants.ActionIngest {
		if workItem.Stage == constants.StageReceive {
			topic = config.FetchWorker.NsqTopic
		} else if workItem.Stage == constants.StageStore {
			topic = config.StoreWorker.NsqTopic
		} else if workItem.Stage == constants.StageRecord {
			topic = config.RecordWorker.NsqTopic
		}
	} else if workItem.Action == constants.ActionFixityCheck {
		topic = config.FixityWorker.NsqTopic
	} else if workItem.Action == constants.ActionRestore {
		topic = config.RestoreWorker.NsqTopic
	} else if workItem.Action == constants.ActionDelete {
		topic = config.FileDeleteWorker.NsqTopic
	} else if workItem.Action == constants.ActionDPN {
		topic = config.DPN.DPNPackageWorker.NsqTopic
	}
	return topic
}

func (aptQueue *APTQueue) GetStats() *stats.APTQueueStats {
	return aptQueue.stats
}

func (aptQueue *APTQueue) printLogHeader() {
	topic := aptQueue.topic
	if aptQueue.topic == "" {
		topic = "ALL"
	}
	aptQueue.Context.MessageLog.Info("apt_queue started with the following params:")
	aptQueue.Context.MessageLog.Info("Topic = %s", topic)
	aptQueue.Context.MessageLog.Info("Enable Stats = %t", aptQueue.statsEnabled)
	aptQueue.Context.MessageLog.Info("Dry Run = %t", aptQueue.dryRun)
}
<commit_msg>PT #158594063: Queue GlacierRestore requests<commit_after>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/network"
	"github.com/APTrust/exchange/stats"
	"net/url"
	"time"
)

const UNKNOWN_TOPIC = "unknown_topic"

type APTQueue struct {
	Context      *context.Context
	NSQClient    *network.NSQClient
	topic        string
	stats        *stats.APTQueueStats
	dryRun       bool
	statsEnabled bool
}

// NewAPTQueue creates a new queue worker to push WorkItems from
// Pharos into NSQ, and marked them as queued. If param topic is
// specified, this will queue items destined for the specified topic;
// otherwise, it will queue items for all topics. If param enableStats
// is true, it will dump stats about what was queued to a JSON file.
// If param dryRun is true, it will log all the items it would have
// queued, without actually pushing anything to NSQ.
func NewAPTQueue(_context *context.Context, topic string, enableStats, dryRun bool) *APTQueue {
	_context.MessageLog.Info("NSQ address: %s", _context.Config.NsqdHttpAddress)
	nsqClient := network.NewNSQClient(_context.Config.NsqdHttpAddress)
	aptQueue := &APTQueue{
		Context:      _context,
		NSQClient:    nsqClient,
		topic:        topic,
		statsEnabled: enableStats,
		dryRun:       dryRun,
	}
	if enableStats {
		aptQueue.stats = stats.NewAPTQueueStats()
	}
	return aptQueue
}

// Run retrieves all unqueued work items from Pharos and pushes
// them into the appropriate NSQ topic.
func (aptQueue *APTQueue) Run() {
	aptQueue.printLogHeader()
	params := url.Values{}
	params.Set("queued", "false")
	params.Set("status", constants.StatusPending)
	params.Set("retry", "true")
	params.Set("node_empty", "true")
	params.Set("page", "1")
	params.Set("per_page", "100")
	for {
		resp := aptQueue.Context.PharosClient.WorkItemList(params)
		aptQueue.Context.MessageLog.Info("GET %s", resp.Request.URL)
		if resp.Error != nil {
			aptQueue.recordError(
				"Error getting WorkItem list from Pharos: %s",
				resp.Error)
		}
		for _, item := range resp.WorkItems() {
			if aptQueue.addToNSQ(item) {
				aptQueue.markAsQueued(item)
			}
		}
		if resp.HasNextPage() == false {
			break
		}
		params = resp.ParamsForNextPage()
	}
}

func (aptQueue *APTQueue) addToNSQ(workItem *models.WorkItem) bool {
	identifier := workItem.Name
	if workItem.ObjectIdentifier != "" {
		identifier = workItem.ObjectIdentifier
	}
	if workItem.GenericFileIdentifier != "" {
		identifier = workItem.GenericFileIdentifier
	}
	topic := aptQueue.getNSQTopic(workItem)
	if aptQueue.topic != "" && topic != aptQueue.topic {
		aptQueue.Context.MessageLog.Info(
			"Skipping WorkItem id %d - %s (%s/%s/%s) because topic would be %s",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status, topic)
		return false
	}
	if topic == UNKNOWN_TOPIC {
		aptQueue.recordError(
			"Unknown topic for WorkItem %d - %s (%s/%s/%s)",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status)
		return false
	}
	if aptQueue.dryRun {
		aptQueue.Context.MessageLog.Info(
			"[DRY RUN] Would add WorkItem id %d - %s (%s/%s/%s) - to %s",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status, topic)
		return false
	}
	err := aptQueue.NSQClient.Enqueue(topic, workItem.Id)
	if err != nil {
		aptQueue.recordError("Error sending WorkItem %d %s (%s/%s/%s) - to %s: %v",
			workItem.Id, identifier, workItem.Action,
			workItem.Stage, workItem.Status, topic, err)
		return false
	}
	aptQueue.Context.MessageLog.Info("Added WorkItem id %d - %s (%s/%s/%s) - to %s",
		workItem.Id, identifier, workItem.Action, workItem.Stage, workItem.Status, topic)
	if aptQueue.stats != nil {
		aptQueue.stats.AddWorkItem(topic, workItem)
	}
	return true
}

func (aptQueue *APTQueue) markAsQueued(workItem *models.WorkItem) *models.WorkItem {
	utcNow := time.Now().UTC()
	workItem.Date = utcNow
	workItem.QueuedAt = &utcNow
	resp := aptQueue.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		aptQueue.recordError("Error setting QueuedAt for WorkItem with id %d: %v",
			workItem.Id, resp.Error)
		return nil
	}
	if resp.Response.StatusCode != 200 {
		aptQueue.processPharosError(resp)
		return nil
	}
	aptQueue.Context.MessageLog.Info("Marked WorkItem id %d (%s/%s/%s) as queued in Pharos",
		workItem.Id, workItem.Action, workItem.Stage, workItem.Status)
	if aptQueue.stats != nil {
		aptQueue.stats.AddItemMarkedAsQueued(workItem)
	}
	return resp.WorkItem()
}

func (aptQueue *APTQueue) processPharosError(resp *network.PharosResponse) {
	respBody := ""
	bytesRead, aptQueuer := resp.RawResponseData()
	if aptQueuer == nil {
		respBody = string(bytesRead)
	} else {
		respBody = fmt.Sprintf("[Could not read response body: %v]", aptQueuer)
	}
	aptQueue.recordError("%s %s returned status code %d. Response body: %s",
		resp.Request.Method, resp.Request.URL, resp.Response.StatusCode, respBody)
}

func (aptQueue *APTQueue) recordError(format string, a ...interface{}) {
	msg := fmt.Sprintf(format, a...)
	if aptQueue.stats != nil {
		aptQueue.stats.AddError(msg)
	}
	aptQueue.Context.MessageLog.Error(msg)
}

func (aptQueue *APTQueue) getNSQTopic(workItem *models.WorkItem) string {
	config := aptQueue.Context.Config
	topic := UNKNOWN_TOPIC
	if workItem.Action == constants.ActionIngest {
		if workItem.Stage == constants.StageReceive {
			topic = config.FetchWorker.NsqTopic
		} else if workItem.Stage == constants.StageStore {
			topic = config.StoreWorker.NsqTopic
		} else if workItem.Stage == constants.StageRecord {
			topic = config.RecordWorker.NsqTopic
		}
	} else if workItem.Action == constants.ActionFixityCheck {
		topic = config.FixityWorker.NsqTopic
	} else if workItem.Action == constants.ActionRestore {
		topic = config.RestoreWorker.NsqTopic
	} else if workItem.Action == constants.ActionGlacierRestore {
		topic = config.GlacierRestoreWorker.NsqTopic
	} else if workItem.Action == constants.ActionDelete {
		topic = config.FileDeleteWorker.NsqTopic
	} else if workItem.Action == constants.ActionDPN {
		topic = config.DPN.DPNPackageWorker.NsqTopic
	}
	return topic
}

func (aptQueue *APTQueue) GetStats() *stats.APTQueueStats {
	return aptQueue.stats
}

func (aptQueue *APTQueue) printLogHeader() {
	topic := aptQueue.topic
	if aptQueue.topic == "" {
		topic = "ALL"
	}
	aptQueue.Context.MessageLog.Info("apt_queue started with the following params:")
	aptQueue.Context.MessageLog.Info("Topic = %s", topic)
	aptQueue.Context.MessageLog.Info("Enable Stats = %t", aptQueue.statsEnabled)
	aptQueue.Context.MessageLog.Info("Dry Run = %t", aptQueue.dryRun)
}
<|endoftext|>


<commit_before>package deployer

import (
	"context"
	"fmt"
	"reflect"

	"github.com/rancher/rancher/pkg/ref"
	"github.com/rancher/types/apis/core/v1"

	"github.com/rancher/norman/controller"
	alertutil "github.com/rancher/rancher/pkg/controllers/user/alert/common"
	"github.com/rancher/rancher/pkg/controllers/user/systemimage"
	monitorutil "github.com/rancher/rancher/pkg/monitoring"
	"github.com/rancher/types/apis/apps/v1beta2"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	projectv3 "github.com/rancher/types/apis/project.cattle.io/v3"

	"github.com/rancher/rancher/pkg/controllers/user/helm/common"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/types/config"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
)

var (
	serviceName = "alerting"
)

type alertService struct {
	clusterName        string
	apps               projectv3.AppInterface
	deployments        v1beta2.DeploymentInterface
	oldClusterAlerts   v3.ClusterAlertInterface
	oldProjectAlerts   v3.ProjectAlertInterface
	clusterAlertGroups v3.ClusterAlertGroupInterface
	projectAlertGroups v3.ProjectAlertGroupInterface
	clusterAlertRules  v3.ClusterAlertRuleInterface
	projectAlertRules  v3.ProjectAlertRuleInterface
	projectLister      v3.ProjectLister
	namespaces         v1.NamespaceInterface
	appDeployer        *appDeployer
}

func init() {
	systemimage.RegisterSystemService(serviceName, &alertService{})
}

func (l *alertService) Init(ctx context.Context, cluster *config.UserContext) {
	ad := &appDeployer{
		appsGetter:       cluster.Management.Project,
		namespaces:       cluster.Management.Core.Namespaces(metav1.NamespaceAll),
		secrets:          cluster.Core.Secrets(metav1.NamespaceAll),
		templateVersions: cluster.Management.Management.TemplateVersions(metav1.NamespaceAll),
	}

	l.clusterName = cluster.ClusterName
	l.deployments = cluster.Apps.Deployments("")
	l.oldClusterAlerts = cluster.Management.Management.ClusterAlerts(cluster.ClusterName)
	l.oldProjectAlerts = cluster.Management.Management.ProjectAlerts(metav1.NamespaceAll)
	l.clusterAlertGroups = cluster.Management.Management.ClusterAlertGroups(cluster.ClusterName)
	l.projectAlertGroups = cluster.Management.Management.ProjectAlertGroups(metav1.NamespaceAll)
	l.clusterAlertRules = cluster.Management.Management.ClusterAlertRules(cluster.ClusterName)
	l.projectAlertRules = cluster.Management.Management.ProjectAlertRules(metav1.NamespaceAll)
	l.projectLister = cluster.Management.Management.Projects(cluster.ClusterName).Controller().Lister()
	l.apps = cluster.Management.Project.Apps(metav1.NamespaceAll)
	l.appDeployer = ad

}

func (l *alertService) Version() (string, error) {
	catalogID := settings.SystemMonitoringCatalogID.Get()
	templateVersionID, err := common.ParseExternalID(catalogID)
	if err != nil {
		return "", fmt.Errorf("get system monitor catalog version failed, %v", err)
	}
	return templateVersionID, nil
}

func (l *alertService) Upgrade(currentVersion string) (string, error) {
	newCatalogID := settings.SystemMonitoringCatalogID.Get()

	NewVersion, err := common.ParseExternalID(newCatalogID)
	if currentVersion == NewVersion {
		return currentVersion, nil
	}

	//migrate legacy
	appName, _ := monitorutil.ClusterAlertManagerInfo()
	oldClusterAlert, err := l.oldClusterAlerts.List(metav1.ListOptions{})
	if err != nil {
		return "", fmt.Errorf("get old cluster alert failed, %s", err)
	}

	for _, v := range oldClusterAlert.Items {
		migrationGroupName := fmt.Sprintf("migrate-group-%s", v.Name)
		groupID := alertutil.GetGroupID(l.clusterName, migrationGroupName)

		name := fmt.Sprintf("migrate-%s", v.Name)
		newClusterRule := &v3.ClusterAlertRule{
			ObjectMeta: metav1.ObjectMeta{
				Name:      name,
				Namespace: l.clusterName,
			},
			Spec: v3.ClusterAlertRuleSpec{
				ClusterName: l.clusterName,
				GroupName:   groupID,
				CommonRuleField: v3.CommonRuleField{
					DisplayName: v.Spec.DisplayName,
					Severity:    v.Spec.Severity,
					TimingField: v3.TimingField{
						GroupIntervalSeconds:  v.Spec.InitialWaitSeconds,
						RepeatIntervalSeconds: v.Spec.RepeatIntervalSeconds,
					},
				},
			},
		}

		if v.Spec.TargetNode != nil {
			newClusterRule.Spec.NodeRule = &v3.NodeRule{
				NodeName:     v.Spec.TargetNode.NodeName,
				Selector:     v.Spec.TargetNode.Selector,
				Condition:    v.Spec.TargetNode.Condition,
				MemThreshold: v.Spec.TargetNode.MemThreshold,
				CPUThreshold: v.Spec.TargetNode.CPUThreshold,
			}
		}

		if v.Spec.TargetEvent != nil {
			newClusterRule.Spec.EventRule = &v3.EventRule{
				EventType:    v.Spec.TargetEvent.EventType,
				ResourceKind: v.Spec.TargetEvent.ResourceKind,
			}
		}

		if v.Spec.TargetSystemService != nil {
			newClusterRule.Spec.SystemServiceRule = &v3.SystemServiceRule{
				Condition: v.Spec.TargetSystemService.Condition,
			}
		}

		oldClusterRule, err := l.clusterAlertRules.Get(newClusterRule.Name, metav1.GetOptions{})
		if err != nil {
			if !apierrors.IsNotFound(err) {
				return "", fmt.Errorf("migrate %s:%s failed, get alert rule failed, %v", v.Namespace, v.Name, err)
			}

			if _, err = l.clusterAlertRules.Create(newClusterRule); err != nil && !apierrors.IsAlreadyExists(err) {
				return "", fmt.Errorf("migrate %s:%s failed, create alert rule failed, %v", v.Namespace, v.Name, err)
			}
		} else {
			updatedClusterRule := oldClusterRule.DeepCopy()
			updatedClusterRule.Spec = newClusterRule.Spec
			if _, err := l.clusterAlertRules.Update(updatedClusterRule); err != nil {
				return "", fmt.Errorf("migrate %s:%s failed, update alert rule failed, %v", v.Namespace, v.Name, err)
			}
		}

		legacyGroup := &v3.ClusterAlertGroup{
			ObjectMeta: metav1.ObjectMeta{
				Name:      migrationGroupName,
				Namespace: l.clusterName,
			},
			Spec: v3.ClusterGroupSpec{
				ClusterName: l.clusterName,
				CommonGroupField: v3.CommonGroupField{
					DisplayName: "Migrate group",
					Description: "Migrate alert from last version",
				},
				Recipients: v.Spec.Recipients,
			},
		}

		legacyGroup, err = l.clusterAlertGroups.Create(legacyGroup)
		if err != nil && !apierrors.IsAlreadyExists(err) {
			return "", fmt.Errorf("migrate failed, create alert group %s:%s failed, %v", l.clusterName, migrationGroupName, err)
		}
	}

	oldProjectAlert, err := l.oldProjectAlerts.List(metav1.ListOptions{})
	if err != nil {
		return "", fmt.Errorf("get old project alert failed, %s", err)
	}

	oldProjectAlertGroup := make(map[string][]*v3.ProjectAlert)
	for _, v := range oldProjectAlert.Items {
		if controller.ObjectInCluster(l.clusterName, v) {
			oldProjectAlertGroup[v.Spec.ProjectName] = append(oldProjectAlertGroup[v.Spec.ProjectName], &v)
		}
	}

	for projectID, oldAlerts := range oldProjectAlertGroup {
		_, projectName := ref.Parse(projectID)

		for _, v := range oldAlerts {
			migrationGroupName := fmt.Sprintf("migrate-group-%s", v.Name)
			groupID := alertutil.GetGroupID(projectName, migrationGroupName)

			migrationRuleName := fmt.Sprintf("migrate-rule-%s", v.Name)
			newProjectRule := &v3.ProjectAlertRule{
				ObjectMeta: metav1.ObjectMeta{
					Name:      migrationRuleName,
					Namespace: projectName,
				},
				Spec: v3.ProjectAlertRuleSpec{
					ProjectName: projectName,
					GroupName:   groupID,
					CommonRuleField: v3.CommonRuleField{
						DisplayName: v.Spec.DisplayName,
						Severity:    v.Spec.Severity,
						TimingField: v3.TimingField{
							GroupIntervalSeconds:  v.Spec.InitialWaitSeconds,
							RepeatIntervalSeconds: v.Spec.RepeatIntervalSeconds,
						},
					},
				},
			}

			if v.Spec.TargetPod != nil {
				newProjectRule.Spec.PodRule = &v3.PodRule{
					PodName:                v.Spec.TargetPod.PodName,
					Condition:              v.Spec.TargetPod.Condition,
					RestartTimes:           v.Spec.TargetPod.RestartTimes,
					RestartIntervalSeconds: v.Spec.TargetPod.RestartIntervalSeconds,
				}
			}

			if v.Spec.TargetWorkload != nil {
				newProjectRule.Spec.WorkloadRule = &v3.WorkloadRule{
					WorkloadID:          v.Spec.TargetWorkload.WorkloadID,
					Selector:            v.Spec.TargetWorkload.Selector,
					AvailablePercentage: v.Spec.TargetWorkload.AvailablePercentage,
				}
			}

			oldProjectRule, err := l.projectAlertRules.GetNamespaced(projectName, newProjectRule.Name, metav1.GetOptions{})
			if err != nil {
				if !apierrors.IsNotFound(err) {
					return "", fmt.Errorf("migrate %s:%s failed, get alert rule failed, %v", v.Namespace, v.Name, err)
				}

				if _, err = l.projectAlertRules.Create(newProjectRule); err != nil && !apierrors.IsAlreadyExists(err) {
					return "", fmt.Errorf("migrate %s:%s failed, create alert rule failed, %v", v.Namespace, v.Name, err)
				}
			} else {
				updatedProjectRule := oldProjectRule.DeepCopy()
				updatedProjectRule.Spec = newProjectRule.Spec
				if _, err := l.projectAlertRules.Update(updatedProjectRule); err != nil {
					return "", fmt.Errorf("migrate %s:%s failed, update alert rule failed, %v", v.Namespace, v.Name, err)
				}
			}

			legacyGroup := &v3.ProjectAlertGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      migrationGroupName,
					Namespace: projectName,
				},
				Spec: v3.ProjectGroupSpec{
					ProjectName: projectID,
					CommonGroupField: v3.CommonGroupField{
						DisplayName: "Migrate group",
						Description: "Migrate alert from last version",
					},
					Recipients: v.Spec.Recipients,
				},
			}

			legacyGroup, err = l.projectAlertGroups.Create(legacyGroup)
			if err != nil && !apierrors.IsAlreadyExists(err) {
				return "", fmt.Errorf("create migrate alert group %s:%s failed, %v", legacyGroup.Namespace, legacyGroup.Name, err)
			}
		}

	}

	//clean up
	set := fields.Set{
		"metadata.namespace": l.clusterName,
	}
	if err := l.oldClusterAlerts.DeleteCollection(&metav1.DeleteOptions{}, metav1.ListOptions{
		FieldSelector: set.AsSelector().String(),
	}); err != nil && !apierrors.IsNotFound(err) {
		return "", fmt.Errorf("migrate failed, delete old cluster alert in namespace %s failded, %v", l.clusterName, err)
	}

	for projectName := range oldProjectAlertGroup {
		if err := l.oldProjectAlerts.DeleteCollection(&metav1.DeleteOptions{}, metav1.ListOptions{
			FieldSelector: fmt.Sprintf("metadata.namespace=%s", projectName),
		}); err != nil && !apierrors.IsNotFound(err) {
			return "", fmt.Errorf("migrate failed, delete old project alert in namespace %s failded, %v", projectName, err)
		}
	}
	//upgrade old app
	defaultSystemProjects, err := l.projectLister.List(metav1.NamespaceAll, labels.Set(systemProjectLabel).AsSelector())
	if err != nil {
		return "", fmt.Errorf("list system project failed, %v", err)
	}

	if len(defaultSystemProjects) == 0 {
		return "", fmt.Errorf("get system project failed")
	}

	systemProject := defaultSystemProjects[0]
	if systemProject == nil {
		return "", fmt.Errorf("get system project failed")
	}

	app, err := l.apps.GetNamespaced(systemProject.Name, appName, metav1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			return NewVersion, nil
		}
		return "", fmt.Errorf("get app %s:%s failed, %v", systemProject.Name, appName, err)
	}
	newApp := app.DeepCopy()
	newApp.Spec.ExternalID = newCatalogID

	if !reflect.DeepEqual(newApp, app) {
		if _, err = l.apps.Update(newApp); err != nil {
			return "", fmt.Errorf("update app %s:%s failed, %v", app.Namespace, app.Name, err)
		}
	}

	return NewVersion, nil
}
<commit_msg>Not delete old alert data when upgrade<commit_after>package deployer

import (
	"context"
	"fmt"
	"reflect"
	"strings"

	"github.com/rancher/rancher/pkg/ref"
	"github.com/rancher/types/apis/core/v1"

	"github.com/rancher/norman/controller"
	alertutil "github.com/rancher/rancher/pkg/controllers/user/alert/common"
	"github.com/rancher/rancher/pkg/controllers/user/systemimage"
	monitorutil "github.com/rancher/rancher/pkg/monitoring"
	"github.com/rancher/types/apis/apps/v1beta2"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	projectv3 "github.com/rancher/types/apis/project.cattle.io/v3"

	"github.com/rancher/rancher/pkg/controllers/user/helm/common"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/types/config"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
)

var (
	serviceName = "alerting"
)

type alertService struct {
	clusterName        string
	apps               projectv3.AppInterface
	deployments        v1beta2.DeploymentInterface
	oldClusterAlerts   v3.ClusterAlertInterface
	oldProjectAlerts   v3.ProjectAlertInterface
	clusterAlertGroups v3.ClusterAlertGroupInterface
	projectAlertGroups v3.ProjectAlertGroupInterface
	clusterAlertRules  v3.ClusterAlertRuleInterface
	projectAlertRules  v3.ProjectAlertRuleInterface
	projectLister      v3.ProjectLister
	namespaces         v1.NamespaceInterface
	appDeployer        *appDeployer
}

func init() {
	systemimage.RegisterSystemService(serviceName, &alertService{})
}

func (l *alertService) Init(ctx context.Context, cluster *config.UserContext) {
	ad := &appDeployer{
		appsGetter:       cluster.Management.Project,
		namespaces:       cluster.Management.Core.Namespaces(metav1.NamespaceAll),
		secrets:          cluster.Core.Secrets(metav1.NamespaceAll),
		templateVersions: cluster.Management.Management.TemplateVersions(metav1.NamespaceAll),
	}

	l.clusterName = cluster.ClusterName
	l.deployments = cluster.Apps.Deployments("")
	l.oldClusterAlerts = cluster.Management.Management.ClusterAlerts(cluster.ClusterName)
	l.oldProjectAlerts = cluster.Management.Management.ProjectAlerts(metav1.NamespaceAll)
	l.clusterAlertGroups = cluster.Management.Management.ClusterAlertGroups(cluster.ClusterName)
	l.projectAlertGroups = cluster.Management.Management.ProjectAlertGroups(metav1.NamespaceAll)
	l.clusterAlertRules = cluster.Management.Management.ClusterAlertRules(cluster.ClusterName)
	l.projectAlertRules = cluster.Management.Management.ProjectAlertRules(metav1.NamespaceAll)
	l.projectLister = cluster.Management.Management.Projects(cluster.ClusterName).Controller().Lister()
	l.apps = cluster.Management.Project.Apps(metav1.NamespaceAll)
	l.appDeployer = ad

}

func (l *alertService) Version() (string, error) {
	catalogID := settings.SystemMonitoringCatalogID.Get()
	templateVersionID, err := common.ParseExternalID(catalogID)
	if err != nil {
		return "", fmt.Errorf("get system monitor catalog version failed, %v", err)
	}
	return templateVersionID, nil
}

func (l *alertService) Upgrade(currentVersion string) (string, error) {
	newCatalogID := settings.SystemMonitoringCatalogID.Get()

	NewVersion, err := common.ParseExternalID(newCatalogID)
	if currentVersion == NewVersion {
		return currentVersion, nil
	}

	appName, _ := monitorutil.ClusterAlertManagerInfo()
	//migrate legacy
	if !strings.Contains(currentVersion, "system-library-rancher-monitoring") {
		if err := l.migrateLegacyClusterAlert(); err != nil {
			return "", err
		}

		if err := l.migrateLegacyProjectAlert(); err != nil {
			return "", err
		}
	}

	//upgrade old app
	defaultSystemProjects, err := l.projectLister.List(metav1.NamespaceAll, labels.Set(systemProjectLabel).AsSelector())
	if err != nil {
		return "", fmt.Errorf("list system project failed, %v", err)
	}

	if len(defaultSystemProjects) == 0 {
		return "", fmt.Errorf("get system project failed")
	}

	systemProject := defaultSystemProjects[0]
	if systemProject == nil {
		return "", fmt.Errorf("get system project failed")
	}
	app, err := l.apps.GetNamespaced(systemProject.Name, appName, metav1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			return NewVersion, nil
		}
		return "", fmt.Errorf("get app %s:%s failed, %v", systemProject.Name, appName, err)
	}
	newApp := app.DeepCopy()
	newApp.Spec.ExternalID = newCatalogID

	if !reflect.DeepEqual(newApp, app) {
		if _, err = l.apps.Update(newApp); err != nil {
			return "", fmt.Errorf("update app %s:%s failed, %v", app.Namespace, app.Name, err)
		}
	}
	return NewVersion, nil
}

func (l *alertService) migrateLegacyClusterAlert() error {
	oldClusterAlert, err := l.oldClusterAlerts.List(metav1.ListOptions{})
	if err != nil {
		return fmt.Errorf("get old cluster alert failed, %s", err)
	}
	for _, v := range oldClusterAlert.Items {
		migrationGroupName := fmt.Sprintf("migrate-group-%s", v.Name)
		groupID := alertutil.GetGroupID(l.clusterName, migrationGroupName)

		name := fmt.Sprintf("migrate-%s", v.Name)
		newClusterRule := &v3.ClusterAlertRule{
			ObjectMeta: metav1.ObjectMeta{
				Name:      name,
				Namespace: l.clusterName,
			},
			Spec: v3.ClusterAlertRuleSpec{
				ClusterName: l.clusterName,
				GroupName:   groupID,
				CommonRuleField: v3.CommonRuleField{
					DisplayName: v.Spec.DisplayName,
					Severity:    v.Spec.Severity,
					TimingField: v3.TimingField{
						GroupIntervalSeconds:  v.Spec.InitialWaitSeconds,
						RepeatIntervalSeconds: v.Spec.RepeatIntervalSeconds,
					},
				},
			},
		}

		if v.Spec.TargetNode != nil {
			newClusterRule.Spec.NodeRule = &v3.NodeRule{
				NodeName:     v.Spec.TargetNode.NodeName,
				Selector:     v.Spec.TargetNode.Selector,
				Condition:    v.Spec.TargetNode.Condition,
				MemThreshold: v.Spec.TargetNode.MemThreshold,
				CPUThreshold: v.Spec.TargetNode.CPUThreshold,
			}
		}

		if v.Spec.TargetEvent != nil {
			newClusterRule.Spec.EventRule = &v3.EventRule{
				EventType:    v.Spec.TargetEvent.EventType,
				ResourceKind: v.Spec.TargetEvent.ResourceKind,
			}
		}

		if v.Spec.TargetSystemService != nil {
			newClusterRule.Spec.SystemServiceRule = &v3.SystemServiceRule{
				Condition: v.Spec.TargetSystemService.Condition,
			}
		}

		oldClusterRule, err := l.clusterAlertRules.Get(newClusterRule.Name, metav1.GetOptions{})
		if err != nil {
			if !apierrors.IsNotFound(err) {
				return fmt.Errorf("migrate %s:%s failed, get alert rule failed, %v", v.Namespace, v.Name, err)
			}

			if _, err = l.clusterAlertRules.Create(newClusterRule); err != nil && !apierrors.IsAlreadyExists(err) {
				return fmt.Errorf("migrate %s:%s failed, create alert rule failed, %v", v.Namespace, v.Name, err)
			}
		} else {
			updatedClusterRule := oldClusterRule.DeepCopy()
			updatedClusterRule.Spec = newClusterRule.Spec
			if _, err := l.clusterAlertRules.Update(updatedClusterRule); err != nil {
				return fmt.Errorf("migrate %s:%s failed, update alert rule failed, %v", v.Namespace, v.Name, err)
			}
		}
		legacyGroup := &v3.ClusterAlertGroup{
			ObjectMeta: metav1.ObjectMeta{
				Name:      migrationGroupName,
				Namespace: l.clusterName,
			},
			Spec: v3.ClusterGroupSpec{
				ClusterName: l.clusterName,
				CommonGroupField: v3.CommonGroupField{
					DisplayName: "Migrate group",
					Description: "Migrate alert from last version",
				},
				Recipients: v.Spec.Recipients,
			},
		}

		legacyGroup, err = l.clusterAlertGroups.Create(legacyGroup)
		if err != nil && !apierrors.IsAlreadyExists(err) {
			return fmt.Errorf("migrate failed, create alert group %s:%s failed, %v", l.clusterName, migrationGroupName, err)
		}
	}
	return nil
}

func (l *alertService) migrateLegacyProjectAlert() error {
	oldProjectAlert, err := l.oldProjectAlerts.List(metav1.ListOptions{})
	if err != nil {
		return fmt.Errorf("get old project alert failed, %s", err)
	}

	oldProjectAlertGroup := make(map[string][]*v3.ProjectAlert)
	for _, v := range oldProjectAlert.Items {
		if controller.ObjectInCluster(l.clusterName, v) {
			oldProjectAlertGroup[v.Spec.ProjectName] = append(oldProjectAlertGroup[v.Spec.ProjectName], &v)
		}
	}

	for projectID, oldAlerts := range oldProjectAlertGroup {
		_, projectName := ref.Parse(projectID)

		for _, v := range oldAlerts {
			migrationGroupName := fmt.Sprintf("migrate-group-%s", v.Name)
			groupID := alertutil.GetGroupID(projectName, migrationGroupName)

			migrationRuleName := fmt.Sprintf("migrate-rule-%s", v.Name)
			newProjectRule := &v3.ProjectAlertRule{
				ObjectMeta: metav1.ObjectMeta{
					Name:      migrationRuleName,
					Namespace: projectName,
				},
				Spec: v3.ProjectAlertRuleSpec{
					ProjectName: projectName,
					GroupName:   groupID,
					CommonRuleField: v3.CommonRuleField{
						DisplayName: v.Spec.DisplayName,
						Severity:    v.Spec.Severity,
						TimingField: v3.TimingField{
							GroupIntervalSeconds:  v.Spec.InitialWaitSeconds,
							RepeatIntervalSeconds: v.Spec.RepeatIntervalSeconds,
						},
					},
				},
			}

			if v.Spec.TargetPod != nil {
				newProjectRule.Spec.PodRule = &v3.PodRule{
					PodName:                v.Spec.TargetPod.PodName,
					Condition:              v.Spec.TargetPod.Condition,
					RestartTimes:           v.Spec.TargetPod.RestartTimes,
					RestartIntervalSeconds: v.Spec.TargetPod.RestartIntervalSeconds,
				}
			}

			if v.Spec.TargetWorkload != nil {
				newProjectRule.Spec.WorkloadRule = &v3.WorkloadRule{
					WorkloadID:          v.Spec.TargetWorkload.WorkloadID,
					Selector:            v.Spec.TargetWorkload.Selector,
					AvailablePercentage: v.Spec.TargetWorkload.AvailablePercentage,
				}
			}

			oldProjectRule, err := l.projectAlertRules.GetNamespaced(projectName, newProjectRule.Name, metav1.GetOptions{})
			if err != nil {
				if !apierrors.IsNotFound(err) {
					return fmt.Errorf("migrate %s:%s failed, get alert rule failed, %v", v.Namespace, v.Name, err)
				}

				if _, err = l.projectAlertRules.Create(newProjectRule); err != nil && !apierrors.IsAlreadyExists(err) {
					return fmt.Errorf("migrate %s:%s failed, create alert rule failed, %v", v.Namespace, v.Name, err)
				}
			} else {
				updatedProjectRule := oldProjectRule.DeepCopy()
				updatedProjectRule.Spec = newProjectRule.Spec
				if _, err := l.projectAlertRules.Update(updatedProjectRule); err != nil {
					return fmt.Errorf("migrate %s:%s failed, update alert rule failed, %v", v.Namespace, v.Name, err)
				}
			}

			legacyGroup := &v3.ProjectAlertGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      migrationGroupName,
					Namespace: projectName,
				},
				Spec: v3.ProjectGroupSpec{
					ProjectName: projectID,
					CommonGroupField: v3.CommonGroupField{
						DisplayName: "Migrate group",
						Description: "Migrate alert from last version",
					},
					Recipients: v.Spec.Recipients,
				},
			}

			legacyGroup, err = l.projectAlertGroups.Create(legacyGroup)
			if err != nil && !apierrors.IsAlreadyExists(err) {
				return fmt.Errorf("create migrate alert group %s:%s failed, %v", legacyGroup.Namespace, legacyGroup.Name, err)
			}
		}
	}
	return nil
}
<|endoftext|>


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

package enumerator // import "go.bug.st/serial.v1/enumerator"

// #cgo LDFLAGS: -framework CoreFoundation -framework IOKit -fconstant-cfstrings
// #include <IOKit/IOKitLib.h>
// #include <CoreFoundation/CoreFoundation.h>
// #include <stdlib.h>
import "C"
import (
	"errors"
	"fmt"
	"unsafe"
)

func nativeGetDetailedPortsList() ([]*PortDetails, error) {
	var ports []*PortDetails

	services, err := getAllServices("IOSerialBSDClient")
	if err != nil {
		return nil, &PortEnumerationError{causedBy: err}
	}
	for _, service := range services {
		defer service.Release()

		port, err := extractPortInfo(C.io_registry_entry_t(service))
		if err != nil {
			return nil, &PortEnumerationError{causedBy: err}
		}
		ports = append(ports, port)
	}
	return ports, nil
}

func extractPortInfo(service C.io_registry_entry_t) (*PortDetails, error) {
	name, err := service.GetStringProperty("IOCalloutDevice")
	if err != nil {
		return nil, fmt.Errorf("Error extracting port info from device: %s", err.Error())
	}
	port := &PortDetails{}
	port.Name = name
	port.IsUSB = false

	usbDevice := service
	for usbDevice.GetClass() != "IOUSBDevice" {
		if usbDevice, err = usbDevice.GetParent("IOService"); err != nil {
			break
		}
	}
	if err == nil {
		// It's an IOUSBDevice
		vid, _ := usbDevice.GetIntProperty("idVendor", C.kCFNumberSInt16Type)
		pid, _ := usbDevice.GetIntProperty("idProduct", C.kCFNumberSInt16Type)
		serialNumber, _ := usbDevice.GetStringProperty("USB Serial Number")
		//product, _ := usbDevice.GetStringProperty("USB Product Name")
		//manufacturer, _ := usbDevice.GetStringProperty("USB Vendor Name")
		//fmt.Println(product + " - " + manufacturer)

		port.IsUSB = true
		port.VID = fmt.Sprintf("%04X", vid)
		port.PID = fmt.Sprintf("%04X", pid)
		port.SerialNumber = serialNumber
	}
	return port, nil
}

func getAllServices(serviceType string) ([]C.io_object_t, error) {
	i, err := getMatchingServices(serviceMatching(serviceType))
	if err != nil {
		return nil, err
	}
	defer i.Release()

	var services []C.io_object_t
	tries := 0
	for tries < 5 {
		// Extract all elements from iterator
		if service, ok := i.Next(); ok {
			services = append(services, service)
			continue
		}
		// If iterator is still valid return the result
		if i.IsValid() {
			return services, nil
		}
		// Otherwise empty the result and retry
		for _, s := range services {
			s.Release()
		}
		services = []C.io_object_t{}
		i.Reset()
		tries++
	}
	// Give up if the iteration continues to fail...
	return nil, fmt.Errorf("IOServiceGetMatchingServices failed, data changed while iterating")
}

// serviceMatching create a matching dictionary that specifies an IOService class match.
func serviceMatching(serviceType string) C.CFMutableDictionaryRef {
	t := C.CString(serviceType)
	defer C.free(unsafe.Pointer(t))
	return C.IOServiceMatching(t)
}

// getMatchingServices look up registered IOService objects that match a matching dictionary.
func getMatchingServices(matcher C.CFMutableDictionaryRef) (C.io_iterator_t, error) {
	var i C.io_iterator_t
	err := C.IOServiceGetMatchingServices(C.kIOMasterPortDefault, matcher, &i)
	if err != C.KERN_SUCCESS {
		return 0, fmt.Errorf("IOServiceGetMatchingServices failed (code %d)", err)
	}
	return i, nil
}

// CFStringRef

func cfStringCreateWithString(s string) C.CFStringRef {
	c := C.CString(s)
	defer C.free(unsafe.Pointer(c))
	return C.CFStringCreateWithCString(
		C.kCFAllocatorDefault, c, C.kCFStringEncodingMacRoman)
}

// io_registry_entry_t

func (me *C.io_registry_entry_t) GetParent(plane string) (C.io_registry_entry_t, error) {
	cPlane := C.CString(plane)
	defer C.free(unsafe.Pointer(cPlane))
	var parent C.io_registry_entry_t
	err := C.IORegistryEntryGetParentEntry(*me, cPlane, &parent)
	if err != 0 {
		return 0, errors.New("No parent device available")
	}
	return parent, nil
}

func (me *C.io_registry_entry_t) GetClass() string {
	obj := (*C.io_object_t)(me)
	return obj.GetClass()
}

func (me *C.io_registry_entry_t) GetStringProperty(key string) (string, error) {
	k := cfStringCreateWithString(key)
	defer C.CFRelease(C.CFTypeRef(k))
	property := C.IORegistryEntryCreateCFProperty(*me, k, C.kCFAllocatorDefault, 0)
	if property == nil {
		return "", errors.New("Property not found: " + key)
	}
	defer C.CFRelease(property)

	if ptr := C.CFStringGetCStringPtr(property, 0); ptr != nil {
		return C.GoString(ptr), nil
	}
	// in certain circumstances CFStringGetCStringPtr may return NULL
	// and we must retrieve the string by copy
	buff := make([]C.char, 1024)
	if C.CFStringGetCString(property, &buff[0], 1024, 0) != C.true {
		return "", fmt.Errorf("Property '%s' can't be converted", key)
	}
	return C.GoString(&buff[0]), nil
}

func (me *C.io_registry_entry_t) GetIntProperty(key string, intType C.CFNumberType) (int, error) {
	k := cfStringCreateWithString(key)
	defer C.CFRelease(C.CFTypeRef(k))
	property := C.IORegistryEntryCreateCFProperty(*me, k, C.kCFAllocatorDefault, 0)
	if property == nil {
		return 0, errors.New("Property not found: " + key)
	}
	defer C.CFRelease(property)
	var res int
	if C.CFNumberGetValue(property, intType, unsafe.Pointer(&res)) != C.true {
		return res, fmt.Errorf("Property '%s' can't be converted or has been truncated", key)
	}
	return res, nil
}

// io_iterator_t

// IsValid checks if an iterator is still valid.
// Some iterators will be made invalid if changes are made to the
// structure they are iterating over. This function checks the iterator
// is still valid and should be called when Next returns zero.
// An invalid iterator can be Reset and the iteration restarted.
func (me *C.io_iterator_t) IsValid() bool {
	return C.IOIteratorIsValid(*me) == C.true
}

func (me *C.io_iterator_t) Reset() {
	C.IOIteratorReset(*me)
}

func (me *C.io_iterator_t) Next() (C.io_object_t, bool) {
	res := C.IOIteratorNext(*me)
	return res, res != 0
}

func (me *C.io_iterator_t) Release() {
	C.IOObjectRelease(C.io_object_t(*me))
}

// io_object_t

func (me *C.io_object_t) Release() {
	C.IOObjectRelease(*me)
}

func (me *C.io_object_t) GetClass() string {
	class := make([]C.char, 1024)
	C.IOObjectGetClass(*me, &class[0])
	return C.GoString(&class[0])
}
<commit_msg>Fixed build of USB enumerator for go compiler >=1.8.3 (@facchinm)<commit_after>//
// Copyright 2014-2017 Cristian Maglie. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package enumerator // import "go.bug.st/serial.v1/enumerator"

// #cgo LDFLAGS: -framework CoreFoundation -framework IOKit -fconstant-cfstrings
// #include <IOKit/IOKitLib.h>
// #include <CoreFoundation/CoreFoundation.h>
// #include <stdlib.h>
import "C"
import (
	"errors"
	"fmt"
	"unsafe"
)

func nativeGetDetailedPortsList() ([]*PortDetails, error) {
	var ports []*PortDetails

	services, err := getAllServices("IOSerialBSDClient")
	if err != nil {
		return nil, &PortEnumerationError{causedBy: err}
	}
	for _, service := range services {
		defer service.Release()

		port, err := extractPortInfo(C.io_registry_entry_t(service))
		if err != nil {
			return nil, &PortEnumerationError{causedBy: err}
		}
		ports = append(ports, port)
	}
	return ports, nil
}

func extractPortInfo(service C.io_registry_entry_t) (*PortDetails, error) {
	name, err := service.GetStringProperty("IOCalloutDevice")
	if err != nil {
		return nil, fmt.Errorf("Error extracting port info from device: %s", err.Error())
	}
	port := &PortDetails{}
	port.Name = name
	port.IsUSB = false

	usbDevice := service
	for usbDevice.GetClass() != "IOUSBDevice" {
		if usbDevice, err = usbDevice.GetParent("IOService"); err != nil {
			break
		}
	}
	if err == nil {
		// It's an IOUSBDevice
		vid, _ := usbDevice.GetIntProperty("idVendor", C.kCFNumberSInt16Type)
		pid, _ := usbDevice.GetIntProperty("idProduct", C.kCFNumberSInt16Type)
		serialNumber, _ := usbDevice.GetStringProperty("USB Serial Number")
		//product, _ := usbDevice.GetStringProperty("USB Product Name")
		//manufacturer, _ := usbDevice.GetStringProperty("USB Vendor Name")
		//fmt.Println(product + " - " + manufacturer)

		port.IsUSB = true
		port.VID = fmt.Sprintf("%04X", vid)
		port.PID = fmt.Sprintf("%04X", pid)
		port.SerialNumber = serialNumber
	}
	return port, nil
}

func getAllServices(serviceType string) ([]C.io_object_t, error) {
	i, err := getMatchingServices(serviceMatching(serviceType))
	if err != nil {
		return nil, err
	}
	defer i.Release()

	var services []C.io_object_t
	tries := 0
	for tries < 5 {
		// Extract all elements from iterator
		if service, ok := i.Next(); ok {
			services = append(services, service)
			continue
		}
		// If iterator is still valid return the result
		if i.IsValid() {
			return services, nil
		}
		// Otherwise empty the result and retry
		for _, s := range services {
			s.Release()
		}
		services = []C.io_object_t{}
		i.Reset()
		tries++
	}
	// Give up if the iteration continues to fail...
	return nil, fmt.Errorf("IOServiceGetMatchingServices failed, data changed while iterating")
}

// serviceMatching create a matching dictionary that specifies an IOService class match.
func serviceMatching(serviceType string) C.CFMutableDictionaryRef {
	t := C.CString(serviceType)
	defer C.free(unsafe.Pointer(t))
	return C.IOServiceMatching(t)
}

// getMatchingServices look up registered IOService objects that match a matching dictionary.
func getMatchingServices(matcher C.CFMutableDictionaryRef) (C.io_iterator_t, error) {
	var i C.io_iterator_t
	err := C.IOServiceGetMatchingServices(C.kIOMasterPortDefault, matcher, &i)
	if err != C.KERN_SUCCESS {
		return 0, fmt.Errorf("IOServiceGetMatchingServices failed (code %d)", err)
	}
	return i, nil
}

// CFStringRef

func cfStringCreateWithString(s string) C.CFStringRef {
	c := C.CString(s)
	defer C.free(unsafe.Pointer(c))
	return C.CFStringCreateWithCString(
		C.kCFAllocatorDefault, c, C.kCFStringEncodingMacRoman)
}

// io_registry_entry_t

func (me *C.io_registry_entry_t) GetParent(plane string) (C.io_registry_entry_t, error) {
	cPlane := C.CString(plane)
	defer C.free(unsafe.Pointer(cPlane))
	var parent C.io_registry_entry_t
	err := C.IORegistryEntryGetParentEntry(*me, cPlane, &parent)
	if err != 0 {
		return 0, errors.New("No parent device available")
	}
	return parent, nil
}

func (me *C.io_registry_entry_t) GetClass() string {
	obj := (*C.io_object_t)(me)
	return obj.GetClass()
}

func (me *C.io_registry_entry_t) GetStringProperty(key string) (string, error) {
	k := cfStringCreateWithString(key)
	defer C.CFRelease(C.CFTypeRef(k))
	property := C.IORegistryEntryCreateCFProperty(*me, k, C.kCFAllocatorDefault, 0)
	if property == nil {
		return "", errors.New("Property not found: " + key)
	}
	defer C.CFRelease(property)

	if ptr := C.CFStringGetCStringPtr((C.CFStringRef)(unsafe.Pointer(property)), 0); ptr != nil {
		return C.GoString(ptr), nil
	}
	// in certain circumstances CFStringGetCStringPtr may return NULL
	// and we must retrieve the string by copy
	buff := make([]C.char, 1024)
	if C.CFStringGetCString((C.CFStringRef)(property), &buff[0], 1024, 0) != C.true {
		return "", fmt.Errorf("Property '%s' can't be converted", key)
	}
	return C.GoString(&buff[0]), nil
}

func (me *C.io_registry_entry_t) GetIntProperty(key string, intType C.CFNumberType) (int, error) {
	k := cfStringCreateWithString(key)
	defer C.CFRelease(C.CFTypeRef(k))
	property := C.IORegistryEntryCreateCFProperty(*me, k, C.kCFAllocatorDefault, 0)
	if property == nil {
		return 0, errors.New("Property not found: " + key)
	}
	defer C.CFRelease(property)
	var res int
	if C.CFNumberGetValue((C.CFNumberRef)(property), intType, unsafe.Pointer(&res)) != C.true {
		return res, fmt.Errorf("Property '%s' can't be converted or has been truncated", key)
	}
	return res, nil
}

// io_iterator_t

// IsValid checks if an iterator is still valid.
// Some iterators will be made invalid if changes are made to the
// structure they are iterating over. This function checks the iterator
// is still valid and should be called when Next returns zero.
// An invalid iterator can be Reset and the iteration restarted.
func (me *C.io_iterator_t) IsValid() bool {
	return C.IOIteratorIsValid(*me) == C.true
}

func (me *C.io_iterator_t) Reset() {
	C.IOIteratorReset(*me)
}

func (me *C.io_iterator_t) Next() (C.io_object_t, bool) {
	res := C.IOIteratorNext(*me)
	return res, res != 0
}

func (me *C.io_iterator_t) Release() {
	C.IOObjectRelease(C.io_object_t(*me))
}

// io_object_t

func (me *C.io_object_t) Release() {
	C.IOObjectRelease(*me)
}

func (me *C.io_object_t) GetClass() string {
	class := make([]C.char, 1024)
	C.IOObjectGetClass(*me, &class[0])
	return C.GoString(&class[0])
}
<|endoftext|>