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

package compiler

import (
	"fmt"
	"sort"

	"github.com/google/syzkaller/pkg/ast"
	"github.com/google/syzkaller/sys"
)

const sizeUnassigned = ^uint64(0)

func (comp *compiler) genResources() []*sys.ResourceDesc {
	var resources []*sys.ResourceDesc
	for _, decl := range comp.desc.Nodes {
		if n, ok := decl.(*ast.Resource); ok {
			resources = append(resources, comp.genResource(n))
		}
	}
	sort.Slice(resources, func(i, j int) bool {
		return resources[i].Name < resources[j].Name
	})
	return resources
}

func (comp *compiler) genResource(n *ast.Resource) *sys.ResourceDesc {
	res := &sys.ResourceDesc{
		Name: n.Name.Name,
	}
	var base *ast.Type
	for n != nil {
		res.Values = append(genIntArray(n.Values), res.Values...)
		res.Kind = append([]string{n.Name.Name}, res.Kind...)
		base = n.Base
		n = comp.resources[n.Base.Ident]
	}
	if len(res.Values) == 0 {
		res.Values = []uint64{0}
	}
	res.Type = comp.genType(base, "", sys.DirIn, false)
	return res
}

func (comp *compiler) genSyscalls() []*sys.Call {
	var calls []*sys.Call
	for _, decl := range comp.desc.Nodes {
		if n, ok := decl.(*ast.Call); ok {
			calls = append(calls, comp.genSyscall(n))
		}
	}
	sort.Slice(calls, func(i, j int) bool {
		return calls[i].Name < calls[j].Name
	})
	return calls
}

func (comp *compiler) genSyscall(n *ast.Call) *sys.Call {
	var ret sys.Type
	if n.Ret != nil {
		ret = comp.genType(n.Ret, "ret", sys.DirOut, true)
	}
	return &sys.Call{
		Name:     n.Name.Name,
		CallName: n.CallName,
		NR:       n.NR,
		Args:     comp.genFieldArray(n.Args, sys.DirIn, true),
		Ret:      ret,
	}
}

func (comp *compiler) genStructDescs(syscalls []*sys.Call) []*sys.KeyedStruct {
	// Calculate struct/union/array sizes, add padding to structs and detach
	// StructDesc's from StructType's. StructType's can be recursive so it's
	// not possible to write them out inline as other types. To break the
	// recursion detach them, and write StructDesc's out as separate array
	// of KeyedStruct's. sys package will reattach them during init.

	padded := make(map[interface{}]bool)
	detach := make(map[**sys.StructDesc]bool)
	var structs []*sys.KeyedStruct
	var rec func(t sys.Type)
	checkStruct := func(key sys.StructKey, descp **sys.StructDesc) bool {
		detach[descp] = true
		desc := *descp
		if padded[desc] {
			return false
		}
		padded[desc] = true
		for _, f := range desc.Fields {
			rec(f)
			if !f.Varlen() && f.Size() == sizeUnassigned {
				// An inner struct is not padded yet.
				// Leave this struct for next iteration.
				delete(padded, desc)
				return false
			}
		}
		structs = append(structs, &sys.KeyedStruct{
			Key:  key,
			Desc: desc,
		})
		return true
	}
	rec = func(t0 sys.Type) {
		switch t := t0.(type) {
		case *sys.PtrType:
			rec(t.Type)
		case *sys.ArrayType:
			if padded[t] {
				return
			}
			rec(t.Type)
			if !t.Type.Varlen() && t.Type.Size() == sizeUnassigned {
				// An inner struct is not padded yet.
				// Leave this array for next iteration.
				return
			}
			padded[t] = true
			t.TypeSize = 0
			if t.Kind == sys.ArrayRangeLen && t.RangeBegin == t.RangeEnd && !t.Type.Varlen() {
				t.TypeSize = t.RangeBegin * t.Type.Size()
			}
		case *sys.StructType:
			if !checkStruct(t.Key, &t.StructDesc) {
				return
			}
			// Add paddings, calculate size, mark bitfields.
			varlen := false
			for _, f := range t.Fields {
				if f.Varlen() {
					varlen = true
				}
			}
			comp.markBitfields(t.Fields)
			packed, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
			t.Fields = comp.addAlignment(t.Fields, varlen, packed, alignAttr)
			t.AlignAttr = alignAttr
			t.TypeSize = 0
			if !varlen {
				for _, f := range t.Fields {
					if !f.BitfieldMiddle() {
						t.TypeSize += f.Size()
					}
				}
			}
		case *sys.UnionType:
			if !checkStruct(t.Key, &t.StructDesc) {
				return
			}
			t.TypeSize = 0
			varlen := comp.parseUnionAttrs(comp.structNodes[t.StructDesc])
			if !varlen {
				for _, fld := range t.Fields {
					if t.TypeSize < fld.Size() {
						t.TypeSize = fld.Size()
					}
				}
			}
		}
	}

	// We have to do this in the loop until we pad nothing new
	// due to recursive structs.
	for {
		start := len(padded)
		for _, c := range syscalls {
			for _, a := range c.Args {
				rec(a)
			}
			if c.Ret != nil {
				rec(c.Ret)
			}
		}
		if start == len(padded) {
			break
		}
	}

	// Detach StructDesc's from StructType's. sys will reattach them again.
	for descp := range detach {
		*descp = nil
	}

	sort.Slice(structs, func(i, j int) bool {
		si, sj := structs[i], structs[j]
		if si.Key.Name != sj.Key.Name {
			return si.Key.Name < sj.Key.Name
		}
		return si.Key.Dir < sj.Key.Dir
	})
	return structs
}

func (comp *compiler) genStructDesc(res *sys.StructDesc, n *ast.Struct, dir sys.Dir) {
	// Leave node for genStructDescs to calculate size/padding.
	comp.structNodes[res] = n
	*res = sys.StructDesc{
		TypeCommon: genCommon(n.Name.Name, "", sizeUnassigned, dir, false),
		Fields:     comp.genFieldArray(n.Fields, dir, false),
	}
}

func (comp *compiler) isStructVarlen(name string) bool {
	if varlen, ok := comp.structVarlen[name]; ok {
		return varlen
	}
	s := comp.structs[name]
	if s.IsUnion && comp.parseUnionAttrs(s) {
		comp.structVarlen[name] = true
		return true
	}
	comp.structVarlen[name] = false // to not hang on recursive types
	varlen := false
	for _, fld := range s.Fields {
		if comp.isVarlen(fld.Type) {
			varlen = true
			break
		}
	}
	comp.structVarlen[name] = varlen
	return varlen
}

func (comp *compiler) markBitfields(fields []sys.Type) {
	var bfOffset uint64
	for i, f := range fields {
		if f.BitfieldLength() == 0 {
			continue
		}
		off, middle := bfOffset, true
		bfOffset += f.BitfieldLength()
		if i == len(fields)-1 || // Last bitfield in a group, if last field of the struct...
			fields[i+1].BitfieldLength() == 0 || // or next field is not a bitfield...
			f.Size() != fields[i+1].Size() || // or next field is of different size...
			bfOffset+fields[i+1].BitfieldLength() > f.Size()*8 { // or next field does not fit into the current group.
			middle, bfOffset = false, 0
		}
		setBitfieldOffset(f, off, middle)
	}
}

func setBitfieldOffset(t0 sys.Type, offset uint64, middle bool) {
	switch t := t0.(type) {
	case *sys.IntType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.ConstType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.LenType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.FlagsType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.ProcType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	default:
		panic(fmt.Sprintf("type %#v can't be a bitfield", t))
	}
}

func (comp *compiler) addAlignment(fields []sys.Type, varlen, packed bool, alignAttr uint64) []sys.Type {
	var newFields []sys.Type
	if packed {
		// If a struct is packed, statically sized and has explicitly set alignment,
		// add a padding at the end.
		newFields = fields
		if !varlen && alignAttr != 0 {
			size := uint64(0)
			for _, f := range fields {
				size += f.Size()
			}
			if tail := size % alignAttr; tail != 0 {
				newFields = append(newFields, genPad(alignAttr-tail))
			}
		}
		return newFields
	}
	var off uint64
	// TODO(dvyukov): this is wrong: if alignAttr!=0, we must use it, not max
	align := alignAttr
	for i, f := range fields {
		if i > 0 && !fields[i-1].BitfieldMiddle() {
			a := comp.typeAlign(f)
			if align < a {
				align = a
			}
			// Append padding if the last field is not a bitfield or it's the last bitfield in a set.
			if off%a != 0 {
				pad := a - off%a
				off += pad
				newFields = append(newFields, genPad(pad))
			}
		}
		newFields = append(newFields, f)
		if !f.BitfieldMiddle() && (i != len(fields)-1 || !f.Varlen()) {
			// Increase offset if the current field is not a bitfield
			// or it's the last bitfield in a set, except when it's
			// the last field in a struct and has variable length.
			off += f.Size()
		}
	}
	if align != 0 && off%align != 0 && !varlen {
		pad := align - off%align
		off += pad
		newFields = append(newFields, genPad(pad))
	}
	return newFields
}

func (comp *compiler) typeAlign(t0 sys.Type) uint64 {
	switch t0.(type) {
	case *sys.IntType, *sys.ConstType, *sys.LenType, *sys.FlagsType, *sys.ProcType,
		*sys.CsumType, *sys.PtrType, *sys.VmaType, *sys.ResourceType:
		return t0.Size()
	case *sys.BufferType:
		return 1
	}

	switch t := t0.(type) {
	case *sys.ArrayType:
		return comp.typeAlign(t.Type)
	case *sys.StructType:
		packed, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
		if alignAttr != 0 {
			return alignAttr // overrided by user attribute
		}
		if packed {
			return 1
		}
		align := uint64(0)
		for _, f := range t.Fields {
			if a := comp.typeAlign(f); align < a {
				align = a
			}
		}
		return align
	case *sys.UnionType:
		align := uint64(0)
		for _, f := range t.Fields {
			if a := comp.typeAlign(f); align < a {
				align = a
			}
		}
		return align
	default:
		panic(fmt.Sprintf("unknown type: %#v", t))
	}
}

func genPad(size uint64) sys.Type {
	return &sys.ConstType{
		IntTypeCommon: genIntCommon(genCommon("pad", "", size, sys.DirIn, false), 0, false),
		IsPad:         true,
	}
}

func (comp *compiler) genField(f *ast.Field, dir sys.Dir, isArg bool) sys.Type {
	return comp.genType(f.Type, f.Name.Name, dir, isArg)
}

func (comp *compiler) genFieldArray(fields []*ast.Field, dir sys.Dir, isArg bool) []sys.Type {
	var res []sys.Type
	for _, f := range fields {
		res = append(res, comp.genField(f, dir, isArg))
	}
	return res
}

func (comp *compiler) genType(t *ast.Type, field string, dir sys.Dir, isArg bool) sys.Type {
	desc, args, base := comp.getArgsBase(t, field, dir, isArg)
	return desc.Gen(comp, t, args, base)
}

func genCommon(name, field string, size uint64, dir sys.Dir, opt bool) sys.TypeCommon {
	return sys.TypeCommon{
		TypeName:   name,
		TypeSize:   size,
		FldName:    field,
		ArgDir:     dir,
		IsOptional: opt,
	}
}

func genIntCommon(com sys.TypeCommon, bitLen uint64, bigEndian bool) sys.IntTypeCommon {
	return sys.IntTypeCommon{
		TypeCommon:  com,
		BigEndian:   bigEndian,
		BitfieldLen: bitLen,
	}
}

func genIntArray(a []*ast.Int) []uint64 {
	r := make([]uint64, len(a))
	for i, v := range a {
		r[i] = v.Value
	}
	return r
}

func genStrArray(a []*ast.String) []string {
	r := make([]string, len(a))
	for i, v := range a {
		r[i] = v.Value
	}
	return r
}
<commit_msg>pkg/compiler: fix alignment calculation bug<commit_after>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package compiler

import (
	"fmt"
	"sort"

	"github.com/google/syzkaller/pkg/ast"
	"github.com/google/syzkaller/sys"
)

const sizeUnassigned = ^uint64(0)

func (comp *compiler) genResources() []*sys.ResourceDesc {
	var resources []*sys.ResourceDesc
	for _, decl := range comp.desc.Nodes {
		if n, ok := decl.(*ast.Resource); ok {
			resources = append(resources, comp.genResource(n))
		}
	}
	sort.Slice(resources, func(i, j int) bool {
		return resources[i].Name < resources[j].Name
	})
	return resources
}

func (comp *compiler) genResource(n *ast.Resource) *sys.ResourceDesc {
	res := &sys.ResourceDesc{
		Name: n.Name.Name,
	}
	var base *ast.Type
	for n != nil {
		res.Values = append(genIntArray(n.Values), res.Values...)
		res.Kind = append([]string{n.Name.Name}, res.Kind...)
		base = n.Base
		n = comp.resources[n.Base.Ident]
	}
	if len(res.Values) == 0 {
		res.Values = []uint64{0}
	}
	res.Type = comp.genType(base, "", sys.DirIn, false)
	return res
}

func (comp *compiler) genSyscalls() []*sys.Call {
	var calls []*sys.Call
	for _, decl := range comp.desc.Nodes {
		if n, ok := decl.(*ast.Call); ok {
			calls = append(calls, comp.genSyscall(n))
		}
	}
	sort.Slice(calls, func(i, j int) bool {
		return calls[i].Name < calls[j].Name
	})
	return calls
}

func (comp *compiler) genSyscall(n *ast.Call) *sys.Call {
	var ret sys.Type
	if n.Ret != nil {
		ret = comp.genType(n.Ret, "ret", sys.DirOut, true)
	}
	return &sys.Call{
		Name:     n.Name.Name,
		CallName: n.CallName,
		NR:       n.NR,
		Args:     comp.genFieldArray(n.Args, sys.DirIn, true),
		Ret:      ret,
	}
}

func (comp *compiler) genStructDescs(syscalls []*sys.Call) []*sys.KeyedStruct {
	// Calculate struct/union/array sizes, add padding to structs and detach
	// StructDesc's from StructType's. StructType's can be recursive so it's
	// not possible to write them out inline as other types. To break the
	// recursion detach them, and write StructDesc's out as separate array
	// of KeyedStruct's. sys package will reattach them during init.

	padded := make(map[interface{}]bool)
	detach := make(map[**sys.StructDesc]bool)
	var structs []*sys.KeyedStruct
	var rec func(t sys.Type)
	checkStruct := func(key sys.StructKey, descp **sys.StructDesc) bool {
		detach[descp] = true
		desc := *descp
		if padded[desc] {
			return false
		}
		padded[desc] = true
		for _, f := range desc.Fields {
			rec(f)
			if !f.Varlen() && f.Size() == sizeUnassigned {
				// An inner struct is not padded yet.
				// Leave this struct for next iteration.
				delete(padded, desc)
				return false
			}
		}
		structs = append(structs, &sys.KeyedStruct{
			Key:  key,
			Desc: desc,
		})
		return true
	}
	rec = func(t0 sys.Type) {
		switch t := t0.(type) {
		case *sys.PtrType:
			rec(t.Type)
		case *sys.ArrayType:
			if padded[t] {
				return
			}
			rec(t.Type)
			if !t.Type.Varlen() && t.Type.Size() == sizeUnassigned {
				// An inner struct is not padded yet.
				// Leave this array for next iteration.
				return
			}
			padded[t] = true
			t.TypeSize = 0
			if t.Kind == sys.ArrayRangeLen && t.RangeBegin == t.RangeEnd && !t.Type.Varlen() {
				t.TypeSize = t.RangeBegin * t.Type.Size()
			}
		case *sys.StructType:
			if !checkStruct(t.Key, &t.StructDesc) {
				return
			}
			// Add paddings, calculate size, mark bitfields.
			varlen := false
			for _, f := range t.Fields {
				if f.Varlen() {
					varlen = true
				}
			}
			comp.markBitfields(t.Fields)
			packed, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
			t.Fields = comp.addAlignment(t.Fields, varlen, packed, alignAttr)
			t.AlignAttr = alignAttr
			t.TypeSize = 0
			if !varlen {
				for _, f := range t.Fields {
					if !f.BitfieldMiddle() {
						t.TypeSize += f.Size()
					}
				}
			}
		case *sys.UnionType:
			if !checkStruct(t.Key, &t.StructDesc) {
				return
			}
			t.TypeSize = 0
			varlen := comp.parseUnionAttrs(comp.structNodes[t.StructDesc])
			if !varlen {
				for _, fld := range t.Fields {
					if t.TypeSize < fld.Size() {
						t.TypeSize = fld.Size()
					}
				}
			}
		}
	}

	// We have to do this in the loop until we pad nothing new
	// due to recursive structs.
	for {
		start := len(padded)
		for _, c := range syscalls {
			for _, a := range c.Args {
				rec(a)
			}
			if c.Ret != nil {
				rec(c.Ret)
			}
		}
		if start == len(padded) {
			break
		}
	}

	// Detach StructDesc's from StructType's. sys will reattach them again.
	for descp := range detach {
		*descp = nil
	}

	sort.Slice(structs, func(i, j int) bool {
		si, sj := structs[i], structs[j]
		if si.Key.Name != sj.Key.Name {
			return si.Key.Name < sj.Key.Name
		}
		return si.Key.Dir < sj.Key.Dir
	})
	return structs
}

func (comp *compiler) genStructDesc(res *sys.StructDesc, n *ast.Struct, dir sys.Dir) {
	// Leave node for genStructDescs to calculate size/padding.
	comp.structNodes[res] = n
	*res = sys.StructDesc{
		TypeCommon: genCommon(n.Name.Name, "", sizeUnassigned, dir, false),
		Fields:     comp.genFieldArray(n.Fields, dir, false),
	}
}

func (comp *compiler) isStructVarlen(name string) bool {
	if varlen, ok := comp.structVarlen[name]; ok {
		return varlen
	}
	s := comp.structs[name]
	if s.IsUnion && comp.parseUnionAttrs(s) {
		comp.structVarlen[name] = true
		return true
	}
	comp.structVarlen[name] = false // to not hang on recursive types
	varlen := false
	for _, fld := range s.Fields {
		if comp.isVarlen(fld.Type) {
			varlen = true
			break
		}
	}
	comp.structVarlen[name] = varlen
	return varlen
}

func (comp *compiler) markBitfields(fields []sys.Type) {
	var bfOffset uint64
	for i, f := range fields {
		if f.BitfieldLength() == 0 {
			continue
		}
		off, middle := bfOffset, true
		bfOffset += f.BitfieldLength()
		if i == len(fields)-1 || // Last bitfield in a group, if last field of the struct...
			fields[i+1].BitfieldLength() == 0 || // or next field is not a bitfield...
			f.Size() != fields[i+1].Size() || // or next field is of different size...
			bfOffset+fields[i+1].BitfieldLength() > f.Size()*8 { // or next field does not fit into the current group.
			middle, bfOffset = false, 0
		}
		setBitfieldOffset(f, off, middle)
	}
}

func setBitfieldOffset(t0 sys.Type, offset uint64, middle bool) {
	switch t := t0.(type) {
	case *sys.IntType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.ConstType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.LenType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.FlagsType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	case *sys.ProcType:
		t.BitfieldOff, t.BitfieldMdl = offset, middle
	default:
		panic(fmt.Sprintf("type %#v can't be a bitfield", t))
	}
}

func (comp *compiler) addAlignment(fields []sys.Type, varlen, packed bool, alignAttr uint64) []sys.Type {
	var newFields []sys.Type
	if packed {
		// If a struct is packed, statically sized and has explicitly set alignment,
		// add a padding at the end.
		newFields = fields
		if !varlen && alignAttr != 0 {
			size := uint64(0)
			for _, f := range fields {
				size += f.Size()
			}
			if tail := size % alignAttr; tail != 0 {
				newFields = append(newFields, genPad(alignAttr-tail))
			}
		}
		return newFields
	}
	var align, off uint64
	for i, f := range fields {
		if i > 0 && !fields[i-1].BitfieldMiddle() {
			a := comp.typeAlign(f)
			if align < a {
				align = a
			}
			// Append padding if the last field is not a bitfield or it's the last bitfield in a set.
			if off%a != 0 {
				pad := a - off%a
				off += pad
				newFields = append(newFields, genPad(pad))
			}
		}
		newFields = append(newFields, f)
		if !f.BitfieldMiddle() && (i != len(fields)-1 || !f.Varlen()) {
			// Increase offset if the current field is not a bitfield
			// or it's the last bitfield in a set, except when it's
			// the last field in a struct and has variable length.
			off += f.Size()
		}
	}
	if alignAttr != 0 {
		align = alignAttr
	}
	if align != 0 && off%align != 0 && !varlen {
		pad := align - off%align
		off += pad
		newFields = append(newFields, genPad(pad))
	}
	return newFields
}

func (comp *compiler) typeAlign(t0 sys.Type) uint64 {
	switch t0.(type) {
	case *sys.IntType, *sys.ConstType, *sys.LenType, *sys.FlagsType, *sys.ProcType,
		*sys.CsumType, *sys.PtrType, *sys.VmaType, *sys.ResourceType:
		return t0.Size()
	case *sys.BufferType:
		return 1
	}

	switch t := t0.(type) {
	case *sys.ArrayType:
		return comp.typeAlign(t.Type)
	case *sys.StructType:
		packed, alignAttr := comp.parseStructAttrs(comp.structNodes[t.StructDesc])
		if alignAttr != 0 {
			return alignAttr // overrided by user attribute
		}
		if packed {
			return 1
		}
		align := uint64(0)
		for _, f := range t.Fields {
			if a := comp.typeAlign(f); align < a {
				align = a
			}
		}
		return align
	case *sys.UnionType:
		align := uint64(0)
		for _, f := range t.Fields {
			if a := comp.typeAlign(f); align < a {
				align = a
			}
		}
		return align
	default:
		panic(fmt.Sprintf("unknown type: %#v", t))
	}
}

func genPad(size uint64) sys.Type {
	return &sys.ConstType{
		IntTypeCommon: genIntCommon(genCommon("pad", "", size, sys.DirIn, false), 0, false),
		IsPad:         true,
	}
}

func (comp *compiler) genField(f *ast.Field, dir sys.Dir, isArg bool) sys.Type {
	return comp.genType(f.Type, f.Name.Name, dir, isArg)
}

func (comp *compiler) genFieldArray(fields []*ast.Field, dir sys.Dir, isArg bool) []sys.Type {
	var res []sys.Type
	for _, f := range fields {
		res = append(res, comp.genField(f, dir, isArg))
	}
	return res
}

func (comp *compiler) genType(t *ast.Type, field string, dir sys.Dir, isArg bool) sys.Type {
	desc, args, base := comp.getArgsBase(t, field, dir, isArg)
	return desc.Gen(comp, t, args, base)
}

func genCommon(name, field string, size uint64, dir sys.Dir, opt bool) sys.TypeCommon {
	return sys.TypeCommon{
		TypeName:   name,
		TypeSize:   size,
		FldName:    field,
		ArgDir:     dir,
		IsOptional: opt,
	}
}

func genIntCommon(com sys.TypeCommon, bitLen uint64, bigEndian bool) sys.IntTypeCommon {
	return sys.IntTypeCommon{
		TypeCommon:  com,
		BigEndian:   bigEndian,
		BitfieldLen: bitLen,
	}
}

func genIntArray(a []*ast.Int) []uint64 {
	r := make([]uint64, len(a))
	for i, v := range a {
		r[i] = v.Value
	}
	return r
}

func genStrArray(a []*ast.String) []string {
	r := make([]string, len(a))
	for i, v := range a {
		r[i] = v.Value
	}
	return r
}
<|endoftext|>


<commit_before>/*******************************************************************************
*
* Copyright 2017 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, you may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*******************************************************************************/

package core

import (
	"fmt"
	"regexp"
	"sort"
	"strconv"

	"github.com/gophercloud/gophercloud"
	"github.com/sapcc/go-bits/logg"
	"github.com/sapcc/limes"
)

//Cluster contains all configuration and runtime information about a single
//cluster. It is passed around a lot in Limes code, mostly for the cluster ID,
//the list of enabled services, and access to the quota and capacity plugins.
type Cluster struct {
	ID                string
	Config            *ClusterConfiguration
	ServiceTypes      []string
	IsServiceShared   map[string]bool
	DiscoveryPlugin   DiscoveryPlugin
	QuotaPlugins      map[string]QuotaPlugin
	CapacityPlugins   map[string]CapacityPlugin
	Authoritative     bool
	QuotaConstraints  *QuotaConstraintSet
	LowPrivilegeRaise struct {
		LimitsForDomains  map[string]map[string]LowPrivilegeRaiseLimit
		LimitsForProjects map[string]map[string]LowPrivilegeRaiseLimit
	}
}

//NewCluster creates a new Cluster instance with the given ID and
//configuration, and also initializes all quota and capacity plugins. Errors
//will be logged when some of the requested plugins cannot be found.
func NewCluster(id string, config *ClusterConfiguration) *Cluster {
	factory, exists := discoveryPluginFactories[config.Discovery.Method]
	if !exists {
		logg.Fatal("setup for cluster %s failed: no suitable discovery plugin found", id)
	}

	c := &Cluster{
		ID:              id,
		Config:          config,
		IsServiceShared: make(map[string]bool),
		DiscoveryPlugin: factory(config.Discovery),
		QuotaPlugins:    make(map[string]QuotaPlugin),
		CapacityPlugins: make(map[string]CapacityPlugin),
		Authoritative:   config.Authoritative,
	}

	for _, srv := range config.Services {
		factory, exists := quotaPluginFactories[srv.Type]
		if !exists {
			logg.Error("skipping service %s: no suitable collector plugin found", srv.Type)
			continue
		}

		scrapeSubresources := map[string]bool{}
		for _, resName := range config.Subresources[srv.Type] {
			scrapeSubresources[resName] = true
		}

		plugin := factory(srv, scrapeSubresources)
		if plugin == nil || plugin.ServiceInfo().Type != srv.Type {
			logg.Error("skipping service %s: failed to initialize collector plugin", srv.Type)
			continue
		}

		c.ServiceTypes = append(c.ServiceTypes, srv.Type)
		c.QuotaPlugins[srv.Type] = plugin
		c.IsServiceShared[srv.Type] = srv.Shared
	}

	scrapeSubcapacities := make(map[string]map[string]bool)
	for serviceType, resourceNames := range config.Subcapacities {
		m := make(map[string]bool)
		for _, resourceName := range resourceNames {
			m[resourceName] = true
		}
		scrapeSubcapacities[serviceType] = m
	}

	for _, capa := range config.Capacitors {
		factory, exists := capacityPluginFactories[capa.ID]
		if !exists {
			logg.Error("skipping capacitor %s: no suitable collector plugin found", capa.ID)
			continue
		}
		plugin := factory(capa, scrapeSubcapacities)
		if plugin == nil || plugin.ID() != capa.ID {
			logg.Error("skipping capacitor %s: failed to initialize collector plugin", capa.ID)
			continue
		}
		c.CapacityPlugins[capa.ID] = plugin
	}

	sort.Strings(c.ServiceTypes) //determinism is useful for unit tests

	return c
}

//Connect calls Connect() on all AuthParameters for this Cluster, thus ensuring
//that all ProviderClient instances are available. It also calls Init() on all
//quota plugins.
//
//It also loads the QuotaConstraints for this cluster, if configured. The
//LowPrivilegeRaise.Limits fields are also initialized here. We also validate
//if Config.ResourceBehavior[].ScalesWith refers to existing resources.
//
//We cannot do any of this earlier because we only know all resources after
//calling Init() on all quota plugins.
func (c *Cluster) Connect() error {
	err := c.Config.Auth.Connect()
	if err != nil {
		return fmt.Errorf("failed to authenticate in cluster %s: %s", c.ID, err.Error())
	}

	for _, srv := range c.Config.Services {
		provider := c.Config.Auth.ProviderClient
		eo := c.Config.Auth.EndpointOpts

		if srv.Auth != nil {
			err := srv.Auth.Connect()
			if err != nil {
				return fmt.Errorf("failed to authenticate for service %s in cluster %s: %s", srv.Type, c.ID, err.Error())
			}
			provider = srv.Auth.ProviderClient
			eo = srv.Auth.EndpointOpts
		}

		err := c.QuotaPlugins[srv.Type].Init(provider, eo)
		if err != nil {
			return fmt.Errorf("failed to initialize service %s in cluster %s: %s", srv.Type, c.ID, err.Error())
		}
	}

	for _, capa := range c.Config.Capacitors {
		provider := c.Config.Auth.ProviderClient
		eo := c.Config.Auth.EndpointOpts

		if capa.Auth != nil {
			err := capa.Auth.Connect()
			if err != nil {
				return fmt.Errorf("failed to authenticate for capacitor %s in cluster %s: %s", capa.ID, c.ID, err.Error())
			}
			provider = capa.Auth.ProviderClient
			eo = capa.Auth.EndpointOpts
		}

		err := c.CapacityPlugins[capa.ID].Init(provider, eo)
		if err != nil {
			return fmt.Errorf("failed to initialize capacitor %s in cluster %s: %s", capa.ID, c.ID, err.Error())
		}
	}

	//load quota constraints
	if c.Config.ConstraintConfigPath != "" && c.QuotaConstraints == nil {
		var errs []error
		c.QuotaConstraints, errs = NewQuotaConstraints(c, c.Config.ConstraintConfigPath)
		if len(errs) > 0 {
			for _, err := range errs {
				logg.Error(err.Error())
			}
			return fmt.Errorf("cannot load quota constraints for cluster %s (see errors above)", c.ID)
		}
	}

	//parse low-privilege raise limits
	c.LowPrivilegeRaise.LimitsForDomains, err = c.parseLowPrivilegeRaiseLimits(
		c.Config.LowPrivilegeRaise.Limits.ForDomains, "domain")
	if err != nil {
		return fmt.Errorf("could not parse low-privilege raise limit: %s", err.Error())
	}
	c.LowPrivilegeRaise.LimitsForProjects, err = c.parseLowPrivilegeRaiseLimits(
		c.Config.LowPrivilegeRaise.Limits.ForProjects, "project")
	if err != nil {
		return fmt.Errorf("could not parse low-privilege raise limit: %s", err.Error())
	}

	//validate scaling relations
	for _, behavior := range c.Config.ResourceBehaviors {
		b := behavior.Compiled
		if b.ScalesWithResourceName == "" {
			continue
		}
		if !c.HasResource(b.ScalesWithServiceType, b.ScalesWithResourceName) {
			return fmt.Errorf(`resources matching "%s" scale with unknown resource "%s/%s"`,
				behavior.FullResourceName, b.ScalesWithServiceType, b.ScalesWithResourceName)
		}
	}

	return nil
}

var percentOfClusterRx = regexp.MustCompile(`^([0-9.]+)\s*% of cluster capacity$`)
var untilPercentOfClusterAssignedRx = regexp.MustCompile(`until ^([0-9.]+)\s*% of cluster capacity is assigned$`)

func (c Cluster) parseLowPrivilegeRaiseLimits(inputs map[string]map[string]string, scopeType string) (map[string]map[string]LowPrivilegeRaiseLimit, error) {
	result := make(map[string]map[string]LowPrivilegeRaiseLimit)
	for srvType, quotaPlugin := range c.QuotaPlugins {
		result[srvType] = make(map[string]LowPrivilegeRaiseLimit)
		for _, res := range quotaPlugin.Resources() {
			limit, exists := inputs[srvType][res.Name]
			if !exists {
				continue
			}

			match := percentOfClusterRx.FindStringSubmatch(limit)
			if match != nil {
				percent, err := strconv.ParseFloat(match[1], 64)
				if err != nil {
					return nil, err
				}
				if percent < 0 || percent > 100 {
					return nil, fmt.Errorf("value out of range: %s%%", match[1])
				}
				result[srvType][res.Name] = LowPrivilegeRaiseLimit{
					PercentOfClusterCapacity: percent,
				}
				continue
			}

			//the "until X% of cluster capacity is assigned" syntax is only allowed for domains
			if scopeType == "domain" {
				match := untilPercentOfClusterAssignedRx.FindStringSubmatch(limit)
				if match != nil {
					percent, err := strconv.ParseFloat(match[1], 64)
					if err != nil {
						return nil, err
					}
					if percent < 0 || percent > 100 {
						return nil, fmt.Errorf("value out of range: %s%%", match[1])
					}
					result[srvType][res.Name] = LowPrivilegeRaiseLimit{
						UntilPercentOfClusterCapacityAssigned: percent,
					}
					continue
				}
			}

			rawValue, err := res.Unit.Parse(limit)
			if err != nil {
				return nil, err
			}
			result[srvType][res.Name] = LowPrivilegeRaiseLimit{
				AbsoluteValue: rawValue,
			}
		}
	}
	return result, nil
}

//ProviderClient returns the gophercloud.ProviderClient for this cluster. This
//returns nil unless Connect() is called first. (This usually happens at
//program startup time for the current cluster.)
func (c *Cluster) ProviderClient() (*gophercloud.ProviderClient, gophercloud.EndpointOpts) {
	return c.Config.Auth.ProviderClient, c.Config.Auth.EndpointOpts
}

//ProviderClientForCapacitor returns the gophercloud.ProviderClient for this
//capacitor. This returns nil unless Connect() is called first. (This usually
//happens at program startup time for the current cluster.)
func (c *Cluster) ProviderClientForCapacitor(capacitorID string) (*gophercloud.ProviderClient, gophercloud.EndpointOpts) {
	for _, capa := range c.Config.Capacitors {
		if capa.ID == capacitorID && capa.Auth != nil {
			return capa.Auth.ProviderClient, capa.Auth.EndpointOpts
		}
	}
	return c.Config.Auth.ProviderClient, c.Config.Auth.EndpointOpts
}

//ProviderClientForService returns the gophercloud.ProviderClient for this
//service. This returns nil unless Connect() is called first. (This usually
//happens at program startup time for the current cluster.)
func (c *Cluster) ProviderClientForService(serviceType string) (*gophercloud.ProviderClient, gophercloud.EndpointOpts) {
	for _, srv := range c.Config.Services {
		if srv.Type == serviceType && srv.Auth != nil {
			return srv.Auth.ProviderClient, srv.Auth.EndpointOpts
		}
	}
	return c.Config.Auth.ProviderClient, c.Config.Auth.EndpointOpts
}

//HasService checks whether the given service is enabled in this cluster.
func (c *Cluster) HasService(serviceType string) bool {
	return c.QuotaPlugins[serviceType] != nil
}

//HasResource checks whether the given service is enabled in this cluster and
//whether it advertises the given resource.
func (c *Cluster) HasResource(serviceType, resourceName string) bool {
	plugin := c.QuotaPlugins[serviceType]
	if plugin == nil {
		return false
	}
	for _, res := range plugin.Resources() {
		if res.Name == resourceName {
			return true
		}
	}
	return false
}

//InfoForResource finds the plugin for the given serviceType and finds within that
//plugin the ResourceInfo for the given resourceName. If the service or
//resource does not exist, an empty ResourceInfo (with .Unit == UnitNone and
//.Category == "") is returned.
func (c *Cluster) InfoForResource(serviceType, resourceName string) limes.ResourceInfo {
	plugin := c.QuotaPlugins[serviceType]
	if plugin == nil {
		return limes.ResourceInfo{Name: resourceName, Unit: limes.UnitNone}
	}
	for _, res := range plugin.Resources() {
		if res.Name == resourceName {
			return res
		}
	}
	return limes.ResourceInfo{Name: resourceName, Unit: limes.UnitNone}
}

//InfoForService finds the plugin for the given serviceType and returns its
//ServiceInfo(), or an empty ServiceInfo (with .Area == "") when no such
//service exists in this cluster.
func (c *Cluster) InfoForService(serviceType string) limes.ServiceInfo {
	plugin := c.QuotaPlugins[serviceType]
	if plugin == nil {
		return limes.ServiceInfo{Type: serviceType}
	}
	return plugin.ServiceInfo()
}

//BehaviorForResource returns the ResourceBehavior for the given resource in
//the given scope. If no special behavior has been configured for this
//resource, or if this resource does not exist, a zero-initialized struct is
//returned.
//
//`scopeName` should be empty for cluster resources, equal to the domain name
//for domain resources, or equal to `$DOMAIN_NAME/$PROJECT_NAME` for project
//resources.
func (c *Cluster) BehaviorForResource(serviceType, resourceName, scopeName string) ResourceBehavior {
	//default behavior
	result := ResourceBehavior{
		MaxBurstMultiplier: c.Config.Bursting.MaxMultiplier,
	}

	//check for specific behavior
	fullName := serviceType + "/" + resourceName
	for _, behaviorConfig := range c.Config.ResourceBehaviors {
		behavior := behaviorConfig.Compiled
		if !behavior.FullResourceNameRx.MatchString(fullName) {
			continue
		}
		if scopeName != "" && behavior.ScopeRx != nil && !behavior.ScopeRx.MatchString(scopeName) {
			continue
		}

		// merge `behavior` into `result`
		if result.MaxBurstMultiplier > behavior.MaxBurstMultiplier {
			result.MaxBurstMultiplier = behavior.MaxBurstMultiplier
		}
		if behavior.OvercommitFactor != 0 {
			result.OvercommitFactor = behavior.OvercommitFactor
		}
		if behavior.ScalingFactor != 0 {
			result.ScalesWithServiceType = behavior.ScalesWithServiceType
			result.ScalesWithResourceName = behavior.ScalesWithResourceName
			result.ScalingFactor = behavior.ScalingFactor
		}
		if len(behavior.Annotations) > 0 && result.Annotations == nil {
			result.Annotations = make(map[string]interface{})
		}
		for k, v := range behavior.Annotations {
			result.Annotations[k] = v
		}
	}

	return result
}
<commit_msg>fix untilPercentOfClusterAssigned Regex<commit_after>/*******************************************************************************
*
* Copyright 2017 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, you may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*******************************************************************************/

package core

import (
	"fmt"
	"regexp"
	"sort"
	"strconv"

	"github.com/gophercloud/gophercloud"
	"github.com/sapcc/go-bits/logg"
	"github.com/sapcc/limes"
)

//Cluster contains all configuration and runtime information about a single
//cluster. It is passed around a lot in Limes code, mostly for the cluster ID,
//the list of enabled services, and access to the quota and capacity plugins.
type Cluster struct {
	ID                string
	Config            *ClusterConfiguration
	ServiceTypes      []string
	IsServiceShared   map[string]bool
	DiscoveryPlugin   DiscoveryPlugin
	QuotaPlugins      map[string]QuotaPlugin
	CapacityPlugins   map[string]CapacityPlugin
	Authoritative     bool
	QuotaConstraints  *QuotaConstraintSet
	LowPrivilegeRaise struct {
		LimitsForDomains  map[string]map[string]LowPrivilegeRaiseLimit
		LimitsForProjects map[string]map[string]LowPrivilegeRaiseLimit
	}
}

//NewCluster creates a new Cluster instance with the given ID and
//configuration, and also initializes all quota and capacity plugins. Errors
//will be logged when some of the requested plugins cannot be found.
func NewCluster(id string, config *ClusterConfiguration) *Cluster {
	factory, exists := discoveryPluginFactories[config.Discovery.Method]
	if !exists {
		logg.Fatal("setup for cluster %s failed: no suitable discovery plugin found", id)
	}

	c := &Cluster{
		ID:              id,
		Config:          config,
		IsServiceShared: make(map[string]bool),
		DiscoveryPlugin: factory(config.Discovery),
		QuotaPlugins:    make(map[string]QuotaPlugin),
		CapacityPlugins: make(map[string]CapacityPlugin),
		Authoritative:   config.Authoritative,
	}

	for _, srv := range config.Services {
		factory, exists := quotaPluginFactories[srv.Type]
		if !exists {
			logg.Error("skipping service %s: no suitable collector plugin found", srv.Type)
			continue
		}

		scrapeSubresources := map[string]bool{}
		for _, resName := range config.Subresources[srv.Type] {
			scrapeSubresources[resName] = true
		}

		plugin := factory(srv, scrapeSubresources)
		if plugin == nil || plugin.ServiceInfo().Type != srv.Type {
			logg.Error("skipping service %s: failed to initialize collector plugin", srv.Type)
			continue
		}

		c.ServiceTypes = append(c.ServiceTypes, srv.Type)
		c.QuotaPlugins[srv.Type] = plugin
		c.IsServiceShared[srv.Type] = srv.Shared
	}

	scrapeSubcapacities := make(map[string]map[string]bool)
	for serviceType, resourceNames := range config.Subcapacities {
		m := make(map[string]bool)
		for _, resourceName := range resourceNames {
			m[resourceName] = true
		}
		scrapeSubcapacities[serviceType] = m
	}

	for _, capa := range config.Capacitors {
		factory, exists := capacityPluginFactories[capa.ID]
		if !exists {
			logg.Error("skipping capacitor %s: no suitable collector plugin found", capa.ID)
			continue
		}
		plugin := factory(capa, scrapeSubcapacities)
		if plugin == nil || plugin.ID() != capa.ID {
			logg.Error("skipping capacitor %s: failed to initialize collector plugin", capa.ID)
			continue
		}
		c.CapacityPlugins[capa.ID] = plugin
	}

	sort.Strings(c.ServiceTypes) //determinism is useful for unit tests

	return c
}

//Connect calls Connect() on all AuthParameters for this Cluster, thus ensuring
//that all ProviderClient instances are available. It also calls Init() on all
//quota plugins.
//
//It also loads the QuotaConstraints for this cluster, if configured. The
//LowPrivilegeRaise.Limits fields are also initialized here. We also validate
//if Config.ResourceBehavior[].ScalesWith refers to existing resources.
//
//We cannot do any of this earlier because we only know all resources after
//calling Init() on all quota plugins.
func (c *Cluster) Connect() error {
	err := c.Config.Auth.Connect()
	if err != nil {
		return fmt.Errorf("failed to authenticate in cluster %s: %s", c.ID, err.Error())
	}

	for _, srv := range c.Config.Services {
		provider := c.Config.Auth.ProviderClient
		eo := c.Config.Auth.EndpointOpts

		if srv.Auth != nil {
			err := srv.Auth.Connect()
			if err != nil {
				return fmt.Errorf("failed to authenticate for service %s in cluster %s: %s", srv.Type, c.ID, err.Error())
			}
			provider = srv.Auth.ProviderClient
			eo = srv.Auth.EndpointOpts
		}

		err := c.QuotaPlugins[srv.Type].Init(provider, eo)
		if err != nil {
			return fmt.Errorf("failed to initialize service %s in cluster %s: %s", srv.Type, c.ID, err.Error())
		}
	}

	for _, capa := range c.Config.Capacitors {
		provider := c.Config.Auth.ProviderClient
		eo := c.Config.Auth.EndpointOpts

		if capa.Auth != nil {
			err := capa.Auth.Connect()
			if err != nil {
				return fmt.Errorf("failed to authenticate for capacitor %s in cluster %s: %s", capa.ID, c.ID, err.Error())
			}
			provider = capa.Auth.ProviderClient
			eo = capa.Auth.EndpointOpts
		}

		err := c.CapacityPlugins[capa.ID].Init(provider, eo)
		if err != nil {
			return fmt.Errorf("failed to initialize capacitor %s in cluster %s: %s", capa.ID, c.ID, err.Error())
		}
	}

	//load quota constraints
	if c.Config.ConstraintConfigPath != "" && c.QuotaConstraints == nil {
		var errs []error
		c.QuotaConstraints, errs = NewQuotaConstraints(c, c.Config.ConstraintConfigPath)
		if len(errs) > 0 {
			for _, err := range errs {
				logg.Error(err.Error())
			}
			return fmt.Errorf("cannot load quota constraints for cluster %s (see errors above)", c.ID)
		}
	}

	//parse low-privilege raise limits
	c.LowPrivilegeRaise.LimitsForDomains, err = c.parseLowPrivilegeRaiseLimits(
		c.Config.LowPrivilegeRaise.Limits.ForDomains, "domain")
	if err != nil {
		return fmt.Errorf("could not parse low-privilege raise limit: %s", err.Error())
	}
	c.LowPrivilegeRaise.LimitsForProjects, err = c.parseLowPrivilegeRaiseLimits(
		c.Config.LowPrivilegeRaise.Limits.ForProjects, "project")
	if err != nil {
		return fmt.Errorf("could not parse low-privilege raise limit: %s", err.Error())
	}

	//validate scaling relations
	for _, behavior := range c.Config.ResourceBehaviors {
		b := behavior.Compiled
		if b.ScalesWithResourceName == "" {
			continue
		}
		if !c.HasResource(b.ScalesWithServiceType, b.ScalesWithResourceName) {
			return fmt.Errorf(`resources matching "%s" scale with unknown resource "%s/%s"`,
				behavior.FullResourceName, b.ScalesWithServiceType, b.ScalesWithResourceName)
		}
	}

	return nil
}

var percentOfClusterRx = regexp.MustCompile(`^([0-9.]+)\s*% of cluster capacity$`)
var untilPercentOfClusterAssignedRx = regexp.MustCompile(`^until ([0-9.]+)\s*% of cluster capacity is assigned$`)

func (c Cluster) parseLowPrivilegeRaiseLimits(inputs map[string]map[string]string, scopeType string) (map[string]map[string]LowPrivilegeRaiseLimit, error) {
	result := make(map[string]map[string]LowPrivilegeRaiseLimit)
	for srvType, quotaPlugin := range c.QuotaPlugins {
		result[srvType] = make(map[string]LowPrivilegeRaiseLimit)
		for _, res := range quotaPlugin.Resources() {
			limit, exists := inputs[srvType][res.Name]
			if !exists {
				continue
			}

			match := percentOfClusterRx.FindStringSubmatch(limit)
			if match != nil {
				percent, err := strconv.ParseFloat(match[1], 64)
				if err != nil {
					return nil, err
				}
				if percent < 0 || percent > 100 {
					return nil, fmt.Errorf("value out of range: %s%%", match[1])
				}
				result[srvType][res.Name] = LowPrivilegeRaiseLimit{
					PercentOfClusterCapacity: percent,
				}
				continue
			}

			//the "until X% of cluster capacity is assigned" syntax is only allowed for domains
			if scopeType == "domain" {
				match := untilPercentOfClusterAssignedRx.FindStringSubmatch(limit)
				if match != nil {
					percent, err := strconv.ParseFloat(match[1], 64)
					if err != nil {
						return nil, err
					}
					if percent < 0 || percent > 100 {
						return nil, fmt.Errorf("value out of range: %s%%", match[1])
					}
					result[srvType][res.Name] = LowPrivilegeRaiseLimit{
						UntilPercentOfClusterCapacityAssigned: percent,
					}
					continue
				}
			}

			rawValue, err := res.Unit.Parse(limit)
			if err != nil {
				return nil, err
			}
			result[srvType][res.Name] = LowPrivilegeRaiseLimit{
				AbsoluteValue: rawValue,
			}
		}
	}
	return result, nil
}

//ProviderClient returns the gophercloud.ProviderClient for this cluster. This
//returns nil unless Connect() is called first. (This usually happens at
//program startup time for the current cluster.)
func (c *Cluster) ProviderClient() (*gophercloud.ProviderClient, gophercloud.EndpointOpts) {
	return c.Config.Auth.ProviderClient, c.Config.Auth.EndpointOpts
}

//ProviderClientForCapacitor returns the gophercloud.ProviderClient for this
//capacitor. This returns nil unless Connect() is called first. (This usually
//happens at program startup time for the current cluster.)
func (c *Cluster) ProviderClientForCapacitor(capacitorID string) (*gophercloud.ProviderClient, gophercloud.EndpointOpts) {
	for _, capa := range c.Config.Capacitors {
		if capa.ID == capacitorID && capa.Auth != nil {
			return capa.Auth.ProviderClient, capa.Auth.EndpointOpts
		}
	}
	return c.Config.Auth.ProviderClient, c.Config.Auth.EndpointOpts
}

//ProviderClientForService returns the gophercloud.ProviderClient for this
//service. This returns nil unless Connect() is called first. (This usually
//happens at program startup time for the current cluster.)
func (c *Cluster) ProviderClientForService(serviceType string) (*gophercloud.ProviderClient, gophercloud.EndpointOpts) {
	for _, srv := range c.Config.Services {
		if srv.Type == serviceType && srv.Auth != nil {
			return srv.Auth.ProviderClient, srv.Auth.EndpointOpts
		}
	}
	return c.Config.Auth.ProviderClient, c.Config.Auth.EndpointOpts
}

//HasService checks whether the given service is enabled in this cluster.
func (c *Cluster) HasService(serviceType string) bool {
	return c.QuotaPlugins[serviceType] != nil
}

//HasResource checks whether the given service is enabled in this cluster and
//whether it advertises the given resource.
func (c *Cluster) HasResource(serviceType, resourceName string) bool {
	plugin := c.QuotaPlugins[serviceType]
	if plugin == nil {
		return false
	}
	for _, res := range plugin.Resources() {
		if res.Name == resourceName {
			return true
		}
	}
	return false
}

//InfoForResource finds the plugin for the given serviceType and finds within that
//plugin the ResourceInfo for the given resourceName. If the service or
//resource does not exist, an empty ResourceInfo (with .Unit == UnitNone and
//.Category == "") is returned.
func (c *Cluster) InfoForResource(serviceType, resourceName string) limes.ResourceInfo {
	plugin := c.QuotaPlugins[serviceType]
	if plugin == nil {
		return limes.ResourceInfo{Name: resourceName, Unit: limes.UnitNone}
	}
	for _, res := range plugin.Resources() {
		if res.Name == resourceName {
			return res
		}
	}
	return limes.ResourceInfo{Name: resourceName, Unit: limes.UnitNone}
}

//InfoForService finds the plugin for the given serviceType and returns its
//ServiceInfo(), or an empty ServiceInfo (with .Area == "") when no such
//service exists in this cluster.
func (c *Cluster) InfoForService(serviceType string) limes.ServiceInfo {
	plugin := c.QuotaPlugins[serviceType]
	if plugin == nil {
		return limes.ServiceInfo{Type: serviceType}
	}
	return plugin.ServiceInfo()
}

//BehaviorForResource returns the ResourceBehavior for the given resource in
//the given scope. If no special behavior has been configured for this
//resource, or if this resource does not exist, a zero-initialized struct is
//returned.
//
//`scopeName` should be empty for cluster resources, equal to the domain name
//for domain resources, or equal to `$DOMAIN_NAME/$PROJECT_NAME` for project
//resources.
func (c *Cluster) BehaviorForResource(serviceType, resourceName, scopeName string) ResourceBehavior {
	//default behavior
	result := ResourceBehavior{
		MaxBurstMultiplier: c.Config.Bursting.MaxMultiplier,
	}

	//check for specific behavior
	fullName := serviceType + "/" + resourceName
	for _, behaviorConfig := range c.Config.ResourceBehaviors {
		behavior := behaviorConfig.Compiled
		if !behavior.FullResourceNameRx.MatchString(fullName) {
			continue
		}
		if scopeName != "" && behavior.ScopeRx != nil && !behavior.ScopeRx.MatchString(scopeName) {
			continue
		}

		// merge `behavior` into `result`
		if result.MaxBurstMultiplier > behavior.MaxBurstMultiplier {
			result.MaxBurstMultiplier = behavior.MaxBurstMultiplier
		}
		if behavior.OvercommitFactor != 0 {
			result.OvercommitFactor = behavior.OvercommitFactor
		}
		if behavior.ScalingFactor != 0 {
			result.ScalesWithServiceType = behavior.ScalesWithServiceType
			result.ScalesWithResourceName = behavior.ScalesWithResourceName
			result.ScalingFactor = behavior.ScalingFactor
		}
		if len(behavior.Annotations) > 0 && result.Annotations == nil {
			result.Annotations = make(map[string]interface{})
		}
		for k, v := range behavior.Annotations {
			result.Annotations[k] = v
		}
	}

	return result
}
<|endoftext|>


<commit_before>package moq

// moqImports are the imports all moq files get.
var moqImports = []string{}

// moqTemplate is the template for mocked code.
var moqTemplate = `// Code generated by moq; DO NOT EDIT.
// github.com/matryer/moq

package {{.PackageName}}
{{- $sourcePackagePrefix := .SourcePackagePrefix}}

import (
{{- range .Imports }}
	"{{.}}"
{{- end }}
)

{{ range $i, $obj := .Objects -}}
var (
{{- range .Methods }}
	lock{{$obj.MockName}}{{.Name}}	sync.RWMutex
{{- end }}
)

// Ensure, that {{.MockName}} does implement {{.InterfaceName}}.
// If this is not the case, regenerate this file with moq.
var _ {{$sourcePackagePrefix}}{{.InterfaceName}} = &{{.MockName}}{}

// {{.MockName}} is a mock implementation of {{$sourcePackagePrefix}}{{.InterfaceName}}.
//
//     func TestSomethingThatUses{{.InterfaceName}}(t *testing.T) {
//
//         // make and configure a mocked {{$sourcePackagePrefix}}{{.InterfaceName}}
//         mocked{{.InterfaceName}} := &{{.MockName}}{ {{ range .Methods }}
//             {{.Name}}Func: func({{ .Arglist }}) {{.ReturnArglist}} {
// 	               panic("mock out the {{.Name}} method")
//             },{{- end }}
//         }
//
//         // use mocked{{.InterfaceName}} in code that requires {{$sourcePackagePrefix}}{{.InterfaceName}}
//         // and then make assertions.
//
//     }
type {{.MockName}} struct {
{{- range .Methods }}
	// {{.Name}}Func mocks the {{.Name}} method.
	{{.Name}}Func func({{ .Arglist }}) {{.ReturnArglist}}
{{ end }}
	// calls tracks calls to the methods.
	calls struct {
{{- range .Methods }}
		// {{ .Name }} holds details about calls to the {{.Name}} method.
		{{ .Name }} []struct {
			{{- range .Params }}
			// {{ .Name | Exported }} is the {{ .Name }} argument value.
			{{ .Name | Exported }} {{ .Type }}
			{{- end }}
		}
{{- end }}
	}
}
{{ range .Methods }}
// {{.Name}} calls {{.Name}}Func.
func (mock *{{$obj.MockName}}) {{.Name}}({{.Arglist}}) {{.ReturnArglist}} {
	if mock.{{.Name}}Func == nil {
		panic("{{$obj.MockName}}.{{.Name}}Func: method is nil but {{$obj.InterfaceName}}.{{.Name}} was just called")
	}
	callInfo := struct {
		{{- range .Params }}
		{{ .Name | Exported }} {{ .Type }}
		{{- end }}
	}{
		{{- range .Params }}
		{{ .Name | Exported }}: {{ .Name }},
		{{- end }}
	}
	lock{{$obj.MockName}}{{.Name}}.Lock()
	mock.calls.{{.Name}} = append(mock.calls.{{.Name}}, callInfo)
	lock{{$obj.MockName}}{{.Name}}.Unlock()
{{- if .ReturnArglist }}
	return mock.{{.Name}}Func({{.ArgCallList}})
{{- else }}
	mock.{{.Name}}Func({{.ArgCallList}})
{{- end }}
}

// {{.Name}}Calls gets all the calls that were made to {{.Name}}.
// Check the length with:
//     len(mocked{{$obj.InterfaceName}}.{{.Name}}Calls())
func (mock *{{$obj.MockName}}) {{.Name}}Calls() []struct {
		{{- range .Params }}
		{{ .Name | Exported }} {{ .Type }}
		{{- end }}
	} {
	var calls []struct {
		{{- range .Params }}
		{{ .Name | Exported }} {{ .Type }}
		{{- end }}
	}
	lock{{$obj.MockName}}{{.Name}}.RLock()
	calls = mock.calls.{{.Name}}
	lock{{$obj.MockName}}{{.Name}}.RUnlock()
	return calls
}
{{ end -}}
{{ end -}}`
<commit_msg>Improve comment on type guard in template (#113)<commit_after>package moq

// moqImports are the imports all moq files get.
var moqImports = []string{}

// moqTemplate is the template for mocked code.
var moqTemplate = `// Code generated by moq; DO NOT EDIT.
// github.com/matryer/moq

package {{.PackageName}}
{{- $sourcePackagePrefix := .SourcePackagePrefix}}

import (
{{- range .Imports }}
	"{{.}}"
{{- end }}
)

{{ range $i, $obj := .Objects -}}
var (
{{- range .Methods }}
	lock{{$obj.MockName}}{{.Name}}	sync.RWMutex
{{- end }}
)

// Ensure, that {{.MockName}} does implement {{$sourcePackagePrefix}}{{.InterfaceName}}.
// If this is not the case, regenerate this file with moq.
var _ {{$sourcePackagePrefix}}{{.InterfaceName}} = &{{.MockName}}{}

// {{.MockName}} is a mock implementation of {{$sourcePackagePrefix}}{{.InterfaceName}}.
//
//     func TestSomethingThatUses{{.InterfaceName}}(t *testing.T) {
//
//         // make and configure a mocked {{$sourcePackagePrefix}}{{.InterfaceName}}
//         mocked{{.InterfaceName}} := &{{.MockName}}{ {{ range .Methods }}
//             {{.Name}}Func: func({{ .Arglist }}) {{.ReturnArglist}} {
// 	               panic("mock out the {{.Name}} method")
//             },{{- end }}
//         }
//
//         // use mocked{{.InterfaceName}} in code that requires {{$sourcePackagePrefix}}{{.InterfaceName}}
//         // and then make assertions.
//
//     }
type {{.MockName}} struct {
{{- range .Methods }}
	// {{.Name}}Func mocks the {{.Name}} method.
	{{.Name}}Func func({{ .Arglist }}) {{.ReturnArglist}}
{{ end }}
	// calls tracks calls to the methods.
	calls struct {
{{- range .Methods }}
		// {{ .Name }} holds details about calls to the {{.Name}} method.
		{{ .Name }} []struct {
			{{- range .Params }}
			// {{ .Name | Exported }} is the {{ .Name }} argument value.
			{{ .Name | Exported }} {{ .Type }}
			{{- end }}
		}
{{- end }}
	}
}
{{ range .Methods }}
// {{.Name}} calls {{.Name}}Func.
func (mock *{{$obj.MockName}}) {{.Name}}({{.Arglist}}) {{.ReturnArglist}} {
	if mock.{{.Name}}Func == nil {
		panic("{{$obj.MockName}}.{{.Name}}Func: method is nil but {{$obj.InterfaceName}}.{{.Name}} was just called")
	}
	callInfo := struct {
		{{- range .Params }}
		{{ .Name | Exported }} {{ .Type }}
		{{- end }}
	}{
		{{- range .Params }}
		{{ .Name | Exported }}: {{ .Name }},
		{{- end }}
	}
	lock{{$obj.MockName}}{{.Name}}.Lock()
	mock.calls.{{.Name}} = append(mock.calls.{{.Name}}, callInfo)
	lock{{$obj.MockName}}{{.Name}}.Unlock()
{{- if .ReturnArglist }}
	return mock.{{.Name}}Func({{.ArgCallList}})
{{- else }}
	mock.{{.Name}}Func({{.ArgCallList}})
{{- end }}
}

// {{.Name}}Calls gets all the calls that were made to {{.Name}}.
// Check the length with:
//     len(mocked{{$obj.InterfaceName}}.{{.Name}}Calls())
func (mock *{{$obj.MockName}}) {{.Name}}Calls() []struct {
		{{- range .Params }}
		{{ .Name | Exported }} {{ .Type }}
		{{- end }}
	} {
	var calls []struct {
		{{- range .Params }}
		{{ .Name | Exported }} {{ .Type }}
		{{- end }}
	}
	lock{{$obj.MockName}}{{.Name}}.RLock()
	calls = mock.calls.{{.Name}}
	lock{{$obj.MockName}}{{.Name}}.RUnlock()
	return calls
}
{{ end -}}
{{ end -}}`
<|endoftext|>


<commit_before>package querier

import (
	"fmt"
	"sync"
	"time"

	"github.com/cortexproject/cortex/pkg/util"
	"github.com/go-kit/kit/log/level"
	"github.com/grafana/loki/pkg/iter"
	"github.com/grafana/loki/pkg/logproto"
	"github.com/pkg/errors"
)

const (
	// keep checking connections with ingesters in duration
	checkConnectionsWithIngestersPeriod = time.Second * 5

	// the size of the channel buffer used to send tailing streams
	// back to the requesting client
	maxBufferedTailResponses = 10

	// the maximum number of entries to return in a TailResponse
	maxEntriesPerTailResponse = 100

	// the maximum number of dropped entries to keep in memory that will be sent along
	// with the next successfully pushed response. Once the dropped entries memory buffer
	// exceed this value, we start skipping dropped entries too.
	maxDroppedEntriesPerTailResponse = 1000
)

type droppedEntry struct {
	Timestamp time.Time
	Labels    string
}

// TailResponse holds response sent by tailer
type TailResponse struct {
	Streams        []logproto.Stream `json:"streams"`
	DroppedEntries []droppedEntry    `json:"dropped_entries"`
}

// Tailer manages complete lifecycle of a tail request
type Tailer struct {
	// openStreamIterator is for streams already open
	openStreamIterator iter.HeapIterator
	streamMtx          sync.Mutex // for synchronizing access to openStreamIterator

	currEntry  logproto.Entry
	currLabels string

	tailDisconnectedIngesters func([]string) (map[string]logproto.Querier_TailClient, error)

	querierTailClients    map[string]logproto.Querier_TailClient // addr -> grpc clients for tailing logs from ingesters
	querierTailClientsMtx sync.Mutex

	stopped         bool
	delayFor        time.Duration
	responseChan    chan *TailResponse
	closeErrChan    chan error
	tailMaxDuration time.Duration

	// if we are not seeing any response from ingester,
	// how long do we want to wait by going into sleep
	waitEntryThrottle time.Duration
}

func (t *Tailer) readTailClients() {
	for addr, querierTailClient := range t.querierTailClients {
		go t.readTailClient(addr, querierTailClient)
	}
}

// keeps sending oldest entry to responseChan. If channel is blocked drop the entry
// When channel is unblocked, send details of dropped entries with current entry
func (t *Tailer) loop() {
	checkConnectionTicker := time.NewTicker(checkConnectionsWithIngestersPeriod)
	defer checkConnectionTicker.Stop()

	tailMaxDurationTicker := time.NewTicker(t.tailMaxDuration)
	defer tailMaxDurationTicker.Stop()

	droppedEntries := make([]droppedEntry, 0)

	for !t.stopped {
		select {
		case <-checkConnectionTicker.C:
			// Try to reconnect dropped ingesters and connect to new ingesters
			if err := t.checkIngesterConnections(); err != nil {
				level.Error(util.Logger).Log("msg", "Error reconnecting to disconnected ingesters", "err", err)
			}
		case <-tailMaxDurationTicker.C:
			if err := t.close(); err != nil {
				level.Error(util.Logger).Log("msg", "Error closing Tailer", "err", err)
			}
			t.closeErrChan <- errors.New("reached tail max duration limit")
			return
		default:
		}

		// Read as much entries as we can (up to the max allowed) and populate the
		// tail response we'll send over the response channel
		tailResponse := new(TailResponse)
		entriesCount := 0

		for ; entriesCount < maxEntriesPerTailResponse && t.next(); entriesCount++ {
			// If the response channel channel is blocked, we drop the current entry directly
			// to save the effort
			if t.isResponseChanBlocked() {
				droppedEntries = dropEntry(droppedEntries, t.currEntry.Timestamp, t.currLabels)
				continue
			}

			tailResponse.Streams = append(tailResponse.Streams, logproto.Stream{
				Labels:  t.currLabels,
				Entries: []logproto.Entry{t.currEntry},
			})
		}

		// If all consumed entries have been dropped because the response channel is blocked
		// we should reiterate on the loop
		if len(tailResponse.Streams) == 0 && entriesCount > 0 {
			continue
		}

		// If no entry has been consumed we should ensure it's not caused by all ingesters
		// connections dropped and then throttle for a while
		if len(tailResponse.Streams) == 0 {
			if len(t.querierTailClients) == 0 {
				// All the connections to ingesters are dropped, try reconnecting or return error
				if err := t.checkIngesterConnections(); err != nil {
					level.Error(util.Logger).Log("Error reconnecting to ingesters", fmt.Sprintf("%v", err))
				} else {
					continue
				}
				if err := t.close(); err != nil {
					level.Error(util.Logger).Log("Error closing Tailer", fmt.Sprintf("%v", err))
				}
				t.closeErrChan <- errors.New("all ingesters closed the connection")
				return
			}

			time.Sleep(t.waitEntryThrottle)
			continue
		}

		// Send the tail response through the response channel without blocking.
		// Drop the entry if the response channel buffer is full.
		if len(droppedEntries) > 0 {
			tailResponse.DroppedEntries = droppedEntries
		}

		select {
		case t.responseChan <- tailResponse:
			if len(droppedEntries) > 0 {
				droppedEntries = make([]droppedEntry, 0)
			}
		default:
			droppedEntries = dropEntries(droppedEntries, tailResponse.Streams)
		}
	}
}

// Checks whether we are connected to all the ingesters to tail the logs.
// Helps in connecting to disconnected ingesters or connecting to new ingesters
func (t *Tailer) checkIngesterConnections() error {
	connectedIngestersAddr := make([]string, 0, len(t.querierTailClients))
	for addr := range t.querierTailClients {
		connectedIngestersAddr = append(connectedIngestersAddr, addr)
	}

	newConnections, err := t.tailDisconnectedIngesters(connectedIngestersAddr)
	if err != nil {
		return err
	}

	if len(newConnections) != 0 {
		for addr, tailClient := range newConnections {
			t.querierTailClients[addr] = tailClient
			go t.readTailClient(addr, tailClient)
		}
	}
	return nil
}

// removes disconnected tail client from map
func (t *Tailer) dropTailClient(addr string) {
	t.querierTailClientsMtx.Lock()
	defer t.querierTailClientsMtx.Unlock()

	delete(t.querierTailClients, addr)
}

// keeps reading streams from grpc connection with ingesters
func (t *Tailer) readTailClient(addr string, querierTailClient logproto.Querier_TailClient) {
	var resp *logproto.TailResponse
	var err error
	defer t.dropTailClient(addr)
	for {
		if t.stopped {
			if err := querierTailClient.CloseSend(); err != nil {
				level.Error(util.Logger).Log("Error closing gprc tail client", fmt.Sprintf("%v", err))
			}
			break
		}
		resp, err = querierTailClient.Recv()
		if err != nil {
			level.Error(util.Logger).Log("Error receiving response from gprc tail client", fmt.Sprintf("%v", err))
			break
		}
		t.pushTailResponseFromIngester(resp)
	}
}

// pushes new streams from ingesters synchronously
func (t *Tailer) pushTailResponseFromIngester(resp *logproto.TailResponse) {
	t.streamMtx.Lock()
	defer t.streamMtx.Unlock()

	t.openStreamIterator.Push(iter.NewStreamIterator(resp.Stream))
}

// finds oldest entry by peeking at open stream iterator.
// Response from ingester is pushed to open stream for further processing
func (t *Tailer) next() bool {
	t.streamMtx.Lock()
	defer t.streamMtx.Unlock()

	if t.openStreamIterator.Len() == 0 || !time.Now().After(t.openStreamIterator.Peek().Add(t.delayFor)) || !t.openStreamIterator.Next() {
		return false
	}

	t.currEntry = t.openStreamIterator.Entry()
	t.currLabels = t.openStreamIterator.Labels()
	return true
}

func (t *Tailer) close() error {
	t.streamMtx.Lock()
	defer t.streamMtx.Unlock()

	t.stopped = true
	return t.openStreamIterator.Close()
}

func (t *Tailer) isResponseChanBlocked() bool {
	// Thread-safety: len() and cap() on a channel are thread-safe. The cap() doesn't
	// change over the time, while len() does.
	return len(t.responseChan) == cap(t.responseChan)
}

func (t *Tailer) getResponseChan() <-chan *TailResponse {
	return t.responseChan
}

func (t *Tailer) getCloseErrorChan() <-chan error {
	return t.closeErrChan
}

func newTailer(
	delayFor time.Duration,
	querierTailClients map[string]logproto.Querier_TailClient,
	historicEntries iter.EntryIterator,
	tailDisconnectedIngesters func([]string) (map[string]logproto.Querier_TailClient, error),
	tailMaxDuration time.Duration,
	waitEntryThrottle time.Duration,
) *Tailer {
	t := Tailer{
		openStreamIterator:        iter.NewHeapIterator([]iter.EntryIterator{historicEntries}, logproto.FORWARD),
		querierTailClients:        querierTailClients,
		delayFor:                  delayFor,
		responseChan:              make(chan *TailResponse, maxBufferedTailResponses),
		closeErrChan:              make(chan error),
		tailDisconnectedIngesters: tailDisconnectedIngesters,
		tailMaxDuration:           tailMaxDuration,
		waitEntryThrottle:         waitEntryThrottle,
	}

	t.readTailClients()
	go t.loop()
	return &t
}

func dropEntry(droppedEntries []droppedEntry, timestamp time.Time, labels string) []droppedEntry {
	if len(droppedEntries) >= maxDroppedEntriesPerTailResponse {
		return droppedEntries
	}

	return append(droppedEntries, droppedEntry{timestamp, labels})
}

func dropEntries(droppedEntries []droppedEntry, streams []logproto.Stream) []droppedEntry {
	for _, stream := range streams {
		for _, entry := range stream.Entries {
			droppedEntries = dropEntry(droppedEntries, entry.Timestamp, entry.Line)
		}
	}

	return droppedEntries
}
<commit_msg>pkg/querier: fix concurrent access to querier tail clients (#944)<commit_after>package querier

import (
	"fmt"
	"sync"
	"time"

	"github.com/cortexproject/cortex/pkg/util"
	"github.com/go-kit/kit/log/level"
	"github.com/grafana/loki/pkg/iter"
	"github.com/grafana/loki/pkg/logproto"
	"github.com/pkg/errors"
)

const (
	// keep checking connections with ingesters in duration
	checkConnectionsWithIngestersPeriod = time.Second * 5

	// the size of the channel buffer used to send tailing streams
	// back to the requesting client
	maxBufferedTailResponses = 10

	// the maximum number of entries to return in a TailResponse
	maxEntriesPerTailResponse = 100

	// the maximum number of dropped entries to keep in memory that will be sent along
	// with the next successfully pushed response. Once the dropped entries memory buffer
	// exceed this value, we start skipping dropped entries too.
	maxDroppedEntriesPerTailResponse = 1000
)

type droppedEntry struct {
	Timestamp time.Time
	Labels    string
}

// TailResponse holds response sent by tailer
type TailResponse struct {
	Streams        []logproto.Stream `json:"streams"`
	DroppedEntries []droppedEntry    `json:"dropped_entries"`
}

// Tailer manages complete lifecycle of a tail request
type Tailer struct {
	// openStreamIterator is for streams already open
	openStreamIterator iter.HeapIterator
	streamMtx          sync.Mutex // for synchronizing access to openStreamIterator

	currEntry  logproto.Entry
	currLabels string

	tailDisconnectedIngesters func([]string) (map[string]logproto.Querier_TailClient, error)

	querierTailClients    map[string]logproto.Querier_TailClient // addr -> grpc clients for tailing logs from ingesters
	querierTailClientsMtx sync.RWMutex

	stopped         bool
	delayFor        time.Duration
	responseChan    chan *TailResponse
	closeErrChan    chan error
	tailMaxDuration time.Duration

	// if we are not seeing any response from ingester,
	// how long do we want to wait by going into sleep
	waitEntryThrottle time.Duration
}

func (t *Tailer) readTailClients() {
	t.querierTailClientsMtx.RLock()
	defer t.querierTailClientsMtx.RUnlock()

	for addr, querierTailClient := range t.querierTailClients {
		go t.readTailClient(addr, querierTailClient)
	}
}

// keeps sending oldest entry to responseChan. If channel is blocked drop the entry
// When channel is unblocked, send details of dropped entries with current entry
func (t *Tailer) loop() {
	checkConnectionTicker := time.NewTicker(checkConnectionsWithIngestersPeriod)
	defer checkConnectionTicker.Stop()

	tailMaxDurationTicker := time.NewTicker(t.tailMaxDuration)
	defer tailMaxDurationTicker.Stop()

	droppedEntries := make([]droppedEntry, 0)

	for !t.stopped {
		select {
		case <-checkConnectionTicker.C:
			// Try to reconnect dropped ingesters and connect to new ingesters
			if err := t.checkIngesterConnections(); err != nil {
				level.Error(util.Logger).Log("msg", "Error reconnecting to disconnected ingesters", "err", err)
			}
		case <-tailMaxDurationTicker.C:
			if err := t.close(); err != nil {
				level.Error(util.Logger).Log("msg", "Error closing Tailer", "err", err)
			}
			t.closeErrChan <- errors.New("reached tail max duration limit")
			return
		default:
		}

		// Read as much entries as we can (up to the max allowed) and populate the
		// tail response we'll send over the response channel
		tailResponse := new(TailResponse)
		entriesCount := 0

		for ; entriesCount < maxEntriesPerTailResponse && t.next(); entriesCount++ {
			// If the response channel channel is blocked, we drop the current entry directly
			// to save the effort
			if t.isResponseChanBlocked() {
				droppedEntries = dropEntry(droppedEntries, t.currEntry.Timestamp, t.currLabels)
				continue
			}

			tailResponse.Streams = append(tailResponse.Streams, logproto.Stream{
				Labels:  t.currLabels,
				Entries: []logproto.Entry{t.currEntry},
			})
		}

		// If all consumed entries have been dropped because the response channel is blocked
		// we should reiterate on the loop
		if len(tailResponse.Streams) == 0 && entriesCount > 0 {
			continue
		}

		// If no entry has been consumed we should ensure it's not caused by all ingesters
		// connections dropped and then throttle for a while
		if len(tailResponse.Streams) == 0 {
			t.querierTailClientsMtx.RLock()
			numClients := len(t.querierTailClients)
			t.querierTailClientsMtx.RUnlock()

			if numClients == 0 {
				// All the connections to ingesters are dropped, try reconnecting or return error
				if err := t.checkIngesterConnections(); err != nil {
					level.Error(util.Logger).Log("Error reconnecting to ingesters", fmt.Sprintf("%v", err))
				} else {
					continue
				}
				if err := t.close(); err != nil {
					level.Error(util.Logger).Log("Error closing Tailer", fmt.Sprintf("%v", err))
				}
				t.closeErrChan <- errors.New("all ingesters closed the connection")
				return
			}

			time.Sleep(t.waitEntryThrottle)
			continue
		}

		// Send the tail response through the response channel without blocking.
		// Drop the entry if the response channel buffer is full.
		if len(droppedEntries) > 0 {
			tailResponse.DroppedEntries = droppedEntries
		}

		select {
		case t.responseChan <- tailResponse:
			if len(droppedEntries) > 0 {
				droppedEntries = make([]droppedEntry, 0)
			}
		default:
			droppedEntries = dropEntries(droppedEntries, tailResponse.Streams)
		}
	}
}

// Checks whether we are connected to all the ingesters to tail the logs.
// Helps in connecting to disconnected ingesters or connecting to new ingesters
func (t *Tailer) checkIngesterConnections() error {
	t.querierTailClientsMtx.Lock()
	defer t.querierTailClientsMtx.Unlock()

	connectedIngestersAddr := make([]string, 0, len(t.querierTailClients))
	for addr := range t.querierTailClients {
		connectedIngestersAddr = append(connectedIngestersAddr, addr)
	}

	newConnections, err := t.tailDisconnectedIngesters(connectedIngestersAddr)
	if err != nil {
		return err
	}

	if len(newConnections) != 0 {
		for addr, tailClient := range newConnections {
			t.querierTailClients[addr] = tailClient
			go t.readTailClient(addr, tailClient)
		}
	}
	return nil
}

// removes disconnected tail client from map
func (t *Tailer) dropTailClient(addr string) {
	t.querierTailClientsMtx.Lock()
	defer t.querierTailClientsMtx.Unlock()

	delete(t.querierTailClients, addr)
}

// keeps reading streams from grpc connection with ingesters
func (t *Tailer) readTailClient(addr string, querierTailClient logproto.Querier_TailClient) {
	var resp *logproto.TailResponse
	var err error
	defer t.dropTailClient(addr)
	for {
		if t.stopped {
			if err := querierTailClient.CloseSend(); err != nil {
				level.Error(util.Logger).Log("Error closing gprc tail client", fmt.Sprintf("%v", err))
			}
			break
		}
		resp, err = querierTailClient.Recv()
		if err != nil {
			level.Error(util.Logger).Log("Error receiving response from gprc tail client", fmt.Sprintf("%v", err))
			break
		}
		t.pushTailResponseFromIngester(resp)
	}
}

// pushes new streams from ingesters synchronously
func (t *Tailer) pushTailResponseFromIngester(resp *logproto.TailResponse) {
	t.streamMtx.Lock()
	defer t.streamMtx.Unlock()

	t.openStreamIterator.Push(iter.NewStreamIterator(resp.Stream))
}

// finds oldest entry by peeking at open stream iterator.
// Response from ingester is pushed to open stream for further processing
func (t *Tailer) next() bool {
	t.streamMtx.Lock()
	defer t.streamMtx.Unlock()

	if t.openStreamIterator.Len() == 0 || !time.Now().After(t.openStreamIterator.Peek().Add(t.delayFor)) || !t.openStreamIterator.Next() {
		return false
	}

	t.currEntry = t.openStreamIterator.Entry()
	t.currLabels = t.openStreamIterator.Labels()
	return true
}

func (t *Tailer) close() error {
	t.streamMtx.Lock()
	defer t.streamMtx.Unlock()

	t.stopped = true
	return t.openStreamIterator.Close()
}

func (t *Tailer) isResponseChanBlocked() bool {
	// Thread-safety: len() and cap() on a channel are thread-safe. The cap() doesn't
	// change over the time, while len() does.
	return len(t.responseChan) == cap(t.responseChan)
}

func (t *Tailer) getResponseChan() <-chan *TailResponse {
	return t.responseChan
}

func (t *Tailer) getCloseErrorChan() <-chan error {
	return t.closeErrChan
}

func newTailer(
	delayFor time.Duration,
	querierTailClients map[string]logproto.Querier_TailClient,
	historicEntries iter.EntryIterator,
	tailDisconnectedIngesters func([]string) (map[string]logproto.Querier_TailClient, error),
	tailMaxDuration time.Duration,
	waitEntryThrottle time.Duration,
) *Tailer {
	t := Tailer{
		openStreamIterator:        iter.NewHeapIterator([]iter.EntryIterator{historicEntries}, logproto.FORWARD),
		querierTailClients:        querierTailClients,
		delayFor:                  delayFor,
		responseChan:              make(chan *TailResponse, maxBufferedTailResponses),
		closeErrChan:              make(chan error),
		tailDisconnectedIngesters: tailDisconnectedIngesters,
		tailMaxDuration:           tailMaxDuration,
		waitEntryThrottle:         waitEntryThrottle,
	}

	t.readTailClients()
	go t.loop()
	return &t
}

func dropEntry(droppedEntries []droppedEntry, timestamp time.Time, labels string) []droppedEntry {
	if len(droppedEntries) >= maxDroppedEntriesPerTailResponse {
		return droppedEntries
	}

	return append(droppedEntries, droppedEntry{timestamp, labels})
}

func dropEntries(droppedEntries []droppedEntry, streams []logproto.Stream) []droppedEntry {
	for _, stream := range streams {
		for _, entry := range stream.Entries {
			droppedEntries = dropEntry(droppedEntries, entry.Timestamp, entry.Line)
		}
	}

	return droppedEntries
}
<|endoftext|>


<commit_before>package assets

import (
	"fmt"
	"io/ioutil"
	"strings"
)

type Position int

const (
	None Position = iota
	Top
	Bottom
)

func (p Position) String() string {
	switch p {
	case None:
		return "None"
	case Top:
		return "Top"
	case Bottom:
		return "Bottom"
	}
	return fmt.Sprintf("invalid position %d", p)
}

type Type int

const (
	TypeOther Type = iota
	TypeCSS
	TypeJavascript
)

func (t Type) String() string {
	switch t {
	case TypeOther:
		return "Other"
	case TypeCSS:
		return "CSS"
	case TypeJavascript:
		return "Javascript"
	}
	return fmt.Sprintf("unknown Type %d", t)
}

func (t Type) Ext() string {
	switch t {
	case TypeCSS:
		return "css"
	case TypeJavascript:
		return "js"
	}
	return ""
}

var (
	parsers = map[string]AssetParser{}
)

type SingleAssetParser func(m *Manager, name string, options Options) ([]*Asset, error)
type AssetParser func(m *Manager, names []string, options Options) ([]*Asset, error)

type Asset struct {
	Name       string
	Type       Type
	Position   Position
	Condition  *Condition
	Attributes Attributes
	HTML       string
}

func (a *Asset) IsRemote() bool {
	name := strings.ToLower(a.Name)
	return strings.HasPrefix(name, "//") || strings.HasPrefix(name, "http://") || strings.HasPrefix(name, "https://")
}

func (a *Asset) IsHTML() bool {
	return a.HTML != ""
}

type Group struct {
	Manager *Manager
	Assets  []*Asset
	Options Options
}

func (g *Group) Names() []string {
	names := make([]string, len(g.Assets))
	for ii, a := range g.Assets {
		names[ii] = a.Name
	}
	return names

}

func (a *Asset) Code(m *Manager) (string, error) {
	f, _, err := m.Load(a.Name)
	if err != nil {
		return "", err
	}
	defer f.Close()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func Parse(m *Manager, name string, names []string, opts Options) (*Group, error) {
	parser := parsers[name]
	if parser == nil {
		return nil, fmt.Errorf("unknown asset type %s", name)
	}
	assets, err := parser(m, names, opts)
	if err != nil {
		return nil, fmt.Errorf("error parsing asset %q: %s", name, err)
	}
	if ifopt := opts["if"]; ifopt != "" {
		cond, err := ParseCondition(ifopt)
		if err != nil {
			return nil, err
		}
		if cond != nil {
			for _, v := range assets {
				if v.Condition == nil {
					v.Condition = cond
				}
			}
		}
	}
	return &Group{
		Manager: m,
		Assets:  assets,
		Options: opts,
	}, nil
}

func SingleParser(parser SingleAssetParser) AssetParser {
	return func(m *Manager, names []string, options Options) ([]*Asset, error) {
		if len(names) != 1 {
			return nil, fmt.Errorf("this asset accepts only one argument (%d given)", len(names))
		}
		return parser(m, names[0], options)
	}
}

func Register(name string, parser AssetParser) {
	parsers[name] = parser
}
<commit_msg>Cosmetic change<commit_after>package assets

import (
	"fmt"
	"io/ioutil"
	"strings"
)

type Position int

const (
	None Position = iota
	Top
	Bottom
)

func (p Position) String() string {
	switch p {
	case None:
		return "None"
	case Top:
		return "Top"
	case Bottom:
		return "Bottom"
	}
	return fmt.Sprintf("invalid position %d", p)
}

type Type int

const (
	TypeOther Type = iota
	TypeCSS
	TypeJavascript
)

func (t Type) String() string {
	switch t {
	case TypeOther:
		return "Other"
	case TypeCSS:
		return "CSS"
	case TypeJavascript:
		return "Javascript"
	}
	return fmt.Sprintf("unknown Type %d", t)
}

func (t Type) Ext() string {
	switch t {
	case TypeCSS:
		return "css"
	case TypeJavascript:
		return "js"
	}
	return ""
}

var (
	parsers = map[string]AssetParser{}
)

type SingleAssetParser func(m *Manager, name string, options Options) ([]*Asset, error)
type AssetParser func(m *Manager, names []string, options Options) ([]*Asset, error)

type Asset struct {
	Name       string
	Type       Type
	Position   Position
	Condition  *Condition
	Attributes Attributes
	HTML       string
}

func (a *Asset) IsRemote() bool {
	name := strings.ToLower(a.Name)
	return strings.HasPrefix(name, "//") || strings.HasPrefix(name, "http://") || strings.HasPrefix(name, "https://")
}

func (a *Asset) IsHTML() bool {
	return a.HTML != ""
}

func (a *Asset) Code(m *Manager) (string, error) {
	f, _, err := m.Load(a.Name)
	if err != nil {
		return "", err
	}
	defer f.Close()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

type Group struct {
	Manager *Manager
	Assets  []*Asset
	Options Options
}

func (g *Group) Names() []string {
	names := make([]string, len(g.Assets))
	for ii, a := range g.Assets {
		names[ii] = a.Name
	}
	return names

}

func Parse(m *Manager, name string, names []string, opts Options) (*Group, error) {
	parser := parsers[name]
	if parser == nil {
		return nil, fmt.Errorf("unknown asset type %s", name)
	}
	assets, err := parser(m, names, opts)
	if err != nil {
		return nil, fmt.Errorf("error parsing asset %q: %s", name, err)
	}
	if ifopt := opts["if"]; ifopt != "" {
		cond, err := ParseCondition(ifopt)
		if err != nil {
			return nil, err
		}
		if cond != nil {
			for _, v := range assets {
				if v.Condition == nil {
					v.Condition = cond
				}
			}
		}
	}
	return &Group{
		Manager: m,
		Assets:  assets,
		Options: opts,
	}, nil
}

func SingleParser(parser SingleAssetParser) AssetParser {
	return func(m *Manager, names []string, options Options) ([]*Asset, error) {
		if len(names) != 1 {
			return nil, fmt.Errorf("this asset accepts only one argument (%d given)", len(names))
		}
		return parser(m, names[0], options)
	}
}

func Register(name string, parser AssetParser) {
	parsers[name] = parser
}
<|endoftext|>


<commit_before>// The `s3` plugin for SHIELD is intended to be a back-end storage
// plugin, wrapping Amazon's Simple Storage Service (S3). It should
// be compatible with other services which emulate the S3 API, offering
// similar storage solutions for private cloud offerings (such as OpenStack
// Swift). However, this plugin has only been tested with Amazon S3.
//
// PLUGIN FEATURES
//
// This plugin implements functionality suitable for use with the following
// SHIELD Job components:
//
//  Target: no
//  Store:  yes
//
// PLUGIN CONFIGURATION
//
// The endpoint configuration passed to this plugin is used to determine
// how to connect to S3, and where to place/retrieve the data once connected.
// your endpoint JSON should look something like this:
//
//    {
//        "s3_host":             "s3.amazonaws.com", # default
//        "access_key_id":       "your-access-key-id",
//        "secret_access_key":   "your-secret-access-key",
//        "skip_ssl_validation":  false,
//        "bucket":              "bucket-name",
//        "prefix":              "/path/inside/bucket/to/place/backup/data",
//        "signature_version":   "4",  # should be 2 or 4. Defaults to 4
//        "socks5_proxy":        "",   # optionally defined SOCKS5 proxy to use for the s3 communications
//    }
//
// `prefix` will default to the empty string, and backups will be placed in the
// root of the bucket.
//
// STORE DETAILS
//
// When storing data, this plugin connects to the S3 service, and uploads the data
// into the specified bucket, using a path/filename with the following format:
//
//    <prefix>/<YYYY>/<MM>/<DD>/<HH-mm-SS>-<UUID>
//
// Upon successful storage, the plugin then returns this filename to SHIELD to use
// as the `store_key` when the data needs to be retrieved, or purged.
//
// RETRIEVE DETAILS
//
// When retrieving data, this plugin connects to the S3 service, and retrieves the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// PURGE DETAILS
//
// When purging data, this plugin connects to the S3 service, and deletes the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// DEPENDENCIES
//
// None.
//
package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"github.com/jhunt/ansi"
	minio "github.com/minio/minio-go"
	"golang.org/x/net/proxy"

	"github.com/starkandwayne/shield/plugin"
)

const (
	DefaultS3Host            = "s3.amazonaws.com"
	DefaultPrefix            = ""
	DefaultSigVersion        = "4"
	DefaultSkipSSLValidation = false
)

func validSigVersion(v string) bool {
	return v == "2" || v == "4"
}

func main() {
	p := S3Plugin{
		Name:    "S3 Backup + Storage Plugin",
		Author:  "Stark & Wayne",
		Version: "0.0.1",
		Features: plugin.PluginFeatures{
			Target: "no",
			Store:  "yes",
		},
	}

	plugin.Run(p)
}

type S3Plugin plugin.PluginInfo

type S3ConnectionInfo struct {
	Host              string
	SkipSSLValidation bool
	AccessKey         string
	SecretKey         string
	Bucket            string
	PathPrefix        string
	SignatureVersion  string
	SOCKS5Proxy       string
}

func (p S3Plugin) Meta() plugin.PluginInfo {
	return plugin.PluginInfo(p)
}

func (p S3Plugin) Validate(endpoint plugin.ShieldEndpoint) error {
	var (
		s    string
		err  error
		fail bool
	)

	s, err = endpoint.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		ansi.Printf("@R{\u2717 s3_host              %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 s3_host}              @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("access_key_id")
	if err != nil {
		ansi.Printf("@R{\u2717 access_key_id        %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 access_key_id}        @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("secret_access_key")
	if err != nil {
		ansi.Printf("@R{\u2717 secret_access_key    %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 secret_access_key}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("bucket")
	if err != nil {
		ansi.Printf("@R{\u2717 bucket               %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 bucket}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		ansi.Printf("@R{\u2717 prefix               %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 prefix}               (none)\n")
	} else {
		ansi.Printf("@G{\u2713 prefix}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("signature_version", DefaultSigVersion)
	if err != nil {
		ansi.Printf("@R{\u2717 signature_version    %s}\n", err)
		fail = true
	} else if !validSigVersion(s) {
		ansi.Printf("@R{\u2717 signature_version    Unexpected signature version '%s' found (expecting '2' or '4')}\n", s)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 signature_version}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("socks5_proxy", "")
	if err != nil {
		ansi.Printf("@R{\u2717 socks5_proxy         %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 socks5_proxy}         (no proxy will be used)\n")
	} else {
		ansi.Printf("@G{\u2713 socks5_proxy}         @C{%s}\n", s)
	}

	tf, err := endpoint.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		ansi.Printf("@R{\u2717 skip_ssl_validation  %s}\n", err)
		fail = true
	} else if tf {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{yes}, SSL will @Y{NOT} be validated\n")
	} else {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{no}, SSL @Y{WILL} be validated\n")
	}

	if fail {
		return fmt.Errorf("s3: invalid configuration")
	}
	return nil
}

func (p S3Plugin) Backup(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Restore(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Store(endpoint plugin.ShieldEndpoint) (string, error) {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return "", err
	}
	client, err := s3.Connect()
	if err != nil {
		return "", err
	}

	path := s3.genBackupPath()
	plugin.DEBUG("Storing data in %s", path)

	// FIXME: should we do something with the size of the write performed?
	_, err = client.PutObject(s3.Bucket, path, os.Stdin, "application/x-gzip")
	if err != nil {
		return "", err
	}

	return path, nil
}

func (p S3Plugin) Retrieve(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	reader, err := client.GetObject(s3.Bucket, file)
	if err != nil {
		return err
	}
	if _, err = io.Copy(os.Stdout, reader); err != nil {
		return err
	}

	err = reader.Close()
	if err != nil {
		return err
	}

	return nil
}

func (p S3Plugin) Purge(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	err = client.RemoveObject(s3.Bucket, file)
	if err != nil {
		return err
	}

	return nil
}

func getS3ConnInfo(e plugin.ShieldEndpoint) (S3ConnectionInfo, error) {
	host, err := e.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	insecure_ssl, err := e.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	key, err := e.StringValue("access_key_id")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	secret, err := e.StringValue("secret_access_key")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	bucket, err := e.StringValue("bucket")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	prefix, err := e.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	sigVer, err := e.StringValueDefault("signature_version", DefaultSigVersion)
	if !validSigVersion(sigVer) {
		return S3ConnectionInfo{}, fmt.Errorf("Invalid `signature_version` specified (`%s`). Expected `2` or `4`", sigVer)
	}

	proxy, err := e.StringValueDefault("socks5_proxy", "")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	return S3ConnectionInfo{
		Host:              host,
		SkipSSLValidation: insecure_ssl,
		AccessKey:         key,
		SecretKey:         secret,
		Bucket:            bucket,
		PathPrefix:        prefix,
		SignatureVersion:  sigVer,
		SOCKS5Proxy:       proxy,
	}, nil
}

func (s3 S3ConnectionInfo) genBackupPath() string {
	t := time.Now()
	year, mon, day := t.Date()
	hour, min, sec := t.Clock()
	uuid := plugin.GenUUID()
	return fmt.Sprintf("%s/%04d/%02d/%02d/%04d-%02d-%02d-%02d%02d%02d-%s", s3.PathPrefix, year, mon, day, year, mon, day, hour, min, sec, uuid)
}

func (s3 S3ConnectionInfo) Connect() (*minio.Client, error) {
	var s3Client *minio.Client
	var err error
	if s3.SignatureVersion == "2" {
		s3Client, err = minio.NewV2(s3.Host, s3.AccessKey, s3.SecretKey, false)
	} else {
		s3Client, err = minio.NewV4(s3.Host, s3.AccessKey, s3.SecretKey, false)
	}
	if err != nil {
		return nil, err
	}

	transport := http.DefaultTransport
	transport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: s3.SkipSSLValidation}
	if s3.SOCKS5Proxy != "" {
		dialer, err := proxy.SOCKS5("tcp", s3.SOCKS5Proxy, nil, proxy.Direct)
		if err != nil {
			fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
			os.Exit(1)
		}
		transport.(*http.Transport).Dial = dialer.Dial
	}

	s3Client.SetCustomTransport(transport)

	return s3Client, nil
}
<commit_msg>drop trailing comma in json docs<commit_after>// The `s3` plugin for SHIELD is intended to be a back-end storage
// plugin, wrapping Amazon's Simple Storage Service (S3). It should
// be compatible with other services which emulate the S3 API, offering
// similar storage solutions for private cloud offerings (such as OpenStack
// Swift). However, this plugin has only been tested with Amazon S3.
//
// PLUGIN FEATURES
//
// This plugin implements functionality suitable for use with the following
// SHIELD Job components:
//
//  Target: no
//  Store:  yes
//
// PLUGIN CONFIGURATION
//
// The endpoint configuration passed to this plugin is used to determine
// how to connect to S3, and where to place/retrieve the data once connected.
// your endpoint JSON should look something like this:
//
//    {
//        "s3_host":             "s3.amazonaws.com", # default
//        "access_key_id":       "your-access-key-id",
//        "secret_access_key":   "your-secret-access-key",
//        "skip_ssl_validation":  false,
//        "bucket":              "bucket-name",
//        "prefix":              "/path/inside/bucket/to/place/backup/data",
//        "signature_version":   "4",  # should be 2 or 4. Defaults to 4
//        "socks5_proxy":        ""    # optionally defined SOCKS5 proxy to use for the s3 communications
//    }
//
// `prefix` will default to the empty string, and backups will be placed in the
// root of the bucket.
//
// STORE DETAILS
//
// When storing data, this plugin connects to the S3 service, and uploads the data
// into the specified bucket, using a path/filename with the following format:
//
//    <prefix>/<YYYY>/<MM>/<DD>/<HH-mm-SS>-<UUID>
//
// Upon successful storage, the plugin then returns this filename to SHIELD to use
// as the `store_key` when the data needs to be retrieved, or purged.
//
// RETRIEVE DETAILS
//
// When retrieving data, this plugin connects to the S3 service, and retrieves the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// PURGE DETAILS
//
// When purging data, this plugin connects to the S3 service, and deletes the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// DEPENDENCIES
//
// None.
//
package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"github.com/jhunt/ansi"
	minio "github.com/minio/minio-go"
	"golang.org/x/net/proxy"

	"github.com/starkandwayne/shield/plugin"
)

const (
	DefaultS3Host            = "s3.amazonaws.com"
	DefaultPrefix            = ""
	DefaultSigVersion        = "4"
	DefaultSkipSSLValidation = false
)

func validSigVersion(v string) bool {
	return v == "2" || v == "4"
}

func main() {
	p := S3Plugin{
		Name:    "S3 Backup + Storage Plugin",
		Author:  "Stark & Wayne",
		Version: "0.0.1",
		Features: plugin.PluginFeatures{
			Target: "no",
			Store:  "yes",
		},
	}

	plugin.Run(p)
}

type S3Plugin plugin.PluginInfo

type S3ConnectionInfo struct {
	Host              string
	SkipSSLValidation bool
	AccessKey         string
	SecretKey         string
	Bucket            string
	PathPrefix        string
	SignatureVersion  string
	SOCKS5Proxy       string
}

func (p S3Plugin) Meta() plugin.PluginInfo {
	return plugin.PluginInfo(p)
}

func (p S3Plugin) Validate(endpoint plugin.ShieldEndpoint) error {
	var (
		s    string
		err  error
		fail bool
	)

	s, err = endpoint.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		ansi.Printf("@R{\u2717 s3_host              %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 s3_host}              @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("access_key_id")
	if err != nil {
		ansi.Printf("@R{\u2717 access_key_id        %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 access_key_id}        @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("secret_access_key")
	if err != nil {
		ansi.Printf("@R{\u2717 secret_access_key    %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 secret_access_key}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("bucket")
	if err != nil {
		ansi.Printf("@R{\u2717 bucket               %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 bucket}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		ansi.Printf("@R{\u2717 prefix               %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 prefix}               (none)\n")
	} else {
		ansi.Printf("@G{\u2713 prefix}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("signature_version", DefaultSigVersion)
	if err != nil {
		ansi.Printf("@R{\u2717 signature_version    %s}\n", err)
		fail = true
	} else if !validSigVersion(s) {
		ansi.Printf("@R{\u2717 signature_version    Unexpected signature version '%s' found (expecting '2' or '4')}\n", s)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 signature_version}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("socks5_proxy", "")
	if err != nil {
		ansi.Printf("@R{\u2717 socks5_proxy         %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 socks5_proxy}         (no proxy will be used)\n")
	} else {
		ansi.Printf("@G{\u2713 socks5_proxy}         @C{%s}\n", s)
	}

	tf, err := endpoint.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		ansi.Printf("@R{\u2717 skip_ssl_validation  %s}\n", err)
		fail = true
	} else if tf {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{yes}, SSL will @Y{NOT} be validated\n")
	} else {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{no}, SSL @Y{WILL} be validated\n")
	}

	if fail {
		return fmt.Errorf("s3: invalid configuration")
	}
	return nil
}

func (p S3Plugin) Backup(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Restore(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Store(endpoint plugin.ShieldEndpoint) (string, error) {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return "", err
	}
	client, err := s3.Connect()
	if err != nil {
		return "", err
	}

	path := s3.genBackupPath()
	plugin.DEBUG("Storing data in %s", path)

	// FIXME: should we do something with the size of the write performed?
	_, err = client.PutObject(s3.Bucket, path, os.Stdin, "application/x-gzip")
	if err != nil {
		return "", err
	}

	return path, nil
}

func (p S3Plugin) Retrieve(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	reader, err := client.GetObject(s3.Bucket, file)
	if err != nil {
		return err
	}
	if _, err = io.Copy(os.Stdout, reader); err != nil {
		return err
	}

	err = reader.Close()
	if err != nil {
		return err
	}

	return nil
}

func (p S3Plugin) Purge(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	err = client.RemoveObject(s3.Bucket, file)
	if err != nil {
		return err
	}

	return nil
}

func getS3ConnInfo(e plugin.ShieldEndpoint) (S3ConnectionInfo, error) {
	host, err := e.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	insecure_ssl, err := e.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	key, err := e.StringValue("access_key_id")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	secret, err := e.StringValue("secret_access_key")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	bucket, err := e.StringValue("bucket")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	prefix, err := e.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	sigVer, err := e.StringValueDefault("signature_version", DefaultSigVersion)
	if !validSigVersion(sigVer) {
		return S3ConnectionInfo{}, fmt.Errorf("Invalid `signature_version` specified (`%s`). Expected `2` or `4`", sigVer)
	}

	proxy, err := e.StringValueDefault("socks5_proxy", "")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	return S3ConnectionInfo{
		Host:              host,
		SkipSSLValidation: insecure_ssl,
		AccessKey:         key,
		SecretKey:         secret,
		Bucket:            bucket,
		PathPrefix:        prefix,
		SignatureVersion:  sigVer,
		SOCKS5Proxy:       proxy,
	}, nil
}

func (s3 S3ConnectionInfo) genBackupPath() string {
	t := time.Now()
	year, mon, day := t.Date()
	hour, min, sec := t.Clock()
	uuid := plugin.GenUUID()
	return fmt.Sprintf("%s/%04d/%02d/%02d/%04d-%02d-%02d-%02d%02d%02d-%s", s3.PathPrefix, year, mon, day, year, mon, day, hour, min, sec, uuid)
}

func (s3 S3ConnectionInfo) Connect() (*minio.Client, error) {
	var s3Client *minio.Client
	var err error
	if s3.SignatureVersion == "2" {
		s3Client, err = minio.NewV2(s3.Host, s3.AccessKey, s3.SecretKey, false)
	} else {
		s3Client, err = minio.NewV4(s3.Host, s3.AccessKey, s3.SecretKey, false)
	}
	if err != nil {
		return nil, err
	}

	transport := http.DefaultTransport
	transport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: s3.SkipSSLValidation}
	if s3.SOCKS5Proxy != "" {
		dialer, err := proxy.SOCKS5("tcp", s3.SOCKS5Proxy, nil, proxy.Direct)
		if err != nil {
			fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
			os.Exit(1)
		}
		transport.(*http.Transport).Dial = dialer.Dial
	}

	s3Client.SetCustomTransport(transport)

	return s3Client, nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"
)

func main() {
	// Simple static webserver:
	log.Fatal(http.ListenAndServe(":8080", http.FileServer(http.Dir("./webroot"))))
}
<commit_msg>Add some notes as to what Go router to use<commit_after>package main

import (
	"log"
	"net/http"
)

// To try out: https://github.com/pressly/chi

func main() {
	// Simple static webserver:
	log.Fatal(http.ListenAndServe(":8080", http.FileServer(http.Dir("./webroot"))))
}
<|endoftext|>