<commit_before>// Copyright 2015/2016 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sysparser

import (
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
)

type Description struct {
	Includes  []string
	Defines   map[string]string
	Syscalls  []Syscall
	Structs   map[string]Struct
	Unnamed   map[string][]string
	Flags     map[string][]string
	Resources map[string]Resource
}

type Syscall struct {
	Name     string
	CallName string
	Args     [][]string
	Ret      []string
}

type Struct struct {
	Name    string
	Flds    [][]string
	IsUnion bool
	Packed  bool
	Varlen  bool
	Align   int
}

type Resource struct {
	Name   string
	Base   string
	Values []string
}

func Parse(in io.Reader) *Description {
	p := newParser(in)
	var includes []string
	defines := make(map[string]string)
	var syscalls []Syscall
	structs := make(map[string]Struct)
	unnamed := make(map[string][]string)
	flags := make(map[string][]string)
	resources := make(map[string]Resource)
	var str *Struct
	for p.Scan() {
		if p.EOF() || p.Char() == '#' {
			continue
		}
		if str != nil {
			// Parsing a struct.
			if p.Char() == '}' || p.Char() == ']' {
				p.Parse(p.Char())
				for _, attr := range parseType1(p, unnamed, flags, "")[1:] {
					if str.IsUnion {
						switch attr {
						case "varlen":
							str.Varlen = true
						default:
							failf("unknown union %v attribute: %v", str.Name, attr)
						}
					} else {
						switch attr {
						case "packed":
							str.Packed = true
						case "align_1":
							str.Align = 1
						case "align_2":
							str.Align = 2
						case "align_4":
							str.Align = 4
						case "align_8":
							str.Align = 8
						default:
							failf("unknown struct %v attribute: %v", str.Name, attr)
						}
					}
				}
				if str.IsUnion {
					if len(str.Flds) <= 1 {
						failf("union %v has only %v fields, need at least 2", str.Name, len(str.Flds))
					}
				}
				fields := make(map[string]bool)
				for _, f := range str.Flds {
					if fields[f[0]] {
						failf("duplicate field %v in struct/union %v", f[0], str.Name)
					}
					fields[f[0]] = true
				}
				structs[str.Name] = *str
				str = nil
			} else {
				p.SkipWs()
				fld := []string{p.Ident()}
				fld = append(fld, parseType(p, unnamed, flags)...)
				str.Flds = append(str.Flds, fld)
			}
		} else {
			name := p.Ident()
			if name == "include" {
				p.Parse('<')
				var include []byte
				for {
					ch := p.Char()
					if ch == '>' {
						break
					}
					p.Parse(ch)
					include = append(include, ch)
				}
				p.Parse('>')
				includes = append(includes, string(include))
			} else if name == "define" {
				key := p.Ident()
				var val []byte
				for !p.EOF() {
					ch := p.Char()
					p.Parse(ch)
					val = append(val, ch)
				}
				if defines[key] != "" {
					failf("%v define is defined multiple times", key)
				}
				defines[key] = fmt.Sprintf("(%s)", val)
			} else if name == "resource" {
				p.SkipWs()
				id := p.Ident()
				p.Parse('[')
				base := p.Ident()
				p.Parse(']')
				var vals []string
				if !p.EOF() && p.Char() == ':' {
					p.Parse(':')
					vals = append(vals, p.Ident())
					for !p.EOF() {
						p.Parse(',')
						vals = append(vals, p.Ident())
					}
				}
				if _, ok := resources[id]; ok {
					failf("resource '%v' is defined multiple times", id)
				}
				if _, ok := structs[id]; ok {
					failf("struct '%v' is redefined as resource", name)
				}
				resources[id] = Resource{id, base, vals}
			} else {
				switch ch := p.Char(); ch {
				case '(':
					// syscall
					p.Parse('(')
					var args [][]string
					for p.Char() != ')' {
						arg := []string{p.Ident()}
						arg = append(arg, parseType(p, unnamed, flags)...)
						args = append(args, arg)
						if p.Char() != ')' {
							p.Parse(',')
						}
					}
					p.Parse(')')
					var ret []string
					if !p.EOF() {
						ret = parseType(p, unnamed, flags)
					}
					callName := name
					if idx := strings.IndexByte(callName, '$'); idx != -1 {
						callName = callName[:idx]
					}
					fields := make(map[string]bool)
					for _, a := range args {
						if fields[a[0]] {
							failf("duplicate arg %v in syscall %v", a[0], name)
						}
						fields[a[0]] = true
					}
					syscalls = append(syscalls, Syscall{name, callName, args, ret})
				case '=':
					// flag
					p.Parse('=')
					vals := []string{p.Ident()}
					for !p.EOF() {
						p.Parse(',')
						vals = append(vals, p.Ident())
					}
					flags[name] = vals
				case '{', '[':
					p.Parse(ch)
					if _, ok := structs[name]; ok {
						failf("struct '%v' is defined multiple times", name)
					}
					if _, ok := resources[name]; ok {
						failf("resource '%v' is redefined as struct", name)
					}
					str = &Struct{Name: name, IsUnion: ch == '['}
				default:
					failf("bad line (%v)", p.Str())
				}
			}
		}
		if !p.EOF() {
			failf("trailing data (%v)", p.Str())
		}
	}
	sort.Sort(syscallArray(syscalls))
	return &Description{
		Includes:  includes,
		Defines:   defines,
		Syscalls:  syscalls,
		Structs:   structs,
		Unnamed:   unnamed,
		Flags:     flags,
		Resources: resources,
	}
}

func isIdentifier(s string) bool {
	for i, c := range s {
		if c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || i > 0 && (c >= '0' && c <= '9') {
			continue
		}
		return false
	}
	return true
}

func parseType(p *parser, unnamed map[string][]string, flags map[string][]string) []string {
	return parseType1(p, unnamed, flags, p.Ident())
}

var (
	unnamedSeq int
	constSeq   int
)

func parseType1(p *parser, unnamed map[string][]string, flags map[string][]string, name string) []string {
	typ := []string{name}
	if !p.EOF() && p.Char() == '[' {
		p.Parse('[')
		for {
			id := p.Ident()
			if p.Char() == '[' {
				inner := parseType1(p, unnamed, flags, id)
				id = fmt.Sprintf("unnamed%v", unnamedSeq)
				unnamedSeq++
				unnamed[id] = inner
			}
			typ = append(typ, id)
			if p.Char() == ']' {
				break
			}
			p.Parse(',')
		}
		p.Parse(']')
	}
	if name == "const" && len(typ) > 1 {
		// Create a fake flag with the const value.
		id := fmt.Sprintf("const_flag_%v", constSeq)
		constSeq++
		flags[id] = typ[1:2]
	}
	if name == "array" && len(typ) > 2 {
		// Create a fake flag with the const value.
		id := fmt.Sprintf("const_flag_%v", constSeq)
		constSeq++
		flags[id] = typ[2:3]
	}
	return typ
}

type syscallArray []Syscall

func (a syscallArray) Len() int           { return len(a) }
func (a syscallArray) Less(i, j int) bool { return a[i].Name < a[j].Name }
func (a syscallArray) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func failf(msg string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, msg+"\n", args...)
	os.Exit(1)
}
<commit_msg>sysparser: check for reserved parent fields<commit_after>// Copyright 2015/2016 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sysparser

import (
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
)

type Description struct {
	Includes  []string
	Defines   map[string]string
	Syscalls  []Syscall
	Structs   map[string]Struct
	Unnamed   map[string][]string
	Flags     map[string][]string
	Resources map[string]Resource
}

type Syscall struct {
	Name     string
	CallName string
	Args     [][]string
	Ret      []string
}

type Struct struct {
	Name    string
	Flds    [][]string
	IsUnion bool
	Packed  bool
	Varlen  bool
	Align   int
}

type Resource struct {
	Name   string
	Base   string
	Values []string
}

func Parse(in io.Reader) *Description {
	p := newParser(in)
	var includes []string
	defines := make(map[string]string)
	var syscalls []Syscall
	structs := make(map[string]Struct)
	unnamed := make(map[string][]string)
	flags := make(map[string][]string)
	resources := make(map[string]Resource)
	var str *Struct
	for p.Scan() {
		if p.EOF() || p.Char() == '#' {
			continue
		}
		if str != nil {
			// Parsing a struct.
			if p.Char() == '}' || p.Char() == ']' {
				p.Parse(p.Char())
				for _, attr := range parseType1(p, unnamed, flags, "")[1:] {
					if str.IsUnion {
						switch attr {
						case "varlen":
							str.Varlen = true
						default:
							failf("unknown union %v attribute: %v", str.Name, attr)
						}
					} else {
						switch attr {
						case "packed":
							str.Packed = true
						case "align_1":
							str.Align = 1
						case "align_2":
							str.Align = 2
						case "align_4":
							str.Align = 4
						case "align_8":
							str.Align = 8
						default:
							failf("unknown struct %v attribute: %v", str.Name, attr)
						}
					}
				}
				if str.IsUnion {
					if len(str.Flds) <= 1 {
						failf("union %v has only %v fields, need at least 2", str.Name, len(str.Flds))
					}
				}
				fields := make(map[string]bool)
				for _, f := range str.Flds {
					if f[0] == "parent" {
						failf("struct/union %v contains reserved field 'parent'", str.Name)
					}
					if fields[f[0]] {
						failf("duplicate field %v in struct/union %v", f[0], str.Name)
					}
					fields[f[0]] = true
				}
				structs[str.Name] = *str
				str = nil
			} else {
				p.SkipWs()
				fld := []string{p.Ident()}
				fld = append(fld, parseType(p, unnamed, flags)...)
				str.Flds = append(str.Flds, fld)
			}
		} else {
			name := p.Ident()
			if name == "include" {
				p.Parse('<')
				var include []byte
				for {
					ch := p.Char()
					if ch == '>' {
						break
					}
					p.Parse(ch)
					include = append(include, ch)
				}
				p.Parse('>')
				includes = append(includes, string(include))
			} else if name == "define" {
				key := p.Ident()
				var val []byte
				for !p.EOF() {
					ch := p.Char()
					p.Parse(ch)
					val = append(val, ch)
				}
				if defines[key] != "" {
					failf("%v define is defined multiple times", key)
				}
				defines[key] = fmt.Sprintf("(%s)", val)
			} else if name == "resource" {
				p.SkipWs()
				id := p.Ident()
				p.Parse('[')
				base := p.Ident()
				p.Parse(']')
				var vals []string
				if !p.EOF() && p.Char() == ':' {
					p.Parse(':')
					vals = append(vals, p.Ident())
					for !p.EOF() {
						p.Parse(',')
						vals = append(vals, p.Ident())
					}
				}
				if _, ok := resources[id]; ok {
					failf("resource '%v' is defined multiple times", id)
				}
				if _, ok := structs[id]; ok {
					failf("struct '%v' is redefined as resource", name)
				}
				resources[id] = Resource{id, base, vals}
			} else {
				switch ch := p.Char(); ch {
				case '(':
					// syscall
					p.Parse('(')
					var args [][]string
					for p.Char() != ')' {
						arg := []string{p.Ident()}
						arg = append(arg, parseType(p, unnamed, flags)...)
						args = append(args, arg)
						if p.Char() != ')' {
							p.Parse(',')
						}
					}
					p.Parse(')')
					var ret []string
					if !p.EOF() {
						ret = parseType(p, unnamed, flags)
					}
					callName := name
					if idx := strings.IndexByte(callName, '$'); idx != -1 {
						callName = callName[:idx]
					}
					fields := make(map[string]bool)
					for _, a := range args {
						if fields[a[0]] {
							failf("duplicate arg %v in syscall %v", a[0], name)
						}
						fields[a[0]] = true
					}
					syscalls = append(syscalls, Syscall{name, callName, args, ret})
				case '=':
					// flag
					p.Parse('=')
					vals := []string{p.Ident()}
					for !p.EOF() {
						p.Parse(',')
						vals = append(vals, p.Ident())
					}
					flags[name] = vals
				case '{', '[':
					p.Parse(ch)
					if _, ok := structs[name]; ok {
						failf("struct '%v' is defined multiple times", name)
					}
					if _, ok := resources[name]; ok {
						failf("resource '%v' is redefined as struct", name)
					}
					str = &Struct{Name: name, IsUnion: ch == '['}
				default:
					failf("bad line (%v)", p.Str())
				}
			}
		}
		if !p.EOF() {
			failf("trailing data (%v)", p.Str())
		}
	}
	sort.Sort(syscallArray(syscalls))
	return &Description{
		Includes:  includes,
		Defines:   defines,
		Syscalls:  syscalls,
		Structs:   structs,
		Unnamed:   unnamed,
		Flags:     flags,
		Resources: resources,
	}
}

func isIdentifier(s string) bool {
	for i, c := range s {
		if c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || i > 0 && (c >= '0' && c <= '9') {
			continue
		}
		return false
	}
	return true
}

func parseType(p *parser, unnamed map[string][]string, flags map[string][]string) []string {
	return parseType1(p, unnamed, flags, p.Ident())
}

var (
	unnamedSeq int
	constSeq   int
)

func parseType1(p *parser, unnamed map[string][]string, flags map[string][]string, name string) []string {
	typ := []string{name}
	if !p.EOF() && p.Char() == '[' {
		p.Parse('[')
		for {
			id := p.Ident()
			if p.Char() == '[' {
				inner := parseType1(p, unnamed, flags, id)
				id = fmt.Sprintf("unnamed%v", unnamedSeq)
				unnamedSeq++
				unnamed[id] = inner
			}
			typ = append(typ, id)
			if p.Char() == ']' {
				break
			}
			p.Parse(',')
		}
		p.Parse(']')
	}
	if name == "const" && len(typ) > 1 {
		// Create a fake flag with the const value.
		id := fmt.Sprintf("const_flag_%v", constSeq)
		constSeq++
		flags[id] = typ[1:2]
	}
	if name == "array" && len(typ) > 2 {
		// Create a fake flag with the const value.
		id := fmt.Sprintf("const_flag_%v", constSeq)
		constSeq++
		flags[id] = typ[2:3]
	}
	return typ
}

type syscallArray []Syscall

func (a syscallArray) Len() int           { return len(a) }
func (a syscallArray) Less(i, j int) bool { return a[i].Name < a[j].Name }
func (a syscallArray) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func failf(msg string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, msg+"\n", args...)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"reflect"
	"time"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/ec2"
)

func resourceAwsSecurityGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSecurityGroupCreate,
		Read:   resourceAwsSecurityGroupRead,
		Update: resourceAwsSecurityGroupUpdate,
		Delete: resourceAwsSecurityGroupDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"ingress": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"from_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"to_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"cidr_blocks": &schema.Schema{
							Type:     schema.TypeList,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},

						"security_groups": &schema.Schema{
							Type:     schema.TypeList,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},
					},
				},
			},

			"owner_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsSecurityGroupCreate(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	securityGroupOpts := ec2.SecurityGroup{
		Name: d.Get("name").(string),
	}

	if v := d.Get("vpc_id"); v != nil {
		securityGroupOpts.VpcId = v.(string)
	}

	if v := d.Get("description"); v != nil {
		securityGroupOpts.Description = v.(string)
	}

	log.Printf(
		"[DEBUG] Security Group create configuration: %#v", securityGroupOpts)
	createResp, err := ec2conn.CreateSecurityGroup(securityGroupOpts)
	if err != nil {
		return fmt.Errorf("Error creating Security Group: %s", err)
	}

	d.SetId(createResp.Id)
	group := createResp.SecurityGroup

	log.Printf("[INFO] Security Group ID: %s", d.Id())

	// Wait for the security group to truly exist
	log.Printf(
		"[DEBUG] Waiting for Security Group (%s) to exist",
		d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{""},
		Target:  "exists",
		Refresh: SGStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for Security Group (%s) to become available: %s",
			d.Id(), err)
	}

	// Expand the "ingress" array to goamz compat []ec2.IPPerm
	ingressRaw := d.Get("ingress")
	if ingressRaw == nil {
		ingressRaw = []interface{}{}
	}
	ingressList := ingressRaw.([]interface{})
	if len(ingressList) > 0 {
		ingressRules := expandIPPerms(ingressList)
		_, err = ec2conn.AuthorizeSecurityGroup(group, ingressRules)
		if err != nil {
			return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
		}
	}

	return resourceAwsSecurityGroupRead(d, meta)
}

func resourceAwsSecurityGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}
	group := sgRaw.(*ec2.SecurityGroupInfo).SecurityGroup

	if d.HasChange("ingress") {
		o, n := d.GetChange("ingress")
		if o == nil {
			o = []interface{}{}
		}
		if n == nil {
			n = []interface{}{}
		}

		oldRules := expandIPPerms(o.([]interface{}))
		newRules := expandIPPerms(n.([]interface{}))

		var add, remove []ec2.IPPerm
		for _, p := range newRules {
			// Check if we have had this rule before
			exists := false
			for _, old := range oldRules {
				if reflect.DeepEqual(old, p) {
					exists = true
					break
				}
			}
			if exists {
				continue
			}
			add = append(add, p)
		}
		for _, p := range oldRules {
			// Check if we have this rule to add
			exists := false
			for _, n := range newRules {
				if reflect.DeepEqual(n, p) {
					exists = true
					break
				}
			}
			if exists {
				continue
			}
			remove = append(remove, p)
		}

		// TODO: We need to handle partial state better in the in-between
		// in this update.

		// Authorize the new rules
		_, err := ec2conn.AuthorizeSecurityGroup(group, add)
		if err != nil {
			return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
		}

		// Revoke the old rules
		_, err = ec2conn.RevokeSecurityGroup(group, remove)
		if err != nil {
			return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
		}
	}

	return nil
}

func resourceAwsSecurityGroupDelete(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	log.Printf("[DEBUG] Security Group destroy: %v", d.Id())

	_, err := ec2conn.DeleteSecurityGroup(ec2.SecurityGroup{Id: d.Id()})
	if err != nil {
		ec2err, ok := err.(*ec2.Error)
		if ok && ec2err.Code == "InvalidGroup.NotFound" {
			return nil
		}
	}

	return err
}

func resourceAwsSecurityGroupRead(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}

	sg := sgRaw.(*ec2.SecurityGroupInfo)

	var deps []terraform.ResourceDependency

	// Gather our ingress rules
	ingressRules := make([]map[string]interface{}, len(sg.IPPerms))
	for i, perm := range sg.IPPerms {
		n := make(map[string]interface{})
		n["from_port"] = perm.FromPort
		n["protocol"] = perm.Protocol
		n["to_port"] = perm.ToPort

		if len(perm.SourceIPs) > 0 {
			n["cidr_blocks"] = perm.SourceIPs
		}

		if len(perm.SourceGroups) > 0 {
			// We depend on other security groups
			for _, v := range perm.SourceGroups {
				deps = append(deps,
					terraform.ResourceDependency{ID: v.Id},
				)
			}

			n["security_groups"] = flattenSecurityGroups(perm.SourceGroups)
		}

		ingressRules[i] = n
	}

	if v := d.Get("vpc_id"); v != nil && v.(string) != "" {
		deps = append(deps,
			terraform.ResourceDependency{ID: v.(string)},
		)
	}

	d.Set("description", sg.Description)
	d.Set("name", sg.Name)
	d.Set("vpc_id", sg.VpcId)
	d.Set("owner_id", sg.OwnerId)
	d.Set("ingress", ingressRules)
	d.SetDependencies(deps)

	return nil
}

// SGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a security group.
func SGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		sgs := []ec2.SecurityGroup{ec2.SecurityGroup{Id: id}}
		resp, err := conn.SecurityGroups(sgs, nil)
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok {
				if ec2err.Code == "InvalidSecurityGroupID.NotFound" ||
					ec2err.Code == "InvalidGroup.NotFound" {
					resp = nil
					err = nil
				}
			}

			if err != nil {
				log.Printf("Error on SGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			return nil, "", nil
		}

		group := &resp.Groups[0]
		return group, "exists", nil
	}
}
<commit_msg>providers/aws: security group ingress rules treated as set [GH-87]<commit_after>package aws

import (
	"bytes"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/ec2"
)

func resourceAwsSecurityGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSecurityGroupCreate,
		Read:   resourceAwsSecurityGroupRead,
		Update: resourceAwsSecurityGroupUpdate,
		Delete: resourceAwsSecurityGroupDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"ingress": &schema.Schema{
				Type:     schema.TypeSet,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"from_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"to_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"cidr_blocks": &schema.Schema{
							Type:     schema.TypeList,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},

						"security_groups": &schema.Schema{
							Type:     schema.TypeList,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},
					},
				},
				Set: resourceAwsSecurityGroupIngressHash,
			},

			"owner_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsSecurityGroupIngressHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["from_port"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["to_port"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["protocol"].(string)))

	if v, ok := m["cidr_blocks"]; ok {
		for _, raw := range v.([]interface{}) {
			buf.WriteString(fmt.Sprintf("%s-", raw.(string)))
		}
	}
	if v, ok := m["security_groups"]; ok {
		for _, raw := range v.([]interface{}) {
			buf.WriteString(fmt.Sprintf("%s-", raw.(string)))
		}
	}

	return hashcode.String(buf.String())
}

func resourceAwsSecurityGroupCreate(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	securityGroupOpts := ec2.SecurityGroup{
		Name: d.Get("name").(string),
	}

	if v := d.Get("vpc_id"); v != nil {
		securityGroupOpts.VpcId = v.(string)
	}

	if v := d.Get("description"); v != nil {
		securityGroupOpts.Description = v.(string)
	}

	log.Printf(
		"[DEBUG] Security Group create configuration: %#v", securityGroupOpts)
	createResp, err := ec2conn.CreateSecurityGroup(securityGroupOpts)
	if err != nil {
		return fmt.Errorf("Error creating Security Group: %s", err)
	}

	d.SetId(createResp.Id)
	group := createResp.SecurityGroup

	log.Printf("[INFO] Security Group ID: %s", d.Id())

	// Wait for the security group to truly exist
	log.Printf(
		"[DEBUG] Waiting for Security Group (%s) to exist",
		d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{""},
		Target:  "exists",
		Refresh: SGStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for Security Group (%s) to become available: %s",
			d.Id(), err)
	}

	// Expand the "ingress" array to goamz compat []ec2.IPPerm
	ingressRaw := d.Get("ingress")
	if ingressRaw == nil {
		ingressRaw = new(schema.Set)
	}
	ingressList := ingressRaw.(*schema.Set).List()
	if len(ingressList) > 0 {
		ingressRules := expandIPPerms(ingressList)
		_, err = ec2conn.AuthorizeSecurityGroup(group, ingressRules)
		if err != nil {
			return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
		}
	}

	return resourceAwsSecurityGroupRead(d, meta)
}

func resourceAwsSecurityGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}
	group := sgRaw.(*ec2.SecurityGroupInfo).SecurityGroup

	if d.HasChange("ingress") {
		o, n := d.GetChange("ingress")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		remove := expandIPPerms(os.Difference(ns).List())
		add := expandIPPerms(ns.Difference(os).List())

		// TODO: We need to handle partial state better in the in-between
		// in this update.
		if len(add) > 0 {
			// Authorize the new rules
			_, err := ec2conn.AuthorizeSecurityGroup(group, add)
			if err != nil {
				return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
			}
		}

		if len(remove) > 0 {
			// Revoke the old rules
			_, err = ec2conn.RevokeSecurityGroup(group, remove)
			if err != nil {
				return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
			}
		}
	}

	return nil
}

func resourceAwsSecurityGroupDelete(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	log.Printf("[DEBUG] Security Group destroy: %v", d.Id())

	_, err := ec2conn.DeleteSecurityGroup(ec2.SecurityGroup{Id: d.Id()})
	if err != nil {
		ec2err, ok := err.(*ec2.Error)
		if ok && ec2err.Code == "InvalidGroup.NotFound" {
			return nil
		}
	}

	return err
}

func resourceAwsSecurityGroupRead(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}

	sg := sgRaw.(*ec2.SecurityGroupInfo)

	var deps []terraform.ResourceDependency

	// Gather our ingress rules
	ingressRules := make([]map[string]interface{}, len(sg.IPPerms))
	for i, perm := range sg.IPPerms {
		n := make(map[string]interface{})
		n["from_port"] = perm.FromPort
		n["protocol"] = perm.Protocol
		n["to_port"] = perm.ToPort

		if len(perm.SourceIPs) > 0 {
			n["cidr_blocks"] = perm.SourceIPs
		}

		if len(perm.SourceGroups) > 0 {
			// We depend on other security groups
			for _, v := range perm.SourceGroups {
				deps = append(deps,
					terraform.ResourceDependency{ID: v.Id},
				)
			}

			n["security_groups"] = flattenSecurityGroups(perm.SourceGroups)
		}

		ingressRules[i] = n
	}

	if v := d.Get("vpc_id"); v != nil && v.(string) != "" {
		deps = append(deps,
			terraform.ResourceDependency{ID: v.(string)},
		)
	}

	d.Set("description", sg.Description)
	d.Set("name", sg.Name)
	d.Set("vpc_id", sg.VpcId)
	d.Set("owner_id", sg.OwnerId)
	d.Set("ingress", ingressRules)
	d.SetDependencies(deps)

	return nil
}

// SGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a security group.
func SGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		sgs := []ec2.SecurityGroup{ec2.SecurityGroup{Id: id}}
		resp, err := conn.SecurityGroups(sgs, nil)
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok {
				if ec2err.Code == "InvalidSecurityGroupID.NotFound" ||
					ec2err.Code == "InvalidGroup.NotFound" {
					resp = nil
					err = nil
				}
			}

			if err != nil {
				log.Printf("Error on SGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			return nil, "", nil
		}

		group := &resp.Groups[0]
		return group, "exists", nil
	}
}
<|endoftext|>


<commit_before>package collector

import (
	"fmt"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/zwopir/osquery_exporter/model"
	"github.com/zwopir/osquery_exporter/osquery"
	"strconv"
	"sync"
)

// singleQueryCollector represents a metric/query definition for a single osquery call
type singleQueryCollector interface {
	String() string
	Id() string
	Query() string
	Desc() *prometheus.Desc
	ValueType() prometheus.ValueType
	Value() string
	Labels() []string
}

// update maps the osquery query result to the singleQueryCollector and updates the provided channel accordingly
func update(sqc singleQueryCollector, result *model.OsqueryResult, ch chan<- prometheus.Metric) error {
	log.Debugf("updating metric %q", sqc.String())
	for _, item := range result.Items {
		value, ok := item[sqc.Value()]
		if !ok {
			return fmt.Errorf("query %q doesn't contain value key %q", sqc.Query(), sqc.Value())
		}
		valueAsFloat, err := strconv.ParseFloat(value, 64)
		if err != nil {
			fmt.Errorf("query %q result %q can't be converted to float", sqc.Query(), value)
		}
		labels := []string{}
		for _, labelIdentifier := range sqc.Labels() {
			if label, ok := item[labelIdentifier]; ok {
				labels = append(labels, label)
			} else {
				return fmt.Errorf("query %q doesn't contain a label key %q", sqc.Query(), labelIdentifier)
			}
		}
		ch <- prometheus.MustNewConstMetric(
			sqc.Desc(),
			sqc.ValueType(),
			valueAsFloat,
			labels...,
		)
	}
	return nil
}

// OsqueryCollector represents a collector that collects metrics from a set of osquery queries. It implements
// prometheus Collector
type OsqueryCollector struct {
	runner         *osquery.OsqueryRunner
	collectors     map[string]singleQueryCollector
	queryDurations *prometheus.SummaryVec
	success        *prometheus.GaugeVec
	resultsets     *prometheus.GaugeVec
}

// NewOsqueryCollector creates an OsQueryCollector from a given osquery-runner and a set of metric definitions
func NewOsqueryCollector(r *osquery.OsqueryRunner, m model.Metrics) *OsqueryCollector {
	collectors := make(map[string]singleQueryCollector)
	for _, c := range m.Counters {
		log.Infof("adding %s to OsqueryCollector", c.String())
		collectors[c.Id()] = c
	}
	for _, cv := range m.CounterVecs {
		log.Infof("adding %s to OsqueryCollector", cv.String())
		collectors[cv.Id()] = cv
	}
	for _, g := range m.Gauges {
		log.Infof("adding %s to OsqueryCollector", g.String())
		collectors[g.Id()] = g
	}
	for _, gv := range m.GaugeVecs {
		log.Infof("adding %s to OsqueryCollector", gv.String())
		collectors[gv.Id()] = gv
	}
	return &OsqueryCollector{
		runner:     r,
		collectors: collectors,
		queryDurations: prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Namespace: "osquery_exporter",
				Name:      "query_duration",
				Help:      "Query duration",
			},
			[]string{"name"}),
		success: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: "osquery_exporter",
				Name:      "query_success",
				Help:      "Query execution status",
			},
			[]string{"name"},
		),
		resultsets: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: "osquery_exporter",
				Name:      "resultsets",
				Help:      "Number of query result sets",
			},
			[]string{"name"},
		),
	}
}

// Describe implements prometheus.Collector
func (c *OsqueryCollector) Describe(ch chan<- *prometheus.Desc) {
	c.queryDurations.Describe(ch)
	c.success.Describe(ch)
	c.resultsets.Describe(ch)
}

// Collect implements prometheus.Collector
func (c *OsqueryCollector) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(c.collectors))
	for _, col := range c.collectors {
		go func(col singleQueryCollector) {
			defer wg.Done()
			result, err := c.runner.Run(col.Query())
			if err != nil {
				log.Errorf("failed to run query %s: %s", col.Query(), err)
				c.success.WithLabelValues(col.String()).Set(0.0)
				return
			}
			c.resultsets.WithLabelValues(col.String()).Set(float64(len(result.Items)))
			err = update(col, result, ch)
			if err != nil {
				log.Warnf("metric %s errors on update: %s", col.String(), err)
				c.success.WithLabelValues(col.String()).Set(0.0)
				return
			}
			log.Debugf("metric %s took %s seconds to run", col.String(), result.Runtime)
			c.queryDurations.WithLabelValues(col.String()).Observe(
				result.Runtime.Seconds(),
			)
			c.success.WithLabelValues(col.String()).Set(1.0)
		}(col)
	}
	c.queryDurations.Collect(ch)
	c.success.Collect(ch)
	c.resultsets.Collect(ch)
	wg.Wait()

}
<commit_msg>fix HTTP/500 collecting error<commit_after>package collector

import (
	"fmt"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/zwopir/osquery_exporter/model"
	"github.com/zwopir/osquery_exporter/osquery"
	"strconv"
	"sync"
)

// singleQueryCollector represents a metric/query definition for a single osquery call
type singleQueryCollector interface {
	String() string
	Id() string
	Query() string
	Desc() *prometheus.Desc
	ValueType() prometheus.ValueType
	Value() string
	Labels() []string
}

// update maps the osquery query result to the singleQueryCollector and updates the provided channel accordingly
func update(sqc singleQueryCollector, result *model.OsqueryResult, ch chan<- prometheus.Metric) error {
	log.Debugf("updating metric %q", sqc.String())
	// metrics with no labels can only accept one result set
	if len(sqc.Labels()) == 0 && len(result.Items) > 1 {
		return fmt.Errorf("metrics with no labels can only accept one result set")
	}
	for _, item := range result.Items {
		value, ok := item[sqc.Value()]
		if !ok {
			return fmt.Errorf("query %q doesn't contain value key %q", sqc.Query(), sqc.Value())
		}
		valueAsFloat, err := strconv.ParseFloat(value, 64)
		if err != nil {
			return fmt.Errorf("query %q result %q can't be converted to float", sqc.Query(), value)
		}
		labels := []string{}
		for _, labelIdentifier := range sqc.Labels() {
			if label, ok := item[labelIdentifier]; ok {
				labels = append(labels, label)
			} else {
				return fmt.Errorf("query %q doesn't contain a label key %q", sqc.Query(), labelIdentifier)
			}
		}
		ch <- prometheus.MustNewConstMetric(
			sqc.Desc(),
			sqc.ValueType(),
			valueAsFloat,
			labels...,
		)
	}
	return nil
}

// OsqueryCollector represents a collector that collects metrics from a set of osquery queries. It implements
// prometheus Collector
type OsqueryCollector struct {
	runner         *osquery.OsqueryRunner
	collectors     map[string]singleQueryCollector
	queryDurations *prometheus.SummaryVec
	success        *prometheus.GaugeVec
	resultsets     *prometheus.GaugeVec
}

// NewOsqueryCollector creates an OsQueryCollector from a given osquery-runner and a set of metric definitions
func NewOsqueryCollector(r *osquery.OsqueryRunner, m model.Metrics) *OsqueryCollector {
	collectors := make(map[string]singleQueryCollector)
	for _, c := range m.Counters {
		log.Infof("adding %s to OsqueryCollector", c.String())
		collectors[c.Id()] = c
	}
	for _, cv := range m.CounterVecs {
		log.Infof("adding %s to OsqueryCollector", cv.String())
		collectors[cv.Id()] = cv
	}
	for _, g := range m.Gauges {
		log.Infof("adding %s to OsqueryCollector", g.String())
		collectors[g.Id()] = g
	}
	for _, gv := range m.GaugeVecs {
		log.Infof("adding %s to OsqueryCollector", gv.String())
		collectors[gv.Id()] = gv
	}
	return &OsqueryCollector{
		runner:     r,
		collectors: collectors,
		queryDurations: prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Namespace: "osquery_exporter",
				Name:      "query_duration",
				Help:      "Query duration",
			},
			[]string{"name"}),
		success: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: "osquery_exporter",
				Name:      "query_success",
				Help:      "Query execution status",
			},
			[]string{"name"},
		),
		resultsets: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: "osquery_exporter",
				Name:      "resultsets",
				Help:      "Number of query result sets",
			},
			[]string{"name"},
		),
	}
}

// Describe implements prometheus.Collector
func (c *OsqueryCollector) Describe(ch chan<- *prometheus.Desc) {
	c.queryDurations.Describe(ch)
	c.success.Describe(ch)
	c.resultsets.Describe(ch)
}

// Collect implements prometheus.Collector
func (c *OsqueryCollector) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(c.collectors))
	for _, col := range c.collectors {
		go func(col singleQueryCollector) {
			defer wg.Done()
			result, err := c.runner.Run(col.Query())
			if err != nil {
				log.Errorf("failed to run query %s: %s", col.Query(), err)
				c.success.WithLabelValues(col.String()).Set(0.0)
				return
			}
			c.resultsets.WithLabelValues(col.String()).Set(float64(len(result.Items)))
			err = update(col, result, ch)
			if err != nil {
				log.Warnf("metric %s errors on update: %s", col.String(), err)
				c.success.WithLabelValues(col.String()).Set(0.0)
				return
			}
			log.Debugf("metric %s took %s seconds to run", col.String(), result.Runtime)
			c.queryDurations.WithLabelValues(col.String()).Observe(
				result.Runtime.Seconds(),
			)
			c.success.WithLabelValues(col.String()).Set(1.0)
		}(col)
	}
	c.queryDurations.Collect(ch)
	c.success.Collect(ch)
	c.resultsets.Collect(ch)
	wg.Wait()

}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present Wenhui Shen <swh@admpub.com>

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

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

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package cmd

import (
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/defaults"
	"github.com/webx-top/echo/middleware/language"

	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/cmd/bootconfig"
	"github.com/admpub/nging/v4/application/handler/setup"
	"github.com/admpub/nging/v4/application/library/config"
	"github.com/admpub/nging/v4/application/library/config/subconfig/sdb"
)

// 静默安装
var InitDBConfig = &sdb.DB{
	Type:     `mysql`, // mysql / sqlite
	User:     `root`,
	Database: `nging`,
	Host:     `127.0.0.1:3306`,
}

var InitInstallConfig = &struct {
	Charset    string
	AdminUser  string
	AdminPass  string
	AdminEmail string
	Language   string // en / zh-cn
}{
	Charset:   sdb.MySQLDefaultCharset,
	AdminUser: `admin`,
	Language:  `zh-cn`,
}

var initCmd = &cobra.Command{
	Use:     "init",
	Short:   "Silent install",
	Example: filepath.Base(os.Args[0]) + " init",
	RunE:    initRunE,
}

func BuildTranslator(c language.Config) echo.Translator {
	c.SetFSFunc(bootconfig.LangFSFunc)
	i18n := language.NewI18n(&c)
	tr := &language.Translate{}
	tr.Reset(InitInstallConfig.Language, i18n)
	return tr
}

func initRunE(cmd *cobra.Command, args []string) error {
	conf, err := config.InitConfig()
	if err != nil {
		return err
	}
	conf.AsDefault()
	ctx := defaults.NewMockContext()

	// 启用多语言支持
	ctx.SetTranslator(BuildTranslator(config.DefaultConfig.Language))

	ctx.Request().Form().Set(`type`, InitDBConfig.Type)
	ctx.Request().Form().Set(`user`, InitDBConfig.User)
	ctx.Request().Form().Set(`host`, InitDBConfig.Host)
	ctx.Request().Form().Set(`password`, InitDBConfig.Password)
	ctx.Request().Form().Set(`database`, InitDBConfig.Database)
	ctx.Request().Form().Set(`prefix`, InitDBConfig.Prefix)
	ctx.Request().Form().Set(`charset`, InitInstallConfig.Charset)
	ctx.Request().Form().Set(`adminUser`, InitInstallConfig.AdminUser)
	ctx.Request().Form().Set(`adminPass`, InitInstallConfig.AdminPass)
	ctx.Request().Form().Set(`adminEmail`, InitInstallConfig.AdminEmail)
	err = setup.Setup(ctx)
	if err == nil {
		log.Okay(`Congratulations, this program has been installed successfully`)
	}
	return err
}

func init() {
	rootCmd.AddCommand(initCmd)

	initCmd.Flags().StringVar(&InitDBConfig.Type, "type", InitDBConfig.Type, "database type")
	initCmd.Flags().StringVar(&InitDBConfig.User, "user", InitDBConfig.User, "database user")
	initCmd.Flags().StringVar(&InitDBConfig.Host, "host", InitDBConfig.Host, "database host")
	initCmd.Flags().StringVar(&InitDBConfig.Password, "password", InitDBConfig.Password, "database password")
	initCmd.Flags().StringVar(&InitDBConfig.Database, "database", InitDBConfig.Database, "database name")
	initCmd.Flags().StringVar(&InitDBConfig.Prefix, "prefix", InitDBConfig.Prefix, "database table prefix")
	initCmd.Flags().StringVar(&InitInstallConfig.Charset, "charset", InitInstallConfig.Charset, "database table charset")
	initCmd.Flags().StringVar(&InitInstallConfig.AdminUser, "adminUser", InitInstallConfig.AdminUser, "administrator name")
	initCmd.Flags().StringVar(&InitInstallConfig.AdminPass, "adminPass", InitInstallConfig.AdminPass, "administrator password")
	initCmd.Flags().StringVar(&InitInstallConfig.AdminEmail, "adminEmail", InitInstallConfig.AdminEmail, "administrator e-mail")
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present Wenhui Shen <swh@admpub.com>

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

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

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package cmd

import (
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/defaults"
	"github.com/webx-top/echo/middleware/language"

	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/cmd/bootconfig"
	"github.com/admpub/nging/v4/application/handler/setup"
	"github.com/admpub/nging/v4/application/library/config"
	"github.com/admpub/nging/v4/application/library/config/subconfig/sdb"
)

// 静默安装
var InitDBConfig = &sdb.DB{
	Type:     `mysql`, // mysql / sqlite
	User:     `root`,
	Database: `nging`,
	Host:     `127.0.0.1:3306`,
}

var InitInstallConfig = &struct {
	Charset    string
	AdminUser  string
	AdminPass  string
	AdminEmail string
	Language   string // en / zh-cn
}{
	Charset:   sdb.MySQLDefaultCharset,
	AdminUser: `admin`,
	Language:  `zh-cn`,
}

var initCmd = &cobra.Command{
	Use:     "init",
	Short:   "Silent install",
	Example: filepath.Base(os.Args[0]) + " init [options]",
	RunE:    initRunE,
}

func BuildTranslator(c language.Config) echo.Translator {
	c.SetFSFunc(bootconfig.LangFSFunc)
	i18n := language.NewI18n(&c)
	tr := &language.Translate{}
	tr.Reset(InitInstallConfig.Language, i18n)
	return tr
}

func initRunE(cmd *cobra.Command, args []string) error {
	conf, err := config.InitConfig()
	if err != nil {
		return err
	}
	conf.AsDefault()
	ctx := defaults.NewMockContext()

	// 启用多语言支持
	ctx.SetTranslator(BuildTranslator(config.DefaultConfig.Language))

	ctx.Request().Form().Set(`type`, InitDBConfig.Type)
	ctx.Request().Form().Set(`user`, InitDBConfig.User)
	ctx.Request().Form().Set(`host`, InitDBConfig.Host)
	ctx.Request().Form().Set(`password`, InitDBConfig.Password)
	ctx.Request().Form().Set(`database`, InitDBConfig.Database)
	ctx.Request().Form().Set(`prefix`, InitDBConfig.Prefix)
	ctx.Request().Form().Set(`charset`, InitInstallConfig.Charset)
	ctx.Request().Form().Set(`adminUser`, InitInstallConfig.AdminUser)
	ctx.Request().Form().Set(`adminPass`, InitInstallConfig.AdminPass)
	ctx.Request().Form().Set(`adminEmail`, InitInstallConfig.AdminEmail)
	err = setup.Setup(ctx)
	if err == nil {
		log.Okay(`Congratulations, this program has been installed successfully`)
	}
	return err
}

func init() {
	rootCmd.AddCommand(initCmd)

	initCmd.Flags().StringVar(&InitDBConfig.Type, "type", InitDBConfig.Type, "database type")
	initCmd.Flags().StringVar(&InitDBConfig.User, "user", InitDBConfig.User, "database user")
	initCmd.Flags().StringVar(&InitDBConfig.Host, "host", InitDBConfig.Host, "database host")
	initCmd.Flags().StringVar(&InitDBConfig.Password, "password", InitDBConfig.Password, "database password")
	initCmd.Flags().StringVar(&InitDBConfig.Database, "database", InitDBConfig.Database, "database name")
	initCmd.Flags().StringVar(&InitDBConfig.Prefix, "prefix", InitDBConfig.Prefix, "database table prefix")
	initCmd.Flags().StringVar(&InitInstallConfig.Charset, "charset", InitInstallConfig.Charset, "database table charset")
	initCmd.Flags().StringVar(&InitInstallConfig.AdminUser, "adminUser", InitInstallConfig.AdminUser, "administrator name")
	initCmd.Flags().StringVar(&InitInstallConfig.AdminPass, "adminPass", InitInstallConfig.AdminPass, "administrator password")
	initCmd.Flags().StringVar(&InitInstallConfig.AdminEmail, "adminEmail", InitInstallConfig.AdminEmail, "administrator e-mail")
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 The Jaeger Authors.
// Copyright (c) 2017 Uber Technologies, 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.

package integration

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/storage/dependencystore"
	"github.com/jaegertracing/jaeger/storage/spanstore"
)

const (
	iterations = 100
)

// StorageIntegration holds components for storage integration test
type StorageIntegration struct {
	SpanWriter       spanstore.Writer
	SpanReader       spanstore.Reader
	DependencyWriter dependencystore.Writer
	DependencyReader dependencystore.Reader
	Fixtures         []*QueryFixtures
	// TODO: remove this flag after all storage plugins returns spanKind with operationNames
	NotSupportSpanKindWithOperation bool
	FixturesPath                    string

	// CleanUp() should ensure that the storage backend is clean before another test.
	// called either before or after each test, and should be idempotent
	CleanUp func() error

	// Refresh() should ensure that the storage backend is up to date before being queried.
	// called between set-up and queries in each test
	Refresh func() error
}

// === SpanStore Integration Tests ===

// QueryFixtures and TraceFixtures are under ./fixtures/queries.json and ./fixtures/traces/*.json respectively.
// Each query fixture includes:
// 	Caption: describes the query we are testing
// 	Query: the query we are testing
//	ExpectedFixture: the trace fixture that we want back from these queries.
// Queries are not necessarily numbered, but since each query requires a service name,
// the service name is formatted "query##-service".
type QueryFixtures struct {
	Caption          string
	Query            *spanstore.TraceQueryParameters
	ExpectedFixtures []string
}

func (s *StorageIntegration) cleanUp(t *testing.T) {
	require.NotNil(t, s.CleanUp, "CleanUp function must be provided")
	require.NoError(t, s.CleanUp())
}

func (s *StorageIntegration) refresh(t *testing.T) {
	require.NotNil(t, s.Refresh, "Refresh function must be provided")
	require.NoError(t, s.Refresh())
}

func (s *StorageIntegration) waitForCondition(t *testing.T, predicate func(t *testing.T) bool) bool {
	for i := 0; i < iterations; i++ {
		t.Logf("Waiting for storage backend to update documents, iteration %d out of %d", i+1, iterations)
		if predicate(t) {
			return true
		}
		time.Sleep(100 * time.Millisecond) // Will wait up to 10 seconds at worst.
	}
	return predicate(t)
}

func (s *StorageIntegration) testGetServices(t *testing.T) {
	defer s.cleanUp(t)

	expected := []string{"example-service-1", "example-service-2", "example-service-3"}
	s.loadParseAndWriteExampleTrace(t)
	s.refresh(t)

	var actual []string
	found := s.waitForCondition(t, func(t *testing.T) bool {
		actual, err := s.SpanReader.GetServices(context.Background())
		require.NoError(t, err)
		return assert.ObjectsAreEqualValues(expected, actual)
	})

	if !assert.True(t, found) {
		t.Log("\t Expected:", expected)
		t.Log("\t Actual  :", actual)
	}
}

func (s *StorageIntegration) testGetLargeSpan(t *testing.T) {
	defer s.cleanUp(t)

	t.Log("Testing Large Trace over 10K ...")
	expected := s.loadParseAndWriteLargeTrace(t)
	expectedTraceID := expected.Spans[0].TraceID
	s.refresh(t)

	var actual *model.Trace
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		actual, err = s.SpanReader.GetTrace(context.Background(), expectedTraceID)
		return err == nil && len(actual.Spans) == len(expected.Spans)
	})
	if !assert.True(t, found) {
		CompareTraces(t, expected, actual)
	}
}

func (s *StorageIntegration) testGetOperations(t *testing.T) {
	defer s.cleanUp(t)

	var expected []spanstore.Operation
	if s.NotSupportSpanKindWithOperation {
		expected = []spanstore.Operation{
			{Name: "example-operation-1"},
			{Name: "example-operation-3"},
			{Name: "example-operation-4"},
		}
	} else {
		expected = []spanstore.Operation{
			{Name: "example-operation-1"},
			{Name: "example-operation-3", SpanKind: "server"},
			{Name: "example-operation-4", SpanKind: "client"},
		}
	}
	s.loadParseAndWriteExampleTrace(t)
	s.refresh(t)

	var actual []spanstore.Operation
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		actual, err = s.SpanReader.GetOperations(context.Background(),
			spanstore.OperationQueryParameters{ServiceName: "example-service-1"})
		require.NoError(t, err)
		return assert.ObjectsAreEqualValues(expected, actual)
	})

	if !assert.True(t, found) {
		t.Log("\t Expected:", expected)
		t.Log("\t Actual  :", actual)
	}
}

func (s *StorageIntegration) testGetTrace(t *testing.T) {
	defer s.cleanUp(t)

	expected := s.loadParseAndWriteExampleTrace(t)
	expectedTraceID := expected.Spans[0].TraceID
	s.refresh(t)

	var actual *model.Trace
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		actual, err = s.SpanReader.GetTrace(context.Background(), expectedTraceID)
		if err != nil {
			t.Log(err)
		}
		return err == nil && len(actual.Spans) == len(expected.Spans)
	})
	if !assert.True(t, found) {
		CompareTraces(t, expected, actual)
	}

	t.Run("NotFound error", func(t *testing.T) {
		fakeTraceID := model.TraceID{High: 0, Low: 0}
		trace, err := s.SpanReader.GetTrace(context.Background(), fakeTraceID)
		assert.Equal(t, spanstore.ErrTraceNotFound, err)
		assert.Nil(t, trace)
	})
}

func (s *StorageIntegration) testFindTraces(t *testing.T) {
	defer s.cleanUp(t)

	fixturesPath := s.FixturesPath
	if s.FixturesPath == "" {
		fixturesPath = "."
	}
	// Note: all cases include ServiceName + StartTime range
	s.Fixtures = append(s.Fixtures, LoadAndParseQueryTestCases(t, fmt.Sprintf("%s/fixtures/queries.json", fixturesPath))...)

	// Each query test case only specifies matching traces, but does not provide counterexamples.
	// To improve coverage we get all possible traces and store all of them before running queries.
	allTraceFixtures := make(map[string]*model.Trace)
	expectedTracesPerTestCase := make([][]*model.Trace, 0, len(s.Fixtures))
	for _, queryTestCase := range s.Fixtures {
		var expected []*model.Trace
		for _, traceFixture := range queryTestCase.ExpectedFixtures {
			trace, ok := allTraceFixtures[traceFixture]
			if !ok {
				trace = s.getTraceFixture(t, traceFixture)
				err := s.writeTrace(t, trace)
				require.NoError(t, err, "Unexpected error when writing trace %s to storage", traceFixture)
				allTraceFixtures[traceFixture] = trace
			}
			expected = append(expected, trace)
		}
		expectedTracesPerTestCase = append(expectedTracesPerTestCase, expected)
	}
	s.refresh(t)
	for i, queryTestCase := range s.Fixtures {
		t.Run(queryTestCase.Caption, func(t *testing.T) {
			expected := expectedTracesPerTestCase[i]
			actual := s.findTracesByQuery(t, queryTestCase.Query, expected)
			CompareSliceOfTraces(t, expected, actual)
		})
	}
}

func (s *StorageIntegration) findTracesByQuery(t *testing.T, query *spanstore.TraceQueryParameters, expected []*model.Trace) []*model.Trace {
	var traces []*model.Trace
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		traces, err = s.SpanReader.FindTraces(context.Background(), query)
		require.NoError(t, err)
		if len(expected) != len(traces) {
			t.Logf("FindTraces: expected: %d, actual: %d", len(expected), len(traces))
			return false
		}
		return true
	})
	require.True(t, found)
	tracesMatch(t, traces, expected)
	return traces
}

func (s *StorageIntegration) writeTrace(t *testing.T, trace *model.Trace) error {
	for _, span := range trace.Spans {
		if err := s.SpanWriter.WriteSpan(context.Background(), span); err != nil {
			return err
		}
	}
	return nil
}

func (s *StorageIntegration) loadParseAndWriteExampleTrace(t *testing.T) *model.Trace {
	trace := s.getTraceFixture(t, "example_trace")
	err := s.writeTrace(t, trace)
	require.NoError(t, err, "Not expecting error when writing example_trace to storage")
	return trace
}

func (s *StorageIntegration) loadParseAndWriteLargeTrace(t *testing.T) *model.Trace {
	trace := s.getTraceFixture(t, "example_trace")
	span := trace.Spans[0]
	spns := make([]*model.Span, 1, 10008)
	trace.Spans = spns
	trace.Spans[0] = span
	for i := 1; i < 10008; i++ {
		s := new(model.Span)
		*s = *span
		s.StartTime = s.StartTime.Add(time.Second * time.Duration(i+1))
		trace.Spans = append(trace.Spans, s)
	}
	err := s.writeTrace(t, trace)
	require.NoError(t, err, "Not expecting error when writing example_trace to storage")
	return trace
}

func (s *StorageIntegration) getTraceFixture(t *testing.T, fixture string) *model.Trace {
	fixturesPath := s.FixturesPath
	if s.FixturesPath == "" {
		fixturesPath = "."
	}
	fileName := fmt.Sprintf("%s/fixtures/traces/%s.json", fixturesPath, fixture)
	return getTraceFixtureExact(t, fileName)
}

func getTraceFixtureExact(t *testing.T, fileName string) *model.Trace {
	var trace model.Trace
	loadAndParseJSONPB(t, fileName, &trace)
	return &trace
}

func loadAndParseJSONPB(t *testing.T, path string, object proto.Message) {
	// #nosec
	inStr, err := os.ReadFile(path)
	require.NoError(t, err, "Not expecting error when loading fixture %s", path)
	err = jsonpb.Unmarshal(bytes.NewReader(correctTime(inStr)), object)
	require.NoError(t, err, "Not expecting error when unmarshaling fixture %s", path)
}

// LoadAndParseQueryTestCases loads and parses query test cases
func LoadAndParseQueryTestCases(t *testing.T, queriesFile string) []*QueryFixtures {
	var queries []*QueryFixtures
	loadAndParseJSON(t, queriesFile, &queries)
	return queries
}

func loadAndParseJSON(t *testing.T, path string, object interface{}) {
	// #nosec
	inStr, err := os.ReadFile(path)
	require.NoError(t, err, "Not expecting error when loading fixture %s", path)
	err = json.Unmarshal(correctTime(inStr), object)
	require.NoError(t, err, "Not expecting error when unmarshaling fixture %s", path)
}

// required, because we want to only query on recent traces, so we replace all the dates with recent dates.
func correctTime(json []byte) []byte {
	jsonString := string(json)
	now := time.Now().UTC()
	today := now.Format("2006-01-02")
	yesterday := now.AddDate(0, 0, -1).Format("2006-01-02")
	retString := strings.Replace(jsonString, "2017-01-26", today, -1)
	retString = strings.Replace(retString, "2017-01-25", yesterday, -1)
	return []byte(retString)
}

func tracesMatch(t *testing.T, actual []*model.Trace, expected []*model.Trace) bool {
	if !assert.Equal(t, len(expected), len(actual), "Expecting certain number of traces") {
		return false
	}
	return assert.Equal(t, spanCount(expected), spanCount(actual), "Expecting certain number of spans")
}

func spanCount(traces []*model.Trace) int {
	var count int
	for _, trace := range traces {
		count += len(trace.Spans)
	}
	return count
}

// === DependencyStore Integration Tests ===

func (s *StorageIntegration) testGetDependencies(t *testing.T) {
	if s.DependencyReader == nil || s.DependencyWriter == nil {
		t.Skipf("Skipping GetDependencies test because dependency reader or writer is nil")
		return
	}

	defer s.cleanUp(t)

	expected := []model.DependencyLink{
		{
			Parent:    "hello",
			Child:     "world",
			CallCount: uint64(1),
		},
		{
			Parent:    "world",
			Child:     "hello",
			CallCount: uint64(3),
		},
	}
	require.NoError(t, s.DependencyWriter.WriteDependencies(time.Now(), expected))
	s.refresh(t)
	actual, err := s.DependencyReader.GetDependencies(context.Background(), time.Now(), 5*time.Minute)
	assert.NoError(t, err)
	assert.EqualValues(t, expected, actual)
}

// IntegrationTestAll runs all integration tests
func (s *StorageIntegration) IntegrationTestAll(t *testing.T) {
	t.Run("GetServices", s.testGetServices)
	t.Run("GetOperations", s.testGetOperations)
	t.Run("GetTrace", s.testGetTrace)
	t.Run("GetLargeSpans", s.testGetLargeSpan)
	t.Run("FindTraces", s.testFindTraces)
	t.Run("GetDependencies", s.testGetDependencies)
}
<commit_msg>fix(integration): no spans in the future (#3538)<commit_after>// Copyright (c) 2019 The Jaeger Authors.
// Copyright (c) 2017 Uber Technologies, 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.

package integration

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/storage/dependencystore"
	"github.com/jaegertracing/jaeger/storage/spanstore"
)

const (
	iterations = 100
)

// StorageIntegration holds components for storage integration test
type StorageIntegration struct {
	SpanWriter       spanstore.Writer
	SpanReader       spanstore.Reader
	DependencyWriter dependencystore.Writer
	DependencyReader dependencystore.Reader
	Fixtures         []*QueryFixtures
	// TODO: remove this flag after all storage plugins returns spanKind with operationNames
	NotSupportSpanKindWithOperation bool
	FixturesPath                    string

	// CleanUp() should ensure that the storage backend is clean before another test.
	// called either before or after each test, and should be idempotent
	CleanUp func() error

	// Refresh() should ensure that the storage backend is up to date before being queried.
	// called between set-up and queries in each test
	Refresh func() error
}

// === SpanStore Integration Tests ===

// QueryFixtures and TraceFixtures are under ./fixtures/queries.json and ./fixtures/traces/*.json respectively.
// Each query fixture includes:
// 	Caption: describes the query we are testing
// 	Query: the query we are testing
//	ExpectedFixture: the trace fixture that we want back from these queries.
// Queries are not necessarily numbered, but since each query requires a service name,
// the service name is formatted "query##-service".
type QueryFixtures struct {
	Caption          string
	Query            *spanstore.TraceQueryParameters
	ExpectedFixtures []string
}

func (s *StorageIntegration) cleanUp(t *testing.T) {
	require.NotNil(t, s.CleanUp, "CleanUp function must be provided")
	require.NoError(t, s.CleanUp())
}

func (s *StorageIntegration) refresh(t *testing.T) {
	require.NotNil(t, s.Refresh, "Refresh function must be provided")
	require.NoError(t, s.Refresh())
}

func (s *StorageIntegration) waitForCondition(t *testing.T, predicate func(t *testing.T) bool) bool {
	for i := 0; i < iterations; i++ {
		t.Logf("Waiting for storage backend to update documents, iteration %d out of %d", i+1, iterations)
		if predicate(t) {
			return true
		}
		time.Sleep(100 * time.Millisecond) // Will wait up to 10 seconds at worst.
	}
	return predicate(t)
}

func (s *StorageIntegration) testGetServices(t *testing.T) {
	defer s.cleanUp(t)

	expected := []string{"example-service-1", "example-service-2", "example-service-3"}
	s.loadParseAndWriteExampleTrace(t)
	s.refresh(t)

	var actual []string
	found := s.waitForCondition(t, func(t *testing.T) bool {
		actual, err := s.SpanReader.GetServices(context.Background())
		require.NoError(t, err)
		return assert.ObjectsAreEqualValues(expected, actual)
	})

	if !assert.True(t, found) {
		t.Log("\t Expected:", expected)
		t.Log("\t Actual  :", actual)
	}
}

func (s *StorageIntegration) testGetLargeSpan(t *testing.T) {
	defer s.cleanUp(t)

	t.Log("Testing Large Trace over 10K ...")
	expected := s.loadParseAndWriteLargeTrace(t)
	expectedTraceID := expected.Spans[0].TraceID
	s.refresh(t)

	var actual *model.Trace
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		actual, err = s.SpanReader.GetTrace(context.Background(), expectedTraceID)
		return err == nil && len(actual.Spans) == len(expected.Spans)
	})
	if !assert.True(t, found) {
		CompareTraces(t, expected, actual)
	}
}

func (s *StorageIntegration) testGetOperations(t *testing.T) {
	defer s.cleanUp(t)

	var expected []spanstore.Operation
	if s.NotSupportSpanKindWithOperation {
		expected = []spanstore.Operation{
			{Name: "example-operation-1"},
			{Name: "example-operation-3"},
			{Name: "example-operation-4"},
		}
	} else {
		expected = []spanstore.Operation{
			{Name: "example-operation-1"},
			{Name: "example-operation-3", SpanKind: "server"},
			{Name: "example-operation-4", SpanKind: "client"},
		}
	}
	s.loadParseAndWriteExampleTrace(t)
	s.refresh(t)

	var actual []spanstore.Operation
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		actual, err = s.SpanReader.GetOperations(context.Background(),
			spanstore.OperationQueryParameters{ServiceName: "example-service-1"})
		require.NoError(t, err)
		return assert.ObjectsAreEqualValues(expected, actual)
	})

	if !assert.True(t, found) {
		t.Log("\t Expected:", expected)
		t.Log("\t Actual  :", actual)
	}
}

func (s *StorageIntegration) testGetTrace(t *testing.T) {
	defer s.cleanUp(t)

	expected := s.loadParseAndWriteExampleTrace(t)
	expectedTraceID := expected.Spans[0].TraceID
	s.refresh(t)

	var actual *model.Trace
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		actual, err = s.SpanReader.GetTrace(context.Background(), expectedTraceID)
		if err != nil {
			t.Log(err)
		}
		return err == nil && len(actual.Spans) == len(expected.Spans)
	})
	if !assert.True(t, found) {
		CompareTraces(t, expected, actual)
	}

	t.Run("NotFound error", func(t *testing.T) {
		fakeTraceID := model.TraceID{High: 0, Low: 0}
		trace, err := s.SpanReader.GetTrace(context.Background(), fakeTraceID)
		assert.Equal(t, spanstore.ErrTraceNotFound, err)
		assert.Nil(t, trace)
	})
}

func (s *StorageIntegration) testFindTraces(t *testing.T) {
	defer s.cleanUp(t)

	fixturesPath := s.FixturesPath
	if s.FixturesPath == "" {
		fixturesPath = "."
	}
	// Note: all cases include ServiceName + StartTime range
	s.Fixtures = append(s.Fixtures, LoadAndParseQueryTestCases(t, fmt.Sprintf("%s/fixtures/queries.json", fixturesPath))...)

	// Each query test case only specifies matching traces, but does not provide counterexamples.
	// To improve coverage we get all possible traces and store all of them before running queries.
	allTraceFixtures := make(map[string]*model.Trace)
	expectedTracesPerTestCase := make([][]*model.Trace, 0, len(s.Fixtures))
	for _, queryTestCase := range s.Fixtures {
		var expected []*model.Trace
		for _, traceFixture := range queryTestCase.ExpectedFixtures {
			trace, ok := allTraceFixtures[traceFixture]
			if !ok {
				trace = s.getTraceFixture(t, traceFixture)
				err := s.writeTrace(t, trace)
				require.NoError(t, err, "Unexpected error when writing trace %s to storage", traceFixture)
				allTraceFixtures[traceFixture] = trace
			}
			expected = append(expected, trace)
		}
		expectedTracesPerTestCase = append(expectedTracesPerTestCase, expected)
	}
	s.refresh(t)
	for i, queryTestCase := range s.Fixtures {
		t.Run(queryTestCase.Caption, func(t *testing.T) {
			expected := expectedTracesPerTestCase[i]
			actual := s.findTracesByQuery(t, queryTestCase.Query, expected)
			CompareSliceOfTraces(t, expected, actual)
		})
	}
}

func (s *StorageIntegration) findTracesByQuery(t *testing.T, query *spanstore.TraceQueryParameters, expected []*model.Trace) []*model.Trace {
	var traces []*model.Trace
	found := s.waitForCondition(t, func(t *testing.T) bool {
		var err error
		traces, err = s.SpanReader.FindTraces(context.Background(), query)
		require.NoError(t, err)
		if len(expected) != len(traces) {
			t.Logf("FindTraces: expected: %d, actual: %d", len(expected), len(traces))
			return false
		}
		return true
	})
	require.True(t, found)
	tracesMatch(t, traces, expected)
	return traces
}

func (s *StorageIntegration) writeTrace(t *testing.T, trace *model.Trace) error {
	for _, span := range trace.Spans {
		if err := s.SpanWriter.WriteSpan(context.Background(), span); err != nil {
			return err
		}
	}
	return nil
}

func (s *StorageIntegration) loadParseAndWriteExampleTrace(t *testing.T) *model.Trace {
	trace := s.getTraceFixture(t, "example_trace")
	err := s.writeTrace(t, trace)
	require.NoError(t, err, "Not expecting error when writing example_trace to storage")
	return trace
}

func (s *StorageIntegration) loadParseAndWriteLargeTrace(t *testing.T) *model.Trace {
	trace := s.getTraceFixture(t, "example_trace")
	span := trace.Spans[0]
	spns := make([]*model.Span, 1, 10008)
	trace.Spans = spns
	trace.Spans[0] = span
	for i := 1; i < 10008; i++ {
		s := new(model.Span)
		*s = *span
		s.StartTime = s.StartTime.Add(time.Second * time.Duration(i+1))
		trace.Spans = append(trace.Spans, s)
	}
	err := s.writeTrace(t, trace)
	require.NoError(t, err, "Not expecting error when writing example_trace to storage")
	return trace
}

func (s *StorageIntegration) getTraceFixture(t *testing.T, fixture string) *model.Trace {
	fixturesPath := s.FixturesPath
	if s.FixturesPath == "" {
		fixturesPath = "."
	}
	fileName := fmt.Sprintf("%s/fixtures/traces/%s.json", fixturesPath, fixture)
	return getTraceFixtureExact(t, fileName)
}

func getTraceFixtureExact(t *testing.T, fileName string) *model.Trace {
	var trace model.Trace
	loadAndParseJSONPB(t, fileName, &trace)
	return &trace
}

func loadAndParseJSONPB(t *testing.T, path string, object proto.Message) {
	// #nosec
	inStr, err := os.ReadFile(path)
	require.NoError(t, err, "Not expecting error when loading fixture %s", path)
	err = jsonpb.Unmarshal(bytes.NewReader(correctTime(inStr)), object)
	require.NoError(t, err, "Not expecting error when unmarshaling fixture %s", path)
}

// LoadAndParseQueryTestCases loads and parses query test cases
func LoadAndParseQueryTestCases(t *testing.T, queriesFile string) []*QueryFixtures {
	var queries []*QueryFixtures
	loadAndParseJSON(t, queriesFile, &queries)
	return queries
}

func loadAndParseJSON(t *testing.T, path string, object interface{}) {
	// #nosec
	inStr, err := os.ReadFile(path)
	require.NoError(t, err, "Not expecting error when loading fixture %s", path)
	err = json.Unmarshal(correctTime(inStr), object)
	require.NoError(t, err, "Not expecting error when unmarshaling fixture %s", path)
}

// required, because we want to only query on recent traces, so we replace all the dates with recent dates.
func correctTime(json []byte) []byte {
	jsonString := string(json)
	now := time.Now().UTC()
	yesterday := now.AddDate(0, 0, -1).Format("2006-01-02")
	twoDaysAgo := now.AddDate(0, 0, -2).Format("2006-01-02")
	retString := strings.Replace(jsonString, "2017-01-26", yesterday, -1)
	retString = strings.Replace(retString, "2017-01-25", twoDaysAgo, -1)
	return []byte(retString)
}

func tracesMatch(t *testing.T, actual []*model.Trace, expected []*model.Trace) bool {
	if !assert.Equal(t, len(expected), len(actual), "Expecting certain number of traces") {
		return false
	}
	return assert.Equal(t, spanCount(expected), spanCount(actual), "Expecting certain number of spans")
}

func spanCount(traces []*model.Trace) int {
	var count int
	for _, trace := range traces {
		count += len(trace.Spans)
	}
	return count
}

// === DependencyStore Integration Tests ===

func (s *StorageIntegration) testGetDependencies(t *testing.T) {
	if s.DependencyReader == nil || s.DependencyWriter == nil {
		t.Skipf("Skipping GetDependencies test because dependency reader or writer is nil")
		return
	}

	defer s.cleanUp(t)

	expected := []model.DependencyLink{
		{
			Parent:    "hello",
			Child:     "world",
			CallCount: uint64(1),
		},
		{
			Parent:    "world",
			Child:     "hello",
			CallCount: uint64(3),
		},
	}
	require.NoError(t, s.DependencyWriter.WriteDependencies(time.Now(), expected))
	s.refresh(t)
	actual, err := s.DependencyReader.GetDependencies(context.Background(), time.Now(), 5*time.Minute)
	assert.NoError(t, err)
	assert.EqualValues(t, expected, actual)
}

// IntegrationTestAll runs all integration tests
func (s *StorageIntegration) IntegrationTestAll(t *testing.T) {
	t.Run("GetServices", s.testGetServices)
	t.Run("GetOperations", s.testGetOperations)
	t.Run("GetTrace", s.testGetTrace)
	t.Run("GetLargeSpans", s.testGetLargeSpan)
	t.Run("FindTraces", s.testFindTraces)
	t.Run("GetDependencies", s.testGetDependencies)
}
<|endoftext|>


<commit_before>package rsync

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/pborman/uuid"

	"github.com/lxc/lxd/lxd/daemon"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/logger"
)

// LocalCopy copies a directory using rsync (with the --devices option).
func LocalCopy(source string, dest string, bwlimit string, xattrs bool) (string, error) {
	err := os.MkdirAll(dest, 0755)
	if err != nil {
		return "", err
	}

	rsyncVerbosity := "-q"
	if daemon.Debug {
		rsyncVerbosity = "-vi"
	}

	if bwlimit == "" {
		bwlimit = "0"
	}

	args := []string{
		"-a",
		"-HA",
		"--sparse",
		"--devices",
		"--delete",
		"--checksum",
		"--numeric-ids",
	}

	if xattrs {
		args = append(args, "--xattrs")
	}

	if bwlimit != "" {
		args = append(args, "--bwlimit", bwlimit)
	}

	args = append(args,
		rsyncVerbosity,
		shared.AddSlash(source),
		dest)
	msg, err := shared.RunCommand("rsync", args...)
	if err != nil {
		runError, ok := err.(shared.RunError)
		if ok {
			exitError, ok := runError.Err.(*exec.ExitError)
			if ok {
				waitStatus := exitError.Sys().(syscall.WaitStatus)
				if waitStatus.ExitStatus() == 24 {
					return msg, nil
				}
			}
		}
		return msg, err
	}

	return msg, nil
}

func sendSetup(name string, path string, bwlimit string, execPath string, features []string) (*exec.Cmd, net.Conn, io.ReadCloser, error) {
	/*
	 * The way rsync works, it invokes a subprocess that does the actual
	 * talking (given to it by a -E argument). Since there isn't an easy
	 * way for us to capture this process' stdin/stdout, we just use netcat
	 * and write to/from a unix socket.
	 *
	 * In principle we don't need this socket. It seems to me that some
	 * clever invocation of rsync --server --sender and usage of that
	 * process' stdin/stdout could work around the need for this socket,
	 * but I couldn't get it to work. Another option would be to look at
	 * the spawned process' first child and read/write from its
	 * stdin/stdout, but that also seemed messy. In any case, this seems to
	 * work just fine.
	 */
	auds := fmt.Sprintf("@lxd/%s", uuid.NewRandom().String())
	// We simply copy a part of the uuid if it's longer than the allowed
	// maximum. That should be safe enough for our purposes.
	if len(auds) > shared.ABSTRACT_UNIX_SOCK_LEN-1 {
		auds = auds[:shared.ABSTRACT_UNIX_SOCK_LEN-1]
	}
	l, err := net.Listen("unix", auds)
	if err != nil {
		return nil, nil, nil, err
	}
	defer l.Close()

	/*
	 * Here, the path /tmp/foo is ignored. Since we specify localhost,
	 * rsync thinks we are syncing to a remote host (in this case, the
	 * other end of the lxd websocket), and so the path specified on the
	 * --server instance of rsync takes precedence.
	 *
	 * Additionally, we use sh -c instead of just calling nc directly
	 * because rsync passes a whole bunch of arguments to the wrapper
	 * command (i.e. the command to run on --server). However, we're
	 * hardcoding that at the other end, so we can just ignore it.
	 */
	rsyncCmd := fmt.Sprintf("sh -c \"%s netcat %s %s\"", execPath, auds, name)
	if bwlimit == "" {
		bwlimit = "0"
	}

	args := []string{
		"-ar",
		"--devices",
		"--numeric-ids",
		"--partial",
		"--sparse",
	}

	if features != nil && len(features) > 0 {
		args = append(args, rsyncFeatureArgs(features)...)
	}

	args = append(args, []string{
		path,
		"localhost:/tmp/foo",
		"-e",
		rsyncCmd,
		"--bwlimit",
		bwlimit}...)

	cmd := exec.Command("rsync", args...)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, nil, err
	}

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

	var conn *net.Conn
	chConn := make(chan *net.Conn, 1)

	go func() {
		conn, err := l.Accept()
		if err != nil {
			chConn <- nil
			return
		}

		chConn <- &conn
	}()

	select {
	case conn = <-chConn:
		if conn == nil {
			cmd.Process.Kill()
			cmd.Wait()
			return nil, nil, nil, fmt.Errorf("Failed to connect to rsync socket")
		}

	case <-time.After(10 * time.Second):
		cmd.Process.Kill()
		cmd.Wait()
		return nil, nil, nil, fmt.Errorf("rsync failed to spawn after 10s")
	}

	return cmd, *conn, stderr, nil
}

// Send sets up the sending half of an rsync, to recursively send the
// directory pointed to by path over the websocket.
func Send(name string, path string, conn io.ReadWriteCloser, tracker *ioprogress.ProgressTracker, features []string, bwlimit string, execPath string) error {
	cmd, netcatConn, stderr, err := sendSetup(name, path, bwlimit, execPath, features)
	if err != nil {
		return err
	}

	// Setup progress tracker.
	readNetcatPipe := io.ReadCloser(netcatConn)
	if tracker != nil {
		readNetcatPipe = &ioprogress.ProgressReader{
			ReadCloser: netcatConn,
			Tracker:    tracker,
		}
	}

	// Forward from netcat to target.
	chCopyNetcat := make(chan error, 1)
	go func() {
		_, err := io.Copy(conn, readNetcatPipe)
		chCopyNetcat <- err
		readNetcatPipe.Close()
		netcatConn.Close()
		conn.Close() // sends barrier message.
	}()

	// Forward from target to netcat.
	writeNetcatPipe := io.WriteCloser(netcatConn)
	chCopyTarget := make(chan error, 1)
	go func() {
		_, err := io.Copy(writeNetcatPipe, conn)
		chCopyTarget <- err
		writeNetcatPipe.Close()
	}()

	// Wait for rsync to complete.
	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		cmd.Process.Kill()
		logger.Errorf("Rsync stderr read failed: %s: %v", path, err)
	}

	err = cmd.Wait()
	errs := []error{}
	chCopyNetcatErr := <-chCopyNetcat
	chCopyTargetErr := <-chCopyTarget

	if err != nil {
		errs = append(errs, err)

		// Try to get more info about the error.
		if chCopyNetcatErr != nil {
			errs = append(errs, chCopyNetcatErr)
		}

		if chCopyTargetErr != nil {
			errs = append(errs, chCopyTargetErr)
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("Rsync send failed: %s, %s: %v (%s)", name, path, errs, string(output))
	}

	return nil
}

// Recv sets up the receiving half of the websocket to rsync (the other
// half set up by rsync.Send), putting the contents in the directory specified
// by path.
func Recv(path string, conn io.ReadWriteCloser, tracker *ioprogress.ProgressTracker, features []string) error {
	args := []string{
		"--server",
		"-vlogDtpre.iLsfx",
		"--numeric-ids",
		"--devices",
		"--partial",
		"--sparse",
	}

	if features != nil && len(features) > 0 {
		args = append(args, rsyncFeatureArgs(features)...)
	}

	args = append(args, []string{".", path}...)

	cmd := exec.Command("rsync", args...)

	// Forward from rsync to source.
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	chCopyRsync := make(chan error, 1)
	go func() {
		_, err := io.Copy(conn, stdout)
		chCopyRsync <- err
		stdout.Close()
		conn.Close() // sends barrier message.
	}()

	// Forward from source to rsync.
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	readSourcePipe := io.ReadCloser(conn)
	if tracker != nil {
		readSourcePipe = &ioprogress.ProgressReader{
			ReadCloser: conn,
			Tracker:    tracker,
		}
	}

	chCopySource := make(chan error, 1)
	go func() {
		_, err := io.Copy(stdin, readSourcePipe)
		chCopySource <- err
		stdin.Close()
	}()

	stderr, err := cmd.StderrPipe()
	if err != nil {
		cmd.Process.Kill()
		logger.Errorf("Rsync stderr read failed: %s: %v", path, err)
	}

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

	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		logger.Errorf("Rsync stderr read failed: %s: %v", path, err)
	}

	err = cmd.Wait()
	errs := []error{}
	chCopyRsyncErr := <-chCopyRsync
	chCopySourceErr := <-chCopySource

	if err != nil {
		errs = append(errs, err)

		// Try to get more info about the error.
		if chCopyRsyncErr != nil {
			errs = append(errs, chCopyRsyncErr)
		}

		if chCopySourceErr != nil {
			errs = append(errs, chCopySourceErr)
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("Rsync receive failed: %s: %v (%s)", path, errs, string(output))
	}

	return nil
}

func rsyncFeatureArgs(features []string) []string {
	args := []string{}
	if shared.StringInSlice("xattrs", features) {
		args = append(args, "--xattrs")
	}

	if shared.StringInSlice("delete", features) {
		args = append(args, "--delete")
	}

	if shared.StringInSlice("compress", features) {
		args = append(args, "--compress")
		args = append(args, "--compress-level=2")
	}

	return args
}
<commit_msg>lxd/rsync: Tweaks Recv's internal synchronisation to avoid race<commit_after>package rsync

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/pborman/uuid"

	"github.com/lxc/lxd/lxd/daemon"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/logger"
)

// LocalCopy copies a directory using rsync (with the --devices option).
func LocalCopy(source string, dest string, bwlimit string, xattrs bool) (string, error) {
	err := os.MkdirAll(dest, 0755)
	if err != nil {
		return "", err
	}

	rsyncVerbosity := "-q"
	if daemon.Debug {
		rsyncVerbosity = "-vi"
	}

	if bwlimit == "" {
		bwlimit = "0"
	}

	args := []string{
		"-a",
		"-HA",
		"--sparse",
		"--devices",
		"--delete",
		"--checksum",
		"--numeric-ids",
	}

	if xattrs {
		args = append(args, "--xattrs")
	}

	if bwlimit != "" {
		args = append(args, "--bwlimit", bwlimit)
	}

	args = append(args,
		rsyncVerbosity,
		shared.AddSlash(source),
		dest)
	msg, err := shared.RunCommand("rsync", args...)
	if err != nil {
		runError, ok := err.(shared.RunError)
		if ok {
			exitError, ok := runError.Err.(*exec.ExitError)
			if ok {
				waitStatus := exitError.Sys().(syscall.WaitStatus)
				if waitStatus.ExitStatus() == 24 {
					return msg, nil
				}
			}
		}
		return msg, err
	}

	return msg, nil
}

func sendSetup(name string, path string, bwlimit string, execPath string, features []string) (*exec.Cmd, net.Conn, io.ReadCloser, error) {
	/*
	 * The way rsync works, it invokes a subprocess that does the actual
	 * talking (given to it by a -E argument). Since there isn't an easy
	 * way for us to capture this process' stdin/stdout, we just use netcat
	 * and write to/from a unix socket.
	 *
	 * In principle we don't need this socket. It seems to me that some
	 * clever invocation of rsync --server --sender and usage of that
	 * process' stdin/stdout could work around the need for this socket,
	 * but I couldn't get it to work. Another option would be to look at
	 * the spawned process' first child and read/write from its
	 * stdin/stdout, but that also seemed messy. In any case, this seems to
	 * work just fine.
	 */
	auds := fmt.Sprintf("@lxd/%s", uuid.NewRandom().String())
	// We simply copy a part of the uuid if it's longer than the allowed
	// maximum. That should be safe enough for our purposes.
	if len(auds) > shared.ABSTRACT_UNIX_SOCK_LEN-1 {
		auds = auds[:shared.ABSTRACT_UNIX_SOCK_LEN-1]
	}
	l, err := net.Listen("unix", auds)
	if err != nil {
		return nil, nil, nil, err
	}
	defer l.Close()

	/*
	 * Here, the path /tmp/foo is ignored. Since we specify localhost,
	 * rsync thinks we are syncing to a remote host (in this case, the
	 * other end of the lxd websocket), and so the path specified on the
	 * --server instance of rsync takes precedence.
	 *
	 * Additionally, we use sh -c instead of just calling nc directly
	 * because rsync passes a whole bunch of arguments to the wrapper
	 * command (i.e. the command to run on --server). However, we're
	 * hardcoding that at the other end, so we can just ignore it.
	 */
	rsyncCmd := fmt.Sprintf("sh -c \"%s netcat %s %s\"", execPath, auds, name)
	if bwlimit == "" {
		bwlimit = "0"
	}

	args := []string{
		"-ar",
		"--devices",
		"--numeric-ids",
		"--partial",
		"--sparse",
	}

	if features != nil && len(features) > 0 {
		args = append(args, rsyncFeatureArgs(features)...)
	}

	args = append(args, []string{
		path,
		"localhost:/tmp/foo",
		"-e",
		rsyncCmd,
		"--bwlimit",
		bwlimit}...)

	cmd := exec.Command("rsync", args...)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, nil, err
	}

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

	var conn *net.Conn
	chConn := make(chan *net.Conn, 1)

	go func() {
		conn, err := l.Accept()
		if err != nil {
			chConn <- nil
			return
		}

		chConn <- &conn
	}()

	select {
	case conn = <-chConn:
		if conn == nil {
			cmd.Process.Kill()
			cmd.Wait()
			return nil, nil, nil, fmt.Errorf("Failed to connect to rsync socket")
		}

	case <-time.After(10 * time.Second):
		cmd.Process.Kill()
		cmd.Wait()
		return nil, nil, nil, fmt.Errorf("rsync failed to spawn after 10s")
	}

	return cmd, *conn, stderr, nil
}

// Send sets up the sending half of an rsync, to recursively send the
// directory pointed to by path over the websocket.
func Send(name string, path string, conn io.ReadWriteCloser, tracker *ioprogress.ProgressTracker, features []string, bwlimit string, execPath string) error {
	cmd, netcatConn, stderr, err := sendSetup(name, path, bwlimit, execPath, features)
	if err != nil {
		return err
	}

	// Setup progress tracker.
	readNetcatPipe := io.ReadCloser(netcatConn)
	if tracker != nil {
		readNetcatPipe = &ioprogress.ProgressReader{
			ReadCloser: netcatConn,
			Tracker:    tracker,
		}
	}

	// Forward from netcat to target.
	chCopyNetcat := make(chan error, 1)
	go func() {
		_, err := io.Copy(conn, readNetcatPipe)
		chCopyNetcat <- err
		readNetcatPipe.Close()
		netcatConn.Close()
		conn.Close() // sends barrier message.
	}()

	// Forward from target to netcat.
	writeNetcatPipe := io.WriteCloser(netcatConn)
	chCopyTarget := make(chan error, 1)
	go func() {
		_, err := io.Copy(writeNetcatPipe, conn)
		chCopyTarget <- err
		writeNetcatPipe.Close()
	}()

	// Wait for rsync to complete.
	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		cmd.Process.Kill()
		logger.Errorf("Rsync stderr read failed: %s: %v", path, err)
	}

	err = cmd.Wait()
	errs := []error{}
	chCopyNetcatErr := <-chCopyNetcat
	chCopyTargetErr := <-chCopyTarget

	if err != nil {
		errs = append(errs, err)

		// Try to get more info about the error.
		if chCopyNetcatErr != nil {
			errs = append(errs, chCopyNetcatErr)
		}

		if chCopyTargetErr != nil {
			errs = append(errs, chCopyTargetErr)
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("Rsync send failed: %s, %s: %v (%s)", name, path, errs, string(output))
	}

	return nil
}

// Recv sets up the receiving half of the websocket to rsync (the other
// half set up by rsync.Send), putting the contents in the directory specified
// by path.
func Recv(path string, conn io.ReadWriteCloser, tracker *ioprogress.ProgressTracker, features []string) error {
	args := []string{
		"--server",
		"-vlogDtpre.iLsfx",
		"--numeric-ids",
		"--devices",
		"--partial",
		"--sparse",
	}

	if features != nil && len(features) > 0 {
		args = append(args, rsyncFeatureArgs(features)...)
	}

	args = append(args, []string{".", path}...)

	cmd := exec.Command("rsync", args...)

	// Forward from rsync to source.
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	chCopyRsync := make(chan error, 1)
	go func() {
		_, err := io.Copy(conn, stdout)
		stdout.Close()
		conn.Close() // sends barrier message.
		chCopyRsync <- err
	}()

	// Forward from source to rsync.
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	readSourcePipe := io.ReadCloser(conn)
	if tracker != nil {
		readSourcePipe = &ioprogress.ProgressReader{
			ReadCloser: conn,
			Tracker:    tracker,
		}
	}

	chCopySource := make(chan error, 1)
	go func() {
		_, err := io.Copy(stdin, readSourcePipe)
		stdin.Close()
		chCopySource <- err
	}()

	stderr, err := cmd.StderrPipe()
	if err != nil {
		cmd.Process.Kill()
		logger.Errorf("Rsync stderr read failed: %s: %v", path, err)
	}

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

	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		logger.Errorf("Rsync stderr read failed: %s: %v", path, err)
	}

	err = cmd.Wait()
	errs := []error{}
	chCopyRsyncErr := <-chCopyRsync
	chCopySourceErr := <-chCopySource

	if err != nil {
		errs = append(errs, err)

		// Try to get more info about the error.
		if chCopyRsyncErr != nil {
			errs = append(errs, chCopyRsyncErr)
		}

		if chCopySourceErr != nil {
			errs = append(errs, chCopySourceErr)
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("Rsync receive failed: %s: %v (%s)", path, errs, string(output))
	}

	return nil
}

func rsyncFeatureArgs(features []string) []string {
	args := []string{}
	if shared.StringInSlice("xattrs", features) {
		args = append(args, "--xattrs")
	}

	if shared.StringInSlice("delete", features) {
		args = append(args, "--delete")
	}

	if shared.StringInSlice("compress", features) {
		args = append(args, "--compress")
		args = append(args, "--compress-level=2")
	}

	return args
}
<|endoftext|>