<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package api

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/cilium/cilium/pkg/fqdn/dns"
	"github.com/cilium/cilium/pkg/fqdn/matchpattern"
)

var (
	// allowedMatchNameChars tests that MatchName contains only valid DNS characters
	allowedMatchNameChars = regexp.MustCompile("^[-a-zA-Z0-9_.]+$")

	// allowedPatternChars tests that the MatchPattern field contains only the
	// characters we want in our wilcard scheme.
	allowedPatternChars = regexp.MustCompile("^[-a-zA-Z0-9_.*]+$") // the * inside the [] is a literal *

	// FQDNMatchNameRegexString is a regex string which matches what's expected
	// in the MatchName field in the FQDNSelector. This should be kept in-sync
	// with the marker comment for validation. There's no way to use a Golang
	// variable in the marker comment, so it's left up to the developer.
	FQDNMatchNameRegexString = `^([-a-zA-Z0-9_]+[.]?)+$`

	// FQDNMatchPatternRegexString is a regex string which matches what's expected
	// in the MatchPattern field in the FQDNSelector. This should be kept in-sync
	// with the marker comment for validation. There's no way to use a Golang
	// variable in the marker comment, so it's left up to the developer.
	FQDNMatchPatternRegexString = `^([-a-zA-Z0-9_*]+[.]?)+$`
)

type FQDNSelector struct {
	// MatchName matches literal DNS names. A trailing "." is automatically added
	// when missing.
	//
	// +kubebuilder:validation:Pattern=`^([-a-zA-Z0-9_]+[.]?)+$`
	MatchName string `json:"matchName,omitempty"`

	// MatchPattern allows using wildcards to match DNS names. All wildcards are
	// case insensitive. The wildcards are:
	// - "*" matches 0 or more DNS valid characters, and may occur anywhere in
	// the pattern. As a special case a "*" as the leftmost character, without a
	// following "." matches all subdomains as well as the name to the right.
	// A trailing "." is automatically added when missing.
	//
	// Examples:
	// `*.cilium.io` matches subomains of cilium at that level
	//   www.cilium.io and blog.cilium.io match, cilium.io and google.com do not
	// `*cilium.io` matches cilium.io and all subdomains 1 level below
	//   www.cilium.io, blog.cilium.io and cilium.io match, google.com does not
	// sub*.cilium.io matches subdomains of cilium where the subdomain component
	// begins with "sub"
	//   sub.cilium.io and subdomain.cilium.io match, www.cilium.io,
	//   blog.cilium.io, cilium.io and google.com do not
	//
	// +kubebuilder:validation:Pattern=`^([-a-zA-Z0-9_*]+[.]?)+$`
	MatchPattern string `json:"matchPattern,omitempty"`
}

func (s *FQDNSelector) String() string {
	const m = "MatchName: "
	const mm = ", MatchPattern: "
	var str strings.Builder
	str.Grow(len(m) + len(mm) + len(s.MatchName) + len(s.MatchPattern))
	str.WriteString(m)
	str.WriteString(s.MatchName)
	str.WriteString(mm)
	str.WriteString(s.MatchPattern)
	return str.String()
}

// sanitize for FQDNSelector is a little wonky. While we do more processing
// when using MatchName the basic requirement is that is a valid regexp. We
// test that it can compile here.
func (s *FQDNSelector) sanitize() error {
	if len(s.MatchName) > 0 && len(s.MatchPattern) > 0 {
		return fmt.Errorf("only one of MatchName or MatchPattern is allowed in an FQDNSelector")
	}
	if len(s.MatchName) > 0 && !allowedMatchNameChars.MatchString(s.MatchName) {
		return fmt.Errorf("Invalid characters in MatchName: \"%s\". Only 0-9, a-z, A-Z and . and - characters are allowed", s.MatchName)
	}

	if len(s.MatchPattern) > 0 && !allowedPatternChars.MatchString(s.MatchPattern) {
		return fmt.Errorf("Invalid characters in MatchPattern: \"%s\". Only 0-9, a-z, A-Z and ., - and * characters are allowed", s.MatchPattern)
	}
	_, err := matchpattern.Validate(s.MatchPattern)
	return err
}

// ToRegex converts the given FQDNSelector to its corresponding regular
// expression. If the MatchName field is set in the selector, it performs all
// needed formatting to ensure that the field is a valid regular expression.
func (s *FQDNSelector) ToRegex() (*regexp.Regexp, error) {
	var preparedMatch string
	if s.MatchName != "" {
		preparedMatch = dns.FQDN(s.MatchName)
	} else {
		preparedMatch = matchpattern.Sanitize(s.MatchPattern)
	}

	regex, err := matchpattern.Validate(preparedMatch)
	return regex, err
}

// PortRuleDNS is a list of allowed DNS lookups.
type PortRuleDNS FQDNSelector

// Sanitize checks that the matchName in the portRule can be compiled as a
// regex. It does not check that a DNS name is a valid DNS name.
func (r *PortRuleDNS) Sanitize() error {
	if len(r.MatchName) > 0 && !allowedMatchNameChars.MatchString(r.MatchName) {
		return fmt.Errorf("Invalid characters in MatchName: \"%s\". Only 0-9, a-z, A-Z and . and - characters are allowed", r.MatchName)
	}

	if len(r.MatchPattern) > 0 && !allowedPatternChars.MatchString(r.MatchPattern) {
		return fmt.Errorf("Invalid characters in MatchPattern: \"%s\". Only 0-9, a-z, A-Z and ., - and * characters are allowed", r.MatchPattern)
	}
	_, err := matchpattern.Validate(r.MatchPattern)
	return err
}

// GetAsEndpointSelectors returns a FQDNSelector as a single EntityNone
// EndpointSelector slice.
// Note that toFQDNs behaves differently than most other rules. The presence of
// any toFQDNs rules means the endpoint must enforce policy, but the IPs are later
// added as toCIDRSet entries and processed as such.
func (s *FQDNSelector) GetAsEndpointSelectors() EndpointSelectorSlice {
	return []EndpointSelector{EndpointSelectorNone}
}

// FQDNSelectorSlice is a wrapper type for []FQDNSelector to make is simpler to
// bind methods.
type FQDNSelectorSlice []FQDNSelector

// GetAsEndpointSelectors will return a single EntityNone if any
// toFQDNs rules exist, and a nil slice otherwise.
func (s FQDNSelectorSlice) GetAsEndpointSelectors() EndpointSelectorSlice {
	for _, rule := range s {
		return rule.GetAsEndpointSelectors()
	}
	return nil
}
<commit_msg>doc: update the api spec for fqdn egress policies code comments.<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package api

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/cilium/cilium/pkg/fqdn/dns"
	"github.com/cilium/cilium/pkg/fqdn/matchpattern"
)

var (
	// allowedMatchNameChars tests that MatchName contains only valid DNS characters
	allowedMatchNameChars = regexp.MustCompile("^[-a-zA-Z0-9_.]+$")

	// allowedPatternChars tests that the MatchPattern field contains only the
	// characters we want in our wilcard scheme.
	allowedPatternChars = regexp.MustCompile("^[-a-zA-Z0-9_.*]+$") // the * inside the [] is a literal *

	// FQDNMatchNameRegexString is a regex string which matches what's expected
	// in the MatchName field in the FQDNSelector. This should be kept in-sync
	// with the marker comment for validation. There's no way to use a Golang
	// variable in the marker comment, so it's left up to the developer.
	FQDNMatchNameRegexString = `^([-a-zA-Z0-9_]+[.]?)+$`

	// FQDNMatchPatternRegexString is a regex string which matches what's expected
	// in the MatchPattern field in the FQDNSelector. This should be kept in-sync
	// with the marker comment for validation. There's no way to use a Golang
	// variable in the marker comment, so it's left up to the developer.
	FQDNMatchPatternRegexString = `^([-a-zA-Z0-9_*]+[.]?)+$`
)

type FQDNSelector struct {
	// MatchName matches literal DNS names. A trailing "." is automatically added
	// when missing.
	//
	// +kubebuilder:validation:Pattern=`^([-a-zA-Z0-9_]+[.]?)+$`
	MatchName string `json:"matchName,omitempty"`

	// MatchPattern allows using wildcards to match DNS names. All wildcards are
	// case insensitive. The wildcards are:
	// - "*" matches 0 or more DNS valid characters, and may occur anywhere in
	// the pattern. As a special case a "*" as the leftmost character, without a
	// following "." matches all subdomains as well as the name to the right.
	// A trailing "." is automatically added when missing.
	//
	// Examples:
	// `*.cilium.io` matches subomains of cilium at that level
	//   www.cilium.io and blog.cilium.io match, cilium.io and google.com do not
	// `*cilium.io` matches cilium.io and all subdomains ends with "cilium.io"
	//   except those containing "." separator, subcilium.io and sub-cilium.io match,
	//   www.cilium.io and blog.cilium.io does not
	// sub*.cilium.io matches subdomains of cilium where the subdomain component
	// begins with "sub"
	//   sub.cilium.io and subdomain.cilium.io match, www.cilium.io,
	//   blog.cilium.io, cilium.io and google.com do not
	//
	// +kubebuilder:validation:Pattern=`^([-a-zA-Z0-9_*]+[.]?)+$`
	MatchPattern string `json:"matchPattern,omitempty"`
}

func (s *FQDNSelector) String() string {
	const m = "MatchName: "
	const mm = ", MatchPattern: "
	var str strings.Builder
	str.Grow(len(m) + len(mm) + len(s.MatchName) + len(s.MatchPattern))
	str.WriteString(m)
	str.WriteString(s.MatchName)
	str.WriteString(mm)
	str.WriteString(s.MatchPattern)
	return str.String()
}

// sanitize for FQDNSelector is a little wonky. While we do more processing
// when using MatchName the basic requirement is that is a valid regexp. We
// test that it can compile here.
func (s *FQDNSelector) sanitize() error {
	if len(s.MatchName) > 0 && len(s.MatchPattern) > 0 {
		return fmt.Errorf("only one of MatchName or MatchPattern is allowed in an FQDNSelector")
	}
	if len(s.MatchName) > 0 && !allowedMatchNameChars.MatchString(s.MatchName) {
		return fmt.Errorf("Invalid characters in MatchName: \"%s\". Only 0-9, a-z, A-Z and . and - characters are allowed", s.MatchName)
	}

	if len(s.MatchPattern) > 0 && !allowedPatternChars.MatchString(s.MatchPattern) {
		return fmt.Errorf("Invalid characters in MatchPattern: \"%s\". Only 0-9, a-z, A-Z and ., - and * characters are allowed", s.MatchPattern)
	}
	_, err := matchpattern.Validate(s.MatchPattern)
	return err
}

// ToRegex converts the given FQDNSelector to its corresponding regular
// expression. If the MatchName field is set in the selector, it performs all
// needed formatting to ensure that the field is a valid regular expression.
func (s *FQDNSelector) ToRegex() (*regexp.Regexp, error) {
	var preparedMatch string
	if s.MatchName != "" {
		preparedMatch = dns.FQDN(s.MatchName)
	} else {
		preparedMatch = matchpattern.Sanitize(s.MatchPattern)
	}

	regex, err := matchpattern.Validate(preparedMatch)
	return regex, err
}

// PortRuleDNS is a list of allowed DNS lookups.
type PortRuleDNS FQDNSelector

// Sanitize checks that the matchName in the portRule can be compiled as a
// regex. It does not check that a DNS name is a valid DNS name.
func (r *PortRuleDNS) Sanitize() error {
	if len(r.MatchName) > 0 && !allowedMatchNameChars.MatchString(r.MatchName) {
		return fmt.Errorf("Invalid characters in MatchName: \"%s\". Only 0-9, a-z, A-Z and . and - characters are allowed", r.MatchName)
	}

	if len(r.MatchPattern) > 0 && !allowedPatternChars.MatchString(r.MatchPattern) {
		return fmt.Errorf("Invalid characters in MatchPattern: \"%s\". Only 0-9, a-z, A-Z and ., - and * characters are allowed", r.MatchPattern)
	}
	_, err := matchpattern.Validate(r.MatchPattern)
	return err
}

// GetAsEndpointSelectors returns a FQDNSelector as a single EntityNone
// EndpointSelector slice.
// Note that toFQDNs behaves differently than most other rules. The presence of
// any toFQDNs rules means the endpoint must enforce policy, but the IPs are later
// added as toCIDRSet entries and processed as such.
func (s *FQDNSelector) GetAsEndpointSelectors() EndpointSelectorSlice {
	return []EndpointSelector{EndpointSelectorNone}
}

// FQDNSelectorSlice is a wrapper type for []FQDNSelector to make is simpler to
// bind methods.
type FQDNSelectorSlice []FQDNSelector

// GetAsEndpointSelectors will return a single EntityNone if any
// toFQDNs rules exist, and a nil slice otherwise.
func (s FQDNSelectorSlice) GetAsEndpointSelectors() EndpointSelectorSlice {
	for _, rule := range s {
		return rule.GetAsEndpointSelectors()
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Implements methods to interact with the official Trade Offer API.

See: https://developer.valvesoftware.com/wiki/Steam_Web_API/IEconService
*/
package tradeoffer

import (
	"encoding/json"
	"github.com/Philipp15b/go-steam/economy/inventory"
	"github.com/Philipp15b/go-steam/steamid"
)

type TradeOfferState uint

const (
	TradeOfferState_Invalid                  TradeOfferState = 1  // Invalid
	TradeOfferState_Active                                   = 2  // This trade offer has been sent, neither party has acted on it yet.
	TradeOfferState_Accepted                                 = 3  // The trade offer was accepted by the recipient and items were exchanged.
	TradeOfferState_Countered                                = 4  // The recipient made a counter offer
	TradeOfferState_Expired                                  = 5  // The trade offer was not accepted before the expiration date
	TradeOfferState_Canceled                                 = 6  // The sender cancelled the offer
	TradeOfferState_Declined                                 = 7  // The recipient declined the offer
	TradeOfferState_InvalidItems                             = 8  // Some of the items in the offer are no longer available (indicated by the missing flag in the output)
	TradeOfferState_CreatedNeedsConfirmation                 = 9  // The offer hasn't been sent yet and is awaiting email/mobile confirmation. The offer is only visible to the sender.
	TradeOfferState_CanceledBySecondFactor                   = 10 // Either party canceled the offer via email/mobile. The offer is visible to both parties, even if the sender canceled it before it was sent.
	TradeOfferState_InEscrow                                 = 11 // The trade has been placed on hold. The items involved in the trade have all been removed from both parties' inventories and will be automatically delivered in the future.
)

type TradeOfferConfirmationMethod uint

const (
	TradeOfferConfirmationMethod_Invalid   TradeOfferConfirmationMethod = 0
	TradeOfferConfirmationMethod_Email                                  = 1
	TradeOfferConfirmationMethod_MobileApp                              = 2
)

type Asset struct {
	AppId      uint32 `json:",string"`
	ContextId  uint64 `json:",string"`
	AssetId    uint64 `json:",string"`
	CurrencyId uint64 `json:",string"`
	ClassId    uint64 `json:",string"`
	InstanceId uint64 `json:",string"`
	Amount     uint64 `json:",string"`
	Missing    bool
}

type TradeOffer struct {
	TradeOfferId       uint64                       `json:",string"`
	TradeId            uint64                       `json:",string"`
	OtherAccountId     uint32                       `json:"accountid_other"`
	OtherSteamId       steamid.SteamId              `json:"-"`
	Message            string                       `json:"message"`
	ExpirationTime     uint32                       `json:"expiraton_time"`
	State              TradeOfferState              `json:"trade_offer_state"`
	ToGive             []*Asset                     `json:"items_to_give"`
	ToReceive          []*Asset                     `json:"items_to_receive"`
	IsOurOffer         bool                         `json:"is_our_offer"`
	TimeCreated        uint32                       `json:"time_created"`
	TimeUpdated        uint32                       `json:"time_updated"`
	EscrowEndDate      uint32                       `json:"escrow_end_date"`
	ConfirmationMethod TradeOfferConfirmationMethod `json:"confirmation_method"`
}

func (t *TradeOffer) UnmarshalJSON(data []byte) error {
	type Alias TradeOffer
	aux := struct {
		*Alias
	}{
		Alias: (*Alias)(t),
	}
	if err := json.Unmarshal(data, &aux); err != nil {
		return err
	}
	if t.OtherAccountId == 0 {
		t.OtherSteamId = steamid.SteamId(0)
		return nil
	}
	t.OtherSteamId = steamid.SteamId(uint64(t.OtherAccountId) + 76561197960265728)
	return nil
}

type TradeOffersResult struct {
	Sent         []*TradeOffer `json:"trade_offers_sent"`
	Received     []*TradeOffer `json:"trade_offers_received"`
	Descriptions []*Description
}

type TradeOfferResult struct {
	Offer        *TradeOffer
	Descriptions []*Description
}
type Description struct {
	AppId      uint32 `json:",string"`
	ClassId    uint64 `json:",string"`
	InstanceId uint64 `json:",string"`

	IconUrl      string `json:"icon_url"`
	IconUrlLarge string `json:"icon_url_large"`

	Name           string
	MarketName     string `json:"market_name"`
	MarketHashName string `json:"market_hash_name"`

	// Colors in hex, for example `B2B2B2`
	NameColor       string `json:"name_color"`
	BackgroundColor string `json:"background_color"`

	Type string

	Tradable                  bool
	Commodity                 bool
	MarketTradableRestriction uint32 `json:"market_tradable_restriction"`

	Descriptions inventory.DescriptionLines
	Actions      []*inventory.Action
}
<commit_msg>Fix tradeoffer descriptions json tags<commit_after>/*
Implements methods to interact with the official Trade Offer API.

See: https://developer.valvesoftware.com/wiki/Steam_Web_API/IEconService
*/
package tradeoffer

import (
	"encoding/json"
	"github.com/Philipp15b/go-steam/economy/inventory"
	"github.com/Philipp15b/go-steam/steamid"
)

type TradeOfferState uint

const (
	TradeOfferState_Invalid                  TradeOfferState = 1  // Invalid
	TradeOfferState_Active                                   = 2  // This trade offer has been sent, neither party has acted on it yet.
	TradeOfferState_Accepted                                 = 3  // The trade offer was accepted by the recipient and items were exchanged.
	TradeOfferState_Countered                                = 4  // The recipient made a counter offer
	TradeOfferState_Expired                                  = 5  // The trade offer was not accepted before the expiration date
	TradeOfferState_Canceled                                 = 6  // The sender cancelled the offer
	TradeOfferState_Declined                                 = 7  // The recipient declined the offer
	TradeOfferState_InvalidItems                             = 8  // Some of the items in the offer are no longer available (indicated by the missing flag in the output)
	TradeOfferState_CreatedNeedsConfirmation                 = 9  // The offer hasn't been sent yet and is awaiting email/mobile confirmation. The offer is only visible to the sender.
	TradeOfferState_CanceledBySecondFactor                   = 10 // Either party canceled the offer via email/mobile. The offer is visible to both parties, even if the sender canceled it before it was sent.
	TradeOfferState_InEscrow                                 = 11 // The trade has been placed on hold. The items involved in the trade have all been removed from both parties' inventories and will be automatically delivered in the future.
)

type TradeOfferConfirmationMethod uint

const (
	TradeOfferConfirmationMethod_Invalid   TradeOfferConfirmationMethod = 0
	TradeOfferConfirmationMethod_Email                                  = 1
	TradeOfferConfirmationMethod_MobileApp                              = 2
)

type Asset struct {
	AppId      uint32 `json:",string"`
	ContextId  uint64 `json:",string"`
	AssetId    uint64 `json:",string"`
	CurrencyId uint64 `json:",string"`
	ClassId    uint64 `json:",string"`
	InstanceId uint64 `json:",string"`
	Amount     uint64 `json:",string"`
	Missing    bool
}

type TradeOffer struct {
	TradeOfferId       uint64                       `json:",string"`
	TradeId            uint64                       `json:",string"`
	OtherAccountId     uint32                       `json:"accountid_other"`
	OtherSteamId       steamid.SteamId              `json:"-"`
	Message            string                       `json:"message"`
	ExpirationTime     uint32                       `json:"expiraton_time"`
	State              TradeOfferState              `json:"trade_offer_state"`
	ToGive             []*Asset                     `json:"items_to_give"`
	ToReceive          []*Asset                     `json:"items_to_receive"`
	IsOurOffer         bool                         `json:"is_our_offer"`
	TimeCreated        uint32                       `json:"time_created"`
	TimeUpdated        uint32                       `json:"time_updated"`
	EscrowEndDate      uint32                       `json:"escrow_end_date"`
	ConfirmationMethod TradeOfferConfirmationMethod `json:"confirmation_method"`
}

func (t *TradeOffer) UnmarshalJSON(data []byte) error {
	type Alias TradeOffer
	aux := struct {
		*Alias
	}{
		Alias: (*Alias)(t),
	}
	if err := json.Unmarshal(data, &aux); err != nil {
		return err
	}
	if t.OtherAccountId == 0 {
		t.OtherSteamId = steamid.SteamId(0)
		return nil
	}
	t.OtherSteamId = steamid.SteamId(uint64(t.OtherAccountId) + 76561197960265728)
	return nil
}

type TradeOffersResult struct {
	Sent         []*TradeOffer `json:"trade_offers_sent"`
	Received     []*TradeOffer `json:"trade_offers_received"`
	Descriptions []*Description
}

type TradeOfferResult struct {
	Offer        *TradeOffer
	Descriptions []*Description
}
type Description struct {
	AppId      uint32 `json:"appid"`
	ClassId    uint64 `json:"classid,string"`
	InstanceId uint64 `json:"instanceid,string"`

	IconUrl      string `json:"icon_url"`
	IconUrlLarge string `json:"icon_url_large"`

	Name           string
	MarketName     string `json:"market_name"`
	MarketHashName string `json:"market_hash_name"`

	// Colors in hex, for example `B2B2B2`
	NameColor       string `json:"name_color"`
	BackgroundColor string `json:"background_color"`

	Type string

	Tradable                  bool   `json:"tradable"`
	Commodity                 bool   `json:"commodity"`
	MarketTradableRestriction uint32 `json:"market_tradable_restriction"`

	Descriptions inventory.DescriptionLines `json:"descriptions"`
	Actions      []*inventory.Action        `json:"actions"`
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"

	"github.com/grafana/metrictank/stats"
	"github.com/raintank/raintank-apps/task-server/api"
	"github.com/raintank/raintank-apps/task-server/event"
	"github.com/raintank/raintank-apps/task-server/manager"
	"github.com/raintank/raintank-apps/task-server/sqlstore"
	tsConfig "github.com/raintank/raintank-apps/task-server/taskserverconfig"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"
)

var (
	GitHash     = "(none)"
	showVersion = flag.Bool("version", false, "print version string")
	logLevel    = flag.Int("log-level", 2, "log level. 0=TRACE|1=DEBUG|2=INFO|3=WARN|4=ERROR|5=CRITICAL|6=FATAL")
	confFile    = flag.String("config", "/etc/raintank/task-server.ini", "configuration file path")

	addr            = flag.String("addr", "localhost:80", "http service address")
	dbType          = flag.String("db-type", "sqlite3", "Database type. sqlite3 or mysql")
	dbConnectString = flag.String("db-connect-str", "file:/tmp/task-server.db?cache=shared&mode=rwc&_loc=Local", "DSN to connect to DB. https://godoc.org/github.com/mattn/go-sqlite3#SQLiteDriver.Open or https://github.com/go-sql-driver/mysql#dsn-data-source-name")

	exchange    = flag.String("exchange", "events", "Rabbitmq Topic Exchange")
	rabbitmqUrl = flag.String("rabbitmq-url", "amqp://guest:guest@localhost:5672/", "rabbitmq Url")

	appAPIKey = flag.String("app-api-key", "app_not_very_secret_key", "API Key for task-server and task-agent communication")
)

var (
	taskServerRunning = stats.NewGauge32("running")
)

func main() {
	flag.Parse()

	var cfile string
	if _, err := os.Stat(*confFile); err == nil {
		cfile = *confFile
	}

	conf, err := globalconf.NewWithOptions(&globalconf.Options{
		Filename:  cfile,
		EnvPrefix: "TASKSERVER_",
	})
	if err != nil {
		panic(fmt.Sprintf("error with configuration file: %s", err))
	}
	tsConfig.ConfigSetup()
	conf.ParseAll()

	log.NewLogger(0, "console", fmt.Sprintf(`{"level": %d, "formatting":true}`, *logLevel))
	// workaround for https://github.com/grafana/grafana/issues/4055
	switch *logLevel {
	case 0:
		log.Level(log.TRACE)
	case 1:
		log.Level(log.DEBUG)
	case 2:
		log.Level(log.INFO)
	case 3:
		log.Level(log.WARN)
	case 4:
		log.Level(log.ERROR)
	case 5:
		log.Level(log.CRITICAL)
	case 6:
		log.Level(log.FATAL)
	}

	if *showVersion {
		fmt.Printf("task-server (built with %s, git hash %s)\n", runtime.Version(), GitHash)
		return
	}

	var hostname, hnErr = os.Hostname()
	if hnErr != nil {
		log.Fatal(4, "failed to get hostname from OS.")
	}

	tsConfig.ConfigProcess(hostname)
	tsConfig.Start()

	// initialize DB
	enableSqlLog := false
	if *logLevel >= int(log.DEBUG) {
		enableSqlLog = true
	}
	sqlstore.NewEngine(*dbType, *dbConnectString, enableSqlLog)

	// delete any stale agentSessions.
	if err := sqlstore.DeleteAgentSessionsByServer(hostname); err != nil {
		panic(err)
	}

	m := api.NewApi(*appAPIKey)
	err = event.Init(*rabbitmqUrl, *exchange)
	if err != nil {
		log.Fatal(4, "failed to init event PubSub. %s", err)
	}

	manager.Init()

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

	log.Info("starting up")
	taskServerRunning.Inc()

	// define our own listner so we can call Close on it
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal(4, err.Error())
	}
	done := make(chan struct{})
	go handleShutdown(done, interrupt, l)
	log.Info("%v", http.Serve(l, m))
	<-done
}

func handleShutdown(done chan struct{}, interrupt chan os.Signal, l net.Listener) {
	<-interrupt
	log.Info("shutdown started.")
	l.Close()
	api.ActiveSockets.CloseAll()
	close(done)
}
<commit_msg>add debug logging for task server api key<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"

	"github.com/grafana/metrictank/stats"
	"github.com/raintank/raintank-apps/task-server/api"
	"github.com/raintank/raintank-apps/task-server/event"
	"github.com/raintank/raintank-apps/task-server/manager"
	"github.com/raintank/raintank-apps/task-server/sqlstore"
	tsConfig "github.com/raintank/raintank-apps/task-server/taskserverconfig"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"
)

var (
	GitHash     = "(none)"
	showVersion = flag.Bool("version", false, "print version string")
	logLevel    = flag.Int("log-level", 2, "log level. 0=TRACE|1=DEBUG|2=INFO|3=WARN|4=ERROR|5=CRITICAL|6=FATAL")
	confFile    = flag.String("config", "/etc/raintank/task-server.ini", "configuration file path")

	addr            = flag.String("addr", "localhost:80", "http service address")
	dbType          = flag.String("db-type", "sqlite3", "Database type. sqlite3 or mysql")
	dbConnectString = flag.String("db-connect-str", "file:/tmp/task-server.db?cache=shared&mode=rwc&_loc=Local", "DSN to connect to DB. https://godoc.org/github.com/mattn/go-sqlite3#SQLiteDriver.Open or https://github.com/go-sql-driver/mysql#dsn-data-source-name")

	exchange    = flag.String("exchange", "events", "Rabbitmq Topic Exchange")
	rabbitmqUrl = flag.String("rabbitmq-url", "amqp://guest:guest@localhost:5672/", "rabbitmq Url")

	appAPIKey = flag.String("app-api-key", "app_not_very_secret_key", "API Key for task-server and task-agent communication")
)

var (
	taskServerRunning = stats.NewGauge32("running")
)

func main() {
	flag.Parse()

	var cfile string
	if _, err := os.Stat(*confFile); err == nil {
		cfile = *confFile
	}

	conf, err := globalconf.NewWithOptions(&globalconf.Options{
		Filename:  cfile,
		EnvPrefix: "TASKSERVER_",
	})
	if err != nil {
		panic(fmt.Sprintf("error with configuration file: %s", err))
	}
	tsConfig.ConfigSetup()
	conf.ParseAll()

	log.NewLogger(0, "console", fmt.Sprintf(`{"level": %d, "formatting":true}`, *logLevel))
	// workaround for https://github.com/grafana/grafana/issues/4055
	switch *logLevel {
	case 0:
		log.Level(log.TRACE)
	case 1:
		log.Level(log.DEBUG)
	case 2:
		log.Level(log.INFO)
	case 3:
		log.Level(log.WARN)
	case 4:
		log.Level(log.ERROR)
	case 5:
		log.Level(log.CRITICAL)
	case 6:
		log.Level(log.FATAL)
	}

	if *showVersion {
		fmt.Printf("task-server (built with %s, git hash %s)\n", runtime.Version(), GitHash)
		return
	}

	var hostname, hnErr = os.Hostname()
	if hnErr != nil {
		log.Fatal(4, "failed to get hostname from OS.")
	}

	tsConfig.ConfigProcess(hostname)
	tsConfig.Start()

	// initialize DB
	enableSqlLog := false
	if *logLevel >= int(log.DEBUG) {
		enableSqlLog = true
	}
	sqlstore.NewEngine(*dbType, *dbConnectString, enableSqlLog)

	// delete any stale agentSessions.
	if err := sqlstore.DeleteAgentSessionsByServer(hostname); err != nil {
		panic(err)
	}

	log.Info("main: using app-api-key: %s", *appAPIKey)

	m := api.NewApi(*appAPIKey)
	err = event.Init(*rabbitmqUrl, *exchange)
	if err != nil {
		log.Fatal(4, "failed to init event PubSub. %s", err)
	}

	manager.Init()

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

	log.Info("starting up")
	taskServerRunning.Inc()

	// define our own listner so we can call Close on it
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal(4, err.Error())
	}
	done := make(chan struct{})
	go handleShutdown(done, interrupt, l)
	log.Info("%v", http.Serve(l, m))
	<-done
}

func handleShutdown(done chan struct{}, interrupt chan os.Signal, l net.Listener) {
	<-interrupt
	log.Info("shutdown started.")
	l.Close()
	api.ActiveSockets.CloseAll()
	close(done)
}
<|endoftext|>


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

package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"io/fs"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"golang.org/x/exp/vulndb/internal/audit"
	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/packages/packagestest"
	"golang.org/x/vulndb/client"
)

// TODO(zpavlinovic): improve integration tests.

// goYamlVuln contains vulnerability info for github.com/go-yaml/yaml package.
var goYamlVuln string = `[{"ID":"GO-2020-0036","Published":"2021-04-14T12:00:00Z","Modified":"2021-04-14T12:00:00Z","Withdrawn":null,"Aliases":["CVE-2019-11254"],"Package":{"Name":"github.com/go-yaml/yaml","Ecosystem":"go"},"Details":"An attacker can craft malicious YAML which will consume significant\nsystem resources when Unmarshalled.\n","Affects":{"Ranges":[{"Type":"SEMVER","Introduced":"","Fixed":"v2.2.8+incompatible"}]},"References":[{"Type":"code review","URL":"https://github.com/go-yaml/yaml/pull/555"},{"Type":"fix","URL":"https://github.com/go-yaml/yaml/commit/53403b58ad1b561927d19068c655246f2db79d48"},{"Type":"misc","URL":"https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=18496"}],"ecosystem_specific":{"Symbols":["yaml_parser_fetch_more_tokens"],"URL":"https://go.googlesource.com/vulndb/+/refs/heads/main/reports/GO-2020-0036.toml"}},{"ID":"GO-2021-0061","Published":"2021-04-14T12:00:00Z","Modified":"2021-04-14T12:00:00Z","Withdrawn":null,"Package":{"Name":"github.com/go-yaml/yaml","Ecosystem":"go"},"Details":"A maliciously crafted input can cause resource exhaustion due to\nalias chasing.\n","Affects":{"Ranges":[{"Type":"SEMVER","Introduced":"","Fixed":"v2.2.3+incompatible"}]},"References":[{"Type":"code review","URL":"https://github.com/go-yaml/yaml/pull/375"},{"Type":"fix","URL":"https://github.com/go-yaml/yaml/commit/bb4e33bf68bf89cad44d386192cbed201f35b241"}],"ecosystem_specific":{"Symbols":["decoder.unmarshal"],"URL":"https://go.googlesource.com/vulndb/+/refs/heads/main/reports/GO-2021-0061.toml"}}]`

// cryptoSSHVuln contains vulnerability info for golang.org/x/crypto/ssh.
var cryptoSSHVuln string = `[{"ID":"GO-2020-0012","Published":"2021-04-14T12:00:00Z","Modified":"2021-04-14T12:00:00Z","Withdrawn":null,"Aliases":["CVE-2020-9283"],"Package":{"Name":"golang.org/x/crypto/ssh","Ecosystem":"go"},"Details":"An attacker can craft an ssh-ed25519 or sk-ssh-ed25519@openssh.com public\nkey, such that the library will panic when trying to verify a signature\nwith it.\n","Affects":{"Ranges":[{"Type":"SEMVER","Introduced":"","Fixed":"v0.0.0-20200220183623-bac4c82f6975"}]},"References":[{"Type":"code review","URL":"https://go-review.googlesource.com/c/crypto/+/220357"},{"Type":"fix","URL":"https://github.com/golang/crypto/commit/bac4c82f69751a6dd76e702d54b3ceb88adab236"},{"Type":"misc","URL":"https://groups.google.com/g/golang-announce/c/3L45YRc91SY"}],"ecosystem_specific":{"Symbols":["parseED25519","ed25519PublicKey.Verify","parseSKEd25519","skEd25519PublicKey.Verify","NewPublicKey"],"URL":"https://go.googlesource.com/vulndb/+/refs/heads/main/reports/GO-2020-0012.toml"}},{"ID":"GO-2020-0013","Published":"2021-04-14T12:00:00Z","Modified":"2021-04-14T12:00:00Z","Withdrawn":null,"Aliases":["CVE-2017-3204"],"Package":{"Name":"golang.org/x/crypto/ssh","Ecosystem":"go"},"Details":"By default host key verification is disabled which allows for\nman-in-the-middle attacks against SSH clients if\n[ClientConfig.HostKeyCallback] is not set.\n","Affects":{"Ranges":[{"Type":"SEMVER","Introduced":"","Fixed":"v0.0.0-20170330155735-e4e2799dd7aa"}]},"References":[{"Type":"code review","URL":"https://go-review.googlesource.com/38701"},{"Type":"fix","URL":"https://github.com/golang/crypto/commit/e4e2799dd7aab89f583e1d898300d96367750991"},{"Type":"misc","URL":"https://github.com/golang/go/issues/19767"},{"Type":"misc","URL":"https://bridge.grumpy-troll.org/2017/04/golang-ssh-security/"}],"ecosystem_specific":{"Symbols":["NewClientConn"],"URL":"https://go.googlesource.com/vulndb/+/refs/heads/main/reports/GO-2020-0013.toml"}}]`

// k8sAPIServerVuln contains vulnerability info for k8s.io/apiextensions-apiserver/pkg/apiserver.
var k8sAPIServerVuln string = `[{"ID":"GO-2021-0062","Published":"2021-04-14T12:00:00Z","Modified":"2021-04-14T12:00:00Z","Withdrawn":null,"Aliases":["CVE-2019-11253"],"Package":{"Name":"k8s.io/apiextensions-apiserver/pkg/apiserver","Ecosystem":"go"},"Details":"A maliciously crafted YAML or JSON message can cause resource\nexhaustion.\n","Affects":{"Ranges":[{"Type":"SEMVER","Introduced":"","Fixed":"v0.17.0"}]},"References":[{"Type":"code review","URL":"https://github.com/kubernetes/kubernetes/pull/83261"},{"Type":"fix","URL":"https://github.com/kubernetes/apiextensions-apiserver/commit/9cfd100448d12f999fbf913ae5d4fef2fcd66871"},{"Type":"misc","URL":"https://github.com/kubernetes/kubernetes/issues/83253"},{"Type":"misc","URL":"https://gist.github.com/bgeesaman/0e0349e94cd22c48bf14d8a9b7d6b8f2"}],"ecosystem_specific":{"Symbols":["NewCustomResourceDefinitionHandler"],"URL":"https://go.googlesource.com/vulndb/+/refs/heads/main/reports/GO-2021-0062.toml"}}]`

// index for dbs containing some entries for each vuln package.
// The timestamp for package is set to random moment in the past.
var index string = `{
	"k8s.io/apiextensions-apiserver/pkg/apiserver": "2021-01-01T12:00:00.000000000-08:00",
	"golang.org/x/crypto/ssh": "2021-01-01T12:00:00.000000000-08:00",
	"github.com/go-yaml/yaml": "2021-01-01T12:00:00.000000000-08:00"
}`

var vulns = map[string]string{
	"github.com/go-yaml/yaml.json":        goYamlVuln,
	"golang.org/x/crypto.json":            cryptoSSHVuln,
	"k8s.io/apiextensions-apiserver.json": k8sAPIServerVuln,
}

// addToLocalDb adds vuln for package p to local db at path db.
func addToLocalDb(db, p, vuln string) error {
	if err := os.MkdirAll(path.Join(db, filepath.Dir(p)), fs.ModePerm); err != nil {
		return err
	}

	f, err := os.Create(path.Join(db, p))
	if err != nil {
		return err
	}
	defer f.Close()

	f.Write([]byte(vuln))
	return nil
}

// addToServerDb adds vuln for package p to localhost server identified by its handler.
func addToServerDb(handler *http.ServeMux, p, vuln string) {
	handler.HandleFunc("/"+p, func(w http.ResponseWriter, req *http.Request) { fmt.Fprint(w, vuln) })
}

// envUpdate updates an environment e by setting the key to value.
func envUpdate(e []string, key, value string) []string {
	var nenv []string
	for _, kv := range e {
		if strings.HasPrefix(kv, key+"=") {
			nenv = append(nenv, key+"="+value)
		} else {
			nenv = append(nenv, kv)
		}
	}
	return nenv
}

// cmd type encapsulating a shell command and its context.
type cmd struct {
	dir  string
	env  []string
	name string
	args []string
}

// execAll executes a sequence of commands cmd. Exits on a first
// encountered error returning the error and the accumulated output.
func execAll(cmds []cmd) ([]byte, error) {
	var out []byte
	for _, c := range cmds {
		o, err := execCmd(c.dir, c.env, c.name, c.args...)
		out = append(out, o...)
		if err != nil {
			return o, err
		}
	}
	return out, nil
}

// execCmd runs the command name with arg in dir location with the env environment.
func execCmd(dir string, env []string, name string, arg ...string) ([]byte, error) {
	cmd := exec.Command(name, arg...)
	cmd.Dir = dir
	cmd.Env = env
	return cmd.CombinedOutput()
}

// finding abstraction of Finding, for test purposes.
type finding struct {
	symbol   string
	traceLen int
}

func testFindings(finds []audit.Finding) []finding {
	var fs []finding
	for _, f := range finds {
		fs = append(fs, finding{symbol: f.Symbol, traceLen: len(f.Trace)})
	}
	return fs
}

func subset(finds1, finds2 []finding) bool {
	fs2 := make(map[finding]bool)
	for _, f := range finds2 {
		fs2[f] = true
	}

	for _, f := range finds1 {
		if !fs2[f] {
			return false
		}
	}
	return true
}

func allFindings(r *audit.Results) []audit.Finding {
	var findings []audit.Finding
	for _, v := range r.Vulnerabilities {
		for _, f := range r.VulnFindings[v.ID] {
			findings = append(findings, f)
		}
	}
	return findings
}

func TestHashicorpVault(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}
	e := packagestest.Export(t, packagestest.Modules, []packagestest.Module{
		{
			Name: "foo",
		},
	})
	defer e.Cleanup()

	hashiVaultOkta := "github.com/hashicorp/vault/builtin/credential/okta"

	// Go get hashicorp-vault okta package v1.6.3.
	env := envUpdate(e.Config.Env, "GOPROXY", "https://proxy.golang.org,direct")
	if out, err := execCmd(e.Config.Dir, env, "go", "get", hashiVaultOkta+"@v1.6.3"); err != nil {
		t.Logf("failed to get %s: %s", hashiVaultOkta+"@v1.6.3", out)
		t.Fatal(err)
	}

	// run goaudit.
	cfg := &packages.Config{
		Mode:  packages.LoadAllSyntax | packages.NeedModule,
		Tests: false,
		Dir:   e.Config.Dir,
	}

	// Create a local filesystem db.
	dbPath := path.Join(e.Config.Dir, "db")
	addToLocalDb(dbPath, "index.json", index)
	// Create a local server db.
	sMux := http.NewServeMux()
	s := http.Server{Addr: ":8080", Handler: sMux}
	go func() { s.ListenAndServe() }()
	defer func() { s.Shutdown(context.Background()) }()
	addToServerDb(sMux, "index.json", index)

	for _, test := range []struct {
		source string
		// list of packages whose vulns should be addded to source
		toAdd []string
		want  []finding
		// "" indicates no cache should be used
		cacheRoot string
	}{
		// test local db without yaml, which should result in no findings.
		{source: "file://" + dbPath, want: nil,
			toAdd: []string{"golang.org/x/crypto.json", "k8s.io/apiextensions-apiserver.json"}},
		// add yaml to the local db, which should produce 2 findings.
		{source: "file://" + dbPath, toAdd: []string{"github.com/go-yaml/yaml.json"},
			want: []finding{
				{"github.com/go-yaml/yaml.decoder.unmarshal", 6},
				{"github.com/go-yaml/yaml.yaml_parser_fetch_more_tokens", 12}},
		},
		// repeat the similar experiment with a server db.
		{source: "http://localhost:8080", toAdd: []string{"k8s.io/apiextensions-apiserver.json"}, want: nil},
		{source: "http://localhost:8080", toAdd: []string{"golang.org/x/crypto.json", "github.com/go-yaml/yaml.json"},
			want: []finding{
				{"github.com/go-yaml/yaml.decoder.unmarshal", 6},
				{"github.com/go-yaml/yaml.yaml_parser_fetch_more_tokens", 12}},
		},
		// repeat the last test but with cache
		{source: "http://localhost:8080", cacheRoot: filepath.Join(e.Config.Dir, "/pkg/mod/cache/download/vulndb"),
			want: []finding{
				{"github.com/go-yaml/yaml.decoder.unmarshal", 6},
				{"github.com/go-yaml/yaml.yaml_parser_fetch_more_tokens", 12}},
		},
	} {
		for _, add := range test.toAdd {
			if strings.HasPrefix(test.source, "file://") {
				addToLocalDb(dbPath, add, vulns[add])
			} else {
				addToServerDb(sMux, add, vulns[add])
			}
		}

		var opts client.Options
		if test.cacheRoot != "" {
			opts.HTTPCache = &fsCache{rootDir: test.cacheRoot}
		}
		dbClient, err := client.NewClient([]string{test.source}, opts)
		if err != nil {
			t.Error(err)
		}

		r, err := run(cfg, []string{hashiVaultOkta}, false, dbClient)
		if err != nil {
			t.Fatal(err)
		}
		if fs := testFindings(allFindings(r)); !subset(test.want, fs) {
			t.Errorf("want %v subset of findings; got %v", test.want, fs)
		}
	}
}

// isSecure checks if http resp was made over a secure connection.
func isSecure(resp *http.Response) bool {
	if resp.TLS == nil {
		return false
	}

	// Check the final URL scheme too for good measure.
	if resp.Request.URL.Scheme != "https" {
		return false
	}

	return true
}

// download fetches the content at url and stores it at destination location.
func download(url, destination string) error {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: false},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if !isSecure(resp) {
		return fmt.Errorf("insecure connection to %s", url)
	}

	out, err := os.Create(destination)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, resp.Body)
	return err
}

// TestKubernetes requires the following system dependencies:
//   - make, tar, unzip, and gcc.
// More information on installing kubernetes: https://github.com/kubernetes/kubernetes.
// Note that the whole installation will require roughly 5GB of disk.
func TestKubernetes(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	// make sure the dependencies are present
	if _, err := exec.LookPath("tar"); err != nil {
		t.Skip("tar needed for this test.")
	}
	if _, err := exec.LookPath("unzip"); err != nil {
		t.Skip("unzip needed for this test.")
	}

	e := packagestest.Export(t, packagestest.Modules, []packagestest.Module{
		{
			Name: "foo",
		},
	})
	defer e.Cleanup()

	// Environments and directories to build and download both k8s and go.
	env := envUpdate(e.Config.Env, "GOPROXY", "https://proxy.golang.org,direct")
	dir := e.Config.Dir
	k8sDir := path.Join(e.Config.Dir, "kubernetes-1.15.11")
	k8sEnv := envUpdate(env, "PATH", path.Join(e.Config.Dir, "go/bin")+":"+os.Getenv("PATH"))

	// Download kubernetes v1.15.11 and the go version 1.12 needed to build it.
	if err := download("https://github.com/kubernetes/kubernetes/archive/v1.15.11.zip", path.Join(dir, "v1.15.11")); err != nil {
		t.Fatal(err)
	}
	goZip := "go1.12.17." + runtime.GOOS + "-" + runtime.GOARCH + ".tar.gz"
	if err := download("https://golang.org/dl/"+goZip, path.Join(dir, goZip)); err != nil {
		t.Fatal(err)
	}

	// Unzip k8s and go, and then build the k8s.
	if out, err := execAll([]cmd{
		{dir, env, "unzip", []string{"v1.15.11"}},
		{dir, env, "tar", []string{"-xf", goZip}},
		{k8sDir, k8sEnv, "make", nil},
	}); err != nil {
		t.Logf("failed to build k8s: %s", out)
		t.Fatal(err)
	}

	// Create a local filesystem db.
	dbPath := path.Join(e.Config.Dir, "db")
	addToLocalDb(dbPath, "index.json", index)
	// Create a local server db.
	sMux := http.NewServeMux()
	s := http.Server{Addr: ":8080", Handler: sMux}
	go func() { s.ListenAndServe() }()
	defer func() { s.Shutdown(context.Background()) }()
	addToServerDb(sMux, "index.json", index)

	// run goaudit.
	cfg := &packages.Config{
		Mode:  packages.LoadAllSyntax | packages.NeedModule,
		Tests: false,
		Dir:   path.Join(e.Config.Dir, "kubernetes-1.15.11"),
	}

	for _, test := range []struct {
		source string
		// list of packages whose vulns should be addded to source
		toAdd []string
		want  []finding
		// "" indicates no cache should be used
		cacheRoot string
	}{
		// test local db with only apiserver vuln, which should result in a single finding.
		{source: "file://" + dbPath, toAdd: []string{"github.com/go-yaml/yaml.json", "k8s.io/apiextensions-apiserver/pkg/apiserver.json"},
			want: []finding{{"k8s.io/apiextensions-apiserver/pkg/apiserver.NewCustomResourceDefinitionHandler", 3}}},
		// add the rest of the vulnerabilites, resulting in more findings.
		{source: "file://" + dbPath, toAdd: []string{"golang.org/x/crypto/ssh.json"},
			want: []finding{
				{"golang.org/x/crypto/ssh.NewPublicKey", 1},
				{"k8s.io/apiextensions-apiserver/pkg/apiserver.NewCustomResourceDefinitionHandler", 3},
				{"golang.org/x/crypto/ssh.NewPublicKey", 4},
				{"golang.org/x/crypto/ssh.parseED25519", 9},
			}},
		// repeat similar experiment with a server db.
		{source: "http://localhost:8080", toAdd: []string{"github.com/go-yaml/yaml.json"}, want: nil},
		{source: "http://localhost:8080", toAdd: []string{"golang.org/x/crypto/ssh.json", "k8s.io/apiextensions-apiserver/pkg/apiserver.json"},
			want: []finding{
				{"golang.org/x/crypto/ssh.NewPublicKey", 1},
				{"k8s.io/apiextensions-apiserver/pkg/apiserver.NewCustomResourceDefinitionHandler", 3},
				{"golang.org/x/crypto/ssh.NewPublicKey", 4},
				{"golang.org/x/crypto/ssh.parseED25519", 9},
			}},
		//repeat the last test but with a cache
		{source: "http://localhost:8080", cacheRoot: filepath.Join(e.Config.Dir, "/pkg/mod/cache/download/vulndb"),
			want: []finding{
				{"golang.org/x/crypto/ssh.NewPublicKey", 1},
				{"k8s.io/apiextensions-apiserver/pkg/apiserver.NewCustomResourceDefinitionHandler", 3},
				{"golang.org/x/crypto/ssh.NewPublicKey", 4},
				{"golang.org/x/crypto/ssh.parseED25519", 9},
			}},
	} {
		for _, add := range test.toAdd {
			if strings.HasPrefix(test.source, "file://") {
				addToLocalDb(dbPath, add, vulns[add])
			} else {
				addToServerDb(sMux, add, vulns[add])
			}
		}

		var opts client.Options
		if test.cacheRoot != "" {
			opts.HTTPCache = &fsCache{rootDir: test.cacheRoot}
		}
		dbClient, err := client.NewClient([]string{test.source}, opts)
		if err != nil {
			t.Error(err)
		}

		r, err := run(cfg, []string{"./..."}, false, dbClient)
		if err != nil {
			t.Fatal(err)
		}
		if fs := testFindings(allFindings(r)); !subset(test.want, fs) {
			t.Errorf("want %v subset of findings; got %v", test.want, fs)
		}
	}
}
<commit_msg>vulndb/govulncheck: remove temporarily integration tests<commit_after><|endoftext|>


<commit_before>/* Copyright (C) 2014 Pivotal Software, Inc.

All rights reserved. This program and the accompanying materials
are made available under the terms of the 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 levo

import (
	"bitbucket.org/pkg/inflect"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"reflect"
	"regexp"
	"strings"
	"text/template"
)

func Lower(input string) string {
	fmt.Print("")
	return strings.ToLower(input)
}

func Upper(input string) string {
	return strings.ToUpper(input)
}

func TestEquality(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func TestInequality(args ...interface{}) bool {
	notEqual := TestEquality(args...)
	return !notEqual
}

func Snakecase(input string) string {
	words := splitByWord(input)
	joined := strings.Join(words, "_")
	return strings.ToLower(joined)
}

func Titlecase(input string) string {
	words := splitByWord(input)
	for index, word := range words {
		if word == "" {
			continue
		}
		words[index] = strings.ToUpper(word[0:1]) + word[1:]
	}
	return strings.Join(words, "")
}

func Camelcase(input string) string {
	if input == "" {
		return ""
	}
	words := splitByWord(input)
	for index, word := range words {
		if word == "" {
			continue
		}
		words[index] = strings.ToUpper(word[0:1]) + word[1:]
	}
	words[0] = strings.ToLower(words[0])
	return strings.Join(words, "")
}

func splitByWord(input string) []string {
	words := splitBySep([]string{input})
	words = splitByLowerToUpper(words)
	words = splitByDoubleUpperToLower(words)
	words = splitByNumber(words)
	output := make([]string, 0)
	for _, word := range words {
		if word != "" {
			output = append(output, word)
		}
	}
	return output
}

func splitBySep(inputWords []string) (outputWords []string) {
	re := regexp.MustCompile("[_ !]")
	outputWords = make([]string, 0)
	for _, inputWord := range inputWords {
		words := re.Split(inputWord, -1)
		for _, word := range words {
			outputWords = append(outputWords, word)
		}
	}
	return
}

func splitByLowerToUpper(inputWords []string) (outputWords []string) {
	re := regexp.MustCompile("[a-z][A-Z]")
	outputWords = make([]string, 0)
	for _, inputWord := range inputWords {
		indexes := re.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 1
			outputWords = append(outputWords, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			outputWords = append(outputWords, inputWord[splitAtIndex:])
		} else {
			outputWords = append(outputWords, inputWord)
		}
	}
	return
}

func splitByDoubleUpperToLower(inputWords []string) (outputWords []string) {
	re := regexp.MustCompile("[A-Z][A-Z][a-z]")
	outputWords = make([]string, 0)
	for _, inputWord := range inputWords {
		indexes := re.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 2
			outputWords = append(outputWords, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			outputWords = append(outputWords, inputWord[splitAtIndex:])
		} else {
			outputWords = append(outputWords, inputWord)
		}
	}
	return
}

func splitByNumber(inputWords []string) (outputWords []string) {
	re1 := regexp.MustCompile("([^0-9][0-9])")
	re2 := regexp.MustCompile("([0-9][^0-9])")
	intermediate := make([]string, 0)
	for _, inputWord := range inputWords {
		indexes := re1.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 1
			intermediate = append(intermediate, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			intermediate = append(intermediate, inputWord[splitAtIndex:])
		} else {
			intermediate = append(intermediate, inputWord)
		}
	}
	outputWords = make([]string, 0)
	for _, inputWord := range intermediate {
		indexes := re2.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 1
			outputWords = append(outputWords, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			outputWords = append(outputWords, inputWord[splitAtIndex:])
		} else {
			outputWords = append(outputWords, inputWord)
		}
	}
	return
}

func Pluralize(input string) string {
	pluralizedString := Snakecase(input)
	pluralizedString = inflect.Pluralize(pluralizedString)
	iIndex := 0
	pIndex := 0
	re := regexp.MustCompile("[A-z0-9]")
	output := ""
	for iIndex < len(input) && pIndex < len(pluralizedString) {
		inputChar := string(input[iIndex])
		pluralChar := string(pluralizedString[pIndex])
		if strings.ToLower(inputChar) == strings.ToLower(pluralChar) {
			output = output + inputChar
			iIndex++
			pIndex++
		} else if pluralChar == "_" {
			pIndex++
		} else if re.MatchString(inputChar) == false {
			output = output + inputChar
			iIndex++
		} else if re.MatchString(pluralChar) {
			output = output + pluralChar
			iIndex++
			pIndex++
		}
	}
	if pIndex <= len(pluralizedString)-1 {
		output = output + pluralizedString[pIndex:]
	}
	return output
}

func PackageToPath(input string) string {
	return strings.Replace(input, ".", "/", -1)
}

func HasListType(model Model) bool {
	for _, property := range model.Properties {
		if property.IsSetType {
			return true
		}
	}
	return false
}

func IsSqliteType(prop ModelProperty) bool {
	_, ok := SqliteTypes()[strings.ToLower(prop.PropertyType)]
	return ok && !prop.IsSetType
}

func ToSqliteType(prop ModelProperty) string {
	if value, ok := SqliteTypes()[strings.ToLower(prop.PropertyType)]; ok {
		return value
	} else {
		return prop.PropertyType
	}
}

func IsJavaType(prop ModelProperty) bool {
	_, ok := JavaTypes()[strings.ToLower(prop.PropertyType)]
	return ok
}

func ToJavaType(prop ModelProperty) string {
	theType := ""
	if javaType, ok := JavaTypes()[strings.ToLower(prop.PropertyType)]; ok {
		theType = javaType
	} else {
		theType = prop.PropertyType
	}

	if prop.IsSetType {
		return "List<" + theType + ">"
	}
	return theType
}

func ToCoreDataType(input string) string {
	return CoreDataTypes()[strings.ToLower(input)]
}

func ToRailsType(prop ModelProperty) string {
	theType := ""
	if railsType, ok := RailsTypes()[strings.ToLower(prop.PropertyType)]; ok {
		theType = railsType
	} else {
		theType = prop.PropertyType
	}

	if prop.IsSetType {
		return "Array"
	}
	return theType
}

func SHA256(data string) string {
	hashWriter := sha1.New()
	io.WriteString(hashWriter, data)
	return hex.EncodeToString(hashWriter.Sum(nil))
}

func Prefix(prefix string, original string) string {
	return Concat(original, prefix)
}

func Suffix(suffix string, original string) string {
	return Concat(original, suffix)
}

func Concat(first string, second string) string {
	return first + second
}

func Truncate(targetLength int, tooLong string) string {
	return string([]byte(tooLong)[0:targetLength])
}

func IdProp(properties []ModelProperty) string {
	for i := 0; i < len(properties); i++ {
		if strings.Contains(properties[i].RemoteIdentifier, "id") {
			underscored := Snakecase(properties[i].RemoteIdentifier)
			return Upper(underscored)
		}
	}
	if len(properties) > 0 {
		underscored := Snakecase(properties[0].RemoteIdentifier)
		return Upper(underscored)
	} else {
		return "UNKNOWN_PROPERTY"
	}
}

func SqliteTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "INTEGER"
	dict["integer"] = "INTEGER"
	dict["short"] = "INTEGER"
	dict["long"] = "INTEGER"
	dict["float"] = "REAL"
	dict["boolean"] = "INTEGER"
	dict["char"] = "TEXT"
	dict["character"] = "TEXT"
	dict["string"] = "TEXT"
	dict["byte"] = "TEXT"
	dict["void"] = "NONE"
	return dict
}

func JavaTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "int"
	dict["integer"] = "int"
	dict["short"] = "short"
	dict["long"] = "long"
	dict["float"] = "float"
	dict["boolean"] = "boolean"
	dict["char"] = "char"
	dict["character"] = "char"
	dict["string"] = "String"
	dict["byte"] = "byte"
	dict["void"] = "NONE"
	return dict
}

func CoreDataTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "Integer 32"
	dict["Integer"] = "Integer 32"
	dict["short"] = "Integer 16"
	dict["short"] = "Integer 16"
	dict["long"] = "Integer 64"
	dict["Long"] = "Integer 64"
	dict["float"] = "Float"
	dict["Float"] = "Float"
	dict["boolean"] = "Boolean"
	dict["Boolean"] = "Boolean"
	dict["char"] = "String"
	dict["Character"] = "String"
	dict["String"] = "String"
	dict["string"] = "String"
	dict["byte"] = "NONE"
	dict["Byte"] = "NONE"
	dict["void"] = "NONE"
	dict["Void"] = "NONE"
	return dict
}

func RailsTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "integer"
	dict["Integer"] = "integer"
	dict["short"] = "integer"
	dict["short"] = "integer"
	dict["long"] = "integer"
	dict["Long"] = "integer"
	dict["float"] = "float"
	dict["Float"] = "float"
	dict["boolean"] = "boolean"
	dict["Boolean"] = "boolean"
	dict["char"] = "string"
	dict["Character"] = "string"
	dict["String"] = "string"
	dict["string"] = "string"
	dict["byte"] = "NONE"
	dict["Byte"] = "NONE"
	dict["void"] = "NONE"
	dict["Void"] = "NONE"
	return dict
}

func addCommonUtilitiesToTemplate(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"eq":        TestEquality,
		"neq":       TestInequality,
		"lower":     Lower,
		"upper":     Upper,
		"pluralize": Pluralize,
		"camelcase": Camelcase,
		"titlecase": Titlecase,
		"snakecase": Snakecase,
		"SHA256":    SHA256,
		"concat":    Concat,
		"truncate":  Truncate,
	})
	return templateObject
}

func addJavaUtilitiesToTemplate(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"hasListType":   HasListType,
		"isSqliteType":  IsSqliteType,
		"toSqliteType":  ToSqliteType,
		"sqliteType":    SqliteTypes,
		"isJavaType":    IsJavaType,
		"toJavaType":    ToJavaType,
		"javaType":      JavaTypes,
		"idProp":        IdProp,
		"packageToPath": PackageToPath,
	})
	return templateObject
}

func addObjectiveCUtilitiesToTempalte(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"toCoreDataType": ToCoreDataType,
	})
	return templateObject
}

func addRailsUitilitiesToTemplate(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"toRailsType": ToRailsType,
	})
	return templateObject
}
<commit_msg>Updated the Objective C template dependancies<commit_after>/* Copyright (C) 2014 Pivotal Software, Inc.

All rights reserved. This program and the accompanying materials
are made available under the terms of the 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 levo

import (
	"bitbucket.org/pkg/inflect"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"reflect"
	"regexp"
	"strings"
	"text/template"
)

func Lower(input string) string {
	fmt.Print("")
	return strings.ToLower(input)
}

func Upper(input string) string {
	return strings.ToUpper(input)
}

func TestEquality(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func TestInequality(args ...interface{}) bool {
	notEqual := TestEquality(args...)
	return !notEqual
}

func Snakecase(input string) string {
	words := splitByWord(input)
	joined := strings.Join(words, "_")
	return strings.ToLower(joined)
}

func Titlecase(input string) string {
	words := splitByWord(input)
	for index, word := range words {
		if word == "" {
			continue
		}
		words[index] = strings.ToUpper(word[0:1]) + word[1:]
	}
	return strings.Join(words, "")
}

func Camelcase(input string) string {
	if input == "" {
		return ""
	}
	words := splitByWord(input)
	for index, word := range words {
		if word == "" {
			continue
		}
		words[index] = strings.ToUpper(word[0:1]) + word[1:]
	}
	words[0] = strings.ToLower(words[0])
	return strings.Join(words, "")
}

func splitByWord(input string) []string {
	words := splitBySep([]string{input})
	words = splitByLowerToUpper(words)
	words = splitByDoubleUpperToLower(words)
	words = splitByNumber(words)
	output := make([]string, 0)
	for _, word := range words {
		if word != "" {
			output = append(output, word)
		}
	}
	return output
}

func splitBySep(inputWords []string) (outputWords []string) {
	re := regexp.MustCompile("[_ !]")
	outputWords = make([]string, 0)
	for _, inputWord := range inputWords {
		words := re.Split(inputWord, -1)
		for _, word := range words {
			outputWords = append(outputWords, word)
		}
	}
	return
}

func splitByLowerToUpper(inputWords []string) (outputWords []string) {
	re := regexp.MustCompile("[a-z][A-Z]")
	outputWords = make([]string, 0)
	for _, inputWord := range inputWords {
		indexes := re.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 1
			outputWords = append(outputWords, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			outputWords = append(outputWords, inputWord[splitAtIndex:])
		} else {
			outputWords = append(outputWords, inputWord)
		}
	}
	return
}

func splitByDoubleUpperToLower(inputWords []string) (outputWords []string) {
	re := regexp.MustCompile("[A-Z][A-Z][a-z]")
	outputWords = make([]string, 0)
	for _, inputWord := range inputWords {
		indexes := re.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 2
			outputWords = append(outputWords, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			outputWords = append(outputWords, inputWord[splitAtIndex:])
		} else {
			outputWords = append(outputWords, inputWord)
		}
	}
	return
}

func splitByNumber(inputWords []string) (outputWords []string) {
	re1 := regexp.MustCompile("([^0-9][0-9])")
	re2 := regexp.MustCompile("([0-9][^0-9])")
	intermediate := make([]string, 0)
	for _, inputWord := range inputWords {
		indexes := re1.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 1
			intermediate = append(intermediate, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			intermediate = append(intermediate, inputWord[splitAtIndex:])
		} else {
			intermediate = append(intermediate, inputWord)
		}
	}
	outputWords = make([]string, 0)
	for _, inputWord := range intermediate {
		indexes := re2.FindAllIndex([]byte(inputWord), -1)
		var splitAtIndex int = -1
		var prevIndex int = 0
		for _, match := range indexes {
			splitAtIndex = match[1] - 1
			outputWords = append(outputWords, inputWord[prevIndex:splitAtIndex])
			prevIndex = splitAtIndex
		}
		if splitAtIndex > -1 {
			outputWords = append(outputWords, inputWord[splitAtIndex:])
		} else {
			outputWords = append(outputWords, inputWord)
		}
	}
	return
}

func Pluralize(input string) string {
	pluralizedString := Snakecase(input)
	pluralizedString = inflect.Pluralize(pluralizedString)
	iIndex := 0
	pIndex := 0
	re := regexp.MustCompile("[A-z0-9]")
	output := ""
	for iIndex < len(input) && pIndex < len(pluralizedString) {
		inputChar := string(input[iIndex])
		pluralChar := string(pluralizedString[pIndex])
		if strings.ToLower(inputChar) == strings.ToLower(pluralChar) {
			output = output + inputChar
			iIndex++
			pIndex++
		} else if pluralChar == "_" {
			pIndex++
		} else if re.MatchString(inputChar) == false {
			output = output + inputChar
			iIndex++
		} else if re.MatchString(pluralChar) {
			output = output + pluralChar
			iIndex++
			pIndex++
		}
	}
	if pIndex <= len(pluralizedString)-1 {
		output = output + pluralizedString[pIndex:]
	}
	return output
}

func PackageToPath(input string) string {
	return strings.Replace(input, ".", "/", -1)
}

func HasListType(model Model) bool {
	for _, property := range model.Properties {
		if property.IsSetType {
			return true
		}
	}
	return false
}

func IsSqliteType(prop ModelProperty) bool {
	_, ok := SqliteTypes()[strings.ToLower(prop.PropertyType)]
	return ok && !prop.IsSetType
}

func ToSqliteType(prop ModelProperty) string {
	if value, ok := SqliteTypes()[strings.ToLower(prop.PropertyType)]; ok {
		return value
	} else {
		return prop.PropertyType
	}
}

func IsJavaType(prop ModelProperty) bool {
	_, ok := JavaTypes()[strings.ToLower(prop.PropertyType)]
	return ok
}

func ToJavaType(prop ModelProperty) string {
	theType := ""
	if javaType, ok := JavaTypes()[strings.ToLower(prop.PropertyType)]; ok {
		theType = javaType
	} else {
		theType = prop.PropertyType
	}

	if prop.IsSetType {
		return "List<" + theType + ">"
	}
	return theType
}

func ToCoreDataType(input string) string {
	return CoreDataTypes()[strings.ToLower(input)]
}

func ToObjectiveCType(input string) string {
	return ObjectiveCTypes()[strings.ToLower(input)]
}

func ToRailsType(prop ModelProperty) string {
	theType := ""
	if railsType, ok := RailsTypes()[strings.ToLower(prop.PropertyType)]; ok {
		theType = railsType
	} else {
		theType = prop.PropertyType
	}

	if prop.IsSetType {
		return "Array"
	}
	return theType
}

func SHA256(data string) string {
	hashWriter := sha1.New()
	io.WriteString(hashWriter, data)
	return hex.EncodeToString(hashWriter.Sum(nil))
}

func Prefix(prefix string, original string) string {
	return Concat(original, prefix)
}

func Suffix(suffix string, original string) string {
	return Concat(original, suffix)
}

func Concat(first string, second string) string {
	return first + second
}

func Truncate(targetLength int, tooLong string) string {
	return string([]byte(tooLong)[0:targetLength])
}

func IdProp(properties []ModelProperty) string {
	for i := 0; i < len(properties); i++ {
		if strings.Contains(properties[i].RemoteIdentifier, "id") {
			underscored := Snakecase(properties[i].RemoteIdentifier)
			return Upper(underscored)
		}
	}
	if len(properties) > 0 {
		underscored := Snakecase(properties[0].RemoteIdentifier)
		return Upper(underscored)
	} else {
		return "UNKNOWN_PROPERTY"
	}
}

func SqliteTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "INTEGER"
	dict["integer"] = "INTEGER"
	dict["short"] = "INTEGER"
	dict["long"] = "INTEGER"
	dict["float"] = "REAL"
	dict["boolean"] = "INTEGER"
	dict["char"] = "TEXT"
	dict["character"] = "TEXT"
	dict["string"] = "TEXT"
	dict["byte"] = "TEXT"
	dict["void"] = "NONE"
	return dict
}

func JavaTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "int"
	dict["integer"] = "int"
	dict["short"] = "short"
	dict["long"] = "long"
	dict["float"] = "float"
	dict["boolean"] = "boolean"
	dict["char"] = "char"
	dict["character"] = "char"
	dict["string"] = "String"
	dict["byte"] = "byte"
	dict["void"] = "NONE"
	return dict
}

func CoreDataTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "Integer 32"
	dict["integer"] = "Integer 32"
	dict["short"] = "Integer 16"
	dict["long"] = "Integer 64"
	dict["float"] = "Float"
	dict["boolean"] = "Boolean"
	dict["char"] = "String"
	dict["character"] = "String"
	dict["string"] = "String"
	dict["date"] = "Date"
	dict["byte"] = "NONE"
	dict["void"] = "NONE"
	return dict
}

func ObjectiveCTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "NSNumber"
	dict["integer"] = "NSNumber"
	dict["short"] = "NSNumber"
	dict["long"] = "NSNumber"
	dict["float"] = "NSNumber"
	dict["boolean"] = "NSNumber"
	dict["char"] = "NSString"
	dict["character"] = "NSString"
	dict["string"] = "NSString"
	dict["date"] = "NSDate"
	dict["byte"] = "NONE"
	dict["void"] = "NONE"
	return dict
}

func RailsTypes() map[string]string {
	var dict = make(map[string]string)
	dict["int"] = "integer"
	dict["Integer"] = "integer"
	dict["short"] = "integer"
	dict["short"] = "integer"
	dict["long"] = "integer"
	dict["Long"] = "integer"
	dict["float"] = "float"
	dict["Float"] = "float"
	dict["boolean"] = "boolean"
	dict["Boolean"] = "boolean"
	dict["char"] = "string"
	dict["Character"] = "string"
	dict["String"] = "string"
	dict["string"] = "string"
	dict["byte"] = "NONE"
	dict["Byte"] = "NONE"
	dict["void"] = "NONE"
	dict["Void"] = "NONE"
	return dict
}

func addCommonUtilitiesToTemplate(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"eq":        TestEquality,
		"neq":       TestInequality,
		"lower":     Lower,
		"upper":     Upper,
		"pluralize": Pluralize,
		"camelcase": Camelcase,
		"titlecase": Titlecase,
		"snakecase": Snakecase,
		"SHA256":    SHA256,
		"concat":    Concat,
		"truncate":  Truncate,
	})
	return templateObject
}

func addJavaUtilitiesToTemplate(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"hasListType":   HasListType,
		"isSqliteType":  IsSqliteType,
		"toSqliteType":  ToSqliteType,
		"sqliteType":    SqliteTypes,
		"isJavaType":    IsJavaType,
		"toJavaType":    ToJavaType,
		"javaType":      JavaTypes,
		"idProp":        IdProp,
		"packageToPath": PackageToPath,
	})
	return templateObject
}

func addObjectiveCUtilitiesToTempalte(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"toCoreDataType": ToCoreDataType,
		"toObjectiveCType" : ToObjectiveCType,
	})
	return templateObject
}

func addRailsUitilitiesToTemplate(templateObject *template.Template) *template.Template {
	templateObject = templateObject.Funcs(template.FuncMap{
		"toRailsType": ToRailsType,
	})
	return templateObject
}
<|endoftext|>


<commit_before>package redlot

import (
	"errors"
	"fmt"
	"strconv"
	"time"

	"github.com/syndtr/goleveldb/leveldb/util"
)

func encodeKvKey(key []byte) (buf []byte) {
	buf = append(buf, typeKV)
	buf = append(buf, key...)
	return
}

func decodeKvKey(buf []byte) (key []byte) {
	if len(buf) < 4 {
		return nil
	}
	return buf[1:]
}

func Get(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return "", errNosArgs
	}

	v, err := db.Get(encodeKvKey(args[0]), nil)
	return string(v), err
}

func Set(args [][]byte) (interface{}, error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	fmt.Printf("SET %s %s\n", args[0], args[1])
	return nil, db.Put(encodeKvKey(args[0]), args[1], nil)
}

func Del(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	return nil, db.Delete(encodeKvKey(args[0]), nil)
}

func Exists(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return int64(-1), errNosArgs
	}

	ret, err := db.Has(encodeKvKey(args[0]), nil)
	if ret {
		return int64(1), err
	}
	return int64(0), err
}

func Expire(args [][]byte) (interface{}, error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	key := encodeKvKey(args[0])
	has, _ := db.Has(key, nil)
	if has {
		ttl := strToInt64(string(args[1]))
		if ttl < 1 {
			return nil, errors.New("TTL must > 0, you set to " + string(args[1]))
		}
		bs := uint64ToBytes(uint64(time.Now().UTC().Unix() + ttl))
		if meta.Put(key, bs, nil) == nil {
			return int64(1), nil
		}
	}
	return int64(0), nil
}

func Setx(args [][]byte) (interface{}, error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	key := encodeKvKey(args[0])

	ttl := strToInt64(string(args[2]))
	if ttl < 1 {
		return nil, errors.New("TTL must > 0, you set to " + string(args[2]))
	}
	bs := uint64ToBytes(uint64(time.Now().UTC().Unix() + ttl))
	meta.Put(key, bs, nil)

	return nil, db.Put(key, args[1], nil)
}

func TTL(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return int64(-1), errNosArgs
	}

	key := encodeKvKey(args[0])
	b, _ := meta.Get(key, nil)
	if len(b) < 1 {
		return int64(-1), nil
	}
	ttl := int64(bytesToUint64(b)) - time.Now().UTC().Unix()
	if ttl < 0 {
		ttl = -1
		meta.Delete(key, nil)
		db.Delete(key, nil)
	}
	return ttl, nil
}

func Keys(args [][]byte) ([]string, error) {
	if len(args) < 3 {
		return []string{}, errNosArgs
	}

	ks := encodeKvKey(args[0])
	ke := encodeKvKey(args[1])
	limit, _ := strconv.Atoi(string(args[2]))

	var keys []string
	iter := db.NewIterator(&util.Range{Start: ks, Limit: ke}, nil)
	for iter.Next() {
		k := decodeKvKey(iter.Key())
		keys = append(keys, string(k))
		limit--
		if limit == 0 {
			break
		}
	}
	iter.Release()
	err := iter.Error()
	return keys, err
}

func Scan(args [][]byte) ([]string, error) {
	if len(args) < 3 {
		return []string{}, errNosArgs
	}

	ks := encodeKvKey(args[0])
	ke := encodeKvKey(args[1])
	limit, _ := strconv.Atoi(string(args[2]))

	var ret []string
	iter := db.NewIterator(&util.Range{Start: ks, Limit: ke}, nil)
	for iter.Next() {
		k := decodeKvKey(iter.Key())
		ret = append(ret, string(k))
		ret = append(ret, string(iter.Value()))
		limit--
		if limit == 0 {
			break
		}
	}
	iter.Release()
	err := iter.Error()
	return ret, err
}
<commit_msg>Comment KV type functions.<commit_after>package redlot

import (
	"errors"
	"fmt"
	"strconv"
	"time"

	"github.com/syndtr/goleveldb/leveldb/util"
)

func encodeKvKey(key []byte) (buf []byte) {
	buf = append(buf, typeKV)
	buf = append(buf, key...)
	return
}

func decodeKvKey(buf []byte) (key []byte) {
	if len(buf) < 4 {
		return nil
	}
	return buf[1:]
}

// Get a value by a key.
// Args: key string
func Get(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return "", errNosArgs
	}

	v, err := db.Get(encodeKvKey(args[0]), nil)
	return string(v), err
}

// Set a value by a key.
// Args: key string, value any
func Set(args [][]byte) (interface{}, error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	fmt.Printf("SET %s %s\n", args[0], args[1])
	return nil, db.Put(encodeKvKey(args[0]), args[1], nil)
}

// Del will delete a value by a key.
// Args: key string
func Del(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	return nil, db.Delete(encodeKvKey(args[0]), nil)
}

// Exists will check key is exists.
// Args: key string
func Exists(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return int64(-1), errNosArgs
	}

	ret, err := db.Has(encodeKvKey(args[0]), nil)
	if ret {
		return int64(1), err
	}
	return int64(0), err
}

// Expire the key after timeout.
// Args: key string, seconds int
func Expire(args [][]byte) (interface{}, error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	key := encodeKvKey(args[0])
	has, _ := db.Has(key, nil)
	if has {
		ttl := strToInt64(string(args[1]))
		if ttl < 1 {
			return nil, errors.New("TTL must > 0, you set to " + string(args[1]))
		}
		bs := uint64ToBytes(uint64(time.Now().UTC().Unix() + ttl))
		if meta.Put(key, bs, nil) == nil {
			return int64(1), nil
		}
	}
	return int64(0), nil
}

// Setx will set a value by the key and expire it after timeout.
// Args: key string, value any, seconds int
func Setx(args [][]byte) (interface{}, error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	key := encodeKvKey(args[0])

	ttl := strToInt64(string(args[2]))
	if ttl < 1 {
		return nil, errors.New("TTL must > 0, you set to " + string(args[2]))
	}
	bs := uint64ToBytes(uint64(time.Now().UTC().Unix() + ttl))
	meta.Put(key, bs, nil)

	return nil, db.Put(key, args[1], nil)
}

// TTL will return the lifetime of the key.
// Args: key string
func TTL(args [][]byte) (interface{}, error) {
	if len(args) < 1 {
		return int64(-1), errNosArgs
	}

	key := encodeKvKey(args[0])
	b, _ := meta.Get(key, nil)
	if len(b) < 1 {
		return int64(-1), nil
	}
	ttl := int64(bytesToUint64(b)) - time.Now().UTC().Unix()
	if ttl < 0 {
		ttl = -1
		meta.Delete(key, nil)
		db.Delete(key, nil)
	}
	return ttl, nil
}

// Keys will list keys in the range.
// Args: start_key string, end_key string, limit_number int
func Keys(args [][]byte) ([]string, error) {
	if len(args) < 3 {
		return []string{}, errNosArgs
	}

	ks := encodeKvKey(args[0])
	ke := encodeKvKey(args[1])
	limit, _ := strconv.Atoi(string(args[2]))

	var keys []string
	iter := db.NewIterator(&util.Range{Start: ks, Limit: ke}, nil)
	for iter.Next() {
		k := decodeKvKey(iter.Key())
		keys = append(keys, string(k))
		limit--
		if limit == 0 {
			break
		}
	}
	iter.Release()
	err := iter.Error()
	return keys, err
}

// Scan will list KV pair that keys in the range.
// Args: start_key string, end_key string, limit_number int
func Scan(args [][]byte) ([]string, error) {
	if len(args) < 3 {
		return []string{}, errNosArgs
	}

	ks := encodeKvKey(args[0])
	ke := encodeKvKey(args[1])
	limit, _ := strconv.Atoi(string(args[2]))

	var ret []string
	iter := db.NewIterator(&util.Range{Start: ks, Limit: ke}, nil)
	for iter.Next() {
		k := decodeKvKey(iter.Key())
		ret = append(ret, string(k))
		ret = append(ret, string(iter.Value()))
		limit--
		if limit == 0 {
			break
		}
	}
	iter.Release()
	err := iter.Error()
	return ret, err
}
<|endoftext|>


<commit_before><commit_msg>option: Use UseAllAttribute() in UseAllAttribute()<commit_after><|endoftext|>


<commit_before>package beater

import (
	"flag"
	"strings"
	"time"

	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/phenomenes/vago"
)

type Varnishbeat struct {
	alive   bool
	client  publisher.Client
	done    chan uint
	period  time.Duration
	varnish *vago.Varnish
}

var logFlag, statsFlag bool

func init() {
	flag.BoolVar(&logFlag, "log", false, "Read data from varnishlog")
	flag.BoolVar(&statsFlag, "stats", false, "Read data from varnishstat")
}

func New() *Varnishbeat {
	return &Varnishbeat{}
}

func (vb *Varnishbeat) HandleFlags(b *beat.Beat) error {
	return nil
}

func (vb *Varnishbeat) Config(b *beat.Beat) error {
	return nil
}

func (vb *Varnishbeat) Setup(b *beat.Beat) error {
	var err error
	vb.varnish, err = vago.Open("")
	if err != nil {
		logp.Err("%s", err)
		return err
	}
	vb.client = b.Publisher.Connect()
	return nil
}

func (vb *Varnishbeat) Run(b *beat.Beat) error {
	vb.alive = true
	var err error
	if logFlag {
		err := vb.exportLog()
		if err != nil {
			logp.Err("%s", err)
		}
	} else {
		vb.period = 10 * time.Second
		ticker := time.NewTicker(vb.period)
		defer ticker.Stop()

		for vb.alive {
			select {
			case <-vb.done:
				return nil
			case <-ticker.C:
			}
			timerStart := time.Now()
			event, err := vb.exportStats()
			if err != nil {
				logp.Err("%s", err)
				break
			}
			vb.client.PublishEvent(event)
			timerEnd := time.Now()
			duration := timerEnd.Sub(timerStart)
			if duration.Nanoseconds() > vb.period.Nanoseconds() {
				logp.Warn("Ignoring tick(s) due to processing taking longer than one period")
			}
		}
	}
	return err
}

func (vb *Varnishbeat) exportStats() (common.MapStr, error) {
	stats := make(common.MapStr)
	for k, v := range vb.varnish.Stats() {
		k1 := strings.Replace(k, ".", "_", -1)
		stats[k1] = v
	}
	event := common.MapStr{
		"@timestamp": common.Time(time.Now()),
		"type":       "stats",
		"stats":      stats,
	}
	return event, nil
}

func (vb *Varnishbeat) exportLog() error {
	tx := make(common.MapStr)
	vb.varnish.Log("", vago.REQ, func(vxid uint32, tag, _type, data string) int {
		if vb.alive == false {
			return -1
		}
		switch _type {
		case "c":
			_type = "client"
		case "b":
			_type = "backend"
		default:
			_type = "ping"
		}
		switch tag {
		case "BereqHeader", "BerespHeader", "ObjHeader", "ReqHeader", "RespHeader":
			header := strings.SplitN(data, ": ", 2)
			k := header[0]
			v := header[1]
			if _, ok := tx[tag]; ok {
				tx[tag].(common.MapStr)[k] = v
			} else {
				tx[tag] = common.MapStr{k: v}
			}
		case "End":
			event := common.MapStr{
				"@timestamp": common.Time(time.Now()),
				"count":      1,
				"type":       _type,
				"vxid":       vxid,
				"tx":         tx,
			}
			vb.client.PublishEvent(event)
			// destroy and re-create the map
			tx = nil
			tx = make(common.MapStr)
		default:
			tx[tag] = data
		}
		return 0
	})
	return nil
}

func (vb *Varnishbeat) Cleanup(b *beat.Beat) error {
	vb.varnish.Close()
	return nil
}

func (vb *Varnishbeat) Stop() {
	vb.alive = false
}
<commit_msg>Properly handle sigint with vago.Stop()<commit_after>package beater

import (
	"flag"
	"strings"
	"time"

	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/phenomenes/vago"
)

type Varnishbeat struct {
	client  publisher.Client
	done    chan uint
	period  time.Duration
	varnish *vago.Varnish
}

var logFlag, statsFlag bool

func init() {
	flag.BoolVar(&logFlag, "log", false, "Read data from varnishlog")
	flag.BoolVar(&statsFlag, "stats", false, "Read data from varnishstat")
}

func New() *Varnishbeat {
	return &Varnishbeat{}
}

func (vb *Varnishbeat) HandleFlags(b *beat.Beat) error {
	return nil
}

func (vb *Varnishbeat) Config(b *beat.Beat) error {
	return nil
}

func (vb *Varnishbeat) Setup(b *beat.Beat) error {
	var err error
	vb.varnish, err = vago.Open("")
	if err != nil {
		logp.Err("%s", err)
		return err
	}
	vb.client = b.Publisher.Connect()
	return nil
}

func (vb *Varnishbeat) Run(b *beat.Beat) error {
	var err error
	if logFlag {
		err := vb.exportLog()
		if err != nil {
			logp.Err("%s", err)
		}
	} else {
		vb.period = 10 * time.Second
		ticker := time.NewTicker(vb.period)
		defer ticker.Stop()

		for {
			select {
			case <-vb.done:
				return nil
			case <-ticker.C:
			}
			timerStart := time.Now()
			event, err := vb.exportStats()
			if err != nil {
				logp.Err("%s", err)
				break
			}
			vb.client.PublishEvent(event)
			timerEnd := time.Now()
			duration := timerEnd.Sub(timerStart)
			if duration.Nanoseconds() > vb.period.Nanoseconds() {
				logp.Warn("Ignoring tick(s) due to processing taking longer than one period")
			}
		}
	}
	return err
}

func (vb *Varnishbeat) exportStats() (common.MapStr, error) {
	stats := make(common.MapStr)
	for k, v := range vb.varnish.Stats() {
		k1 := strings.Replace(k, ".", "_", -1)
		stats[k1] = v
	}
	event := common.MapStr{
		"@timestamp": common.Time(time.Now()),
		"type":       "stats",
		"stats":      stats,
	}
	return event, nil
}

func (vb *Varnishbeat) exportLog() error {
	tx := make(common.MapStr)
	vb.varnish.Log("", vago.REQ, func(vxid uint32, tag, _type, data string) int {
		switch _type {
		case "c":
			_type = "client"
		case "b":
			_type = "backend"
		default:
			_type = "ping"
		}
		switch tag {
		case "BereqHeader", "BerespHeader", "ObjHeader", "ReqHeader", "RespHeader":
			header := strings.SplitN(data, ": ", 2)
			k := header[0]
			v := header[1]
			if _, ok := tx[tag]; ok {
				tx[tag].(common.MapStr)[k] = v
			} else {
				tx[tag] = common.MapStr{k: v}
			}
		case "End":
			event := common.MapStr{
				"@timestamp": common.Time(time.Now()),
				"count":      1,
				"type":       _type,
				"vxid":       vxid,
				"tx":         tx,
			}
			vb.client.PublishEvent(event)
			// destroy and re-create the map
			tx = nil
			tx = make(common.MapStr)
		default:
			tx[tag] = data
		}
		return 0
	})
	return nil
}

func (vb *Varnishbeat) Cleanup(b *beat.Beat) error {
	vb.varnish.Close()
	return nil
}

func (vb *Varnishbeat) Stop() {
	vb.varnish.Stop()
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/output"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/olivere/elastic.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 5
	esMaxConns         = 20
	esRecvBuffer       = 100
	esSendBuffer       = 100
)

type Indexer struct {
	bulkService       *elastic.BulkService
	indexPrefix       string
	indexType         string
	RateCounter       *ratecounter.RateCounter
	lastDisplayUpdate time.Time
}

type Config struct {
	Hosts       []string `json:"hosts"`
	IndexPrefix string   `json:"index"`
	IndexType   string   `json:"indexType"`
}

type ESServer struct {
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
}

func init() {
	output.Register("elasticsearch", &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	})
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func (i *Indexer) flush() error {
	numEvents := i.bulkService.NumberOfActions()

	if numEvents > 0 {
		if time.Now().Sub(i.lastDisplayUpdate) >= time.Duration(1*time.Second) {
			log.Printf("Flushing %d event(s) to Elasticsearch, current rate: %d/s", numEvents, i.RateCounter.Rate())
			i.lastDisplayUpdate = time.Now()
		}

		_, err := i.bulkService.Do()

		if err != nil {
			log.Printf("Unable to flush events: %s", err)
		}

		return err
	}

	return nil
}

func (i *Indexer) index(ev *buffer.Event) error {
	doc := *ev.Text
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkService.Add(request)
	i.RateCounter.Incr(1)

	numEvents := i.bulkService.NumberOfActions()

	if numEvents < esSendBuffer {
		return nil
	}

	return i.flush()
}

func (e *ESServer) Init(config json.RawMessage, b buffer.Sender) error {
	var esConfig *Config
	if err := json.Unmarshal(config, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	return nil
}

func readInputChannel(idx *Indexer, receiveChan chan *buffer.Event) {
	// Drain the channel only if we have room
	if idx.bulkService.NumberOfActions() < esSendBuffer {
		select {
		case ev := <-receiveChan:
			idx.index(ev)
		}
	} else {
		log.Printf("Internal Elasticsearch buffer is full, waiting")
		time.Sleep(1 * time.Second)
	}
}

func (es *ESServer) Start() error {
	var client *elastic.Client
	var err error

	for {
		client, err = elastic.NewClient(elastic.SetURL(es.hosts...))
		if err != nil {
			log.Printf("Error starting Elasticsearch: %s, will retry", err)
			time.Sleep(2 * time.Second)
			continue
		}

		break
	}

	log.Printf("Connected to Elasticsearch")

	service := elastic.NewBulkService(client)

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	// Create indexer
	idx := &Indexer{service, es.config.IndexPrefix, es.config.IndexType, rateCounter, time.Now()}

	// Loop events and publish to elasticsearch
	tick := time.NewTicker(time.Duration(esFlushInterval) * time.Second)

	for {
		readInputChannel(idx, receiveChan)

		select {
		case <-tick.C:
			idx.flush()
		case <-es.term:
			tick.Stop()
			log.Println("Elasticsearch received term signal")
			break
		}
	}

	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<commit_msg>Fix bug where Elasticsearch events would not be flushed quickly<commit_after>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/output"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/olivere/elastic.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 5
	esMaxConns         = 20
	esRecvBuffer       = 100
	esSendBuffer       = 100
)

type Indexer struct {
	bulkService       *elastic.BulkService
	indexPrefix       string
	indexType         string
	RateCounter       *ratecounter.RateCounter
	lastDisplayUpdate time.Time
}

type Config struct {
	Hosts       []string `json:"hosts"`
	IndexPrefix string   `json:"index"`
	IndexType   string   `json:"indexType"`
}

type ESServer struct {
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
}

func init() {
	output.Register("elasticsearch", &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	})
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func (i *Indexer) flush() error {
	numEvents := i.bulkService.NumberOfActions()

	if numEvents > 0 {
		if time.Now().Sub(i.lastDisplayUpdate) >= time.Duration(1*time.Second) {
			log.Printf("Flushing %d event(s) to Elasticsearch, current rate: %d/s", numEvents, i.RateCounter.Rate())
			i.lastDisplayUpdate = time.Now()
		}

		_, err := i.bulkService.Do()

		if err != nil {
			log.Printf("Unable to flush events: %s", err)
		}

		return err
	}

	return nil
}

func (i *Indexer) index(ev *buffer.Event) error {
	doc := *ev.Text
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkService.Add(request)
	i.RateCounter.Incr(1)

	numEvents := i.bulkService.NumberOfActions()

	if numEvents < esSendBuffer {
		return nil
	}

	return i.flush()
}

func (e *ESServer) Init(config json.RawMessage, b buffer.Sender) error {
	var esConfig *Config
	if err := json.Unmarshal(config, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	return nil
}

func readInputChannel(idx *Indexer, receiveChan chan *buffer.Event) {
	// Drain the channel only if we have room
	if idx.bulkService.NumberOfActions() < esSendBuffer {
		select {
		case ev := <-receiveChan:
			idx.index(ev)
		}
	} else {
		log.Printf("Internal Elasticsearch buffer is full, waiting")
		time.Sleep(1 * time.Second)
	}
}

func (es *ESServer) Start() error {
	var client *elastic.Client
	var err error

	for {
		client, err = elastic.NewClient(elastic.SetURL(es.hosts...))
		if err != nil {
			log.Printf("Error starting Elasticsearch: %s, will retry", err)
			time.Sleep(2 * time.Second)
			continue
		}

		break
	}

	log.Printf("Connected to Elasticsearch")

	service := elastic.NewBulkService(client)

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	// Create indexer
	idx := &Indexer{service, es.config.IndexPrefix, es.config.IndexType, rateCounter, time.Now()}

	// Loop events and publish to elasticsearch
	tick := time.NewTicker(time.Duration(esFlushInterval) * time.Second)

	for {
		readInputChannel(idx, receiveChan)

		if len(tick.C) > 0 || len(es.term) >0 {
			select {
			case <-tick.C:
				idx.flush()
			case <-es.term:
				tick.Stop()
				log.Println("Elasticsearch received term signal")
				break
			}
		}
	}

	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<|endoftext|>


<commit_before><commit_msg>reuseport: close Listener in case of error (#94)<commit_after><|endoftext|>