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

package datasrcs

import (
	"context"
	"errors"
	"fmt"

	"github.com/OWASP/Amass/v3/config"
	"github.com/OWASP/Amass/v3/eventbus"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/systems"
	luaurl "github.com/cjoudrey/gluaurl"
	lua "github.com/yuin/gopher-lua"
	luajson "layeh.com/gopher-json"
)

// Script is the Service that handles access to the Script data source.
type Script struct {
	requests.BaseService

	SourceType string
	luaState   *lua.LState
	// Script callback functions
	start      lua.LValue
	stop       lua.LValue
	vertical   lua.LValue
	horizontal lua.LValue
	address    lua.LValue
	asn        lua.LValue
	resolved   lua.LValue
	subdomain  lua.LValue
}

// NewScript returns he object initialized, but not yet started.
func NewScript(script string, sys systems.System) *Script {
	s := new(Script)
	L := s.newLuaState(sys.Config())
	s.luaState = L

	// Load the script
	err := L.DoString(script)
	if err != nil {
		msg := fmt.Sprintf("Script: Failed to load script: %v", err)

		sys.Config().Log.Print(msg)
		return nil
	}

	// Pull the script type from the script
	s.SourceType, err = s.scriptType()
	if err != nil {
		msg := fmt.Sprintf("Script: Failed to obtain the script type: %v", err)

		sys.Config().Log.Print(msg)
		return nil
	}

	// Pull the script name from the script
	name, err := s.scriptName()
	if err != nil {
		msg := fmt.Sprintf("Script: Failed to obtain the script name: %v", err)

		sys.Config().Log.Print(msg)
		return nil
	}
	s.BaseService = *requests.NewBaseService(s, name)

	// Acquire API authentication info and make it global in the script
	s.registerAPIKey(L, sys.Config())
	// Save references to the callbacks defined within the script
	s.getScriptCallbacks()
	return s
}

// Setup the Lua state with desired constraints and access to necessary functionality.
func (s *Script) newLuaState(cfg *config.Config) *lua.LState {
	L := lua.NewState()

	L.PreloadModule("url", luaurl.Loader)
	L.PreloadModule("json", luajson.Loader)
	L.SetGlobal("log", L.NewFunction(s.log))
	L.SetGlobal("find", L.NewFunction(s.find))
	L.SetGlobal("submatch", L.NewFunction(s.submatch))
	L.SetGlobal("active", L.NewFunction(s.active))
	L.SetGlobal("newname", L.NewFunction(s.newName))
	L.SetGlobal("newaddr", L.NewFunction(s.newAddr))
	L.SetGlobal("newasn", L.NewFunction(s.newASN))
	L.SetGlobal("associated", L.NewFunction(s.associated))
	L.SetGlobal("inscope", L.NewFunction(s.inScope))
	L.SetGlobal("request", L.NewFunction(s.request))
	L.SetGlobal("scrape", L.NewFunction(s.scrape))
	L.SetGlobal("crawl", L.NewFunction(s.crawl))
	L.SetGlobal("outputdir", L.NewFunction(s.outputdir))
	L.SetGlobal("setratelimit", L.NewFunction(s.setRateLimit))
	L.SetGlobal("checkratelimit", L.NewFunction(s.checkRateLimit))
	L.SetGlobal("subdomainre", lua.LString(dns.AnySubdomainRegexString()))
	return L
}

// Fetch provided API authentication information and provide to the script as a global table.
func (s *Script) registerAPIKey(L *lua.LState, cfg *config.Config) {
	api := cfg.GetAPIKey(s.String())
	if api == nil {
		return
	}

	tb := L.NewTable()
	if api.Username != "" {
		tb.RawSetString("username", lua.LString(api.Username))
	}
	if api.Password != "" {
		tb.RawSetString("password", lua.LString(api.Password))
	}
	if api.Key != "" {
		tb.RawSetString("key", lua.LString(api.Key))
	}
	if api.Secret != "" {
		tb.RawSetString("secret", lua.LString(api.Secret))
	}

	L.SetGlobal("api", tb)
}

// Save references to the script functions that serve as callbacks for Amass events.
func (s *Script) getScriptCallbacks() {
	L := s.luaState

	s.start = L.GetGlobal("start")
	s.stop = L.GetGlobal("stop")
	s.vertical = L.GetGlobal("vertical")
	s.horizontal = L.GetGlobal("horizontal")
	s.address = L.GetGlobal("address")
	s.asn = L.GetGlobal("asn")
	s.resolved = L.GetGlobal("resolved")
	s.subdomain = L.GetGlobal("subdomain")
}

// Acquires the script name of the script by accessing the global variable.
func (s *Script) scriptName() (string, error) {
	L := s.luaState

	lv := L.GetGlobal("name")
	if lv.Type() == lua.LTNil {
		return "", errors.New("Script does not contain the 'name' global")
	}

	if str, ok := lv.(lua.LString); ok {
		return string(str), nil
	}

	return "", errors.New("The script global 'name' is not a string")
}

// Acquires the script type of the script by accessing the global variable.
func (s *Script) scriptType() (string, error) {
	L := s.luaState

	lv := L.GetGlobal("type")
	if lv.Type() == lua.LTNil {
		return "", errors.New("Script does not contain the 'type' global")
	}

	if str, ok := lv.(lua.LString); ok {
		return string(str), nil
	}

	return "", errors.New("The script global 'type' is not a string")
}

// Type implements the Service interface.
func (s *Script) Type() string {
	return s.SourceType
}

// OnStart implements the Service interface.
func (s *Script) OnStart() error {
	s.BaseService.OnStart()

	L := s.luaState
	if s.start.Type() == lua.LTNil {
		return nil
	}

	L.CallByParam(lua.P{
		Fn:      s.start,
		NRet:    0,
		Protect: true,
	})
	return nil
}

// OnStop implements the Service interface.
func (s *Script) OnStop() error {
	defer s.luaState.Close()

	L := s.luaState
	if s.stop.Type() == lua.LTNil {
		return nil
	}

	L.CallByParam(lua.P{
		Fn:      s.stop,
		NRet:    0,
		Protect: true,
	})
	return nil
}

// OnDNSRequest implements the Service interface.
func (s *Script) OnDNSRequest(ctx context.Context, req *requests.DNSRequest) {
	L := s.luaState

	if s.vertical.Type() == lua.LTNil || req == nil || req.Domain == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())
	bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
		fmt.Sprintf("Querying %s for %s subdomains", s.String(), req.Domain))

	L.CallByParam(lua.P{
		Fn:      s.vertical,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Domain))
}

// OnResolved implements the Service interface.
func (s *Script) OnResolved(ctx context.Context, req *requests.DNSRequest) {
	L := s.luaState

	if s.resolved.Type() == lua.LTNil || req == nil || req.Name == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.resolved,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Name))
}

// OnSubdomainDiscovered implements the Service interface.
func (s *Script) OnSubdomainDiscovered(ctx context.Context, req *requests.DNSRequest, times int) {
	L := s.luaState

	if s.subdomain.Type() == lua.LTNil || req == nil || req.Name == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.subdomain,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LNumber(times))
}

// OnAddrRequest implements the Service interface.
func (s *Script) OnAddrRequest(ctx context.Context, req *requests.AddrRequest) {
	L := s.luaState

	if s.address.Type() == lua.LTNil || req == nil || req.Address == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.address,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Address))
}

// OnASNRequest implements the Service interface.
func (s *Script) OnASNRequest(ctx context.Context, req *requests.ASNRequest) {
	L := s.luaState

	if s.asn.Type() == lua.LTNil || req == nil || req.Address == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.asn,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Address))
}

// OnWhoisRequest implements the Service interface.
func (s *Script) OnWhoisRequest(ctx context.Context, req *requests.WhoisRequest) {
	L := s.luaState

	if s.horizontal.Type() == lua.LTNil {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.horizontal,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Domain))
}
<commit_msg>bug fix related to enum termination<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package datasrcs

import (
	"context"
	"errors"
	"fmt"

	"github.com/OWASP/Amass/v3/config"
	"github.com/OWASP/Amass/v3/eventbus"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/systems"
	luaurl "github.com/cjoudrey/gluaurl"
	lua "github.com/yuin/gopher-lua"
	luajson "layeh.com/gopher-json"
)

// Script is the Service that handles access to the Script data source.
type Script struct {
	requests.BaseService

	SourceType string
	luaState   *lua.LState
	// Script callback functions
	start      lua.LValue
	stop       lua.LValue
	vertical   lua.LValue
	horizontal lua.LValue
	address    lua.LValue
	asn        lua.LValue
	resolved   lua.LValue
	subdomain  lua.LValue
}

// NewScript returns he object initialized, but not yet started.
func NewScript(script string, sys systems.System) *Script {
	s := new(Script)
	L := s.newLuaState(sys.Config())
	s.luaState = L

	// Load the script
	err := L.DoString(script)
	if err != nil {
		msg := fmt.Sprintf("Script: Failed to load script: %v", err)

		sys.Config().Log.Print(msg)
		return nil
	}

	// Pull the script type from the script
	s.SourceType, err = s.scriptType()
	if err != nil {
		msg := fmt.Sprintf("Script: Failed to obtain the script type: %v", err)

		sys.Config().Log.Print(msg)
		return nil
	}

	// Pull the script name from the script
	name, err := s.scriptName()
	if err != nil {
		msg := fmt.Sprintf("Script: Failed to obtain the script name: %v", err)

		sys.Config().Log.Print(msg)
		return nil
	}
	s.BaseService = *requests.NewBaseService(s, name)

	// Acquire API authentication info and make it global in the script
	s.registerAPIKey(L, sys.Config())
	// Save references to the callbacks defined within the script
	s.getScriptCallbacks()
	return s
}

// Setup the Lua state with desired constraints and access to necessary functionality.
func (s *Script) newLuaState(cfg *config.Config) *lua.LState {
	L := lua.NewState()

	L.PreloadModule("url", luaurl.Loader)
	L.PreloadModule("json", luajson.Loader)
	L.SetGlobal("log", L.NewFunction(s.log))
	L.SetGlobal("find", L.NewFunction(s.find))
	L.SetGlobal("submatch", L.NewFunction(s.submatch))
	L.SetGlobal("active", L.NewFunction(s.active))
	L.SetGlobal("newname", L.NewFunction(s.newName))
	L.SetGlobal("newaddr", L.NewFunction(s.newAddr))
	L.SetGlobal("newasn", L.NewFunction(s.newASN))
	L.SetGlobal("associated", L.NewFunction(s.associated))
	L.SetGlobal("inscope", L.NewFunction(s.inScope))
	L.SetGlobal("request", L.NewFunction(s.request))
	L.SetGlobal("scrape", L.NewFunction(s.scrape))
	L.SetGlobal("crawl", L.NewFunction(s.crawl))
	L.SetGlobal("outputdir", L.NewFunction(s.outputdir))
	L.SetGlobal("setratelimit", L.NewFunction(s.setRateLimit))
	L.SetGlobal("checkratelimit", L.NewFunction(s.checkRateLimit))
	L.SetGlobal("subdomainre", lua.LString(dns.AnySubdomainRegexString()))
	return L
}

// Fetch provided API authentication information and provide to the script as a global table.
func (s *Script) registerAPIKey(L *lua.LState, cfg *config.Config) {
	api := cfg.GetAPIKey(s.String())
	if api == nil {
		return
	}

	tb := L.NewTable()
	if api.Username != "" {
		tb.RawSetString("username", lua.LString(api.Username))
	}
	if api.Password != "" {
		tb.RawSetString("password", lua.LString(api.Password))
	}
	if api.Key != "" {
		tb.RawSetString("key", lua.LString(api.Key))
	}
	if api.Secret != "" {
		tb.RawSetString("secret", lua.LString(api.Secret))
	}

	L.SetGlobal("api", tb)
}

// Save references to the script functions that serve as callbacks for Amass events.
func (s *Script) getScriptCallbacks() {
	L := s.luaState

	s.start = L.GetGlobal("start")
	s.stop = L.GetGlobal("stop")
	s.vertical = L.GetGlobal("vertical")
	s.horizontal = L.GetGlobal("horizontal")
	s.address = L.GetGlobal("address")
	s.asn = L.GetGlobal("asn")
	s.resolved = L.GetGlobal("resolved")
	s.subdomain = L.GetGlobal("subdomain")
}

// Acquires the script name of the script by accessing the global variable.
func (s *Script) scriptName() (string, error) {
	L := s.luaState

	lv := L.GetGlobal("name")
	if lv.Type() == lua.LTNil {
		return "", errors.New("Script does not contain the 'name' global")
	}

	if str, ok := lv.(lua.LString); ok {
		return string(str), nil
	}

	return "", errors.New("The script global 'name' is not a string")
}

// Acquires the script type of the script by accessing the global variable.
func (s *Script) scriptType() (string, error) {
	L := s.luaState

	lv := L.GetGlobal("type")
	if lv.Type() == lua.LTNil {
		return "", errors.New("Script does not contain the 'type' global")
	}

	if str, ok := lv.(lua.LString); ok {
		return string(str), nil
	}

	return "", errors.New("The script global 'type' is not a string")
}

// Type implements the Service interface.
func (s *Script) Type() string {
	return s.SourceType
}

// OnStart implements the Service interface.
func (s *Script) OnStart() error {
	s.BaseService.OnStart()

	L := s.luaState
	if s.start.Type() == lua.LTNil {
		return nil
	}

	L.CallByParam(lua.P{
		Fn:      s.start,
		NRet:    0,
		Protect: true,
	})
	return nil
}

// OnStop implements the Service interface.
func (s *Script) OnStop() error {
	defer s.luaState.Close()

	L := s.luaState
	if s.stop.Type() == lua.LTNil {
		return nil
	}

	L.CallByParam(lua.P{
		Fn:      s.stop,
		NRet:    0,
		Protect: true,
	})
	return nil
}

// OnDNSRequest implements the Service interface.
func (s *Script) OnDNSRequest(ctx context.Context, req *requests.DNSRequest) {
	L := s.luaState

	if s.vertical.Type() == lua.LTNil || req == nil || req.Domain == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())
	bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
		fmt.Sprintf("Querying %s for %s subdomains", s.String(), req.Domain))

	L.CallByParam(lua.P{
		Fn:      s.vertical,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Domain))
}

// OnResolved implements the Service interface.
func (s *Script) OnResolved(ctx context.Context, req *requests.DNSRequest) {
	L := s.luaState

	if s.resolved.Type() == lua.LTNil || req == nil || req.Name == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	L.CallByParam(lua.P{
		Fn:      s.resolved,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Name))
}

// OnSubdomainDiscovered implements the Service interface.
func (s *Script) OnSubdomainDiscovered(ctx context.Context, req *requests.DNSRequest, times int) {
	L := s.luaState

	if s.subdomain.Type() == lua.LTNil || req == nil || req.Name == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	L.CallByParam(lua.P{
		Fn:      s.subdomain,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LNumber(times))
}

// OnAddrRequest implements the Service interface.
func (s *Script) OnAddrRequest(ctx context.Context, req *requests.AddrRequest) {
	L := s.luaState

	if s.address.Type() == lua.LTNil || req == nil || req.Address == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.address,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Address))
}

// OnASNRequest implements the Service interface.
func (s *Script) OnASNRequest(ctx context.Context, req *requests.ASNRequest) {
	L := s.luaState

	if s.asn.Type() == lua.LTNil || req == nil || req.Address == "" {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.asn,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Address))
}

// OnWhoisRequest implements the Service interface.
func (s *Script) OnWhoisRequest(ctx context.Context, req *requests.WhoisRequest) {
	L := s.luaState

	if s.horizontal.Type() == lua.LTNil {
		return
	}

	cfg := ctx.Value(requests.ContextConfig).(*config.Config)
	bus := ctx.Value(requests.ContextEventBus).(*eventbus.EventBus)
	if cfg == nil || bus == nil {
		return
	}

	s.CheckRateLimit()
	bus.Publish(requests.SetActiveTopic, eventbus.PriorityCritical, s.String())

	L.CallByParam(lua.P{
		Fn:      s.horizontal,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Domain))
}
<|endoftext|>


<commit_before>package datastore

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gocql/gocql"

	"github.com/jeffpierce/cassabon/config"
	"github.com/jeffpierce/cassabon/logging"
	"github.com/jeffpierce/cassabon/middleware"
)

// rollup contains the accumulated metrics data for a path.
type rollup struct {
	expr  string    // The text form of the path expression, to locate the definition
	count []uint64  // The number of data points accumulated (for averaging)
	value []float64 // One rollup per window definition
}

// runlist contains the paths to be written for an expression, and when to write the rollups.
type runlist struct {
	nextWriteTime []time.Time        // The next write time for each rollup bucket
	path          map[string]*rollup // The rollup data for each path matched by the expression
}

type StoreManager struct {

	// Rollup configuration.
	// Note: Does not reload on SIGHUP.
	rollupPriority []string                    // First matched expression wins
	rollup         map[string]config.RollupDef // Rollup processing definitions by path expression

	// Timer management.
	setTimeout chan time.Duration // Write a duration to this to get a notification on timeout channel
	timeout    chan struct{}      // Timeout notifications arrive on this channel

	// Database connection.
	dbClient *gocql.Session

	// Rollup data.
	byPath map[string]*rollup  // Stats, by path, for rollup accumulation
	byExpr map[string]*runlist // Stats, by path within expression, for rollup processing
}

func (sm *StoreManager) Init() {

	// Copy in the configuration (requires hard restart to refresh).
	sm.rollupPriority = config.G.RollupPriority
	sm.rollup = config.G.Rollup

	// Initialize private objects.
	sm.setTimeout = make(chan time.Duration, 0)
	sm.timeout = make(chan struct{}, 1)

	// Start the persistent goroutines.
	config.G.OnExitWG.Add(2)
	go sm.timer()
	go sm.run()

	// Kick off the timer.
	sm.setTimeout <- time.Second
}

func (sm *StoreManager) Start() {
}

func (sm *StoreManager) resetRollupData() {

	// Initialize rollup data structures.
	sm.byPath = make(map[string]*rollup)
	sm.byExpr = make(map[string]*runlist)
	baseTime := time.Now()
	for expr, rollupdef := range sm.rollup {
		// For each expression, provide a place to record all the paths that it matches.
		rl := new(runlist)
		rl.nextWriteTime = make([]time.Time, len(rollupdef.Windows))
		rl.path = make(map[string]*rollup)
		// Establish the next time boundary on which each write will take place.
		for i, v := range rollupdef.Windows {
			rl.nextWriteTime[i] = nextTimeBoundary(baseTime, v.Window)
		}
		sm.byExpr[expr] = rl
	}
}

func (sm *StoreManager) populateSchema() {
	// Keyspace exists since we have a successful dbClient connection, create tables if they do not exist
	for _, table := range config.G.RollupTables {
		var ttlfloat float64
		ttl := strings.Split(table, "_")[1]
		ttlfloat, _ = strconv.ParseFloat(ttl, 64)
		query := fmt.Sprintf(
			`CREATE TABLE IF NOT EXISTS %s (path text, timestamp timestamp, stat double, PRIMARY KEY (path, timestamp))
			WITH COMPACT STORAGE
			  AND CLUSTERING ORDER BY (timestamp ASC)
			  AND compaction = {'class': 'org.apache.cassandra.db.compaction.DateTieredCompactionStrategy'}
			  AND compression = {'sstable_compression': 'org.apache.cassandra.io.compress.LZ4Compressor'}
			  AND dclocal_read_repair_chance = 0.1
			  AND default_time_to_live = %v
			  AND gc_grace_seconds = 864000
			  AND memtable_flush_period_in_ms = 0
			  AND read_repair_chance = 0.0
			  AND speculative_retry = '99.0PERCENTILE';`, table, int(ttlfloat*1.1))

		config.G.Log.System.LogDebug(query)

		if err := sm.dbClient.Query(query).Exec(); err != nil {
			config.G.Log.System.LogFatal("Could not configure cassabon keyspace, error is %v", err.Error())
		}
	}
}

func (sm *StoreManager) run() {

	// Perform first-time initialization of rollup data accumulation structures.
	sm.resetRollupData()

	// Open connection to the Cassandra database here, so we can defer the close.
	var err error
	config.G.Log.System.LogDebug("StoreManager initializing Cassandra client")
	sm.dbClient, err = middleware.CassandraSession(
		config.G.Cassandra.Hosts,
		config.G.Cassandra.Port,
		"cassabon",
	)
	if err != nil {
		// Without Cassandra client we can't do our job, so log, whine, and crash.
		config.G.Log.System.LogFatal("StoreManager unable to connect to Cassandra at %v, port %s: %v",
			config.G.Cassandra.Hosts, config.G.Cassandra.Port, err)
	}

	defer sm.dbClient.Close()
	config.G.Log.System.LogDebug("StoreManager Cassandra client initialized")

	config.G.Log.System.LogDebug("StoreManager Cassandra Keyspace configuration starting...")
	sm.populateSchema()

	for {
		select {
		case <-config.G.OnPeerChangeReq:
			config.G.Log.System.LogDebug("StoreManager::run received PEERCHANGE message")
			sm.flush(true)
			sm.resetRollupData()
			config.G.OnPeerChangeRsp <- struct{}{} // Unblock sender
		case <-config.G.OnExit:
			config.G.Log.System.LogDebug("StoreManager::run received QUIT message")
			sm.flush(true)
			config.G.OnExitWG.Done()
			return
		case metric := <-config.G.Channels.DataStore:
			sm.accumulate(metric)
		case <-sm.timeout:
			sm.flush(false)
		}
	}
}

// timer sends a message on the "timeout" channel after the specified duration.
func (sm *StoreManager) timer() {
	for {
		select {
		case <-config.G.OnExit:
			config.G.Log.System.LogDebug("StoreManager::timer received QUIT message")
			config.G.OnExitWG.Done()
			return
		case duration := <-sm.setTimeout:
			// Block in this state until a new entry is received.
			select {
			case <-config.G.OnExit:
				// Nothing; do handling above on next iteration.
			case <-time.After(duration):
				select {
				case sm.timeout <- struct{}{}:
					// Timeout sent.
				default:
					// Do not block.
				}
			}
		}
	}
}

// accumulate records a metric according to the rollup definitions.
func (sm *StoreManager) accumulate(metric config.CarbonMetric) {
	config.G.Log.System.LogDebug("StoreManager::accumulate %s=%v", metric.Path, metric.Value)

	// Locate the metric in the map.
	var currentRollup *rollup
	var found bool
	if currentRollup, found = sm.byPath[metric.Path]; !found {

		// Determine which expression matches this path.
		var expr string
		for _, expr = range sm.rollupPriority {
			if expr != config.ROLLUP_CATCHALL {
				if sm.rollup[expr].Expression.MatchString(metric.Path) {
					break
				}
			}
			// Catchall always appears last, and is therefore the default value.
		}

		// Initialize, and insert the new rollup into both maps.
		currentRollup = new(rollup)
		currentRollup.expr = expr
		currentRollup.count = make([]uint64, len(sm.rollup[expr].Windows))
		currentRollup.value = make([]float64, len(sm.rollup[expr].Windows))
		sm.byPath[metric.Path] = currentRollup
		sm.byExpr[expr].path[metric.Path] = currentRollup

		// Send the entry off for writing to the path index.
		config.G.Channels.IndexStore <- metric
	}

	// Apply the incoming metric to each rollup bucket.
	switch sm.rollup[currentRollup.expr].Method {
	case config.AVERAGE:
		for i, v := range currentRollup.value {
			currentRollup.value[i] = (v*float64(currentRollup.count[i]) + metric.Value) /
				float64(currentRollup.count[i]+1)
		}
	case config.MAX:
		for i, v := range currentRollup.value {
			if v < metric.Value {
				currentRollup.value[i] = metric.Value
			}
		}
	case config.MIN:
		for i, v := range currentRollup.value {
			if v > metric.Value || currentRollup.count[i] == 0 {
				currentRollup.value[i] = metric.Value
			}
		}
	case config.SUM:
		for i, v := range currentRollup.value {
			currentRollup.value[i] = v + metric.Value
		}
	}

	// Note that we added a data point into each bucket.
	for i, _ := range currentRollup.count {
		currentRollup.count[i]++
	}
}

// flush persists the accumulated metrics to the database.
func (sm *StoreManager) flush(terminating bool) {
	config.G.Log.System.LogDebug("StoreManager::flush terminating=%v", terminating)

	// Report the current length of the list of unique paths seen.
	logging.Statsd.Client.Gauge("path.count", int64(len(sm.byPath)), 1.0)

	// Use a consistent current time for all tests in this cycle.
	baseTime := time.Now()

	// Use a reasonable default value for setting the next timer delay.
	nextFlush := baseTime.Add(time.Minute)

	// Walk the set of expressions, looking for closed rollup windows.
	for expr, rl := range sm.byExpr {

		// Inspect each rollup window defined for this expression.
		for i, windowEnd := range rl.nextWriteTime {

			// If the window has closed, process and clear the data.
			if windowEnd.Before(baseTime) {

				// Iterate over all the paths that match the current expression.
				for path, rollup := range rl.path {

					// Has any data accumulated while the window was open?
					if rollup.count[i] > 0 {
						// TODO: Write the data to persistent storage.
						config.G.Log.System.LogInfo("Write expr=%s win=%v ret=%v ts=%v path=%s value=%.4f",
							expr,
							sm.rollup[expr].Windows[i].Window,
							sm.rollup[expr].Windows[i].Retention,
							windowEnd.Format("15:04:05.000"), // Window end time
							path,
							rollup.value[i])

						sm.write(path, windowEnd, rollup.value[i], sm.rollup[expr].Windows[i].Table)
					}

					// Ensure the bucket is empty for the next open window.
					rollup.count[i] = 0
					rollup.value[i] = 0
				}

				// Set a new window closing time for the just-cleared window.
				rl.nextWriteTime[i] = nextTimeBoundary(baseTime, sm.rollup[expr].Windows[i].Window)
			}

			// If terminating, write out all remaining data, stamped with current time.
			if terminating {

				// Iterate over all the paths that match the current expression.
				for path, rollup := range rl.path {

					// Has any data accumulated while the window was open?
					if rollup.count[i] > 0 {
						// TODO: Write the data to persistent storage.
						config.G.Log.System.LogInfo("Write expr=%s win=%v ret=%v ts=%v path=%s value=%.4f",
							expr,
							sm.rollup[expr].Windows[i].Window,
							sm.rollup[expr].Windows[i].Retention,
							baseTime.Format("15:04:05.000"), // Current time, window end is in future
							path,
							rollup.value[i])

						sm.write(path, baseTime, rollup.value[i], sm.rollup[expr].Windows[i].Table)
					}

					// Ensure the bucket is empty for the next open window.
					rollup.count[i] = 0
					rollup.value[i] = 0
				}

				// Set a new window closing time for the just-cleared window.
				rl.nextWriteTime[i] = nextTimeBoundary(baseTime, sm.rollup[expr].Windows[i].Window)
			}

			// ASSERT: rl.nextWriteTime[i] time is in the future (later than baseTime).

			// Adjust the timer delay downwards if this window closing time is
			// earlier than all others seen so far.
			if nextFlush.After(rl.nextWriteTime[i]) {
				nextFlush = rl.nextWriteTime[i]
			}
		}
	}

	// Set a timer to expire when the earliest future window closing occurs.
	if !terminating {

		// Convert window closing time to a duration, and do a sanity check.
		delay := nextFlush.Sub(baseTime)
		if delay.Nanoseconds() < 0 {
			delay = time.Millisecond
		}

		// Perform a non-blocking write to the timeout channel.
		select {
		case sm.setTimeout <- delay:
			// Notification sent
		default:
			// Do not block if channel is at capacity
		}
	}
}

// flush persists the accumulated metrics to the database.
func (sm *StoreManager) write(path string, ts time.Time, value float64, table string) {
	query := fmt.Sprintf(`INSERT INTO %s (path, timestamp, stat) VALUES (?, ?, ?)`, table)
	if err := sm.dbClient.Query(query, path, ts, value).Exec(); err != nil {
		// Could not write to Cassandra cluster...we should scream loudly about this.  Possibly a failure case?.
		config.G.Log.System.LogError("Unable to write stats to Cassandra cluster, error is %s", err.Error())
	}
}
<commit_msg>Fix compile error: import not used<commit_after>package datastore

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/gocql/gocql"

	"github.com/jeffpierce/cassabon/config"
	"github.com/jeffpierce/cassabon/logging"
	"github.com/jeffpierce/cassabon/middleware"
)

// rollup contains the accumulated metrics data for a path.
type rollup struct {
	expr  string    // The text form of the path expression, to locate the definition
	count []uint64  // The number of data points accumulated (for averaging)
	value []float64 // One rollup per window definition
}

// runlist contains the paths to be written for an expression, and when to write the rollups.
type runlist struct {
	nextWriteTime []time.Time        // The next write time for each rollup bucket
	path          map[string]*rollup // The rollup data for each path matched by the expression
}

type StoreManager struct {

	// Rollup configuration.
	// Note: Does not reload on SIGHUP.
	rollupPriority []string                    // First matched expression wins
	rollup         map[string]config.RollupDef // Rollup processing definitions by path expression

	// Timer management.
	setTimeout chan time.Duration // Write a duration to this to get a notification on timeout channel
	timeout    chan struct{}      // Timeout notifications arrive on this channel

	// Database connection.
	dbClient *gocql.Session

	// Rollup data.
	byPath map[string]*rollup  // Stats, by path, for rollup accumulation
	byExpr map[string]*runlist // Stats, by path within expression, for rollup processing
}

func (sm *StoreManager) Init() {

	// Copy in the configuration (requires hard restart to refresh).
	sm.rollupPriority = config.G.RollupPriority
	sm.rollup = config.G.Rollup

	// Initialize private objects.
	sm.setTimeout = make(chan time.Duration, 0)
	sm.timeout = make(chan struct{}, 1)

	// Start the persistent goroutines.
	config.G.OnExitWG.Add(2)
	go sm.timer()
	go sm.run()

	// Kick off the timer.
	sm.setTimeout <- time.Second
}

func (sm *StoreManager) Start() {
}

func (sm *StoreManager) resetRollupData() {

	// Initialize rollup data structures.
	sm.byPath = make(map[string]*rollup)
	sm.byExpr = make(map[string]*runlist)
	baseTime := time.Now()
	for expr, rollupdef := range sm.rollup {
		// For each expression, provide a place to record all the paths that it matches.
		rl := new(runlist)
		rl.nextWriteTime = make([]time.Time, len(rollupdef.Windows))
		rl.path = make(map[string]*rollup)
		// Establish the next time boundary on which each write will take place.
		for i, v := range rollupdef.Windows {
			rl.nextWriteTime[i] = nextTimeBoundary(baseTime, v.Window)
		}
		sm.byExpr[expr] = rl
	}
}

func (sm *StoreManager) populateSchema() {
	// Keyspace exists since we have a successful dbClient connection, create tables if they do not exist
	for _, table := range config.G.RollupTables {
		var ttlfloat float64
		ttl := strings.Split(table, "_")[1]
		ttlfloat, _ = strconv.ParseFloat(ttl, 64)
		query := fmt.Sprintf(
			`CREATE TABLE IF NOT EXISTS %s (path text, timestamp timestamp, stat double, PRIMARY KEY (path, timestamp))
			WITH COMPACT STORAGE
			  AND CLUSTERING ORDER BY (timestamp ASC)
			  AND compaction = {'class': 'org.apache.cassandra.db.compaction.DateTieredCompactionStrategy'}
			  AND compression = {'sstable_compression': 'org.apache.cassandra.io.compress.LZ4Compressor'}
			  AND dclocal_read_repair_chance = 0.1
			  AND default_time_to_live = %v
			  AND gc_grace_seconds = 864000
			  AND memtable_flush_period_in_ms = 0
			  AND read_repair_chance = 0.0
			  AND speculative_retry = '99.0PERCENTILE';`, table, int(ttlfloat*1.1))

		config.G.Log.System.LogDebug(query)

		if err := sm.dbClient.Query(query).Exec(); err != nil {
			config.G.Log.System.LogFatal("Could not configure cassabon keyspace, error is %v", err.Error())
		}
	}
}

func (sm *StoreManager) run() {

	// Perform first-time initialization of rollup data accumulation structures.
	sm.resetRollupData()

	// Open connection to the Cassandra database here, so we can defer the close.
	var err error
	config.G.Log.System.LogDebug("StoreManager initializing Cassandra client")
	sm.dbClient, err = middleware.CassandraSession(
		config.G.Cassandra.Hosts,
		config.G.Cassandra.Port,
		"cassabon",
	)
	if err != nil {
		// Without Cassandra client we can't do our job, so log, whine, and crash.
		config.G.Log.System.LogFatal("StoreManager unable to connect to Cassandra at %v, port %s: %v",
			config.G.Cassandra.Hosts, config.G.Cassandra.Port, err)
	}

	defer sm.dbClient.Close()
	config.G.Log.System.LogDebug("StoreManager Cassandra client initialized")

	config.G.Log.System.LogDebug("StoreManager Cassandra Keyspace configuration starting...")
	sm.populateSchema()

	for {
		select {
		case <-config.G.OnPeerChangeReq:
			config.G.Log.System.LogDebug("StoreManager::run received PEERCHANGE message")
			sm.flush(true)
			sm.resetRollupData()
			config.G.OnPeerChangeRsp <- struct{}{} // Unblock sender
		case <-config.G.OnExit:
			config.G.Log.System.LogDebug("StoreManager::run received QUIT message")
			sm.flush(true)
			config.G.OnExitWG.Done()
			return
		case metric := <-config.G.Channels.DataStore:
			sm.accumulate(metric)
		case <-sm.timeout:
			sm.flush(false)
		}
	}
}

// timer sends a message on the "timeout" channel after the specified duration.
func (sm *StoreManager) timer() {
	for {
		select {
		case <-config.G.OnExit:
			config.G.Log.System.LogDebug("StoreManager::timer received QUIT message")
			config.G.OnExitWG.Done()
			return
		case duration := <-sm.setTimeout:
			// Block in this state until a new entry is received.
			select {
			case <-config.G.OnExit:
				// Nothing; do handling above on next iteration.
			case <-time.After(duration):
				select {
				case sm.timeout <- struct{}{}:
					// Timeout sent.
				default:
					// Do not block.
				}
			}
		}
	}
}

// accumulate records a metric according to the rollup definitions.
func (sm *StoreManager) accumulate(metric config.CarbonMetric) {
	config.G.Log.System.LogDebug("StoreManager::accumulate %s=%v", metric.Path, metric.Value)

	// Locate the metric in the map.
	var currentRollup *rollup
	var found bool
	if currentRollup, found = sm.byPath[metric.Path]; !found {

		// Determine which expression matches this path.
		var expr string
		for _, expr = range sm.rollupPriority {
			if expr != config.ROLLUP_CATCHALL {
				if sm.rollup[expr].Expression.MatchString(metric.Path) {
					break
				}
			}
			// Catchall always appears last, and is therefore the default value.
		}

		// Initialize, and insert the new rollup into both maps.
		currentRollup = new(rollup)
		currentRollup.expr = expr
		currentRollup.count = make([]uint64, len(sm.rollup[expr].Windows))
		currentRollup.value = make([]float64, len(sm.rollup[expr].Windows))
		sm.byPath[metric.Path] = currentRollup
		sm.byExpr[expr].path[metric.Path] = currentRollup

		// Send the entry off for writing to the path index.
		config.G.Channels.IndexStore <- metric
	}

	// Apply the incoming metric to each rollup bucket.
	switch sm.rollup[currentRollup.expr].Method {
	case config.AVERAGE:
		for i, v := range currentRollup.value {
			currentRollup.value[i] = (v*float64(currentRollup.count[i]) + metric.Value) /
				float64(currentRollup.count[i]+1)
		}
	case config.MAX:
		for i, v := range currentRollup.value {
			if v < metric.Value {
				currentRollup.value[i] = metric.Value
			}
		}
	case config.MIN:
		for i, v := range currentRollup.value {
			if v > metric.Value || currentRollup.count[i] == 0 {
				currentRollup.value[i] = metric.Value
			}
		}
	case config.SUM:
		for i, v := range currentRollup.value {
			currentRollup.value[i] = v + metric.Value
		}
	}

	// Note that we added a data point into each bucket.
	for i, _ := range currentRollup.count {
		currentRollup.count[i]++
	}
}

// flush persists the accumulated metrics to the database.
func (sm *StoreManager) flush(terminating bool) {
	config.G.Log.System.LogDebug("StoreManager::flush terminating=%v", terminating)

	// Report the current length of the list of unique paths seen.
	logging.Statsd.Client.Gauge("path.count", int64(len(sm.byPath)), 1.0)

	// Use a consistent current time for all tests in this cycle.
	baseTime := time.Now()

	// Use a reasonable default value for setting the next timer delay.
	nextFlush := baseTime.Add(time.Minute)

	// Walk the set of expressions, looking for closed rollup windows.
	for expr, rl := range sm.byExpr {

		// Inspect each rollup window defined for this expression.
		for i, windowEnd := range rl.nextWriteTime {

			// If the window has closed, process and clear the data.
			if windowEnd.Before(baseTime) {

				// Iterate over all the paths that match the current expression.
				for path, rollup := range rl.path {

					// Has any data accumulated while the window was open?
					if rollup.count[i] > 0 {
						// TODO: Write the data to persistent storage.
						config.G.Log.System.LogInfo("Write expr=%s win=%v ret=%v ts=%v path=%s value=%.4f",
							expr,
							sm.rollup[expr].Windows[i].Window,
							sm.rollup[expr].Windows[i].Retention,
							windowEnd.Format("15:04:05.000"), // Window end time
							path,
							rollup.value[i])

						sm.write(path, windowEnd, rollup.value[i], sm.rollup[expr].Windows[i].Table)
					}

					// Ensure the bucket is empty for the next open window.
					rollup.count[i] = 0
					rollup.value[i] = 0
				}

				// Set a new window closing time for the just-cleared window.
				rl.nextWriteTime[i] = nextTimeBoundary(baseTime, sm.rollup[expr].Windows[i].Window)
			}

			// If terminating, write out all remaining data, stamped with current time.
			if terminating {

				// Iterate over all the paths that match the current expression.
				for path, rollup := range rl.path {

					// Has any data accumulated while the window was open?
					if rollup.count[i] > 0 {
						// TODO: Write the data to persistent storage.
						config.G.Log.System.LogInfo("Write expr=%s win=%v ret=%v ts=%v path=%s value=%.4f",
							expr,
							sm.rollup[expr].Windows[i].Window,
							sm.rollup[expr].Windows[i].Retention,
							baseTime.Format("15:04:05.000"), // Current time, window end is in future
							path,
							rollup.value[i])

						sm.write(path, baseTime, rollup.value[i], sm.rollup[expr].Windows[i].Table)
					}

					// Ensure the bucket is empty for the next open window.
					rollup.count[i] = 0
					rollup.value[i] = 0
				}

				// Set a new window closing time for the just-cleared window.
				rl.nextWriteTime[i] = nextTimeBoundary(baseTime, sm.rollup[expr].Windows[i].Window)
			}

			// ASSERT: rl.nextWriteTime[i] time is in the future (later than baseTime).

			// Adjust the timer delay downwards if this window closing time is
			// earlier than all others seen so far.
			if nextFlush.After(rl.nextWriteTime[i]) {
				nextFlush = rl.nextWriteTime[i]
			}
		}
	}

	// Set a timer to expire when the earliest future window closing occurs.
	if !terminating {

		// Convert window closing time to a duration, and do a sanity check.
		delay := nextFlush.Sub(baseTime)
		if delay.Nanoseconds() < 0 {
			delay = time.Millisecond
		}

		// Perform a non-blocking write to the timeout channel.
		select {
		case sm.setTimeout <- delay:
			// Notification sent
		default:
			// Do not block if channel is at capacity
		}
	}
}

// flush persists the accumulated metrics to the database.
func (sm *StoreManager) write(path string, ts time.Time, value float64, table string) {
	query := fmt.Sprintf(`INSERT INTO %s (path, timestamp, stat) VALUES (?, ?, ?)`, table)
	if err := sm.dbClient.Query(query, path, ts, value).Exec(); err != nil {
		// Could not write to Cassandra cluster...we should scream loudly about this.  Possibly a failure case?.
		config.G.Log.System.LogError("Unable to write stats to Cassandra cluster, error is %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package buffalo

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/gobuffalo/buffalo/render"
	"github.com/gorilla/schema"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
)

// DefaultContext is, as its name implies, a default
// implementation of the Context interface.
type DefaultContext struct {
	response    http.ResponseWriter
	request     *http.Request
	params      url.Values
	logger      Logger
	session     *Session
	contentType string
	notFound    http.Handler
	data        map[string]interface{}
}

// Response returns the original Response for the request.
func (d *DefaultContext) Response() http.ResponseWriter {
	return d.response
}

// Request returns the original Request.
func (d *DefaultContext) Request() *http.Request {
	return d.request
}

// Params returns all of the parameters for the request,
// including both named params and query string parameters.
// These parameters are automatically available in templates
// as "{{.params}}".
func (d *DefaultContext) Params() ParamValues {
	return d.params
}

// Logger returns the Logger for this context.
func (d *DefaultContext) Logger() Logger {
	return d.logger
}

// Param returns a param, either named or query string,
// based on the key.
func (d *DefaultContext) Param(key string) string {
	return d.Params().Get(key)
}

// ParamInt tries to convert the requested parameter to
// an int. It will  return an error if there is a problem.
func (d *DefaultContext) ParamInt(key string) (int, error) {
	k := d.Params().Get(key)
	i, err := strconv.Atoi(k)
	return i, errors.WithMessage(err, fmt.Sprintf("could not convert %s to an int", k))
}

// Set a value onto the Context. Any value set onto the Context
// will be automatically available in templates.
func (d *DefaultContext) Set(key string, value interface{}) {
	d.data[key] = value
}

// Get a value that was previous set onto the Context.
func (d *DefaultContext) Get(key string) interface{} {
	return d.data[key]
}

// Session for the associated Request.
func (d *DefaultContext) Session() *Session {
	return d.session
}

// Render a status code and render.Renderer to the associated Response.
// The request parameters will be made available to the render.Renderer
// "{{.params}}". Any values set onto the Context will also automatically
// be made available to the render.Renderer. To render "no content" pass
// in a nil render.Renderer.
func (d *DefaultContext) Render(status int, rr render.Renderer) error {
	now := time.Now()
	defer func() {
		d.LogField("render", time.Now().Sub(now))
	}()
	if rr != nil {
		data := d.data
		pp := map[string]string{}
		for k, v := range d.params {
			pp[k] = v[0]
		}
		data["params"] = pp
		bb := &bytes.Buffer{}
		err := rr.Render(bb, data)
		if err != nil {
			return err
		}
		d.Response().Header().Set("Content-Type", rr.ContentType())
		d.Response().WriteHeader(status)
		_, err = io.Copy(d.Response(), bb)
		return err
	}
	d.Response().WriteHeader(status)
	return nil
}

// Bind the interface to the request.Body. The type of binding
// is dependent on the "Content-Type" for the request. If the type
// is "application/json" it will use "json.NewDecoder". If the type
// is "application/xml" it will use "xml.NewDecoder". The default
// binder is "http://www.gorillatoolkit.org/pkg/schema".
func (d *DefaultContext) Bind(value interface{}) error {
	switch strings.ToLower(d.Request().Header.Get("Content-Type")) {
	case "application/json", "text/json", "json":
		return json.NewDecoder(d.Request().Body).Decode(value)
	case "application/xml", "text/xml", "xml":
		return xml.NewDecoder(d.Request().Body).Decode(value)
	default:
		err := d.Request().ParseForm()
		if err != nil {
			return err
		}
		dec := schema.NewDecoder()
		dec.IgnoreUnknownKeys(true)
		dec.ZeroEmpty(true)
		return dec.Decode(value, d.Request().PostForm)
	}
}

// LogField adds the key/value pair onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogField(key string, value interface{}) {
	d.logger = d.logger.WithField(key, value)
}

// LogFields adds the key/value pairs onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogFields(values map[string]interface{}) {
	d.logger = d.logger.WithFields(values)
}

func (d *DefaultContext) Error(status int, err error) error {
	if status == 404 {
		req := d.Request()
		req.URL.Query().Set("error", err.Error())
		d.notFound.ServeHTTP(d.Response(), req)
		return nil
	}
	err = errors.WithStack(err)
	d.Logger().Error(err)
	msg := fmt.Sprintf("%+v", err)
	d.Response().WriteHeader(status)

	ct := d.Request().Header.Get("Content-Type")
	switch strings.ToLower(ct) {
	case "application/json", "text/json", "json":
		err = json.NewEncoder(d.Response()).Encode(map[string]interface{}{
			"error": msg,
			"code":  status,
		})
	case "application/xml", "text/xml", "xml":
	default:
		_, err = d.Response().Write([]byte(fmt.Sprintf("
%+v
", msg)))
	}
	return err
}

// Websocket returns an upgraded github.com/gorilla/websocket.Conn
// that can then be used to work with websockets easily.
func (d *DefaultContext) Websocket() (*websocket.Conn, error) {
	return defaultUpgrader.Upgrade(d.Response(), d.Request(), nil)
}

// Redirect a request with the given status to the given URL.
func (d *DefaultContext) Redirect(status int, url string, args ...interface{}) error {
	http.Redirect(d.Response(), d.Request(), fmt.Sprintf(url, args...), status)
	return nil
}

// Data contains all the values set through Get/Set.
func (d *DefaultContext) Data() map[string]interface{} {
	return d.data
}

var defaultUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}
<commit_msg>Remove an extra space<commit_after>package buffalo

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/gobuffalo/buffalo/render"
	"github.com/gorilla/schema"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
)

// DefaultContext is, as its name implies, a default
// implementation of the Context interface.
type DefaultContext struct {
	response    http.ResponseWriter
	request     *http.Request
	params      url.Values
	logger      Logger
	session     *Session
	contentType string
	notFound    http.Handler
	data        map[string]interface{}
}

// Response returns the original Response for the request.
func (d *DefaultContext) Response() http.ResponseWriter {
	return d.response
}

// Request returns the original Request.
func (d *DefaultContext) Request() *http.Request {
	return d.request
}

// Params returns all of the parameters for the request,
// including both named params and query string parameters.
// These parameters are automatically available in templates
// as "{{.params}}".
func (d *DefaultContext) Params() ParamValues {
	return d.params
}

// Logger returns the Logger for this context.
func (d *DefaultContext) Logger() Logger {
	return d.logger
}

// Param returns a param, either named or query string,
// based on the key.
func (d *DefaultContext) Param(key string) string {
	return d.Params().Get(key)
}

// ParamInt tries to convert the requested parameter to
// an int. It will return an error if there is a problem.
func (d *DefaultContext) ParamInt(key string) (int, error) {
	k := d.Params().Get(key)
	i, err := strconv.Atoi(k)
	return i, errors.WithMessage(err, fmt.Sprintf("could not convert %s to an int", k))
}

// Set a value onto the Context. Any value set onto the Context
// will be automatically available in templates.
func (d *DefaultContext) Set(key string, value interface{}) {
	d.data[key] = value
}

// Get a value that was previous set onto the Context.
func (d *DefaultContext) Get(key string) interface{} {
	return d.data[key]
}

// Session for the associated Request.
func (d *DefaultContext) Session() *Session {
	return d.session
}

// Render a status code and render.Renderer to the associated Response.
// The request parameters will be made available to the render.Renderer
// "{{.params}}". Any values set onto the Context will also automatically
// be made available to the render.Renderer. To render "no content" pass
// in a nil render.Renderer.
func (d *DefaultContext) Render(status int, rr render.Renderer) error {
	now := time.Now()
	defer func() {
		d.LogField("render", time.Now().Sub(now))
	}()
	if rr != nil {
		data := d.data
		pp := map[string]string{}
		for k, v := range d.params {
			pp[k] = v[0]
		}
		data["params"] = pp
		bb := &bytes.Buffer{}
		err := rr.Render(bb, data)
		if err != nil {
			return err
		}
		d.Response().Header().Set("Content-Type", rr.ContentType())
		d.Response().WriteHeader(status)
		_, err = io.Copy(d.Response(), bb)
		return err
	}
	d.Response().WriteHeader(status)
	return nil
}

// Bind the interface to the request.Body. The type of binding
// is dependent on the "Content-Type" for the request. If the type
// is "application/json" it will use "json.NewDecoder". If the type
// is "application/xml" it will use "xml.NewDecoder". The default
// binder is "http://www.gorillatoolkit.org/pkg/schema".
func (d *DefaultContext) Bind(value interface{}) error {
	switch strings.ToLower(d.Request().Header.Get("Content-Type")) {
	case "application/json", "text/json", "json":
		return json.NewDecoder(d.Request().Body).Decode(value)
	case "application/xml", "text/xml", "xml":
		return xml.NewDecoder(d.Request().Body).Decode(value)
	default:
		err := d.Request().ParseForm()
		if err != nil {
			return err
		}
		dec := schema.NewDecoder()
		dec.IgnoreUnknownKeys(true)
		dec.ZeroEmpty(true)
		return dec.Decode(value, d.Request().PostForm)
	}
}

// LogField adds the key/value pair onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogField(key string, value interface{}) {
	d.logger = d.logger.WithField(key, value)
}

// LogFields adds the key/value pairs onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogFields(values map[string]interface{}) {
	d.logger = d.logger.WithFields(values)
}

func (d *DefaultContext) Error(status int, err error) error {
	if status == 404 {
		req := d.Request()
		req.URL.Query().Set("error", err.Error())
		d.notFound.ServeHTTP(d.Response(), req)
		return nil
	}
	err = errors.WithStack(err)
	d.Logger().Error(err)
	msg := fmt.Sprintf("%+v", err)
	d.Response().WriteHeader(status)

	ct := d.Request().Header.Get("Content-Type")
	switch strings.ToLower(ct) {
	case "application/json", "text/json", "json":
		err = json.NewEncoder(d.Response()).Encode(map[string]interface{}{
			"error": msg,
			"code":  status,
		})
	case "application/xml", "text/xml", "xml":
	default:
		_, err = d.Response().Write([]byte(fmt.Sprintf("
%+v
", msg)))
	}
	return err
}

// Websocket returns an upgraded github.com/gorilla/websocket.Conn
// that can then be used to work with websockets easily.
func (d *DefaultContext) Websocket() (*websocket.Conn, error) {
	return defaultUpgrader.Upgrade(d.Response(), d.Request(), nil)
}

// Redirect a request with the given status to the given URL.
func (d *DefaultContext) Redirect(status int, url string, args ...interface{}) error {
	http.Redirect(d.Response(), d.Request(), fmt.Sprintf(url, args...), status)
	return nil
}

// Data contains all the values set through Get/Set.
func (d *DefaultContext) Data() map[string]interface{} {
	return d.data
}

var defaultUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}
<|endoftext|>


<commit_before>package database

import _ "github.com/go-sql-driver/mysql"

const stmtSelectUserByNumber = "SELECT name, cash, dream FROM users WHERE number = ?"

type user struct {
	name         string
	dreamBalance int
	cashBalance  int
}

// SelectUserByNumber function selects user info via collector number
func SelectUserByNumber(num int) (user, error) {
	stmt, err := db.Prepare(stmtSelectUserByNumber)
	if err != nil {
		return &user, err
	}

	err = stmt.QueryRow(1).Scan(&user)
	if err != nil {
		return &user, err
	}
	return &user, nil
}
<commit_msg>Database cleaned up, needs testing<commit_after>package database

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
)

const stmtSelectUserByNumber = "SELECT name, cash, dream FROM users WHERE number = ?"

// User database structure
type User struct {
	name         string
	dreamBalance int
	cashBalance  int
}

// SelectUserByNumber function selects user info via collector number
func SelectUserByNumber(num int) (*User, error) {
	tempUser := User{}
	stmt, err := db.Prepare(stmtSelectUserByNumber)
	if err != nil {
		return &tempUser, err
	}
	defer stmt.Close()

	err = stmt.QueryRow(1).Scan(&tempUser)
	if err != nil {
		return &tempUser, err
	}
	return &tempUser, nil
}
<|endoftext|>


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

// Package utils contains helper utils for osconfig_tests.

package packagemanagement

import (
	"fmt"

	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/utils"
)

func getPackageInstallStartupScript(pkgManager, packageName string) string {
	var ss string

	switch pkgManager {
	case "apt":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/dpkg-query -s %s`\n" +
			"if [[ $isinstalled =~ \"Status: install ok installed\" ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5;\n" +
			"done;\n"

		ss = fmt.Sprintf(ss, utils.InstallOSConfigDeb, packageName, packageInstalledString, packageNotInstalledString)

	case "yum":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/rpmquery -a %s`\n" +
			"if [[ $isinstalled =~ ^cowsay-* ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5\n" +
			"done\n"
		ss = fmt.Sprintf(ss, utils.InstallOSConfigYumEL7, packageName, packageInstalledString, packageNotInstalledString)

	default:
		panic(fmt.Sprintf("invalid package manager: %s", pkgManager))
	}

	return ss
}

func getPackageRemovalStartupScript(pkgManager, packageName string) string {
	var ss string

	switch pkgManager {
	case "apt":
		ss = "%s\n" +
			"sudo apt-get -y install %s\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"could not install package\"\n" +
			"exit 1\n" +
			"fi\n" +
			"systemctl restart google-osconfig-agent\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"Error restarting google-osconfig-agent\"\n" +
			"exit 1\n" +
			"fi\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=\"$(/usr/bin/dpkg-query -s %s 2>&1 )\"\n" +
			"if [[ $isinstalled =~ \"package '%s' is not installed\" ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5;\n" +
			"done;\n"

		ss = fmt.Sprintf(ss, utils.InstallOSConfigDeb, packageName, packageNotInstalledString, packageInstalledString)

	case "yum":
		ss = "%s\n" +
			"sudo yum -y install %s\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"could not install package\"\n" +
			"exit 1\n" +
			"fi\n" +
			"systemctl restart google-osconfig-agent\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"Error restarting google-osconfig-agent\"\n" +
			"exit 1\n" +
			"fi\n" + "while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/rpmquery -a %s`\n" +
			"if [[ $isinstalled =~ ^cowsay-* ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5\n" +
			"done\n"
		ss = fmt.Sprintf(ss, utils.InstallOSConfigYumEL7, packageName, packageName, packageInstalledString, packageNotInstalledString)

	default:
		panic(fmt.Sprintf("invalid package manager: %s", pkgManager))
	}

	return ss
}

func getPackageInstallRemovalStartupScript(pkgManager, packageName string) string {
	var ss string

	switch pkgManager {
	case "apt":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/dpkg-query -s %s`\n" +
			"if [[ $isinstalled =~ \"package '%s' is not installed\" ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5;\n" +
			"done;\n"

		ss = fmt.Sprintf(ss, utils.InstallOSConfigDeb, packageName, packageName, packageNotInstalledString, packageInstalledString)

	case "yum":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/rpmquery -a %s`\n" +
			"if [[ $isinstalled =~ ^cowsay-* ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5\n" +
			"done\n"
		ss = fmt.Sprintf(ss, utils.InstallOSConfigYumEL7, packageName, packageName, packageInstalledString, packageNotInstalledString)

	default:
		panic(fmt.Sprintf("invalid package manager: %s", pkgManager))
	}

	return ss
}
<commit_msg>Fix broken and flaky tests (#700)<commit_after>//  Copyright 2019 Google Inc. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

// Package utils contains helper utils for osconfig_tests.

package packagemanagement

import (
	"fmt"

	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/utils"
)

func getPackageInstallStartupScript(pkgManager, packageName string) string {
	var ss string

	switch pkgManager {
	case "apt":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/dpkg-query -s %s`\n" +
			"if [[ $isinstalled =~ \"Status: install ok installed\" ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5;\n" +
			"done;\n"

		ss = fmt.Sprintf(ss, utils.InstallOSConfigDeb, packageName, packageInstalledString, packageNotInstalledString)

	case "yum":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/rpmquery -a %s`\n" +
			"if [[ $isinstalled =~ ^cowsay-* ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5\n" +
			"done\n"
		ss = fmt.Sprintf(ss, utils.InstallOSConfigYumEL7, packageName, packageInstalledString, packageNotInstalledString)

	default:
		panic(fmt.Sprintf("invalid package manager: %s", pkgManager))
	}

	return ss
}

func getPackageRemovalStartupScript(pkgManager, packageName string) string {
	var ss string

	switch pkgManager {
	case "apt":
		ss = "%s\n" +
			"n=0\n" +
			"while ! apt-get -y install %s; do\n" +
			"if [[ n -gt 3 ]]; then\n" +
			"echo \"could not install package\"\n" +
			"exit 1\n" +
			"fi\n" +
			"n=$[$n+1]\n" +
			"sleep 5\n" +
			"done\n" +
			"systemctl restart google-osconfig-agent\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"Error restarting google-osconfig-agent\"\n" +
			"exit 1\n" +
			"fi\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=\"$(/usr/bin/dpkg-query -s %s 2>&1 )\"\n" +
			"if [[ $isinstalled =~ \"package '%s' is not installed\" ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5;\n" +
			"done;\n"

		ss = fmt.Sprintf(ss, utils.InstallOSConfigDeb, packageName, packageName, packageName, packageNotInstalledString, packageInstalledString)

	case "yum":
		ss = "%s\n" +
			"yum -y install %s\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"could not install package\"\n" +
			"exit 1\n" +
			"fi\n" +
			"systemctl restart google-osconfig-agent\n" +
			"if [[ $? != 0 ]]; then\n" +
			"echo \"Error restarting google-osconfig-agent\"\n" +
			"exit 1\n" +
			"fi\n" + "while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/rpmquery -a %s`\n" +
			"if [[ $isinstalled =~ ^%s-* ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5\n" +
			"done\n"
		ss = fmt.Sprintf(ss, utils.InstallOSConfigYumEL7, packageName, packageName, packageName, packageInstalledString, packageNotInstalledString)

	default:
		panic(fmt.Sprintf("invalid package manager: %s", pkgManager))
	}

	return ss
}

func getPackageInstallRemovalStartupScript(pkgManager, packageName string) string {
	var ss string

	switch pkgManager {
	case "apt":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=\"$(/usr/bin/dpkg-query -s %s 2>&1 )\"\n" +
			"if [[ $isinstalled =~ \"package '%s' is not installed\" ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5;\n" +
			"done;\n"

		ss = fmt.Sprintf(ss, utils.InstallOSConfigDeb, packageName, packageName, packageNotInstalledString, packageInstalledString)

	case "yum":
		ss = "%s\n" +
			"while true;\n" +
			"do\n" +
			"isinstalled=`/usr/bin/rpmquery -a %s`\n" +
			"if [[ $isinstalled =~ ^cowsay-* ]]; then\n" +
			"echo \"%s\"\n" +
			"else\n" +
			"echo \"%s\"\n" +
			"fi\n" +
			"sleep 5\n" +
			"done\n"
		ss = fmt.Sprintf(ss, utils.InstallOSConfigYumEL7, packageName, packageInstalledString, packageNotInstalledString)

	default:
		panic(fmt.Sprintf("invalid package manager: %s", pkgManager))
	}

	return ss
}
<|endoftext|>


<commit_before>package testing

import (
	"net/http"
	"net/http/httptest"

	"github.com/admpub/log"

	"github.com/webx-top/echo/engine"
	"github.com/webx-top/echo/engine/standard"
)

// Request testing
func Request(method, path string, handler engine.Handler, reqRewrite ...func(*http.Request)) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	if len(reqRewrite) > 0 && reqRewrite[0] != nil {
		reqRewrite[0](req)
	}
	rec := httptest.NewRecorder()

	handler.ServeHTTP(WrapRequest(req), WrapResponse(req, rec))
	//rec.Code, rec.Body.String(),rec.Header
	return rec
}

func WrapRequest(req *http.Request) engine.Request {
	return standard.NewRequest(req)
}

func WrapResponse(req *http.Request, rw http.ResponseWriter) engine.Response {
	return standard.NewResponse(rw, req, log.New().Sync())
}
<commit_msg>update<commit_after>package testing

import (
	"net/http"
	"net/http/httptest"

	"github.com/admpub/log"

	"github.com/webx-top/echo/engine"
	"github.com/webx-top/echo/engine/standard"
)

// Request testing
func Request(method, path string, handler engine.Handler, reqRewrite ...func(*http.Request)) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	if len(reqRewrite) > 0 && reqRewrite[0] != nil {
		reqRewrite[0](req)
	}
	rec := httptest.NewRecorder()

	handler.ServeHTTP(WrapRequest(req), WrapResponse(req, rec))
	//rec.Code, rec.Body.String(),rec.Header
	return rec
}

func NewStdRequest(method, path string) *http.Request {
	req, _ := http.NewRequest(method, path, nil)
	return req
}

func NewStdResponse() http.ResponseWriter {
	return httptest.NewRecorder()
}

func NewRequestAndResponse(method, path string) (engine.Request, engine.Response) {
	req := NewStdRequest(method, path)
	return WrapRequest(req), WrapResponse(req, NewStdResponse())
}

func WrapRequest(req *http.Request) engine.Request {
	return standard.NewRequest(req)
}

func WrapResponse(req *http.Request, rw http.ResponseWriter) engine.Response {
	return standard.NewResponse(rw, req, log.New().Sync())
}
<|endoftext|>


<commit_before>package goproxy

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"net"
	"strings"
	"sync"
	"time"
)

// MaxSerialNumber is the upper boundary that is used to create unique serial
// numbers for the certificate. This can be any unsigned integer up to 20
// bytes (2^(8*20)-1).
var MaxSerialNumber = big.NewInt(0).SetBytes(bytes.Repeat([]byte{255}, 20))

func getWildcardHost(host string) string {
	first := strings.Index(host, ".")
	if first <= 0 {
		return host
	}
	last := strings.LastIndex(host, ".")
	if last == first {
		// root domain, no wildcard
		return host
	}
	return "*" + host[first:]
}

// Config is a set of configuration values that are used to build TLS configs
// capable of MITM.
type GoproxyConfig struct {
	Root   *x509.Certificate
	capriv interface{}

	priv  *rsa.PrivateKey
	keyID []byte

	validity time.Duration

	certmu sync.RWMutex

	*tls.Config
}

// NewConfig creates a MITM config using the CA certificate and
// private key to generate on-the-fly certificates.
func NewConfig(ca *x509.Certificate, privateKey interface{}) (*GoproxyConfig, error) {
	roots := x509.NewCertPool()
	roots.AddCert(ca)

	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	pub := priv.Public()

	// Subject Key Identifier support for end entity certificate.
	// https://www.ietf.org/rfc/rfc3280.txt (section 4.2.1.2)
	pkixpub, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return nil, err
	}
	h := sha1.New()
	h.Write(pkixpub)
	keyID := h.Sum(nil)

	tlsConfig := &GoproxyConfig{
		Root:     ca,
		capriv:   privateKey,
		priv:     priv,
		keyID:    keyID,
		validity: time.Hour * 24 * 3650,

		Config: &tls.Config{
			InsecureSkipVerify: true,
			Certificates:       make([]tls.Certificate, 0),
			MinVersion:         tls.VersionTLS12,
			RootCAs:            roots,
			NameToCertificate:  make(map[string]*tls.Certificate),
		},
	}

	return tlsConfig, nil
}

func (c *GoproxyConfig) cert(hostname string) error {
	// Remove the port if it exists.
	host, _, err := net.SplitHostPort(getWildcardHost(hostname))
	if err == nil {
		hostname = host
	}

	// Remove the port if it exists.
	c.certmu.RLock()
	tlsc, ok := c.NameToCertificate[hostname]
	c.certmu.RUnlock()

	if ok {
		// Check validity of the certificate for hostname match, expiry, etc. In
		// particular, if the cached certificate has expired, create a new one.
		if _, err := tlsc.Leaf.Verify(x509.VerifyOptions{
			DNSName: hostname,
			Roots:   c.RootCAs,
		}); err == nil {
			return nil
		}
	}

	serial, err := rand.Int(rand.Reader, MaxSerialNumber)
	if err != nil {
		return err
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   hostname,
			Organization: []string{"StopLight"},
		},
		SubjectKeyId:          c.keyID,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		NotBefore:             time.Now().Add(-c.validity),
		NotAfter:              time.Now().Add(c.validity),
	}

	if ip := net.ParseIP(hostname); ip != nil {
		tmpl.IPAddresses = []net.IP{ip}
	} else {
		tmpl.DNSNames = []string{hostname}
	}

	raw, err := x509.CreateCertificate(rand.Reader, tmpl, c.Root, c.priv.Public(), c.capriv)
	if err != nil {
		return err
	}

	// Parse certificate bytes so that we have a leaf certificate.
	x509c, err := x509.ParseCertificate(raw)
	if err != nil {
		return err
	}

	tlsc = &tls.Certificate{
		Certificate: [][]byte{raw, c.Root.Raw},
		PrivateKey:  c.priv,
		Leaf:        x509c,
	}

	c.certmu.Lock()
	c.NameToCertificate[hostname] = tlsc
	c.Certificates = append(c.Certificates, *tlsc)
	c.certmu.Unlock()

	return nil
}
<commit_msg>set the min ssl version to tls10<commit_after>package goproxy

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"net"
	"strings"
	"sync"
	"time"
)

// MaxSerialNumber is the upper boundary that is used to create unique serial
// numbers for the certificate. This can be any unsigned integer up to 20
// bytes (2^(8*20)-1).
var MaxSerialNumber = big.NewInt(0).SetBytes(bytes.Repeat([]byte{255}, 20))

func getWildcardHost(host string) string {
	first := strings.Index(host, ".")
	if first <= 0 {
		return host
	}
	last := strings.LastIndex(host, ".")
	if last == first {
		// root domain, no wildcard
		return host
	}
	return "*" + host[first:]
}

// Config is a set of configuration values that are used to build TLS configs
// capable of MITM.
type GoproxyConfig struct {
	Root   *x509.Certificate
	capriv interface{}

	priv  *rsa.PrivateKey
	keyID []byte

	validity time.Duration

	certmu sync.RWMutex

	*tls.Config
}

// NewConfig creates a MITM config using the CA certificate and
// private key to generate on-the-fly certificates.
func NewConfig(ca *x509.Certificate, privateKey interface{}) (*GoproxyConfig, error) {
	roots := x509.NewCertPool()
	roots.AddCert(ca)

	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	pub := priv.Public()

	// Subject Key Identifier support for end entity certificate.
	// https://www.ietf.org/rfc/rfc3280.txt (section 4.2.1.2)
	pkixpub, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return nil, err
	}
	h := sha1.New()
	h.Write(pkixpub)
	keyID := h.Sum(nil)

	tlsConfig := &GoproxyConfig{
		Root:     ca,
		capriv:   privateKey,
		priv:     priv,
		keyID:    keyID,
		validity: time.Hour * 24 * 3650,

		Config: &tls.Config{
			InsecureSkipVerify: true,
			Certificates:       make([]tls.Certificate, 0),
			MinVersion:         tls.VersionTLS10,
			RootCAs:            roots,
			NameToCertificate:  make(map[string]*tls.Certificate),
		},
	}

	return tlsConfig, nil
}

func (c *GoproxyConfig) cert(hostname string) error {
	// Remove the port if it exists.
	host, _, err := net.SplitHostPort(getWildcardHost(hostname))
	if err == nil {
		hostname = host
	}

	// Remove the port if it exists.
	c.certmu.RLock()
	tlsc, ok := c.NameToCertificate[hostname]
	c.certmu.RUnlock()

	if ok {
		// Check validity of the certificate for hostname match, expiry, etc. In
		// particular, if the cached certificate has expired, create a new one.
		if _, err := tlsc.Leaf.Verify(x509.VerifyOptions{
			DNSName: hostname,
			Roots:   c.RootCAs,
		}); err == nil {
			return nil
		}
	}

	serial, err := rand.Int(rand.Reader, MaxSerialNumber)
	if err != nil {
		return err
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   hostname,
			Organization: []string{"StopLight"},
		},
		SubjectKeyId:          c.keyID,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		NotBefore:             time.Now().Add(-c.validity),
		NotAfter:              time.Now().Add(c.validity),
	}

	if ip := net.ParseIP(hostname); ip != nil {
		tmpl.IPAddresses = []net.IP{ip}
	} else {
		tmpl.DNSNames = []string{hostname}
	}

	raw, err := x509.CreateCertificate(rand.Reader, tmpl, c.Root, c.priv.Public(), c.capriv)
	if err != nil {
		return err
	}

	// Parse certificate bytes so that we have a leaf certificate.
	x509c, err := x509.ParseCertificate(raw)
	if err != nil {
		return err
	}

	tlsc = &tls.Certificate{
		Certificate: [][]byte{raw, c.Root.Raw},
		PrivateKey:  c.priv,
		Leaf:        x509c,
	}

	c.certmu.Lock()
	c.NameToCertificate[hostname] = tlsc
	c.Certificates = append(c.Certificates, *tlsc)
	c.certmu.Unlock()

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"context"
	"github.com/google/mako/go/quickstore"
	"log"
	"net"
	"sort"
	"sync"
	"time"

	"google.golang.org/grpc"

	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"

	pb "knative.dev/eventing/test/test_images/latencymako/event_state"
	"knative.dev/pkg/test/mako"
)

const maxRcvMsgSize = 1024 * 1024 * 100

// thread-safe events recording map
type eventsRecord struct {
	sync.RWMutex
	*pb.EventsRecord
}

type aggregatorExecutor struct {
	// thread-safe events recording maps
	sentEvents     *eventsRecord
	acceptedEvents *eventsRecord
	failedEvents   *eventsRecord
	receivedEvents *eventsRecord

	// channel to notify the main goroutine that an events record has been received
	notifyEventsReceived chan struct{}

	// GRPC server
	listener net.Listener
	server   *grpc.Server
}

func newAggregatorExecutor(lis net.Listener) testExecutor {
	executor := &aggregatorExecutor{
		listener:             lis,
		notifyEventsReceived: make(chan struct{}),
	}

	// --- Create GRPC server

	s := grpc.NewServer(grpc.MaxRecvMsgSize(maxRcvMsgSize))
	pb.RegisterEventsRecorderServer(s, executor)
	executor.server = s

	// --- Initialize records maps

	executor.sentEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_SENT,
		Events: make(map[string]*timestamp.Timestamp),
	}}
	executor.acceptedEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_ACCEPTED,
		Events: make(map[string]*timestamp.Timestamp),
	}}
	executor.failedEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_FAILED,
		Events: make(map[string]*timestamp.Timestamp),
	}}
	executor.receivedEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_RECEIVED,
		Events: make(map[string]*timestamp.Timestamp),
	}}

	return executor
}

func (ex *aggregatorExecutor) Run(ctx context.Context) {
	// --- Configure mako

	printf("Configuring Mako")

	// Use the benchmark key created
	ctx, q, qclose, err := mako.Setup(ctx)
	if err != nil {
		fatalf("Failed to setup mako: %v", err)
	}

	// Use a fresh context here so that our RPC to terminate the sidecar
	// isn't subject to our timeout (or we won't shut it down when we time out)
	defer qclose(context.Background())

	// Wrap fatalf in a helper or our sidecar will live forever.
	fatalf = func(f string, args ...interface{}) {
		qclose(context.Background())
		fatalf(f, args...)
	}

	// --- Run GRPC events receiver

	printf("Starting events recorder server")

	go func() {
		if err := ex.server.Serve(ex.listener); err != nil {
			fatalf("Failed to serve: %v", err)
		}
	}()
	go func() {
		<-ctx.Done()
		printf("Terminating events recorder server")
		ex.server.GracefulStop()
	}()

	printf("Expecting %d events records", expectRecords)
	ex.waitForEvents()
	printf("Received all expected events records")

	ex.server.GracefulStop()

	// --- Publish latencies

	printf("%-15s: %d", "Sent count", len(ex.sentEvents.Events))
	printf("%-15s: %d", "Accepted count", len(ex.acceptedEvents.Events))
	printf("%-15s: %d", "Failed count", len(ex.failedEvents.Events))
	printf("%-15s: %d", "Received count", len(ex.receivedEvents.Events))

	printf("Publishing latencies")

	// count errors
	var publishErrorCount int
	var deliverErrorCount int

	for sentID := range ex.sentEvents.Events {
		timestampSentProto := ex.sentEvents.Events[sentID]
		timestampSent, _ := ptypes.Timestamp(timestampSentProto)

		timestampAcceptedProto, accepted := ex.acceptedEvents.Events[sentID]
		timestampAccepted, _ := ptypes.Timestamp(timestampAcceptedProto)

		timestampReceivedProto, received := ex.receivedEvents.Events[sentID]
		timestampReceived, _ := ptypes.Timestamp(timestampReceivedProto)

		if !accepted {
			errMsg := "Failed on broker"
			if _, failed := ex.failedEvents.Events[sentID]; !failed {
				// TODO(antoineco): should never happen, check whether the failed map makes any sense
				errMsg = "Event not accepted but missing from failed map"
			}

			deliverErrorCount++

			if qerr := q.AddError(mako.XTime(timestampSent), errMsg); qerr != nil {
				log.Printf("ERROR AddError: %v", qerr)
			}
			continue
		}

		sendLatency := timestampAccepted.Sub(timestampSent)
		// Uncomment to get CSV directly from this container log
		//fmt.Printf("%f,%d,\n", mako.XTime(timestampSent), sendLatency.Nanoseconds())
		// TODO mako accepts float64, which imo could lead to losing some precision on local tests. It should accept int64
		if qerr := q.AddSamplePoint(mako.XTime(timestampSent), map[string]float64{"pl": sendLatency.Seconds()}); qerr != nil {
			log.Printf("ERROR AddSamplePoint: %v", qerr)
		}

		if !received {
			publishErrorCount++

			if qerr := q.AddError(mako.XTime(timestampSent), "Event not delivered"); qerr != nil {
				log.Printf("ERROR AddError: %v", qerr)
			}
			continue
		}

		e2eLatency := timestampReceived.Sub(timestampSent)
		// Uncomment to get CSV directly from this container log
		//fmt.Printf("%f,,%d\n", mako.XTime(timestampSent), e2eLatency.Nanoseconds())
		// TODO mako accepts float64, which imo could lead to losing some precision on local tests. It should accept int64
		if qerr := q.AddSamplePoint(mako.XTime(timestampSent), map[string]float64{"dl": e2eLatency.Seconds()}); qerr != nil {
			log.Printf("ERROR AddSamplePoint: %v", qerr)
		}
	}

	// --- Publish throughput

	printf("Publishing throughputs")

	sentTimestamps := eventsToTimestampsArray(&ex.sentEvents.Events)
	err = publishThpt(sentTimestamps, q, "st")
	if err != nil {
		log.Printf("ERROR AddSamplePoint: %v", err)
	}

	receivedTimestamps := eventsToTimestampsArray(&ex.receivedEvents.Events)
	err = publishThpt(receivedTimestamps, q, "dt")
	if err != nil {
		log.Printf("ERROR AddSamplePoint: %v", err)
	}

	// --- Publish error counts as aggregate metrics

	printf("Publishing aggregates")

	q.AddRunAggregate("pe", float64(publishErrorCount))
	q.AddRunAggregate("de", float64(deliverErrorCount))

	printf("Store to mako")

	if out, err := q.Store(); err != nil {
		fatalf("Failed to store data: %v\noutput: %v", err, out)
	}

	printf("Aggregation completed")
}

func eventsToTimestampsArray(events *map[string]*timestamp.Timestamp) []time.Time {
	values := make([]time.Time, 0, len(*events))
	for _, v := range *events {
		t, _ := ptypes.Timestamp(v)
		values = append(values, t)
	}
	sort.Slice(values, func(x, y int) bool { return values[x].Before(values[y]) })
	return values
}

func publishThpt(timestamps []time.Time, q *quickstore.Quickstore, metricName string) error {
	for i, t := range timestamps[1:] {
		var thpt uint
		j := i - 1
		for j >= 0 && t.Sub(timestamps[j]) <= time.Second {
			thpt++
			j--
		}
		if qerr := q.AddSamplePoint(mako.XTime(t), map[string]float64{metricName: float64(thpt)}); qerr != nil {
			return qerr
		}
	}
	return nil
}

// waitForEvents blocks until the expected number of events records has been received.
func (ex *aggregatorExecutor) waitForEvents() {
	for receivedRecords := uint(0); receivedRecords < expectRecords; receivedRecords++ {
		<-ex.notifyEventsReceived
	}
}

// RecordSentEvents implements event_state.EventsRecorder
func (ex *aggregatorExecutor) RecordEvents(_ context.Context, in *pb.EventsRecordList) (*pb.RecordReply, error) {
	defer func() {
		ex.notifyEventsReceived <- struct{}{}
	}()

	for _, recIn := range in.Items {
		recType := recIn.GetType()

		var rec *eventsRecord

		switch recType {
		case pb.EventsRecord_SENT:
			rec = ex.sentEvents
		case pb.EventsRecord_ACCEPTED:
			rec = ex.acceptedEvents
		case pb.EventsRecord_FAILED:
			rec = ex.failedEvents
		case pb.EventsRecord_RECEIVED:
			rec = ex.receivedEvents
		default:
			printf("Ignoring events record of type %s", recType)
			continue
		}

		printf("-> Recording %d %s events", uint64(len(recIn.Events)), recType)

		func() {
			rec.Lock()
			defer rec.Unlock()
			for id, t := range recIn.Events {
				if _, exists := rec.Events[id]; exists {
					log.Printf("!! Found duplicate %s event ID %s", recType, id)
					continue
				}
				rec.Events[id] = t
			}
		}()
	}

	return &pb.RecordReply{Count: uint32(len(in.Items))}, nil
}
<commit_msg>golang format tools (#1929)<commit_after>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"context"
	"log"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/google/mako/go/quickstore"

	"google.golang.org/grpc"

	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"

	pb "knative.dev/eventing/test/test_images/latencymako/event_state"
	"knative.dev/pkg/test/mako"
)

const maxRcvMsgSize = 1024 * 1024 * 100

// thread-safe events recording map
type eventsRecord struct {
	sync.RWMutex
	*pb.EventsRecord
}

type aggregatorExecutor struct {
	// thread-safe events recording maps
	sentEvents     *eventsRecord
	acceptedEvents *eventsRecord
	failedEvents   *eventsRecord
	receivedEvents *eventsRecord

	// channel to notify the main goroutine that an events record has been received
	notifyEventsReceived chan struct{}

	// GRPC server
	listener net.Listener
	server   *grpc.Server
}

func newAggregatorExecutor(lis net.Listener) testExecutor {
	executor := &aggregatorExecutor{
		listener:             lis,
		notifyEventsReceived: make(chan struct{}),
	}

	// --- Create GRPC server

	s := grpc.NewServer(grpc.MaxRecvMsgSize(maxRcvMsgSize))
	pb.RegisterEventsRecorderServer(s, executor)
	executor.server = s

	// --- Initialize records maps

	executor.sentEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_SENT,
		Events: make(map[string]*timestamp.Timestamp),
	}}
	executor.acceptedEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_ACCEPTED,
		Events: make(map[string]*timestamp.Timestamp),
	}}
	executor.failedEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_FAILED,
		Events: make(map[string]*timestamp.Timestamp),
	}}
	executor.receivedEvents = &eventsRecord{EventsRecord: &pb.EventsRecord{
		Type:   pb.EventsRecord_RECEIVED,
		Events: make(map[string]*timestamp.Timestamp),
	}}

	return executor
}

func (ex *aggregatorExecutor) Run(ctx context.Context) {
	// --- Configure mako

	printf("Configuring Mako")

	// Use the benchmark key created
	ctx, q, qclose, err := mako.Setup(ctx)
	if err != nil {
		fatalf("Failed to setup mako: %v", err)
	}

	// Use a fresh context here so that our RPC to terminate the sidecar
	// isn't subject to our timeout (or we won't shut it down when we time out)
	defer qclose(context.Background())

	// Wrap fatalf in a helper or our sidecar will live forever.
	fatalf = func(f string, args ...interface{}) {
		qclose(context.Background())
		fatalf(f, args...)
	}

	// --- Run GRPC events receiver

	printf("Starting events recorder server")

	go func() {
		if err := ex.server.Serve(ex.listener); err != nil {
			fatalf("Failed to serve: %v", err)
		}
	}()
	go func() {
		<-ctx.Done()
		printf("Terminating events recorder server")
		ex.server.GracefulStop()
	}()

	printf("Expecting %d events records", expectRecords)
	ex.waitForEvents()
	printf("Received all expected events records")

	ex.server.GracefulStop()

	// --- Publish latencies

	printf("%-15s: %d", "Sent count", len(ex.sentEvents.Events))
	printf("%-15s: %d", "Accepted count", len(ex.acceptedEvents.Events))
	printf("%-15s: %d", "Failed count", len(ex.failedEvents.Events))
	printf("%-15s: %d", "Received count", len(ex.receivedEvents.Events))

	printf("Publishing latencies")

	// count errors
	var publishErrorCount int
	var deliverErrorCount int

	for sentID := range ex.sentEvents.Events {
		timestampSentProto := ex.sentEvents.Events[sentID]
		timestampSent, _ := ptypes.Timestamp(timestampSentProto)

		timestampAcceptedProto, accepted := ex.acceptedEvents.Events[sentID]
		timestampAccepted, _ := ptypes.Timestamp(timestampAcceptedProto)

		timestampReceivedProto, received := ex.receivedEvents.Events[sentID]
		timestampReceived, _ := ptypes.Timestamp(timestampReceivedProto)

		if !accepted {
			errMsg := "Failed on broker"
			if _, failed := ex.failedEvents.Events[sentID]; !failed {
				// TODO(antoineco): should never happen, check whether the failed map makes any sense
				errMsg = "Event not accepted but missing from failed map"
			}

			deliverErrorCount++

			if qerr := q.AddError(mako.XTime(timestampSent), errMsg); qerr != nil {
				log.Printf("ERROR AddError: %v", qerr)
			}
			continue
		}

		sendLatency := timestampAccepted.Sub(timestampSent)
		// Uncomment to get CSV directly from this container log
		//fmt.Printf("%f,%d,\n", mako.XTime(timestampSent), sendLatency.Nanoseconds())
		// TODO mako accepts float64, which imo could lead to losing some precision on local tests. It should accept int64
		if qerr := q.AddSamplePoint(mako.XTime(timestampSent), map[string]float64{"pl": sendLatency.Seconds()}); qerr != nil {
			log.Printf("ERROR AddSamplePoint: %v", qerr)
		}

		if !received {
			publishErrorCount++

			if qerr := q.AddError(mako.XTime(timestampSent), "Event not delivered"); qerr != nil {
				log.Printf("ERROR AddError: %v", qerr)
			}
			continue
		}

		e2eLatency := timestampReceived.Sub(timestampSent)
		// Uncomment to get CSV directly from this container log
		//fmt.Printf("%f,,%d\n", mako.XTime(timestampSent), e2eLatency.Nanoseconds())
		// TODO mako accepts float64, which imo could lead to losing some precision on local tests. It should accept int64
		if qerr := q.AddSamplePoint(mako.XTime(timestampSent), map[string]float64{"dl": e2eLatency.Seconds()}); qerr != nil {
			log.Printf("ERROR AddSamplePoint: %v", qerr)
		}
	}

	// --- Publish throughput

	printf("Publishing throughputs")

	sentTimestamps := eventsToTimestampsArray(&ex.sentEvents.Events)
	err = publishThpt(sentTimestamps, q, "st")
	if err != nil {
		log.Printf("ERROR AddSamplePoint: %v", err)
	}

	receivedTimestamps := eventsToTimestampsArray(&ex.receivedEvents.Events)
	err = publishThpt(receivedTimestamps, q, "dt")
	if err != nil {
		log.Printf("ERROR AddSamplePoint: %v", err)
	}

	// --- Publish error counts as aggregate metrics

	printf("Publishing aggregates")

	q.AddRunAggregate("pe", float64(publishErrorCount))
	q.AddRunAggregate("de", float64(deliverErrorCount))

	printf("Store to mako")

	if out, err := q.Store(); err != nil {
		fatalf("Failed to store data: %v\noutput: %v", err, out)
	}

	printf("Aggregation completed")
}

func eventsToTimestampsArray(events *map[string]*timestamp.Timestamp) []time.Time {
	values := make([]time.Time, 0, len(*events))
	for _, v := range *events {
		t, _ := ptypes.Timestamp(v)
		values = append(values, t)
	}
	sort.Slice(values, func(x, y int) bool { return values[x].Before(values[y]) })
	return values
}

func publishThpt(timestamps []time.Time, q *quickstore.Quickstore, metricName string) error {
	for i, t := range timestamps[1:] {
		var thpt uint
		j := i - 1
		for j >= 0 && t.Sub(timestamps[j]) <= time.Second {
			thpt++
			j--
		}
		if qerr := q.AddSamplePoint(mako.XTime(t), map[string]float64{metricName: float64(thpt)}); qerr != nil {
			return qerr
		}
	}
	return nil
}

// waitForEvents blocks until the expected number of events records has been received.
func (ex *aggregatorExecutor) waitForEvents() {
	for receivedRecords := uint(0); receivedRecords < expectRecords; receivedRecords++ {
		<-ex.notifyEventsReceived
	}
}

// RecordSentEvents implements event_state.EventsRecorder
func (ex *aggregatorExecutor) RecordEvents(_ context.Context, in *pb.EventsRecordList) (*pb.RecordReply, error) {
	defer func() {
		ex.notifyEventsReceived <- struct{}{}
	}()

	for _, recIn := range in.Items {
		recType := recIn.GetType()

		var rec *eventsRecord

		switch recType {
		case pb.EventsRecord_SENT:
			rec = ex.sentEvents
		case pb.EventsRecord_ACCEPTED:
			rec = ex.acceptedEvents
		case pb.EventsRecord_FAILED:
			rec = ex.failedEvents
		case pb.EventsRecord_RECEIVED:
			rec = ex.receivedEvents
		default:
			printf("Ignoring events record of type %s", recType)
			continue
		}

		printf("-> Recording %d %s events", uint64(len(recIn.Events)), recType)

		func() {
			rec.Lock()
			defer rec.Unlock()
			for id, t := range recIn.Events {
				if _, exists := rec.Events[id]; exists {
					log.Printf("!! Found duplicate %s event ID %s", recType, id)
					continue
				}
				rec.Events[id] = t
			}
		}()
	}

	return &pb.RecordReply{Count: uint32(len(in.Items))}, nil
}
<|endoftext|>


<commit_before>package logwriter_test

import (
	"database/sql"

	testdb "github.com/erikstmartin/go-testdb"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/cloudfoundry-incubator/galera-healthcheck/cluster-health-logger/logwriter"
	"os"
	"io/ioutil"
)

var (
	logFile *os.File
)

var _ = Describe("Cluster Health Logger", func() {

	BeforeEach(func() {
		logFile, _ = ioutil.TempFile(os.TempDir(), "logFile")

	})

	AfterEach(func() {
		os.Remove(logFile.Name())
	})

	Context("when the log file does not exist", func() {
		BeforeEach(func() {
			os.Remove(logFile.Name())
		})

		It("writes headers to the file", func() {
			logWriter := logWriterTestHelper(logFile.Name())
			ts := "happy-time"
			logWriter.Write(ts)
			contents, err := ioutil.ReadFile(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
			contentsStr := string(contents)
			Expect(contentsStr).To(Equal("timestamp,a,b,c,d,e,f,g,h,i\nhappy-time,1,2,3,4,5,6,7,8,9\n"))
		})
	})

	Context("when the log file exists", func() {

		It("writes only the rows to the file", func() {
			logWriter := logWriterTestHelper(logFile.Name())
			ts := "happy-time"
			logWriter.Write(ts)
			contents, err := ioutil.ReadFile(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
			contentsStr := string(contents)
			Expect(contentsStr).To(Equal("happy-time,1,2,3,4,5,6,7,8,9\n"))
		})

		It("writes a new line", func() {
			logWriter := logWriterTestHelper(logFile.Name())
			ts1 := "happy-time"
			logWriter.Write(ts1)
			ts2 := "sad-time"
			logWriter.Write(ts2)
			contents, err := ioutil.ReadFile(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
			contentsStr := string(contents)
			Expect(contentsStr).To(Equal("happy-time,1,2,3,4,5,6,7,8,9\nsad-time,1,2,3,4,5,6,7,8,9\n"))
		})

	})
})

func logWriterTestHelper(filePath string) logwriter.LogWriter {
	db, _ := sql.Open("testdb", "")

	sql := "SHOW STATUS WHERE Variable_name IN ('wsrep_ready','wsrep_cluster_conf_id','wsrep_cluster_status','wsrep_connected','wsrep_local_state_comment','wsrep_local_recv_queue_avg','wsrep_flow_control_paused','wsrep_cert_deps_distance','wsrep_local_send_queue_avg')"
	columns := []string{"Variable_name", "Value"}
	result := "a,1\nb,2\nc,3\nd,4\ne,5\nf,6\ng,7\nh,8\ni,9"
	testdb.StubQuery(sql, testdb.RowsFromCSVString(columns, result))

	return logwriter.New(db, filePath)
}
<commit_msg>Handle errors in the cluster-health tests<commit_after>package logwriter_test

import (
	"database/sql"

	testdb "github.com/erikstmartin/go-testdb"

	"io/ioutil"
	"os"

	"github.com/cloudfoundry-incubator/galera-healthcheck/cluster-health-logger/logwriter"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var (
	logFile *os.File
	err     error
)

var _ = Describe("Cluster Health Logger", func() {

	BeforeEach(func() {
		logFile, err = ioutil.TempFile(os.TempDir(), "cluster-health-logger")
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		err = os.Remove(logFile.Name())
		Expect(err).ToNot(HaveOccurred())
	})

	Context("when the log file does not exist", func() {
		BeforeEach(func() {
			err = os.Remove(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
		})

		It("writes headers to the file", func() {
			logWriter := logWriterTestHelper(logFile.Name())
			ts := "happy-time"
			logWriter.Write(ts)
			contents, err := ioutil.ReadFile(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
			contentsStr := string(contents)
			Expect(contentsStr).To(Equal("timestamp,a,b,c,d,e,f,g,h,i\nhappy-time,1,2,3,4,5,6,7,8,9\n"))
		})
	})

	Context("when the log file exists", func() {

		It("writes only the rows to the file", func() {
			logWriter := logWriterTestHelper(logFile.Name())
			ts := "happy-time"
			logWriter.Write(ts)
			contents, err := ioutil.ReadFile(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
			contentsStr := string(contents)
			Expect(contentsStr).To(Equal("happy-time,1,2,3,4,5,6,7,8,9\n"))
		})

		It("writes a new line", func() {
			logWriter := logWriterTestHelper(logFile.Name())
			ts1 := "happy-time"
			logWriter.Write(ts1)
			ts2 := "sad-time"
			logWriter.Write(ts2)
			contents, err := ioutil.ReadFile(logFile.Name())
			Expect(err).ToNot(HaveOccurred())
			contentsStr := string(contents)
			Expect(contentsStr).To(Equal("happy-time,1,2,3,4,5,6,7,8,9\nsad-time,1,2,3,4,5,6,7,8,9\n"))
		})

	})
})

func logWriterTestHelper(filePath string) logwriter.LogWriter {
	db, err := sql.Open("testdb", "")
	Expect(err).ToNot(HaveOccurred())

	sql := "SHOW STATUS WHERE Variable_name IN ('wsrep_ready','wsrep_cluster_conf_id','wsrep_cluster_status','wsrep_connected','wsrep_local_state_comment','wsrep_local_recv_queue_avg','wsrep_flow_control_paused','wsrep_cert_deps_distance','wsrep_local_send_queue_avg')"
	columns := []string{"Variable_name", "Value"}
	result := "a,1\nb,2\nc,3\nd,4\ne,5\nf,6\ng,7\nh,8\ni,9"
	testdb.StubQuery(sql, testdb.RowsFromCSVString(columns, result))

	return logwriter.New(db, filePath)
}
<|endoftext|>