<commit_before>package cfg

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

	"github.com/scrapli/scrapligo/logging"

	"github.com/scrapli/scrapligo/driver/base"
	"github.com/scrapli/scrapligo/driver/network"
)

type eosPatterns struct {
	globalCommentLinePattern *regexp.Regexp
	bannerPattern            *regexp.Regexp
	endPattern               *regexp.Regexp
}

var (
	eosPatternsInstance     *eosPatterns //nolint:gochecknoglobals
	eosPAtternsInstanceOnce sync.Once    //nolint:gochecknoglobals
)

func getEOSPatterns() *eosPatterns {
	eosPAtternsInstanceOnce.Do(func() {
		eosPatternsInstance = &eosPatterns{
			globalCommentLinePattern: regexp.MustCompile(`(?im)^! .*$`),
			bannerPattern:            regexp.MustCompile(`(?ims)^banner.*EOF$`),
			endPattern:               regexp.MustCompile(`end$`),
		}
	})

	return eosPatternsInstance
}

type EOSCfg struct {
	conn              *network.Driver
	VersionPattern    *regexp.Regexp
	configCommandMap  map[string]string
	configSessionName string
}

// NewEOSCfg return a cfg instance setup for an Arista EOS device.
func NewEOSCfg(
	conn *network.Driver,
	options ...Option,
) (*Cfg, error) {
	options = append([]Option{WithConfigSources([]string{"running", "startup"})}, options...)

	c, err := newCfg(conn, options...)
	if err != nil {
		return nil, err
	}

	c.Platform = &EOSCfg{
		conn:           conn,
		VersionPattern: regexp.MustCompile(`(?i)\d+\.\d+\.[a-z0-9\-]+(\.\d+[a-z]?)?`),
		configCommandMap: map[string]string{
			"running": "show running-config",
			"startup": "show startup-config",
		},
	}

	err = setPlatformOptions(c.Platform, options...)
	if err != nil {
		return nil, err
	}

	return c, nil
}

func (p *EOSCfg) ClearConfigSession() {
	p.configSessionName = ""
}

// GetVersion get the version from the device.
func (p *EOSCfg) GetVersion() (string, []*base.Response, error) {
	versionResult, err := p.conn.SendCommand("show version | i Software image version")
	if err != nil {
		return "", nil, err
	}

	return p.VersionPattern.FindString(
			versionResult.Result,
		), []*base.Response{
			versionResult,
		}, nil
}

// GetConfig get the configuration of a source datastore from the device.
func (p *EOSCfg) GetConfig(source string) (string, []*base.Response, error) {
	cmd, err := getConfigCommand(p.configCommandMap, source)
	if err != nil {
		return "", nil, err
	}

	configResult, err := p.conn.SendCommand(cmd)

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

	return configResult.Result, []*base.Response{configResult}, nil
}

func (p *EOSCfg) prepareConfigPayloads(config string) (stdConfig, eagerConfig string) {
	patterns := getEOSPatterns()

	// remove comment lines
	config = patterns.globalCommentLinePattern.ReplaceAllString(config, "!")

	// remove "end" at the end of the config - if its present it will drop scrapli out
	// of the config session which we do not want
	config = patterns.endPattern.ReplaceAllString(config, "!")

	// find all sections that need to be "eagerly" sent; remove those sections from the "normal"
	// config, then join all the eager sections into a single string
	eagerSections := patterns.bannerPattern.FindStringSubmatch(config)
	eagerConfig = strings.Join(eagerSections, "\n")

	for _, section := range eagerSections {
		config = strings.Replace(config, section, "!", -1)
	}

	return config, eagerConfig
}

// RegisterConfigSession register a configuration session in EOS.
func (p *EOSCfg) RegisterConfigSession(sessionName string) error {
	_, ok := p.conn.PrivilegeLevels[sessionName]

	if ok {
		return ErrConfigSessionAlreadyExists
	}

	sessionPrompt := regexp.QuoteMeta(sessionName[:6])
	sessionPromptPattern := fmt.Sprintf(
		`(?im)^[\w.\-@()/:\s]{1,63}\(config\-s\-%s[\w.\-@_/:]{0,32}\)#\s?$`,
		sessionPrompt,
	)

	sessionPrivilegeLevel := &base.PrivilegeLevel{
		Pattern:        sessionPromptPattern,
		Name:           sessionName,
		PreviousPriv:   execPrivLevel,
		Deescalate:     "end",
		Escalate:       fmt.Sprintf("configure session %s", sessionName),
		EscalateAuth:   false,
		EscalatePrompt: "",
	}

	p.conn.PrivilegeLevels[sessionName] = sessionPrivilegeLevel
	p.conn.UpdatePrivilegeLevels()

	return nil
}

func (p *EOSCfg) loadConfig(
	stdConfig, eagerConfig string,
	replace bool,
) ([]*base.Response, error) {
	var scrapliResponses []*base.Response

	if replace {
		rollbackCleanConfigResult, rollbackErr := p.conn.SendConfig("rollback clean-config",
			base.WithDesiredPrivilegeLevel(p.configSessionName))
		if rollbackErr != nil {
			return scrapliResponses, rollbackErr
		}

		scrapliResponses = append(scrapliResponses, rollbackCleanConfigResult)
	}

	configResult, stdConfigErr := p.conn.SendConfig(
		stdConfig,
		base.WithDesiredPrivilegeLevel(p.configSessionName),
	)
	if stdConfigErr != nil || configResult.Failed != nil {
		return scrapliResponses, stdConfigErr
	}

	scrapliResponses = append(scrapliResponses, configResult)

	eagerResult, eagerConfigErr := p.conn.SendConfig(
		eagerConfig,
		base.WithDesiredPrivilegeLevel(p.configSessionName),
		base.WithSendEager(true),
	)

	if eagerConfigErr != nil {
		return scrapliResponses, eagerConfigErr
	}

	if eagerResult.Failed != nil {
		return scrapliResponses, eagerConfigErr
	}

	scrapliResponses = append(scrapliResponses, eagerResult)

	return scrapliResponses, nil
}

// LoadConfig load a candidate configuration.
func (p *EOSCfg) LoadConfig(
	config string,
	replace bool,
	options *OperationOptions,
) ([]*base.Response, error) {
	// options are unused for eos load config
	_ = options

	stdConfig, eagerConfig := p.prepareConfigPayloads(config)

	if p.configSessionName == "" {
		p.configSessionName = fmt.Sprintf("scrapli_cfg_%d", time.Now().Unix())

		logging.LogDebug(
			FormatLogMessage(
				p.conn,
				"debug",
				fmt.Sprintf("configuration session name will be %s", p.configSessionName),
			),
		)

		err := p.RegisterConfigSession(p.configSessionName)
		if err != nil {
			return nil, err
		}
	}

	return p.loadConfig(stdConfig, eagerConfig, replace)
}

// AbortConfig abort the loaded candidate configuration.
func (p *EOSCfg) AbortConfig() ([]*base.Response, error) {
	var scrapliResponses []*base.Response

	err := p.conn.AcquirePriv(p.configSessionName)
	if err != nil {
		return scrapliResponses, err
	}

	_, err = p.conn.Channel.SendInput("abort", false, false, p.conn.Channel.TimeoutOps)
	if err != nil {
		return scrapliResponses, err
	}

	p.conn.CurrentPriv = execPrivLevel

	return scrapliResponses, nil
}

// CommitConfig commit the loaded candidate configuration.
func (p *EOSCfg) CommitConfig(source string) ([]*base.Response, error) {
	if source != runningConfig {
		logging.LogDebug(
			FormatLogMessage(
				p.conn,
				"warning",
				"eos only supports committing to running config, running config is automatically copied to "+
					"startup during commit operation",
			),
		)
	}

	var scrapliResponses []*base.Response

	commands := []string{
		fmt.Sprintf("configure session %s commit", p.configSessionName),
		"copy running-config startup-config",
	}

	m, err := p.conn.SendCommands(commands)
	if err != nil {
		return scrapliResponses, err
	}

	scrapliResponses = append(scrapliResponses, m.Responses...)

	return scrapliResponses, err
}

func (p *EOSCfg) normalizeSourceAndCandidateConfigs(
	sourceConfig, candidateConfig string,
) (normalizedSourceConfig, normalizedCandidateConfig string) {
	patterns := getEOSPatterns()

	// Remove all comment lines from both the source and candidate configs -- this is only done
	// here pre-diff, so we dont modify the user provided candidate config which can totally have
	// those comment lines - we only remove "global" (top level) comments though... user comments
	// attached to interfaces and the stuff will remain
	normalizedSourceConfig = patterns.globalCommentLinePattern.ReplaceAllString(sourceConfig, "")
	normalizedSourceConfig = strings.Replace(normalizedSourceConfig, "\n\n", "\n", -1)

	normalizedCandidateConfig = patterns.globalCommentLinePattern.ReplaceAllString(
		candidateConfig,
		"",
	)
	normalizedCandidateConfig = strings.Replace(normalizedCandidateConfig, "\n\n", "\n", -1)

	return normalizedSourceConfig, normalizedCandidateConfig
}

// DiffConfig diff the candidate configuration against a source config.
func (p *EOSCfg) DiffConfig(
	source, candidateConfig string,
) (responses []*base.Response,
	normalizedSourceConfig,
	normalizedCandidateConfig,
	deviceDiff string, err error) {
	if source != runningConfig {
		logging.LogDebug(
			FormatLogMessage(
				p.conn,
				"warning",
				"eos only supports diffing against the running config",
			),
		)
	}

	var scrapliResponses []*base.Response

	diffResult, err := p.conn.SendConfig(
		"show session-config diffs",
		base.WithDesiredPrivilegeLevel(p.configSessionName),
	)
	if err != nil {
		return scrapliResponses, "", "", "", err
	}

	scrapliResponses = append(scrapliResponses, diffResult)

	if diffResult.Failed != nil {
		logging.LogError(
			FormatLogMessage(
				p.conn,
				"error",
				"failed generating diff for config session",
			),
		)

		return scrapliResponses, "", "", "", nil
	}

	deviceDiff = diffResult.Result

	sourceConfig, getConfigR, err := p.GetConfig(source)
	if err != nil {
		return scrapliResponses, "", "", "", nil
	}

	scrapliResponses = append(scrapliResponses, getConfigR[0])

	if getConfigR[0].Failed != nil {
		logging.LogError(
			FormatLogMessage(
				p.conn,
				"error",
				"failed fetching source config for diff comparison",
			),
		)

		return scrapliResponses, "", "", "", nil
	}

	normalizedSourceConfig, normalizedCandidateConfig = p.normalizeSourceAndCandidateConfigs(
		sourceConfig,
		candidateConfig,
	)

	return scrapliResponses, normalizedSourceConfig, normalizedCandidateConfig, deviceDiff, nil
}
<commit_msg>fix typo<commit_after>package cfg

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

	"github.com/scrapli/scrapligo/logging"

	"github.com/scrapli/scrapligo/driver/base"
	"github.com/scrapli/scrapligo/driver/network"
)

type eosPatterns struct {
	globalCommentLinePattern *regexp.Regexp
	bannerPattern            *regexp.Regexp
	endPattern               *regexp.Regexp
}

var (
	eosPatternsInstance     *eosPatterns //nolint:gochecknoglobals
	eosPatternsInstanceOnce sync.Once    //nolint:gochecknoglobals
)

func getEOSPatterns() *eosPatterns {
	eosPatternsInstanceOnce.Do(func() {
		eosPatternsInstance = &eosPatterns{
			globalCommentLinePattern: regexp.MustCompile(`(?im)^! .*$`),
			bannerPattern:            regexp.MustCompile(`(?ims)^banner.*EOF$`),
			endPattern:               regexp.MustCompile(`end$`),
		}
	})

	return eosPatternsInstance
}

type EOSCfg struct {
	conn              *network.Driver
	VersionPattern    *regexp.Regexp
	configCommandMap  map[string]string
	configSessionName string
}

// NewEOSCfg return a cfg instance setup for an Arista EOS device.
func NewEOSCfg(
	conn *network.Driver,
	options ...Option,
) (*Cfg, error) {
	options = append([]Option{WithConfigSources([]string{"running", "startup"})}, options...)

	c, err := newCfg(conn, options...)
	if err != nil {
		return nil, err
	}

	c.Platform = &EOSCfg{
		conn:           conn,
		VersionPattern: regexp.MustCompile(`(?i)\d+\.\d+\.[a-z0-9\-]+(\.\d+[a-z]?)?`),
		configCommandMap: map[string]string{
			"running": "show running-config",
			"startup": "show startup-config",
		},
	}

	err = setPlatformOptions(c.Platform, options...)
	if err != nil {
		return nil, err
	}

	return c, nil
}

func (p *EOSCfg) ClearConfigSession() {
	p.configSessionName = ""
}

// GetVersion get the version from the device.
func (p *EOSCfg) GetVersion() (string, []*base.Response, error) {
	versionResult, err := p.conn.SendCommand("show version | i Software image version")
	if err != nil {
		return "", nil, err
	}

	return p.VersionPattern.FindString(
			versionResult.Result,
		), []*base.Response{
			versionResult,
		}, nil
}

// GetConfig get the configuration of a source datastore from the device.
func (p *EOSCfg) GetConfig(source string) (string, []*base.Response, error) {
	cmd, err := getConfigCommand(p.configCommandMap, source)
	if err != nil {
		return "", nil, err
	}

	configResult, err := p.conn.SendCommand(cmd)

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

	return configResult.Result, []*base.Response{configResult}, nil
}

func (p *EOSCfg) prepareConfigPayloads(config string) (stdConfig, eagerConfig string) {
	patterns := getEOSPatterns()

	// remove comment lines
	config = patterns.globalCommentLinePattern.ReplaceAllString(config, "!")

	// remove "end" at the end of the config - if its present it will drop scrapli out
	// of the config session which we do not want
	config = patterns.endPattern.ReplaceAllString(config, "!")

	// find all sections that need to be "eagerly" sent; remove those sections from the "normal"
	// config, then join all the eager sections into a single string
	eagerSections := patterns.bannerPattern.FindStringSubmatch(config)
	eagerConfig = strings.Join(eagerSections, "\n")

	for _, section := range eagerSections {
		config = strings.Replace(config, section, "!", -1)
	}

	return config, eagerConfig
}

// RegisterConfigSession register a configuration session in EOS.
func (p *EOSCfg) RegisterConfigSession(sessionName string) error {
	_, ok := p.conn.PrivilegeLevels[sessionName]

	if ok {
		return ErrConfigSessionAlreadyExists
	}

	sessionPrompt := regexp.QuoteMeta(sessionName[:6])
	sessionPromptPattern := fmt.Sprintf(
		`(?im)^[\w.\-@()/:\s]{1,63}\(config\-s\-%s[\w.\-@_/:]{0,32}\)#\s?$`,
		sessionPrompt,
	)

	sessionPrivilegeLevel := &base.PrivilegeLevel{
		Pattern:        sessionPromptPattern,
		Name:           sessionName,
		PreviousPriv:   execPrivLevel,
		Deescalate:     "end",
		Escalate:       fmt.Sprintf("configure session %s", sessionName),
		EscalateAuth:   false,
		EscalatePrompt: "",
	}

	p.conn.PrivilegeLevels[sessionName] = sessionPrivilegeLevel
	p.conn.UpdatePrivilegeLevels()

	return nil
}

func (p *EOSCfg) loadConfig(
	stdConfig, eagerConfig string,
	replace bool,
) ([]*base.Response, error) {
	var scrapliResponses []*base.Response

	if replace {
		rollbackCleanConfigResult, rollbackErr := p.conn.SendConfig("rollback clean-config",
			base.WithDesiredPrivilegeLevel(p.configSessionName))
		if rollbackErr != nil {
			return scrapliResponses, rollbackErr
		}

		scrapliResponses = append(scrapliResponses, rollbackCleanConfigResult)
	}

	configResult, stdConfigErr := p.conn.SendConfig(
		stdConfig,
		base.WithDesiredPrivilegeLevel(p.configSessionName),
	)
	if stdConfigErr != nil || configResult.Failed != nil {
		return scrapliResponses, stdConfigErr
	}

	scrapliResponses = append(scrapliResponses, configResult)

	eagerResult, eagerConfigErr := p.conn.SendConfig(
		eagerConfig,
		base.WithDesiredPrivilegeLevel(p.configSessionName),
		base.WithSendEager(true),
	)

	if eagerConfigErr != nil {
		return scrapliResponses, eagerConfigErr
	}

	if eagerResult.Failed != nil {
		return scrapliResponses, eagerConfigErr
	}

	scrapliResponses = append(scrapliResponses, eagerResult)

	return scrapliResponses, nil
}

// LoadConfig load a candidate configuration.
func (p *EOSCfg) LoadConfig(
	config string,
	replace bool,
	options *OperationOptions,
) ([]*base.Response, error) {
	// options are unused for eos load config
	_ = options

	stdConfig, eagerConfig := p.prepareConfigPayloads(config)

	if p.configSessionName == "" {
		p.configSessionName = fmt.Sprintf("scrapli_cfg_%d", time.Now().Unix())

		logging.LogDebug(
			FormatLogMessage(
				p.conn,
				"debug",
				fmt.Sprintf("configuration session name will be %s", p.configSessionName),
			),
		)

		err := p.RegisterConfigSession(p.configSessionName)
		if err != nil {
			return nil, err
		}
	}

	return p.loadConfig(stdConfig, eagerConfig, replace)
}

// AbortConfig abort the loaded candidate configuration.
func (p *EOSCfg) AbortConfig() ([]*base.Response, error) {
	var scrapliResponses []*base.Response

	err := p.conn.AcquirePriv(p.configSessionName)
	if err != nil {
		return scrapliResponses, err
	}

	_, err = p.conn.Channel.SendInput("abort", false, false, p.conn.Channel.TimeoutOps)
	if err != nil {
		return scrapliResponses, err
	}

	p.conn.CurrentPriv = execPrivLevel

	return scrapliResponses, nil
}

// CommitConfig commit the loaded candidate configuration.
func (p *EOSCfg) CommitConfig(source string) ([]*base.Response, error) {
	if source != runningConfig {
		logging.LogDebug(
			FormatLogMessage(
				p.conn,
				"warning",
				"eos only supports committing to running config, running config is automatically copied to "+
					"startup during commit operation",
			),
		)
	}

	var scrapliResponses []*base.Response

	commands := []string{
		fmt.Sprintf("configure session %s commit", p.configSessionName),
		"copy running-config startup-config",
	}

	m, err := p.conn.SendCommands(commands)
	if err != nil {
		return scrapliResponses, err
	}

	scrapliResponses = append(scrapliResponses, m.Responses...)

	return scrapliResponses, err
}

func (p *EOSCfg) normalizeSourceAndCandidateConfigs(
	sourceConfig, candidateConfig string,
) (normalizedSourceConfig, normalizedCandidateConfig string) {
	patterns := getEOSPatterns()

	// Remove all comment lines from both the source and candidate configs -- this is only done
	// here pre-diff, so we dont modify the user provided candidate config which can totally have
	// those comment lines - we only remove "global" (top level) comments though... user comments
	// attached to interfaces and the stuff will remain
	normalizedSourceConfig = patterns.globalCommentLinePattern.ReplaceAllString(sourceConfig, "")
	normalizedSourceConfig = strings.Replace(normalizedSourceConfig, "\n\n", "\n", -1)

	normalizedCandidateConfig = patterns.globalCommentLinePattern.ReplaceAllString(
		candidateConfig,
		"",
	)
	normalizedCandidateConfig = strings.Replace(normalizedCandidateConfig, "\n\n", "\n", -1)

	return normalizedSourceConfig, normalizedCandidateConfig
}

// DiffConfig diff the candidate configuration against a source config.
func (p *EOSCfg) DiffConfig(
	source, candidateConfig string,
) (responses []*base.Response,
	normalizedSourceConfig,
	normalizedCandidateConfig,
	deviceDiff string, err error) {
	if source != runningConfig {
		logging.LogDebug(
			FormatLogMessage(
				p.conn,
				"warning",
				"eos only supports diffing against the running config",
			),
		)
	}

	var scrapliResponses []*base.Response

	diffResult, err := p.conn.SendConfig(
		"show session-config diffs",
		base.WithDesiredPrivilegeLevel(p.configSessionName),
	)
	if err != nil {
		return scrapliResponses, "", "", "", err
	}

	scrapliResponses = append(scrapliResponses, diffResult)

	if diffResult.Failed != nil {
		logging.LogError(
			FormatLogMessage(
				p.conn,
				"error",
				"failed generating diff for config session",
			),
		)

		return scrapliResponses, "", "", "", nil
	}

	deviceDiff = diffResult.Result

	sourceConfig, getConfigR, err := p.GetConfig(source)
	if err != nil {
		return scrapliResponses, "", "", "", nil
	}

	scrapliResponses = append(scrapliResponses, getConfigR[0])

	if getConfigR[0].Failed != nil {
		logging.LogError(
			FormatLogMessage(
				p.conn,
				"error",
				"failed fetching source config for diff comparison",
			),
		)

		return scrapliResponses, "", "", "", nil
	}

	normalizedSourceConfig, normalizedCandidateConfig = p.normalizeSourceAndCandidateConfigs(
		sourceConfig,
		candidateConfig,
	)

	return scrapliResponses, normalizedSourceConfig, normalizedCandidateConfig, deviceDiff, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The logr Authors.

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 main

import (
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/go-logr/logr"
)

// tabLogSink is a sample logr.LogSink that logs to stderr.
// It's terribly inefficient, and is only a basic example.
type tabLogSink struct {
	name      string
	keyValues map[string]interface{}
	writer    *tabwriter.Writer
}

var _ logr.LogSink = tabLogSink{}

func (_ tabLogSink) Init(info logr.RuntimeInfo) {
}

func (_ tabLogSink) Enabled(level int) bool {
	return true
}

func (l tabLogSink) Info(level int, msg string, kvs ...interface{}) {
	fmt.Fprintf(l.writer, "%s\t%s\t", l.name, msg)
	for k, v := range l.keyValues {
		fmt.Fprintf(l.writer, "%s: %+v  ", k, v)
	}
	for i := 0; i < len(kvs); i += 2 {
		fmt.Fprintf(l.writer, "%s: %+v  ", kvs[i], kvs[i+1])
	}
	fmt.Fprintf(l.writer, "\n")
	l.writer.Flush()
}

func (l tabLogSink) Error(err error, msg string, kvs ...interface{}) {
	kvs = append(kvs, "error", err)
	l.Info(0, msg, kvs...)
}

func (l tabLogSink) WithName(name string) logr.LogSink {
	return tabLogSink{
		name:      l.name + "." + name,
		keyValues: l.keyValues,
		writer:    l.writer,
	}
}

func (l tabLogSink) WithValues(kvs ...interface{}) logr.LogSink {
	newMap := make(map[string]interface{}, len(l.keyValues)+len(kvs)/2)
	for k, v := range l.keyValues {
		newMap[k] = v
	}
	for i := 0; i < len(kvs); i += 2 {
		newMap[kvs[i].(string)] = kvs[i+1]
	}
	return tabLogSink{
		name:      l.name,
		keyValues: newMap,
		writer:    l.writer,
	}
}

// NewTabLogger is the main entry-point to this implementation.  App developers
// call this somewhere near main() and thenceforth only deal with logr.Logger.
func NewTabLogger() logr.Logger {
	sink := tabLogSink{
		writer: tabwriter.NewWriter(os.Stderr, 40, 8, 2, '\t', 0),
	}
	return logr.New(sink)
}
<commit_msg>Make tab logger Init() use a pointer receiver<commit_after>/*
Copyright 2019 The logr Authors.

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 main

import (
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/go-logr/logr"
)

// tabLogSink is a sample logr.LogSink that logs to stderr.
// It's terribly inefficient, and is only a basic example.
type tabLogSink struct {
	name      string
	keyValues map[string]interface{}
	writer    *tabwriter.Writer
}

var _ logr.LogSink = &tabLogSink{}

// Note that Init usually takes a pointer so it can modify the receiver to save
// runtime info.
func (_ *tabLogSink) Init(info logr.RuntimeInfo) {
}

func (_ tabLogSink) Enabled(level int) bool {
	return true
}

func (l tabLogSink) Info(level int, msg string, kvs ...interface{}) {
	fmt.Fprintf(l.writer, "%s\t%s\t", l.name, msg)
	for k, v := range l.keyValues {
		fmt.Fprintf(l.writer, "%s: %+v  ", k, v)
	}
	for i := 0; i < len(kvs); i += 2 {
		fmt.Fprintf(l.writer, "%s: %+v  ", kvs[i], kvs[i+1])
	}
	fmt.Fprintf(l.writer, "\n")
	l.writer.Flush()
}

func (l tabLogSink) Error(err error, msg string, kvs ...interface{}) {
	kvs = append(kvs, "error", err)
	l.Info(0, msg, kvs...)
}

func (l tabLogSink) WithName(name string) logr.LogSink {
	return &tabLogSink{
		name:      l.name + "." + name,
		keyValues: l.keyValues,
		writer:    l.writer,
	}
}

func (l tabLogSink) WithValues(kvs ...interface{}) logr.LogSink {
	newMap := make(map[string]interface{}, len(l.keyValues)+len(kvs)/2)
	for k, v := range l.keyValues {
		newMap[k] = v
	}
	for i := 0; i < len(kvs); i += 2 {
		newMap[kvs[i].(string)] = kvs[i+1]
	}
	return &tabLogSink{
		name:      l.name,
		keyValues: newMap,
		writer:    l.writer,
	}
}

// NewTabLogger is the main entry-point to this implementation.  App developers
// call this somewhere near main() and thenceforth only deal with logr.Logger.
func NewTabLogger() logr.Logger {
	sink := &tabLogSink{
		writer: tabwriter.NewWriter(os.Stderr, 40, 8, 2, '\t', 0),
	}
	return logr.New(sink)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/artyom/autoflags"
	"github.com/ttacon/chalk"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] ip1 ip2 ...\n", os.Args[0])
		flag.PrintDefaults()
	}
}

type Config struct {
	Concurrency int    `flag:"n,concurrent ssh sessions"`
	Command     string `flag:"cmd,command to run"`
	Login       string `flag:"l,login to use"`
	Port        int    `flag:"p,default port"`
	StdinFile   string `flag:"stdin,REGULAR (no piping!) file to pass to stdin of remote command"`
	DumpFiles   bool   `flag:"logs,save stdout/stderr to separate per-host logs"`
	StdoutFmt   string `flag:"logs.stdout,format of stdout per-host log name"`
	StderrFmt   string `flag:"logs.stderr,format of stderr per-host log name"`
}

func main() {
	log.SetFlags(0)
	conf := Config{
		Concurrency: 100,
		Login:       "root",
		Port:        22,
		StdoutFmt:   "/tmp/%s.stdout",
		StderrFmt:   "/tmp/%s.stderr",
	}
	autoflags.Define(&conf)
	flag.Parse()
	if conf.Concurrency < 1 {
		conf.Concurrency = 1
	}
	hosts := flag.Args()
	if len(hosts) == 0 || len(conf.Command) == 0 {
		flag.Usage()
		return
	}
	if conf.DumpFiles {
		if conf.StdoutFmt == conf.StderrFmt {
			log.Fatal("file format for stdout and stderr should differ")
		}
		if !strings.Contains(conf.StdoutFmt, `%s`) {
			log.Fatal("invalid format for stdout log name")
		}
		if !strings.Contains(conf.StderrFmt, `%s`) {
			log.Fatal("invalid format for stderr log name")
		}
	}
	var sshAgent agent.Agent
	agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		log.Fatal(err)
	}
	sshAgent = agent.NewClient(agentConn)
	defer agentConn.Close()

	signers, err := sshAgent.Signers()
	if err != nil {
		log.Fatal(err)
	}

	config := &ssh.ClientConfig{
		User: conf.Login,
		Auth: []ssh.AuthMethod{ssh.PublicKeys(signers...)},
	}

	var wg sync.WaitGroup

	stdout := make(chan string, 10)
	stderr := make(chan string, 10)
	go func() {
		for {
			stdoutPrefix, stderrPrefix := ".", "E"
			if isTerminal(os.Stdout) {
				stdoutPrefix = chalk.Green.Color(stdoutPrefix)
			}
			if isTerminal(os.Stderr) {
				stderrPrefix = chalk.Red.Color(stderrPrefix)
			}
			select {
			case item := <-stdout:
				fmt.Fprintln(os.Stdout, stdoutPrefix, item)
			case item := <-stderr:
				fmt.Fprintln(os.Stderr, stderrPrefix, item)
			}
		}
	}()

	limit := make(chan struct{}, conf.Concurrency)

	uniqueHosts := make(map[string]struct{})
	for _, host := range hosts {
		uniqueHosts[host] = struct{}{}
	}

	for host := range uniqueHosts {
		limit <- struct{}{}
		wg.Add(1)
		go func(host string) {
			defer wg.Done()
			defer func() { <-limit }()
			switch err := RemoteCommand(host, conf, config, stdout, stderr); {
			case err == nil && conf.DumpFiles:
				fmt.Println(host, "processed")
			case err == nil:
			default:
				log.Println(host, err)
			}
		}(host)
	}
	wg.Wait()
}

func RemoteCommand(host string, conf Config, config *ssh.ClientConfig, stdout, stderr chan<- string) error {
	addr := host
	if !strings.ContainsRune(addr, ':') {
		addr = fmt.Sprintf("%s:%d", addr, conf.Port)
	}
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	if conf.StdinFile != "" {
		f, err := os.Open(conf.StdinFile)
		if err != nil {
			return err
		}
		defer f.Close()
		st, err := f.Stat()
		if err != nil {
			return err
		}
		if !st.Mode().IsRegular() {
			return errors.New("file passed to stdin is not a regular file")
		}
		session.Stdin = f
	}

	switch {
	case conf.DumpFiles:
		stdoutLog, err := os.Create(filepath.Clean(fmt.Sprintf(conf.StdoutFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stdoutLog)
		session.Stdout = stdoutLog
		stderrLog, err := os.Create(filepath.Clean(fmt.Sprintf(conf.StderrFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stderrLog)
		session.Stderr = stderrLog
	default:
		stdoutPipe, err := session.StdoutPipe()
		if err != nil {
			return err
		}
		stderrPipe, err := session.StderrPipe()
		if err != nil {
			return err
		}

		go pipeFeeder(host, stdoutPipe, stdout)
		go pipeFeeder(host, stderrPipe, stderr)
	}

	return session.Run(conf.Command)
}

func pipeFeeder(prefix string, pipe io.Reader, sink chan<- string) {
	scanner := bufio.NewScanner(pipe)
	for scanner.Scan() {
		sink <- prefix + "\t" + scanner.Text()
	}
	if err := scanner.Err(); err != nil {
		log.Print("string scanner error", err)
	}
}

func closeAndRemoveIfAt0(f *os.File) {
	if n, err := f.Seek(0, os.SEEK_CUR); err == nil && n == 0 {
		os.Remove(f.Name())
	}
	f.Close()
}

func isTerminal(f *os.File) bool {
	st, err := f.Stat()
	if err != nil {
		return false
	}
	return st.Mode()&os.ModeDevice != 0
}
<commit_msg>Support remote connections via socks5 proxy<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/artyom/autoflags"
	"github.com/ttacon/chalk"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/net/proxy"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] ip1 ip2 ...\n", os.Args[0])
		flag.PrintDefaults()
	}
}

type Config struct {
	Concurrency int    `flag:"n,concurrent ssh sessions"`
	Command     string `flag:"cmd,command to run"`
	Login       string `flag:"l,login to use"`
	Port        int    `flag:"p,default port"`
	StdinFile   string `flag:"stdin,REGULAR (no piping!) file to pass to stdin of remote command"`
	DumpFiles   bool   `flag:"logs,save stdout/stderr to separate per-host logs"`
	StdoutFmt   string `flag:"logs.stdout,format of stdout per-host log name"`
	StderrFmt   string `flag:"logs.stderr,format of stderr per-host log name"`
}

func main() {
	log.SetFlags(0)
	conf := Config{
		Concurrency: 100,
		Login:       "root",
		Port:        22,
		StdoutFmt:   "/tmp/%s.stdout",
		StderrFmt:   "/tmp/%s.stderr",
	}
	autoflags.Define(&conf)
	flag.Parse()
	if conf.Concurrency < 1 {
		conf.Concurrency = 1
	}
	hosts := flag.Args()
	if len(hosts) == 0 || len(conf.Command) == 0 {
		flag.Usage()
		return
	}
	if conf.DumpFiles {
		if conf.StdoutFmt == conf.StderrFmt {
			log.Fatal("file format for stdout and stderr should differ")
		}
		if !strings.Contains(conf.StdoutFmt, `%s`) {
			log.Fatal("invalid format for stdout log name")
		}
		if !strings.Contains(conf.StderrFmt, `%s`) {
			log.Fatal("invalid format for stderr log name")
		}
	}
	var sshAgent agent.Agent
	agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		log.Fatal(err)
	}
	sshAgent = agent.NewClient(agentConn)
	defer agentConn.Close()

	signers, err := sshAgent.Signers()
	if err != nil {
		log.Fatal(err)
	}

	config := &ssh.ClientConfig{
		User: conf.Login,
		Auth: []ssh.AuthMethod{ssh.PublicKeys(signers...)},
	}

	var wg sync.WaitGroup

	stdout := make(chan string, 10)
	stderr := make(chan string, 10)
	go func() {
		for {
			stdoutPrefix, stderrPrefix := ".", "E"
			if isTerminal(os.Stdout) {
				stdoutPrefix = chalk.Green.Color(stdoutPrefix)
			}
			if isTerminal(os.Stderr) {
				stderrPrefix = chalk.Red.Color(stderrPrefix)
			}
			select {
			case item := <-stdout:
				fmt.Fprintln(os.Stdout, stdoutPrefix, item)
			case item := <-stderr:
				fmt.Fprintln(os.Stderr, stderrPrefix, item)
			}
		}
	}()

	limit := make(chan struct{}, conf.Concurrency)

	uniqueHosts := make(map[string]struct{})
	for _, host := range hosts {
		uniqueHosts[host] = struct{}{}
	}

	for host := range uniqueHosts {
		limit <- struct{}{}
		wg.Add(1)
		go func(host string) {
			defer wg.Done()
			defer func() { <-limit }()
			switch err := RemoteCommand(host, conf, config, stdout, stderr); {
			case err == nil && conf.DumpFiles:
				fmt.Println(host, "processed")
			case err == nil:
			default:
				log.Println(host, err)
			}
		}(host)
	}
	wg.Wait()
}

func RemoteCommand(host string, conf Config, config *ssh.ClientConfig, stdout, stderr chan<- string) error {
	addr := host
	if !strings.ContainsRune(addr, ':') {
		addr = fmt.Sprintf("%s:%d", addr, conf.Port)
	}
	client, err := sshDial("tcp", addr, config)
	if err != nil {
		return err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	if conf.StdinFile != "" {
		f, err := os.Open(conf.StdinFile)
		if err != nil {
			return err
		}
		defer f.Close()
		st, err := f.Stat()
		if err != nil {
			return err
		}
		if !st.Mode().IsRegular() {
			return errors.New("file passed to stdin is not a regular file")
		}
		session.Stdin = f
	}

	switch {
	case conf.DumpFiles:
		stdoutLog, err := os.Create(filepath.Clean(fmt.Sprintf(conf.StdoutFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stdoutLog)
		session.Stdout = stdoutLog
		stderrLog, err := os.Create(filepath.Clean(fmt.Sprintf(conf.StderrFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stderrLog)
		session.Stderr = stderrLog
	default:
		stdoutPipe, err := session.StdoutPipe()
		if err != nil {
			return err
		}
		stderrPipe, err := session.StderrPipe()
		if err != nil {
			return err
		}

		go pipeFeeder(host, stdoutPipe, stdout)
		go pipeFeeder(host, stderrPipe, stderr)
	}

	return session.Run(conf.Command)
}

func pipeFeeder(prefix string, pipe io.Reader, sink chan<- string) {
	scanner := bufio.NewScanner(pipe)
	for scanner.Scan() {
		sink <- prefix + "\t" + scanner.Text()
	}
	if err := scanner.Err(); err != nil {
		log.Print("string scanner error", err)
	}
}

func closeAndRemoveIfAt0(f *os.File) {
	if n, err := f.Seek(0, os.SEEK_CUR); err == nil && n == 0 {
		os.Remove(f.Name())
	}
	f.Close()
}

func isTerminal(f *os.File) bool {
	st, err := f.Stat()
	if err != nil {
		return false
	}
	return st.Mode()&os.ModeDevice != 0
}

func sshDial(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	conn, err := proxyDialer.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		return nil, err
	}
	return ssh.NewClient(c, chans, reqs), nil
}

var proxyDialer = proxy.FromEnvironment()
<|endoftext|>


<commit_before>package citrixadc

import (
	"github.com/chiradeep/go-nitro/config/cs"
	"github.com/chiradeep/go-nitro/netscaler"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"

	"fmt"
	"log"
)

func resourceCitrixAdcCspolicy() *schema.Resource {
	return &schema.Resource{
		SchemaVersion: 1,
		Create:        createCspolicyFunc,
		Read:          readCspolicyFunc,
		Update:        updateCspolicyFunc,
		Delete:        deleteCspolicyFunc,
		Schema: map[string]*schema.Schema{
			"action": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"domain": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"logaction": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"newname": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"policyname": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"rule": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"url": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"csvserver": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"targetlbvserver": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"forcenew_id_set": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				ForceNew: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},
			"priority": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
		},
	}
}

func createCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In createCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client

	csvserver := d.Get("csvserver").(string)
	targetlbvserver, lbok := d.GetOk("targetlbvserver")
	priority, pok := d.GetOk("priority")
	action, aok := d.GetOk("action")
	_, dok := d.GetOk("domain")
	_, uok := d.GetOk("url")
	_, rok := d.GetOk("rule")

	if lbok && rok && !pok {
		return fmt.Errorf("[ERROR] netscaler-provider: Priority needs to be specified if target lb vserver and rule is specified")
	}
	if !lbok && !pok {
		return fmt.Errorf("[ERROR] netscaler-provider: Priority needs to be specified if target lb vserver is not specified")
	}
	if !lbok && !aok {
		return fmt.Errorf("[ERROR] netscaler-provider: Action or targetlbvserver needs to be specified")
	}
	if aok {
		actionExists := client.ResourceExists(netscaler.Csaction.Type(), action.(string))
		if !actionExists {
			return fmt.Errorf("[ERROR] netscaler-provider: Specified Action %s does not exist", action.(string))
		}
		if !rok {
			return fmt.Errorf("[ERROR] netscaler-provider: Action  %s specified without rule", action.(string))
		}
		if dok || uok {
			return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify url or domain when action  %s is specified", action.(string))
		}
	}
	if uok && dok {
		return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify both url and domain ")
	}
	if rok && (uok || dok) {
		return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify both rule and domain or url ")
	}
	if (uok || dok) && pok {
		return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify both priority and domain or url ")
	}

	var cspolicyName string
	if v, ok := d.GetOk("policyname"); ok {
		cspolicyName = v.(string)
	} else {
		cspolicyName = resource.PrefixedUniqueId("tf-cspolicy-")
		d.Set("name", cspolicyName)
	}
	cspolicy := cs.Cspolicy{
		Policyname: d.Get("policyname").(string),
		Action:     d.Get("action").(string),
		Domain:     d.Get("domain").(string),
		Logaction:  d.Get("logaction").(string),
		Rule:       d.Get("rule").(string),
		Url:        d.Get("url").(string),
	}

	_, err := client.AddResource(netscaler.Cspolicy.Type(), cspolicyName, &cspolicy)
	if err != nil {
		return err
	}

	binding := cs.Csvservercspolicybinding{
		Name:            csvserver,
		Policyname:      cspolicyName,
		Targetlbvserver: targetlbvserver.(string),
		Priority:        priority.(int),
	}

	if !lbok {
		binding = cs.Csvservercspolicybinding{
			Name:       csvserver,
			Policyname: cspolicyName,
			Priority:   priority.(int),
		}
	}

	err = client.BindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, &binding)
	if err != nil {
		d.SetId("")
		err2 := client.DeleteResource(netscaler.Cspolicy.Type(), cspolicyName)
		if err2 != nil {
			return fmt.Errorf("[ERROR] netscaler-provider:  Failed to undo add cspolicy after bind cspolicy %s to Csvserver failed err=%v", cspolicyName, err2)
		}
		return fmt.Errorf("[ERROR] netscaler-provider:  Failed to bind cspolicy %s to Csvserver, err=%v", cspolicyName, err)
	}
	d.SetId(cspolicyName)
	err = readCspolicyFunc(d, meta)
	if err != nil {
		log.Printf("[ERROR] netscaler-provider:  ?? we just created this cspolicy but we can't read it ?? %s", cspolicyName)
		return nil
	}
	return nil
}

func readCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In readCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client
	cspolicyName := d.Id()
	log.Printf("[DEBUG] netscaler-provider:  Reading cspolicy state %s", cspolicyName)
	data, err := client.FindResource(netscaler.Cspolicy.Type(), cspolicyName)
	if err != nil {
		log.Printf("[WARN] netscaler-provider:  Clearing cspolicy state %s", cspolicyName)
		d.SetId("")
		return nil
	}
	d.Set("name", data["name"])
	d.Set("action", data["action"])
	d.Set("domain", data["domain"])
	d.Set("logaction", data["logaction"])
	d.Set("rule", data["rule"])
	d.Set("url", data["url"])

	//read the csvserver binding and update
	bindings, err := client.FindAllBoundResources(netscaler.Cspolicy.Type(), cspolicyName, netscaler.Csvserver.Type())
	if err != nil {
		log.Printf("[WARN] netscaler-provider: cspolicy binding to csvserver error %s", cspolicyName)
		return nil
	}
	var boundCsvserver string
	for _, binding := range bindings {
		log.Printf("[TRACE] netscaler-provider: csvserver_cspolicy binding %v", binding)
		csv, ok := binding["domain"]
		if ok {
			boundCsvserver = csv.(string)
			break
		}
	}
	log.Printf("[TRACE] netscaler-provider: boundCsvserver %v", boundCsvserver)
	d.Set("csvserver", boundCsvserver)

	return nil

}

func updateCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In updateCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client
	cspolicyName := d.Get("policyname").(string)
	csvserver := d.Get("csvserver").(string)

	cspolicy := cs.Cspolicy{
		Policyname: d.Get("policyname").(string),
	}
	hasChange := false
	lbvserverChanged := false
	priorityChanged := false

	if d.HasChange("action") {
		log.Printf("[DEBUG] netscaler-provider: Action has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Action = d.Get("action").(string)
		hasChange = true
	}
	if d.HasChange("domain") {
		log.Printf("[DEBUG] netscaler-provider: Domain has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Domain = d.Get("domain").(string)
		hasChange = true
	}
	if d.HasChange("logaction") {
		log.Printf("[DEBUG] netscaler-provider: Logaction has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Logaction = d.Get("logaction").(string)
		hasChange = true
	}
	if d.HasChange("newname") {
		log.Printf("[DEBUG] netscaler-provider: Newname has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Newname = d.Get("newname").(string)
		hasChange = true
	}
	if d.HasChange("rule") {
		log.Printf("[DEBUG] netscaler-provider: Rule has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Rule = d.Get("rule").(string)
		hasChange = true
	}
	if d.HasChange("url") {
		log.Printf("[DEBUG] netscaler-provider: Url has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Url = d.Get("url").(string)
		hasChange = true
	}
	if d.HasChange("priority") {
		log.Printf("[DEBUG] netscaler-provider: Priority has changed for cspolicy %s, starting update", cspolicyName)
		priorityChanged = true
	}

	if d.HasChange("targetlbvserver") {
		log.Printf("[DEBUG] netscaler-provider: targetlbvserver has changed for cspolicy %s, starting update", cspolicyName)
		lbvserverChanged = true
	}

	if lbvserverChanged || priorityChanged {
		//Binding has to be updated
		//First we unbind from cs vserver
		err := client.UnbindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, "policyname")
		if err != nil {
			return fmt.Errorf("[ERROR] netscaler-provider: Error unbinding cspolicy from csvserver %s", cspolicyName)
		}
		log.Printf("[DEBUG] netscaler-provider: cspolicy has been unbound from csvserver for cspolicy %s ", cspolicyName)
	}

	if hasChange {
		_, err := client.UpdateResource(netscaler.Cspolicy.Type(), cspolicyName, &cspolicy)
		if err != nil {
			return fmt.Errorf("[ERROR] netscaler-provider: Error updating cspolicy %s", cspolicyName)
		}
		log.Printf("[DEBUG] netscaler-provider: cspolicy has been updated  cspolicy %s ", cspolicyName)
	}

	if lbvserverChanged || priorityChanged {
		//Binding has to be updated
		//rebind
		targetlbvserver, lbok := d.GetOk("targetlbvserver")
		priority, pok := d.GetOk("priority")

		if !pok && lbok {
			return fmt.Errorf("[ERROR] netscaler-provider: Need to specify priority if lbvserver is specified")
		}

		binding := cs.Csvservercspolicybinding{
			Name:            csvserver,
			Policyname:      cspolicyName,
			Targetlbvserver: targetlbvserver.(string),
			Priority:        priority.(int),
		}
		err := client.BindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, &binding)
		if err != nil {
			return fmt.Errorf("[ERROR] netscaler-provider: Failed to bind new cspolicy to Csvserver")
		}
		log.Printf("[DEBUG] netscaler-provider: cspolicy has been bound to csvserver  cspolicy %s csvserver %s", cspolicyName, csvserver)
	}

	return readCspolicyFunc(d, meta)
}

func deleteCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In deleteCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client
	cspolicyName := d.Id()
	csvserver := d.Get("csvserver").(string)

	//First we unbind from cs vserver if necessary
	err := client.UnbindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, "policyname")
	if err != nil {
		return fmt.Errorf("[ERROR] netscaler-provider: Error unbinding cspolicy \"%s\" from csvserver \"%v\": %v ", cspolicyName, csvserver, err.Error())
	}
	err = client.DeleteResource(netscaler.Cspolicy.Type(), cspolicyName)
	if err != nil {
		return fmt.Errorf("[ERROR] netscaler-provider: Error  deleting cspolicy %s", cspolicyName)
	}

	d.SetId("")

	return nil
}
<commit_msg>Implement import operation for resource cspolicy<commit_after>package citrixadc

import (
	"github.com/chiradeep/go-nitro/config/cs"
	"github.com/chiradeep/go-nitro/netscaler"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"

	"fmt"
	"log"
)

func resourceCitrixAdcCspolicy() *schema.Resource {
	return &schema.Resource{
		SchemaVersion: 1,
		Create:        createCspolicyFunc,
		Read:          readCspolicyFunc,
		Update:        updateCspolicyFunc,
		Delete:        deleteCspolicyFunc,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},
		Schema: map[string]*schema.Schema{
			"action": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"domain": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"logaction": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"newname": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"policyname": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"rule": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"url": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"csvserver": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"targetlbvserver": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"forcenew_id_set": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				ForceNew: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},
			"priority": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
		},
	}
}

func createCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In createCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client

	csvserver := d.Get("csvserver").(string)
	targetlbvserver, lbok := d.GetOk("targetlbvserver")
	priority, pok := d.GetOk("priority")
	action, aok := d.GetOk("action")
	_, dok := d.GetOk("domain")
	_, uok := d.GetOk("url")
	_, rok := d.GetOk("rule")

	if lbok && rok && !pok {
		return fmt.Errorf("[ERROR] netscaler-provider: Priority needs to be specified if target lb vserver and rule is specified")
	}
	if !lbok && !pok {
		return fmt.Errorf("[ERROR] netscaler-provider: Priority needs to be specified if target lb vserver is not specified")
	}
	if !lbok && !aok {
		return fmt.Errorf("[ERROR] netscaler-provider: Action or targetlbvserver needs to be specified")
	}
	if aok {
		actionExists := client.ResourceExists(netscaler.Csaction.Type(), action.(string))
		if !actionExists {
			return fmt.Errorf("[ERROR] netscaler-provider: Specified Action %s does not exist", action.(string))
		}
		if !rok {
			return fmt.Errorf("[ERROR] netscaler-provider: Action  %s specified without rule", action.(string))
		}
		if dok || uok {
			return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify url or domain when action  %s is specified", action.(string))
		}
	}
	if uok && dok {
		return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify both url and domain ")
	}
	if rok && (uok || dok) {
		return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify both rule and domain or url ")
	}
	if (uok || dok) && pok {
		return fmt.Errorf("[ERROR] netscaler-provider: Cannot specify both priority and domain or url ")
	}

	var cspolicyName string
	if v, ok := d.GetOk("policyname"); ok {
		cspolicyName = v.(string)
	} else {
		cspolicyName = resource.PrefixedUniqueId("tf-cspolicy-")
		d.Set("name", cspolicyName)
	}
	cspolicy := cs.Cspolicy{
		Policyname: d.Get("policyname").(string),
		Action:     d.Get("action").(string),
		Domain:     d.Get("domain").(string),
		Logaction:  d.Get("logaction").(string),
		Rule:       d.Get("rule").(string),
		Url:        d.Get("url").(string),
	}

	_, err := client.AddResource(netscaler.Cspolicy.Type(), cspolicyName, &cspolicy)
	if err != nil {
		return err
	}

	binding := cs.Csvservercspolicybinding{
		Name:            csvserver,
		Policyname:      cspolicyName,
		Targetlbvserver: targetlbvserver.(string),
		Priority:        priority.(int),
	}

	if !lbok {
		binding = cs.Csvservercspolicybinding{
			Name:       csvserver,
			Policyname: cspolicyName,
			Priority:   priority.(int),
		}
	}

	err = client.BindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, &binding)
	if err != nil {
		d.SetId("")
		err2 := client.DeleteResource(netscaler.Cspolicy.Type(), cspolicyName)
		if err2 != nil {
			return fmt.Errorf("[ERROR] netscaler-provider:  Failed to undo add cspolicy after bind cspolicy %s to Csvserver failed err=%v", cspolicyName, err2)
		}
		return fmt.Errorf("[ERROR] netscaler-provider:  Failed to bind cspolicy %s to Csvserver, err=%v", cspolicyName, err)
	}
	d.SetId(cspolicyName)
	err = readCspolicyFunc(d, meta)
	if err != nil {
		log.Printf("[ERROR] netscaler-provider:  ?? we just created this cspolicy but we can't read it ?? %s", cspolicyName)
		return nil
	}
	return nil
}

func readCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In readCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client
	cspolicyName := d.Id()
	log.Printf("[DEBUG] netscaler-provider:  Reading cspolicy state %s", cspolicyName)
	data, err := client.FindResource(netscaler.Cspolicy.Type(), cspolicyName)
	if err != nil {
		log.Printf("[WARN] netscaler-provider:  Clearing cspolicy state %s", cspolicyName)
		d.SetId("")
		return nil
	}
	d.Set("name", data["name"])
	d.Set("action", data["action"])
	d.Set("domain", data["domain"])
	d.Set("logaction", data["logaction"])
	d.Set("rule", data["rule"])
	d.Set("url", data["url"])

	//read the csvserver binding and update
	bindings, err := client.FindAllBoundResources(netscaler.Cspolicy.Type(), cspolicyName, netscaler.Csvserver.Type())
	if err != nil {
		log.Printf("[WARN] netscaler-provider: cspolicy binding to csvserver error %s", cspolicyName)
		return nil
	}
	var boundCsvserver string
	for _, binding := range bindings {
		log.Printf("[TRACE] netscaler-provider: csvserver_cspolicy binding %v", binding)
		csv, ok := binding["domain"]
		if ok {
			boundCsvserver = csv.(string)
			break
		}
	}
	log.Printf("[TRACE] netscaler-provider: boundCsvserver %v", boundCsvserver)
	d.Set("csvserver", boundCsvserver)

	return nil

}

func updateCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In updateCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client
	cspolicyName := d.Get("policyname").(string)
	csvserver := d.Get("csvserver").(string)

	cspolicy := cs.Cspolicy{
		Policyname: d.Get("policyname").(string),
	}
	hasChange := false
	lbvserverChanged := false
	priorityChanged := false

	if d.HasChange("action") {
		log.Printf("[DEBUG] netscaler-provider: Action has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Action = d.Get("action").(string)
		hasChange = true
	}
	if d.HasChange("domain") {
		log.Printf("[DEBUG] netscaler-provider: Domain has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Domain = d.Get("domain").(string)
		hasChange = true
	}
	if d.HasChange("logaction") {
		log.Printf("[DEBUG] netscaler-provider: Logaction has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Logaction = d.Get("logaction").(string)
		hasChange = true
	}
	if d.HasChange("newname") {
		log.Printf("[DEBUG] netscaler-provider: Newname has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Newname = d.Get("newname").(string)
		hasChange = true
	}
	if d.HasChange("rule") {
		log.Printf("[DEBUG] netscaler-provider: Rule has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Rule = d.Get("rule").(string)
		hasChange = true
	}
	if d.HasChange("url") {
		log.Printf("[DEBUG] netscaler-provider: Url has changed for cspolicy %s, starting update", cspolicyName)
		cspolicy.Url = d.Get("url").(string)
		hasChange = true
	}
	if d.HasChange("priority") {
		log.Printf("[DEBUG] netscaler-provider: Priority has changed for cspolicy %s, starting update", cspolicyName)
		priorityChanged = true
	}

	if d.HasChange("targetlbvserver") {
		log.Printf("[DEBUG] netscaler-provider: targetlbvserver has changed for cspolicy %s, starting update", cspolicyName)
		lbvserverChanged = true
	}

	if lbvserverChanged || priorityChanged {
		//Binding has to be updated
		//First we unbind from cs vserver
		err := client.UnbindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, "policyname")
		if err != nil {
			return fmt.Errorf("[ERROR] netscaler-provider: Error unbinding cspolicy from csvserver %s", cspolicyName)
		}
		log.Printf("[DEBUG] netscaler-provider: cspolicy has been unbound from csvserver for cspolicy %s ", cspolicyName)
	}

	if hasChange {
		_, err := client.UpdateResource(netscaler.Cspolicy.Type(), cspolicyName, &cspolicy)
		if err != nil {
			return fmt.Errorf("[ERROR] netscaler-provider: Error updating cspolicy %s", cspolicyName)
		}
		log.Printf("[DEBUG] netscaler-provider: cspolicy has been updated  cspolicy %s ", cspolicyName)
	}

	if lbvserverChanged || priorityChanged {
		//Binding has to be updated
		//rebind
		targetlbvserver, lbok := d.GetOk("targetlbvserver")
		priority, pok := d.GetOk("priority")

		if !pok && lbok {
			return fmt.Errorf("[ERROR] netscaler-provider: Need to specify priority if lbvserver is specified")
		}

		binding := cs.Csvservercspolicybinding{
			Name:            csvserver,
			Policyname:      cspolicyName,
			Targetlbvserver: targetlbvserver.(string),
			Priority:        priority.(int),
		}
		err := client.BindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, &binding)
		if err != nil {
			return fmt.Errorf("[ERROR] netscaler-provider: Failed to bind new cspolicy to Csvserver")
		}
		log.Printf("[DEBUG] netscaler-provider: cspolicy has been bound to csvserver  cspolicy %s csvserver %s", cspolicyName, csvserver)
	}

	return readCspolicyFunc(d, meta)
}

func deleteCspolicyFunc(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] netscaler-provider: In deleteCspolicyFunc")
	client := meta.(*NetScalerNitroClient).client
	cspolicyName := d.Id()
	csvserver := d.Get("csvserver").(string)

	//First we unbind from cs vserver if necessary
	err := client.UnbindResource(netscaler.Csvserver.Type(), csvserver, netscaler.Cspolicy.Type(), cspolicyName, "policyname")
	if err != nil {
		return fmt.Errorf("[ERROR] netscaler-provider: Error unbinding cspolicy \"%s\" from csvserver \"%v\": %v ", cspolicyName, csvserver, err.Error())
	}
	err = client.DeleteResource(netscaler.Cspolicy.Type(), cspolicyName)
	if err != nil {
		return fmt.Errorf("[ERROR] netscaler-provider: Error  deleting cspolicy %s", cspolicyName)
	}

	d.SetId("")

	return nil
}
<|endoftext|>


<commit_before>package check

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	kafka "github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	"github.com/pkg/errors"
	"github.com/samuel/go-zookeeper/zk"
)

func (check *HealthCheck) connect(firstConnection bool, stop <-chan struct{}) error {
	var createHealthTopicIfMissing = firstConnection
	var createReplicationTopicIfMissing = firstConnection
	ticker := time.NewTicker(check.config.retryInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			if err := check.tryConnectOnce(&createHealthTopicIfMissing, &createReplicationTopicIfMissing); err == nil {
				return nil
			}
		case <-stop:
			return errors.New("connect was asked to stop")
		}
	}
}

func (check *HealthCheck) tryConnectOnce(createBrokerTopic, createReplicationTopic *bool) error {
	pauseTime := check.config.retryInterval
	// connect to kafka cluster
	connectString := fmt.Sprintf("localhost:%d", check.config.brokerPort)
	err := check.broker.Dial(connectString, check.brokerConfig())
	if err != nil {
		log.Printf("unable to connect to broker, retrying in %s (%s)", pauseTime.String(), err)
		return err
	}

	metadata, err := check.broker.Metadata()
	if err != nil {
		return errors.Wrap(err, "failure retrieving metadata")
	}

	check.partitionID, err = check.findPartitionID(check.config.topicName, true, createBrokerTopic, metadata)
	if err != nil {
		log.Printf("%s retrying in %s", err.Error(), pauseTime)
		check.broker.Close()
		return err
	}

	check.replicationPartitionID, err = check.findPartitionID(check.config.replicationTopicName, false, createReplicationTopic, metadata)
	if err != nil {
		log.Printf("%s retrying in %s", err.Error(), pauseTime)
		check.broker.Close()
		return err
	}

	consumer, err := check.broker.Consumer(check.consumerConfig())
	if err != nil {
		log.Printf("unable to create consumer, retrying in %s: %s", pauseTime.String(), err)
		check.broker.Close()
		return err
	}

	producer, err := check.broker.Producer(check.producerConfig())
	if err != nil {
		log.Printf("unable to create producer, retrying in %s: %s", pauseTime.String(), err)
		check.broker.Close()
		return err
	}

	check.consumer = consumer
	check.producer = producer
	return nil
}

func (check *HealthCheck) findPartitionID(topicName string, forHealthCheck bool, createIfMissing *bool, metadata *kafka.MetadataResponse) (int32, error) {
	brokerID := int32(check.config.brokerID)

	if !brokerExists(brokerID, metadata) {
		return 0, fmt.Errorf("unable to find broker %d in metadata", brokerID)
	}

	topic, ok := findTopic(topicName, metadata)

	if ok {
		for _, partition := range topic.Partitions {
			if forHealthCheck && partition.Leader != brokerID {
				continue
			}
			if !contains(partition.Replicas, brokerID) {
				continue
			}

			log.Printf(`found partition id %d for broker %d in topic "%s"`, partition.ID, brokerID, topicName)
			return partition.ID, nil
		}
	}

	if *createIfMissing {
		err := check.createTopic(topicName, forHealthCheck)
		if err != nil {
			return 0, errors.Wrapf(err, `unable to create topic "%s"`, topicName)
		}
		log.Printf(`topic "%s" created`, topicName)
		*createIfMissing = false
		return 0, errors.New("topic created, try again")
	}

	if ok {
		return 0, fmt.Errorf(`Unable to find broker's parition in topic "%s" in metadata`, topicName)
	}
	return 0, fmt.Errorf(`Unable to find broker's topic "%s" in metadata`, topicName)
}

func findTopic(name string, metadata *kafka.MetadataResponse) (*kafka.TopicMetadata, bool) {
	for _, topic := range metadata.Topics {
		if topic.Name == name {
			return topic, true
		}
	}

	return nil, false
}

func brokerExists(brokerID int32, metadata *kafka.MetadataResponse) bool {
	for _, broker := range metadata.Brokers {
		if broker.ID() == brokerID {
			return true
		}
	}
	return false
}

func zookeeperEnsembleAndChroot(connectString string) (ensemble []string, chroot string) {
	result := strings.Split(connectString, "/")
	switch len(result) {
	case 1:
		ensemble = strings.Split(result[0], ",")
		chroot = ""
	default:
		ensemble = strings.Split(result[0], ",")
		chroot = "/" + strings.Join(result[1:], "/")
		if strings.HasSuffix(chroot, "/") {
			chroot = chroot[:len(chroot)-1]
		}
	}
	return
}

func (check *HealthCheck) createTopic(name string, forHealthCheck bool) (err error) {
	log.Printf("connecting to ZooKeeper ensemble %s", check.config.zookeeperConnect)
	connectString, chroot := zookeeperEnsembleAndChroot(check.config.zookeeperConnect)
	zkConn := check.zookeeper

	if _, err = zkConn.Connect(connectString, 10*time.Second); err != nil {
		return
	}
	defer zkConn.Close()

	topicPath := chroot + "/config/topics/" + name

	exists := false
	if !forHealthCheck {
		exists, _, err = zkConn.Exists(topicPath)
		if err != nil {
			return
		}
	}

	brokerID := int32(check.config.brokerID)

	if !exists {
		topicConfig := `{"version":1,"config":{"delete.retention.ms":"10000",` +
			`"cleanup.policy":"delete","compression.type":"uncompressed"}}`
		log.Infof(`creating topic "%s" configuration node`, name)

		if err = createZkNode(zkConn, topicPath, topicConfig, false); err != nil {
			return
		}

		partitionAssignment := fmt.Sprintf(`{"version":1,"partitions":{"0":[%d]}}`, brokerID)
		log.Infof(`creating topic "%s" partition assignment node`, name)

		if err = createZkNode(zkConn, chroot+"/brokers/topics/"+name, partitionAssignment, false); err != nil {
			return
		}
	}

	if !forHealthCheck {
		err = maybeExpandReplicationTopic(zkConn, brokerID, check.replicationPartitionID, name, chroot)
	}

	return

}

func maybeExpandReplicationTopic(zk ZkConnection, brokerID, partitionID int32, topicName, chroot string) error {
	topic := ZkTopic{Name: topicName}
	err := zkPartitions(&topic, zk, topicName, chroot)
	if err != nil {
		return errors.Wrap(err, "Unable to determine if replication topic should be expanded")
	}

	replicas, ok := topic.Partitions[strconv.Itoa(int(partitionID))]
	if !ok {
		return fmt.Errorf(`Cannot find partition with ID %d in topic "%s"`, partitionID, topicName)
	}

	if !contains(replicas, brokerID) {
		log.Info("Expanding replication check topic to include broker ", brokerID)
		replicas = append(replicas, brokerID)

		return reassignPartition(zk, partitionID, replicas, topicName, chroot)
	}
	return nil
}

func reassignPartition(zk ZkConnection, partitionID int32, replicas []int32, topicName, chroot string) (err error) {

	repeat := true
	for repeat {
		time.Sleep(1 * time.Second)
		exists, _, rpErr := zk.Exists(chroot + "/admin/reassign_partitions")
		if rpErr != nil {
			log.Warn("Error while checking if reassign_partitions node exists", rpErr)
		}
		repeat = exists || rpErr != nil
	}

	var replicasStr []string
	for _, ID := range replicas {
		replicasStr = append(replicasStr, fmt.Sprintf("%d", ID))
	}

	reassign := fmt.Sprintf(`{"version":1,"partitions":[{"topic":"%s","partition":%d,"replicas":[%s]}]}`,
		topicName, partitionID, strings.Join(replicasStr, ","))

	repeat = true
	for repeat {
		log.Info("Creating reassign partition node")
		err = createZkNode(zk, chroot+"/admin/reassign_partitions", reassign, true)
		if err != nil {
			log.Warn("Error while creating reassignment node", err)
		}
		repeat = err != nil
	}

	return
}

func createZkNode(zookeeper ZkConnection, path string, content string, failIfExists bool) error {
	nodeExists, _, err := zookeeper.Exists(path)

	if err != nil {
		return err
	}

	if nodeExists {
		if failIfExists {
			return fmt.Errorf("node %s cannot be created, exists already", path)
		}
		return nil
	}

	log.Println("creating node", path)
	flags := int32(0) // permanent node.
	acl := zk.WorldACL(zk.PermAll)
	_, err = zookeeper.Create(path, []byte(content), flags, acl)
	return err
}

func (check *HealthCheck) closeConnection(deleteTopicIfPresent bool) {
	if deleteTopicIfPresent {
		log.Infof("connecting to ZooKeeper ensemble %s", check.config.zookeeperConnect)
		connectString, chroot := zookeeperEnsembleAndChroot(check.config.zookeeperConnect)

		zkConn := check.zookeeper
		_, err := zkConn.Connect(connectString, 10*time.Second)
		if err != nil {
			return
		}
		defer zkConn.Close()

		check.deleteTopic(zkConn, chroot, check.config.topicName, check.partitionID)
		check.deleteTopic(zkConn, chroot, check.config.replicationTopicName, check.replicationPartitionID)
	}
	check.broker.Close()
}

func (check *HealthCheck) deleteTopic(zkConn ZkConnection, chroot, name string, partitionID int32) error {
	topic := ZkTopic{Name: name}
	err := zkPartitions(&topic, zkConn, name, chroot)
	if err != nil {
		return err
	}

	replicas, ok := topic.Partitions[strconv.Itoa(int(partitionID))]
	if !ok {
		return fmt.Errorf(`Cannot find partition with ID %d in topic "%s"`, partitionID, name)
	}

	brokerID := int32(check.config.brokerID)
	if len(replicas) > 1 {
		log.Info("Shrinking replication check topic to exclude broker ", brokerID)
		replicas = delAll(replicas, brokerID)
		return reassignPartition(zkConn, partitionID, replicas, name, chroot)
	}

	delTopicPath := chroot + "/admin/delete_topics/" + name

	err = createZkNode(check.zookeeper, delTopicPath, "", true)
	if err != nil {
		return err
	}
	return check.waitForTopicDeletion(delTopicPath)
}

func (check *HealthCheck) waitForTopicDeletion(topicPath string) error {
	for {
		exists, _, err := check.zookeeper.Exists(topicPath)
		if err != nil {
			return err
		}
		if !exists {
			return nil
		}
		time.Sleep(check.config.retryInterval)
	}
}

func (check *HealthCheck) reconnect(stop <-chan struct{}) error {
	check.closeConnection(false)
	return check.connect(false, stop)
}
<commit_msg>set createIfMissing<commit_after>package check

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	kafka "github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	"github.com/pkg/errors"
	"github.com/samuel/go-zookeeper/zk"
)

func (check *HealthCheck) connect(firstConnection bool, stop <-chan struct{}) error {
	var createHealthTopicIfMissing = firstConnection
	var createReplicationTopicIfMissing = firstConnection
	ticker := time.NewTicker(check.config.retryInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			if err := check.tryConnectOnce(&createHealthTopicIfMissing, &createReplicationTopicIfMissing); err == nil {
				return nil
			}
		case <-stop:
			return errors.New("connect was asked to stop")
		}
	}
}

func (check *HealthCheck) tryConnectOnce(createBrokerTopic, createReplicationTopic *bool) error {
	pauseTime := check.config.retryInterval
	// connect to kafka cluster
	connectString := fmt.Sprintf("localhost:%d", check.config.brokerPort)
	err := check.broker.Dial(connectString, check.brokerConfig())
	if err != nil {
		log.Printf("unable to connect to broker, retrying in %s (%s)", pauseTime.String(), err)
		return err
	}

	metadata, err := check.broker.Metadata()
	if err != nil {
		return errors.Wrap(err, "failure retrieving metadata")
	}

	check.partitionID, err = check.findPartitionID(check.config.topicName, true, createBrokerTopic, metadata)
	if err != nil {
		log.Printf("%s retrying in %s", err.Error(), pauseTime)
		check.broker.Close()
		return err
	}

	check.replicationPartitionID, err = check.findPartitionID(check.config.replicationTopicName, false, createReplicationTopic, metadata)
	if err != nil {
		log.Printf("%s retrying in %s", err.Error(), pauseTime)
		check.broker.Close()
		return err
	}

	consumer, err := check.broker.Consumer(check.consumerConfig())
	if err != nil {
		log.Printf("unable to create consumer, retrying in %s: %s", pauseTime.String(), err)
		check.broker.Close()
		return err
	}

	producer, err := check.broker.Producer(check.producerConfig())
	if err != nil {
		log.Printf("unable to create producer, retrying in %s: %s", pauseTime.String(), err)
		check.broker.Close()
		return err
	}

	check.consumer = consumer
	check.producer = producer
	return nil
}

func (check *HealthCheck) findPartitionID(topicName string, forHealthCheck bool, createIfMissing *bool, metadata *kafka.MetadataResponse) (int32, error) {
	brokerID := int32(check.config.brokerID)

	if !brokerExists(brokerID, metadata) {
		return 0, fmt.Errorf("unable to find broker %d in metadata", brokerID)
	}

	topic, ok := findTopic(topicName, metadata)

	if ok {
		for _, partition := range topic.Partitions {
			if forHealthCheck && partition.Leader != brokerID {
				continue
			}
			if !contains(partition.Replicas, brokerID) {
				continue
			}

			log.Printf(`found partition id %d for broker %d in topic "%s"`, partition.ID, brokerID, topicName)
			return partition.ID, nil
		}
	}

	if *createIfMissing {
		err := check.createTopic(topicName, forHealthCheck)
		if err != nil {
			return 0, errors.Wrapf(err, `unable to create topic "%s"`, topicName)
		}
		log.Printf(`topic "%s" created`, topicName)
		*createIfMissing = false
		return 0, errors.New("topic created, try again")
	}

	if ok {
		return 0, fmt.Errorf(`Unable to find broker's parition in topic "%s" in metadata`, topicName)
	}

	*createIfMissing = true
	return 0, fmt.Errorf(`Unable to find broker's topic "%s" in metadata`, topicName)
}

func findTopic(name string, metadata *kafka.MetadataResponse) (*kafka.TopicMetadata, bool) {
	for _, topic := range metadata.Topics {
		if topic.Name == name {
			return topic, true
		}
	}

	return nil, false
}

func brokerExists(brokerID int32, metadata *kafka.MetadataResponse) bool {
	for _, broker := range metadata.Brokers {
		if broker.ID() == brokerID {
			return true
		}
	}
	return false
}

func zookeeperEnsembleAndChroot(connectString string) (ensemble []string, chroot string) {
	result := strings.Split(connectString, "/")
	switch len(result) {
	case 1:
		ensemble = strings.Split(result[0], ",")
		chroot = ""
	default:
		ensemble = strings.Split(result[0], ",")
		chroot = "/" + strings.Join(result[1:], "/")
		if strings.HasSuffix(chroot, "/") {
			chroot = chroot[:len(chroot)-1]
		}
	}
	return
}

func (check *HealthCheck) createTopic(name string, forHealthCheck bool) (err error) {
	log.Printf("connecting to ZooKeeper ensemble %s", check.config.zookeeperConnect)
	connectString, chroot := zookeeperEnsembleAndChroot(check.config.zookeeperConnect)
	zkConn := check.zookeeper

	if _, err = zkConn.Connect(connectString, 10*time.Second); err != nil {
		return
	}
	defer zkConn.Close()

	topicPath := chroot + "/config/topics/" + name

	exists := false
	if !forHealthCheck {
		exists, _, err = zkConn.Exists(topicPath)
		if err != nil {
			return
		}
	}

	brokerID := int32(check.config.brokerID)

	if !exists {
		topicConfig := `{"version":1,"config":{"delete.retention.ms":"10000",` +
			`"cleanup.policy":"delete","compression.type":"uncompressed"}}`
		log.Infof(`creating topic "%s" configuration node`, name)

		if err = createZkNode(zkConn, topicPath, topicConfig, false); err != nil {
			return
		}

		partitionAssignment := fmt.Sprintf(`{"version":1,"partitions":{"0":[%d]}}`, brokerID)
		log.Infof(`creating topic "%s" partition assignment node`, name)

		if err = createZkNode(zkConn, chroot+"/brokers/topics/"+name, partitionAssignment, false); err != nil {
			return
		}
	}

	if !forHealthCheck {
		err = maybeExpandReplicationTopic(zkConn, brokerID, check.replicationPartitionID, name, chroot)
	}

	return

}

func maybeExpandReplicationTopic(zk ZkConnection, brokerID, partitionID int32, topicName, chroot string) error {
	topic := ZkTopic{Name: topicName}
	err := zkPartitions(&topic, zk, topicName, chroot)
	if err != nil {
		return errors.Wrap(err, "Unable to determine if replication topic should be expanded")
	}

	replicas, ok := topic.Partitions[strconv.Itoa(int(partitionID))]
	if !ok {
		return fmt.Errorf(`Cannot find partition with ID %d in topic "%s"`, partitionID, topicName)
	}

	if !contains(replicas, brokerID) {
		log.Info("Expanding replication check topic to include broker ", brokerID)
		replicas = append(replicas, brokerID)

		return reassignPartition(zk, partitionID, replicas, topicName, chroot)
	}
	return nil
}

func reassignPartition(zk ZkConnection, partitionID int32, replicas []int32, topicName, chroot string) (err error) {

	repeat := true
	for repeat {
		time.Sleep(1 * time.Second)
		exists, _, rpErr := zk.Exists(chroot + "/admin/reassign_partitions")
		if rpErr != nil {
			log.Warn("Error while checking if reassign_partitions node exists", rpErr)
		}
		repeat = exists || rpErr != nil
	}

	var replicasStr []string
	for _, ID := range replicas {
		replicasStr = append(replicasStr, fmt.Sprintf("%d", ID))
	}

	reassign := fmt.Sprintf(`{"version":1,"partitions":[{"topic":"%s","partition":%d,"replicas":[%s]}]}`,
		topicName, partitionID, strings.Join(replicasStr, ","))

	repeat = true
	for repeat {
		log.Info("Creating reassign partition node")
		err = createZkNode(zk, chroot+"/admin/reassign_partitions", reassign, true)
		if err != nil {
			log.Warn("Error while creating reassignment node", err)
		}
		repeat = err != nil
	}

	return
}

func createZkNode(zookeeper ZkConnection, path string, content string, failIfExists bool) error {
	nodeExists, _, err := zookeeper.Exists(path)

	if err != nil {
		return err
	}

	if nodeExists {
		if failIfExists {
			return fmt.Errorf("node %s cannot be created, exists already", path)
		}
		return nil
	}

	log.Println("creating node", path)
	flags := int32(0) // permanent node.
	acl := zk.WorldACL(zk.PermAll)
	_, err = zookeeper.Create(path, []byte(content), flags, acl)
	return err
}

func (check *HealthCheck) closeConnection(deleteTopicIfPresent bool) {
	if deleteTopicIfPresent {
		log.Infof("connecting to ZooKeeper ensemble %s", check.config.zookeeperConnect)
		connectString, chroot := zookeeperEnsembleAndChroot(check.config.zookeeperConnect)

		zkConn := check.zookeeper
		_, err := zkConn.Connect(connectString, 10*time.Second)
		if err != nil {
			return
		}
		defer zkConn.Close()

		check.deleteTopic(zkConn, chroot, check.config.topicName, check.partitionID)
		check.deleteTopic(zkConn, chroot, check.config.replicationTopicName, check.replicationPartitionID)
	}
	check.broker.Close()
}

func (check *HealthCheck) deleteTopic(zkConn ZkConnection, chroot, name string, partitionID int32) error {
	topic := ZkTopic{Name: name}
	err := zkPartitions(&topic, zkConn, name, chroot)
	if err != nil {
		return err
	}

	replicas, ok := topic.Partitions[strconv.Itoa(int(partitionID))]
	if !ok {
		return fmt.Errorf(`Cannot find partition with ID %d in topic "%s"`, partitionID, name)
	}

	brokerID := int32(check.config.brokerID)
	if len(replicas) > 1 {
		log.Info("Shrinking replication check topic to exclude broker ", brokerID)
		replicas = delAll(replicas, brokerID)
		return reassignPartition(zkConn, partitionID, replicas, name, chroot)
	}

	delTopicPath := chroot + "/admin/delete_topics/" + name

	err = createZkNode(check.zookeeper, delTopicPath, "", true)
	if err != nil {
		return err
	}
	return check.waitForTopicDeletion(delTopicPath)
}

func (check *HealthCheck) waitForTopicDeletion(topicPath string) error {
	for {
		exists, _, err := check.zookeeper.Exists(topicPath)
		if err != nil {
			return err
		}
		if !exists {
			return nil
		}
		time.Sleep(check.config.retryInterval)
	}
}

func (check *HealthCheck) reconnect(stop <-chan struct{}) error {
	check.closeConnection(false)
	return check.connect(false, stop)
}
<|endoftext|>


<commit_before>package discovery

import (
	"errors"
	"io"
	"io/ioutil"
	golog "log"
	"net"
	"sync"
	"time"

	"github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/hashicorp/mdns"
	ma "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	manet "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr-net"

	"github.com/ipfs/go-ipfs/p2p/host"
	"github.com/ipfs/go-ipfs/p2p/peer"
	u "github.com/ipfs/go-ipfs/util"
)

var log = u.Logger("mdns")

const ServiceTag = "discovery.ipfs.io"

type Service interface {
	io.Closer
	RegisterNotifee(Notifee)
	UnregisterNotifee(Notifee)
}

type Notifee interface {
	HandlePeerFound(peer.PeerInfo)
}

type mdnsService struct {
	server  *mdns.Server
	service *mdns.MDNSService
	host    host.Host

	lk       sync.Mutex
	notifees []Notifee
	interval time.Duration
}

func getDialableListenAddr(ph host.Host) (*net.TCPAddr, error) {
	for _, addr := range ph.Addrs() {
		if manet.IsIPLoopback(addr) {
			continue
		}
		na, err := manet.ToNetAddr(addr)
		if err != nil {
			continue
		}
		tcp, ok := na.(*net.TCPAddr)
		if ok {
			return tcp, nil
		}
	}
	return nil, errors.New("failed to find good external addr from peerhost")
}

func NewMdnsService(peerhost host.Host, interval time.Duration) (Service, error) {

	// TODO: dont let mdns use logging...
	golog.SetOutput(ioutil.Discard)

	var ipaddrs []net.IP
	port := 4001

	addr, err := getDialableListenAddr(peerhost)
	if err != nil {
		log.Warning(err)
	} else {
		ipaddrs = []net.IP{addr.IP}
		port = addr.Port
	}

	myid := peerhost.ID().Pretty()

	info := []string{myid}
	service, err := mdns.NewMDNSService(myid, ServiceTag, "", "", port, ipaddrs, info)
	if err != nil {
		return nil, err
	}

	// Create the mDNS server, defer shutdown
	server, err := mdns.NewServer(&mdns.Config{Zone: service})
	if err != nil {
		return nil, err
	}

	s := &mdnsService{
		server:   server,
		service:  service,
		host:     peerhost,
		interval: interval,
	}

	go s.pollForEntries()

	return s, nil
}

func (m *mdnsService) Close() error {
	return m.server.Shutdown()
}

func (m *mdnsService) pollForEntries() {
	ticker := time.NewTicker(m.interval)
	for {
		select {
		case <-ticker.C:
			entriesCh := make(chan *mdns.ServiceEntry, 16)
			go func() {
				for entry := range entriesCh {
					m.handleEntry(entry)
				}
			}()

			qp := mdns.QueryParam{}
			qp.Domain = "local"
			qp.Entries = entriesCh
			qp.Service = ServiceTag
			qp.Timeout = time.Second * 5

			err := mdns.Query(&qp)
			if err != nil {
				log.Error("mdns lookup error: ", err)
			}
			close(entriesCh)
		}
	}
}

func (m *mdnsService) handleEntry(e *mdns.ServiceEntry) {
	mpeer, err := peer.IDB58Decode(e.Info)
	if err != nil {
		log.Warning("Error parsing peer ID from mdns entry: ", err)
		return
	}

	if mpeer == m.host.ID() {
		return
	}

	maddr, err := manet.FromNetAddr(&net.TCPAddr{
		IP:   e.AddrV4,
		Port: e.Port,
	})
	if err != nil {
		log.Warning("Error parsing multiaddr from mdns entry: ", err)
		return
	}

	pi := peer.PeerInfo{
		ID:    mpeer,
		Addrs: []ma.Multiaddr{maddr},
	}

	m.lk.Lock()
	for _, n := range m.notifees {
		n.HandlePeerFound(pi)
	}
	m.lk.Unlock()
}

func (m *mdnsService) RegisterNotifee(n Notifee) {
	m.lk.Lock()
	m.notifees = append(m.notifees, n)
	m.lk.Unlock()
}

func (m *mdnsService) UnregisterNotifee(n Notifee) {
	m.lk.Lock()
	found := -1
	for i, notif := range m.notifees {
		if notif == n {
			found = i
			break
		}
	}
	if found != -1 {
		m.notifees = append(m.notifees[:found], m.notifees[found+1:]...)
	}
	m.lk.Unlock()
}
<commit_msg>use ALL THE ADDRESSES!<commit_after>package discovery

import (
	"errors"
	"io"
	"io/ioutil"
	golog "log"
	"net"
	"sync"
	"time"

	"github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/hashicorp/mdns"
	ma "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	manet "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr-net"

	"github.com/ipfs/go-ipfs/p2p/host"
	"github.com/ipfs/go-ipfs/p2p/peer"
	u "github.com/ipfs/go-ipfs/util"
)

var log = u.Logger("mdns")

const ServiceTag = "discovery.ipfs.io"

type Service interface {
	io.Closer
	RegisterNotifee(Notifee)
	UnregisterNotifee(Notifee)
}

type Notifee interface {
	HandlePeerFound(peer.PeerInfo)
}

type mdnsService struct {
	server  *mdns.Server
	service *mdns.MDNSService
	host    host.Host

	lk       sync.Mutex
	notifees []Notifee
	interval time.Duration
}

func getDialableListenAddrs(ph host.Host) ([]*net.TCPAddr, error) {
	var out []*net.TCPAddr
	for _, addr := range ph.Addrs() {
		na, err := manet.ToNetAddr(addr)
		if err != nil {
			continue
		}
		tcp, ok := na.(*net.TCPAddr)
		if ok {
			out = append(out, tcp)
		}
	}
	if len(out) == 0 {
		return nil, errors.New("failed to find good external addr from peerhost")
	}
	return out, nil
}

func NewMdnsService(peerhost host.Host, interval time.Duration) (Service, error) {

	// TODO: dont let mdns use logging...
	golog.SetOutput(ioutil.Discard)

	var ipaddrs []net.IP
	port := 4001

	addrs, err := getDialableListenAddrs(peerhost)
	if err != nil {
		log.Warning(err)
	} else {
		port = addrs[0].Port
		for _, a := range addrs {
			ipaddrs = append(ipaddrs, a.IP)
		}
	}

	myid := peerhost.ID().Pretty()

	info := []string{myid}
	service, err := mdns.NewMDNSService(myid, ServiceTag, "", "", port, ipaddrs, info)
	if err != nil {
		return nil, err
	}

	// Create the mDNS server, defer shutdown
	server, err := mdns.NewServer(&mdns.Config{Zone: service})
	if err != nil {
		return nil, err
	}

	s := &mdnsService{
		server:   server,
		service:  service,
		host:     peerhost,
		interval: interval,
	}

	go s.pollForEntries()

	return s, nil
}

func (m *mdnsService) Close() error {
	return m.server.Shutdown()
}

func (m *mdnsService) pollForEntries() {
	ticker := time.NewTicker(m.interval)
	for {
		select {
		case <-ticker.C:
			entriesCh := make(chan *mdns.ServiceEntry, 16)
			go func() {
				for entry := range entriesCh {
					m.handleEntry(entry)
				}
			}()

			qp := mdns.QueryParam{}
			qp.Domain = "local"
			qp.Entries = entriesCh
			qp.Service = ServiceTag
			qp.Timeout = time.Second * 5

			err := mdns.Query(&qp)
			if err != nil {
				log.Error("mdns lookup error: ", err)
			}
			close(entriesCh)
		}
	}
}

func (m *mdnsService) handleEntry(e *mdns.ServiceEntry) {
	mpeer, err := peer.IDB58Decode(e.Info)
	if err != nil {
		log.Warning("Error parsing peer ID from mdns entry: ", err)
		return
	}

	if mpeer == m.host.ID() {
		return
	}

	maddr, err := manet.FromNetAddr(&net.TCPAddr{
		IP:   e.AddrV4,
		Port: e.Port,
	})
	if err != nil {
		log.Warning("Error parsing multiaddr from mdns entry: ", err)
		return
	}

	pi := peer.PeerInfo{
		ID:    mpeer,
		Addrs: []ma.Multiaddr{maddr},
	}

	m.lk.Lock()
	for _, n := range m.notifees {
		n.HandlePeerFound(pi)
	}
	m.lk.Unlock()
}

func (m *mdnsService) RegisterNotifee(n Notifee) {
	m.lk.Lock()
	m.notifees = append(m.notifees, n)
	m.lk.Unlock()
}

func (m *mdnsService) UnregisterNotifee(n Notifee) {
	m.lk.Lock()
	found := -1
	for i, notif := range m.notifees {
		if notif == n {
			found = i
			break
		}
	}
	if found != -1 {
		m.notifees = append(m.notifees[:found], m.notifees[found+1:]...)
	}
	m.lk.Unlock()
}
<|endoftext|>


<commit_before>package reedsolomon

import (
	"bytes"
	"math/rand"
	"testing"
	"runtime"
	"sync"
)

//------------

func TestEncode(t *testing.T) {
	size := 50000
	r, err := New(10, 3)
	if err != nil {
		t.Fatal(err)
	}
	dp := NewMatrix(13, size)
	rand.Seed(0)
	for s := 0; s < 13; s++ {
		fillRandom(dp[s])
	}
	err = r.Encode(dp)
	if err != nil {
		t.Fatal(err)
	}
	badDP := NewMatrix(13, 100)
	badDP[0] = make([]byte, 1)
	err = r.Encode(badDP)
	if err != ErrShardSize {
		t.Errorf("expected %v, got %v", ErrShardSize, err)
	}
}

// test low, high table work
func TestVerifyEncode(t *testing.T) {
	r, err := New(5, 5)
	if err != nil {
		t.Fatal(err)
	}
	shards := [][]byte{
		{0, 1},
		{4, 5},
		{2, 3},
		{6, 7},
		{8, 9},
		{0, 0},
		{0, 0},
		{0, 0},
		{0, 0},
		{0, 0},
	}
	r.Encode(shards)
	if shards[5][0] != 97 || shards[5][1] != 64 {
		t.Fatal("shard 5 mismatch")
	}
	if shards[6][0] != 173 || shards[6][1] != 3 {
		t.Fatal("shard 6 mismatch")
	}
	if shards[7][0] != 218 || shards[7][1] != 14 {
		t.Fatal("shard 7 mismatch")
	}
	if shards[8][0] != 107 || shards[8][1] != 35 {
		t.Fatal("shard 8 mismatch")
	}
	if shards[9][0] != 110 || shards[9][1] != 177 {
		t.Fatal("shard 9 mismatch")
	}
}

func TestASM(t *testing.T) {
	d := 10
	p := 4
	size := 65 * 1024
	r, err := New(d, p)
	if err != nil {
		t.Fatal(err)
	}
	// asm
	dp := NewMatrix(d+p, size)
	rand.Seed(0)
	for i := 0; i < d; i++ {
		fillRandom(dp[i])
	}
	err = r.Encode(dp)
	if err != nil {
		t.Fatal(err)
	}
	// mulTable
	mDP := NewMatrix(d+p, size)
	for i := 0; i < d; i++ {
		mDP[i] = dp[i]
	}
	err = r.noasmEncode(mDP)
	if err != nil {
		t.Fatal(err)
	}
	for i, asm := range dp {
		if !bytes.Equal(asm, mDP[i]) {
			t.Fatal("verify asm failed, no match noasm version; shards: ", i)
		}
	}
}

func BenchmarkEncode10x4x1M(b *testing.B) {
	benchmarkEncode(b, 10, 4, 1024*1024)
}

func BenchmarkEncode10x4x4M(b *testing.B) {
	benchmarkEncode(b, 10, 4, 4*1024*1024)
}

func BenchmarkEncode10x4x16M(b *testing.B) {
	benchmarkEncode(b, 10, 4, 16*1024*1024)
}

func BenchmarkEncode17x3x1M(b *testing.B) {
	benchmarkEncode(b, 17, 3, 1024*1024)
}

func BenchmarkEncode17x3x4M(b *testing.B) {
	benchmarkEncode(b, 17, 3, 4*1024*1024)
}

func BenchmarkEncode17x3x16M(b *testing.B) {
	benchmarkEncode(b, 17, 3, 16*1024*1024)
}

func benchmarkEncode(b *testing.B, data, parity, size int) {
	r, err := New(data, parity)
	if err != nil {
		b.Fatal(err)
	}
	dp := NewMatrix(data+parity, size)
	rand.Seed(0)
	for i := 0; i < data; i++ {
		fillRandom(dp[i])
	}
	b.SetBytes(int64(size * data))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r.Encode(dp)
	}
}

func fillRandom(p []byte) {
	for i := 0; i < len(p); i += 7 {
		val := rand.Int63()
		for j := 0; i+j < len(p) && j < 7; j++ {
			p[i+j] = byte(val)
			val >>= 8
		}
	}
}

func (r *rs) noasmEncode(dp matrix) error {
	// check args
	if len(dp) != r.shards {
		return ErrTooFewShards
	}
	_, err := checkShardSize(dp)
	if err != nil {
		return err
	}
	// encoding
	input := dp[0:r.data]
	output := dp[r.data:]
	noasmRunner(r.gen, input, output, r.data, r.parity)
	return nil
}

func noasmRunner(gen, input, output matrix, numData, numParity int) {
	for i := 0; i < numData; i++ {
		in := input[i]
		for oi := 0; oi < numParity; oi++ {
			if i == 0 {
				noasmGfVectMul(gen[oi][i], in[:], output[oi][:])
			} else {
				noasmGfVectMulXor(gen[oi][i], in[:], output[oi][:])
			}
		}
	}
}

func noasmGfVectMul(c byte, in, out []byte) {
	mt := mulTable[c]
	for i := 0; i < len(in); i++ {
		out[i] = mt[in[i]]
	}
}

func noasmGfVectMulXor(c byte, in, out []byte) {
	mt := mulTable[c]
	for i := 0; i < len(in); i++ {
		out[i] ^= mt[in[i]]
	}
}

func BenchmarkEncode28x4x16M(b *testing.B) {
	benchmarkEncode(b, 28, 4, 16776168)
}

func BenchmarkEncode14x10x16M(b *testing.B) {
	benchmarkEncode(b, 14, 10, 16776168)
}

func BenchmarkEncode28x4x16M_ConCurrency(b *testing.B) {
	benchmarkEncode_ConCurrency(b, 28, 4, 16776168)
}

func BenchmarkEncode14x10x16M_ConCurrency(b *testing.B) {
	benchmarkEncode_ConCurrency(b, 14, 10, 16776168)
}

func benchmarkEncode_ConCurrency(b *testing.B, data, parity, size int) {
	count := runtime.NumCPU()
	Instances := make([]*rs, count)
	dps := make([]matrix, count)
	for i := 0; i < count; i++ {
		r, err := New(data, parity)
		if err != nil {
			b.Fatal(err)
		}
		Instances[i] = r
	}
	for i := 0; i < count; i++ {
		dps[i] = NewMatrix(data+parity, size)
		rand.Seed(0)
		for j := 0; j < data; j++ {
			fillRandom(dps[i][j])
		}
	}

	runtime.GOMAXPROCS(count)
	b.SetBytes(int64(size * data * count))
	b.ResetTimer()
	var g sync.WaitGroup
	for i := 0; i < count; i++ {
		g.Add(1)
		go func(i int) {
			Instances[i].Encode(dps[i])
			g.Done()
		}(i)
	}
	g.Wait()
}<commit_msg>add test<commit_after>package reedsolomon

import (
	"bytes"
	"math/rand"
	"testing"
	"runtime"
	"sync"
	"fmt"
)

//------------

func TestEncode(t *testing.T) {
	size := 50000
	r, err := New(10, 3)
	if err != nil {
		t.Fatal(err)
	}
	dp := NewMatrix(13, size)
	rand.Seed(0)
	for s := 0; s < 13; s++ {
		fillRandom(dp[s])
	}
	err = r.Encode(dp)
	if err != nil {
		t.Fatal(err)
	}
	badDP := NewMatrix(13, 100)
	badDP[0] = make([]byte, 1)
	err = r.Encode(badDP)
	if err != ErrShardSize {
		t.Errorf("expected %v, got %v", ErrShardSize, err)
	}
}

// test low, high table work
func TestVerifyEncode(t *testing.T) {
	r, err := New(5, 5)
	if err != nil {
		t.Fatal(err)
	}
	shards := [][]byte{
		{0, 1},
		{4, 5},
		{2, 3},
		{6, 7},
		{8, 9},
		{0, 0},
		{0, 0},
		{0, 0},
		{0, 0},
		{0, 0},
	}
	r.Encode(shards)
	if shards[5][0] != 97 || shards[5][1] != 64 {
		t.Fatal("shard 5 mismatch")
	}
	if shards[6][0] != 173 || shards[6][1] != 3 {
		t.Fatal("shard 6 mismatch")
	}
	if shards[7][0] != 218 || shards[7][1] != 14 {
		t.Fatal("shard 7 mismatch")
	}
	if shards[8][0] != 107 || shards[8][1] != 35 {
		t.Fatal("shard 8 mismatch")
	}
	if shards[9][0] != 110 || shards[9][1] != 177 {
		t.Fatal("shard 9 mismatch")
	}
}

func TestASM(t *testing.T) {
	d := 10
	p := 4
	size := 65 * 1024
	r, err := New(d, p)
	if err != nil {
		t.Fatal(err)
	}
	// asm
	dp := NewMatrix(d+p, size)
	rand.Seed(0)
	for i := 0; i < d; i++ {
		fillRandom(dp[i])
	}
	err = r.Encode(dp)
	if err != nil {
		t.Fatal(err)
	}
	// mulTable
	mDP := NewMatrix(d+p, size)
	for i := 0; i < d; i++ {
		mDP[i] = dp[i]
	}
	err = r.noasmEncode(mDP)
	if err != nil {
		t.Fatal(err)
	}
	for i, asm := range dp {
		if !bytes.Equal(asm, mDP[i]) {
			t.Fatal("verify asm failed, no match noasm version; shards: ", i)
		}
	}
}

func BenchmarkEncode10x4x1M(b *testing.B) {
	benchmarkEncode(b, 10, 4, 1024*1024)
}

func BenchmarkEncode10x4x4M(b *testing.B) {
	benchmarkEncode(b, 10, 4, 4*1024*1024)
}

func BenchmarkEncode10x4x16M(b *testing.B) {
	benchmarkEncode(b, 10, 4, 16*1024*1024)
}

func BenchmarkEncode17x3x1M(b *testing.B) {
	benchmarkEncode(b, 17, 3, 1024*1024)
}

func BenchmarkEncode17x3x4M(b *testing.B) {
	benchmarkEncode(b, 17, 3, 4*1024*1024)
}

func BenchmarkEncode17x3x16M(b *testing.B) {
	benchmarkEncode(b, 17, 3, 16*1024*1024)
}

func benchmarkEncode(b *testing.B, data, parity, size int) {
	r, err := New(data, parity)
	if err != nil {
		b.Fatal(err)
	}
	dp := NewMatrix(data+parity, size)
	rand.Seed(0)
	for i := 0; i < data; i++ {
		fillRandom(dp[i])
	}
	b.SetBytes(int64(size * data))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r.Encode(dp)
	}
}

func fillRandom(p []byte) {
	for i := 0; i < len(p); i += 7 {
		val := rand.Int63()
		for j := 0; i+j < len(p) && j < 7; j++ {
			p[i+j] = byte(val)
			val >>= 8
		}
	}
}

func (r *rs) noasmEncode(dp matrix) error {
	// check args
	if len(dp) != r.shards {
		return ErrTooFewShards
	}
	_, err := checkShardSize(dp)
	if err != nil {
		return err
	}
	// encoding
	input := dp[0:r.data]
	output := dp[r.data:]
	noasmRunner(r.gen, input, output, r.data, r.parity)
	return nil
}

func noasmRunner(gen, input, output matrix, numData, numParity int) {
	for i := 0; i < numData; i++ {
		in := input[i]
		for oi := 0; oi < numParity; oi++ {
			if i == 0 {
				noasmGfVectMul(gen[oi][i], in[:], output[oi][:])
			} else {
				noasmGfVectMulXor(gen[oi][i], in[:], output[oi][:])
			}
		}
	}
}

func noasmGfVectMul(c byte, in, out []byte) {
	mt := mulTable[c]
	for i := 0; i < len(in); i++ {
		out[i] = mt[in[i]]
	}
}

func noasmGfVectMulXor(c byte, in, out []byte) {
	mt := mulTable[c]
	for i := 0; i < len(in); i++ {
		out[i] ^= mt[in[i]]
	}
}

func BenchmarkEncode28x4x16M(b *testing.B) {
	benchmarkEncode(b, 28, 4, 16776168)
}

func BenchmarkEncode14x10x16M(b *testing.B) {
	benchmarkEncode(b, 14, 10, 16776168)
}

func BenchmarkEncode28x4x16M_ConCurrency(b *testing.B) {
	benchmarkEncode_ConCurrency(b, 28, 4, 16776168)
}

func BenchmarkEncode14x10x16M_ConCurrency(b *testing.B) {
	benchmarkEncode_ConCurrency(b, 14, 10, 16776168)
}

func benchmarkEncode_ConCurrency(b *testing.B, data, parity, size int) {
	count := runtime.NumCPU()
	fmt.Println("-------", count)
	Instances := make([]*rs, count)
	dps := make([]matrix, count)
	for i := 0; i < count; i++ {
		r, err := New(data, parity)
		if err != nil {
			b.Fatal(err)
		}
		Instances[i] = r
	}
	for i := 0; i < count; i++ {
		dps[i] = NewMatrix(data+parity, size)
		rand.Seed(0)
		for j := 0; j < data; j++ {
			fillRandom(dps[i][j])
		}
	}

	runtime.GOMAXPROCS(count)
	b.SetBytes(int64(size * data * count))
	b.ResetTimer()
	var g sync.WaitGroup
	for i := 0; i < count; i++ {
		g.Add(1)
		go func(i int) {
			Instances[i].Encode(dps[i])
			g.Done()
		}(i)
	}
	g.Wait()
}<|endoftext|>


<commit_before>package python

import (
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/kr/fs"

	"sourcegraph.com/sourcegraph/srclib/unit"
)

func Scan(srcdir string, repoURI string, repoSubdir string) ([]*unit.SourceUnit, error) {
	if units, isSpecial := specialUnits[repoURI]; isSpecial {
		return units, nil
	}

	cmd := exec.Command("pydep-run.py", "list", srcdir)

	cmd.Stderr = os.Stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	var pkgs []*pkgInfo
	if err := json.NewDecoder(stdout).Decode(&pkgs); err != nil {
		return nil, err
	}

	if err := cmd.Wait(); err != nil {
		return nil, err
	}

	units := make([]*unit.SourceUnit, len(pkgs))
	for i, pkg := range pkgs {
		units[i] = pkg.SourceUnit()
		units[i].Files = pythonSourceFiles(pkg.RootDir)
	}

	return units, nil
}

func pythonSourceFiles(dir string) (files []string) {
	walker := fs.Walk(dir)
	for walker.Step() {
		if err := walker.Err(); err == nil && !walker.Stat().IsDir() && filepath.Ext(walker.Path()) == ".py" {
			file, _ := filepath.Rel(dir, walker.Path())
			files = append(files, file)
		}
	}
	return
}
<commit_msg>scanner with requirements detection<commit_after>package python

import (
	"encoding/json"
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/kr/fs"

	"sourcegraph.com/sourcegraph/srclib/unit"
)

func Scan(srcdir string, repoURI string, repoSubdir string) ([]*unit.SourceUnit, error) {
	if units, isSpecial := specialUnits[repoURI]; isSpecial {
		return units, nil
	}

	cmd := exec.Command("pydep-run.py", "list", srcdir)

	cmd.Stderr = os.Stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	var pkgs []*pkgInfo
	if err := json.NewDecoder(stdout).Decode(&pkgs); err != nil {
		return nil, err
	}

	if err := cmd.Wait(); err != nil {
		return nil, err
	}

	units := make([]*unit.SourceUnit, len(pkgs))
	for i, pkg := range pkgs {
		units[i] = pkg.SourceUnit()
		units[i].Files = pythonSourceFiles(pkg.RootDir)

		reqs, err := requirements(pkg.RootDir)
		if err != nil {
			return nil, err
		}
		reqs_ := make([]interface{}, len(reqs))
		for i, req := range reqs {
			reqs_[i] = req
		}
		units[i].Dependencies = reqs_
	}

	return units, nil
}

func requirements(unitDir string) ([]requirement, error) {
	depCmd := exec.Command("pydep-run.py", "dep", unitDir)
	depCmd.Stderr = os.Stderr
	b, err := depCmd.Output()
	if err != nil {
		return nil, err
	}

	var reqs []requirement
	err = json.Unmarshal(b, &reqs)
	if err != nil {
		return nil, err
	}
	reqs, ignoredReqs := pruneReqs(reqs)
	if len(ignoredReqs) > 0 {
		ignoredKeys := make([]string, len(ignoredReqs))
		for r, req := range ignoredReqs {
			ignoredKeys[r] = req.Key
		}
		log.Printf("(warn) ignoring dependencies %v because repo URL absent", ignoredKeys)
	}

	return reqs, nil
}

// Get all python source files under dir
func pythonSourceFiles(dir string) (files []string) {
	walker := fs.Walk(dir)
	for walker.Step() {
		if err := walker.Err(); err == nil && !walker.Stat().IsDir() && filepath.Ext(walker.Path()) == ".py" {
			file, _ := filepath.Rel(dir, walker.Path())
			files = append(files, file)
		}
	}
	return
}

// Remove unresolvable requirements (i.e., requirements with no clone URL)
func pruneReqs(reqs []requirement) (kept, ignored []requirement) {
	for _, req := range reqs {
		if req.RepoURL != "" { // cannot resolve dependencies with no clone URL
			kept = append(kept, req)
		} else {
			ignored = append(ignored, req)
		}
	}
	return
}
<|endoftext|>