<commit_before>package zenodb

import (
	"fmt"
	"time"

	"github.com/getlantern/goexpr"
	"github.com/getlantern/zenodb/expr"
)

// sequence represents a time-ordered sequence of accumulator states in
// descending time order. The first 8 bytes are the timestamp at which the
// sequence starts, and after that each n bytes are a floating point value for
// the next interval in the sequence, where n is determined by the type of
// accumulator.
type sequence []byte

func newSequence(periodWidth int, numPeriods int) sequence {
	return make(sequence, width64bits+numPeriods*periodWidth)
}

func (seq sequence) start() time.Time {
	if seq == nil {
		return zeroTime
	}
	return timeFromBytes(seq)
}

func (seq sequence) setStart(t time.Time) {
	binaryEncoding.PutUint64(seq, uint64(t.UnixNano()))
}

func (seq sequence) numPeriods(periodWidth int) int {
	if seq == nil {
		return 0
	}
	return seq.dataLength() / periodWidth
}

// length without start time
func (seq sequence) dataLength() int {
	return len(seq) - width64bits
}

func (seq sequence) valueAtTime(t time.Time, e expr.Expr, resolution time.Duration) (float64, bool) {
	if seq == nil {
		return 0, false
	}
	start := seq.start()
	if t.After(start) {
		return 0, false
	}
	period := int(start.Sub(t) / resolution)
	return seq.ValueAt(period, e)
}

func (seq sequence) ValueAt(period int, e expr.Expr) (float64, bool) {
	if seq == nil {
		return 0, false
	}
	if period < 0 {
		return 0, false
	}
	return seq.valueAtOffset(period*e.EncodedWidth(), e)
}

func (seq sequence) dataAt(period int, e expr.Expr) ([]byte, bool) {
	if seq == nil {
		return nil, false
	}
	if period < 0 {
		return nil, false
	}
	return seq.dataAtOffset(period*e.EncodedWidth(), e)
}

func (seq sequence) valueAtOffset(offset int, e expr.Expr) (float64, bool) {
	if seq == nil {
		return 0, false
	}
	offset = offset + width64bits
	if offset >= len(seq) {
		return 0, false
	}
	val, wasSet, _ := e.Get(seq[offset:])
	return val, wasSet
}

func (seq sequence) dataAtOffset(offset int, e expr.Expr) ([]byte, bool) {
	if seq == nil {
		return nil, false
	}
	offset = offset + width64bits
	if offset >= len(seq) {
		return nil, false
	}
	return seq[offset : offset+e.EncodedWidth()], true
}

func (seq sequence) updateValueAtTime(t time.Time, resolution time.Duration, e expr.Expr, params expr.Params, metadata goexpr.Params) {
	start := seq.start()
	period := int(start.Sub(t) / resolution)
	seq.updateValueAt(period, e, params, metadata)
}

func (seq sequence) updateValueAt(period int, e expr.Expr, params expr.Params, metadata goexpr.Params) {
	seq.updateValueAtOffset(period*e.EncodedWidth(), e, params, metadata)
}

func (seq sequence) mergeValueAt(period int, e expr.Expr, other []byte, metadata goexpr.Params) {
	seq.mergeValueAtOffset(period*e.EncodedWidth(), e, other, metadata)
}

func (seq sequence) subMergeValueAt(period int, e expr.Expr, subMerge expr.SubMerge, other []byte, metadata goexpr.Params) {
	seq.subMergeValueAtOffset(period*e.EncodedWidth(), e, subMerge, other, metadata)
}

func (seq sequence) updateValueAtOffset(offset int, e expr.Expr, params expr.Params, metadata goexpr.Params) {
	offset = offset + width64bits
	e.Update(seq[offset:], params, metadata)
}

func (seq sequence) mergeValueAtOffset(offset int, e expr.Expr, other []byte, metadata goexpr.Params) {
	offset = offset + width64bits
	orig := seq[offset:]
	e.Merge(orig, orig, other, metadata)
}

func (seq sequence) subMergeValueAtOffset(offset int, e expr.Expr, subMerge expr.SubMerge, other []byte, metadata goexpr.Params) {
	offset = offset + width64bits
	orig := seq[offset:]
	subMerge(orig, other, metadata)
}

func (seq sequence) update(tsp tsparams, metadata goexpr.Params, e expr.Expr, resolution time.Duration, truncateBefore time.Time) sequence {
	ts, params := tsp.timeAndParams()
	return seq.updateValue(ts, params, metadata, e, resolution, truncateBefore)
}

func (seq sequence) updateValue(ts time.Time, params expr.Params, metadata goexpr.Params, e expr.Expr, resolution time.Duration, truncateBefore time.Time) sequence {
	periodWidth := e.EncodedWidth()

	if log.IsTraceEnabled() {
		log.Tracef("Updating sequence starting at %v to %v at %v, truncating before %v", seq.start().In(time.UTC), params, ts.In(time.UTC), truncateBefore.In(time.UTC))
	}

	if !ts.After(truncateBefore) {
		log.Trace("New value falls outside of truncation range, just truncate existing sequence")
		if len(seq) == 0 {
			return nil
		}
		return seq.truncate(periodWidth, resolution, truncateBefore)
	}

	sequenceEmpty := len(seq) == 0
	var start time.Time
	var gapPeriods int
	var maxPeriods int
	if !sequenceEmpty {
		start = seq.start()
		gapPeriods = int(ts.Sub(start) / resolution)
		maxPeriods = int(ts.Sub(truncateBefore) / resolution)
	}
	if sequenceEmpty || start.Before(truncateBefore) || gapPeriods > maxPeriods {
		log.Trace("Creating new sequence")
		out := make(sequence, width64bits+periodWidth)
		out.setStart(ts)
		out.updateValueAt(0, e, params, metadata)
		return out
	}

	if ts.After(start) {
		log.Trace("Prepending to sequence")
		numPeriods := seq.numPeriods(periodWidth) + gapPeriods
		origEnd := len(seq)
		if numPeriods > maxPeriods {
			log.Trace("Truncating existing sequence")
			numPeriods = maxPeriods
			origEnd = width64bits + periodWidth*(numPeriods-gapPeriods)
		}
		out := newSequence(periodWidth, numPeriods)
		copy(out[width64bits+gapPeriods*periodWidth:], seq[width64bits:origEnd])
		out.setStart(ts)
		out.updateValueAt(0, e, params, metadata)
		return out
	}

	log.Trace("Updating existing entry on sequence")
	out := seq
	period := int(start.Sub(ts) / resolution)
	offset := period * periodWidth
	if offset+periodWidth >= len(seq) {
		// Grow seq
		out = make(sequence, offset+width64bits+periodWidth)
		copy(out, seq)
	}
	out.updateValueAtOffset(offset, e, params, metadata)
	return out
}

func (seq sequence) merge(other sequence, e expr.Expr, resolution time.Duration, truncateBefore time.Time) sequence {
	if seq == nil {
		return other
	}
	if other == nil {
		return seq
	}

	sa := seq
	sb := other
	startA := sa.start()
	startB := sb.start()
	if startB.After(startA) {
		// Switch
		sa, startA, sb, startB = sb, startB, sa, startA
	}

	if startB.Before(truncateBefore) {
		return sa
	}

	encodedWidth := e.EncodedWidth()
	aPeriods := sa.numPeriods(encodedWidth)
	bPeriods := sb.numPeriods(encodedWidth)
	endA := startA.Add(-1 * time.Duration(aPeriods) * resolution)
	endB := startB.Add(-1 * time.Duration(bPeriods) * resolution)
	end := endB
	if endA.Before(endB) {
		end = endA
	}
	totalPeriods := int(startA.Sub(end) / resolution)

	out := make(sequence, width64bits+totalPeriods*encodedWidth)
	sout := out

	// Set start
	copy(sout, sa[:width64bits])
	sout = sout[width64bits:]
	sa = sa[width64bits:]
	sb = sb[width64bits:]

	// Handle starting window with no overlap
	leadEnd := startB
	if startB.Before(endA) {
		leadEnd = endA
	}
	leadNoOverlapPeriods := int(startA.Sub(leadEnd) / resolution)
	if leadNoOverlapPeriods > 0 {
		l := leadNoOverlapPeriods * encodedWidth
		copy(sout, sa[:l])
		sout = sout[l:]
		sa = sa[l:]
	}

	if startB.After(endA) {
		// Handle middle window with overlap
		overlapPeriods := 0
		if endB.After(endA) {
			overlapPeriods = int(startA.Sub(endB) / resolution)
		} else {
			overlapPeriods = int(startA.Sub(endA) / resolution)
		}
		overlapPeriods -= leadNoOverlapPeriods
		for i := 0; i < overlapPeriods; i++ {
			sout, sa, sb = e.Merge(sout, sa, sb, nil)
		}
	} else if startB.Before(endA) {
		// Handle gap
		gapPeriods := int(endA.Sub(startB) / resolution)
		gap := gapPeriods * encodedWidth
		sout = sout[gap:]
	}

	// Handle end window with no overlap
	if endA.Before(endB) {
		copy(sout, sa)
	} else if endB.Before(endA) {
		copy(sout, sb)
	}

	return out
}

func (seq sequence) truncate(periodWidth int, resolution time.Duration, truncateBefore time.Time) sequence {
	if seq == nil {
		return nil
	}
	maxPeriods := int(seq.start().Sub(truncateBefore) / resolution)
	if maxPeriods <= 0 {
		// Entire sequence falls outside of truncation range
		return nil
	}
	maxLength := width64bits + maxPeriods*periodWidth
	if maxLength >= len(seq) {
		return seq
	}
	return seq[:maxLength]
}

func (seq sequence) String(e expr.Expr) string {
	if seq == nil {
		return ""
	}

	values := ""

	numPeriods := seq.numPeriods(e.EncodedWidth())
	for i := 0; i < numPeriods; i++ {
		if i > 0 {
			values += " "
		}
		val, _ := seq.ValueAt(i, e)
		values += fmt.Sprint(val)
	}
	return fmt.Sprintf("%v at %v: %v", e, seq.start(), values)
}
<commit_msg>Fixed panic in seq.Truncate<commit_after>package zenodb

import (
	"fmt"
	"time"

	"github.com/getlantern/goexpr"
	"github.com/getlantern/zenodb/expr"
)

// sequence represents a time-ordered sequence of accumulator states in
// descending time order. The first 8 bytes are the timestamp at which the
// sequence starts, and after that each n bytes are a floating point value for
// the next interval in the sequence, where n is determined by the type of
// accumulator.
type sequence []byte

func newSequence(periodWidth int, numPeriods int) sequence {
	return make(sequence, width64bits+numPeriods*periodWidth)
}

func (seq sequence) start() time.Time {
	if seq == nil {
		return zeroTime
	}
	return timeFromBytes(seq)
}

func (seq sequence) setStart(t time.Time) {
	binaryEncoding.PutUint64(seq, uint64(t.UnixNano()))
}

func (seq sequence) numPeriods(periodWidth int) int {
	if seq == nil {
		return 0
	}
	return seq.dataLength() / periodWidth
}

// length without start time
func (seq sequence) dataLength() int {
	return len(seq) - width64bits
}

func (seq sequence) valueAtTime(t time.Time, e expr.Expr, resolution time.Duration) (float64, bool) {
	if seq == nil {
		return 0, false
	}
	start := seq.start()
	if t.After(start) {
		return 0, false
	}
	period := int(start.Sub(t) / resolution)
	return seq.ValueAt(period, e)
}

func (seq sequence) ValueAt(period int, e expr.Expr) (float64, bool) {
	if seq == nil {
		return 0, false
	}
	if period < 0 {
		return 0, false
	}
	return seq.valueAtOffset(period*e.EncodedWidth(), e)
}

func (seq sequence) dataAt(period int, e expr.Expr) ([]byte, bool) {
	if seq == nil {
		return nil, false
	}
	if period < 0 {
		return nil, false
	}
	return seq.dataAtOffset(period*e.EncodedWidth(), e)
}

func (seq sequence) valueAtOffset(offset int, e expr.Expr) (float64, bool) {
	if seq == nil {
		return 0, false
	}
	offset = offset + width64bits
	if offset >= len(seq) {
		return 0, false
	}
	val, wasSet, _ := e.Get(seq[offset:])
	return val, wasSet
}

func (seq sequence) dataAtOffset(offset int, e expr.Expr) ([]byte, bool) {
	if seq == nil {
		return nil, false
	}
	offset = offset + width64bits
	if offset >= len(seq) {
		return nil, false
	}
	return seq[offset : offset+e.EncodedWidth()], true
}

func (seq sequence) updateValueAtTime(t time.Time, resolution time.Duration, e expr.Expr, params expr.Params, metadata goexpr.Params) {
	start := seq.start()
	period := int(start.Sub(t) / resolution)
	seq.updateValueAt(period, e, params, metadata)
}

func (seq sequence) updateValueAt(period int, e expr.Expr, params expr.Params, metadata goexpr.Params) {
	seq.updateValueAtOffset(period*e.EncodedWidth(), e, params, metadata)
}

func (seq sequence) mergeValueAt(period int, e expr.Expr, other []byte, metadata goexpr.Params) {
	seq.mergeValueAtOffset(period*e.EncodedWidth(), e, other, metadata)
}

func (seq sequence) subMergeValueAt(period int, e expr.Expr, subMerge expr.SubMerge, other []byte, metadata goexpr.Params) {
	seq.subMergeValueAtOffset(period*e.EncodedWidth(), e, subMerge, other, metadata)
}

func (seq sequence) updateValueAtOffset(offset int, e expr.Expr, params expr.Params, metadata goexpr.Params) {
	offset = offset + width64bits
	e.Update(seq[offset:], params, metadata)
}

func (seq sequence) mergeValueAtOffset(offset int, e expr.Expr, other []byte, metadata goexpr.Params) {
	offset = offset + width64bits
	orig := seq[offset:]
	e.Merge(orig, orig, other, metadata)
}

func (seq sequence) subMergeValueAtOffset(offset int, e expr.Expr, subMerge expr.SubMerge, other []byte, metadata goexpr.Params) {
	offset = offset + width64bits
	orig := seq[offset:]
	subMerge(orig, other, metadata)
}

func (seq sequence) update(tsp tsparams, metadata goexpr.Params, e expr.Expr, resolution time.Duration, truncateBefore time.Time) sequence {
	ts, params := tsp.timeAndParams()
	return seq.updateValue(ts, params, metadata, e, resolution, truncateBefore)
}

func (seq sequence) updateValue(ts time.Time, params expr.Params, metadata goexpr.Params, e expr.Expr, resolution time.Duration, truncateBefore time.Time) sequence {
	periodWidth := e.EncodedWidth()

	if log.IsTraceEnabled() {
		log.Tracef("Updating sequence starting at %v to %v at %v, truncating before %v", seq.start().In(time.UTC), params, ts.In(time.UTC), truncateBefore.In(time.UTC))
	}

	if !ts.After(truncateBefore) {
		log.Trace("New value falls outside of truncation range, just truncate existing sequence")
		if len(seq) == 0 {
			return nil
		}
		return seq.truncate(periodWidth, resolution, truncateBefore)
	}

	sequenceEmpty := len(seq) == 0
	var start time.Time
	var gapPeriods int
	var maxPeriods int
	if !sequenceEmpty {
		start = seq.start()
		gapPeriods = int(ts.Sub(start) / resolution)
		maxPeriods = int(ts.Sub(truncateBefore) / resolution)
	}
	if sequenceEmpty || start.Before(truncateBefore) || gapPeriods > maxPeriods {
		log.Trace("Creating new sequence")
		out := make(sequence, width64bits+periodWidth)
		out.setStart(ts)
		out.updateValueAt(0, e, params, metadata)
		return out
	}

	if ts.After(start) {
		log.Trace("Prepending to sequence")
		numPeriods := seq.numPeriods(periodWidth) + gapPeriods
		origEnd := len(seq)
		if numPeriods > maxPeriods {
			log.Trace("Truncating existing sequence")
			numPeriods = maxPeriods
			origEnd = width64bits + periodWidth*(numPeriods-gapPeriods)
		}
		out := newSequence(periodWidth, numPeriods)
		copy(out[width64bits+gapPeriods*periodWidth:], seq[width64bits:origEnd])
		out.setStart(ts)
		out.updateValueAt(0, e, params, metadata)
		return out
	}

	log.Trace("Updating existing entry on sequence")
	out := seq
	period := int(start.Sub(ts) / resolution)
	offset := period * periodWidth
	if offset+periodWidth >= len(seq) {
		// Grow seq
		out = make(sequence, offset+width64bits+periodWidth)
		copy(out, seq)
	}
	out.updateValueAtOffset(offset, e, params, metadata)
	return out
}

func (seq sequence) merge(other sequence, e expr.Expr, resolution time.Duration, truncateBefore time.Time) sequence {
	if seq == nil {
		return other
	}
	if other == nil {
		return seq
	}

	sa := seq
	sb := other
	startA := sa.start()
	startB := sb.start()
	if startB.After(startA) {
		// Switch
		sa, startA, sb, startB = sb, startB, sa, startA
	}

	if startB.Before(truncateBefore) {
		return sa
	}

	encodedWidth := e.EncodedWidth()
	aPeriods := sa.numPeriods(encodedWidth)
	bPeriods := sb.numPeriods(encodedWidth)
	endA := startA.Add(-1 * time.Duration(aPeriods) * resolution)
	endB := startB.Add(-1 * time.Duration(bPeriods) * resolution)
	end := endB
	if endA.Before(endB) {
		end = endA
	}
	totalPeriods := int(startA.Sub(end) / resolution)

	out := make(sequence, width64bits+totalPeriods*encodedWidth)
	sout := out

	// Set start
	copy(sout, sa[:width64bits])
	sout = sout[width64bits:]
	sa = sa[width64bits:]
	sb = sb[width64bits:]

	// Handle starting window with no overlap
	leadEnd := startB
	if startB.Before(endA) {
		leadEnd = endA
	}
	leadNoOverlapPeriods := int(startA.Sub(leadEnd) / resolution)
	if leadNoOverlapPeriods > 0 {
		l := leadNoOverlapPeriods * encodedWidth
		copy(sout, sa[:l])
		sout = sout[l:]
		sa = sa[l:]
	}

	if startB.After(endA) {
		// Handle middle window with overlap
		overlapPeriods := 0
		if endB.After(endA) {
			overlapPeriods = int(startA.Sub(endB) / resolution)
		} else {
			overlapPeriods = int(startA.Sub(endA) / resolution)
		}
		overlapPeriods -= leadNoOverlapPeriods
		for i := 0; i < overlapPeriods; i++ {
			sout, sa, sb = e.Merge(sout, sa, sb, nil)
		}
	} else if startB.Before(endA) {
		// Handle gap
		gapPeriods := int(endA.Sub(startB) / resolution)
		gap := gapPeriods * encodedWidth
		sout = sout[gap:]
	}

	// Handle end window with no overlap
	if endA.Before(endB) {
		copy(sout, sa)
	} else if endB.Before(endA) {
		copy(sout, sb)
	}

	return out
}

func (seq sequence) truncate(periodWidth int, resolution time.Duration, truncateBefore time.Time) sequence {
	if len(seq) == 0 {
		return nil
	}
	maxPeriods := int(seq.start().Sub(truncateBefore) / resolution)
	if maxPeriods <= 0 {
		// Entire sequence falls outside of truncation range
		return nil
	}
	maxLength := width64bits + maxPeriods*periodWidth
	if maxLength >= len(seq) {
		return seq
	}
	return seq[:maxLength]
}

func (seq sequence) String(e expr.Expr) string {
	if seq == nil {
		return ""
	}

	values := ""

	numPeriods := seq.numPeriods(e.EncodedWidth())
	for i := 0; i < numPeriods; i++ {
		if i > 0 {
			values += " "
		}
		val, _ := seq.ValueAt(i, e)
		values += fmt.Sprint(val)
	}
	return fmt.Sprintf("%v at %v: %v", e, seq.start(), values)
}
<|endoftext|>


<commit_before>/*
Open Source Initiative OSI - The MIT License (MIT):Licensing

The MIT License (MIT)
Copyright (c) 2013 Ralph Caraveo (deckarep@gmail.com)

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

// Package mapset implements a simple and generic set collection.
// Items stored within it are unordered and unique
// It supports typical set operations: membership testing, intersection, union, difference and symmetric difference

package mapset

import "fmt"
import "strings"

// The primary type that represents a set
type Set struct {
	set map[interface{}]_placeHolder
}

type _placeHolder struct{}

// Creates and returns a pointer to an empty set.
func NewSet() *Set {
	return &Set{make(map[interface{}]_placeHolder)}
}

// Adds an item to the current set if it doesn't already exist in the set.
func (set *Set) Add(i interface{}) bool {
	_, found := set.set[i]
	set.set[i] = _placeHolder{}
	return !found //False if it existed already
}

// Determines if a given item is already in the set.
func (set *Set) Contains(i interface{}) bool {
	if _, found := set.set[i]; found {
		return found //true if it existed already
	}
	return false
}

// Determines if every item in the other set is in this set.
func (set *Set) IsSubset(other *Set) bool {
	for key, _ := range set.set {
		if !other.Contains(key) {
			return false
		}
	}
	return true
}

// Determines if every item of this set is in the other set.
func (set *Set) IsSuperset(other *Set) bool {
	for key, _ := range other.set {
		if !set.Contains(key) {
			return false
		}
	}
	return true
}

// Returns a new set with all items in both sets.
func (set *Set) Union(other *Set) *Set {
	if set != nil && other != nil {
		unionedSet := NewSet()

		for key, _ := range set.set {
			unionedSet.Add(key)
		}
		for key, _ := range other.set {
			unionedSet.Add(key)
		}
		return unionedSet
	}
	return nil
}

// Returns a new set with items that exist only in both sets.
func (set *Set) Intersect(other *Set) *Set {
	if set != nil && other != nil {
		intersectedSet := NewSet()
		var smallerSet *Set = nil
		var largerSet *Set = nil

		//figure out the smaller of the two sets and loop on that one as an optimization.
		if set.Size() < other.Size() {
			smallerSet = set
			largerSet = other
		} else {
			smallerSet = other
			largerSet = set
		}

		for key, _ := range smallerSet.set {
			if largerSet.Contains(key) {
				intersectedSet.Add(key)
			}
		}
		return intersectedSet
	}
	return nil
}

// Returns a new set with items in the current set but not in the other set
func (set *Set) Difference(other *Set) *Set {
	if set != nil && other != nil {
		differencedSet := NewSet()

		for key, _ := range set.set {
			if !other.Contains(key) {
				differencedSet.Add(key)
			}
		}

		return differencedSet
	}
	return nil
}

// Returns a new set with items in the current set or the other set but not in both.
func (set *Set) SymmetricDifference(other *Set) *Set {
	if set != nil && other != nil {
		aDiff := set.Difference(other)
		bDiff := other.Difference(set)

		symDifferencedSet := aDiff.Union(bDiff)

		return symDifferencedSet
	}
	return nil
}

// Clears the entire set to be the empty set.
func (set *Set) Clear() {
	set.set = make(map[interface{}]_placeHolder)
}

// Allows the removal of a single item in the set.
func (set *Set) Remove(i interface{}) {
	delete(set.set, i)
}

// Size returns the how many items are currently in the set.
func (set *Set) Size() int {
	return len(set.set)
}

// Equal determines if two sets are equal to each other.
// If they both are the same size and have the same items they are considered equal.
// Order of items is not relevent for sets to be equal.
func (set *Set) Equal(other *Set) bool {
	if set != nil && other != nil {
		if !(set.Size() == other.Size()) {
			return false
		} else {
			for key, _ := range set.set {
				if !other.Contains(key) {
					return false
				}
			}
			return true
		}
	}
	return false
}

// Provides a convenient string representation of the current state of the set.
func (set *Set) String() string {
	items := make([]string, 0, len(set.set))

	for key := range set.set {
		items = append(items, fmt.Sprintf("%v", key))
	}
	return fmt.Sprintf("Set{%s}", strings.Join(items, ", "))
}
<commit_msg>The Set type is now simply an alias instead of a struct with a map, it basically is a map now<commit_after>/*
Open Source Initiative OSI - The MIT License (MIT):Licensing

The MIT License (MIT)
Copyright (c) 2013 Ralph Caraveo (deckarep@gmail.com)

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

// Package mapset implements a simple and generic set collection.
// Items stored within it are unordered and unique
// It supports typical set operations: membership testing, intersection, union, difference and symmetric difference

package mapset

import "fmt"
import "strings"

// The primary type that represents a set
type Set map[interface{}]_placeHolder

type _placeHolder struct{}

// Creates and returns a pointer to an empty set.
func NewSet() Set {
	return make(map[interface{}]_placeHolder)
}

// Adds an item to the current set if it doesn't already exist in the set.
func (set Set) Add(i interface{}) bool {
	_, found := set[i]
	set[i] = _placeHolder{}
	return !found //False if it existed already
}

// Determines if a given item is already in the set.
func (set Set) Contains(i interface{}) bool {
	if _, found := set[i]; found {
		return found //true if it existed already
	}
	return false
}

// Determines if every item in the other set is in this set.
func (set Set) IsSubset(other Set) bool {
	for key, _ := range set {
		if !other.Contains(key) {
			return false
		}
	}
	return true
}

// Determines if every item of this set is in the other set.
func (set Set) IsSuperset(other Set) bool {
	for key, _ := range other {
		if !set.Contains(key) {
			return false
		}
	}
	return true
}

// Returns a new set with all items in both sets.
func (set Set) Union(other Set) Set {
	if set != nil && other != nil {
		unionedSet := NewSet()

		for key, _ := range set {
			unionedSet.Add(key)
		}
		for key, _ := range other {
			unionedSet.Add(key)
		}
		return unionedSet
	}
	return nil
}

// Returns a new set with items that exist only in both sets.
func (set Set) Intersect(other Set) Set {
	if set != nil && other != nil {
		intersectedSet := NewSet()
		var smallerSet Set = nil
		var largerSet Set = nil

		//figure out the smaller of the two sets and loop on that one as an optimization.
		if set.Size() < other.Size() {
			smallerSet = set
			largerSet = other
		} else {
			smallerSet = other
			largerSet = set
		}

		for key, _ := range smallerSet {
			if largerSet.Contains(key) {
				intersectedSet.Add(key)
			}
		}
		return intersectedSet
	}
	return nil
}

// Returns a new set with items in the current set but not in the other set
func (set Set) Difference(other Set) Set {
	if set != nil && other != nil {
		differencedSet := NewSet()

		for key, _ := range set {
			if !other.Contains(key) {
				differencedSet.Add(key)
			}
		}

		return differencedSet
	}
	return nil
}

// Returns a new set with items in the current set or the other set but not in both.
func (set Set) SymmetricDifference(other Set) Set {
	if set != nil && other != nil {
		aDiff := set.Difference(other)
		bDiff := other.Difference(set)

		symDifferencedSet := aDiff.Union(bDiff)

		return symDifferencedSet
	}
	return nil
}

// Clears the entire set to be the empty set.
func (set *Set) Clear() {
	*set = make(map[interface{}]_placeHolder)
}

// Allows the removal of a single item in the set.
func (set Set) Remove(i interface{}) {
	delete(set, i)
}

// Size returns the how many items are currently in the set.
func (set Set) Size() int {
	return len(set)
}

// Equal determines if two sets are equal to each other.
// If they both are the same size and have the same items they are considered equal.
// Order of items is not relevent for sets to be equal.
func (set Set) Equal(other Set) bool {
	if set != nil && other != nil {
		if !(set.Size() == other.Size()) {
			return false
		} else {
			for key, _ := range set {
				if !other.Contains(key) {
					return false
				}
			}
			return true
		}
	}
	return false
}

// Provides a convenient string representation of the current state of the set.
func (set Set) String() string {
	items := make([]string, 0, len(set))

	for key := range set {
		items = append(items, fmt.Sprintf("%v", key))
	}
	return fmt.Sprintf("Set{%s}", strings.Join(items, ", "))
}
<|endoftext|>


<commit_before>// +build linux

package linux

import (
	"github.com/mackerelio/go-osstat/memory"
	"github.com/mackerelio/golib/logging"
	"github.com/mackerelio/mackerel-agent/metrics"
)

/*
MemoryGenerator collect memory usage

`memory.{metric}`: using memory size[KiB] retrieved from /proc/meminfo

metric = "total", "free", "buffers", "cached", "active", "inactive", "swap_cached", "swap_total", "swap_free"

Metrics "used" is calculated here like (total - free - buffers - cached) for ease.
This calculation may be going to be done in server side in the future.

graph: stacks `memory.{metric}`
*/
type MemoryGenerator struct {
}

var memoryLogger = logging.GetLogger("metrics.memory")

// Generate memory values
func (g *MemoryGenerator) Generate() (metrics.Values, error) {
	mem, err := memory.Get()
	if err != nil {
		memoryLogger.Errorf("failed to get memory statistics: %s", err)
		return nil, err
	}

	ret := map[string]float64{
		"memory.total":       float64(mem.Total),
		"memory.used":        float64(mem.Used),
		"memory.active":      float64(mem.Active),
		"memory.inactive":    float64(mem.Inactive),
		"memory.swap_total":  float64(mem.SwapTotal),
		"memory.swap_cached": float64(mem.SwapCached),
		"memory.swap_free":   float64(mem.SwapFree),
	}

	if mem.MemAvailableEnabled {
		ret["memory.mem_available"] = float64(mem.Available)
	} else {
		ret["memory.buffers"] = float64(mem.Buffers)
		ret["memory.cached"] = float64(mem.Cacned)
		ret["memory.free"] = float64(mem.Free)
	}

	return metrics.Values(ret), nil
}
<commit_msg>fix typo<commit_after>// +build linux

package linux

import (
	"github.com/mackerelio/go-osstat/memory"
	"github.com/mackerelio/golib/logging"
	"github.com/mackerelio/mackerel-agent/metrics"
)

/*
MemoryGenerator collect memory usage

`memory.{metric}`: using memory size[KiB] retrieved from /proc/meminfo

metric = "total", "free", "buffers", "cached", "active", "inactive", "swap_cached", "swap_total", "swap_free"

Metrics "used" is calculated here like (total - free - buffers - cached) for ease.
This calculation may be going to be done in server side in the future.

graph: stacks `memory.{metric}`
*/
type MemoryGenerator struct {
}

var memoryLogger = logging.GetLogger("metrics.memory")

// Generate memory values
func (g *MemoryGenerator) Generate() (metrics.Values, error) {
	mem, err := memory.Get()
	if err != nil {
		memoryLogger.Errorf("failed to get memory statistics: %s", err)
		return nil, err
	}

	ret := map[string]float64{
		"memory.total":       float64(mem.Total),
		"memory.used":        float64(mem.Used),
		"memory.active":      float64(mem.Active),
		"memory.inactive":    float64(mem.Inactive),
		"memory.swap_total":  float64(mem.SwapTotal),
		"memory.swap_cached": float64(mem.SwapCached),
		"memory.swap_free":   float64(mem.SwapFree),
	}

	if mem.MemAvailableEnabled {
		ret["memory.mem_available"] = float64(mem.Available)
	} else {
		ret["memory.buffers"] = float64(mem.Buffers)
		ret["memory.cached"] = float64(mem.Cached)
		ret["memory.free"] = float64(mem.Free)
	}

	return metrics.Values(ret), nil
}
<|endoftext|>


<commit_before>package middleware

import (
	"github.com/gin-gonic/gin"
	"strings"
	"time"

	u "github.com/techjanitor/pram-get/utils"
)

var analyticsWorker *requestWorker

type requestWorker struct {
	queue chan RequestType
}

// requesttype holds the data we want to capture
type RequestType struct {
	Ib        string
	Ip        string
	User      uint
	Path      string
	ItemKey   string
	ItemValue string
	Status    int
	Latency   time.Duration
	Cached    bool
}

func init() {
	// make worker channel
	analyticsWorker = &requestWorker{
		make(chan RequestType, 64),
	}

	// Get Database handle
	db, err := u.GetDb()
	if err != nil {
		return
	}

	go func() {

		// prepare query for analytics table
		ps1, err := db.Prepare("INSERT INTO analytics (ib_id, user_id, request_ip, request_path, request_status, request_latency, request_itemkey, request_itemvalue, request_cached, request_time) VALUES (?,?,?,?,?,?,?,?,?,NOW())")
		if err != nil {
			return
		}

		// range through tasks channel
		for request := range analyticsWorker.queue {

			// input data
			_, err = ps1.Exec(request.Ib, request.User, request.Ip, request.Path, request.Status, request.Latency, request.ItemKey, request.ItemValue, request.Cached)
			if err != nil {
				return
			}

		}

	}()

}

// Analytics will log requests in the database
func Analytics() gin.HandlerFunc {
	return func(c *gin.Context) {
		req := c.Request
		// get request path
		path := req.URL.Path
		// Start timer
		start := time.Now()

		// Process request
		c.Next()

		// Stop timer
		end := time.Now()
		// get request latency
		latency := end.Sub(start)

		// get userdata from session middleware
		userdata := c.MustGet("userdata").(u.User)

		// get cached state from cache middleware
		cached := c.MustGet("cached").(bool)

		// get the ib
		ib := c.Param("ib")

		// Trim leading / from path and split
		params := strings.Split(strings.Trim(path, "/"), "/")

		// Make key from path
		key := itemKey{}
		key.generateKey(params...)

		if !skipKey(params[0]) {

			// set our data
			request := RequestType{
				Ib:        ib,
				Ip:        c.ClientIP(),
				User:      userdata.Id,
				Path:      path,
				Status:    c.Writer.Status(),
				Latency:   latency,
				ItemKey:   key.Key,
				ItemValue: key.Value,
				Cached:    cached,
			}

			analyticsWorker.queue <- request

		}

	}
}

type itemKey struct {
	Key   string
	Value string
}

// Will take the params from the request and turn them into a key
func (r *itemKey) generateKey(params ...string) {

	switch {
	case len(params) <= 2:
		r.Key = params[0]
		r.Value = "1"
	case len(params) >= 3:
		r.Key = params[0]
		r.Value = params[2]
	}

	return

}

// Check if key should be expired
func skipKey(key string) bool {

	keyList := map[string]bool{
		"taginfo":  true,
		"tags":     true,
		"post":     true,
		"tagtypes": true,
		"pram":     true,
	}

	if keyList[strings.ToLower(key)] {
		return true
	}

	return false

}
<commit_msg>test async analytics submission<commit_after>package middleware

import (
	"github.com/gin-gonic/gin"
	"strings"
	"time"

	u "github.com/techjanitor/pram-get/utils"
)

var analyticsWorker *requestWorker

type requestWorker struct {
	queue chan RequestType
}

// requesttype holds the data we want to capture
type RequestType struct {
	Ib        string
	Ip        string
	User      uint
	Path      string
	ItemKey   string
	ItemValue string
	Status    int
	Latency   time.Duration
	Cached    bool
}

func init() {
	// make worker channel
	analyticsWorker = &requestWorker{
		make(chan RequestType, 64),
	}

	go func() {

		// Get Database handle
		db, err := u.GetDb()
		if err != nil {
			return
		}

		// prepare query for analytics table
		ps1, err := db.Prepare("INSERT INTO analytics (ib_id, user_id, request_ip, request_path, request_status, request_latency, request_itemkey, request_itemvalue, request_cached, request_time) VALUES (?,?,?,?,?,?,?,?,?,NOW())")
		if err != nil {
			return
		}

		// range through tasks channel
		for request := range analyticsWorker.queue {

			// input data
			_, err = ps1.Exec(request.Ib, request.User, request.Ip, request.Path, request.Status, request.Latency, request.ItemKey, request.ItemValue, request.Cached)
			if err != nil {
				return
			}

		}

	}()

}

// Analytics will log requests in the database
func Analytics() gin.HandlerFunc {
	return func(c *gin.Context) {
		req := c.Request
		// get request path
		path := req.URL.Path
		// Start timer
		start := time.Now()

		// Process request
		c.Next()

		// Stop timer
		end := time.Now()
		// get request latency
		latency := end.Sub(start)

		// get userdata from session middleware
		userdata := c.MustGet("userdata").(u.User)

		// get cached state from cache middleware
		cached := c.MustGet("cached").(bool)

		// get the ib
		ib := c.Param("ib")

		// Trim leading / from path and split
		params := strings.Split(strings.Trim(path, "/"), "/")

		// Make key from path
		key := itemKey{}
		key.generateKey(params...)

		if !skipKey(params[0]) {

			// set our data
			request := RequestType{
				Ib:        ib,
				Ip:        c.ClientIP(),
				User:      userdata.Id,
				Path:      path,
				Status:    c.Writer.Status(),
				Latency:   latency,
				ItemKey:   key.Key,
				ItemValue: key.Value,
				Cached:    cached,
			}

			analyticsWorker.queue <- request

		}

	}
}

type itemKey struct {
	Key   string
	Value string
}

// Will take the params from the request and turn them into a key
func (r *itemKey) generateKey(params ...string) {

	switch {
	case len(params) <= 2:
		r.Key = params[0]
		r.Value = "1"
	case len(params) >= 3:
		r.Key = params[0]
		r.Value = params[2]
	}

	return

}

// Check if key should be expired
func skipKey(key string) bool {

	keyList := map[string]bool{
		"taginfo":  true,
		"tags":     true,
		"post":     true,
		"tagtypes": true,
		"pram":     true,
	}

	if keyList[strings.ToLower(key)] {
		return true
	}

	return false

}
<|endoftext|>


<commit_before>// Package gzip provides a simple middleware layer that performs
// gzip compression on the response.
package gzip

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strings"

	"github.com/mholt/caddy/middleware"
)

// Gzip is a middleware type which gzips HTTP responses. It is
// imperative that any handler which writes to a gzipped response
// specifies the Content-Type, otherwise some clients will assume
// application/x-gzip and try to download a file.
type Gzip struct {
	Next    middleware.Handler
	Configs []Config
}

// Config holds the configuration for Gzip middleware
type Config struct {
	RequestFilters  []RequestFilter
	ResponseFilters []ResponseFilter
	Level           int // Compression level
}

// ServeHTTP serves a gzipped response if the client supports it.
func (g Gzip) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		return g.Next.ServeHTTP(w, r)
	}

outer:
	for _, c := range g.Configs {

		// Check request filters to determine if gzipping is permitted for this request
		for _, filter := range c.RequestFilters {
			if !filter.ShouldCompress(r) {
				continue outer
			}
		}

		// Delete this header so gzipping is not repeated later in the chain
		r.Header.Del("Accept-Encoding")

		// gzipWriter modifies underlying writer at init,
		// use a discard writer instead to leave ResponseWriter in
		// original form.
		gzipWriter, err := newWriter(c, ioutil.Discard)
		if err != nil {
			// should not happen
			return http.StatusInternalServerError, err
		}
		defer gzipWriter.Close()
		gz := &gzipResponseWriter{Writer: gzipWriter, ResponseWriter: w}

		var rw http.ResponseWriter
		// if no response filter is used
		if len(c.ResponseFilters) == 0 {
			// replace discard writer with ResponseWriter
			gzipWriter.Reset(w)
			rw = gz
		} else {
			// wrap gzip writer with ResponseFilterWriter
			rw = NewResponseFilterWriter(c.ResponseFilters, gz)
		}

		// Any response in forward middleware will now be compressed
		status, err := g.Next.ServeHTTP(rw, r)

		// If there was an error that remained unhandled, we need
		// to send something back before gzipWriter gets closed at
		// the return of this method!
		if status >= 400 {
			gz.Header().Set("Content-Type", "text/plain") // very necessary
			gz.WriteHeader(status)
			fmt.Fprintf(gz, "%d %s", status, http.StatusText(status))
			return 0, err
		}
		return status, err
	}

	// no matching filter
	return g.Next.ServeHTTP(w, r)
}

// newWriter create a new Gzip Writer based on the compression level.
// If the level is valid (i.e. between 1 and 9), it uses the level.
// Otherwise, it uses default compression level.
func newWriter(c Config, w io.Writer) (*gzip.Writer, error) {
	if c.Level >= gzip.BestSpeed && c.Level <= gzip.BestCompression {
		return gzip.NewWriterLevel(w, c.Level)
	}
	return gzip.NewWriter(w), nil
}

// gzipResponeWriter wraps the underlying Write method
// with a gzip.Writer to compress the output.
type gzipResponseWriter struct {
	io.Writer
	http.ResponseWriter
	statusCodeWritten bool
}

// WriteHeader wraps the underlying WriteHeader method to prevent
// problems with conflicting headers from proxied backends. For
// example, a backend system that calculates Content-Length would
// be wrong because it doesn't know it's being gzipped.
func (w *gzipResponseWriter) WriteHeader(code int) {
	w.Header().Del("Content-Length")
	w.Header().Set("Content-Encoding", "gzip")
	w.Header().Add("Vary", "Accept-Encoding")
	w.ResponseWriter.WriteHeader(code)
	w.statusCodeWritten = true
}

// Write wraps the underlying Write method to do compression.
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
	if w.Header().Get("Content-Type") == "" {
		w.Header().Set("Content-Type", http.DetectContentType(b))
	}
	if !w.statusCodeWritten {
		w.WriteHeader(http.StatusOK)
	}
	n, err := w.Writer.Write(b)
	return n, err
}

// Hijack implements http.Hijacker. It simply wraps the underlying
// ResponseWriter's Hijack method if there is one, or returns an error.
func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := w.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, fmt.Errorf("not a Hijacker")
}

// CloseNotify implements http.CloseNotifier.
// It just inherits the underlying ResponseWriter's CloseNotify method.
func (w *gzipResponseWriter) CloseNotify() <-chan bool {
	if cn, ok := w.ResponseWriter.(http.CloseNotifier); ok {
		return cn.CloseNotify()
	}
	panic("not a CloseNotifier")
}
<commit_msg>gzipResponseWriter should also be a Flusher<commit_after>// Package gzip provides a simple middleware layer that performs
// gzip compression on the response.
package gzip

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strings"

	"github.com/mholt/caddy/middleware"
)

// Gzip is a middleware type which gzips HTTP responses. It is
// imperative that any handler which writes to a gzipped response
// specifies the Content-Type, otherwise some clients will assume
// application/x-gzip and try to download a file.
type Gzip struct {
	Next    middleware.Handler
	Configs []Config
}

// Config holds the configuration for Gzip middleware
type Config struct {
	RequestFilters  []RequestFilter
	ResponseFilters []ResponseFilter
	Level           int // Compression level
}

// ServeHTTP serves a gzipped response if the client supports it.
func (g Gzip) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		return g.Next.ServeHTTP(w, r)
	}

outer:
	for _, c := range g.Configs {

		// Check request filters to determine if gzipping is permitted for this request
		for _, filter := range c.RequestFilters {
			if !filter.ShouldCompress(r) {
				continue outer
			}
		}

		// Delete this header so gzipping is not repeated later in the chain
		r.Header.Del("Accept-Encoding")

		// gzipWriter modifies underlying writer at init,
		// use a discard writer instead to leave ResponseWriter in
		// original form.
		gzipWriter, err := newWriter(c, ioutil.Discard)
		if err != nil {
			// should not happen
			return http.StatusInternalServerError, err
		}
		defer gzipWriter.Close()
		gz := &gzipResponseWriter{Writer: gzipWriter, ResponseWriter: w}

		var rw http.ResponseWriter
		// if no response filter is used
		if len(c.ResponseFilters) == 0 {
			// replace discard writer with ResponseWriter
			gzipWriter.Reset(w)
			rw = gz
		} else {
			// wrap gzip writer with ResponseFilterWriter
			rw = NewResponseFilterWriter(c.ResponseFilters, gz)
		}

		// Any response in forward middleware will now be compressed
		status, err := g.Next.ServeHTTP(rw, r)

		// If there was an error that remained unhandled, we need
		// to send something back before gzipWriter gets closed at
		// the return of this method!
		if status >= 400 {
			gz.Header().Set("Content-Type", "text/plain") // very necessary
			gz.WriteHeader(status)
			fmt.Fprintf(gz, "%d %s", status, http.StatusText(status))
			return 0, err
		}
		return status, err
	}

	// no matching filter
	return g.Next.ServeHTTP(w, r)
}

// newWriter create a new Gzip Writer based on the compression level.
// If the level is valid (i.e. between 1 and 9), it uses the level.
// Otherwise, it uses default compression level.
func newWriter(c Config, w io.Writer) (*gzip.Writer, error) {
	if c.Level >= gzip.BestSpeed && c.Level <= gzip.BestCompression {
		return gzip.NewWriterLevel(w, c.Level)
	}
	return gzip.NewWriter(w), nil
}

// gzipResponeWriter wraps the underlying Write method
// with a gzip.Writer to compress the output.
type gzipResponseWriter struct {
	io.Writer
	http.ResponseWriter
	statusCodeWritten bool
}

// WriteHeader wraps the underlying WriteHeader method to prevent
// problems with conflicting headers from proxied backends. For
// example, a backend system that calculates Content-Length would
// be wrong because it doesn't know it's being gzipped.
func (w *gzipResponseWriter) WriteHeader(code int) {
	w.Header().Del("Content-Length")
	w.Header().Set("Content-Encoding", "gzip")
	w.Header().Add("Vary", "Accept-Encoding")
	w.ResponseWriter.WriteHeader(code)
	w.statusCodeWritten = true
}

// Write wraps the underlying Write method to do compression.
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
	if w.Header().Get("Content-Type") == "" {
		w.Header().Set("Content-Type", http.DetectContentType(b))
	}
	if !w.statusCodeWritten {
		w.WriteHeader(http.StatusOK)
	}
	n, err := w.Writer.Write(b)
	return n, err
}

// Hijack implements http.Hijacker. It simply wraps the underlying
// ResponseWriter's Hijack method if there is one, or returns an error.
func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := w.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, fmt.Errorf("not a Hijacker")
}

// Flush implements http.Flusher. It simply wraps the underlying
// ResponseWriter's Flush method if there is one, or panics.
func (w *gzipResponseWriter) Flush() {
	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	} else {
		panic("not a Flusher") // should be recovered at the beginning of middleware stack
	}
}

// CloseNotify implements http.CloseNotifier.
// It just inherits the underlying ResponseWriter's CloseNotify method.
func (w *gzipResponseWriter) CloseNotify() <-chan bool {
	if cn, ok := w.ResponseWriter.(http.CloseNotifier); ok {
		return cn.CloseNotify()
	}
	panic("not a CloseNotifier")
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package manager

import (
	"sort"
	"sync"

	"v.io/v23/context"
	"v.io/v23/naming"
	"v.io/x/ref/runtime/internal/flow/conn"
)

// ConnCache is a cache of Conns keyed by (protocol, address) and (routingID).
// Multiple goroutines can invoke methods on the ConnCache simultaneously.
// TODO(suharshs): We should periodically look for closed connections and remove them.
type ConnCache struct {
	mu            *sync.Mutex
	cond          *sync.Cond
	addrCache     map[string]*connEntry           // keyed by (protocol, address)
	ridCache      map[naming.RoutingID]*connEntry // keyed by naming.RoutingID
	started       map[string]bool                 // keyed by (protocol, address)
	unmappedConns map[*connEntry]bool             // list of connEntries replaced by other entries
}

type connEntry struct {
	conn          *conn.Conn
	rid           naming.RoutingID
	addrKey       string
	blessingNames []string
	proxy         bool
}

func NewConnCache() *ConnCache {
	mu := &sync.Mutex{}
	cond := sync.NewCond(mu)
	return &ConnCache{
		mu:            mu,
		cond:          cond,
		addrCache:     make(map[string]*connEntry),
		ridCache:      make(map[naming.RoutingID]*connEntry),
		started:       make(map[string]bool),
		unmappedConns: make(map[*connEntry]bool),
	}
}

// Insert adds conn to the cache, keyed by both (protocol, address) and (routingID)
// An error will be returned iff the cache has been closed.
func (c *ConnCache) Insert(conn *conn.Conn, protocol, address string, proxy bool) error {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return NewErrCacheClosed(nil)
	}
	ep := conn.RemoteEndpoint()
	k := key(protocol, address)
	entry := &connEntry{
		conn:    conn,
		rid:     ep.RoutingID(),
		addrKey: k,
		proxy:   proxy,
	}
	if !entry.proxy {
		entry.blessingNames = ep.BlessingNames()
	}
	if old := c.ridCache[entry.rid]; old != nil {
		c.unmappedConns[old] = true
	}
	c.addrCache[k] = entry
	c.ridCache[entry.rid] = entry
	return nil
}

// InsertWithRoutingID adds conn to the cache keyed only by conn's RoutingID.
func (c *ConnCache) InsertWithRoutingID(conn *conn.Conn, proxy bool) error {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return NewErrCacheClosed(nil)
	}
	ep := conn.RemoteEndpoint()
	entry := &connEntry{
		conn:  conn,
		rid:   ep.RoutingID(),
		proxy: proxy,
	}
	if !entry.proxy {
		entry.blessingNames = ep.BlessingNames()
	}
	if old := c.ridCache[entry.rid]; old != nil {
		c.unmappedConns[old] = true
	}
	c.ridCache[entry.rid] = entry
	return nil
}

func (c *ConnCache) Find(protocol, address string, rid naming.RoutingID, blessingNames []string) (*conn.Conn, error) {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return nil, NewErrCacheClosed(nil)
	}
	if rid != naming.NullRoutingID {
		if entry := c.removeUndialable(c.ridCache[rid], blessingNames); entry != nil {
			return entry, nil
		}
	}
	k := key(protocol, address)
	return c.removeUndialable(c.addrCache[k], blessingNames), nil
}

// ReservedFind returns a Conn based on the input remoteEndpoint.
// nil is returned if there is no such Conn.
//
// ReservedFind will return an error iff the cache is closed.
// If no error is returned, the caller is required to call Unreserve, to avoid
// deadlock. The (protocol, address) provided to Unreserve must be the same as
// the arguments provided to ReservedFind.
// All new ReservedFind calls for the (protocol, address) will Block
// until the corresponding Unreserve call is made.
func (c *ConnCache) ReservedFind(protocol, address string, rid naming.RoutingID, blessingNames []string) (*conn.Conn, error) {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return nil, NewErrCacheClosed(nil)
	}
	if rid != naming.NullRoutingID {
		if entry := c.removeUndialable(c.ridCache[rid], blessingNames); entry != nil {
			return entry, nil
		}
	}
	k := key(protocol, address)
	for c.started[k] {
		c.cond.Wait()
	}
	c.started[k] = true
	return c.removeUndialable(c.addrCache[k], blessingNames), nil
}

// Unreserve marks the status of the (protocol, address) as no longer started, and
// broadcasts waiting threads.
func (c *ConnCache) Unreserve(protocol, address string) {
	c.mu.Lock()
	delete(c.started, key(protocol, address))
	c.cond.Broadcast()
	c.mu.Unlock()
}

// Close marks the ConnCache as closed and closes all Conns in the cache.
func (c *ConnCache) Close(ctx *context.T) {
	defer c.mu.Unlock()
	c.mu.Lock()
	c.addrCache, c.started = nil, nil
	err := NewErrCacheClosed(ctx)
	for _, d := range c.ridCache {
		d.conn.Close(ctx, err)
	}
	for d := range c.unmappedConns {
		d.conn.Close(ctx, err)
	}
}

// removeUndialable filters connections that are closed, lameducked, or non-proxied
// connections whose blessings dont match blessingNames.
func (c *ConnCache) removeUndialable(e *connEntry, blessingNames []string) *conn.Conn {
	if e == nil {
		return nil
	}
	if status := e.conn.Status(); status >= conn.Closing || e.conn.RemoteLameDuck() {
		delete(c.addrCache, e.addrKey)
		delete(c.ridCache, e.rid)
		if status < conn.Closing {
			c.unmappedConns[e] = true
		}
		return nil
	}
	if !e.proxy && len(e.blessingNames) > 0 && !matchBlessings(e.blessingNames, blessingNames) {
		return nil
	}
	return e.conn
}

// KillConnections will close and remove num LRU Conns in the cache.
// If connections are already closed they will be removed from the cache.
// This is useful when the manager is approaching system FD limits.
// If num is greater than the number of connections in the cache, all cached
// connections will be closed and removed.
// KillConnections returns an error iff the cache is closed.
func (c *ConnCache) KillConnections(ctx *context.T, num int) error {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return NewErrCacheClosed(ctx)
	}
	err := NewErrConnKilledToFreeResources(ctx)
	pq := make(connEntries, 0, len(c.ridCache))
	for _, e := range c.ridCache {
		if c.removeUndialable(e, nil) == nil {
			continue
		}
		if e.conn.IsEncapsulated() {
			// Killing a proxied connection doesn't save us any FD resources, just memory.
			continue
		}
		pq = append(pq, e)
	}
	for d := range c.unmappedConns {
		if status := d.conn.Status(); status == conn.Closed {
			delete(c.unmappedConns, d)
			continue
		}
		if d.conn.IsEncapsulated() {
			continue
		}
		pq = append(pq, d)
	}
	sort.Sort(pq)
	for i := 0; i < num; i++ {
		d := pq[i]
		d.conn.Close(ctx, err)
		delete(c.addrCache, d.addrKey)
		delete(c.ridCache, d.rid)
		delete(c.unmappedConns, d)
	}
	return nil
}

func (c *ConnCache) EnterLameDuckMode(ctx *context.T) {
	c.mu.Lock()
	n := len(c.ridCache) + len(c.unmappedConns)
	conns := make([]*conn.Conn, 0, n)
	for _, e := range c.ridCache {
		conns = append(conns, e.conn)
	}
	for d := range c.unmappedConns {
		conns = append(conns, d.conn)
	}
	c.mu.Unlock()

	waitfor := make([]chan struct{}, 0, n)
	for _, c := range conns {
		waitfor = append(waitfor, c.EnterLameDuck(ctx))
	}
	for _, w := range waitfor {
		<-w
	}
}

// matchBlessings return true if the intersection of a and b is not empty.
func matchBlessings(a, b []string) bool {
	m := make(map[string]bool, len(a))
	for _, i := range a {
		m[i] = true
	}
	for _, j := range b {
		if m[j] {
			return true
		}
	}
	return false
}

// TODO(suharshs): If sorting the connections becomes too slow, switch to
// container/heap instead of sorting all the connections.
type connEntries []*connEntry

func (c connEntries) Len() int {
	return len(c)
}

func (c connEntries) Less(i, j int) bool {
	return c[i].conn.LastUsedTime().Before(c[j].conn.LastUsedTime())
}

func (c connEntries) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

func key(protocol, address string) string {
	return protocol + "," + address
}
<commit_msg>flow/manager: Fix conncache reservedfind to respect when cache is closed.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package manager

import (
	"sort"
	"sync"

	"v.io/v23/context"
	"v.io/v23/naming"
	"v.io/x/ref/runtime/internal/flow/conn"
)

// ConnCache is a cache of Conns keyed by (protocol, address) and (routingID).
// Multiple goroutines can invoke methods on the ConnCache simultaneously.
// TODO(suharshs): We should periodically look for closed connections and remove them.
type ConnCache struct {
	mu            *sync.Mutex
	cond          *sync.Cond
	addrCache     map[string]*connEntry           // keyed by (protocol, address)
	ridCache      map[naming.RoutingID]*connEntry // keyed by naming.RoutingID
	started       map[string]bool                 // keyed by (protocol, address)
	unmappedConns map[*connEntry]bool             // list of connEntries replaced by other entries
}

type connEntry struct {
	conn          *conn.Conn
	rid           naming.RoutingID
	addrKey       string
	blessingNames []string
	proxy         bool
}

func NewConnCache() *ConnCache {
	mu := &sync.Mutex{}
	cond := sync.NewCond(mu)
	return &ConnCache{
		mu:            mu,
		cond:          cond,
		addrCache:     make(map[string]*connEntry),
		ridCache:      make(map[naming.RoutingID]*connEntry),
		started:       make(map[string]bool),
		unmappedConns: make(map[*connEntry]bool),
	}
}

// Insert adds conn to the cache, keyed by both (protocol, address) and (routingID)
// An error will be returned iff the cache has been closed.
func (c *ConnCache) Insert(conn *conn.Conn, protocol, address string, proxy bool) error {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return NewErrCacheClosed(nil)
	}
	ep := conn.RemoteEndpoint()
	k := key(protocol, address)
	entry := &connEntry{
		conn:    conn,
		rid:     ep.RoutingID(),
		addrKey: k,
		proxy:   proxy,
	}
	if !entry.proxy {
		entry.blessingNames = ep.BlessingNames()
	}
	if old := c.ridCache[entry.rid]; old != nil {
		c.unmappedConns[old] = true
	}
	c.addrCache[k] = entry
	c.ridCache[entry.rid] = entry
	return nil
}

// InsertWithRoutingID adds conn to the cache keyed only by conn's RoutingID.
func (c *ConnCache) InsertWithRoutingID(conn *conn.Conn, proxy bool) error {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return NewErrCacheClosed(nil)
	}
	ep := conn.RemoteEndpoint()
	entry := &connEntry{
		conn:  conn,
		rid:   ep.RoutingID(),
		proxy: proxy,
	}
	if !entry.proxy {
		entry.blessingNames = ep.BlessingNames()
	}
	if old := c.ridCache[entry.rid]; old != nil {
		c.unmappedConns[old] = true
	}
	c.ridCache[entry.rid] = entry
	return nil
}

func (c *ConnCache) Find(protocol, address string, rid naming.RoutingID, blessingNames []string) (*conn.Conn, error) {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return nil, NewErrCacheClosed(nil)
	}
	if rid != naming.NullRoutingID {
		if entry := c.removeUndialable(c.ridCache[rid], blessingNames); entry != nil {
			return entry, nil
		}
	}
	k := key(protocol, address)
	return c.removeUndialable(c.addrCache[k], blessingNames), nil
}

// ReservedFind returns a Conn based on the input remoteEndpoint.
// nil is returned if there is no such Conn.
//
// ReservedFind will return an error iff the cache is closed.
// If no error is returned, the caller is required to call Unreserve, to avoid
// deadlock. The (protocol, address) provided to Unreserve must be the same as
// the arguments provided to ReservedFind.
// All new ReservedFind calls for the (protocol, address) will Block
// until the corresponding Unreserve call is made.
func (c *ConnCache) ReservedFind(protocol, address string, rid naming.RoutingID, blessingNames []string) (*conn.Conn, error) {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return nil, NewErrCacheClosed(nil)
	}
	if rid != naming.NullRoutingID {
		if entry := c.removeUndialable(c.ridCache[rid], blessingNames); entry != nil {
			return entry, nil
		}
	}
	k := key(protocol, address)
	for c.started[k] {
		c.cond.Wait()
		if c.addrCache == nil {
			return nil, NewErrCacheClosed(nil)
		}
	}
	c.started[k] = true
	return c.removeUndialable(c.addrCache[k], blessingNames), nil
}

// Unreserve marks the status of the (protocol, address) as no longer started, and
// broadcasts waiting threads.
func (c *ConnCache) Unreserve(protocol, address string) {
	c.mu.Lock()
	delete(c.started, key(protocol, address))
	c.cond.Broadcast()
	c.mu.Unlock()
}

// Close marks the ConnCache as closed and closes all Conns in the cache.
func (c *ConnCache) Close(ctx *context.T) {
	defer c.mu.Unlock()
	c.mu.Lock()
	c.addrCache, c.started = nil, nil
	err := NewErrCacheClosed(ctx)
	for _, d := range c.ridCache {
		d.conn.Close(ctx, err)
	}
	for d := range c.unmappedConns {
		d.conn.Close(ctx, err)
	}
}

// removeUndialable filters connections that are closed, lameducked, or non-proxied
// connections whose blessings dont match blessingNames.
func (c *ConnCache) removeUndialable(e *connEntry, blessingNames []string) *conn.Conn {
	if e == nil {
		return nil
	}
	if status := e.conn.Status(); status >= conn.Closing || e.conn.RemoteLameDuck() {
		delete(c.addrCache, e.addrKey)
		delete(c.ridCache, e.rid)
		if status < conn.Closing {
			c.unmappedConns[e] = true
		}
		return nil
	}
	if !e.proxy && len(e.blessingNames) > 0 && !matchBlessings(e.blessingNames, blessingNames) {
		return nil
	}
	return e.conn
}

// KillConnections will close and remove num LRU Conns in the cache.
// If connections are already closed they will be removed from the cache.
// This is useful when the manager is approaching system FD limits.
// If num is greater than the number of connections in the cache, all cached
// connections will be closed and removed.
// KillConnections returns an error iff the cache is closed.
func (c *ConnCache) KillConnections(ctx *context.T, num int) error {
	defer c.mu.Unlock()
	c.mu.Lock()
	if c.addrCache == nil {
		return NewErrCacheClosed(ctx)
	}
	err := NewErrConnKilledToFreeResources(ctx)
	pq := make(connEntries, 0, len(c.ridCache))
	for _, e := range c.ridCache {
		if c.removeUndialable(e, nil) == nil {
			continue
		}
		if e.conn.IsEncapsulated() {
			// Killing a proxied connection doesn't save us any FD resources, just memory.
			continue
		}
		pq = append(pq, e)
	}
	for d := range c.unmappedConns {
		if status := d.conn.Status(); status == conn.Closed {
			delete(c.unmappedConns, d)
			continue
		}
		if d.conn.IsEncapsulated() {
			continue
		}
		pq = append(pq, d)
	}
	sort.Sort(pq)
	for i := 0; i < num; i++ {
		d := pq[i]
		d.conn.Close(ctx, err)
		delete(c.addrCache, d.addrKey)
		delete(c.ridCache, d.rid)
		delete(c.unmappedConns, d)
	}
	return nil
}

func (c *ConnCache) EnterLameDuckMode(ctx *context.T) {
	c.mu.Lock()
	n := len(c.ridCache) + len(c.unmappedConns)
	conns := make([]*conn.Conn, 0, n)
	for _, e := range c.ridCache {
		conns = append(conns, e.conn)
	}
	for d := range c.unmappedConns {
		conns = append(conns, d.conn)
	}
	c.mu.Unlock()

	waitfor := make([]chan struct{}, 0, n)
	for _, c := range conns {
		waitfor = append(waitfor, c.EnterLameDuck(ctx))
	}
	for _, w := range waitfor {
		<-w
	}
}

// matchBlessings return true if the intersection of a and b is not empty.
func matchBlessings(a, b []string) bool {
	m := make(map[string]bool, len(a))
	for _, i := range a {
		m[i] = true
	}
	for _, j := range b {
		if m[j] {
			return true
		}
	}
	return false
}

// TODO(suharshs): If sorting the connections becomes too slow, switch to
// container/heap instead of sorting all the connections.
type connEntries []*connEntry

func (c connEntries) Len() int {
	return len(c)
}

func (c connEntries) Less(i, j int) bool {
	return c[i].conn.LastUsedTime().Before(c[j].conn.LastUsedTime())
}

func (c connEntries) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

func key(protocol, address string) string {
	return protocol + "," + address
}
<|endoftext|>


<commit_before>package connection

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/deweysasser/locksmith/data"
	"github.com/deweysasser/locksmith/output"
	"sync"
	"time"
)

type AWSConnection struct {
	Type, Profile string
}

func (a *AWSConnection) String() string {
	return fmt.Sprintf("aws://%s", a.Profile)
}

type userMap map[string]*iam.User

func (a *AWSConnection) Fetch() (keys <-chan data.Key, accounts <-chan data.Account) {
	output.Debug("Fetching from aws", a.Profile)
	cKeys := make(chan data.Key)
	cAccounts := make(chan data.Account)

	sharedCredentials := credentials.NewSharedCredentials("", a.Profile)
	region := aws.String("us-east-1")

	wg := sync.WaitGroup{}

	wg.Add(1)
	go func() {
		defer wg.Done()

		if sess, err := session.NewSession(&aws.Config{
			Region:      region,
			Credentials: sharedCredentials,
		}); err == nil {

			e := ec2.New(sess)

			if arn, err := a.fetchAccountInfo(sess, cAccounts); err == nil {

				wg.Add(1)
				go func() {
					defer wg.Done()
					usermap := a.fetchAccounts(sess, cAccounts, cKeys)
					a.fetchAccessKeys(sess, cAccounts, cKeys, usermap)
				}()
				wg.Add(1)
				go func() {
					defer wg.Done()
					a.fetchAccountInfo(sess, cAccounts)
				}()
				if dro, err := e.DescribeRegions(&ec2.DescribeRegionsInput{}); err == nil {
					for _, r := range dro.Regions {
						wg.Add(1)
						go func(region *string) {
							defer wg.Done()
							keymap := a.fetchKeyPairs(arn, region, sharedCredentials, cKeys, cAccounts)
							a.fetchInstances(region, sharedCredentials, cAccounts, keymap)
						}(r.RegionName)

					}
				} else {
					output.Error(a, "failed to lookup EC2 regions")
				}
			}
		}

	}()

	go func() {
		wg.Wait()
		defer close(cKeys)
		defer close(cAccounts)
	}()

	return cKeys, cAccounts
}

func (a *AWSConnection) fetchAccountInfo(sess *session.Session, accounts chan<- data.Account) (data.AWSAccountID, error) {
	s := sts.New(sess)
	i := iam.New(sess)

	if out, err := s.GetCallerIdentity(&sts.GetCallerIdentityInput{}); err == nil {
		arn := data.AWSAccountID(*out.Account)
		if iout, err := i.ListAccountAliases(&iam.ListAccountAliasesInput{}); err == nil {
			aliases := make([]string, len(iout.AccountAliases))
			for _, a := range iout.AccountAliases {
				aliases = append(aliases, *a)
			}
			accounts <- data.NewAWSAccount(arn, a.Id(), []data.KeyBinding{}, aliases...)
		} else {
			output.Warn("Failed to get accont alias for", a.Profile, "account", arn)
			accounts <- data.NewAWSAccount(arn, a.Id(), []data.KeyBinding{})
		}
		return arn, nil
	} else {
		output.Error("Failed to get account identity for", a.Profile)
		return data.AWSAccountID(""), err
	}
}

func (a *AWSConnection) fetchAccessKeys(sess *session.Session, accounts chan<- data.Account, keys chan<- data.Key, usermap userMap) {

	i := iam.New(sess)

	if lako, err := i.ListAccessKeys(&iam.ListAccessKeysInput{}); err == nil {
		for _, md := range lako.AccessKeyMetadata {
			output.Debug("Found acces key", *md.AccessKeyId)
			userName := *md.UserName
			keys <- data.NewAwsKey(*md.AccessKeyId, *md.CreateDate, *md.Status == "Active", userName, *usermap[userName].Arn)
			//accounts <- data.NewIAMAccount(usermap[*md.UserName], a.Id(), md)
			accounts <- data.NewIAMAccountFromKey(md, usermap[userName], a.Id())
		}
	} else {
		output.Error(a, "failed to list IAM users")
	}
}

func (a *AWSConnection) fetchAccounts(sess *session.Session, accounts chan<- data.Account, keys chan<- data.Key) userMap {
	usermap := make(userMap)
	i := iam.New(sess)

	if r, err := i.ListUsers(&iam.ListUsersInput{}); err == nil {
		for _, user := range r.Users {
			usermap[*user.UserName] = user
		}
	} else {
		output.Error(a, "failed to list IAM users")
	}

	return usermap
}

func (a *AWSConnection) fetchInstances(region *string, sharedCredentials *credentials.Credentials, cAccounts chan data.Account, keymap map[string]data.ID) {
	output.Debug(a, "fetching instances from", *region)
	if sess, err := session.NewSession(&aws.Config{
		Region:      region,
		Credentials: sharedCredentials,
	}); err == nil {
		e := ec2.New(sess)
		//ec2.DescribeInstancesOutput{}.Reservations[0].Instances[0].k
		if dio, err := e.DescribeInstances(&ec2.DescribeInstancesInput{}); err == nil {
			output.Debug(a, *region, "reservations:", len(dio.Reservations))
			for _, res := range dio.Reservations {
				output.Debug(a, *region, "instances:", len(res.Instances))
				for _, instance := range res.Instances {
					keyID := keymap[*instance.KeyName]
					keys := []data.KeyBinding{
						{
							KeyID:    keyID,
							Location: data.INSTANCE_ROOT_CREDENTIALS,
						},
					}

					acct := data.NewAWSInstanceAccount(instance, a.Id(), keys)
					output.Debug("Found instance account", acct)
					cAccounts <- acct
				}
			}
		} else {
			output.Error(a, "Failed to fetch instances")
		}

	} else {
		output.Error(a, "Failed to connect")
	}
}

func (a *AWSConnection) fetchKeyPairs(arn data.AWSAccountID, region *string, sharedCredentials *credentials.Credentials, cKeys chan data.Key, cAccounts chan data.Account) (keymap map[string]data.ID) {
	output.Debug(a, "fetching key pairs from", *region)
	keymap = make(map[string]data.ID)

	if sess, err := session.NewSession(&aws.Config{
		Region:      region,
		Credentials: sharedCredentials,
	}); err == nil {
		e := ec2.New(sess)

		if out, err := e.DescribeKeyPairs(&ec2.DescribeKeyPairsInput{}); err == nil {
			bindings := make([]data.KeyBinding, 0)
			for _, p := range out.KeyPairs {
				fp := p.KeyFingerprint
				name := p.KeyName
				key := data.NewSSHKeyFromFingerprint(*name, time.Now(), data.ID(*fp))
				cKeys <- key
				bindings = append(bindings, data.KeyBinding{KeyID: data.ID(*fp), Name: *name})
				keymap[*name] = data.ID(*fp)
			}

			acct := data.NewAWSAccount(arn, a.Id(), bindings)

			cAccounts <- acct
		} else {
			output.Warn(a.String()+":", "Failed to find key pairs:", err)
		}

	} else {
		output.Error("Failed to connect to", a)
	}

	return
}

func (a *AWSConnection) Id() data.ID {
	return data.ID(a.Profile)
}
<commit_msg>Better error messages in aws connection<commit_after>package connection

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/deweysasser/locksmith/data"
	"github.com/deweysasser/locksmith/output"
	"sync"
	"time"
)

type AWSConnection struct {
	Type, Profile string
}

func (a *AWSConnection) String() string {
	return fmt.Sprintf("aws://%s", a.Profile)
}

type userMap map[string]*iam.User

func (a *AWSConnection) Fetch() (keys <-chan data.Key, accounts <-chan data.Account) {
	output.Debug("Fetching from aws", a.Profile)
	cKeys := make(chan data.Key)
	cAccounts := make(chan data.Account)

	sharedCredentials := credentials.NewSharedCredentials("", a.Profile)
	region := aws.String("us-east-1")

	wg := sync.WaitGroup{}

	wg.Add(1)
	go func() {
		defer wg.Done()

		if sess, err := session.NewSession(&aws.Config{
			Region:      region,
			Credentials: sharedCredentials,
		}); err == nil {

			e := ec2.New(sess)

			// TODO:  make this code still fetch keys and instances even if fetching the account ARN fails
			if arn, err := a.fetchAccountInfo(sess, cAccounts); err == nil {

				wg.Add(1)
				go func() {
					defer wg.Done()
					usermap := a.fetchAccounts(sess, cAccounts, cKeys)
					a.fetchAccessKeys(sess, cAccounts, cKeys, usermap)
				}()
				wg.Add(1)
				go func() {
					defer wg.Done()
					a.fetchAccountInfo(sess, cAccounts)
				}()
				if dro, err := e.DescribeRegions(&ec2.DescribeRegionsInput{}); err == nil {
					for _, r := range dro.Regions {
						wg.Add(1)
						go func(region *string) {
							defer wg.Done()
							keymap := a.fetchKeyPairs(arn, region, sharedCredentials, cKeys, cAccounts)
							a.fetchInstances(region, sharedCredentials, cAccounts, keymap)
						}(r.RegionName)
					}
				} else {
					output.Error(a, "failed to lookup EC2 regions")
				}
			} else {
				output.Error("Failed to get account identity for", a.Profile, ":", err)
			}
		} else {
			output.Error("Failed to create AWS BOTO session:", err)
		}

	}()

	go func() {
		wg.Wait()
		defer close(cKeys)
		defer close(cAccounts)
	}()

	return cKeys, cAccounts
}

func (a *AWSConnection) fetchAccountInfo(sess *session.Session, accounts chan<- data.Account) (data.AWSAccountID, error) {
	s := sts.New(sess)
	i := iam.New(sess)

	if out, err := s.GetCallerIdentity(&sts.GetCallerIdentityInput{}); err == nil {
		arn := data.AWSAccountID(*out.Account)
		if iout, err := i.ListAccountAliases(&iam.ListAccountAliasesInput{}); err == nil {
			aliases := make([]string, len(iout.AccountAliases))
			for _, a := range iout.AccountAliases {
				aliases = append(aliases, *a)
			}
			accounts <- data.NewAWSAccount(arn, a.Id(), []data.KeyBinding{}, aliases...)
		} else {
			output.Warn("Failed to get accont alias for", a.Profile, "account", arn)
			accounts <- data.NewAWSAccount(arn, a.Id(), []data.KeyBinding{})
		}
		return arn, nil
	} else {
		return data.AWSAccountID(""), err
	}
}

func (a *AWSConnection) fetchAccessKeys(sess *session.Session, accounts chan<- data.Account, keys chan<- data.Key, usermap userMap) {

	i := iam.New(sess)

	if lako, err := i.ListAccessKeys(&iam.ListAccessKeysInput{}); err == nil {
		for _, md := range lako.AccessKeyMetadata {
			output.Debug("Found acces key", *md.AccessKeyId)
			userName := *md.UserName
			keys <- data.NewAwsKey(*md.AccessKeyId, *md.CreateDate, *md.Status == "Active", userName, *usermap[userName].Arn)
			//accounts <- data.NewIAMAccount(usermap[*md.UserName], a.Id(), md)
			accounts <- data.NewIAMAccountFromKey(md, usermap[userName], a.Id())
		}
	} else {
		output.Error(a, "failed to list IAM users")
	}
}

func (a *AWSConnection) fetchAccounts(sess *session.Session, accounts chan<- data.Account, keys chan<- data.Key) userMap {
	usermap := make(userMap)
	i := iam.New(sess)

	if r, err := i.ListUsers(&iam.ListUsersInput{}); err == nil {
		for _, user := range r.Users {
			usermap[*user.UserName] = user
		}
	} else {
		output.Error(a, "failed to list IAM users")
	}

	return usermap
}

func (a *AWSConnection) fetchInstances(region *string, sharedCredentials *credentials.Credentials, cAccounts chan data.Account, keymap map[string]data.ID) {
	output.Debug(a, "fetching instances from", *region)
	if sess, err := session.NewSession(&aws.Config{
		Region:      region,
		Credentials: sharedCredentials,
	}); err == nil {
		e := ec2.New(sess)
		//ec2.DescribeInstancesOutput{}.Reservations[0].Instances[0].k
		if dio, err := e.DescribeInstances(&ec2.DescribeInstancesInput{}); err == nil {
			output.Debug(a, *region, "reservations:", len(dio.Reservations))
			for _, res := range dio.Reservations {
				output.Debug(a, *region, "instances:", len(res.Instances))
				for _, instance := range res.Instances {
					keyID := keymap[*instance.KeyName]
					keys := []data.KeyBinding{
						{
							KeyID:    keyID,
							Location: data.INSTANCE_ROOT_CREDENTIALS,
						},
					}

					acct := data.NewAWSInstanceAccount(instance, a.Id(), keys)
					output.Debug("Found instance account", acct)
					cAccounts <- acct
				}
			}
		} else {
			output.Error(a, "Failed to fetch instances")
		}

	} else {
		output.Error(a, "Failed to connect")
	}
}

func (a *AWSConnection) fetchKeyPairs(arn data.AWSAccountID, region *string, sharedCredentials *credentials.Credentials, cKeys chan data.Key, cAccounts chan data.Account) (keymap map[string]data.ID) {
	output.Debug(a, "fetching key pairs from", *region)
	keymap = make(map[string]data.ID)

	if sess, err := session.NewSession(&aws.Config{
		Region:      region,
		Credentials: sharedCredentials,
	}); err == nil {
		e := ec2.New(sess)

		if out, err := e.DescribeKeyPairs(&ec2.DescribeKeyPairsInput{}); err == nil {
			bindings := make([]data.KeyBinding, 0)
			for _, p := range out.KeyPairs {
				fp := p.KeyFingerprint
				name := p.KeyName
				key := data.NewSSHKeyFromFingerprint(*name, time.Now(), data.ID(*fp))
				cKeys <- key
				bindings = append(bindings, data.KeyBinding{KeyID: data.ID(*fp), Name: *name})
				keymap[*name] = data.ID(*fp)
			}

			acct := data.NewAWSAccount(arn, a.Id(), bindings)

			cAccounts <- acct
		} else {
			output.Warn(a.String()+":", "Failed to find key pairs:", err)
		}

	} else {
		output.Error("Failed to connect to", a)
	}

	return
}

func (a *AWSConnection) Id() data.ID {
	return data.ID(a.Profile)
}
<|endoftext|>


<commit_before>package stream

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/mjibson/mog/codec"
	"github.com/mjibson/mog/codec/mpa"
	"github.com/mjibson/mog/protocol"
	"golang.org/x/oauth2"
)

func init() {
	protocol.Register("stream", []string{"URL"}, New, reflect.TypeOf(&Stream{}))
	gob.Register(new(Stream))
}

func New(params []string, token *oauth2.Token) (protocol.Instance, error) {
	if len(params) != 1 {
		return nil, fmt.Errorf("expected one parameter")
	}
	addr, err := url.Parse(params[0])
	if err != nil {
		return nil, err
	}
	u, name := tryPLS(params[0])
	if name == "" {
		name = params[0]
	}
	s := Stream{
		Orig: params[0],
		URL:  u,
		Name: name,
		Host: addr.Host,
	}
	resp, err := s.get()
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	return &s, nil
}

// tryPLS checks if u is a URL to a .pls file. If it is, it returns the
// first File entry of as target and first Title entry as name.
func tryPLS(u string) (target, name string) {
	target = u
	resp, err := http.Get(u)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	sc := bufio.NewScanner(io.LimitReader(resp.Body, 1024))
	i := 0
	for sc.Scan() {
		if i > 5 {
			break
		}
		i++
		sp := strings.SplitN(sc.Text(), "=", 2)
		if len(sp) != 2 {
			continue
		}
		if strings.HasPrefix(sp[0], "File") {
			_, err := url.Parse(sp[1])
			if err != nil {
				return
			}
			target = sp[1]
		} else if strings.HasPrefix(sp[0], "Title") {
			name = sp[1]
		}
		if target != u && name != "" {
			return
		}
	}
	return
}

type Stream struct {
	Orig           string
	URL            string
	Host           string
	Name           string
	metaint, count int
	body           io.ReadCloser
	title          string
	songtitle      string
}

type dialer struct {
	*net.Dialer
}

type conn struct {
	net.Conn
	read bool
}

// Read modifies the first line of an ICY stream response,
// if needed, to conform to Go's HTTP version requirements:
// http://golang.org/pkg/net/http/#ParseHTTPVersion.
func (c *conn) Read(b []byte) (n int, err error) {
	if !c.read {
		const headerICY = "ICY"
		const headerHTTP = "HTTP/1.1"
		// Hold 5 bytes because "HTTP/1.1" is 5 bytes longer than "ICY".
		n, err := c.Conn.Read(b[:len(b)+len(headerICY)-len(headerHTTP)])
		if bytes.HasPrefix(b, []byte(headerICY)) {
			copy(b[len(headerHTTP):], b[len(headerICY):])
			copy(b, []byte(headerHTTP))
		}
		c.read = true
		return n, err
	}
	return c.Conn.Read(b)
}

func (d *dialer) Dial(network, address string) (net.Conn, error) {
	c, err := d.Dialer.Dial(network, address)
	cn := conn{
		Conn: c,
	}
	return &cn, err
}

var client = &http.Client{
	Transport: &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&dialer{
			Dialer: &net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			},
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
	},
}

func (s *Stream) get() (*http.Response, error) {
	req, err := http.NewRequest("GET", s.URL, nil)
	if err != nil {
		panic(err)
		log.Fatal(err)
	}
	req.Header.Add("Icy-MetaData", "1")
	log.Println("stream open", req.URL)
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("stream status: %v", resp.Status)
	}
	s.metaint, err = strconv.Atoi(resp.Header.Get("Icy-Metaint"))
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (s *Stream) info() *codec.SongInfo {
	return &codec.SongInfo{
		Title: s.Name,
		Album: s.Host,
	}
}

func (s *Stream) Key() string {
	return s.Orig
}

func (s *Stream) List() (protocol.SongList, error) {
	return protocol.SongList{
		codec.ID(s.URL): s.info(),
	}, nil
}

func (s *Stream) Refresh() (protocol.SongList, error) {
	return s.List()
}

func (s *Stream) Info(codec.ID) (*codec.SongInfo, error) {
	i := s.info()
	i.SongTitle = s.songtitle
	return i, nil
}

func (s *Stream) GetSong(codec.ID) (codec.Song, error) {
	return mpa.NewSong(s.reader())
}

func (s *Stream) reader() codec.Reader {
	return func() (io.ReadCloser, int64, error) {
		resp, err := s.get()
		if err != nil {
			return nil, 0, err
		}
		s.Close()
		s.body = resp.Body
		s.songtitle = ""
		return s, 0, nil
	}
}

var titleRE = regexp.MustCompile("StreamTitle='(.*?)';")

func (s *Stream) Read(p []byte) (n int, err error) {
	if s.metaint == 0 {
		return s.body.Read(p)
	}
	l := s.metaint - s.count
	if len(p) > l {
		p = p[:l]
	}
	n, err = s.body.Read(p)
	s.count += n
	if s.count == s.metaint {
		s.count = 0
		mlen := make([]byte, 1)
		if _, err := io.ReadFull(s.body, mlen); err != nil {
			return n, err
		}
		meta := make([]byte, int(mlen[0])*16)
		if _, err := io.ReadFull(s.body, meta); err != nil {
			return n, err
		}
		matches := titleRE.FindSubmatch(meta)
		if len(matches) == 2 {
			s.songtitle = string(matches[1])
		}
	}
	return
}

func (s *Stream) Close() error {
	var err error
	if s.body != nil {
		err = s.body.Close()
	}
	s.body = nil
	s.count = 0
	s.title = ""
	return err
}
<commit_msg>Refetch PLS on stream refresh<commit_after>package stream

import (
	"bufio"
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/mjibson/mog/codec"
	"github.com/mjibson/mog/codec/mpa"
	"github.com/mjibson/mog/protocol"
	"golang.org/x/oauth2"
)

func init() {
	protocol.Register("stream", []string{"URL"}, New, reflect.TypeOf(&Stream{}))
	gob.Register(new(Stream))
}

func New(params []string, token *oauth2.Token) (protocol.Instance, error) {
	if len(params) != 1 {
		return nil, fmt.Errorf("expected one parameter")
	}
	addr, err := url.Parse(params[0])
	if err != nil {
		return nil, err
	}
	s := Stream{
		Orig: params[0],
		Host: addr.Host,
	}
	s.Refresh()
	resp, err := s.get()
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	return &s, nil
}

// tryPLS checks if u is a URL to a .pls file. If it is, it returns the
// first File entry of as target and first Title entry as name.
func tryPLS(u string) (target, name string) {
	target = u
	resp, err := http.Get(u)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	sc := bufio.NewScanner(io.LimitReader(resp.Body, 1024))
	i := 0
	for sc.Scan() {
		if i > 5 {
			break
		}
		i++
		sp := strings.SplitN(sc.Text(), "=", 2)
		if len(sp) != 2 {
			continue
		}
		if strings.HasPrefix(sp[0], "File") {
			_, err := url.Parse(sp[1])
			if err != nil {
				return
			}
			target = sp[1]
		} else if strings.HasPrefix(sp[0], "Title") {
			name = sp[1]
		}
		if target != u && name != "" {
			return
		}
	}
	return
}

type Stream struct {
	Orig           string
	URL            string
	Host           string
	Name           string
	metaint, count int
	body           io.ReadCloser
	title          string
	songtitle      string
}

type dialer struct {
	*net.Dialer
}

type conn struct {
	net.Conn
	read bool
}

// Read modifies the first line of an ICY stream response,
// if needed, to conform to Go's HTTP version requirements:
// http://golang.org/pkg/net/http/#ParseHTTPVersion.
func (c *conn) Read(b []byte) (n int, err error) {
	if !c.read {
		const headerICY = "ICY"
		const headerHTTP = "HTTP/1.1"
		// Hold 5 bytes because "HTTP/1.1" is 5 bytes longer than "ICY".
		n, err := c.Conn.Read(b[:len(b)+len(headerICY)-len(headerHTTP)])
		if bytes.HasPrefix(b, []byte(headerICY)) {
			copy(b[len(headerHTTP):], b[len(headerICY):])
			copy(b, []byte(headerHTTP))
		}
		c.read = true
		return n, err
	}
	return c.Conn.Read(b)
}

func (d *dialer) Dial(network, address string) (net.Conn, error) {
	c, err := d.Dialer.Dial(network, address)
	cn := conn{
		Conn: c,
	}
	return &cn, err
}

var client = &http.Client{
	Transport: &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&dialer{
			Dialer: &net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			},
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
	},
}

func (s *Stream) get() (*http.Response, error) {
	req, err := http.NewRequest("GET", s.URL, nil)
	if err != nil {
		panic(err)
		log.Fatal(err)
	}
	req.Header.Add("Icy-MetaData", "1")
	log.Println("stream open", req.URL)
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("stream status: %v", resp.Status)
	}
	s.metaint, err = strconv.Atoi(resp.Header.Get("Icy-Metaint"))
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (s *Stream) info() *codec.SongInfo {
	return &codec.SongInfo{
		Title: s.Name,
		Album: s.Host,
	}
}

func (s *Stream) Key() string {
	return s.Orig
}

func (s *Stream) List() (protocol.SongList, error) {
	return protocol.SongList{
		codec.ID(s.URL): s.info(),
	}, nil
}

func (s *Stream) Refresh() (protocol.SongList, error) {
	u, name := tryPLS(s.Orig)
	if name == "" {
		name = s.Orig
	}
	s.URL = u
	s.Name = name
	return s.List()
}

func (s *Stream) Info(codec.ID) (*codec.SongInfo, error) {
	i := s.info()
	i.SongTitle = s.songtitle
	return i, nil
}

func (s *Stream) GetSong(codec.ID) (codec.Song, error) {
	return mpa.NewSong(s.reader())
}

func (s *Stream) reader() codec.Reader {
	return func() (io.ReadCloser, int64, error) {
		resp, err := s.get()
		if err != nil {
			return nil, 0, err
		}
		s.Close()
		s.body = resp.Body
		s.songtitle = ""
		return s, 0, nil
	}
}

var titleRE = regexp.MustCompile("StreamTitle='(.*?)';")

func (s *Stream) Read(p []byte) (n int, err error) {
	if s.metaint == 0 {
		return s.body.Read(p)
	}
	l := s.metaint - s.count
	if len(p) > l {
		p = p[:l]
	}
	n, err = s.body.Read(p)
	s.count += n
	if s.count == s.metaint {
		s.count = 0
		mlen := make([]byte, 1)
		if _, err := io.ReadFull(s.body, mlen); err != nil {
			return n, err
		}
		meta := make([]byte, int(mlen[0])*16)
		if _, err := io.ReadFull(s.body, meta); err != nil {
			return n, err
		}
		matches := titleRE.FindSubmatch(meta)
		if len(matches) == 2 {
			s.songtitle = string(matches[1])
		}
	}
	return
}

func (s *Stream) Close() error {
	var err error
	if s.body != nil {
		err = s.body.Close()
	}
	s.body = nil
	s.count = 0
	s.title = ""
	return err
}
<|endoftext|>