<commit_before>package stripe

import (
	"encoding/json"
)

// ChargeFraudUserReport is the list of allowed values for reporting fraud.
type ChargeFraudUserReport string

// List of values that ChargeFraudUserReport can take.
const (
	ChargeFraudUserReportFraudulent ChargeFraudUserReport = "fraudulent"
	ChargeFraudUserReportSafe       ChargeFraudUserReport = "safe"
)

// ChargeFraudStripeReport is the list of allowed values for reporting fraud.
type ChargeFraudStripeReport string

// List of values that ChargeFraudStripeReport can take.
const (
	ChargeFraudStripeReportFraudulent ChargeFraudStripeReport = "fraudulent"
)

// ChargeLevel3LineItemsParams is the set of parameters that represent a line item on level III data.
type ChargeLevel3LineItemsParams struct {
	DiscountAmount     *int64  `form:"discount_amount"`
	ProductCode        *string `form:"product_code"`
	ProductDescription *string `form:"product_description"`
	Quantity           *int64  `form:"quantity"`
	TaxAmount          *int64  `form:"tax_amount"`
	UnitCost           *int64  `form:"unit_cost"`
}

// ChargeLevel3Params is the set of parameters that can be used for the Level III data.
type ChargeLevel3Params struct {
	CustomerReference  *string                        `form:"customer_reference"`
	LineItems          []*ChargeLevel3LineItemsParams `form:"line_items"`
	MerchantReference  *string                        `form:"merchant_reference"`
	ShippingAddressZip *string                        `form:"shipping_address_zip"`
	ShippingFromZip    *string                        `form:"shipping_from_zip"`
	ShippingAmount     *int64                         `form:"shipping_amount"`
}

// ChargeParams is the set of parameters that can be used when creating or updating a charge.
type ChargeParams struct {
	Params              `form:"*"`
	Amount              *int64                 `form:"amount"`
	ApplicationFee      *int64                 `form:"application_fee"`
	Capture             *bool                  `form:"capture"`
	Currency            *string                `form:"currency"`
	Customer            *string                `form:"customer"`
	Description         *string                `form:"description"`
	Destination         *DestinationParams     `form:"destination"`
	ExchangeRate        *float64               `form:"exchange_rate"`
	FraudDetails        *FraudDetailsParams    `form:"fraud_details"`
	Level3              *ChargeLevel3Params    `form:"level3"`
	OnBehalfOf          *string                `form:"on_behalf_of"`
	ReceiptEmail        *string                `form:"receipt_email"`
	Shipping            *ShippingDetailsParams `form:"shipping"`
	Source              *SourceParams          `form:"*"` // SourceParams has custom encoding so brought to top level with "*"
	StatementDescriptor *string                `form:"statement_descriptor"`
	TransferGroup       *string                `form:"transfer_group"`
}

// ShippingDetailsParams is the structure containing shipping information as parameters
type ShippingDetailsParams struct {
	Address        *AddressParams `form:"address"`
	Carrier        *string        `form:"carrier"`
	Name           *string        `form:"name"`
	Phone          *string        `form:"phone"`
	TrackingNumber *string        `form:"tracking_number"`
}

// SetSource adds valid sources to a ChargeParams object,
// returning an error for unsupported sources.
func (p *ChargeParams) SetSource(sp interface{}) error {
	source, err := SourceParamsFor(sp)
	p.Source = source
	return err
}

// DestinationParams describes the parameters available for the destination hash when creating a charge.
type DestinationParams struct {
	Account *string `form:"account"`
	Amount  *int64  `form:"amount"`
}

// FraudDetailsParams provides information on the fraud details for a charge.
type FraudDetailsParams struct {
	UserReport *string `form:"user_report"`
}

// ChargeListParams is the set of parameters that can be used when listing charges.
type ChargeListParams struct {
	ListParams    `form:"*"`
	Created       *int64            `form:"created"`
	CreatedRange  *RangeQueryParams `form:"created"`
	Customer      *string           `form:"customer"`
	TransferGroup *string           `form:"transfer_group"`
}

// CaptureParams is the set of parameters that can be used when capturing a charge.
type CaptureParams struct {
	Params              `form:"*"`
	Amount              *int64   `form:"amount"`
	ApplicationFee      *int64   `form:"application_fee"`
	ExchangeRate        *float64 `form:"exchange_rate"`
	ReceiptEmail        *string  `form:"receipt_email"`
	StatementDescriptor *string  `form:"statement_descriptor"`
}

// ChargeLevel3LineItem represents a line item on level III data.
// This is in private beta and would be empty for most integrations
type ChargeLevel3LineItem struct {
	DiscountAmount     int64  `json:"discount_amount"`
	ProductCode        string `json:"product_code"`
	ProductDescription string `json:"product_description"`
	Quantity           int64  `json:"quantity"`
	TaxAmount          int64  `json:"tax_amount"`
	UnitCost           int64  `json:"unit_cost"`
}

// ChargeLevel3 represents the Level III data.
// This is in private beta and would be empty for most integrations
type ChargeLevel3 struct {
	CustomerReference  string                  `json:"customer_reference"`
	LineItems          []*ChargeLevel3LineItem `json:"line_items"`
	MerchantReference  string                  `json:"merchant_reference"`
	ShippingAddressZip string                  `json:"shipping_address_zip"`
	ShippingFromZip    string                  `json:"shipping_from_zip"`
	ShippingAmount     int64                   `json:"shipping_amount"`
}

// Charge is the resource representing a Stripe charge.
// For more details see https://stripe.com/docs/api#charges.
type Charge struct {
	Amount              int64               `json:"amount"`
	AmountRefunded      int64               `json:"amount_refunded"`
	Application         *Application        `json:"application"`
	ApplicationFee      *ApplicationFee     `json:"application_fee"`
	BalanceTransaction  *BalanceTransaction `json:"balance_transaction"`
	Captured            bool                `json:"captured"`
	Created             int64               `json:"created"`
	Currency            Currency            `json:"currency"`
	Customer            *Customer           `json:"customer"`
	Description         string              `json:"description"`
	Destination         *Account            `json:"destination"`
	Dispute             *Dispute            `json:"dispute"`
	FailureCode         string              `json:"failure_code"`
	FailureMessage      string              `json:"failure_message"`
	FraudDetails        *FraudDetails       `json:"fraud_details"`
	ID                  string              `json:"id"`
	Invoice             *Invoice            `json:"invoice"`
	Level3              ChargeLevel3        `json:"level3"`
	Livemode            bool                `json:"livemode"`
	Metadata            map[string]string   `json:"metadata"`
	Outcome             *ChargeOutcome      `json:"outcome"`
	Paid                bool                `json:"paid"`
	ReceiptEmail        string              `json:"receipt_email"`
	ReceiptNumber       string              `json:"receipt_number"`
	Refunded            bool                `json:"refunded"`
	Refunds             *RefundList         `json:"refunds"`
	Review              *Review             `json:"review"`
	Shipping            *ShippingDetails    `json:"shipping"`
	Source              *PaymentSource      `json:"source"`
	SourceTransfer      *Transfer           `json:"source_transfer"`
	StatementDescriptor string              `json:"statement_descriptor"`
	Status              string              `json:"status"`
	Transfer            *Transfer           `json:"transfer"`
	TransferGroup       string              `json:"transfer_group"`
}

// UnmarshalJSON handles deserialization of a charge.
// This custom unmarshaling is needed because the resulting
// property may be an ID or the full struct if it was expanded.
func (c *Charge) UnmarshalJSON(data []byte) error {
	if id, ok := ParseID(data); ok {
		c.ID = id
		return nil
	}

	type charge Charge
	var v charge
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	*c = Charge(v)
	return nil
}

// ChargeList is a list of charges as retrieved from a list endpoint.
type ChargeList struct {
	ListMeta
	Data []*Charge `json:"data"`
}

// FraudDetails is the structure detailing fraud status.
type FraudDetails struct {
	UserReport   ChargeFraudUserReport   `json:"user_report"`
	StripeReport ChargeFraudStripeReport `json:"stripe_report"`
}

// ChargeOutcomeRule tells you the Radar rule that blocked the charge, if any.
type ChargeOutcomeRule struct {
	Action    string `json:"action"`
	ID        string `json:"id"`
	Predicate string `json:"predicate"`
}

// ChargeOutcome is the charge's outcome that details whether a payment
// was accepted and why.
type ChargeOutcome struct {
	NetworkStatus string             `json:"network_status"`
	Reason        string             `json:"reason"`
	RiskLevel     string             `json:"risk_level"`
	RiskScore     int64              `json:"risk_score"`
	Rule          *ChargeOutcomeRule `json:"rule"`
	SellerMessage string             `json:"seller_message"`
	Type          string             `json:"type"`
}

// ShippingDetails is the structure containing shipping information.
type ShippingDetails struct {
	Address        *Address `json:"address"`
	Carrier        string   `json:"carrier"`
	Name           string   `json:"name"`
	Phone          string   `json:"phone"`
	TrackingNumber string   `json:"tracking_number"`
}

var depth = -1

// UnmarshalJSON handles deserialization of a ChargeOutcomeRule.
// This custom unmarshaling is needed because the resulting
// property may be an id or the full struct if it was expanded.
func (c *ChargeOutcomeRule) UnmarshalJSON(data []byte) error {
	if id, ok := ParseID(data); ok {
		c.ID = id
		return nil
	}

	type chargeOutcomeRule ChargeOutcomeRule
	var v chargeOutcomeRule
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	*c = ChargeOutcomeRule(v)
	return nil
}
<commit_msg>Add support for authorization_code on the Charge resource<commit_after>package stripe

import (
	"encoding/json"
)

// ChargeFraudUserReport is the list of allowed values for reporting fraud.
type ChargeFraudUserReport string

// List of values that ChargeFraudUserReport can take.
const (
	ChargeFraudUserReportFraudulent ChargeFraudUserReport = "fraudulent"
	ChargeFraudUserReportSafe       ChargeFraudUserReport = "safe"
)

// ChargeFraudStripeReport is the list of allowed values for reporting fraud.
type ChargeFraudStripeReport string

// List of values that ChargeFraudStripeReport can take.
const (
	ChargeFraudStripeReportFraudulent ChargeFraudStripeReport = "fraudulent"
)

// ChargeLevel3LineItemsParams is the set of parameters that represent a line item on level III data.
type ChargeLevel3LineItemsParams struct {
	DiscountAmount     *int64  `form:"discount_amount"`
	ProductCode        *string `form:"product_code"`
	ProductDescription *string `form:"product_description"`
	Quantity           *int64  `form:"quantity"`
	TaxAmount          *int64  `form:"tax_amount"`
	UnitCost           *int64  `form:"unit_cost"`
}

// ChargeLevel3Params is the set of parameters that can be used for the Level III data.
type ChargeLevel3Params struct {
	CustomerReference  *string                        `form:"customer_reference"`
	LineItems          []*ChargeLevel3LineItemsParams `form:"line_items"`
	MerchantReference  *string                        `form:"merchant_reference"`
	ShippingAddressZip *string                        `form:"shipping_address_zip"`
	ShippingFromZip    *string                        `form:"shipping_from_zip"`
	ShippingAmount     *int64                         `form:"shipping_amount"`
}

// ChargeParams is the set of parameters that can be used when creating or updating a charge.
type ChargeParams struct {
	Params              `form:"*"`
	Amount              *int64                 `form:"amount"`
	ApplicationFee      *int64                 `form:"application_fee"`
	Capture             *bool                  `form:"capture"`
	Currency            *string                `form:"currency"`
	Customer            *string                `form:"customer"`
	Description         *string                `form:"description"`
	Destination         *DestinationParams     `form:"destination"`
	ExchangeRate        *float64               `form:"exchange_rate"`
	FraudDetails        *FraudDetailsParams    `form:"fraud_details"`
	Level3              *ChargeLevel3Params    `form:"level3"`
	OnBehalfOf          *string                `form:"on_behalf_of"`
	ReceiptEmail        *string                `form:"receipt_email"`
	Shipping            *ShippingDetailsParams `form:"shipping"`
	Source              *SourceParams          `form:"*"` // SourceParams has custom encoding so brought to top level with "*"
	StatementDescriptor *string                `form:"statement_descriptor"`
	TransferGroup       *string                `form:"transfer_group"`
}

// ShippingDetailsParams is the structure containing shipping information as parameters
type ShippingDetailsParams struct {
	Address        *AddressParams `form:"address"`
	Carrier        *string        `form:"carrier"`
	Name           *string        `form:"name"`
	Phone          *string        `form:"phone"`
	TrackingNumber *string        `form:"tracking_number"`
}

// SetSource adds valid sources to a ChargeParams object,
// returning an error for unsupported sources.
func (p *ChargeParams) SetSource(sp interface{}) error {
	source, err := SourceParamsFor(sp)
	p.Source = source
	return err
}

// DestinationParams describes the parameters available for the destination hash when creating a charge.
type DestinationParams struct {
	Account *string `form:"account"`
	Amount  *int64  `form:"amount"`
}

// FraudDetailsParams provides information on the fraud details for a charge.
type FraudDetailsParams struct {
	UserReport *string `form:"user_report"`
}

// ChargeListParams is the set of parameters that can be used when listing charges.
type ChargeListParams struct {
	ListParams    `form:"*"`
	Created       *int64            `form:"created"`
	CreatedRange  *RangeQueryParams `form:"created"`
	Customer      *string           `form:"customer"`
	TransferGroup *string           `form:"transfer_group"`
}

// CaptureParams is the set of parameters that can be used when capturing a charge.
type CaptureParams struct {
	Params              `form:"*"`
	Amount              *int64   `form:"amount"`
	ApplicationFee      *int64   `form:"application_fee"`
	ExchangeRate        *float64 `form:"exchange_rate"`
	ReceiptEmail        *string  `form:"receipt_email"`
	StatementDescriptor *string  `form:"statement_descriptor"`
}

// ChargeLevel3LineItem represents a line item on level III data.
// This is in private beta and would be empty for most integrations
type ChargeLevel3LineItem struct {
	DiscountAmount     int64  `json:"discount_amount"`
	ProductCode        string `json:"product_code"`
	ProductDescription string `json:"product_description"`
	Quantity           int64  `json:"quantity"`
	TaxAmount          int64  `json:"tax_amount"`
	UnitCost           int64  `json:"unit_cost"`
}

// ChargeLevel3 represents the Level III data.
// This is in private beta and would be empty for most integrations
type ChargeLevel3 struct {
	CustomerReference  string                  `json:"customer_reference"`
	LineItems          []*ChargeLevel3LineItem `json:"line_items"`
	MerchantReference  string                  `json:"merchant_reference"`
	ShippingAddressZip string                  `json:"shipping_address_zip"`
	ShippingFromZip    string                  `json:"shipping_from_zip"`
	ShippingAmount     int64                   `json:"shipping_amount"`
}

// Charge is the resource representing a Stripe charge.
// For more details see https://stripe.com/docs/api#charges.
type Charge struct {
	Amount              int64               `json:"amount"`
	AmountRefunded      int64               `json:"amount_refunded"`
	Application         *Application        `json:"application"`
	ApplicationFee      *ApplicationFee     `json:"application_fee"`
	AuthorizationCode   string              `json:"authorization_code"`
	BalanceTransaction  *BalanceTransaction `json:"balance_transaction"`
	Captured            bool                `json:"captured"`
	Created             int64               `json:"created"`
	Currency            Currency            `json:"currency"`
	Customer            *Customer           `json:"customer"`
	Description         string              `json:"description"`
	Destination         *Account            `json:"destination"`
	Dispute             *Dispute            `json:"dispute"`
	FailureCode         string              `json:"failure_code"`
	FailureMessage      string              `json:"failure_message"`
	FraudDetails        *FraudDetails       `json:"fraud_details"`
	ID                  string              `json:"id"`
	Invoice             *Invoice            `json:"invoice"`
	Level3              ChargeLevel3        `json:"level3"`
	Livemode            bool                `json:"livemode"`
	Metadata            map[string]string   `json:"metadata"`
	Outcome             *ChargeOutcome      `json:"outcome"`
	Paid                bool                `json:"paid"`
	ReceiptEmail        string              `json:"receipt_email"`
	ReceiptNumber       string              `json:"receipt_number"`
	Refunded            bool                `json:"refunded"`
	Refunds             *RefundList         `json:"refunds"`
	Review              *Review             `json:"review"`
	Shipping            *ShippingDetails    `json:"shipping"`
	Source              *PaymentSource      `json:"source"`
	SourceTransfer      *Transfer           `json:"source_transfer"`
	StatementDescriptor string              `json:"statement_descriptor"`
	Status              string              `json:"status"`
	Transfer            *Transfer           `json:"transfer"`
	TransferGroup       string              `json:"transfer_group"`
}

// UnmarshalJSON handles deserialization of a charge.
// This custom unmarshaling is needed because the resulting
// property may be an ID or the full struct if it was expanded.
func (c *Charge) UnmarshalJSON(data []byte) error {
	if id, ok := ParseID(data); ok {
		c.ID = id
		return nil
	}

	type charge Charge
	var v charge
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	*c = Charge(v)
	return nil
}

// ChargeList is a list of charges as retrieved from a list endpoint.
type ChargeList struct {
	ListMeta
	Data []*Charge `json:"data"`
}

// FraudDetails is the structure detailing fraud status.
type FraudDetails struct {
	UserReport   ChargeFraudUserReport   `json:"user_report"`
	StripeReport ChargeFraudStripeReport `json:"stripe_report"`
}

// ChargeOutcomeRule tells you the Radar rule that blocked the charge, if any.
type ChargeOutcomeRule struct {
	Action    string `json:"action"`
	ID        string `json:"id"`
	Predicate string `json:"predicate"`
}

// ChargeOutcome is the charge's outcome that details whether a payment
// was accepted and why.
type ChargeOutcome struct {
	NetworkStatus string             `json:"network_status"`
	Reason        string             `json:"reason"`
	RiskLevel     string             `json:"risk_level"`
	RiskScore     int64              `json:"risk_score"`
	Rule          *ChargeOutcomeRule `json:"rule"`
	SellerMessage string             `json:"seller_message"`
	Type          string             `json:"type"`
}

// ShippingDetails is the structure containing shipping information.
type ShippingDetails struct {
	Address        *Address `json:"address"`
	Carrier        string   `json:"carrier"`
	Name           string   `json:"name"`
	Phone          string   `json:"phone"`
	TrackingNumber string   `json:"tracking_number"`
}

var depth = -1

// UnmarshalJSON handles deserialization of a ChargeOutcomeRule.
// This custom unmarshaling is needed because the resulting
// property may be an id or the full struct if it was expanded.
func (c *ChargeOutcomeRule) UnmarshalJSON(data []byte) error {
	if id, ok := ParseID(data); ok {
		c.ID = id
		return nil
	}

	type chargeOutcomeRule ChargeOutcomeRule
	var v chargeOutcomeRule
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	*c = ChargeOutcomeRule(v)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"go/parser"
	"go/token"
	"log"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

const (
	bitflags = 1 << iota
	jsonMarshalIsString
	jsonUnmarshalIsString
	xmlMarshalIsString
	xmlUnmarshalIsString
)

type EnumData struct {
	Package string
	File    string
	Reprs   []*EnumRepr
}

type EnumRepr struct {
	Name string

	flags    uint
	FlagSep  string
	iterName string
	unique   string
	Fields   []*FieldRepr
}

type FieldRepr struct {
	Name        string
	String      string
	Description string
	Value       int64
}

func (self *EnumData) DoFile(file string) error {
	fset := token.NewFileSet()

	f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	self.Package = f.Name.Name
	if len(self.Package) == 0 {
		return fmt.Errorf("No package Name")
	}

	var dir, filename = filepath.Split(file)

	self.File = filepath.Join(dir, "enum____"+filename)

	for _, cg := range f.Comments {
		self.doComment(cg.Text())
	}

	if err := self.generateCode(); err != nil {
		return err
	}

	return nil
}

func (self *EnumRepr) GetUniqueName() string {
	if self.unique == "" {
		self.unique = "value_" + strconv.FormatInt(rand.Int63(), 36)
	}
	return self.unique
}

func (self *EnumRepr) JsonMarshalIsString() bool {
	return self.flags&jsonMarshalIsString == jsonMarshalIsString
}

func (self *EnumRepr) JsonUnmarshalIsString() bool {
	return self.flags&jsonUnmarshalIsString == jsonUnmarshalIsString
}

func (self *EnumRepr) GetIterName() string {
	if len(self.iterName) == 0 {
		return self.Name + "Values"
	}
	return self.iterName
}

func (repr *EnumRepr) GetIntType() string {
	var bf = repr.flags&bitflags == bitflags
	var ln = int64(len(repr.Fields))

	switch {
	case (bf && ln <= 8) || ln < 256:
		return "uint8"
	case (bf && ln <= 16) || ln < 65536:
		return "uint16"
	case (bf && ln <= 32) || ln < 4294967296:
		return "uint32"
	}
	return "uint64"
}

func (self *EnumData) checkValidity(flgs, flds, errs bool, repr *EnumRepr) {
	const warnMarshal = "WARNING: %s marshal and unmarshal do not match for %q\n"

	if !flgs || !flds || errs {
		if errs {
			log.Println("Enums with errors are excluded")
		} else {
			log.Println("Incomplete definition. Both flags and Fields are required.")
		}
		if len(self.Reprs) > 0 {
			self.Reprs = self.Reprs[0 : len(self.Reprs)-1]
		}

	} else if repr != nil {
		// Warn if marshaling to string and unmarshaling to number, or vice versa.
		if (repr.flags&jsonMarshalIsString == 0) !=
			(repr.flags&jsonUnmarshalIsString == 0) {
			log.Printf(warnMarshal, "JSON", repr.Name)
		}
		if (repr.flags&xmlMarshalIsString == 0) !=
			(repr.flags&xmlUnmarshalIsString == 0) {
			log.Printf(warnMarshal, "XML", repr.Name)
		}
	}
}

func (self *EnumData) doComment(cgText string) {
	s := bufio.NewScanner(strings.NewReader(cgText))

	firstPass := true
	doFlags, doFields, hasErrors := true, true, false

	var repr *EnumRepr

	for line := skipEmptyLines(s); len(line) > 0; line = skipEmptyLines(s) {
		line = strings.TrimSpace(line)

		if strings.HasPrefix(line, "@enum") {
			firstPass = false
			self.checkValidity(doFlags, doFields, hasErrors, repr)

			repr = &EnumRepr{
				FlagSep: ",", // Go ahead and set the default even if not needed
			}
			self.Reprs = append(self.Reprs, repr)

			doFlags, doFields, hasErrors = false, false, false

			line = line[5:] // strip out the @enum

			if len(strings.TrimSpace(line)) == 0 {
				continue

			} else if unicode.IsSpace(rune(line[0])) {
				line = strings.TrimSpace(line)

			} else {
				hasErrors = true
				continue
			}
		} else if firstPass {
			return // comment group didn't start with @enum

		} else if hasErrors {
			continue
		}

		if strings.HasPrefix(line, "--") {
			if doFields {
				log.Println("Flags must come before enum variant definitions")
				hasErrors = true
				continue
			}
			doFlags = true

			if repr.setFlags(line) == false {
				hasErrors = true
			}

		} else { // Get the field definitions
			if !doFlags {
				log.Println("No flags were defined before the first enum variant")
				hasErrors = true
				continue
			}
			doFields = true

			if repr.setField(line) == false {
				hasErrors = true
			}
		}
	}

	self.checkValidity(doFlags, doFields, hasErrors, repr)
}

const unexpectedValue = "Unexpected value %q for %q\n"

func (self *EnumRepr) setFlags(flags string) bool {
	var Name, value string
	var found, foundEqual bool

	for len(flags) > 0 {
		if flags, found = trashUntil(strings.TrimSpace(flags), "--", true); !found {
			if len(strings.TrimSpace(flags)) > 0 {
				log.Printf(`Expected "--", but found: %q`, flags)
				return false
			}
			break
		}

		flags, Name = getIdent(flags)

		switch strings.ToLower(Name) {

		case "name": // Set the base Name for the enum
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}

			if len(self.Name) > 0 {
				log.Printf("Name is already set: %q, but found: %q\n", self.Name, value)
				return false
			}

			self.Name = value

		case "bitflags": // The enum values are to be bitflags
			self.flags |= bitflags

		case "bitflag_separator": // The separator used when joining bitflags
			if flags, self.FlagSep, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}

		case "iterator_name": // Custom Name for Array of values
			if flags, self.iterName, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}

		case "json": // Set type of JSON marshaler and unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, jsonMarshalIsString|jsonUnmarshalIsString)

		case "xml": // Set type of XML marshaler and unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, xmlMarshalIsString|xmlUnmarshalIsString)

		case "json_marshaler": // Set type of JSON marshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, jsonMarshalIsString)

		case "json_unmarshaler": // Set type of JSON unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, jsonUnmarshalIsString)

		case "xml_marshaler": // Set type of XML marshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, xmlMarshalIsString)

		case "xml_unmarshaler": // Set type of XML unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, xmlUnmarshalIsString)

		default:
			log.Printf("Unknown flag %q\n", Name)
		}
	}
	return true
}

func (self *EnumRepr) setField(field string) bool {
	var f = FieldRepr{
		Value: -1,
	}

	field, f.Name = getIdent(strings.TrimSpace(field))

	if len(f.Name) == 0 {
		log.Println("Field Name is empty")
		return false
	}

	var Name string
	var foundEqual, found bool

	for len(field) > 0 {
		if field, found = trashUntil(strings.TrimSpace(field), "--", true); !found {
			if len(strings.TrimSpace(field)) > 0 {
				log.Printf(`Expected "--", but found: %q`, field)
			}
			break
		}

		field, Name = getIdent(field)

		switch strings.ToLower(Name) {

		case "string": // The string representation of the field
			if field, f.String, foundEqual = getValue(Name, field); !foundEqual {
				return false
			}

		case "description": // The description of the field
			if field, f.Description, foundEqual = getValue(Name, field); !foundEqual {
				return false
			}

		case "value": // Custom value for the field
			var v string

			if field, v, foundEqual = getValue(Name, field); !foundEqual {
				return false
			}
			if n, err := strconv.ParseUint(v, 10, 32); err != nil {
				log.Printf("%q is not a valid uint\n", v)
				return false
			} else {
				f.Value = int64(n)
			}
		default:
			log.Printf("Unknown flag %q\n", Name)
		}
	}

	if self.flags&bitflags == bitflags && f.Value != -1 {
		log.Println("bitflag enums may not have a custom --value setting")
		return false
	}

	if len(f.String) == 0 {
		f.String = f.Name
	}
	if len(f.Description) == 0 {
		f.Description = f.String
	}

	if f.Value == -1 {
		if self.flags&bitflags == bitflags {
			f.Value = 1 << uint(len(self.Fields))
		} else {
			// TODO: Make sure there are no Custom number conflicts
			f.Value = int64(len(self.Fields) + 1)
		}
	}

	self.Fields = append(self.Fields, &f)

	return true
}

func (self *EnumRepr) setMarshalingFlags(Name, value string, flags uint) {
	switch strings.ToLower(value) {
	case "string":
		self.flags |= flags
	case "value":
		self.flags &^= flags
	default:
		log.Printf(unexpectedValue, value, Name)
	}
}

func (self *EnumRepr) IsBitflag() bool {
	return self.flags&bitflags == bitflags
}

func trashUntil(source, search string, exclude bool) (string, bool) {
	var idx = strings.Index(source, search)
	if idx == -1 {
		return source, false
	}

	if idx > 0 {
		log.Printf("Expected %q, but found: %q\n", search, source[0:idx])
	}

	if exclude { // caller wants the leading search string removed
		idx += len(search)
	}
	return source[idx:], true
}

func getIdent(source string) (_ string, word string) {
	var i = 0
	for j, c := range source {
		if ('a' <= c && c <= 'z') || c == '_' || ('A' <= c && c <= 'Z') ||
			(j > 0 && '0' <= c && c <= '9') {
			word += string(c)
			i += utf8.RuneLen(c)
		} else {
			break
		}
	}
	return source[i:], word
}

func getValue(Name, source string) (_, value string, foundEqual bool) {
	if len(source) == 0 {
		return "", "", false
	}

	source, foundEqual = trashUntil(strings.TrimSpace(source), "=", true)
	if !foundEqual {
		log.Printf(`Expected "=" after %q.\n`, Name)
		return source, "", foundEqual
	}

	source = strings.TrimSpace(source)

	var idx = 0

	if source[0] == '"' || source[0] == '\'' {
		idx = strings.IndexByte(source[1:], source[0])
		if idx == -1 {
			return source, "", false // Expected closing quote
		} else {
			idx += 1 // Because we started searching on the second character
		}
		return source[idx+1:], source[1:idx], foundEqual
	} else { // Get unquoted value
		for _, r := range source {
			if unicode.IsSpace(r) {
				break
			}
			idx += utf8.RuneLen(r)
		}
		return source[idx:], source[0:idx], foundEqual
	}
}

// Returns an empty string when scanner is complete
func skipEmptyLines(s *bufio.Scanner) string {
	for s.Scan() {
		t := strings.TrimSpace(s.Text())

		if err := s.Err(); err != nil {
			fmt.Println(err)
		}

		if len(t) == 0 {
			continue
		}
		return t
	}
	return ""
}
<commit_msg>Renamed marshaler flag to marshal; same with unmarshal<commit_after>package main

import (
	"bufio"
	"fmt"
	"go/parser"
	"go/token"
	"log"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

const (
	bitflags = 1 << iota
	jsonMarshalIsString
	jsonUnmarshalIsString
	xmlMarshalIsString
	xmlUnmarshalIsString
)

type EnumData struct {
	Package string
	File    string
	Reprs   []*EnumRepr
}

type EnumRepr struct {
	Name string

	flags    uint
	FlagSep  string
	iterName string
	unique   string
	Fields   []*FieldRepr
}

type FieldRepr struct {
	Name        string
	String      string
	Description string
	Value       int64
}

func (self *EnumData) DoFile(file string) error {
	fset := token.NewFileSet()

	f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	self.Package = f.Name.Name
	if len(self.Package) == 0 {
		return fmt.Errorf("No package Name")
	}

	var dir, filename = filepath.Split(file)

	self.File = filepath.Join(dir, "enum____"+filename)

	for _, cg := range f.Comments {
		self.doComment(cg.Text())
	}

	if err := self.generateCode(); err != nil {
		return err
	}

	return nil
}

func (self *EnumRepr) GetUniqueName() string {
	if self.unique == "" {
		self.unique = "value_" + strconv.FormatInt(rand.Int63(), 36)
	}
	return self.unique
}

func (self *EnumRepr) JsonMarshalIsString() bool {
	return self.flags&jsonMarshalIsString == jsonMarshalIsString
}

func (self *EnumRepr) JsonUnmarshalIsString() bool {
	return self.flags&jsonUnmarshalIsString == jsonUnmarshalIsString
}

func (self *EnumRepr) GetIterName() string {
	if len(self.iterName) == 0 {
		return self.Name + "Values"
	}
	return self.iterName
}

func (repr *EnumRepr) GetIntType() string {
	var bf = repr.flags&bitflags == bitflags
	var ln = int64(len(repr.Fields))

	switch {
	case (bf && ln <= 8) || ln < 256:
		return "uint8"
	case (bf && ln <= 16) || ln < 65536:
		return "uint16"
	case (bf && ln <= 32) || ln < 4294967296:
		return "uint32"
	}
	return "uint64"
}

func (self *EnumData) checkValidity(flgs, flds, errs bool, repr *EnumRepr) {
	const warnMarshal = "WARNING: %s marshal and unmarshal do not match for %q\n"

	if !flgs || !flds || errs {
		if errs {
			log.Println("Enums with errors are excluded")
		} else {
			log.Println("Incomplete definition. Both flags and Fields are required.")
		}
		if len(self.Reprs) > 0 {
			self.Reprs = self.Reprs[0 : len(self.Reprs)-1]
		}

	} else if repr != nil {
		// Warn if marshaling to string and unmarshaling to number, or vice versa.
		if (repr.flags&jsonMarshalIsString == 0) !=
			(repr.flags&jsonUnmarshalIsString == 0) {
			log.Printf(warnMarshal, "JSON", repr.Name)
		}
		if (repr.flags&xmlMarshalIsString == 0) !=
			(repr.flags&xmlUnmarshalIsString == 0) {
			log.Printf(warnMarshal, "XML", repr.Name)
		}
	}
}

func (self *EnumData) doComment(cgText string) {
	s := bufio.NewScanner(strings.NewReader(cgText))

	firstPass := true
	doFlags, doFields, hasErrors := true, true, false

	var repr *EnumRepr

	for line := skipEmptyLines(s); len(line) > 0; line = skipEmptyLines(s) {
		line = strings.TrimSpace(line)

		if strings.HasPrefix(line, "@enum") {
			firstPass = false
			self.checkValidity(doFlags, doFields, hasErrors, repr)

			repr = &EnumRepr{
				FlagSep: ",", // Go ahead and set the default even if not needed
			}
			self.Reprs = append(self.Reprs, repr)

			doFlags, doFields, hasErrors = false, false, false

			line = line[5:] // strip out the @enum

			if len(strings.TrimSpace(line)) == 0 {
				continue

			} else if unicode.IsSpace(rune(line[0])) {
				line = strings.TrimSpace(line)

			} else {
				hasErrors = true
				continue
			}
		} else if firstPass {
			return // comment group didn't start with @enum

		} else if hasErrors {
			continue
		}

		if strings.HasPrefix(line, "--") {
			if doFields {
				log.Println("Flags must come before enum variant definitions")
				hasErrors = true
				continue
			}
			doFlags = true

			if repr.setFlags(line) == false {
				hasErrors = true
			}

		} else { // Get the field definitions
			if !doFlags {
				log.Println("No flags were defined before the first enum variant")
				hasErrors = true
				continue
			}
			doFields = true

			if repr.setField(line) == false {
				hasErrors = true
			}
		}
	}

	self.checkValidity(doFlags, doFields, hasErrors, repr)
}

const unexpectedValue = "Unexpected value %q for %q\n"

func (self *EnumRepr) setFlags(flags string) bool {
	var Name, value string
	var found, foundEqual bool

	for len(flags) > 0 {
		if flags, found = trashUntil(strings.TrimSpace(flags), "--", true); !found {
			if len(strings.TrimSpace(flags)) > 0 {
				log.Printf(`Expected "--", but found: %q`, flags)
				return false
			}
			break
		}

		flags, Name = getIdent(flags)

		switch strings.ToLower(Name) {

		case "name": // Set the base Name for the enum
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}

			if len(self.Name) > 0 {
				log.Printf("Name is already set: %q, but found: %q\n", self.Name, value)
				return false
			}

			self.Name = value

		case "bitflags": // The enum values are to be bitflags
			self.flags |= bitflags

		case "bitflag_separator": // The separator used when joining bitflags
			if flags, self.FlagSep, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}

		case "iterator_name": // Custom Name for Array of values
			if flags, self.iterName, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}

		case "json": // Set type of JSON marshaler and unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, jsonMarshalIsString|jsonUnmarshalIsString)

		case "xml": // Set type of XML marshaler and unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, xmlMarshalIsString|xmlUnmarshalIsString)

		case "json_marshal": // Set type of JSON marshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, jsonMarshalIsString)

		case "json_unmarshal": // Set type of JSON unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, jsonUnmarshalIsString)

		case "xml_marshal": // Set type of XML marshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, xmlMarshalIsString)

		case "xml_unmarshal": // Set type of XML unmarshaler
			if flags, value, foundEqual = getValue(Name, flags); !foundEqual {
				return false
			}
			self.setMarshalingFlags(Name, value, xmlUnmarshalIsString)

		default:
			log.Printf("Unknown flag %q\n", Name)
		}
	}
	return true
}

func (self *EnumRepr) setField(field string) bool {
	var f = FieldRepr{
		Value: -1,
	}

	field, f.Name = getIdent(strings.TrimSpace(field))

	if len(f.Name) == 0 {
		log.Println("Field Name is empty")
		return false
	}

	var Name string
	var foundEqual, found bool

	for len(field) > 0 {
		if field, found = trashUntil(strings.TrimSpace(field), "--", true); !found {
			if len(strings.TrimSpace(field)) > 0 {
				log.Printf(`Expected "--", but found: %q`, field)
			}
			break
		}

		field, Name = getIdent(field)

		switch strings.ToLower(Name) {

		case "string": // The string representation of the field
			if field, f.String, foundEqual = getValue(Name, field); !foundEqual {
				return false
			}

		case "description": // The description of the field
			if field, f.Description, foundEqual = getValue(Name, field); !foundEqual {
				return false
			}

		case "value": // Custom value for the field
			var v string

			if field, v, foundEqual = getValue(Name, field); !foundEqual {
				return false
			}
			if n, err := strconv.ParseUint(v, 10, 32); err != nil {
				log.Printf("%q is not a valid uint\n", v)
				return false
			} else {
				f.Value = int64(n)
			}
		default:
			log.Printf("Unknown flag %q\n", Name)
		}
	}

	if self.flags&bitflags == bitflags && f.Value != -1 {
		log.Println("bitflag enums may not have a custom --value setting")
		return false
	}

	if len(f.String) == 0 {
		f.String = f.Name
	}
	if len(f.Description) == 0 {
		f.Description = f.String
	}

	if f.Value == -1 {
		if self.flags&bitflags == bitflags {
			f.Value = 1 << uint(len(self.Fields))
		} else {
			// TODO: Make sure there are no Custom number conflicts
			f.Value = int64(len(self.Fields) + 1)
		}
	}

	self.Fields = append(self.Fields, &f)

	return true
}

func (self *EnumRepr) setMarshalingFlags(Name, value string, flags uint) {
	switch strings.ToLower(value) {
	case "string":
		self.flags |= flags
	case "value":
		self.flags &^= flags
	default:
		log.Printf(unexpectedValue, value, Name)
	}
}

func (self *EnumRepr) IsBitflag() bool {
	return self.flags&bitflags == bitflags
}

func trashUntil(source, search string, exclude bool) (string, bool) {
	var idx = strings.Index(source, search)
	if idx == -1 {
		return source, false
	}

	if idx > 0 {
		log.Printf("Expected %q, but found: %q\n", search, source[0:idx])
	}

	if exclude { // caller wants the leading search string removed
		idx += len(search)
	}
	return source[idx:], true
}

func getIdent(source string) (_ string, word string) {
	var i = 0
	for j, c := range source {
		if ('a' <= c && c <= 'z') || c == '_' || ('A' <= c && c <= 'Z') ||
			(j > 0 && '0' <= c && c <= '9') {
			word += string(c)
			i += utf8.RuneLen(c)
		} else {
			break
		}
	}
	return source[i:], word
}

func getValue(Name, source string) (_, value string, foundEqual bool) {
	if len(source) == 0 {
		return "", "", false
	}

	source, foundEqual = trashUntil(strings.TrimSpace(source), "=", true)
	if !foundEqual {
		log.Printf(`Expected "=" after %q.\n`, Name)
		return source, "", foundEqual
	}

	source = strings.TrimSpace(source)

	var idx = 0

	if source[0] == '"' || source[0] == '\'' {
		idx = strings.IndexByte(source[1:], source[0])
		if idx == -1 {
			return source, "", false // Expected closing quote
		} else {
			idx += 1 // Because we started searching on the second character
		}
		return source[idx+1:], source[1:idx], foundEqual
	} else { // Get unquoted value
		for _, r := range source {
			if unicode.IsSpace(r) {
				break
			}
			idx += utf8.RuneLen(r)
		}
		return source[idx:], source[0:idx], foundEqual
	}
}

// Returns an empty string when scanner is complete
func skipEmptyLines(s *bufio.Scanner) string {
	for s.Scan() {
		t := strings.TrimSpace(s.Text())

		if err := s.Err(); err != nil {
			fmt.Println(err)
		}

		if len(t) == 0 {
			continue
		}
		return t
	}
	return ""
}
<|endoftext|>


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

/*
	Package iothrottler implements application IO throttling.
*/

package iothrottler

import (
	"errors"
	"io"
	"log"
	"math"
	"net"
	"time"
)

// The Bandwidth type represents a bandwidth quantity in bytes per second.
// Sub-byte per seconds values are not supported
type Bandwidth int64

const (
	// Bytes per second
	BytesPerSecond Bandwidth = 1
	// Kilobits per second
	Kbps = BytesPerSecond * (1024 / 8)
	// Megabits per second
	Mbps = Kbps * 1024
	// Gigabits per second
	Gbps = Mbps * 1024
	// Unlimited bandwidth
	Unlimited = math.MaxInt64
)

// A pool for throttling IO
type IOThrottlerPool struct {
	// A channel for setting the pools bandwith
	bandwidthSettingChan chan Bandwidth
	// A channel for allocating bandwidth
	bandwidthAllocatorChan chan Bandwidth
	// A channel for returning unused bandwidth to server
	bandwidthFreeChan chan Bandwidth
	// A channel for getting a count of the clients
	// A pool only accumulates bandwidth if the pool is non-empty
	clientCountChan chan int64
	// A channel for getting pool release messages
	releasePoolChan chan bool
}

// Construct a new IO throttling pool
// The bandwidth for this pool will be limited to 'bandwidth'
func NewIOThrottlerPool(bandwidth Bandwidth) *IOThrottlerPool {

	pool := &IOThrottlerPool{make(chan Bandwidth), make(chan Bandwidth), make(chan Bandwidth), make(chan int64), make(chan bool)}

	go throttlerPoolDriver(pool)

	pool.bandwidthSettingChan <- bandwidth

	return pool
}

func throttlerPoolDriver(pool *IOThrottlerPool) {

	// These will all be recalculated as soon as the bandwidth is set
	clientCount := int64(0)
	currentBandwidth := Bandwidth(0)
	totalbandwidth := Bandwidth(0)
	allocationSize := Bandwidth(0)
	timeout := time.NewTicker(time.Second)
	var thisBandwidthAllocatorChan chan Bandwidth = nil

	recalculateAllocationSize := func() {
		if currentBandwidth == Unlimited {
			totalbandwidth = Unlimited
		}

		if totalbandwidth == Unlimited {
			allocationSize = Unlimited
		} else {

			// Calculate how much bandwidth each consumer will get
			// We divvy the available bandwidth among the existing
			// clients but leave a bit of room in case more clients
			// connect in the mean time. This greatly improves
			// performance
			if clientCount == 0 {
				allocationSize = totalbandwidth / 2
			} else {
				allocationSize = totalbandwidth / Bandwidth(clientCount*2)
			}

			// Even if we have a negative totalbandwidth we never want to
			// allocate negative bandwidth to members of our pool
			if allocationSize < 0 {
				allocationSize = 0
			}

			// If we do have some bandwidth make sure we at least allocate 1 byte
			if allocationSize == 0 && totalbandwidth > 0 {
				allocationSize = 1
			}
		}

		if allocationSize > 0 {
			// Since we have bandwidth to allocate we can select on
			// the bandwidth allocator chan
			thisBandwidthAllocatorChan = pool.bandwidthAllocatorChan
		} else {
			// We've allocate all out bandwidth so we need to wait for
			// more
			thisBandwidthAllocatorChan = nil
		}
	}

	for {
		select {
		// Release the pool
		case release := <-pool.releasePoolChan:
			if release {
				close(pool.bandwidthAllocatorChan)
				close(pool.bandwidthFreeChan)
				// Don't close the clientCountChan it's not needed and it
				// complicates the code (two different functions need to recover
				// the panic if it's closed
				pool.releasePoolChan <- true
				close(pool.releasePoolChan)
				close(pool.clientCountChan)
				return
			}

		// Register a new client
		case increment := <-pool.clientCountChan:
			// We got our first client
			// We start the timer as soon as we get our first client
			if clientCount == 0 {
				timeout.Reset(time.Second)
			}
			clientCount += increment
			// Our last client left so stop the timer
			if clientCount == 0 {
				timeout.Stop()
			}
			recalculateAllocationSize()

		// Set the new bandwidth
		case newBandwidth := <-pool.bandwidthSettingChan:
			// If we've accumulated more bandwidth then the new amount we
			// truncate the totalbandwidth to the new set amount. This is
			// important if the totalbandwidth is much larger than the
			// new bandwidth value we could end up not really respecting the
			// new bandwidth setting. An extreme example of this is if the
			// old bandwidth was set to Unlimited (totalbandwidth would be
			// Unlimited)
			//
			// If the totalbandwidth is less than the new bandwidth setting
			// we want to bring it up to the new bandwidth value so clients
			// can immediately use the new available bandwidth
			totalbandwidth = newBandwidth

			// Update the current bandwidth
			currentBandwidth = newBandwidth

			recalculateAllocationSize()

		// Allocate some bandwidth
		case thisBandwidthAllocatorChan <- allocationSize:
			if Unlimited != totalbandwidth {
				totalbandwidth -= allocationSize

				recalculateAllocationSize()
			}

		// Get unused bandwidth back from client
		case returnSize := <-pool.bandwidthFreeChan:
			if Unlimited != totalbandwidth {
				totalbandwidth += returnSize
			}

			recalculateAllocationSize()

		// Get more bandwidth to allocate
		case <-timeout.C:
			if clientCount > 0 {
				if Unlimited != totalbandwidth {
					// Get a new allotment of bandwidth
					totalbandwidth += currentBandwidth

					recalculateAllocationSize()
				}
			}
		}
	}
}

// Release the IOThrottlerPool all bandwidth
func (pool *IOThrottlerPool) ReleasePool() {
	// If pool.releasePoolChan is already closed (called ReleasePool more than
	// once) then sending to it will panic so just swallow the panic
	defer func() {
		recover()
	}()
	pool.releasePoolChan <- true
	<-pool.releasePoolChan
}

// Sets the IOThrottlerPool's bandwith rate
func (pool *IOThrottlerPool) SetBandwidth(bandwith Bandwidth) {
	pool.bandwidthSettingChan <- bandwith
}

// Returns the first error or nil if neither are errors
func orErrors(er0, er1 error) error {
	if er0 != nil {
		return er0
	}
	return er1
}

/*
 * Updates the client count for a pool return error on failure
 */
func twiddleClientCount(p *IOThrottlerPool, change int64) (err error) {
	// When the pool has been released the server closes clientCountChan
	// so our channel send will panic. We want to set the return error
	defer func() {
		if r := recover(); r != nil {
			err = errors.New("Pool has been released")
		}
	}()
	// Update the client count
	p.clientCountChan <- change

	return nil
}

// A ReadCloser that will respect the bandwidth limitations of the IOThrottlerPool
type throttledReadCloser struct {
	origReadCloser io.ReadCloser
	pool           *IOThrottlerPool
}

// A WriteCloser that will respect the bandwidth limitations of the IOThrottlerPool
type throttledWriteCloser struct {
	origWriteCloser io.WriteCloser
	pool            *IOThrottlerPool
}

// A ReadWriteCloser that will respect the bandwidth limitations of the IOThrottlerPool
type throttledReadWriteCloser struct {
	throttledReadCloser
	throttledWriteCloser
}

// Read method for the throttledReadCloser
func (t *throttledReadCloser) Read(b []byte) (int, error) {
	// Get some bandwidth
	allocation, ok := <-t.pool.bandwidthAllocatorChan
	if !ok {
		// Pool has been released
		return 0, errors.New("Pool has been released")
	}

	// Calculate how much we can read
	toRead := Bandwidth(len(b))
	if allocation < Bandwidth(len(b)) {
		toRead = allocation
	}

	// Do the limited read
	n, err := t.origReadCloser.Read(b[:toRead])

	// Free up what we didn't use
	if Bandwidth(n) < allocation && allocation != Unlimited {
		t.pool.bandwidthFreeChan <- allocation - Bandwidth(n)
	}

	return n, err
}

// Write method for the throttledWriteCloser
func (t *throttledWriteCloser) Write(data []byte) (int, error) {

	// Write must either write len(data) bytes or return an error
	allocation := Bandwidth(0)

	for allocation < Bandwidth(len(data)) {
		// Get some bandwidth
		thisAllocation, ok := <-t.pool.bandwidthAllocatorChan
		if !ok {
			// Pool has been released
			return 0, errors.New("Pool has been released")
		}
		allocation += thisAllocation
	}

	// Do the write
	n, err := t.origWriteCloser.Write(data)

	// Free up what we didn't use
	if Bandwidth(n) < allocation && allocation != Unlimited {
		t.pool.bandwidthFreeChan <- allocation - Bandwidth(n)
	}

	return n, err
}

// Close method for the throttledReadCloser
func (t *throttledReadCloser) Close() error {
	// Unregister with the pool
	err := twiddleClientCount(t.pool, -1)

	return orErrors(err, t.origReadCloser.Close())
}

// Close method for the throttledWriteCloser
func (t *throttledWriteCloser) Close() error {
	// Unregister with the pool
	err := twiddleClientCount(t.pool, -1)

	return orErrors(err, t.origWriteCloser.Close())
}

// Close method for the throttledReadWriteCloser
func (t *throttledReadWriteCloser) Close() error {
	// In this case we really have two copies of all the data
	// It really doesn't matter which we use as the reader and writer hold the
	// same data

	// Unregister with the pool
	err := twiddleClientCount(t.throttledReadCloser.pool, -1)

	return orErrors(err, t.throttledReadCloser.origReadCloser.Close())
}

// Add a io.ReadCloser to the pool. The returned io.ReadCloser shares the
// IOThrottlerPool's bandwidth with other items in the pool.
func (p *IOThrottlerPool) AddReader(reader io.ReadCloser) (io.ReadCloser, error) {
	// Register with the pool
	err := twiddleClientCount(p, 1)
	if err != nil {
		return nil, err
	}

	return &throttledReadCloser{reader, p}, nil
}

// Add a io.WriteCloser to the pool. The returned io.WriteCloser shares the
// IOThrottlerPool's bandwidth with other items in the pool.
func (p *IOThrottlerPool) AddWriter(writer io.WriteCloser) (io.WriteCloser, error) {
	// Register with the pool
	err := twiddleClientCount(p, 1)
	if err != nil {
		return nil, err
	}

	return &throttledWriteCloser{writer, p}, nil
}

// Add a io.ReadWriteCloser to the pool. The returned io.ReadWriteCloser shares the
// IOThrottlerPool's bandwidth with other items in the pool.
func (p *IOThrottlerPool) AddReadWriter(readWriter io.ReadWriteCloser) (io.ReadWriteCloser, error) {
	// Register with the pool
	err := twiddleClientCount(p, 1)
	if err != nil {
		return nil, err
	}

	return &throttledReadWriteCloser{throttledReadCloser{readWriter, p},
		throttledWriteCloser{readWriter, p}}, err
}

// Add a net.Conn to the pool. The returned net.Conn shares the
// IOThrottlerPool's bandwidth with other items in the pool.
type throttledConn struct {
	throttledReadWriteCloser
	originalConn net.Conn
}

// Implements the net.Conn LocalAddr method
func (c *throttledConn) LocalAddr() net.Addr {
	return c.originalConn.LocalAddr()
}

// Implements the net.Conn RemoteAddr method
func (c *throttledConn) RemoteAddr() net.Addr {
	return c.originalConn.RemoteAddr()
}

// Implements the net.Conn SetDeadline method
func (c *throttledConn) SetDeadline(t time.Time) error {
	return c.originalConn.SetDeadline(t)
}

// Implements the net.Conn SetReadDeadline method
func (c *throttledConn) SetReadDeadline(t time.Time) error {
	return c.originalConn.SetReadDeadline(t)
}

// Implements the net.Conn SetWriteDeadline method
func (c *throttledConn) SetWriteDeadline(t time.Time) error {
	return c.originalConn.SetWriteDeadline(t)
}

// Restrict the network connection to the bandwidth limitations of the IOThrottlerPool
func (p *IOThrottlerPool) AddConn(conn net.Conn) (net.Conn, error) {

	rwCloser, err := p.AddReadWriter(conn)
	if err != nil {
		return nil, err
	}
	throttledRWC, ok := rwCloser.(*throttledReadWriteCloser)
	if !ok {
		log.Fatalf("Programming error, expecting *throttledReadWriteCloser but got %v", rwCloser)
	}

	return &throttledConn{*throttledRWC, conn}, nil
}
<commit_msg>Don't start the time until we get our first client<commit_after>// Copyright 2012 Evan Farrer. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
	Package iothrottler implements application IO throttling.
*/

package iothrottler

import (
	"errors"
	"io"
	"log"
	"math"
	"net"
	"time"
)

// The Bandwidth type represents a bandwidth quantity in bytes per second.
// Sub-byte per seconds values are not supported
type Bandwidth int64

const (
	// Bytes per second
	BytesPerSecond Bandwidth = 1
	// Kilobits per second
	Kbps = BytesPerSecond * (1024 / 8)
	// Megabits per second
	Mbps = Kbps * 1024
	// Gigabits per second
	Gbps = Mbps * 1024
	// Unlimited bandwidth
	Unlimited = math.MaxInt64
)

// A pool for throttling IO
type IOThrottlerPool struct {
	// A channel for setting the pools bandwith
	bandwidthSettingChan chan Bandwidth
	// A channel for allocating bandwidth
	bandwidthAllocatorChan chan Bandwidth
	// A channel for returning unused bandwidth to server
	bandwidthFreeChan chan Bandwidth
	// A channel for getting a count of the clients
	// A pool only accumulates bandwidth if the pool is non-empty
	clientCountChan chan int64
	// A channel for getting pool release messages
	releasePoolChan chan bool
}

// Construct a new IO throttling pool
// The bandwidth for this pool will be limited to 'bandwidth'
func NewIOThrottlerPool(bandwidth Bandwidth) *IOThrottlerPool {

	pool := &IOThrottlerPool{make(chan Bandwidth), make(chan Bandwidth), make(chan Bandwidth), make(chan int64), make(chan bool)}

	go throttlerPoolDriver(pool)

	pool.bandwidthSettingChan <- bandwidth

	return pool
}

func throttlerPoolDriver(pool *IOThrottlerPool) {

	// These will all be recalculated as soon as the bandwidth is set
	clientCount := int64(0)
	currentBandwidth := Bandwidth(0)
	totalbandwidth := Bandwidth(0)
	allocationSize := Bandwidth(0)
	timeout := time.NewTicker(time.Second)
	var thisBandwidthAllocatorChan chan Bandwidth = nil

	// Start the timer until we get the first client
	timeout.Stop()

	recalculateAllocationSize := func() {
		if currentBandwidth == Unlimited {
			totalbandwidth = Unlimited
		}

		if totalbandwidth == Unlimited {
			allocationSize = Unlimited
		} else {

			// Calculate how much bandwidth each consumer will get
			// We divvy the available bandwidth among the existing
			// clients but leave a bit of room in case more clients
			// connect in the mean time. This greatly improves
			// performance
			if clientCount == 0 {
				allocationSize = totalbandwidth / 2
			} else {
				allocationSize = totalbandwidth / Bandwidth(clientCount*2)
			}

			// Even if we have a negative totalbandwidth we never want to
			// allocate negative bandwidth to members of our pool
			if allocationSize < 0 {
				allocationSize = 0
			}

			// If we do have some bandwidth make sure we at least allocate 1 byte
			if allocationSize == 0 && totalbandwidth > 0 {
				allocationSize = 1
			}
		}

		if allocationSize > 0 {
			// Since we have bandwidth to allocate we can select on
			// the bandwidth allocator chan
			thisBandwidthAllocatorChan = pool.bandwidthAllocatorChan
		} else {
			// We've allocate all out bandwidth so we need to wait for
			// more
			thisBandwidthAllocatorChan = nil
		}
	}

	for {
		select {
		// Release the pool
		case release := <-pool.releasePoolChan:
			if release {
				close(pool.bandwidthAllocatorChan)
				close(pool.bandwidthFreeChan)
				// Don't close the clientCountChan it's not needed and it
				// complicates the code (two different functions need to recover
				// the panic if it's closed
				pool.releasePoolChan <- true
				close(pool.releasePoolChan)
				close(pool.clientCountChan)
				return
			}

		// Register a new client
		case increment := <-pool.clientCountChan:
			// We got our first client
			// We start the timer as soon as we get our first client
			if clientCount == 0 {
				timeout.Reset(time.Second)
			}
			clientCount += increment
			// Our last client left so stop the timer
			if clientCount == 0 {
				timeout.Stop()
			}
			recalculateAllocationSize()

		// Set the new bandwidth
		case newBandwidth := <-pool.bandwidthSettingChan:
			// If we've accumulated more bandwidth then the new amount we
			// truncate the totalbandwidth to the new set amount. This is
			// important if the totalbandwidth is much larger than the
			// new bandwidth value we could end up not really respecting the
			// new bandwidth setting. An extreme example of this is if the
			// old bandwidth was set to Unlimited (totalbandwidth would be
			// Unlimited)
			//
			// If the totalbandwidth is less than the new bandwidth setting
			// we want to bring it up to the new bandwidth value so clients
			// can immediately use the new available bandwidth
			totalbandwidth = newBandwidth

			// Update the current bandwidth
			currentBandwidth = newBandwidth

			recalculateAllocationSize()

		// Allocate some bandwidth
		case thisBandwidthAllocatorChan <- allocationSize:
			if Unlimited != totalbandwidth {
				totalbandwidth -= allocationSize

				recalculateAllocationSize()
			}

		// Get unused bandwidth back from client
		case returnSize := <-pool.bandwidthFreeChan:
			if Unlimited != totalbandwidth {
				totalbandwidth += returnSize
			}

			recalculateAllocationSize()

		// Get more bandwidth to allocate
		case <-timeout.C:
			if clientCount > 0 {
				if Unlimited != totalbandwidth {
					// Get a new allotment of bandwidth
					totalbandwidth += currentBandwidth

					recalculateAllocationSize()
				}
			}
		}
	}
}

// Release the IOThrottlerPool all bandwidth
func (pool *IOThrottlerPool) ReleasePool() {
	// If pool.releasePoolChan is already closed (called ReleasePool more than
	// once) then sending to it will panic so just swallow the panic
	defer func() {
		recover()
	}()
	pool.releasePoolChan <- true
	<-pool.releasePoolChan
}

// Sets the IOThrottlerPool's bandwith rate
func (pool *IOThrottlerPool) SetBandwidth(bandwith Bandwidth) {
	pool.bandwidthSettingChan <- bandwith
}

// Returns the first error or nil if neither are errors
func orErrors(er0, er1 error) error {
	if er0 != nil {
		return er0
	}
	return er1
}

/*
 * Updates the client count for a pool return error on failure
 */
func twiddleClientCount(p *IOThrottlerPool, change int64) (err error) {
	// When the pool has been released the server closes clientCountChan
	// so our channel send will panic. We want to set the return error
	defer func() {
		if r := recover(); r != nil {
			err = errors.New("Pool has been released")
		}
	}()
	// Update the client count
	p.clientCountChan <- change

	return nil
}

// A ReadCloser that will respect the bandwidth limitations of the IOThrottlerPool
type throttledReadCloser struct {
	origReadCloser io.ReadCloser
	pool           *IOThrottlerPool
}

// A WriteCloser that will respect the bandwidth limitations of the IOThrottlerPool
type throttledWriteCloser struct {
	origWriteCloser io.WriteCloser
	pool            *IOThrottlerPool
}

// A ReadWriteCloser that will respect the bandwidth limitations of the IOThrottlerPool
type throttledReadWriteCloser struct {
	throttledReadCloser
	throttledWriteCloser
}

// Read method for the throttledReadCloser
func (t *throttledReadCloser) Read(b []byte) (int, error) {
	// Get some bandwidth
	allocation, ok := <-t.pool.bandwidthAllocatorChan
	if !ok {
		// Pool has been released
		return 0, errors.New("Pool has been released")
	}

	// Calculate how much we can read
	toRead := Bandwidth(len(b))
	if allocation < Bandwidth(len(b)) {
		toRead = allocation
	}

	// Do the limited read
	n, err := t.origReadCloser.Read(b[:toRead])

	// Free up what we didn't use
	if Bandwidth(n) < allocation && allocation != Unlimited {
		t.pool.bandwidthFreeChan <- allocation - Bandwidth(n)
	}

	return n, err
}

// Write method for the throttledWriteCloser
func (t *throttledWriteCloser) Write(data []byte) (int, error) {

	// Write must either write len(data) bytes or return an error
	allocation := Bandwidth(0)

	for allocation < Bandwidth(len(data)) {
		// Get some bandwidth
		thisAllocation, ok := <-t.pool.bandwidthAllocatorChan
		if !ok {
			// Pool has been released
			return 0, errors.New("Pool has been released")
		}
		allocation += thisAllocation
	}

	// Do the write
	n, err := t.origWriteCloser.Write(data)

	// Free up what we didn't use
	if Bandwidth(n) < allocation && allocation != Unlimited {
		t.pool.bandwidthFreeChan <- allocation - Bandwidth(n)
	}

	return n, err
}

// Close method for the throttledReadCloser
func (t *throttledReadCloser) Close() error {
	// Unregister with the pool
	err := twiddleClientCount(t.pool, -1)

	return orErrors(err, t.origReadCloser.Close())
}

// Close method for the throttledWriteCloser
func (t *throttledWriteCloser) Close() error {
	// Unregister with the pool
	err := twiddleClientCount(t.pool, -1)

	return orErrors(err, t.origWriteCloser.Close())
}

// Close method for the throttledReadWriteCloser
func (t *throttledReadWriteCloser) Close() error {
	// In this case we really have two copies of all the data
	// It really doesn't matter which we use as the reader and writer hold the
	// same data

	// Unregister with the pool
	err := twiddleClientCount(t.throttledReadCloser.pool, -1)

	return orErrors(err, t.throttledReadCloser.origReadCloser.Close())
}

// Add a io.ReadCloser to the pool. The returned io.ReadCloser shares the
// IOThrottlerPool's bandwidth with other items in the pool.
func (p *IOThrottlerPool) AddReader(reader io.ReadCloser) (io.ReadCloser, error) {
	// Register with the pool
	err := twiddleClientCount(p, 1)
	if err != nil {
		return nil, err
	}

	return &throttledReadCloser{reader, p}, nil
}

// Add a io.WriteCloser to the pool. The returned io.WriteCloser shares the
// IOThrottlerPool's bandwidth with other items in the pool.
func (p *IOThrottlerPool) AddWriter(writer io.WriteCloser) (io.WriteCloser, error) {
	// Register with the pool
	err := twiddleClientCount(p, 1)
	if err != nil {
		return nil, err
	}

	return &throttledWriteCloser{writer, p}, nil
}

// Add a io.ReadWriteCloser to the pool. The returned io.ReadWriteCloser shares the
// IOThrottlerPool's bandwidth with other items in the pool.
func (p *IOThrottlerPool) AddReadWriter(readWriter io.ReadWriteCloser) (io.ReadWriteCloser, error) {
	// Register with the pool
	err := twiddleClientCount(p, 1)
	if err != nil {
		return nil, err
	}

	return &throttledReadWriteCloser{throttledReadCloser{readWriter, p},
		throttledWriteCloser{readWriter, p}}, err
}

// Add a net.Conn to the pool. The returned net.Conn shares the
// IOThrottlerPool's bandwidth with other items in the pool.
type throttledConn struct {
	throttledReadWriteCloser
	originalConn net.Conn
}

// Implements the net.Conn LocalAddr method
func (c *throttledConn) LocalAddr() net.Addr {
	return c.originalConn.LocalAddr()
}

// Implements the net.Conn RemoteAddr method
func (c *throttledConn) RemoteAddr() net.Addr {
	return c.originalConn.RemoteAddr()
}

// Implements the net.Conn SetDeadline method
func (c *throttledConn) SetDeadline(t time.Time) error {
	return c.originalConn.SetDeadline(t)
}

// Implements the net.Conn SetReadDeadline method
func (c *throttledConn) SetReadDeadline(t time.Time) error {
	return c.originalConn.SetReadDeadline(t)
}

// Implements the net.Conn SetWriteDeadline method
func (c *throttledConn) SetWriteDeadline(t time.Time) error {
	return c.originalConn.SetWriteDeadline(t)
}

// Restrict the network connection to the bandwidth limitations of the IOThrottlerPool
func (p *IOThrottlerPool) AddConn(conn net.Conn) (net.Conn, error) {

	rwCloser, err := p.AddReadWriter(conn)
	if err != nil {
		return nil, err
	}
	throttledRWC, ok := rwCloser.(*throttledReadWriteCloser)
	if !ok {
		log.Fatalf("Programming error, expecting *throttledReadWriteCloser but got %v", rwCloser)
	}

	return &throttledConn{*throttledRWC, conn}, nil
}
<|endoftext|>


<commit_before>package felixcheck

import (
	"fmt"
	"net"
	"strings"
	"time"

	"net/http"
	"net/url"

	"github.com/aleasoluciones/gosnmpquerier"
	"github.com/aleasoluciones/simpleamqp"
	"github.com/bigdatadev/goryman"
	"github.com/tatsushid/go-fastping"
)

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

const (
	sysName     = "1.3.6.1.2.1.1.5.0"
	maxPingTime = 4 * time.Second
)

type CheckFunction func() goryman.Event

func (f CheckFunction) Tags(tags ...string) CheckFunction {
	return func() goryman.Event {
		result := f()
		result.Tags = tags
		return result
	}
}

func (f CheckFunction) Ttl(ttl float32) CheckFunction {
	return func() goryman.Event {
		result := f()
		result.Ttl = ttl
		return result
	}
}

func NewPingChecker(host, service, ip string) CheckFunction {
	return func() goryman.Event {
		var retRtt time.Duration = 0
		var result goryman.Event = goryman.Event{Host: host, Service: service, State: "critical"}

		p := fastping.NewPinger()
		p.MaxRTT = maxPingTime
		ra, err := net.ResolveIPAddr("ip4:icmp", ip)
		if err != nil {
			result.Description = err.Error()
		}

		p.AddIPAddr(ra)
		p.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
			result.State = "ok"
			result.Metric = float32(retRtt.Nanoseconds() / 1e6)
		}

		err = p.Run()
		if err != nil {
			result.Description = err.Error()
		}
		return result
	}
}

type TcpCheckerConf struct {
	retries   int
	timeout   time.Duration
	retrytime time.Duration
}

var DefaultTcpCheckConf = TcpCheckerConf{
	retries:   3,
	timeout:   2 * time.Second,
	retrytime: 1 * time.Second,
}

func NewTcpPortChecker(host, service, ip string, port int, conf TcpCheckerConf) CheckFunction {
	return func() goryman.Event {
		var err error
		var conn net.Conn

		for attempt := 0; attempt < conf.retries; attempt++ {
			conn, err = net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), conf.timeout)
			if err == nil {
				conn.Close()
				return goryman.Event{Host: host, Service: service, State: "ok"}
			}
			time.Sleep(conf.retrytime)
		}
		return goryman.Event{Host: host, Service: service, State: "critical"}
	}
}

func NewHttpChecker(host, service, url string, expectedStatusCode int) CheckFunction {
	return func() goryman.Event {
		var t1 = time.Now()

		response, err := http.Get(url)
		milliseconds := float32((time.Now().Sub(t1)).Nanoseconds() / 1e6)
		result := goryman.Event{Host: host, Service: service, State: "critical", Metric: milliseconds}
		if err != nil {
			result.Description = err.Error()
		} else {
			if response.Body != nil {
				defer response.Body.Close()
			}
			if response.StatusCode == expectedStatusCode {
				result.State = "ok"
			} else {
				result.Description = fmt.Sprintf("Response %d", response.StatusCode)
			}
		}
		return result
	}
}

type SnmpCheckerConf struct {
	retries    int
	timeout    time.Duration
	oidToCheck string
}

var DefaultSnmpCheckConf = SnmpCheckerConf{
	retries:    1,
	timeout:    1 * time.Second,
	oidToCheck: sysName,
}

func NewSnmpChecker(host, service, ip, community string, conf SnmpCheckerConf, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		_, err := snmpQuerier.Get(ip, community, []string{conf.oidToCheck}, conf.timeout, conf.retries)
		if err == nil {
			return goryman.Event{Host: host, Service: service, State: "ok", Description: err.Error()}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewC4CMTSTempChecker(host, service, ip, community string, maxAllowedTemp int, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		result, err := snmpQuerier.Walk(ip, community, "1.3.6.1.4.1.4998.1.1.10.1.4.2.1.29", 2*time.Second, 1)

		if err == nil {
			max := 0
			for _, r := range result {
				if r.Value.(int) != 999 && r.Value.(int) > max {
					max = r.Value.(int)
				}
			}
			var state string = "critical"
			if max < maxAllowedTemp {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(max)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func getMaxValueFromSnmpWalk(oid, ip, community string, snmpQuerier gosnmpquerier.SyncQuerier) (uint, error) {
	result, err := snmpQuerier.Walk(ip, community, oid, 2*time.Second, 1)
	if err == nil {
		max := uint(0)
		for _, r := range result {
			if r.Value.(uint) > max {
				max = r.Value.(uint)
			}
		}
		return max, nil
	} else {
		return 0, err
	}
}

func NewJuniperTempChecker(host, service, ip, community string, maxAllowedTemp uint, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		max, err := getMaxValueFromSnmpWalk("1.3.6.1.4.1.2636.3.1.13.1.7", ip, community, snmpQuerier)
		if err == nil {
			var state string = "critical"
			if max < maxAllowedTemp {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(max)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewJuniperCpuChecker(host, service, ip, community string, maxAllowedTemp uint, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		max, err := getMaxValueFromSnmpWalk("1.3.6.1.4.1.2636.3.1.13.1.8", ip, community, snmpQuerier)
		if err == nil {
			var state string = "critical"
			if max < maxAllowedTemp {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(max)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewRabbitMQQueueLenCheck(host, service, amqpuri, queue string, max int) CheckFunction {
	return func() goryman.Event {
		queueInfo, err := simpleamqp.NewAmqpManagement(amqpuri).QueueInfo(queue)
		if err == nil {
			var state string = "critical"
			if queueInfo.Messages < max {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(queueInfo.Messages)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewMysqlConnectionCheck(host, service, mysqluri string) CheckFunction {
	return func() goryman.Event {
		u, err := url.Parse(mysqluri)
		if err != nil {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}

		password, _ := u.User.Password()
		hostAndPort := u.Host
		if !strings.Contains(hostAndPort, ":") {
			hostAndPort = hostAndPort + ":3306"
		}
		con, err := sql.Open("mysql", u.User.Username()+":"+password+"@"+"tcp("+hostAndPort+")"+u.Path)
		defer con.Close()
		if err != nil {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
		q := `select CURTIME()`
		row := con.QueryRow(q)
		var date string
		err = row.Scan(&date)
		if err != nil {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
		return goryman.Event{Host: host, Service: service, State: "ok"}
	}
}
<commit_msg>Spike generic check with moving mean calculation<commit_after>package felixcheck

import (
	"fmt"
	"net"
	"strings"
	"time"

	"container/ring"
	"net/http"
	"net/url"

	"github.com/aleasoluciones/gosnmpquerier"
	"github.com/aleasoluciones/simpleamqp"
	"github.com/bigdatadev/goryman"
	"github.com/tatsushid/go-fastping"
)

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

const (
	sysName     = "1.3.6.1.2.1.1.5.0"
	maxPingTime = 4 * time.Second
)

type CheckFunction func() goryman.Event

func (f CheckFunction) Tags(tags ...string) CheckFunction {
	return func() goryman.Event {
		result := f()
		result.Tags = tags
		return result
	}
}

func (f CheckFunction) Ttl(ttl float32) CheckFunction {
	return func() goryman.Event {
		result := f()
		result.Ttl = ttl
		return result
	}
}

func NewPingChecker(host, service, ip string) CheckFunction {
	return func() goryman.Event {
		var retRtt time.Duration = 0
		var result goryman.Event = goryman.Event{Host: host, Service: service, State: "critical"}

		p := fastping.NewPinger()
		p.MaxRTT = maxPingTime
		ra, err := net.ResolveIPAddr("ip4:icmp", ip)
		if err != nil {
			result.Description = err.Error()
		}

		p.AddIPAddr(ra)
		p.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
			result.State = "ok"
			result.Metric = float32(retRtt.Nanoseconds() / 1e6)
		}

		err = p.Run()
		if err != nil {
			result.Description = err.Error()
		}
		return result
	}
}

type TcpCheckerConf struct {
	retries   int
	timeout   time.Duration
	retrytime time.Duration
}

var DefaultTcpCheckConf = TcpCheckerConf{
	retries:   3,
	timeout:   2 * time.Second,
	retrytime: 1 * time.Second,
}

func NewTcpPortChecker(host, service, ip string, port int, conf TcpCheckerConf) CheckFunction {
	return func() goryman.Event {
		var err error
		var conn net.Conn

		for attempt := 0; attempt < conf.retries; attempt++ {
			conn, err = net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), conf.timeout)
			if err == nil {
				conn.Close()
				return goryman.Event{Host: host, Service: service, State: "ok"}
			}
			time.Sleep(conf.retrytime)
		}
		return goryman.Event{Host: host, Service: service, State: "critical"}
	}
}

func NewHttpChecker(host, service, url string, expectedStatusCode int) CheckFunction {
	return func() goryman.Event {
		var t1 = time.Now()

		response, err := http.Get(url)
		milliseconds := float32((time.Now().Sub(t1)).Nanoseconds() / 1e6)
		result := goryman.Event{Host: host, Service: service, State: "critical", Metric: milliseconds}
		if err != nil {
			result.Description = err.Error()
		} else {
			if response.Body != nil {
				defer response.Body.Close()
			}
			if response.StatusCode == expectedStatusCode {
				result.State = "ok"
			} else {
				result.Description = fmt.Sprintf("Response %d", response.StatusCode)
			}
		}
		return result
	}
}

type SnmpCheckerConf struct {
	retries    int
	timeout    time.Duration
	oidToCheck string
}

var DefaultSnmpCheckConf = SnmpCheckerConf{
	retries:    1,
	timeout:    1 * time.Second,
	oidToCheck: sysName,
}

func NewSnmpChecker(host, service, ip, community string, conf SnmpCheckerConf, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		_, err := snmpQuerier.Get(ip, community, []string{conf.oidToCheck}, conf.timeout, conf.retries)
		if err == nil {
			return goryman.Event{Host: host, Service: service, State: "ok", Description: err.Error()}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewC4CMTSTempChecker(host, service, ip, community string, maxAllowedTemp int, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		result, err := snmpQuerier.Walk(ip, community, "1.3.6.1.4.1.4998.1.1.10.1.4.2.1.29", 2*time.Second, 1)

		if err == nil {
			max := 0
			for _, r := range result {
				if r.Value.(int) != 999 && r.Value.(int) > max {
					max = r.Value.(int)
				}
			}
			var state string = "critical"
			if max < maxAllowedTemp {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(max)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func getMaxValueFromSnmpWalk(oid, ip, community string, snmpQuerier gosnmpquerier.SyncQuerier) (uint, error) {
	result, err := snmpQuerier.Walk(ip, community, oid, 2*time.Second, 1)
	if err == nil {
		max := uint(0)
		for _, r := range result {
			if r.Value.(uint) > max {
				max = r.Value.(uint)
			}
		}
		return max, nil
	} else {
		return 0, err
	}
}

func NewJuniperTempChecker(host, service, ip, community string, maxAllowedTemp uint, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		max, err := getMaxValueFromSnmpWalk("1.3.6.1.4.1.2636.3.1.13.1.7", ip, community, snmpQuerier)
		if err == nil {
			var state string = "critical"
			if max < maxAllowedTemp {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(max)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewJuniperCpuChecker(host, service, ip, community string, maxAllowedTemp uint, snmpQuerier gosnmpquerier.SyncQuerier) CheckFunction {
	return func() goryman.Event {
		max, err := getMaxValueFromSnmpWalk("1.3.6.1.4.1.2636.3.1.13.1.8", ip, community, snmpQuerier)
		if err == nil {
			var state string = "critical"
			if max < maxAllowedTemp {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(max)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewRabbitMQQueueLenCheck(host, service, amqpuri, queue string, max int) CheckFunction {
	return func() goryman.Event {
		queueInfo, err := simpleamqp.NewAmqpManagement(amqpuri).QueueInfo(queue)
		if err == nil {
			var state string = "critical"
			if queueInfo.Messages < max {
				state = "ok"
			}
			return goryman.Event{Host: host, Service: service, State: state, Metric: float32(queueInfo.Messages)}
		} else {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
	}
}

func NewMysqlConnectionCheck(host, service, mysqluri string) CheckFunction {
	return func() goryman.Event {
		u, err := url.Parse(mysqluri)
		if err != nil {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}

		password, _ := u.User.Password()
		hostAndPort := u.Host
		if !strings.Contains(hostAndPort, ":") {
			hostAndPort = hostAndPort + ":3306"
		}
		con, err := sql.Open("mysql", u.User.Username()+":"+password+"@"+"tcp("+hostAndPort+")"+u.Path)
		defer con.Close()
		if err != nil {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
		q := `select CURTIME()`
		row := con.QueryRow(q)
		var date string
		err = row.Scan(&date)
		if err != nil {
			return goryman.Event{Host: host, Service: service, State: "critical", Description: err.Error()}
		}
		return goryman.Event{Host: host, Service: service, State: "ok"}
	}
}

type ObtainMetricFunction func() float32
type CalculateStateFunction func(float32, float32) string

func NewGenericCheck(host, service string, metricFunc ObtainMetricFunction, stateFunc CalculateStateFunction, windowSize int) CheckFunction {
	window := ring.New(windowSize)

	return func() goryman.Event {
		value := metricFunc()
		var state string = stateFunc(windowMean(window), value)
		window.Value = value
		window = window.Next()
		return goryman.Event{Host: host, Service: service, State: state, Metric: value}
	}
}

func windowMean(window *ring.Ring) float32 {
	numPoints := 0
	sum := float32(0)

	window.Do(func(x interface{}) {
		switch x.(type) {
		case float32:
			sum = sum + x.(float32)
			numPoints = numPoints + 1
		}
	})

	if numPoints == 0 {
		return 0
	}
	return sum / float32(numPoints)
}
<|endoftext|>


<commit_before>package ipnets

import (
	"net"
	"testing"
)

func TestIpnets(t *testing.T) {
	stringNets := "192.168.3.4,172.17.4.127/24 , 10.0.0.1/30, fe81::20f:75fa:fe1a:1af2/64,fe99::20f:75fa:fe1a:1af2"
	nets := new(IPNets)
	// test parsing
	if err := nets.Set(stringNets); err != nil {
		t.Fatal("Failed to parse networks from string: ", err)
	}
	if l := len(*nets); l != 5 {
		t.Fatal("The length of IPNets should be 5, not ", l)
	}
	// test individual nets
	for i, proposed := range []string{"192.168.3.4/32", "172.17.4.0/24", "10.0.0.0/30", "fe81::/64", "fe99::20f:75fa:fe1a:1af2/128"} {
		if parsed := (*nets)[i].String(); parsed != proposed {
			t.Errorf("The element %d parsed as %v, but should be %v", i, parsed, proposed)
		}
	}
	// test networks matching
	for _, s := range []string{"192.168.3.4", "172.17.4.0", "172.17.4.128", "172.17.4.255", "10.0.0.1", "fe81::ff", "fe99::20f:75fa:fe1a:1af2"} {
		if ip := net.ParseIP(s); !nets.Contains(ip) {
			t.Errorf("%s should match the IPNets %v", ip, nets)
		}
	}
	// test not matching networks
	for _, s := range []string{"192.168.3.5", "172.17.3.255", "172.17.5.0", "10.0.0.4", "fe82::ff", "fe99::20f:75fa:fe1a:1af3"} {
		if ip := net.ParseIP(s); nets.Contains(ip) {
			t.Errorf("%s should not match the IPNets %v", ip, nets)
		}
	}
	// test ivalid CIDR:
	if err := nets.Set("10.0.0.0/99"); err == nil {
		t.Error("Should fail while parsing invalid CIDR 10.0.0.0/99")
	}
}
<commit_msg>Achieved 100% test coverage by "testing" Strings()<commit_after>package ipnets

import (
	"net"
	"testing"
)

func TestIpnets(t *testing.T) {
	stringNets := "192.168.3.4,172.17.4.127/24 , 10.0.0.1/30, fe81::20f:75fa:fe1a:1af2/64,fe99::20f:75fa:fe1a:1af2"
	nets := new(IPNets)
	// test parsing
	if err := nets.Set(stringNets); err != nil {
		t.Fatal("Failed to parse networks from string: ", err)
	}
	if l := len(*nets); l != 5 {
		t.Fatal("The length of IPNets should be 5, not ", l)
	}
	// make 100% coverage by "testing" String()
	_ = nets.String()
	// test individual nets
	for i, proposed := range []string{"192.168.3.4/32", "172.17.4.0/24", "10.0.0.0/30", "fe81::/64", "fe99::20f:75fa:fe1a:1af2/128"} {
		if parsed := (*nets)[i].String(); parsed != proposed {
			t.Errorf("The element %d parsed as %v, but should be %v", i, parsed, proposed)
		}
	}
	// test networks matching
	for _, s := range []string{"192.168.3.4", "172.17.4.0", "172.17.4.128", "172.17.4.255", "10.0.0.1", "fe81::ff", "fe99::20f:75fa:fe1a:1af2"} {
		if ip := net.ParseIP(s); !nets.Contains(ip) {
			t.Errorf("%s should match the IPNets %v", ip, nets)
		}
	}
	// test not matching networks
	for _, s := range []string{"192.168.3.5", "172.17.3.255", "172.17.5.0", "10.0.0.4", "fe82::ff", "fe99::20f:75fa:fe1a:1af3"} {
		if ip := net.ParseIP(s); nets.Contains(ip) {
			t.Errorf("%s should not match the IPNets %v", ip, nets)
		}
	}
	// test ivalid CIDR:
	if err := nets.Set("10.0.0.0/99"); err == nil {
		t.Error("Should fail while parsing invalid CIDR 10.0.0.0/99")
	}
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 graphx

import (
	"fmt"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
)

// NamedEdge is a named MultiEdge.
type NamedEdge struct {
	Name string
	Edge *graph.MultiEdge
}

// NamedScope is a named Scope.
type NamedScope struct {
	Name  string
	Scope *graph.Scope
}

// ScopeTree is a convenient representation of the Scope-structure as a tree.
// Each ScopeTree may also be a subtree of another ScopeTree. The tree structure
// must respect the underlying Scope structure, i.e., if Scope 'a' has a parent
// 'b' then the ScopeTree for 'b' must have the ScopeTree for 'a' as a child.
type ScopeTree struct {
	// Scope is the named scope at the root of the (sub)tree.
	Scope NamedScope
	// Edges are the named edges directly under this scope.
	Edges []NamedEdge

	// Children are the scopes directly under this scope.
	Children []*ScopeTree
}

// NewScopeTree computes the ScopeTree for a set of edges.
func NewScopeTree(edges []*graph.MultiEdge) *ScopeTree {
	t := newTreeBuilder()
	for _, edge := range edges {
		t.addEdge(edge)
	}
	return t.root
}

// treeBuilder is a builder of a ScopeTree from any set of edges and
// scopes from the same graph.
type treeBuilder struct {
	root    *ScopeTree
	id2tree map[int]*ScopeTree
}

func newTreeBuilder() *treeBuilder {
	return &treeBuilder{
		id2tree: make(map[int]*ScopeTree),
	}
}

func (t *treeBuilder) addEdge(edge *graph.MultiEdge) {
	tree := t.addScope(edge.Scope())
	tree.Edges = append(tree.Edges, NamedEdge{Name: edge.Name(), Edge: edge})
}

func (t *treeBuilder) addScope(s *graph.Scope) *ScopeTree {
	if tree, exists := t.id2tree[s.ID()]; exists {
		return tree
	}

	tree := &ScopeTree{Scope: NamedScope{Name: s.Label, Scope: s}}
	t.id2tree[s.ID()] = tree

	if s.Parent == nil {
		t.root = tree
	} else {
		parent := t.addScope(s.Parent)
		parent.Children = append(parent.Children, tree)
	}
	return tree
}

// EnsureUniqueNames ensures that each name is unique within each ScopeTree
// recursively. Any conflict is resolved by adding '1, '2, etc to the name.
func EnsureUniqueNames(tree *ScopeTree) {
	seen := make(map[string]bool)
	for _, edge := range tree.Edges {
		edge.Name = findFreeName(seen, edge.Name)
		seen[edge.Name] = true
	}

	for _, s := range tree.Children {
		EnsureUniqueNames(s)

		s.Scope.Name = findFreeName(seen, s.Scope.Name)
		seen[s.Scope.Name] = true
	}
}

func findFreeName(seen map[string]bool, name string) string {
	if !seen[name] {
		return name
	}
	for i := 1; ; i++ {
		next := fmt.Sprintf("%v'%v", name, i)
		if !seen[next] {
			return next
		}
	}
}
<commit_msg>CR: [BEAM-3287] Use model pipelines in Go SDK<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 graphx

import (
	"fmt"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
)

// NamedEdge is a named MultiEdge.
type NamedEdge struct {
	Name string
	Edge *graph.MultiEdge
}

// NamedScope is a named Scope.
type NamedScope struct {
	Name  string
	Scope *graph.Scope
}

// ScopeTree is a convenient representation of the Scope-structure as a tree.
// Each ScopeTree may also be a subtree of another ScopeTree. The tree structure
// respects the underlying Scope structure, i.e., if Scope 'a' has a parent
// 'b' then the ScopeTree for 'b' must have the ScopeTree for 'a' as a child.
type ScopeTree struct {
	// Scope is the named scope at the root of the (sub)tree.
	Scope NamedScope
	// Edges are the named edges directly under this scope.
	Edges []NamedEdge

	// Children are the scopes directly under this scope.
	Children []*ScopeTree
}

// NewScopeTree computes the ScopeTree for a set of edges.
func NewScopeTree(edges []*graph.MultiEdge) *ScopeTree {
	t := newTreeBuilder()
	for _, edge := range edges {
		t.addEdge(edge)
	}
	return t.root
}

// treeBuilder is a builder of a ScopeTree from any set of edges and
// scopes from the same graph.
type treeBuilder struct {
	root    *ScopeTree
	id2tree map[int]*ScopeTree
}

func newTreeBuilder() *treeBuilder {
	return &treeBuilder{
		id2tree: make(map[int]*ScopeTree),
	}
}

func (t *treeBuilder) addEdge(edge *graph.MultiEdge) {
	tree := t.addScope(edge.Scope())
	tree.Edges = append(tree.Edges, NamedEdge{Name: edge.Name(), Edge: edge})
}

func (t *treeBuilder) addScope(s *graph.Scope) *ScopeTree {
	if tree, exists := t.id2tree[s.ID()]; exists {
		return tree
	}

	tree := &ScopeTree{Scope: NamedScope{Name: s.Label, Scope: s}}
	t.id2tree[s.ID()] = tree

	if s.Parent == nil {
		t.root = tree
	} else {
		parent := t.addScope(s.Parent)
		parent.Children = append(parent.Children, tree)
	}
	return tree
}

// EnsureUniqueNames ensures that each name is unique within each ScopeTree
// recursively. Any conflict is resolved by adding '1, '2, etc to the name.
func EnsureUniqueNames(tree *ScopeTree) {
	seen := make(map[string]bool)
	for _, edge := range tree.Edges {
		edge.Name = findFreeName(seen, edge.Name)
		seen[edge.Name] = true
	}

	for _, s := range tree.Children {
		EnsureUniqueNames(s)

		s.Scope.Name = findFreeName(seen, s.Scope.Name)
		seen[s.Scope.Name] = true
	}
}

func findFreeName(seen map[string]bool, name string) string {
	if !seen[name] {
		return name
	}
	for i := 1; ; i++ {
		next := fmt.Sprintf("%v'%v", name, i)
		if !seen[next] {
			return next
		}
	}
}
<|endoftext|>


<commit_before>package api

import (
	"net/http"
	"os"

	"github.com/go-chi/chi"
	"github.com/go-chi/chi/middleware"
	"github.com/goware/cors"
	"github.com/murdinc/terminal"
	"github.com/skratchdot/open-golang/open"
)

// StartAPI Starts the API listener on port 8081
func StartAPI(withDashboard bool) error {
	r := chi.NewRouter()

	cors := cors.New(cors.Options{
		AllowedOrigins:   []string{"*"},
		AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE"},
		AllowCredentials: true,
	})

	r.Use(cors.Handler)
	r.Use(middleware.Recoverer)
	r.Use(middleware.StripSlashes)
	//r.Use(middleware.Logger)

	r.Route("/api", func(r chi.Router) {
		r.Route("/dashboard", func(r chi.Router) {
			r.Route("/widgets", func(r chi.Router) {
				r.Get("/", getWidgets)
				r.Get("/events", getEvents)
				r.Get("/feed/:feedName", getFeed)
				r.Get("/options", getWidgetOptions)
				r.Get("/names", getWidgetNames)
				r.Get("/name/:widgetName", getWidgetByName)
				r.Put("/name/:widgetName", putWidget)
				r.Delete("/name/:widgetName", deleteWidget)
			})
		})
		r.Route("/assets", func(r chi.Router) {
			r.Route("/:assetType", func(r chi.Router) {
				r.Get("/", getAssets)
			})
		})
		r.Route("/classes", func(r chi.Router) {
			r.Get("/export", exportClasses)
			//r.Get("/import", importClasses) // TODO
			r.Route("/:classType", func(r chi.Router) {
				r.Get("/", getClasses)
				r.Get("/options", getClassOptions)
				r.Get("/names", getClassNames)
				r.Get("/name/:className", getClassByName)
				r.Put("/name/:className", putClass)
				r.Delete("/name/:className", deleteClass)
			})
		})
	})

	if withDashboard {
		src := "/usr/local/awsmDashboard"

		_, err := os.Stat(src)
		if os.IsNotExist(err) {
			// TODO install the dashboard automatically
			terminal.Notice("No awsm dashboard found!")
			terminal.Notice("Please install the dashboard by running the following command:")
			terminal.Notice("curl -s http://dl.sudoba.sh/get/awsmDashboard | sh")
			return err
		}

		r.Get("/*", func(w http.ResponseWriter, r *http.Request) {
			path := r.URL.Path

			if path == "/" {
				path = "/index.html"
			} else if _, err := os.Stat(src + path); os.IsNotExist(err) {
				path = "/index.html"
			}

			terminal.Information("Serving file: " + src + path)

			// return file
			http.ServeFile(w, r, src+path)

		})

		open.Start("http://localhost:8081")
	}

	return http.ListenAndServe(":8081", r)
}
<commit_msg>updates api route definitions to latest method<commit_after>package api

import (
	"net/http"
	"os"

	"github.com/go-chi/chi"
	"github.com/go-chi/chi/middleware"
	"github.com/goware/cors"
	"github.com/murdinc/terminal"
	"github.com/skratchdot/open-golang/open"
)

// StartAPI Starts the API listener on port 8081
func StartAPI(withDashboard bool) error {
	r := chi.NewRouter()

	cors := cors.New(cors.Options{
		AllowedOrigins:   []string{"*"},
		AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE"},
		AllowCredentials: true,
	})

	r.Use(cors.Handler)
	r.Use(middleware.Recoverer)
	r.Use(middleware.StripSlashes)
	//r.Use(middleware.Logger)

	r.Route("/api", func(r chi.Router) {
		r.Route("/dashboard", func(r chi.Router) {
			r.Route("/widgets", func(r chi.Router) {
				r.Get("/", getWidgets)
				r.Get("/events", getEvents)
				r.Get("/feed/{feedName}", getFeed)
				r.Get("/options", getWidgetOptions)
				r.Get("/names", getWidgetNames)
				r.Get("/name/{widgetName}", getWidgetByName)
				r.Put("/name/{widgetName}", putWidget)
				r.Delete("/name/{widgetName}", deleteWidget)
			})
		})
		r.Route("/assets", func(r chi.Router) {
			r.Route("/{assetType}", func(r chi.Router) {
				r.Get("/", getAssets)
			})
		})
		r.Route("/classes", func(r chi.Router) {
			r.Get("/export", exportClasses)
			//r.Get("/import", importClasses) // TODO
			r.Route("/{classType}", func(r chi.Router) {
				r.Get("/", getClasses)
				r.Get("/options", getClassOptions)
				r.Get("/names", getClassNames)
				r.Get("/name/{className}", getClassByName)
				r.Put("/name/{className}", putClass)
				r.Delete("/name/{className}", deleteClass)
			})
		})
	})

	if withDashboard {
		src := "/usr/local/awsmDashboard"

		_, err := os.Stat(src)
		if os.IsNotExist(err) {
			// TODO install the dashboard automatically
			terminal.Notice("No awsm dashboard found!")
			terminal.Notice("Please install the dashboard by running the following command:")
			terminal.Notice("curl -s http://dl.sudoba.sh/get/awsmDashboard | sh")
			return err
		}

		r.Get("/*", func(w http.ResponseWriter, r *http.Request) {
			path := r.URL.Path

			if path == "/" {
				path = "/index.html"
			} else if _, err := os.Stat(src + path); os.IsNotExist(err) {
				path = "/index.html"
			}

			terminal.Information("Serving file: " + src + path)

			// return file
			http.ServeFile(w, r, src+path)

		})

		open.Start("http://localhost:8081")
	}

	return http.ListenAndServe(":8081", r)
}
<|endoftext|>