<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package plugin

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/getgauge/common"
	"github.com/getgauge/gauge/config"
	"github.com/getgauge/gauge/conn"
	"github.com/getgauge/gauge/gauge_messages"
	"github.com/getgauge/gauge/logger"
	"github.com/getgauge/gauge/manifest"
	"github.com/getgauge/gauge/plugin/pluginInfo"
	"github.com/getgauge/gauge/reporter"
	"github.com/getgauge/gauge/version"
	"github.com/golang/protobuf/proto"
)

type pluginScope string

const (
	executionScope          pluginScope = "execution"
	docScope                pluginScope = "documentation"
	pluginConnectionPortEnv             = "plugin_connection_port"
	debugEnv                            = "debugging"
)

type plugin struct {
	mutex      *sync.Mutex
	connection net.Conn
	pluginCmd  *exec.Cmd
	descriptor *pluginDescriptor
	killTimer  *time.Timer
}

func isProcessRunning(p *plugin) bool {
	p.mutex.Lock()
	ps := p.pluginCmd.ProcessState
	p.mutex.Unlock()
	return ps == nil || !ps.Exited()
}

func (p *plugin) rejuvenate() error {
	if p.killTimer == nil {
		return fmt.Errorf("timer is uninitialized. Perhaps kill is not yet invoked")
	}
	logger.Debugf(true, "Extending the plugin_kill_timeout for %s", p.descriptor.ID)
	p.killTimer.Reset(config.PluginConnectionTimeout())
	return nil
}

func (p *plugin) kill(wg *sync.WaitGroup) error {
	defer wg.Done()
	if isProcessRunning(p) {
		defer p.connection.Close()
		p.killTimer = time.NewTimer(config.PluginConnectionTimeout())
		conn.SendProcessKillMessage(p.connection)

		exited := make(chan bool, 1)
		go func() {
			for {
				if isProcessRunning(p) {
					time.Sleep(100 * time.Millisecond)
				} else {
					exited <- true
					return
				}
			}
		}()
		select {
		case <-exited:
			if !p.killTimer.Stop() {
				<-p.killTimer.C
			}
			logger.Debugf(true, "Plugin [%s] with pid [%d] has exited", p.descriptor.Name, p.pluginCmd.Process.Pid)
		case <-p.killTimer.C:
			logger.Warningf(true, "Plugin [%s] with pid [%d] did not exit after %.2f seconds. Forcefully killing it.", p.descriptor.Name, p.pluginCmd.Process.Pid, config.PluginKillTimeout().Seconds())
			err := p.pluginCmd.Process.Kill()
			if err != nil {
				logger.Warningf(true, "Error while killing plugin %s : %s ", p.descriptor.Name, err.Error())
			}
			return err
		}
	}
	return nil
}

func IsPluginInstalled(pluginName, pluginVersion string) bool {
	pluginsInstallDir, err := common.GetPluginsInstallDir(pluginName)
	if err != nil {
		return false
	}

	thisPluginDir := filepath.Join(pluginsInstallDir, pluginName)
	if !common.DirExists(thisPluginDir) {
		return false
	}

	if pluginVersion != "" {
		pluginJSON := filepath.Join(thisPluginDir, pluginVersion, common.PluginJSONFile)
		if common.FileExists(pluginJSON) {
			return true
		}
		return false
	}
	return true
}

func getPluginJSONPath(pluginName, pluginVersion string) (string, error) {
	if !IsPluginInstalled(pluginName, pluginVersion) {
		plugin := strings.TrimSpace(fmt.Sprintf("%s %s", pluginName, pluginVersion))
		return "", fmt.Errorf("Plugin %s is not installed", plugin)
	}

	pluginInstallDir, err := GetInstallDir(pluginName, "")
	if err != nil {
		return "", err
	}
	return filepath.Join(pluginInstallDir, common.PluginJSONFile), nil
}

func GetPluginDescriptor(pluginID, pluginVersion string) (*pluginDescriptor, error) {
	pluginJSON, err := getPluginJSONPath(pluginID, pluginVersion)
	if err != nil {
		return nil, err
	}
	return GetPluginDescriptorFromJSON(pluginJSON)
}

func GetPluginDescriptorFromJSON(pluginJSON string) (*pluginDescriptor, error) {
	pluginJSONContents, err := common.ReadFileContents(pluginJSON)
	if err != nil {
		return nil, err
	}
	var pd pluginDescriptor
	if err = json.Unmarshal([]byte(pluginJSONContents), &pd); err != nil {
		return nil, fmt.Errorf("%s: %s", pluginJSON, err.Error())
	}
	pd.pluginPath = filepath.Dir(pluginJSON)

	return &pd, nil
}

func StartPlugin(pd *pluginDescriptor, action pluginScope) (*plugin, error) {
	command := []string{}
	switch runtime.GOOS {
	case "windows":
		command = pd.Command.Windows
		break
	case "darwin":
		command = pd.Command.Darwin
		break
	default:
		command = pd.Command.Linux
		break
	}
	if len(command) == 0 {
		return nil, fmt.Errorf("Platform specific command not specified: %s.", runtime.GOOS)
	}

	cmd, err := common.ExecuteCommand(command, pd.pluginPath, reporter.Current(), reporter.Current())

	if err != nil {
		return nil, err
	}
	var mutex = &sync.Mutex{}
	go func() {
		pState, _ := cmd.Process.Wait()
		mutex.Lock()
		cmd.ProcessState = pState
		mutex.Unlock()
	}()
	plugin := &plugin{pluginCmd: cmd, descriptor: pd, mutex: mutex}
	return plugin, nil
}

func SetEnvForPlugin(action pluginScope, pd *pluginDescriptor, manifest *manifest.Manifest, pluginEnvVars map[string]string) error {
	pluginEnvVars[fmt.Sprintf("%s_action", pd.ID)] = string(action)
	pluginEnvVars["test_language"] = manifest.Language
	if err := setEnvironmentProperties(pluginEnvVars); err != nil {
		return err
	}
	return nil
}

func setEnvironmentProperties(properties map[string]string) error {
	for k, v := range properties {
		if err := common.SetEnvVariable(k, v); err != nil {
			return err
		}
	}
	return nil
}

func IsPluginAdded(manifest *manifest.Manifest, descriptor *pluginDescriptor) bool {
	for _, pluginID := range manifest.Plugins {
		if pluginID == descriptor.ID {
			return true
		}
	}
	return false
}

func startPluginsForExecution(manifest *manifest.Manifest) (Handler, []string) {
	var warnings []string
	handler := &GaugePlugins{}
	envProperties := make(map[string]string)

	for _, pluginID := range manifest.Plugins {
		pd, err := GetPluginDescriptor(pluginID, "")
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("Unable to start plugin %s. %s. To install, run `gauge install %s`.", pluginID, err.Error(), pluginID))
			continue
		}
		compatibilityErr := version.CheckCompatibility(version.CurrentGaugeVersion, &pd.GaugeVersionSupport)
		if compatibilityErr != nil {
			warnings = append(warnings, fmt.Sprintf("Compatible %s plugin version to current Gauge version %s not found", pd.Name, version.CurrentGaugeVersion))
			continue
		}
		if pd.hasScope(executionScope) {
			gaugeConnectionHandler, err := conn.NewGaugeConnectionHandler(0, &keepAliveHandler{ph: handler})
			if err != nil {
				warnings = append(warnings, err.Error())
				continue
			}
			envProperties[pluginConnectionPortEnv] = strconv.Itoa(gaugeConnectionHandler.ConnectionPortNumber())
			prop, err := common.GetGaugeConfiguration()
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Unable to read Gauge configuration. %s", err.Error()))
				continue
			}
			envProperties["plugin_kill_timeout"] = prop["plugin_kill_timeout"]
			err = SetEnvForPlugin(executionScope, pd, manifest, envProperties)
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Error setting environment for plugin %s %s. %s", pd.Name, pd.Version, err.Error()))
				continue
			}

			plugin, err := StartPlugin(pd, executionScope)
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Error starting plugin %s %s. %s", pd.Name, pd.Version, err.Error()))
				continue
			}
			pluginConnection, err := gaugeConnectionHandler.AcceptConnection(config.PluginConnectionTimeout(), make(chan error))
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Error starting plugin %s %s. Failed to connect to plugin. %s", pd.Name, pd.Version, err.Error()))
				plugin.pluginCmd.Process.Kill()
				continue
			}
			plugin.connection = pluginConnection
			handler.addPlugin(pluginID, plugin)
		}

	}
	return handler, warnings
}

func GenerateDoc(pluginName string, specDirs []string, port int) {
	pd, err := GetPluginDescriptor(pluginName, "")
	if err != nil {
		logger.Fatalf(true, "Error starting plugin %s. Failed to get plugin.json. %s. To install, run `gauge install %s`.", pluginName, err.Error(), pluginName)
	}
	if err := version.CheckCompatibility(version.CurrentGaugeVersion, &pd.GaugeVersionSupport); err != nil {
		logger.Fatalf(true, "Compatible %s plugin version to current Gauge version %s not found", pd.Name, version.CurrentGaugeVersion)
	}
	if !pd.hasScope(docScope) {
		logger.Fatalf(true, "Invalid plugin name: %s, this plugin cannot generate documentation.", pd.Name)
	}
	var sources []string
	for _, src := range specDirs {
		path, _ := filepath.Abs(src)
		sources = append(sources, path)
	}
	os.Setenv("GAUGE_SPEC_DIRS", strings.Join(sources, "||"))
	os.Setenv("GAUGE_PROJECT_ROOT", config.ProjectRoot)
	os.Setenv(common.APIPortEnvVariableName, strconv.Itoa(port))
	p, err := StartPlugin(pd, docScope)
	if err != nil {
		logger.Fatalf(true, "Error starting plugin %s %s. %s", pd.Name, pd.Version, err.Error())
	}
	for isProcessRunning(p) {
	}
}

func (p *plugin) sendMessage(message *gauge_messages.Message) error {
	messageID := common.GetUniqueID()
	message.MessageId = messageID
	messageBytes, err := proto.Marshal(message)
	if err != nil {
		return err
	}
	err = conn.Write(p.connection, messageBytes)
	if err != nil {
		return fmt.Errorf("[Warning] Failed to send message to plugin: %s  %s", p.descriptor.ID, err.Error())
	}
	return nil
}

func StartPlugins(manifest *manifest.Manifest) Handler {
	pluginHandler, warnings := startPluginsForExecution(manifest)
	logger.HandleWarningMessages(true, warnings)
	return pluginHandler
}

func PluginsWithoutScope() (infos []pluginInfo.PluginInfo) {
	if plugins, err := pluginInfo.GetAllInstalledPluginsWithVersion(); err == nil {
		for _, p := range plugins {
			pd, err := GetPluginDescriptor(p.Name, p.Version.String())
			if err == nil && !pd.hasAnyScope() {
				infos = append(infos, p)
			}
		}
	}
	return
}

// GetInstallDir returns the install directory of given plugin and a given version.
func GetInstallDir(pluginName, version string) (string, error) {
	allPluginsInstallDir, err := common.GetPluginsInstallDir(pluginName)
	if err != nil {
		return "", err
	}
	pluginDir := filepath.Join(allPluginsInstallDir, pluginName)
	if version != "" {
		pluginDir = filepath.Join(pluginDir, version)
	} else {
		latestPlugin, err := pluginInfo.GetLatestInstalledPlugin(pluginDir)
		if err != nil {
			return "", err
		}
		pluginDir = latestPlugin.Path
	}
	return pluginDir, nil
}

func GetLanguageJSONFilePath(language string) (string, error) {
	languageInstallDir, err := GetInstallDir(language, "")
	if err != nil {
		return "", err
	}
	languageJSON := filepath.Join(languageInstallDir, fmt.Sprintf("%s.json", language))
	if !common.FileExists(languageJSON) {
		return "", fmt.Errorf("Failed to find the implementation for: %s. %s does not exist.", language, languageJSON)
	}

	return languageJSON, nil
}

func IsLanguagePlugin(plugin string) bool {
	if _, err := GetLanguageJSONFilePath(plugin); err != nil {
		return false
	}
	return true
}

func QueryParams() string {
	return fmt.Sprintf("?l=%s&p=%s&o=%s&a=%s", language(), plugins(), runtime.GOOS, runtime.GOARCH)
}

func language() string {
	if config.ProjectRoot == "" {
		return ""
	}
	m, err := manifest.ProjectManifest()
	if err != nil {
		return ""
	}
	return m.Language
}

func plugins() string {
	pluginInfos, err := pluginInfo.GetAllInstalledPluginsWithVersion()
	if err != nil {
		return ""
	}
	var plugins []string
	for _, p := range pluginInfos {
		plugins = append(plugins, p.Name)
	}
	return strings.Join(plugins, ",")
}
<commit_msg>Fixed plugin kill timeout, #1176<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package plugin

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/getgauge/common"
	"github.com/getgauge/gauge/config"
	"github.com/getgauge/gauge/conn"
	"github.com/getgauge/gauge/gauge_messages"
	"github.com/getgauge/gauge/logger"
	"github.com/getgauge/gauge/manifest"
	"github.com/getgauge/gauge/plugin/pluginInfo"
	"github.com/getgauge/gauge/reporter"
	"github.com/getgauge/gauge/version"
	"github.com/golang/protobuf/proto"
)

type pluginScope string

const (
	executionScope          pluginScope = "execution"
	docScope                pluginScope = "documentation"
	pluginConnectionPortEnv             = "plugin_connection_port"
	debugEnv                            = "debugging"
)

type plugin struct {
	mutex      *sync.Mutex
	connection net.Conn
	pluginCmd  *exec.Cmd
	descriptor *pluginDescriptor
	killTimer  *time.Timer
}

func isProcessRunning(p *plugin) bool {
	p.mutex.Lock()
	ps := p.pluginCmd.ProcessState
	p.mutex.Unlock()
	return ps == nil || !ps.Exited()
}

func (p *plugin) rejuvenate() error {
	if p.killTimer == nil {
		return fmt.Errorf("timer is uninitialized. Perhaps kill is not yet invoked")
	}
	logger.Debugf(true, "Extending the plugin_kill_timeout for %s", p.descriptor.ID)
	p.killTimer.Reset(config.PluginConnectionTimeout())
	return nil
}

func (p *plugin) kill(wg *sync.WaitGroup) error {
	defer wg.Done()
	if isProcessRunning(p) {
		defer p.connection.Close()
		p.killTimer = time.NewTimer(config.PluginKillTimeout())
		conn.SendProcessKillMessage(p.connection)

		exited := make(chan bool, 1)
		go func() {
			for {
				if isProcessRunning(p) {
					time.Sleep(100 * time.Millisecond)
				} else {
					exited <- true
					return
				}
			}
		}()
		select {
		case <-exited:
			if !p.killTimer.Stop() {
				<-p.killTimer.C
			}
			logger.Debugf(true, "Plugin [%s] with pid [%d] has exited", p.descriptor.Name, p.pluginCmd.Process.Pid)
		case <-p.killTimer.C:
			logger.Warningf(true, "Plugin [%s] with pid [%d] did not exit after %.2f seconds. Forcefully killing it.", p.descriptor.Name, p.pluginCmd.Process.Pid, config.PluginKillTimeout().Seconds())
			err := p.pluginCmd.Process.Kill()
			if err != nil {
				logger.Warningf(true, "Error while killing plugin %s : %s ", p.descriptor.Name, err.Error())
			}
			return err
		}
	}
	return nil
}

func IsPluginInstalled(pluginName, pluginVersion string) bool {
	pluginsInstallDir, err := common.GetPluginsInstallDir(pluginName)
	if err != nil {
		return false
	}

	thisPluginDir := filepath.Join(pluginsInstallDir, pluginName)
	if !common.DirExists(thisPluginDir) {
		return false
	}

	if pluginVersion != "" {
		pluginJSON := filepath.Join(thisPluginDir, pluginVersion, common.PluginJSONFile)
		if common.FileExists(pluginJSON) {
			return true
		}
		return false
	}
	return true
}

func getPluginJSONPath(pluginName, pluginVersion string) (string, error) {
	if !IsPluginInstalled(pluginName, pluginVersion) {
		plugin := strings.TrimSpace(fmt.Sprintf("%s %s", pluginName, pluginVersion))
		return "", fmt.Errorf("Plugin %s is not installed", plugin)
	}

	pluginInstallDir, err := GetInstallDir(pluginName, "")
	if err != nil {
		return "", err
	}
	return filepath.Join(pluginInstallDir, common.PluginJSONFile), nil
}

func GetPluginDescriptor(pluginID, pluginVersion string) (*pluginDescriptor, error) {
	pluginJSON, err := getPluginJSONPath(pluginID, pluginVersion)
	if err != nil {
		return nil, err
	}
	return GetPluginDescriptorFromJSON(pluginJSON)
}

func GetPluginDescriptorFromJSON(pluginJSON string) (*pluginDescriptor, error) {
	pluginJSONContents, err := common.ReadFileContents(pluginJSON)
	if err != nil {
		return nil, err
	}
	var pd pluginDescriptor
	if err = json.Unmarshal([]byte(pluginJSONContents), &pd); err != nil {
		return nil, fmt.Errorf("%s: %s", pluginJSON, err.Error())
	}
	pd.pluginPath = filepath.Dir(pluginJSON)

	return &pd, nil
}

func StartPlugin(pd *pluginDescriptor, action pluginScope) (*plugin, error) {
	command := []string{}
	switch runtime.GOOS {
	case "windows":
		command = pd.Command.Windows
		break
	case "darwin":
		command = pd.Command.Darwin
		break
	default:
		command = pd.Command.Linux
		break
	}
	if len(command) == 0 {
		return nil, fmt.Errorf("Platform specific command not specified: %s.", runtime.GOOS)
	}

	cmd, err := common.ExecuteCommand(command, pd.pluginPath, reporter.Current(), reporter.Current())

	if err != nil {
		return nil, err
	}
	var mutex = &sync.Mutex{}
	go func() {
		pState, _ := cmd.Process.Wait()
		mutex.Lock()
		cmd.ProcessState = pState
		mutex.Unlock()
	}()
	plugin := &plugin{pluginCmd: cmd, descriptor: pd, mutex: mutex}
	return plugin, nil
}

func SetEnvForPlugin(action pluginScope, pd *pluginDescriptor, manifest *manifest.Manifest, pluginEnvVars map[string]string) error {
	pluginEnvVars[fmt.Sprintf("%s_action", pd.ID)] = string(action)
	pluginEnvVars["test_language"] = manifest.Language
	if err := setEnvironmentProperties(pluginEnvVars); err != nil {
		return err
	}
	return nil
}

func setEnvironmentProperties(properties map[string]string) error {
	for k, v := range properties {
		if err := common.SetEnvVariable(k, v); err != nil {
			return err
		}
	}
	return nil
}

func IsPluginAdded(manifest *manifest.Manifest, descriptor *pluginDescriptor) bool {
	for _, pluginID := range manifest.Plugins {
		if pluginID == descriptor.ID {
			return true
		}
	}
	return false
}

func startPluginsForExecution(manifest *manifest.Manifest) (Handler, []string) {
	var warnings []string
	handler := &GaugePlugins{}
	envProperties := make(map[string]string)

	for _, pluginID := range manifest.Plugins {
		pd, err := GetPluginDescriptor(pluginID, "")
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("Unable to start plugin %s. %s. To install, run `gauge install %s`.", pluginID, err.Error(), pluginID))
			continue
		}
		compatibilityErr := version.CheckCompatibility(version.CurrentGaugeVersion, &pd.GaugeVersionSupport)
		if compatibilityErr != nil {
			warnings = append(warnings, fmt.Sprintf("Compatible %s plugin version to current Gauge version %s not found", pd.Name, version.CurrentGaugeVersion))
			continue
		}
		if pd.hasScope(executionScope) {
			gaugeConnectionHandler, err := conn.NewGaugeConnectionHandler(0, &keepAliveHandler{ph: handler})
			if err != nil {
				warnings = append(warnings, err.Error())
				continue
			}
			envProperties[pluginConnectionPortEnv] = strconv.Itoa(gaugeConnectionHandler.ConnectionPortNumber())
			prop, err := common.GetGaugeConfiguration()
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Unable to read Gauge configuration. %s", err.Error()))
				continue
			}
			envProperties["plugin_kill_timeout"] = prop["plugin_kill_timeout"]
			err = SetEnvForPlugin(executionScope, pd, manifest, envProperties)
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Error setting environment for plugin %s %s. %s", pd.Name, pd.Version, err.Error()))
				continue
			}

			plugin, err := StartPlugin(pd, executionScope)
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Error starting plugin %s %s. %s", pd.Name, pd.Version, err.Error()))
				continue
			}
			pluginConnection, err := gaugeConnectionHandler.AcceptConnection(config.PluginConnectionTimeout(), make(chan error))
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("Error starting plugin %s %s. Failed to connect to plugin. %s", pd.Name, pd.Version, err.Error()))
				plugin.pluginCmd.Process.Kill()
				continue
			}
			plugin.connection = pluginConnection
			handler.addPlugin(pluginID, plugin)
		}

	}
	return handler, warnings
}

func GenerateDoc(pluginName string, specDirs []string, port int) {
	pd, err := GetPluginDescriptor(pluginName, "")
	if err != nil {
		logger.Fatalf(true, "Error starting plugin %s. Failed to get plugin.json. %s. To install, run `gauge install %s`.", pluginName, err.Error(), pluginName)
	}
	if err := version.CheckCompatibility(version.CurrentGaugeVersion, &pd.GaugeVersionSupport); err != nil {
		logger.Fatalf(true, "Compatible %s plugin version to current Gauge version %s not found", pd.Name, version.CurrentGaugeVersion)
	}
	if !pd.hasScope(docScope) {
		logger.Fatalf(true, "Invalid plugin name: %s, this plugin cannot generate documentation.", pd.Name)
	}
	var sources []string
	for _, src := range specDirs {
		path, _ := filepath.Abs(src)
		sources = append(sources, path)
	}
	os.Setenv("GAUGE_SPEC_DIRS", strings.Join(sources, "||"))
	os.Setenv("GAUGE_PROJECT_ROOT", config.ProjectRoot)
	os.Setenv(common.APIPortEnvVariableName, strconv.Itoa(port))
	p, err := StartPlugin(pd, docScope)
	if err != nil {
		logger.Fatalf(true, "Error starting plugin %s %s. %s", pd.Name, pd.Version, err.Error())
	}
	for isProcessRunning(p) {
	}
}

func (p *plugin) sendMessage(message *gauge_messages.Message) error {
	messageID := common.GetUniqueID()
	message.MessageId = messageID
	messageBytes, err := proto.Marshal(message)
	if err != nil {
		return err
	}
	err = conn.Write(p.connection, messageBytes)
	if err != nil {
		return fmt.Errorf("[Warning] Failed to send message to plugin: %s  %s", p.descriptor.ID, err.Error())
	}
	return nil
}

func StartPlugins(manifest *manifest.Manifest) Handler {
	pluginHandler, warnings := startPluginsForExecution(manifest)
	logger.HandleWarningMessages(true, warnings)
	return pluginHandler
}

func PluginsWithoutScope() (infos []pluginInfo.PluginInfo) {
	if plugins, err := pluginInfo.GetAllInstalledPluginsWithVersion(); err == nil {
		for _, p := range plugins {
			pd, err := GetPluginDescriptor(p.Name, p.Version.String())
			if err == nil && !pd.hasAnyScope() {
				infos = append(infos, p)
			}
		}
	}
	return
}

// GetInstallDir returns the install directory of given plugin and a given version.
func GetInstallDir(pluginName, version string) (string, error) {
	allPluginsInstallDir, err := common.GetPluginsInstallDir(pluginName)
	if err != nil {
		return "", err
	}
	pluginDir := filepath.Join(allPluginsInstallDir, pluginName)
	if version != "" {
		pluginDir = filepath.Join(pluginDir, version)
	} else {
		latestPlugin, err := pluginInfo.GetLatestInstalledPlugin(pluginDir)
		if err != nil {
			return "", err
		}
		pluginDir = latestPlugin.Path
	}
	return pluginDir, nil
}

func GetLanguageJSONFilePath(language string) (string, error) {
	languageInstallDir, err := GetInstallDir(language, "")
	if err != nil {
		return "", err
	}
	languageJSON := filepath.Join(languageInstallDir, fmt.Sprintf("%s.json", language))
	if !common.FileExists(languageJSON) {
		return "", fmt.Errorf("Failed to find the implementation for: %s. %s does not exist.", language, languageJSON)
	}

	return languageJSON, nil
}

func IsLanguagePlugin(plugin string) bool {
	if _, err := GetLanguageJSONFilePath(plugin); err != nil {
		return false
	}
	return true
}

func QueryParams() string {
	return fmt.Sprintf("?l=%s&p=%s&o=%s&a=%s", language(), plugins(), runtime.GOOS, runtime.GOARCH)
}

func language() string {
	if config.ProjectRoot == "" {
		return ""
	}
	m, err := manifest.ProjectManifest()
	if err != nil {
		return ""
	}
	return m.Language
}

func plugins() string {
	pluginInfos, err := pluginInfo.GetAllInstalledPluginsWithVersion()
	if err != nil {
		return ""
	}
	var plugins []string
	for _, p := range pluginInfos {
		plugins = append(plugins, p.Name)
	}
	return strings.Join(plugins, ",")
}
<|endoftext|>


<commit_before>// Package plugin provides some types and functions common among plugin.
package plugin

import (
	"context"
	"errors"
	"fmt"

	"github.com/miekg/dns"
	ot "github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
)

type (
	// Plugin is a middle layer which represents the traditional
	// idea of plugin: it chains one Handler to the next by being
	// passed the next Handler in the chain.
	Plugin func(Handler) Handler

	// Handler is like dns.Handler except ServeDNS may return an rcode
	// and/or error.
	//
	// If ServeDNS writes to the response body, it should return a status
	// code. If the status code is not one of the following:
	//
	// * SERVFAIL (dns.RcodeServerFailure)
	//
	// * REFUSED (dns.RecodeRefused)
	//
	// * FORMERR (dns.RcodeFormatError)
	//
	// * NOTIMP (dns.RcodeNotImplemented)
	//
	// CoreDNS assumes *no* reply has yet been written. All other response
	// codes signal other handlers above it that the response message is
	// already written, and that they should not write to it also.
	//
	// If ServeDNS encounters an error, it should return the error value
	// so it can be logged by designated error-handling plugin.
	//
	// If writing a response after calling another ServeDNS method, the
	// returned rcode SHOULD be used when writing the response.
	//
	// If handling errors after calling another ServeDNS method, the
	// returned error value SHOULD be logged or handled accordingly.
	//
	// Otherwise, return values should be propagated down the plugin
	// chain by returning them unchanged.
	Handler interface {
		ServeDNS(context.Context, dns.ResponseWriter, *dns.Msg) (int, error)
		Name() string
	}

	// HandlerFunc is a convenience type like dns.HandlerFunc, except
	// ServeDNS returns an rcode and an error. See Handler
	// documentation for more information.
	HandlerFunc func(context.Context, dns.ResponseWriter, *dns.Msg) (int, error)
)

// ServeDNS implements the Handler interface.
func (f HandlerFunc) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	return f(ctx, w, r)
}

// Name implements the Handler interface.
func (f HandlerFunc) Name() string { return "handlerfunc" }

// Error returns err with 'plugin/name: ' prefixed to it.
func Error(name string, err error) error { return fmt.Errorf("%s/%s: %s", "plugin", name, err) }

// NextOrFailure calls next.ServeDNS when next is not nill, otherwise it will return, a ServerFailure and a nil error.
func NextOrFailure(name string, next Handler, ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) { // nolint: golint
	if next != nil {
		if span := ot.SpanFromContext(ctx); span != nil {
			child := span.Tracer().StartSpan(next.Name(), ot.ChildOf(span.Context()))
			defer child.Finish()
			ctx = ot.ContextWithSpan(ctx, child)
		}
		return next.ServeDNS(ctx, w, r)
	}

	return dns.RcodeServerFailure, Error(name, errors.New("no next plugin found"))
}

// ClientWrite returns true if the response has been written to the client.
// Each plugin to adhire to this protocol.
func ClientWrite(rcode int) bool {
	switch rcode {
	case dns.RcodeServerFailure:
		fallthrough
	case dns.RcodeRefused:
		fallthrough
	case dns.RcodeFormatError:
		fallthrough
	case dns.RcodeNotImplemented:
		return false
	}
	return true
}

// Namespace is the namespace used for the metrics.
const Namespace = "coredns"

// TimeBuckets is based on Prometheus client_golang prometheus.DefBuckets
var TimeBuckets = prometheus.ExponentialBuckets(0.00025, 2, 16) // from 0.25ms to 8 seconds

// ErrOnce is returned when a plugin doesn't support multiple setups per server.
var ErrOnce = errors.New("this plugin can only be used once per Server Block")

// ServerCtx is the context key to pass server address context to the plugins handling the request.
type ServerCtx struct{}
<commit_msg>fix typo (#2272)<commit_after>// Package plugin provides some types and functions common among plugin.
package plugin

import (
	"context"
	"errors"
	"fmt"

	"github.com/miekg/dns"
	ot "github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
)

type (
	// Plugin is a middle layer which represents the traditional
	// idea of plugin: it chains one Handler to the next by being
	// passed the next Handler in the chain.
	Plugin func(Handler) Handler

	// Handler is like dns.Handler except ServeDNS may return an rcode
	// and/or error.
	//
	// If ServeDNS writes to the response body, it should return a status
	// code. If the status code is not one of the following:
	//
	// * SERVFAIL (dns.RcodeServerFailure)
	//
	// * REFUSED (dns.RecodeRefused)
	//
	// * FORMERR (dns.RcodeFormatError)
	//
	// * NOTIMP (dns.RcodeNotImplemented)
	//
	// CoreDNS assumes *no* reply has yet been written. All other response
	// codes signal other handlers above it that the response message is
	// already written, and that they should not write to it also.
	//
	// If ServeDNS encounters an error, it should return the error value
	// so it can be logged by designated error-handling plugin.
	//
	// If writing a response after calling another ServeDNS method, the
	// returned rcode SHOULD be used when writing the response.
	//
	// If handling errors after calling another ServeDNS method, the
	// returned error value SHOULD be logged or handled accordingly.
	//
	// Otherwise, return values should be propagated down the plugin
	// chain by returning them unchanged.
	Handler interface {
		ServeDNS(context.Context, dns.ResponseWriter, *dns.Msg) (int, error)
		Name() string
	}

	// HandlerFunc is a convenience type like dns.HandlerFunc, except
	// ServeDNS returns an rcode and an error. See Handler
	// documentation for more information.
	HandlerFunc func(context.Context, dns.ResponseWriter, *dns.Msg) (int, error)
)

// ServeDNS implements the Handler interface.
func (f HandlerFunc) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	return f(ctx, w, r)
}

// Name implements the Handler interface.
func (f HandlerFunc) Name() string { return "handlerfunc" }

// Error returns err with 'plugin/name: ' prefixed to it.
func Error(name string, err error) error { return fmt.Errorf("%s/%s: %s", "plugin", name, err) }

// NextOrFailure calls next.ServeDNS when next is not nil, otherwise it will return, a ServerFailure and a nil error.
func NextOrFailure(name string, next Handler, ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) { // nolint: golint
	if next != nil {
		if span := ot.SpanFromContext(ctx); span != nil {
			child := span.Tracer().StartSpan(next.Name(), ot.ChildOf(span.Context()))
			defer child.Finish()
			ctx = ot.ContextWithSpan(ctx, child)
		}
		return next.ServeDNS(ctx, w, r)
	}

	return dns.RcodeServerFailure, Error(name, errors.New("no next plugin found"))
}

// ClientWrite returns true if the response has been written to the client.
// Each plugin to adhire to this protocol.
func ClientWrite(rcode int) bool {
	switch rcode {
	case dns.RcodeServerFailure:
		fallthrough
	case dns.RcodeRefused:
		fallthrough
	case dns.RcodeFormatError:
		fallthrough
	case dns.RcodeNotImplemented:
		return false
	}
	return true
}

// Namespace is the namespace used for the metrics.
const Namespace = "coredns"

// TimeBuckets is based on Prometheus client_golang prometheus.DefBuckets
var TimeBuckets = prometheus.ExponentialBuckets(0.00025, 2, 16) // from 0.25ms to 8 seconds

// ErrOnce is returned when a plugin doesn't support multiple setups per server.
var ErrOnce = errors.New("this plugin can only be used once per Server Block")

// ServerCtx is the context key to pass server address context to the plugins handling the request.
type ServerCtx struct{}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"strings"
	"testing"
)

func IsSeparator(s string) bool {
	r, _ := separator(s)
	return (r > 0)
}

func IsNotSeparator() bool {
	r, _ := separator("false")
	return (r == ',')
}

func TestGetSeparator(t *testing.T) {
	if !IsSeparator(",") {
		t.Error(`IsSeparator(",") = false`)
	}
	if !IsNotSeparator() {
		t.Error(`IsNotSeparator() = false`)
	}
}

func TestCsvInputNew(t *testing.T) {
	trdsql := trdsqlNew()
	file, err := tableFileOpen("testdata/test.csv")
	if err != nil {
		t.Error(err)
	}
	_, err = trdsql.csvInputNew(file)
	if err != nil {
		t.Error(`csvInputNew error`)
	}
}

func TestCsvEmptyNew(t *testing.T) {
	trdsql := trdsqlNew()
	const csvStream = ``
	s := strings.NewReader(csvStream)
	r, err := trdsql.csvInputNew(s)
	if err != nil {
		t.Error(err)
	}
	_, err = r.FirstRead()
	if err.Error() != "EOF" {
		t.Error(err)
	}
}

func TestCsvHeaderNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,h2
	v1,v2`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	header, _ := r.FirstRead()
	if header[0] != "h1" || header[1] != "h2" {
		t.Error("invalid header")
	}
}

func TestCsvEmptyColumnHeaderNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,
	v1,v2`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	header, _ := r.FirstRead()
	if header[0] != "h1" || header[1] != "c2" {
		t.Error("invalid header")
	}
}

func TestCsvEmptyColumnRowNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,h2
	,v2`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	_, err := r.FirstRead()
	if err != nil {
		t.Error(err)
	}
	record := make([]interface{}, 2)
	record, _ = r.RowRead(record)
	if record[0] != "" || record[1] != "v2" {
		t.Error("invalid value")
	}
}

func TestCsvColumnDifferenceNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,h2,h3
	v1,v2,v3
	x1,x2
	z1`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	_, err := r.FirstRead()
	if err != nil {
		t.Error(err)
	}
	record := make([]interface{}, 3)
	for {
		record, err = r.RowRead(record)
		if err == io.EOF {
			break
		} else if err != nil {
			t.Error(err)
		}
		if len(record) != 3 {
			t.Error("row difference")
		}
	}
}

func TestCsvNoInputNew(t *testing.T) {
	trdsql := trdsqlNew()
	file, err := tableFileOpen("nofile")
	if err == nil {
		t.Error(`Should error`)
	}
	_, err = trdsql.csvInputNew(file)
	if err != nil {
		t.Error(`csvInputNew error`)
	}
}

func TestCsvOutNew(t *testing.T) {
	trdsql := trdsqlNew()
	out := trdsql.csvOutNew()
	if out == nil {
		t.Error(`csvOut error`)
	}
}
<commit_msg>fix set of delimiters.<commit_after>package main

import (
	"io"
	"strings"
	"testing"
)

func IsSeparator(s string) bool {
	r, _ := separator(s)
	return (r > 0)
}

func IsNotSeparator() bool {
	r, _ := separator("false")
	return (r == ',')
}

func TestGetSeparator(t *testing.T) {
	if !IsSeparator(",") {
		t.Error(`IsSeparator(",") = false`)
	}
	if !IsNotSeparator() {
		t.Error(`IsNotSeparator() = false`)
	}
}

func TestCsvInputNew(t *testing.T) {
	trdsql := trdsqlNew()
	file, err := tableFileOpen("testdata/test.csv")
	if err != nil {
		t.Error(err)
	}
	_, err = trdsql.csvInputNew(file)
	if err != nil {
		t.Error(`csvInputNew error`)
	}
}

func TestCsvEmptyNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inSep = ","
	const csvStream = ``
	s := strings.NewReader(csvStream)
	r, err := trdsql.csvInputNew(s)
	if err != nil {
		t.Error(err)
	}
	_, err = r.FirstRead()
	if err.Error() != "EOF" {
		t.Error(err)
	}
}

func TestCsvHeaderNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,h2
	v1,v2`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	header, _ := r.FirstRead()
	if header[0] != "h1" || header[1] != "h2" {
		t.Error("invalid header")
	}
}

func TestCsvEmptyColumnHeaderNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,
	v1,v2`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	header, _ := r.FirstRead()
	if header[0] != "h1" || header[1] != "c2" {
		t.Error("invalid header")
	}
}

func TestCsvEmptyColumnRowNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,h2
	,v2`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	_, err := r.FirstRead()
	if err != nil {
		t.Error(err)
	}
	record := make([]interface{}, 2)
	record, _ = r.RowRead(record)
	if record[0] != "" || record[1] != "v2" {
		t.Error("invalid value")
	}
}

func TestCsvColumnDifferenceNew(t *testing.T) {
	trdsql := trdsqlNew()
	trdsql.inHeader = true
	trdsql.inSep = ","
	csvStream := `h1,h2,h3
	v1,v2,v3
	x1,x2
	z1`
	s := strings.NewReader(csvStream)
	r, _ := trdsql.csvInputNew(s)
	_, err := r.FirstRead()
	if err != nil {
		t.Error(err)
	}
	record := make([]interface{}, 3)
	for {
		record, err = r.RowRead(record)
		if err == io.EOF {
			break
		} else if err != nil {
			t.Error(err)
		}
		if len(record) != 3 {
			t.Error("row difference")
		}
	}
}

func TestCsvNoInputNew(t *testing.T) {
	trdsql := trdsqlNew()
	file, err := tableFileOpen("nofile")
	if err == nil {
		t.Error(`Should error`)
	}
	_, err = trdsql.csvInputNew(file)
	if err != nil {
		t.Error(`csvInputNew error`)
	}
}

func TestCsvOutNew(t *testing.T) {
	trdsql := trdsqlNew()
	out := trdsql.csvOutNew()
	if out == nil {
		t.Error(`csvOut error`)
	}
}
<|endoftext|>


<commit_before>package goalone

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"hash"
	"sync"
	"time"
)

// Options that can be configured and passed to New()
type Options struct {
	// hash algorithm to use when signing tokens, ex. sha1.New
	Algorithm func() hash.Hash

	// Epoch to use for Timestamps, when signing/parsing Tokens
	Epoch int64

	// Should the Sign method add a timestamp to tokens?
	Timestamp bool
}

// Sword is a Wooden Sword to be used for protection, because it's dangerous out
// there... Also, it is the main struct used to sign and unsign data using this
// package.
type Sword struct {
	sync.Mutex
	hash  hash.Hash
	dirty bool

	Options
}

// ErrInvalidSignature is returned by Unsign when the provided token's
// signatuire is not valid.
var ErrInvalidSignature = errors.New("invalid signature")

// ErrShortToken is returned by Unsign when the provided token's length
// is too short to be a vlaid token.
var ErrShortToken = errors.New("token is too small to be valid")

// New takes a secret key and returns a new Sword.  If no Options are provided
// then minimal defaults will be used.
func New(key []byte, o *Options) *Sword {

	if key == nil {
		return &Sword{}
	}

	if o == nil {
		return &Sword{hash: hmac.New(sha1.New, key)}
	}

	s := &Sword{Options: *o}
	if s.Algorithm == nil {
		s.hash = hmac.New(sha1.New, key)
	} else {
		s.hash = hmac.New(s.Algorithm, key)
	}

	return s
}

// Sign signs data and returns []byte in the format `data.signature`.
func (s *Sword) Sign(data []byte) []byte {

	// Build the payload
	el := base64.RawURLEncoding.EncodedLen(s.hash.Size())
	var t []byte

	if s.Timestamp {
		now := time.Now().UTC().Unix() - s.Epoch
		ts := encodeUint64(uint64(now))
		t = make([]byte, 0, len(data)+len(ts)+el+2)
		t = append(t, data...)
		t = append(t, '.')
		t = append(t, ts...)
	} else {
		t = make([]byte, 0, len(data)+el+1)
		t = append(t, data...)
	}

	// Now lets lock the hash and create the signature
	s.Lock()
	if s.dirty {
		s.hash.Reset()
	}
	s.dirty = true
	s.hash.Write(t)
	h := s.hash.Sum(nil)
	s.Unlock()

	// Append signature to token
	t = append(t, '.')
	tl := len(t)
	t = t[0 : tl+el]
	base64.RawURLEncoding.Encode(t[tl:], h)

	// Return the token to the caller
	return t
}

// Unsign validates a signature and if successful returns the data portion of
// the []byte. If unsuccessful it will return an error and nil for the data.
func (s *Sword) Unsign(token []byte) ([]byte, error) {

	tl := len(token)
	el := base64.RawURLEncoding.EncodedLen(s.hash.Size())

	// A token must be at least el+2 bytes long to be valid.
	if tl < el+2 {
		return nil, ErrShortToken
	}

	// Now lets lock the hash and create the signature
	s.Lock()
	if s.dirty {
		s.hash.Reset()
	}
	s.dirty = true
	s.hash.Write(token[0 : tl-(el+1)])
	h := s.hash.Sum(nil)
	s.Unlock()

	// Encode hash into dst
	dst := make([]byte, el)
	base64.RawURLEncoding.Encode(dst, h)

	if subtle.ConstantTimeCompare(token[tl-el:], dst) != 1 {
		return nil, ErrInvalidSignature
	}

	return token[0 : tl-(el+1)], nil
}

func encodeUint64(i uint64) []byte {

	// covert int into bytes
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, i)

	// encode bytes to base64, striping off leading 0's
	buf2 := make([]byte, 11)
	base64.RawURLEncoding.Encode(buf2, bytes.TrimLeft(buf, "\x00"))

	// return result, removing training `\x00`
	return bytes.TrimRight(buf2, "\x00")
}

func decodeUint64(b []byte) uint64 {

	l := len(b) * 6 / 8 // cause DecodedLen has a bug // https://github.com/golang/go/commit/87151c82b68023e4224b016a6a66ead2c4b8ece7
	buf := make([]byte, 8)

	base64.RawURLEncoding.Decode(buf[8-l:], b)
	return binary.BigEndian.Uint64(buf)
}
<commit_msg>Few comment changes<commit_after>package goalone

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"hash"
	"sync"
	"time"
)

// Options that can be configured and passed to New()
type Options struct {
	// hash algorithm to use when signing tokens, ex. sha1.New
	Algorithm func() hash.Hash

	// Epoch to use for Timestamps, when signing/parsing Tokens
	Epoch int64

	// Should the Sign method add a timestamp to tokens?
	Timestamp bool
}

// Sword is a magical Wooden Sword to be used for protection, because it's dangerous out
// there... Also, it is the main struct used to sign and unsign data using this
// package.
type Sword struct {
	sync.Mutex
	hash  hash.Hash
	dirty bool

	Options
}

// ErrInvalidSignature is returned by Unsign when the provided token's
// signatuire is not valid.
var ErrInvalidSignature = errors.New("invalid signature")

// ErrShortToken is returned by Unsign when the provided token's length
// is too short to be a vlaid token.
var ErrShortToken = errors.New("token is too small to be valid")

// New takes a secret key and returns a new Sword.  If no Options are provided
// then minimal defaults will be used.
func New(key []byte, o *Options) *Sword {

	if key == nil {
		return &Sword{}
	}

	if o == nil {
		return &Sword{hash: hmac.New(sha1.New, key)}
	}

	s := &Sword{Options: *o}
	if s.Algorithm == nil {
		s.hash = hmac.New(sha1.New, key)
	} else {
		s.hash = hmac.New(s.Algorithm, key)
	}

	return s
}

// Sign signs data and returns []byte in the format `data.signature`. Optionally
// add a timestamp and return in the format `data.timestamp.signature`
func (s *Sword) Sign(data []byte) []byte {

	// Build the payload
	el := base64.RawURLEncoding.EncodedLen(s.hash.Size())
	var t []byte

	if s.Timestamp {
		now := time.Now().UTC().Unix() - s.Epoch
		ts := encodeUint64(uint64(now))
		t = make([]byte, 0, len(data)+len(ts)+el+2)
		t = append(t, data...)
		t = append(t, '.')
		t = append(t, ts...)
	} else {
		t = make([]byte, 0, len(data)+el+1)
		t = append(t, data...)
	}

	// Now lets lock the hash and create the signature
	s.Lock()
	if s.dirty {
		s.hash.Reset()
	}
	s.dirty = true
	s.hash.Write(t)
	h := s.hash.Sum(nil)
	s.Unlock()

	// Append signature to token
	t = append(t, '.')
	tl := len(t)
	t = t[0 : tl+el]
	base64.RawURLEncoding.Encode(t[tl:], h)

	// Return the token to the caller
	return t
}

// Unsign validates a signature and if successful returns the data portion of
// the []byte. If unsuccessful it will return an error and nil for the data.
func (s *Sword) Unsign(token []byte) ([]byte, error) {

	tl := len(token)
	el := base64.RawURLEncoding.EncodedLen(s.hash.Size())

	// A token must be at least el+2 bytes long to be valid.
	if tl < el+2 {
		return nil, ErrShortToken
	}

	// Now lets lock the hash and create the signature
	s.Lock()
	if s.dirty {
		s.hash.Reset()
	}
	s.dirty = true
	s.hash.Write(token[0 : tl-(el+1)])
	h := s.hash.Sum(nil)
	s.Unlock()

	// Encode hash into dst
	dst := make([]byte, el)
	base64.RawURLEncoding.Encode(dst, h)

	if subtle.ConstantTimeCompare(token[tl-el:], dst) != 1 {
		return nil, ErrInvalidSignature
	}

	return token[0 : tl-(el+1)], nil
}

func encodeUint64(i uint64) []byte {

	// covert int into bytes
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, i)

	// encode bytes to base64, striping off leading 0's
	buf2 := make([]byte, 11)
	base64.RawURLEncoding.Encode(buf2, bytes.TrimLeft(buf, "\x00"))

	// return result, removing training `\x00`
	return bytes.TrimRight(buf2, "\x00")
}

func decodeUint64(b []byte) uint64 {

	l := len(b) * 6 / 8 // cause DecodedLen has a bug // https://github.com/golang/go/commit/87151c82b68023e4224b016a6a66ead2c4b8ece7
	buf := make([]byte, 8)

	base64.RawURLEncoding.Decode(buf[8-l:], b)
	return binary.BigEndian.Uint64(buf)
}
<|endoftext|>


<commit_before>package gocache

import(
    "fmt"
    "net/http"
    "errors"
    "os"
    "strings"
    "log"
    "encoding/json"
    "bytes"
    "io/ioutil"
    "path"
    "io"
)

type Entry struct{
    Response *http.Response
    Content string
}

func main(){
    fmt.Println("Testing web cache");
    url := "http://blog.golang.org/error-handling-and-go"
    Get(url)
}

func Get(url string) (resp *http.Response, err error) {

    resp, err = getFromCache(url)

    if err != nil{
        //log.Println(url,"could not be served out of cache.", err)
        return getFromWeb(url)
    }

    //log.Print("Request to ", url, " served out of cache")
    return
}

// Fetch contents of url from web and write to cache
func getFromWeb(url string) (resp *http.Response, err error){
    resp, err = http.Get(url)

    if err != nil{
        return
    }
    var resp2 http.Response
    resp2 = *resp

    defer resp.Body.Close()

    // Draining the body

    body, err := ioutil.ReadAll(resp.Body)

    if err != nil {
        log.Print("Reading response body went bad. ", err);

    }

    writeToCache(url, body, &resp2)

    resp.Body = nopCloser{bytes.NewBufferString(string(body))}

    return
}

func writeToCache(url string, body []byte, resp *http.Response){

    cacheEntry := generateCacheEntry(resp, body)

    b, err := json.Marshal(cacheEntry)

    if err != nil {
        log.Println("Could not marshal response ", err)
        return
    }

    filename := getFilePath(url)

    // Set up any directory needed to write the file.
    // e.g. vg.no/first/second/file will be stored as
    // cache/first/second/file

    err = createDirectories(filename)
    if err != nil{
        pe, _ := err.(*os.PathError)

        if ! strings.Contains(pe.Error(),"file exists") {
            log.Println("Could not create directories ", err)
            return
        }
    }

    // If the file hasn't got an extension, set it to .json
    name := strings.Split(filename,"/")
    fn := name[len(name)-1]

    if len(strings.Split(fn, ".")) < 2{
        log.Println("No extension", filename)
        filename = filename + ".json"
    }

    // Create file
    file, err := os.Create(filename)
    if err != nil {
        log.Println("could not create cache file ", err)
        return
    }

    // Close and check for error on exit
    defer func() {
        if err := file.Close(); err != nil {
            log.Println("Could not close file ", err)
        }
    }()

    // Write file
    _, err = file.Write(b)
    if err != nil{
        log.Println("Could not write to cache file ", err)
    }

    //log.Println(url, "successfully written to cache");
}

func generateCacheEntry(resp *http.Response, body []byte ) Entry {
    Response := resp
    Content := string(body)
    entry := Entry{Response, Content}

    // Cannot marshal the body from get go
    entry.Response.Body = nil

    return entry

}

// Try to fetch contents of url from cache
func getFromCache(url string) (resp *http.Response, err error) {

    filename := getFilePath(url)

    file, err := os.Open(filename)
    if err != nil{
        err = errors.New("File '"+filename+"' not Found")
        return
    }

    entry, err := readFromFile(file)

    if err != nil{
        return
    }

    resp = entry.GenerateHttpResponse()

    return

}

// Read cache entry from file
func readFromFile(file *os.File) (entry *Entry, err error) {
    fileinfo, err := file.Stat()
    var size int
    size = int(fileinfo.Size())

    buf := make([]byte, size)

    _, err = file.Read(buf)

    if err != nil{
        //log.Println("Reading file got:", err, "after",n,"bytes");
        return entry, errors.New("Reading file went wrong")
    }

    // Must trim buffer before unmarshaling it. This is because of
    // the unmarshaling failing if entire buffer is returned
    buf = bytes.Trim(buf[0:], string(0))

    entry = new(Entry)
    err = json.Unmarshal(buf, entry)

    if err != nil {
        log.Print("Unmarshaling gone wrong: ", err)
        return entry, errors.New("Unmarshaling gone wrong")
    }

    return
}

func (entry *Entry) Print () {

    //n := bytes.Index(entry.Content, []byte{0})
    //content := string(entry.Content[:n])
    log.Print("Content: ", entry.Content)

}

func (entry *Entry) GenerateHttpResponse() (resp *http.Response){

    resp = entry.Response
    resp.Body = nopCloser{bytes.NewBufferString(entry.Content)}

    return resp

}

func getFilePath(url string) (dir string) {
	urlTokens := strings.Split(url, "/")
    // 2 because we need to strip away 'http:', ' '
    strippedUrl := urlTokens[2:]
	dir = "cache/"+strings.Join(strippedUrl,"/")
	return
}

func createDirectories(filename string) error {

    filepath := path.Dir(filename)
    directories := strings.Split(filepath, "/")

    p := ""
    for i := range directories {
        p += directories[i] + "/"
        err := os.Mkdir(p, 0755)

        if err != nil {
            pe, _ := err.(*os.PathError)

            // if folder exists, continue to the next one
            if ! strings.Contains(pe.Error(),"file exists") {
                log.Println("Mkdir failed miserably: ", err)
                return err
            }
        }
    }

    return nil
}

/* Below are from:
    https://groups.google.com/forum/#!topic/golang-nuts/J-Y4LtdGNSw
*/

type nopCloser struct {
    io.Reader
}

func (nopCloser) Close() error {
    return nil
}
<commit_msg>don't need this<commit_after>package gocache

import(
    "fmt"
    "net/http"
    "errors"
    "os"
    "strings"
    "log"
    "encoding/json"
    "bytes"
    "io/ioutil"
    "path"
    "io"
)

type Entry struct{
    Response *http.Response
    Content string
}

func main(){
    fmt.Println("Testing web cache");
    url := "http://blog.golang.org/error-handling-and-go"
    Get(url)
}

func Get(url string) (resp *http.Response, err error) {

    resp, err = getFromCache(url)

    if err != nil{
        //log.Println(url,"could not be served out of cache.", err)
        return getFromWeb(url)
    }

    //log.Print("Request to ", url, " served out of cache")
    return
}

// Fetch contents of url from web and write to cache
func getFromWeb(url string) (resp *http.Response, err error){
    resp, err = http.Get(url)

    if err != nil{
        return
    }
    var resp2 http.Response
    resp2 = *resp

    defer resp.Body.Close()

    // Draining the body

    body, err := ioutil.ReadAll(resp.Body)

    if err != nil {
        log.Print("Reading response body went bad. ", err);

    }

    writeToCache(url, body, &resp2)

    resp.Body = nopCloser{bytes.NewBufferString(string(body))}

    return
}

func writeToCache(url string, body []byte, resp *http.Response){

    cacheEntry := generateCacheEntry(resp, body)

    b, err := json.Marshal(cacheEntry)

    if err != nil {
        log.Println("Could not marshal response ", err)
        return
    }

    filename := getFilePath(url)

    // Set up any directory needed to write the file.
    // e.g. vg.no/first/second/file will be stored as
    // cache/first/second/file

    err = createDirectories(filename)
    if err != nil{
        pe, _ := err.(*os.PathError)

        if ! strings.Contains(pe.Error(),"file exists") {
            log.Println("Could not create directories ", err)
            return
        }
    }

    // If the file hasn't got an extension, set it to .json
    name := strings.Split(filename,"/")
    fn := name[len(name)-1]

    if len(strings.Split(fn, ".")) < 2{
        filename = filename + ".json"
    }

    // Create file
    file, err := os.Create(filename)
    if err != nil {
        log.Println("could not create cache file ", err)
        return
    }

    // Close and check for error on exit
    defer func() {
        if err := file.Close(); err != nil {
            log.Println("Could not close file ", err)
        }
    }()

    // Write file
    _, err = file.Write(b)
    if err != nil{
        log.Println("Could not write to cache file ", err)
    }

    //log.Println(url, "successfully written to cache");
}

func generateCacheEntry(resp *http.Response, body []byte ) Entry {
    Response := resp
    Content := string(body)
    entry := Entry{Response, Content}

    // Cannot marshal the body from get go
    entry.Response.Body = nil

    return entry

}

// Try to fetch contents of url from cache
func getFromCache(url string) (resp *http.Response, err error) {

    filename := getFilePath(url)

    file, err := os.Open(filename)
    if err != nil{
        err = errors.New("File '"+filename+"' not Found")
        return
    }

    entry, err := readFromFile(file)

    if err != nil{
        return
    }

    resp = entry.GenerateHttpResponse()

    return

}

// Read cache entry from file
func readFromFile(file *os.File) (entry *Entry, err error) {
    fileinfo, err := file.Stat()
    var size int
    size = int(fileinfo.Size())

    buf := make([]byte, size)

    _, err = file.Read(buf)

    if err != nil{
        //log.Println("Reading file got:", err, "after",n,"bytes");
        return entry, errors.New("Reading file went wrong")
    }

    // Must trim buffer before unmarshaling it. This is because of
    // the unmarshaling failing if entire buffer is returned
    buf = bytes.Trim(buf[0:], string(0))

    entry = new(Entry)
    err = json.Unmarshal(buf, entry)

    if err != nil {
        log.Print("Unmarshaling gone wrong: ", err)
        return entry, errors.New("Unmarshaling gone wrong")
    }

    return
}

func (entry *Entry) Print () {

    //n := bytes.Index(entry.Content, []byte{0})
    //content := string(entry.Content[:n])
    log.Print("Content: ", entry.Content)

}

func (entry *Entry) GenerateHttpResponse() (resp *http.Response){

    resp = entry.Response
    resp.Body = nopCloser{bytes.NewBufferString(entry.Content)}

    return resp

}

func getFilePath(url string) (dir string) {
	urlTokens := strings.Split(url, "/")
    // 2 because we need to strip away 'http:', ' '
    strippedUrl := urlTokens[2:]
	dir = "cache/"+strings.Join(strippedUrl,"/")
	return
}

func createDirectories(filename string) error {

    filepath := path.Dir(filename)
    directories := strings.Split(filepath, "/")

    p := ""
    for i := range directories {
        p += directories[i] + "/"
        err := os.Mkdir(p, 0755)

        if err != nil {
            pe, _ := err.(*os.PathError)

            // if folder exists, continue to the next one
            if ! strings.Contains(pe.Error(),"file exists") {
                log.Println("Mkdir failed miserably: ", err)
                return err
            }
        }
    }

    return nil
}

/* Below are from:
    https://groups.google.com/forum/#!topic/golang-nuts/J-Y4LtdGNSw
*/

type nopCloser struct {
    io.Reader
}

func (nopCloser) Close() error {
    return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
)

var (
	ScanOptions = []string{"openscap"}
)

// MultiStringVar is implementing flag.Value
type MultiStringVar struct {
	Values []string
}

func (sv *MultiStringVar) Set(s string) error {
	sv.Values = append(sv.Values, s)
	return nil
}

func (sv *MultiStringVar) String() string {
	return fmt.Sprintf("%v", sv.Values)
}

// ImageInspectorOptions is the main inspector implementation and holds the configuration
// for an image inspector.
type ImageInspectorOptions struct {
	// URI contains the location of the docker daemon socket to connect to.
	URI string
	// Image contains the docker image to inspect.
	Image string
	// DstPath is the destination path for image files.
	DstPath string
	// Serve holds the host and port for where to serve the image with webdav.
	Serve string
	// Chroot controls whether or not a chroot is excuted when serving the image with webdav.
	Chroot bool
	// DockerCfg is the location of the docker config file.
	DockerCfg MultiStringVar
	// Username is the username for authenticating to the docker registry.
	Username string
	// PasswordFile is the location of the file containing the password for authentication to the
	// docker registry.
	PasswordFile string
	// ScanType is the type of the scan to be done on the inspected image
	ScanType string
	// ScanResultsDir is the directory that will contain the results of the scan
	ScanResultsDir string
}

// NewDefaultImageInspectorOptions provides a new ImageInspectorOptions with default values.
func NewDefaultImageInspectorOptions() *ImageInspectorOptions {
	return &ImageInspectorOptions{
		URI:            "unix:///var/run/docker.sock",
		Image:          "",
		DstPath:        "",
		Serve:          "",
		Chroot:         false,
		DockerCfg:      MultiStringVar{[]string{}},
		Username:       "",
		PasswordFile:   "",
		ScanType:       "",
		ScanResultsDir: "",
	}
}

// Validate performs validation on the field settings.
func (i *ImageInspectorOptions) Validate() error {
	if len(i.URI) == 0 {
		return fmt.Errorf("Docker socket connection must be specified")
	}
	if len(i.Image) == 0 {
		return fmt.Errorf("Docker image to inspect must be specified")
	}
	if len(i.DockerCfg.Values) > 0 && len(i.Username) > 0 {
		return fmt.Errorf("Only specify dockercfg file or username/password pair for authentication")
	}
	if len(i.Username) > 0 && len(i.PasswordFile) == 0 {
		return fmt.Errorf("Please specify password for the username")
	}
	if len(i.Serve) == 0 && i.Chroot {
		return fmt.Errorf("Change root can be used only when serving the image through webdav")
	}
	if len(i.ScanResultsDir) > 0 && len(i.ScanType) == 0 {
		return fmt.Errorf("scan-result-dir can be used only when spacifing scan-type")
	}
	if len(i.ScanResultsDir) > 0 {
		fi, err := os.Stat(i.ScanResultsDir)
		if err == nil && !fi.IsDir() {
			return fmt.Errorf("%s is not a directory", i.ScanResultsDir)
		}
	}
	for _, fl := range append(i.DockerCfg.Values, i.PasswordFile) {
		if len(fl) > 0 {
			if _, err := os.Stat(fl); os.IsNotExist(err) {
				return fmt.Errorf("%s does not exists", fl)
			}
		}
	}
	if len(i.ScanType) > 0 {
		var found bool = false
		for _, opt := range ScanOptions {
			if i.ScanType == opt {
				found = true
				break
			}
		}
		if !found {
			return fmt.Errorf("%s is not one of the available scan-type's which are %v", i.ScanType, ScanOptions)
		}
	}
	return nil
}
<commit_msg>fix scan type error message<commit_after>package cmd

import (
	"fmt"
	"os"
)

var (
	ScanOptions = []string{"openscap"}
)

// MultiStringVar is implementing flag.Value
type MultiStringVar struct {
	Values []string
}

func (sv *MultiStringVar) Set(s string) error {
	sv.Values = append(sv.Values, s)
	return nil
}

func (sv *MultiStringVar) String() string {
	return fmt.Sprintf("%v", sv.Values)
}

// ImageInspectorOptions is the main inspector implementation and holds the configuration
// for an image inspector.
type ImageInspectorOptions struct {
	// URI contains the location of the docker daemon socket to connect to.
	URI string
	// Image contains the docker image to inspect.
	Image string
	// DstPath is the destination path for image files.
	DstPath string
	// Serve holds the host and port for where to serve the image with webdav.
	Serve string
	// Chroot controls whether or not a chroot is excuted when serving the image with webdav.
	Chroot bool
	// DockerCfg is the location of the docker config file.
	DockerCfg MultiStringVar
	// Username is the username for authenticating to the docker registry.
	Username string
	// PasswordFile is the location of the file containing the password for authentication to the
	// docker registry.
	PasswordFile string
	// ScanType is the type of the scan to be done on the inspected image
	ScanType string
	// ScanResultsDir is the directory that will contain the results of the scan
	ScanResultsDir string
}

// NewDefaultImageInspectorOptions provides a new ImageInspectorOptions with default values.
func NewDefaultImageInspectorOptions() *ImageInspectorOptions {
	return &ImageInspectorOptions{
		URI:            "unix:///var/run/docker.sock",
		Image:          "",
		DstPath:        "",
		Serve:          "",
		Chroot:         false,
		DockerCfg:      MultiStringVar{[]string{}},
		Username:       "",
		PasswordFile:   "",
		ScanType:       "",
		ScanResultsDir: "",
	}
}

// Validate performs validation on the field settings.
func (i *ImageInspectorOptions) Validate() error {
	if len(i.URI) == 0 {
		return fmt.Errorf("Docker socket connection must be specified")
	}
	if len(i.Image) == 0 {
		return fmt.Errorf("Docker image to inspect must be specified")
	}
	if len(i.DockerCfg.Values) > 0 && len(i.Username) > 0 {
		return fmt.Errorf("Only specify dockercfg file or username/password pair for authentication")
	}
	if len(i.Username) > 0 && len(i.PasswordFile) == 0 {
		return fmt.Errorf("Please specify password for the username")
	}
	if len(i.Serve) == 0 && i.Chroot {
		return fmt.Errorf("Change root can be used only when serving the image through webdav")
	}
	if len(i.ScanResultsDir) > 0 && len(i.ScanType) == 0 {
		return fmt.Errorf("scan-result-dir can be used only when spacifing scan-type")
	}
	if len(i.ScanResultsDir) > 0 {
		fi, err := os.Stat(i.ScanResultsDir)
		if err == nil && !fi.IsDir() {
			return fmt.Errorf("%s is not a directory", i.ScanResultsDir)
		}
	}
	for _, fl := range append(i.DockerCfg.Values, i.PasswordFile) {
		if len(fl) > 0 {
			if _, err := os.Stat(fl); os.IsNotExist(err) {
				return fmt.Errorf("%s does not exists", fl)
			}
		}
	}
	if len(i.ScanType) > 0 {
		var found bool = false
		for _, opt := range ScanOptions {
			if i.ScanType == opt {
				found = true
				break
			}
		}
		if !found {
			return fmt.Errorf("%s is not one of the available scan-types which are %v", i.ScanType, ScanOptions)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors All rights reserved.

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 hook

import (
	"context"
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"k8s.io/git-sync/pkg/logging"
)

var (
	hookRunCount = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "git_sync_hook_run_count_total",
		Help: "How many hook runs completed, partitioned by name and state (success, error)",
	}, []string{"name", "status"})
)

// Describes what a Hook needs to implement, run by HookRunner
type Hook interface {
	// Describes hook
	Name() string
	// Function that called by HookRunner
	Do(ctx context.Context, hash string) error
}

type hookData struct {
	ch    chan struct{}
	mutex sync.Mutex
	hash  string
}

// NewHookData returns a new HookData
func NewHookData() *hookData {
	return &hookData{
		ch: make(chan struct{}, 1),
	}
}

func (d *hookData) events() chan struct{} {
	return d.ch
}

func (d *hookData) get() string {
	d.mutex.Lock()
	defer d.mutex.Unlock()
	return d.hash
}

func (d *hookData) set(newHash string) {
	d.mutex.Lock()
	defer d.mutex.Unlock()
	d.hash = newHash
}

func (d *hookData) send(newHash string) {
	d.set(newHash)

	// Non-blocking write.  If the channel is full, the consumer will see the
	// newest value.  If the channel was not full, the consumer will get another
	// event.
	select {
	case d.ch <- struct{}{}:
	default:
	}
}

// NewHookRunner returns a new HookRunner
func NewHookRunner(hook Hook, backoff time.Duration, data *hookData, log *logging.Logger, hasSucceededOnce chan bool) *HookRunner {
	return &HookRunner{hook: hook, backoff: backoff, data: data, logger: log, hasCompletedOnce: hasSucceededOnce}
}

// HookRunner struct
type HookRunner struct {
	// Hook to run and check
	hook Hook
	// Backoff for failed hooks
	backoff time.Duration
	// Holds the data as it crosses from producer to consumer.
	data *hookData
	// Logger
	logger *logging.Logger
	// hasCompletedOnce is used to send true if and only if first run executed
	// successfully and false otherwise to some receiver.  Should be
	// initialised to a buffered channel of size 1.
	// Is only meant for use within sendCompletedOnceMessageIfApplicable.
	hasCompletedOnce chan bool
}

// Send sends hash to hookdata
func (r *HookRunner) Send(hash string) {
	r.data.send(hash)
}

// Run waits for trigger events from the channel, and run hook when triggered
func (r *HookRunner) Run(ctx context.Context) {
	var lastHash string
	prometheus.MustRegister(hookRunCount)

	// Wait for trigger from hookData.Send
	for range r.data.events() {
		// Retry in case of error
		for {
			// Always get the latest value, in case we fail-and-retry and the
			// value changed in the meantime.  This means that we might not send
			// every single hash.
			hash := r.data.get()
			if hash == lastHash {
				break
			}

			if err := r.hook.Do(ctx, hash); err != nil {
				r.logger.Error(err, "hook failed")
				updateHookRunCountMetric(r.hook.Name(), "error")
				// don't want to sleep unnecessarily terminating anyways
				r.sendCompletedOnceMessageIfApplicable(false)
				time.Sleep(r.backoff)
			} else {
				updateHookRunCountMetric(r.hook.Name(), "success")
				lastHash = hash
				r.sendCompletedOnceMessageIfApplicable(true)
				break
			}
		}
	}
}

// If hasCompletedOnce is nil, does nothing. Otherwise, forwards the caller
// provided success status (as a boolean) of HookRunner to receivers of
// hasCompletedOnce, closes said chanel, and terminates this goroutine.
// Using this function to write to hasCompletedOnce ensures it is only ever
// written to once.
func (r *HookRunner) sendCompletedOnceMessageIfApplicable(completedSuccessfully bool) {
	if r.hasCompletedOnce != nil {
		r.hasCompletedOnce <- completedSuccessfully
		close(r.hasCompletedOnce)
		runtime.Goexit()
	}
}

// WaitForCompletion waits for HookRunner to send completion message to
// calling thread and returns either true if HookRunner executed successfully
// and some error otherwise.
// Assumes that r.hasCompletedOnce is not nil, but if it is, returns an error
func (r *HookRunner) WaitForCompletion() error {
	// Make sure function should be called
	if r.hasCompletedOnce == nil {
		return fmt.Errorf("HookRunner.WaitForCompletion called on async runner")
	}

	// Perform wait on HookRunner
	hookRunnerFinishedSuccessfully := <-r.hasCompletedOnce
	if !hookRunnerFinishedSuccessfully {
		return fmt.Errorf("hook completed with error")
	}

	return nil
}

func updateHookRunCountMetric(name, status string) {
	hookRunCount.WithLabelValues(name, status).Inc()
}
<commit_msg>Don't double-register the hook metric<commit_after>/*
Copyright 2021 The Kubernetes Authors All rights reserved.

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 hook

import (
	"context"
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"k8s.io/git-sync/pkg/logging"
)

var (
	hookRunCount = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "git_sync_hook_run_count_total",
		Help: "How many hook runs completed, partitioned by name and state (success, error)",
	}, []string{"name", "status"})
)

func init() {
	prometheus.MustRegister(hookRunCount)
}

// Describes what a Hook needs to implement, run by HookRunner
type Hook interface {
	// Describes hook
	Name() string
	// Function that called by HookRunner
	Do(ctx context.Context, hash string) error
}

type hookData struct {
	ch    chan struct{}
	mutex sync.Mutex
	hash  string
}

// NewHookData returns a new HookData
func NewHookData() *hookData {
	return &hookData{
		ch: make(chan struct{}, 1),
	}
}

func (d *hookData) events() chan struct{} {
	return d.ch
}

func (d *hookData) get() string {
	d.mutex.Lock()
	defer d.mutex.Unlock()
	return d.hash
}

func (d *hookData) set(newHash string) {
	d.mutex.Lock()
	defer d.mutex.Unlock()
	d.hash = newHash
}

func (d *hookData) send(newHash string) {
	d.set(newHash)

	// Non-blocking write.  If the channel is full, the consumer will see the
	// newest value.  If the channel was not full, the consumer will get another
	// event.
	select {
	case d.ch <- struct{}{}:
	default:
	}
}

// NewHookRunner returns a new HookRunner
func NewHookRunner(hook Hook, backoff time.Duration, data *hookData, log *logging.Logger, hasSucceededOnce chan bool) *HookRunner {
	return &HookRunner{hook: hook, backoff: backoff, data: data, logger: log, hasCompletedOnce: hasSucceededOnce}
}

// HookRunner struct
type HookRunner struct {
	// Hook to run and check
	hook Hook
	// Backoff for failed hooks
	backoff time.Duration
	// Holds the data as it crosses from producer to consumer.
	data *hookData
	// Logger
	logger *logging.Logger
	// hasCompletedOnce is used to send true if and only if first run executed
	// successfully and false otherwise to some receiver.  Should be
	// initialised to a buffered channel of size 1.
	// Is only meant for use within sendCompletedOnceMessageIfApplicable.
	hasCompletedOnce chan bool
}

// Send sends hash to hookdata
func (r *HookRunner) Send(hash string) {
	r.data.send(hash)
}

// Run waits for trigger events from the channel, and run hook when triggered
func (r *HookRunner) Run(ctx context.Context) {
	var lastHash string

	// Wait for trigger from hookData.Send
	for range r.data.events() {
		// Retry in case of error
		for {
			// Always get the latest value, in case we fail-and-retry and the
			// value changed in the meantime.  This means that we might not send
			// every single hash.
			hash := r.data.get()
			if hash == lastHash {
				break
			}

			if err := r.hook.Do(ctx, hash); err != nil {
				r.logger.Error(err, "hook failed")
				updateHookRunCountMetric(r.hook.Name(), "error")
				// don't want to sleep unnecessarily terminating anyways
				r.sendCompletedOnceMessageIfApplicable(false)
				time.Sleep(r.backoff)
			} else {
				updateHookRunCountMetric(r.hook.Name(), "success")
				lastHash = hash
				r.sendCompletedOnceMessageIfApplicable(true)
				break
			}
		}
	}
}

// If hasCompletedOnce is nil, does nothing. Otherwise, forwards the caller
// provided success status (as a boolean) of HookRunner to receivers of
// hasCompletedOnce, closes said chanel, and terminates this goroutine.
// Using this function to write to hasCompletedOnce ensures it is only ever
// written to once.
func (r *HookRunner) sendCompletedOnceMessageIfApplicable(completedSuccessfully bool) {
	if r.hasCompletedOnce != nil {
		r.hasCompletedOnce <- completedSuccessfully
		close(r.hasCompletedOnce)
		runtime.Goexit()
	}
}

// WaitForCompletion waits for HookRunner to send completion message to
// calling thread and returns either true if HookRunner executed successfully
// and some error otherwise.
// Assumes that r.hasCompletedOnce is not nil, but if it is, returns an error
func (r *HookRunner) WaitForCompletion() error {
	// Make sure function should be called
	if r.hasCompletedOnce == nil {
		return fmt.Errorf("HookRunner.WaitForCompletion called on async runner")
	}

	// Perform wait on HookRunner
	hookRunnerFinishedSuccessfully := <-r.hasCompletedOnce
	if !hookRunnerFinishedSuccessfully {
		return fmt.Errorf("hook completed with error")
	}

	return nil
}

func updateHookRunCountMetric(name, status string) {
	hookRunCount.WithLabelValues(name, status).Inc()
}
<|endoftext|>


<commit_before>package ice

import (
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/pions/pkg/stun"
	"github.com/pions/webrtc/internal/util"
	"github.com/pkg/errors"
)

// Unknown defines default public constant to use for "enum" like struct
// comparisons when no value was defined.
const Unknown = iota

// OutboundCallback is the user defined Callback that is called when ICE traffic needs to sent
type OutboundCallback func(raw []byte, local *stun.TransportAddr, remote *net.UDPAddr)

func newCandidatePair(local, remote Candidate) CandidatePair {
	return CandidatePair{
		remote: remote,
		local:  local,
	}
}

// CandidatePair represents a combination of a local and remote candidate
type CandidatePair struct {
	// lastUpdateTime ?
	remote Candidate
	local  Candidate
}

// GetAddrs returns network addresses for the candidate pair
func (c CandidatePair) GetAddrs() (local *stun.TransportAddr, remote *net.UDPAddr) {
	return &stun.TransportAddr{
			IP:   net.ParseIP(c.local.GetBase().Address),
			Port: c.local.GetBase().Port,
		}, &net.UDPAddr{
			IP:   net.ParseIP(c.remote.GetBase().Address),
			Port: c.remote.GetBase().Port,
		}
}

// Agent represents the ICE agent
type Agent struct {
	sync.RWMutex

	outboundCallback OutboundCallback
	iceNotifier      func(ConnectionState)

	tieBreaker      uint64
	connectionState ConnectionState
	gatheringState  GatheringState

	haveStarted   bool
	isControlling bool
	taskLoopChan  chan bool

	LocalUfrag      string
	LocalPwd        string
	LocalCandidates []Candidate

	remoteUfrag      string
	remotePwd        string
	remoteCandidates []Candidate

	selectedPair CandidatePair
	validPairs   []CandidatePair
}

const (
	agentTickerBaseInterval = 3 * time.Second
	stunTimeout             = 10 * time.Second
)

// NewAgent creates a new Agent
func NewAgent(outboundCallback OutboundCallback, iceNotifier func(ConnectionState)) *Agent {
	return &Agent{
		outboundCallback: outboundCallback,
		iceNotifier:      iceNotifier,

		tieBreaker:      rand.New(rand.NewSource(time.Now().UnixNano())).Uint64(),
		gatheringState:  GatheringStateComplete, // TODO trickle-ice
		connectionState: ConnectionStateNew,

		LocalUfrag: util.RandSeq(16),
		LocalPwd:   util.RandSeq(32),
	}
}

// Start starts the agent
func (a *Agent) Start(isControlling bool, remoteUfrag, remotePwd string) error {
	a.Lock()
	defer a.Unlock()

	if a.haveStarted {
		return errors.Errorf("Attempted to start agent twice")
	} else if remoteUfrag == "" {
		return errors.Errorf("remoteUfrag is empty")
	} else if remotePwd == "" {
		return errors.Errorf("remotePwd is empty")
	}

	a.isControlling = isControlling
	a.remoteUfrag = remoteUfrag
	a.remotePwd = remotePwd

	go a.agentTaskLoop()
	return nil
}

func (a *Agent) pingCandidate(local, remote Candidate) {
	if a.isControlling {
		msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, stun.GenerateTransactionId(),
			&stun.Username{Username: a.remoteUfrag + ":" + a.LocalUfrag},
			&stun.UseCandidate{},
			&stun.IceControlling{TieBreaker: a.tieBreaker},
			&stun.Priority{Priority: uint32(local.GetBase().Priority(HostCandidatePreference, 1))},
			&stun.MessageIntegrity{
				Key: []byte(a.remotePwd),
			},
			&stun.Fingerprint{},
		)
		if err != nil {
			fmt.Println(err)
			return
		}

		a.outboundCallback(msg.Pack(), &stun.TransportAddr{
			IP:   net.ParseIP(local.GetBase().Address),
			Port: local.GetBase().Port,
		}, &net.UDPAddr{
			IP:   net.ParseIP(remote.GetBase().Address),
			Port: remote.GetBase().Port,
		})
	} else {
		msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, stun.GenerateTransactionId(),
			&stun.Username{Username: a.remoteUfrag + ":" + a.LocalUfrag},
			&stun.UseCandidate{},
			&stun.IceControlled{TieBreaker: a.tieBreaker},
			&stun.Priority{Priority: uint32(local.GetBase().Priority(HostCandidatePreference, 1))},
			&stun.MessageIntegrity{
				Key: []byte(a.remotePwd),
			},
			&stun.Fingerprint{},
		)
		if err != nil {
			fmt.Println(err)
			return
		}

		a.outboundCallback(msg.Pack(), &stun.TransportAddr{
			IP:   net.ParseIP(local.GetBase().Address),
			Port: local.GetBase().Port,
		}, &net.UDPAddr{
			IP:   net.ParseIP(remote.GetBase().Address),
			Port: remote.GetBase().Port,
		})
	}

}

func (a *Agent) updateConnectionState(newState ConnectionState) {
	a.connectionState = newState
	// Call handler async since we may be holding the agent lock
	// and the handler may also require it
	go a.iceNotifier(a.connectionState)
}

func (a *Agent) setValidPair(local, remote Candidate, selected bool) {
	p := newCandidatePair(local, remote)

	if selected {
		a.selectedPair = p
		a.validPairs = nil
		// TODO: only set state to connected on selecting final pair?
		a.updateConnectionState(ConnectionStateConnected)
	} else {
		// keep track of pairs with succesfull bindings since any of them
		// can be used for communication until the final pair is selected:
		// https://tools.ietf.org/html/draft-ietf-ice-rfc5245bis-20#section-12
		a.validPairs = append(a.validPairs, p)
	}
}

func (a *Agent) agentTaskLoop() {
	// TODO this should be dynamic, and grow when the connection is stable
	t := time.NewTicker(agentTickerBaseInterval)
	a.updateConnectionState(ConnectionStateChecking)

	assertSelectedPairValid := func() bool {
		if a.selectedPair.remote == nil || a.selectedPair.local == nil {
			return false
		} else if time.Since(a.selectedPair.remote.GetBase().LastSeen) > stunTimeout {
			a.selectedPair.remote = nil
			a.selectedPair.local = nil
			a.updateConnectionState(ConnectionStateDisconnected)
			return false
		}

		return true
	}

	for {
		select {
		case <-t.C:
			a.Lock()
			if a.isControlling {
				if assertSelectedPairValid() {
					a.Unlock()
					continue
				}

				for _, localCandidate := range a.LocalCandidates {
					for _, remoteCandidate := range a.remoteCandidates {
						a.pingCandidate(localCandidate, remoteCandidate)
					}
				}
			} else {
				if assertSelectedPairValid() {
					a.Unlock()
					continue
				}

				for _, localCandidate := range a.LocalCandidates {
					for _, remoteCandidate := range a.remoteCandidates {
						a.pingCandidate(localCandidate, remoteCandidate)
					}
				}
			}
			a.Unlock()
		case <-a.taskLoopChan:
			t.Stop()
			return
		}
	}
}

// AddRemoteCandidate adds a new remote candidate
func (a *Agent) AddRemoteCandidate(c Candidate) {
	a.Lock()
	defer a.Unlock()
	a.remoteCandidates = append(a.remoteCandidates, c)
}

// AddLocalCandidate adds a new local candidate
func (a *Agent) AddLocalCandidate(c Candidate) {
	a.Lock()
	defer a.Unlock()
	a.LocalCandidates = append(a.LocalCandidates, c)
}

// Close cleans up the Agent
func (a *Agent) Close() {
	if a.taskLoopChan != nil {
		close(a.taskLoopChan)
	}
}

func isCandidateMatch(c Candidate, testAddress string, testPort int) bool {
	if c.GetBase().Address == testAddress && c.GetBase().Port == testPort {
		return true
	}

	switch c := c.(type) {
	case *CandidateSrflx:
		if c.RemoteAddress == testAddress && c.RemotePort == testPort {
			return true
		}
	}

	return false
}

func getTransportAddrCandidate(candidates []Candidate, addr *stun.TransportAddr) Candidate {
	for _, c := range candidates {
		if isCandidateMatch(c, addr.IP.String(), addr.Port) {
			return c
		}
	}
	return nil
}

func getUDPAddrCandidate(candidates []Candidate, addr *net.UDPAddr) Candidate {
	for _, c := range candidates {
		if isCandidateMatch(c, addr.IP.String(), addr.Port) {
			return c
		}
	}
	return nil
}

func (a *Agent) sendBindingSuccess(m *stun.Message, local *stun.TransportAddr, remote *net.UDPAddr) {
	if out, err := stun.Build(stun.ClassSuccessResponse, stun.MethodBinding, m.TransactionID,
		&stun.XorMappedAddress{
			XorAddress: stun.XorAddress{
				IP:   remote.IP,
				Port: remote.Port,
			},
		},
		&stun.MessageIntegrity{
			Key: []byte(a.LocalPwd),
		},
		&stun.Fingerprint{},
	); err != nil {
		fmt.Printf("Failed to handle inbound ICE from: %s to: %s error: %s", local.String(), remote.String(), err.Error())
	} else {
		a.outboundCallback(out.Pack(), local, remote)
	}

}

func (a *Agent) handleInboundControlled(m *stun.Message, local *stun.TransportAddr, remote *net.UDPAddr, localCandidate, remoteCandidate Candidate) {
	if _, isControlled := m.GetOneAttribute(stun.AttrIceControlled); isControlled && !a.isControlling {
		fmt.Println("inbound isControlled && a.isControlling == false")
		return
	}

	_, useCandidateFound := m.GetOneAttribute(stun.AttrUseCandidate)
	a.setValidPair(localCandidate, remoteCandidate, useCandidateFound)

	a.sendBindingSuccess(m, local, remote)
}

func (a *Agent) handleInboundControlling(m *stun.Message, local *stun.TransportAddr, remote *net.UDPAddr, localCandidate, remoteCandidate Candidate) {
	if _, isControlling := m.GetOneAttribute(stun.AttrIceControlling); isControlling && a.isControlling {
		fmt.Println("inbound isControlling && a.isControlling == true")
		return
	} else if _, useCandidate := m.GetOneAttribute(stun.AttrUseCandidate); useCandidate && a.isControlling {
		fmt.Println("useCandidate && a.isControlling == true")
		return
	}

	final := m.Class == stun.ClassSuccessResponse && m.Method == stun.MethodBinding
	a.setValidPair(localCandidate, remoteCandidate, final)

	if !final {
		a.sendBindingSuccess(m, local, remote)
	}
}

// HandleInbound processes traffic from a remote candidate
func (a *Agent) HandleInbound(buf []byte, local *stun.TransportAddr, remote *net.UDPAddr) {
	a.Lock()
	defer a.Unlock()

	localCandidate := getTransportAddrCandidate(a.LocalCandidates, local)
	if localCandidate == nil {
		// TODO debug
		// fmt.Printf("Could not find local candidate for %s:%d ", local.IP.String(), local.Value)
		return
	}

	remoteCandidate := getUDPAddrCandidate(a.remoteCandidates, remote)
	if remoteCandidate == nil {
		// TODO debug
		// fmt.Printf("Could not find remote candidate for %s:%d ", remote.IP.String(), remote.Value)
		return
	}
	remoteCandidate.GetBase().LastSeen = time.Now()

	m, err := stun.NewMessage(buf)
	if err != nil {
		fmt.Println(fmt.Sprintf("Failed to handle decode ICE from: %s to: %s error: %s", local.String(), remote.String(), err.Error()))
		return
	}

	if a.isControlling {
		a.handleInboundControlling(m, local, remote, localCandidate, remoteCandidate)
	} else {
		a.handleInboundControlled(m, local, remote, localCandidate, remoteCandidate)
	}

}

// SelectedPair gets the current selected pair's Addresses (or returns nil)
func (a *Agent) SelectedPair() (local *stun.TransportAddr, remote *net.UDPAddr) {
	a.RLock()
	defer a.RUnlock()

	if a.selectedPair.remote == nil || a.selectedPair.local == nil {
		for _, p := range a.validPairs {
			return p.GetAddrs()
		}
		return nil, nil
	}

	return a.selectedPair.GetAddrs()
}
<commit_msg>fix gofmt<commit_after>package ice

import (
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/pions/pkg/stun"
	"github.com/pions/webrtc/internal/util"
	"github.com/pkg/errors"
)

// Unknown defines default public constant to use for "enum" like struct
// comparisons when no value was defined.
const Unknown = iota

// OutboundCallback is the user defined Callback that is called when ICE traffic needs to sent
type OutboundCallback func(raw []byte, local *stun.TransportAddr, remote *net.UDPAddr)

func newCandidatePair(local, remote Candidate) CandidatePair {
	return CandidatePair{
		remote: remote,
		local:  local,
	}
}

// CandidatePair represents a combination of a local and remote candidate
type CandidatePair struct {
	// lastUpdateTime ?
	remote Candidate
	local  Candidate
}

// GetAddrs returns network addresses for the candidate pair
func (c CandidatePair) GetAddrs() (local *stun.TransportAddr, remote *net.UDPAddr) {
	return &stun.TransportAddr{
			IP:   net.ParseIP(c.local.GetBase().Address),
			Port: c.local.GetBase().Port,
		}, &net.UDPAddr{
			IP:   net.ParseIP(c.remote.GetBase().Address),
			Port: c.remote.GetBase().Port,
		}
}

// Agent represents the ICE agent
type Agent struct {
	sync.RWMutex

	outboundCallback OutboundCallback
	iceNotifier      func(ConnectionState)

	tieBreaker      uint64
	connectionState ConnectionState
	gatheringState  GatheringState

	haveStarted   bool
	isControlling bool
	taskLoopChan  chan bool

	LocalUfrag      string
	LocalPwd        string
	LocalCandidates []Candidate

	remoteUfrag      string
	remotePwd        string
	remoteCandidates []Candidate

	selectedPair CandidatePair
	validPairs   []CandidatePair
}

const (
	agentTickerBaseInterval = 3 * time.Second
	stunTimeout             = 10 * time.Second
)

// NewAgent creates a new Agent
func NewAgent(outboundCallback OutboundCallback, iceNotifier func(ConnectionState)) *Agent {
	return &Agent{
		outboundCallback: outboundCallback,
		iceNotifier:      iceNotifier,

		tieBreaker:      rand.New(rand.NewSource(time.Now().UnixNano())).Uint64(),
		gatheringState:  GatheringStateComplete, // TODO trickle-ice
		connectionState: ConnectionStateNew,

		LocalUfrag: util.RandSeq(16),
		LocalPwd:   util.RandSeq(32),
	}
}

// Start starts the agent
func (a *Agent) Start(isControlling bool, remoteUfrag, remotePwd string) error {
	a.Lock()
	defer a.Unlock()

	if a.haveStarted {
		return errors.Errorf("Attempted to start agent twice")
	} else if remoteUfrag == "" {
		return errors.Errorf("remoteUfrag is empty")
	} else if remotePwd == "" {
		return errors.Errorf("remotePwd is empty")
	}

	a.isControlling = isControlling
	a.remoteUfrag = remoteUfrag
	a.remotePwd = remotePwd

	go a.agentTaskLoop()
	return nil
}

func (a *Agent) pingCandidate(local, remote Candidate) {
	if a.isControlling {
		msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, stun.GenerateTransactionId(),
			&stun.Username{Username: a.remoteUfrag + ":" + a.LocalUfrag},
			&stun.UseCandidate{},
			&stun.IceControlling{TieBreaker: a.tieBreaker},
			&stun.Priority{Priority: uint32(local.GetBase().Priority(HostCandidatePreference, 1))},
			&stun.MessageIntegrity{
				Key: []byte(a.remotePwd),
			},
			&stun.Fingerprint{},
		)
		if err != nil {
			fmt.Println(err)
			return
		}

		a.outboundCallback(msg.Pack(), &stun.TransportAddr{
			IP:   net.ParseIP(local.GetBase().Address),
			Port: local.GetBase().Port,
		}, &net.UDPAddr{
			IP:   net.ParseIP(remote.GetBase().Address),
			Port: remote.GetBase().Port,
		})
	} else {
		msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, stun.GenerateTransactionId(),
			&stun.Username{Username: a.remoteUfrag + ":" + a.LocalUfrag},
			&stun.UseCandidate{},
			&stun.IceControlled{TieBreaker: a.tieBreaker},
			&stun.Priority{Priority: uint32(local.GetBase().Priority(HostCandidatePreference, 1))},
			&stun.MessageIntegrity{
				Key: []byte(a.remotePwd),
			},
			&stun.Fingerprint{},
		)
		if err != nil {
			fmt.Println(err)
			return
		}

		a.outboundCallback(msg.Pack(), &stun.TransportAddr{
			IP:   net.ParseIP(local.GetBase().Address),
			Port: local.GetBase().Port,
		}, &net.UDPAddr{
			IP:   net.ParseIP(remote.GetBase().Address),
			Port: remote.GetBase().Port,
		})
	}
}

func (a *Agent) updateConnectionState(newState ConnectionState) {
	a.connectionState = newState
	// Call handler async since we may be holding the agent lock
	// and the handler may also require it
	go a.iceNotifier(a.connectionState)
}

func (a *Agent) setValidPair(local, remote Candidate, selected bool) {
	p := newCandidatePair(local, remote)

	if selected {
		a.selectedPair = p
		a.validPairs = nil
		// TODO: only set state to connected on selecting final pair?
		a.updateConnectionState(ConnectionStateConnected)
	} else {
		// keep track of pairs with succesfull bindings since any of them
		// can be used for communication until the final pair is selected:
		// https://tools.ietf.org/html/draft-ietf-ice-rfc5245bis-20#section-12
		a.validPairs = append(a.validPairs, p)
	}
}

func (a *Agent) agentTaskLoop() {
	// TODO this should be dynamic, and grow when the connection is stable
	t := time.NewTicker(agentTickerBaseInterval)
	a.updateConnectionState(ConnectionStateChecking)

	assertSelectedPairValid := func() bool {
		if a.selectedPair.remote == nil || a.selectedPair.local == nil {
			return false
		} else if time.Since(a.selectedPair.remote.GetBase().LastSeen) > stunTimeout {
			a.selectedPair.remote = nil
			a.selectedPair.local = nil
			a.updateConnectionState(ConnectionStateDisconnected)
			return false
		}

		return true
	}

	for {
		select {
		case <-t.C:
			a.Lock()
			if a.isControlling {
				if assertSelectedPairValid() {
					a.Unlock()
					continue
				}

				for _, localCandidate := range a.LocalCandidates {
					for _, remoteCandidate := range a.remoteCandidates {
						a.pingCandidate(localCandidate, remoteCandidate)
					}
				}
			} else {
				if assertSelectedPairValid() {
					a.Unlock()
					continue
				}

				for _, localCandidate := range a.LocalCandidates {
					for _, remoteCandidate := range a.remoteCandidates {
						a.pingCandidate(localCandidate, remoteCandidate)
					}
				}
			}
			a.Unlock()
		case <-a.taskLoopChan:
			t.Stop()
			return
		}
	}
}

// AddRemoteCandidate adds a new remote candidate
func (a *Agent) AddRemoteCandidate(c Candidate) {
	a.Lock()
	defer a.Unlock()
	a.remoteCandidates = append(a.remoteCandidates, c)
}

// AddLocalCandidate adds a new local candidate
func (a *Agent) AddLocalCandidate(c Candidate) {
	a.Lock()
	defer a.Unlock()
	a.LocalCandidates = append(a.LocalCandidates, c)
}

// Close cleans up the Agent
func (a *Agent) Close() {
	if a.taskLoopChan != nil {
		close(a.taskLoopChan)
	}
}

func isCandidateMatch(c Candidate, testAddress string, testPort int) bool {
	if c.GetBase().Address == testAddress && c.GetBase().Port == testPort {
		return true
	}

	switch c := c.(type) {
	case *CandidateSrflx:
		if c.RemoteAddress == testAddress && c.RemotePort == testPort {
			return true
		}
	}

	return false
}

func getTransportAddrCandidate(candidates []Candidate, addr *stun.TransportAddr) Candidate {
	for _, c := range candidates {
		if isCandidateMatch(c, addr.IP.String(), addr.Port) {
			return c
		}
	}
	return nil
}

func getUDPAddrCandidate(candidates []Candidate, addr *net.UDPAddr) Candidate {
	for _, c := range candidates {
		if isCandidateMatch(c, addr.IP.String(), addr.Port) {
			return c
		}
	}
	return nil
}

func (a *Agent) sendBindingSuccess(m *stun.Message, local *stun.TransportAddr, remote *net.UDPAddr) {
	if out, err := stun.Build(stun.ClassSuccessResponse, stun.MethodBinding, m.TransactionID,
		&stun.XorMappedAddress{
			XorAddress: stun.XorAddress{
				IP:   remote.IP,
				Port: remote.Port,
			},
		},
		&stun.MessageIntegrity{
			Key: []byte(a.LocalPwd),
		},
		&stun.Fingerprint{},
	); err != nil {
		fmt.Printf("Failed to handle inbound ICE from: %s to: %s error: %s", local.String(), remote.String(), err.Error())
	} else {
		a.outboundCallback(out.Pack(), local, remote)
	}

}

func (a *Agent) handleInboundControlled(m *stun.Message, local *stun.TransportAddr, remote *net.UDPAddr, localCandidate, remoteCandidate Candidate) {
	if _, isControlled := m.GetOneAttribute(stun.AttrIceControlled); isControlled && !a.isControlling {
		fmt.Println("inbound isControlled && a.isControlling == false")
		return
	}

	_, useCandidateFound := m.GetOneAttribute(stun.AttrUseCandidate)
	a.setValidPair(localCandidate, remoteCandidate, useCandidateFound)

	a.sendBindingSuccess(m, local, remote)
}

func (a *Agent) handleInboundControlling(m *stun.Message, local *stun.TransportAddr, remote *net.UDPAddr, localCandidate, remoteCandidate Candidate) {
	if _, isControlling := m.GetOneAttribute(stun.AttrIceControlling); isControlling && a.isControlling {
		fmt.Println("inbound isControlling && a.isControlling == true")
		return
	} else if _, useCandidate := m.GetOneAttribute(stun.AttrUseCandidate); useCandidate && a.isControlling {
		fmt.Println("useCandidate && a.isControlling == true")
		return
	}

	final := m.Class == stun.ClassSuccessResponse && m.Method == stun.MethodBinding
	a.setValidPair(localCandidate, remoteCandidate, final)

	if !final {
		a.sendBindingSuccess(m, local, remote)
	}
}

// HandleInbound processes traffic from a remote candidate
func (a *Agent) HandleInbound(buf []byte, local *stun.TransportAddr, remote *net.UDPAddr) {
	a.Lock()
	defer a.Unlock()

	localCandidate := getTransportAddrCandidate(a.LocalCandidates, local)
	if localCandidate == nil {
		// TODO debug
		// fmt.Printf("Could not find local candidate for %s:%d ", local.IP.String(), local.Value)
		return
	}

	remoteCandidate := getUDPAddrCandidate(a.remoteCandidates, remote)
	if remoteCandidate == nil {
		// TODO debug
		// fmt.Printf("Could not find remote candidate for %s:%d ", remote.IP.String(), remote.Value)
		return
	}
	remoteCandidate.GetBase().LastSeen = time.Now()

	m, err := stun.NewMessage(buf)
	if err != nil {
		fmt.Println(fmt.Sprintf("Failed to handle decode ICE from: %s to: %s error: %s", local.String(), remote.String(), err.Error()))
		return
	}

	if a.isControlling {
		a.handleInboundControlling(m, local, remote, localCandidate, remoteCandidate)
	} else {
		a.handleInboundControlled(m, local, remote, localCandidate, remoteCandidate)
	}

}

// SelectedPair gets the current selected pair's Addresses (or returns nil)
func (a *Agent) SelectedPair() (local *stun.TransportAddr, remote *net.UDPAddr) {
	a.RLock()
	defer a.RUnlock()

	if a.selectedPair.remote == nil || a.selectedPair.local == nil {
		for _, p := range a.validPairs {
			return p.GetAddrs()
		}
		return nil, nil
	}

	return a.selectedPair.GetAddrs()
}
<|endoftext|>