<commit_before>package internal

import (
	"bytes"
	"errors"
	"regexp"
	"strings"

	"gopkg.in/yaml.v2"
	"square/vis/metrics-indexer/api"
)

var defaultRegex = "[^.]+"

var (
	// ErrInvalidPattern is returned when an invalid rule pattern is provided.
	ErrInvalidPattern = errors.New("Invalid Pattern")
	// ErrInvalidMetricKey is retruned when an invalid metric key is provided.
	ErrInvalidMetricKey = errors.New("Invalid Metric Key")
	// ErrInvalidCustomRegex is retruned when the custom regex is invalid.
	ErrInvalidCustomRegex = errors.New("Invalid Custom Regex")
	// ErrMissingTag is returned during the reverse mapping, when a tag is not provided.
	ErrMissingTag = errors.New("Missing Tag")
	// ErrCannotInterpolate is returned when the tag interpolation fails.
	ErrCannotInterpolate = errors.New("Cannot Interpolate")
	// ErrNoMatch is returned when the conversion to tagged metric fails.
	ErrNoMatch = errors.New("No Match")
)

// RawRule is the input provided by the YAML file to specify the rul.
type RawRule struct {
	Pattern          string            `yaml:"pattern"`
	MetricKeyPattern string            `yaml:"metric_key"`
	Regex            map[string]string `yaml:"regex,omitempty"`
}

// RawRules is list of RawRule
type RawRules struct {
	RawRules []RawRule `yaml:"rules"`
}

// Rule is a sanitized version of RawRule. Only valid rules
// can be converted to Rule.
type Rule struct {
	raw                 RawRule
	regex               *regexp.Regexp
	graphitePatternTags []string // tags extracted from the raw graphite string, in the order of appearance.
	metricKeyTags       []string // tags extracted from MetricKey, in the order of appearance.
}

// RuleSet is a sanitized version of RawRules.
// Rules are matched sequentially until a correct one is matched.
type RuleSet struct {
	rules []Rule
}

// Compile a given RawRule into a regex and exposed tagset.
func Compile(rule RawRule) (Rule, error) {
	if len(rule.MetricKeyPattern) == 0 {
		return Rule{}, ErrInvalidMetricKey
	}
	graphitePatternTags := extractTags(rule.Pattern)
	if graphitePatternTags == nil {
		return Rule{}, ErrInvalidPattern
	}
	metricKeyTags := extractTags(string(rule.MetricKeyPattern))
	if !isSubset(metricKeyTags, graphitePatternTags) {
		return Rule{}, ErrInvalidMetricKey
	}
	if metricKeyTags == nil {
		return Rule{}, ErrInvalidMetricKey
	}
	if !rule.checkTagRegexes() {
		return Rule{}, ErrInvalidCustomRegex
	}
	regex := rule.toRegexp()
	if regex == nil {
		return Rule{}, ErrInvalidPattern
	}
	if regex.NumSubexp() != len(graphitePatternTags) {
		return Rule{}, ErrInvalidCustomRegex
	}
	return Rule{
		rule,
		regex,
		graphitePatternTags,
		metricKeyTags,
	}, nil
}

// MatchRule sees if a given graphite string matches the rule, and if so, returns the generated tag.
func (rule Rule) MatchRule(input string) (api.TaggedMetric, bool) {
	matches := rule.regex.FindStringSubmatch(input)
	if matches == nil {
		return api.TaggedMetric{}, false
	}
	tagSet := api.NewTagSet()
	for index, tagValue := range matches {
		if index == 0 {
			continue
		}
		tagKey := rule.graphitePatternTags[index-1]
		tagSet[tagKey] = tagValue
	}
	interpolatedKey, err := interpolateTags(rule.raw.MetricKeyPattern, tagSet)
	if err != nil {
		return api.TaggedMetric{}, false
	}
	// removed tags that are part of the generated metric key.
	for _, metricKeyTag := range rule.metricKeyTags {
		if _, containsKey := tagSet[metricKeyTag]; containsKey {
			delete(tagSet, metricKeyTag)
		}
	}
	return api.TaggedMetric{
		api.MetricKey(interpolatedKey),
		tagSet,
	}, true
}

// ToGraphiteName transforms the given tagged metric back to its graphite metric.
func (rule Rule) ToGraphiteName(taggedMetric api.TaggedMetric) (api.GraphiteMetric, error) {
	interpolatedKey, err := interpolateTags(rule.raw.MetricKeyPattern, taggedMetric.TagSet)
	if err != nil {
		return "", err
	}
	if interpolatedKey != string(taggedMetric.MetricKey) {
		return "", ErrCannotInterpolate
	}
	interpolated, err := interpolateTags(rule.raw.Pattern, taggedMetric.TagSet)
	if err != nil {
		return "", err
	}
	return api.GraphiteMetric(interpolated), nil
}

// MatchRule sees if a given graphite string matches
// any of the specified rules.
func (ruleSet RuleSet) MatchRule(input string) (api.TaggedMetric, bool) {
	for _, rule := range ruleSet.rules {
		value, matched := rule.MatchRule(input)
		if matched {
			return value, matched
		}
	}
	return api.TaggedMetric{}, false
}

// ToGraphiteName transforms the given tagged metric back to its graphite name,
// checking against all the rules.
func (ruleSet RuleSet) ToGraphiteName(taggedMetric api.TaggedMetric) (api.GraphiteMetric, error) {
	for _, rule := range ruleSet.rules {
		reversed, err := rule.ToGraphiteName(taggedMetric)
		if err == nil {
			return reversed, nil
		}
	}
	return "", ErrCannotInterpolate
}

// checkTagRegexes sees if any of the custom regular expressions are invalid.
func (rule RawRule) checkTagRegexes() bool {
	for _, regex := range rule.Regex {
		compiled, err := regexp.Compile(regex)
		if err != nil {
			return false
		}
		if compiled.NumSubexp() > 0 {
			return false // do not allow subexpressions.
		}
	}
	return true
}

func (rule RawRule) toRegexp() *regexp.Regexp {
	splitted := strings.Split(rule.Pattern, "%")
	buffer := new(bytes.Buffer)
	if len(splitted)%2 == 0 {
		// invalid pattern - even number of parts mean odd number of %.
		return nil
	}
	buffer.WriteString("^")
	for index, token := range splitted {
		if isTagPortion(index) {
			regex, contains := rule.Regex[token]
			if !contains {
				// use the defuault regex
				regex = defaultRegex
			}
			// wrap the regex in parenthesis, so that it can be matched.
			// if the regex contains matching groups, bad things will happen.
			buffer.WriteString("(" + regex + ")")
		} else {
			buffer.WriteString(regexp.QuoteMeta(token))
		}
	}
	buffer.WriteString("$")
	compiled, err := regexp.Compile(buffer.String())
	if err != nil {
		return nil
	}
	return compiled
}

func extractTags(pattern string) []string {
	if len(pattern) == 0 {
		return nil // empty pattern is not allowed.
	}
	splitted := strings.Split(pattern, "%")
	if len(splitted)%2 == 0 {
		return nil // invalid tags.
	}

	result := make([]string, len(splitted)/2, len(splitted)/2)
	for index, token := range splitted {
		if isTagPortion(index) {
			if len(token) == 0 {
				return nil // no empty tag
			}
			result[index/2] = token
		}
	}

	// check for duplicates.
	exists := make(map[string]bool)
	for _, token := range result {
		if exists[token] {
			return nil // no duplicate
		}
		exists[token] = true
	}
	return result
}

// LoadYAML loads a RuleSet from the byte array of the YAML file.
func LoadYAML(input []byte) (RuleSet, error) {
	rawRules := RawRules{}
	if err := yaml.Unmarshal(input, &rawRules); err != nil {
		return RuleSet{}, err
	}
	rules := make([]Rule, len(rawRules.RawRules))
	for index, rawRule := range rawRules.RawRules {
		rule, err := Compile(rawRule)
		if err != nil {
			return RuleSet{}, err
		}
		rules[index] = rule
	}
	return RuleSet{rules}, nil
}

// check if setA is subset of setB.
func isSubset(setA, setB []string) bool {
	set := make(map[string]bool)
	for _, v := range setB {
		set[v] = true
	}
	for _, v := range setA {
		if !set[v] {
			return false
		}
	}
	return true
}

func isTagPortion(index int) bool {
	return index%2 == 1
}

func interpolateTags(pattern string, tagSet api.TagSet) (string, error) {
	if !strings.Contains(pattern, "%") {
		return pattern, nil // short circuit when there are no tags to interpolate.
	}
	splitted := strings.Split(pattern, "%")
	buffer := new(bytes.Buffer)
	for index, token := range splitted {
		if isTagPortion(index) {
			tagValue, hasTag := tagSet[token]
			if hasTag {
				buffer.WriteString(tagValue)
			} else {
				return "", ErrMissingTag
			}
		} else {
			buffer.WriteString(token)
		}
	}
	return buffer.String(), nil
}
<commit_msg>better comment<commit_after>package internal

import (
	"bytes"
	"errors"
	"regexp"
	"strings"

	"gopkg.in/yaml.v2"
	"square/vis/metrics-indexer/api"
)

var defaultRegex = "[^.]+"

var (
	// ErrInvalidPattern is returned when an invalid rule pattern is provided.
	ErrInvalidPattern = errors.New("Invalid Pattern")
	// ErrInvalidMetricKey is retruned when an invalid metric key is provided.
	ErrInvalidMetricKey = errors.New("Invalid Metric Key")
	// ErrInvalidCustomRegex is retruned when the custom regex is invalid.
	ErrInvalidCustomRegex = errors.New("Invalid Custom Regex")
	// ErrMissingTag is returned during the reverse mapping, when a tag is not provided.
	ErrMissingTag = errors.New("Missing Tag")
	// ErrCannotInterpolate is returned when the tag interpolation fails.
	ErrCannotInterpolate = errors.New("Cannot Interpolate")
	// ErrNoMatch is returned when the conversion to tagged metric fails.
	ErrNoMatch = errors.New("No Match")
)

// RawRule is the input provided by the YAML file to specify the rul.
type RawRule struct {
	Pattern          string            `yaml:"pattern"`
	MetricKeyPattern string            `yaml:"metric_key"`
	Regex            map[string]string `yaml:"regex,omitempty"`
}

// RawRules is list of RawRule
type RawRules struct {
	RawRules []RawRule `yaml:"rules"`
}

// Rule is a sanitized version of RawRule. Only valid rules
// can be converted to Rule.
type Rule struct {
	raw                 RawRule
	regex               *regexp.Regexp
	graphitePatternTags []string // tags extracted from the raw graphite string, in the order of appearance.
	metricKeyTags       []string // tags extracted from MetricKey, in the order of appearance.
}

// RuleSet is a sanitized version of RawRules.
// Rules are matched sequentially until a correct one is matched.
type RuleSet struct {
	rules []Rule
}

// Compile a given RawRule into a regex and exposed tagset.
func Compile(rule RawRule) (Rule, error) {
	if len(rule.MetricKeyPattern) == 0 {
		return Rule{}, ErrInvalidMetricKey
	}
	graphitePatternTags := extractTags(rule.Pattern)
	if graphitePatternTags == nil {
		return Rule{}, ErrInvalidPattern
	}
	metricKeyTags := extractTags(string(rule.MetricKeyPattern))
	if !isSubset(metricKeyTags, graphitePatternTags) {
		return Rule{}, ErrInvalidMetricKey
	}
	if metricKeyTags == nil {
		return Rule{}, ErrInvalidMetricKey
	}
	if !rule.checkTagRegexes() {
		return Rule{}, ErrInvalidCustomRegex
	}
	regex := rule.toRegexp()
	if regex == nil {
		return Rule{}, ErrInvalidPattern
	}
	if regex.NumSubexp() != len(graphitePatternTags) {
		return Rule{}, ErrInvalidCustomRegex
	}
	return Rule{
		rule,
		regex,
		graphitePatternTags,
		metricKeyTags,
	}, nil
}

// MatchRule sees if a given graphite string matches the rule, and if so, returns the generated tag.
func (rule Rule) MatchRule(input string) (api.TaggedMetric, bool) {
	matches := rule.regex.FindStringSubmatch(input)
	if matches == nil {
		return api.TaggedMetric{}, false
	}
	tagSet := api.NewTagSet()
	for index, tagValue := range matches {
		if index == 0 {
			continue
		}
		tagKey := rule.graphitePatternTags[index-1]
		tagSet[tagKey] = tagValue
	}
	interpolatedKey, err := interpolateTags(rule.raw.MetricKeyPattern, tagSet)
	if err != nil {
		return api.TaggedMetric{}, false
	}
	// Do not output tags appearing in both graphite metric & metric key.
	// for exmaple, if graphite metric is
	//   `foo.%a%.%b%`
	// and metric key is
	//   `bar.%b%`
	// the resulting tag set should only contain {a} after the matching
	// because %b% is already encoded.
	for _, metricKeyTag := range rule.metricKeyTags {
		if _, containsKey := tagSet[metricKeyTag]; containsKey {
			delete(tagSet, metricKeyTag)
		}
	}
	return api.TaggedMetric{
		api.MetricKey(interpolatedKey),
		tagSet,
	}, true
}

// ToGraphiteName transforms the given tagged metric back to its graphite metric.
func (rule Rule) ToGraphiteName(taggedMetric api.TaggedMetric) (api.GraphiteMetric, error) {
	interpolatedKey, err := interpolateTags(rule.raw.MetricKeyPattern, taggedMetric.TagSet)
	if err != nil {
		return "", err
	}
	if interpolatedKey != string(taggedMetric.MetricKey) {
		return "", ErrCannotInterpolate
	}
	interpolated, err := interpolateTags(rule.raw.Pattern, taggedMetric.TagSet)
	if err != nil {
		return "", err
	}
	return api.GraphiteMetric(interpolated), nil
}

// MatchRule sees if a given graphite string matches
// any of the specified rules.
func (ruleSet RuleSet) MatchRule(input string) (api.TaggedMetric, bool) {
	for _, rule := range ruleSet.rules {
		value, matched := rule.MatchRule(input)
		if matched {
			return value, matched
		}
	}
	return api.TaggedMetric{}, false
}

// ToGraphiteName transforms the given tagged metric back to its graphite name,
// checking against all the rules.
func (ruleSet RuleSet) ToGraphiteName(taggedMetric api.TaggedMetric) (api.GraphiteMetric, error) {
	for _, rule := range ruleSet.rules {
		reversed, err := rule.ToGraphiteName(taggedMetric)
		if err == nil {
			return reversed, nil
		}
	}
	return "", ErrCannotInterpolate
}

// checkTagRegexes sees if any of the custom regular expressions are invalid.
func (rule RawRule) checkTagRegexes() bool {
	for _, regex := range rule.Regex {
		compiled, err := regexp.Compile(regex)
		if err != nil {
			return false
		}
		if compiled.NumSubexp() > 0 {
			return false // do not allow subexpressions.
		}
	}
	return true
}

func (rule RawRule) toRegexp() *regexp.Regexp {
	splitted := strings.Split(rule.Pattern, "%")
	buffer := new(bytes.Buffer)
	if len(splitted)%2 == 0 {
		// invalid pattern - even number of parts mean odd number of %.
		return nil
	}
	buffer.WriteString("^")
	for index, token := range splitted {
		if isTagPortion(index) {
			regex, contains := rule.Regex[token]
			if !contains {
				// use the defuault regex
				regex = defaultRegex
			}
			// wrap the regex in parenthesis, so that it can be matched.
			// if the regex contains matching groups, bad things will happen.
			buffer.WriteString("(" + regex + ")")
		} else {
			buffer.WriteString(regexp.QuoteMeta(token))
		}
	}
	buffer.WriteString("$")
	compiled, err := regexp.Compile(buffer.String())
	if err != nil {
		return nil
	}
	return compiled
}

func extractTags(pattern string) []string {
	if len(pattern) == 0 {
		return nil // empty pattern is not allowed.
	}
	splitted := strings.Split(pattern, "%")
	if len(splitted)%2 == 0 {
		return nil // invalid tags.
	}

	result := make([]string, len(splitted)/2, len(splitted)/2)
	for index, token := range splitted {
		if isTagPortion(index) {
			if len(token) == 0 {
				return nil // no empty tag
			}
			result[index/2] = token
		}
	}

	// check for duplicates.
	exists := make(map[string]bool)
	for _, token := range result {
		if exists[token] {
			return nil // no duplicate
		}
		exists[token] = true
	}
	return result
}

// LoadYAML loads a RuleSet from the byte array of the YAML file.
func LoadYAML(input []byte) (RuleSet, error) {
	rawRules := RawRules{}
	if err := yaml.Unmarshal(input, &rawRules); err != nil {
		return RuleSet{}, err
	}
	rules := make([]Rule, len(rawRules.RawRules))
	for index, rawRule := range rawRules.RawRules {
		rule, err := Compile(rawRule)
		if err != nil {
			return RuleSet{}, err
		}
		rules[index] = rule
	}
	return RuleSet{rules}, nil
}

// check if setA is subset of setB.
func isSubset(setA, setB []string) bool {
	set := make(map[string]bool)
	for _, v := range setB {
		set[v] = true
	}
	for _, v := range setA {
		if !set[v] {
			return false
		}
	}
	return true
}

func isTagPortion(index int) bool {
	return index%2 == 1
}

func interpolateTags(pattern string, tagSet api.TagSet) (string, error) {
	if !strings.Contains(pattern, "%") {
		return pattern, nil // short circuit when there are no tags to interpolate.
	}
	splitted := strings.Split(pattern, "%")
	buffer := new(bytes.Buffer)
	for index, token := range splitted {
		if isTagPortion(index) {
			tagValue, hasTag := tagSet[token]
			if hasTag {
				buffer.WriteString(tagValue)
			} else {
				return "", ErrMissingTag
			}
		} else {
			buffer.WriteString(token)
		}
	}
	return buffer.String(), nil
}
<|endoftext|>


<commit_before>package awsrules

import (
	"fmt"
	"log"

	"github.com/hashicorp/hcl2/hcl"
	"github.com/wata727/tflint/tflint"
)

// AwsInstancePreviousTypeRule checks whether the resource uses previous generation instance type
type AwsInstancePreviousTypeRule struct {
	resourceType          string
	attributeName         string
	previousInstanceTypes map[string]bool
}

// NewAwsInstancePreviousTypeRule returns new rule with default attributes
func NewAwsInstancePreviousTypeRule() *AwsInstancePreviousTypeRule {
	return &AwsInstancePreviousTypeRule{
		resourceType:  "aws_instance",
		attributeName: "instance_type",
		previousInstanceTypes: map[string]bool{
			"t1.micro":    true,
			"m1.small":    true,
			"m1.medium":   true,
			"m1.large":    true,
			"m1.xlarge":   true,
			"c1.medium":   true,
			"c1.xlarge":   true,
			"cc2.8xlarge": true,
			"cg1.4xlarge": true,
			"m2.xlarge":   true,
			"m2.2xlarge":  true,
			"m2.4xlarge":  true,
			"cr1.8xlarge": true,
			"hi1.4xlarge": true,
			"hs1.8xlarge": true,
		},
	}
}

// Name returns the rule name
func (r *AwsInstancePreviousTypeRule) Name() string {
	return "aws_instance_previous_type"
}

// Enabled returns whether the rule is enabled by default
func (r *AwsInstancePreviousTypeRule) Enabled() bool {
	return true
}

// Severity returns the rule severity
func (r *AwsInstancePreviousTypeRule) Severity() string {
	return tflint.WARNING
}

// Link returns the rule reference link
func (r *AwsInstancePreviousTypeRule) Link() string {
	return tflint.ReferenceLink(r.Name())
}

// Check checks whether the resource's `instance_type` is included in the list of previous generation instance type
func (r *AwsInstancePreviousTypeRule) Check(runner *tflint.Runner) error {
	log.Printf("[TRACE] Check `%s` rule for `%s` runner", r.Name(), runner.TFConfigPath())

	return runner.WalkResourceAttributes(r.resourceType, r.attributeName, func(attribute *hcl.Attribute) error {
		var instanceType string
		err := runner.EvaluateExpr(attribute.Expr, &instanceType)

		return runner.EnsureNoError(err, func() error {
			if r.previousInstanceTypes[instanceType] {
				runner.EmitIssue(
					r,
					fmt.Sprintf("\"%s\" is previous generation instance type.", instanceType),
					attribute.Expr.Range(),
				)
			}
			return nil
		})
	})
}
<commit_msg>Update EC2 instance type list<commit_after>package awsrules

import (
	"fmt"
	"log"

	"github.com/hashicorp/hcl2/hcl"
	"github.com/wata727/tflint/tflint"
)

// AwsInstancePreviousTypeRule checks whether the resource uses previous generation instance type
type AwsInstancePreviousTypeRule struct {
	resourceType          string
	attributeName         string
	previousInstanceTypes map[string]bool
}

// NewAwsInstancePreviousTypeRule returns new rule with default attributes
func NewAwsInstancePreviousTypeRule() *AwsInstancePreviousTypeRule {
	return &AwsInstancePreviousTypeRule{
		resourceType:  "aws_instance",
		attributeName: "instance_type",
		previousInstanceTypes: map[string]bool{
			"t1.micro":    true,
			"m1.small":    true,
			"m1.medium":   true,
			"m1.large":    true,
			"m1.xlarge":   true,
			"m3.medium":   true,
			"m3.large":    true,
			"m3.xlarge":   true,
			"m3.2xlarge":  true,
			"c1.medium":   true,
			"c1.xlarge":   true,
			"cc2.8xlarge": true,
			"cg1.4xlarge": true,
			"c3.large":    true,
			"c3.xlarge":   true,
			"c3.2xlarge":  true,
			"c3.4xlarge":  true,
			"c3.8xlarge":  true,
			"g2.2xlarge":  true,
			"g2.8xlarge":  true,
			"m2.xlarge":   true,
			"m2.2xlarge":  true,
			"m2.4xlarge":  true,
			"cr1.8xlarge": true,
			"r3.large":    true,
			"r3.xlarge":   true,
			"r3.2xlarge":  true,
			"r3.4xlarge":  true,
			"r3.8xlarge":  true,
			"i2.xlarge":   true,
			"i2.2xlarge":  true,
			"i2.4xlarge":  true,
			"i2.8xlarge":  true,
			"hi1.4xlarge": true,
			"hs1.8xlarge": true,
		},
	}
}

// Name returns the rule name
func (r *AwsInstancePreviousTypeRule) Name() string {
	return "aws_instance_previous_type"
}

// Enabled returns whether the rule is enabled by default
func (r *AwsInstancePreviousTypeRule) Enabled() bool {
	return true
}

// Severity returns the rule severity
func (r *AwsInstancePreviousTypeRule) Severity() string {
	return tflint.WARNING
}

// Link returns the rule reference link
func (r *AwsInstancePreviousTypeRule) Link() string {
	return tflint.ReferenceLink(r.Name())
}

// Check checks whether the resource's `instance_type` is included in the list of previous generation instance type
func (r *AwsInstancePreviousTypeRule) Check(runner *tflint.Runner) error {
	log.Printf("[TRACE] Check `%s` rule for `%s` runner", r.Name(), runner.TFConfigPath())

	return runner.WalkResourceAttributes(r.resourceType, r.attributeName, func(attribute *hcl.Attribute) error {
		var instanceType string
		err := runner.EvaluateExpr(attribute.Expr, &instanceType)

		return runner.EnsureNoError(err, func() error {
			if r.previousInstanceTypes[instanceType] {
				runner.EmitIssue(
					r,
					fmt.Sprintf("\"%s\" is previous generation instance type.", instanceType),
					attribute.Expr.Range(),
				)
			}
			return nil
		})
	})
}
<|endoftext|>


<commit_before>package coremain

// Various CoreDNS constants.
const (
	CoreVersion = "1.6.4"
	coreName    = "CoreDNS"
	serverType  = "dns"
)
<commit_msg>dump version to 1.6.5 (#3425)<commit_after>package coremain

// Various CoreDNS constants.
const (
	CoreVersion = "1.6.5"
	coreName    = "CoreDNS"
	serverType  = "dns"
)
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/system"
)

const version = "0.2.1"

func main() {
	var printVersion bool
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")

	var ignoreFailure bool
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")

	var file string
	flag.StringVar(&file, "from-file", "", "Read user-data from provided file")

	var url string
	flag.StringVar(&url, "from-url", "", "Download user-data from provided url")

	var workspace string
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")

	var sshKeyName string
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")

	flag.Parse()

	if printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	if file != "" && url != "" {
		fmt.Println("Provide one of --from-file or --from-url")
		os.Exit(1)
	}

	var ds datasource.Datasource
	if file != "" {
		ds = datasource.NewLocalFile(file)
	} else if url != "" {
		ds = datasource.NewMetadataService(url)
	} else {
		fmt.Println("Provide one of --from-file or --from-url")
		os.Exit(1)
	}

	log.Printf("Fetching user-data from datasource of type %q", ds.Type())
	userdata, err := ds.Fetch()
	if err != nil {
		log.Printf("Failed fetching user-data from datasource: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	if len(userdata) == 0 {
		log.Printf("No user data to handle, exiting.")
		os.Exit(0)
	}

	parsed, err := ParseUserData(userdata)
	if err != nil {
		log.Printf("Failed parsing user-data: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	env := initialize.NewEnvironment("/", workspace)
	err = initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		log.Fatalf("Failed preparing workspace: %v", err)
	}

	switch t := parsed.(type) {
	case initialize.CloudConfig:
		err = initialize.Apply(t, env)
	case system.Script:
		var path string
		path, err = initialize.PersistScriptInWorkspace(t, env.Workspace())
		if err == nil {
			var name string
			name, err = system.ExecuteScript(path)
			initialize.PersistUnitNameInWorkspace(name, workspace)
		}
	}

	if err != nil {
		log.Fatalf("Failed resolving user-data: %v", err)
	}
}

func ParseUserData(contents []byte) (interface{}, error) {
	bytereader := bytes.NewReader(contents)
	bufreader := bufio.NewReader(bytereader)
	header, _ := bufreader.ReadString('\n')

	if strings.HasPrefix(header, "#!") {
		log.Printf("Parsing user-data as script")
		return system.Script(contents), nil

	} else if header == "#cloud-config\n" {
		log.Printf("Parsing user-data as cloud-config")
		cfg, err := initialize.NewCloudConfig(contents)
		if err != nil {
			log.Fatal(err.Error())
		}
		return *cfg, nil
	} else {
		return nil, fmt.Errorf("Unrecognized user-data header: %s", header)
	}
}
<commit_msg>chore(release): Bump version to v0.2.1+git<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/system"
)

const version = "0.2.1+git"

func main() {
	var printVersion bool
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")

	var ignoreFailure bool
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")

	var file string
	flag.StringVar(&file, "from-file", "", "Read user-data from provided file")

	var url string
	flag.StringVar(&url, "from-url", "", "Download user-data from provided url")

	var workspace string
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")

	var sshKeyName string
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")

	flag.Parse()

	if printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	if file != "" && url != "" {
		fmt.Println("Provide one of --from-file or --from-url")
		os.Exit(1)
	}

	var ds datasource.Datasource
	if file != "" {
		ds = datasource.NewLocalFile(file)
	} else if url != "" {
		ds = datasource.NewMetadataService(url)
	} else {
		fmt.Println("Provide one of --from-file or --from-url")
		os.Exit(1)
	}

	log.Printf("Fetching user-data from datasource of type %q", ds.Type())
	userdata, err := ds.Fetch()
	if err != nil {
		log.Printf("Failed fetching user-data from datasource: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	if len(userdata) == 0 {
		log.Printf("No user data to handle, exiting.")
		os.Exit(0)
	}

	parsed, err := ParseUserData(userdata)
	if err != nil {
		log.Printf("Failed parsing user-data: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	env := initialize.NewEnvironment("/", workspace)
	err = initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		log.Fatalf("Failed preparing workspace: %v", err)
	}

	switch t := parsed.(type) {
	case initialize.CloudConfig:
		err = initialize.Apply(t, env)
	case system.Script:
		var path string
		path, err = initialize.PersistScriptInWorkspace(t, env.Workspace())
		if err == nil {
			var name string
			name, err = system.ExecuteScript(path)
			initialize.PersistUnitNameInWorkspace(name, workspace)
		}
	}

	if err != nil {
		log.Fatalf("Failed resolving user-data: %v", err)
	}
}

func ParseUserData(contents []byte) (interface{}, error) {
	bytereader := bytes.NewReader(contents)
	bufreader := bufio.NewReader(bytereader)
	header, _ := bufreader.ReadString('\n')

	if strings.HasPrefix(header, "#!") {
		log.Printf("Parsing user-data as script")
		return system.Script(contents), nil

	} else if header == "#cloud-config\n" {
		log.Printf("Parsing user-data as cloud-config")
		cfg, err := initialize.NewCloudConfig(contents)
		if err != nil {
			log.Fatal(err.Error())
		}
		return *cfg, nil
	} else {
		return nil, fmt.Errorf("Unrecognized user-data header: %s", header)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Google Inc.
//
// 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.

// stable-api demonstrates a bar that exercises barista's stable API.
package main

import (
	"os/exec"
	"syscall"
	"time"

	"github.com/dustin/go-humanize"

	"github.com/soumya92/barista"
	"github.com/soumya92/barista/bar"
	"github.com/soumya92/barista/timing"
)

type simpleClockModule struct {
	format   string
	interval time.Duration
}

func (s simpleClockModule) Stream(sink bar.Sink) {
	sch := timing.NewScheduler()
	for {
		now := timing.Now()
		sink.Output(bar.TextSegment(now.Format(s.format)))
		next := now.Add(s.interval).Truncate(s.interval)
		<-sch.At(next).Tick()
	}
}

type diskSpaceModule string

func (d diskSpaceModule) Stream(sink bar.Sink) {
	sch := timing.NewScheduler().Every(5 * time.Second)
	for {
		var stat syscall.Statfs_t
		err := syscall.Statfs(string(d), &stat)
		if sink.Error(err) {
			return
		}
		sink.Output(bar.TextSegment(
			humanize.IBytes(stat.Bavail * uint64(stat.Bsize)),
		))
		<-sch.Tick()
	}
}

func (d diskSpaceModule) Click(e bar.Event) {
	exec.Command("xdg-open", string(d)).Run()
}

func main() {
	panic(barista.Run(
		diskSpaceModule("/home"),
		simpleClockModule{"Mon Jan 02", time.Hour},
		simpleClockModule{"15:04:05", time.Second},
	))
}
<commit_msg>Update stable-api sample to reflect changed event handling<commit_after>// Copyright 2018 Google Inc.
//
// 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.

// stable-api demonstrates a bar that exercises barista's stable API.
package main

import (
	"syscall"
	"time"

	"github.com/dustin/go-humanize"

	"github.com/soumya92/barista"
	"github.com/soumya92/barista/bar"
	"github.com/soumya92/barista/timing"
)

type simpleClockModule struct {
	format   string
	interval time.Duration
}

func (s simpleClockModule) Stream(sink bar.Sink) {
	sch := timing.NewScheduler()
	for {
		now := timing.Now()
		sink.Output(bar.TextSegment(now.Format(s.format)))
		next := now.Add(s.interval).Truncate(s.interval)
		<-sch.At(next).Tick()
	}
}

type diskSpaceModule string

func (d diskSpaceModule) Stream(sink bar.Sink) {
	sch := timing.NewScheduler().Every(5 * time.Second)
	for {
		var stat syscall.Statfs_t
		err := syscall.Statfs(string(d), &stat)
		if sink.Error(err) {
			return
		}
		sink.Output(bar.TextSegment(
			humanize.IBytes(stat.Bavail * uint64(stat.Bsize)),
		))
		<-sch.Tick()
	}
}

func main() {
	panic(barista.Run(
		diskSpaceModule("/home"),
		simpleClockModule{"Mon Jan 02", time.Hour},
		simpleClockModule{"15:04:05", time.Second},
	))
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save_test

import (
	"testing"

	"golang.org/x/net/context"

	. "github.com/jacobsa/ogletest"
)

func TestFileSystemVisitor(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FileSystemVisitorTest struct {
	ctx context.Context
	dir string
}

var _ SetUpInterface = &FileSystemVisitorTest{}
var _ TearDownInterface = &FileSystemVisitorTest{}

func init() { RegisterTestSuite(&FileSystemVisitorTest{}) }

func (t *FileSystemVisitorTest) SetUp(ti *TestInfo) {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) TearDown() {
	AssertFalse(true, "TODO")
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FileSystemVisitorTest) NonExistentPath() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) NotADirectory() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) VisitRootNode() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) VisitNonRootNode() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Files() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Directories() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Symlinks() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Devices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) CharDevices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) NamedPipes() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Sockets() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Exclusions() {
	AssertFalse(true, "TODO")
}
<commit_msg>Set up and tear down.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save_test

import (
	"io/ioutil"
	"os"
	"testing"

	"golang.org/x/net/context"

	. "github.com/jacobsa/ogletest"
)

func TestFileSystemVisitor(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FileSystemVisitorTest struct {
	ctx context.Context

	// A temporary directory that is cleaned up at the end of the test.
	dir string
}

var _ SetUpInterface = &FileSystemVisitorTest{}
var _ TearDownInterface = &FileSystemVisitorTest{}

func init() { RegisterTestSuite(&FileSystemVisitorTest{}) }

func (t *FileSystemVisitorTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Create the temporary directory.
	var err error
	t.dir, err = ioutil.TempDir("", "file_system_visistor_test")
	AssertEq(nil, err)
}

func (t *FileSystemVisitorTest) TearDown() {
	var err error

	err = os.RemoveAll(t.dir)
	AssertEq(nil, err)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FileSystemVisitorTest) NonExistentPath() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) NotADirectory() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) VisitRootNode() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) VisitNonRootNode() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Files() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Directories() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Symlinks() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Devices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) CharDevices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) NamedPipes() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Sockets() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Exclusions() {
	AssertFalse(true, "TODO")
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save_test

import (
	"io/ioutil"
	"os"
	"path"
	"sort"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/save"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestFileSystemVisitor(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type PathAndFileInfoSlice []save.PathAndFileInfo

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

func (p PathAndFileInfoSlice) Less(i, j int) bool {
	return p[i].Path < p[j].Path
}

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

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FileSystemVisitorTest struct {
	ctx context.Context

	// A temporary directory that is cleaned up at the end of the test. This is
	// the base path with which the visitor is configured.
	dir string

	// The channel into which the visitor writes. Configured with a large buffer.
	output chan save.PathAndFileInfo

	visitor graph.Visitor
}

var _ SetUpInterface = &FileSystemVisitorTest{}
var _ TearDownInterface = &FileSystemVisitorTest{}

func init() { RegisterTestSuite(&FileSystemVisitorTest{}) }

func (t *FileSystemVisitorTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.output = make(chan save.PathAndFileInfo, 10e3)

	// Create the base directory.
	var err error
	t.dir, err = ioutil.TempDir("", "file_system_visistor_test")
	AssertEq(nil, err)

	// And the visitor.
	t.visitor = save.NewFileSystemVisitor(t.dir, t.output)
}

func (t *FileSystemVisitorTest) TearDown() {
	var err error

	// Clean up the junk we left in the file system.
	err = os.RemoveAll(t.dir)
	AssertEq(nil, err)
}

// Consume the output, returning a slice sorted by path.
func (t *FileSystemVisitorTest) sortOutput() (output PathAndFileInfoSlice) {
	for o := range t.output {
		output = append(output, o)
	}

	sort.Sort(output)
	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FileSystemVisitorTest) NonExistentPath() {
	const node = "foo/bar"

	_, err := t.visitor.Visit(t.ctx, node)
	ExpectThat(err, Error(HasSubstr(node)))
	ExpectThat(err, Error(HasSubstr("not found")))
}

func (t *FileSystemVisitorTest) NotADirectory() {
	const node = "foo"
	var err error

	// Create a file.
	err = ioutil.WriteFile(path.Join(t.dir, node), []byte{}, 0500)
	AssertEq(nil, err)

	// Attempt to visit it.
	_, err = t.visitor.Visit(t.ctx, node)
	ExpectThat(err, Error(HasSubstr(node)))
	ExpectThat(err, Error(HasSubstr("TODO")))
}

func (t *FileSystemVisitorTest) VisitRootNode() {
	var err error

	// Create two children.
	err = ioutil.WriteFile(path.Join(t.dir, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(t.dir, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the root.
	_, err = t.visitor.Visit(t.ctx, "")
	AssertEq(nil, err)

	// Check the output.
	output := t.sortOutput()
	AssertEq(2, len(output))
	ExpectEq(path.Join(t.dir, "bar"), output[0].Path)
	ExpectEq(path.Join(t.dir, "foo"), output[1].Path)
}

func (t *FileSystemVisitorTest) VisitNonRootNode() {
	var err error

	// Make a few levels of sub-directories.
	d := path.Join(t.dir, "sub/dirs")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create two children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(d, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the root.
	_, err = t.visitor.Visit(t.ctx, "sub/dirs")
	AssertEq(nil, err)

	// Check the output.
	output := t.sortOutput()
	AssertEq(2, len(output))
	ExpectEq(path.Join(t.dir, "bar"), output[0].Path)
	ExpectEq(path.Join(t.dir, "foo"), output[1].Path)
}

func (t *FileSystemVisitorTest) Files() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Directories() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Symlinks() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Devices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) CharDevices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) NamedPipes() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Sockets() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Exclusions() {
	AssertFalse(true, "TODO")
}
<commit_msg>FileSystemVisitorTest.Files<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save_test

import (
	"io/ioutil"
	"os"
	"path"
	"sort"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/save"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestFileSystemVisitor(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type PathAndFileInfoSlice []save.PathAndFileInfo

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

func (p PathAndFileInfoSlice) Less(i, j int) bool {
	return p[i].Path < p[j].Path
}

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

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FileSystemVisitorTest struct {
	ctx context.Context

	// A temporary directory that is cleaned up at the end of the test. This is
	// the base path with which the visitor is configured.
	dir string

	// The channel into which the visitor writes. Configured with a large buffer.
	output chan save.PathAndFileInfo

	visitor graph.Visitor
}

var _ SetUpInterface = &FileSystemVisitorTest{}
var _ TearDownInterface = &FileSystemVisitorTest{}

func init() { RegisterTestSuite(&FileSystemVisitorTest{}) }

func (t *FileSystemVisitorTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.output = make(chan save.PathAndFileInfo, 10e3)

	// Create the base directory.
	var err error
	t.dir, err = ioutil.TempDir("", "file_system_visistor_test")
	AssertEq(nil, err)

	// And the visitor.
	t.visitor = save.NewFileSystemVisitor(t.dir, t.output)
}

func (t *FileSystemVisitorTest) TearDown() {
	var err error

	// Clean up the junk we left in the file system.
	err = os.RemoveAll(t.dir)
	AssertEq(nil, err)
}

// Consume the output, returning a slice sorted by path.
func (t *FileSystemVisitorTest) sortOutput() (output PathAndFileInfoSlice) {
	for o := range t.output {
		output = append(output, o)
	}

	sort.Sort(output)
	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FileSystemVisitorTest) NonExistentPath() {
	const node = "foo/bar"

	_, err := t.visitor.Visit(t.ctx, node)
	ExpectThat(err, Error(HasSubstr(node)))
	ExpectThat(err, Error(HasSubstr("not found")))
}

func (t *FileSystemVisitorTest) NotADirectory() {
	const node = "foo"
	var err error

	// Create a file.
	err = ioutil.WriteFile(path.Join(t.dir, node), []byte{}, 0500)
	AssertEq(nil, err)

	// Attempt to visit it.
	_, err = t.visitor.Visit(t.ctx, node)
	ExpectThat(err, Error(HasSubstr(node)))
	ExpectThat(err, Error(HasSubstr("TODO")))
}

func (t *FileSystemVisitorTest) VisitRootNode() {
	var err error

	// Create two children.
	err = ioutil.WriteFile(path.Join(t.dir, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(t.dir, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the root.
	_, err = t.visitor.Visit(t.ctx, "")
	AssertEq(nil, err)

	// Check the output.
	output := t.sortOutput()
	AssertEq(2, len(output))
	ExpectEq(path.Join(t.dir, "bar"), output[0].Path)
	ExpectEq(path.Join(t.dir, "foo"), output[1].Path)
}

func (t *FileSystemVisitorTest) VisitNonRootNode() {
	var err error

	// Make a few levels of sub-directories.
	d := path.Join(t.dir, "sub/dirs")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create two children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(d, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the root.
	_, err = t.visitor.Visit(t.ctx, "sub/dirs")
	AssertEq(nil, err)

	// Check the output.
	output := t.sortOutput()
	AssertEq(2, len(output))
	ExpectEq(path.Join(t.dir, "bar"), output[0].Path)
	ExpectEq(path.Join(t.dir, "foo"), output[1].Path)
}

func (t *FileSystemVisitorTest) Files() {
	var err error
	var pfi save.PathAndFileInfo

	// Create two children.
	err = ioutil.WriteFile(path.Join(t.dir, "foo"), []byte("taco"), 0400)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(t.dir, "bar"), []byte("burrito"), 0400)
	AssertEq(nil, err)

	// Visit the root.
	adjacent, err := t.visitor.Visit(t.ctx, "")

	AssertEq(nil, err)
	ExpectThat(adjacent, ElementsAre())

	// Check the output.
	output := t.sortOutput()
	AssertEq(2, len(output))

	pfi = output[0]
	ExpectEq(path.Join(t.dir, "bar"), pfi.Path)
	ExpectEq("bar", pfi.Info.Name())
	ExpectEq(len("burrito"), pfi.Info.Size())

	pfi = output[1]
	ExpectEq(path.Join(t.dir, "foo"), pfi.Path)
	ExpectEq("foo", pfi.Info.Name())
	ExpectEq(len("taco"), pfi.Info.Size())
}

func (t *FileSystemVisitorTest) Directories() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Symlinks() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Devices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) CharDevices() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) NamedPipes() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Sockets() {
	AssertFalse(true, "TODO")
}

func (t *FileSystemVisitorTest) Exclusions() {
	AssertFalse(true, "TODO")
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.
//
// Usage: generate_catalog SRC_REPO_DIR/ DEST_MD_DIR/
//
// The command will create a README.md file under DEST_MD_DIR/ containing a table
// of collected functions. Source files for the catalog will also appear in
// DEST_MD_DIR/
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"golang.org/x/mod/semver"
	"gopkg.in/yaml.v2"
)

var (
	branchesToSkip = []string{
		"sops/v0.2",
		// TODO: Remove this line after port this PR to release branch sops/v0.3
		"sops/v0.3",
	}
)

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintf(os.Stderr, "Usage: generate-catalog SRC_REPO_DIR/ DEST_MD_DIR/\n")
		os.Exit(1)
	}
	source, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	dest, err := filepath.Abs(os.Args[2])
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	branches, err := getBranches()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	functions := getFunctions(branches, source, dest)
	curatedFns, contribFns, err := writeFunctionIndex(functions, source, dest)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	err = writeExampleIndex(curatedFns, source, dest)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	err = writeExampleIndex(contribFns, source, filepath.Join(dest, "contrib"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}

type function struct {
	FunctionName      string
	ImagePath         string
	VersionToExamples map[string]map[string]example
	LatestVersion     string
	Path              string
	Description       string
	Tags              string
	Gcp               bool
}

type example struct {
	LocalExamplePath  string
	RemoteExamplePath string
	RemoteSourcePath  string
}

type metadata struct {
	Image              string
	Description        string
	Tags               []string
	SourceUrl          string   `yaml:"sourceURL"`
	ExamplePackageUrls []string `yaml:"examplePackageURLs"`
	Hidden             bool
}

var (
	// Match start of a version such as v1.9.1
	branchSemverPrefix = regexp.MustCompile(`[-\w]*\/(v\d*\.\d*)`)
	functionDirPrefix  = regexp.MustCompile(`.+/functions/`)
)

func getBranches() ([]string, error) {
	verBranches := make([]string, 0)

	var buf bytes.Buffer
	cmd := exec.Command("git", "branch", "-a")
	cmd.Stdout = &buf
	err := cmd.Run()
	if err != nil {
		return verBranches, err
	}

	for _, branch := range strings.Split(buf.String(), "\n") {
		skipCurrentBranch := false
		for _, branchToSkip := range branchesToSkip {
			if strings.Contains(branch, branchToSkip) {
				skipCurrentBranch = true
				break
			}
		}
		if skipCurrentBranch {
			continue
		}

		if branchSemverPrefix.MatchString(branch) {
			verBranches = append(verBranches, strings.TrimSpace(branch))
		}
	}
	return verBranches, err
}

func getFunctions(branches []string, source string, dest string) []function {
	functions := make(map[string]function)
	for _, b := range branches {
		segments := strings.Split(b, "/")
		funcName := segments[len(segments)-2]
		minorVersion := segments[len(segments)-1]
		relativeFuncPath, err := getRelativeFunctionPath(source, funcName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			continue
		}
		versionDest := filepath.Join(dest, funcName, minorVersion)
		if strings.Contains(relativeFuncPath, "contrib") {
			versionDest = filepath.Join(dest, "contrib", funcName, minorVersion)
		}

		// Functions with the hidden field enabled should not be processed.
		metadataPath := strings.TrimSpace(fmt.Sprintf("%v:%v", b, filepath.Join(relativeFuncPath, "metadata.yaml")))
		md, err := getMetadata(metadataPath)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error getting metadata for %q from %q: %v\n", funcName, b, err)
			os.Exit(1)
		}
		if md.Hidden {
			continue
		}
		err = copyExamples(b, md.ExamplePackageUrls, versionDest, minorVersion)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error getting examples for %q from %q: %v\n", funcName, b, err)
			os.Exit(1)
		}

		err = copyReadme(b, funcName, relativeFuncPath, versionDest)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error getting README for %q from %q: %v\n", funcName, b, err)
			os.Exit(1)
		}

		f := functions[funcName]
		f.FunctionName = funcName
		if f.VersionToExamples == nil {
			f.VersionToExamples = make(map[string]map[string]example)
		}

		functions[funcName] = parseMetadata(f, md, minorVersion, versionDest)
	}

	flattenedFunctions := make([]function, 0)
	for _, f := range functions {
		flattenedFunctions = append(flattenedFunctions, f)
	}
	sort.Slice(flattenedFunctions, func(i int, j int) bool {
		return flattenedFunctions[i].FunctionName < flattenedFunctions[j].FunctionName
	})
	return flattenedFunctions
}

func copyExamples(b string, exampleSources []string, versionDest, minorVersion string) error {
	// Copy examples for the function's version to a temporary directory.
	tempDir, err := ioutil.TempDir("", "examples")
	if err != nil {
		return err
	}

	// Prepare destination for versioned examples.
	err = os.MkdirAll(versionDest, 0744)
	if err != nil {
		return err
	}

	for _, exampleSource := range exampleSources {
		splitedPaths := strings.SplitN(exampleSource, minorVersion+string(filepath.Separator), 2)
		if len(splitedPaths) != 2 {
			return fmt.Errorf("expect 2 substring after spliting %q by %q", exampleSource, minorVersion+string(filepath.Separator))
		}
		relativePath := splitedPaths[1]
		// Fetch example into temporary directory.
		cmd := exec.Command("git", fmt.Sprintf("--work-tree=%v", tempDir), "checkout", b, "--", relativePath)
		err = cmd.Run()
		if err != nil {
			return fmt.Errorf("Error running %v: %v", cmd, err)
		}

		exampleName := filepath.Base(relativePath)
		src := filepath.Join(tempDir, relativePath)
		dest := filepath.Join(versionDest, exampleName)

		// Move example content to the site's example directory.
		if err = os.RemoveAll(dest); err != nil {
			return err
		}
		err = os.Rename(src, dest)
		if err != nil {
			return err
		}

	}

	return nil
}

func copyReadme(b string, funcName string, relativeFuncPath string, versionDest string) error {
	// Copy README for the function's version to the function's directory.
	tempDir, err := ioutil.TempDir("", "functions")
	if err != nil {
		return err
	}
	cmd := exec.Command("git", fmt.Sprintf("--work-tree=%v", tempDir), "checkout", b, "--", filepath.Join(relativeFuncPath, "README.md"))
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("Error running %v: %v", cmd, err)
	}

	// Find the README in the function's directory.
	m, err := filepath.Glob(filepath.Join(tempDir, "functions", "*", funcName, "README.md"))
	if err != nil {
		return err
	}

	if m == nil {
		m, err = filepath.Glob(filepath.Join(tempDir, "contrib", "functions", "*", funcName, "README.md"))
		if err != nil {
			return err
		}
	}

	// Move the README to the destination directory.
	err = os.Rename(m[0], filepath.Join(versionDest, "README.md"))
	if err != nil {
		return err
	}

	return nil
}

func getMetadata(metadataPath string) (metadata, error) {
	var buf bytes.Buffer
	var md metadata
	// Get the content of metadata.yaml from the appropriate release branch.
	cmd := exec.Command("git", "cat-file", "blob", metadataPath)
	cmd.Stdout = &buf
	err := cmd.Run()
	if err != nil {
		return md, err
	}

	yaml.Unmarshal(buf.Bytes(), &md)
	return md, nil
}

func parseMetadata(f function, md metadata, version string, versionDest string) function {

	// Add examples to version map.
	if f.VersionToExamples[version] == nil {
		f.VersionToExamples[version] = make(map[string]example)
	}
	e := f.VersionToExamples[version]
	for _, exUrl := range md.ExamplePackageUrls {
		exUrlSegments := strings.Split(exUrl, "/")
		exName := exUrlSegments[len(exUrlSegments)-1]
		ex := example{
			RemoteSourcePath:  md.SourceUrl,
			RemoteExamplePath: exUrl,
			LocalExamplePath:  filepath.Join(versionDest, exName),
		}
		e[exName] = ex
	}
	f.VersionToExamples[version] = e
	if semver.Compare(f.LatestVersion, version) == 1 {
		return f
	}

	// If this is the latest version,
	// update latest version, default path, description and rags.
	f.LatestVersion = version
	f.Path = versionDest
	f.Description = md.Description
	functionTags := make([]string, 0)
	for _, tag := range md.Tags {
		normalizedTag := strings.ToLower(tag)
		if normalizedTag == "gcp" {
			f.Gcp = true
		} else {
			functionTags = append(functionTags, normalizedTag)
		}
	}
	sort.Strings(functionTags)
	f.Tags = strings.Join(functionTags, ", ")
	f.ImagePath = md.Image

	return f
}

func getRelativeFunctionPath(source string, funcName string) (string, error) {
	// Find the directory for the function's source.
	sourcePattern := filepath.Join(source, "functions", "*", funcName)
	m, err := filepath.Glob(sourcePattern)
	if err != nil {
		return "", err
	}
	if m != nil {
		return functionDirPrefix.ReplaceAllString(m[0], "functions/"), nil
	}

	contribPattern := filepath.Join(source, "contrib", "functions", "*", funcName)
	m, err = filepath.Glob(contribPattern)
	if err != nil {
		return  "", err
	}
	if m == nil {
		return  "", fmt.Errorf("Could not find a function with the following name: %v", funcName)
	}
	return functionDirPrefix.ReplaceAllString(m[0], "contrib/functions/"), nil
}

func writeFunctionIndex(functions []function, source string, dest string) ([]function, []function, error) {
	out := []string{"# Curated Functions Catalog", ""}
	var contribOut []string

	genericFunctions := make([]function, 0)
	gcp := make([]function, 0)
	contribFunctions := make([]function, 0)
	for _, f := range functions {
		if strings.Contains(f.ImagePath, "contrib") {
			contribFunctions = append(contribFunctions, f)
		} else {
			if f.Gcp {
				gcp = append(gcp, f)
			} else {
				genericFunctions = append(genericFunctions, f)
			}
		}
	}

	out = append(out, getFunctionTable(genericFunctions, source)...)

	if len(gcp) > 0 {
		out = append(out, "", "## GCP Functions", "")
		out = append(out, getFunctionTable(gcp, source)...)
	}

	if len(contribFunctions) > 0 {
		contribOut = append(contribOut, "# Contrib Functions Catalog", "")
		contribOut = append(contribOut, getFunctionTable(contribFunctions, source)...)
	}

	o := strings.Join(out, "\n")
	if err := ioutil.WriteFile(filepath.Join(dest, "README.md"), []byte(o), 0744); err != nil {
		return nil, nil, err
	}
	co := strings.Join(contribOut, "\n")
	if err := ioutil.WriteFile(filepath.Join(dest, "contrib", "README.md"), []byte(co), 0744); err != nil {
		return nil, nil, err
	}
	return append(genericFunctions, gcp...), contribFunctions, nil
}

func getFunctionTable(functions []function, source string) []string {
	out := []string{"| Name | Description | Tags |", "| ---- | ----------- | ---- |"}
	for _, f := range functions {
		functionEntry := fmt.Sprintf("| [%v](%v/) | %v | %v |", f.FunctionName, strings.Replace(f.Path, filepath.Join(source, "site"), "", 1), f.Description, f.Tags)
		out = append(out, functionEntry)
	}
	return out
}

func writeExampleIndex(functions []function, source string, dest string) error {
	// Key a function's version's examples by the function's name -> version
	functionVersionMap := make(map[string]map[string]map[string]example)
	for _, f := range functions {
		vToE := make(map[string]map[string]example)
		for v, examples := range f.VersionToExamples {
			exampleToPaths := make(map[string]example)
			for exName, ex := range examples {
				e := ex
				e.LocalExamplePath = strings.Replace(ex.LocalExamplePath, filepath.Join(source, "site"), "", 1)
				exampleToPaths[exName] = e
			}
			vToE[v] = exampleToPaths
		}
		functionVersionMap[f.FunctionName] = vToE
	}

	funcJson, err := json.Marshal(functionVersionMap)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(dest, "catalog.json"), funcJson, 0600)
	return err
}
<commit_msg>chore: don't skip sops/v0.3 branch (#598)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.
//
// Usage: generate_catalog SRC_REPO_DIR/ DEST_MD_DIR/
//
// The command will create a README.md file under DEST_MD_DIR/ containing a table
// of collected functions. Source files for the catalog will also appear in
// DEST_MD_DIR/
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"golang.org/x/mod/semver"
	"gopkg.in/yaml.v2"
)

var (
	branchesToSkip = []string{
		"sops/v0.2",
	}
)

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintf(os.Stderr, "Usage: generate-catalog SRC_REPO_DIR/ DEST_MD_DIR/\n")
		os.Exit(1)
	}
	source, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	dest, err := filepath.Abs(os.Args[2])
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	branches, err := getBranches()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	functions := getFunctions(branches, source, dest)
	curatedFns, contribFns, err := writeFunctionIndex(functions, source, dest)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	err = writeExampleIndex(curatedFns, source, dest)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	err = writeExampleIndex(contribFns, source, filepath.Join(dest, "contrib"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}

type function struct {
	FunctionName      string
	ImagePath         string
	VersionToExamples map[string]map[string]example
	LatestVersion     string
	Path              string
	Description       string
	Tags              string
	Gcp               bool
}

type example struct {
	LocalExamplePath  string
	RemoteExamplePath string
	RemoteSourcePath  string
}

type metadata struct {
	Image              string
	Description        string
	Tags               []string
	SourceUrl          string   `yaml:"sourceURL"`
	ExamplePackageUrls []string `yaml:"examplePackageURLs"`
	Hidden             bool
}

var (
	// Match start of a version such as v1.9.1
	branchSemverPrefix = regexp.MustCompile(`[-\w]*\/(v\d*\.\d*)`)
	functionDirPrefix  = regexp.MustCompile(`.+/functions/`)
)

func getBranches() ([]string, error) {
	verBranches := make([]string, 0)

	var buf bytes.Buffer
	cmd := exec.Command("git", "branch", "-a")
	cmd.Stdout = &buf
	err := cmd.Run()
	if err != nil {
		return verBranches, err
	}

	for _, branch := range strings.Split(buf.String(), "\n") {
		skipCurrentBranch := false
		for _, branchToSkip := range branchesToSkip {
			if strings.Contains(branch, branchToSkip) {
				skipCurrentBranch = true
				break
			}
		}
		if skipCurrentBranch {
			continue
		}

		if branchSemverPrefix.MatchString(branch) {
			verBranches = append(verBranches, strings.TrimSpace(branch))
		}
	}
	return verBranches, err
}

func getFunctions(branches []string, source string, dest string) []function {
	functions := make(map[string]function)
	for _, b := range branches {
		segments := strings.Split(b, "/")
		funcName := segments[len(segments)-2]
		minorVersion := segments[len(segments)-1]
		relativeFuncPath, err := getRelativeFunctionPath(source, funcName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			continue
		}
		versionDest := filepath.Join(dest, funcName, minorVersion)
		if strings.Contains(relativeFuncPath, "contrib") {
			versionDest = filepath.Join(dest, "contrib", funcName, minorVersion)
		}

		// Functions with the hidden field enabled should not be processed.
		metadataPath := strings.TrimSpace(fmt.Sprintf("%v:%v", b, filepath.Join(relativeFuncPath, "metadata.yaml")))
		md, err := getMetadata(metadataPath)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error getting metadata for %q from %q: %v\n", funcName, b, err)
			os.Exit(1)
		}
		if md.Hidden {
			continue
		}
		err = copyExamples(b, md.ExamplePackageUrls, versionDest, minorVersion)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error getting examples for %q from %q: %v\n", funcName, b, err)
			os.Exit(1)
		}

		err = copyReadme(b, funcName, relativeFuncPath, versionDest)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error getting README for %q from %q: %v\n", funcName, b, err)
			os.Exit(1)
		}

		f := functions[funcName]
		f.FunctionName = funcName
		if f.VersionToExamples == nil {
			f.VersionToExamples = make(map[string]map[string]example)
		}

		functions[funcName] = parseMetadata(f, md, minorVersion, versionDest)
	}

	flattenedFunctions := make([]function, 0)
	for _, f := range functions {
		flattenedFunctions = append(flattenedFunctions, f)
	}
	sort.Slice(flattenedFunctions, func(i int, j int) bool {
		return flattenedFunctions[i].FunctionName < flattenedFunctions[j].FunctionName
	})
	return flattenedFunctions
}

func copyExamples(b string, exampleSources []string, versionDest, minorVersion string) error {
	// Copy examples for the function's version to a temporary directory.
	tempDir, err := ioutil.TempDir("", "examples")
	if err != nil {
		return err
	}

	// Prepare destination for versioned examples.
	err = os.MkdirAll(versionDest, 0744)
	if err != nil {
		return err
	}

	for _, exampleSource := range exampleSources {
		splitedPaths := strings.SplitN(exampleSource, minorVersion+string(filepath.Separator), 2)
		if len(splitedPaths) != 2 {
			return fmt.Errorf("expect 2 substring after spliting %q by %q", exampleSource, minorVersion+string(filepath.Separator))
		}
		relativePath := splitedPaths[1]
		// Fetch example into temporary directory.
		cmd := exec.Command("git", fmt.Sprintf("--work-tree=%v", tempDir), "checkout", b, "--", relativePath)
		err = cmd.Run()
		if err != nil {
			return fmt.Errorf("Error running %v: %v", cmd, err)
		}

		exampleName := filepath.Base(relativePath)
		src := filepath.Join(tempDir, relativePath)
		dest := filepath.Join(versionDest, exampleName)

		// Move example content to the site's example directory.
		if err = os.RemoveAll(dest); err != nil {
			return err
		}
		err = os.Rename(src, dest)
		if err != nil {
			return err
		}

	}

	return nil
}

func copyReadme(b string, funcName string, relativeFuncPath string, versionDest string) error {
	// Copy README for the function's version to the function's directory.
	tempDir, err := ioutil.TempDir("", "functions")
	if err != nil {
		return err
	}
	cmd := exec.Command("git", fmt.Sprintf("--work-tree=%v", tempDir), "checkout", b, "--", filepath.Join(relativeFuncPath, "README.md"))
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("Error running %v: %v", cmd, err)
	}

	// Find the README in the function's directory.
	m, err := filepath.Glob(filepath.Join(tempDir, "functions", "*", funcName, "README.md"))
	if err != nil {
		return err
	}

	if m == nil {
		m, err = filepath.Glob(filepath.Join(tempDir, "contrib", "functions", "*", funcName, "README.md"))
		if err != nil {
			return err
		}
	}

	// Move the README to the destination directory.
	err = os.Rename(m[0], filepath.Join(versionDest, "README.md"))
	if err != nil {
		return err
	}

	return nil
}

func getMetadata(metadataPath string) (metadata, error) {
	var buf bytes.Buffer
	var md metadata
	// Get the content of metadata.yaml from the appropriate release branch.
	cmd := exec.Command("git", "cat-file", "blob", metadataPath)
	cmd.Stdout = &buf
	err := cmd.Run()
	if err != nil {
		return md, err
	}

	yaml.Unmarshal(buf.Bytes(), &md)
	return md, nil
}

func parseMetadata(f function, md metadata, version string, versionDest string) function {

	// Add examples to version map.
	if f.VersionToExamples[version] == nil {
		f.VersionToExamples[version] = make(map[string]example)
	}
	e := f.VersionToExamples[version]
	for _, exUrl := range md.ExamplePackageUrls {
		exUrlSegments := strings.Split(exUrl, "/")
		exName := exUrlSegments[len(exUrlSegments)-1]
		ex := example{
			RemoteSourcePath:  md.SourceUrl,
			RemoteExamplePath: exUrl,
			LocalExamplePath:  filepath.Join(versionDest, exName),
		}
		e[exName] = ex
	}
	f.VersionToExamples[version] = e
	if semver.Compare(f.LatestVersion, version) == 1 {
		return f
	}

	// If this is the latest version,
	// update latest version, default path, description and rags.
	f.LatestVersion = version
	f.Path = versionDest
	f.Description = md.Description
	functionTags := make([]string, 0)
	for _, tag := range md.Tags {
		normalizedTag := strings.ToLower(tag)
		if normalizedTag == "gcp" {
			f.Gcp = true
		} else {
			functionTags = append(functionTags, normalizedTag)
		}
	}
	sort.Strings(functionTags)
	f.Tags = strings.Join(functionTags, ", ")
	f.ImagePath = md.Image

	return f
}

func getRelativeFunctionPath(source string, funcName string) (string, error) {
	// Find the directory for the function's source.
	sourcePattern := filepath.Join(source, "functions", "*", funcName)
	m, err := filepath.Glob(sourcePattern)
	if err != nil {
		return "", err
	}
	if m != nil {
		return functionDirPrefix.ReplaceAllString(m[0], "functions/"), nil
	}

	contribPattern := filepath.Join(source, "contrib", "functions", "*", funcName)
	m, err = filepath.Glob(contribPattern)
	if err != nil {
		return "", err
	}
	if m == nil {
		return "", fmt.Errorf("Could not find a function with the following name: %v", funcName)
	}
	return functionDirPrefix.ReplaceAllString(m[0], "contrib/functions/"), nil
}

func writeFunctionIndex(functions []function, source string, dest string) ([]function, []function, error) {
	out := []string{"# Curated Functions Catalog", ""}
	var contribOut []string

	genericFunctions := make([]function, 0)
	gcp := make([]function, 0)
	contribFunctions := make([]function, 0)
	for _, f := range functions {
		if strings.Contains(f.ImagePath, "contrib") {
			contribFunctions = append(contribFunctions, f)
		} else {
			if f.Gcp {
				gcp = append(gcp, f)
			} else {
				genericFunctions = append(genericFunctions, f)
			}
		}
	}

	out = append(out, getFunctionTable(genericFunctions, source)...)

	if len(gcp) > 0 {
		out = append(out, "", "## GCP Functions", "")
		out = append(out, getFunctionTable(gcp, source)...)
	}

	if len(contribFunctions) > 0 {
		contribOut = append(contribOut, "# Contrib Functions Catalog", "")
		contribOut = append(contribOut, getFunctionTable(contribFunctions, source)...)
	}

	o := strings.Join(out, "\n")
	if err := ioutil.WriteFile(filepath.Join(dest, "README.md"), []byte(o), 0744); err != nil {
		return nil, nil, err
	}
	co := strings.Join(contribOut, "\n")
	if err := ioutil.WriteFile(filepath.Join(dest, "contrib", "README.md"), []byte(co), 0744); err != nil {
		return nil, nil, err
	}
	return append(genericFunctions, gcp...), contribFunctions, nil
}

func getFunctionTable(functions []function, source string) []string {
	out := []string{"| Name | Description | Tags |", "| ---- | ----------- | ---- |"}
	for _, f := range functions {
		functionEntry := fmt.Sprintf("| [%v](%v/) | %v | %v |", f.FunctionName, strings.Replace(f.Path, filepath.Join(source, "site"), "", 1), f.Description, f.Tags)
		out = append(out, functionEntry)
	}
	return out
}

func writeExampleIndex(functions []function, source string, dest string) error {
	// Key a function's version's examples by the function's name -> version
	functionVersionMap := make(map[string]map[string]map[string]example)
	for _, f := range functions {
		vToE := make(map[string]map[string]example)
		for v, examples := range f.VersionToExamples {
			exampleToPaths := make(map[string]example)
			for exName, ex := range examples {
				e := ex
				e.LocalExamplePath = strings.Replace(ex.LocalExamplePath, filepath.Join(source, "site"), "", 1)
				exampleToPaths[exName] = e
			}
			vToE[v] = exampleToPaths
		}
		functionVersionMap[f.FunctionName] = vToE
	}

	funcJson, err := json.Marshal(functionVersionMap)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(dest, "catalog.json"), funcJson, 0600)
	return err
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2017 VMware, Inc. All Rights Reserved.

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

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

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

package simulator

import (
	"github.com/vmware/govmomi/vim25/methods"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

type HostNetworkSystem struct {
	mo.HostNetworkSystem

	Host *mo.HostSystem
}

func NewHostNetworkSystem(host *mo.HostSystem) *HostNetworkSystem {
	return &HostNetworkSystem{
		Host: host,
		HostNetworkSystem: mo.HostNetworkSystem{
			NetworkInfo: &types.HostNetworkInfo{
				Vswitch: []types.HostVirtualSwitch{
					{
						Name:      "vSwitch0",
						Portgroup: []string{"VM Network"},
					},
				},
			},
		},
	}
}

func (s *HostNetworkSystem) folder() *Folder {
	f := Map.getEntityDatacenter(s.Host).NetworkFolder
	return Map.Get(f).(*Folder)
}

func (s *HostNetworkSystem) AddVirtualSwitch(c *types.AddVirtualSwitch) soap.HasFault {
	r := &methods.AddVirtualSwitchBody{}

	for _, vswitch := range s.NetworkInfo.Vswitch {
		if vswitch.Name == c.VswitchName {
			r.Fault_ = Fault("", &types.AlreadyExists{Name: c.VswitchName})
			return r
		}
	}

	s.NetworkInfo.Vswitch = append(s.NetworkInfo.Vswitch, types.HostVirtualSwitch{
		Name: c.VswitchName,
	})

	r.Res = &types.AddVirtualSwitchResponse{}

	return r
}

func (s *HostNetworkSystem) RemoveVirtualSwitch(c *types.RemoveVirtualSwitch) soap.HasFault {
	r := &methods.RemoveVirtualSwitchBody{}

	vs := s.NetworkInfo.Vswitch

	for i, v := range vs {
		if v.Name == c.VswitchName {
			s.NetworkInfo.Vswitch = append(vs[:i], vs[i+1:]...)
			r.Res = &types.RemoveVirtualSwitchResponse{}
			return r
		}
	}

	r.Fault_ = Fault("", &types.NotFound{})

	return r
}

func (s *HostNetworkSystem) AddPortGroup(c *types.AddPortGroup) soap.HasFault {
	var vswitch *types.HostVirtualSwitch

	r := &methods.AddPortGroupBody{}

	if c.Portgrp.Name == "" {
		r.Fault_ = Fault("", &types.InvalidArgument{InvalidProperty: "name"})
		return r
	}

	for i := range s.NetworkInfo.Vswitch {
		if s.NetworkInfo.Vswitch[i].Name == c.Portgrp.VswitchName {
			vswitch = &s.NetworkInfo.Vswitch[i]
			break
		}
	}

	if vswitch == nil {
		r.Fault_ = Fault("", &types.NotFound{})
		return r
	}

	network := &mo.Network{}
	network.Name = c.Portgrp.Name
	network.Entity().Name = network.Name

	folder := s.folder()

	if obj := Map.FindByName(c.Portgrp.Name, folder.ChildEntity); obj != nil {
		r.Fault_ = Fault("", &types.DuplicateName{
			Name:   c.Portgrp.Name,
			Object: obj.Reference(),
		})

		return r
	}

	folder.putChild(network)

	vswitch.Portgroup = append(vswitch.Portgroup, c.Portgrp.Name)
	r.Res = &types.AddPortGroupResponse{}

	return r
}

func (s *HostNetworkSystem) RemovePortGroup(c *types.RemovePortGroup) soap.HasFault {
	var vswitch *types.HostVirtualSwitch

	r := &methods.RemovePortGroupBody{}

	for i, v := range s.NetworkInfo.Vswitch {
		for j, pg := range v.Portgroup {
			if pg == c.PgName {
				vswitch = &s.NetworkInfo.Vswitch[i]
				vswitch.Portgroup = append(vswitch.Portgroup[:j], vswitch.Portgroup[j+1:]...)
			}
		}
	}

	if vswitch == nil {
		r.Fault_ = Fault("", &types.NotFound{})
		return r
	}

	folder := s.folder()
	e := Map.FindByName(c.PgName, folder.ChildEntity)
	folder.removeChild(e.Reference())

	r.Res = &types.RemovePortGroupResponse{}

	return r
}
<commit_msg>vcsim: Add UpdateNetworkConfig to HostNetworkSystem<commit_after>/*
Copyright (c) 2017 VMware, Inc. All Rights Reserved.

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

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

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

package simulator

import (
	"github.com/vmware/govmomi/vim25/methods"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

type HostNetworkSystem struct {
	mo.HostNetworkSystem

	Host *mo.HostSystem
}

func NewHostNetworkSystem(host *mo.HostSystem) *HostNetworkSystem {
	return &HostNetworkSystem{
		Host: host,
		HostNetworkSystem: mo.HostNetworkSystem{
			NetworkInfo: &types.HostNetworkInfo{
				Vswitch: []types.HostVirtualSwitch{
					{
						Name:      "vSwitch0",
						Portgroup: []string{"VM Network"},
					},
				},
			},
		},
	}
}

func (s *HostNetworkSystem) folder() *Folder {
	f := Map.getEntityDatacenter(s.Host).NetworkFolder
	return Map.Get(f).(*Folder)
}

func (s *HostNetworkSystem) AddVirtualSwitch(c *types.AddVirtualSwitch) soap.HasFault {
	r := &methods.AddVirtualSwitchBody{}

	for _, vswitch := range s.NetworkInfo.Vswitch {
		if vswitch.Name == c.VswitchName {
			r.Fault_ = Fault("", &types.AlreadyExists{Name: c.VswitchName})
			return r
		}
	}

	s.NetworkInfo.Vswitch = append(s.NetworkInfo.Vswitch, types.HostVirtualSwitch{
		Name: c.VswitchName,
	})

	r.Res = &types.AddVirtualSwitchResponse{}

	return r
}

func (s *HostNetworkSystem) RemoveVirtualSwitch(c *types.RemoveVirtualSwitch) soap.HasFault {
	r := &methods.RemoveVirtualSwitchBody{}

	vs := s.NetworkInfo.Vswitch

	for i, v := range vs {
		if v.Name == c.VswitchName {
			s.NetworkInfo.Vswitch = append(vs[:i], vs[i+1:]...)
			r.Res = &types.RemoveVirtualSwitchResponse{}
			return r
		}
	}

	r.Fault_ = Fault("", &types.NotFound{})

	return r
}

func (s *HostNetworkSystem) AddPortGroup(c *types.AddPortGroup) soap.HasFault {
	var vswitch *types.HostVirtualSwitch

	r := &methods.AddPortGroupBody{}

	if c.Portgrp.Name == "" {
		r.Fault_ = Fault("", &types.InvalidArgument{InvalidProperty: "name"})
		return r
	}

	for i := range s.NetworkInfo.Vswitch {
		if s.NetworkInfo.Vswitch[i].Name == c.Portgrp.VswitchName {
			vswitch = &s.NetworkInfo.Vswitch[i]
			break
		}
	}

	if vswitch == nil {
		r.Fault_ = Fault("", &types.NotFound{})
		return r
	}

	network := &mo.Network{}
	network.Name = c.Portgrp.Name
	network.Entity().Name = network.Name

	folder := s.folder()

	if obj := Map.FindByName(c.Portgrp.Name, folder.ChildEntity); obj != nil {
		r.Fault_ = Fault("", &types.DuplicateName{
			Name:   c.Portgrp.Name,
			Object: obj.Reference(),
		})

		return r
	}

	folder.putChild(network)

	vswitch.Portgroup = append(vswitch.Portgroup, c.Portgrp.Name)
	r.Res = &types.AddPortGroupResponse{}

	return r
}

func (s *HostNetworkSystem) RemovePortGroup(c *types.RemovePortGroup) soap.HasFault {
	var vswitch *types.HostVirtualSwitch

	r := &methods.RemovePortGroupBody{}

	for i, v := range s.NetworkInfo.Vswitch {
		for j, pg := range v.Portgroup {
			if pg == c.PgName {
				vswitch = &s.NetworkInfo.Vswitch[i]
				vswitch.Portgroup = append(vswitch.Portgroup[:j], vswitch.Portgroup[j+1:]...)
			}
		}
	}

	if vswitch == nil {
		r.Fault_ = Fault("", &types.NotFound{})
		return r
	}

	folder := s.folder()
	e := Map.FindByName(c.PgName, folder.ChildEntity)
	folder.removeChild(e.Reference())

	r.Res = &types.RemovePortGroupResponse{}

	return r
}

func (s *HostNetworkSystem) UpdateNetworkConfig(req *types.UpdateNetworkConfig) soap.HasFault {
	s.NetworkConfig = &req.Config

	return &methods.UpdateNetworkConfigBody{
		Res: &types.UpdateNetworkConfigResponse{
			Returnval: types.HostNetworkConfigResult{},
		},
	}
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"

	"github.com/micro/go-micro/codec"
	raw "github.com/micro/go-micro/codec/bytes"
	"github.com/micro/go-micro/codec/grpc"
	"github.com/micro/go-micro/codec/json"
	"github.com/micro/go-micro/codec/jsonrpc"
	"github.com/micro/go-micro/codec/proto"
	"github.com/micro/go-micro/codec/protorpc"
	"github.com/micro/go-micro/transport"
	"github.com/pkg/errors"
)

type rpcCodec struct {
	socket transport.Socket
	codec  codec.Codec
	first  bool

	req *transport.Message
	buf *readWriteCloser
}

type readWriteCloser struct {
	wbuf *bytes.Buffer
	rbuf *bytes.Buffer
}

var (
	DefaultContentType = "application/protobuf"

	DefaultCodecs = map[string]codec.NewCodec{
		"application/grpc":         grpc.NewCodec,
		"application/grpc+json":    grpc.NewCodec,
		"application/grpc+proto":   grpc.NewCodec,
		"application/json":         json.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/protobuf":     proto.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": raw.NewCodec,
	}

	// TODO: remove legacy codec list
	defaultCodecs = map[string]codec.NewCodec{
		"application/json":         jsonrpc.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/protobuf":     protorpc.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": protorpc.NewCodec,
	}
)

func (rwc *readWriteCloser) Read(p []byte) (n int, err error) {
	return rwc.rbuf.Read(p)
}

func (rwc *readWriteCloser) Write(p []byte) (n int, err error) {
	return rwc.wbuf.Write(p)
}

func (rwc *readWriteCloser) Close() error {
	rwc.rbuf.Reset()
	rwc.wbuf.Reset()
	return nil
}

func getHeader(hdr string, md map[string]string) string {
	if hd := md[hdr]; len(hd) > 0 {
		return hd
	}
	return md["X-"+hdr]
}

func getHeaders(m *codec.Message) {
	get := func(hdr, v string) string {
		if len(v) > 0 {
			return v
		}

		if hd := m.Header[hdr]; len(hd) > 0 {
			return hd
		}

		// old
		return m.Header["X-"+hdr]
	}

	m.Id = get("Micro-Id", m.Id)
	m.Error = get("Micro-Error", m.Error)
	m.Endpoint = get("Micro-Endpoint", m.Endpoint)
	m.Method = get("Micro-Method", m.Method)
	m.Target = get("Micro-Service", m.Target)

	// TODO: remove this cruft
	if len(m.Endpoint) == 0 {
		m.Endpoint = m.Method
	}
}

func setHeaders(m, r *codec.Message) {
	set := func(hdr, v string) {
		if len(v) == 0 {
			return
		}
		m.Header[hdr] = v
		m.Header["X-"+hdr] = v
	}

	// set headers
	set("Micro-Id", r.Id)
	set("Micro-Service", r.Target)
	set("Micro-Method", r.Method)
	set("Micro-Endpoint", r.Endpoint)
	set("Micro-Error", r.Error)
}

// setupProtocol sets up the old protocol
func setupProtocol(msg *transport.Message) codec.NewCodec {
	service := getHeader("Micro-Service", msg.Header)
	method := getHeader("Micro-Method", msg.Header)
	endpoint := getHeader("Micro-Endpoint", msg.Header)
	protocol := getHeader("Micro-Protocol", msg.Header)
	target := getHeader("Micro-Target", msg.Header)

	// if the protocol exists (mucp) do nothing
	if len(protocol) > 0 {
		return nil
	}

	// if no service/method/endpoint then it's the old protocol
	if len(service) == 0 && len(method) == 0 && len(endpoint) == 0 {
		return defaultCodecs[msg.Header["Content-Type"]]
	}

	// old target method specified
	if len(target) > 0 {
		return defaultCodecs[msg.Header["Content-Type"]]
	}

	// no method then set to endpoint
	if len(method) == 0 {
		msg.Header["Micro-Method"] = endpoint
	}

	// no endpoint then set to method
	if len(endpoint) == 0 {
		msg.Header["Micro-Endpoint"] = method
	}

	return nil
}

func newRpcCodec(req *transport.Message, socket transport.Socket, c codec.NewCodec) codec.Codec {
	rwc := &readWriteCloser{
		rbuf: bytes.NewBuffer(req.Body),
		wbuf: bytes.NewBuffer(nil),
	}
	r := &rpcCodec{
		first:  true,
		buf:    rwc,
		codec:  c(rwc),
		req:    req,
		socket: socket,
	}
	return r
}

func (c *rpcCodec) ReadHeader(r *codec.Message, t codec.MessageType) error {
	// the initial message
	m := codec.Message{
		Header: c.req.Header,
		Body:   c.req.Body,
	}

	// if its a follow on request read it
	if !c.first {
		var tm transport.Message

		// read off the socket
		if err := c.socket.Recv(&tm); err != nil {
			return err
		}
		// reset the read buffer
		c.buf.rbuf.Reset()

		// write the body to the buffer
		if _, err := c.buf.rbuf.Write(tm.Body); err != nil {
			return err
		}

		// set the message header
		m.Header = tm.Header
		// set the message body
		m.Body = tm.Body

		// set req
		c.req = &tm
	}

	// no longer first read
	c.first = false

	// set some internal things
	getHeaders(&m)

	// read header via codec
	if err := c.codec.ReadHeader(&m, codec.Request); err != nil {
		return err
	}

	// set message
	*r = m

	return nil
}

func (c *rpcCodec) ReadBody(b interface{}) error {
	// don't read empty body
	if len(c.req.Body) == 0 {
		return nil
	}
	// read raw data
	if v, ok := b.(*raw.Frame); ok {
		v.Data = c.req.Body
		return nil
	}
	// decode the usual way
	return c.codec.ReadBody(b)
}

func (c *rpcCodec) Write(r *codec.Message, b interface{}) error {
	c.buf.wbuf.Reset()

	// create a new message
	m := &codec.Message{
		Target:   r.Target,
		Method:   r.Method,
		Endpoint: r.Endpoint,
		Id:       r.Id,
		Error:    r.Error,
		Type:     r.Type,
		Header:   r.Header,
	}

	if m.Header == nil {
		m.Header = map[string]string{}
	}

	setHeaders(m, r)

	// the body being sent
	var body []byte

	// is it a raw frame?
	if v, ok := b.(*raw.Frame); ok {
		body = v.Data
		// if we have encoded data just send it
	} else if len(r.Body) > 0 {
		body = r.Body
		// write the body to codec
	} else if err := c.codec.Write(m, b); err != nil {
		c.buf.wbuf.Reset()

		// write an error if it failed
		m.Error = errors.Wrapf(err, "Unable to encode body").Error()
		m.Header["X-Micro-Error"] = m.Error
		m.Header["Micro-Error"] = m.Error
		// no body to write
		if err := c.codec.Write(m, nil); err != nil {
			return err
		}
	} else {
		// set the body
		body = c.buf.wbuf.Bytes()
	}

	// Set content type if theres content
	if len(body) > 0 {
		m.Header["Content-Type"] = c.req.Header["Content-Type"]
	}

	// send on the socket
	return c.socket.Send(&transport.Message{
		Header: m.Header,
		Body:   body,
	})
}

func (c *rpcCodec) Close() error {
	c.buf.Close()
	c.codec.Close()
	return c.socket.Close()
}

func (c *rpcCodec) String() string {
	return "rpc"
}
<commit_msg>fallback for 0.14.0 and older<commit_after>package server

import (
	"bytes"

	"github.com/micro/go-micro/codec"
	raw "github.com/micro/go-micro/codec/bytes"
	"github.com/micro/go-micro/codec/grpc"
	"github.com/micro/go-micro/codec/json"
	"github.com/micro/go-micro/codec/jsonrpc"
	"github.com/micro/go-micro/codec/proto"
	"github.com/micro/go-micro/codec/protorpc"
	"github.com/micro/go-micro/transport"
	"github.com/pkg/errors"
)

type rpcCodec struct {
	socket transport.Socket
	codec  codec.Codec
	first  bool

	req *transport.Message
	buf *readWriteCloser
}

type readWriteCloser struct {
	wbuf *bytes.Buffer
	rbuf *bytes.Buffer
}

var (
	DefaultContentType = "application/protobuf"

	DefaultCodecs = map[string]codec.NewCodec{
		"application/grpc":         grpc.NewCodec,
		"application/grpc+json":    grpc.NewCodec,
		"application/grpc+proto":   grpc.NewCodec,
		"application/json":         json.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/protobuf":     proto.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": raw.NewCodec,
	}

	// TODO: remove legacy codec list
	defaultCodecs = map[string]codec.NewCodec{
		"application/json":         jsonrpc.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/protobuf":     protorpc.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": protorpc.NewCodec,
	}
)

func (rwc *readWriteCloser) Read(p []byte) (n int, err error) {
	return rwc.rbuf.Read(p)
}

func (rwc *readWriteCloser) Write(p []byte) (n int, err error) {
	return rwc.wbuf.Write(p)
}

func (rwc *readWriteCloser) Close() error {
	rwc.rbuf.Reset()
	rwc.wbuf.Reset()
	return nil
}

func getHeader(hdr string, md map[string]string) string {
	if hd := md[hdr]; len(hd) > 0 {
		return hd
	}
	return md["X-"+hdr]
}

func getHeaders(m *codec.Message) {
	get := func(hdr, v string) string {
		if len(v) > 0 {
			return v
		}

		if hd := m.Header[hdr]; len(hd) > 0 {
			return hd
		}

		// old
		return m.Header["X-"+hdr]
	}

	m.Id = get("Micro-Id", m.Id)
	m.Error = get("Micro-Error", m.Error)
	m.Endpoint = get("Micro-Endpoint", m.Endpoint)
	m.Method = get("Micro-Method", m.Method)
	m.Target = get("Micro-Service", m.Target)

	// TODO: remove this cruft
	if len(m.Endpoint) == 0 {
		m.Endpoint = m.Method
	}
}

func setHeaders(m, r *codec.Message) {
	set := func(hdr, v string) {
		if len(v) == 0 {
			return
		}
		m.Header[hdr] = v
		m.Header["X-"+hdr] = v
	}

	// set headers
	set("Micro-Id", r.Id)
	set("Micro-Service", r.Target)
	set("Micro-Method", r.Method)
	set("Micro-Endpoint", r.Endpoint)
	set("Micro-Error", r.Error)
}

// setupProtocol sets up the old protocol
func setupProtocol(msg *transport.Message) codec.NewCodec {
	service := getHeader("Micro-Service", msg.Header)
	method := getHeader("Micro-Method", msg.Header)
	endpoint := getHeader("Micro-Endpoint", msg.Header)
	protocol := getHeader("Micro-Protocol", msg.Header)
	target := getHeader("Micro-Target", msg.Header)

	// if the protocol exists (mucp) do nothing
	if len(protocol) > 0 {
		return nil
	}

	// if no service/method/endpoint then it's the old protocol
	if len(service) == 0 && len(method) == 0 && len(endpoint) == 0 {
		return defaultCodecs[msg.Header["Content-Type"]]
	}

	// old target method specified
	if len(target) > 0 {
		return defaultCodecs[msg.Header["Content-Type"]]
	}

	// no method then set to endpoint
	if len(method) == 0 {
		msg.Header["Micro-Method"] = endpoint
	}

	// no endpoint then set to method
	if len(endpoint) == 0 {
		msg.Header["Micro-Endpoint"] = method
	}

	return nil
}

func newRpcCodec(req *transport.Message, socket transport.Socket, c codec.NewCodec) codec.Codec {
	rwc := &readWriteCloser{
		rbuf: bytes.NewBuffer(req.Body),
		wbuf: bytes.NewBuffer(nil),
	}
	r := &rpcCodec{
		first:  true,
		buf:    rwc,
		codec:  c(rwc),
		req:    req,
		socket: socket,
	}
	return r
}

func (c *rpcCodec) ReadHeader(r *codec.Message, t codec.MessageType) error {
	// the initial message
	m := codec.Message{
		Header: c.req.Header,
		Body:   c.req.Body,
	}

	// if its a follow on request read it
	if !c.first {
		var tm transport.Message

		// read off the socket
		if err := c.socket.Recv(&tm); err != nil {
			return err
		}
		// reset the read buffer
		c.buf.rbuf.Reset()

		// write the body to the buffer
		if _, err := c.buf.rbuf.Write(tm.Body); err != nil {
			return err
		}

		// set the message header
		m.Header = tm.Header
		// set the message body
		m.Body = tm.Body

		// set req
		c.req = &tm
	}

	// no longer first read
	c.first = false

	// set some internal things
	getHeaders(&m)

	// read header via codec
	if err := c.codec.ReadHeader(&m, codec.Request); err != nil {
		return err
	}

	// fallback for 0.14 and older
	if len(m.Endpoint) == 0 {
		m.Endpoint = m.Method
	}

	// set message
	*r = m

	return nil
}

func (c *rpcCodec) ReadBody(b interface{}) error {
	// don't read empty body
	if len(c.req.Body) == 0 {
		return nil
	}
	// read raw data
	if v, ok := b.(*raw.Frame); ok {
		v.Data = c.req.Body
		return nil
	}
	// decode the usual way
	return c.codec.ReadBody(b)
}

func (c *rpcCodec) Write(r *codec.Message, b interface{}) error {
	c.buf.wbuf.Reset()

	// create a new message
	m := &codec.Message{
		Target:   r.Target,
		Method:   r.Method,
		Endpoint: r.Endpoint,
		Id:       r.Id,
		Error:    r.Error,
		Type:     r.Type,
		Header:   r.Header,
	}

	if m.Header == nil {
		m.Header = map[string]string{}
	}

	setHeaders(m, r)

	// the body being sent
	var body []byte

	// is it a raw frame?
	if v, ok := b.(*raw.Frame); ok {
		body = v.Data
		// if we have encoded data just send it
	} else if len(r.Body) > 0 {
		body = r.Body
		// write the body to codec
	} else if err := c.codec.Write(m, b); err != nil {
		c.buf.wbuf.Reset()

		// write an error if it failed
		m.Error = errors.Wrapf(err, "Unable to encode body").Error()
		m.Header["X-Micro-Error"] = m.Error
		m.Header["Micro-Error"] = m.Error
		// no body to write
		if err := c.codec.Write(m, nil); err != nil {
			return err
		}
	} else {
		// set the body
		body = c.buf.wbuf.Bytes()
	}

	// Set content type if theres content
	if len(body) > 0 {
		m.Header["Content-Type"] = c.req.Header["Content-Type"]
	}

	// send on the socket
	return c.socket.Send(&transport.Message{
		Header: m.Header,
		Body:   body,
	})
}

func (c *rpcCodec) Close() error {
	c.buf.Close()
	c.codec.Close()
	return c.socket.Close()
}

func (c *rpcCodec) String() string {
	return "rpc"
}
<|endoftext|>