<commit_before>// MIT License
//
// Copyright (c) 2019 Stefan Wichmann
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	hue "github.com/stefanwichmann/go.hue"
)

// HueBridge represents the Philips Hue bridge in
// your system.
// It is used to communicate with all devices.
type HueBridge struct {
	bridge   hue.Bridge
	BridgeIP string
	Username string
	Version  int
}

const hueBridgeAppName = "kelvin"

// InitializeBridge creates and returns an initialized HueBridge.
// If you have a valid configuration this will be used. Otherwise a local
// discovery will be started, followed by a user registration on your bridge.
func (bridge *HueBridge) InitializeBridge(configuration *Configuration) error {
	err := bridge.discover(configuration.Bridge.IP)
	if err != nil {
		return err
	}
	configuration.Bridge.IP = bridge.BridgeIP

	if configuration.Bridge.Username != "" {
		log.Debugf("⌘ Found bridge username in configuration: %s", configuration.Bridge.Username)
		bridge.Username = configuration.Bridge.Username
	} else {
		log.Debugf("⌘ No username found in bridge configuration. Starting registration...")
		err := bridge.register()
		if err != nil {
			return err
		}
		log.Debugf("⌘ Saving new username in bridge configuration: %s", bridge.Username)
		configuration.Bridge.Username = bridge.Username
	}

	log.Debugf("⌘ Connecting to bridge %s with username %s", bridge.BridgeIP, bridge.Username)
	err = bridge.connect()
	if err != nil {
		return err
	}
	log.Println("⌘ Connection to bridge established")
	bridge.validateSofwareVersion()

	err = bridge.populateSchedule(configuration)
	return err
}

// Lights return all known lights on your bridge.
func (bridge *HueBridge) Lights() ([]*Light, error) {
	var lights []*Light
	hueLights, err := bridge.bridge.GetAllLights()
	if err != nil {
		return lights, err
	}

	for _, hueLight := range hueLights {
		var light Light
		light.ID, err = strconv.Atoi(hueLight.Id)
		if err != nil {
			return lights, err
		}

		light.HueLight.HueLight = *hueLight
		light.HueLight.initialize(hueLight.Attributes)
		light.Name = light.HueLight.Name
		light.Reachable = light.HueLight.Reachable
		light.On = light.HueLight.On

		lights = append(lights, &light)
	}

	sort.Slice(lights, func(i, j int) bool { return lights[i].ID < lights[j].ID })
	return lights, nil
}

// LightStates returns the current state for lights on the bridge
func (bridge *HueBridge) LightStates() (map[int]hue.LightAttributes, error) {
	var states = make(map[int]hue.LightAttributes)
	hueLights, err := bridge.bridge.GetAllLights()
	if err != nil {
		return states, err
	}

	for _, hueLight := range hueLights {
		lightID, err := strconv.Atoi(hueLight.Id)
		if err != nil {
			return states, err
		}

		states[lightID] = hueLight.Attributes
	}

	return states, nil
}

func (bridge *HueBridge) discover(ip string) error {
	if ip != "" {
		// we have a known IP address. Validate if it points to a reachable bridge
		bridge.BridgeIP = ip
		err := bridge.validateBridge()
		if err != nil {
			return err
		}
		return nil
	}
	log.Debugf("⌘ Starting bridge discovery")
	bridges, err := hue.DiscoverBridges(false)
	if err != nil {
		bridge.BridgeIP = ""
		return err
	}
	if len(bridges) == 0 {
		bridge.BridgeIP = ""
		return errors.New("Bridge discovery failed. Please configure manually in config.json")
	}
	for _, candidate := range bridges {
		bridge.BridgeIP = candidate.IpAddr
		err := bridge.validateBridge()
		if err == nil {
			log.Printf("⌘ Found bridge at %s", bridge.BridgeIP)
			return nil
		}
	}
	bridge.BridgeIP = ""
	return errors.New("Bridge discovery failed. Please configure manually in config.json")
}

func (bridge *HueBridge) register() error {
	if bridge.BridgeIP == "" {
		return errors.New("Registration at bridge not possible because no IP is configured. Start discovery first or enter manually")
	}

	bridge.bridge = *hue.NewBridge(bridge.BridgeIP, "")
	log.Printf("⌘ Starting user registration.")
	log.Warningf("⌘ PLEASE PUSH THE BLUE BUTTON ON YOUR HUE BRIDGE")
	for {
		time.Sleep(5 * time.Second)

		// try user creation, will fail if the button wasn't pressed.
		err := bridge.bridge.CreateUser(hueBridgeAppName)
		if err != nil {
			return err
		}

		if bridge.bridge.Username != "" {
			// registration successful
			bridge.Username = bridge.bridge.Username
			log.Printf("⌘ User registration successful.")
			return nil
		}
	}
}

func (bridge *HueBridge) connect() error {
	if bridge.BridgeIP == "" {
		return errors.New("No bridge IP configured")
	}

	if bridge.Username == "" {
		return errors.New("No username on bridge configured")
	}
	bridge.bridge = *hue.NewBridge(bridge.BridgeIP, bridge.Username)

	// Test bridge
	configuration, err := bridge.bridge.Configuration()
	if err != nil {
		return err
	}

	// Enable HTTPS if supported
	// TODO HTTPS supported on Model BSB001?
	swversion, err := strconv.Atoi(configuration.SoftwareVersion)
	if err != nil {
		return err
	}
	if configuration.ModelId == "BSB002" && swversion >= 1802201122 {
		bridge.bridge.EnableHTTPS(true)
	}

	return nil
}

func (bridge *HueBridge) populateSchedule(configuration *Configuration) error {
	if len(configuration.Schedules) == 0 {
		return errors.New("Configuration does not contain any schedules to populate")
	}

	// Do we have associated lights?
	for _, schedule := range configuration.Schedules {
		if len(schedule.AssociatedDeviceIDs) > 0 {
			log.Debugf("⌘ Configuration contains at least one schedule with associated lights.")
			return nil // At least one schedule is configured
		}
	}

	// No schedule has associated lights
	log.Debugf("⌘ Configuration contains no schedule with associated lights. Initializing first schedule with all lights.")
	lights, err := bridge.Lights()
	if err != nil {
		return err
	}
	var lightIDs []int
	for _, light := range lights {
		lightIDs = append(lightIDs, light.ID)
	}
	configuration.Schedules[0].AssociatedDeviceIDs = lightIDs
	return nil
}

func (bridge *HueBridge) validateSofwareVersion() {
	configuration, err := bridge.bridge.Configuration()
	if err != nil {
		log.Warningf("⌘ Could not validate bridge software version: %v", err)
		return
	}

	swversion, err := strconv.Atoi(configuration.SoftwareVersion)
	if err != nil {
		log.Warningf("⌘ Could not validate bridge software version: %v", err)
		return
	}
	log.Debugf("⌘ Bridge is running software version %s", configuration.SoftwareVersion)

	if (bridge.Version == 1 && swversion < 1041302) || (bridge.Version == 2 && swversion < 1901181309) {
		log.Warningf("⌘ Your hue bridge is running an old software version. Please update using the hue app to ensure Kelvin will run smoothly.")
	} else {
		log.Debugf("⌘ Bridge software is up to date")
	}
}

func (bridge *HueBridge) validateBridge() error {
	if bridge.BridgeIP == "" {
		return errors.New("No bridge configured. Could not validate")
	}
	resp, err := http.Get("http://" + bridge.BridgeIP + "/description.xml")
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return fmt.Errorf("Could not read bridge description: %v", err)
	}

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("Could not read bridge description: %v", err)
	}
	if strings.Contains(string(data), "<modelNumber>929000226503</modelNumber>") {
		bridge.Version = 1
		return nil
	}
	if strings.Contains(string(data), "<modelNumber>BSB002</modelNumber>") {
		bridge.Version = 2
		return nil
	}
	return fmt.Errorf("Bridge validation failed")
}
<commit_msg>Improve behavior on inital setup, update lastest known bridge version<commit_after>// MIT License
//
// Copyright (c) 2019 Stefan Wichmann
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	hue "github.com/stefanwichmann/go.hue"
)

// HueBridge represents the Philips Hue bridge in
// your system.
// It is used to communicate with all devices.
type HueBridge struct {
	bridge   hue.Bridge
	BridgeIP string
	Username string
	Version  int
}

const hueBridgeAppName = "kelvin"

// InitializeBridge creates and returns an initialized HueBridge.
// If you have a valid configuration this will be used. Otherwise a local
// discovery will be started, followed by a user registration on your bridge.
func (bridge *HueBridge) InitializeBridge(configuration *Configuration) error {
	err := bridge.discover(configuration.Bridge.IP)
	if err != nil {
		return err
	}
	configuration.Bridge.IP = bridge.BridgeIP

	if configuration.Bridge.Username != "" {
		log.Debugf("⌘ Found bridge username in configuration: %s", configuration.Bridge.Username)
		bridge.Username = configuration.Bridge.Username
	} else {
		log.Debugf("⌘ No username found in bridge configuration. Starting registration...")
		err := bridge.register()
		if err != nil {
			return err
		}
		log.Debugf("⌘ Saving new username in bridge configuration: %s", bridge.Username)
		configuration.Bridge.Username = bridge.Username
	}

	log.Debugf("⌘ Connecting to bridge %s with username %s", bridge.BridgeIP, bridge.Username)
	err = bridge.connect()
	if err != nil {
		return err
	}
	log.Println("⌘ Connection to bridge established")
	bridge.validateSofwareVersion()

	err = bridge.populateSchedule(configuration)
	return err
}

// Lights return all known lights on your bridge.
func (bridge *HueBridge) Lights() ([]*Light, error) {
	var lights []*Light
	hueLights, err := bridge.bridge.GetAllLights()
	if err != nil {
		return lights, err
	}

	for _, hueLight := range hueLights {
		var light Light
		light.ID, err = strconv.Atoi(hueLight.Id)
		if err != nil {
			return lights, err
		}

		light.HueLight.HueLight = *hueLight
		light.HueLight.initialize(hueLight.Attributes)
		light.Name = light.HueLight.Name
		light.Reachable = light.HueLight.Reachable
		light.On = light.HueLight.On

		lights = append(lights, &light)
	}

	sort.Slice(lights, func(i, j int) bool { return lights[i].ID < lights[j].ID })
	return lights, nil
}

// LightStates returns the current state for lights on the bridge
func (bridge *HueBridge) LightStates() (map[int]hue.LightAttributes, error) {
	var states = make(map[int]hue.LightAttributes)
	hueLights, err := bridge.bridge.GetAllLights()
	if err != nil {
		return states, err
	}

	for _, hueLight := range hueLights {
		lightID, err := strconv.Atoi(hueLight.Id)
		if err != nil {
			return states, err
		}

		states[lightID] = hueLight.Attributes
	}

	return states, nil
}

func (bridge *HueBridge) discover(ip string) error {
	if ip != "" {
		// we have a known IP address. Validate if it points to a reachable bridge
		bridge.BridgeIP = ip
		err := bridge.validateBridge()
		if err != nil {
			return err
		}
		return nil
	}
	log.Debugf("⌘ Starting bridge discovery")
	bridges, err := hue.DiscoverBridges(false)
	if err != nil {
		bridge.BridgeIP = ""
		return err
	}
	if len(bridges) == 0 {
		bridge.BridgeIP = ""
		return errors.New("Bridge discovery failed. Please configure manually in config.json")
	}
	for _, candidate := range bridges {
		bridge.BridgeIP = candidate.IpAddr
		err := bridge.validateBridge()
		if err == nil {
			log.Printf("⌘ Found bridge at %s", bridge.BridgeIP)
			return nil
		}
	}
	bridge.BridgeIP = ""
	return errors.New("Bridge discovery failed. Please configure manually in config.json")
}

func (bridge *HueBridge) register() error {
	if bridge.BridgeIP == "" {
		return errors.New("Registration at bridge not possible because no IP is configured. Start discovery first or enter manually")
	}

	bridge.bridge = *hue.NewBridge(bridge.BridgeIP, "")
	log.Printf("⌘ Starting user registration.")
	log.Warningf("⌘ PLEASE PUSH THE BLUE BUTTON ON YOUR HUE BRIDGE")
	for {
		time.Sleep(5 * time.Second)

		// try user creation, will fail if the button wasn't pressed.
		err := bridge.bridge.CreateUser(hueBridgeAppName)
		if err != nil {
			log.Debugf("⌘ Button wasn't pressed yet. Waiting...")
			continue
		}

		if bridge.bridge.Username != "" {
			// registration successful
			bridge.Username = bridge.bridge.Username
			log.Printf("⌘ User registration successful.")
			return nil
		}
	}
}

func (bridge *HueBridge) connect() error {
	if bridge.BridgeIP == "" {
		return errors.New("No bridge IP configured")
	}

	if bridge.Username == "" {
		return errors.New("No username on bridge configured")
	}
	bridge.bridge = *hue.NewBridge(bridge.BridgeIP, bridge.Username)

	// Test bridge
	configuration, err := bridge.bridge.Configuration()
	if err != nil {
		return err
	}

	// Enable HTTPS if supported
	// TODO HTTPS supported on Model BSB001?
	swversion, err := strconv.Atoi(configuration.SoftwareVersion)
	if err != nil {
		return err
	}
	if configuration.ModelId == "BSB002" && swversion >= 1802201122 {
		bridge.bridge.EnableHTTPS(true)
	}

	return nil
}

func (bridge *HueBridge) populateSchedule(configuration *Configuration) error {
	if len(configuration.Schedules) == 0 {
		return errors.New("Configuration does not contain any schedules to populate")
	}

	// Do we have associated lights?
	for _, schedule := range configuration.Schedules {
		if len(schedule.AssociatedDeviceIDs) > 0 {
			log.Debugf("⌘ Configuration contains at least one schedule with associated lights.")
			return nil // At least one schedule is configured
		}
	}

	// No schedule has associated lights
	log.Debugf("⌘ Configuration contains no schedule with associated lights. Initializing first schedule with all lights.")
	lights, err := bridge.Lights()
	if err != nil {
		return err
	}
	var lightIDs []int
	for _, light := range lights {
		lightIDs = append(lightIDs, light.ID)
	}
	configuration.Schedules[0].AssociatedDeviceIDs = lightIDs
	return nil
}

func (bridge *HueBridge) validateSofwareVersion() {
	configuration, err := bridge.bridge.Configuration()
	if err != nil {
		log.Warningf("⌘ Could not validate bridge software version: %v", err)
		return
	}

	swversion, err := strconv.Atoi(configuration.SoftwareVersion)
	if err != nil {
		log.Warningf("⌘ Could not validate bridge software version: %v", err)
		return
	}
	log.Debugf("⌘ Bridge is running software version %s", configuration.SoftwareVersion)

	if (bridge.Version == 1 && swversion < 1041302) || (bridge.Version == 2 && swversion < 1931069120) {
		log.Warningf("⌘ Your hue bridge is running an old software version. Please update using the hue app to ensure Kelvin will run smoothly.")
	} else {
		log.Debugf("⌘ Bridge software is up to date")
	}
}

func (bridge *HueBridge) validateBridge() error {
	if bridge.BridgeIP == "" {
		return errors.New("No bridge configured. Could not validate")
	}
	resp, err := http.Get("http://" + bridge.BridgeIP + "/description.xml")
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return fmt.Errorf("Could not read bridge description: %v", err)
	}

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("Could not read bridge description: %v", err)
	}
	if strings.Contains(string(data), "<modelNumber>929000226503</modelNumber>") {
		bridge.Version = 1
		return nil
	}
	if strings.Contains(string(data), "<modelNumber>BSB002</modelNumber>") {
		bridge.Version = 2
		return nil
	}
	return fmt.Errorf("Bridge validation failed")
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"net/http"
	"net/url"

	"github.com/docker/docker/api/types"
	"golang.org/x/net/context"
)

// PluginInstall installs a plugin
func (cli *Client) PluginInstall(ctx context.Context, name string, options types.PluginInstallOptions) error {
	// FIXME(vdemeester) name is a ref, we might want to parse/validate it here.
	query := url.Values{}
	query.Set("name", name)
	resp, err := cli.tryPluginPull(ctx, query, options.RegistryAuth)
	if resp.statusCode == http.StatusUnauthorized && options.PrivilegeFunc != nil {
		newAuthHeader, privilegeErr := options.PrivilegeFunc()
		if privilegeErr != nil {
			ensureReaderClosed(resp)
			return privilegeErr
		}
		resp, err = cli.tryPluginPull(ctx, query, newAuthHeader)
	}
	if err != nil {
		ensureReaderClosed(resp)
		return err
	}
	var privileges types.PluginPrivileges
	if err := json.NewDecoder(resp.body).Decode(&privileges); err != nil {
		ensureReaderClosed(resp)
		return err
	}
	ensureReaderClosed(resp)

	if !options.AcceptAllPermissions && options.AcceptPermissionsFunc != nil && len(privileges) > 0 {
		accept, err := options.AcceptPermissionsFunc(privileges)
		if err != nil {
			return err
		}
		if !accept {
			resp, _ := cli.delete(ctx, "/plugins/"+name, nil, nil)
			ensureReaderClosed(resp)
			return pluginPermissionDenied{name}
		}
	}
	if options.Disabled {
		return nil
	}
	return cli.PluginEnable(ctx, name)
}

func (cli *Client) tryPluginPull(ctx context.Context, query url.Values, registryAuth string) (serverResponse, error) {
	headers := map[string][]string{"X-Registry-Auth": {registryAuth}}
	return cli.post(ctx, "/plugins/pull", query, nil, headers)
}
<commit_msg>support settings in docker plugins install<commit_after>package client

import (
	"encoding/json"
	"net/http"
	"net/url"

	"github.com/docker/docker/api/types"
	"golang.org/x/net/context"
)

// PluginInstall installs a plugin
func (cli *Client) PluginInstall(ctx context.Context, name string, options types.PluginInstallOptions) error {
	// FIXME(vdemeester) name is a ref, we might want to parse/validate it here.
	query := url.Values{}
	query.Set("name", name)
	resp, err := cli.tryPluginPull(ctx, query, options.RegistryAuth)
	if resp.statusCode == http.StatusUnauthorized && options.PrivilegeFunc != nil {
		newAuthHeader, privilegeErr := options.PrivilegeFunc()
		if privilegeErr != nil {
			ensureReaderClosed(resp)
			return privilegeErr
		}
		resp, err = cli.tryPluginPull(ctx, query, newAuthHeader)
	}
	if err != nil {
		ensureReaderClosed(resp)
		return err
	}
	var privileges types.PluginPrivileges
	if err := json.NewDecoder(resp.body).Decode(&privileges); err != nil {
		ensureReaderClosed(resp)
		return err
	}
	ensureReaderClosed(resp)

	if !options.AcceptAllPermissions && options.AcceptPermissionsFunc != nil && len(privileges) > 0 {
		accept, err := options.AcceptPermissionsFunc(privileges)
		if err != nil {
			return err
		}
		if !accept {
			resp, _ := cli.delete(ctx, "/plugins/"+name, nil, nil)
			ensureReaderClosed(resp)
			return pluginPermissionDenied{name}
		}
	}

	if len(options.Args) > 0 {
		if err := cli.PluginSet(ctx, name, options.Args); err != nil {
			return err
		}
	}

	if options.Disabled {
		return nil
	}

	return cli.PluginEnable(ctx, name)
}

func (cli *Client) tryPluginPull(ctx context.Context, query url.Values, registryAuth string) (serverResponse, error) {
	headers := map[string][]string{"X-Registry-Auth": {registryAuth}}
	return cli.post(ctx, "/plugins/pull", query, nil, headers)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package containeranalysis

import (
	"fmt"
	"strings"

	"github.com/google/go-containerregistry/pkg/name"

	gen "cloud.google.com/go/devtools/containeranalysis/apiv1alpha1"
	kritisv1beta1 "github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/metadata"
	"golang.org/x/net/context"
	"google.golang.org/api/iterator"
	containeranalysispb "google.golang.org/genproto/googleapis/devtools/containeranalysis/v1alpha1"
)

const (
	PkgVulnerability = "PACKAGE_VULNERABILITY"
	PageSize         = int32(100)
)

// The ContainerAnalysis struct implements MetadataFetcher Interface.
type ContainerAnalysis struct {
	client *gen.Client
	ctx    context.Context
}

func NewContainerAnalysisClient() (*ContainerAnalysis, error) {
	ctx := context.Background()
	client, err := gen.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	return &ContainerAnalysis{
		client: client,
		ctx:    ctx,
	}, nil
}

// GetVulnerabilites gets Package Vulnerabilities Occurrences for a specified image.
func (c ContainerAnalysis) GetVulnerabilities(containerImage string) ([]metadata.Vulnerability, error) {
	// Make sure container image is a GCR image
	ref, err := name.ParseReference(containerImage, name.WeakValidation)
	if err != nil {
		return nil, err
	}
	if !isRegistryGCR(ref.Context().RegistryStr()) {
		return nil, fmt.Errorf("%s is not a valid image hosted in GCR", containerImage)
	}
	project := strings.Split(containerImage, "/")[1]

	req := &containeranalysispb.ListOccurrencesRequest{
		Filter:   fmt.Sprintf("resource_url=%q AND kind=%q", fmt.Sprintf("https://%s", containerImage), PkgVulnerability),
		PageSize: PageSize,
		Parent:   fmt.Sprintf("projects/%s", project),
	}
	it := c.client.ListOccurrences(c.ctx, req)
	vulnz := []metadata.Vulnerability{}
	for {
		occ, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		vulnz = append(vulnz, GetVulnerabilityFromOccurence(occ))
	}
	return vulnz, nil
}

func GetVulnerabilityFromOccurence(occ *containeranalysispb.Occurrence) metadata.Vulnerability {
	vulnDetails := occ.GetDetails().(*containeranalysispb.Occurrence_VulnerabilityDetails).VulnerabilityDetails
	hasFixAvailable := isFixAvaliable(vulnDetails.GetPackageIssue())
	vulnerability := metadata.Vulnerability{
		Severity:        containeranalysispb.VulnerabilityType_Severity_name[int32(vulnDetails.Severity)],
		HasFixAvailable: hasFixAvailable,
		CVE:             occ.GetNoteName(),
	}
	return vulnerability
}

func isFixAvaliable(pis []*containeranalysispb.VulnerabilityType_PackageIssue) bool {
	for _, pi := range pis {
		if pi.GetFixedLocation().GetVersion().Kind == containeranalysispb.VulnerabilityType_Version_MAXIMUM {
			// If FixedLocation.Version.Kind = MAXIMUM then no fix is available. Return false
			return false
		}
	}
	return true
}

func isRegistryGCR(r string) bool {
	registry := strings.Split(r, ".")
	if len(registry) < 2 {
		return false
	}
	if registry[len(registry)-2] != "gcr" || registry[len(registry)-1] != "io" {
		return false
	}
	return true
}

func getProjectFromNoteReference(ref string) (string, error) {
	str := strings.Split(ref, "/")
	if len(str) < 3 {
		return "", fmt.Errorf("Invalid Note Reference. Should be in format <api>/projects/<project_id")
	}
	return str[2], nil
}

func (c ContainerAnalysis) CreateAttestationNote(aa kritisv1beta1.AttestationAuthority) error {
	noteProject, err := getProjectFromNoteReference(aa.NoteReference)
	if err != nil {
		return err
	}
	aaNote := &containeranalysispb.AttestationAuthority{
		Hint: &containeranalysispb.AttestationAuthority_AttestationAuthorityHint{
			HumanReadableName: aa.Name,
		},
	}
	note := containeranalysispb.Note{
		Name:             fmt.Sprintf("projects/%s/notes/%s", noteProject, aa.Name),
		ShortDescription: fmt.Sprintf("Image Policy Security Attestor"),
		LongDescription:  fmt.Sprintf("Image Policy Security Attestor deployed in %s namespace", aa.Namespace),
		NoteType: &containeranalysispb.Note_AttestationAuthority{
			AttestationAuthority: aaNote,
		},
	}

	req := &containeranalysispb.CreateNoteRequest{
		Note:   &note,
		NoteId: aa.Name,
		Parent: fmt.Sprintf("projects/%s", noteProject),
	}
	_, err = c.client.CreateNote(c.ctx, req)
	return err
}

func (c ContainerAnalysis) GetAttestationNote(aa kritisv1beta1.AttestationAuthority) (*containeranalysispb.Note, error) {
	noteProject, err := getProjectFromNotReference(aa.NoteReference)
	if err != nil {
		return nil, err
	}
	req := &containeranalysispb.GetNoteRequest{
		Name: fmt.Sprintf("projects/%s/notes/%s", noteProject, aa.Name),
	}
	resp, err := c.client.GetNote(c.ctx, req)
	return resp, nil
}

// This is used for Testing.
func (c ContainerAnalysis) DeleteAttestationNote(aa kritisv1beta1.AttestationAuthority) error {
	noteProject, err := getProjectFromNotReference(aa.NoteReference)
	if err != nil {
		return err
	}
	req := &containeranalysispb.DeleteNoteRequest{
		Name: fmt.Sprintf("projects/%s/notes/%s", noteProject, aa.Name),
	}
	return c.client.DeleteNote(c.ctx, req)
}
<commit_msg>fix error<commit_after>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package containeranalysis

import (
	"fmt"
	"strings"

	"github.com/google/go-containerregistry/pkg/name"

	gen "cloud.google.com/go/devtools/containeranalysis/apiv1alpha1"
	kritisv1beta1 "github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/metadata"
	"golang.org/x/net/context"
	"google.golang.org/api/iterator"
	containeranalysispb "google.golang.org/genproto/googleapis/devtools/containeranalysis/v1alpha1"
)

const (
	PkgVulnerability = "PACKAGE_VULNERABILITY"
	PageSize         = int32(100)
)

// The ContainerAnalysis struct implements MetadataFetcher Interface.
type ContainerAnalysis struct {
	client *gen.Client
	ctx    context.Context
}

func NewContainerAnalysisClient() (*ContainerAnalysis, error) {
	ctx := context.Background()
	client, err := gen.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	return &ContainerAnalysis{
		client: client,
		ctx:    ctx,
	}, nil
}

// GetVulnerabilites gets Package Vulnerabilities Occurrences for a specified image.
func (c ContainerAnalysis) GetVulnerabilities(containerImage string) ([]metadata.Vulnerability, error) {
	// Make sure container image is a GCR image
	ref, err := name.ParseReference(containerImage, name.WeakValidation)
	if err != nil {
		return nil, err
	}
	if !isRegistryGCR(ref.Context().RegistryStr()) {
		return nil, fmt.Errorf("%s is not a valid image hosted in GCR", containerImage)
	}
	project := strings.Split(containerImage, "/")[1]

	req := &containeranalysispb.ListOccurrencesRequest{
		Filter:   fmt.Sprintf("resource_url=%q AND kind=%q", fmt.Sprintf("https://%s", containerImage), PkgVulnerability),
		PageSize: PageSize,
		Parent:   fmt.Sprintf("projects/%s", project),
	}
	it := c.client.ListOccurrences(c.ctx, req)
	vulnz := []metadata.Vulnerability{}
	for {
		occ, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		vulnz = append(vulnz, GetVulnerabilityFromOccurence(occ))
	}
	return vulnz, nil
}

func GetVulnerabilityFromOccurence(occ *containeranalysispb.Occurrence) metadata.Vulnerability {
	vulnDetails := occ.GetDetails().(*containeranalysispb.Occurrence_VulnerabilityDetails).VulnerabilityDetails
	hasFixAvailable := isFixAvaliable(vulnDetails.GetPackageIssue())
	vulnerability := metadata.Vulnerability{
		Severity:        containeranalysispb.VulnerabilityType_Severity_name[int32(vulnDetails.Severity)],
		HasFixAvailable: hasFixAvailable,
		CVE:             occ.GetNoteName(),
	}
	return vulnerability
}

func isFixAvaliable(pis []*containeranalysispb.VulnerabilityType_PackageIssue) bool {
	for _, pi := range pis {
		if pi.GetFixedLocation().GetVersion().Kind == containeranalysispb.VulnerabilityType_Version_MAXIMUM {
			// If FixedLocation.Version.Kind = MAXIMUM then no fix is available. Return false
			return false
		}
	}
	return true
}

func isRegistryGCR(r string) bool {
	registry := strings.Split(r, ".")
	if len(registry) < 2 {
		return false
	}
	if registry[len(registry)-2] != "gcr" || registry[len(registry)-1] != "io" {
		return false
	}
	return true
}

func getProjectFromNoteReference(ref string) (string, error) {
	str := strings.Split(ref, "/")
	if len(str) < 3 {
		return "", fmt.Errorf("Invalid Note Reference. Should be in format <api>/projects/<project_id")
	}
	return str[2], nil
}

func (c ContainerAnalysis) CreateAttestationNote(aa kritisv1beta1.AttestationAuthority) error {
	noteProject, err := getProjectFromNoteReference(aa.NoteReference)
	if err != nil {
		return err
	}
	aaNote := &containeranalysispb.AttestationAuthority{
		Hint: &containeranalysispb.AttestationAuthority_AttestationAuthorityHint{
			HumanReadableName: aa.Name,
		},
	}
	note := containeranalysispb.Note{
		Name:             fmt.Sprintf("projects/%s/notes/%s", noteProject, aa.Name),
		ShortDescription: fmt.Sprintf("Image Policy Security Attestor"),
		LongDescription:  fmt.Sprintf("Image Policy Security Attestor deployed in %s namespace", aa.Namespace),
		NoteType: &containeranalysispb.Note_AttestationAuthority{
			AttestationAuthority: aaNote,
		},
	}

	req := &containeranalysispb.CreateNoteRequest{
		Note:   &note,
		NoteId: aa.Name,
		Parent: fmt.Sprintf("projects/%s", noteProject),
	}
	_, err = c.client.CreateNote(c.ctx, req)
	return err
}

func (c ContainerAnalysis) GetAttestationNote(aa kritisv1beta1.AttestationAuthority) (*containeranalysispb.Note, error) {
	noteProject, err := getProjectFromNoteReference(aa.NoteReference)
	if err != nil {
		return nil, err
	}
	req := &containeranalysispb.GetNoteRequest{
		Name: fmt.Sprintf("projects/%s/notes/%s", noteProject, aa.Name),
	}
	resp, err := c.client.GetNote(c.ctx, req)
	return resp, nil
}

// This is used for Testing.
func (c ContainerAnalysis) DeleteAttestationNote(aa kritisv1beta1.AttestationAuthority) error {
	noteProject, err := getProjectFromNoteReference(aa.NoteReference)
	if err != nil {
		return err
	}
	req := &containeranalysispb.DeleteNoteRequest{
		Name: fmt.Sprintf("projects/%s/notes/%s", noteProject, aa.Name),
	}
	return c.client.DeleteNote(c.ctx, req)
}
<|endoftext|>


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

package server

import (
	"bytes"
	"fmt"
	"os"
	"time"

	"github.com/osrg/gobgp/config"
	"github.com/osrg/gobgp/packet/bgp"
	"github.com/osrg/gobgp/packet/mrt"
	"github.com/osrg/gobgp/table"
	log "github.com/sirupsen/logrus"
)

const (
	MIN_ROTATION_INTERVAL = 60
	MIN_DUMP_INTERVAL     = 60
)

type mrtWriter struct {
	dead             chan struct{}
	s                *BgpServer
	c                *config.MrtConfig
	file             *os.File
	rotationInterval uint64
	dumpInterval     uint64
}

func (m *mrtWriter) Stop() {
	close(m.dead)
}

func (m *mrtWriter) loop() error {
	ops := []WatchOption{}
	switch m.c.DumpType {
	case config.MRT_TYPE_UPDATES:
		ops = append(ops, WatchUpdate(false))
	case config.MRT_TYPE_TABLE:
		if len(m.c.TableName) > 0 {
			ops = append(ops, WatchTableName(m.c.TableName))
		}
	}
	w := m.s.Watch(ops...)
	rotator := func() *time.Ticker {
		if m.rotationInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.rotationInterval))
	}()
	dump := func() *time.Ticker {
		if m.dumpInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.dumpInterval))
	}()

	defer func() {
		if m.file != nil {
			m.file.Close()
		}
		if m.rotationInterval != 0 {
			rotator.Stop()
		}
		if m.dumpInterval != 0 {
			dump.Stop()
		}
		w.Stop()
	}()

	for {
		serialize := func(ev WatchEvent) []*mrt.MRTMessage {
			msg := make([]*mrt.MRTMessage, 0, 1)
			switch e := ev.(type) {
			case *WatchEventUpdate:
				if e.Init {
					return nil
				}
				mp := mrt.NewBGP4MPMessage(e.PeerAS, e.LocalAS, 0, e.PeerAddress.String(), e.LocalAddress.String(), e.FourBytesAs, nil)
				mp.BGPMessagePayload = e.Payload
				isAddPath := e.Neighbor.IsAddPathReceiveEnabled(e.PathList[0].GetRouteFamily())
				subtype := mrt.MESSAGE
				switch {
				case isAddPath && e.FourBytesAs:
					subtype = mrt.MESSAGE_AS4_ADDPATH
				case isAddPath:
					subtype = mrt.MESSAGE_ADDPATH
				case e.FourBytesAs:
					subtype = mrt.MESSAGE_AS4
				}
				if bm, err := mrt.NewMRTMessage(uint32(e.Timestamp.Unix()), mrt.BGP4MP, subtype, mp); err != nil {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Data":  e,
						"Error": err,
					}).Warnf("Failed to create MRT BGP4MP message (subtype %d)", subtype)
				} else {
					msg = append(msg, bm)
				}
			case *WatchEventTable:
				t := uint32(time.Now().Unix())
				peers := make([]*mrt.Peer, 0, len(e.Neighbor))
				neighborMap := make(map[string]*config.Neighbor)
				for _, pconf := range e.Neighbor {
					peers = append(peers, mrt.NewPeer(pconf.State.RemoteRouterId, pconf.State.NeighborAddress, pconf.Config.PeerAs, true))
					neighborMap[pconf.State.NeighborAddress] = pconf
				}
				if bm, err := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, mrt.PEER_INDEX_TABLE, mrt.NewPeerIndexTable(e.RouterId, "", peers)); err != nil {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Data":  e,
						"Error": err,
					}).Warnf("Failed to create MRT TABLE_DUMPv2 message (subtype %d)", mrt.PEER_INDEX_TABLE)
					break
				} else {
					msg = append(msg, bm)
				}

				idx := func(p *table.Path) uint16 {
					for i, pconf := range e.Neighbor {
						if p.GetSource().Address.String() == pconf.State.NeighborAddress {
							return uint16(i)
						}
					}
					return uint16(len(e.Neighbor))
				}

				subtype := func(p *table.Path, isAddPath bool) mrt.MRTSubTypeTableDumpv2 {
					t := mrt.RIB_GENERIC
					switch p.GetRouteFamily() {
					case bgp.RF_IPv4_UC:
						t = mrt.RIB_IPV4_UNICAST
					case bgp.RF_IPv4_MC:
						t = mrt.RIB_IPV4_MULTICAST
					case bgp.RF_IPv6_UC:
						t = mrt.RIB_IPV6_UNICAST
					case bgp.RF_IPv6_MC:
						t = mrt.RIB_IPV6_MULTICAST
					}
					if isAddPath {
						// Shift non-additional-path version to *_ADDPATH
						t += 6
					}
					return t
				}

				seq := uint32(0)
				appendTableDumpMsg := func(path *table.Path, entries []*mrt.RibEntry, isAddPath bool) {
					st := subtype(path, isAddPath)
					if bm, err := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, st, mrt.NewRib(seq, path.GetNlri(), entries)); err != nil {
						log.WithFields(log.Fields{
							"Topic": "mrt",
							"Data":  e,
							"Error": err,
						}).Warnf("Failed to create MRT TABLE_DUMPv2 message (subtype %d)", st)
					} else {
						msg = append(msg, bm)
						seq++
					}
				}
				for _, pathList := range e.PathList {
					entries := make([]*mrt.RibEntry, 0, len(pathList))
					entriesAddPath := make([]*mrt.RibEntry, 0, len(pathList))
					for _, path := range pathList {
						if path.IsLocal() {
							continue
						}
						isAddPath := false
						if neighbor, ok := neighborMap[path.GetSource().Address.String()]; ok {
							isAddPath = neighbor.IsAddPathReceiveEnabled(path.GetRouteFamily())
						}
						if !isAddPath {
							entries = append(entries, mrt.NewRibEntry(idx(path), uint32(path.GetTimestamp().Unix()), 0, path.GetPathAttrs(), false))
						} else {
							entriesAddPath = append(entriesAddPath, mrt.NewRibEntry(idx(path), uint32(path.GetTimestamp().Unix()), path.GetNlri().PathIdentifier(), path.GetPathAttrs(), true))
						}
					}
					if len(entries) > 0 {
						appendTableDumpMsg(pathList[0], entries, false)
					}
					if len(entriesAddPath) > 0 {
						appendTableDumpMsg(pathList[0], entriesAddPath, true)
					}
				}
			}
			return msg
		}

		drain := func(ev WatchEvent) {
			events := make([]WatchEvent, 0, 1+len(w.Event()))
			if ev != nil {
				events = append(events, ev)
			}

			for len(w.Event()) > 0 {
				events = append(events, <-w.Event())
			}

			w := func(buf []byte) {
				if _, err := m.file.Write(buf); err == nil {
					m.file.Sync()
				} else {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Error": err,
					}).Warn("Can't write to destination MRT file")
				}
			}

			var b bytes.Buffer
			for _, e := range events {
				for _, m := range serialize(e) {
					if buf, err := m.Serialize(); err != nil {
						log.WithFields(log.Fields{
							"Topic": "mrt",
							"Data":  e,
							"Error": err,
						}).Warn("Failed to serialize event")
					} else {
						b.Write(buf)
						if b.Len() > 1*1000*1000 {
							w(b.Bytes())
							b.Reset()
						}
					}
				}
			}
			if b.Len() > 0 {
				w(b.Bytes())
			}
		}
		rotate := func() {
			m.file.Close()
			file, err := mrtFileOpen(m.c.FileName, m.rotationInterval)
			if err == nil {
				m.file = file
			} else {
				log.WithFields(log.Fields{
					"Topic": "mrt",
					"Error": err,
				}).Warn("can't rotate MRT file")
			}
		}

		select {
		case <-m.dead:
			drain(nil)
			return nil
		case e := <-w.Event():
			drain(e)
			if m.c.DumpType == config.MRT_TYPE_TABLE && m.rotationInterval != 0 {
				rotate()
			}
		case <-rotator.C:
			if m.c.DumpType == config.MRT_TYPE_UPDATES {
				rotate()
			} else {
				w.Generate(WATCH_EVENT_TYPE_TABLE)
			}
		case <-dump.C:
			w.Generate(WATCH_EVENT_TYPE_TABLE)
		}
	}
}

func mrtFileOpen(filename string, interval uint64) (*os.File, error) {
	realname := filename
	if interval != 0 {
		realname = time.Now().Format(filename)
	}
	log.WithFields(log.Fields{
		"Topic":         "mrt",
		"Filename":      realname,
		"Dump Interval": interval,
	}).Debug("Setting new MRT destination file")

	i := len(realname)
	for i > 0 && os.IsPathSeparator(realname[i-1]) {
		// skip trailing path separators
		i--
	}
	j := i

	for j > 0 && !os.IsPathSeparator(realname[j-1]) {
		j--
	}

	if j > 0 {
		if err := os.MkdirAll(realname[0:j-1], 0755); err != nil {
			log.WithFields(log.Fields{
				"Topic": "mrt",
				"Error": err,
			}).Warn("can't create MRT destination directory")
			return nil, err
		}
	}

	file, err := os.OpenFile(realname, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
	if err != nil {
		log.WithFields(log.Fields{
			"Topic": "mrt",
			"Error": err,
		}).Warn("can't create MRT destination file")
	}
	return file, err
}

func newMrtWriter(s *BgpServer, c *config.MrtConfig, rInterval, dInterval uint64) (*mrtWriter, error) {
	file, err := mrtFileOpen(c.FileName, rInterval)
	if err != nil {
		return nil, err
	}
	m := mrtWriter{
		s:                s,
		c:                c,
		file:             file,
		rotationInterval: rInterval,
		dumpInterval:     dInterval,
	}
	go m.loop()
	return &m, nil
}

type mrtManager struct {
	bgpServer *BgpServer
	writer    map[string]*mrtWriter
}

func (m *mrtManager) enable(c *config.MrtConfig) error {
	if _, ok := m.writer[c.FileName]; ok {
		return fmt.Errorf("%s already exists", c.FileName)
	}

	rInterval := c.RotationInterval
	dInterval := c.DumpInterval

	setRotationMin := func() {
		if rInterval < MIN_ROTATION_INTERVAL {
			log.WithFields(log.Fields{
				"Topic": "MRT",
			}).Infof("minimum mrt rotation interval is %d seconds", MIN_ROTATION_INTERVAL)
			rInterval = MIN_ROTATION_INTERVAL
		}
	}

	if c.DumpType == config.MRT_TYPE_TABLE {
		if rInterval == 0 {
			if dInterval < MIN_DUMP_INTERVAL {
				log.WithFields(log.Fields{
					"Topic": "MRT",
				}).Infof("minimum mrt dump interval is %d seconds", MIN_DUMP_INTERVAL)
				dInterval = MIN_DUMP_INTERVAL
			}
		} else if dInterval == 0 {
			setRotationMin()
		} else {
			return fmt.Errorf("can't specify both intervals in the table dump type")
		}
	} else if c.DumpType == config.MRT_TYPE_UPDATES {
		// ignore the dump interval
		dInterval = 0
		if len(c.TableName) > 0 {
			return fmt.Errorf("can't specify the table name with the update dump type")
		}
		setRotationMin()
	}

	w, err := newMrtWriter(m.bgpServer, c, rInterval, dInterval)
	if err == nil {
		m.writer[c.FileName] = w
	}
	return err
}

func (m *mrtManager) disable(c *config.MrtConfig) error {
	if w, ok := m.writer[c.FileName]; !ok {
		return fmt.Errorf("%s doesn't exists", c.FileName)
	} else {
		w.Stop()
		delete(m.writer, c.FileName)
	}
	return nil
}

func newMrtManager(s *BgpServer) *mrtManager {
	return &mrtManager{
		bgpServer: s,
		writer:    make(map[string]*mrtWriter),
	}
}
<commit_msg>mrt: Enable to dump locally generated routes<commit_after>// Copyright (C) 2016 Nippon Telegraph and Telephone Corporation.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"fmt"
	"os"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/osrg/gobgp/config"
	"github.com/osrg/gobgp/packet/bgp"
	"github.com/osrg/gobgp/packet/mrt"
	"github.com/osrg/gobgp/table"
)

const (
	MIN_ROTATION_INTERVAL = 60
	MIN_DUMP_INTERVAL     = 60
)

type mrtWriter struct {
	dead             chan struct{}
	s                *BgpServer
	c                *config.MrtConfig
	file             *os.File
	rotationInterval uint64
	dumpInterval     uint64
}

func (m *mrtWriter) Stop() {
	close(m.dead)
}

func (m *mrtWriter) loop() error {
	ops := []WatchOption{}
	switch m.c.DumpType {
	case config.MRT_TYPE_UPDATES:
		ops = append(ops, WatchUpdate(false))
	case config.MRT_TYPE_TABLE:
		if len(m.c.TableName) > 0 {
			ops = append(ops, WatchTableName(m.c.TableName))
		}
	}
	w := m.s.Watch(ops...)
	rotator := func() *time.Ticker {
		if m.rotationInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.rotationInterval))
	}()
	dump := func() *time.Ticker {
		if m.dumpInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.dumpInterval))
	}()

	defer func() {
		if m.file != nil {
			m.file.Close()
		}
		if m.rotationInterval != 0 {
			rotator.Stop()
		}
		if m.dumpInterval != 0 {
			dump.Stop()
		}
		w.Stop()
	}()

	for {
		serialize := func(ev WatchEvent) []*mrt.MRTMessage {
			msg := make([]*mrt.MRTMessage, 0, 1)
			switch e := ev.(type) {
			case *WatchEventUpdate:
				if e.Init {
					return nil
				}
				mp := mrt.NewBGP4MPMessage(e.PeerAS, e.LocalAS, 0, e.PeerAddress.String(), e.LocalAddress.String(), e.FourBytesAs, nil)
				mp.BGPMessagePayload = e.Payload
				isAddPath := e.Neighbor.IsAddPathReceiveEnabled(e.PathList[0].GetRouteFamily())
				subtype := mrt.MESSAGE
				switch {
				case isAddPath && e.FourBytesAs:
					subtype = mrt.MESSAGE_AS4_ADDPATH
				case isAddPath:
					subtype = mrt.MESSAGE_ADDPATH
				case e.FourBytesAs:
					subtype = mrt.MESSAGE_AS4
				}
				if bm, err := mrt.NewMRTMessage(uint32(e.Timestamp.Unix()), mrt.BGP4MP, subtype, mp); err != nil {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Data":  e,
						"Error": err,
					}).Warnf("Failed to create MRT BGP4MP message (subtype %d)", subtype)
				} else {
					msg = append(msg, bm)
				}
			case *WatchEventTable:
				t := uint32(time.Now().Unix())

				peers := make([]*mrt.Peer, 1, len(e.Neighbor)+1)
				// Adding dummy Peer record for locally generated routes
				peers[0] = mrt.NewPeer("0.0.0.0", "0.0.0.0", 0, true)
				neighborMap := make(map[string]*config.Neighbor)
				for _, pconf := range e.Neighbor {
					peers = append(peers, mrt.NewPeer(pconf.State.RemoteRouterId, pconf.State.NeighborAddress, pconf.Config.PeerAs, true))
					neighborMap[pconf.State.NeighborAddress] = pconf
				}

				if bm, err := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, mrt.PEER_INDEX_TABLE, mrt.NewPeerIndexTable(e.RouterId, "", peers)); err != nil {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Data":  e,
						"Error": err,
					}).Warnf("Failed to create MRT TABLE_DUMPv2 message (subtype %d)", mrt.PEER_INDEX_TABLE)
					break
				} else {
					msg = append(msg, bm)
				}

				idx := func(p *table.Path) uint16 {
					for i, pconf := range e.Neighbor {
						if p.GetSource().Address.String() == pconf.State.NeighborAddress {
							return uint16(i)
						}
					}
					return uint16(len(e.Neighbor))
				}

				subtype := func(p *table.Path, isAddPath bool) mrt.MRTSubTypeTableDumpv2 {
					t := mrt.RIB_GENERIC
					switch p.GetRouteFamily() {
					case bgp.RF_IPv4_UC:
						t = mrt.RIB_IPV4_UNICAST
					case bgp.RF_IPv4_MC:
						t = mrt.RIB_IPV4_MULTICAST
					case bgp.RF_IPv6_UC:
						t = mrt.RIB_IPV6_UNICAST
					case bgp.RF_IPv6_MC:
						t = mrt.RIB_IPV6_MULTICAST
					}
					if isAddPath {
						// Shift non-additional-path version to *_ADDPATH
						t += 6
					}
					return t
				}

				seq := uint32(0)
				appendTableDumpMsg := func(path *table.Path, entries []*mrt.RibEntry, isAddPath bool) {
					st := subtype(path, isAddPath)
					if bm, err := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, st, mrt.NewRib(seq, path.GetNlri(), entries)); err != nil {
						log.WithFields(log.Fields{
							"Topic": "mrt",
							"Data":  e,
							"Error": err,
						}).Warnf("Failed to create MRT TABLE_DUMPv2 message (subtype %d)", st)
					} else {
						msg = append(msg, bm)
						seq++
					}
				}
				for _, pathList := range e.PathList {
					entries := make([]*mrt.RibEntry, 0, len(pathList))
					entriesAddPath := make([]*mrt.RibEntry, 0, len(pathList))
					for _, path := range pathList {
						isAddPath := false
						if path.IsLocal() {
							isAddPath = true
						} else if neighbor, ok := neighborMap[path.GetSource().Address.String()]; ok {
							isAddPath = neighbor.IsAddPathReceiveEnabled(path.GetRouteFamily())
						}
						if !isAddPath {
							entries = append(entries, mrt.NewRibEntry(idx(path), uint32(path.GetTimestamp().Unix()), 0, path.GetPathAttrs(), false))
						} else {
							entriesAddPath = append(entriesAddPath, mrt.NewRibEntry(idx(path), uint32(path.GetTimestamp().Unix()), path.GetNlri().PathIdentifier(), path.GetPathAttrs(), true))
						}
					}
					if len(entries) > 0 {
						appendTableDumpMsg(pathList[0], entries, false)
					}
					if len(entriesAddPath) > 0 {
						appendTableDumpMsg(pathList[0], entriesAddPath, true)
					}
				}
			}
			return msg
		}

		drain := func(ev WatchEvent) {
			events := make([]WatchEvent, 0, 1+len(w.Event()))
			if ev != nil {
				events = append(events, ev)
			}

			for len(w.Event()) > 0 {
				events = append(events, <-w.Event())
			}

			w := func(buf []byte) {
				if _, err := m.file.Write(buf); err == nil {
					m.file.Sync()
				} else {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Error": err,
					}).Warn("Can't write to destination MRT file")
				}
			}

			var b bytes.Buffer
			for _, e := range events {
				for _, m := range serialize(e) {
					if buf, err := m.Serialize(); err != nil {
						log.WithFields(log.Fields{
							"Topic": "mrt",
							"Data":  e,
							"Error": err,
						}).Warn("Failed to serialize event")
					} else {
						b.Write(buf)
						if b.Len() > 1*1000*1000 {
							w(b.Bytes())
							b.Reset()
						}
					}
				}
			}
			if b.Len() > 0 {
				w(b.Bytes())
			}
		}
		rotate := func() {
			m.file.Close()
			file, err := mrtFileOpen(m.c.FileName, m.rotationInterval)
			if err == nil {
				m.file = file
			} else {
				log.WithFields(log.Fields{
					"Topic": "mrt",
					"Error": err,
				}).Warn("can't rotate MRT file")
			}
		}

		select {
		case <-m.dead:
			drain(nil)
			return nil
		case e := <-w.Event():
			drain(e)
			if m.c.DumpType == config.MRT_TYPE_TABLE && m.rotationInterval != 0 {
				rotate()
			}
		case <-rotator.C:
			if m.c.DumpType == config.MRT_TYPE_UPDATES {
				rotate()
			} else {
				w.Generate(WATCH_EVENT_TYPE_TABLE)
			}
		case <-dump.C:
			w.Generate(WATCH_EVENT_TYPE_TABLE)
		}
	}
}

func mrtFileOpen(filename string, interval uint64) (*os.File, error) {
	realname := filename
	if interval != 0 {
		realname = time.Now().Format(filename)
	}
	log.WithFields(log.Fields{
		"Topic":         "mrt",
		"Filename":      realname,
		"Dump Interval": interval,
	}).Debug("Setting new MRT destination file")

	i := len(realname)
	for i > 0 && os.IsPathSeparator(realname[i-1]) {
		// skip trailing path separators
		i--
	}
	j := i

	for j > 0 && !os.IsPathSeparator(realname[j-1]) {
		j--
	}

	if j > 0 {
		if err := os.MkdirAll(realname[0:j-1], 0755); err != nil {
			log.WithFields(log.Fields{
				"Topic": "mrt",
				"Error": err,
			}).Warn("can't create MRT destination directory")
			return nil, err
		}
	}

	file, err := os.OpenFile(realname, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
	if err != nil {
		log.WithFields(log.Fields{
			"Topic": "mrt",
			"Error": err,
		}).Warn("can't create MRT destination file")
	}
	return file, err
}

func newMrtWriter(s *BgpServer, c *config.MrtConfig, rInterval, dInterval uint64) (*mrtWriter, error) {
	file, err := mrtFileOpen(c.FileName, rInterval)
	if err != nil {
		return nil, err
	}
	m := mrtWriter{
		s:                s,
		c:                c,
		file:             file,
		rotationInterval: rInterval,
		dumpInterval:     dInterval,
	}
	go m.loop()
	return &m, nil
}

type mrtManager struct {
	bgpServer *BgpServer
	writer    map[string]*mrtWriter
}

func (m *mrtManager) enable(c *config.MrtConfig) error {
	if _, ok := m.writer[c.FileName]; ok {
		return fmt.Errorf("%s already exists", c.FileName)
	}

	rInterval := c.RotationInterval
	dInterval := c.DumpInterval

	setRotationMin := func() {
		if rInterval < MIN_ROTATION_INTERVAL {
			log.WithFields(log.Fields{
				"Topic": "MRT",
			}).Infof("minimum mrt rotation interval is %d seconds", MIN_ROTATION_INTERVAL)
			rInterval = MIN_ROTATION_INTERVAL
		}
	}

	if c.DumpType == config.MRT_TYPE_TABLE {
		if rInterval == 0 {
			if dInterval < MIN_DUMP_INTERVAL {
				log.WithFields(log.Fields{
					"Topic": "MRT",
				}).Infof("minimum mrt dump interval is %d seconds", MIN_DUMP_INTERVAL)
				dInterval = MIN_DUMP_INTERVAL
			}
		} else if dInterval == 0 {
			setRotationMin()
		} else {
			return fmt.Errorf("can't specify both intervals in the table dump type")
		}
	} else if c.DumpType == config.MRT_TYPE_UPDATES {
		// ignore the dump interval
		dInterval = 0
		if len(c.TableName) > 0 {
			return fmt.Errorf("can't specify the table name with the update dump type")
		}
		setRotationMin()
	}

	w, err := newMrtWriter(m.bgpServer, c, rInterval, dInterval)
	if err == nil {
		m.writer[c.FileName] = w
	}
	return err
}

func (m *mrtManager) disable(c *config.MrtConfig) error {
	if w, ok := m.writer[c.FileName]; !ok {
		return fmt.Errorf("%s doesn't exists", c.FileName)
	} else {
		w.Stop()
		delete(m.writer, c.FileName)
	}
	return nil
}

func newMrtManager(s *BgpServer) *mrtManager {
	return &mrtManager{
		bgpServer: s,
		writer:    make(map[string]*mrtWriter),
	}
}
<|endoftext|>


<commit_before>package util

import (
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"

	log "github.com/Sirupsen/logrus"
)

var (
	//The number of goroutines waiting for a close signal
	closeNumber  = int32(0)
	closeCounter = make(chan bool, 1)
	closeExiter  = sync.Once{}
)

type Closeable interface {
	Close()
}

// CloseOnExit closes a resource when the program is exiting.
func CloseOnExit(closeable Closeable) {
	atomic.AddInt32(&closeNumber, 1)
	closeExiter.Do(func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)

		go func() {
			<-c
			log.Warn("Exiting...")
			cnum := atomic.LoadInt32(&closeNumber)
			for cnum > 0 {
				<-closeCounter
				cnum = atomic.LoadInt32(&closeNumber)
			}
			log.Debug("bye!")
			os.Exit(0)
		}()
	})

	if closeable == nil {
		return
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				atomic.AddInt32(&closeNumber, -1)
				closeCounter <- false
			}
		}()
		<-c
		closeable.Close()
		atomic.AddInt32(&closeNumber, -1)
		closeCounter <- true
	}()
}

// SendCloseSignal sends the program the terminate signal so all items waiting for a
// CloseOnExit will complete.
func SendCloseSignal() {
	syscall.Kill(syscall.Getpid(), syscall.SIGTERM)
}
<commit_msg>Switched close to WaitGroup<commit_after>package util

import (
	"os"
	"os/signal"
	"sync"
	"syscall"

	log "github.com/Sirupsen/logrus"
)

var (
	closeWaiter = sync.WaitGroup{}
	closeExiter = sync.Once{}
)

type Closeable interface {
	Close()
}

// CloseOnExit closes a resource when the program is exiting.
func CloseOnExit(closeable Closeable) {
	closeWaiter.Add(1)
	closeExiter.Do(func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)

		go func() {
			<-c
			log.Warn("Exiting...")
			closeWaiter.Wait()
			log.Debug("bye!")
			os.Exit(0)
		}()
	})

	if closeable == nil {
		return
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-c
		closeable.Close()
		closeWaiter.Done()
	}()
}

// SendCloseSignal sends the program the terminate signal so all items waiting for a
// CloseOnExit will complete.
func SendCloseSignal() {
	syscall.Kill(syscall.Getpid(), syscall.SIGTERM)
}
<|endoftext|>


<commit_before>// +build windows

package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"time"
	"unsafe"

	"github.com/go-ole/go-ole"
	"github.com/moutend/go-wav"
	"github.com/moutend/go-wca"
)

var version = "latest"
var revision = "latest"

type FilenameFlag struct {
	Value string
}

func (f *FilenameFlag) Set(value string) (err error) {
	if !strings.HasSuffix(value, ".wav") {
		err = fmt.Errorf("specify WAVE audio file (*.wav)")
		return
	}
	f.Value = value
	return
}

func (f *FilenameFlag) String() string {
	return f.Value
}

func main() {
	var err error
	if err = run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func run(args []string) (err error) {
	var filenameFlag FilenameFlag
	var versionFlag bool
	var audio = &wav.File{}
	var file []byte

	f := flag.NewFlagSet(args[0], flag.ExitOnError)
	f.Var(&filenameFlag, "input", "Specify WAVE format audio (e.g. music.wav)")
	f.Var(&filenameFlag, "i", "Alias of --input")
	f.BoolVar(&versionFlag, "version", false, "Show version")
	f.Parse(args[1:])

	if versionFlag {
		fmt.Printf("%s-%s\n", version, revision)
		return
	}
	if filenameFlag.Value == "" {
		return
	}
	if file, err = ioutil.ReadFile(filenameFlag.Value); err != nil {
		return
	}
	if err = wav.Unmarshal(file, audio); err != nil {
		return
	}

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		select {
		case <-signalChan:
			fmt.Println("Interrupted by SIGINT")
			cancel()
		}
		return
	}()

	if err = renderSharedTimerDriven(ctx, audio); err != nil {
		return
	}
	fmt.Println("Successfully done")
	return
}

func renderSharedTimerDriven(ctx context.Context, audio *wav.File) (err error) {
	if err = ole.CoInitializeEx(0, ole.COINIT_APARTMENTTHREADED); err != nil {
		return
	}
	defer ole.CoUninitialize()

	var de *wca.IMMDeviceEnumerator
	if err = wca.CoCreateInstance(wca.CLSID_MMDeviceEnumerator, 0, wca.CLSCTX_ALL, wca.IID_IMMDeviceEnumerator, &de); err != nil {
		return
	}
	defer de.Release()

	var mmd *wca.IMMDevice
	if err = de.GetDefaultAudioEndpoint(wca.ERender, wca.EConsole, &mmd); err != nil {
		return
	}
	defer mmd.Release()

	var ps *wca.IPropertyStore
	if err = mmd.OpenPropertyStore(wca.STGM_READ, &ps); err != nil {
		return
	}
	defer ps.Release()

	var pv wca.PROPVARIANT
	if err = ps.GetValue(&wca.PKEY_Device_FriendlyName, &pv); err != nil {
		return
	}
	fmt.Printf("Rendering audio to: %s\n", pv.String())

	var ac *wca.IAudioClient
	if err = mmd.Activate(wca.IID_IAudioClient, wca.CLSCTX_ALL, nil, &ac); err != nil {
		return
	}
	defer ac.Release()

	var wfx *wca.WAVEFORMATEX
	if err = ac.GetMixFormat(&wfx); err != nil {
		return
	}
	defer ole.CoTaskMemFree(uintptr(unsafe.Pointer(wfx)))

	wfx.WFormatTag = 1
	wfx.NSamplesPerSec = uint32(audio.SamplesPerSec())
	wfx.WBitsPerSample = uint16(audio.BitsPerSample())
	wfx.NChannels = uint16(audio.Channels())
	wfx.NBlockAlign = uint16(audio.BlockAlign())
	wfx.NAvgBytesPerSec = uint32(audio.AvgBytesPerSec())
	wfx.CbSize = 0

	fmt.Println("--------")
	fmt.Printf("Format: PCM %d bit signed integer\n", wfx.WBitsPerSample)
	fmt.Printf("Rate: %d Hz\n", wfx.NSamplesPerSec)
	fmt.Printf("Channels: %d\n", wfx.NChannels)
	fmt.Println("--------")

	var defaultPeriod wca.REFERENCE_TIME
	var minimumPeriod wca.REFERENCE_TIME
	var latency time.Duration
	if err = ac.GetDevicePeriod(&defaultPeriod, &minimumPeriod); err != nil {
		return
	}
	latency = time.Duration(int(defaultPeriod) * 100)

	fmt.Println("Default period: ", defaultPeriod)
	fmt.Println("Minimum period: ", minimumPeriod)
	fmt.Println("Latency: ", latency)

	if err = ac.Initialize(wca.AUDCLNT_SHAREMODE_SHARED, 0, defaultPeriod, 0, wfx, nil); err != nil {
		return
	}

	var bufferFrameSize uint32
	if err = ac.GetBufferSize(&bufferFrameSize); err != nil {
		return
	}
	fmt.Printf("Allocated buffer size: %d\n", bufferFrameSize)

	var arc *wca.IAudioRenderClient
	if err = ac.GetService(wca.IID_IAudioRenderClient, &arc); err != nil {
		return
	}
	defer arc.Release()

	if err = ac.Start(); err != nil {
		return
	}

	fmt.Println("Start rendering audio with shared-timer-driven mode")
	fmt.Println("Press Ctrl-C to stop rendering")

	var input = audio.Bytes()
	var data *byte
	var offset int
	var padding uint32
	var availableFrameSize uint32
	var b *byte
	var isPlaying bool = true

	for {
		if !isPlaying {
			break
		}
		select {
		case <-ctx.Done():
			isPlaying = false
			break
		default:
			if offset >= audio.Length() {
				isPlaying = false
				break
			}
			if err = ac.GetCurrentPadding(&padding); err != nil {
				return
			}
			availableFrameSize = bufferFrameSize - padding
			if availableFrameSize == 0 {
				continue
			}
			if err = arc.GetBuffer(availableFrameSize, &data); err != nil {
				return
			}

			start := unsafe.Pointer(data)
			lim := int(availableFrameSize) * int(wfx.NBlockAlign)
			remaining := audio.Length() - offset
			if remaining < lim {
				lim = remaining
			}
			for n := 0; n < lim; n++ {
				b = (*byte)(unsafe.Pointer(uintptr(start) + uintptr(n)))
				*b = input[offset+n]
			}
			offset += lim
			if err = arc.ReleaseBuffer(availableFrameSize, 0); err != nil {
				return
			}
			time.Sleep(latency)
		}
	}

	time.Sleep(latency)

	return ac.Stop()
}
<commit_msg>Chore<commit_after>// +build windows

package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"time"
	"unsafe"

	"github.com/go-ole/go-ole"
	"github.com/moutend/go-wav"
	"github.com/moutend/go-wca"
)

var version = "latest"
var revision = "latest"

type FilenameFlag struct {
	Value string
}

func (f *FilenameFlag) Set(value string) (err error) {
	if !strings.HasSuffix(value, ".wav") {
		err = fmt.Errorf("specify WAVE audio file (*.wav)")
		return
	}
	f.Value = value
	return
}

func (f *FilenameFlag) String() string {
	return f.Value
}

func main() {
	var err error
	if err = run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func run(args []string) (err error) {
	var filenameFlag FilenameFlag
	var versionFlag bool
	var audio = &wav.File{}
	var file []byte

	f := flag.NewFlagSet(args[0], flag.ExitOnError)
	f.Var(&filenameFlag, "input", "Specify WAVE format audio (e.g. music.wav)")
	f.Var(&filenameFlag, "i", "Alias of --input")
	f.BoolVar(&versionFlag, "version", false, "Show version")
	f.Parse(args[1:])

	if versionFlag {
		fmt.Printf("%s-%s\n", version, revision)
		return
	}
	if filenameFlag.Value == "" {
		return
	}
	if file, err = ioutil.ReadFile(filenameFlag.Value); err != nil {
		return
	}
	if err = wav.Unmarshal(file, audio); err != nil {
		return
	}

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		select {
		case <-signalChan:
			fmt.Println("Interrupted by SIGINT")
			cancel()
		}
		return
	}()

	if err = renderSharedTimerDriven(ctx, audio); err != nil {
		return
	}
	fmt.Println("Successfully done")
	return
}

func renderSharedTimerDriven(ctx context.Context, audio *wav.File) (err error) {
	if err = ole.CoInitializeEx(0, ole.COINIT_APARTMENTTHREADED); err != nil {
		return
	}
	defer ole.CoUninitialize()

	var de *wca.IMMDeviceEnumerator
	if err = wca.CoCreateInstance(wca.CLSID_MMDeviceEnumerator, 0, wca.CLSCTX_ALL, wca.IID_IMMDeviceEnumerator, &de); err != nil {
		return
	}
	defer de.Release()

	var mmd *wca.IMMDevice
	if err = de.GetDefaultAudioEndpoint(wca.ERender, wca.EConsole, &mmd); err != nil {
		return
	}
	defer mmd.Release()

	var ps *wca.IPropertyStore
	if err = mmd.OpenPropertyStore(wca.STGM_READ, &ps); err != nil {
		return
	}
	defer ps.Release()

	var pv wca.PROPVARIANT
	if err = ps.GetValue(&wca.PKEY_Device_FriendlyName, &pv); err != nil {
		return
	}
	fmt.Printf("Rendering audio to: %s\n", pv.String())

	var ac *wca.IAudioClient
	if err = mmd.Activate(wca.IID_IAudioClient, wca.CLSCTX_ALL, nil, &ac); err != nil {
		return
	}
	defer ac.Release()

	var wfx *wca.WAVEFORMATEX
	if err = ac.GetMixFormat(&wfx); err != nil {
		return
	}
	defer ole.CoTaskMemFree(uintptr(unsafe.Pointer(wfx)))

	wfx.WFormatTag = 1
	wfx.NSamplesPerSec = uint32(audio.SamplesPerSec())
	wfx.WBitsPerSample = uint16(audio.BitsPerSample())
	wfx.NChannels = uint16(audio.Channels())
	wfx.NBlockAlign = uint16(audio.BlockAlign())
	wfx.NAvgBytesPerSec = uint32(audio.AvgBytesPerSec())
	wfx.CbSize = 0

	fmt.Println("--------")
	fmt.Printf("Format: PCM %d bit signed integer\n", wfx.WBitsPerSample)
	fmt.Printf("Rate: %d Hz\n", wfx.NSamplesPerSec)
	fmt.Printf("Channels: %d\n", wfx.NChannels)
	fmt.Println("--------")

	var defaultPeriod wca.REFERENCE_TIME
	var minimumPeriod wca.REFERENCE_TIME
	var latency time.Duration
	if err = ac.GetDevicePeriod(&defaultPeriod, &minimumPeriod); err != nil {
		return
	}
	latency = time.Duration(int(defaultPeriod) * 100)

	fmt.Println("Default period: ", defaultPeriod)
	fmt.Println("Minimum period: ", minimumPeriod)
	fmt.Println("Latency: ", latency)

	if err = ac.Initialize(wca.AUDCLNT_SHAREMODE_SHARED, 0, defaultPeriod, 0, wfx, nil); err != nil {
		return
	}

	var bufferFrameSize uint32
	if err = ac.GetBufferSize(&bufferFrameSize); err != nil {
		return
	}
	fmt.Printf("Allocated buffer size: %d\n", bufferFrameSize)

	var arc *wca.IAudioRenderClient
	if err = ac.GetService(wca.IID_IAudioRenderClient, &arc); err != nil {
		return
	}
	defer arc.Release()

	if err = ac.Start(); err != nil {
		return
	}

	fmt.Println("Start rendering audio with shared-timer-driven mode")
	fmt.Println("Press Ctrl-C to stop rendering")

	var input = audio.Bytes()
	var data *byte
	var offset int
	var padding uint32
	var availableFrameSize uint32
	var b *byte
	var isPlaying bool = true

	for {
		if !isPlaying {
			break
		}
		select {
		case <-ctx.Done():
			isPlaying = false
			break
		default:
			if offset >= audio.Length() {
				isPlaying = false
				break
			}
			if err = ac.GetCurrentPadding(&padding); err != nil {
				return
			}
			availableFrameSize = bufferFrameSize - padding
			if availableFrameSize == 0 {
				continue
			}
			if err = arc.GetBuffer(availableFrameSize, &data); err != nil {
				return
			}

			start := unsafe.Pointer(data)
			lim := int(availableFrameSize) * int(wfx.NBlockAlign)
			remaining := audio.Length() - offset
			if remaining < lim {
				lim = remaining
			}
			for n := 0; n < lim; n++ {
				b = (*byte)(unsafe.Pointer(uintptr(start) + uintptr(n)))
				*b = input[offset+n]
			}
			offset += lim
			if err = arc.ReleaseBuffer(availableFrameSize, 0); err != nil {
				return
			}
			time.Sleep(latency)
		}
	}

// Render samples remaining in buffer.
	time.Sleep(latency)

	return ac.Stop()
}
<|endoftext|>


<commit_before>package digests

import (
	"fmt"
	"sort"
	"strings"

	"github.com/ion-channel/ionic/rulesets"
	"github.com/ion-channel/ionic/scanner"
	"github.com/ion-channel/ionic/scans"
)

const (
	differenceIndex = iota
	virusFoundIndex
	criticalVulnerabilitiesIndex
	highVulnerabilitiesIndex
	totalVulnerabilitiesIndex
	uniqueVulnerabilitiesIndex
	licensesIndex
	filesScannedIndex
	directDependencyIndex
	transitiveDependencyIndex
	dependencyOutdatedIndex
	compilersIndex
	containerImagesIndex
	noVersionIndex
	languagesIndex
	uniqueCommittersIndex
	codeCoverageIndex
)

// NewDigests takes an applied ruleset and returns the relevant digests derived
// from all the evaluations in it, and any errors it encounters.
func NewDigests(appliedRuleset *rulesets.AppliedRulesetSummary, statuses []scanner.ScanStatus, experimental bool) ([]Digest, error) {
	ds := make([]Digest, 0)
	errs := make([]string, 0, 0)

	for i := range statuses {
		s := statuses[i]

		var e *scans.Evaluation
		if appliedRuleset != nil && appliedRuleset.RuleEvaluationSummary != nil {
			for i := range appliedRuleset.RuleEvaluationSummary.Ruleresults {
				if appliedRuleset.RuleEvaluationSummary.Ruleresults[i].ID == s.ID {
					e = &appliedRuleset.RuleEvaluationSummary.Ruleresults[i]
					break
				}
			}
		}

		d, err := _newDigests(&s, e, experimental)
		if err != nil {
			errs = append(errs, fmt.Sprintf("failed to make digest(s) from scan: %v", err.Error()))
			continue
		}

		if d != nil {
			ds = append(ds, d...)
		}
	}

	sort.Slice(ds, func(i, j int) bool { return ds[i].Index < ds[j].Index })

	if len(errs) > 0 {
		return ds, fmt.Errorf("failed to make some digests: %v", strings.Join(errs, "; "))
	}

	return ds, nil
}

func _newDigests(status *scanner.ScanStatus, eval *scans.Evaluation, experimental bool) ([]Digest, error) {
	if eval != nil {
		err := eval.Translate()
		if err != nil {
			return nil, fmt.Errorf("evaluation translate error: %v", err.Error())
		}
	}

	switch strings.ToLower(status.Name) {
	case "ecosystems":
		return ecosystemsDigests(status, eval)

	case "dependency":
		return dependencyDigests(status, eval)

	case "vulnerability":
		return vulnerabilityDigests(status, eval)

	case "virus":
		return virusDigests(status, eval)

	case "community":
		return communityDigests(status, eval)

	case "license":
		return licenseDigests(status, eval)

	case "external_coverage", "code_coverage", "coverage":
		return coveragDigests(status, eval)

	case "difference":
		return differenceDigests(status, eval)

	case "buildsystems":
		if experimental {
			return buildsystemsDigests(status, eval)
		}
		return nil, nil

	case "about_yml", "file_type":
		return nil, nil

	default:
		return nil, fmt.Errorf("Couldn't figure out how to map '%v' to a digest", status.Name)
	}
}
<commit_msg>reorder tiles<commit_after>package digests

import (
	"fmt"
	"sort"
	"strings"

	"github.com/ion-channel/ionic/rulesets"
	"github.com/ion-channel/ionic/scanner"
	"github.com/ion-channel/ionic/scans"
)

const (
	differenceIndex = iota
	virusFoundIndex
	criticalVulnerabilitiesIndex
	highVulnerabilitiesIndex
	totalVulnerabilitiesIndex
	uniqueVulnerabilitiesIndex
	licensesIndex
	filesScannedIndex
	directDependencyIndex
	transitiveDependencyIndex
	dependencyOutdatedIndex
	noVersionIndex
	compilersIndex
	containerImagesIndex
	languagesIndex
	uniqueCommittersIndex
	codeCoverageIndex
)

// NewDigests takes an applied ruleset and returns the relevant digests derived
// from all the evaluations in it, and any errors it encounters.
func NewDigests(appliedRuleset *rulesets.AppliedRulesetSummary, statuses []scanner.ScanStatus, experimental bool) ([]Digest, error) {
	ds := make([]Digest, 0)
	errs := make([]string, 0, 0)

	for i := range statuses {
		s := statuses[i]

		var e *scans.Evaluation
		if appliedRuleset != nil && appliedRuleset.RuleEvaluationSummary != nil {
			for i := range appliedRuleset.RuleEvaluationSummary.Ruleresults {
				if appliedRuleset.RuleEvaluationSummary.Ruleresults[i].ID == s.ID {
					e = &appliedRuleset.RuleEvaluationSummary.Ruleresults[i]
					break
				}
			}
		}

		d, err := _newDigests(&s, e, experimental)
		if err != nil {
			errs = append(errs, fmt.Sprintf("failed to make digest(s) from scan: %v", err.Error()))
			continue
		}

		if d != nil {
			ds = append(ds, d...)
		}
	}

	sort.Slice(ds, func(i, j int) bool { return ds[i].Index < ds[j].Index })

	if len(errs) > 0 {
		return ds, fmt.Errorf("failed to make some digests: %v", strings.Join(errs, "; "))
	}

	return ds, nil
}

func _newDigests(status *scanner.ScanStatus, eval *scans.Evaluation, experimental bool) ([]Digest, error) {
	if eval != nil {
		err := eval.Translate()
		if err != nil {
			return nil, fmt.Errorf("evaluation translate error: %v", err.Error())
		}
	}

	switch strings.ToLower(status.Name) {
	case "ecosystems":
		return ecosystemsDigests(status, eval)

	case "dependency":
		return dependencyDigests(status, eval)

	case "vulnerability":
		return vulnerabilityDigests(status, eval)

	case "virus":
		return virusDigests(status, eval)

	case "community":
		return communityDigests(status, eval)

	case "license":
		return licenseDigests(status, eval)

	case "external_coverage", "code_coverage", "coverage":
		return coveragDigests(status, eval)

	case "difference":
		return differenceDigests(status, eval)

	case "buildsystems":
		if experimental {
			return buildsystemsDigests(status, eval)
		}
		return nil, nil

	case "about_yml", "file_type":
		return nil, nil

	default:
		return nil, fmt.Errorf("Couldn't figure out how to map '%v' to a digest", status.Name)
	}
}
<|endoftext|>


<commit_before>package discover

import (
	"errors"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/flynn/rpcplus"
)

var WaitTimeoutSecs = 10

type Service struct {
	Created int
	Name    string
	Host    string
	Port    string
	Addr    string
	Attrs   map[string]string
}

type ServiceSet struct {
	services   map[string]*Service
	filters    map[string]string
	watches    map[chan *ServiceUpdate]struct{}
	serMutex   sync.Mutex
	filMutex   sync.Mutex
	watchMutex sync.Mutex
	call       *rpcplus.Call
}

func copyService(service *Service) *Service {
	s := *service
	s.Attrs = make(map[string]string, len(service.Attrs))
	for k, v := range service.Attrs {
		s.Attrs[k] = v
	}
	return &s
}

func makeServiceSet(call *rpcplus.Call) *ServiceSet {
	return &ServiceSet{
		services: make(map[string]*Service),
		filters:  make(map[string]string),
		watches:  make(map[chan *ServiceUpdate]struct{}),
		call:     call,
	}
}

func (s *ServiceSet) bind(updates chan *ServiceUpdate) chan struct{} {
	// current is an event when enough service updates have been
	// received to bring us to "current" state (when subscribed)
	current := make(chan struct{})
	go func() {
		isCurrent := false
		for update := range updates {
			if update.Addr == "" && update.Name == "" && !isCurrent {
				close(current)
				isCurrent = true
				continue
			}
			if s.filters != nil && !s.matchFilters(update.Attrs) {
				continue
			}

			s.serMutex.Lock()
			if update.Online {
				if _, exists := s.services[update.Addr]; !exists {
					host, port, _ := net.SplitHostPort(update.Addr)
					s.services[update.Addr] = &Service{
						Name: update.Name,
						Addr: update.Addr,
						Host: host,
						Port: port,
					}
				}
				s.services[update.Addr].Attrs = update.Attrs
			} else {
				if _, exists := s.services[update.Addr]; exists {
					delete(s.services, update.Addr)
				} else {
					s.serMutex.Unlock()
					continue
				}
			}
			s.serMutex.Unlock()

			s.updateWatches(update)
		}
		//s.closeWatches()
	}()
	return current
}

func (s *ServiceSet) updateWatches(update *ServiceUpdate) {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	for ch := range s.watches {
		ch <- update
	}
}

func (s *ServiceSet) closeWatches() {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	for ch := range s.watches {
		close(ch)
	}
}

func (s *ServiceSet) matchFilters(attrs map[string]string) bool {
	s.filMutex.Lock()
	defer s.filMutex.Unlock()
	for key, value := range s.filters {
		if attrs[key] != value {
			return false
		}
	}
	return true
}

// deprecated
func (s *ServiceSet) Online() []*Service {
	return s.Services()
}

func (s *ServiceSet) Services() []*Service {
	s.serMutex.Lock()
	defer s.serMutex.Unlock()
	list := make([]*Service, 0, len(s.services))
	for _, service := range s.services {
		list = append(list, copyService(service))
	}
	return list
}

// deprecated
func (s *ServiceSet) OnlineAddrs() []string {
	return s.Addrs()
}

func (s *ServiceSet) Addrs() []string {
	list := make([]string, 0, len(s.services))
	for _, service := range s.Services() {
		list = append(list, service.Addr)
	}
	return list
}

func (s *ServiceSet) Select(attrs map[string]string) []*Service {
	s.serMutex.Lock()
	defer s.serMutex.Unlock()
	list := make([]*Service, 0, len(s.services))
outer:
	for _, service := range s.services {
		for key, value := range attrs {
			if service.Attrs[key] != value {
				continue outer
			}
		}
		list = append(list, service)
	}
	return list
}

func (s *ServiceSet) Filter(attrs map[string]string) {
	s.filMutex.Lock()
	s.filters = attrs
	s.filMutex.Unlock()
	s.serMutex.Lock()
	defer s.serMutex.Unlock()
	for key, service := range s.services {
		if !s.matchFilters(service.Attrs) {
			delete(s.services, key)
		}
	}
}

func (s *ServiceSet) Watch(ch chan *ServiceUpdate, bringCurrent bool) {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	s.watches[ch] = struct{}{}
	if bringCurrent {
		go func() {
			s.serMutex.Lock()
			defer s.serMutex.Unlock()
			for _, service := range s.services {
				ch <- &ServiceUpdate{
					Name:   service.Name,
					Addr:   service.Addr,
					Online: true,
					Attrs:  service.Attrs,
				}
			}
		}()
	}
}

func (s *ServiceSet) Unwatch(ch chan *ServiceUpdate) {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	delete(s.watches, ch)
}

func (s *ServiceSet) Wait() (*ServiceUpdate, error) {
	updateCh := make(chan *ServiceUpdate, 1)
	s.Watch(updateCh, true)
	defer s.Unwatch(updateCh)
	timeout := make(chan struct{})
	go func() {
		time.Sleep(time.Duration(WaitTimeoutSecs) * time.Second)
		timeout <- struct{}{}
	}()
	select {
	case update := <-updateCh:
		return update, nil
	case <-timeout:
		return nil, errors.New("discover: wait timeout exceeded")
	}
}

func (s *ServiceSet) Close() error {
	err := s.call.CloseStream()
	if err != nil {
		return err
	}
	return nil
}

type Client struct {
	client     *rpcplus.Client
	heartbeats map[string]chan struct{}
	hbMutex    sync.Mutex
}

func NewClient() (*Client, error) {
	addr := os.Getenv("DISCOVERD")
	if addr == "" {
		addr = "127.0.0.1:1111"
	}
	return NewClientUsingAddress(addr)
}

func NewClientUsingAddress(addr string) (*Client, error) {
	client, err := rpcplus.DialHTTP("tcp", addr)
	return &Client{
		client:     client,
		heartbeats: make(map[string]chan struct{}),
	}, err
}

func pickMostPublicIp() string {
	// TODO: prefer non 10.0.0.0, 172.16.0.0, and 192.168.0.0
	addrs, _ := net.InterfaceAddrs()
	var ip string
	for _, addr := range addrs {
		ip = strings.SplitN(addr.String(), "/", 2)[0]
		if !strings.Contains(ip, "::") && ip != "127.0.0.1" {
			return ip
		}
	}
	return ip
}

func (c *Client) ServiceSet(name string) (*ServiceSet, error) {
	updates := make(chan *ServiceUpdate)
	call := c.client.StreamGo("Agent.Subscribe", &Args{
		Name: name,
	}, updates)
	set := makeServiceSet(call)
	<-set.bind(updates)
	return set, nil
}

func (c *Client) Services(name string) ([]*Service, error) {
	set, err := c.ServiceSet(name)
	if err != nil {
		return nil, err
	}
	if len(set.Services()) == 0 {
		_, err = set.Wait()
		if err != nil {
			return nil, err
		}
	}
	set.Close()
	return set.Services(), nil

}

func (c *Client) Register(name, port string, attributes map[string]string) error {
	return c.RegisterWithHost(name, pickMostPublicIp(), port, attributes)
}

func (c *Client) RegisterWithHost(name, host, port string, attributes map[string]string) error {
	args := &Args{
		Name:  name,
		Addr:  net.JoinHostPort(host, port),
		Attrs: attributes,
	}
	var ret struct{}
	err := c.client.Call("Agent.Register", args, &ret)
	if err != nil {
		return errors.New("discover: register failed: " + err.Error())
	}
	done := make(chan struct{})
	c.hbMutex.Lock()
	c.heartbeats[args.Addr] = done
	c.hbMutex.Unlock()
	go func() {
		ticker := time.NewTicker(HeartbeatIntervalSecs * time.Second) // TODO: add jitter
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				// TODO: log error here
				c.client.Call("Agent.Heartbeat", &Args{
					Name: name,
					Addr: args.Addr,
				}, &struct{}{})
			case <-done:
				return
			}
		}
	}()
	return nil
}

func (c *Client) Unregister(name, port string) error {
	return c.UnregisterWithHost(name, pickMostPublicIp(), port)
}

func (c *Client) UnregisterWithHost(name, host, port string) error {
	args := &Args{
		Name: name,
		Addr: net.JoinHostPort(host, port),
	}
	c.hbMutex.Lock()
	close(c.heartbeats[args.Addr])
	delete(c.heartbeats, args.Addr)
	c.hbMutex.Unlock()
	err := c.client.Call("Agent.Unregister", args, &struct{}{})
	if err != nil {
		return errors.New("discover: unregister failed: " + err.Error())
	}
	return nil
}
<commit_msg>discoverd: uncomment closeWatchers<commit_after>package discover

import (
	"errors"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/flynn/rpcplus"
)

var WaitTimeoutSecs = 10

type Service struct {
	Created int
	Name    string
	Host    string
	Port    string
	Addr    string
	Attrs   map[string]string
}

type ServiceSet struct {
	services   map[string]*Service
	filters    map[string]string
	watches    map[chan *ServiceUpdate]struct{}
	serMutex   sync.Mutex
	filMutex   sync.Mutex
	watchMutex sync.Mutex
	call       *rpcplus.Call
}

func copyService(service *Service) *Service {
	s := *service
	s.Attrs = make(map[string]string, len(service.Attrs))
	for k, v := range service.Attrs {
		s.Attrs[k] = v
	}
	return &s
}

func makeServiceSet(call *rpcplus.Call) *ServiceSet {
	return &ServiceSet{
		services: make(map[string]*Service),
		filters:  make(map[string]string),
		watches:  make(map[chan *ServiceUpdate]struct{}),
		call:     call,
	}
}

func (s *ServiceSet) bind(updates chan *ServiceUpdate) chan struct{} {
	// current is an event when enough service updates have been
	// received to bring us to "current" state (when subscribed)
	current := make(chan struct{})
	go func() {
		isCurrent := false
		for update := range updates {
			if update.Addr == "" && update.Name == "" && !isCurrent {
				close(current)
				isCurrent = true
				continue
			}
			if s.filters != nil && !s.matchFilters(update.Attrs) {
				continue
			}

			s.serMutex.Lock()
			if update.Online {
				if _, exists := s.services[update.Addr]; !exists {
					host, port, _ := net.SplitHostPort(update.Addr)
					s.services[update.Addr] = &Service{
						Name: update.Name,
						Addr: update.Addr,
						Host: host,
						Port: port,
					}
				}
				s.services[update.Addr].Attrs = update.Attrs
			} else {
				if _, exists := s.services[update.Addr]; exists {
					delete(s.services, update.Addr)
				} else {
					s.serMutex.Unlock()
					continue
				}
			}
			s.serMutex.Unlock()

			s.updateWatches(update)
		}
		s.closeWatches()
	}()
	return current
}

func (s *ServiceSet) updateWatches(update *ServiceUpdate) {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	for ch := range s.watches {
		ch <- update
	}
}

func (s *ServiceSet) closeWatches() {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	for ch := range s.watches {
		close(ch)
	}
}

func (s *ServiceSet) matchFilters(attrs map[string]string) bool {
	s.filMutex.Lock()
	defer s.filMutex.Unlock()
	for key, value := range s.filters {
		if attrs[key] != value {
			return false
		}
	}
	return true
}

// deprecated
func (s *ServiceSet) Online() []*Service {
	return s.Services()
}

func (s *ServiceSet) Services() []*Service {
	s.serMutex.Lock()
	defer s.serMutex.Unlock()
	list := make([]*Service, 0, len(s.services))
	for _, service := range s.services {
		list = append(list, copyService(service))
	}
	return list
}

// deprecated
func (s *ServiceSet) OnlineAddrs() []string {
	return s.Addrs()
}

func (s *ServiceSet) Addrs() []string {
	list := make([]string, 0, len(s.services))
	for _, service := range s.Services() {
		list = append(list, service.Addr)
	}
	return list
}

func (s *ServiceSet) Select(attrs map[string]string) []*Service {
	s.serMutex.Lock()
	defer s.serMutex.Unlock()
	list := make([]*Service, 0, len(s.services))
outer:
	for _, service := range s.services {
		for key, value := range attrs {
			if service.Attrs[key] != value {
				continue outer
			}
		}
		list = append(list, service)
	}
	return list
}

func (s *ServiceSet) Filter(attrs map[string]string) {
	s.filMutex.Lock()
	s.filters = attrs
	s.filMutex.Unlock()
	s.serMutex.Lock()
	defer s.serMutex.Unlock()
	for key, service := range s.services {
		if !s.matchFilters(service.Attrs) {
			delete(s.services, key)
		}
	}
}

func (s *ServiceSet) Watch(ch chan *ServiceUpdate, bringCurrent bool) {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	s.watches[ch] = struct{}{}
	if bringCurrent {
		go func() {
			s.serMutex.Lock()
			defer s.serMutex.Unlock()
			for _, service := range s.services {
				ch <- &ServiceUpdate{
					Name:   service.Name,
					Addr:   service.Addr,
					Online: true,
					Attrs:  service.Attrs,
				}
			}
		}()
	}
}

func (s *ServiceSet) Unwatch(ch chan *ServiceUpdate) {
	s.watchMutex.Lock()
	defer s.watchMutex.Unlock()
	delete(s.watches, ch)
}

func (s *ServiceSet) Wait() (*ServiceUpdate, error) {
	updateCh := make(chan *ServiceUpdate, 1)
	s.Watch(updateCh, true)
	defer s.Unwatch(updateCh)
	timeout := make(chan struct{})
	go func() {
		time.Sleep(time.Duration(WaitTimeoutSecs) * time.Second)
		timeout <- struct{}{}
	}()
	select {
	case update := <-updateCh:
		return update, nil
	case <-timeout:
		return nil, errors.New("discover: wait timeout exceeded")
	}
}

func (s *ServiceSet) Close() error {
	err := s.call.CloseStream()
	if err != nil {
		return err
	}
	return nil
}

type Client struct {
	client     *rpcplus.Client
	heartbeats map[string]chan struct{}
	hbMutex    sync.Mutex
}

func NewClient() (*Client, error) {
	addr := os.Getenv("DISCOVERD")
	if addr == "" {
		addr = "127.0.0.1:1111"
	}
	return NewClientUsingAddress(addr)
}

func NewClientUsingAddress(addr string) (*Client, error) {
	client, err := rpcplus.DialHTTP("tcp", addr)
	return &Client{
		client:     client,
		heartbeats: make(map[string]chan struct{}),
	}, err
}

func pickMostPublicIp() string {
	// TODO: prefer non 10.0.0.0, 172.16.0.0, and 192.168.0.0
	addrs, _ := net.InterfaceAddrs()
	var ip string
	for _, addr := range addrs {
		ip = strings.SplitN(addr.String(), "/", 2)[0]
		if !strings.Contains(ip, "::") && ip != "127.0.0.1" {
			return ip
		}
	}
	return ip
}

func (c *Client) ServiceSet(name string) (*ServiceSet, error) {
	updates := make(chan *ServiceUpdate)
	call := c.client.StreamGo("Agent.Subscribe", &Args{
		Name: name,
	}, updates)
	set := makeServiceSet(call)
	<-set.bind(updates)
	return set, nil
}

func (c *Client) Services(name string) ([]*Service, error) {
	set, err := c.ServiceSet(name)
	if err != nil {
		return nil, err
	}
	if len(set.Services()) == 0 {
		_, err = set.Wait()
		if err != nil {
			return nil, err
		}
	}
	set.Close()
	return set.Services(), nil

}

func (c *Client) Register(name, port string, attributes map[string]string) error {
	return c.RegisterWithHost(name, pickMostPublicIp(), port, attributes)
}

func (c *Client) RegisterWithHost(name, host, port string, attributes map[string]string) error {
	args := &Args{
		Name:  name,
		Addr:  net.JoinHostPort(host, port),
		Attrs: attributes,
	}
	var ret struct{}
	err := c.client.Call("Agent.Register", args, &ret)
	if err != nil {
		return errors.New("discover: register failed: " + err.Error())
	}
	done := make(chan struct{})
	c.hbMutex.Lock()
	c.heartbeats[args.Addr] = done
	c.hbMutex.Unlock()
	go func() {
		ticker := time.NewTicker(HeartbeatIntervalSecs * time.Second) // TODO: add jitter
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				// TODO: log error here
				c.client.Call("Agent.Heartbeat", &Args{
					Name: name,
					Addr: args.Addr,
				}, &struct{}{})
			case <-done:
				return
			}
		}
	}()
	return nil
}

func (c *Client) Unregister(name, port string) error {
	return c.UnregisterWithHost(name, pickMostPublicIp(), port)
}

func (c *Client) UnregisterWithHost(name, host, port string) error {
	args := &Args{
		Name: name,
		Addr: net.JoinHostPort(host, port),
	}
	c.hbMutex.Lock()
	close(c.heartbeats[args.Addr])
	delete(c.heartbeats, args.Addr)
	c.hbMutex.Unlock()
	err := c.client.Call("Agent.Unregister", args, &struct{}{})
	if err != nil {
		return errors.New("discover: unregister failed: " + err.Error())
	}
	return nil
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2016 Jamie Alquiza
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package tachymeter

import (
	"fmt"
	"math"
	"sort"
	"sync/atomic"
	"time"
)

// Satisfy sort for timeSlice.
// Sorts in increasing order of duration.

func (p timeSlice) Len() int {
	return len(p)
}

func (p timeSlice) Less(i, j int) bool {
	return int64(p[i]) < int64(p[j])
}

func (p timeSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

// Calc summarizes Tachymeter sample data
// and returns it in the form of a *Metrics.
func (m *Tachymeter) Calc() *Metrics {
	metrics := &Metrics{}
	if m.Count == 0 {
		return metrics
	}

	m.Lock()

	metrics.Samples = int(math.Min(float64(atomic.LoadUint64(&m.Count)), float64(m.Size)))
	metrics.Count = int(atomic.LoadUint64(&m.Count))
	times := make(timeSlice, metrics.Samples)
	copy(times, m.Times[:metrics.Samples])
	sort.Sort(times)

	metrics.Time.Cumulative = calcTimeCumulative(times)
	var rateTime float64
	if m.WallTime != 0 {
		rateTime = float64(metrics.Count) / float64(m.WallTime)
	} else {
		rateTime = float64(metrics.Samples) / float64(metrics.Time.Cumulative)
	}

	metrics.Rate.Second = rateTime * 1e9

	m.Unlock()

	metrics.Time.Avg = calcAvg(times, metrics.Samples)
	metrics.Time.P50 = times[len(times)/2]
	metrics.Time.P75 = calcP(times, 0.75)
	metrics.Time.P95 = calcP(times, 0.95)
	metrics.Time.P99 = calcP(times, 0.99)
	metrics.Time.P999 = calcP(times, 0.999)
	metrics.Time.Long5p = calcLong5p(times)
	metrics.Time.Short5p = calcShort5p(times)
	metrics.Time.Max = times[metrics.Samples-1]
	metrics.Time.Min = times[0]
	metrics.Time.Range = metrics.Time.Max - metrics.Time.Min

	metrics.Histogram = calcHgram(m.HBuckets, times, metrics.Time.Min, metrics.Time.Max, metrics.Time.Range)

	return metrics
}

// calcHgram returns a histogram of event durations t in b buckets.
// A histogram bucket is a map["low-high duration"]count of events that
// fall within the low / high range.
func calcHgram(b int, t timeSlice, low, max, r time.Duration) []map[string]int {
	// Interval is the time range / n buckets.
	interval := time.Duration(int64(r) / int64(b))
	high := low + interval
	hgram := []map[string]int{}
	pos := 1 // Bucket position.

	bstring := fmt.Sprintf("%s-%s", low, high)
	bucket := map[string]int{}

	for _, v := range t {
		// If v fits in the current bucket,
		// increment the bucket count.
		if v <= high {
			bucket[bstring]++
		} else {
			// If not, prepare the next bucket.
			hgram = append(hgram, bucket)
			bucket = map[string]int{}

			// Update the high/low range values.
			low = high + time.Nanosecond

			high += interval
			// if we're going into the
			// last bucket, set high to max.
			if pos == b-1 {
				high = max
			}

			bstring = fmt.Sprintf("%s - %s", low, high)

			// The value didn't fit in the previous
			// bucket, so the new bucket count should
			// be incremented.
			bucket[bstring]++

			pos++
		}
	}

	hgram = append(hgram, bucket)

	return hgram
}

// These should be self-explanatory:

func calcTimeCumulative(d []time.Duration) time.Duration {
	var total time.Duration
	for _, t := range d {
		total += t
	}

	return total
}

func calcAvg(d []time.Duration, c int) time.Duration {
	var total time.Duration
	for _, t := range d {
		total += t
	}
	return time.Duration(int(total) / c)
}

func calcP(d []time.Duration, p float64) time.Duration {
	return d[int(float64(len(d))*p+0.5)-1]
}

func calcLong5p(d []time.Duration) time.Duration {
	set := d[int(float64(len(d))*0.95+0.5):]

	if len(set) <= 1 {
		return d[len(d)-1]
	}

	var t time.Duration
	var i int
	for _, n := range set {
		t += n
		i++
	}

	return time.Duration(int(t) / i)
}

func calcShort5p(d []time.Duration) time.Duration {
	set := d[:int(float64(len(d))*0.05+0.5)]

	if len(set) <= 1 {
		return d[0]
	}

	var t time.Duration
	var i int
	for _, n := range set {
		t += n
		i++
	}

	return time.Duration(int(t) / i)
}
<commit_msg>histogram key consistency<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2016 Jamie Alquiza
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package tachymeter

import (
	"fmt"
	"math"
	"sort"
	"sync/atomic"
	"time"
)

// Satisfy sort for timeSlice.
// Sorts in increasing order of duration.

func (p timeSlice) Len() int {
	return len(p)
}

func (p timeSlice) Less(i, j int) bool {
	return int64(p[i]) < int64(p[j])
}

func (p timeSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

// Calc summarizes Tachymeter sample data
// and returns it in the form of a *Metrics.
func (m *Tachymeter) Calc() *Metrics {
	metrics := &Metrics{}
	if m.Count == 0 {
		return metrics
	}

	m.Lock()

	metrics.Samples = int(math.Min(float64(atomic.LoadUint64(&m.Count)), float64(m.Size)))
	metrics.Count = int(atomic.LoadUint64(&m.Count))
	times := make(timeSlice, metrics.Samples)
	copy(times, m.Times[:metrics.Samples])
	sort.Sort(times)

	metrics.Time.Cumulative = calcTimeCumulative(times)
	var rateTime float64
	if m.WallTime != 0 {
		rateTime = float64(metrics.Count) / float64(m.WallTime)
	} else {
		rateTime = float64(metrics.Samples) / float64(metrics.Time.Cumulative)
	}

	metrics.Rate.Second = rateTime * 1e9

	m.Unlock()

	metrics.Time.Avg = calcAvg(times, metrics.Samples)
	metrics.Time.P50 = times[len(times)/2]
	metrics.Time.P75 = calcP(times, 0.75)
	metrics.Time.P95 = calcP(times, 0.95)
	metrics.Time.P99 = calcP(times, 0.99)
	metrics.Time.P999 = calcP(times, 0.999)
	metrics.Time.Long5p = calcLong5p(times)
	metrics.Time.Short5p = calcShort5p(times)
	metrics.Time.Max = times[metrics.Samples-1]
	metrics.Time.Min = times[0]
	metrics.Time.Range = metrics.Time.Max - metrics.Time.Min

	metrics.Histogram = calcHgram(m.HBuckets, times, metrics.Time.Min, metrics.Time.Max, metrics.Time.Range)

	return metrics
}

// calcHgram returns a histogram of event durations t in b buckets.
// A histogram bucket is a map["low-high duration"]count of events that
// fall within the low / high range.
func calcHgram(b int, t timeSlice, low, max, r time.Duration) []map[string]int {
	// Interval is the time range / n buckets.
	interval := time.Duration(int64(r) / int64(b))
	high := low + interval
	hgram := []map[string]int{}
	pos := 1 // Bucket position.

	bstring := fmt.Sprintf("%s - %s", low, high)
	bucket := map[string]int{}

	for _, v := range t {
		// If v fits in the current bucket,
		// increment the bucket count.
		if v <= high {
			bucket[bstring]++
		} else {
			// If not, prepare the next bucket.
			hgram = append(hgram, bucket)
			bucket = map[string]int{}

			// Update the high/low range values.
			low = high + time.Nanosecond

			high += interval
			// if we're going into the
			// last bucket, set high to max.
			if pos == b-1 {
				high = max
			}

			bstring = fmt.Sprintf("%s - %s", low, high)

			// The value didn't fit in the previous
			// bucket, so the new bucket count should
			// be incremented.
			bucket[bstring]++

			pos++
		}
	}

	hgram = append(hgram, bucket)

	return hgram
}

// These should be self-explanatory:

func calcTimeCumulative(d []time.Duration) time.Duration {
	var total time.Duration
	for _, t := range d {
		total += t
	}

	return total
}

func calcAvg(d []time.Duration, c int) time.Duration {
	var total time.Duration
	for _, t := range d {
		total += t
	}
	return time.Duration(int(total) / c)
}

func calcP(d []time.Duration, p float64) time.Duration {
	return d[int(float64(len(d))*p+0.5)-1]
}

func calcLong5p(d []time.Duration) time.Duration {
	set := d[int(float64(len(d))*0.95+0.5):]

	if len(set) <= 1 {
		return d[len(d)-1]
	}

	var t time.Duration
	var i int
	for _, n := range set {
		t += n
		i++
	}

	return time.Duration(int(t) / i)
}

func calcShort5p(d []time.Duration) time.Duration {
	set := d[:int(float64(len(d))*0.05+0.5)]

	if len(set) <= 1 {
		return d[0]
	}

	var t time.Duration
	var i int
	for _, n := range set {
		t += n
		i++
	}

	return time.Duration(int(t) / i)
}
<|endoftext|>