<commit_before>package enum

//TreeEnum is a special type of Enum where the list of values also have a tree
//structure that can be interrogated. TreeEnums always have 0 map to "" as the
//root value.
type TreeEnum interface {
	Enum

	//IsLeaf returns true if the given value in the enum represents a leaf (as
	//opposed to branch)
	IsLeaf(val int) bool

	//Children returns all of the val beneath this branch that are direct
	//descendents, either including or excluding non-leaf nodes.
	Children(node int, includeBranches bool) []int

	//Descendants returns all enumvals beneath this point, recursively.
	Descendants(node int, includeBranches bool) []int

	NewImmutableTreeVal(val int) (ImmutableTreeVal, error)
	NewTreeVal() TreeVal

	MustNewImmutableTreeVal(val int) ImmutableTreeVal
	MustNewTreeVal(val int) TreeVal
}

//TreeValGetters is the collection of methods that TreeVals have beyodn normal
//Vals. It is factored out into a separate interface to clarify how
//ImmutableTreeVal and TreeVal differ from their non-treeval types.
type TreeValGetters interface {
	//IsLeaf is a convenience for val.Enum().TreeEnum().IsLeaf(val.Value())
	IsLeaf() bool

	//Children is a convenience for val.Enum().TreeEnum().Children(val).
	Children(includeBranches bool) []int

	//Descendants is a convenience for val.Enum().TreeEnum().Descendents(val).
	Descendants(includeBranches bool) []int

	//NodeString returns the name of this specific node, whereas String
	//returns the fully qualified name. So whereas String() might return
	//"Normal - Default - Save Item", NodeString() will return "Default".
	NodeString() string
}

//ImmutableTreeVal is a value from a TreeEnum.
type ImmutableTreeVal interface {
	ImmutableVal
	TreeValGetters
}

//TreeVal is a value from a tree enum.
type TreeVal interface {
	Val
	TreeValGetters
}

func (e *enum) IsLeaf(val int) bool {
	if e.tree == nil {
		return false
	}

	return len(e.tree[val]) == 0
}

func (e *enum) Children(node int, includeBranches bool) []int {
	if e.tree == nil {
		return nil
	}

	var result []int

	for _, val := range e.tree[node] {
		if !includeBranches && !e.IsLeaf(val) {
			continue
		}
		result = append(result, val)
	}

	return result
}

func (e *enum) Descendants(node int, includeBranches bool) []int {
	if e.tree == nil {
		return nil
	}

	if e.IsLeaf(node) {
		return []int{node}
	}

	var result []int

	if includeBranches {
		result = []int{node}
	}

	for _, val := range e.Children(node, true) {
		result = append(result, e.Descendants(val, includeBranches)...)
	}

	return result

}

func (e *enum) NewImmutableTreeVal(val int) (ImmutableTreeVal, error) {
	v := e.NewTreeVal()
	if err := v.SetValue(val); err != nil {
		return nil, err
	}
	return v, nil
}

func (e *enum) NewTreeVal() TreeVal {
	return &variable{
		e,
		e.DefaultValue(),
	}
}

func (e *enum) MustNewImmutableTreeVal(val int) ImmutableTreeVal {
	return e.MustNewTreeVal(val)
}

func (e *enum) MustNewTreeVal(val int) TreeVal {
	v := e.NewTreeVal()
	if err := v.SetValue(val); err != nil {
		panic("Couldn't set string value: " + err.Error())
	}
	return v
}

func (v *variable) IsLeaf() bool {
	return v.Enum().TreeEnum().IsLeaf(v.Value())
}

func (v *variable) Children(includeBranches bool) []int {
	return v.Enum().TreeEnum().Children(v.Value(), includeBranches)
}

func (v *variable) Descendants(includeBranches bool) []int {
	return v.Enum().TreeEnum().Descendants(v.Value(), includeBranches)
}

func (v *variable) NodeString() string {
	//TODO: actually implement this properly
	return v.String()
}
<commit_msg>Add TreeEnum.Parent() (inefficiently) and convenience on TreeVal. Part of #628.<commit_after>package enum

//TreeEnum is a special type of Enum where the list of values also have a tree
//structure that can be interrogated. TreeEnums always have 0 map to "" as the
//root value.
type TreeEnum interface {
	Enum

	//IsLeaf returns true if the given value in the enum represents a leaf (as
	//opposed to branch)
	IsLeaf(val int) bool

	//Parent returns the value of the node who is the direct parent. The root
	//will return itself.
	Parent(val int) int

	//Children returns all of the val beneath this branch that are direct
	//descendents, either including or excluding non-leaf nodes.
	Children(node int, includeBranches bool) []int

	//Descendants returns all enumvals beneath this point, recursively.
	Descendants(node int, includeBranches bool) []int

	NewImmutableTreeVal(val int) (ImmutableTreeVal, error)
	NewTreeVal() TreeVal

	MustNewImmutableTreeVal(val int) ImmutableTreeVal
	MustNewTreeVal(val int) TreeVal
}

//TreeValGetters is the collection of methods that TreeVals have beyodn normal
//Vals. It is factored out into a separate interface to clarify how
//ImmutableTreeVal and TreeVal differ from their non-treeval types.
type TreeValGetters interface {
	//IsLeaf is a convenience for val.Enum().TreeEnum().IsLeaf(val.Value())
	IsLeaf() bool

	//Parent is a convenience for val.Enum().TreeEnum().Parent(val).
	Parent() int

	//Children is a convenience for val.Enum().TreeEnum().Children(val).
	Children(includeBranches bool) []int

	//Descendants is a convenience for val.Enum().TreeEnum().Descendents(val).
	Descendants(includeBranches bool) []int

	//NodeString returns the name of this specific node, whereas String
	//returns the fully qualified name. So whereas String() might return
	//"Normal - Default - Save Item", NodeString() will return "Default".
	NodeString() string
}

//ImmutableTreeVal is a value from a TreeEnum.
type ImmutableTreeVal interface {
	ImmutableVal
	TreeValGetters
}

//TreeVal is a value from a tree enum.
type TreeVal interface {
	Val
	TreeValGetters
}

func (e *enum) IsLeaf(val int) bool {
	if e.tree == nil {
		return false
	}

	return len(e.tree[val]) == 0
}

func (e *enum) Parent(val int) int {
	if val == 0 {
		return 0
	}

	//TODO: more efficient implementation using the node -> parent map.

	for parent, children := range e.tree {
		for _, child := range children {
			if child == val {
				return parent
			}
		}
	}
	//This shouldn't happen if the value is actually in the tree.
	return -1
}

func (e *enum) Children(node int, includeBranches bool) []int {
	if e.tree == nil {
		return nil
	}

	var result []int

	for _, val := range e.tree[node] {
		if !includeBranches && !e.IsLeaf(val) {
			continue
		}
		result = append(result, val)
	}

	return result
}

func (e *enum) Descendants(node int, includeBranches bool) []int {
	if e.tree == nil {
		return nil
	}

	if e.IsLeaf(node) {
		return []int{node}
	}

	var result []int

	if includeBranches {
		result = []int{node}
	}

	for _, val := range e.Children(node, true) {
		result = append(result, e.Descendants(val, includeBranches)...)
	}

	return result

}

func (e *enum) NewImmutableTreeVal(val int) (ImmutableTreeVal, error) {
	v := e.NewTreeVal()
	if err := v.SetValue(val); err != nil {
		return nil, err
	}
	return v, nil
}

func (e *enum) NewTreeVal() TreeVal {
	return &variable{
		e,
		e.DefaultValue(),
	}
}

func (e *enum) MustNewImmutableTreeVal(val int) ImmutableTreeVal {
	return e.MustNewTreeVal(val)
}

func (e *enum) MustNewTreeVal(val int) TreeVal {
	v := e.NewTreeVal()
	if err := v.SetValue(val); err != nil {
		panic("Couldn't set string value: " + err.Error())
	}
	return v
}

func (v *variable) IsLeaf() bool {
	return v.Enum().TreeEnum().IsLeaf(v.Value())
}

func (v *variable) Parent() int {
	return v.Enum().TreeEnum().Parent(v.Value())
}

func (v *variable) Children(includeBranches bool) []int {
	return v.Enum().TreeEnum().Children(v.Value(), includeBranches)
}

func (v *variable) Descendants(includeBranches bool) []int {
	return v.Enum().TreeEnum().Descendants(v.Value(), includeBranches)
}

func (v *variable) NodeString() string {
	//TODO: actually implement this properly
	return v.String()
}
<|endoftext|>


<commit_before>// +build ignore

// Enumerate unique keys from key/values found in the whois responses.
// To use: go run enumerate.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"

	"github.com/domainr/whois"
	"github.com/domainr/whoistest"
	"github.com/wsxiaoys/terminal/color"
)

var (
	keys = make(map[string]string)
)

func main() {
	flag.Parse()
	if err := main1(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func main1() error {
	fns, err := whoistest.ResponseFiles()
	if err != nil {
		return err
	}
	for _, fn := range fns {
		res, err := whois.ReadMIMEFile(fn)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error reading response file %s: %s\n", fn, err)
			continue
		}
		if res.MediaType != "text/plain" {
			continue
		}
		scan(res)
	}

	sorted := make([]string, 0, len(keys))
	for k, _ := range keys {
		sorted = append(sorted, k)
	}
	sort.Strings(sorted)

	for _, k := range sorted {
		color.Printf("@{|c}%- 40s  @{|.}%s\n", k, keys[k])
	}
	color.Printf("@{|w}%d unique keys found\n", len(keys))

	return nil
}

var (
	reEmptyLine   = regexp.MustCompile(`^\s*$`)
	reKey         = `([^,a-z\:\],][^\:\]]{0,39}\S|[a-z-]{3,40})`
	reBareKey     = regexp.MustCompile(`^\s*` + reKey + `\s*\:\s*$`)
	reKeyValue    = regexp.MustCompile(`^\s*` + reKey + `\s*\:\s*(.*\S)\s*$`)
	reAltKey      = regexp.MustCompile(`^\s*\[` + reKey + `\]\s*$`)
	reAltKeyValue = regexp.MustCompile(`^\s*\[` + reKey + `\]\s*(.*\S)\s*$`)
	reBareValue   = regexp.MustCompile(`^      \s+(.*\S)\s*$`)
	reNotice      = regexp.MustCompile(strings.Join([]string{
		`^%`,                // whois.de, whois.registro.br
		`^# `,               // whois.kr
		`^\[ .+ \]$`,        // whois.jprs.jp
		`^>>>.+<<<$`,        // Database last updated...
		`^[^\:]+https?\://`, // Line with an URL
		`^NOTE: |^NOTICE: |^TERMS OF USE: `,
	}, "|"))
)

func scan(res *whois.Response) {
	r, err := res.Reader()
	if err != nil {
		return
	}
	line := 0
	s := bufio.NewScanner(r)
	for s.Scan() {
		line++
		color.Printf("@{|.}% 4d  ", line)

		// Get next line
		text := s.Text()

		// Notices and empty lines
		if reEmptyLine.MatchString(text) {
			color.Printf("@{|w}EMPTY\n")
			continue
		}
		if m := reNotice.FindStringSubmatch(text); m != nil {
			color.Printf("@{|w}%- 16s  %s\n", "NOTICE", text)
			continue
		}

		// Keys and values
		if m := reAltKeyValue.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%- 40s @{w}%s\n", "ALT_KEY_VALUE", m[1], m[2])
			continue
		}
		if m := reAltKey.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%s\n", "BARE_ALT_KEY", m[1])
			continue
		}
		if m := reKeyValue.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%- 40s @{w}%s\n", "KEY_VALUE", m[1], m[2])
			continue
		}
		if m := reBareKey.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%s\n", "BARE_KEY", m[1])
			continue
		}
		if m := reBareValue.FindStringSubmatch(text); m != nil {
			color.Printf("@{|w}%- 16s  @{c}%- 40s @{w}%s\n", "BARE_VALUE", "", m[1])
			continue
		}

		// Unknown
		color.Printf("@{|.}%- 16s  @{|.}%s\n", "UNKNOWN", text)
	}

	fmt.Printf("\n")
}

func addKey(k, host string) string {
	k = transformKey(k)
	if _, ok := keys[k]; !ok {
		keys[k] = host
	} else if !strings.Contains(keys[k], host) {
		keys[k] = keys[k] + "  " + host
	}
	return k
}

var (
	reStrip = regexp.MustCompile(`[[:punct:]]`)
	reSpace = regexp.MustCompile(`\s+`)
)

func transformKey(k string) string {
	k = strings.ToUpper(k)
	k = reStrip.ReplaceAllLiteralString(k, " ")
	k = strings.TrimSpace(k)
	k = reSpace.ReplaceAllLiteralString(k, "_")
	return k
}
<commit_msg>Print response filename.<commit_after>// +build ignore

// Enumerate unique keys from key/values found in the whois responses.
// To use: go run enumerate.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"

	"github.com/domainr/whois"
	"github.com/domainr/whoistest"
	"github.com/wsxiaoys/terminal/color"
)

var (
	keys = make(map[string]string)
)

func main() {
	flag.Parse()
	if err := main1(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func main1() error {
	wd, err := os.Getwd()
	if err != nil {
		return err
	}
	wd = wd + "/"

	fns, err := whoistest.ResponseFiles()
	if err != nil {
		return err
	}

	for _, fn := range fns {
		res, err := whois.ReadMIMEFile(fn)
		if err != nil {
			color.Fprintf(os.Stderr, "@{|r}Error reading response file %s: %s\n", fn, err)
			continue
		}
		if res.MediaType != "text/plain" {
			continue
		}
		color.Printf("@{|g}%s\n", strings.TrimPrefix(fn, wd))
		scan(res)
	}

	sorted := make([]string, 0, len(keys))
	for k, _ := range keys {
		sorted = append(sorted, k)
	}
	sort.Strings(sorted)

	for _, k := range sorted {
		color.Printf("@{|c}%- 40s  @{|.}%s\n", k, keys[k])
	}
	color.Printf("@{|w}%d unique keys found\n", len(keys))

	return nil
}

var (
	reEmptyLine   = regexp.MustCompile(`^\s*$`)
	reKey         = `([^,a-z\:\],][^\:\]]{0,39}\S|[a-z-]{3,40})`
	reBareKey     = regexp.MustCompile(`^\s*` + reKey + `\s*\:\s*$`)
	reKeyValue    = regexp.MustCompile(`^\s*` + reKey + `\s*\:\s*(.*\S)\s*$`)
	reAltKey      = regexp.MustCompile(`^\s*\[` + reKey + `\]\s*$`)
	reAltKeyValue = regexp.MustCompile(`^\s*\[` + reKey + `\]\s*(.*\S)\s*$`)
	reBareValue   = regexp.MustCompile(`^      \s+(.*\S)\s*$`)
	reNotice      = regexp.MustCompile(strings.Join([]string{
		`^%`,                // whois.de, whois.registro.br
		`^# `,               // whois.kr
		`^\[ .+ \]$`,        // whois.jprs.jp
		`^>>>.+<<<$`,        // Database last updated...
		`^[^\:]+https?\://`, // Line with an URL
		`^NOTE: |^NOTICE: |^TERMS OF USE: `,
	}, "|"))
)

func scan(res *whois.Response) {
	r, err := res.Reader()
	if err != nil {
		return
	}
	line := 0
	s := bufio.NewScanner(r)
	for s.Scan() {
		line++
		color.Printf("@{|.}% 4d  ", line)

		// Get next line
		text := s.Text()

		// Notices and empty lines
		if reEmptyLine.MatchString(text) {
			color.Printf("@{|w}EMPTY\n")
			continue
		}
		if m := reNotice.FindStringSubmatch(text); m != nil {
			color.Printf("@{|w}%- 16s  %s\n", "NOTICE", text)
			continue
		}

		// Keys and values
		if m := reAltKeyValue.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%- 40s @{w}%s\n", "ALT_KEY_VALUE", m[1], m[2])
			continue
		}
		if m := reAltKey.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%s\n", "BARE_ALT_KEY", m[1])
			continue
		}
		if m := reKeyValue.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%- 40s @{w}%s\n", "KEY_VALUE", m[1], m[2])
			continue
		}
		if m := reBareKey.FindStringSubmatch(text); m != nil {
			addKey(m[1], res.Host)
			color.Printf("@{|w}%- 16s  @{c}%s\n", "BARE_KEY", m[1])
			continue
		}
		if m := reBareValue.FindStringSubmatch(text); m != nil {
			color.Printf("@{|w}%- 16s  @{c}%- 40s @{w}%s\n", "BARE_VALUE", "", m[1])
			continue
		}

		// Unknown
		color.Printf("@{|.}%- 16s  @{|.}%s\n", "UNKNOWN", text)
	}

	fmt.Printf("\n")
}

func addKey(k, host string) string {
	k = transformKey(k)
	if _, ok := keys[k]; !ok {
		keys[k] = host
	} else if !strings.Contains(keys[k], host) {
		keys[k] = keys[k] + "  " + host
	}
	return k
}

var (
	reStrip = regexp.MustCompile(`[[:punct:]]`)
	reSpace = regexp.MustCompile(`\s+`)
)

func transformKey(k string) string {
	k = strings.ToUpper(k)
	k = reStrip.ReplaceAllLiteralString(k, " ")
	k = strings.TrimSpace(k)
	k = reSpace.ReplaceAllLiteralString(k, "_")
	return k
}
<|endoftext|>


<commit_before>package envconfig

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"
)

var (
	ErrUnexportedField = errors.New("envconfig: unexported field")
)

// Unmarshaler is the interface implemented by objects that can unmarshal
// a environment variable string of themselves.
type Unmarshaler interface {
	Unmarshal(s string) error
}

// Init reads the configuration from environment variables and populates the conf object. conf must be a pointer
func Init(conf interface{}) error {
	return InitWithPrefix(conf, "")
}

// InitWithPrefix reads the configuration from environment variables and populates the conf object. conf must be a pointer.
// Each key read will be prefixed with the prefix string.
func InitWithPrefix(conf interface{}, prefix string) error {
	value := reflect.ValueOf(conf)
	if value.Kind() != reflect.Ptr {
		return errors.New("envconfig: value is not a pointer")
	}

	elem := value.Elem()

	switch elem.Kind() {
	case reflect.Ptr:
		elem.Set(reflect.New(elem.Type().Elem()))
		return readStruct(elem.Elem(), prefix, false)
	case reflect.Struct:
		return readStruct(elem, prefix, false)
	default:
		return errors.New("envconfig: invalid value kind, only works on structs")
	}
}

type tag struct {
	customName string
	optional   bool
}

func parseTag(s string) *tag {
	tag := &tag{}

	tokens := strings.Split(s, ",")
	if len(tokens) == 0 {
		return tag
	}

	for _, v := range tokens {
		if v == "optional" {
			tag.optional = true
		} else {
			tag.customName = v
		}
	}

	return tag
}

func readStruct(value reflect.Value, parentName string, optional bool) (err error) {
	for i := 0; i < value.NumField(); i++ {
		field := value.Field(i)
		name := value.Type().Field(i).Name

		tag := parseTag(value.Type().Field(i).Tag.Get("envconfig"))

		var combinedName string
		if tag.customName != "" {
			combinedName = tag.customName
		} else {
			combinedName = combineName(parentName, name)
		}

	doRead:
		switch field.Kind() {
		case reflect.Ptr:
			// it's a pointer, create a new value and restart the switch
			field.Set(reflect.New(field.Type().Elem()))
			field = field.Elem()
			goto doRead
		case reflect.Struct:
			err = readStruct(field, combinedName, optional || tag.optional)
		default:
			err = setField(field, combinedName, tag.customName != "", optional || tag.optional)
		}

		if err != nil {
			return err
		}
	}

	return
}

func setField(value reflect.Value, name string, customName, optional bool) (err error) {
	str, err := readValue(name, customName, optional)
	if err != nil {
		return err
	}

	switch value.Kind() {
	case reflect.Slice:
		return setSliceField(value, str)
	default:
		return parseValue(value, str)
	}
}

func setSliceField(value reflect.Value, str string) error {
	elType := value.Type().Elem()
	tnz := newSliceTokenizer(str)

	slice := reflect.MakeSlice(value.Type(), value.Len(), value.Cap())

	for tnz.scan() {
		token := tnz.text()

		el := reflect.New(elType).Elem()

		if err := parseValue(el, token); err != nil {
			return err
		}

		slice = reflect.Append(slice, el)
	}

	value.Set(slice)

	return tnz.Err()
}

var (
	durationType    = reflect.TypeOf(new(time.Duration)).Elem()
	unmarshalerType = reflect.TypeOf(new(Unmarshaler)).Elem()
)

func isDurationField(t reflect.Type) bool {
	return t.AssignableTo(durationType)
}

func isUnmarshaler(t reflect.Type) bool {
	return t.Implements(unmarshalerType) || reflect.PtrTo(t).Implements(unmarshalerType)
}

func parseValue(v reflect.Value, str string) (err error) {
	if !v.CanSet() {
		return ErrUnexportedField
	}

	vtype := v.Type()

	// Special case for Unmarshaler
	if isUnmarshaler(vtype) {
		return parseWithUnmarshaler(v, str)
	}

	// Special case for time.Duration
	if isDurationField(vtype) {
		return parseDuration(v, str)
	}

	kind := vtype.Kind()
	switch kind {
	case reflect.Bool:
		err = parseBoolValue(v, str)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		err = parseIntValue(v, str)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		err = parseUintValue(v, str)
	case reflect.Float32, reflect.Float64:
		err = parseFloatValue(v, str)
	case reflect.Ptr:
		v.Set(reflect.New(vtype.Elem()))
		return parseValue(v.Elem(), str)
	case reflect.String:
		v.SetString(str)
	case reflect.Struct:
		err = parseStruct(v, str)
	default:
		return fmt.Errorf("envconfig: slice element type %v not supported", kind)
	}

	return
}

func parseWithUnmarshaler(v reflect.Value, str string) error {
	var u Unmarshaler
	vtype := v.Type()

	if vtype.Implements(unmarshalerType) {
		u = v.Interface().(Unmarshaler)
	} else if reflect.PtrTo(vtype).Implements(unmarshalerType) {
		// We know the interface has a pointer receiver, but our value might not be a pointer, so we get one
		if v.Kind() != reflect.Ptr {
			u = v.Addr().Interface().(Unmarshaler)
		} else {
			u = v.Interface().(Unmarshaler)
		}
	}

	return u.Unmarshal(str)
}

func parseDuration(v reflect.Value, str string) error {
	d, err := time.ParseDuration(str)
	if err != nil {
		return err
	}

	v.SetInt(int64(d))

	return nil
}

func parseStruct(value reflect.Value, token string) error {
	tokens := strings.Split(token[1:len(token)-1], ",")
	if len(tokens) == 0 {
		return errors.New("envconfig: struct token should not be empty")
	}

	if len(tokens) != value.NumField() {
		return fmt.Errorf("envconfig: struct token has %d fields but struct has %d",
			len(tokens), value.NumField(),
		)
	}

	for i := 0; i < value.NumField(); i++ {
		field := value.Field(i)
		t := tokens[i]

		if err := parseValue(field, t); err != nil {
			return err
		}
	}

	return nil
}

func parseBoolValue(v reflect.Value, str string) error {
	val, err := strconv.ParseBool(str)
	if err != nil {
		return err
	}
	v.SetBool(val)

	return nil
}

func parseIntValue(v reflect.Value, str string) error {
	val, err := strconv.ParseInt(str, 10, 64)
	if err != nil {
		return err
	}
	v.SetInt(val)

	return nil
}

func parseUintValue(v reflect.Value, str string) error {
	val, err := strconv.ParseUint(str, 10, 64)
	if err != nil {
		return err
	}
	v.SetUint(val)

	return nil
}

func parseFloatValue(v reflect.Value, str string) error {
	val, err := strconv.ParseFloat(str, 64)
	if err != nil {
		return err
	}
	v.SetFloat(val)

	return nil
}

func nameToKey(name string) string {
	s := strings.Replace(name, ".", "_", -1)
	s = strings.ToUpper(s)

	return s
}

func combineName(parentName, name string) string {
	if parentName == "" {
		return name
	}

	return parentName + "." + name
}

func readValue(name string, customName, optional bool) (string, error) {
	key := name
	if !customName {
		key = nameToKey(name)
	}

	str := os.Getenv(key)
	if str != "" {
		return str, nil
	}

	if optional {
		return str, nil
	}

	return "", fmt.Errorf("envconfig: key %v not found", key)
}
<commit_msg>optional bool should not throw exception if empty<commit_after>package envconfig

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"
)

var (
	ErrUnexportedField = errors.New("envconfig: unexported field")
)

// Unmarshaler is the interface implemented by objects that can unmarshal
// a environment variable string of themselves.
type Unmarshaler interface {
	Unmarshal(s string) error
}

// Init reads the configuration from environment variables and populates the conf object. conf must be a pointer
func Init(conf interface{}) error {
	return InitWithPrefix(conf, "")
}

// InitWithPrefix reads the configuration from environment variables and populates the conf object. conf must be a pointer.
// Each key read will be prefixed with the prefix string.
func InitWithPrefix(conf interface{}, prefix string) error {
	value := reflect.ValueOf(conf)
	if value.Kind() != reflect.Ptr {
		return errors.New("envconfig: value is not a pointer")
	}

	elem := value.Elem()

	switch elem.Kind() {
	case reflect.Ptr:
		elem.Set(reflect.New(elem.Type().Elem()))
		return readStruct(elem.Elem(), prefix, false)
	case reflect.Struct:
		return readStruct(elem, prefix, false)
	default:
		return errors.New("envconfig: invalid value kind, only works on structs")
	}
}

type tag struct {
	customName string
	optional   bool
}

func parseTag(s string) *tag {
	tag := &tag{}

	tokens := strings.Split(s, ",")
	if len(tokens) == 0 {
		return tag
	}

	for _, v := range tokens {
		if v == "optional" {
			tag.optional = true
		} else {
			tag.customName = v
		}
	}

	return tag
}

func readStruct(value reflect.Value, parentName string, optional bool) (err error) {
	for i := 0; i < value.NumField(); i++ {
		field := value.Field(i)
		name := value.Type().Field(i).Name

		tag := parseTag(value.Type().Field(i).Tag.Get("envconfig"))

		var combinedName string
		if tag.customName != "" {
			combinedName = tag.customName
		} else {
			combinedName = combineName(parentName, name)
		}

	doRead:
		switch field.Kind() {
		case reflect.Ptr:
			// it's a pointer, create a new value and restart the switch
			field.Set(reflect.New(field.Type().Elem()))
			field = field.Elem()
			goto doRead
		case reflect.Struct:
			err = readStruct(field, combinedName, optional || tag.optional)
		default:
			err = setField(field, combinedName, tag.customName != "", optional || tag.optional)
		}

		if err != nil {
			return err
		}
	}

	return
}

func setField(value reflect.Value, name string, customName, optional bool) (err error) {
	str, err := readValue(name, customName, optional)
	if err != nil {
		return err
	}
	if len(str) == 0 && optional {
		return nil
	}

	switch value.Kind() {
	case reflect.Slice:
		return setSliceField(value, str)
	default:
		return parseValue(value, str)
	}
}

func setSliceField(value reflect.Value, str string) error {
	elType := value.Type().Elem()
	tnz := newSliceTokenizer(str)

	slice := reflect.MakeSlice(value.Type(), value.Len(), value.Cap())

	for tnz.scan() {
		token := tnz.text()

		el := reflect.New(elType).Elem()

		if err := parseValue(el, token); err != nil {
			return err
		}

		slice = reflect.Append(slice, el)
	}

	value.Set(slice)

	return tnz.Err()
}

var (
	durationType    = reflect.TypeOf(new(time.Duration)).Elem()
	unmarshalerType = reflect.TypeOf(new(Unmarshaler)).Elem()
)

func isDurationField(t reflect.Type) bool {
	return t.AssignableTo(durationType)
}

func isUnmarshaler(t reflect.Type) bool {
	return t.Implements(unmarshalerType) || reflect.PtrTo(t).Implements(unmarshalerType)
}

func parseValue(v reflect.Value, str string) (err error) {
	if !v.CanSet() {
		return ErrUnexportedField
	}

	vtype := v.Type()

	// Special case for Unmarshaler
	if isUnmarshaler(vtype) {
		return parseWithUnmarshaler(v, str)
	}

	// Special case for time.Duration
	if isDurationField(vtype) {
		return parseDuration(v, str)
	}

	kind := vtype.Kind()
	switch kind {
	case reflect.Bool:
		err = parseBoolValue(v, str)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		err = parseIntValue(v, str)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		err = parseUintValue(v, str)
	case reflect.Float32, reflect.Float64:
		err = parseFloatValue(v, str)
	case reflect.Ptr:
		v.Set(reflect.New(vtype.Elem()))
		return parseValue(v.Elem(), str)
	case reflect.String:
		v.SetString(str)
	case reflect.Struct:
		err = parseStruct(v, str)
	default:
		return fmt.Errorf("envconfig: slice element type %v not supported", kind)
	}

	return
}

func parseWithUnmarshaler(v reflect.Value, str string) error {
	var u Unmarshaler
	vtype := v.Type()

	if vtype.Implements(unmarshalerType) {
		u = v.Interface().(Unmarshaler)
	} else if reflect.PtrTo(vtype).Implements(unmarshalerType) {
		// We know the interface has a pointer receiver, but our value might not be a pointer, so we get one
		if v.Kind() != reflect.Ptr {
			u = v.Addr().Interface().(Unmarshaler)
		} else {
			u = v.Interface().(Unmarshaler)
		}
	}

	return u.Unmarshal(str)
}

func parseDuration(v reflect.Value, str string) error {
	d, err := time.ParseDuration(str)
	if err != nil {
		return err
	}

	v.SetInt(int64(d))

	return nil
}

func parseStruct(value reflect.Value, token string) error {
	tokens := strings.Split(token[1:len(token)-1], ",")
	if len(tokens) == 0 {
		return errors.New("envconfig: struct token should not be empty")
	}

	if len(tokens) != value.NumField() {
		return fmt.Errorf("envconfig: struct token has %d fields but struct has %d",
			len(tokens), value.NumField(),
		)
	}

	for i := 0; i < value.NumField(); i++ {
		field := value.Field(i)
		t := tokens[i]

		if err := parseValue(field, t); err != nil {
			return err
		}
	}

	return nil
}

func parseBoolValue(v reflect.Value, str string) error {
	val, err := strconv.ParseBool(str)
	if err != nil {
		return err
	}
	v.SetBool(val)

	return nil
}

func parseIntValue(v reflect.Value, str string) error {
	val, err := strconv.ParseInt(str, 10, 64)
	if err != nil {
		return err
	}
	v.SetInt(val)

	return nil
}

func parseUintValue(v reflect.Value, str string) error {
	val, err := strconv.ParseUint(str, 10, 64)
	if err != nil {
		return err
	}
	v.SetUint(val)

	return nil
}

func parseFloatValue(v reflect.Value, str string) error {
	val, err := strconv.ParseFloat(str, 64)
	if err != nil {
		return err
	}
	v.SetFloat(val)

	return nil
}

func nameToKey(name string) string {
	s := strings.Replace(name, ".", "_", -1)
	s = strings.ToUpper(s)

	return s
}

func combineName(parentName, name string) string {
	if parentName == "" {
		return name
	}

	return parentName + "." + name
}

func readValue(name string, customName, optional bool) (string, error) {
	key := name
	if !customName {
		key = nameToKey(name)
	}

	str := os.Getenv(key)
	if str != "" {
		return str, nil
	}

	if optional {
		return str, nil
	}

	return "", fmt.Errorf("envconfig: key %v not found", key)
}
<|endoftext|>


<commit_before>package middlewares

import (
	"net/http"
	"strconv"
	"strings"

	build "github.com/cozy/cozy-stack/pkg/config"
	"github.com/labstack/echo/v4"
	"github.com/mssola/user_agent"
)

// Some constants for the browser names
const (
	InternetExplorer = "Internet Explorer"
	Edge             = "Edge"
	Firefox          = "Firefox"
	Chrome           = "Chrome"
	Chromium         = "Chromium"
	Opera            = "Opera"
	Safari           = "Safari"
	Android          = "Android"
)

// browser is a struct with a name and a minimal version
type browser struct {
	name       string
	minVersion *int
}

// We don't support Edge before version 17, because some webapps (like Drive)
// needs URLSearchParams.
var minEdgeVersion = 17

// We don't support Firefox before version 52
var minFirefoxVersion = 52

var rules = []browser{
	{
		name: InternetExplorer,
	},
	{
		name:       Edge,
		minVersion: &minEdgeVersion,
	},
	{
		name:       Firefox,
		minVersion: &minFirefoxVersion,
	},
}

// CheckUserAgent is a middleware that shows an HTML page of error when a
// browser that is not supported try to load a webapp.
func CheckUserAgent(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		ua := user_agent.New(c.Request().UserAgent())
		browser, rawVersion := ua.Browser()
		acceptHeader := c.Request().Header.Get(echo.HeaderAccept)

		if strings.Contains(acceptHeader, echo.MIMETextHTML) {
			for _, rule := range rules {
				if browser == rule.name {
					version, ok := getMajorVersion(rawVersion)
					if !ok || (rule.minVersion != nil && !(version < *rule.minVersion)) {
						return next(c)
					}

					instance := GetInstance(c)
					return c.Render(http.StatusOK, "compat.html", echo.Map{
						"Domain":      instance.ContextualDomain(),
						"ContextName": instance.ContextName,
						"Locale":      instance.Locale,
						"Favicon":     Favicon(instance),
					})
				}
			}
		}
		return next(c)
	}
}

// getMajorVersion returns the major version of a browser
// 12 => 12
// 12.13 => 12
func getMajorVersion(rawVersion string) (int, bool) {
	splitted := strings.SplitN(rawVersion, ".", 2)
	v, err := strconv.Atoi(splitted[0])
	if err != nil {
		return -1, false
	}
	return v, true
}

// CryptoPolyfill returns true if the browser can't use its window.crypto API
// to hash the password with PBKDF2. It is the case for Edge, but also for
// Chrome in development mode, because this API is only available in secure
// more (HTTPS).
func CryptoPolyfill(c echo.Context) bool {
	ua := user_agent.New(c.Request().UserAgent())
	browser, _ := ua.Browser()
	if browser == Edge {
		return true
	}
	if build.IsDevRelease() {
		// XXX electron is seen as Safari
		return browser == Chrome || browser == Chromium || browser == Opera || browser == Safari || browser == Android
	}
	return false
}
<commit_msg>Block old versions of Safari (#2310)<commit_after>package middlewares

import (
	"net/http"
	"strconv"
	"strings"

	build "github.com/cozy/cozy-stack/pkg/config"
	"github.com/labstack/echo/v4"
	"github.com/mssola/user_agent"
)

// Some constants for the browser names
const (
	InternetExplorer = "Internet Explorer"
	Edge             = "Edge"
	Firefox          = "Firefox"
	Chrome           = "Chrome"
	Chromium         = "Chromium"
	Opera            = "Opera"
	Safari           = "Safari"
	Android          = "Android"
)

// browser is a struct with a name and a minimal version
type browser struct {
	name       string
	minVersion *int
}

// We don't support Edge before version 17, because some webapps (like Drive)
// needs URLSearchParams.
var minEdgeVersion = 17

// We don't support Firefox before version 52
var minFirefoxVersion = 52

// We don't support Safari before version 11, as window.crypto is not
// available.
var minSafariVersion = 11

var rules = []browser{
	{
		name: InternetExplorer,
	},
	{
		name:       Edge,
		minVersion: &minEdgeVersion,
	},
	{
		name:       Firefox,
		minVersion: &minFirefoxVersion,
	},
	{
		name:       Safari,
		minVersion: &minSafariVersion,
	},
}

// CheckUserAgent is a middleware that shows an HTML page of error when a
// browser that is not supported try to load a webapp.
func CheckUserAgent(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		ua := user_agent.New(c.Request().UserAgent())
		browser, rawVersion := ua.Browser()
		acceptHeader := c.Request().Header.Get(echo.HeaderAccept)

		if strings.Contains(acceptHeader, echo.MIMETextHTML) {
			for _, rule := range rules {
				if browser == rule.name {
					version, ok := getMajorVersion(rawVersion)
					if !ok || (rule.minVersion != nil && !(version < *rule.minVersion)) {
						return next(c)
					}

					instance := GetInstance(c)
					return c.Render(http.StatusOK, "compat.html", echo.Map{
						"Domain":      instance.ContextualDomain(),
						"ContextName": instance.ContextName,
						"Locale":      instance.Locale,
						"Favicon":     Favicon(instance),
					})
				}
			}
		}
		return next(c)
	}
}

// getMajorVersion returns the major version of a browser
// 12 => 12
// 12.13 => 12
func getMajorVersion(rawVersion string) (int, bool) {
	splitted := strings.SplitN(rawVersion, ".", 2)
	v, err := strconv.Atoi(splitted[0])
	if err != nil {
		return -1, false
	}
	return v, true
}

// CryptoPolyfill returns true if the browser can't use its window.crypto API
// to hash the password with PBKDF2. It is the case for Edge, but also for
// Chrome in development mode, because this API is only available in secure
// more (HTTPS).
func CryptoPolyfill(c echo.Context) bool {
	ua := user_agent.New(c.Request().UserAgent())
	browser, _ := ua.Browser()
	if browser == Edge {
		return true
	}
	if build.IsDevRelease() {
		// XXX electron is seen as Safari
		return browser == Chrome || browser == Chromium || browser == Opera || browser == Safari || browser == Android
	}
	return false
}
<|endoftext|>


<commit_before>package etcd

import (
	"bytes"
	"encoding/gob"
	"net/http"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/coreos/etcd/etcdserver/stats"
	"github.com/coreos/etcd/pkg/types"
	etcdraft "github.com/coreos/etcd/raft"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/coreos/etcd/rafthttp"
	"github.com/relab/raft"
	"github.com/relab/raft/commonpb"
)

type tag struct {
	UID  uint64
	Data []byte
}

type future struct {
	res chan raft.Result
}

func (f *future) ResultCh() <-chan raft.Result {
	return f.res
}

// Wrapper wraps an etcd/raft.Node and implements relab/raft.Raft and
// etcd/rafthttp.Raft.
type Wrapper struct {
	n         etcdraft.Node
	sm        raft.StateMachine
	storage   *etcdraft.MemoryStorage
	transport *rafthttp.Transport
	logger    logrus.FieldLogger

	heartbeat time.Duration

	uid       uint64
	propLock  sync.Mutex
	proposals map[uint64]chan<- raft.Result

	apply chan []raftpb.Entry
}

func (w *Wrapper) encodeProposal(data []byte) (uint64, []byte, error) {
	uid := atomic.AddUint64(&w.uid, 1)
	t := &tag{
		UID:  uid,
		Data: data,
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	err := enc.Encode(t)
	if err != nil {
		return 0, nil, err
	}

	return uid, b.Bytes(), nil
}

func (w *Wrapper) newFuture(uid uint64) raft.Future {
	res := make(chan raft.Result, 1)
	future := &future{res}

	w.propLock.Lock()
	w.proposals[uid] = res
	w.propLock.Unlock()

	return future
}

// Hack: Returns nil if commit fails to decode. This when etcd append the
// initial configuration to the log. Fix: Make state machine return
// clientID/clientSeq along with result, we can then use that to figure out who
// to respond to instead of using the tag struct. Etcd uses the same hack
// internally so we'll just leave it be for now.
func (w *Wrapper) decodeCommit(commit []byte) *tag {
	b := bytes.NewBuffer(commit)
	dec := gob.NewDecoder(b)

	var t tag
	err := dec.Decode(&t)

	if err != nil {
		return nil
	}

	return &t
}

func NewRaft(logger logrus.FieldLogger,
	sm raft.StateMachine, storage *etcdraft.MemoryStorage, cfg *etcdraft.Config,
	peers []etcdraft.Peer, heartbeat time.Duration,
) *Wrapper {
	w := &Wrapper{
		sm:        sm,
		storage:   storage,
		heartbeat: heartbeat,
		proposals: make(map[uint64]chan<- raft.Result),
		logger:    logger,
		apply:     make(chan []raftpb.Entry, 2048),
	}
	w.n = etcdraft.StartNode(cfg, append(peers, etcdraft.Peer{ID: cfg.ID}))

	ss := &stats.ServerStats{}
	ss.Initialize()

	transport := &rafthttp.Transport{
		ID:          types.ID(cfg.ID),
		ClusterID:   1,
		Raft:        w,
		ServerStats: ss,
		LeaderStats: stats.NewLeaderStats(strconv.FormatUint(cfg.ID, 10)),
		ErrorC:      make(chan error),
	}
	err := transport.Start()

	if err != nil {
		panic("start transport: " + err.Error())
	}

	for i, peer := range peers {
		w.logger.WithFields(logrus.Fields{
			"i":       i,
			"id":      peer.ID,
			"context": string(peer.Context),
		}).Warnln("Add peer")

		transport.AddPeer(types.ID(peer.ID), []string{string(peer.Context)})
	}

	w.transport = transport

	go w.runSM()
	go w.run()
	return w
}

func (w *Wrapper) Handler() http.Handler {
	return w.transport.Handler()
}

func (w *Wrapper) ProposeCmd(ctx context.Context, req []byte) (raft.Future, error) {
	uid, prop, err := w.encodeProposal(req)

	if err != nil {
		return nil, err
	}

	future := w.newFuture(uid)
	err = w.n.Propose(ctx, prop)

	if err != nil {
		w.propLock.Lock()
		delete(w.proposals, uid)
		w.propLock.Unlock()
		return nil, err
	}

	return future, nil
}

func (w *Wrapper) ReadCmd(context.Context, []byte) (raft.Future, error) {
	w.logger.Warnln("ProposeRead")
	return nil, nil
}

func (w *Wrapper) ProposeConf(ctx context.Context, req *commonpb.ReconfRequest) (raft.Future, error) {
	cc := raftpb.ConfChange{
		Type:    raftpb.ConfChangeType(req.ReconfType),
		ID:      req.ServerID, // ?
		NodeID:  req.ServerID,
		Context: []byte(""), // ?
	}

	err := w.n.ProposeConfChange(ctx, cc)

	w.logger.WithError(err).WithFields(logrus.Fields{
		"req": req,
		"cc":  cc,
	}).Warnln("ProposeConf")

	if err != nil {
		return nil, err
	}

	panic("ProposeConf not implemented")
}

func (w *Wrapper) run() {
	s := time.NewTicker(w.heartbeat)

	for {
		select {
		case <-s.C:
			w.n.Tick()
		case rd := <-w.n.Ready():
			if rd.SoftState != nil {
				rmetrics.leader.Set(float64(rd.Lead))
				w.propLock.Lock()
				if rd.RaftState != etcdraft.StateLeader && len(w.proposals) > 0 {
					for uid, respCh := range w.proposals {
						respCh <- raft.Result{
							Value: raft.ErrNotLeader{Leader: rd.Lead},
						}
						delete(w.proposals, uid)
					}
				}
				w.propLock.Unlock()
			}
			w.storage.Append(rd.Entries)
			if !etcdraft.IsEmptyHardState(rd.HardState) {
				w.storage.SetHardState(rd.HardState)
			}
			if !etcdraft.IsEmptySnap(rd.Snapshot) {
				w.storage.ApplySnapshot(rd.Snapshot)
			}
			w.n.Advance()
			w.transport.Send(rd.Messages)
			w.apply <- rd.CommittedEntries
		}
	}
}

func (w *Wrapper) runSM() {
	for entries := range w.apply {
		tags := make(map[uint64]*tag)

		for _, entry := range entries {
			tags[entry.Index] = w.decodeCommit(entry.Data)
		}

		chs := make(map[uint64]chan<- raft.Result)

		w.propLock.Lock()
		for index, tag := range tags {
			if tag == nil {
				continue
			}
			chs[index] = w.proposals[tag.UID]
		}
		w.propLock.Unlock()

		for _, entry := range entries {
			rmetrics.commitIndex.Set(float64(entry.Index))

			t := tags[entry.Index]
			ch := chs[entry.Index]

			switch entry.Type {
			case raftpb.EntryNormal:
				w.handleNormal(&entry, t, ch)
			case raftpb.EntryConfChange:
				w.logger.WithField(
					"entry", etcdraft.DescribeEntry(entry, nil),
				).Warnln("Committed conf change entry")

				w.handleConfChange(&entry, t, ch)
			}

			// TODO Respond to future.
		}

		w.propLock.Lock()
		for _, tag := range tags {
			if tag == nil {
				continue
			}
			delete(w.proposals, tag.UID)
		}
		w.propLock.Unlock()
	}
}

func (w *Wrapper) handleNormal(entry *raftpb.Entry, t *tag, respCh chan<- raft.Result) {
	if t == nil {
		return
	}

	res := w.sm.Apply(&commonpb.Entry{
		Term:      entry.Term,
		Index:     entry.Index,
		EntryType: commonpb.EntryNormal,
		Data:      t.Data,
	})

	if respCh == nil {
		return
	}

	respCh <- raft.Result{
		Index: entry.Index,
		Value: res,
	}
}

func (w *Wrapper) handleConfChange(entry *raftpb.Entry, t *tag, respCh chan<- raft.Result) {
	var cc raftpb.ConfChange
	err := cc.Unmarshal(entry.Data)

	if err != nil {
		panic("unmarshal conf change: " + err.Error())
	}

	w.logger.WithField("cc", cc).Warnln("Applying conf change")
	w.n.ApplyConfChange(cc)
	switch cc.Type {
	case raftpb.ConfChangeAddNode:
		if len(cc.Context) > 0 {
			w.transport.AddPeer(types.ID(cc.NodeID), []string{string(cc.Context)})
		}
	case raftpb.ConfChangeRemoveNode:
		tid := types.ID(cc.NodeID)
		if tid == w.transport.ID {
			w.logger.Warnln("Shutting down")
			return
		}
		w.transport.RemovePeer(tid)
	}

	if t == nil {
		return
	}

	// Inform state machine about new configuration.
	w.sm.Apply(&commonpb.Entry{
		Term:      entry.Term,
		Index:     entry.Index,
		EntryType: commonpb.EntryReconf,
		Data:      t.Data,
	})

	// TODO Respond / Cleanup future.
}

func (w *Wrapper) Process(ctx context.Context, m raftpb.Message) error {
	return w.n.Step(ctx, m)
}
func (w *Wrapper) IsIDRemoved(id uint64) bool                               { return false }
func (w *Wrapper) ReportUnreachable(id uint64)                              {}
func (w *Wrapper) ReportSnapshot(id uint64, status etcdraft.SnapshotStatus) {}
<commit_msg>etcd/raft.go: Append uid to byte slice rather than wrap it<commit_after>package etcd

import (
	"encoding/binary"
	"net/http"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/coreos/etcd/etcdserver/stats"
	"github.com/coreos/etcd/pkg/types"
	etcdraft "github.com/coreos/etcd/raft"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/coreos/etcd/rafthttp"
	"github.com/relab/raft"
	"github.com/relab/raft/commonpb"
)

type future struct {
	res chan raft.Result
}

func (f *future) ResultCh() <-chan raft.Result {
	return f.res
}

// Wrapper wraps an etcd/raft.Node and implements relab/raft.Raft and
// etcd/rafthttp.Raft.
type Wrapper struct {
	n         etcdraft.Node
	sm        raft.StateMachine
	storage   *etcdraft.MemoryStorage
	transport *rafthttp.Transport
	logger    logrus.FieldLogger

	heartbeat time.Duration

	uid       uint64
	propLock  sync.Mutex
	proposals map[uint64]chan<- raft.Result

	apply chan []raftpb.Entry
}

func (w *Wrapper) newFuture(uid uint64) raft.Future {
	res := make(chan raft.Result, 1)
	future := &future{res}

	w.propLock.Lock()
	w.proposals[uid] = res
	w.propLock.Unlock()

	return future
}

func NewRaft(logger logrus.FieldLogger,
	sm raft.StateMachine, storage *etcdraft.MemoryStorage, cfg *etcdraft.Config,
	peers []etcdraft.Peer, heartbeat time.Duration,
) *Wrapper {
	w := &Wrapper{
		sm:        sm,
		storage:   storage,
		heartbeat: heartbeat,
		proposals: make(map[uint64]chan<- raft.Result),
		logger:    logger,
		apply:     make(chan []raftpb.Entry, 2048),
	}
	w.n = etcdraft.StartNode(cfg, append(peers, etcdraft.Peer{ID: cfg.ID}))

	ss := &stats.ServerStats{}
	ss.Initialize()

	transport := &rafthttp.Transport{
		ID:          types.ID(cfg.ID),
		ClusterID:   1,
		Raft:        w,
		ServerStats: ss,
		LeaderStats: stats.NewLeaderStats(strconv.FormatUint(cfg.ID, 10)),
		ErrorC:      make(chan error),
	}
	err := transport.Start()

	if err != nil {
		panic("start transport: " + err.Error())
	}

	for i, peer := range peers {
		w.logger.WithFields(logrus.Fields{
			"i":       i,
			"id":      peer.ID,
			"context": string(peer.Context),
		}).Warnln("Add peer")

		transport.AddPeer(types.ID(peer.ID), []string{string(peer.Context)})
	}

	w.transport = transport

	go w.runSM()
	go w.run()
	return w
}

func (w *Wrapper) Handler() http.Handler {
	return w.transport.Handler()
}

func (w *Wrapper) ProposeCmd(ctx context.Context, req []byte) (raft.Future, error) {
	uid := atomic.AddUint64(&w.uid, 1)
	b := make([]byte, 9)
	binary.LittleEndian.PutUint64(b, uid)
	b[8] = 0x9
	req = append(req, b...)

	future := w.newFuture(uid)

	if err := w.n.Propose(ctx, req); err != nil {
		w.propLock.Lock()
		delete(w.proposals, uid)
		w.propLock.Unlock()
		return nil, err
	}

	return future, nil
}

func (w *Wrapper) ReadCmd(context.Context, []byte) (raft.Future, error) {
	w.logger.Warnln("ProposeRead")
	return nil, nil
}

func (w *Wrapper) ProposeConf(ctx context.Context, req *commonpb.ReconfRequest) (raft.Future, error) {
	cc := raftpb.ConfChange{
		Type:    raftpb.ConfChangeType(req.ReconfType),
		ID:      req.ServerID, // ?
		NodeID:  req.ServerID,
		Context: []byte(""), // ?
	}

	err := w.n.ProposeConfChange(ctx, cc)

	w.logger.WithError(err).WithFields(logrus.Fields{
		"req": req,
		"cc":  cc,
	}).Warnln("ProposeConf")

	if err != nil {
		return nil, err
	}

	panic("ProposeConf not implemented")
}

func (w *Wrapper) run() {
	s := time.NewTicker(w.heartbeat)

	for {
		select {
		case <-s.C:
			w.n.Tick()
		case rd := <-w.n.Ready():
			if rd.SoftState != nil {
				rmetrics.leader.Set(float64(rd.Lead))
				w.propLock.Lock()
				if rd.RaftState != etcdraft.StateLeader && len(w.proposals) > 0 {
					for uid, respCh := range w.proposals {
						respCh <- raft.Result{
							Value: raft.ErrNotLeader{Leader: rd.Lead},
						}
						delete(w.proposals, uid)
					}
				}
				w.propLock.Unlock()
			}
			w.storage.Append(rd.Entries)
			if !etcdraft.IsEmptyHardState(rd.HardState) {
				w.storage.SetHardState(rd.HardState)
			}
			if !etcdraft.IsEmptySnap(rd.Snapshot) {
				w.storage.ApplySnapshot(rd.Snapshot)
			}
			w.n.Advance()
			w.transport.Send(rd.Messages)
			w.apply <- rd.CommittedEntries
		}
	}
}

func (w *Wrapper) runSM() {
	for entries := range w.apply {
		for _, entry := range entries {
			rmetrics.commitIndex.Set(float64(entry.Index))

			switch entry.Type {
			case raftpb.EntryNormal:
				w.handleNormal(&entry)
			case raftpb.EntryConfChange:
				w.logger.WithField(
					"entry", etcdraft.DescribeEntry(entry, nil),
				).Warnln("Committed conf change entry")

				w.handleConfChange(&entry)
			}
		}
	}
}

func (w *Wrapper) handleNormal(entry *raftpb.Entry) {
	if len(entry.Data) == 0 {
		return
	}

	uid := binary.LittleEndian.Uint64(entry.Data[len(entry.Data)-9 : len(entry.Data)-1])
	data := entry.Data[:len(entry.Data)-9]

	res := w.sm.Apply(&commonpb.Entry{
		Term:      entry.Term,
		Index:     entry.Index,
		EntryType: commonpb.EntryNormal,
		Data:      data,
	})

	w.propLock.Lock()
	ch, ok := w.proposals[uid]
	w.propLock.Unlock()

	if !ok {
		return
	}

	ch <- raft.Result{
		Index: entry.Index,
		Value: res,
	}

	w.propLock.Lock()
	delete(w.proposals, uid)
	w.propLock.Unlock()
}

func (w *Wrapper) handleConfChange(entry *raftpb.Entry) {
	var remove int
	if entry.Data[len(entry.Data)-1] == 0x9 {
		remove = 9
	}
	data := entry.Data[:len(entry.Data)-remove]
	var cc raftpb.ConfChange
	err := cc.Unmarshal(data)

	if err != nil {
		panic("unmarshal conf change: " + err.Error())
	}

	w.logger.WithField("cc", cc).Warnln("Applying conf change")
	w.n.ApplyConfChange(cc)
	switch cc.Type {
	case raftpb.ConfChangeAddNode:
		if len(cc.Context) > 0 {
			w.transport.AddPeer(types.ID(cc.NodeID), []string{string(cc.Context)})
		}
	case raftpb.ConfChangeRemoveNode:
		tid := types.ID(cc.NodeID)
		if tid == w.transport.ID {
			w.logger.Warnln("Shutting down")
			return
		}
		w.transport.RemovePeer(tid)
	}

	// Inform state machine about new configuration.
	w.sm.Apply(&commonpb.Entry{
		Term:      entry.Term,
		Index:     entry.Index,
		EntryType: commonpb.EntryReconf,
		Data:      data,
	})

	// TODO Respond / Cleanup future.
}

func (w *Wrapper) Process(ctx context.Context, m raftpb.Message) error {
	return w.n.Step(ctx, m)
}
func (w *Wrapper) IsIDRemoved(id uint64) bool                               { return false }
func (w *Wrapper) ReportUnreachable(id uint64)                              {}
func (w *Wrapper) ReportSnapshot(id uint64, status etcdraft.SnapshotStatus) {}
<|endoftext|>


<commit_before>package reception

import (
	etcd "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

type (
	keysAPI interface {
		Get(path string) (*etcd.Response, error)
		Set(path, value string, opts *etcd.SetOptions) error
		Watcher(path string) etcd.Watcher
	}

	etcdShim struct {
		api etcd.KeysAPI
	}
)

func (s *etcdShim) Get(path string) (*etcd.Response, error) {
	return s.api.Get(context.Background(), path, &etcd.GetOptions{
		Recursive: true,
	})
}

func (s *etcdShim) Set(path, value string, opts *etcd.SetOptions) error {
	_, err := s.api.Set(context.Background(), path, value, opts)
	return err
}

func (s *etcdShim) Watcher(path string) etcd.Watcher {
	return s.api.Watcher(path, &etcd.WatcherOptions{Recursive: true})
}
<commit_msg>Update one remaining context import.<commit_after>package reception

import (
	"context"

	etcd "github.com/coreos/etcd/client"
)

type (
	keysAPI interface {
		Get(path string) (*etcd.Response, error)
		Set(path, value string, opts *etcd.SetOptions) error
		Watcher(path string) etcd.Watcher
	}

	etcdShim struct {
		api etcd.KeysAPI
	}
)

func (s *etcdShim) Get(path string) (*etcd.Response, error) {
	return s.api.Get(context.Background(), path, &etcd.GetOptions{
		Recursive: true,
	})
}

func (s *etcdShim) Set(path, value string, opts *etcd.SetOptions) error {
	_, err := s.api.Set(context.Background(), path, value, opts)
	return err
}

func (s *etcdShim) Watcher(path string) etcd.Watcher {
	return s.api.Watcher(path, &etcd.WatcherOptions{Recursive: true})
}
<|endoftext|>


<commit_before>package server

import (
	"net"
	"sync"

	"google.golang.org/grpc"

	"github.com/barakmich/agro"
	"github.com/barakmich/agro/models"
)

const (
	defaultInodeReplication = 3
)

type distributor struct {
	mut     sync.RWMutex
	blocks  agro.BlockStore
	inodes  agro.INodeStore
	srv     *server
	client  *distClient
	lis     net.Listener
	grpcSrv *grpc.Server

	ring           agro.Ring
	closed         bool
	rebalancerChan chan struct{}
}

func newDistributor(srv *server, addr string, listen bool) (*distributor, error) {
	var err error
	d := &distributor{
		blocks: srv.blocks,
		inodes: srv.inodes,
		srv:    srv,
	}
	if listen {
		d.lis, err = net.Listen("tcp", addr)
		if err != nil {
			return nil, err
		}
		d.grpcSrv = grpc.NewServer()
		models.RegisterAgroStorageServer(d.grpcSrv, d)
		go d.grpcSrv.Serve(d.lis)
		srv.BeginHeartbeat()
	}
	// Set up the rebalancer
	d.ring, err = d.srv.mds.GetRing()
	if err != nil {
		return nil, err
	}
	d.rebalancerChan = make(chan struct{})
	go d.rebalancer(d.rebalancerChan)
	d.client = newDistClient(d)
	return d, nil
}

func (d *distributor) UUID() string {
	return d.srv.mds.UUID()
}

func (d *distributor) inodeReplication() int {
	r, err := d.srv.mds.GlobalMetadata()
	if err != nil {
		clog.Error(err)
		return defaultInodeReplication
	}
	return r.INodeReplication
}

func (d *distributor) Close() error {
	if d.closed {
		return nil
	}
	close(d.rebalancerChan)
	if d.lis != nil {
		d.grpcSrv.Stop()
		err := d.lis.Close()
		if err != nil {
			return err
		}
	}
	err := d.inodes.Close()
	if err != nil {
		return err
	}
	err = d.blocks.Close()
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>fix typo<commit_after>package server

import (
	"net"
	"sync"

	"google.golang.org/grpc"

	"github.com/barakmich/agro"
	"github.com/barakmich/agro/models"
)

const (
	defaultINodeReplication = 3
)

type distributor struct {
	mut     sync.RWMutex
	blocks  agro.BlockStore
	inodes  agro.INodeStore
	srv     *server
	client  *distClient
	lis     net.Listener
	grpcSrv *grpc.Server

	ring           agro.Ring
	closed         bool
	rebalancerChan chan struct{}
}

func newDistributor(srv *server, addr string, listen bool) (*distributor, error) {
	var err error
	d := &distributor{
		blocks: srv.blocks,
		inodes: srv.inodes,
		srv:    srv,
	}
	if listen {
		d.lis, err = net.Listen("tcp", addr)
		if err != nil {
			return nil, err
		}
		d.grpcSrv = grpc.NewServer()
		models.RegisterAgroStorageServer(d.grpcSrv, d)
		go d.grpcSrv.Serve(d.lis)
		srv.BeginHeartbeat()
	}
	// Set up the rebalancer
	d.ring, err = d.srv.mds.GetRing()
	if err != nil {
		return nil, err
	}
	d.rebalancerChan = make(chan struct{})
	go d.rebalancer(d.rebalancerChan)
	d.client = newDistClient(d)
	return d, nil
}

func (d *distributor) UUID() string {
	return d.srv.mds.UUID()
}

func (d *distributor) inodeReplication() int {
	r, err := d.srv.mds.GlobalMetadata()
	if err != nil {
		clog.Error(err)
		return defaultINodeReplication
	}
	return r.INodeReplication
}

func (d *distributor) Close() error {
	if d.closed {
		return nil
	}
	close(d.rebalancerChan)
	if d.lis != nil {
		d.grpcSrv.Stop()
		err := d.lis.Close()
		if err != nil {
			return err
		}
	}
	err := d.inodes.Close()
	if err != nil {
		return err
	}
	err = d.blocks.Close()
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package uploads

import (
	"time"

	"math"

	"fmt"

	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/domain"
	"github.com/willf/bitset"
)

var _ = fmt.Println

// A IDRequest requests a new upload id be created for
// the given parameters.
type IDRequest struct {
	User        string
	DirectoryID string
	ProjectID   string
	FileName    string
	FileSize    int64
	Checksum    string
	ChunkSize   int32
	FileMTime   time.Time
	Host        string
	Birthtime   time.Time
}

// IDService creates new upload requests
type IDService interface {
	ID(req IDRequest) (*schema.Upload, error)
	Delete(requestID, user string) error
	UploadsForProject(projectID, user string) ([]schema.Upload, error)
}

// idService implements the IDService interface using
// the dai services.
type idService struct {
	dirs        dai.Dirs
	projects    dai.Projects
	uploads     dai.Uploads
	files       dai.Files
	access      domain.Access
	fops        file.Operations
	tracker     tracker
	requestPath requestPath
}

// NewIDService creates a new idService. It uses db.RSessionMust() to get
// a session to connect to the database. It will panic if it cannot connect to
// the database.
func NewIDService() *idService {
	session := db.RSessionMust()
	access := domain.NewAccess(dai.NewRProjects(session), dai.NewRFiles(session), dai.NewRUsers(session))
	return &idService{
		dirs:        dai.NewRDirs(session),
		projects:    dai.NewRProjects(session),
		uploads:     dai.NewRUploads(session),
		files:       dai.NewRFiles(session),
		access:      access,
		fops:        file.OS,
		tracker:     requestBlockTracker,
		requestPath: &mcdirRequestPath{},
	}
}

// NewIDServiceUsingSession creates a new idService that connects to the database using
// the given session.
func NewIDServiceUsingSession(session *r.Session) *idService {
	access := domain.NewAccess(dai.NewRProjects(session), dai.NewRFiles(session), dai.NewRUsers(session))
	return &idService{
		dirs:        dai.NewRDirs(session),
		projects:    dai.NewRProjects(session),
		uploads:     dai.NewRUploads(session),
		files:       dai.NewRFiles(session),
		access:      access,
		fops:        file.OS,
		tracker:     requestBlockTracker,
		requestPath: &mcdirRequestPath{},
	}
}

// ID will create a new Upload request or return an existing one.
func (s *idService) ID(req IDRequest) (*schema.Upload, error) {
	///files, err := s.files.AllByChecksum(req.Checksum)
	proj, err := s.getProj(req.ProjectID, req.User)
	if err != nil {
		fmt.Println("getProj err", err)
		return nil, err
	}

	dir, err := s.getDir(req.DirectoryID, proj.ID, req.User)
	if err != nil {
		fmt.Println("getDir err", err)
		return nil, err
	}

	searchParams := dai.UploadSearch{
		ProjectID:   proj.ID,
		DirectoryID: dir.ID,
		FileName:    req.FileName,
		Checksum:    req.Checksum,
	}

	if existingUpload, err := s.uploads.Search(searchParams); err == nil {
		// Found existing
		if bset := s.tracker.getBlocks(existingUpload.ID); bset != nil {
			existingUpload.File.Blocks = bset
		}
		return existingUpload, nil
	}

	n := uint(numBlocks(req.FileSize, req.ChunkSize))
	upload := schema.CUpload().
		Owner(req.User).
		Project(req.ProjectID, proj.Name).
		ProjectOwner(proj.Owner).
		Directory(req.DirectoryID, dir.Name).
		Host(req.Host).
		FName(req.FileName).
		FSize(req.FileSize).
		FChecksum(req.Checksum).
		FRemoteMTime(req.FileMTime).
		FBlocks(bitset.New(n)).
		Create()
	u, err := s.uploads.Insert(&upload)
	if err != nil {
		fmt.Println("insert err", err)
		return nil, err
	}

	if err := s.initUpload(u.ID, req.FileSize, req.ChunkSize); err != nil {
		fmt.Println("initUpload err", err)
		s.uploads.Delete(u.ID)
		return nil, err
	}
	return u, nil
}

// getProj retrieves the project with the given projectID. It checks that the
// given user has access to that project.
func (s *idService) getProj(projectID, user string) (*schema.Project, error) {
	project, err := s.projects.ByID(projectID)
	switch {
	case err != nil:
		return nil, err
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return project, nil
	}
}

// getDir retrieves the directory with the given directoryID. It checks access to the
// directory and validates that the directory exists in the given project.
func (s *idService) getDir(directoryID, projectID, user string) (*schema.Directory, error) {
	dir, err := s.dirs.ByID(directoryID)
	switch {
	case err != nil:
		return nil, err
	case !s.projects.HasDirectory(projectID, directoryID):
		return nil, app.ErrInvalid
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return dir, nil
	}
}

// initUpload
func (s *idService) initUpload(id string, fileSize int64, chunkSize int32) error {
	if err := s.requestPath.mkdirFromID(id); err != nil {
		return err
	}

	s.tracker.load(id, numBlocks(fileSize, chunkSize))
	return nil
}

// numBlocks
func numBlocks(fileSize int64, chunkSize int32) int {
	// round up to nearest number of blocks
	d := float64(fileSize) / float64(chunkSize)
	n := int(math.Ceil(d))
	return n
}

// Delete will delete the given requestID if the user has access
// to delete that request. Owners of requests can delete their
// own requests. Project owners can delete any request, even if
// they don't own the request.
func (s *idService) Delete(requestID, user string) error {
	upload, err := s.uploads.ByID(requestID)
	switch {
	case err != nil:
		return err
	case !s.access.AllowedByOwner(upload.ProjectID, user):
		return app.ErrNoAccess
	default:
		if err := s.uploads.Delete(requestID); err != nil {
			return err
		}
		// Delete the directory where chunks were being written
		s.fops.RemoveAll(app.MCDir.UploadDir(requestID))
		return nil
	}
}

// ListForProject will return all the uploads associated with a project.
func (s *idService) UploadsForProject(projectID, user string) ([]schema.Upload, error) {
	// getProj will validate the project and access.
	_, err := s.getProj(projectID, user)
	switch {
	case err == app.ErrNotFound:
		// Invalid project
		return nil, app.ErrInvalid
	case err != nil:
		return nil, err
	default:
		return s.uploads.ForProject(projectID)
	}
}
<commit_msg>Refactor ID service to include checks for existing file uploads.<commit_after>package uploads

import (
	"time"

	"math"

	"fmt"

	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/domain"
	"github.com/willf/bitset"
)

var _ = fmt.Println

// A IDRequest requests a new upload id be created for
// the given parameters.
type IDRequest struct {
	User        string
	DirectoryID string
	ProjectID   string
	FileName    string
	FileSize    int64
	Checksum    string
	ChunkSize   int32
	FileMTime   time.Time
	Host        string
	Birthtime   time.Time
}

// IDService creates new upload requests
type IDService interface {
	ID(req IDRequest) (*schema.Upload, error)
	Delete(requestID, user string) error
	UploadsForProject(projectID, user string) ([]schema.Upload, error)
}

// idService implements the IDService interface using
// the dai services.
type idService struct {
	dirs        dai.Dirs
	projects    dai.Projects
	uploads     dai.Uploads
	files       dai.Files
	access      domain.Access
	fops        file.Operations
	tracker     tracker
	requestPath requestPath
}

// NewIDService creates a new idService. It uses db.RSessionMust() to get
// a session to connect to the database. It will panic if it cannot connect to
// the database.
func NewIDService() *idService {
	session := db.RSessionMust()
	access := domain.NewAccess(dai.NewRProjects(session), dai.NewRFiles(session), dai.NewRUsers(session))
	return &idService{
		dirs:        dai.NewRDirs(session),
		projects:    dai.NewRProjects(session),
		uploads:     dai.NewRUploads(session),
		files:       dai.NewRFiles(session),
		access:      access,
		fops:        file.OS,
		tracker:     requestBlockTracker,
		requestPath: &mcdirRequestPath{},
	}
}

// NewIDServiceUsingSession creates a new idService that connects to the database using
// the given session.
func NewIDServiceUsingSession(session *r.Session) *idService {
	access := domain.NewAccess(dai.NewRProjects(session), dai.NewRFiles(session), dai.NewRUsers(session))
	return &idService{
		dirs:        dai.NewRDirs(session),
		projects:    dai.NewRProjects(session),
		uploads:     dai.NewRUploads(session),
		files:       dai.NewRFiles(session),
		access:      access,
		fops:        file.OS,
		tracker:     requestBlockTracker,
		requestPath: &mcdirRequestPath{},
	}
}

func (s *idService) ID2(req IDRequest) (*schema.Upload, error) {
	var (
		err  error
		proj *schema.Project
		dir  *schema.Directory
	)

	if proj, err = s.getProjectValidatingAccess(req.ProjectID, req.User); err != nil {
		return nil, err
	}

	if dir, err = s.getDirectoryValidatingAccess(req.DirectoryID, req.ProjectID, req.User); err != nil {
		return nil, err
	}

	upload, err := s.findExisting(req, proj, dir)
	switch {
	case err == app.ErrNotFound:
		return s.createNewUploadRequest(req, proj, dir)
	case err != nil:
		return nil, err
	default:
		return upload, nil
	}
}

func (s *idService) findExisting(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	if file, err := s.files.ByChecksum(req); err == nil {
		return s.createUploadFromFile(req, file)
	}
	return s.findMatchingUploadRequest(req)
}

func (s *idService) createUploadFromFile(req IDRequest, file *schema.File) (*schema.Upload, error) {
	// We have the real uploaded file. There are 4 cases:
	//   1. This file matches the project, directory and name we want to upload, so
	//      we create a finished upload request because there is nothing to do.
	//
	//   2. This file matches the project and directory, but not name, so we create
	//      a new file entry pointing at this one and return a finished upload request.
	//
	//   3. This entry is in a different project and/or directory and there isn't a matching
	//      one in the given project/directory, so create a new file entry pointing at this
	//      one and return a finished upload request.
	//
	//   4. This entry is in a different project and/or directory and there is a matching
	//      one in the given project/directory. Return a finished upload request.
	return nil, nil
}

func (s *idService) findMatchingUploadRequest(req IDRequest) (*schema.Upload, error) {
	searchParams := dai.UploadSearch{
		ProjectID:   req.ProjectID,
		DirectoryID: req.DirectoryID,
		FileName:    req.FileName,
		Checksum:    req.Checksum,
	}

	if existingUpload, err := s.uploads.Search(searchParams); err == nil {
		// Found existing
		if bset := s.tracker.getBlocks(existingUpload.ID); bset != nil {
			existingUpload.File.Blocks = bset
		}
		return existingUpload, nil
	}
	return nil, app.ErrNotFound
}

func (s *idService) createNewUploadRequest(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	n := uint(numBlocks(req.FileSize, req.ChunkSize))
	upload := schema.CUpload().
		Owner(req.User).
		Project(req.ProjectID, proj.Name).
		ProjectOwner(proj.Owner).
		Directory(req.DirectoryID, dir.Name).
		Host(req.Host).
		FName(req.FileName).
		FSize(req.FileSize).
		FChecksum(req.Checksum).
		FRemoteMTime(req.FileMTime).
		FBlocks(bitset.New(n)).
		Create()
	u, err := s.uploads.Insert(&upload)
	if err != nil {
		fmt.Println("insert err", err)
		return nil, err
	}

	if err := s.initUpload(u.ID, req.FileSize, req.ChunkSize); err != nil {
		fmt.Println("initUpload err", err)
		s.uploads.Delete(u.ID)
		return nil, err
	}
	return u, nil
}

// ID will create a new Upload request or return an existing one.
func (s *idService) ID(req IDRequest) (*schema.Upload, error) {
	///files, err := s.files.AllByChecksum(req.Checksum)
	proj, err := s.getProjectValidatingAccess(req.ProjectID, req.User)
	if err != nil {
		fmt.Println("getProj err", err)
		return nil, err
	}

	dir, err := s.getDirectoryValidatingAccess(req.DirectoryID, proj.ID, req.User)
	if err != nil {
		fmt.Println("getDir err", err)
		return nil, err
	}

	searchParams := dai.UploadSearch{
		ProjectID:   proj.ID,
		DirectoryID: dir.ID,
		FileName:    req.FileName,
		Checksum:    req.Checksum,
	}

	if existingUpload, err := s.uploads.Search(searchParams); err == nil {
		// Found existing
		if bset := s.tracker.getBlocks(existingUpload.ID); bset != nil {
			existingUpload.File.Blocks = bset
		}
		return existingUpload, nil
	}

	n := uint(numBlocks(req.FileSize, req.ChunkSize))
	upload := schema.CUpload().
		Owner(req.User).
		Project(req.ProjectID, proj.Name).
		ProjectOwner(proj.Owner).
		Directory(req.DirectoryID, dir.Name).
		Host(req.Host).
		FName(req.FileName).
		FSize(req.FileSize).
		FChecksum(req.Checksum).
		FRemoteMTime(req.FileMTime).
		FBlocks(bitset.New(n)).
		Create()
	u, err := s.uploads.Insert(&upload)
	if err != nil {
		fmt.Println("insert err", err)
		return nil, err
	}

	if err := s.initUpload(u.ID, req.FileSize, req.ChunkSize); err != nil {
		fmt.Println("initUpload err", err)
		s.uploads.Delete(u.ID)
		return nil, err
	}
	return u, nil
}

func (s *idService) getAndValidateDependencies(req IDRequest) (*schema.Project, *schema.Directory, error) {
	proj, err := s.getProjectValidatingAccess(req.ProjectID, req.User)
	if err != nil {
		fmt.Println("getProj err", err)
		return nil, nil, err
	}

	dir, err := s.getDirectoryValidatingAccess(req.DirectoryID, proj.ID, req.User)
	if err != nil {
		fmt.Println("getDir err", err)
		return nil, nil, err
	}

	return proj, dir, nil
}

// getProj retrieves the project with the given projectID. It checks that the
// given user has access to that project.
func (s *idService) getProjectValidatingAccess(projectID, user string) (*schema.Project, error) {
	project, err := s.projects.ByID(projectID)
	switch {
	case err != nil:
		return nil, err
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return project, nil
	}
}

// getDir retrieves the directory with the given directoryID. It checks access to the
// directory and validates that the directory exists in the given project.
func (s *idService) getDirectoryValidatingAccess(directoryID, projectID, user string) (*schema.Directory, error) {
	dir, err := s.dirs.ByID(directoryID)
	switch {
	case err != nil:
		return nil, err
	case !s.projects.HasDirectory(projectID, directoryID):
		return nil, app.ErrInvalid
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return dir, nil
	}
}

// initUpload
func (s *idService) initUpload(id string, fileSize int64, chunkSize int32) error {
	if err := s.requestPath.mkdirFromID(id); err != nil {
		return err
	}

	s.tracker.load(id, numBlocks(fileSize, chunkSize))
	return nil
}

// numBlocks
func numBlocks(fileSize int64, chunkSize int32) int {
	// round up to nearest number of blocks
	d := float64(fileSize) / float64(chunkSize)
	n := int(math.Ceil(d))
	return n
}

// Delete will delete the given requestID if the user has access
// to delete that request. Owners of requests can delete their
// own requests. Project owners can delete any request, even if
// they don't own the request.
func (s *idService) Delete(requestID, user string) error {
	upload, err := s.uploads.ByID(requestID)
	switch {
	case err != nil:
		return err
	case !s.access.AllowedByOwner(upload.ProjectID, user):
		return app.ErrNoAccess
	default:
		if err := s.uploads.Delete(requestID); err != nil {
			return err
		}
		// Delete the directory where chunks were being written
		s.fops.RemoveAll(app.MCDir.UploadDir(requestID))
		return nil
	}
}

// ListForProject will return all the uploads associated with a project.
func (s *idService) UploadsForProject(projectID, user string) ([]schema.Upload, error) {
	// getProj will validate the project and access.
	_, err := s.getProjectValidatingAccess(projectID, user)
	switch {
	case err == app.ErrNotFound:
		// Invalid project
		return nil, app.ErrInvalid
	case err != nil:
		return nil, err
	default:
		return s.uploads.ForProject(projectID)
	}
}
<|endoftext|>


<commit_before>package couchbase

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"hash/crc32"
	"math/rand"
	"strconv"
	"strings"
	"time"

	cb "github.com/couchbaselabs/go-couchbase"
	"github.com/couchbaselabs/tuqtng/ast"
	"github.com/couchbaselabs/tuqtng/catalog"
)

type ddocJSON struct {
	cb.DDocJSON
	IndexOn       []string `json:"indexOn"`
	IndexChecksum int      `json:"indexChecksum"`
}

func newViewIndex(name string, on catalog.IndexKey, bkt *bucket) (*viewIndex, error) {

	doc, err := newDesignDoc(name, on)
	if err != nil {
		return nil, err
	}

	inst := viewIndex{
		name:   name,
		using:  catalog.VIEW,
		on:     on,
		ddoc:   doc,
		bucket: bkt,
	}

	err = inst.putDesignDoc()
	if err != nil {
		return nil, err
	}

	err = inst.WaitForIndex()
	if err != nil {
		return nil, err
	}

	return &inst, nil
}

func (vi *viewIndex) String() string {
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("name: %v ", vi.name))
	buf.WriteString(fmt.Sprintf("on: %v ", vi.on))
	buf.WriteString(fmt.Sprintf("using: %v ", vi.using))
	buf.WriteString(fmt.Sprintf("ddoc: %v ", *vi.ddoc))
	buf.WriteString(fmt.Sprintf("bucket: %v ", *vi.bucket))
	return buf.String()
}

func newDesignDoc(idxname string, on catalog.IndexKey) (*designdoc, error) {
	var doc designdoc

	doc.name = "ddl_" + idxname
	doc.viewname = idxname

	err := generateMap(on, &doc)
	if err != nil {
		return nil, err
	}

	err = generateReduce(on, &doc)
	if err != nil {
		return nil, err
	}

	return &doc, nil
}

func loadViewIndexes(b *bucket) ([]*catalog.Index, error) {

	rows, err := b.cbbucket.GetDDocs()
	if err != nil {
		return nil, err
	}

	inames := make([]string, 0, len(rows.Rows))
	for _, row := range rows.Rows {
		cdoc := row.DDoc
		id := cdoc.Meta["id"].(string)
		if !strings.HasPrefix(id, "_design/ddl_") {
			continue
		}
		iname := strings.TrimPrefix(id, "_design/ddl_")
		inames = append(inames, iname)
	}

	indexes := make([]*catalog.Index, 0, len(inames))
	for _, iname := range inames {
		ddname := "ddl_" + iname
		jdoc, err := getDesignDoc(b, ddname)
		if err != nil {
			return nil, err
		}
		jview, ok := jdoc.Views[iname]
		if !ok {
			return nil, errors.New("Missing view for index " + iname)
		}

		exprlist := make([]ast.Expression, 0, len(jdoc.IndexOn))

		for _, ser := range jdoc.IndexOn {
			// HACK - remove this when Unmarshall supports META()
			if iname == PRIMARY_INDEX {
				meta := ast.NewFunctionCall("meta", ast.FunctionArgExpressionList{})
				mdid := ast.NewDotMemberOperator(meta, ast.NewProperty("id"))
				exprlist = append(exprlist, mdid)
			} else {
				expr, err := ast.UnmarshalExpression([]byte(ser))
				if err != nil {
					return nil, errors.New("Cannot unmarshal expression for index " + iname)
				}
				exprlist = append(exprlist, expr)
			}
		}
		if len(exprlist) != len(jdoc.IndexOn) {
			continue
		}

		ddoc := designdoc{
			name:     ddname,
			viewname: iname,
			mapfn:    jview.Map,
			reducefn: jview.Reduce,
		}
		if ddoc.checksum() != jdoc.IndexChecksum {
			return nil, errors.New("Warning - checksum failed on index " + iname)
		}

		var index catalog.Index

		if iname == PRIMARY_INDEX {
			index = &primaryIndex{
				viewIndex{
					name:   iname,
					bucket: b,
					using:  catalog.VIEW,
					ddoc:   &ddoc,
					on:     exprlist,
				},
			}
			indexes = append(indexes, &index)
		} else {
			index = &viewIndex{
				name:   iname,
				bucket: b,
				using:  catalog.VIEW,
				ddoc:   &ddoc,
				on:     exprlist,
			}
			indexes = append(indexes, &index)
		}
	}

	return indexes, nil
}

func newPrimaryIndex(b *bucket) (*primaryIndex, error) {
	ddoc := newPrimaryDDoc()
	meta := ast.NewFunctionCall("meta", ast.FunctionArgExpressionList{})
	mdid := ast.NewDotMemberOperator(meta, ast.NewProperty("id"))
	inst := primaryIndex{
		viewIndex{
			name:   PRIMARY_INDEX,
			using:  catalog.VIEW,
			on:     catalog.IndexKey{mdid},
			ddoc:   ddoc,
			bucket: b,
		},
	}

	err := inst.putDesignDoc()
	if err != nil {
		return nil, err
	}

	err = inst.WaitForIndex()
	if err != nil {
		return nil, err
	}

	return &inst, nil
}

func newAllDocsIndex(b *bucket) *primaryIndex {
	meta := ast.NewFunctionCall("meta", ast.FunctionArgExpressionList{})
	mdid := ast.NewDotMemberOperator(meta, ast.NewProperty("id"))
	ddoc := designdoc{name: "", viewname: "_all_docs"}
	idx := primaryIndex{
		viewIndex{
			name:   ALLDOCS_INDEX,
			using:  catalog.VIEW,
			on:     catalog.IndexKey{mdid},
			ddoc:   &ddoc,
			bucket: b,
		},
	}
	return &idx
}

func newPrimaryDDoc() *designdoc {
	var doc designdoc
	line := strings.Replace(templPrimary, "$rnd", strconv.Itoa(int(rand.Int31())), -1)
	line = strings.Replace(line, "$string", strconv.Itoa(TYPE_STRING), -1)
	doc.mapfn = line
	doc.reducefn = ""
	doc.name = "ddl_" + PRIMARY_INDEX
	doc.viewname = PRIMARY_INDEX
	return &doc
}

func generateMap(on catalog.IndexKey, doc *designdoc) error {

	buf := new(bytes.Buffer)

	fmt.Fprintln(buf, templStart)
	line := strings.Replace(templFunctions, "$null", strconv.Itoa(TYPE_NULL), -1)
	line = strings.Replace(line, "$boolean", strconv.Itoa(TYPE_BOOLEAN), -1)
	line = strings.Replace(line, "$number", strconv.Itoa(TYPE_NUMBER), -1)
	line = strings.Replace(line, "$string", strconv.Itoa(TYPE_STRING), -1)
	line = strings.Replace(line, "$array", strconv.Itoa(TYPE_ARRAY), -1)
	line = strings.Replace(line, "$object", strconv.Itoa(TYPE_OBJECT), -1)
	fmt.Fprintln(buf, line)

	keylist := new(bytes.Buffer)
	for idx, expr := range on {

		walker := NewWalker()
		_, err := walker.Visit(expr)
		if err != nil {
			return err
		}

		jvar := fmt.Sprintf("key%v", idx+1)
		line := strings.Replace(templExpr, "$var", jvar, -1)
		line = strings.Replace(line, "$path", walker.JS(), -1)
		fmt.Fprint(buf, line)

		if idx > 0 {
			fmt.Fprint(keylist, ", ")
		}
		fmt.Fprint(keylist, jvar)
	}

	line = strings.Replace(templKey, "$keylist", keylist.String(), -1)

	fmt.Fprint(buf, line)
	fmt.Fprint(buf, templEmit)

	line = strings.Replace(templEnd, "$rnd", strconv.Itoa(int(rand.Int31())), -1)
	fmt.Fprint(buf, line)

	doc.mapfn = buf.String()
	return nil
}

func generateReduce(on catalog.IndexKey, doc *designdoc) error {
	doc.reducefn = ""
	return nil
}

func (idx *viewIndex) putDesignDoc() error {
	var view cb.ViewDefinition
	view.Map = idx.ddoc.mapfn

	var put ddocJSON
	put.Views = make(map[string]cb.ViewDefinition)
	put.Views[idx.name] = view
	put.IndexChecksum = idx.ddoc.checksum()

	put.IndexOn = make([]string, len(idx.on))
	for idx, expr := range idx.on {
		ser, err := json.Marshal(expr)
		if err != nil {
			return err
		}
		put.IndexOn[idx] = string(ser)
	}

	if err := idx.bucket.cbbucket.PutDDoc(idx.DDocName(), &put); err != nil {
		return err
	}

	var saved *ddocJSON = nil
	var err error = nil

	// give the PUT some time to register
	for i := 0; i < 3; i++ {
		if i > 1 {
			time.Sleep(time.Duration(i*3) * time.Second)
		}

		saved, err = getDesignDoc(idx.bucket, idx.DDocName())
		if err == nil {
			break
		}
	}

	if err != nil {
		return errors.New("Creating index '" + idx.name + "' failed: " + err.Error())
	}

	if saved.IndexChecksum != idx.ddoc.checksum() {
		return errors.New("Checksum mismatch after creating index '" + idx.name + "'")
	}

	return nil
}

func (ddoc *designdoc) checksum() int {
	mapSum := crc32.ChecksumIEEE([]byte(ddoc.mapfn))
	reduceSum := crc32.ChecksumIEEE([]byte(ddoc.reducefn))
	return int(mapSum + reduceSum)
}

func getDesignDoc(b *bucket, ddocname string) (*ddocJSON, error) {
	var ddoc ddocJSON
	err := b.cbbucket.GetDDoc(ddocname, &ddoc)
	if err != nil {
		return nil, err
	}
	return &ddoc, nil
}

func (idx *viewIndex) DropViewIndex() error {
	if err := idx.bucket.cbbucket.DeleteDDoc(idx.ddoc.name); err != nil {
		return err
	}
	return nil
}

func (idx *viewIndex) WaitForIndex() error {
	var err error
	// if we have got this far, very likely any errors are
	// due to index not yet being noticed by the system.
	for i := 0; i < 3; i++ {
		if i > 1 {
			time.Sleep(time.Duration(i*3) * time.Second)
		}
		_, err = idx.bucket.cbbucket.View(
			idx.ddoc.name,
			idx.ddoc.viewname,
			map[string]interface{}{
				"start_key": []interface{}{"thing"},
				"end_key":   []interface{}{"thing", map[string]string{}},
				"stale":     false,
			})
		if err == nil {
			break
		}
	}
	return err
}

// AST to JS conversion
type JsStatement struct {
	js bytes.Buffer
}

func NewWalker() *JsStatement {
	var js JsStatement
	return &js
}

func (this *JsStatement) JS() string {
	return this.js.String()
}

// inorder traversal of the AST to get JS expression out of it
func (this *JsStatement) Visit(e ast.Expression) (ast.Expression, error) {
	switch expr := e.(type) {

	case *ast.DotMemberOperator:
		if this.js.Len() == 0 {
			this.js.WriteString("doc.")
		}
		_, err := expr.Left.Accept(this)
		if err != nil {
			return nil, err
		}
		this.js.WriteString(".")
		_, err = expr.Right.Accept(this)
		if err != nil {
			return nil, err
		}

	case *ast.BracketMemberOperator:
		if this.js.Len() == 0 {
			this.js.WriteString("doc.")
		}
		_, err := expr.Left.Accept(this)
		if err != nil {
			return nil, err
		}
		this.js.WriteString("[")
		_, err = expr.Right.Accept(this)
		if err != nil {
			return nil, err
		}
		this.js.WriteString("]")

	case *ast.Property:
		if this.js.Len() == 0 {
			this.js.WriteString("doc.")
		}
		this.js.WriteString(expr.Path)

	case *ast.LiteralNumber:
		this.js.WriteString(fmt.Sprintf("%v", expr.Val))

	case *ast.LiteralString:
		this.js.WriteString(expr.Val)

	default:
		return e, errors.New("Expression is not supported by indexing currently: " + e.String())

	}
	return e, nil
}
<commit_msg>Rename DDocJSON to DDoc to mirror go-couchbase change 88b6e00<commit_after>package couchbase

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"hash/crc32"
	"math/rand"
	"strconv"
	"strings"
	"time"

	cb "github.com/couchbaselabs/go-couchbase"
	"github.com/couchbaselabs/tuqtng/ast"
	"github.com/couchbaselabs/tuqtng/catalog"
)

type ddocJSON struct {
	cb.DDoc
	IndexOn       []string `json:"indexOn"`
	IndexChecksum int      `json:"indexChecksum"`
}

func newViewIndex(name string, on catalog.IndexKey, bkt *bucket) (*viewIndex, error) {

	doc, err := newDesignDoc(name, on)
	if err != nil {
		return nil, err
	}

	inst := viewIndex{
		name:   name,
		using:  catalog.VIEW,
		on:     on,
		ddoc:   doc,
		bucket: bkt,
	}

	err = inst.putDesignDoc()
	if err != nil {
		return nil, err
	}

	err = inst.WaitForIndex()
	if err != nil {
		return nil, err
	}

	return &inst, nil
}

func (vi *viewIndex) String() string {
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("name: %v ", vi.name))
	buf.WriteString(fmt.Sprintf("on: %v ", vi.on))
	buf.WriteString(fmt.Sprintf("using: %v ", vi.using))
	buf.WriteString(fmt.Sprintf("ddoc: %v ", *vi.ddoc))
	buf.WriteString(fmt.Sprintf("bucket: %v ", *vi.bucket))
	return buf.String()
}

func newDesignDoc(idxname string, on catalog.IndexKey) (*designdoc, error) {
	var doc designdoc

	doc.name = "ddl_" + idxname
	doc.viewname = idxname

	err := generateMap(on, &doc)
	if err != nil {
		return nil, err
	}

	err = generateReduce(on, &doc)
	if err != nil {
		return nil, err
	}

	return &doc, nil
}

func loadViewIndexes(b *bucket) ([]*catalog.Index, error) {

	rows, err := b.cbbucket.GetDDocs()
	if err != nil {
		return nil, err
	}

	inames := make([]string, 0, len(rows.Rows))
	for _, row := range rows.Rows {
		cdoc := row.DDoc
		id := cdoc.Meta["id"].(string)
		if !strings.HasPrefix(id, "_design/ddl_") {
			continue
		}
		iname := strings.TrimPrefix(id, "_design/ddl_")
		inames = append(inames, iname)
	}

	indexes := make([]*catalog.Index, 0, len(inames))
	for _, iname := range inames {
		ddname := "ddl_" + iname
		jdoc, err := getDesignDoc(b, ddname)
		if err != nil {
			return nil, err
		}
		jview, ok := jdoc.Views[iname]
		if !ok {
			return nil, errors.New("Missing view for index " + iname)
		}

		exprlist := make([]ast.Expression, 0, len(jdoc.IndexOn))

		for _, ser := range jdoc.IndexOn {
			// HACK - remove this when Unmarshall supports META()
			if iname == PRIMARY_INDEX {
				meta := ast.NewFunctionCall("meta", ast.FunctionArgExpressionList{})
				mdid := ast.NewDotMemberOperator(meta, ast.NewProperty("id"))
				exprlist = append(exprlist, mdid)
			} else {
				expr, err := ast.UnmarshalExpression([]byte(ser))
				if err != nil {
					return nil, errors.New("Cannot unmarshal expression for index " + iname)
				}
				exprlist = append(exprlist, expr)
			}
		}
		if len(exprlist) != len(jdoc.IndexOn) {
			continue
		}

		ddoc := designdoc{
			name:     ddname,
			viewname: iname,
			mapfn:    jview.Map,
			reducefn: jview.Reduce,
		}
		if ddoc.checksum() != jdoc.IndexChecksum {
			return nil, errors.New("Warning - checksum failed on index " + iname)
		}

		var index catalog.Index

		if iname == PRIMARY_INDEX {
			index = &primaryIndex{
				viewIndex{
					name:   iname,
					bucket: b,
					using:  catalog.VIEW,
					ddoc:   &ddoc,
					on:     exprlist,
				},
			}
			indexes = append(indexes, &index)
		} else {
			index = &viewIndex{
				name:   iname,
				bucket: b,
				using:  catalog.VIEW,
				ddoc:   &ddoc,
				on:     exprlist,
			}
			indexes = append(indexes, &index)
		}
	}

	return indexes, nil
}

func newPrimaryIndex(b *bucket) (*primaryIndex, error) {
	ddoc := newPrimaryDDoc()
	meta := ast.NewFunctionCall("meta", ast.FunctionArgExpressionList{})
	mdid := ast.NewDotMemberOperator(meta, ast.NewProperty("id"))
	inst := primaryIndex{
		viewIndex{
			name:   PRIMARY_INDEX,
			using:  catalog.VIEW,
			on:     catalog.IndexKey{mdid},
			ddoc:   ddoc,
			bucket: b,
		},
	}

	err := inst.putDesignDoc()
	if err != nil {
		return nil, err
	}

	err = inst.WaitForIndex()
	if err != nil {
		return nil, err
	}

	return &inst, nil
}

func newAllDocsIndex(b *bucket) *primaryIndex {
	meta := ast.NewFunctionCall("meta", ast.FunctionArgExpressionList{})
	mdid := ast.NewDotMemberOperator(meta, ast.NewProperty("id"))
	ddoc := designdoc{name: "", viewname: "_all_docs"}
	idx := primaryIndex{
		viewIndex{
			name:   ALLDOCS_INDEX,
			using:  catalog.VIEW,
			on:     catalog.IndexKey{mdid},
			ddoc:   &ddoc,
			bucket: b,
		},
	}
	return &idx
}

func newPrimaryDDoc() *designdoc {
	var doc designdoc
	line := strings.Replace(templPrimary, "$rnd", strconv.Itoa(int(rand.Int31())), -1)
	line = strings.Replace(line, "$string", strconv.Itoa(TYPE_STRING), -1)
	doc.mapfn = line
	doc.reducefn = ""
	doc.name = "ddl_" + PRIMARY_INDEX
	doc.viewname = PRIMARY_INDEX
	return &doc
}

func generateMap(on catalog.IndexKey, doc *designdoc) error {

	buf := new(bytes.Buffer)

	fmt.Fprintln(buf, templStart)
	line := strings.Replace(templFunctions, "$null", strconv.Itoa(TYPE_NULL), -1)
	line = strings.Replace(line, "$boolean", strconv.Itoa(TYPE_BOOLEAN), -1)
	line = strings.Replace(line, "$number", strconv.Itoa(TYPE_NUMBER), -1)
	line = strings.Replace(line, "$string", strconv.Itoa(TYPE_STRING), -1)
	line = strings.Replace(line, "$array", strconv.Itoa(TYPE_ARRAY), -1)
	line = strings.Replace(line, "$object", strconv.Itoa(TYPE_OBJECT), -1)
	fmt.Fprintln(buf, line)

	keylist := new(bytes.Buffer)
	for idx, expr := range on {

		walker := NewWalker()
		_, err := walker.Visit(expr)
		if err != nil {
			return err
		}

		jvar := fmt.Sprintf("key%v", idx+1)
		line := strings.Replace(templExpr, "$var", jvar, -1)
		line = strings.Replace(line, "$path", walker.JS(), -1)
		fmt.Fprint(buf, line)

		if idx > 0 {
			fmt.Fprint(keylist, ", ")
		}
		fmt.Fprint(keylist, jvar)
	}

	line = strings.Replace(templKey, "$keylist", keylist.String(), -1)

	fmt.Fprint(buf, line)
	fmt.Fprint(buf, templEmit)

	line = strings.Replace(templEnd, "$rnd", strconv.Itoa(int(rand.Int31())), -1)
	fmt.Fprint(buf, line)

	doc.mapfn = buf.String()
	return nil
}

func generateReduce(on catalog.IndexKey, doc *designdoc) error {
	doc.reducefn = ""
	return nil
}

func (idx *viewIndex) putDesignDoc() error {
	var view cb.ViewDefinition
	view.Map = idx.ddoc.mapfn

	var put ddocJSON
	put.Views = make(map[string]cb.ViewDefinition)
	put.Views[idx.name] = view
	put.IndexChecksum = idx.ddoc.checksum()

	put.IndexOn = make([]string, len(idx.on))
	for idx, expr := range idx.on {
		ser, err := json.Marshal(expr)
		if err != nil {
			return err
		}
		put.IndexOn[idx] = string(ser)
	}

	if err := idx.bucket.cbbucket.PutDDoc(idx.DDocName(), &put); err != nil {
		return err
	}

	var saved *ddocJSON = nil
	var err error = nil

	// give the PUT some time to register
	for i := 0; i < 3; i++ {
		if i > 1 {
			time.Sleep(time.Duration(i*3) * time.Second)
		}

		saved, err = getDesignDoc(idx.bucket, idx.DDocName())
		if err == nil {
			break
		}
	}

	if err != nil {
		return errors.New("Creating index '" + idx.name + "' failed: " + err.Error())
	}

	if saved.IndexChecksum != idx.ddoc.checksum() {
		return errors.New("Checksum mismatch after creating index '" + idx.name + "'")
	}

	return nil
}

func (ddoc *designdoc) checksum() int {
	mapSum := crc32.ChecksumIEEE([]byte(ddoc.mapfn))
	reduceSum := crc32.ChecksumIEEE([]byte(ddoc.reducefn))
	return int(mapSum + reduceSum)
}

func getDesignDoc(b *bucket, ddocname string) (*ddocJSON, error) {
	var ddoc ddocJSON
	err := b.cbbucket.GetDDoc(ddocname, &ddoc)
	if err != nil {
		return nil, err
	}
	return &ddoc, nil
}

func (idx *viewIndex) DropViewIndex() error {
	if err := idx.bucket.cbbucket.DeleteDDoc(idx.ddoc.name); err != nil {
		return err
	}
	return nil
}

func (idx *viewIndex) WaitForIndex() error {
	var err error
	// if we have got this far, very likely any errors are
	// due to index not yet being noticed by the system.
	for i := 0; i < 3; i++ {
		if i > 1 {
			time.Sleep(time.Duration(i*3) * time.Second)
		}
		_, err = idx.bucket.cbbucket.View(
			idx.ddoc.name,
			idx.ddoc.viewname,
			map[string]interface{}{
				"start_key": []interface{}{"thing"},
				"end_key":   []interface{}{"thing", map[string]string{}},
				"stale":     false,
			})
		if err == nil {
			break
		}
	}
	return err
}

// AST to JS conversion
type JsStatement struct {
	js bytes.Buffer
}

func NewWalker() *JsStatement {
	var js JsStatement
	return &js
}

func (this *JsStatement) JS() string {
	return this.js.String()
}

// inorder traversal of the AST to get JS expression out of it
func (this *JsStatement) Visit(e ast.Expression) (ast.Expression, error) {
	switch expr := e.(type) {

	case *ast.DotMemberOperator:
		if this.js.Len() == 0 {
			this.js.WriteString("doc.")
		}
		_, err := expr.Left.Accept(this)
		if err != nil {
			return nil, err
		}
		this.js.WriteString(".")
		_, err = expr.Right.Accept(this)
		if err != nil {
			return nil, err
		}

	case *ast.BracketMemberOperator:
		if this.js.Len() == 0 {
			this.js.WriteString("doc.")
		}
		_, err := expr.Left.Accept(this)
		if err != nil {
			return nil, err
		}
		this.js.WriteString("[")
		_, err = expr.Right.Accept(this)
		if err != nil {
			return nil, err
		}
		this.js.WriteString("]")

	case *ast.Property:
		if this.js.Len() == 0 {
			this.js.WriteString("doc.")
		}
		this.js.WriteString(expr.Path)

	case *ast.LiteralNumber:
		this.js.WriteString(fmt.Sprintf("%v", expr.Val))

	case *ast.LiteralString:
		this.js.WriteString(expr.Val)

	default:
		return e, errors.New("Expression is not supported by indexing currently: " + e.String())

	}
	return e, nil
}
<|endoftext|>


<commit_before>package antibodylib_test

import (
	"bytes"
	"io/ioutil"
	"os"
	"runtime"
	"strings"
	"testing"

	"github.com/getantibody/antibody/antibodylib"
	"github.com/stretchr/testify/assert"
)

func TestAntibody(t *testing.T) {
	assert := assert.New(t)
	home := home()
	bundles := []string{
		"# comments also are allowed",
		"caarlos0/ports kind:path # comment at the end of the line",
		"caarlos0/jvm kind:path branch:gh-pages",
		"caarlos0/zsh-open-pr     kind:zsh",
		"",
		"        ",
		"  # trick play",
		"/tmp kind:path",
	}
	sh, err := antibodylib.New(
		home,
		bytes.NewBufferString(strings.Join(bundles, "\n")),
		runtime.NumCPU(),
	).Bundle()
	assert.NoError(err)
	files, err := ioutil.ReadDir(home)
	assert.NoError(err)
	assert.Len(files, 3)
	assert.Contains(sh, `export PATH="/tmp:$PATH"`)
	assert.Contains(sh, `export PATH="`+home+`/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-ports:$PATH"`)
	assert.Contains(sh, `export PATH="`+home+`/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-jvm:$PATH"`)
	assert.Contains(sh, `source `+home+`/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-zsh-open-pr/git-open-pr.plugin.zsh`)
}

func TestAntibodyError(t *testing.T) {
	assert := assert.New(t)
	home := home()
	bundles := bytes.NewBufferString("invalid-repo")
	sh, err := antibodylib.New(home, bundles, runtime.NumCPU()).Bundle()
	assert.Error(err)
	assert.Empty(sh)
}

func TestHome(t *testing.T) {
	assert.Contains(t, antibodylib.Home(), "antibody")
}

func TestHomeFromEnvironmentVariable(t *testing.T) {
	assert.NoError(t, os.Setenv("ANTIBODY_HOME", "/tmp"))
	assert.Equal(t, "/tmp", antibodylib.Home())
}

func home() string {
	home, err := ioutil.TempDir(os.TempDir(), "antibody")
	if err != nil {
		panic(err.Error())
	}
	return home
}
<commit_msg>chore: added more tests<commit_after>package antibodylib_test

import (
	"bytes"
	"io/ioutil"
	"os"
	"runtime"
	"strings"
	"testing"

	"github.com/getantibody/antibody/antibodylib"
	"github.com/stretchr/testify/assert"
)

func TestAntibody(t *testing.T) {
	assert := assert.New(t)
	home := home()
	bundles := []string{
		"# comments also are allowed",
		"caarlos0/ports kind:path # comment at the end of the line",
		"caarlos0/jvm kind:path branch:gh-pages",
		"caarlos0/zsh-open-pr     kind:zsh",
		"",
		"        ",
		"  # trick play",
		"/tmp kind:path",
	}
	sh, err := antibodylib.New(
		home,
		bytes.NewBufferString(strings.Join(bundles, "\n")),
		runtime.NumCPU(),
	).Bundle()
	assert.NoError(err)
	files, err := ioutil.ReadDir(home)
	assert.NoError(err)
	assert.Len(files, 3)
	assert.Contains(sh, `export PATH="/tmp:$PATH"`)
	assert.Contains(sh, `export PATH="`+home+`/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-ports:$PATH"`)
	assert.Contains(sh, `export PATH="`+home+`/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-jvm:$PATH"`)
	assert.Contains(sh, `source `+home+`/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-zsh-open-pr/git-open-pr.plugin.zsh`)
}

func TestAntibodyError(t *testing.T) {
	assert := assert.New(t)
	home := home()
	bundles := bytes.NewBufferString("invalid-repo")
	sh, err := antibodylib.New(home, bundles, runtime.NumCPU()).Bundle()
	assert.Error(err)
	assert.Empty(sh)
}

func TestMultipleRepositories(t *testing.T) {
	var assert = assert.New(t)
	home := home()
	bundles := []string{
		"# this block is in alphabetic order",
		"caarlos0/git-add-remote kind:path",
		"caarlos0/jvm",
		"caarlos0/ports kind:path",
		"caarlos0/zsh-git-fetch-merge kind:path",
		"caarlos0/zsh-git-sync kind:path",
		"caarlos0/zsh-mkc",
		"caarlos0/zsh-open-pr kind:path",
		"mafredri/zsh-async",
		"rupa/z",
		"Tarrasch/zsh-bd",
		"wbinglee/zsh-wakatime",
		"zsh-users/zsh-completions",
		"zsh-users/zsh-autosuggestions",
		"",
		"# these should be at last!",
		"sindresorhus/pure",
		"zsh-users/zsh-syntax-highlighting",
		"zsh-users/zsh-history-substring-search",
	}
	sh, err := antibodylib.New(
		home,
		bytes.NewBufferString(strings.Join(bundles, "\n")),
		runtime.NumCPU(),
	).Bundle()
	assert.NoError(err)
	assert.Len(strings.Split(sh, "\n"), 16)
}

func TestHome(t *testing.T) {
	assert.Contains(t, antibodylib.Home(), "antibody")
}

func TestHomeFromEnvironmentVariable(t *testing.T) {
	assert.NoError(t, os.Setenv("ANTIBODY_HOME", "/tmp"))
	assert.Equal(t, "/tmp", antibodylib.Home())
}

func home() string {
	home, err := ioutil.TempDir(os.TempDir(), "antibody")
	if err != nil {
		panic(err.Error())
	}
	return home
}
<|endoftext|>