<commit_before>package elastigo

import "encoding/json"

func Aggregate(name string) *AggregateDsl {
	return &AggregateDsl{Name: name}
}

type AggregateDsl struct {
	Name          string
	TypeName      string
	Type          interface{}
	Filters       *FilterWrap              `json:"filters,omitempty"`
	AggregatesVal map[string]*AggregateDsl `json:"aggregations,omitempty"`
}

type FieldAggregate struct {
	Field string `json:"field"`
}

/**
 * Aggregates accepts n "sub-aggregates" to be applied to this aggregate
 *
 * agg := Aggregate("user").Term("user_id")
 * agg.Aggregates(
 *   Aggregate("total_spent").Sum("price"),
 *   Aggregate("total_saved").Sum("discount"),
 * )
 */
func (d *AggregateDsl) Aggregates(aggs ...*AggregateDsl) *AggregateDsl {
	if len(aggs) < 1 {
		return d
	}
	if len(d.AggregatesVal) == 0 {
		d.AggregatesVal = make(map[string]*AggregateDsl)
	}

	for _, agg := range aggs {
		d.AggregatesVal[agg.Name] = agg
	}
	return d
}

func (d *AggregateDsl) Min(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "min"
	return d
}

func (d *AggregateDsl) Max(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "max"
	return d
}

func (d *AggregateDsl) Sum(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "sum"
	return d
}

func (d *AggregateDsl) Avg(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "avg"
	return d
}

func (d *AggregateDsl) Stats(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "stats"
	return d
}

func (d *AggregateDsl) ExtendedStats(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "extended_stats"
	return d
}

func (d *AggregateDsl) ValueCount(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "value_count"
	return d
}

func (d *AggregateDsl) Percentiles(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "percentiles"
	return d
}

type Cardinality struct {
	Field              string  `json:"field"`
	PrecisionThreshold float64 `json:"precision_threshold,omitempty"`
	Rehash             bool    `json:"rehash,omitempty"`
}

/**
 * Cardinality(
 *	 "field_name",
 *	 true,
 *   0,
 * )
 */
func (d *AggregateDsl) Cardinality(field string, rehash bool, threshold int) *AggregateDsl {
	c := Cardinality{Field: field}

	// Only set if it's false, since the default is true
	if !rehash {
		c.Rehash = false
	}

	if threshold > 0 {
		c.PrecisionThreshold = float64(threshold)
	}
	d.Type = c
	d.TypeName = "cardinality"
	return d
}

func (d *AggregateDsl) Global() *AggregateDsl {
	d.Type = struct{}{}
	d.TypeName = "global"
	return d
}

func (d *AggregateDsl) Filter(filters ...interface{}) *AggregateDsl {

	if len(filters) == 0 {
		return d
	}

	if d.Filters == nil {
		d.Filters = NewFilterWrap()
	}

	d.Filters.addFilters(filters)
	return d
}

func (d *AggregateDsl) Missing(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "missing"
	return d
}

func (d *AggregateDsl) Terms(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "terms"
	return d
}

func (d *AggregateDsl) SignificantTerms(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "significant_terms"
	return d
}

type Histogram struct {
	Field    string  `json:"field"`
	Interval float64 `json:"interval"`
}

func (d *AggregateDsl) Histogram(field string, interval int) *AggregateDsl {
	d.Type = Histogram{
		Field:    field,
		Interval: float64(interval),
	}
	d.TypeName = "histogram"
	return d
}

type DateHistogram struct {
	Field    string `json:"field"`
	Interval string `json:"interval"`
}

func (d *AggregateDsl) DateHistogram(field, interval string) *AggregateDsl {
	d.Type = DateHistogram{
		Field:    field,
		Interval: interval,
	}
	d.TypeName = "date_histogram"
	return d
}

func (d *AggregateDsl) MarshalJSON() ([]byte, error) {
	return json.Marshal(d.toMap())
}

func (d *AggregateDsl) toMap() map[string]interface{} {
	root := map[string]interface{}{}

	if d.Type != nil {
		root[d.TypeName] = d.Type
	}
	aggregates := d.aggregatesMap()

	if d.Filters != nil {
		root["filter"] = d.Filters
	}

	if len(aggregates) > 0 {
		root["aggregations"] = aggregates
	}
	return root

}
func (d *AggregateDsl) aggregatesMap() map[string]interface{} {
	root := map[string]interface{}{}

	if len(d.AggregatesVal) > 0 {
		for _, agg := range d.AggregatesVal {
			root[agg.Name] = agg.toMap()
		}
	}
	return root
}
<commit_msg>allow setting of terms aggregation size<commit_after>package elastigo

import "encoding/json"

func Aggregate(name string) *AggregateDsl {
	return &AggregateDsl{Name: name}
}

type AggregateDsl struct {
	Name          string
	TypeName      string
	Type          interface{}
	Filters       *FilterWrap              `json:"filters,omitempty"`
	AggregatesVal map[string]*AggregateDsl `json:"aggregations,omitempty"`
}

type FieldAggregate struct {
	Field string `json:"field"`
	Size  int    `json:"size,omitempty"`
}

/**
 * Aggregates accepts n "sub-aggregates" to be applied to this aggregate
 *
 * agg := Aggregate("user").Term("user_id")
 * agg.Aggregates(
 *   Aggregate("total_spent").Sum("price"),
 *   Aggregate("total_saved").Sum("discount"),
 * )
 */
func (d *AggregateDsl) Aggregates(aggs ...*AggregateDsl) *AggregateDsl {
	if len(aggs) < 1 {
		return d
	}
	if len(d.AggregatesVal) == 0 {
		d.AggregatesVal = make(map[string]*AggregateDsl)
	}

	for _, agg := range aggs {
		d.AggregatesVal[agg.Name] = agg
	}
	return d
}

func (d *AggregateDsl) Min(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "min"
	return d
}

func (d *AggregateDsl) Max(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "max"
	return d
}

func (d *AggregateDsl) Sum(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "sum"
	return d
}

func (d *AggregateDsl) Avg(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "avg"
	return d
}

func (d *AggregateDsl) Stats(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "stats"
	return d
}

func (d *AggregateDsl) ExtendedStats(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "extended_stats"
	return d
}

func (d *AggregateDsl) ValueCount(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "value_count"
	return d
}

func (d *AggregateDsl) Percentiles(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "percentiles"
	return d
}

type Cardinality struct {
	Field              string  `json:"field"`
	PrecisionThreshold float64 `json:"precision_threshold,omitempty"`
	Rehash             bool    `json:"rehash,omitempty"`
}

/**
 * Cardinality(
 *	 "field_name",
 *	 true,
 *   0,
 * )
 */
func (d *AggregateDsl) Cardinality(field string, rehash bool, threshold int) *AggregateDsl {
	c := Cardinality{Field: field}

	// Only set if it's false, since the default is true
	if !rehash {
		c.Rehash = false
	}

	if threshold > 0 {
		c.PrecisionThreshold = float64(threshold)
	}
	d.Type = c
	d.TypeName = "cardinality"
	return d
}

func (d *AggregateDsl) Global() *AggregateDsl {
	d.Type = struct{}{}
	d.TypeName = "global"
	return d
}

func (d *AggregateDsl) Filter(filters ...interface{}) *AggregateDsl {

	if len(filters) == 0 {
		return d
	}

	if d.Filters == nil {
		d.Filters = NewFilterWrap()
	}

	d.Filters.addFilters(filters)
	return d
}

func (d *AggregateDsl) Missing(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "missing"
	return d
}

func (d *AggregateDsl) Terms(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "terms"
	return d
}

func (d *AggregateDsl) TermsWithSize(field string, size int) *AggregateDsl {
	d.Type = FieldAggregate{Field: field, Size: size}
	d.TypeName = "terms"
	return d
}

func (d *AggregateDsl) SignificantTerms(field string) *AggregateDsl {
	d.Type = FieldAggregate{Field: field}
	d.TypeName = "significant_terms"
	return d
}

type Histogram struct {
	Field    string  `json:"field"`
	Interval float64 `json:"interval"`
}

func (d *AggregateDsl) Histogram(field string, interval int) *AggregateDsl {
	d.Type = Histogram{
		Field:    field,
		Interval: float64(interval),
	}
	d.TypeName = "histogram"
	return d
}

type DateHistogram struct {
	Field    string `json:"field"`
	Interval string `json:"interval"`
}

func (d *AggregateDsl) DateHistogram(field, interval string) *AggregateDsl {
	d.Type = DateHistogram{
		Field:    field,
		Interval: interval,
	}
	d.TypeName = "date_histogram"
	return d
}

func (d *AggregateDsl) MarshalJSON() ([]byte, error) {
	return json.Marshal(d.toMap())
}

func (d *AggregateDsl) toMap() map[string]interface{} {
	root := map[string]interface{}{}

	if d.Type != nil {
		root[d.TypeName] = d.Type
	}
	aggregates := d.aggregatesMap()

	if d.Filters != nil {
		root["filter"] = d.Filters
	}

	if len(aggregates) > 0 {
		root["aggregations"] = aggregates
	}
	return root

}

func (d *AggregateDsl) aggregatesMap() map[string]interface{} {
	root := map[string]interface{}{}

	if len(d.AggregatesVal) > 0 {
		for _, agg := range d.AggregatesVal {
			root[agg.Name] = agg.toMap()
		}
	}
	return root
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014, The Tor Project, Inc.
 * See LICENSE for licensing information
 */

package upnp

import (
	"bufio"
	"bytes"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"math"
	"net"
	"net/http"
	"strconv"
	"syscall"
)

const maxMappingDuration = 604800

// The people who made this abomination of a protocol used SOAP.  Presumably
// the "right" way to do this is to use an existing SOAP client, but Go does
// not have such a thing.

type soapEnvelope struct {
	XMLName       xml.Name `xml:"http://schemas.xmlsoap.org/soap/envelope/ Envelope"`
	EncodingStyle string   `xml:"http://schemas.xmlsoap.org/soap/envelope/ encodingStyle,attr"`
	Body          soapBody `xml:"http://schemas.xmlsoap.org/soap/envelope/ Body"`
}

type soapBody struct {
	Fault                              *soapFault             `xml:"Fault"`
	GetExternalIPAddressResponse       *getExtIPResponse      `xml:"GetExternalIPAddressResponse"`
	GetGenericPortMappingEntryResponse *getGenPMapEntResponse `xml:"GetGenericPortMappingEntryResponse"`
}

type soapFault struct {
	FaultCode   string      `xml:"faultcode"`
	FaultString string      `xml:"faultstring"`
	Detail      *soapDetail `xml:"detail"`
}

type soapDetail struct {
	UPnPError *upnpError `xml:"UPnPError"`
}

type upnpError struct {
	ErrorCode        int    `xml:"errorCode"`
	ErrorDescription string `xml:"errorDescription"`
}

type getExtIPResponse struct {
	IP string `xml:"NewExternalIPAddress"`
}

type getGenPMapEntResponse struct {
	RemoteHost             string `xml:"NewRemoteHost"`
	ExternalPort           int    `xml:"NewExternalPort"`
	Protocol               string `xml:"NewProtocol"`
	InternalPort           int    `xml:"NewInternalPort"`
	InternalClient         string `xml:"NewInternalClient"`
	Enabled                int    `xml:"NewEnabled"`
	PortMappingDescription string `xml:"NewPortMappingDescription"`
	LeaseDuration          int    `xml:"NewLeaseDuration"`
}

func (f *soapFault) String() string {
	if f.Detail.UPnPError != nil {
		return fmt.Sprintf("upnp error: %d - %s", f.Detail.UPnPError.ErrorCode, f.Detail.UPnPError.ErrorDescription)
	}
	return fmt.Sprintf("fault: %s - %s", f.FaultCode, f.FaultString)
}

func (c *Client) issueSoapRequest(actionName, argsXML string) (*soapBody, error) {
	// Apparently a lot of routers puke horribly on XML that's well-formed but
	// not exactly what they expect, so requests are crafted by hand.  At a
	// future time when more than 2 requests need to be supported, revisit.
	const header = xml.Header +
		"<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" " +
		"s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">" +
		"<s:Body>"
	const footer = "</s:Body></s:Envelope>"

	actionOpen := "<u:" + actionName + " xmlns:u=\"" + c.ctrl.urn.String() + "\">"
	actionClose := "</u:" + actionName + ">"
	body := []byte(header + actionOpen + argsXML + actionClose + footer)
	soapAction := "\"" + c.ctrl.urn.String() + "#" + actionName + "\""

	c.Vlogf("soap: issuing %s\n", actionName)

	// miniupnpd (used by a lot of routers) can't handle chunked transfer
	// encoding at all and just passes the raw body to it's XML parser.  This
	// is all sorts of garbage and violates RFC 2616.
	reqBuf := bytes.NewBuffer(body)
	req, err := http.NewRequest("POST", c.ctrl.url.String(), bufio.NewReader(reqBuf))
	if err != nil {
		return nil, err
	}
	req.ContentLength = int64(len(body))
	req.TransferEncoding = []string{"identity"}
	req.Header.Set("Content-Type", "text/xml; charset=\"utf-8\"")
	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("SOAPAction", soapAction)

	httpTransport := &http.Transport{DisableKeepAlives: true, DisableCompression: true}
	httpClient := &http.Client{Transport: httpTransport}

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	respEnvelope := &soapEnvelope{}
	if err = xml.Unmarshal(body, respEnvelope); err != nil {
		return nil, err
	}
	if respEnvelope.Body.Fault != nil {
		return nil, fmt.Errorf("soap: %s", respEnvelope.Body.Fault)
	}
	if resp.StatusCode != http.StatusOK {
		// Yes, this is at the end because the SOAP Fault gives more useful
		// diagnostics than "500 Internal Server Error".
		return nil, fmt.Errorf("soap: request failed with status: %s", resp.Status)
	}
	return &respEnvelope.Body, nil
}

// GetExternalIPAddress queries the router's external IP address.
func (c *Client) GetExternalIPAddress() (net.IP, error) {

	respBody, err := c.issueSoapRequest("GetExternalIPAddress", "")
	if err != nil {
		return nil, err
	}

	if respBody.GetExternalIPAddressResponse != nil {
		ip := net.ParseIP(respBody.GetExternalIPAddressResponse.IP)
		if ip != nil {
			return ip, nil
		}
	}
	return nil, fmt.Errorf("igd: GetExternalIPAddress() failed")
}

// GetListOfPortMappings queries the router for the list of port forwarding
// entries.
func (c *Client) GetListOfPortMappings() ([]string, error) {
	// Sad panda, GetListOfPortMappings requires IGD2 or later, so emulate it
	// with GetGenericPortMappingEntry.  Theoretically if the number of entries
	// changes during this process we would need to start over from the
	// begining, but we don't monitor events so we can't tell.

	resps := make([]string, 0)
	for idx := 0; idx < math.MaxUint16; idx++ {
		argsXML := "<NewPortMappingIndex>" + strconv.FormatUint(uint64(idx), 10) + "</NewPortMappingIndex>"
		respBody, err := c.issueSoapRequest("GetGenericPortMappingEntry", argsXML)
		if err != nil {
			// Probably SpecifiedArrayIndexInvalid. (XXX: Check?)
			c.Vlogf("igd: GetGenericPortMappingEntry returned: %s\n", err)
			break
		}
		if respBody.GetGenericPortMappingEntryResponse != nil {
			// Too long, much too long.
			r := respBody.GetGenericPortMappingEntryResponse
			remoteHost := r.RemoteHost
			if remoteHost == "" {
				remoteHost = "0.0.0.0"
			}
			s := fmt.Sprintf("'%s' %s:%d <-> %s:%d %s (%d sec)",
				r.PortMappingDescription,
				r.InternalClient,
				r.InternalPort,
				remoteHost,
				r.ExternalPort,
				r.Protocol,
				r.LeaseDuration)
			c.Vlogf("%s\n", s)
			resps = append(resps, s)
		}
	}
	return resps, nil
}

// AddPortMapping adds a new TCP/IP port mapping.  The internal IP address of
// the client is used as the destination.  Per the UPnP spec, duration can
// range from 0 to 604800, with the behavior on 0 changing depending on the
// version of the spec.
func (c *Client) AddPortMapping(descr string, internalPort, externalPort, duration int) error {
	if duration > maxMappingDuration {
		return syscall.ERANGE
	}

	c.Vlogf("AddPortMapping: '%s' %s:%d <-> 0.0.0.0:%d (%d sec)\n", descr, c.internalAddr, internalPort, externalPort, duration)

	argsXML := "<NewRemoteHost></NewRemoteHost>" +
		"<NewExternalPort>" + strconv.FormatUint(uint64(externalPort), 10) + "</NewExternalPort>" +
		"<NewProtocol>TCP</NewProtocol>" +
		"<NewInternalPort>" + strconv.FormatUint(uint64(internalPort), 10) + "</NewInternalPort>" +
		"<NewInternalClient>" + c.internalAddr.String() + "</NewInternalClient>" +
		"<NewEnabled>1</NewEnabled>" +
		"<NewPortMappingDescription>" + descr + "</NewPortMappingDescription>" +
		"<NewLeaseDuration>" + strconv.FormatUint(uint64(duration), 10) + "</NewLeaseDuration>"

	// HTTP 200 means that things worked.  The response isn't interesting
	// enough to warrant parsing.
	_, err := c.issueSoapRequest("AddPortMapping", argsXML)
	if err != nil {
		c.Vlogf("igd: AddPortMapping failed: %s\n", err)
		return err
	}
	return nil
}

// DeletePortMapping removes an existing TCP/IP port forwarding entry
// between clientIP:internalPort and 0.0.0.0:externalPort.
func (c *Client) DeletePortMapping(internalPort, externalPort int) error {
	c.Vlogf("DeletePortMapping: %s:%d <-> 0.0.0.0:%d\n", c.internalAddr, internalPort, externalPort)

	argsXML := "<NewRemoteHost></NewRemoteHost>" +
		"<NewExternalPort>" + strconv.FormatUint(uint64(externalPort), 10) + "</NewExternalPort>" +
		"<NewProtocol>TCP</NewProtocol>"

	// HTTP 200 means that things worked.  The response isn't interesting
	// enough to warrant parsing.
	_, err := c.issueSoapRequest("DeletePortMapping", argsXML)
	if err != nil {
		c.Vlogf("igd: DeletePortMapping failed: %s\n", err)
		return err
	}
	return nil
}
<commit_msg>Allocating space for resps is not always necessary<commit_after>/*
 * Copyright (c) 2014, The Tor Project, Inc.
 * See LICENSE for licensing information
 */

package upnp

import (
	"bufio"
	"bytes"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"math"
	"net"
	"net/http"
	"strconv"
	"syscall"
)

const maxMappingDuration = 604800

// The people who made this abomination of a protocol used SOAP.  Presumably
// the "right" way to do this is to use an existing SOAP client, but Go does
// not have such a thing.

type soapEnvelope struct {
	XMLName       xml.Name `xml:"http://schemas.xmlsoap.org/soap/envelope/ Envelope"`
	EncodingStyle string   `xml:"http://schemas.xmlsoap.org/soap/envelope/ encodingStyle,attr"`
	Body          soapBody `xml:"http://schemas.xmlsoap.org/soap/envelope/ Body"`
}

type soapBody struct {
	Fault                              *soapFault             `xml:"Fault"`
	GetExternalIPAddressResponse       *getExtIPResponse      `xml:"GetExternalIPAddressResponse"`
	GetGenericPortMappingEntryResponse *getGenPMapEntResponse `xml:"GetGenericPortMappingEntryResponse"`
}

type soapFault struct {
	FaultCode   string      `xml:"faultcode"`
	FaultString string      `xml:"faultstring"`
	Detail      *soapDetail `xml:"detail"`
}

type soapDetail struct {
	UPnPError *upnpError `xml:"UPnPError"`
}

type upnpError struct {
	ErrorCode        int    `xml:"errorCode"`
	ErrorDescription string `xml:"errorDescription"`
}

type getExtIPResponse struct {
	IP string `xml:"NewExternalIPAddress"`
}

type getGenPMapEntResponse struct {
	RemoteHost             string `xml:"NewRemoteHost"`
	ExternalPort           int    `xml:"NewExternalPort"`
	Protocol               string `xml:"NewProtocol"`
	InternalPort           int    `xml:"NewInternalPort"`
	InternalClient         string `xml:"NewInternalClient"`
	Enabled                int    `xml:"NewEnabled"`
	PortMappingDescription string `xml:"NewPortMappingDescription"`
	LeaseDuration          int    `xml:"NewLeaseDuration"`
}

func (f *soapFault) String() string {
	if f.Detail.UPnPError != nil {
		return fmt.Sprintf("upnp error: %d - %s", f.Detail.UPnPError.ErrorCode, f.Detail.UPnPError.ErrorDescription)
	}
	return fmt.Sprintf("fault: %s - %s", f.FaultCode, f.FaultString)
}

func (c *Client) issueSoapRequest(actionName, argsXML string) (*soapBody, error) {
	// Apparently a lot of routers puke horribly on XML that's well-formed but
	// not exactly what they expect, so requests are crafted by hand.  At a
	// future time when more than 2 requests need to be supported, revisit.
	const header = xml.Header +
		"<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" " +
		"s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">" +
		"<s:Body>"
	const footer = "</s:Body></s:Envelope>"

	actionOpen := "<u:" + actionName + " xmlns:u=\"" + c.ctrl.urn.String() + "\">"
	actionClose := "</u:" + actionName + ">"
	body := []byte(header + actionOpen + argsXML + actionClose + footer)
	soapAction := "\"" + c.ctrl.urn.String() + "#" + actionName + "\""

	c.Vlogf("soap: issuing %s\n", actionName)

	// miniupnpd (used by a lot of routers) can't handle chunked transfer
	// encoding at all and just passes the raw body to it's XML parser.  This
	// is all sorts of garbage and violates RFC 2616.
	reqBuf := bytes.NewBuffer(body)
	req, err := http.NewRequest("POST", c.ctrl.url.String(), bufio.NewReader(reqBuf))
	if err != nil {
		return nil, err
	}
	req.ContentLength = int64(len(body))
	req.TransferEncoding = []string{"identity"}
	req.Header.Set("Content-Type", "text/xml; charset=\"utf-8\"")
	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("SOAPAction", soapAction)

	httpTransport := &http.Transport{DisableKeepAlives: true, DisableCompression: true}
	httpClient := &http.Client{Transport: httpTransport}

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	respEnvelope := &soapEnvelope{}
	if err = xml.Unmarshal(body, respEnvelope); err != nil {
		return nil, err
	}
	if respEnvelope.Body.Fault != nil {
		return nil, fmt.Errorf("soap: %s", respEnvelope.Body.Fault)
	}
	if resp.StatusCode != http.StatusOK {
		// Yes, this is at the end because the SOAP Fault gives more useful
		// diagnostics than "500 Internal Server Error".
		return nil, fmt.Errorf("soap: request failed with status: %s", resp.Status)
	}
	return &respEnvelope.Body, nil
}

// GetExternalIPAddress queries the router's external IP address.
func (c *Client) GetExternalIPAddress() (net.IP, error) {

	respBody, err := c.issueSoapRequest("GetExternalIPAddress", "")
	if err != nil {
		return nil, err
	}

	if respBody.GetExternalIPAddressResponse != nil {
		ip := net.ParseIP(respBody.GetExternalIPAddressResponse.IP)
		if ip != nil {
			return ip, nil
		}
	}
	return nil, fmt.Errorf("igd: GetExternalIPAddress() failed")
}

// GetListOfPortMappings queries the router for the list of port forwarding
// entries.
func (c *Client) GetListOfPortMappings() ([]string, error) {
	// Sad panda, GetListOfPortMappings requires IGD2 or later, so emulate it
	// with GetGenericPortMappingEntry.  Theoretically if the number of entries
	// changes during this process we would need to start over from the
	// begining, but we don't monitor events so we can't tell.

	var resps []string
	for idx := 0; idx < math.MaxUint16; idx++ {
		argsXML := "<NewPortMappingIndex>" + strconv.FormatUint(uint64(idx), 10) + "</NewPortMappingIndex>"
		respBody, err := c.issueSoapRequest("GetGenericPortMappingEntry", argsXML)
		if err != nil {
			// Probably SpecifiedArrayIndexInvalid. (XXX: Check?)
			c.Vlogf("igd: GetGenericPortMappingEntry returned: %s\n", err)
			break
		}
		if respBody.GetGenericPortMappingEntryResponse != nil {
			// Too long, much too long.
			r := respBody.GetGenericPortMappingEntryResponse
			remoteHost := r.RemoteHost
			if remoteHost == "" {
				remoteHost = "0.0.0.0"
			}
			s := fmt.Sprintf("'%s' %s:%d <-> %s:%d %s (%d sec)",
				r.PortMappingDescription,
				r.InternalClient,
				r.InternalPort,
				remoteHost,
				r.ExternalPort,
				r.Protocol,
				r.LeaseDuration)
			c.Vlogf("%s\n", s)
			resps = append(resps, s)
		}
	}
	return resps, nil
}

// AddPortMapping adds a new TCP/IP port mapping.  The internal IP address of
// the client is used as the destination.  Per the UPnP spec, duration can
// range from 0 to 604800, with the behavior on 0 changing depending on the
// version of the spec.
func (c *Client) AddPortMapping(descr string, internalPort, externalPort, duration int) error {
	if duration > maxMappingDuration {
		return syscall.ERANGE
	}

	c.Vlogf("AddPortMapping: '%s' %s:%d <-> 0.0.0.0:%d (%d sec)\n", descr, c.internalAddr, internalPort, externalPort, duration)

	argsXML := "<NewRemoteHost></NewRemoteHost>" +
		"<NewExternalPort>" + strconv.FormatUint(uint64(externalPort), 10) + "</NewExternalPort>" +
		"<NewProtocol>TCP</NewProtocol>" +
		"<NewInternalPort>" + strconv.FormatUint(uint64(internalPort), 10) + "</NewInternalPort>" +
		"<NewInternalClient>" + c.internalAddr.String() + "</NewInternalClient>" +
		"<NewEnabled>1</NewEnabled>" +
		"<NewPortMappingDescription>" + descr + "</NewPortMappingDescription>" +
		"<NewLeaseDuration>" + strconv.FormatUint(uint64(duration), 10) + "</NewLeaseDuration>"

	// HTTP 200 means that things worked.  The response isn't interesting
	// enough to warrant parsing.
	_, err := c.issueSoapRequest("AddPortMapping", argsXML)
	if err != nil {
		c.Vlogf("igd: AddPortMapping failed: %s\n", err)
		return err
	}
	return nil
}

// DeletePortMapping removes an existing TCP/IP port forwarding entry
// between clientIP:internalPort and 0.0.0.0:externalPort.
func (c *Client) DeletePortMapping(internalPort, externalPort int) error {
	c.Vlogf("DeletePortMapping: %s:%d <-> 0.0.0.0:%d\n", c.internalAddr, internalPort, externalPort)

	argsXML := "<NewRemoteHost></NewRemoteHost>" +
		"<NewExternalPort>" + strconv.FormatUint(uint64(externalPort), 10) + "</NewExternalPort>" +
		"<NewProtocol>TCP</NewProtocol>"

	// HTTP 200 means that things worked.  The response isn't interesting
	// enough to warrant parsing.
	_, err := c.issueSoapRequest("DeletePortMapping", argsXML)
	if err != nil {
		c.Vlogf("igd: DeletePortMapping failed: %s\n", err)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// The amazonebs package contains a packer.Builder implementation that
// builds AMIs for Amazon EC2.
//
// In general, there are two types of AMIs that can be created: ebs-backed or
// instance-store. This builder _only_ builds ebs-backed images.
package amazonebs

import (
	"encoding/json"
	"github.com/mitchellh/packer/packer"
	"log"
)

type config struct {
	AccessKey string `json:"access_key"`
	Region    string
	SecretKey string `json:"secret_key"`
	SourceAmi string `json:"source_ami"`
}

type Builder struct {
	config config
}

func (b *Builder) Prepare(raw interface{}) (err error) {
	// Marshal and unmarshal the raw configuration as a way to get it
	// into our "config" struct.
	// TODO: Use the reflection package and provide this as an API for
	// better error messages
	jsonBytes, err := json.Marshal(raw)
	if err != nil {
		return
	}

	err = json.Unmarshal(jsonBytes, &b.config)
	if err != nil {
		return
	}

	log.Printf("Config: %+v\n", b.config)

	// TODO: Validate the configuration
	return
}

func (*Builder) Run(b packer.Build, ui packer.Ui) {
	ui.Say("Building an AWS image...\n")
}
<commit_msg>builder/amazonebs: Can make an AMI. WHAT!<commit_after>// The amazonebs package contains a packer.Builder implementation that
// builds AMIs for Amazon EC2.
//
// In general, there are two types of AMIs that can be created: ebs-backed or
// instance-store. This builder _only_ builds ebs-backed images.
package amazonebs

import (
	"encoding/json"
	"fmt"
	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/packer/packer"
	"log"
	"time"
)

type config struct {
	AccessKey string `json:"access_key"`
	Region    string
	SecretKey string `json:"secret_key"`
	SourceAmi string `json:"source_ami"`
}

type Builder struct {
	config config
}

func (b *Builder) Prepare(raw interface{}) (err error) {
	// Marshal and unmarshal the raw configuration as a way to get it
	// into our "config" struct.
	// TODO: Use the reflection package and provide this as an API for
	// better error messages
	jsonBytes, err := json.Marshal(raw)
	if err != nil {
		return
	}

	err = json.Unmarshal(jsonBytes, &b.config)
	if err != nil {
		return
	}

	log.Printf("Config: %+v\n", b.config)

	// TODO: Validate the configuration
	return
}

func (b *Builder) Run(build packer.Build, ui packer.Ui) {
	auth := aws.Auth{b.config.AccessKey, b.config.SecretKey}
	region := aws.Regions[b.config.Region]
	ec2conn := ec2.New(auth, region)

	runOpts := &ec2.RunInstances{
		ImageId: b.config.SourceAmi,
		InstanceType: "m1.small",
		MinCount: 0,
		MaxCount: 0,
	}

	ui.Say("Launching a source AWS instance...\n")
	runResp, err := ec2conn.RunInstances(runOpts)
	if err != nil {
		ui.Error("%s\n", err.Error())
		return
	}

	instance := &runResp.Instances[0]
	log.Printf("instance id: %s\n", instance.InstanceId)
	ui.Say("Waiting for instance to become ready...\n")
	err = waitForState(ec2conn, instance, "running")
	if err != nil {
		ui.Error("%s\n", err.Error())
		return
	}

	// Stop the instance so we can create an AMI from it
	_, err = ec2conn.StopInstances(instance.InstanceId)
	if err != nil {
		ui.Error("%s\n", err.Error())
		return
	}

	// Wait for the instance to actual stop
	// TODO: Handle diff source states, i.e. this force state sucks
	ui.Say("Waiting for the instance to stop...\n")
	instance.State.Name = "stopping"
	err = waitForState(ec2conn, instance, "stopped")
	if err != nil {
		ui.Error("%s\n", err.Error())
		return
	}

	// Create the image
	ui.Say("Creating the AMI...\n")
	createOpts := &ec2.CreateImage{
		InstanceId: instance.InstanceId,
		Name: "ZIMAGINE",
	}

	createResp, err := ec2conn.CreateImage(createOpts)
	if err != nil {
		ui.Error("%s\n", err.Error())
		return
	}

	ui.Say("AMI: %s\n", createResp.ImageId)

	// Make sure we clean up the instance by terminating it, no matter what
	defer func() {
		// TODO: error handling
		ui.Say("Terminating the source AWS instance...\n")
		ec2conn.TerminateInstances([]string{instance.InstanceId})
	}()
}

func waitForState(ec2conn *ec2.EC2, i *ec2.Instance, target string) (err error) {
	log.Printf("Waiting for instance state to become: %s\n", target)

	original := i.State.Name
	for i.State.Name == original {
		var resp *ec2.InstancesResp
		resp, err = ec2conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
		if err != nil {
			return
		}

		i = &resp.Reservations[0].Instances[0]

		time.Sleep(2 * time.Second)
	}

	if i.State.Name != target {
		return fmt.Errorf("unexpected target state '%s', wanted '%s'", i.State.Name, target)
	}

	return
}
<|endoftext|>


<commit_before>package githop

import (
	"sort"
	"time"

	"github.com/google/go-github/github"
)

type RepoDigest struct {
	Repo    *github.Repository
	Commits []github.RepositoryCommit
}

// sort.Interface implementation for sorting RepoDigests.
type ByRepoFullName []*RepoDigest

func (a ByRepoFullName) Len() int           { return len(a) }
func (a ByRepoFullName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByRepoFullName) Less(i, j int) bool { return *a[i].Repo.FullName < *a[j].Repo.FullName }

type Digest struct {
	User        *github.User
	StartTime   time.Time
	EndTime     time.Time
	RepoDigests []*RepoDigest
}

func newDigest(githubClient *github.Client) (*Digest, error) {
	user, _, err := githubClient.Users.Get("")
	if err != nil {
		return nil, err
	}

	// The username parameter must be left blank so that we can get all of the
	// repositories the user has access to, not just ones that they own.
	repos, _, err := githubClient.Repositories.List("", nil)
	if err != nil {
		return nil, err
	}

	orgs, _, err := githubClient.Organizations.List("", nil)
	if err != nil {
		return nil, err
	}
	for _, org := range orgs {
		orgRepos, _, err := githubClient.Repositories.ListByOrg(*org.Login, nil)
		if err != nil {
			return nil, err
		}
		newRepos := make([]github.Repository, len(repos)+len(orgRepos))
		copy(newRepos, repos)
		copy(newRepos[len(repos):], orgRepos)
		repos = newRepos
	}

	now := time.Now()
	digestStartTime := time.Date(now.Year()-1, now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
	digestEndTime := digestStartTime.AddDate(0, 0, 1)

	// Only look at repos that may have activity in the digest interval.
	var digestRepos []github.Repository
	for _, repo := range repos {
		if repo.CreatedAt.Before(digestEndTime) && repo.PushedAt.After(digestStartTime) {
			digestRepos = append(digestRepos, repo)
		}
	}
	repos = digestRepos
	digest := &Digest{
		User:        user,
		RepoDigests: make([]*RepoDigest, 0, len(repos)),
		StartTime:   digestStartTime,
		EndTime:     digestEndTime,
	}
	err = digest.fetch(repos, githubClient)
	return digest, err
}

func (digest *Digest) fetch(repos []github.Repository, githubClient *github.Client) error {
	type RepoDigestResponse struct {
		repoDigest *RepoDigest
		err        error
	}
	ch := make(chan *RepoDigestResponse)
	for _, repo := range repos {
		go func(repo github.Repository) {
			commits, _, err := githubClient.Repositories.ListCommits(
				*repo.Owner.Login,
				*repo.Name,
				&github.CommitsListOptions{
					Author: *digest.User.Login,
					Since:  digest.StartTime,
					Until:  digest.EndTime,
				})
			if err != nil {
				ch <- &RepoDigestResponse{nil, err}
			} else {
				ch <- &RepoDigestResponse{&RepoDigest{&repo, commits}, nil}
			}
		}(repo)
	}
	for i := 0; i < len(repos); i++ {
		select {
		case r := <-ch:
			if r.err != nil {
				return r.err
			}
			digest.RepoDigests = append(digest.RepoDigests, r.repoDigest)
		}
	}
	sort.Sort(ByRepoFullName(digest.RepoDigests))
	return nil
}
<commit_msg>Omit repo digests that don't have any commits in that interval.<commit_after>package githop

import (
	"sort"
	"time"

	"github.com/google/go-github/github"
)

type RepoDigest struct {
	Repo    *github.Repository
	Commits []github.RepositoryCommit
}

// sort.Interface implementation for sorting RepoDigests.
type ByRepoFullName []*RepoDigest

func (a ByRepoFullName) Len() int           { return len(a) }
func (a ByRepoFullName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByRepoFullName) Less(i, j int) bool { return *a[i].Repo.FullName < *a[j].Repo.FullName }

type Digest struct {
	User        *github.User
	StartTime   time.Time
	EndTime     time.Time
	RepoDigests []*RepoDigest
}

func newDigest(githubClient *github.Client) (*Digest, error) {
	user, _, err := githubClient.Users.Get("")
	if err != nil {
		return nil, err
	}

	// The username parameter must be left blank so that we can get all of the
	// repositories the user has access to, not just ones that they own.
	repos, _, err := githubClient.Repositories.List("", nil)
	if err != nil {
		return nil, err
	}

	orgs, _, err := githubClient.Organizations.List("", nil)
	if err != nil {
		return nil, err
	}
	for _, org := range orgs {
		orgRepos, _, err := githubClient.Repositories.ListByOrg(*org.Login, nil)
		if err != nil {
			return nil, err
		}
		newRepos := make([]github.Repository, len(repos)+len(orgRepos))
		copy(newRepos, repos)
		copy(newRepos[len(repos):], orgRepos)
		repos = newRepos
	}

	now := time.Now()
	digestStartTime := time.Date(now.Year()-1, now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
	digestEndTime := digestStartTime.AddDate(0, 0, 1)

	// Only look at repos that may have activity in the digest interval.
	var digestRepos []github.Repository
	for _, repo := range repos {
		if repo.CreatedAt.Before(digestEndTime) && repo.PushedAt.After(digestStartTime) {
			digestRepos = append(digestRepos, repo)
		}
	}
	repos = digestRepos
	digest := &Digest{
		User:        user,
		RepoDigests: make([]*RepoDigest, 0, len(repos)),
		StartTime:   digestStartTime,
		EndTime:     digestEndTime,
	}
	err = digest.fetch(repos, githubClient)
	return digest, err
}

func (digest *Digest) fetch(repos []github.Repository, githubClient *github.Client) error {
	type RepoDigestResponse struct {
		repoDigest *RepoDigest
		err        error
	}
	ch := make(chan *RepoDigestResponse)
	for _, repo := range repos {
		go func(repo github.Repository) {
			commits, _, err := githubClient.Repositories.ListCommits(
				*repo.Owner.Login,
				*repo.Name,
				&github.CommitsListOptions{
					Author: *digest.User.Login,
					Since:  digest.StartTime,
					Until:  digest.EndTime,
				})
			if err != nil {
				ch <- &RepoDigestResponse{nil, err}
			} else {
				ch <- &RepoDigestResponse{&RepoDigest{&repo, commits}, nil}
			}
		}(repo)
	}
	for i := 0; i < len(repos); i++ {
		select {
		case r := <-ch:
			if r.err != nil {
				return r.err
			}
			if len(r.repoDigest.Commits) > 0 {
				digest.RepoDigests = append(digest.RepoDigests, r.repoDigest)
			}
		}
	}
	sort.Sort(ByRepoFullName(digest.RepoDigests))
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 - Olivier Wulveryck

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

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

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

import (
	"fmt"
)

// PropertyDefinition as described in Appendix 5.7:
// A property definition defines a named, typed value and related data
// that can be associated with an entity defined in this specification
// (e.g., Node Types, Relation ship Types, Capability Types, etc.).
// Properties are used by template authors to provide input values to
// TOSCA entities which indicate their “desired state” when they are instantiated.
// The value of a property can be retrieved using the
// get_property function within TOSCA Service Templates
type PropertyDefinition struct {
	Type        string      `yaml:"type" json:"type"`                                   // The required data type for the property
	Description string      `yaml:"description,omitempty" json:"description,omitempty"` // The optional description for the property.
	Required    bool        `yaml:"required,omitempty" json:"required,omitempty"`       // An optional key that declares a property as required ( true) or not ( false) Default: true
	Default     string      `yaml:"default,omitempty" json:"default,omitempty"`
	Status      Status      `yaml:"status,omitempty" json:"status,omitempty"`
	Constraints Constraints `yaml:"constraints,omitempty,flow" json:"constraints,omitempty"`
	EntrySchema string      `yaml:"entry_schema,omitempty" json:"entry_schema,omitempty"`
}

// A Property assignment is always a map, but the key may be value
type PropertyAssignment map[string][]string

func (p *PropertyAssignment) MarshalYAML() (interface{}, error) {
	for k, v := range *p {
		if k == "value" {
			if len(v) != 1 {
				return nil, fmt.Errorf("too many values")
			}
			return v[0], nil
		}
	}
	return p, nil
}

func (p *PropertyAssignment) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	*p = make(map[string][]string, 1)
	if err := unmarshal(&s); err == nil {
		(*p)["value"] = append([]string{}, s)
		return nil
	}
	var m map[string]string
	if err := unmarshal(&m); err != nil {
		for k, v := range m {
			(*p)[k] = append([]string{}, v)
		}
		return nil
	}
	var mm map[string][]string
	if err := unmarshal(&mm); err != nil {
		for k, v := range mm {
			(*p)[k] = v
		}
		return nil
	}
	var res interface{}
	unmarshal(&res)
	return fmt.Errorf("Cannot parse Property %v", res)
}
<commit_msg>[Fix] Bug in the PropertyAssignment's Unmarshaling The error check was incorrect<commit_after>/*
Copyright 2015 - Olivier Wulveryck

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

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

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

import (
	"fmt"
)

// PropertyDefinition as described in Appendix 5.7:
// A property definition defines a named, typed value and related data
// that can be associated with an entity defined in this specification
// (e.g., Node Types, Relation ship Types, Capability Types, etc.).
// Properties are used by template authors to provide input values to
// TOSCA entities which indicate their “desired state” when they are instantiated.
// The value of a property can be retrieved using the
// get_property function within TOSCA Service Templates
type PropertyDefinition struct {
	Type        string      `yaml:"type" json:"type"`                                   // The required data type for the property
	Description string      `yaml:"description,omitempty" json:"description,omitempty"` // The optional description for the property.
	Required    bool        `yaml:"required,omitempty" json:"required,omitempty"`       // An optional key that declares a property as required ( true) or not ( false) Default: true
	Default     string      `yaml:"default,omitempty" json:"default,omitempty"`
	Status      Status      `yaml:"status,omitempty" json:"status,omitempty"`
	Constraints Constraints `yaml:"constraints,omitempty,flow" json:"constraints,omitempty"`
	EntrySchema string      `yaml:"entry_schema,omitempty" json:"entry_schema,omitempty"`
}

// A Property assignment is always a map, but the key may be value
type PropertyAssignment map[string][]string

func (p *PropertyAssignment) MarshalYAML() (interface{}, error) {
	for k, v := range *p {
		if k == "value" {
			if len(v) != 1 {
				return nil, fmt.Errorf("too many values")
			}
			return v[0], nil
		}
	}
	return p, nil
}

func (p *PropertyAssignment) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	*p = make(map[string][]string, 1)
	if err := unmarshal(&s); err == nil {
		(*p)["value"] = append([]string{}, s)
		return nil
	}
	var m map[string]string
	if err := unmarshal(&m); err == nil {
		for k, v := range m {
			(*p)[k] = append([]string{}, v)
		}
		return nil
	}
	var mm map[string][]string
	if err := unmarshal(&mm); err == nil {
		for k, v := range mm {
			(*p)[k] = v
		}
		return nil
	}
	var res interface{}
	unmarshal(&res)
	return fmt.Errorf("Cannot parse Property %v", res)
}
<|endoftext|>


<commit_before>package libgobuster

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

	multierror "github.com/hashicorp/go-multierror"
)

const (
	// ModeDir represents -m dir
	ModeDir = "dir"
	// ModeDNS represents -m dns
	ModeDNS = "dns"
)

// Options helds all options that can be passed to libgobuster
type Options struct {
	Extensions        string
	ExtensionsParsed  []string
	Mode              string
	Password          string
	StatusCodes       string
	StatusCodesParsed intSet
	Threads           int
	URL               string
	UserAgent         string
	Username          string
	Wordlist          string
	Proxy             string
	Cookies           string
	Timeout           time.Duration
	FollowRedirect    bool
	IncludeLength     bool
	NoStatus          bool
	Expanded          bool
	Quiet             bool
	ShowIPs           bool
	ShowCNAME         bool
	InsecureSSL       bool
	WildcardForced    bool
	Verbose           bool
	UseSlash          bool
}

// NewOptions returns a new initialized Options object
func NewOptions() *Options {
	return &Options{
		StatusCodesParsed: intSet{Set: map[int]bool{}},
	}
}

// Validate validates the given options
func (opt *Options) validate() *multierror.Error {
	var errorList *multierror.Error

	if strings.ToLower(opt.Mode) != ModeDir && strings.ToLower(opt.Mode) != ModeDNS {
		errorList = multierror.Append(errorList, fmt.Errorf("Mode (-m): Invalid value: %s", opt.Mode))
	}

	if opt.Threads < 0 {
		errorList = multierror.Append(errorList, fmt.Errorf("Threads (-t): Invalid value: %d", opt.Threads))
	}

	if opt.Wordlist == "" {
		errorList = multierror.Append(errorList, fmt.Errorf("WordList (-w): Must be specified"))
	} else if opt.Wordlist == "-" {
		// STDIN
	} else if _, err := os.Stat(opt.Wordlist); os.IsNotExist(err) {
		errorList = multierror.Append(errorList, fmt.Errorf("Wordlist (-w): File does not exist: %s", opt.Wordlist))
	}

	if opt.URL == "" {
		errorList = multierror.Append(errorList, fmt.Errorf("Url/Domain (-u): Must be specified"))
	}

	if opt.StatusCodes != "" {
		if err := opt.parseStatusCodes(); err != nil {
			errorList = multierror.Append(errorList, err)
		}
	}

	if opt.Extensions != "" {
		if err := opt.parseExtensions(); err != nil {
			errorList = multierror.Append(errorList, err)
		}
	}

	if !strings.HasSuffix(opt.URL, "/") {
		opt.URL = fmt.Sprintf("%s/", opt.URL)
	}

	if opt.Mode == ModeDir {
		if err := opt.validateDirMode(); err != nil {
			errorList = multierror.Append(errorList, err)
		}
	}

	return errorList
}

// ParseExtensions parses the extensions provided as a comma seperated list
func (opt *Options) parseExtensions() error {
	if opt.Extensions == "" {
		return fmt.Errorf("invalid extension string provided")
	}

	opt.ExtensionsParsed = append(opt.ExtensionsParsed, strings.Split(opt.Extensions, ",")...)
	return nil
}

// ParseStatusCodes parses the status codes provided as a comma seperated list
func (opt *Options) parseStatusCodes() error {
	if opt.StatusCodes == "" {
		return fmt.Errorf("invalid status code string provided")
	}

	for _, c := range strings.Split(opt.StatusCodes, ",") {
		i, err := strconv.Atoi(c)
		if err != nil {
			return fmt.Errorf("invalid status code given: %s", c)
		}
		opt.StatusCodesParsed.Add(i)
	}
	return nil
}

func (opt *Options) validateDirMode() error {
	// bail out if we are not in dir mode
	if opt.Mode != ModeDir {
		return nil
	}
	if !strings.HasPrefix(opt.URL, "http") {
		// check to see if a port was specified
		re := regexp.MustCompile(`^[^/]+:(\d+)`)
		match := re.FindStringSubmatch(opt.URL)

		if len(match) < 2 {
			// no port, default to http on 80
			opt.URL = fmt.Sprintf("http://%s", opt.URL)
		} else {
			port, err := strconv.Atoi(match[1])
			if err != nil || (port != 80 && port != 443) {
				return fmt.Errorf("url scheme not specified")
			} else if port == 80 {
				opt.URL = fmt.Sprintf("http://%s", opt.URL)
			} else {
				opt.URL = fmt.Sprintf("https://%s", opt.URL)
			}
		}
	}

	if opt.Username != "" && opt.Password == "" {
		return fmt.Errorf("username was provided but password is missing")
	}

	return nil
}
<commit_msg>remove leading dot from extensions<commit_after>package libgobuster

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

	multierror "github.com/hashicorp/go-multierror"
)

const (
	// ModeDir represents -m dir
	ModeDir = "dir"
	// ModeDNS represents -m dns
	ModeDNS = "dns"
)

// Options helds all options that can be passed to libgobuster
type Options struct {
	Extensions        string
	ExtensionsParsed  []string
	Mode              string
	Password          string
	StatusCodes       string
	StatusCodesParsed intSet
	Threads           int
	URL               string
	UserAgent         string
	Username          string
	Wordlist          string
	Proxy             string
	Cookies           string
	Timeout           time.Duration
	FollowRedirect    bool
	IncludeLength     bool
	NoStatus          bool
	Expanded          bool
	Quiet             bool
	ShowIPs           bool
	ShowCNAME         bool
	InsecureSSL       bool
	WildcardForced    bool
	Verbose           bool
	UseSlash          bool
}

// NewOptions returns a new initialized Options object
func NewOptions() *Options {
	return &Options{
		StatusCodesParsed: intSet{Set: map[int]bool{}},
	}
}

// Validate validates the given options
func (opt *Options) validate() *multierror.Error {
	var errorList *multierror.Error

	if strings.ToLower(opt.Mode) != ModeDir && strings.ToLower(opt.Mode) != ModeDNS {
		errorList = multierror.Append(errorList, fmt.Errorf("Mode (-m): Invalid value: %s", opt.Mode))
	}

	if opt.Threads < 0 {
		errorList = multierror.Append(errorList, fmt.Errorf("Threads (-t): Invalid value: %d", opt.Threads))
	}

	if opt.Wordlist == "" {
		errorList = multierror.Append(errorList, fmt.Errorf("WordList (-w): Must be specified"))
	} else if opt.Wordlist == "-" {
		// STDIN
	} else if _, err := os.Stat(opt.Wordlist); os.IsNotExist(err) {
		errorList = multierror.Append(errorList, fmt.Errorf("Wordlist (-w): File does not exist: %s", opt.Wordlist))
	}

	if opt.URL == "" {
		errorList = multierror.Append(errorList, fmt.Errorf("Url/Domain (-u): Must be specified"))
	}

	if opt.StatusCodes != "" {
		if err := opt.parseStatusCodes(); err != nil {
			errorList = multierror.Append(errorList, err)
		}
	}

	if opt.Extensions != "" {
		if err := opt.parseExtensions(); err != nil {
			errorList = multierror.Append(errorList, err)
		}
	}

	if !strings.HasSuffix(opt.URL, "/") {
		opt.URL = fmt.Sprintf("%s/", opt.URL)
	}

	if opt.Mode == ModeDir {
		if err := opt.validateDirMode(); err != nil {
			errorList = multierror.Append(errorList, err)
		}
	}

	return errorList
}

// ParseExtensions parses the extensions provided as a comma seperated list
func (opt *Options) parseExtensions() error {
	if opt.Extensions == "" {
		return fmt.Errorf("invalid extension string provided")
	}

	exts := strings.Split(opt.Extensions, ",")
	for _, e := range exts {
		// remove leading . from extensions
		opt.ExtensionsParsed = append(opt.ExtensionsParsed, strings.TrimPrefix(e, "."))
	}
	return nil
}

// ParseStatusCodes parses the status codes provided as a comma seperated list
func (opt *Options) parseStatusCodes() error {
	if opt.StatusCodes == "" {
		return fmt.Errorf("invalid status code string provided")
	}

	for _, c := range strings.Split(opt.StatusCodes, ",") {
		i, err := strconv.Atoi(c)
		if err != nil {
			return fmt.Errorf("invalid status code given: %s", c)
		}
		opt.StatusCodesParsed.Add(i)
	}
	return nil
}

func (opt *Options) validateDirMode() error {
	// bail out if we are not in dir mode
	if opt.Mode != ModeDir {
		return nil
	}
	if !strings.HasPrefix(opt.URL, "http") {
		// check to see if a port was specified
		re := regexp.MustCompile(`^[^/]+:(\d+)`)
		match := re.FindStringSubmatch(opt.URL)

		if len(match) < 2 {
			// no port, default to http on 80
			opt.URL = fmt.Sprintf("http://%s", opt.URL)
		} else {
			port, err := strconv.Atoi(match[1])
			if err != nil || (port != 80 && port != 443) {
				return fmt.Errorf("url scheme not specified")
			} else if port == 80 {
				opt.URL = fmt.Sprintf("http://%s", opt.URL)
			} else {
				opt.URL = fmt.Sprintf("https://%s", opt.URL)
			}
		}
	}

	if opt.Username != "" && opt.Password == "" {
		return fmt.Errorf("username was provided but password is missing")
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import "testing"

var inodePathTests = []struct {
	path  string
	inode uint64
}{
	{"/", 1251674434},
	{"/foo/bar", 902704296},
}

func TestInodePath(t *testing.T) {
	for i, test := range inodePathTests {
		inode := inodePath(test.path)

		if inode != test.inode {
			t.Errorf("test %d: wrong inode returned, want %d, got %d", i, test.inode, inode)
		}
	}
}
<commit_msg>Fix inode tests<commit_after>package main

import (
	"testing"

	"github.com/jacobsa/fuse/fuseops"
)

var inodePathTests = []struct {
	path  string
	inode fuseops.InodeID
}{
	{"/", 1251674434},
	{"/foo/bar", 902704296},
}

func TestInodePath(t *testing.T) {
	for i, test := range inodePathTests {
		inode := inodePath(test.path)

		if inode != test.inode {
			t.Errorf("test %d: wrong inode returned, want %d, got %d", i, test.inode, inode)
		}
	}
}
<|endoftext|>


<commit_before>package test

import (
	"testing"
	"time"

	"golang.org/x/net/websocket"

	"github.com/neptulon/neptulon"
	"github.com/neptulon/neptulon/middleware"
	"github.com/neptulon/randstr"
)

type echoMsg struct {
	Message string `json:"message"`
}

var (
	msg1 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
	msg2 = "In sit amet lectus felis, at pellentesque turpis."
	msg3 = "Nunc urna enim, cursus varius aliquet ac, imperdiet eget tellus."
	msg4 = randstr.Get(45 * 1000)       // 0.45 MB
	msg5 = randstr.Get(5 * 1000 * 1000) // 5.0 MB
)

func TestEchoWithoutTestHelpers(t *testing.T) {
	s := neptulon.NewServer("127.0.0.1:3001")
	go s.ListenAndServe()
	time.Sleep(time.Millisecond * 30)
	defer s.Close()

	s.MiddlewareFunc(func(ctx *neptulon.ReqCtx) error {
		t.Log("Request received:", ctx.Method)
		ctx.Res = "response-wow!"
		return ctx.Next()
	})

	origin := "http://127.0.0.1"
	url := "ws://127.0.0.1:3001"
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		t.Fatal(err)
	}
	if err := websocket.JSON.Send(ws, map[string]string{"id": "123", "method": "test"}); err != nil {
		t.Fatal(err)
	}
	var res interface{}
	if err := websocket.JSON.Receive(ws, &res); err != nil {
		t.Fatal(err)
	}
	t.Log("Got response:", res)

	if err := ws.Close(); err != nil {
		t.Fatal(err)
	}

	time.Sleep(time.Millisecond * 30)
}

func TestEcho(t *testing.T) {
	sh := NewServerHelper(t)
	rout := middleware.NewRouter()
	sh.Server.MiddlewareFunc(middleware.Logger)
	sh.Server.Middleware(rout)
	rout.Request("echo", middleware.Echo)
	defer sh.ListenAndServe().CloseWait()

	ch := sh.GetConnHelper()
	defer ch.Connect().CloseWait()

	m := "Hello!"
	ch.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
		var msg echoMsg
		if err := ctx.Result(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Message != m {
			t.Fatalf("expected: %v got: %v", m, msg.Message)
		}
		return nil
	})
}

func TestMessages(t *testing.T) {
	// todo: verify all message echoes from small to big
}

func TestBidirectional(t *testing.T) {
	// todo: test simultaneous read/writes
}

func TestTLS(t *testing.T) {
	// todo: ...
}

func TestError(t *testing.T) {
	sh := NewServerHelper(t)
	sh.Server.MiddlewareFunc(middleware.Logger)
	sh.Server.MiddlewareFunc(func(ctx *neptulon.ReqCtx) error {
		ctx.Err = &neptulon.ResError{
			Code:    1234,
			Message: "much error",
			Data:    map[string]string{"keykey": "valuevalue"},
		}
		return ctx.Next()
	})
	defer sh.ListenAndServe().CloseWait()

	ch := sh.GetConnHelper()
	defer ch.Connect().CloseWait()

	ch.SendRequest("testerror", nil, func(ctx *neptulon.ResCtx) error {

		return nil
	})
}
<commit_msg>add TestError impl<commit_after>package test

import (
	"testing"
	"time"

	"golang.org/x/net/websocket"

	"github.com/neptulon/neptulon"
	"github.com/neptulon/neptulon/middleware"
	"github.com/neptulon/randstr"
)

type echoMsg struct {
	Message string `json:"message"`
}

var (
	msg1 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
	msg2 = "In sit amet lectus felis, at pellentesque turpis."
	msg3 = "Nunc urna enim, cursus varius aliquet ac, imperdiet eget tellus."
	msg4 = randstr.Get(45 * 1000)       // 0.45 MB
	msg5 = randstr.Get(5 * 1000 * 1000) // 5.0 MB
)

func TestEchoWithoutTestHelpers(t *testing.T) {
	s := neptulon.NewServer("127.0.0.1:3001")
	go s.ListenAndServe()
	time.Sleep(time.Millisecond * 30)
	defer s.Close()

	s.MiddlewareFunc(func(ctx *neptulon.ReqCtx) error {
		t.Log("Request received:", ctx.Method)
		ctx.Res = "response-wow!"
		return ctx.Next()
	})

	origin := "http://127.0.0.1"
	url := "ws://127.0.0.1:3001"
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		t.Fatal(err)
	}
	if err := websocket.JSON.Send(ws, map[string]string{"id": "123", "method": "test"}); err != nil {
		t.Fatal(err)
	}
	var res interface{}
	if err := websocket.JSON.Receive(ws, &res); err != nil {
		t.Fatal(err)
	}
	t.Log("Got response:", res)

	if err := ws.Close(); err != nil {
		t.Fatal(err)
	}

	time.Sleep(time.Millisecond * 30)
}

func TestEcho(t *testing.T) {
	sh := NewServerHelper(t)
	rout := middleware.NewRouter()
	sh.Server.MiddlewareFunc(middleware.Logger)
	sh.Server.Middleware(rout)
	rout.Request("echo", middleware.Echo)
	defer sh.ListenAndServe().CloseWait()

	ch := sh.GetConnHelper()
	defer ch.Connect().CloseWait()

	m := "Hello!"
	ch.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
		var msg echoMsg
		if err := ctx.Result(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Message != m {
			t.Fatalf("expected: %v got: %v", m, msg.Message)
		}
		return nil
	})
}

func TestMessages(t *testing.T) {
	// todo: verify all message echoes from small to big
}

func TestBidirectional(t *testing.T) {
	// todo: test simultaneous read/writes
}

func TestTLS(t *testing.T) {
	// todo: ...
}

func TestError(t *testing.T) {
	sh := NewServerHelper(t)
	sh.Server.MiddlewareFunc(middleware.Logger)
	sh.Server.MiddlewareFunc(func(ctx *neptulon.ReqCtx) error {
		ctx.Err = &neptulon.ResError{
			Code:    1234,
			Message: "much error",
			Data:    map[string]string{"keykey": "valuevalue"},
		}
		return ctx.Next()
	})
	defer sh.ListenAndServe().CloseWait()

	ch := sh.GetConnHelper()
	defer ch.Connect().CloseWait()

	ch.SendRequest("testerror", nil, func(ctx *neptulon.ResCtx) error {
		var v map[string]string
		if ctx.Success {
			t.Error("expected to get error response")
		}
		if ctx.Result(&v) == nil {
			t.Error("did not expect to get any result for expected error response")
		}
		if ctx.ErrorCode != 1234 {
			t.Errorf("expected error code %v got %v", 1234, ctx.ErrorCode)
		}
		if ctx.ErrorMessage != "much error" {
			t.Errorf("expected error message %v got %v", "much error", ctx.ErrorMessage)
		}
		if ctx.ErrorData(&v) != nil || v["keykey"] != "valuevalue" {
			t.Errorf("expected error data %v got %v or errored during deserialization", "valuevalue", v["keykey"])
		}

		return nil
	})
}
<|endoftext|>


<commit_before>// Package forward implements a forwarding proxy. It caches an upstream net.Conn for some time, so if the same
// client returns the upstream's Conn will be precached. Depending on how you benchmark this looks to be
// 50% faster than just opening a new connection for every client. It works with UDP and TCP and uses
// inband healthchecking.
package forward

import (
	"context"
	"io"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/coredns/coredns/request"

	"github.com/miekg/dns"
)

// limitTimeout is a utility function to auto-tune timeout values
// average observed time is moved towards the last observed delay moderated by a weight
// next timeout to use will be the double of the computed average, limited by min and max frame.
func limitTimeout(currentAvg *int64, minValue time.Duration, maxValue time.Duration) time.Duration {
	rt := time.Duration(atomic.LoadInt64(currentAvg))
	if rt < minValue {
		return minValue
	}
	if rt < maxValue/2 {
		return 2 * rt
	}
	return maxValue
}

func averageTimeout(currentAvg *int64, observedDuration time.Duration, weight int64) {
	dt := time.Duration(atomic.LoadInt64(currentAvg))
	atomic.AddInt64(currentAvg, int64(observedDuration-dt)/weight)
}

func (t *Transport) dialTimeout() time.Duration {
	return limitTimeout(&t.avgDialTime, minDialTimeout, maxDialTimeout)
}

func (t *Transport) updateDialTimeout(newDialTime time.Duration) {
	averageTimeout(&t.avgDialTime, newDialTime, cumulativeAvgWeight)
}

// Dial dials the address configured in transport, potentially reusing a connection or creating a new one.
func (t *Transport) Dial(proto string) (*persistConn, bool, error) {
	// If tls has been configured; use it.
	if t.tlsConfig != nil {
		proto = "tcp-tls"
	}

	t.dial <- proto
	pc := <-t.ret

	if pc != nil {
		ConnCacheHitsCount.WithLabelValues(t.addr, proto).Add(1)
		return pc, true, nil
	}
	ConnCacheMissesCount.WithLabelValues(t.addr, proto).Add(1)

	reqTime := time.Now()
	timeout := t.dialTimeout()
	if proto == "tcp-tls" {
		conn, err := dns.DialTimeoutWithTLS("tcp", t.addr, t.tlsConfig, timeout)
		t.updateDialTimeout(time.Since(reqTime))
		return &persistConn{c: conn}, false, err
	}
	conn, err := dns.DialTimeout(proto, t.addr, timeout)
	t.updateDialTimeout(time.Since(reqTime))
	return &persistConn{c: conn}, false, err
}

// Connect selects an upstream, sends the request and waits for a response.
func (p *Proxy) Connect(ctx context.Context, state request.Request, opts options) (*dns.Msg, error) {
	start := time.Now()

	proto := ""
	switch {
	case opts.forceTCP: // TCP flag has precedence over UDP flag
		proto = "tcp"
	case opts.preferUDP:
		proto = "udp"
	default:
		proto = state.Proto()
	}

	pc, cached, err := p.transport.Dial(proto)
	if err != nil {
		return nil, err
	}

	// Set buffer size correctly for this client.
	pc.c.UDPSize = uint16(state.Size())
	if pc.c.UDPSize < 512 {
		pc.c.UDPSize = 512
	}

	pc.c.SetWriteDeadline(time.Now().Add(maxTimeout))
	if err := pc.c.WriteMsg(state.Req); err != nil {
		pc.c.Close() // not giving it back
		if err == io.EOF && cached {
			return nil, ErrCachedClosed
		}
		return nil, err
	}

	var ret *dns.Msg
	pc.c.SetReadDeadline(time.Now().Add(readTimeout))
	for {
		ret, err = pc.c.ReadMsg()
		if err != nil {
			pc.c.Close() // not giving it back
			if err == io.EOF && cached {
				return nil, ErrCachedClosed
			}
			return ret, err
		}
		// drop out-of-order responses
		if state.Req.Id == ret.Id {
			break
		}
	}

	p.transport.Yield(pc)

	rc, ok := dns.RcodeToString[ret.Rcode]
	if !ok {
		rc = strconv.Itoa(ret.Rcode)
	}

	RequestCount.WithLabelValues(p.addr).Add(1)
	RcodeCount.WithLabelValues(rc, p.addr).Add(1)
	RequestDuration.WithLabelValues(p.addr, rc).Observe(time.Since(start).Seconds())

	return ret, nil
}

const cumulativeAvgWeight = 4
<commit_msg>plugin/forward: Use new msg.Id for upstream queries (#4841)<commit_after>// Package forward implements a forwarding proxy. It caches an upstream net.Conn for some time, so if the same
// client returns the upstream's Conn will be precached. Depending on how you benchmark this looks to be
// 50% faster than just opening a new connection for every client. It works with UDP and TCP and uses
// inband healthchecking.
package forward

import (
	"context"
	"io"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/coredns/coredns/request"

	"github.com/miekg/dns"
)

// limitTimeout is a utility function to auto-tune timeout values
// average observed time is moved towards the last observed delay moderated by a weight
// next timeout to use will be the double of the computed average, limited by min and max frame.
func limitTimeout(currentAvg *int64, minValue time.Duration, maxValue time.Duration) time.Duration {
	rt := time.Duration(atomic.LoadInt64(currentAvg))
	if rt < minValue {
		return minValue
	}
	if rt < maxValue/2 {
		return 2 * rt
	}
	return maxValue
}

func averageTimeout(currentAvg *int64, observedDuration time.Duration, weight int64) {
	dt := time.Duration(atomic.LoadInt64(currentAvg))
	atomic.AddInt64(currentAvg, int64(observedDuration-dt)/weight)
}

func (t *Transport) dialTimeout() time.Duration {
	return limitTimeout(&t.avgDialTime, minDialTimeout, maxDialTimeout)
}

func (t *Transport) updateDialTimeout(newDialTime time.Duration) {
	averageTimeout(&t.avgDialTime, newDialTime, cumulativeAvgWeight)
}

// Dial dials the address configured in transport, potentially reusing a connection or creating a new one.
func (t *Transport) Dial(proto string) (*persistConn, bool, error) {
	// If tls has been configured; use it.
	if t.tlsConfig != nil {
		proto = "tcp-tls"
	}

	t.dial <- proto
	pc := <-t.ret

	if pc != nil {
		ConnCacheHitsCount.WithLabelValues(t.addr, proto).Add(1)
		return pc, true, nil
	}
	ConnCacheMissesCount.WithLabelValues(t.addr, proto).Add(1)

	reqTime := time.Now()
	timeout := t.dialTimeout()
	if proto == "tcp-tls" {
		conn, err := dns.DialTimeoutWithTLS("tcp", t.addr, t.tlsConfig, timeout)
		t.updateDialTimeout(time.Since(reqTime))
		return &persistConn{c: conn}, false, err
	}
	conn, err := dns.DialTimeout(proto, t.addr, timeout)
	t.updateDialTimeout(time.Since(reqTime))
	return &persistConn{c: conn}, false, err
}

// Connect selects an upstream, sends the request and waits for a response.
func (p *Proxy) Connect(ctx context.Context, state request.Request, opts options) (*dns.Msg, error) {
	start := time.Now()

	proto := ""
	switch {
	case opts.forceTCP: // TCP flag has precedence over UDP flag
		proto = "tcp"
	case opts.preferUDP:
		proto = "udp"
	default:
		proto = state.Proto()
	}

	pc, cached, err := p.transport.Dial(proto)
	if err != nil {
		return nil, err
	}

	// Set buffer size correctly for this client.
	pc.c.UDPSize = uint16(state.Size())
	if pc.c.UDPSize < 512 {
		pc.c.UDPSize = 512
	}

	pc.c.SetWriteDeadline(time.Now().Add(maxTimeout))
	// records the origin Id before upstream.
	originId := state.Req.Id
	state.Req.Id = dns.Id()
	defer func(){
		state.Req.Id = originId
	}()

	if err := pc.c.WriteMsg(state.Req); err != nil {
		pc.c.Close() // not giving it back
		if err == io.EOF && cached {
			return nil, ErrCachedClosed
		}
		return nil, err
	}

	var ret *dns.Msg
	pc.c.SetReadDeadline(time.Now().Add(readTimeout))
	for {
		ret, err = pc.c.ReadMsg()
		if err != nil {
			pc.c.Close() // not giving it back
			if err == io.EOF && cached {
				return nil, ErrCachedClosed
			}
			// recovery the origin Id after upstream.
			if ret != nil{
				ret.Id = originId
			}
			return ret, err
		}
		// drop out-of-order responses
		if state.Req.Id == ret.Id {
			break
		}
	}
	// recovery the origin Id after upstream.
	ret.Id = originId

	p.transport.Yield(pc)

	rc, ok := dns.RcodeToString[ret.Rcode]
	if !ok {
		rc = strconv.Itoa(ret.Rcode)
	}

	RequestCount.WithLabelValues(p.addr).Add(1)
	RcodeCount.WithLabelValues(rc, p.addr).Add(1)
	RequestDuration.WithLabelValues(p.addr, rc).Observe(time.Since(start).Seconds())

	return ret, nil
}

const cumulativeAvgWeight = 4
<|endoftext|>


<commit_before>package plugin

import (
	log "github.com/Sirupsen/logrus"
	. "github.com/talbright/go-curator"

	"errors"
	"sync"
)

var ErrNoAvailableWorkers = errors.New("no workers available to process work")
var ErrWorkerNotFound = errors.New("worker not found")
var ErrWorkNotFound = errors.New("work not found")
var ErrWorkerAlreadyExists = errors.New("worker already exists")
var ErrUnableToRemoveWorker = errors.New("could not remove worker")

type WorkSupervisorType interface {
	AddWorker(worker *Worker) error
	RemoveWorker(worker *Worker) error
	AddWork(node *Znode) error
	RemoveWork(node *Znode) error
	Load() error
}

type WorkSupervisor struct {
	work         *Work
	workers      *WorkerList
	workersLock  *sync.Mutex
	nextWorker   int
	client       *Client
	Logger       *log.Logger
	LogComponent string
}

func NewWorkSupervisor(client *Client, workPath string) *WorkSupervisor {
	return &WorkSupervisor{
		client:       client,
		work:         NewWork(client, workPath),
		workers:      NewWorkerList(),
		workersLock:  &sync.Mutex{},
		Logger:       log.StandardLogger(),
		LogComponent: "work_supervisor",
	}
}

func (s *WorkSupervisor) entry() *log.Entry {
	return s.Logger.WithField("component", s.LogComponent)
}

func (s *WorkSupervisor) AddWorker(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "AddWorker")
	entry.Info("adding worker")
	w := NewWorker(s.client, node.Path)
	if !s.workers.Add(w) {
		return ErrWorkerAlreadyExists
	}
	w.Children.LoadCache()
	w.Children.Clear()
	totalWork := s.work.Children.Size()
	workerCount := s.workers.Size()
	nodesToAdd := make([]Znode, 0)
	entry.WithFields(log.Fields{
		"totalWork":   totalWork,
		"workerCount": workerCount,
	}).Debug("preliminary calculations for work distribution")
	if workerCount > 1 && totalWork > 0 {
		workPerWorker := int(totalWork / workerCount)
		workRemainder := totalWork % workerCount
		if workRemainder != 0 {
			workPerWorker = workPerWorker - int(workRemainder/(workerCount-1))
		}
		entry.WithFields(log.Fields{
			"workerPerWorker": workPerWorker,
			"remainder":       workRemainder,
		}).Debug("post calculations for work distribution")
		if workPerWorker > 0 {
			for _, worker := range s.workers.ToSlice() {
				if worker.Id() == w.Id() {
					continue
				}
				var shiftAmount int
				if workRemainder == 0 {
					shiftAmount = worker.Children.Size() - workPerWorker
				} else {
					shiftAmount = workPerWorker
				}
				capacityLeft := worker.Children.Size() - shiftAmount
				if capacityLeft <= 0 {
					continue
				}
				if capacityLeft < workPerWorker {
					shiftAmount = worker.Children.Size() % workPerWorker
				}
				sliced := worker.ShiftWork(shiftAmount)
				entry.WithFields(log.Fields{
					"sliced": sliced,
					"worker": worker,
				}).Debug("work removed from existing worker")
				nodesToAdd = append(nodesToAdd, sliced...)
				if len(nodesToAdd) > workPerWorker {
					break
				}
			}
		}
	} else {
		nodesToAdd = s.work.Children.ToSlice()
	}
	entry.WithFields(log.Fields{
		"appended": nodesToAdd,
		"worker":   w,
	}).Debug("preparing to add nodes to worker")
	w.UnshiftWork(nodesToAdd)
	return nil
}

func (s *WorkSupervisor) RemoveWorker(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "RemoveWorker")
	entry.Info("removing worker")
	_, worker := s.workers.FindById(node.Path)
	if worker == nil {
		return ErrWorkerNotFound
	}
	children := worker.Children.ToSlice()
	if err = worker.Children.Clear(); err != nil {
		entry.WithError(err).WithField("worker", worker).Warn("unable to clear cache")
	}
	if !s.workers.Remove(worker) {
		return ErrUnableToRemoveWorker
	}
	for _, node := range children {
		newNode := node
		if assignErr := s.AssignWork(&newNode); assignErr != nil {
			entry.WithError(assignErr).WithField("node", node).WithField("worker", worker).Warn("unable to assign work to worker")
		}
	}
	return err
}

func (s *WorkSupervisor) AddWork(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "AddWork")
	entry.Info("add work")
	if err = s.work.Children.Add(node); err == nil {
		if err = s.AssignWork(node); err != nil {
			entry.WithError(err).WithField("node", node).Warn("unable to assign work to worker")
		}
	}
	return err
}

func (s *WorkSupervisor) RemoveWork(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "RemoveWork")
	entry.Info("remove work")
	if err = s.work.Children.Remove(node); err == nil {
		_, worker := s.workers.Find(func(i int, w *Worker) (found bool) {
			if w.Children.Contains(node.Name) {
				err = w.Children.Remove(node)
				found = true
			}
			return found
		})
		if err != nil {
			return err
		}
		if worker == nil {
			return ErrNoAvailableWorkers
		}
	}
	if err == ErrNodeNotInCache {
		err = ErrWorkNotFound
	}
	return err
}

func (s *WorkSupervisor) Load() (err error) {
	entry := s.entry().WithField("method", "Load")
	entry.Info("loading cache")
	if err := s.work.Children.LoadCache(); err != nil {
		entry.WithError(err).Error("error loading work cache")
	}
	for _, w := range s.workers.ToSlice() {
		if err := w.Children.LoadCache(); err != nil {
			entry.WithField("worker", w).WithError(err).Error("error loading worker cache")
		}
		w.Children.Clear()
	}
	for _, n := range s.work.Children.ToSlice() {
		s.AddWork(n.DeepCopy())
	}
	return nil
}

func (s *WorkSupervisor) AssignWork(work *Znode) error {
	entry := s.entry().WithField("work", work).WithField("method", "AssignWork")
	entry.Info("assigning work")
	s.workersLock.Lock()
	defer s.workersLock.Unlock()
	workerCount := s.workers.Size()
	if workerCount < 1 {
		s.nextWorker = 0
		return ErrNoAvailableWorkers
	}
	if s.nextWorker >= workerCount {
		s.nextWorker = 0
	}
	worker := s.workers.At(s.nextWorker)
	if worker == nil {
		return ErrNoAvailableWorkers
	}
	s.nextWorker++
	err := worker.Children.Add(work)
	entry.WithFields(log.Fields{
		"worker":      worker,
		"workerCount": workerCount,
		"index":       s.nextWorker,
		"children":    worker.Children.ToSlice(),
		"error":       err,
	}).Debug("post assignment info")
	return err
}

func (s *WorkSupervisor) GetWorkers() *WorkerList {
	return s.workers
}

func (s *WorkSupervisor) GetWork() *Work {
	return s.work
}
<commit_msg>Improve balancing. Deleted code reduces the amount of work that needs to shifted needlessly.<commit_after>package plugin

import (
	log "github.com/Sirupsen/logrus"
	. "github.com/talbright/go-curator"

	"errors"
	"sync"
)

var ErrNoAvailableWorkers = errors.New("no workers available to process work")
var ErrWorkerNotFound = errors.New("worker not found")
var ErrWorkNotFound = errors.New("work not found")
var ErrWorkerAlreadyExists = errors.New("worker already exists")
var ErrUnableToRemoveWorker = errors.New("could not remove worker")

type WorkSupervisorType interface {
	AddWorker(worker *Worker) error
	RemoveWorker(worker *Worker) error
	AddWork(node *Znode) error
	RemoveWork(node *Znode) error
	Load() error
}

type WorkSupervisor struct {
	work         *Work
	workers      *WorkerList
	workersLock  *sync.Mutex
	nextWorker   int
	client       *Client
	Logger       *log.Logger
	LogComponent string
}

func NewWorkSupervisor(client *Client, workPath string) *WorkSupervisor {
	return &WorkSupervisor{
		client:       client,
		work:         NewWork(client, workPath),
		workers:      NewWorkerList(),
		workersLock:  &sync.Mutex{},
		Logger:       log.StandardLogger(),
		LogComponent: "work_supervisor",
	}
}

func (s *WorkSupervisor) entry() *log.Entry {
	return s.Logger.WithField("component", s.LogComponent)
}

func (s *WorkSupervisor) AddWorker(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "AddWorker")
	entry.Info("adding worker")
	w := NewWorker(s.client, node.Path)
	if !s.workers.Add(w) {
		return ErrWorkerAlreadyExists
	}
	w.Children.LoadCache()
	w.Children.Clear()
	totalWork := s.work.Children.Size()
	workerCount := s.workers.Size()
	nodesToAdd := make([]Znode, 0)
	entry.WithFields(log.Fields{
		"totalWork":   totalWork,
		"workerCount": workerCount,
	}).Debug("preliminary calculations for work distribution")
	if workerCount > 1 && totalWork > 0 {
		workPerWorker := int(totalWork / workerCount)
		workRemainder := totalWork % workerCount
		entry.WithFields(log.Fields{
			"workerPerWorker": workPerWorker,
			"remainder":       workRemainder,
		}).Debug("post calculations for work distribution")
		if workPerWorker > 0 {
			for _, worker := range s.workers.ToSlice() {
				if worker.Id() == w.Id() {
					continue
				}
				var shiftAmount int
				if workRemainder == 0 {
					shiftAmount = worker.Children.Size() - workPerWorker
				} else {
					shiftAmount = workPerWorker
				}
				capacityLeft := worker.Children.Size() - shiftAmount
				if capacityLeft <= 0 {
					continue
				}
				if capacityLeft < workPerWorker {
					shiftAmount = worker.Children.Size() % workPerWorker
				}
				sliced := worker.ShiftWork(shiftAmount)
				entry.WithFields(log.Fields{
					"sliced": sliced,
					"worker": worker,
				}).Debug("work removed from existing worker")
				nodesToAdd = append(nodesToAdd, sliced...)
				if len(nodesToAdd) > workPerWorker {
					break
				}
			}
		}
	} else {
		nodesToAdd = s.work.Children.ToSlice()
	}
	entry.WithFields(log.Fields{
		"appended": nodesToAdd,
		"worker":   w,
	}).Debug("preparing to add nodes to worker")
	w.UnshiftWork(nodesToAdd)
	return nil
}

func (s *WorkSupervisor) RemoveWorker(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "RemoveWorker")
	entry.Info("removing worker")
	_, worker := s.workers.FindById(node.Path)
	if worker == nil {
		return ErrWorkerNotFound
	}
	children := worker.Children.ToSlice()
	if err = worker.Children.Clear(); err != nil {
		entry.WithError(err).WithField("worker", worker).Warn("unable to clear cache")
	}
	if !s.workers.Remove(worker) {
		return ErrUnableToRemoveWorker
	}
	for _, node := range children {
		newNode := node
		if assignErr := s.AssignWork(&newNode); assignErr != nil {
			entry.WithError(assignErr).WithField("node", node).WithField("worker", worker).Warn("unable to assign work to worker")
		}
	}
	return err
}

func (s *WorkSupervisor) AddWork(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "AddWork")
	entry.Info("add work")
	if err = s.work.Children.Add(node); err == nil {
		if err = s.AssignWork(node); err != nil {
			entry.WithError(err).WithField("node", node).Warn("unable to assign work to worker")
		}
	}
	return err
}

func (s *WorkSupervisor) RemoveWork(node *Znode) (err error) {
	entry := s.entry().WithField("node", node.Spew()).WithField("method", "RemoveWork")
	entry.Info("remove work")
	if err = s.work.Children.Remove(node); err == nil {
		_, worker := s.workers.Find(func(i int, w *Worker) (found bool) {
			if w.Children.Contains(node.Name) {
				err = w.Children.Remove(node)
				found = true
			}
			return found
		})
		if err != nil {
			return err
		}
		if worker == nil {
			return ErrNoAvailableWorkers
		}
	}
	if err == ErrNodeNotInCache {
		err = ErrWorkNotFound
	}
	return err
}

func (s *WorkSupervisor) Load() (err error) {
	entry := s.entry().WithField("method", "Load")
	entry.Info("loading cache")
	if err := s.work.Children.LoadCache(); err != nil {
		entry.WithError(err).Error("error loading work cache")
	}
	for _, w := range s.workers.ToSlice() {
		if err := w.Children.LoadCache(); err != nil {
			entry.WithField("worker", w).WithError(err).Error("error loading worker cache")
		}
		w.Children.Clear()
	}
	for _, n := range s.work.Children.ToSlice() {
		s.AddWork(n.DeepCopy())
	}
	return nil
}

func (s *WorkSupervisor) AssignWork(work *Znode) error {
	entry := s.entry().WithField("work", work).WithField("method", "AssignWork")
	entry.Info("assigning work")
	s.workersLock.Lock()
	defer s.workersLock.Unlock()
	workerCount := s.workers.Size()
	if workerCount < 1 {
		s.nextWorker = 0
		return ErrNoAvailableWorkers
	}
	if s.nextWorker >= workerCount {
		s.nextWorker = 0
	}
	worker := s.workers.At(s.nextWorker)
	if worker == nil {
		return ErrNoAvailableWorkers
	}
	s.nextWorker++
	err := worker.Children.Add(work)
	entry.WithFields(log.Fields{
		"worker":      worker,
		"workerCount": workerCount,
		"index":       s.nextWorker,
		"children":    worker.Children.ToSlice(),
		"error":       err,
	}).Debug("post assignment info")
	return err
}

func (s *WorkSupervisor) GetWorkers() *WorkerList {
	return s.workers
}

func (s *WorkSupervisor) GetWork() *Work {
	return s.work
}
<|endoftext|>


<commit_before>package gutscript

// vim:list:

import "testing"
import "gutscript/codegen/phpcodegen"
import "io/ioutil"
import "errors"
import "path/filepath"
import "fmt"

var parserInputs = []string{
	`pi = 3.1415926`,
	`pi = 3.1415926 + 4`,
	`pi = 4455
a = 4 * (2 + 3)
b = a`,

	// XXX: need trailing line
	`x = 3 * -2`,
	`if a > 10
    a = 10
`,
	`if a > 10
    a = 10
else
    a = 0
`,
	`if a > 10
    a = 10
    b = 20
else
    a = 0
    b = 0
`,
	// test if && single else if
	`if a > 10
    a = 10
    b = 20
elseif b > 10
    a = 1
    b = 1
`,
	// test if, elseif && else
	`if a > 10
    a = 10
    b = 20
elseif b > 10
    a = 1
    b = 1
else
    a = 0
    b = 0
`,
	`foo :: ->
    a = 10
    return a
`,
}

func LexFile(srcFile string) error {
	bytes, err := ioutil.ReadFile(srcFile)
	if err != nil {
		return err
	}
	// code := string(bytes)
	lexer := GutsLex{
		input: string(bytes),
		start: 0,
		pos:   0,
		items: make(chan *GutsSymType, 100),
	}
	go lexer.run()
	dumpLexItems(lexer.items)
	return nil
}

func CompileFile(srcFile string) (string, error) {
	bytes, err := ioutil.ReadFile(srcFile)
	if err != nil {
		return "", err
	}
	// code := string(bytes)
	lexer := GutsLex{
		input: string(bytes),
		start: 0,
		pos:   0,
		items: make(chan *GutsSymType, 100),
	}
	go lexer.run()
	parser := GutsParser{}
	if parser.Parse(&lexer) == 1 {
		return "", errors.New("syntax error")
	}
	lexer.close()
	visitor := phpcodegen.Visitor{}
	fmt.Printf("AST: %#v\n", parser.Val.val)
	return visitor.Visit(parser.Val.val), nil
}

func TestCompileFile(t *testing.T) {
	srcFiles, err := filepath.Glob("tests/*.guts")
	if err != nil {
		t.Fatal(err)
	}
	for _, srcFile := range srcFiles {
		t.Log("Lexing", srcFile)
		LexFile(srcFile)

		t.Log("Compiling", srcFile)
		out, err := CompileFile(srcFile)
		if err != nil {
			t.Fatalf("Compilation of %s failed: %s", srcFile, err)
		}
		t.Log(out)
		_ = out
	}
}

func TestCodeGen(t *testing.T) {
	for _, input := range parserInputs {
		t.Log(input)
		lexer := GutsLex{
			input: input,
			start: 0,
			pos:   0,
			items: make(chan *GutsSymType, 100),
		}
		go lexer.run()
		parser := GutsParser{}
		if parser.Parse(&lexer) == 1 {
			t.Fatal("syntax error")
		}
		lexer.close()
		t.Logf("AST: %#v", parser.Val.val)
		visitor := phpcodegen.Visitor{}
		t.Log(visitor.Visit(parser.Val.val))
	}
}

func TestParser(t *testing.T) {
	for _, input := range parserInputs {
		t.Log(input)
		lexer := GutsLex{
			input: input,
			start: 0,
			pos:   0,
			items: make(chan *GutsSymType, 100),
		}
		go lexer.run()
		parser := GutsParser{}
		if parser.Parse(&lexer) == 1 {
			t.Fatal("syntax error")
		}
		lexer.close()
	}
}
<commit_msg>Remove CodeGen Test function since we have TestCompile()<commit_after>package gutscript

// vim:list:

import "testing"
import "gutscript/codegen/phpcodegen"
import "io/ioutil"
import "errors"
import "path/filepath"
import "fmt"

var parserInputs = []string{
	`pi = 3.1415926`,
	`pi = 3.1415926 + 4`,
	`pi = 4455
a = 4 * (2 + 3)
b = a`,

	// XXX: need trailing line
	`x = 3 * -2`,
	`if a > 10
    a = 10
`,
	`if a > 10
    a = 10
else
    a = 0
`,
	`if a > 10
    a = 10
    b = 20
else
    a = 0
    b = 0
`,
	// test if && single else if
	`if a > 10
    a = 10
    b = 20
elseif b > 10
    a = 1
    b = 1
`,
	// test if, elseif && else
	`if a > 10
    a = 10
    b = 20
elseif b > 10
    a = 1
    b = 1
else
    a = 0
    b = 0
`,
	`foo :: ->
    a = 10
    return a
`,
}

func LexFile(srcFile string) error {
	bytes, err := ioutil.ReadFile(srcFile)
	if err != nil {
		return err
	}
	// code := string(bytes)
	lexer := GutsLex{
		input: string(bytes),
		start: 0,
		pos:   0,
		items: make(chan *GutsSymType, 100),
	}
	go lexer.run()
	dumpLexItems(lexer.items)
	return nil
}

func CompileFile(srcFile string) (string, error) {
	bytes, err := ioutil.ReadFile(srcFile)
	if err != nil {
		return "", err
	}
	// code := string(bytes)
	lexer := GutsLex{
		input: string(bytes),
		start: 0,
		pos:   0,
		items: make(chan *GutsSymType, 100),
	}
	go lexer.run()
	parser := GutsParser{}
	if parser.Parse(&lexer) == 1 {
		return "", errors.New("syntax error")
	}
	lexer.close()
	visitor := phpcodegen.Visitor{}
	fmt.Printf("AST: %#v\n", parser.Val.val)
	return visitor.Visit(parser.Val.val), nil
}

func TestCompileFile(t *testing.T) {
	srcFiles, err := filepath.Glob("tests/*.guts")
	if err != nil {
		t.Fatal(err)
	}
	for _, srcFile := range srcFiles {
		t.Log("Lexing", srcFile)
		LexFile(srcFile)

		t.Log("Compiling", srcFile)
		out, err := CompileFile(srcFile)
		if err != nil {
			t.Fatalf("Compilation of %s failed: %s", srcFile, err)
		}
		t.Log(out)
		_ = out
	}
}

func TestParser(t *testing.T) {
	for _, input := range parserInputs {
		t.Log(input)
		lexer := GutsLex{
			input: input,
			start: 0,
			pos:   0,
			items: make(chan *GutsSymType, 100),
		}
		go lexer.run()
		parser := GutsParser{}
		if parser.Parse(&lexer) == 1 {
			t.Fatal("syntax error")
		}
		lexer.close()
	}
}
<|endoftext|>