<commit_before>package shared

import (
	"bytes"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	"github.com/sneakybeaky/aws-volumes/shared/iface"
	"github.com/sneakybeaky/aws-volumes/shared/log"
)

// VolumeTagPrefix prefixes the name of a tag describing an allocated volume
const VolumeTagPrefix = "volume_"

// DetachVolumesTag when set to a true value signals volumes can be detached
const DetachVolumesTag = "detach_volumes"

// EC2Instance provides metadata about an EC2 instance.
type EC2Instance struct {
	svc      ec2iface.EC2API
	metadata iface.Metadata
}

func NewEC2Instance(metadata iface.Metadata, svc ec2iface.EC2API) *EC2Instance {

	return &EC2Instance{
		svc:      svc,
		metadata: metadata,
	}

}

func (e EC2Instance) Tags() ([]*ec2.TagDescription, error) {

	instanceid, err := e.metadata.InstanceID()

	if err != nil {
		return nil, err
	}

	params := &ec2.DescribeTagsInput{
		Filters: []*ec2.Filter{
			{ // Required
				Name: aws.String("resource-id"),
				Values: []*string{
					aws.String(instanceid), // Required
				},
			},
		},
		MaxResults: aws.Int64(1000),
	}
	resp, err := e.svc.DescribeTags(params)

	if err != nil {
		return nil, err
	}

	return resp.Tags, nil
}

func (e EC2Instance) AllocatedVolumes() ([]*AllocatedVolume, error) {
	var allocated []*AllocatedVolume

	tags, err := e.Tags()

	if err != nil {
		return allocated, err
	}

	for _, tag := range tags {
		if strings.HasPrefix(*tag.Key, VolumeTagPrefix) {

			key := *tag.Key
			device := key[len(VolumeTagPrefix):]
			allocated = append(allocated, NewAllocatedVolume(*tag.Value, device, *tag.ResourceId, e.svc))
		}
	}

	return allocated, nil
}

func (e EC2Instance) DetachVolumes() error {

	tags, err := e.Tags()

	if err != nil {
		return err
	}

	for _, tag := range tags {
		if *tag.Key == DetachVolumesTag {

			detachVolumes, _ := strconv.ParseBool(*tag.Value)

			if detachVolumes {
				e.applyToVolumes(detachVolume)
			} else {
				log.Debug.Printf("Tag '%s' value is '%s' - not detaching volumes", DetachVolumesTag, *tag.Value)
			}

			break

		}
	}

	return nil
}

func (e EC2Instance) AttachVolumes() {
	e.applyToVolumes(attachVolume)
}

func (e EC2Instance) ShowVolumesInfo() {
	e.applyToVolumes(showVolumeInfo)
}

var attachVolume = func(volume *AllocatedVolume) {

	if err := volume.Attach(); err != nil {
		log.Error.Printf("Unable to attach volume : %s\n", err)
	}
}

var detachVolume = func(volume *AllocatedVolume) {

	if err := volume.Detach(); err != nil {
		log.Error.Printf("Unable to detach volume : %s\n", err)
	}
}

var showVolumeInfo = func(volume *AllocatedVolume) {
	buf := new(bytes.Buffer)

	if err := volume.Info(buf); err != nil {
		log.Error.Printf("Unable to get info for volume : %s\n", err)
		return
	}
	os.Stdout.WriteString(buf.String())

}

func (e EC2Instance) applyToVolumes(action func(volume *AllocatedVolume)) {
	if volumes, err := e.AllocatedVolumes(); err != nil {
		log.Error.Printf("Unable to find allocated volumes : %s", err)
	} else {

		var wg sync.WaitGroup

		for _, volume := range volumes {

			wg.Add(1)
			go func(action func(volume *AllocatedVolume), volume *AllocatedVolume) {

				defer wg.Done()
				action(volume)

			}(action, volume)

		}

		wg.Wait()
	}

}
<commit_msg>Alerts when not detaching volumes because of tag logic. Fixes #2<commit_after>package shared

import (
	"bytes"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	"github.com/sneakybeaky/aws-volumes/shared/iface"
	"github.com/sneakybeaky/aws-volumes/shared/log"
)

// VolumeTagPrefix prefixes the name of a tag describing an allocated volume
const VolumeTagPrefix = "volume_"

// DetachVolumesTag when set to a true value signals volumes can be detached
const DetachVolumesTag = "detach_volumes"

// EC2Instance provides metadata about an EC2 instance.
type EC2Instance struct {
	svc      ec2iface.EC2API
	metadata iface.Metadata
}

func NewEC2Instance(metadata iface.Metadata, svc ec2iface.EC2API) *EC2Instance {

	return &EC2Instance{
		svc:      svc,
		metadata: metadata,
	}

}

func (e EC2Instance) Tags() ([]*ec2.TagDescription, error) {

	instanceid, err := e.metadata.InstanceID()

	if err != nil {
		return nil, err
	}

	params := &ec2.DescribeTagsInput{
		Filters: []*ec2.Filter{
			{ // Required
				Name: aws.String("resource-id"),
				Values: []*string{
					aws.String(instanceid), // Required
				},
			},
		},
		MaxResults: aws.Int64(1000),
	}
	resp, err := e.svc.DescribeTags(params)

	if err != nil {
		return nil, err
	}

	return resp.Tags, nil
}

func (e EC2Instance) AllocatedVolumes() ([]*AllocatedVolume, error) {
	var allocated []*AllocatedVolume

	tags, err := e.Tags()

	if err != nil {
		return allocated, err
	}

	for _, tag := range tags {
		if strings.HasPrefix(*tag.Key, VolumeTagPrefix) {

			key := *tag.Key
			device := key[len(VolumeTagPrefix):]
			allocated = append(allocated, NewAllocatedVolume(*tag.Value, device, *tag.ResourceId, e.svc))
		}
	}

	return allocated, nil
}

func (e EC2Instance) DetachVolumes() error {

	tags, err := e.Tags()

	if err != nil {
		return err
	}

	for _, tag := range tags {
		if *tag.Key == DetachVolumesTag {

			detachVolumes, _ := strconv.ParseBool(*tag.Value)

			if detachVolumes {
				e.applyToVolumes(detachVolume)
			} else {
				log.Info.Printf("Tag '%s' value is '%s' - not detaching volumes", DetachVolumesTag, *tag.Value)
			}

			break

		}
	}

	return nil
}

func (e EC2Instance) AttachVolumes() {
	e.applyToVolumes(attachVolume)
}

func (e EC2Instance) ShowVolumesInfo() {
	e.applyToVolumes(showVolumeInfo)
}

var attachVolume = func(volume *AllocatedVolume) {

	if err := volume.Attach(); err != nil {
		log.Error.Printf("Unable to attach volume : %s\n", err)
	}
}

var detachVolume = func(volume *AllocatedVolume) {

	if err := volume.Detach(); err != nil {
		log.Error.Printf("Unable to detach volume : %s\n", err)
	}
}

var showVolumeInfo = func(volume *AllocatedVolume) {
	buf := new(bytes.Buffer)

	if err := volume.Info(buf); err != nil {
		log.Error.Printf("Unable to get info for volume : %s\n", err)
		return
	}
	os.Stdout.WriteString(buf.String())

}

func (e EC2Instance) applyToVolumes(action func(volume *AllocatedVolume)) {
	if volumes, err := e.AllocatedVolumes(); err != nil {
		log.Error.Printf("Unable to find allocated volumes : %s", err)
	} else {

		var wg sync.WaitGroup

		for _, volume := range volumes {

			wg.Add(1)
			go func(action func(volume *AllocatedVolume), volume *AllocatedVolume) {

				defer wg.Done()
				action(volume)

			}(action, volume)

		}

		wg.Wait()
	}

}
<|endoftext|>


<commit_before>package logpeck

import (
	"log"
	"net/http"
	"net/http/httputil"
)

func NewAddTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Printf("AddTaskHandler")
		r_str, _ := httputil.DumpRequest(r, false)
		log.Printf("Request len[%d], body[%s]", len(r_str), r_str)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("AddTaskHandler Success\n"))
	}
}

func NewUpdateTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Printf("UpdateTaskHandler")
		r_str, _ := httputil.DumpRequest(r, false)
		log.Printf("Request len[%d], body[%s]", len(r_str), r_str)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("UpdateTaskHandler Success\n"))
	}
}

func NewPauseTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Printf("PauseTaskHandler")
		r_str, _ := httputil.DumpRequest(r, false)
		log.Printf("Request len[%d], body[%s]", len(r_str), r_str)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("PauseTaskHandler Success\n"))
	}
}

func NewRemoveTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Printf("RemoveTaskHandler")
		r_str, _ := httputil.DumpRequest(r, false)
		log.Printf("Request len[%d], body[%s]", len(r_str), r_str)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("RemoveTaskHandler Success\n"))
	}
}
<commit_msg>update http handler<commit_after>package logpeck

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
)

func logRequest(r *http.Request, prefix string) {
	r_str, _ := httputil.DumpRequest(r, true)
	log.Printf("[%s] req_len[%d] req[%s]", prefix, len(r_str), r_str)
}

func NewAddTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logRequest(r, "AddTaskHandler")
		defer r.Body.Close()

		var config PeckTaskConfig
		raw, _ := ioutil.ReadAll(r.Body)
		err := json.Unmarshal(raw, &config)
		if err != nil {
			log.Printf("Parse PeckTaskConfig error, %s", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Bad Request\n"))
			return
		}

		err = pecker.AddPeckTask(&config)
		if err != nil {
			log.Printf("AddTaskConfig error, %s", err)
			w.WriteHeader(http.StatusNotAcceptable)
			w.Write([]byte("Add failed, " + err.Error() + "\n"))
			return
		}

		err = db.SaveConfig(&config)
		if err != nil {
			log.Printf("AddTaskConfig error, save config error, %s", err)
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error() + "\n"))
			return
		}

		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK\n"))
		return
	}
}

func NewUpdateTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logRequest(r, "UpdateTaskHandler")
		defer r.Body.Close()

		var config PeckTaskConfig
		raw, _ := ioutil.ReadAll(r.Body)
		err := json.Unmarshal(raw, &config)
		if err != nil {
			log.Printf("Parse PeckTaskConfig error, %s", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Bad Request\n"))
			return
		}

		err = pecker.UpdatePeckTask(&config)
		if err != nil {
			w.WriteHeader(http.StatusNotAcceptable)
			w.Write([]byte("Update failed, " + err.Error() + "\n"))
			return
		}

		err = db.SaveConfig(&config)
		if err != nil {
			log.Printf("UpdateTaskConfig error, save config error, %s", err)
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error() + "\n"))
			return
		}

		w.WriteHeader(http.StatusOK)
		w.Write([]byte("UpdateTaskHandler Success\n"))
	}
}

func NewStartTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logRequest(r, "StartTaskHandler")
		defer r.Body.Close()

		var config PeckTaskConfig
		raw, _ := ioutil.ReadAll(r.Body)
		err := json.Unmarshal(raw, &config)
		if err != nil {
			log.Printf("Start PeckTaskConfig error, %s", err)
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Bad Request\n"))
			return
		}

		config_w, c_err := db.GetConfig(&config)
		if c_err != nil {
			w.WriteHeader(http.StatusNotAcceptable)
			w.Write([]byte("Start failed, " + c_err.Error() + "\n"))
			return
		}

		err = pecker.StartPeckTask(config_w)
		if err != nil {
			w.WriteHeader(http.StatusNotAcceptable)
			w.Write([]byte("Update failed, " + err.Error() + "\n"))
			return
		}

		err = db.SaveConfig(&config)
		if err != nil {
			log.Printf("UpdateTaskConfig error, save config error, %s", err)
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error() + "\n"))
			return
		}

		w.WriteHeader(http.StatusOK)
		w.Write([]byte("PauseTaskHandler Success\n"))
	}
}

func NewPauseTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logRequest(r, "PauseTaskHandler")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("PauseTaskHandler Success\n"))
	}
}

func NewRemoveTaskHandler(pecker *Pecker, db *DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logRequest(r, "RemoveTaskHandler")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("RemoveTaskHandler Success\n"))
	}
}
<|endoftext|>


<commit_before>package itchio

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/mitchellh/mapstructure"
	"github.com/pkg/errors"
)

func getDebugLevel() int64 {
	val, err := strconv.ParseInt(os.Getenv("GO_ITCHIO_DEBUG"), 10, 64)
	if err != nil {
		return 0
	}
	return val
}

var (
	debugLevel   = getDebugLevel()
	logRequests  = debugLevel >= 1
	dumpAPICalls = debugLevel >= 2
)

// Get performs an HTTP GET request to the API
func (c *Client) Get(ctx context.Context, url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	return c.Do(req)
}

// GetResponse performs an HTTP GET request and parses the API response.
func (c *Client) GetResponse(ctx context.Context, url string, dst interface{}) error {
	resp, err := c.Get(ctx, url)
	if err != nil {
		return errors.WithStack(err)
	}

	err = ParseAPIResponse(dst, resp)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// PostForm performs an HTTP POST request to the API, with url-encoded parameters
func (c *Client) PostForm(ctx context.Context, url string, data url.Values) (*http.Response, error) {
	req, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return c.Do(req)
}

// PostFormResponse performs an HTTP POST request to the API *and* parses the API response.
func (c *Client) PostFormResponse(ctx context.Context, url string, data url.Values, dst interface{}) error {
	resp, err := c.PostForm(ctx, url, data)
	if err != nil {
		return errors.WithStack(err)
	}

	err = ParseAPIResponse(dst, resp)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Do performs a request (any method). It takes care of JWT or API key
// authentication, sets the proper user agent, has built-in retry,
func (c *Client) Do(req *http.Request) (*http.Response, error) {
	c.Limiter.Wait(req.Context())
	req.Header.Add("Authorization", c.Key)
	req.Header.Set("User-Agent", c.UserAgent)
	req.Header.Set("Accept-Language", c.AcceptedLanguage)
	req.Header.Set("Accept", "application/vnd.itch.v2")

	var res *http.Response
	var err error

	if logRequests {
		fmt.Fprintf(os.Stderr, "%s %s [request]\n", req.Method, req.URL)
	}

	if dumpAPICalls {
		for k, vv := range req.Header {
			for _, v := range vv {
				fmt.Fprintf(os.Stderr, "[request] %s: %s\n", k, v)
			}
		}
	}

	retryPatterns := append(c.RetryPatterns, time.Millisecond)

	for _, sleepTime := range retryPatterns {
		if c.onOutgoingRequest != nil {
			c.onOutgoingRequest(req)
		}
		res, err = c.HTTPClient.Do(req)
		if err != nil {
			if strings.Contains(err.Error(), "TLS handshake timeout") {
				time.Sleep(sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond)
				continue
			}
			return nil, err
		}

		if res.StatusCode == 503 {
			res.Body.Close()

			// Rate limited, try again according to patterns.
			// following https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload#exp-backoff to the letter
			actualSleepTime := sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond
			if c.onRateLimited != nil {
				c.onRateLimited(req, res)
			}
			if logRequests {
				fmt.Fprintf(os.Stderr, "%s %s [rate limited, sleeping %v]\n", req.Method, req.URL, actualSleepTime)
			}
			time.Sleep(actualSleepTime)
			continue
		}

		break
	}

	return res, err
}

// MakePath crafts an API url from our configured base URL
func (c *Client) MakePath(format string, a ...interface{}) string {
	return c.MakeValuesPath(nil, format, a...)
}

// MakeValuesPath crafts an API url from our configured base URL
func (c *Client) MakeValuesPath(values url.Values, format string, a ...interface{}) string {
	base := strings.Trim(c.BaseURL, "/")
	subPath := strings.Trim(fmt.Sprintf(format, a...), "/")
	path := fmt.Sprintf("%s/%s", base, subPath)
	if len(values) == 0 {
		return path
	}
	return fmt.Sprintf("%s?%s", path, values.Encode())
}

func asHTTPCodeError(res *http.Response) error {
	if res.StatusCode/100 != 2 {
		err := fmt.Errorf("Server error: HTTP %s for %s", res.Status, res.Request.URL.Path)
		return err
	}
	return nil
}

// ParseAPIResponse unmarshals an HTTP response into one of out response
// data structures
func ParseAPIResponse(dst interface{}, res *http.Response) error {
	if res == nil || res.Body == nil {
		return fmt.Errorf("No response from server")
	}

	bodyReader := res.Body
	defer bodyReader.Close()

	body, err := ioutil.ReadAll(bodyReader)
	if err != nil {
		return errors.WithStack(err)
	}

	if dumpAPICalls {
		fmt.Fprintf(os.Stderr, "[response] %s\n", string(body))
	}

	intermediate := make(map[string]interface{})

	err = json.NewDecoder(bytes.NewReader(body)).Decode(&intermediate)
	if err != nil {
		if he := asHTTPCodeError(res); he != nil {
			return he
		}

		msg := fmt.Sprintf("JSON decode error: %s\n\nBody: %s\n\n", err.Error(), string(body))
		return errors.New(msg)
	}

	if errorsField, ok := intermediate["errors"]; ok {
		if errorsList, ok := errorsField.([]interface{}); ok {
			var messages []string
			for _, el := range errorsList {
				if errorMessage, ok := el.(string); ok {
					messages = append(messages, errorMessage)
				}
			}
			if len(messages) > 0 {
				return &APIError{Messages: messages, StatusCode: res.StatusCode, Path: res.Request.URL.Path}
			}
		}
	}

	if he := asHTTPCodeError(res); he != nil {
		return he
	}

	intermediate = camelifyMap(intermediate)

	if dumpAPICalls {
		enc := json.NewEncoder(os.Stderr)
		enc.SetIndent("[intermediate] ", "  ")
		enc.Encode(intermediate)
	}

	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
		TagName: "json",
		Result:  dst,
		// see https://github.com/itchio/itch/issues/1549
		WeaklyTypedInput: true,
		DecodeHook: mapstructure.ComposeDecodeHookFunc(
			mapstructure.StringToTimeHookFunc(time.RFC3339Nano),
			GameHookFunc,
			UploadHookFunc,
		),
	})
	if err != nil {
		return errors.WithStack(err)
	}

	err = decoder.Decode(intermediate)
	if err != nil {
		msg := fmt.Sprintf("mapstructure decode error: %s\n\nBody: %#v\n\n", err.Error(), intermediate)
		return errors.New(msg)
	}

	if res.StatusCode != 200 {
		return errors.Errorf("HTTP %v", res.StatusCode)
	}

	return nil
}

// FindBuildFile looks for an uploaded file of the right type
// in a list of file. Returns nil if it can't find one.
func FindBuildFile(fileType BuildFileType, files []*BuildFile) *BuildFile {
	for _, f := range files {
		if f.Type == fileType && f.State == BuildFileStateUploaded {
			return f
		}
	}

	return nil
}

// FindBuildFileEx looks for an uploaded file of the right type
// and subtype in a list of file. Returns nil if it can't find one.
func FindBuildFileEx(fileType BuildFileType, fileSubType BuildFileSubType, files []*BuildFile) *BuildFile {
	for _, f := range files {
		if f.Type == fileType && f.SubType == fileSubType && f.State == BuildFileStateUploaded {
			return f
		}
	}

	return nil
}
<commit_msg>Call limiter *inside* retry loop<commit_after>package itchio

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/mitchellh/mapstructure"
	"github.com/pkg/errors"
)

func getDebugLevel() int64 {
	val, err := strconv.ParseInt(os.Getenv("GO_ITCHIO_DEBUG"), 10, 64)
	if err != nil {
		return 0
	}
	return val
}

var (
	debugLevel   = getDebugLevel()
	logRequests  = debugLevel >= 1
	dumpAPICalls = debugLevel >= 2
)

// Get performs an HTTP GET request to the API
func (c *Client) Get(ctx context.Context, url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	return c.Do(req)
}

// GetResponse performs an HTTP GET request and parses the API response.
func (c *Client) GetResponse(ctx context.Context, url string, dst interface{}) error {
	resp, err := c.Get(ctx, url)
	if err != nil {
		return errors.WithStack(err)
	}

	err = ParseAPIResponse(dst, resp)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// PostForm performs an HTTP POST request to the API, with url-encoded parameters
func (c *Client) PostForm(ctx context.Context, url string, data url.Values) (*http.Response, error) {
	req, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return c.Do(req)
}

// PostFormResponse performs an HTTP POST request to the API *and* parses the API response.
func (c *Client) PostFormResponse(ctx context.Context, url string, data url.Values, dst interface{}) error {
	resp, err := c.PostForm(ctx, url, data)
	if err != nil {
		return errors.WithStack(err)
	}

	err = ParseAPIResponse(dst, resp)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Do performs a request (any method). It takes care of JWT or API key
// authentication, sets the proper user agent, has built-in retry,
func (c *Client) Do(req *http.Request) (*http.Response, error) {
	req.Header.Add("Authorization", c.Key)
	req.Header.Set("User-Agent", c.UserAgent)
	req.Header.Set("Accept-Language", c.AcceptedLanguage)
	req.Header.Set("Accept", "application/vnd.itch.v2")

	var res *http.Response
	var err error

	if logRequests {
		fmt.Fprintf(os.Stderr, "%s %s [request]\n", req.Method, req.URL)
	}

	if dumpAPICalls {
		for k, vv := range req.Header {
			for _, v := range vv {
				fmt.Fprintf(os.Stderr, "[request] %s: %s\n", k, v)
			}
		}
	}

	retryPatterns := append(c.RetryPatterns, time.Millisecond)

	for _, sleepTime := range retryPatterns {
		c.Limiter.Wait(req.Context())
		if c.onOutgoingRequest != nil {
			c.onOutgoingRequest(req)
		}
		res, err = c.HTTPClient.Do(req)
		if err != nil {
			if strings.Contains(err.Error(), "TLS handshake timeout") {
				time.Sleep(sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond)
				continue
			}
			return nil, err
		}

		if res.StatusCode == 503 {
			res.Body.Close()

			// Rate limited, try again according to patterns.
			// following https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload#exp-backoff to the letter
			actualSleepTime := sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond
			if c.onRateLimited != nil {
				c.onRateLimited(req, res)
			}
			if logRequests {
				fmt.Fprintf(os.Stderr, "%s %s [rate limited, sleeping %v]\n", req.Method, req.URL, actualSleepTime)
			}
			time.Sleep(actualSleepTime)
			continue
		}

		break
	}

	return res, err
}

// MakePath crafts an API url from our configured base URL
func (c *Client) MakePath(format string, a ...interface{}) string {
	return c.MakeValuesPath(nil, format, a...)
}

// MakeValuesPath crafts an API url from our configured base URL
func (c *Client) MakeValuesPath(values url.Values, format string, a ...interface{}) string {
	base := strings.Trim(c.BaseURL, "/")
	subPath := strings.Trim(fmt.Sprintf(format, a...), "/")
	path := fmt.Sprintf("%s/%s", base, subPath)
	if len(values) == 0 {
		return path
	}
	return fmt.Sprintf("%s?%s", path, values.Encode())
}

func asHTTPCodeError(res *http.Response) error {
	if res.StatusCode/100 != 2 {
		err := fmt.Errorf("Server error: HTTP %s for %s", res.Status, res.Request.URL.Path)
		return err
	}
	return nil
}

// ParseAPIResponse unmarshals an HTTP response into one of out response
// data structures
func ParseAPIResponse(dst interface{}, res *http.Response) error {
	if res == nil || res.Body == nil {
		return fmt.Errorf("No response from server")
	}

	bodyReader := res.Body
	defer bodyReader.Close()

	body, err := ioutil.ReadAll(bodyReader)
	if err != nil {
		return errors.WithStack(err)
	}

	if dumpAPICalls {
		fmt.Fprintf(os.Stderr, "[response] %s\n", string(body))
	}

	intermediate := make(map[string]interface{})

	err = json.NewDecoder(bytes.NewReader(body)).Decode(&intermediate)
	if err != nil {
		if he := asHTTPCodeError(res); he != nil {
			return he
		}

		msg := fmt.Sprintf("JSON decode error: %s\n\nBody: %s\n\n", err.Error(), string(body))
		return errors.New(msg)
	}

	if errorsField, ok := intermediate["errors"]; ok {
		if errorsList, ok := errorsField.([]interface{}); ok {
			var messages []string
			for _, el := range errorsList {
				if errorMessage, ok := el.(string); ok {
					messages = append(messages, errorMessage)
				}
			}
			if len(messages) > 0 {
				return &APIError{Messages: messages, StatusCode: res.StatusCode, Path: res.Request.URL.Path}
			}
		}
	}

	if he := asHTTPCodeError(res); he != nil {
		return he
	}

	intermediate = camelifyMap(intermediate)

	if dumpAPICalls {
		enc := json.NewEncoder(os.Stderr)
		enc.SetIndent("[intermediate] ", "  ")
		enc.Encode(intermediate)
	}

	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
		TagName: "json",
		Result:  dst,
		// see https://github.com/itchio/itch/issues/1549
		WeaklyTypedInput: true,
		DecodeHook: mapstructure.ComposeDecodeHookFunc(
			mapstructure.StringToTimeHookFunc(time.RFC3339Nano),
			GameHookFunc,
			UploadHookFunc,
		),
	})
	if err != nil {
		return errors.WithStack(err)
	}

	err = decoder.Decode(intermediate)
	if err != nil {
		msg := fmt.Sprintf("mapstructure decode error: %s\n\nBody: %#v\n\n", err.Error(), intermediate)
		return errors.New(msg)
	}

	if res.StatusCode != 200 {
		return errors.Errorf("HTTP %v", res.StatusCode)
	}

	return nil
}

// FindBuildFile looks for an uploaded file of the right type
// in a list of file. Returns nil if it can't find one.
func FindBuildFile(fileType BuildFileType, files []*BuildFile) *BuildFile {
	for _, f := range files {
		if f.Type == fileType && f.State == BuildFileStateUploaded {
			return f
		}
	}

	return nil
}

// FindBuildFileEx looks for an uploaded file of the right type
// and subtype in a list of file. Returns nil if it can't find one.
func FindBuildFileEx(fileType BuildFileType, fileSubType BuildFileSubType, files []*BuildFile) *BuildFile {
	for _, f := range files {
		if f.Type == fileType && f.SubType == fileSubType && f.State == BuildFileStateUploaded {
			return f
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

var shellGodownloader = `#!/bin/sh
set -e
#  Code generated by godownloader. DO NOT EDIT.
#

usage() {
  this=$1
  cat <<EOF
$this: download go binaries for {{ $.Release.GitHub.Owner }}/{{ $.Release.GitHub.Name }}

Usage: $this [version]
  where [version] is 'latest' or a version number from
  https://github.com/{{ $.Release.GitHub.Owner }}/{{ $.Release.GitHub.Name }}/releases

Generated by godownloader
 https://github.com/goreleaser/godownloader

EOF
}
` + shellfn + `
is_supported() {
  platform=$1
  found=0
  case "$platform" in
	  {{- range $goos := $.Build.Goos }}{{ range $goarch := $.Build.Goarch }}
	  {{ $goos }}/{{ $goarch }}) found=1 ;;
	  {{- end }}{{ end }}
  esac
  {{- if $.Build.Ignore }}
  case "$platform" in
    {{- range $ignore := $.Build.Ignore }}
    {{ $ignore.Goos }}/{{ $ignore.Goarch }}) found=0 ;;
    {{- end }}
  esac
  {{- end }}
  return $found
}

OWNER={{ $.Release.GitHub.Owner }}
REPO={{ $.Release.GitHub.Name }}
BINARY={{ .Build.Binary }}
FORMAT={{ .Archive.Format }}
BINDIR=${BINDIR:-./bin}

uname_os_check
uname_arch_check

VERSION=$1
case "${VERSION}" in
 latest)
    VERSION=""
    ;;
 -h|-?|*help*)
   usage "$0"
   exit 1
   ;;
esac

# check to see if request version is supported
# based on goreleaser.yml file

PREFIX="$OWNER/$REPO"

if [ -z "${VERSION}" ]; then
  echo "$PREFIX: checking GitHub for latest version"
  VERSION=$(github_last_release "$OWNER/$REPO")
fi
# if version starts with 'v', remove it
VERSION=${VERSION#v}

OS=$(uname_os)
ARCH=$(uname_arch)

# change format (tar.gz or zip) based on ARCH
{{- with .Archive.FormatOverrides }}
case ${ARCH} in
{{- range . }}
{{ .Goos }}) FORMAT={{ .Format }} ;;
esac
{{- end }}
{{- end }}

# adjust archive name based on OS
{{- with .Archive.Replacements }}
case ${OS} in
{{- range $k, $v := . }}
{{ $k }}) OS={{ $v }} ;;
{{- end }}
esac

# adjust archive name based on ARCH
case ${ARCH} in
{{- range $k, $v := . }}
{{ $k }}) ARCH={{ $v }} ;;
{{- end }}
esac
{{- end }}

echo "$PREFIX: found version ${VERSION} for ${OS}/${ARCH}"

{{ .Archive.NameTemplate }}
TARBALL=${NAME}.${FORMAT}
TARBALL_URL=https://github.com/${OWNER}/${REPO}/releases/download/v${VERSION}/${TARBALL}
CHECKSUM=${REPO}_checksums.txt
CHECKSUM_URL=https://github.com/${OWNER}/${REPO}/releases/download/v${VERSION}/${CHECKSUM}

# this function wraps all the destructive operations
# if a curl|bash cuts off the end of the script due to
# network, either nothing will happen or will syntax error
# out preventing half-done work
execute() {
  TMPDIR=$(mktmpdir)
  echo "$PREFIX: downloading ${TARBALL_URL}"
  http_download "${TMPDIR}/${TARBALL}" "${TARBALL_URL}"

  echo "$PREFIX: verifying checksums"
  http_download "${TMPDIR}/${CHECKSUM}" "${CHECKSUM_URL}"
  hash_sha256_verify "${TMPDIR}/${TARBALL}" "${TMPDIR}/${CHECKSUM}"

  (cd "${TMPDIR}" && untar "${TARBALL}")
  install -d "${BINDIR}"
  install "${TMPDIR}/${BINARY}" "${BINDIR}/"
  echo "$PREFIX: installed as ${BINDIR}/${BINARY}"
}

execute
`
<commit_msg>Issue #20 - warn if platform is not supported<commit_after>package main

var shellGodownloader = `#!/bin/sh
set -e
#  Code generated by godownloader. DO NOT EDIT.
#

usage() {
  this=$1
  cat <<EOF
$this: download go binaries for {{ $.Release.GitHub.Owner }}/{{ $.Release.GitHub.Name }}

Usage: $this [version]
  where [version] is 'latest' or a version number from
  https://github.com/{{ $.Release.GitHub.Owner }}/{{ $.Release.GitHub.Name }}/releases

Generated by godownloader
 https://github.com/goreleaser/godownloader

EOF
}
` + shellfn + `
is_supported_platform() {
  platform=$1
  found=1
  case "$platform" in
    {{- range $goos := $.Build.Goos }}{{ range $goarch := $.Build.Goarch }}
    {{ $goos }}/{{ $goarch }}) found=0 ;;
    {{- end }}{{ end }}
  esac
  {{- if $.Build.Ignore }}
  case "$platform" in
    {{- range $ignore := $.Build.Ignore }}
    {{ $ignore.Goos }}/{{ $ignore.Goarch }}) found=1 ;;
    {{- end }}
  esac
  {{- end }}
  return $found
}

OWNER={{ $.Release.GitHub.Owner }}
REPO={{ $.Release.GitHub.Name }}
BINARY={{ .Build.Binary }}
FORMAT={{ .Archive.Format }}
BINDIR=${BINDIR:-./bin}

VERSION=$1
case "${VERSION}" in
 latest)
    VERSION=""
    ;;
 -h|-?|*help*)
   usage "$0"
   exit 1
   ;;
esac

uname_os_check
uname_arch_check

OS=$(uname_os)
ARCH=$(uname_arch)
PREFIX="$OWNER/$REPO"
PLATFORM="${OS}/${ARCH}"
if  is_supported_platform "$PLATFORM"; then
  # optional logging goes here
  true
else
  echo "${PREFIX}: platform $PLATFORM is not supported.  Make sure this script is up-to-date and file request at https://github.com/${PREFIX}/issues/new"
  exit 1
fi

if [ -z "${VERSION}" ]; then
  echo "$PREFIX: checking GitHub for latest version"
  VERSION=$(github_last_release "$OWNER/$REPO")
fi
# if version starts with 'v', remove it
VERSION=${VERSION#v}

# change format (tar.gz or zip) based on ARCH
{{- with .Archive.FormatOverrides }}
case ${ARCH} in
{{- range . }}
{{ .Goos }}) FORMAT={{ .Format }} ;;
esac
{{- end }}
{{- end }}

# adjust archive name based on OS
{{- with .Archive.Replacements }}
case ${OS} in
{{- range $k, $v := . }}
{{ $k }}) OS={{ $v }} ;;
{{- end }}
esac

# adjust archive name based on ARCH
case ${ARCH} in
{{- range $k, $v := . }}
{{ $k }}) ARCH={{ $v }} ;;
{{- end }}
esac
{{- end }}

echo "$PREFIX: found version ${VERSION} for ${OS}/${ARCH}"

{{ .Archive.NameTemplate }}
TARBALL=${NAME}.${FORMAT}
TARBALL_URL=https://github.com/${OWNER}/${REPO}/releases/download/v${VERSION}/${TARBALL}
CHECKSUM=${REPO}_checksums.txt
CHECKSUM_URL=https://github.com/${OWNER}/${REPO}/releases/download/v${VERSION}/${CHECKSUM}

# this function wraps all the destructive operations
# if a curl|bash cuts off the end of the script due to
# network, either nothing will happen or will syntax error
# out preventing half-done work
execute() {
  TMPDIR=$(mktmpdir)
  echo "$PREFIX: downloading ${TARBALL_URL}"
  http_download "${TMPDIR}/${TARBALL}" "${TARBALL_URL}"

  echo "$PREFIX: verifying checksums"
  http_download "${TMPDIR}/${CHECKSUM}" "${CHECKSUM_URL}"
  hash_sha256_verify "${TMPDIR}/${TARBALL}" "${TMPDIR}/${CHECKSUM}"

  (cd "${TMPDIR}" && untar "${TARBALL}")
  install -d "${BINDIR}"
  install "${TMPDIR}/${BINARY}" "${BINDIR}/"
  echo "$PREFIX: installed as ${BINDIR}/${BINARY}"
}

execute
`
<|endoftext|>


<commit_before>package beater

import (
	"sync"
	"time"

	cfg "github.com/elastic/beats/filebeat/config"
	"github.com/elastic/beats/filebeat/input"
	"github.com/elastic/beats/libbeat/logp"
)

var debugf = logp.MakeDebug("spooler")

// channelSize is the number of events Channel can buffer before blocking will occur.
const channelSize = 16

// Spooler aggregates the events and sends the aggregated data to the publisher.
type Spooler struct {
	Channel chan *input.FileEvent // Channel is the input to the Spooler.

	// Config
	idleTimeout time.Duration // How often to flush the spooler if spoolSize is not reached.
	spoolSize   uint64        // Maximum number of events that are stored before a flush occurs.

	exit          chan struct{}             // Channel used to signal shutdown.
	nextFlushTime time.Time                 // Scheduled time of the next flush.
	publisher     chan<- []*input.FileEvent // Channel used to publish events.
	spool         []*input.FileEvent        // FileEvents being held by the Spooler.
	wg            sync.WaitGroup            // WaitGroup used to control the shutdown.
}

// NewSpooler creates and returns a new Spooler. The returned Spooler must be
// started by calling Start before it can be used.
func NewSpooler(
	config cfg.FilebeatConfig,
	publisher chan<- []*input.FileEvent,
) *Spooler {
	spoolSize := config.SpoolSize
	if spoolSize <= 0 {
		spoolSize = cfg.DefaultSpoolSize
		debugf("Spooler will use the default spool_size of %d", spoolSize)
	}

	idleTimeout := config.IdleTimeout
	if idleTimeout <= 0 {
		idleTimeout = cfg.DefaultIdleTimeout
		debugf("Spooler will use the default idle_timeout of %s", idleTimeout)
	}

	return &Spooler{
		Channel:       make(chan *input.FileEvent, channelSize),
		idleTimeout:   idleTimeout,
		spoolSize:     spoolSize,
		exit:          make(chan struct{}),
		nextFlushTime: time.Now().Add(idleTimeout),
		publisher:     publisher,
		spool:         make([]*input.FileEvent, 0, spoolSize),
	}
}

// Start starts the Spooler. Stop must be called to stop the Spooler.
func (s *Spooler) Start() {
	s.wg.Add(1)
	go s.run()
}

// run queues events that it reads from Channel and flushes them when either the
// queue reaches its capacity (which is spoolSize) or a timeout period elapses.
func (s *Spooler) run() {
	defer s.wg.Done()

	ticker := time.NewTicker(s.idleTimeout / 2)

	logp.Info("Starting spooler: spool_size: %v; idle_timeout: %s",
		s.spoolSize, s.idleTimeout)

loop:
	for {
		select {
		case <-s.exit:
			ticker.Stop()
			break loop
		case event := <-s.Channel:
			if event != nil {
				s.queue(event)
			}
		case <-ticker.C:
			s.timedFlush()
		}
	}

	// Drain any events that may remain in Channel.
	for e := range s.Channel {
		s.queue(e)
	}
	debugf("Flushing events from spooler at shutdown")
	s.flush()
}

// Stop stops this Spooler. This method blocks until all events have been
// flushed to the publisher. The method should only be invoked one time after
// Start has been invoked.
func (s *Spooler) Stop() {
	logp.Info("Stopping spooler")

	// Signal to the run method that it should stop.
	close(s.exit)

	// Stop accepting writes. Any events in the channel will be flushed.
	close(s.Channel)

	// Wait for the flush to complete.
	s.wg.Wait()
	debugf("Spooler has stopped")
}

// queue queues a single event to be spooled. If the queue reaches spoolSize
// while calling this method then all events in the queue will be flushed to
// the publisher.
func (s *Spooler) queue(event *input.FileEvent) {
	s.spool = append(s.spool, event)
	if len(s.spool) == cap(s.spool) {
		debugf("Flushing spooler because spooler full. Events flushed: %v", len(s.spool))
		s.flush()
	}
}

// timedFlush flushes the events in the queue if a flush has not occurred
// for a period of time greater than idleTimeout.
func (s *Spooler) timedFlush() {
	if time.Now().After(s.nextFlushTime) {
		debugf("Flushing spooler because of timeout. Events flushed: %v", len(s.spool))
		s.flush()
	}
}

// flush flushes all events to the publisher.
func (s *Spooler) flush() {
	if len(s.spool) > 0 {
		// copy buffer
		tmpCopy := make([]*input.FileEvent, len(s.spool))
		copy(tmpCopy, s.spool)

		// clear buffer
		s.spool = s.spool[:0]

		// send
		s.publisher <- tmpCopy
	}
	s.nextFlushTime = time.Now().Add(s.idleTimeout)
}
<commit_msg>Fix spooler shutdown blocking<commit_after>package beater

import (
	"sync"
	"time"

	cfg "github.com/elastic/beats/filebeat/config"
	"github.com/elastic/beats/filebeat/input"
	"github.com/elastic/beats/libbeat/logp"
)

var debugf = logp.MakeDebug("spooler")

// channelSize is the number of events Channel can buffer before blocking will occur.
const channelSize = 16

// Spooler aggregates the events and sends the aggregated data to the publisher.
type Spooler struct {
	Channel chan *input.FileEvent // Channel is the input to the Spooler.

	// Config
	idleTimeout time.Duration // How often to flush the spooler if spoolSize is not reached.
	spoolSize   uint64        // Maximum number of events that are stored before a flush occurs.

	exit          chan struct{}             // Channel used to signal shutdown.
	nextFlushTime time.Time                 // Scheduled time of the next flush.
	publisher     chan<- []*input.FileEvent // Channel used to publish events.
	spool         []*input.FileEvent        // FileEvents being held by the Spooler.
	wg            sync.WaitGroup            // WaitGroup used to control the shutdown.
}

// NewSpooler creates and returns a new Spooler. The returned Spooler must be
// started by calling Start before it can be used.
func NewSpooler(
	config cfg.FilebeatConfig,
	publisher chan<- []*input.FileEvent,
) *Spooler {
	spoolSize := config.SpoolSize
	if spoolSize <= 0 {
		spoolSize = cfg.DefaultSpoolSize
		debugf("Spooler will use the default spool_size of %d", spoolSize)
	}

	idleTimeout := config.IdleTimeout
	if idleTimeout <= 0 {
		idleTimeout = cfg.DefaultIdleTimeout
		debugf("Spooler will use the default idle_timeout of %s", idleTimeout)
	}

	return &Spooler{
		Channel:       make(chan *input.FileEvent, channelSize),
		idleTimeout:   idleTimeout,
		spoolSize:     spoolSize,
		exit:          make(chan struct{}),
		nextFlushTime: time.Now().Add(idleTimeout),
		publisher:     publisher,
		spool:         make([]*input.FileEvent, 0, spoolSize),
	}
}

// Start starts the Spooler. Stop must be called to stop the Spooler.
func (s *Spooler) Start() {
	s.wg.Add(1)
	go s.run()
}

// run queues events that it reads from Channel and flushes them when either the
// queue reaches its capacity (which is spoolSize) or a timeout period elapses.
func (s *Spooler) run() {
	defer s.wg.Done()

	ticker := time.NewTicker(s.idleTimeout / 2)

	logp.Info("Starting spooler: spool_size: %v; idle_timeout: %s",
		s.spoolSize, s.idleTimeout)

loop:
	for {
		select {
		case <-s.exit:
			ticker.Stop()
			break loop
		case event := <-s.Channel:
			if event != nil {
				s.queue(event)
			}
		case <-ticker.C:
			s.timedFlush()
		}
	}

	// Drain any events that may remain in Channel.
	for e := range s.Channel {
		s.queue(e)
	}
	debugf("Flushing events from spooler at shutdown")
	s.flush()
}

// Stop stops this Spooler. This method blocks until all events have been
// flushed to the publisher. The method should only be invoked one time after
// Start has been invoked.
func (s *Spooler) Stop() {
	logp.Info("Stopping spooler")

	// Signal to the run method that it should stop.
	close(s.exit)

	// Stop accepting writes. Any events in the channel will be flushed.
	close(s.Channel)

	// Wait for the flush to complete.
	s.wg.Wait()
	debugf("Spooler has stopped")
}

// queue queues a single event to be spooled. If the queue reaches spoolSize
// while calling this method then all events in the queue will be flushed to
// the publisher.
func (s *Spooler) queue(event *input.FileEvent) {
	s.spool = append(s.spool, event)
	if len(s.spool) == cap(s.spool) {
		debugf("Flushing spooler because spooler full. Events flushed: %v", len(s.spool))
		s.flush()
	}
}

// timedFlush flushes the events in the queue if a flush has not occurred
// for a period of time greater than idleTimeout.
func (s *Spooler) timedFlush() {
	if time.Now().After(s.nextFlushTime) {
		debugf("Flushing spooler because of timeout. Events flushed: %v", len(s.spool))
		s.flush()
	}
}

// flush flushes all events to the publisher.
func (s *Spooler) flush() {
	if len(s.spool) > 0 {
		// copy buffer
		tmpCopy := make([]*input.FileEvent, len(s.spool))
		copy(tmpCopy, s.spool)

		// clear buffer
		s.spool = s.spool[:0]

		select {
		case <-s.exit:
		case s.publisher <- tmpCopy: // send
		}
	}
	s.nextFlushTime = time.Now().Add(s.idleTimeout)
}
<|endoftext|>


<commit_before>package socket

import (
	"testing"
)

func TestPacketString(t *testing.T) {
	var p = NewPacket(nil)
	p.SetSeq(21)
	p.SetPtype(3)
	p.SetSize(300)
	p.SetBody(map[string]int{"a": 1})
	p.SetUri("uri/b")
	p.SetBodyCodec(5)
	p.Meta().Set("key", "value")
	t.Logf("%%s:%s", p.String())
	t.Logf("%%v:%v", p)
	t.Logf("%%#v:%#v", p)
	t.Logf("%%+v:%+v", p)
}
<commit_msg>Update TestPacketString<commit_after>package socket

import (
	"testing"
)

func TestPacketString(t *testing.T) {
	var p = NewPacket()
	p.SetSeq(21)
	p.XferPipe().Append('g')
	p.SetPtype(3)
	p.SetSize(300)
	p.SetBody(map[string]int{"a": 1})
	p.SetUri("uri/b")
	p.SetBodyCodec(5)
	p.Meta().Set("key", "value")
	t.Logf("%%s:%s", p.String())
	t.Logf("%%v:%v", p)
	t.Logf("%%#v:%#v", p)
	t.Logf("%%+v:%+v", p)
}
<|endoftext|>


<commit_before>package multiaddr

import "testing"

func expectVarint(t *testing.T, x, expected int) {
	size := VarintSize(x)
	if size != expected {
		t.Fatalf("expected varintsize of %d to be %d, got %d", x, expected, size)
	}
}

func TestVarintSize(t *testing.T) {
	expectVarint(t, (1<<7)-1, 1)
	expectVarint(t, 0, 1)
	expectVarint(t, 1<<7, 2)
}
<commit_msg>test: test all varints less than 2**16 againt VarintSize<commit_after>package multiaddr

import (
	"encoding/binary"
	"testing"
)

func checkVarint(t *testing.T, x int) {
	buf := make([]byte, binary.MaxVarintLen64)
	expected := binary.PutUvarint(buf, uint64(x))

	size := VarintSize(x)
	if size != expected {
		t.Fatalf("expected varintsize of %d to be %d, got %d", x, expected, size)
	}
}

func TestVarintSize(t *testing.T) {
	max := 1 << 16
	for x := 0; x < max; x++ {
		checkVarint(t, x)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Pulcy.
//
// 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 vault

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/hashicorp/vault/api"
	"github.com/juju/errgo"
	"github.com/mitchellh/go-homedir"
	"github.com/op/go-logging"
)

type VaultConfig struct {
	VaultAddr   string
	VaultCACert string
	VaultCAPath string
}

type Vault struct {
	vaultClient *api.Client
	log         *logging.Logger
}

func NewVault(srvCfg VaultConfig, log *logging.Logger) (*Vault, error) {
	// Create a vault client
	config := api.DefaultConfig()
	if err := config.ReadEnvironment(); err != nil {
		return nil, maskAny(err)
	}
	if srvCfg.VaultAddr != "" {
		config.Address = srvCfg.VaultAddr
	}
	if srvCfg.VaultCACert != "" || srvCfg.VaultCAPath != "" {
		var newCertPool *x509.CertPool
		var err error
		if srvCfg.VaultCACert != "" {
			log.Debugf("Loading CA cert: %s", srvCfg.VaultCACert)
			newCertPool, err = api.LoadCACert(srvCfg.VaultCACert)
		} else {
			log.Debugf("Loading CA certs from: %s", srvCfg.VaultCAPath)
			newCertPool, err = api.LoadCAPath(srvCfg.VaultCAPath)
		}
		if err != nil {
			return nil, maskAny(err)
		}
		clientTLSConfig := config.HttpClient.Transport.(*http.Transport).TLSClientConfig
		clientTLSConfig.RootCAs = newCertPool
	}
	client, err := api.NewClient(config)
	if err != nil {
		return nil, maskAny(err)
	}

	return &Vault{
		log:         log,
		vaultClient: client,
	}, nil

}

type GithubLoginData struct {
	GithubToken     string
	GithubTokenPath string
	Mount           string // defaults to "github"
}

// GithubLogin performs a standard Github authentication and initializes the vaultClient with the resulting token.
func (s *Vault) GithubLogin(data GithubLoginData) error {
	// Read token
	var err error
	data.GithubToken, err = s.readGithubToken(data)
	if err != nil {
		return maskAny(err)
	}
	// Perform login
	s.vaultClient.ClearToken()
	logical := s.vaultClient.Logical()
	loginData := make(map[string]interface{})
	loginData["token"] = data.GithubToken
	if data.Mount == "" {
		data.Mount = "github"
	}
	path := fmt.Sprintf("auth/%s/login", data.Mount)
	if loginSecret, err := logical.Write(path, loginData); err != nil {
		return maskAny(err)
	} else if loginSecret.Auth == nil {
		return maskAny(errgo.WithCausef(nil, VaultError, "missing authentication in secret response"))
	} else {
		// Use token
		s.vaultClient.SetToken(loginSecret.Auth.ClientToken)
	}

	// We're done
	return nil
}

func (s *Vault) readGithubToken(data GithubLoginData) (string, error) {
	if data.GithubToken != "" {
		return data.GithubToken, nil
	}
	if data.GithubTokenPath == "" {
		return "", maskAny(errgo.WithCausef(nil, InvalidArgumentError, "No github token path set"))
	}
	path, err := homedir.Expand(data.GithubTokenPath)
	if err != nil {
		return "", maskAny(err)
	}
	raw, err := ioutil.ReadFile(path)
	if err != nil {
		return "", maskAny(err)
	}
	return strings.TrimSpace(string(raw)), nil
}

// extractSecret extracts a secret based on given variables
// Call a login method before calling this method.
func (s *Vault) Extract(secretPath, secretField string) (string, error) {
	if secretPath == "" {
		return "", maskAny(errgo.WithCausef(nil, InvalidArgumentError, "path not set"))
	}
	if secretField == "" {
		return "", maskAny(errgo.WithCausef(nil, InvalidArgumentError, "field not set"))
	}

	// Load secret
	s.log.Infof("Read %s#%s", secretPath, secretField)
	secret, err := s.vaultClient.Logical().Read(secretPath)
	if err != nil {
		return "", maskAny(errgo.WithCausef(nil, VaultError, "error reading %s: %s", secretPath, err))
	}
	if secret == nil {
		return "", maskAny(errgo.WithCausef(nil, VaultError, "no value found at %s", secretPath))
	}

	if value, ok := secret.Data[secretField]; !ok {
		return "", maskAny(errgo.WithCausef(nil, VaultError, "no field '%s' found at %s", secretField, secretPath))
	} else {
		return value.(string), nil
	}
}
<commit_msg>Fixed SSL certificate error regarding IP addresses missing from certificate<commit_after>// Copyright (c) 2016 Pulcy.
//
// 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 vault

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"

	"github.com/hashicorp/vault/api"
	"github.com/juju/errgo"
	"github.com/mitchellh/go-homedir"
	"github.com/op/go-logging"
)

type VaultConfig struct {
	VaultAddr   string
	VaultCACert string
	VaultCAPath string
}

type Vault struct {
	vaultClient *api.Client
	log         *logging.Logger
}

func NewVault(srvCfg VaultConfig, log *logging.Logger) (*Vault, error) {
	// Create a vault client
	config := api.DefaultConfig()
	if err := config.ReadEnvironment(); err != nil {
		return nil, maskAny(err)
	}
	var serverName string
	if srvCfg.VaultAddr != "" {
		log.Debugf("Setting vault address to %s", srvCfg.VaultAddr)
		config.Address = srvCfg.VaultAddr
		url, err := url.Parse(config.Address)
		if err != nil {
			return nil, maskAny(err)
		}
		host, _, err := net.SplitHostPort(url.Host)
		if err != nil {
			return nil, maskAny(err)
		}
		serverName = host
	}
	if srvCfg.VaultCACert != "" || srvCfg.VaultCAPath != "" {
		var newCertPool *x509.CertPool
		var err error
		if srvCfg.VaultCACert != "" {
			log.Debugf("Loading CA cert: %s", srvCfg.VaultCACert)
			newCertPool, err = api.LoadCACert(srvCfg.VaultCACert)
		} else {
			log.Debugf("Loading CA certs from: %s", srvCfg.VaultCAPath)
			newCertPool, err = api.LoadCAPath(srvCfg.VaultCAPath)
		}
		if err != nil {
			return nil, maskAny(err)
		}
		clientTLSConfig := config.HttpClient.Transport.(*http.Transport).TLSClientConfig
		clientTLSConfig.RootCAs = newCertPool
		clientTLSConfig.ServerName = serverName
	}
	client, err := api.NewClient(config)
	if err != nil {
		return nil, maskAny(err)
	}

	return &Vault{
		log:         log,
		vaultClient: client,
	}, nil

}

type GithubLoginData struct {
	GithubToken     string
	GithubTokenPath string
	Mount           string // defaults to "github"
}

// GithubLogin performs a standard Github authentication and initializes the vaultClient with the resulting token.
func (s *Vault) GithubLogin(data GithubLoginData) error {
	// Read token
	var err error
	data.GithubToken, err = s.readGithubToken(data)
	if err != nil {
		return maskAny(err)
	}
	// Perform login
	s.vaultClient.ClearToken()
	logical := s.vaultClient.Logical()
	loginData := make(map[string]interface{})
	loginData["token"] = data.GithubToken
	if data.Mount == "" {
		data.Mount = "github"
	}
	path := fmt.Sprintf("auth/%s/login", data.Mount)
	if loginSecret, err := logical.Write(path, loginData); err != nil {
		return maskAny(err)
	} else if loginSecret.Auth == nil {
		return maskAny(errgo.WithCausef(nil, VaultError, "missing authentication in secret response"))
	} else {
		// Use token
		s.vaultClient.SetToken(loginSecret.Auth.ClientToken)
	}

	// We're done
	return nil
}

func (s *Vault) readGithubToken(data GithubLoginData) (string, error) {
	if data.GithubToken != "" {
		return data.GithubToken, nil
	}
	if data.GithubTokenPath == "" {
		return "", maskAny(errgo.WithCausef(nil, InvalidArgumentError, "No github token path set"))
	}
	path, err := homedir.Expand(data.GithubTokenPath)
	if err != nil {
		return "", maskAny(err)
	}
	raw, err := ioutil.ReadFile(path)
	if err != nil {
		return "", maskAny(err)
	}
	return strings.TrimSpace(string(raw)), nil
}

// extractSecret extracts a secret based on given variables
// Call a login method before calling this method.
func (s *Vault) Extract(secretPath, secretField string) (string, error) {
	if secretPath == "" {
		return "", maskAny(errgo.WithCausef(nil, InvalidArgumentError, "path not set"))
	}
	if secretField == "" {
		return "", maskAny(errgo.WithCausef(nil, InvalidArgumentError, "field not set"))
	}

	// Load secret
	s.log.Infof("Read %s#%s", secretPath, secretField)
	secret, err := s.vaultClient.Logical().Read(secretPath)
	if err != nil {
		return "", maskAny(errgo.WithCausef(nil, VaultError, "error reading %s: %s", secretPath, err))
	}
	if secret == nil {
		return "", maskAny(errgo.WithCausef(nil, VaultError, "no value found at %s", secretPath))
	}

	if value, ok := secret.Data[secretField]; !ok {
		return "", maskAny(errgo.WithCausef(nil, VaultError, "no field '%s' found at %s", secretField, secretPath))
	} else {
		return value.(string), nil
	}
}
<|endoftext|>


<commit_before>package vault

import (
	"io/ioutil"
	"log"
	"os"

	vaultapi "github.com/hashicorp/vault/api"
	"k8s.io/client-go/rest"
)

// Client is a Vault client with Kubernetes support and token automatic renewing
type Client struct {
	client       *vaultapi.Client
	logical      *vaultapi.Logical
	tokenRenewer *vaultapi.Renewer
}

// NewClient creates a new Vault client
func NewClient(role string) (*Client, error) {
	return NewClientWithConfig(vaultapi.DefaultConfig(), role)
}

// NewClientWithConfig creates a new Vault client with custom configuration
func NewClientWithConfig(config *vaultapi.Config, role string) (*Client, error) {
	client, err := vaultapi.NewClient(config)
	if err != nil {
		return nil, err
	}
	logical := client.Logical()
	var tokenRenewer *vaultapi.Renewer

	if client.Token() == "" {

		token, err := ioutil.ReadFile(os.Getenv("HOME") + "/.vault-token")

		if err == nil {

			client.SetToken(string(token))

		} else {
			// If VAULT_TOKEN or ~/.vault-token wasn't provided let's suppose
			// we are in Kubernetes and try to get one with the ServiceAccount token

			k8sconfig, err := rest.InClusterConfig()
			if err != nil {
				return nil, err
			}

			initialTokenArrived := make(chan string, 1)

			go func() {
				for {
					data := map[string]interface{}{"jwt": k8sconfig.BearerToken, "role": role}
					secret, err := logical.Write("auth/kubernetes/login", data)
					if err != nil {
						log.Println(err.Error())
						continue
					}

					log.Println("Received new Vault token")

					// Set the first token from the response
					client.SetToken(secret.Auth.ClientToken)

					initialTokenArrived <- secret.LeaseID

					// Start the renewing process
					tokenRenewer, err = client.NewRenewer(&vaultapi.RenewerInput{Secret: secret})
					if err != nil {
						log.Println(err.Error())
						continue
					}

					go tokenRenewer.Renew()

					runRenewChecker(tokenRenewer)
				}
			}()

			<-initialTokenArrived
		}
	}

	return &Client{client: client, logical: logical, tokenRenewer: tokenRenewer}, nil
}

func runRenewChecker(tokenRenewer *vaultapi.Renewer) {
	for {
		select {
		case err := <-tokenRenewer.DoneCh():
			if err != nil {
				log.Println("Renew error:", err.Error())
				return
			}
		case renewal := <-tokenRenewer.RenewCh():
			log.Printf("Successfully renewed at: %s", renewal.RenewedAt)
		}
	}
}

// Vault returns the underlying hashicorp Vault client
func (client *Client) Vault() *vaultapi.Client {
	return client.client
}

// Close stops the token renewing process of this client
func (client *Client) Close() {
	if client.tokenRenewer != nil {
		log.Println("Stopped Vault tokenRenewer")
		client.tokenRenewer.Stop()
	}
}
<commit_msg>Replace old tokenRenewer in case of a new one is created<commit_after>package vault

import (
	"io/ioutil"
	"log"
	"os"

	vaultapi "github.com/hashicorp/vault/api"
	"k8s.io/client-go/rest"
)

// Client is a Vault client with Kubernetes support and token automatic renewing
type Client struct {
	client       *vaultapi.Client
	logical      *vaultapi.Logical
	tokenRenewer *vaultapi.Renewer
}

// NewClient creates a new Vault client
func NewClient(role string) (*Client, error) {
	return NewClientWithConfig(vaultapi.DefaultConfig(), role)
}

// NewClientWithConfig creates a new Vault client with custom configuration
func NewClientWithConfig(config *vaultapi.Config, role string) (*Client, error) {
	rawClient, err := vaultapi.NewClient(config)
	if err != nil {
		return nil, err
	}
	logical := rawClient.Logical()
	var tokenRenewer *vaultapi.Renewer

	client := &Client{client: rawClient, logical: logical}

	if rawClient.Token() == "" {

		token, err := ioutil.ReadFile(os.Getenv("HOME") + "/.vault-token")

		if err == nil {

			rawClient.SetToken(string(token))

		} else {
			// If VAULT_TOKEN or ~/.vault-token wasn't provided let's suppose
			// we are in Kubernetes and try to get one with the ServiceAccount token

			k8sconfig, err := rest.InClusterConfig()
			if err != nil {
				return nil, err
			}

			initialTokenArrived := make(chan string, 1)

			go func() {
				for {
					data := map[string]interface{}{"jwt": k8sconfig.BearerToken, "role": role}
					secret, err := logical.Write("auth/kubernetes/login", data)
					if err != nil {
						log.Println(err.Error())
						continue
					}

					log.Println("Received new Vault token")

					// Set the first token from the response
					rawClient.SetToken(secret.Auth.ClientToken)

					initialTokenArrived <- secret.LeaseID

					// Start the renewing process
					tokenRenewer, err = rawClient.NewRenewer(&vaultapi.RenewerInput{Secret: secret})
					if err != nil {
						log.Println(err.Error())
						continue
					}
					client.tokenRenewer = tokenRenewer

					go tokenRenewer.Renew()

					runRenewChecker(tokenRenewer)
				}
			}()

			<-initialTokenArrived
		}
	}

	return client, nil
}

func runRenewChecker(tokenRenewer *vaultapi.Renewer) {
	for {
		select {
		case err := <-tokenRenewer.DoneCh():
			if err != nil {
				log.Println("Renew error:", err.Error())
				return
			}
		case renewal := <-tokenRenewer.RenewCh():
			log.Printf("Successfully renewed at: %s", renewal.RenewedAt)
		}
	}
}

// Vault returns the underlying hashicorp Vault client
func (client *Client) Vault() *vaultapi.Client {
	return client.client
}

// Close stops the token renewing process of this client
func (client *Client) Close() {
	if client.tokenRenewer != nil {
		log.Println("Stopped Vault tokenRenewer")
		client.tokenRenewer.Stop()
	}
}
<|endoftext|>


<commit_before>package schema

import (
	"encoding/json"

	"github.com/zclconf/go-cty/cty"
	ctyjson "github.com/zclconf/go-cty/cty/json"

	"github.com/hashicorp/terraform/configs/configschema"
	"github.com/hashicorp/terraform/terraform"
)

// DiffFromValues takes the current state and desired state as cty.Values and
// derives a terraform.InstanceDiff to give to the legacy providers. This is
// used to take the states provided by the new ApplyResourceChange method and
// convert them to a state+diff required for the legacy Apply method.
func DiffFromValues(prior, planned cty.Value, res *Resource) (*terraform.InstanceDiff, error) {
	return diffFromValues(prior, planned, res, nil)
}

// diffFromValues takes an additional CustomizeDiffFunc, so we can generate our
// test fixtures from the legacy tests. In the new provider protocol the diff
// only needs to be created for the apply operation, and any customizations
// have already been done.
func diffFromValues(prior, planned cty.Value, res *Resource, cust CustomizeDiffFunc) (*terraform.InstanceDiff, error) {
	instanceState, err := res.ShimInstanceStateFromValue(prior)
	if err != nil {
		return nil, err
	}

	configSchema := res.CoreConfigSchema()

	cfg := terraform.NewResourceConfigShimmed(planned, configSchema)

	diff, err := schemaMap(res.Schema).Diff(instanceState, cfg, cust, nil, false)
	if err != nil {
		return nil, err
	}

	return diff, err
}

// ApplyDiff takes a cty.Value state and applies a terraform.InstanceDiff to
// get a new cty.Value state. This is used to convert the diff returned from
// the legacy provider Diff method to the state required for the new
// PlanResourceChange method.
func ApplyDiff(base cty.Value, d *terraform.InstanceDiff, schema *configschema.Block) (cty.Value, error) {
	return d.ApplyToValue(base, schema)
}

// StateValueToJSONMap converts a cty.Value to generic JSON map via the cty JSON
// encoding.
func StateValueToJSONMap(val cty.Value, ty cty.Type) (map[string]interface{}, error) {
	js, err := ctyjson.Marshal(val, ty)
	if err != nil {
		return nil, err
	}

	var m map[string]interface{}
	if err := json.Unmarshal(js, &m); err != nil {
		return nil, err
	}

	return m, nil
}

// JSONMapToStateValue takes a generic json map[string]interface{} and converts it
// to the specific type, ensuring that the values conform to the schema.
func JSONMapToStateValue(m map[string]interface{}, block *configschema.Block) (cty.Value, error) {
	var val cty.Value

	js, err := json.Marshal(m)
	if err != nil {
		return val, err
	}

	val, err = ctyjson.Unmarshal(js, block.ImpliedType())
	if err != nil {
		return val, err
	}

	return block.CoerceValue(val)
}

// StateValueFromInstanceState converts a terraform.InstanceState to a
// cty.Value as described by the provided cty.Type, and maintains the resource
// ID as the "id" attribute.
func StateValueFromInstanceState(is *terraform.InstanceState, ty cty.Type) (cty.Value, error) {
	return is.AttrsAsObjectValue(ty)
}

// LegacyResourceSchema takes a *Resource and returns a deep copy with 0.12 specific
// features removed. This is used by the shims to get a configschema that
// directly matches the structure of the schema.Resource.
func LegacyResourceSchema(r *Resource) *Resource {
	if r == nil {
		return nil
	}
	// start with a shallow copy
	newResource := new(Resource)
	*newResource = *r
	newResource.Schema = map[string]*Schema{}

	for k, s := range r.Schema {
		newResource.Schema[k] = LegacySchema(s)
	}

	return newResource
}

// LegacySchema takes a *Schema and returns a deep copy with 0.12 specific
// features removed. This is used by the shims to get a configschema that
// directly matches the structure of the schema.Resource.
func LegacySchema(s *Schema) *Schema {
	if s == nil {
		return nil
	}
	// start with a shallow copy
	newSchema := new(Schema)
	*newSchema = *s
	newSchema.ConfigMode = SchemaConfigModeAuto
	newSchema.PromoteSingle = false
	newSchema.SkipCoreTypeCheck = false

	switch e := newSchema.Elem.(type) {
	case *Schema:
		newSchema.Elem = LegacySchema(e)
	case *Resource:
		newSchema.Elem = LegacyResourceSchema(e)
	}

	return newSchema
}
<commit_msg>PromoteSingle is used in 0.11 mode<commit_after>package schema

import (
	"encoding/json"

	"github.com/zclconf/go-cty/cty"
	ctyjson "github.com/zclconf/go-cty/cty/json"

	"github.com/hashicorp/terraform/configs/configschema"
	"github.com/hashicorp/terraform/terraform"
)

// DiffFromValues takes the current state and desired state as cty.Values and
// derives a terraform.InstanceDiff to give to the legacy providers. This is
// used to take the states provided by the new ApplyResourceChange method and
// convert them to a state+diff required for the legacy Apply method.
func DiffFromValues(prior, planned cty.Value, res *Resource) (*terraform.InstanceDiff, error) {
	return diffFromValues(prior, planned, res, nil)
}

// diffFromValues takes an additional CustomizeDiffFunc, so we can generate our
// test fixtures from the legacy tests. In the new provider protocol the diff
// only needs to be created for the apply operation, and any customizations
// have already been done.
func diffFromValues(prior, planned cty.Value, res *Resource, cust CustomizeDiffFunc) (*terraform.InstanceDiff, error) {
	instanceState, err := res.ShimInstanceStateFromValue(prior)
	if err != nil {
		return nil, err
	}

	configSchema := res.CoreConfigSchema()

	cfg := terraform.NewResourceConfigShimmed(planned, configSchema)

	diff, err := schemaMap(res.Schema).Diff(instanceState, cfg, cust, nil, false)
	if err != nil {
		return nil, err
	}

	return diff, err
}

// ApplyDiff takes a cty.Value state and applies a terraform.InstanceDiff to
// get a new cty.Value state. This is used to convert the diff returned from
// the legacy provider Diff method to the state required for the new
// PlanResourceChange method.
func ApplyDiff(base cty.Value, d *terraform.InstanceDiff, schema *configschema.Block) (cty.Value, error) {
	return d.ApplyToValue(base, schema)
}

// StateValueToJSONMap converts a cty.Value to generic JSON map via the cty JSON
// encoding.
func StateValueToJSONMap(val cty.Value, ty cty.Type) (map[string]interface{}, error) {
	js, err := ctyjson.Marshal(val, ty)
	if err != nil {
		return nil, err
	}

	var m map[string]interface{}
	if err := json.Unmarshal(js, &m); err != nil {
		return nil, err
	}

	return m, nil
}

// JSONMapToStateValue takes a generic json map[string]interface{} and converts it
// to the specific type, ensuring that the values conform to the schema.
func JSONMapToStateValue(m map[string]interface{}, block *configschema.Block) (cty.Value, error) {
	var val cty.Value

	js, err := json.Marshal(m)
	if err != nil {
		return val, err
	}

	val, err = ctyjson.Unmarshal(js, block.ImpliedType())
	if err != nil {
		return val, err
	}

	return block.CoerceValue(val)
}

// StateValueFromInstanceState converts a terraform.InstanceState to a
// cty.Value as described by the provided cty.Type, and maintains the resource
// ID as the "id" attribute.
func StateValueFromInstanceState(is *terraform.InstanceState, ty cty.Type) (cty.Value, error) {
	return is.AttrsAsObjectValue(ty)
}

// LegacyResourceSchema takes a *Resource and returns a deep copy with 0.12 specific
// features removed. This is used by the shims to get a configschema that
// directly matches the structure of the schema.Resource.
func LegacyResourceSchema(r *Resource) *Resource {
	if r == nil {
		return nil
	}
	// start with a shallow copy
	newResource := new(Resource)
	*newResource = *r
	newResource.Schema = map[string]*Schema{}

	for k, s := range r.Schema {
		newResource.Schema[k] = LegacySchema(s)
	}

	return newResource
}

// LegacySchema takes a *Schema and returns a deep copy with 0.12 specific
// features removed. This is used by the shims to get a configschema that
// directly matches the structure of the schema.Resource.
func LegacySchema(s *Schema) *Schema {
	if s == nil {
		return nil
	}
	// start with a shallow copy
	newSchema := new(Schema)
	*newSchema = *s
	newSchema.ConfigMode = SchemaConfigModeAuto
	newSchema.SkipCoreTypeCheck = false

	switch e := newSchema.Elem.(type) {
	case *Schema:
		newSchema.Elem = LegacySchema(e)
	case *Resource:
		newSchema.Elem = LegacyResourceSchema(e)
	}

	return newSchema
}
<|endoftext|>


<commit_before>// Package configutil provides an interface for loading and validating configuration
// data from YAML files.
//
// Other YAML files could be included via the following directive:
//
// production.yaml:
// extends: base.yaml
//
// There is no multiple inheritance supported. Dependency tree suppossed to
// form a directed acyclic graph(DAG).
//
//
// Values from multiple configurations within the same hierarchy are deep merged
//
// Package support multiple configuraton directories potentially having multiple files
// with the the same name. In this the we just follow the path in extends and load all
// the file accroding to the relative directory, i.e configA: base.yaml production.yaml(extends base.yaml), configB: base.yaml the load sequance
// will be the following: configA(base.yaml), configA(production.yaml)
//
// Note regarding configuration merging:
//   Array defined in YAML will be overriden based on load sequence.
//   e.g. in the base.yaml:
//        sports:
//           - football
//        in the development.yaml:
//        extends: base.yaml
//        sports:
//           - basketball
//        after the merge:
//        sports:
//           - basketball  // only keep the latest one
//
//   Map defined in YAML will be merged together based on load sequence.
//   e.g. in the base.yaml:
//        sports:
//           football: true
//        in the development.yaml:
//        extends: base.yaml
//        sports:
//           basketball: true
//        after the merge:
//        sports:  // combine all the map fields
//           football: true
//           basketball: true
//
package configutil

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"

	"code.uber.internal/infra/kraken/utils/stringset"

	"gopkg.in/validator.v2"
	"gopkg.in/yaml.v2"
)

const (
	configDirKey = "UBER_CONFIG_DIR"
)

const (
	configDir       = "config"
	secretsFile     = "secrets.yaml"
	configSeparator = ":"
)

// ErrNoFilesToLoad is returned when you attemp to call LoadFiles with no file paths
var ErrNoFilesToLoad = errors.New("attempt to load configuration with no files")

// ErrCycleRef is returned when there are circular dependencies detected in
// configuraiton files extending each other
var ErrCycleRef = errors.New("cyclic reference in configuration extends detected")

// Extends define a keywoword in config for extending a base configuration file
type Extends struct {
	Extends string `yaml:"extends"`
}

// ValidationError is the returned when a configuration fails to pass validation
type ValidationError struct {
	errorMap validator.ErrorMap
}

// ErrForField returns the validation error for the given field
func (e ValidationError) ErrForField(name string) error {
	return e.errorMap[name]
}

func (e ValidationError) Error() string {
	var w bytes.Buffer

	fmt.Fprintf(&w, "validation failed")
	for f, err := range e.errorMap {
		fmt.Fprintf(&w, "   %s: %v\n", f, err)
	}

	return w.String()
}

// FilterCandidates filters candidate config files into only the ones that exist
func FilterCandidates(fname string) ([]string, error) {
	realConfigDirs := []string{configDir}
	// Allow overriding the directory config is loaded from, useful for tests
	// inside subdirectories when the config/ dir is in the top-level of a project.
	if configRoot := os.Getenv(configDirKey); configRoot != "" {
		realConfigDirs = strings.Split(configRoot, configSeparator)
	}
	return filterCandidatesFromDirs(fname, realConfigDirs)
}

func readExtend(configFile string) (string, error) {
	var cfg Extends

	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return "", err
	}

	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return "", err
	}
	return cfg.Extends, nil
}

func getCandidate(fname string, dirs []string) string {
	candidate := ""
	for _, realConfigDir := range dirs {
		configFile := path.Join(realConfigDir, fname)
		if _, err := os.Stat(configFile); err == nil {
			candidate = configFile
		}
	}
	return candidate
}

func filterCandidatesFromDirs(fname string, dirs []string) ([]string, error) {

	var paths []string
	cSet := make(stringset.Set)

	// Go through all the 'extends' hierarchy until
	// there is no base anymore or some reference cycles
	// been detected.
	cSet.Add(fname)

	candidate := getCandidate(fname, dirs)
	fmt.Fprintf(os.Stderr, "candidate: %s\n", candidate)
	if candidate == "" {
		return nil, fmt.Errorf("file %s not found in %s", fname, dirs)
	}

	for {
		extends, err := readExtend(candidate)
		if err != nil {
			return nil, err
		}

		paths = append([]string{candidate}, paths...)
		if extends != "" {
			// prevent cycle references
			if !cSet.Has(extends) {
				candidate = path.Join(filepath.Dir(candidate), extends)
				cSet.Add(extends)
			} else {
				return nil, ErrCycleRef
			}
		} else {
			break
		}
	}

	// append secrets
	candidate = getCandidate(secretsFile, dirs)
	if candidate != "" {
		paths = append(paths, candidate)
	}

	return paths, nil
}

// Load loads configuration based on environment variables
func Load(fname string, config interface{}) error {
	candidates, err := FilterCandidates(fname)

	if err != nil {
		return err
	}

	return LoadFiles(config, candidates...)
}

// LoadFile loads configuration based on config directory
// where the input file is located
func LoadFile(fname string, config interface{}) error {
	candidates, err := filterCandidatesFromDirs(
		filepath.Base(fname), []string{filepath.Dir(fname)})

	if err != nil {
		return err
	}

	return LoadFiles(config, candidates...)
}

// LoadFiles loads a list of files, deep-merging values.
// This function is exposed for using from tests.  For production it's recommended
// to use the default resolution and the Load() method.
func LoadFiles(config interface{}, fnames ...string) error {
	if len(fnames) == 0 {
		return ErrNoFilesToLoad
	}
	for _, fname := range fnames {
		data, err := ioutil.ReadFile(fname)
		if err != nil {
			return err
		}

		if err := yaml.Unmarshal(data, config); err != nil {
			return err
		}
	}

	// Validate on the merged config at the end.
	if err := validator.Validate(config); err != nil {
		return ValidationError{
			errorMap: err.(validator.ErrorMap),
		}
	}
	return nil
}
<commit_msg>Fix proxy config<commit_after>// Package configutil provides an interface for loading and validating configuration
// data from YAML files.
//
// Other YAML files could be included via the following directive:
//
// production.yaml:
// extends: base.yaml
//
// There is no multiple inheritance supported. Dependency tree suppossed to
// form a directed acyclic graph(DAG).
//
//
// Values from multiple configurations within the same hierarchy are deep merged
//
// Package support multiple configuraton directories potentially having multiple files
// with the the same name. In this the we just follow the path in extends and load all
// the file accroding to the relative directory, i.e configA: base.yaml production.yaml(extends base.yaml), configB: base.yaml the load sequance
// will be the following: configA(base.yaml), configA(production.yaml)
//
// Note regarding configuration merging:
//   Array defined in YAML will be overriden based on load sequence.
//   e.g. in the base.yaml:
//        sports:
//           - football
//        in the development.yaml:
//        extends: base.yaml
//        sports:
//           - basketball
//        after the merge:
//        sports:
//           - basketball  // only keep the latest one
//
//   Map defined in YAML will be merged together based on load sequence.
//   e.g. in the base.yaml:
//        sports:
//           football: true
//        in the development.yaml:
//        extends: base.yaml
//        sports:
//           basketball: true
//        after the merge:
//        sports:  // combine all the map fields
//           football: true
//           basketball: true
//
package configutil

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"

	"code.uber.internal/infra/kraken/utils/stringset"

	"gopkg.in/validator.v2"
	"gopkg.in/yaml.v2"
)

const (
	configDirKey = "UBER_CONFIG_DIR"
)

const (
	configDir       = "config"
	secretsFile     = "secrets.yaml"
	configSeparator = ":"
)

// ErrNoFilesToLoad is returned when you attemp to call LoadFiles with no file paths
var ErrNoFilesToLoad = errors.New("attempt to load configuration with no files")

// ErrCycleRef is returned when there are circular dependencies detected in
// configuraiton files extending each other
var ErrCycleRef = errors.New("cyclic reference in configuration extends detected")

// Extends define a keywoword in config for extending a base configuration file
type Extends struct {
	Extends string `yaml:"extends"`
}

// ValidationError is the returned when a configuration fails to pass validation
type ValidationError struct {
	errorMap validator.ErrorMap
}

// ErrForField returns the validation error for the given field
func (e ValidationError) ErrForField(name string) error {
	return e.errorMap[name]
}

func (e ValidationError) Error() string {
	var w bytes.Buffer

	fmt.Fprintf(&w, "validation failed")
	for f, err := range e.errorMap {
		fmt.Fprintf(&w, "   %s: %v\n", f, err)
	}

	return w.String()
}

// FilterCandidates filters candidate config files into only the ones that exist
func FilterCandidates(fname string) ([]string, error) {
	realConfigDirs := []string{configDir}
	// Allow overriding the directory config is loaded from, useful for tests
	// inside subdirectories when the config/ dir is in the top-level of a project.
	if configRoot := os.Getenv(configDirKey); configRoot != "" {
		realConfigDirs = strings.Split(configRoot, configSeparator)
	}
	return filterCandidatesFromDirs(fname, realConfigDirs)
}

func readExtend(configFile string) (string, error) {
	var cfg Extends

	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return "", err
	}

	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return "", fmt.Errorf("unmarshal %s: %s", configFile, err)
	}
	return cfg.Extends, nil
}

func getCandidate(fname string, dirs []string) string {
	candidate := ""
	for _, realConfigDir := range dirs {
		configFile := path.Join(realConfigDir, fname)
		if _, err := os.Stat(configFile); err == nil {
			candidate = configFile
		}
	}
	return candidate
}

func filterCandidatesFromDirs(fname string, dirs []string) ([]string, error) {

	var paths []string
	cSet := make(stringset.Set)

	// Go through all the 'extends' hierarchy until
	// there is no base anymore or some reference cycles
	// been detected.
	cSet.Add(fname)

	candidate := getCandidate(fname, dirs)
	fmt.Fprintf(os.Stderr, "candidate: %s\n", candidate)
	if candidate == "" {
		return nil, fmt.Errorf("file %s not found in %s", fname, dirs)
	}

	for {
		extends, err := readExtend(candidate)
		if err != nil {
			return nil, err
		}

		paths = append([]string{candidate}, paths...)
		if extends != "" {
			// prevent cycle references
			if !cSet.Has(extends) {
				candidate = path.Join(filepath.Dir(candidate), extends)
				cSet.Add(extends)
			} else {
				return nil, ErrCycleRef
			}
		} else {
			break
		}
	}

	// append secrets
	candidate = getCandidate(secretsFile, dirs)
	if candidate != "" {
		paths = append(paths, candidate)
	}

	return paths, nil
}

// Load loads configuration based on environment variables
func Load(fname string, config interface{}) error {
	candidates, err := FilterCandidates(fname)

	if err != nil {
		return err
	}

	return LoadFiles(config, candidates...)
}

// LoadFile loads configuration based on config directory
// where the input file is located
func LoadFile(fname string, config interface{}) error {
	candidates, err := filterCandidatesFromDirs(
		filepath.Base(fname), []string{filepath.Dir(fname)})

	if err != nil {
		return err
	}

	return LoadFiles(config, candidates...)
}

// LoadFiles loads a list of files, deep-merging values.
// This function is exposed for using from tests.  For production it's recommended
// to use the default resolution and the Load() method.
func LoadFiles(config interface{}, fnames ...string) error {
	if len(fnames) == 0 {
		return ErrNoFilesToLoad
	}
	for _, fname := range fnames {
		data, err := ioutil.ReadFile(fname)
		if err != nil {
			return err
		}

		if err := yaml.Unmarshal(data, config); err != nil {
			return fmt.Errorf("unmarshal %s: %s", fname, err)
		}
	}

	// Validate on the merged config at the end.
	if err := validator.Validate(config); err != nil {
		return ValidationError{
			errorMap: err.(validator.ErrorMap),
		}
	}
	return nil
}
<|endoftext|>


<commit_before>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    ipfix.go
//: details: ipfix decoders handler
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: 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 main

import (
	"bytes"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/VerizonDigital/vflow/ipfix"
	"github.com/VerizonDigital/vflow/producer"
)

// IPFIX represents IPFIX collector
type IPFIX struct {
	port    int
	addr    string
	workers int
	stop    bool
	stats   IPFIXStats
	pool    chan chan struct{}
}

// IPFIXUDPMsg represents IPFIX UDP data
type IPFIXUDPMsg struct {
	raddr *net.UDPAddr
	body  []byte
}

// IPFIXStats represents IPFIX stats
type IPFIXStats struct {
	UDPQueue       int
	UDPMirrorQueue int
	MessageQueue   int
	UDPCount       uint64
	DecodedCount   uint64
	MQErrorCount   uint64
	Workers        int32
}

var (
	ipfixUDPCh         = make(chan IPFIXUDPMsg, 1000)
	ipfixMCh           = make(chan IPFIXUDPMsg, 1000)
	ipfixMQCh          = make(chan []byte, 1000)
	ipfixMirrorEnabled bool

	// templates memory cache
	mCache ipfix.MemCache

	// ipfix udp payload pool
	ipfixBuffer = &sync.Pool{
		New: func() interface{} {
			return make([]byte, opts.IPFIXUDPSize)
		},
	}
)

// NewIPFIX constructs IPFIX
func NewIPFIX() *IPFIX {
	return &IPFIX{
		port:    opts.IPFIXPort,
		workers: opts.IPFIXWorkers,
		pool:    make(chan chan struct{}, maxWorkers),
	}
}

func (i *IPFIX) run() {
	// exit if the ipfix is disabled
	if !opts.IPFIXEnabled {
		logger.Println("ipfix has been disabled")
		return
	}

	hostPort := net.JoinHostPort(i.addr, strconv.Itoa(i.port))
	udpAddr, _ := net.ResolveUDPAddr("udp", hostPort)

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		logger.Fatal(err)
	}

	atomic.AddInt32(&i.stats.Workers, int32(i.workers))
	for n := 0; n < i.workers; n++ {
		go func() {
			wQuit := make(chan struct{})
			i.pool <- wQuit
			i.ipfixWorker(wQuit)
		}()
	}

	logger.Printf("ipfix is running (workers#: %d)", i.workers)

	mCache = ipfix.GetCache(opts.IPFIXTplCacheFile)
	go ipfix.RPC(mCache, &ipfix.RPCConfig{
		Enabled: opts.IPFIXRPCEnabled,
		Logger:  logger,
	})

	go mirrorIPFIXDispatcher(ipfixMCh)

	go func() {
		p := producer.NewProducer(opts.MQName)

		p.MQConfigFile = opts.MQConfigFile
		p.MQErrorCount = &i.stats.MQErrorCount
		p.Logger = logger
		p.Chan = ipfixMQCh
		p.Topic = opts.IPFIXTopic

		if err := p.Run(); err != nil {
			logger.Fatal(err)
		}
	}()

	go func() {
		if !opts.DynWorkers {
			logger.Println("IPFIX dynamic worker disabled")
			return
		}

		i.dynWorkers()
	}()

	for !i.stop {
		b := ipfixBuffer.Get().([]byte)
		conn.SetReadDeadline(time.Now().Add(1e9))
		n, raddr, err := conn.ReadFromUDP(b)
		if err != nil {
			continue
		}
		atomic.AddUint64(&i.stats.UDPCount, 1)
		ipfixUDPCh <- IPFIXUDPMsg{raddr, b[:n]}
	}

}

func (i *IPFIX) shutdown() {
	// exit if the ipfix is disabled
	if !opts.IPFIXEnabled {
		logger.Println("ipfix disabled")
		return
	}

	// stop reading from UDP listener
	i.stop = true
	logger.Println("stopping ipfix service gracefully ...")
	time.Sleep(1 * time.Second)

	// dump the templates to storage
	if err := mCache.Dump(opts.IPFIXTplCacheFile); err != nil {
		logger.Println("couldn't not dump template", err)
	}

	// logging and close UDP channel
	logger.Println("ipfix has been shutdown")
	close(ipfixUDPCh)
}

func (i *IPFIX) ipfixWorker(wQuit chan struct{}) {
	var (
		decodedMsg *ipfix.Message
		mirror     IPFIXUDPMsg
		msg        = IPFIXUDPMsg{body: ipfixBuffer.Get().([]byte)}
		buf        = new(bytes.Buffer)
		err        error
		ok         bool
		b          []byte
	)

LOOP:
	for {

		ipfixBuffer.Put(msg.body[:opts.IPFIXUDPSize])
		buf.Reset()

		select {
		case <-wQuit:
			break LOOP
		case msg, ok = <-ipfixUDPCh:
			if !ok {
				break LOOP
			}
		}

		if opts.Verbose {
			logger.Printf("rcvd ipfix data from: %s, size: %d bytes",
				msg.raddr, len(msg.body))
		}

		if ipfixMirrorEnabled {
			mirror.body = ipfixBuffer.Get().([]byte)
			mirror.raddr = msg.raddr
			mirror.body = append(mirror.body[:0], msg.body...)

			select {
			case ipfixMCh <- mirror:
			default:
			}
		}

		d := ipfix.NewDecoder(msg.raddr.IP, msg.body)
		if decodedMsg, err = d.Decode(mCache); err != nil {
			logger.Println(err)
			continue
		}

		atomic.AddUint64(&i.stats.DecodedCount, 1)

		if len(decodedMsg.DataSets) > 0 {
			b, err = decodedMsg.JSONMarshal(buf)
			if err != nil {
				logger.Println(err)
				continue
			}

			select {
			case ipfixMQCh <- append([]byte{}, b...):
			default:
			}

			if opts.Verbose {
				logger.Println(string(b))
			}
		}

	}
}

func (i *IPFIX) status() *IPFIXStats {
	return &IPFIXStats{
		UDPQueue:       len(ipfixUDPCh),
		UDPMirrorQueue: len(ipfixMCh),
		MessageQueue:   len(ipfixMQCh),
		UDPCount:       atomic.LoadUint64(&i.stats.UDPCount),
		DecodedCount:   atomic.LoadUint64(&i.stats.DecodedCount),
		MQErrorCount:   atomic.LoadUint64(&i.stats.MQErrorCount),
		Workers:        atomic.LoadInt32(&i.stats.Workers),
	}
}

func (i *IPFIX) dynWorkers() {
	var load, nSeq, newWorkers, workers, n int

	tick := time.Tick(120 * time.Second)

	for {
		<-tick
		load = 0

		for n = 0; n < 30; n++ {
			time.Sleep(1 * time.Second)
			load += len(ipfixUDPCh)
		}

		if load > 15 {

			switch {
			case load > 300:
				newWorkers = 100
			case load > 200:
				newWorkers = 60
			case load > 100:
				newWorkers = 40
			default:
				newWorkers = 30
			}

			workers = int(atomic.LoadInt32(&i.stats.Workers))
			if workers+newWorkers > maxWorkers {
				logger.Println("ipfix :: max out workers")
				continue
			}

			for n = 0; n < newWorkers; n++ {
				go func() {
					atomic.AddInt32(&i.stats.Workers, 1)
					wQuit := make(chan struct{})
					i.pool <- wQuit
					i.ipfixWorker(wQuit)
				}()
			}

		}

		if load == 0 {
			nSeq++
		} else {
			nSeq = 0
			continue
		}

		if nSeq > 15 {
			for n = 0; n < 10; n++ {
				if len(i.pool) > i.workers {
					atomic.AddInt32(&i.stats.Workers, -1)
					wQuit := <-i.pool
					close(wQuit)
				}
			}

			nSeq = 0
		}
	}
}
<commit_msg>prevent continue loop once we have at least one data set<commit_after>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    ipfix.go
//: details: ipfix decoders handler
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: 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 main

import (
	"bytes"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/VerizonDigital/vflow/ipfix"
	"github.com/VerizonDigital/vflow/producer"
)

// IPFIX represents IPFIX collector
type IPFIX struct {
	port    int
	addr    string
	workers int
	stop    bool
	stats   IPFIXStats
	pool    chan chan struct{}
}

// IPFIXUDPMsg represents IPFIX UDP data
type IPFIXUDPMsg struct {
	raddr *net.UDPAddr
	body  []byte
}

// IPFIXStats represents IPFIX stats
type IPFIXStats struct {
	UDPQueue       int
	UDPMirrorQueue int
	MessageQueue   int
	UDPCount       uint64
	DecodedCount   uint64
	MQErrorCount   uint64
	Workers        int32
}

var (
	ipfixUDPCh         = make(chan IPFIXUDPMsg, 1000)
	ipfixMCh           = make(chan IPFIXUDPMsg, 1000)
	ipfixMQCh          = make(chan []byte, 1000)
	ipfixMirrorEnabled bool

	// templates memory cache
	mCache ipfix.MemCache

	// ipfix udp payload pool
	ipfixBuffer = &sync.Pool{
		New: func() interface{} {
			return make([]byte, opts.IPFIXUDPSize)
		},
	}
)

// NewIPFIX constructs IPFIX
func NewIPFIX() *IPFIX {
	return &IPFIX{
		port:    opts.IPFIXPort,
		workers: opts.IPFIXWorkers,
		pool:    make(chan chan struct{}, maxWorkers),
	}
}

func (i *IPFIX) run() {
	// exit if the ipfix is disabled
	if !opts.IPFIXEnabled {
		logger.Println("ipfix has been disabled")
		return
	}

	hostPort := net.JoinHostPort(i.addr, strconv.Itoa(i.port))
	udpAddr, _ := net.ResolveUDPAddr("udp", hostPort)

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		logger.Fatal(err)
	}

	atomic.AddInt32(&i.stats.Workers, int32(i.workers))
	for n := 0; n < i.workers; n++ {
		go func() {
			wQuit := make(chan struct{})
			i.pool <- wQuit
			i.ipfixWorker(wQuit)
		}()
	}

	logger.Printf("ipfix is running (workers#: %d)", i.workers)

	mCache = ipfix.GetCache(opts.IPFIXTplCacheFile)
	go ipfix.RPC(mCache, &ipfix.RPCConfig{
		Enabled: opts.IPFIXRPCEnabled,
		Logger:  logger,
	})

	go mirrorIPFIXDispatcher(ipfixMCh)

	go func() {
		p := producer.NewProducer(opts.MQName)

		p.MQConfigFile = opts.MQConfigFile
		p.MQErrorCount = &i.stats.MQErrorCount
		p.Logger = logger
		p.Chan = ipfixMQCh
		p.Topic = opts.IPFIXTopic

		if err := p.Run(); err != nil {
			logger.Fatal(err)
		}
	}()

	go func() {
		if !opts.DynWorkers {
			logger.Println("IPFIX dynamic worker disabled")
			return
		}

		i.dynWorkers()
	}()

	for !i.stop {
		b := ipfixBuffer.Get().([]byte)
		conn.SetReadDeadline(time.Now().Add(1e9))
		n, raddr, err := conn.ReadFromUDP(b)
		if err != nil {
			continue
		}
		atomic.AddUint64(&i.stats.UDPCount, 1)
		ipfixUDPCh <- IPFIXUDPMsg{raddr, b[:n]}
	}

}

func (i *IPFIX) shutdown() {
	// exit if the ipfix is disabled
	if !opts.IPFIXEnabled {
		logger.Println("ipfix disabled")
		return
	}

	// stop reading from UDP listener
	i.stop = true
	logger.Println("stopping ipfix service gracefully ...")
	time.Sleep(1 * time.Second)

	// dump the templates to storage
	if err := mCache.Dump(opts.IPFIXTplCacheFile); err != nil {
		logger.Println("couldn't not dump template", err)
	}

	// logging and close UDP channel
	logger.Println("ipfix has been shutdown")
	close(ipfixUDPCh)
}

func (i *IPFIX) ipfixWorker(wQuit chan struct{}) {
	var (
		decodedMsg *ipfix.Message
		mirror     IPFIXUDPMsg
		msg        = IPFIXUDPMsg{body: ipfixBuffer.Get().([]byte)}
		buf        = new(bytes.Buffer)
		err        error
		ok         bool
		b          []byte
	)

LOOP:
	for {

		ipfixBuffer.Put(msg.body[:opts.IPFIXUDPSize])
		buf.Reset()

		select {
		case <-wQuit:
			break LOOP
		case msg, ok = <-ipfixUDPCh:
			if !ok {
				break LOOP
			}
		}

		if opts.Verbose {
			logger.Printf("rcvd ipfix data from: %s, size: %d bytes",
				msg.raddr, len(msg.body))
		}

		if ipfixMirrorEnabled {
			mirror.body = ipfixBuffer.Get().([]byte)
			mirror.raddr = msg.raddr
			mirror.body = append(mirror.body[:0], msg.body...)

			select {
			case ipfixMCh <- mirror:
			default:
			}
		}

		d := ipfix.NewDecoder(msg.raddr.IP, msg.body)
		if decodedMsg, err = d.Decode(mCache); err != nil {
			logger.Println(err)
			// in case ipfix message header couldn't decode
			if decodedMsg == nil {
				continue
			}
		}

		atomic.AddUint64(&i.stats.DecodedCount, 1)

		if len(decodedMsg.DataSets) > 0 {
			b, err = decodedMsg.JSONMarshal(buf)
			if err != nil {
				logger.Println(err)
				continue
			}

			select {
			case ipfixMQCh <- append([]byte{}, b...):
			default:
			}

			if opts.Verbose {
				logger.Println(string(b))
			}
		}

	}
}

func (i *IPFIX) status() *IPFIXStats {
	return &IPFIXStats{
		UDPQueue:       len(ipfixUDPCh),
		UDPMirrorQueue: len(ipfixMCh),
		MessageQueue:   len(ipfixMQCh),
		UDPCount:       atomic.LoadUint64(&i.stats.UDPCount),
		DecodedCount:   atomic.LoadUint64(&i.stats.DecodedCount),
		MQErrorCount:   atomic.LoadUint64(&i.stats.MQErrorCount),
		Workers:        atomic.LoadInt32(&i.stats.Workers),
	}
}

func (i *IPFIX) dynWorkers() {
	var load, nSeq, newWorkers, workers, n int

	tick := time.Tick(120 * time.Second)

	for {
		<-tick
		load = 0

		for n = 0; n < 30; n++ {
			time.Sleep(1 * time.Second)
			load += len(ipfixUDPCh)
		}

		if load > 15 {

			switch {
			case load > 300:
				newWorkers = 100
			case load > 200:
				newWorkers = 60
			case load > 100:
				newWorkers = 40
			default:
				newWorkers = 30
			}

			workers = int(atomic.LoadInt32(&i.stats.Workers))
			if workers+newWorkers > maxWorkers {
				logger.Println("ipfix :: max out workers")
				continue
			}

			for n = 0; n < newWorkers; n++ {
				go func() {
					atomic.AddInt32(&i.stats.Workers, 1)
					wQuit := make(chan struct{})
					i.pool <- wQuit
					i.ipfixWorker(wQuit)
				}()
			}

		}

		if load == 0 {
			nSeq++
		} else {
			nSeq = 0
			continue
		}

		if nSeq > 15 {
			for n = 0; n < 10; n++ {
				if len(i.pool) > i.workers {
					atomic.AddInt32(&i.stats.Workers, -1)
					wQuit := <-i.pool
					close(wQuit)
				}
			}

			nSeq = 0
		}
	}
}
<|endoftext|>