<commit_before>package ParquetType

import (
	"github.com/xitongsys/parquet-go/parquet"
	"testing"
)

func TestNameToBaseType(t *testing.T) {
	testData := []struct {
		Name string
		Type parquet.Type
	}{
		{"BOOLEAN", parquet.Type_BOOLEAN},
		{"INT32", parquet.Type_INT32},
		{"INT64", parquet.Type_INT64},
		{"INT96", parquet.Type_INT96},
		{"FLOAT", parquet.Type_FLOAT},
		{"DOUBLE", parquet.Type_DOUBLE},
		{"BYTE_ARRAY", parquet.Type_BYTE_ARRAY},
		{"FIXED_LEN_BYTE_ARRAY", parquet.Type_FIXED_LEN_BYTE_ARRAY},
	}

	for _, data := range testData {
		res := NameToBaseType(data.Name)
		if res != data.Type {
			t.Errorf("NameToBaseType err, expect %v, get %v", data.Type, res)
		}
	}
}

func TestNameToConvertedType(t *testing.T) {
	testData := []struct {
		Name string
		Type parquet.ConvertedType
	}{
		{"UTF8", parquet.ConvertedType_UTF8},
		{"INT_8", parquet.ConvertedType_INT_8},
		{"INT_16", parquet.ConvertedType_INT_16},
		{"INT_32", parquet.ConvertedType_INT_32},
		{"INT_64", parquet.ConvertedType_INT_64},
		{"UINT_8", parquet.ConvertedType_UINT_8},
		{"UINT_16", parquet.ConvertedType_UINT_16},
		{"UINT_32", parquet.ConvertedType_UINT_32},
		{"UINT_64", parquet.ConvertedType_UINT_64},
		{"DATE", parquet.ConvertedType_DATE},
		{"TIME_MILLIS", parquet.ConvertedType_TIME_MILLIS},
		{"TIME_MICROS", parquet.ConvertedType_TIME_MICROS},
		{"TIMESTAMP_MILLIS", parquet.ConvertedType_TIMESTAMP_MILLIS},
		{"TIMESTAMP_MICROS", parquet.ConvertedType_TIMESTAMP_MICROS},
		{"INTERVAL", parquet.ConvertedType_INTERVAL},
		{"DECIMAL", parquet.ConvertedType_DECIMAL},
	}

	for _, data := range testData {
		res := NameToConvertedType(data.Name)
		if res != data.Type {
			t.Errorf("NameToBaseType err, expect %v, get %v", data.Type, res)
		}
	}
}

func TestIsBaseType(t *testing.T) {
	testData := []struct {
		Name string
		Flag bool
	}{
		{"BOOLEAN", true},
		{"INT32", true},
		{"INT64", true},
		{"INT96", true},
		{"FLOAT", true},
		{"DOUBLE", true},
		{"BYTE_ARRAY", true},
		{"FIXED_LEN_BYTE_ARRAY", true},
		{"UTF8", false},
		{"INT_8", false},
		{"INT_16", false},
		{"INT_32", false},
		{"INT_64", false},
		{"UINT_8", false},
		{"UINT_16", false},
		{"UINT_32", false},
		{"UINT_64", false},
		{"DATE", false},
		{"TIME_MILLIS", false},
		{"TIME_MICROS", false},
		{"TIMESTAMP_MILLIS", false},
		{"TIMESTAMP_MICROS", false},
		{"INTERAL", false},
		{"DECIMAL", false},
	}

	for _, data := range testData {
		res := IsBaseType(data.Name)
		if res != data.Flag {
			t.Errorf("IsBaseType err")
		}
	}
}
<commit_msg>adding ut<commit_after>package ParquetType

import (
	"github.com/xitongsys/parquet-go/parquet"
	"testing"
)

func TestNameToBaseType(t *testing.T) {
	testData := []struct {
		Name string
		Type parquet.Type
	}{
		{"BOOLEAN", parquet.Type_BOOLEAN},
		{"INT32", parquet.Type_INT32},
		{"INT64", parquet.Type_INT64},
		{"INT96", parquet.Type_INT96},
		{"FLOAT", parquet.Type_FLOAT},
		{"DOUBLE", parquet.Type_DOUBLE},
		{"BYTE_ARRAY", parquet.Type_BYTE_ARRAY},
		{"FIXED_LEN_BYTE_ARRAY", parquet.Type_FIXED_LEN_BYTE_ARRAY},
	}

	for _, data := range testData {
		res := NameToBaseType(data.Name)
		if res != data.Type {
			t.Errorf("NameToBaseType err, expect %v, get %v", data.Type, res)
		}
	}
}

func TestNameToConvertedType(t *testing.T) {
	testData := []struct {
		Name string
		Type parquet.ConvertedType
	}{
		{"UTF8", parquet.ConvertedType_UTF8},
		{"INT_8", parquet.ConvertedType_INT_8},
		{"INT_16", parquet.ConvertedType_INT_16},
		{"INT_32", parquet.ConvertedType_INT_32},
		{"INT_64", parquet.ConvertedType_INT_64},
		{"UINT_8", parquet.ConvertedType_UINT_8},
		{"UINT_16", parquet.ConvertedType_UINT_16},
		{"UINT_32", parquet.ConvertedType_UINT_32},
		{"UINT_64", parquet.ConvertedType_UINT_64},
		{"DATE", parquet.ConvertedType_DATE},
		{"TIME_MILLIS", parquet.ConvertedType_TIME_MILLIS},
		{"TIME_MICROS", parquet.ConvertedType_TIME_MICROS},
		{"TIMESTAMP_MILLIS", parquet.ConvertedType_TIMESTAMP_MILLIS},
		{"TIMESTAMP_MICROS", parquet.ConvertedType_TIMESTAMP_MICROS},
		{"INTERVAL", parquet.ConvertedType_INTERVAL},
		{"DECIMAL", parquet.ConvertedType_DECIMAL},
	}

	for _, data := range testData {
		res := NameToConvertedType(data.Name)
		if res != data.Type {
			t.Errorf("NameToBaseType err, expect %v, get %v", data.Type, res)
		}
	}
}

func TestIsBaseType(t *testing.T) {
	testData := []struct {
		Name string
		Flag bool
	}{
		{"BOOLEAN", true},
		{"INT32", true},
		{"INT64", true},
		{"INT96", true},
		{"FLOAT", true},
		{"DOUBLE", true},
		{"BYTE_ARRAY", true},
		{"FIXED_LEN_BYTE_ARRAY", true},
		{"UTF8", false},
		{"INT_8", false},
		{"INT_16", false},
		{"INT_32", false},
		{"INT_64", false},
		{"UINT_8", false},
		{"UINT_16", false},
		{"UINT_32", false},
		{"UINT_64", false},
		{"DATE", false},
		{"TIME_MILLIS", false},
		{"TIME_MICROS", false},
		{"TIMESTAMP_MILLIS", false},
		{"TIMESTAMP_MICROS", false},
		{"INTERAL", false},
		{"DECIMAL", false},
	}

	for _, data := range testData {
		res := IsBaseType(data.Name)
		if res != data.Flag {
			t.Errorf("IsBaseType err")
		}
	}
}

func TestStrToParquetType(t *testing.T) {
	testData := []struct {
		StrData     string
		Type        string
		ParquetData interface{}
	}{
		{"false", "BOOLEAN", BOOLEAN(false)},
		{"1", "INT32", INT32(1)},
		{"0", "INT64", INT64(0)},
		{"012345678901", "INT96", INT96("012345678901")},
		{"0.1", "FLOAT", FLOAT(0.1)},
		{"0.1", "DOUBLE", DOUBLE(0.1)},
		{"abc bcd", "BYTE_ARRAY", BYTE_ARRAY("abc bcd")},
		{"abc bcd", "FIXED_LEN_BYTE_ARRAY", FIXED_LEN_BYTE_ARRAY("abc bcd")},
	}

	for _, data := range testData {
		res := StrToParquetType(data.StrData, data.Type)
		if res != data.ParquetData {
			t.Errorf("StrToParquetType err, expect %v, get %v", data.ParquetData, res)
		}
	}
}
<|endoftext|>


<commit_before>package memory

import (
	"sort"

	"github.com/grafana/metrictank/expr/tagquery"
	"github.com/grafana/metrictank/schema"
	log "github.com/sirupsen/logrus"
)

type expressionFilter struct {
	// expr is the expression based on which this filter has been generated
	expr tagquery.Expression

	// test is a filter function which takes certain descriptive properties of a
	// MetricDefinition and returns a tagquery.FilterDecision type indicating
	// whether the MD satisfies this expression or not
	// this decision is made based on the metric tag index, the meta tag index
	// does not get take into account
	testByMetricTags tagquery.MetricDefinitionFilter

	// testByMetaTags is a filter function which has been generated from the meta
	// records that match this filter's expression.
	// it looks at certain descriptive properties of a MetricDefinition and decides
	// whether this metric definition satisfies the given expression based on the
	// meta tag index, it does not take the metric tag index into account
	testByMetaTags tagquery.MetricDefinitionFilter

	// the default decision which should be applied if none of test & testByMetaTags
	// have come to a conclusive decision
	defaultDecision tagquery.FilterDecision
}

// idFilter contains one or many filter functions which are used to filter metrics
// by the expressions from which its filters have been generated.
// once initialized it can be used concurrently by many worker routines
type idFilter struct {
	ctx     *TagQueryContext
	filters []expressionFilter
}

// newIdFilter takes a set of expressions and a tag query context, then it generates
// various filter functions from the expressions which are going to be used to decide
// whether a given metric matches the provided expressions
func newIdFilter(expressions tagquery.Expressions, ctx *TagQueryContext) *idFilter {
	res := idFilter{
		ctx:     ctx,
		filters: make([]expressionFilter, len(expressions)),
	}

	for i, expr := range expressions {
		res.filters[i] = expressionFilter{
			expr:             expr,
			testByMetricTags: expr.GetMetricDefinitionFilter(ctx.index.idHasTag),
			defaultDecision:  expr.GetDefaultDecision(),
		}

		if !ctx.useMetaTagIndex() {
			continue
		}

		// this is a performacnce optimization:
		// some expressions indicate that they'll likely result in a smaller result set
		// if they get inverted.
		// f.e. a!=b becomes a=b)
		// in this case we want to get the inverted set of meta records that matches
		// "expr", because when generating a off of this smaller result set less meta
		// record expressions will need to be checked, which will improve the filter
		// speed. if the meta record filter set has been generated from an inverted
		// expression, then its result will need to be inverted again to get the
		// correct result.
		// f.e. if a!=b previously has been inverted to a=b, and for a given MD a=b
		// results in true, then a!=b would result in false
		invertSetOfMetaRecords := expr.ResultIsSmallerWhenInverted()

		// if no meta records match this expression, then we don't need to generate
		// a meta record filter for it
		metaRecordIds := ctx.metaTagIndex.getMetaRecordIdsByExpression(expr, invertSetOfMetaRecords)
		if len(metaRecordIds) == 0 {
			continue
		}

		records := make([]tagquery.MetaTagRecord, 0, len(metaRecordIds))
		for _, id := range metaRecordIds {
			record, ok := ctx.metaTagRecords.getMetaRecordById(id)
			if !ok {
				corruptIndex.Inc()
				log.Errorf("TagQueryContext: Tried to lookup a meta tag record id that does not exist, index is corrupted")
				continue
			}

			records = append(records, record)
		}

		// if we don't use an inverted set of meta records, then we check if
		// all meta records involved in a meta tag filter use the "=" operator.
		// if this is the case then it is cheaper to build a set of acceptable tags
		// based on the meta record expressions and just check whether they are present
		// in a metric that gets filtered, compared to doing a full tag index lookup
		// to check whether a metric has one of the necessary meta tags associated
		// with it.
		onlyEqualOperators, singleExprPerRecord := viableOptimizations(records)

		if onlyEqualOperators {
			// there are two different ways how we optimize for the case where all expressions
			// of all involved meta records are using the "=" operator. the first and fastest
			// way can only be used if each involved meta record only has one single expression,
			// otherwise we use the second way which is a bit more expensive but it also works
			// if some of the involved meta records have multiple expressions.
			if singleExprPerRecord {
				res.filters[i].testByMetaTags = metaRecordFilterBySetOfValidValues(records, invertSetOfMetaRecords)
			} else {
				res.filters[i].testByMetaTags = metaRecordFilterBySetOfValidValueSets(records, invertSetOfMetaRecords)
			}
		} else {
			metaRecordFilters := make([]tagquery.MetricDefinitionFilter, 0, len(records))
			for recordIdx := range records {
				metaRecordFilters = append(metaRecordFilters, records[recordIdx].GetMetricDefinitionFilter(ctx.index.idHasTag))
			}

			if invertSetOfMetaRecords {
				res.filters[i].testByMetaTags = metaRecordFilterInverted(metaRecordFilters, res.filters[i].defaultDecision)
			} else {
				res.filters[i].testByMetaTags = metaRecordFilterNormal(metaRecordFilters, res.filters[i].defaultDecision)
			}
		}
	}

	return &res
}

// viableOptimizations looks at a set of meta tag records and decides whether two possible
// optimizations can be applied when filtering by these records. it returns two bools to
// indicate which optimizations are or are not viable.
// if invertSetOfMetaRecords is true then none of these optimizations can be used.
//
// * the first bool refers to the optimization for sets of records which all have only one
//   expression and this expression is using the equal operator.
// * the second bool refers to the optimization for sets of records which all only have
//   expressions using the equal operator, but there may be more than one per record.
func viableOptimizations(records []tagquery.MetaTagRecord) (bool, bool) {
	singleExprPerRecord := true
	for recordIdx := range records {
		for exprIdx := range records[recordIdx].Expressions {
			if records[recordIdx].Expressions[exprIdx].GetOperator() != tagquery.EQUAL {
				return false, false
			}
		}
		if len(records[recordIdx].Expressions) > 1 {
			singleExprPerRecord = false
		}
	}

	return true, singleExprPerRecord
}

// metaRecordFilterBySetOfValidValues creates a filter function to filter by a meta tag
// which only involves meta records of which each only has exactly one expression and that
// expression is using the "=" operator. this is quite a narrow scenario, but since it is
// a very common use case it makes sense to optimize for it.
func metaRecordFilterBySetOfValidValues(records []tagquery.MetaTagRecord, invertFilter bool) tagquery.MetricDefinitionFilter {
	// we first build a set of valid tags and names.
	// since we know that each of the involved meta records uses exactly one expression
	// which is using the "=" operator we know that if a given metric's name matches a
	// value in validNames or if one of its tags matches a value in validValues then this
	// is sufficient to let it pass the filter.
	validValues := make(map[string]struct{})
	validNames := make(map[string]struct{})
	for _, record := range records {
		if len(record.Expressions) < 1 {
			corruptIndex.Inc()
			log.Errorf("meta-record-filter-by-set-of-valid-values: received a meta record without any associated expressions")
			continue
		}

		if record.Expressions[0].GetKey() == "name" {
			validNames[record.Expressions[0].GetValue()] = struct{}{}
		} else {
			validValues[record.Expressions[0].GetKey()+"="+record.Expressions[0].GetValue()] = struct{}{}
		}
	}

	resultOnHit := tagquery.Pass
	if invertFilter {
		resultOnHit = tagquery.Fail
	}

	return func(_ schema.MKey, name string, tags []string) tagquery.FilterDecision {
		for i := range tags {
			if _, ok := validValues[tags[i]]; ok {
				return resultOnHit
			}
		}
		if _, ok := validNames[name]; ok {
			return resultOnHit
		}
		return tagquery.None
	}
}

// metaRecordFilterBySetOfValidValueSets creates a filter function to filter by a meta tag
// which only involves meta records of which all expressions are only using the "=" operator,
// it is ok if one meta record uses multiple such expressions.
func metaRecordFilterBySetOfValidValueSets(records []tagquery.MetaTagRecord, invertFilter bool) tagquery.MetricDefinitionFilter {
	// we first build a set of tag and name value combinations of which each is sufficient
	// to pass the generated filter when a metric contains all values of one of these
	// combinations
	validValueSets := make([]struct {
		name string
		tags []string
	}, len(records))
	for i := range records {
		validValueSets[i].tags = make([]string, 0, len(records[i].Expressions))
		for j := range records[i].Expressions {
			if records[i].Expressions[j].GetKey() == "name" {
				validValueSets[i].name = records[i].Expressions[j].GetValue()
			} else {
				validValueSets[i].tags = append(validValueSets[i].tags, records[i].Expressions[j].GetKey()+"="+records[i].Expressions[j].GetValue())
			}
		}

		// the tags must be sorted because that's a requirement for sliceContainsElements()
		sort.Strings(validValueSets[i].tags)
	}

	resultOnHit := tagquery.Pass
	if invertFilter {
		resultOnHit = tagquery.Fail
	}

	return func(_ schema.MKey, name string, tags []string) tagquery.FilterDecision {
		// iterate over the acceptable value combinations and check if one matches this metric
		for _, validValueSet := range validValueSets {
			if len(validValueSet.name) > 0 {
				if name != validValueSet.name {
					continue
				}
			}

			if sliceContainsElements(validValueSet.tags, tags) {
				return resultOnHit
			}
		}

		return tagquery.None
	}
}

// sliceContainsElements returns true if the elements in the slice "find"
// are all present int the slice "in". It requires both slices to be sorted
func sliceContainsElements(find, in []string) bool {
	var findIdx, inIdx int
	for {
		if findIdx == len(find) {
			return true
		}

		if inIdx == len(in) {
			return false
		}

		if find[findIdx] == in[inIdx] {
			findIdx++
			inIdx++
			continue
		}

		if find[findIdx] < in[inIdx] {
			return false
		}

		inIdx++
	}
}

func metaRecordFilterInverted(metaRecordFilters []tagquery.MetricDefinitionFilter, defaultDecision tagquery.FilterDecision) tagquery.MetricDefinitionFilter {
	return func(id schema.MKey, name string, tags []string) tagquery.FilterDecision {
		for _, metaRecordFilter := range metaRecordFilters {
			decision := metaRecordFilter(id, name, tags)
			if decision == tagquery.None {
				decision = defaultDecision
			}

			if decision == tagquery.Pass {
				return tagquery.Fail
			}
		}

		return tagquery.Pass
	}
}

func metaRecordFilterNormal(metaRecordFilters []tagquery.MetricDefinitionFilter, defaultDecision tagquery.FilterDecision) tagquery.MetricDefinitionFilter {
	return func(id schema.MKey, name string, tags []string) tagquery.FilterDecision {
		for _, metaRecordFilter := range metaRecordFilters {
			decision := metaRecordFilter(id, name, tags)
			if decision == tagquery.None {
				decision = defaultDecision
			}

			if decision == tagquery.Pass {
				return tagquery.Pass
			}
		}

		return tagquery.Fail
	}
}

// matches takes descriptive properties of a metric definition and runs them through all
// filters required to come to a decision whether this metric definition should be part
// of the result or not.
// it uses the filter functions that have previously been generated when this instance
// of idFilter was instantiated
func (f *idFilter) matches(id schema.MKey, name string, tags []string) bool {
	for i := range f.filters {
		decision := f.filters[i].testByMetricTags(id, name, tags)

		if decision == tagquery.Pass {
			continue
		}

		if decision == tagquery.Fail {
			return false
		}

		if f.filters[i].testByMetaTags != nil {
			decision = f.filters[i].testByMetaTags(id, name, tags)
		}

		if decision == tagquery.None {
			decision = f.filters[i].defaultDecision
		}

		if decision == tagquery.Pass {
			continue
		}

		return false
	}

	return true
}
<commit_msg>update comments<commit_after>package memory

import (
	"sort"

	"github.com/grafana/metrictank/expr/tagquery"
	"github.com/grafana/metrictank/schema"
	log "github.com/sirupsen/logrus"
)

type expressionFilter struct {
	// expr is the expression based on which this filter has been generated
	expr tagquery.Expression

	// test is a filter function which takes certain descriptive properties of a
	// MetricDefinition and returns a tagquery.FilterDecision type indicating
	// whether the MD satisfies this expression or not
	// this decision is made based on the metric tag index, the meta tag index
	// does not get take into account
	testByMetricTags tagquery.MetricDefinitionFilter

	// testByMetaTags is a filter function which has been generated from the meta
	// records that match this filter's expression.
	// it looks at certain descriptive properties of a MetricDefinition and decides
	// whether this metric definition satisfies the given expression based on the
	// meta tag index, it does not take the metric tag index into account
	testByMetaTags tagquery.MetricDefinitionFilter

	// the default decision which should be applied if none of test & testByMetaTags
	// have come to a conclusive decision
	defaultDecision tagquery.FilterDecision
}

// idFilter contains one or many filter functions which are used to filter metrics
// by the expressions from which its filters have been generated.
// once initialized it can be used concurrently by many worker routines
type idFilter struct {
	ctx     *TagQueryContext
	filters []expressionFilter
}

// newIdFilter takes a set of expressions and a tag query context, then it generates
// various filter functions from the expressions which are going to be used to decide
// whether a given metric matches the provided expressions
func newIdFilter(expressions tagquery.Expressions, ctx *TagQueryContext) *idFilter {
	res := idFilter{
		ctx:     ctx,
		filters: make([]expressionFilter, len(expressions)),
	}

	for i, expr := range expressions {
		res.filters[i] = expressionFilter{
			expr:             expr,
			testByMetricTags: expr.GetMetricDefinitionFilter(ctx.index.idHasTag),
			defaultDecision:  expr.GetDefaultDecision(),
		}

		if !ctx.useMetaTagIndex() {
			continue
		}

		// this is a performacnce optimization:
		// some expressions indicate that they'll likely result in a smaller result set
		// if they get inverted.
		// f.e. a!=b becomes a=b)
		// in this case we want to get the inverted set of meta records that matches
		// "expr", because when generating a off of this smaller result set less meta
		// record expressions will need to be checked, which will improve the filter
		// speed. if the meta record filter set has been generated from an inverted
		// expression, then its result will need to be inverted again to get the
		// correct result.
		// f.e. if a!=b previously has been inverted to a=b, and for a given MD a=b
		// results in true, then a!=b would result in false
		invertSetOfMetaRecords := expr.ResultIsSmallerWhenInverted()

		// if no meta records match this expression, then we don't need to generate
		// a meta record filter for it
		metaRecordIds := ctx.metaTagIndex.getMetaRecordIdsByExpression(expr, invertSetOfMetaRecords)
		if len(metaRecordIds) == 0 {
			continue
		}

		records := make([]tagquery.MetaTagRecord, 0, len(metaRecordIds))
		for _, id := range metaRecordIds {
			record, ok := ctx.metaTagRecords.getMetaRecordById(id)
			if !ok {
				corruptIndex.Inc()
				log.Errorf("TagQueryContext: Tried to lookup a meta tag record id that does not exist, index is corrupted")
				continue
			}

			records = append(records, record)
		}

		// if a query only involves meta tags of which all underlying expressions
		// use the "=" operator, then it is cheaper to build a set of acceptable tags
		// based on the meta record expressions and just check whether they are present
		// in a metric that gets filtered, compared to doing a full tag index lookup
		// to check whether a metric has one of the necessary meta tags associated
		// with it.
		onlyEqualOperators, singleExprPerRecord := viableOptimizations(records)

		if onlyEqualOperators {
			// there are two different ways how we optimize for the case where all expressions
			// of all involved meta records are using the "=" operator. the first and fastest
			// way can only be used if each involved meta record only has one single expression,
			// otherwise we use the second way which is a bit more expensive but it also works
			// if some of the involved meta records have multiple expressions.
			if singleExprPerRecord {
				res.filters[i].testByMetaTags = metaRecordFilterBySetOfValidValues(records, invertSetOfMetaRecords)
			} else {
				res.filters[i].testByMetaTags = metaRecordFilterBySetOfValidValueSets(records, invertSetOfMetaRecords)
			}
		} else {
			metaRecordFilters := make([]tagquery.MetricDefinitionFilter, 0, len(records))
			for recordIdx := range records {
				metaRecordFilters = append(metaRecordFilters, records[recordIdx].GetMetricDefinitionFilter(ctx.index.idHasTag))
			}

			if invertSetOfMetaRecords {
				res.filters[i].testByMetaTags = metaRecordFilterInverted(metaRecordFilters, res.filters[i].defaultDecision)
			} else {
				res.filters[i].testByMetaTags = metaRecordFilterNormal(metaRecordFilters, res.filters[i].defaultDecision)
			}
		}
	}

	return &res
}

// viableOptimizations looks at a set of meta tag records and decides whether two possible
// optimizations can be applied when filtering by these records. it returns two bools to
// indicate which optimizations are or are not viable.
//
// * the first bool refers to the optimization for sets of records which all have only one
//   expression and this expression is using the equal operator.
// * the second bool refers to the optimization for sets of records which all only have
//   expressions using the equal operator, but there may be more than one per record.
func viableOptimizations(records []tagquery.MetaTagRecord) (bool, bool) {
	singleExprPerRecord := true
	for recordIdx := range records {
		for exprIdx := range records[recordIdx].Expressions {
			if records[recordIdx].Expressions[exprIdx].GetOperator() != tagquery.EQUAL {
				return false, false
			}
		}
		if len(records[recordIdx].Expressions) > 1 {
			singleExprPerRecord = false
		}
	}

	return true, singleExprPerRecord
}

// metaRecordFilterBySetOfValidValues creates a filter function to filter by a meta tag
// which only involves meta records of which each only has exactly one expression and that
// expression is using the "=" operator. this is quite a narrow scenario, but since it is
// a very common use case it makes sense to optimize for it.
// The invertFilter bool flips the filter logic so that instead of removing metrics which
// do not have a meta tag it filters metrics which do have a meta tag.
func metaRecordFilterBySetOfValidValues(records []tagquery.MetaTagRecord, invertFilter bool) tagquery.MetricDefinitionFilter {
	// we first build a set of valid tags and names.
	// since we know that each of the involved meta records uses exactly one expression
	// which is using the "=" operator we know that if a given metric's name matches a
	// value in validNames or if one of its tags matches a value in validValues then this
	// is sufficient to let it pass the filter.
	validValues := make(map[string]struct{})
	validNames := make(map[string]struct{})
	for _, record := range records {
		if len(record.Expressions) < 1 {
			corruptIndex.Inc()
			log.Errorf("meta-record-filter-by-set-of-valid-values: received a meta record without any associated expressions")
			continue
		}

		if record.Expressions[0].GetKey() == "name" {
			validNames[record.Expressions[0].GetValue()] = struct{}{}
		} else {
			validValues[record.Expressions[0].GetKey()+"="+record.Expressions[0].GetValue()] = struct{}{}
		}
	}

	resultOnHit := tagquery.Pass
	if invertFilter {
		resultOnHit = tagquery.Fail
	}

	return func(_ schema.MKey, name string, tags []string) tagquery.FilterDecision {
		for i := range tags {
			if _, ok := validValues[tags[i]]; ok {
				return resultOnHit
			}
		}
		if _, ok := validNames[name]; ok {
			return resultOnHit
		}
		return tagquery.None
	}
}

// metaRecordFilterBySetOfValidValueSets creates a filter function to filter by a meta tag
// which only involves meta records of which all expressions are only using the "=" operator,
// it is ok if one meta record uses multiple such expressions.
// The invertFilter bool flips the filter logic so that instead of removing metrics which
// do not have a meta tag it filters metrics which do have a meta tag.
func metaRecordFilterBySetOfValidValueSets(records []tagquery.MetaTagRecord, invertFilter bool) tagquery.MetricDefinitionFilter {
	// we first build a set of tag and name value combinations of which each is sufficient
	// to pass the generated filter when a metric contains all values of one of these
	// combinations
	validValueSets := make([]struct {
		name string
		tags []string
	}, len(records))
	for i := range records {
		validValueSets[i].tags = make([]string, 0, len(records[i].Expressions))
		for j := range records[i].Expressions {
			if records[i].Expressions[j].GetKey() == "name" {
				validValueSets[i].name = records[i].Expressions[j].GetValue()
			} else {
				validValueSets[i].tags = append(validValueSets[i].tags, records[i].Expressions[j].GetKey()+"="+records[i].Expressions[j].GetValue())
			}
		}

		// the tags must be sorted because that's a requirement for sliceContainsElements()
		sort.Strings(validValueSets[i].tags)
	}

	resultOnHit := tagquery.Pass
	if invertFilter {
		resultOnHit = tagquery.Fail
	}

	return func(_ schema.MKey, name string, tags []string) tagquery.FilterDecision {
		// iterate over the acceptable value combinations and check if one matches this metric
		for _, validValueSet := range validValueSets {
			if len(validValueSet.name) > 0 {
				if name != validValueSet.name {
					continue
				}
			}

			if sliceContainsElements(validValueSet.tags, tags) {
				return resultOnHit
			}
		}

		return tagquery.None
	}
}

// sliceContainsElements returns true if the elements in the slice "find"
// are all present int the slice "in". It requires both slices to be sorted
func sliceContainsElements(find, in []string) bool {
	var findIdx, inIdx int
	for {
		if findIdx == len(find) {
			return true
		}

		if inIdx == len(in) {
			return false
		}

		if find[findIdx] == in[inIdx] {
			findIdx++
			inIdx++
			continue
		}

		if find[findIdx] < in[inIdx] {
			return false
		}

		inIdx++
	}
}

func metaRecordFilterInverted(metaRecordFilters []tagquery.MetricDefinitionFilter, defaultDecision tagquery.FilterDecision) tagquery.MetricDefinitionFilter {
	return func(id schema.MKey, name string, tags []string) tagquery.FilterDecision {
		for _, metaRecordFilter := range metaRecordFilters {
			decision := metaRecordFilter(id, name, tags)
			if decision == tagquery.None {
				decision = defaultDecision
			}

			if decision == tagquery.Pass {
				return tagquery.Fail
			}
		}

		return tagquery.Pass
	}
}

func metaRecordFilterNormal(metaRecordFilters []tagquery.MetricDefinitionFilter, defaultDecision tagquery.FilterDecision) tagquery.MetricDefinitionFilter {
	return func(id schema.MKey, name string, tags []string) tagquery.FilterDecision {
		for _, metaRecordFilter := range metaRecordFilters {
			decision := metaRecordFilter(id, name, tags)
			if decision == tagquery.None {
				decision = defaultDecision
			}

			if decision == tagquery.Pass {
				return tagquery.Pass
			}
		}

		return tagquery.Fail
	}
}

// matches takes descriptive properties of a metric definition and runs them through all
// filters required to come to a decision whether this metric definition should be part
// of the result or not.
// it uses the filter functions that have previously been generated when this instance
// of idFilter was instantiated
func (f *idFilter) matches(id schema.MKey, name string, tags []string) bool {
	for i := range f.filters {
		decision := f.filters[i].testByMetricTags(id, name, tags)

		if decision == tagquery.Pass {
			continue
		}

		if decision == tagquery.Fail {
			return false
		}

		if f.filters[i].testByMetaTags != nil {
			decision = f.filters[i].testByMetaTags(id, name, tags)
		}

		if decision == tagquery.None {
			decision = f.filters[i].defaultDecision
		}

		if decision == tagquery.Pass {
			continue
		}

		return false
	}

	return true
}
<|endoftext|>


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

package ssh

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/Unknwon/com"
	"golang.org/x/crypto/ssh"

	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/setting"
)

func cleanCommand(cmd string) string {
	i := strings.Index(cmd, "git")
	if i == -1 {
		return cmd
	}
	return cmd[i:]
}

func handleServerConn(keyID string, chans <-chan ssh.NewChannel) {
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		ch, reqs, err := newChan.Accept()
		if err != nil {
			log.Error(3, "Error accepting channel: %v", err)
			continue
		}

		go func(in <-chan *ssh.Request) {
			defer ch.Close()
			for req := range in {
				payload := cleanCommand(string(req.Payload))
				switch req.Type {
				case "env":
					args := strings.Split(strings.Replace(payload, "\x00", "", -1), "\v")
					if len(args) != 2 {
						log.Warn("Invalid env arguments: '%#v'", args)
						continue
					}
					args[0] = strings.TrimLeft(args[0], "\x04")
					_, _, err := com.ExecCmdBytes("env", args[0]+"="+args[1])
					if err != nil {
						log.Error(3, "env: %v", err)
						return
					}
				case "exec":
					cmdName := strings.TrimLeft(payload, "'()")
					os.Setenv("SSH_ORIGINAL_COMMAND", cmdName)
					log.Trace("Payload: %v", cmdName)

					args := []string{"serv", "key-" + keyID, "--config=" + setting.CustomConf}
					log.Trace("Arguments: %v", args)
					cmd := exec.Command(setting.AppPath, args...)

					stdout, err := cmd.StdoutPipe()
					if err != nil {
						log.Error(3, "StdoutPipe: %v", err)
						return
					}
					stderr, err := cmd.StderrPipe()
					if err != nil {
						log.Error(3, "StderrPipe: %v", err)
						return
					}
					input, err := cmd.StdinPipe()
					if err != nil {
						log.Error(3, "StdinPipe: %v", err)
						return
					}

					// FIXME: check timeout
					if err = cmd.Start(); err != nil {
						log.Error(3, "Start: %v", err)
						return
					}

					go io.Copy(input, ch)
					io.Copy(ch, stdout)
					io.Copy(ch.Stderr(), stderr)

					if err = cmd.Wait(); err != nil {
						log.Error(3, "Wait: %v", err)
						return
					}

					ch.SendRequest("exit-status", false, []byte{0, 0, 0, 0})
					return
				default:
				}
			}
		}(reqs)
	}
}

func listen(config *ssh.ServerConfig, port int) {
	listener, err := net.Listen("tcp", "0.0.0.0:"+com.ToStr(port))
	if err != nil {
		panic(err)
	}
	for {
		// Once a ServerConfig has been configured, connections can be accepted.
		conn, err := listener.Accept()
		if err != nil {
			log.Error(3, "Error accepting incoming connection: %v", err)
			continue
		}
		// Before use, a handshake must be performed on the incoming net.Conn.
		sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
		if err != nil {
			log.Error(3, "Error on handshaking: %v", err)
			continue
		}

		log.Trace("Connection from %s (%s)", sConn.RemoteAddr(), sConn.ClientVersion())
		// The incoming Request channel must be serviced.
		go ssh.DiscardRequests(reqs)
		go handleServerConn(sConn.Permissions.Extensions["key-id"], chans)
	}
}

// Listen starts a SSH server listens on given port.
func Listen(port int) {
	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			pkey, err := models.SearchPublicKeyByContent(strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key))))
			if err != nil {
				log.Error(3, "SearchPublicKeyByContent: %v", err)
				return nil, err
			}
			return &ssh.Permissions{Extensions: map[string]string{"key-id": com.ToStr(pkey.ID)}}, nil
		},
	}

	keyPath := filepath.Join(setting.AppDataPath, "ssh/gogs.rsa")
	if !com.IsExist(keyPath) {
		os.MkdirAll(filepath.Dir(keyPath), os.ModePerm)
		_, stderr, err := com.ExecCmd("ssh-keygen", "-f", keyPath, "-t", "rsa", "-N", "")
		if err != nil {
			panic(fmt.Sprintf("Fail to generate private key: %v - %s", err, stderr))
		}
		log.Trace("New private key is generateed: %s", keyPath)
	}

	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		panic("Fail to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		panic("Fail to parse private key")
	}
	config.AddHostKey(private)

	go listen(config, port)
}
<commit_msg>Fix #2431 - handle requests waiting for reply<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package ssh

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/Unknwon/com"
	"golang.org/x/crypto/ssh"

	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/setting"
)

func cleanCommand(cmd string) string {
	i := strings.Index(cmd, "git")
	if i == -1 {
		return cmd
	}
	return cmd[i:]
}

func handleServerConn(keyID string, chans <-chan ssh.NewChannel) {
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		ch, reqs, err := newChan.Accept()
		if err != nil {
			log.Error(3, "Error accepting channel: %v", err)
			continue
		}

		go func(in <-chan *ssh.Request) {
			defer ch.Close()
			for req := range in {
				payload := cleanCommand(string(req.Payload))
				switch req.Type {
				case "env":
					args := strings.Split(strings.Replace(payload, "\x00", "", -1), "\v")
					if len(args) != 2 {
						log.Warn("Invalid env arguments: '%#v'", args)
						continue
					}
					args[0] = strings.TrimLeft(args[0], "\x04")
					_, _, err := com.ExecCmdBytes("env", args[0]+"="+args[1])
					if err != nil {
						log.Error(3, "env: %v", err)
						return
					}
				case "exec":
					cmdName := strings.TrimLeft(payload, "'()")
					os.Setenv("SSH_ORIGINAL_COMMAND", cmdName)
					log.Trace("Payload: %v", cmdName)

					args := []string{"serv", "key-" + keyID, "--config=" + setting.CustomConf}
					log.Trace("Arguments: %v", args)
					cmd := exec.Command(setting.AppPath, args...)

					stdout, err := cmd.StdoutPipe()
					if err != nil {
						log.Error(3, "StdoutPipe: %v", err)
						return
					}
					stderr, err := cmd.StderrPipe()
					if err != nil {
						log.Error(3, "StderrPipe: %v", err)
						return
					}
					input, err := cmd.StdinPipe()
					if err != nil {
						log.Error(3, "StdinPipe: %v", err)
						return
					}

					// FIXME: check timeout
					if err = cmd.Start(); err != nil {
						log.Error(3, "Start: %v", err)
						return
					}

					req.Reply(true, nil);
					go io.Copy(input, ch)
					io.Copy(ch, stdout)
					io.Copy(ch.Stderr(), stderr)

					if err = cmd.Wait(); err != nil {
						log.Error(3, "Wait: %v", err)
						return
					}

					ch.SendRequest("exit-status", false, []byte{0, 0, 0, 0})
					return
				default:
				}
			}
		}(reqs)
	}
}

func listen(config *ssh.ServerConfig, port int) {
	listener, err := net.Listen("tcp", "0.0.0.0:"+com.ToStr(port))
	if err != nil {
		panic(err)
	}
	for {
		// Once a ServerConfig has been configured, connections can be accepted.
		conn, err := listener.Accept()
		if err != nil {
			log.Error(3, "Error accepting incoming connection: %v", err)
			continue
		}
		// Before use, a handshake must be performed on the incoming net.Conn.
		sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
		if err != nil {
			log.Error(3, "Error on handshaking: %v", err)
			continue
		}

		log.Trace("Connection from %s (%s)", sConn.RemoteAddr(), sConn.ClientVersion())
		// The incoming Request channel must be serviced.
		go ssh.DiscardRequests(reqs)
		go handleServerConn(sConn.Permissions.Extensions["key-id"], chans)
	}
}

// Listen starts a SSH server listens on given port.
func Listen(port int) {
	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			pkey, err := models.SearchPublicKeyByContent(strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key))))
			if err != nil {
				log.Error(3, "SearchPublicKeyByContent: %v", err)
				return nil, err
			}
			return &ssh.Permissions{Extensions: map[string]string{"key-id": com.ToStr(pkey.ID)}}, nil
		},
	}

	keyPath := filepath.Join(setting.AppDataPath, "ssh/gogs.rsa")
	if !com.IsExist(keyPath) {
		os.MkdirAll(filepath.Dir(keyPath), os.ModePerm)
		_, stderr, err := com.ExecCmd("ssh-keygen", "-f", keyPath, "-t", "rsa", "-N", "")
		if err != nil {
			panic(fmt.Sprintf("Fail to generate private key: %v - %s", err, stderr))
		}
		log.Trace("New private key is generateed: %s", keyPath)
	}

	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		panic("Fail to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		panic("Fail to parse private key")
	}
	config.AddHostKey(private)

	go listen(config, port)
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2017 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package ssh

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/Unknwon/com"
	"golang.org/x/crypto/ssh"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
)

func cleanCommand(cmd string) string {
	i := strings.Index(cmd, "git")
	if i == -1 {
		return cmd
	}
	return cmd[i:]
}

func handleServerConn(keyID string, chans <-chan ssh.NewChannel) {
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		ch, reqs, err := newChan.Accept()
		if err != nil {
			log.Error("Error accepting channel: %v", err)
			continue
		}

		go func(in <-chan *ssh.Request) {
			defer ch.Close()
			for req := range in {
				payload := cleanCommand(string(req.Payload))
				switch req.Type {
				case "env":
					args := strings.Split(strings.Replace(payload, "\x00", "", -1), "\v")
					if len(args) != 2 {
						log.Warn("SSH: Invalid env arguments: '%#v'", args)
						continue
					}
					args[0] = strings.TrimLeft(args[0], "\x04")
					_, _, err := com.ExecCmdBytes("env", args[0]+"="+args[1])
					if err != nil {
						log.Error("env: %v", err)
						return
					}
				case "exec":
					cmdName := strings.TrimLeft(payload, "'()")
					log.Trace("SSH: Payload: %v", cmdName)

					args := []string{"serv", "key-" + keyID, "--config=" + setting.CustomConf}
					log.Trace("SSH: Arguments: %v", args)
					cmd := exec.Command(setting.AppPath, args...)
					cmd.Env = append(
						os.Environ(),
						"SSH_ORIGINAL_COMMAND="+cmdName,
						"SKIP_MINWINSVC=1",
					)

					stdout, err := cmd.StdoutPipe()
					if err != nil {
						log.Error("SSH: StdoutPipe: %v", err)
						return
					}
					stderr, err := cmd.StderrPipe()
					if err != nil {
						log.Error("SSH: StderrPipe: %v", err)
						return
					}
					input, err := cmd.StdinPipe()
					if err != nil {
						log.Error("SSH: StdinPipe: %v", err)
						return
					}

					// FIXME: check timeout
					if err = cmd.Start(); err != nil {
						log.Error("SSH: Start: %v", err)
						return
					}

					req.Reply(true, nil)
					go io.Copy(input, ch)
					io.Copy(ch, stdout)
					io.Copy(ch.Stderr(), stderr)

					if err = cmd.Wait(); err != nil {
						log.Error("SSH: Wait: %v", err)
						return
					}

					ch.SendRequest("exit-status", false, []byte{0, 0, 0, 0})
					return
				default:
				}
			}
		}(reqs)
	}
}

func listen(config *ssh.ServerConfig, host string, port int) {
	listener, err := net.Listen("tcp", host+":"+com.ToStr(port))
	if err != nil {
		log.Fatal("Failed to start SSH server: %v", err)
	}
	for {
		// Once a ServerConfig has been configured, connections can be accepted.
		conn, err := listener.Accept()
		if err != nil {
			log.Error("SSH: Error accepting incoming connection: %v", err)
			continue
		}

		// Before use, a handshake must be performed on the incoming net.Conn.
		// It must be handled in a separate goroutine,
		// otherwise one user could easily block entire loop.
		// For example, user could be asked to trust server key fingerprint and hangs.
		go func() {
			log.Trace("SSH: Handshaking for %s", conn.RemoteAddr())
			sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
			if err != nil {
				if err == io.EOF {
					log.Warn("SSH: Handshaking with %s was terminated: %v", conn.RemoteAddr(), err)
				} else {
					log.Error("SSH: Error on handshaking with %s: %v", conn.RemoteAddr(), err)
				}
				return
			}

			log.Trace("SSH: Connection from %s (%s)", sConn.RemoteAddr(), sConn.ClientVersion())
			// The incoming Request channel must be serviced.
			go ssh.DiscardRequests(reqs)
			go handleServerConn(sConn.Permissions.Extensions["key-id"], chans)
		}()
	}
}

// Listen starts a SSH server listens on given port.
func Listen(host string, port int, ciphers []string, keyExchanges []string, macs []string) {
	config := &ssh.ServerConfig{
		Config: ssh.Config{
			Ciphers:      ciphers,
			KeyExchanges: keyExchanges,
			MACs:         macs,
		},
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			pkey, err := models.SearchPublicKeyByContent(strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key))))
			if err != nil {
				log.Error("SearchPublicKeyByContent: %v", err)
				return nil, err
			}
			return &ssh.Permissions{Extensions: map[string]string{"key-id": com.ToStr(pkey.ID)}}, nil
		},
	}

	keyPath := filepath.Join(setting.AppDataPath, "ssh/gogs.rsa")
	if !com.IsExist(keyPath) {
		filePath := filepath.Dir(keyPath)

		if err := os.MkdirAll(filePath, os.ModePerm); err != nil {
			log.Error("Failed to create dir %s: %v", filePath, err)
		}

		err := GenKeyPair(keyPath)
		if err != nil {
			log.Fatal("Failed to generate private key: %v", err)
		}
		log.Trace("SSH: New private key is generateed: %s", keyPath)
	}

	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		log.Fatal("SSH: Failed to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		log.Fatal("SSH: Failed to parse private key")
	}
	config.AddHostKey(private)

	go listen(config, host, port)
}

// GenKeyPair make a pair of public and private keys for SSH access.
// Public key is encoded in the format for inclusion in an OpenSSH authorized_keys file.
// Private Key generated is PEM encoded
func GenKeyPair(keyPath string) error {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}

	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	f, err := os.OpenFile(keyPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	if err := pem.Encode(f, privateKeyPEM); err != nil {
		return err
	}

	// generate public key
	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return err
	}

	public := ssh.MarshalAuthorizedKey(pub)
	p, err := os.OpenFile(keyPath+".pub", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer p.Close()
	_, err = p.Write(public)
	return err
}
<commit_msg>internal/ssh: ignore env command totally (#6825)<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2017 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package ssh

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/Unknwon/com"
	"golang.org/x/crypto/ssh"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
)

func cleanCommand(cmd string) string {
	i := strings.Index(cmd, "git")
	if i == -1 {
		return cmd
	}
	return cmd[i:]
}

func handleServerConn(keyID string, chans <-chan ssh.NewChannel) {
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		ch, reqs, err := newChan.Accept()
		if err != nil {
			log.Error("Error accepting channel: %v", err)
			continue
		}

		go func(in <-chan *ssh.Request) {
			defer ch.Close()
			for req := range in {
				payload := cleanCommand(string(req.Payload))
				switch req.Type {
				case "exec":
					cmdName := strings.TrimLeft(payload, "'()")
					log.Trace("SSH: Payload: %v", cmdName)

					args := []string{"serv", "key-" + keyID, "--config=" + setting.CustomConf}
					log.Trace("SSH: Arguments: %v", args)
					cmd := exec.Command(setting.AppPath, args...)
					cmd.Env = append(
						os.Environ(),
						"SSH_ORIGINAL_COMMAND="+cmdName,
						"SKIP_MINWINSVC=1",
					)

					stdout, err := cmd.StdoutPipe()
					if err != nil {
						log.Error("SSH: StdoutPipe: %v", err)
						return
					}
					stderr, err := cmd.StderrPipe()
					if err != nil {
						log.Error("SSH: StderrPipe: %v", err)
						return
					}
					input, err := cmd.StdinPipe()
					if err != nil {
						log.Error("SSH: StdinPipe: %v", err)
						return
					}

					// FIXME: check timeout
					if err = cmd.Start(); err != nil {
						log.Error("SSH: Start: %v", err)
						return
					}

					req.Reply(true, nil)
					go io.Copy(input, ch)
					io.Copy(ch, stdout)
					io.Copy(ch.Stderr(), stderr)

					if err = cmd.Wait(); err != nil {
						log.Error("SSH: Wait: %v", err)
						return
					}

					ch.SendRequest("exit-status", false, []byte{0, 0, 0, 0})
					return
				default:
				}
			}
		}(reqs)
	}
}

func listen(config *ssh.ServerConfig, host string, port int) {
	listener, err := net.Listen("tcp", host+":"+com.ToStr(port))
	if err != nil {
		log.Fatal("Failed to start SSH server: %v", err)
	}
	for {
		// Once a ServerConfig has been configured, connections can be accepted.
		conn, err := listener.Accept()
		if err != nil {
			log.Error("SSH: Error accepting incoming connection: %v", err)
			continue
		}

		// Before use, a handshake must be performed on the incoming net.Conn.
		// It must be handled in a separate goroutine,
		// otherwise one user could easily block entire loop.
		// For example, user could be asked to trust server key fingerprint and hangs.
		go func() {
			log.Trace("SSH: Handshaking for %s", conn.RemoteAddr())
			sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
			if err != nil {
				if err == io.EOF {
					log.Warn("SSH: Handshaking with %s was terminated: %v", conn.RemoteAddr(), err)
				} else {
					log.Error("SSH: Error on handshaking with %s: %v", conn.RemoteAddr(), err)
				}
				return
			}

			log.Trace("SSH: Connection from %s (%s)", sConn.RemoteAddr(), sConn.ClientVersion())
			// The incoming Request channel must be serviced.
			go ssh.DiscardRequests(reqs)
			go handleServerConn(sConn.Permissions.Extensions["key-id"], chans)
		}()
	}
}

// Listen starts a SSH server listens on given port.
func Listen(host string, port int, ciphers []string, keyExchanges []string, macs []string) {
	config := &ssh.ServerConfig{
		Config: ssh.Config{
			Ciphers:      ciphers,
			KeyExchanges: keyExchanges,
			MACs:         macs,
		},
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			pkey, err := models.SearchPublicKeyByContent(strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key))))
			if err != nil {
				log.Error("SearchPublicKeyByContent: %v", err)
				return nil, err
			}
			return &ssh.Permissions{Extensions: map[string]string{"key-id": com.ToStr(pkey.ID)}}, nil
		},
	}

	keyPath := filepath.Join(setting.AppDataPath, "ssh/gogs.rsa")
	if !com.IsExist(keyPath) {
		filePath := filepath.Dir(keyPath)

		if err := os.MkdirAll(filePath, os.ModePerm); err != nil {
			log.Error("Failed to create dir %s: %v", filePath, err)
		}

		err := GenKeyPair(keyPath)
		if err != nil {
			log.Fatal("Failed to generate private key: %v", err)
		}
		log.Trace("SSH: New private key is generateed: %s", keyPath)
	}

	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		log.Fatal("SSH: Failed to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		log.Fatal("SSH: Failed to parse private key")
	}
	config.AddHostKey(private)

	go listen(config, host, port)
}

// GenKeyPair make a pair of public and private keys for SSH access.
// Public key is encoded in the format for inclusion in an OpenSSH authorized_keys file.
// Private Key generated is PEM encoded
func GenKeyPair(keyPath string) error {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}

	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	f, err := os.OpenFile(keyPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	if err := pem.Encode(f, privateKeyPEM); err != nil {
		return err
	}

	// generate public key
	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return err
	}

	public := ssh.MarshalAuthorizedKey(pub)
	p, err := os.OpenFile(keyPath+".pub", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer p.Close()
	_, err = p.Write(public)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"net/http"
	"net/http/pprof"
	"os"

	"github.com/ViaQ/logerr/log"
	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
	// to ensure that exec-entrypoint and run can make use of them.
	_ "k8s.io/client-go/plugin/pkg/client/auth"

	lokiv1beta1 "github.com/grafana/loki-operator/api/v1beta1"
	"github.com/grafana/loki-operator/controllers"
	"github.com/grafana/loki-operator/internal/manifests"
	"github.com/grafana/loki-operator/internal/metrics"
	configv1 "github.com/openshift/api/config/v1"
	routev1 "github.com/openshift/api/route/v1"
	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"

	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	// +kubebuilder:scaffold:imports
)

var scheme = runtime.NewScheme()

func init() {
	utilruntime.Must(clientgoscheme.AddToScheme(scheme))

	utilruntime.Must(lokiv1beta1.AddToScheme(scheme))
	// +kubebuilder:scaffold:scheme
}

func main() {
	var (
		metricsAddr              string
		enableLeaderElection     bool
		probeAddr                string
		enableCertSigning        bool
		enableServiceMonitors    bool
		enableTLSServiceMonitors bool
		enableGateway            bool
		enableGatewayRoute       bool
	)

	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
	flag.BoolVar(&enableLeaderElection, "leader-elect", false,
		"Enable leader election for controller manager. "+
			"Enabling this will ensure there is only one active controller manager.")
	flag.BoolVar(&enableCertSigning, "with-cert-signing-service", false,
		"Enables features in an Openshift cluster.")
	flag.BoolVar(&enableServiceMonitors, "with-service-monitors", false, "Enables service monitoring")
	flag.BoolVar(&enableTLSServiceMonitors, "with-tls-service-monitors", false,
		"Enables loading of a prometheus service monitor.")
	flag.BoolVar(&enableGateway, "with-lokistack-gateway", false,
		"Enables the manifest creation for the entire lokistack-gateway.")
	flag.BoolVar(&enableGatewayRoute, "with-lokistack-gateway-route", false,
		"Enables the usage of Route for the lokistack-gateway instead of Ingress (OCP Only!)")
	flag.Parse()

	log.Init("loki-operator")
	ctrl.SetLogger(log.GetLogger())

	if enableServiceMonitors || enableTLSServiceMonitors {
		utilruntime.Must(monitoringv1.AddToScheme(scheme))
	}

	if enableGateway {
		utilruntime.Must(configv1.AddToScheme(scheme))

		if enableGatewayRoute {
			utilruntime.Must(routev1.AddToScheme(scheme))
		}
	}

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:                 scheme,
		MetricsBindAddress:     metricsAddr,
		Port:                   9443,
		HealthProbeBindAddress: probeAddr,
		LeaderElection:         enableLeaderElection,
		LeaderElectionID:       "e3716011.grafana.com",
	})
	if err != nil {
		log.Error(err, "unable to start manager")
		os.Exit(1)
	}

	featureFlags := manifests.FeatureFlags{
		EnableCertificateSigningService: enableCertSigning,
		EnableServiceMonitors:           enableServiceMonitors,
		EnableTLSServiceMonitorConfig:   enableTLSServiceMonitors,
		EnableGateway:                   enableGateway,
		EnableGatewayRoute:              enableGatewayRoute,
	}

	if err = (&controllers.LokiStackReconciler{
		Client: mgr.GetClient(),
		Log:    log.WithName("controllers").WithName("LokiStack"),
		Scheme: mgr.GetScheme(),
		Flags:  featureFlags,
	}).SetupWithManager(mgr); err != nil {
		log.Error(err, "unable to create controller", "controller", "LokiStack")
		os.Exit(1)
	}
	// +kubebuilder:scaffold:builder

	if err = mgr.AddHealthzCheck("health", healthz.Ping); err != nil {
		log.Error(err, "unable to set up health check")
		os.Exit(1)
	}
	if err = mgr.AddReadyzCheck("check", healthz.Ping); err != nil {
		log.Error(err, "unable to set up ready check")
		os.Exit(1)
	}

	log.Info("registering metrics")
	metrics.RegisterMetricCollectors()

	log.Info("Registring profiling endpoints.")
	err = registerProfiler(mgr)
	if err != nil {
		log.Error(err, "failed to register extra pprof handler")
		os.Exit(1)
	}

	log.Info("starting manager")
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		log.Error(err, "problem running manager")
		os.Exit(1)
	}
}

func registerProfiler(m ctrl.Manager) error {
	endpoints := map[string]http.HandlerFunc{
		"/debug/pprof/":        pprof.Index,
		"/debug/pprof/cmdline": pprof.Cmdline,
		"/debug/pprof/profile": pprof.Profile,
		"/debug/pprof/symbol":  pprof.Symbol,
		"/debug/pprof/trace":   pprof.Trace,
	}

	for path, handler := range endpoints {
		err := m.AddMetricsExtraHandler(path, handler)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Fix typo 'Registring'->'Registering' (#4979)<commit_after>package main

import (
	"flag"
	"net/http"
	"net/http/pprof"
	"os"

	"github.com/ViaQ/logerr/log"
	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
	// to ensure that exec-entrypoint and run can make use of them.
	_ "k8s.io/client-go/plugin/pkg/client/auth"

	lokiv1beta1 "github.com/grafana/loki-operator/api/v1beta1"
	"github.com/grafana/loki-operator/controllers"
	"github.com/grafana/loki-operator/internal/manifests"
	"github.com/grafana/loki-operator/internal/metrics"
	configv1 "github.com/openshift/api/config/v1"
	routev1 "github.com/openshift/api/route/v1"
	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"

	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	// +kubebuilder:scaffold:imports
)

var scheme = runtime.NewScheme()

func init() {
	utilruntime.Must(clientgoscheme.AddToScheme(scheme))

	utilruntime.Must(lokiv1beta1.AddToScheme(scheme))
	// +kubebuilder:scaffold:scheme
}

func main() {
	var (
		metricsAddr              string
		enableLeaderElection     bool
		probeAddr                string
		enableCertSigning        bool
		enableServiceMonitors    bool
		enableTLSServiceMonitors bool
		enableGateway            bool
		enableGatewayRoute       bool
	)

	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
	flag.BoolVar(&enableLeaderElection, "leader-elect", false,
		"Enable leader election for controller manager. "+
			"Enabling this will ensure there is only one active controller manager.")
	flag.BoolVar(&enableCertSigning, "with-cert-signing-service", false,
		"Enables features in an Openshift cluster.")
	flag.BoolVar(&enableServiceMonitors, "with-service-monitors", false, "Enables service monitoring")
	flag.BoolVar(&enableTLSServiceMonitors, "with-tls-service-monitors", false,
		"Enables loading of a prometheus service monitor.")
	flag.BoolVar(&enableGateway, "with-lokistack-gateway", false,
		"Enables the manifest creation for the entire lokistack-gateway.")
	flag.BoolVar(&enableGatewayRoute, "with-lokistack-gateway-route", false,
		"Enables the usage of Route for the lokistack-gateway instead of Ingress (OCP Only!)")
	flag.Parse()

	log.Init("loki-operator")
	ctrl.SetLogger(log.GetLogger())

	if enableServiceMonitors || enableTLSServiceMonitors {
		utilruntime.Must(monitoringv1.AddToScheme(scheme))
	}

	if enableGateway {
		utilruntime.Must(configv1.AddToScheme(scheme))

		if enableGatewayRoute {
			utilruntime.Must(routev1.AddToScheme(scheme))
		}
	}

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:                 scheme,
		MetricsBindAddress:     metricsAddr,
		Port:                   9443,
		HealthProbeBindAddress: probeAddr,
		LeaderElection:         enableLeaderElection,
		LeaderElectionID:       "e3716011.grafana.com",
	})
	if err != nil {
		log.Error(err, "unable to start manager")
		os.Exit(1)
	}

	featureFlags := manifests.FeatureFlags{
		EnableCertificateSigningService: enableCertSigning,
		EnableServiceMonitors:           enableServiceMonitors,
		EnableTLSServiceMonitorConfig:   enableTLSServiceMonitors,
		EnableGateway:                   enableGateway,
		EnableGatewayRoute:              enableGatewayRoute,
	}

	if err = (&controllers.LokiStackReconciler{
		Client: mgr.GetClient(),
		Log:    log.WithName("controllers").WithName("LokiStack"),
		Scheme: mgr.GetScheme(),
		Flags:  featureFlags,
	}).SetupWithManager(mgr); err != nil {
		log.Error(err, "unable to create controller", "controller", "LokiStack")
		os.Exit(1)
	}
	// +kubebuilder:scaffold:builder

	if err = mgr.AddHealthzCheck("health", healthz.Ping); err != nil {
		log.Error(err, "unable to set up health check")
		os.Exit(1)
	}
	if err = mgr.AddReadyzCheck("check", healthz.Ping); err != nil {
		log.Error(err, "unable to set up ready check")
		os.Exit(1)
	}

	log.Info("registering metrics")
	metrics.RegisterMetricCollectors()

	log.Info("Registering profiling endpoints.")
	err = registerProfiler(mgr)
	if err != nil {
		log.Error(err, "failed to register extra pprof handler")
		os.Exit(1)
	}

	log.Info("starting manager")
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		log.Error(err, "problem running manager")
		os.Exit(1)
	}
}

func registerProfiler(m ctrl.Manager) error {
	endpoints := map[string]http.HandlerFunc{
		"/debug/pprof/":        pprof.Index,
		"/debug/pprof/cmdline": pprof.Cmdline,
		"/debug/pprof/profile": pprof.Profile,
		"/debug/pprof/symbol":  pprof.Symbol,
		"/debug/pprof/trace":   pprof.Trace,
	}

	for path, handler := range endpoints {
		err := m.AddMetricsExtraHandler(path, handler)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/jmcvetta/neoism"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"
)

func main() {

	var err error
	db, err = neoism.Connect("http://localhost:7474/db/data")
	if err != nil {
		panic(err)
	}

	ensureIndexes(db)

	writeQueue = make(chan organisation, 2048)

	port := 8082

	m := mux.NewRouter()
	http.Handle("/", m)

	m.HandleFunc("/organisations/{uuid}", idWriteHandler).Methods("PUT")
	m.HandleFunc("/organisations/", allWriteHandler).Methods("PUT")

	go func() {
		log.Printf("listening on %d", port)
		if err := http.ListenAndServe(fmt.Sprintf(":%d", port), nil); err != nil {
			log.Printf("web stuff failed: %v\n", err)
		}
	}()

	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		orgWriteLoop()
		wg.Done()
	}()

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	// wait for ctrl-c
	<-c
	close(writeQueue)
	wg.Wait()
	println("exiting")

}

func ensureIndexes(db *neoism.Database) {

	indexes, err := db.Indexes("Organisation")
	if err != nil {
		panic(err)
	}
	if len(indexes) == 0 {
		if _, err := db.CreateIndex("Organisation", "uuid"); err != nil {
			panic(err)
		}
		if _, err := db.CreateIndex("Industry", "uuid"); err != nil {
			panic(err)
		}
	}
}

var db *neoism.Database

var writeQueue chan organisation

func orgWriteLoop() {
	var qs []*neoism.CypherQuery

	timer := time.NewTimer(1 * time.Second)

	defer println("asdasd")
	for {
		select {
		case o, ok := <-writeQueue:
			if !ok {
				return
			}
			for _, q := range toQueries(o) {
				qs = append(qs, q)
			}
			if len(qs) < 1024 {
				timer.Reset(1 * time.Second)
				continue
			}
		case <-timer.C:
		}
		if len(qs) > 0 {
			fmt.Printf("writing batch of %d\n", len(qs))
			err := db.CypherBatch(qs)
			if err != nil {
				panic(err)
			}
			fmt.Printf("wrote batch of %d\n", len(qs))
			qs = qs[0:0]
			timer.Stop()
		}
	}
}

func allWriteHandler(w http.ResponseWriter, r *http.Request) {

	dec := json.NewDecoder(r.Body)

	for {
		var o organisation
		err := dec.Decode(&o)
		if err == io.ErrUnexpectedEOF {
			println("eof")
			return
		}
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		writeQueue <- o
	}

	w.WriteHeader(http.StatusAccepted)
}

func idWriteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	UUID := vars["uuid"]

	var o organisation
	dec := json.NewDecoder(r.Body)
	err := dec.Decode(&o)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if o.UUID != UUID {
		fmt.Printf("%v\n", o)
		http.Error(w, fmt.Sprintf("id does not match: %v %v", o.UUID, UUID), http.StatusBadRequest)
		return
	}

	writeQueue <- o

	w.WriteHeader(http.StatusAccepted)
}

func toQueries(o organisation) []*neoism.CypherQuery {
	props := toProps(o)

	var queries []*neoism.CypherQuery

	queries = append(queries, &neoism.CypherQuery{
		Statement: `
			MERGE (n:Organisation {uuid: {uuid}})
			SET n = {allProps}
			RETURN n
		`,
		Parameters: map[string]interface{}{
			"uuid":     o.UUID,
			"allProps": props,
		},
	})

	t := string(o.Type)
	if t != "Organisation" && t != "" {
		queries = append(queries, &neoism.CypherQuery{
			Statement: fmt.Sprintf("MERGE (n:Organisation {uuid: {uuid}}) SET n :%s", t),
			Parameters: map[string]interface{}{
				"uuid": o.UUID,
			},
		})
	}

	if o.ParentOrganisation != "" {
		queries = append(queries, &neoism.CypherQuery{
			Statement: `
			MERGE (n:Organisation {uuid: {uuid}})
			MERGE (p:Organisation {uuid: {puuid}})
			MERGE (n)-[r:SUB_ORG_OF]->(p)
		`,
			Parameters: map[string]interface{}{
				"uuid":  o.UUID,
				"puuid": o.ParentOrganisation,
			},
		})
	}

	if o.IndustryClassification != "" {
		queries = append(queries, &neoism.CypherQuery{
			Statement: `
			MERGE (n:Organisation {uuid: {uuid}})
			MERGE (ic:Industry {uuid: {icuuid}})
			MERGE (n)-[r:IN_INDUSTRY]->(ic)
		`,
			Parameters: map[string]interface{}{
				"uuid":   o.UUID,
				"icuuid": o.IndustryClassification,
			},
		})

	}

	return queries
}

func toProps(o organisation) neoism.Props {
	p := map[string]interface{}{
		"uuid": o.UUID,
	}

	if o.Extinct == true {
		p["extinct"] = true
	}
	if o.FormerNames != nil && len(o.FormerNames) != 0 {
		p["formerNames"] = o.FormerNames
	}
	if o.HiddenLabel != "" {
		p["hiddenLabel"] = o.HiddenLabel
	}
	if o.LegalName != "" {
		p["legalName"] = o.LegalName
	}
	if o.LocalNames != nil && len(o.LocalNames) != 0 {
		p["localNames"] = o.LocalNames
	}
	if o.ProperName != "" {
		p["properName"] = o.ProperName
	}
	if o.ShortName != "" {
		p["shortName"] = o.ShortName
	}
	if o.TradeNames != nil && len(o.TradeNames) != 0 {
		p["tradeNames"] = o.TradeNames
	}
	p["uuid"] = o.UUID

	return neoism.Props(p)
}
<commit_msg>paramater for url<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/jmcvetta/neoism"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"
)

func main() {

	neoUrl := os.Getenv("NEO_URL")
	if neoUrl == "" {
		log.Println("no $NEO_URL set, defaulting to local")
		neoUrl = "http://localhost:7474/db/data"
	}
	log.Printf("connecting to %s\n", neoUrl)

	var err error
	db, err = neoism.Connect(neoUrl)
	if err != nil {
		panic(err)
	}

	ensureIndexes(db)

	writeQueue = make(chan organisation, 2048)

	port := 8082

	m := mux.NewRouter()
	http.Handle("/", m)

	m.HandleFunc("/organisations/{uuid}", idWriteHandler).Methods("PUT")
	m.HandleFunc("/organisations/", allWriteHandler).Methods("PUT")

	go func() {
		log.Printf("listening on %d", port)
		if err := http.ListenAndServe(fmt.Sprintf(":%d", port), nil); err != nil {
			log.Printf("web stuff failed: %v\n", err)
		}
	}()

	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		orgWriteLoop()
		wg.Done()
	}()

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	// wait for ctrl-c
	<-c
	close(writeQueue)
	wg.Wait()
	println("exiting")

}

func ensureIndexes(db *neoism.Database) {

	indexes, err := db.Indexes("Organisation")
	if err != nil {
		panic(err)
	}
	if len(indexes) == 0 {
		if _, err := db.CreateIndex("Organisation", "uuid"); err != nil {
			panic(err)
		}
		if _, err := db.CreateIndex("Industry", "uuid"); err != nil {
			panic(err)
		}
	}
}

var db *neoism.Database

var writeQueue chan organisation

func orgWriteLoop() {
	var qs []*neoism.CypherQuery

	timer := time.NewTimer(1 * time.Second)

	defer println("asdasd")
	for {
		select {
		case o, ok := <-writeQueue:
			if !ok {
				return
			}
			for _, q := range toQueries(o) {
				qs = append(qs, q)
			}
			if len(qs) < 1024 {
				timer.Reset(1 * time.Second)
				continue
			}
		case <-timer.C:
		}
		if len(qs) > 0 {
			fmt.Printf("writing batch of %d\n", len(qs))
			err := db.CypherBatch(qs)
			if err != nil {
				panic(err)
			}
			fmt.Printf("wrote batch of %d\n", len(qs))
			qs = qs[0:0]
			timer.Stop()
		}
	}
}

func allWriteHandler(w http.ResponseWriter, r *http.Request) {

	dec := json.NewDecoder(r.Body)

	for {
		var o organisation
		err := dec.Decode(&o)
		if err == io.ErrUnexpectedEOF {
			println("eof")
			return
		}
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		writeQueue <- o
	}

	w.WriteHeader(http.StatusAccepted)
}

func idWriteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	UUID := vars["uuid"]

	var o organisation
	dec := json.NewDecoder(r.Body)
	err := dec.Decode(&o)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if o.UUID != UUID {
		fmt.Printf("%v\n", o)
		http.Error(w, fmt.Sprintf("id does not match: %v %v", o.UUID, UUID), http.StatusBadRequest)
		return
	}

	writeQueue <- o

	w.WriteHeader(http.StatusAccepted)
}

func toQueries(o organisation) []*neoism.CypherQuery {
	props := toProps(o)

	var queries []*neoism.CypherQuery

	queries = append(queries, &neoism.CypherQuery{
		Statement: `
			MERGE (n:Organisation {uuid: {uuid}})
			SET n = {allProps}
			RETURN n
		`,
		Parameters: map[string]interface{}{
			"uuid":     o.UUID,
			"allProps": props,
		},
	})

	t := string(o.Type)
	if t != "Organisation" && t != "" {
		queries = append(queries, &neoism.CypherQuery{
			Statement: fmt.Sprintf("MERGE (n:Organisation {uuid: {uuid}}) SET n :%s", t),
			Parameters: map[string]interface{}{
				"uuid": o.UUID,
			},
		})
	}

	if o.ParentOrganisation != "" {
		queries = append(queries, &neoism.CypherQuery{
			Statement: `
			MERGE (n:Organisation {uuid: {uuid}})
			MERGE (p:Organisation {uuid: {puuid}})
			MERGE (n)-[r:SUB_ORG_OF]->(p)
		`,
			Parameters: map[string]interface{}{
				"uuid":  o.UUID,
				"puuid": o.ParentOrganisation,
			},
		})
	}

	if o.IndustryClassification != "" {
		queries = append(queries, &neoism.CypherQuery{
			Statement: `
			MERGE (n:Organisation {uuid: {uuid}})
			MERGE (ic:Industry {uuid: {icuuid}})
			MERGE (n)-[r:IN_INDUSTRY]->(ic)
		`,
			Parameters: map[string]interface{}{
				"uuid":   o.UUID,
				"icuuid": o.IndustryClassification,
			},
		})

	}

	return queries
}

func toProps(o organisation) neoism.Props {
	p := map[string]interface{}{
		"uuid": o.UUID,
	}

	if o.Extinct == true {
		p["extinct"] = true
	}
	if o.FormerNames != nil && len(o.FormerNames) != 0 {
		p["formerNames"] = o.FormerNames
	}
	if o.HiddenLabel != "" {
		p["hiddenLabel"] = o.HiddenLabel
	}
	if o.LegalName != "" {
		p["legalName"] = o.LegalName
	}
	if o.LocalNames != nil && len(o.LocalNames) != 0 {
		p["localNames"] = o.LocalNames
	}
	if o.ProperName != "" {
		p["properName"] = o.ProperName
	}
	if o.ShortName != "" {
		p["shortName"] = o.ShortName
	}
	if o.TradeNames != nil && len(o.TradeNames) != 0 {
		p["tradeNames"] = o.TradeNames
	}
	p["uuid"] = o.UUID

	return neoism.Props(p)
}
<|endoftext|>


<commit_before>package ostent

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
)

// NewErrorLog creates a logger and returns a func to defer.
func NewErrorLog() (*log.Logger, func() error) {
	logger := logrus.New()
	logger.Formatter = &logrus.TextFormatter{
		FullTimestamp: true,
	}
	errlog := logger.Writer()
	return log.New(errlog, "", 0), errlog.Close
}

type Access struct {
	TaggedBin bool
	Logger    *logrus.Logger
	Hosts     struct {
		Mutex sync.Mutex
		Hosts map[string]struct{}
	}
}

func NewAccess(taggedbin bool) *Access {
	logger := logrus.New()
	/* logger.Formatter = &logrus.TextFormatter{
		// DisableTimestamp:true,
		FullTimestamp:   true,
		TimestampFormat: "02/Jan/2006:15:04:05 -0700",
	} // */
	logger.Formatter = &AccessFormatter{}
	a := &Access{
		TaggedBin: taggedbin,
		Logger:    logger,
	}
	a.Hosts.Hosts = make(map[string]struct{})
	return a
}

func (a *Access) Constructor(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		start := time.Now()
		rsp := &Responding{ResponseWriter: w}
		handler.ServeHTTP(rsp, req)

		if entry := a.Entry(start, *rsp, req); a.TaggedBin && rsp.StatusGood() {
			a.InfoForGood(entry, req)
		} else {
			entry.Info("")
		}
	})
}

func (a *Access) Entry(start time.Time, rsp Responding, req *http.Request) *logrus.Entry {
	since := ZEROTIME.Add(time.Since(start)).Format("5.0000s")
	host := RemoteHost(req)

	uri := req.URL.Path // OR req.RequestURI ??
	if req.Form != nil && len(req.Form) > 0 {
		uri += "?" + req.Form.Encode()
	}
	return a.Logger.WithFields(logrus.Fields{
		"code":      rsp.Status,
		"duration":  since,
		"host":      host,
		"method":    req.Method,
		"proto":     req.Proto,
		"referer":   req.Header.Get("Referer"),
		"size":      rsp.Size,
		"uri":       uri,
		"useragent": req.Header.Get("User-Agent"),
	})
}

func (a *Access) InfoForGood(entry *logrus.Entry, req *http.Request) {
	if host := RemoteHost(req); !a.Seen(host) {
		entry.Data["comment"] = fmt.Sprintf(
			";last info-logged successful request from %s", host)
		entry.Info("")
	} else {
		entry.Debug("")
	}
}

func (a *Access) Seen(host string) bool {
	a.Hosts.Mutex.Lock()
	defer a.Hosts.Mutex.Unlock()
	if _, ok := a.Hosts.Hosts[host]; ok {
		return true
	}
	a.Hosts.Hosts[host] = struct{}{}
	return false
}

func RemoteHost(req *http.Request) string {
	host, _, err := net.SplitHostPort(req.RemoteAddr)
	if err != nil {
		host = req.RemoteAddr
	}
	return host
}

var ZEROTIME, _ = time.Parse("15:04:05", "00:00:00")

type AccessFormatter struct{}

func (af *AccessFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	f := entry.Data
	comment, okc := f["comment"]
	if !okc {
		comment = ""
	} else if s, ok := comment.(string); ok && s != "" {
		comment = "\t" + s
	}
	echo := func(v interface{}) interface{} {
		if s, ok := v.(string); ok && s == "" {
			return "-"
		}
		return v
	}
	return []byte(fmt.Sprintf("%s - - [%s] %#v %d %d %#v %#v\t;%s%s\n",
		f["host"],
		entry.Time.Format("02/Jan/2006:15:04:05 -0700"),
		fmt.Sprintf("%s %s %s", f["method"], f["uri"], f["proto"]),
		f["code"],
		f["size"],
		echo(f["referer"]),
		echo(f["useragent"]),
		f["duration"],
		comment)), nil
}

type Responding struct {
	http.ResponseWriter
	http.Flusher // ?
	Status       int
	Size         int
}

func (r Responding) StatusGood() bool {
	var good = []int{
		http.StatusSwitchingProtocols, // 101
		http.StatusOK,                 // 200
		http.StatusNotModified,        // 304
	}
	for _, g := range good {
		if r.Status == g {
			return true
		}
	}
	return false
}

func (r *Responding) Flush() {
	if f, ok := r.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (r *Responding) WriteHeader(status int) {
	r.ResponseWriter.WriteHeader(status)
	r.Status = status
}

func (r *Responding) Write(text []byte) (int, error) {
	if r.Status == 0 { // generic approach to Write-ing before WriteHeader call
		r.WriteHeader(http.StatusOK)
	}
	size, err := r.ResponseWriter.Write(text)
	if err == nil {
		r.Size += size
	}
	return size, err
}

func (r *Responding) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := r.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, fmt.Errorf("Responding's ResponseWriter doesn't support the Hijacker interface")
}
<commit_msg>log cleanup<commit_after>package ostent

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
)

// NewErrorLog creates a logger and returns a func to defer.
func NewErrorLog() (*log.Logger, func() error) {
	xlog := logrus.New() // into os.Stderr
	xlog.Formatter = &logrus.TextFormatter{FullTimestamp: true}
	// , TimestampFormat: "02/Jan/2006:15:04:05 -0700",
	wlog := xlog.Writer()
	return log.New(wlog, "", 0), wlog.Close
}

type Access struct {
	TaggedBin bool
	Log       *logrus.Logger
	Hosts     struct {
		Mutex sync.Mutex
		Hosts map[string]struct{}
	}
}

func NewAccess(taggedbin bool) *Access {
	a := &Access{TaggedBin: taggedbin}
	a.Hosts.Hosts = make(map[string]struct{})
	a.Log = logrus.New() // into os.Stderr
	a.Log.Formatter = &AccessFormat{}
	return a
}

func (a *Access) Constructor(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		start := time.Now()
		rsp := &Responding{ResponseWriter: w}
		handler.ServeHTTP(rsp, req)
		a.DoLog(start, req, rsp)
	})
}

// DoLog logs the arguments specifics.
func (a *Access) DoLog(start time.Time, req *http.Request, rsp *Responding) {
	since := ZeroTime.Add(time.Since(start))
	host, _, err := net.SplitHostPort(req.RemoteAddr)
	if err != nil {
		host = req.RemoteAddr
	}
	uri := req.URL.Path // OR req.RequestURI ??
	if req.Form != nil && len(req.Form) > 0 {
		uri += "?" + req.Form.Encode()
	}
	// msg is unused by formatter
	msg, entry := "msg", a.Log.WithFields(logrus.Fields{
		// every value must be string for formatter
		"host":      host,
		"method":    req.Method,
		"uri":       uri,
		"proto":     req.Proto,
		"code":      fmt.Sprintf("%d", rsp.Status),
		"size":      fmt.Sprintf("%d", rsp.Size),
		"referer":   req.Header.Get("Referer"),
		"useragent": req.Header.Get("User-Agent"),
		"duration":  since.Format("5.0000s"),
	})
	if !a.TaggedBin || !rsp.StatusGood() {
		entry.Info(msg)
		return
	}
	if a.Seen(host) {
		entry.Debug(msg)
		return
	}
	entry.Data["comment"] = fmt.Sprintf(
		";last logged successful request from %s", host)
	entry.Info(msg)
}

// Seen returns whether host has been recorded.
// The record is created if it did not exist.
func (a *Access) Seen(host string) bool {
	a.Hosts.Mutex.Lock()
	defer a.Hosts.Mutex.Unlock()
	if _, ok := a.Hosts.Hosts[host]; ok {
		return true
	}
	a.Hosts.Hosts[host] = struct{}{}
	return false
}

// ZeroTime is zero time for formatting duration from it.
var ZeroTime, _ = time.Parse("15:04:05", "00:00:00")

// AccessFormat is a dummy struct with Format method.
type AccessFormat struct{}

// Format conforms to logrus.Formatter interface.
func (af *AccessFormat) Format(entry *logrus.Entry) ([]byte, error) {
	buf := bytes.NewBuffer(make([]byte, 256))
	WriteField(buf, entry.Data["host"], " - - [")
	WriteField(buf, entry.Time.Format("02/Jan/2006:15:04:05 -0700"), "] ")
	WriteField(buf, entry.Data["method"])
	WriteField(buf, entry.Data["uri"])
	WriteField(buf, entry.Data["proto"])
	WriteField(buf, entry.Data["code"])
	WriteField(buf, entry.Data["size"])
	WriteField(buf, strconv.Quote(entry.Data["referer"].(string)))
	WriteField(buf, strconv.Quote(entry.Data["useragent"].(string)), "\t;")
	WriteField(buf, entry.Data["duration"], "")
	if comment, ok := entry.Data["comment"]; ok {
		WriteField(buf, "\t", comment.(string))
	}
	WriteField(buf, "\n", "")
	return buf.Bytes(), nil
}

// WriteField is a helper to write val and optional posts[0] or a whitespace.
// The name relate to use with logrus.Entry.Data fields, no other Field relation.
func WriteField(buf *bytes.Buffer, val interface{}, posts ...string) {
	if s := val.(string); s == "" {
		buf.WriteString("-")
	} else if s == `""` {
		buf.WriteString(`"-"`)
	} else {
		buf.WriteString(s)
	}
	if len(posts) == 0 {
		buf.WriteString(" ")
	} else { // just [0] is written
		buf.WriteString(posts[0])
	}
}

type Responding struct {
	http.ResponseWriter
	http.Flusher // ?
	Status       int
	Size         int
}

func (r Responding) StatusGood() bool {
	var good = []int{
		http.StatusSwitchingProtocols, // 101
		http.StatusOK,                 // 200
		http.StatusNotModified,        // 304
	}
	for _, g := range good {
		if r.Status == g {
			return true
		}
	}
	return false
}

func (r *Responding) Flush() {
	if f, ok := r.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (r *Responding) WriteHeader(status int) {
	r.ResponseWriter.WriteHeader(status)
	r.Status = status
}

func (r *Responding) Write(text []byte) (int, error) {
	if r.Status == 0 { // generic approach to Write-ing before WriteHeader call
		r.WriteHeader(http.StatusOK)
	}
	size, err := r.ResponseWriter.Write(text)
	if err == nil {
		r.Size += size
	}
	return size, err
}

func (r *Responding) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := r.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, fmt.Errorf("Responding's ResponseWriter doesn't support the Hijacker interface")
}
<|endoftext|>


<commit_before>package stages

import (
	"fmt"
	"reflect"

	"github.com/vito/booklit"
	"github.com/vito/booklit/ast"
)

type Evaluate struct {
	Plugins []booklit.Plugin

	Result booklit.Content
}

func (eval *Evaluate) VisitString(str ast.String) error {
	eval.Result = booklit.Append(eval.Result, booklit.String(str))
	return nil
}

func (eval *Evaluate) VisitSequence(seq ast.Sequence) error {
	for _, node := range seq {
		err := node.Visit(eval)
		if err != nil {
			return err
		}
	}

	return nil
}

func (eval *Evaluate) VisitParagraph(node ast.Paragraph) error {
	previous := eval.Result

	para := booklit.Paragraph{}
	for _, sentence := range node {
		eval.Result = nil

		err := sentence.Visit(eval)
		if err != nil {
			return err
		}

		if eval.Result != nil {
			para = append(para, eval.Result)
		}
	}

	eval.Result = nil

	if len(para) == 0 {
		// paragraph resulted in no content (e.g. an invoke with no return value)
		eval.Result = previous
		return nil
	}

	if len(para) == 1 && !para[0].IsSentence() {
		// paragraph resulted in block content (e.g. a section)
		eval.Result = booklit.Append(previous, para[0])
		return nil
	}

	eval.Result = booklit.Append(previous, para)

	return nil
}

func (eval *Evaluate) VisitPreformatted(node ast.Preformatted) error {
	previous := eval.Result

	pre := booklit.Preformatted{}
	for _, sentence := range node {
		eval.Result = nil

		err := sentence.Visit(eval)
		if err != nil {
			return err
		}

		if eval.Result != nil {
			pre = append(pre, eval.Result)
		}
	}

	eval.Result = booklit.Append(previous, pre)

	return nil
}

func (eval *Evaluate) VisitInvoke(invoke ast.Invoke) error {
	methodName := invoke.Method()

	var method reflect.Value
	for _, p := range eval.Plugins {
		value := reflect.ValueOf(p)
		method = value.MethodByName(methodName)
		if method.IsValid() {
			break
		}
	}

	if !method.IsValid() {
		return fmt.Errorf("undefined method: %s", method)
	}

	rawArgs := invoke.Arguments

	argc := method.Type().NumIn()
	if method.Type().IsVariadic() {
		argc--

		if len(rawArgs) < argc {
			return fmt.Errorf("argument count mismatch for %s: given %d, need at least %d", methodName, len(rawArgs), argc)
		}
	} else {
		if len(rawArgs) != argc {
			return fmt.Errorf("argument count mismatch for %s: given %d, need %d", methodName, len(rawArgs), argc)
		}
	}

	argv := make([]reflect.Value, argc)
	for i := 0; i < argc; i++ {
		t := method.Type().In(i)
		arg, err := eval.convert(t, rawArgs[i])
		if err != nil {
			return err
		}

		argv[i] = arg
	}

	if method.Type().IsVariadic() {
		variadic := rawArgs[argc:]
		variadicType := method.Type().In(argc)

		subType := variadicType.Elem()
		for _, varg := range variadic {
			arg, err := eval.convert(subType, varg)
			if err != nil {
				return err
			}

			argv = append(argv, arg)
		}
	}

	result := method.Call(argv)
	switch len(result) {
	case 0:
		return nil
	case 1:
		last := result[0]
		switch v := last.Interface().(type) {
		case nil:
		case error:
			return v
		case booklit.Content:
			eval.Result = booklit.Append(eval.Result, v)
		default:
			return fmt.Errorf("unknown return type: %T", v)
		}
	case 2:
		first := result[0]
		switch v := first.Interface().(type) {
		case booklit.Content:
			eval.Result = booklit.Append(eval.Result, v)
		default:
			return fmt.Errorf("unknown first return type: %T", v)
		}

		last := result[1]
		switch v := last.Interface().(type) {
		case nil:
		case error:
			return v
		default:
			return fmt.Errorf("unknown second return type: %T", v)
		}
	default:
		return fmt.Errorf("expected 0-2 return values from %s, got %d", methodName, len(result))
	}

	return nil
}

func (eval Evaluate) convert(to reflect.Type, node ast.Node) (reflect.Value, error) {
	iface := reflect.New(to).Interface()

	if _, ok := iface.(*ast.Node); ok {
		return reflect.ValueOf(node), nil
	}

	argEval := &Evaluate{
		Plugins: eval.Plugins,
	}

	err := node.Visit(argEval)
	if err != nil {
		return reflect.ValueOf(nil), err
	}

	content := argEval.Result

	switch iface.(type) {
	case *string:
		return reflect.ValueOf(content.String()), nil
	case *booklit.Content:
		return reflect.ValueOf(content), nil
	default:
		return reflect.ValueOf(nil), fmt.Errorf("unsupported argument type: %T", iface)
	}
}
<commit_msg>say function name, not method name, in errors<commit_after>package stages

import (
	"fmt"
	"reflect"

	"github.com/vito/booklit"
	"github.com/vito/booklit/ast"
)

type Evaluate struct {
	Plugins []booklit.Plugin

	Result booklit.Content
}

func (eval *Evaluate) VisitString(str ast.String) error {
	eval.Result = booklit.Append(eval.Result, booklit.String(str))
	return nil
}

func (eval *Evaluate) VisitSequence(seq ast.Sequence) error {
	for _, node := range seq {
		err := node.Visit(eval)
		if err != nil {
			return err
		}
	}

	return nil
}

func (eval *Evaluate) VisitParagraph(node ast.Paragraph) error {
	previous := eval.Result

	para := booklit.Paragraph{}
	for _, sentence := range node {
		eval.Result = nil

		err := sentence.Visit(eval)
		if err != nil {
			return err
		}

		if eval.Result != nil {
			para = append(para, eval.Result)
		}
	}

	eval.Result = nil

	if len(para) == 0 {
		// paragraph resulted in no content (e.g. an invoke with no return value)
		eval.Result = previous
		return nil
	}

	if len(para) == 1 && !para[0].IsSentence() {
		// paragraph resulted in block content (e.g. a section)
		eval.Result = booklit.Append(previous, para[0])
		return nil
	}

	eval.Result = booklit.Append(previous, para)

	return nil
}

func (eval *Evaluate) VisitPreformatted(node ast.Preformatted) error {
	previous := eval.Result

	pre := booklit.Preformatted{}
	for _, sentence := range node {
		eval.Result = nil

		err := sentence.Visit(eval)
		if err != nil {
			return err
		}

		if eval.Result != nil {
			pre = append(pre, eval.Result)
		}
	}

	eval.Result = booklit.Append(previous, pre)

	return nil
}

func (eval *Evaluate) VisitInvoke(invoke ast.Invoke) error {
	methodName := invoke.Method()

	var method reflect.Value
	for _, p := range eval.Plugins {
		value := reflect.ValueOf(p)
		method = value.MethodByName(methodName)
		if method.IsValid() {
			break
		}
	}

	if !method.IsValid() {
		return fmt.Errorf("undefined booklit function: %s", invoke.Function)
	}

	rawArgs := invoke.Arguments

	argc := method.Type().NumIn()
	if method.Type().IsVariadic() {
		argc--

		if len(rawArgs) < argc {
			return fmt.Errorf("argument count mismatch for %s: given %d, need at least %d", invoke.Function, len(rawArgs), argc)
		}
	} else {
		if len(rawArgs) != argc {
			return fmt.Errorf("argument count mismatch for %s: given %d, need %d", invoke.Function, len(rawArgs), argc)
		}
	}

	argv := make([]reflect.Value, argc)
	for i := 0; i < argc; i++ {
		t := method.Type().In(i)
		arg, err := eval.convert(t, rawArgs[i])
		if err != nil {
			return err
		}

		argv[i] = arg
	}

	if method.Type().IsVariadic() {
		variadic := rawArgs[argc:]
		variadicType := method.Type().In(argc)

		subType := variadicType.Elem()
		for _, varg := range variadic {
			arg, err := eval.convert(subType, varg)
			if err != nil {
				return err
			}

			argv = append(argv, arg)
		}
	}

	result := method.Call(argv)
	switch len(result) {
	case 0:
		return nil
	case 1:
		last := result[0]
		switch v := last.Interface().(type) {
		case nil:
		case error:
			return v
		case booklit.Content:
			eval.Result = booklit.Append(eval.Result, v)
		default:
			return fmt.Errorf("unknown return type: %T", v)
		}
	case 2:
		first := result[0]
		switch v := first.Interface().(type) {
		case booklit.Content:
			eval.Result = booklit.Append(eval.Result, v)
		default:
			return fmt.Errorf("unknown first return type: %T", v)
		}

		last := result[1]
		switch v := last.Interface().(type) {
		case nil:
		case error:
			return v
		default:
			return fmt.Errorf("unknown second return type: %T", v)
		}
	default:
		return fmt.Errorf("expected 0-2 return values from %s, got %d", invoke.Function, len(result))
	}

	return nil
}

func (eval Evaluate) convert(to reflect.Type, node ast.Node) (reflect.Value, error) {
	iface := reflect.New(to).Interface()

	if _, ok := iface.(*ast.Node); ok {
		return reflect.ValueOf(node), nil
	}

	argEval := &Evaluate{
		Plugins: eval.Plugins,
	}

	err := node.Visit(argEval)
	if err != nil {
		return reflect.ValueOf(nil), err
	}

	content := argEval.Result

	switch iface.(type) {
	case *string:
		return reflect.ValueOf(content.String()), nil
	case *booklit.Content:
		return reflect.ValueOf(content), nil
	default:
		return reflect.ValueOf(nil), fmt.Errorf("unsupported argument type: %T", iface)
	}
}
<|endoftext|>


<commit_before>package fusefrontend

import (
	"fmt"
	"log"
	"sync"
	"syscall"
	"time"

	"github.com/rfjakob/gocryptfs/internal/nametransform"
	"github.com/rfjakob/gocryptfs/internal/tlog"
)

const (
	// Number of entries in the dirCache. Three entries work well for two
	// (probably also three) parallel tar extracts (hit rate around 92%).
	// Keep in sync with test_helpers.maxCacheFds !
	// TODO: How to share this constant without causing an import cycle?
	dirCacheSize = 3
	// Enable Lookup/Store/Clear debug messages
	enableDebugMessages = false
	// Enable hit rate statistics printing
	enableStats = false
)

type dirCacheEntryStruct struct {
	// relative plaintext path to the directory
	dirRelPath string
	// fd to the directory (opened with O_PATH!)
	fd int
	// content of gocryptfs.diriv in this directory
	iv []byte
}

func (e *dirCacheEntryStruct) Clear() {
	// An earlier clear may have already closed the fd, or the cache
	// has never been filled (fd is 0 in that case).
	// Note: package ensurefds012, imported from main, guarantees that dirCache
	// can never get fds 0,1,2.
	if e.fd > 0 {
		err := syscall.Close(e.fd)
		if err != nil {
			tlog.Warn.Printf("dirCache.Clear: Close failed: %v", err)
		}
	}
	e.fd = -1
	e.dirRelPath = ""
	e.iv = nil
}

type dirCacheStruct struct {
	sync.Mutex
	// Cache entries
	entries [dirCacheSize]dirCacheEntryStruct
	// Where to store the next entry (index into entries)
	nextIndex int
	// On the first Lookup(), the expire thread is started, and this flag is set
	// to true.
	expireThreadRunning bool
	// Hit rate stats. Evaluated and reset by the expire thread.
	lookups uint64
	hits    uint64
}

// Clear clears the cache contents.
func (d *dirCacheStruct) Clear() {
	d.Lock()
	defer d.Unlock()
	for i := range d.entries {
		d.entries[i].Clear()
	}
}

// Store the entry in the cache. The passed "fd" will be Dup()ed, and the caller
// can close their copy at will.
func (d *dirCacheStruct) Store(dirRelPath string, fd int, iv []byte) {
	// Note: package ensurefds012, imported from main, guarantees that dirCache
	// can never get fds 0,1,2.
	if fd <= 0 || len(iv) != nametransform.DirIVLen {
		log.Panicf("Store sanity check failed: fd=%d len=%d", fd, len(iv))
	}
	d.Lock()
	defer d.Unlock()
	e := &d.entries[d.nextIndex]
	// Round-robin works well enough
	d.nextIndex = (d.nextIndex + 1) % dirCacheSize
	// Close the old fd
	e.Clear()
	fd2, err := syscall.Dup(fd)
	if err != nil {
		tlog.Warn.Printf("dirCache.Store: Dup failed: %v", err)
		return
	}
	d.dbg("Store: %q %d %x\n", dirRelPath, fd2, iv)
	e.fd = fd2
	e.dirRelPath = dirRelPath
	e.iv = iv
	// expireThread is started on the first Lookup()
	if !d.expireThreadRunning {
		d.expireThreadRunning = true
		go d.expireThread()
	}
}

// Lookup checks if relPath is in the cache, and returns an (fd, iv) pair.
// It returns (-1, nil) if not found. The fd is internally Dup()ed and the
// caller must close it when done.
func (d *dirCacheStruct) Lookup(dirRelPath string) (fd int, iv []byte) {
	d.Lock()
	defer d.Unlock()
	if enableStats {
		d.lookups++
	}
	for i := range d.entries {
		e := &d.entries[i]
		if e.fd <= 0 {
			// Cache slot is empty
			continue
		}
		if dirRelPath != e.dirRelPath {
			// Not the right path
			continue
		}
		var err error
		fd, err = syscall.Dup(e.fd)
		if err != nil {
			tlog.Warn.Printf("dirCache.Lookup: Dup failed: %v", err)
			return -1, nil
		}
		iv = e.iv
		break
	}
	if fd == 0 {
		d.dbg("Lookup %q: miss\n", dirRelPath)
		return -1, nil
	}
	if enableStats {
		d.hits++
	}
	if fd <= 0 || len(iv) != nametransform.DirIVLen {
		log.Panicf("Lookup sanity check failed: fd=%d len=%d", fd, len(iv))
	}
	d.dbg("Lookup %q: hit %d %x\n", dirRelPath, fd, iv)
	return fd, iv
}

// expireThread is started on the first Lookup()
func (d *dirCacheStruct) expireThread() {
	for {
		time.Sleep(1 * time.Second)
		d.Clear()
		if enableStats {
			d.Lock()
			lookups := d.lookups
			hits := d.hits
			d.lookups = 0
			d.hits = 0
			d.Unlock()
			if lookups > 0 {
				fmt.Printf("dirCache: hits=%3d lookups=%3d, rate=%3d%%\n", hits, lookups, (hits*100)/lookups)
			}
		}
	}
}

// dbg prints a debug message. Usually disabled.
func (d *dirCacheStruct) dbg(format string, a ...interface{}) {
	if enableDebugMessages {
		fmt.Printf(format, a...)
	}
}
<commit_msg>dircache: improve debug messages<commit_after>package fusefrontend

import (
	"fmt"
	"log"
	"sync"
	"syscall"
	"time"

	"github.com/rfjakob/gocryptfs/internal/nametransform"
	"github.com/rfjakob/gocryptfs/internal/tlog"
)

const (
	// Number of entries in the dirCache. Three entries work well for two
	// (probably also three) parallel tar extracts (hit rate around 92%).
	// Keep in sync with test_helpers.maxCacheFds !
	// TODO: How to share this constant without causing an import cycle?
	dirCacheSize = 3
	// Enable Lookup/Store/Clear debug messages
	enableDebugMessages = false
	// Enable hit rate statistics printing
	enableStats = false

	pathFmt = "%-40q"
)

type dirCacheEntryStruct struct {
	// relative plaintext path to the directory
	dirRelPath string
	// fd to the directory (opened with O_PATH!)
	fd int
	// content of gocryptfs.diriv in this directory
	iv []byte
}

func (e *dirCacheEntryStruct) Clear() {
	// An earlier clear may have already closed the fd, or the cache
	// has never been filled (fd is 0 in that case).
	// Note: package ensurefds012, imported from main, guarantees that dirCache
	// can never get fds 0,1,2.
	if e.fd > 0 {
		err := syscall.Close(e.fd)
		if err != nil {
			tlog.Warn.Printf("dirCache.Clear: Close failed: %v", err)
		}
	}
	e.fd = -1
	e.dirRelPath = ""
	e.iv = nil
}

type dirCacheStruct struct {
	sync.Mutex
	// Cache entries
	entries [dirCacheSize]dirCacheEntryStruct
	// Where to store the next entry (index into entries)
	nextIndex int
	// On the first Lookup(), the expire thread is started, and this flag is set
	// to true.
	expireThreadRunning bool
	// Hit rate stats. Evaluated and reset by the expire thread.
	lookups uint64
	hits    uint64
}

// Clear clears the cache contents.
func (d *dirCacheStruct) Clear() {
	d.dbg("Clear\n")
	d.Lock()
	defer d.Unlock()
	for i := range d.entries {
		d.entries[i].Clear()
	}
}

// Store the entry in the cache. The passed "fd" will be Dup()ed, and the caller
// can close their copy at will.
func (d *dirCacheStruct) Store(dirRelPath string, fd int, iv []byte) {
	// Note: package ensurefds012, imported from main, guarantees that dirCache
	// can never get fds 0,1,2.
	if fd <= 0 || len(iv) != nametransform.DirIVLen {
		log.Panicf("Store sanity check failed: fd=%d len=%d", fd, len(iv))
	}
	d.Lock()
	defer d.Unlock()
	e := &d.entries[d.nextIndex]
	// Round-robin works well enough
	d.nextIndex = (d.nextIndex + 1) % dirCacheSize
	// Close the old fd
	e.Clear()
	fd2, err := syscall.Dup(fd)
	if err != nil {
		tlog.Warn.Printf("dirCache.Store: Dup failed: %v", err)
		return
	}
	d.dbg("Store  "+pathFmt+" fd=%d iv=%x\n", dirRelPath, fd2, iv)
	e.fd = fd2
	e.dirRelPath = dirRelPath
	e.iv = iv
	// expireThread is started on the first Lookup()
	if !d.expireThreadRunning {
		d.expireThreadRunning = true
		go d.expireThread()
	}
}

// Lookup checks if relPath is in the cache, and returns an (fd, iv) pair.
// It returns (-1, nil) if not found. The fd is internally Dup()ed and the
// caller must close it when done.
func (d *dirCacheStruct) Lookup(dirRelPath string) (fd int, iv []byte) {
	d.Lock()
	defer d.Unlock()
	if enableStats {
		d.lookups++
	}
	var e *dirCacheEntryStruct
	for i := range d.entries {
		e = &d.entries[i]
		if e.fd <= 0 {
			// Cache slot is empty
			continue
		}
		if dirRelPath != e.dirRelPath {
			// Not the right path
			continue
		}
		var err error
		fd, err = syscall.Dup(e.fd)
		if err != nil {
			tlog.Warn.Printf("dirCache.Lookup: Dup failed: %v", err)
			return -1, nil
		}
		iv = e.iv
		break
	}
	if fd == 0 {
		d.dbg("Lookup "+pathFmt+" miss\n", dirRelPath)
		return -1, nil
	}
	if enableStats {
		d.hits++
	}
	if fd <= 0 || len(iv) != nametransform.DirIVLen {
		log.Panicf("Lookup sanity check failed: fd=%d len=%d", fd, len(iv))
	}
	d.dbg("Lookup "+pathFmt+" hit fd=%d dup=%d iv=%x\n", dirRelPath, e.fd, fd, iv)
	return fd, iv
}

// expireThread is started on the first Lookup()
func (d *dirCacheStruct) expireThread() {
	for {
		time.Sleep(1 * time.Second)
		d.Clear()
		if enableStats {
			d.Lock()
			lookups := d.lookups
			hits := d.hits
			d.lookups = 0
			d.hits = 0
			d.Unlock()
			if lookups > 0 {
				fmt.Printf("dirCache: hits=%3d lookups=%3d, rate=%3d%%\n", hits, lookups, (hits*100)/lookups)
			}
		}
	}
}

// dbg prints a debug message. Usually disabled.
func (d *dirCacheStruct) dbg(format string, a ...interface{}) {
	if enableDebugMessages {
		fmt.Printf(format, a...)
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"errors"
	"io"

	"github.com/appc/acserver/aci"
	"github.com/appc/acserver/upload"
)

var (
	ErrGPGPubKeyNotProvided = errors.New("GPG Public Key not provided")
)

type Storage interface {
	GetGPGPubKey() ([]byte, error)
	ListACIs() ([]aci.Aci, error)
	UploadACI(upload.Upload, io.Reader) error
	UploadASC(upload.Upload, io.Reader) error
	FinishUpload(upload.Upload) error
	CancelUpload(upload.Upload) error
}
<commit_msg>Add the DownloadACI method to the storage interface<commit_after>package storage

import (
	"errors"
	"io"

	"github.com/appc/acserver/aci"
	"github.com/appc/acserver/upload"
)

var (
	ErrGPGPubKeyNotProvided = errors.New("GPG Public Key not provided")
)

type Storage interface {
	GetGPGPubKey() ([]byte, error)
	ListACIs() ([]aci.Aci, error)
	DownloadACI(string) (io.ReadSeeker, error)
	UploadACI(upload.Upload, io.Reader) error
	UploadASC(upload.Upload, io.Reader) error
	FinishUpload(upload.Upload) error
	CancelUpload(upload.Upload) error
}
<|endoftext|>