<commit_before>// Copyright 2017-2021 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 scripting

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"time"

	"github.com/OWASP/Amass/v3/config"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/systems"
	"github.com/caffix/eventbus"
	"github.com/caffix/service"
	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 {
	service.BaseService

	SourceType string
	sys        systems.System
	luaState   *lua.LState
	// Script callback functions
	start      lua.LValue
	stop       lua.LValue
	check      lua.LValue
	vertical   lua.LValue
	horizontal lua.LValue
	address    lua.LValue
	asn        lua.LValue
	resolved   lua.LValue
	subdomain  lua.LValue
	// Regexp to match any subdomain name
	subre   *regexp.Regexp
	seconds int
	cancel  context.CancelFunc
}

// NewScript returns he object initialized, but not yet started.
func NewScript(script string, sys systems.System) *Script {
	re, err := regexp.Compile(dns.AnySubdomainRegexString())
	if err != nil {
		return nil
	}

	ctx, cancel := context.WithCancel(context.Background())

	s := &Script{
		sys:    sys,
		subre:  re,
		cancel: cancel,
	}

	L := s.newLuaState(sys.Config())
	s.luaState = L
	L.SetContext(ctx)

	// Load the script
	if err := L.DoString(script); 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 = *service.NewBaseService(s, name)

	// 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("config", L.NewFunction(s.config))
	L.SetGlobal("datasrc_config", L.NewFunction(s.dataSourceConfig))
	L.SetGlobal("brute_wordlist", L.NewFunction(s.bruteWordlist))
	L.SetGlobal("alt_wordlist", L.NewFunction(s.altWordlist))
	L.SetGlobal("log", L.NewFunction(s.log))
	L.SetGlobal("find", L.NewFunction(s.find))
	L.SetGlobal("submatch", L.NewFunction(s.submatch))
	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("obtain_response", L.NewFunction(s.obtainResponse))
	L.SetGlobal("cache_response", L.NewFunction(s.cacheResponse))
	L.SetGlobal("subdomainre", lua.LString(dns.AnySubdomainRegexString()))
	return L
}

// 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.check = L.GetGlobal("check")
	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")
}

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

// OnStart implements the Service interface.
func (s *Script) OnStart() error {
	L := s.luaState
	if s.start.Type() == lua.LTNil {
		return nil
	}

	err := L.CallByParam(lua.P{
		Fn:      s.start,
		NRet:    0,
		Protect: true,
	})
	if err != nil {
		s.sys.Config().Log.Print(fmt.Sprintf("%s: start callback: %v", s.String(), err))
	}

	if s.seconds > 0 {
		s.SetRateLimit(1)
	}

	return s.checkConfig()
}

// OnStop implements the Service interface.
func (s *Script) OnStop() error {
	defer func() {
		s.cancel()
		time.Sleep(2 * time.Second)
		s.luaState.Close()
	}()

	var err error
	L := s.luaState
	if s.stop.Type() != lua.LTNil {
		err = L.CallByParam(lua.P{
			Fn:      s.stop,
			NRet:    0,
			Protect: true,
		})
		if err != nil {
			err = fmt.Errorf("%s: stop callback: %v", s.String(), err)

			s.sys.Config().Log.Print(err.Error())
		}
	}

	return err
}

func (s *Script) checkConfig() error {
	L := s.luaState

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

	err := L.CallByParam(lua.P{
		Fn:      s.check,
		NRet:    1,
		Protect: true,
	})
	if err != nil {
		estr := fmt.Sprintf("%s: check callback: %v", s.String(), err)

		s.sys.Config().Log.Print(estr)
		return errors.New(estr)
	}

	ret := L.Get(-1)
	L.Pop(1)

	passed, ok := ret.(lua.LBool)
	if ok && bool(passed) {
		return nil
	}

	estr := fmt.Sprintf("%s: check callback failed for the configuration", s.String())
	s.sys.Config().Log.Print(estr)
	return errors.New(estr)
}

// OnRequest implements the Service interface.
func (s *Script) OnRequest(ctx context.Context, args service.Args) {
	var check bool

	switch req := args.(type) {
	case *requests.DNSRequest:
		if s.vertical.Type() != lua.LTNil && req != nil && req.Domain != "" {
			s.dnsRequest(ctx, req)
			check = true
		}
	case *requests.ResolvedRequest:
		if s.resolved.Type() != lua.LTNil && req != nil && req.Name != "" && len(req.Records) > 0 {
			s.resolvedRequest(ctx, req)
			check = true
		}
	case *requests.SubdomainRequest:
		if s.subdomain.Type() != lua.LTNil && req != nil && req.Name != "" {
			s.subdomainRequest(ctx, req)
			check = true
		}
	case *requests.AddrRequest:
		if s.address.Type() != lua.LTNil && req != nil && req.Address != "" {
			s.addrRequest(ctx, req)
			check = true
		}
	case *requests.ASNRequest:
		if s.asn.Type() != lua.LTNil && req != nil && (req.Address != "" || req.ASN != 0) {
			s.asnRequest(ctx, req)
			check = true
		}
	case *requests.WhoisRequest:
		if s.horizontal.Type() != lua.LTNil {
			s.whoisRequest(ctx, req)
			check = true
		}
	}

	if check {
		numRateLimitChecks(s, s.seconds)
	}
}

func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: vertical callback: %v", s.String(), err))
	}
}

func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

	records := L.NewTable()
	for _, rec := range req.Records {
		tb := L.NewTable()

		tb.RawSetString("rrname", lua.LString(rec.Name))
		tb.RawSetString("rrtype", lua.LNumber(rec.Type))
		tb.RawSetString("rrdata", lua.LString(rec.Data))
		records.Append(tb)
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: resolved callback: %v", s.String(), err))
	}
}

func (s *Script) subdomainRequest(ctx context.Context, req *requests.SubdomainRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: subdomain callback: %v", s.String(), err))
	}
}

func (s *Script) addrRequest(ctx context.Context, req *requests.AddrRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: address callback: %v", s.String(), err))
	}
}

func (s *Script) asnRequest(ctx context.Context, req *requests.ASNRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: asn callback: %v", s.String(), err))
	}
}

func (s *Script) whoisRequest(ctx context.Context, req *requests.WhoisRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: horizontal callback: %v", s.String(), err))
	}
}
<commit_msg>Increased script request processing<commit_after>// Copyright 2017-2021 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 scripting

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"time"

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

// Script callback functions
type callbacks struct {
	Start      lua.LValue
	Stop       lua.LValue
	Check      lua.LValue
	Vertical   lua.LValue
	Horizontal lua.LValue
	Address    lua.LValue
	Asn        lua.LValue
	Resolved   lua.LValue
	Subdomain  lua.LValue
}

// Script is the Service that handles access to the Script data source.
type Script struct {
	service.BaseService
	SourceType string
	sys        systems.System
	luaState   *lua.LState
	cbs        *callbacks
	subre      *regexp.Regexp
	seconds    int
	cancel     context.CancelFunc
}

// NewScript returns he object initialized, but not yet started.
func NewScript(script string, sys systems.System) *Script {
	re, err := regexp.Compile(dns.AnySubdomainRegexString())
	if err != nil {
		return nil
	}

	ctx, cancel := context.WithCancel(context.Background())

	s := &Script{
		sys:    sys,
		subre:  re,
		cancel: cancel,
	}

	L := s.newLuaState(sys.Config())
	s.luaState = L
	L.SetContext(ctx)

	// Load the script
	if err := L.DoString(script); 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 = *service.NewBaseService(s, name)

	// Save references to the callbacks defined within the script
	s.assignCallbacks()
	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("config", L.NewFunction(s.config))
	L.SetGlobal("datasrc_config", L.NewFunction(s.dataSourceConfig))
	L.SetGlobal("brute_wordlist", L.NewFunction(s.bruteWordlist))
	L.SetGlobal("alt_wordlist", L.NewFunction(s.altWordlist))
	L.SetGlobal("log", L.NewFunction(s.log))
	L.SetGlobal("find", L.NewFunction(s.find))
	L.SetGlobal("submatch", L.NewFunction(s.submatch))
	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("obtain_response", L.NewFunction(s.obtainResponse))
	L.SetGlobal("cache_response", L.NewFunction(s.cacheResponse))
	L.SetGlobal("subdomainre", lua.LString(dns.AnySubdomainRegexString()))
	return L
}

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

	s.cbs = &callbacks{
		Start:      L.GetGlobal("start"),
		Stop:       L.GetGlobal("stop"),
		Check:      L.GetGlobal("check"),
		Vertical:   L.GetGlobal("vertical"),
		Horizontal: L.GetGlobal("horizontal"),
		Address:    L.GetGlobal("address"),
		Asn:        L.GetGlobal("asn"),
		Resolved:   L.GetGlobal("resolved"),
		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")
}

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

// OnStart implements the Service interface.
func (s *Script) OnStart() error {
	L := s.luaState
	if s.cbs.Start.Type() == lua.LTNil {
		return nil
	}

	err := L.CallByParam(lua.P{
		Fn:      s.cbs.Start,
		NRet:    0,
		Protect: true,
	})
	if err != nil {
		s.sys.Config().Log.Print(fmt.Sprintf("%s: start callback: %v", s.String(), err))
	}

	if s.seconds > 0 {
		s.SetRateLimit(1)
	}

	return s.checkConfig()
}

// OnStop implements the Service interface.
func (s *Script) OnStop() error {
	defer func() {
		s.cancel()
		time.Sleep(2 * time.Second)
		s.luaState.Close()
	}()

	var err error
	L := s.luaState
	if s.cbs.Stop.Type() != lua.LTNil {
		err = L.CallByParam(lua.P{
			Fn:      s.cbs.Stop,
			NRet:    0,
			Protect: true,
		})
		if err != nil {
			err = fmt.Errorf("%s: stop callback: %v", s.String(), err)

			s.sys.Config().Log.Print(err.Error())
		}
	}

	return err
}

func (s *Script) checkConfig() error {
	L := s.luaState

	if s.cbs.Check.Type() == lua.LTNil {
		return nil
	}

	err := L.CallByParam(lua.P{
		Fn:      s.cbs.Check,
		NRet:    1,
		Protect: true,
	})
	if err != nil {
		estr := fmt.Sprintf("%s: check callback: %v", s.String(), err)

		s.sys.Config().Log.Print(estr)
		return errors.New(estr)
	}

	ret := L.Get(-1)
	L.Pop(1)

	passed, ok := ret.(lua.LBool)
	if ok && bool(passed) {
		return nil
	}

	estr := fmt.Sprintf("%s: check callback failed for the configuration", s.String())
	s.sys.Config().Log.Print(estr)
	return errors.New(estr)
}

// OnRequest implements the Service interface.
func (s *Script) OnRequest(ctx context.Context, args service.Args) {
	switch req := args.(type) {
	case *requests.DNSRequest:
		if s.cbs.Vertical.Type() != lua.LTNil && req != nil && req.Domain != "" {
			s.dnsRequest(ctx, req)
		}
	case *requests.ResolvedRequest:
		if s.cbs.Resolved.Type() != lua.LTNil && req != nil && req.Name != "" && len(req.Records) > 0 {
			s.resolvedRequest(ctx, req)
		}
	case *requests.SubdomainRequest:
		if s.cbs.Subdomain.Type() != lua.LTNil && req != nil && req.Name != "" {
			s.subdomainRequest(ctx, req)
		}
	case *requests.AddrRequest:
		if s.cbs.Address.Type() != lua.LTNil && req != nil && req.Address != "" {
			s.addrRequest(ctx, req)
		}
	case *requests.ASNRequest:
		if s.cbs.Asn.Type() != lua.LTNil && req != nil && (req.Address != "" || req.ASN != 0) {
			s.asnRequest(ctx, req)
		}
	case *requests.WhoisRequest:
		if s.cbs.Horizontal.Type() != lua.LTNil {
			s.whoisRequest(ctx, req)
		}
	}
}

func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: vertical callback: %v", s.String(), err))
	}
}

func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

	records := L.NewTable()
	for _, rec := range req.Records {
		tb := L.NewTable()

		tb.RawSetString("rrname", lua.LString(rec.Name))
		tb.RawSetString("rrtype", lua.LNumber(rec.Type))
		tb.RawSetString("rrdata", lua.LString(rec.Data))
		records.Append(tb)
	}

	err = L.CallByParam(lua.P{
		Fn:      s.cbs.Resolved,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LString(req.Domain), records)

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: resolved callback: %v", s.String(), err))
	}
}

func (s *Script) subdomainRequest(ctx context.Context, req *requests.SubdomainRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

	err = L.CallByParam(lua.P{
		Fn:      s.cbs.Subdomain,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LString(req.Domain), lua.LNumber(req.Times))

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: subdomain callback: %v", s.String(), err))
	}
}

func (s *Script) addrRequest(ctx context.Context, req *requests.AddrRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: address callback: %v", s.String(), err))
	}
}

func (s *Script) asnRequest(ctx context.Context, req *requests.ASNRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

	err = L.CallByParam(lua.P{
		Fn:      s.cbs.Asn,
		NRet:    0,
		Protect: true,
	}, s.contextToUserData(ctx), lua.LString(req.Address), lua.LNumber(req.ASN))

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: asn callback: %v", s.String(), err))
	}
}

func (s *Script) whoisRequest(ctx context.Context, req *requests.WhoisRequest) {
	L := s.luaState

	if err := checkContextExpired(ctx); err != nil {
		return
	}

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		return
	}

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

	if err != nil {
		bus.Publish(requests.LogTopic, eventbus.PriorityHigh,
			fmt.Sprintf("%s: horizontal callback: %v", s.String(), err))
	}
}
<|endoftext|>


<commit_before>package consul

import (
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/consul/agent/pool"
	"github.com/hashicorp/consul/agent/router"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/lib"
	"github.com/hashicorp/serf/serf"
	"golang.org/x/time/rate"
)

const (
	// clientRPCConnMaxIdle controls how long we keep an idle connection
	// open to a server.  127s was chosen as the first prime above 120s
	// (arbitrarily chose to use a prime) with the intent of reusing
	// connections who are used by once-a-minute cron(8) jobs *and* who
	// use a 60s jitter window (e.g. in vixie cron job execution can
	// drift by up to 59s per job, or 119s for a once-a-minute cron job).
	clientRPCConnMaxIdle = 127 * time.Second

	// clientMaxStreams controls how many idle streams we keep
	// open to a server
	clientMaxStreams = 32

	// serfEventBacklog is the maximum number of unprocessed Serf Events
	// that will be held in queue before new serf events block.  A
	// blocking serf event queue is a bad thing.
	serfEventBacklog = 256

	// serfEventBacklogWarning is the threshold at which point log
	// warnings will be emitted indicating a problem when processing serf
	// events.
	serfEventBacklogWarning = 200
)

// Client is Consul client which uses RPC to communicate with the
// services for service discovery, health checking, and DC forwarding.
type Client struct {
	config *Config

	// Connection pool to consul servers
	connPool *pool.ConnPool

	// routers is responsible for the selection and maintenance of
	// Consul servers this agent uses for RPC requests
	routers *router.Manager

	// rpcLimiter is used to rate limit the total number of RPCs initiated
	// from an agent.
	rpcLimiter *rate.Limiter

	// eventCh is used to receive events from the
	// serf cluster in the datacenter
	eventCh chan serf.Event

	// Logger uses the provided LogOutput
	logger *log.Logger

	// serf is the Serf cluster maintained inside the DC
	// which contains all the DC nodes
	serf *serf.Serf

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex

	// embedded struct to hold all the enterprise specific data
	EnterpriseClient
}

// NewClient is used to construct a new Consul client from the
// configuration, potentially returning an error
func NewClient(config *Config) (*Client, error) {
	return NewClientLogger(config, nil)
}

func NewClientLogger(config *Config, logger *log.Logger) (*Client, error) {
	// Check the protocol version
	if err := config.CheckProtocolVersion(); err != nil {
		return nil, err
	}

	// Check for a data directory!
	if config.DataDir == "" {
		return nil, fmt.Errorf("Config must provide a DataDir")
	}

	// Sanity check the ACLs
	if err := config.CheckACL(); err != nil {
		return nil, err
	}

	// Ensure we have a log output
	if config.LogOutput == nil {
		config.LogOutput = os.Stderr
	}

	// Create the tls Wrapper
	tlsWrap, err := config.tlsConfig().OutgoingTLSWrapper()
	if err != nil {
		return nil, err
	}

	// Create a logger
	if logger == nil {
		logger = log.New(config.LogOutput, "", log.LstdFlags)
	}

	connPool := &pool.ConnPool{
		SrcAddr:    config.RPCSrcAddr,
		LogOutput:  config.LogOutput,
		MaxTime:    clientRPCConnMaxIdle,
		MaxStreams: clientMaxStreams,
		TLSWrapper: tlsWrap,
		ForceTLS:   config.VerifyOutgoing,
	}

	// Create client
	c := &Client{
		config:     config,
		connPool:   connPool,
		rpcLimiter: rate.NewLimiter(config.RPCRate, config.RPCMaxBurst),
		eventCh:    make(chan serf.Event, serfEventBacklog),
		logger:     logger,
		shutdownCh: make(chan struct{}),
	}

	if err := c.initEnterprise(); err != nil {
		c.Shutdown()
		return nil, err
	}

	// Initialize the LAN Serf
	c.serf, err = c.setupSerf(config.SerfLANConfig,
		c.eventCh, serfLANSnapshot)
	if err != nil {
		c.Shutdown()
		return nil, fmt.Errorf("Failed to start lan serf: %v", err)
	}

	// Start maintenance task for servers
	c.routers = router.New(c.logger, c.shutdownCh, c.serf, c.connPool)
	go c.routers.Start()

	// Start LAN event handlers after the router is complete since the event
	// handlers depend on the router and the router depends on Serf.
	go c.lanEventHandler()

	if err := c.startEnterprise(); err != nil {
		c.Shutdown()
		return nil, err
	}

	return c, nil
}

// Shutdown is used to shutdown the client
func (c *Client) Shutdown() error {
	c.logger.Printf("[INFO] consul: shutting down client")
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if c.shutdown {
		return nil
	}

	c.shutdown = true
	close(c.shutdownCh)

	if c.serf != nil {
		c.serf.Shutdown()
	}

	// Close the connection pool
	c.connPool.Shutdown()
	return nil
}

// Leave is used to prepare for a graceful shutdown
func (c *Client) Leave() error {
	c.logger.Printf("[INFO] consul: client starting leave")

	// Leave the LAN pool
	if c.serf != nil {
		if err := c.serf.Leave(); err != nil {
			c.logger.Printf("[ERR] consul: Failed to leave LAN Serf cluster: %v", err)
		}
	}
	return nil
}

// JoinLAN is used to have Consul client join the inner-DC pool
// The target address should be another node inside the DC
// listening on the Serf LAN address
func (c *Client) JoinLAN(addrs []string) (int, error) {
	return c.serf.Join(addrs, true)
}

// LocalMember is used to return the local node
func (c *Client) LocalMember() serf.Member {
	return c.serf.LocalMember()
}

// LANMembers is used to return the members of the LAN cluster
func (c *Client) LANMembers() []serf.Member {
	return c.serf.Members()
}

// LANMembersAllSegments returns members from all segments.
func (c *Client) LANMembersAllSegments() ([]serf.Member, error) {
	return c.serf.Members(), nil
}

// LANSegmentMembers only returns our own segment's members, because clients
// can't be in multiple segments.
func (c *Client) LANSegmentMembers(segment string) ([]serf.Member, error) {
	if segment == c.config.Segment {
		return c.LANMembers(), nil
	}

	return nil, fmt.Errorf("segment %q not found", segment)
}

// RemoveFailedNode is used to remove a failed node from the cluster
func (c *Client) RemoveFailedNode(node string) error {
	return c.serf.RemoveFailedNode(node)
}

// KeyManagerLAN returns the LAN Serf keyring manager
func (c *Client) KeyManagerLAN() *serf.KeyManager {
	return c.serf.KeyManager()
}

// Encrypted determines if gossip is encrypted
func (c *Client) Encrypted() bool {
	return c.serf.EncryptionEnabled()
}

// RPC is used to forward an RPC call to a consul server, or fail if no servers
func (c *Client) RPC(method string, args interface{}, reply interface{}) error {
	// This is subtle but we start measuring the time on the client side
	// right at the time of the first request, vs. on the first retry as
	// is done on the server side inside forward(). This is because the
	// servers may already be applying the RPCHoldTimeout up there, so by
	// starting the timer here we won't potentially double up the delay.
	// TODO (slackpad) Plumb a deadline here with a context.
	firstCheck := time.Now()

TRY:
	server := c.routers.FindServer()
	if server == nil {
		return structs.ErrNoServers
	}

	// Enforce the RPC limit.
	metrics.IncrCounter([]string{"client", "rpc"}, 1)
	if !c.rpcLimiter.Allow() {
		metrics.IncrCounter([]string{"client", "rpc", "exceeded"}, 1)
		return structs.ErrRPCRateExceeded
	}

	// Make the request.
	rpcErr := c.connPool.RPC(c.config.Datacenter, server.Addr, server.Version, method, server.UseTLS, args, reply)
	if rpcErr == nil {
		return nil
	}

	// Move off to another server, and see if we can retry.
	c.logger.Printf("[ERR] consul: %q RPC failed to server %s: %v", method, server.Addr, rpcErr)
	metrics.IncrCounterWithLabels([]string{"client", "rpc", "failed"}, 1, []metrics.Label{{Name: "server", Value: server.ID}})
	c.routers.NotifyFailedServer(server)
	if retry := canRetry(args, rpcErr); !retry {
		return rpcErr
	}

	// We can wait a bit and retry!
	if time.Since(firstCheck) < c.config.RPCHoldTimeout {
		jitter := lib.RandomStagger(c.config.RPCHoldTimeout / jitterFraction)
		select {
		case <-time.After(jitter):
			goto TRY
		case <-c.shutdownCh:
		}
	}
	return rpcErr
}

// SnapshotRPC sends the snapshot request to one of the servers, reading from
// the streaming input and writing to the streaming output depending on the
// operation.
func (c *Client) SnapshotRPC(args *structs.SnapshotRequest, in io.Reader, out io.Writer,
	replyFn structs.SnapshotReplyFn) error {
	server := c.routers.FindServer()
	if server == nil {
		return structs.ErrNoServers
	}

	// Enforce the RPC limit.
	metrics.IncrCounter([]string{"client", "rpc"}, 1)
	if !c.rpcLimiter.Allow() {
		metrics.IncrCounter([]string{"client", "rpc", "exceeded"}, 1)
		return structs.ErrRPCRateExceeded
	}

	// Request the operation.
	var reply structs.SnapshotResponse
	snap, err := SnapshotRPC(c.connPool, c.config.Datacenter, server.Addr, server.UseTLS, args, in, &reply)
	if err != nil {
		return err
	}
	defer func() {
		if err := snap.Close(); err != nil {
			c.logger.Printf("[WARN] consul: Failed closing snapshot stream: %v", err)
		}
	}()

	// Let the caller peek at the reply.
	if replyFn != nil {
		if err := replyFn(&reply); err != nil {
			return nil
		}
	}

	// Stream the snapshot.
	if out != nil {
		if _, err := io.Copy(out, snap); err != nil {
			return fmt.Errorf("failed to stream snapshot: %v", err)
		}
	}

	return nil
}

// Stats is used to return statistics for debugging and insight
// for various sub-systems
func (c *Client) Stats() map[string]map[string]string {
	numServers := c.routers.NumServers()

	toString := func(v uint64) string {
		return strconv.FormatUint(v, 10)
	}
	stats := map[string]map[string]string{
		"consul": map[string]string{
			"server":        "false",
			"known_servers": toString(uint64(numServers)),
		},
		"serf_lan": c.serf.Stats(),
		"runtime":  runtimeStats(),
	}

	for outerKey, outerValue := range c.enterpriseStats() {
		if _, ok := stats[outerKey]; ok {
			for innerKey, innerValue := range outerValue {
				stats[outerKey][innerKey] = innerValue
			}
		} else {
			stats[outerKey] = outerValue
		}
	}

	return stats
}

// GetLANCoordinate returns the network coordinate of the current node, as
// maintained by Serf.
func (c *Client) GetLANCoordinate() (lib.CoordinateSet, error) {
	lan, err := c.serf.GetCoordinate()
	if err != nil {
		return nil, err
	}

	cs := lib.CoordinateSet{c.config.Segment: lan}
	return cs, nil
}
<commit_msg>Attach server.Name label to client.rpc.failed<commit_after>package consul

import (
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/consul/agent/pool"
	"github.com/hashicorp/consul/agent/router"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/lib"
	"github.com/hashicorp/serf/serf"
	"golang.org/x/time/rate"
)

const (
	// clientRPCConnMaxIdle controls how long we keep an idle connection
	// open to a server.  127s was chosen as the first prime above 120s
	// (arbitrarily chose to use a prime) with the intent of reusing
	// connections who are used by once-a-minute cron(8) jobs *and* who
	// use a 60s jitter window (e.g. in vixie cron job execution can
	// drift by up to 59s per job, or 119s for a once-a-minute cron job).
	clientRPCConnMaxIdle = 127 * time.Second

	// clientMaxStreams controls how many idle streams we keep
	// open to a server
	clientMaxStreams = 32

	// serfEventBacklog is the maximum number of unprocessed Serf Events
	// that will be held in queue before new serf events block.  A
	// blocking serf event queue is a bad thing.
	serfEventBacklog = 256

	// serfEventBacklogWarning is the threshold at which point log
	// warnings will be emitted indicating a problem when processing serf
	// events.
	serfEventBacklogWarning = 200
)

// Client is Consul client which uses RPC to communicate with the
// services for service discovery, health checking, and DC forwarding.
type Client struct {
	config *Config

	// Connection pool to consul servers
	connPool *pool.ConnPool

	// routers is responsible for the selection and maintenance of
	// Consul servers this agent uses for RPC requests
	routers *router.Manager

	// rpcLimiter is used to rate limit the total number of RPCs initiated
	// from an agent.
	rpcLimiter *rate.Limiter

	// eventCh is used to receive events from the
	// serf cluster in the datacenter
	eventCh chan serf.Event

	// Logger uses the provided LogOutput
	logger *log.Logger

	// serf is the Serf cluster maintained inside the DC
	// which contains all the DC nodes
	serf *serf.Serf

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex

	// embedded struct to hold all the enterprise specific data
	EnterpriseClient
}

// NewClient is used to construct a new Consul client from the
// configuration, potentially returning an error
func NewClient(config *Config) (*Client, error) {
	return NewClientLogger(config, nil)
}

func NewClientLogger(config *Config, logger *log.Logger) (*Client, error) {
	// Check the protocol version
	if err := config.CheckProtocolVersion(); err != nil {
		return nil, err
	}

	// Check for a data directory!
	if config.DataDir == "" {
		return nil, fmt.Errorf("Config must provide a DataDir")
	}

	// Sanity check the ACLs
	if err := config.CheckACL(); err != nil {
		return nil, err
	}

	// Ensure we have a log output
	if config.LogOutput == nil {
		config.LogOutput = os.Stderr
	}

	// Create the tls Wrapper
	tlsWrap, err := config.tlsConfig().OutgoingTLSWrapper()
	if err != nil {
		return nil, err
	}

	// Create a logger
	if logger == nil {
		logger = log.New(config.LogOutput, "", log.LstdFlags)
	}

	connPool := &pool.ConnPool{
		SrcAddr:    config.RPCSrcAddr,
		LogOutput:  config.LogOutput,
		MaxTime:    clientRPCConnMaxIdle,
		MaxStreams: clientMaxStreams,
		TLSWrapper: tlsWrap,
		ForceTLS:   config.VerifyOutgoing,
	}

	// Create client
	c := &Client{
		config:     config,
		connPool:   connPool,
		rpcLimiter: rate.NewLimiter(config.RPCRate, config.RPCMaxBurst),
		eventCh:    make(chan serf.Event, serfEventBacklog),
		logger:     logger,
		shutdownCh: make(chan struct{}),
	}

	if err := c.initEnterprise(); err != nil {
		c.Shutdown()
		return nil, err
	}

	// Initialize the LAN Serf
	c.serf, err = c.setupSerf(config.SerfLANConfig,
		c.eventCh, serfLANSnapshot)
	if err != nil {
		c.Shutdown()
		return nil, fmt.Errorf("Failed to start lan serf: %v", err)
	}

	// Start maintenance task for servers
	c.routers = router.New(c.logger, c.shutdownCh, c.serf, c.connPool)
	go c.routers.Start()

	// Start LAN event handlers after the router is complete since the event
	// handlers depend on the router and the router depends on Serf.
	go c.lanEventHandler()

	if err := c.startEnterprise(); err != nil {
		c.Shutdown()
		return nil, err
	}

	return c, nil
}

// Shutdown is used to shutdown the client
func (c *Client) Shutdown() error {
	c.logger.Printf("[INFO] consul: shutting down client")
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if c.shutdown {
		return nil
	}

	c.shutdown = true
	close(c.shutdownCh)

	if c.serf != nil {
		c.serf.Shutdown()
	}

	// Close the connection pool
	c.connPool.Shutdown()
	return nil
}

// Leave is used to prepare for a graceful shutdown
func (c *Client) Leave() error {
	c.logger.Printf("[INFO] consul: client starting leave")

	// Leave the LAN pool
	if c.serf != nil {
		if err := c.serf.Leave(); err != nil {
			c.logger.Printf("[ERR] consul: Failed to leave LAN Serf cluster: %v", err)
		}
	}
	return nil
}

// JoinLAN is used to have Consul client join the inner-DC pool
// The target address should be another node inside the DC
// listening on the Serf LAN address
func (c *Client) JoinLAN(addrs []string) (int, error) {
	return c.serf.Join(addrs, true)
}

// LocalMember is used to return the local node
func (c *Client) LocalMember() serf.Member {
	return c.serf.LocalMember()
}

// LANMembers is used to return the members of the LAN cluster
func (c *Client) LANMembers() []serf.Member {
	return c.serf.Members()
}

// LANMembersAllSegments returns members from all segments.
func (c *Client) LANMembersAllSegments() ([]serf.Member, error) {
	return c.serf.Members(), nil
}

// LANSegmentMembers only returns our own segment's members, because clients
// can't be in multiple segments.
func (c *Client) LANSegmentMembers(segment string) ([]serf.Member, error) {
	if segment == c.config.Segment {
		return c.LANMembers(), nil
	}

	return nil, fmt.Errorf("segment %q not found", segment)
}

// RemoveFailedNode is used to remove a failed node from the cluster
func (c *Client) RemoveFailedNode(node string) error {
	return c.serf.RemoveFailedNode(node)
}

// KeyManagerLAN returns the LAN Serf keyring manager
func (c *Client) KeyManagerLAN() *serf.KeyManager {
	return c.serf.KeyManager()
}

// Encrypted determines if gossip is encrypted
func (c *Client) Encrypted() bool {
	return c.serf.EncryptionEnabled()
}

// RPC is used to forward an RPC call to a consul server, or fail if no servers
func (c *Client) RPC(method string, args interface{}, reply interface{}) error {
	// This is subtle but we start measuring the time on the client side
	// right at the time of the first request, vs. on the first retry as
	// is done on the server side inside forward(). This is because the
	// servers may already be applying the RPCHoldTimeout up there, so by
	// starting the timer here we won't potentially double up the delay.
	// TODO (slackpad) Plumb a deadline here with a context.
	firstCheck := time.Now()

TRY:
	server := c.routers.FindServer()
	if server == nil {
		return structs.ErrNoServers
	}

	// Enforce the RPC limit.
	metrics.IncrCounter([]string{"client", "rpc"}, 1)
	if !c.rpcLimiter.Allow() {
		metrics.IncrCounter([]string{"client", "rpc", "exceeded"}, 1)
		return structs.ErrRPCRateExceeded
	}

	// Make the request.
	rpcErr := c.connPool.RPC(c.config.Datacenter, server.Addr, server.Version, method, server.UseTLS, args, reply)
	if rpcErr == nil {
		return nil
	}

	// Move off to another server, and see if we can retry.
	c.logger.Printf("[ERR] consul: %q RPC failed to server %s: %v", method, server.Addr, rpcErr)
	metrics.IncrCounterWithLabels([]string{"client", "rpc", "failed"}, 1, []metrics.Label{{Name: "server", Value: server.Name}})
	c.routers.NotifyFailedServer(server)
	if retry := canRetry(args, rpcErr); !retry {
		return rpcErr
	}

	// We can wait a bit and retry!
	if time.Since(firstCheck) < c.config.RPCHoldTimeout {
		jitter := lib.RandomStagger(c.config.RPCHoldTimeout / jitterFraction)
		select {
		case <-time.After(jitter):
			goto TRY
		case <-c.shutdownCh:
		}
	}
	return rpcErr
}

// SnapshotRPC sends the snapshot request to one of the servers, reading from
// the streaming input and writing to the streaming output depending on the
// operation.
func (c *Client) SnapshotRPC(args *structs.SnapshotRequest, in io.Reader, out io.Writer,
	replyFn structs.SnapshotReplyFn) error {
	server := c.routers.FindServer()
	if server == nil {
		return structs.ErrNoServers
	}

	// Enforce the RPC limit.
	metrics.IncrCounter([]string{"client", "rpc"}, 1)
	if !c.rpcLimiter.Allow() {
		metrics.IncrCounter([]string{"client", "rpc", "exceeded"}, 1)
		return structs.ErrRPCRateExceeded
	}

	// Request the operation.
	var reply structs.SnapshotResponse
	snap, err := SnapshotRPC(c.connPool, c.config.Datacenter, server.Addr, server.UseTLS, args, in, &reply)
	if err != nil {
		return err
	}
	defer func() {
		if err := snap.Close(); err != nil {
			c.logger.Printf("[WARN] consul: Failed closing snapshot stream: %v", err)
		}
	}()

	// Let the caller peek at the reply.
	if replyFn != nil {
		if err := replyFn(&reply); err != nil {
			return nil
		}
	}

	// Stream the snapshot.
	if out != nil {
		if _, err := io.Copy(out, snap); err != nil {
			return fmt.Errorf("failed to stream snapshot: %v", err)
		}
	}

	return nil
}

// Stats is used to return statistics for debugging and insight
// for various sub-systems
func (c *Client) Stats() map[string]map[string]string {
	numServers := c.routers.NumServers()

	toString := func(v uint64) string {
		return strconv.FormatUint(v, 10)
	}
	stats := map[string]map[string]string{
		"consul": map[string]string{
			"server":        "false",
			"known_servers": toString(uint64(numServers)),
		},
		"serf_lan": c.serf.Stats(),
		"runtime":  runtimeStats(),
	}

	for outerKey, outerValue := range c.enterpriseStats() {
		if _, ok := stats[outerKey]; ok {
			for innerKey, innerValue := range outerValue {
				stats[outerKey][innerKey] = innerValue
			}
		} else {
			stats[outerKey] = outerValue
		}
	}

	return stats
}

// GetLANCoordinate returns the network coordinate of the current node, as
// maintained by Serf.
func (c *Client) GetLANCoordinate() (lib.CoordinateSet, error) {
	lan, err := c.serf.GetCoordinate()
	if err != nil {
		return nil, err
	}

	cs := lib.CoordinateSet{c.config.Segment: lan}
	return cs, nil
}
<|endoftext|>


<commit_before>package schema

// TODO: copy disco schema here.

// Meta contains the archive and parse metadata.
type Meta struct {
	FileName  string `json:"task_filename, string" bigquery:"task_filename"`
	TestName  string `json:"test_id, string" bigquery:"test_id"`
	ParseTime int64  `json:"parse_time, int64" bigquery:"parse_time"`
}

// Sample is an individual measurement taken by DISCO.
type Sample struct {
	Timestamp int64   `json:"timestamp, int64" bigquery:"timestamp"`
	Value     float32 `json:"value, float32" bigquery:"value"`
}

// SwitchStats represents a row of data taken from the raw DISCO export file.
type SwitchStats struct {
	Meta       Meta     `json:"meta" bigquery:"meta"`
	Sample     []Sample `json:"sample" bigquery:"sample"`
	Metric     string   `json:"metric" bigquery:"metric"`
	Hostname   string   `json:"hostname" bigquery:"hostname"`
	Experiment string   `json:"experiment" bigquery:"experiment"`
}
<commit_msg>Add SwitchStats row size estimation.<commit_after>package schema

// TODO: copy disco schema here.

// Meta contains the archive and parse metadata.
type Meta struct {
	FileName  string `json:"task_filename, string" bigquery:"task_filename"`
	TestName  string `json:"test_id, string" bigquery:"test_id"`
	ParseTime int64  `json:"parse_time, int64" bigquery:"parse_time"`
}

// Sample is an individual measurement taken by DISCO.
type Sample struct {
	Timestamp int64   `json:"timestamp, int64" bigquery:"timestamp"`
	Value     float32 `json:"value, float32" bigquery:"value"`
}

// SwitchStats represents a row of data taken from the raw DISCO export file.
type SwitchStats struct {
	Meta       Meta     `json:"meta" bigquery:"meta"`
	Sample     []Sample `json:"sample" bigquery:"sample"`
	Metric     string   `json:"metric" bigquery:"metric"`
	Hostname   string   `json:"hostname" bigquery:"hostname"`
	Experiment string   `json:"experiment" bigquery:"experiment"`
}

// Size estimates the number of bytes in the SwitchStats object.
func (s *SwitchStats) Size() int {
	return (len(s.Meta.FileName) + len(s.Meta.TestName) + 8 +
		12*len(s.Sample) + len(s.Metric) + len(s.Hostname) + len(s.Experiment))
}
<|endoftext|>


<commit_before>// Copyright 2014 The Transporter Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// seed is a reimagining of the seed mongo to mongo tool

package main

import (
	"fmt"
	"os"
	"time"

	"github.com/compose/transporter/pkg/events"
	"github.com/compose/transporter/pkg/transporter"
)

var (
	sourceUri string     = os.Getenv("SOURCE_MONGO_URL")
	destUri string       = os.Getenv("DESTINATION_MONGO_URL")
	sourceNS string      = os.Getenv("SOURCE_NS")
	destinationNS string = os.Getenv("DEST_NS")
	tail	             = true
	debug       	     = true
)

func main() {
	source :=
		transporter.NewNode("source", "mongo", map[string]interface{}{"uri": fmt.Println(sourceUri), "namespace": fmt.Println(sourceNS), "tail": fmt.Println(tail)}).
			Add(transporter.NewNode("out", "mongo", map[string]interface{}{"uri": fmt.Println(destUri), "namespace": fmt.Println(destinationNS)}))

	if fmt.Println(debug) == true {
		source.Add(transporter.NewNode("out", "file", map[string]interface{}{"uri": "stdout://"}))
	}

	pipeline, err := transporter.NewPipeline(source, events.NewLogEmitter(), 1*time.Second)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	pipeline.Run()
}
<commit_msg>Var changes<commit_after>// Copyright 2014 The Transporter Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// seed is a reimagining of the seed mongo to mongo tool

package main

import (
	"fmt"
	"os"
	"time"

	"github.com/compose/transporter/pkg/events"
	"github.com/compose/transporter/pkg/transporter"
)

var (
	sourceUri string     = os.Getenv("SOURCE_MONGO_URL")
	destUri string       = os.Getenv("DESTINATION_MONGO_URL")
	sourceNS string      = os.Getenv("SOURCE_NS")
	destinationNS string = os.Getenv("DEST_NS")
	tail	             = true
	debug       	     = true
)

func main() {
	source :=
		transporter.NewNode("source", "mongo", map[string]interface{}{"uri": sourceUri, "namespace": sourceNS, "tail": tail}).
			Add(transporter.NewNode("out", "mongo", map[string]interface{}{"uri": destUri, "namespace": destinationNS}))

	if debug == true {
		source.Add(transporter.NewNode("out", "file", map[string]interface{}{"uri": "stdout://"}))
	}

	pipeline, err := transporter.NewPipeline(source, events.NewLogEmitter(), 1*time.Second)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	pipeline.Run()
}
<|endoftext|>


<commit_before><commit_msg>trivial fix<commit_after><|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/hanwen/termite/cba"
	"github.com/hanwen/termite/termite"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"os/user"
	"runtime"
	"runtime/pprof"
	"syscall"
)

var _ = log.Printf

func handleStop(daemon *termite.Worker) {
	for {
		sig := <-signal.Incoming
		switch sig.(os.UnixSignal) {
		case syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT:
			log.Println("got signal: ", sig)
			req := termite.ShutdownRequest{}
			rep := termite.ShutdownResponse{}
			daemon.Shutdown(&req, &rep)
		case syscall.SIGHUP:
			daemon.Report()
		}
	}
}

func main() {
	cachedir := flag.String("cachedir", "/var/cache/termite/worker-cache", "content cache")
	tmpdir := flag.String("tmpdir", "/var/tmp",
		"where to create FUSE mounts; should be on same partition as cachedir.")
	secretFile := flag.String("secret", "secret.txt", "file containing password.")
	port := flag.Int("port", 1232, "Start of port to try.")
	portRetry := flag.Int("port-retry", 10, "How many other ports to try.")
	coordinator := flag.String("coordinator", "", "Where to register the worker.")
	jobs := flag.Int("jobs", 1, "Max number of jobs to run.")
	reapcount := flag.Int("reap-count", 1, "Number of jobs per filesystem.")
	userFlag := flag.String("user", "nobody", "Run as this user.")
	memcache := flag.Int("filecache", 1024, "number of <128k files to cache in memory")
	logfile := flag.String("logfile", "", "Output log file to use.")
	paranoia := flag.Bool("paranoia", false, "Check attribute cache.")
	cpus := flag.Int("cpus", 1, "Number of CPUs to use.")
	heap := flag.Int("heap-size", 0, "Maximum heap size in MB.")
	cpuprofile := flag.String("profile", "", "File to write profile output to.")
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if os.Geteuid() != 0 {
		log.Fatal("This program must run as root")
	}
	secret, err := ioutil.ReadFile(*secretFile)
	if err != nil {
		log.Fatal("ReadFile", err)
	}

	if *logfile != "" {
		f, err := os.OpenFile(*logfile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			log.Fatal("Could not open log file.", err)
		}
		log.Println("Log output to", *logfile)
		log.SetOutput(f)
	} else {
		log.SetPrefix("W")
	}

	opts := termite.WorkerOptions{
		Secret:      secret,
		TempDir:     *tmpdir,
		Jobs:        *jobs,
		Paranoia:    *paranoia,
		ReapCount:   *reapcount,
		LogFileName: *logfile,
		ContentCacheOptions: cba.ContentCacheOptions{
			Dir:      *cachedir,
			MemCount: *memcache,
		},
		HeapLimit:   uint64(*heap) * (1 << 20),
		Coordinator: *coordinator,
		Port:        *port,
		PortRetry:   *portRetry,
	}
	if os.Geteuid() == 0 {
		nobody, err := user.Lookup(*userFlag)
		if err != nil {
			log.Fatalf("can't lookup %q: %v", *userFlag, err)
		}
		opts.User = nobody
	}

	daemon := termite.NewWorker(&opts)
	if *cpus > 0 {
		runtime.GOMAXPROCS(*cpus)
	}
	log.Printf("%s on %d CPUs", termite.Version(), runtime.GOMAXPROCS(0))
	go handleStop(daemon)
	daemon.RunWorkerServer()
}
<commit_msg>Make ^C be aggressive for worker.<commit_after>package main

import (
	"flag"
	"github.com/hanwen/termite/cba"
	"github.com/hanwen/termite/termite"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"os/user"
	"runtime"
	"runtime/pprof"
	"syscall"
)

var _ = log.Printf

func handleStop(daemon *termite.Worker) {
	for {
		sig := <-signal.Incoming
		switch sig.(os.UnixSignal) {
		case syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT:
			log.Println("got signal: ", sig)
			req := termite.ShutdownRequest{Kill: true}
			rep := termite.ShutdownResponse{}
			daemon.Shutdown(&req, &rep)
		case syscall.SIGHUP:
			daemon.Report()
		}
	}
}

func main() {
	cachedir := flag.String("cachedir", "/var/cache/termite/worker-cache", "content cache")
	tmpdir := flag.String("tmpdir", "/var/tmp",
		"where to create FUSE mounts; should be on same partition as cachedir.")
	secretFile := flag.String("secret", "secret.txt", "file containing password.")
	port := flag.Int("port", 1232, "Start of port to try.")
	portRetry := flag.Int("port-retry", 10, "How many other ports to try.")
	coordinator := flag.String("coordinator", "", "Where to register the worker.")
	jobs := flag.Int("jobs", 1, "Max number of jobs to run.")
	reapcount := flag.Int("reap-count", 1, "Number of jobs per filesystem.")
	userFlag := flag.String("user", "nobody", "Run as this user.")
	memcache := flag.Int("filecache", 1024, "number of <128k files to cache in memory")
	logfile := flag.String("logfile", "", "Output log file to use.")
	paranoia := flag.Bool("paranoia", false, "Check attribute cache.")
	cpus := flag.Int("cpus", 1, "Number of CPUs to use.")
	heap := flag.Int("heap-size", 0, "Maximum heap size in MB.")
	cpuprofile := flag.String("profile", "", "File to write profile output to.")
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if os.Geteuid() != 0 {
		log.Fatal("This program must run as root")
	}
	secret, err := ioutil.ReadFile(*secretFile)
	if err != nil {
		log.Fatal("ReadFile", err)
	}

	if *logfile != "" {
		f, err := os.OpenFile(*logfile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			log.Fatal("Could not open log file.", err)
		}
		log.Println("Log output to", *logfile)
		log.SetOutput(f)
	} else {
		log.SetPrefix("W")
	}

	opts := termite.WorkerOptions{
		Secret:      secret,
		TempDir:     *tmpdir,
		Jobs:        *jobs,
		Paranoia:    *paranoia,
		ReapCount:   *reapcount,
		LogFileName: *logfile,
		ContentCacheOptions: cba.ContentCacheOptions{
			Dir:      *cachedir,
			MemCount: *memcache,
		},
		HeapLimit:   uint64(*heap) * (1 << 20),
		Coordinator: *coordinator,
		Port:        *port,
		PortRetry:   *portRetry,
	}
	if os.Geteuid() == 0 {
		nobody, err := user.Lookup(*userFlag)
		if err != nil {
			log.Fatalf("can't lookup %q: %v", *userFlag, err)
		}
		opts.User = nobody
	}

	daemon := termite.NewWorker(&opts)
	if *cpus > 0 {
		runtime.GOMAXPROCS(*cpus)
	}
	log.Printf("%s on %d CPUs", termite.Version(), runtime.GOMAXPROCS(0))
	go handleStop(daemon)
	daemon.RunWorkerServer()
}
<|endoftext|>


<commit_before><commit_msg>remove public_key com deployment<commit_after><|endoftext|>


<commit_before>// Copyright 2020 CoreOS, Inc.
//
// 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 files

import (
	"github.com/coreos/ignition/v2/tests/register"
	"github.com/coreos/ignition/v2/tests/types"
)

func init() {
	register.Register(register.PositiveTest, CreateInstantiatedService())
	register.Register(register.PositiveTest, CreateEmptyDropinsUnit())
}

func CreateInstantiatedService() types.Test {
	name := "instantiated.unit.create"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": { "version": "$version" },
		"systemd": {
			"units": [
			  {
				"name": "echo@.service",
				"contents": "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/bin/echo %i\nRemainAfterExit=yes\n[Install]\nWantedBy=multi-user.target\n"
			  },
			  {
				"enabled": true,
				"name": "echo@bar.service"
			  },
			  {
				"enabled": true,
				"name": "echo@foo.service"
			  }
			]
		  }
		}`
	configMinVersion := "3.0.0"
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "echo@.service",
				Directory: "etc/systemd/system",
			},
			Contents: "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/bin/echo %i\nRemainAfterExit=yes\n[Install]\nWantedBy=multi-user.target\n",
		},
		{
			Node: types.Node{
				Name:      "20-ignition.preset",
				Directory: "etc/systemd/system-preset",
			},
			Contents: "enable echo@.service bar foo\n",
		},
	})

	return types.Test{
		Name:             name,
		In:               in,
		Out:              out,
		Config:           config,
		ConfigMinVersion: configMinVersion,
	}
}

// CreateEmptyDropinsUnit writes an empty dropin to the disk.
func CreateEmptyDropinsUnit() types.Test {
	name := "unit.create.emptydropin"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": { "version": "$version" },
		"systemd": {
			"units": [
			  {
				"name": "zincati.service",
				"dropins": [
					{
						"name": "empty.conf",
						"contents": ""
					}
				]
			  }
			]
		  }
		}`
	configMinVersion := "3.0.0"
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "empty.conf",
				Directory: "etc/systemd/system/zincati.service.d",
			},
			Contents: "",
		},
	})

	return types.Test{
		Name:             name,
		In:               in,
		Out:              out,
		Config:           config,
		ConfigMinVersion: configMinVersion,
	}
}
<commit_msg>tests/positive: Add tests for masking and unmasking units<commit_after>// Copyright 2020 CoreOS, Inc.
//
// 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 files

import (
	"github.com/coreos/ignition/v2/tests/register"
	"github.com/coreos/ignition/v2/tests/types"
)

func init() {
	register.Register(register.PositiveTest, CreateInstantiatedService())
	register.Register(register.PositiveTest, CreateEmptyDropinsUnit())
	register.Register(register.PositiveTest, TestUnmaskUnit())
	register.Register(register.PositiveTest, TestMaskUnit())
}

func CreateInstantiatedService() types.Test {
	name := "instantiated.unit.create"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": { "version": "$version" },
		"systemd": {
			"units": [
			  {
				"name": "echo@.service",
				"contents": "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/bin/echo %i\nRemainAfterExit=yes\n[Install]\nWantedBy=multi-user.target\n"
			  },
			  {
				"enabled": true,
				"name": "echo@bar.service"
			  },
			  {
				"enabled": true,
				"name": "echo@foo.service"
			  }
			]
		  }
		}`
	configMinVersion := "3.0.0"
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "echo@.service",
				Directory: "etc/systemd/system",
			},
			Contents: "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/bin/echo %i\nRemainAfterExit=yes\n[Install]\nWantedBy=multi-user.target\n",
		},
		{
			Node: types.Node{
				Name:      "20-ignition.preset",
				Directory: "etc/systemd/system-preset",
			},
			Contents: "enable echo@.service bar foo\n",
		},
	})

	return types.Test{
		Name:             name,
		In:               in,
		Out:              out,
		Config:           config,
		ConfigMinVersion: configMinVersion,
	}
}

// CreateEmptyDropinsUnit writes an empty dropin to the disk.
func CreateEmptyDropinsUnit() types.Test {
	name := "unit.create.emptydropin"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": { "version": "$version" },
		"systemd": {
			"units": [
			  {
				"name": "zincati.service",
				"dropins": [
					{
						"name": "empty.conf",
						"contents": ""
					}
				]
			  }
			]
		  }
		}`
	configMinVersion := "3.0.0"
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "empty.conf",
				Directory: "etc/systemd/system/zincati.service.d",
			},
			Contents: "",
		},
	})

	return types.Test{
		Name:             name,
		In:               in,
		Out:              out,
		Config:           config,
		ConfigMinVersion: configMinVersion,
	}
}

// TestUnmaskUnit tests unmasking systemd units. In the case that a
// systemd unit is masked already it will unmask it. In the case it is
// not masked it will do nothing. In the code below we have foo.service,
// which is masked in the input, and bar.service, which is not masked in
// the input. We test that foo.service is properly unmasked and that the
// unmasking of bar.service is skipped (because it wasn't masked to begin
// with) and that bar.service gets enabled.
func TestUnmaskUnit() types.Test {
	name := "unit.unmask"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": { "version": "$version" },
		"systemd": {
			"units": [
				{
					"name": "foo.service",
					"mask": false
				},
				{
					"name": "bar.service",
					"enabled": true,
					"mask": false
				}
			]
        }
		}`
	configMinVersion := "3.0.0"
	in[0].Partitions.AddLinks("ROOT", []types.Link{
		{
			Node: types.Node{
				Directory: "etc/systemd/system",
				Name:      "foo.service",
			},
			Target: "/dev/null",
			Hard:   false,
		},
	})
	in[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "foo.service",
				Directory: "usr/lib/systemd/system",
			},
			Contents: "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/usr/bin/true\n[Install]\nWantedBy=multi-user.target\n",
		},
		{
			Node: types.Node{
				Name:      "bar.service",
				Directory: "usr/lib/systemd/system",
			},
			Contents: "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/usr/bin/true\n[Install]\nWantedBy=multi-user.target\n",
		},
	})
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "20-ignition.preset",
				Directory: "etc/systemd/system-preset",
			},
			Contents: "enable bar.service\n",
		},
	})
	out[0].Partitions.AddRemovedNodes("ROOT", []types.Node{
		{
			Directory: "etc/systemd/system",
			Name:      "foo.service",
		},
		{
			Directory: "etc/systemd/system",
			Name:      "bar.service",
		},
	})

	return types.Test{
		Name:             name,
		In:               in,
		Out:              out,
		Config:           config,
		ConfigMinVersion: configMinVersion,
	}
}

// TestMaskUnit tests masking systemd units. In this case foo.service
// verified to be masked by checking it's symlinked to /dev/null.
// We'll also mask a non-existent bar.service.
func TestMaskUnit() types.Test {
	name := "unit.mask"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": { "version": "$version" },
		"systemd": {
			"units": [
				{
					"name": "foo.service",
					"mask": true
				},
				{
					"name": "bar.service",
					"mask": true
				}
			]
        }
		}`
	configMinVersion := "3.0.0"
	in[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "foo.service",
				Directory: "usr/lib/systemd/system",
			},
			Contents: "[Unit]\nDescription=f\n[Service]\nType=oneshot\nExecStart=/usr/bin/true\n[Install]\nWantedBy=multi-user.target\n",
		},
	})
	out[0].Partitions.AddLinks("ROOT", []types.Link{
		{
			Node: types.Node{
				Directory: "etc/systemd/system",
				Name:      "foo.service",
			},
			Target: "/dev/null",
			Hard:   false,
		},
		{
			Node: types.Node{
				Directory: "etc/systemd/system",
				Name:      "bar.service",
			},
			Target: "/dev/null",
			Hard:   false,
		},
	})

	return types.Test{
		Name:             name,
		In:               in,
		Out:              out,
		Config:           config,
		ConfigMinVersion: configMinVersion,
	}
}
<|endoftext|>


<commit_before><commit_msg>rewrite<commit_after><|endoftext|>


<commit_before>package integration_test

import (
	"fmt"
	"strings"
	"time"

	"github.com/stretchr/testify/suite"

	"github.com/telepresenceio/telepresence/v2/integration_test/itest"
)

type helmSuite struct {
	itest.Suite
	itest.HelmAndService
	mgrSpace2 string
	appSpace2 string
}

func init() {
	itest.AddHelmAndServiceSuite("-1", "echo", func(h itest.HelmAndService) suite.TestingSuite {
		s := &helmSuite{Suite: itest.Suite{Harness: h}, HelmAndService: h}
		suffix := itest.GetGlobalHarness(h.HarnessContext()).Suffix()
		s.appSpace2, s.mgrSpace2 = itest.AppAndMgrNSName(suffix + "-2")
		return s
	})
}

func (s *helmSuite) SetupSuite() {
	s.Suite.SetupSuite()
	ctx := s.Context()
	itest.CreateNamespaces(ctx, s.appSpace2, s.mgrSpace2)
	itest.ApplyEchoService(ctx, s.ServiceName(), s.appSpace2, 80)
}

func (s *helmSuite) TearDownSuite() {
	itest.DeleteNamespaces(s.Context(), s.appSpace2, s.mgrSpace2)
}

func (s *helmSuite) Test_HelmCanInterceptInManagedNamespace() {
	ctx := s.Context()
	defer itest.TelepresenceOk(ctx, "leave", s.ServiceName()+"-"+s.AppNamespace())

	stdout := itest.TelepresenceOk(ctx, "intercept", "--namespace", s.AppNamespace(), "--mount", "false", s.ServiceName(), "--port", "9090")
	s.Contains(stdout, "Using Deployment "+s.ServiceName())
	stdout = itest.TelepresenceOk(ctx, "list", "--namespace", s.AppNamespace(), "--intercepts")
	s.Contains(stdout, s.ServiceName()+": intercepted")
}

func (s *helmSuite) Test_HelmCannotInterceptInUnmanagedNamespace() {
	ctx := s.Context()
	_, stderr, err := itest.Telepresence(ctx, "intercept", "--namespace", s.appSpace2, "--mount", "false", s.ServiceName(), "--port", "9090")
	s.Error(err)
	s.Contains(stderr, "Failed to establish intercept")
}

func (s *helmSuite) Test_HelmWebhookInjectsInManagedNamespace() {
	ctx := s.Context()
	s.ApplyApp(ctx, "echo-auto-inject", "deploy/echo-auto-inject")
	defer s.DeleteSvcAndWorkload(ctx, "deploy", "echo-auto-inject")

	s.Eventually(func() bool {
		stdout, _, err := itest.Telepresence(ctx, "list", "--namespace", s.AppNamespace(), "--agents")
		return err == nil && strings.Contains(stdout, "echo-auto-inject: ready to intercept (traffic-agent already installed)")
	},
		20*time.Second, // waitFor
		2*time.Second,  // polling interval
	)
}

func (s *helmSuite) Test_HelmWebhookDoesntInjectInUnmanagedNamespace() {
	ctx := s.Context()
	itest.ApplyApp(ctx, "echo-auto-inject", s.appSpace2, "deploy/echo-auto-inject")
	defer itest.DeleteSvcAndWorkload(ctx, "deploy", "echo-auto-inject", s.appSpace2)

	s.Never(func() bool {
		stdout, _, err := itest.Telepresence(ctx, "list", "--namespace", s.appSpace2, "--agents")
		return err == nil && strings.Contains(stdout, "echo-auto-inject: ready to intercept (traffic-agent already installed)")
	},
		10*time.Second, // waitFor
		2*time.Second,  // polling interval
	)
}

func (s *helmSuite) Test_HelmMultipleInstalls() {
	ctx := s.Context()
	svc := s.ServiceName()
	itest.TelepresenceQuitOk(ctx)

	ctx = itest.WithEnv(ctx, map[string]string{"TELEPRESENCE_MANAGER_NAMESPACE": s.mgrSpace2})
	s.Run("Installs Successfully", func() {
		s.NoError(s.InstallTrafficManager(ctx, s.mgrSpace2, s.appSpace2))
	})

	s.Run("Can be connected to", func() {
		stdout := itest.TelepresenceOk(ctx, "connect")
		s.Contains(stdout, "Connected to context")
		itest.TelepresenceOk(ctx, "loglevel", "trace")
		s.Eventually(func() bool {
			return itest.Run(ctx, "curl", "--silent", "--connect-timeout", "1", fmt.Sprintf("%s.%s", svc, s.appSpace2)) == nil
		}, 7*time.Second, 1*time.Second)
	})

	s.Run("Can intercept", func() {
		defer itest.TelepresenceOk(ctx, "leave", fmt.Sprintf("%s-%s", svc, s.appSpace2))

		stdout := itest.TelepresenceOk(ctx, "intercept", "--namespace", s.appSpace2, "--mount", "false", svc, "--port", "9090")
		s.Contains(stdout, "Using Deployment "+svc)
		stdout = itest.TelepresenceOk(ctx, "list", "--namespace", s.appSpace2, "--intercepts")
		s.Contains(stdout, svc+": intercepted")
	})

	s.Run("Uninstalls Successfully", func() {
		s.UninstallTrafficManager(ctx, s.mgrSpace2)
	})
}

func (s *helmSuite) Test_CollidingInstalls() {
	s.Error(s.InstallTrafficManager(s.Context(), s.mgrSpace2, s.AppNamespace(), s.appSpace2))
}
<commit_msg>Fix helm intercept test that fails with other error now.<commit_after>package integration_test

import (
	"fmt"
	"strings"
	"time"

	"github.com/stretchr/testify/suite"

	"github.com/telepresenceio/telepresence/v2/integration_test/itest"
)

type helmSuite struct {
	itest.Suite
	itest.HelmAndService
	mgrSpace2 string
	appSpace2 string
}

func init() {
	itest.AddHelmAndServiceSuite("-1", "echo", func(h itest.HelmAndService) suite.TestingSuite {
		s := &helmSuite{Suite: itest.Suite{Harness: h}, HelmAndService: h}
		suffix := itest.GetGlobalHarness(h.HarnessContext()).Suffix()
		s.appSpace2, s.mgrSpace2 = itest.AppAndMgrNSName(suffix + "-2")
		return s
	})
}

func (s *helmSuite) SetupSuite() {
	s.Suite.SetupSuite()
	ctx := s.Context()
	itest.CreateNamespaces(ctx, s.appSpace2, s.mgrSpace2)
	itest.ApplyEchoService(ctx, s.ServiceName(), s.appSpace2, 80)
}

func (s *helmSuite) TearDownSuite() {
	itest.DeleteNamespaces(s.Context(), s.appSpace2, s.mgrSpace2)
}

func (s *helmSuite) Test_HelmCanInterceptInManagedNamespace() {
	ctx := s.Context()
	defer itest.TelepresenceOk(ctx, "leave", s.ServiceName()+"-"+s.AppNamespace())

	stdout := itest.TelepresenceOk(ctx, "intercept", "--namespace", s.AppNamespace(), "--mount", "false", s.ServiceName(), "--port", "9090")
	s.Contains(stdout, "Using Deployment "+s.ServiceName())
	stdout = itest.TelepresenceOk(ctx, "list", "--namespace", s.AppNamespace(), "--intercepts")
	s.Contains(stdout, s.ServiceName()+": intercepted")
}

func (s *helmSuite) Test_HelmCannotInterceptInUnmanagedNamespace() {
	ctx := s.Context()
	_, stderr, err := itest.Telepresence(ctx, "intercept", "--namespace", s.appSpace2, "--mount", "false", s.ServiceName(), "--port", "9090")
	s.Error(err)
	s.Contains(stderr, `deployments.apps "echo" is forbidden`)
}

func (s *helmSuite) Test_HelmWebhookInjectsInManagedNamespace() {
	ctx := s.Context()
	s.ApplyApp(ctx, "echo-auto-inject", "deploy/echo-auto-inject")
	defer s.DeleteSvcAndWorkload(ctx, "deploy", "echo-auto-inject")

	s.Eventually(func() bool {
		stdout, _, err := itest.Telepresence(ctx, "list", "--namespace", s.AppNamespace(), "--agents")
		return err == nil && strings.Contains(stdout, "echo-auto-inject: ready to intercept (traffic-agent already installed)")
	},
		20*time.Second, // waitFor
		2*time.Second,  // polling interval
	)
}

func (s *helmSuite) Test_HelmWebhookDoesntInjectInUnmanagedNamespace() {
	ctx := s.Context()
	itest.ApplyApp(ctx, "echo-auto-inject", s.appSpace2, "deploy/echo-auto-inject")
	defer itest.DeleteSvcAndWorkload(ctx, "deploy", "echo-auto-inject", s.appSpace2)

	s.Never(func() bool {
		stdout, _, err := itest.Telepresence(ctx, "list", "--namespace", s.appSpace2, "--agents")
		return err == nil && strings.Contains(stdout, "echo-auto-inject: ready to intercept (traffic-agent already installed)")
	},
		10*time.Second, // waitFor
		2*time.Second,  // polling interval
	)
}

func (s *helmSuite) Test_HelmMultipleInstalls() {
	ctx := s.Context()
	svc := s.ServiceName()
	itest.TelepresenceQuitOk(ctx)

	ctx = itest.WithEnv(ctx, map[string]string{"TELEPRESENCE_MANAGER_NAMESPACE": s.mgrSpace2})
	s.Run("Installs Successfully", func() {
		s.NoError(s.InstallTrafficManager(ctx, s.mgrSpace2, s.appSpace2))
	})

	s.Run("Can be connected to", func() {
		stdout := itest.TelepresenceOk(ctx, "connect")
		s.Contains(stdout, "Connected to context")
		itest.TelepresenceOk(ctx, "loglevel", "trace")
		s.Eventually(func() bool {
			return itest.Run(ctx, "curl", "--silent", "--connect-timeout", "1", fmt.Sprintf("%s.%s", svc, s.appSpace2)) == nil
		}, 7*time.Second, 1*time.Second)
	})

	s.Run("Can intercept", func() {
		defer itest.TelepresenceOk(ctx, "leave", fmt.Sprintf("%s-%s", svc, s.appSpace2))

		stdout := itest.TelepresenceOk(ctx, "intercept", "--namespace", s.appSpace2, "--mount", "false", svc, "--port", "9090")
		s.Contains(stdout, "Using Deployment "+svc)
		stdout = itest.TelepresenceOk(ctx, "list", "--namespace", s.appSpace2, "--intercepts")
		s.Contains(stdout, svc+": intercepted")
	})

	s.Run("Uninstalls Successfully", func() {
		s.UninstallTrafficManager(ctx, s.mgrSpace2)
	})
}

func (s *helmSuite) Test_CollidingInstalls() {
	s.Error(s.InstallTrafficManager(s.Context(), s.mgrSpace2, s.AppNamespace(), s.appSpace2))
}
<|endoftext|>


<commit_before>package agent

import (
	"log"
	"math"
	"time"

	"github.com/influxdata/telegraf"
)

func NewAccumulator(
	metrics chan telegraf.Metric,
) *accumulator {
	acc := accumulator{}
	acc.metrics = metrics
	return &acc
}

type accumulator struct {
	metrics chan telegraf.Metric

	debug bool
}

// Add of telegraf.Accumulator interface.
func (ac *accumulator) Add(measurement string, value interface{}, tags map[string]string, t ...time.Time) {
	ac.AddFields(measurement, map[string]interface{}{"value": value}, tags, t...)
}

// AddFields of telegraf.Accumulator interface.
func (ac *accumulator) AddFields(measurement string, fields map[string]interface{}, tags map[string]string, t ...time.Time) {
	if measurement == "" || len(fields) == 0 {
		return
	}
	if tags == nil {
		tags = make(map[string]string)
	}
	values := make(map[string]interface{})
	for k, v := range fields {

		// Validate uint64 and float64 fields
		switch val := v.(type) {
		case uint64:
			// InfluxDB does not support writing uint64
			if val < uint64(9223372036854775808) {
				values[k] = int64(val)
			} else {
				values[k] = int64(9223372036854775807)
			}
			continue
		case float64:
			// NaNs are invalid values in influxdb, skip measurement
			if math.IsNaN(val) || math.IsInf(val, 0) {
				if false { // if ac.debug TODO
					log.Printf(
						"Measurement [%s] field [%s] has a NaN or Inf field, skipping",
						measurement, k)
				}
				continue
			}
		}

		values[k] = v
	}
	if len(values) == 0 {
		return
	}

	var ts time.Time
	if len(t) > 0 {
		ts = t[0]
	} else {
		ts = time.Now()
	}
	// timestamp = timestamp.Round(ac.precision) // TODO

	m, err := telegraf.NewMetric(measurement, tags, values, ts)
	if err != nil {
		log.Printf("Error adding point [%s]: %s\n", measurement, err.Error())
		return
	}
	// if ac.trace { fmt.Println("> " + m.String()) } // TODO
	ac.metrics <- m
}

func (ac accumulator) Debug() bool {
	return ac.debug
}

func (ac *accumulator) SetDebug(on bool) {
	ac.debug = on
}

func (ac *accumulator) SetPrecision(precision, interval time.Duration) {
	// TODO
}

func (ac *accumulator) DisablePrecision() {
	// TODO
}

// func (ac *accumulator) AddError(err error) { log.Printf("Error in input: %s", err) }
<commit_msg>fix compliance with telegraf.Accumulator interface<commit_after>package agent

import (
	"log"
	"math"
	"time"

	"github.com/influxdata/telegraf"
)

func NewAccumulator(
	metrics chan telegraf.Metric,
) *accumulator {
	acc := accumulator{}
	acc.metrics = metrics
	return &acc
}

type accumulator struct {
	metrics chan telegraf.Metric

	debug bool
}

// Add of telegraf.Accumulator interface.
func (ac *accumulator) Add(measurement string, value interface{}, tags map[string]string, t ...time.Time) {
	ac.AddFields(measurement, map[string]interface{}{"value": value}, tags, t...)
}

// AddFields of telegraf.Accumulator interface.
func (ac *accumulator) AddFields(measurement string, fields map[string]interface{}, tags map[string]string, t ...time.Time) {
	if measurement == "" || len(fields) == 0 {
		return
	}
	if tags == nil {
		tags = make(map[string]string)
	}
	values := make(map[string]interface{})
	for k, v := range fields {

		// Validate uint64 and float64 fields
		switch val := v.(type) {
		case uint64:
			// InfluxDB does not support writing uint64
			if val < uint64(9223372036854775808) {
				values[k] = int64(val)
			} else {
				values[k] = int64(9223372036854775807)
			}
			continue
		case float64:
			// NaNs are invalid values in influxdb, skip measurement
			if math.IsNaN(val) || math.IsInf(val, 0) {
				if false { // if ac.debug TODO
					log.Printf(
						"Measurement [%s] field [%s] has a NaN or Inf field, skipping",
						measurement, k)
				}
				continue
			}
		}

		values[k] = v
	}
	if len(values) == 0 {
		return
	}

	var ts time.Time
	if len(t) > 0 {
		ts = t[0]
	} else {
		ts = time.Now()
	}
	// timestamp = timestamp.Round(ac.precision) // TODO

	m, err := telegraf.NewMetric(measurement, tags, values, ts)
	if err != nil {
		log.Printf("Error adding point [%s]: %s\n", measurement, err.Error())
		return
	}
	// if ac.trace { fmt.Println("> " + m.String()) } // TODO
	ac.metrics <- m
}

func (ac accumulator) Debug() bool {
	return ac.debug
}

func (ac *accumulator) SetDebug(on bool) {
	ac.debug = on
}

func (ac *accumulator) SetPrecision(precision, interval time.Duration) {
	// TODO
}

func (ac *accumulator) DisablePrecision() {
	// TODO
}

// TODO AddError is not part of telegraf.Accumulator interface as of 1.0.0-beta3
func (ac *accumulator) AddError(err error) { log.Printf("Error in input: %s", err) }
<|endoftext|>


<commit_before>package handler

import (
	"context"
	"errors"
	"os"
	"testing"

	"github.com/stretchr/testify/require"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	pb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
)

func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn) (int32, error) {
	return func(ctx context.Context, client *grpc.ClientConn) (int32, error) {
		require.NotNil(t, ctx)
		require.NotNil(t, client)

		return 0, err
	}
}

func TestRunGitalyCommand(t *testing.T) {
	cmd := GitalyCommand{
		Config:  &config.Config{},
		Address: "tcp://localhost:9999",
	}

	err := cmd.RunGitalyCommand(makeHandler(t, nil))
	require.NoError(t, err)

	expectedErr := errors.New("error")
	err = cmd.RunGitalyCommand(makeHandler(t, expectedErr))
	require.Equal(t, err, expectedErr)
}

func TestMissingGitalyAddress(t *testing.T) {
	cmd := GitalyCommand{Config: &config.Config{}}

	err := cmd.RunGitalyCommand(makeHandler(t, nil))
	require.EqualError(t, err, "no gitaly_address given")
}

func TestGetConnMetadata(t *testing.T) {
	tests := []struct {
		name string
		gc   *GitalyCommand
		want map[string]string
	}{
		{
			name: "gitaly_feature_flags",
			gc: &GitalyCommand{
				Config:  &config.Config{},
				Address: "tcp://localhost:9999",
				Features: map[string]string{
					"gitaly-feature-cache_invalidator":        "true",
					"other-ff":                                "true",
					"gitaly-feature-inforef_uploadpack_cache": "false",
				},
			},
			want: map[string]string{
				"gitaly-feature-cache_invalidator":        "true",
				"gitaly-feature-inforef_uploadpack_cache": "false",
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			conn, err := getConn(tt.gc)
			require.NoError(t, err)

			md, exists := metadata.FromOutgoingContext(conn.ctx)
			require.True(t, exists)
			require.Equal(t, len(tt.want), md.Len())

			for k, v := range tt.want {
				values := md.Get(k)
				require.Equal(t, 1, len(values))
				require.Equal(t, v, values[0])
			}

		})
	}
}

func TestPrepareContext(t *testing.T) {
	tests := []struct {
		name             string
		gc               *GitalyCommand
		sshConnectionEnv string
		repo             *pb.Repository
		response         *accessverifier.Response
		want             map[string]string
	}{
		{
			name: "client_identity",
			gc: &GitalyCommand{
				Config:  &config.Config{},
				Address: "tcp://localhost:9999",
			},
			sshConnectionEnv: "10.0.0.1 1234 127.0.0.1 5678",
			repo: &pb.Repository{
				StorageName:                   "default",
				RelativePath:                  "@hashed/5f/9c/5f9c4ab08cac7457e9111a30e4664920607ea2c115a1433d7be98e97e64244ca.git",
				GitObjectDirectory:            "path/to/git_object_directory",
				GitAlternateObjectDirectories: []string{"path/to/git_alternate_object_directory"},
				GlRepository:                  "project-26",
				GlProjectPath:                 "group/private",
			},
			response: &accessverifier.Response{
				UserId:   "6",
				Username: "jane.doe",
			},
			want: map[string]string{
				"remote_ip": "10.0.0.1",
				"user_id":   "6",
				"username":  "jane.doe",
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()

			orig := os.Getenv("SSH_CONNECTION")
			defer os.Setenv("SSH_CONNECTION", orig)
			os.Setenv("SSH_CONNECTION", tt.sshConnectionEnv)

			ctx, cancel := tt.gc.PrepareContext(ctx, tt.repo, tt.response, "protocol")
			defer cancel()

			md, exists := metadata.FromOutgoingContext(ctx)
			require.True(t, exists)
			require.Equal(t, len(tt.want), md.Len())

			for k, v := range tt.want {
				values := md.Get(k)
				require.Equal(t, 1, len(values))
				require.Equal(t, v, values[0])
			}

		})
	}
}
<commit_msg>use testhelper for SSH_CONNECTION<commit_after>package handler

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/require"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	pb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
)

func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn) (int32, error) {
	return func(ctx context.Context, client *grpc.ClientConn) (int32, error) {
		require.NotNil(t, ctx)
		require.NotNil(t, client)

		return 0, err
	}
}

func TestRunGitalyCommand(t *testing.T) {
	cmd := GitalyCommand{
		Config:  &config.Config{},
		Address: "tcp://localhost:9999",
	}

	err := cmd.RunGitalyCommand(makeHandler(t, nil))
	require.NoError(t, err)

	expectedErr := errors.New("error")
	err = cmd.RunGitalyCommand(makeHandler(t, expectedErr))
	require.Equal(t, err, expectedErr)
}

func TestMissingGitalyAddress(t *testing.T) {
	cmd := GitalyCommand{Config: &config.Config{}}

	err := cmd.RunGitalyCommand(makeHandler(t, nil))
	require.EqualError(t, err, "no gitaly_address given")
}

func TestGetConnMetadata(t *testing.T) {
	tests := []struct {
		name string
		gc   *GitalyCommand
		want map[string]string
	}{
		{
			name: "gitaly_feature_flags",
			gc: &GitalyCommand{
				Config:  &config.Config{},
				Address: "tcp://localhost:9999",
				Features: map[string]string{
					"gitaly-feature-cache_invalidator":        "true",
					"other-ff":                                "true",
					"gitaly-feature-inforef_uploadpack_cache": "false",
				},
			},
			want: map[string]string{
				"gitaly-feature-cache_invalidator":        "true",
				"gitaly-feature-inforef_uploadpack_cache": "false",
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			conn, err := getConn(tt.gc)
			require.NoError(t, err)

			md, exists := metadata.FromOutgoingContext(conn.ctx)
			require.True(t, exists)
			require.Equal(t, len(tt.want), md.Len())

			for k, v := range tt.want {
				values := md.Get(k)
				require.Equal(t, 1, len(values))
				require.Equal(t, v, values[0])
			}

		})
	}
}

func TestPrepareContext(t *testing.T) {
	tests := []struct {
		name             string
		gc               *GitalyCommand
		sshConnectionEnv string
		repo             *pb.Repository
		response         *accessverifier.Response
		want             map[string]string
	}{
		{
			name: "client_identity",
			gc: &GitalyCommand{
				Config:  &config.Config{},
				Address: "tcp://localhost:9999",
			},
			sshConnectionEnv: "10.0.0.1 1234 127.0.0.1 5678",
			repo: &pb.Repository{
				StorageName:                   "default",
				RelativePath:                  "@hashed/5f/9c/5f9c4ab08cac7457e9111a30e4664920607ea2c115a1433d7be98e97e64244ca.git",
				GitObjectDirectory:            "path/to/git_object_directory",
				GitAlternateObjectDirectories: []string{"path/to/git_alternate_object_directory"},
				GlRepository:                  "project-26",
				GlProjectPath:                 "group/private",
			},
			response: &accessverifier.Response{
				UserId:   "6",
				Username: "jane.doe",
			},
			want: map[string]string{
				"remote_ip": "10.0.0.1",
				"user_id":   "6",
				"username":  "jane.doe",
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cleanup, err := testhelper.Setenv("SSH_CONNECTION", tt.sshConnectionEnv)
			require.NoError(t, err)
			defer cleanup()

			ctx := context.Background()

			ctx, cancel := tt.gc.PrepareContext(ctx, tt.repo, tt.response, "protocol")
			defer cancel()

			md, exists := metadata.FromOutgoingContext(ctx)
			require.True(t, exists)
			require.Equal(t, len(tt.want), md.Len())

			for k, v := range tt.want {
				values := md.Get(k)
				require.Equal(t, 1, len(values))
				require.Equal(t, v, values[0])
			}

		})
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Hajime Hoshi
//
// 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.

// +build js

package opengl

import (
	"errors"
	"fmt"
	"github.com/gopherjs/gopherjs/js"
	"github.com/gopherjs/webgl"
)

type Texture js.Object
type Framebuffer js.Object
type Shader js.Object
type Program js.Object

type context struct {
	gl *webgl.Context
}

func NewContext(gl *webgl.Context) *Context {
	c := &Context{
		Nearest:            FilterType(gl.NEAREST),
		Linear:             FilterType(gl.LINEAR),
		VertexShader:       ShaderType(gl.VERTEX_SHADER),
		FragmentShader:     ShaderType(gl.FRAGMENT_SHADER),
		ArrayBuffer:        BufferType(gl.ARRAY_BUFFER),
		ElementArrayBuffer: BufferType(gl.ELEMENT_ARRAY_BUFFER),
		DynamicDraw:        BufferUsageType(gl.DYNAMIC_DRAW),
		StaticDraw:         BufferUsageType(gl.STATIC_DRAW),
	}
	c.gl = gl
	c.init()
	return c
}

func (c *Context) init() {
	gl := c.gl
	// Textures' pixel formats are alpha premultiplied.
	gl.Enable(gl.BLEND)
	gl.BlendFunc(gl.ONE, gl.ONE_MINUS_SRC_ALPHA)
}

func (c *Context) NewTexture(width, height int, pixels []uint8, filter FilterType) (Texture, error) {
	gl := c.gl
	t := gl.CreateTexture()
	if t.IsNull() {
		return nil, errors.New("glGenTexture failed")
	}
	gl.PixelStorei(gl.UNPACK_ALIGNMENT, 4)
	gl.BindTexture(gl.TEXTURE_2D, t)

	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, int(filter))
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, int(filter))

	// void texImage2D(GLenum target, GLint level, GLenum internalformat,
	//     GLsizei width, GLsizei height, GLint border, GLenum format,
	//     GLenum type, ArrayBufferView? pixels);
	var p interface{}
	if pixels != nil {
		p = pixels
	}
	gl.Call("texImage2D", gl.TEXTURE_2D, 0, gl.RGBA, width, height, 0, gl.RGBA, gl.UNSIGNED_BYTE, p)

	return Texture(t), nil
}

func (c *Context) TexturePixels(t Texture, width, height int) ([]uint8, error) {
	gl := c.gl
	// TODO: Use glGetTexLevelParameteri and GL_TEXTURE_WIDTH?
	pixels := js.Global.Get("Uint8Array").New(4 * width * height)
	gl.BindTexture(gl.TEXTURE_2D, t)
	gl.ReadPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, pixels)
	if e := gl.GetError(); e != gl.NO_ERROR {
		// TODO: Use glu.ErrorString
		return nil, errors.New(fmt.Sprintf("gl error: %d", e))
	}
	return pixels.Interface().([]uint8), nil
}

func (c *Context) BindTexture(t Texture) {
	gl := c.gl
	gl.BindTexture(gl.TEXTURE_2D, t)
}

func (c *Context) DeleteTexture(t Texture) {
	gl := c.gl
	gl.DeleteTexture(t)
}

func (c *Context) GlslHighpSupported() bool {
	gl := c.gl
	return gl.Call("getShaderPrecisionFormat", gl.FRAGMENT_SHADER, gl.HIGH_FLOAT).Get("precision").Int() != 0
}

func (c *Context) NewFramebuffer(texture Texture) (Framebuffer, error) {
	gl := c.gl
	f := gl.CreateFramebuffer()
	gl.BindFramebuffer(gl.FRAMEBUFFER, f)

	gl.FramebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0)
	if gl.CheckFramebufferStatus(gl.FRAMEBUFFER) != gl.FRAMEBUFFER_COMPLETE {
		return nil, errors.New("creating framebuffer failed")
	}

	return Framebuffer(f), nil
}

func (c *Context) SetViewport(f Framebuffer, width, height int) error {
	gl := c.gl
	gl.Flush()
	gl.BindFramebuffer(gl.FRAMEBUFFER, f)
	err := gl.CheckFramebufferStatus(gl.FRAMEBUFFER)
	if err != gl.FRAMEBUFFER_COMPLETE {
		if gl.GetError() != 0 {
			return errors.New(fmt.Sprintf("glBindFramebuffer failed: %d", gl.GetError()))
		}
		return errors.New("glBindFramebuffer failed: the context is different?")
	}
	gl.Viewport(0, 0, width, height)
	return nil
}

func (c *Context) FillFramebuffer(f Framebuffer, r, g, b, a float64) error {
	gl := c.gl
	gl.ClearColor(float32(r), float32(g), float32(b), float32(a))
	gl.Clear(gl.COLOR_BUFFER_BIT)
	return nil
}

func (c *Context) DeleteFramebuffer(f Framebuffer) {
	gl := c.gl
	gl.DeleteFramebuffer(f)
}

func (c *Context) NewShader(shaderType ShaderType, source string) (Shader, error) {
	gl := c.gl
	s := gl.CreateShader(int(shaderType))
	if s.IsNull() {
		println(gl.GetError())
		return nil, errors.New("glCreateShader failed")
	}

	gl.ShaderSource(s, source)
	gl.CompileShader(s)

	if !gl.GetShaderParameterb(s, gl.COMPILE_STATUS) {
		log := gl.GetShaderInfoLog(s)
		return nil, errors.New(fmt.Sprintf("shader compile failed: %s", log))
	}
	return Shader(s), nil
}

func (c *Context) DeleteShader(s Shader) {
	gl := c.gl
	gl.DeleteShader(s)
}

func (c *Context) NewProgram(shaders []Shader) (Program, error) {
	gl := c.gl
	p := gl.CreateProgram()
	if p.IsNull() {
		return nil, errors.New("glCreateProgram failed")
	}

	for _, shader := range shaders {
		gl.AttachShader(p, shader)
	}
	gl.LinkProgram(p)
	if !gl.GetProgramParameterb(p, gl.LINK_STATUS) {
		return nil, errors.New("program error")
	}
	return Program(p), nil
}

func (c *Context) UseProgram(p Program) {
	gl := c.gl
	gl.UseProgram(p)
}

func (c *Context) Uniform(p Program, location string, v interface{}) {
	gl := c.gl
	l := gl.GetUniformLocation(p, location)
	switch v := v.(type) {
	case int:
		gl.Uniform1i(l, v)
	case [4]float32:
		gl.Call("uniform4fv", l, v[:])
	case [16]float32:
		gl.UniformMatrix4fv(l, false, v[:])
	}
}

func (c *Context) VertexAttribPointer(p Program, location string, stride int, v uintptr) {
	gl := c.gl
	l := gl.GetAttribLocation(p, location)
	gl.VertexAttribPointer(l, 2, gl.FLOAT, false, stride, int(v))
}

func (c *Context) EnableVertexAttribArray(p Program, location string) {
	gl := c.gl
	l := gl.GetAttribLocation(p, location)
	gl.EnableVertexAttribArray(l)
}

func (c *Context) DisableVertexAttribArray(p Program, location string) {
	gl := c.gl
	l := gl.GetAttribLocation(p, location)
	gl.DisableVertexAttribArray(l)
}

func (c *Context) NewBuffer(bufferType BufferType, v interface{}, bufferUsageType BufferUsageType) {
	gl := c.gl
	b := gl.CreateBuffer()
	gl.BindBuffer(int(bufferType), b)
	// TODO: Can we assign integer to gl.bufferData?
	gl.BufferData(int(bufferType), v, int(bufferUsageType))
}

func (c *Context) BufferSubData(bufferType BufferType, data []float32) {
	gl := c.gl
	const float32Size = 4
	gl.BufferSubData(int(bufferType), 0, data)
}

func (c *Context) DrawElements(len int) {
	gl := c.gl
	gl.DrawElements(gl.TRIANGLES, len, gl.UNSIGNED_SHORT, 0)
}

func (c *Context) Flush() {
	gl := c.gl
	gl.Flush()
}
<commit_msg>Fix BindFramebuffer to pass JS null object explicitly<commit_after>// Copyright 2014 Hajime Hoshi
//
// 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.

// +build js

package opengl

import (
	"errors"
	"fmt"
	"github.com/gopherjs/gopherjs/js"
	"github.com/gopherjs/webgl"
)

type Texture js.Object
type Framebuffer js.Object
type Shader js.Object
type Program js.Object

type context struct {
	gl *webgl.Context
}

// TODO: Is there any better way to get null value?
var nullVal = js.Global.Call("eval", "null")

func NewContext(gl *webgl.Context) *Context {
	c := &Context{
		Nearest:            FilterType(gl.NEAREST),
		Linear:             FilterType(gl.LINEAR),
		VertexShader:       ShaderType(gl.VERTEX_SHADER),
		FragmentShader:     ShaderType(gl.FRAGMENT_SHADER),
		ArrayBuffer:        BufferType(gl.ARRAY_BUFFER),
		ElementArrayBuffer: BufferType(gl.ELEMENT_ARRAY_BUFFER),
		DynamicDraw:        BufferUsageType(gl.DYNAMIC_DRAW),
		StaticDraw:         BufferUsageType(gl.STATIC_DRAW),
	}
	c.gl = gl
	c.init()
	return c
}

func (c *Context) init() {
	gl := c.gl
	// Textures' pixel formats are alpha premultiplied.
	gl.Enable(gl.BLEND)
	gl.BlendFunc(gl.ONE, gl.ONE_MINUS_SRC_ALPHA)
}

func (c *Context) NewTexture(width, height int, pixels []uint8, filter FilterType) (Texture, error) {
	gl := c.gl
	t := gl.CreateTexture()
	if t.IsNull() {
		return nil, errors.New("glGenTexture failed")
	}
	gl.PixelStorei(gl.UNPACK_ALIGNMENT, 4)
	gl.BindTexture(gl.TEXTURE_2D, t)

	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, int(filter))
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, int(filter))

	// void texImage2D(GLenum target, GLint level, GLenum internalformat,
	//     GLsizei width, GLsizei height, GLint border, GLenum format,
	//     GLenum type, ArrayBufferView? pixels);
	var p interface{}
	if pixels != nil {
		p = pixels
	}
	gl.Call("texImage2D", gl.TEXTURE_2D, 0, gl.RGBA, width, height, 0, gl.RGBA, gl.UNSIGNED_BYTE, p)

	return Texture(t), nil
}

func (c *Context) TexturePixels(t Texture, width, height int) ([]uint8, error) {
	gl := c.gl
	// TODO: Use glGetTexLevelParameteri and GL_TEXTURE_WIDTH?
	pixels := js.Global.Get("Uint8Array").New(4 * width * height)
	gl.BindTexture(gl.TEXTURE_2D, t)
	gl.ReadPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, pixels)
	if e := gl.GetError(); e != gl.NO_ERROR {
		// TODO: Use glu.ErrorString
		return nil, errors.New(fmt.Sprintf("gl error: %d", e))
	}
	return pixels.Interface().([]uint8), nil
}

func (c *Context) BindTexture(t Texture) {
	gl := c.gl
	gl.BindTexture(gl.TEXTURE_2D, t)
}

func (c *Context) DeleteTexture(t Texture) {
	gl := c.gl
	gl.DeleteTexture(t)
}

func (c *Context) GlslHighpSupported() bool {
	gl := c.gl
	return gl.Call("getShaderPrecisionFormat", gl.FRAGMENT_SHADER, gl.HIGH_FLOAT).Get("precision").Int() != 0
}

func (c *Context) NewFramebuffer(texture Texture) (Framebuffer, error) {
	gl := c.gl
	f := gl.CreateFramebuffer()
	gl.BindFramebuffer(gl.FRAMEBUFFER, f)

	gl.FramebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0)
	if gl.CheckFramebufferStatus(gl.FRAMEBUFFER) != gl.FRAMEBUFFER_COMPLETE {
		return nil, errors.New("creating framebuffer failed")
	}

	return Framebuffer(f), nil
}

func (c *Context) SetViewport(f Framebuffer, width, height int) error {
	gl := c.gl
	gl.Flush()
	if f != nil {
		gl.BindFramebuffer(gl.FRAMEBUFFER, f)
	} else {
		gl.BindFramebuffer(gl.FRAMEBUFFER, nullVal)
	}
	err := gl.CheckFramebufferStatus(gl.FRAMEBUFFER)
	if err != gl.FRAMEBUFFER_COMPLETE {
		if gl.GetError() != 0 {
			return errors.New(fmt.Sprintf("glBindFramebuffer failed: %d", gl.GetError()))
		}
		return errors.New("glBindFramebuffer failed: the context is different?")
	}
	gl.Viewport(0, 0, width, height)
	return nil
}

func (c *Context) FillFramebuffer(f Framebuffer, r, g, b, a float64) error {
	gl := c.gl
	gl.ClearColor(float32(r), float32(g), float32(b), float32(a))
	gl.Clear(gl.COLOR_BUFFER_BIT)
	return nil
}

func (c *Context) DeleteFramebuffer(f Framebuffer) {
	gl := c.gl
	gl.DeleteFramebuffer(f)
}

func (c *Context) NewShader(shaderType ShaderType, source string) (Shader, error) {
	gl := c.gl
	s := gl.CreateShader(int(shaderType))
	if s.IsNull() {
		println(gl.GetError())
		return nil, errors.New("glCreateShader failed")
	}

	gl.ShaderSource(s, source)
	gl.CompileShader(s)

	if !gl.GetShaderParameterb(s, gl.COMPILE_STATUS) {
		log := gl.GetShaderInfoLog(s)
		return nil, errors.New(fmt.Sprintf("shader compile failed: %s", log))
	}
	return Shader(s), nil
}

func (c *Context) DeleteShader(s Shader) {
	gl := c.gl
	gl.DeleteShader(s)
}

func (c *Context) NewProgram(shaders []Shader) (Program, error) {
	gl := c.gl
	p := gl.CreateProgram()
	if p.IsNull() {
		return nil, errors.New("glCreateProgram failed")
	}

	for _, shader := range shaders {
		gl.AttachShader(p, shader)
	}
	gl.LinkProgram(p)
	if !gl.GetProgramParameterb(p, gl.LINK_STATUS) {
		return nil, errors.New("program error")
	}
	return Program(p), nil
}

func (c *Context) UseProgram(p Program) {
	gl := c.gl
	gl.UseProgram(p)
}

func (c *Context) Uniform(p Program, location string, v interface{}) {
	gl := c.gl
	l := gl.GetUniformLocation(p, location)
	switch v := v.(type) {
	case int:
		gl.Uniform1i(l, v)
	case [4]float32:
		gl.Call("uniform4fv", l, v[:])
	case [16]float32:
		gl.UniformMatrix4fv(l, false, v[:])
	}
}

func (c *Context) VertexAttribPointer(p Program, location string, stride int, v uintptr) {
	gl := c.gl
	l := gl.GetAttribLocation(p, location)
	gl.VertexAttribPointer(l, 2, gl.FLOAT, false, stride, int(v))
}

func (c *Context) EnableVertexAttribArray(p Program, location string) {
	gl := c.gl
	l := gl.GetAttribLocation(p, location)
	gl.EnableVertexAttribArray(l)
}

func (c *Context) DisableVertexAttribArray(p Program, location string) {
	gl := c.gl
	l := gl.GetAttribLocation(p, location)
	gl.DisableVertexAttribArray(l)
}

func (c *Context) NewBuffer(bufferType BufferType, v interface{}, bufferUsageType BufferUsageType) {
	gl := c.gl
	b := gl.CreateBuffer()
	gl.BindBuffer(int(bufferType), b)
	// TODO: Can we assign integer to gl.bufferData?
	gl.BufferData(int(bufferType), v, int(bufferUsageType))
}

func (c *Context) BufferSubData(bufferType BufferType, data []float32) {
	gl := c.gl
	const float32Size = 4
	gl.BufferSubData(int(bufferType), 0, data)
}

func (c *Context) DrawElements(len int) {
	gl := c.gl
	gl.DrawElements(gl.TRIANGLES, len, gl.UNSIGNED_SHORT, 0)
}

func (c *Context) Flush() {
	gl := c.gl
	gl.Flush()
}
<|endoftext|>