<commit_before>// This file is part of Ringot.
/*
Copyright 2016 tSU-RooT <tsu.root@gmail.com>

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 (
	"fmt"
	"github.com/ChimeraCoder/anaconda"
)

const (
	loadMax = 50
)

type conversationview struct {
	*tweetview

	loadPreviousTweetCh chan *anaconda.Tweet
}

func newConversationview() *conversationview {
	return &conversationview{
		tweetview:           newTweetview(),
		loadPreviousTweetCh: make(chan *anaconda.Tweet),
	}
}

func (cv *conversationview) setTopTweet(ts tweetstatus) {
	if ts.Empty || ts.Content == nil {
		return
	}
	cv.tweets = []tweetstatus{ts}

	count := 1
	tweet := ts.Content
	for {
		id := tweet.InReplyToStatusID
		if id == 0 {
			break
		} else if t, ok := tweetmap.get(id); ok {
			cv.tweets = append(cv.tweets, tweetstatus{Content: t})
			tweet = t
		} else {
			break
		}
		count++
		if count >= loadMax {
			break
		}
	}
}

func (cv *conversationview) addPreviousTweet(ts tweetstatus) {
	cv.tweets = append(cv.tweets, ts)
}

func (cv *conversationview) loadTweet() {
	count := 1
	if len(cv.tweets) < 1 {
		return
	}
	tweet := cv.tweets[len(cv.tweets)-1].Content
	for {
		id := tweet.InReplyToStatusID
		if id == 0 {
			break
		} else if t, ok := tweetmap.get(id); ok {
			cv.loadPreviousTweetCh <- t
			tweet = t
		} else {
			t, err := api.GetTweet(id, nil)
			if err != nil {
				changeBufferState(fmt.Sprintf("Err:Load Tweet(ID:%d)", id))
			}
			cv.loadPreviousTweetCh <- &t
			tweet = &t

		}
		count++
		if count >= loadMax {
			break
		}
	}
}
<commit_msg>Fix handling err at Conversation View<commit_after>// This file is part of Ringot.
/*
Copyright 2016 tSU-RooT <tsu.root@gmail.com>

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 (
	"fmt"
	"github.com/ChimeraCoder/anaconda"
)

const (
	loadMax = 50
)

type conversationview struct {
	*tweetview

	loadPreviousTweetCh chan *anaconda.Tweet
}

func newConversationview() *conversationview {
	return &conversationview{
		tweetview:           newTweetview(),
		loadPreviousTweetCh: make(chan *anaconda.Tweet),
	}
}

func (cv *conversationview) setTopTweet(ts tweetstatus) {
	if ts.Empty || ts.Content == nil {
		return
	}
	cv.tweets = []tweetstatus{ts}

	count := 1
	tweet := ts.Content
	for {
		id := tweet.InReplyToStatusID
		if id == 0 {
			break
		} else if t, ok := tweetmap.get(id); ok {
			cv.tweets = append(cv.tweets, tweetstatus{Content: t})
			tweet = t
		} else {
			break
		}
		count++
		if count >= loadMax {
			break
		}
	}
}

func (cv *conversationview) addPreviousTweet(ts tweetstatus) {
	cv.tweets = append(cv.tweets, ts)
}

func (cv *conversationview) loadTweet() {
	count := 1
	if len(cv.tweets) < 1 {
		return
	}
	tweet := cv.tweets[len(cv.tweets)-1].Content
	for {
		id := tweet.InReplyToStatusID
		if id == 0 {
			break
		} else if t, ok := tweetmap.get(id); ok {
			cv.loadPreviousTweetCh <- t
			tweet = t
		} else {
			t, err := api.GetTweet(id, nil)
			if err != nil {
				changeBufferState(fmt.Sprintf("Err:Load Tweet(ID:%d)", id))
				break
			}
			cv.loadPreviousTweetCh <- &t
			tweet = &t

		}
		count++
		if count >= loadMax {
			break
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/Sirupsen/logrus"
	"github.com/TykTechnologies/goverify"
	"github.com/TykTechnologies/tykcommon"

	"archive/zip"
	"bytes"
	"crypto/md5"
	b64 "encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
)

// Bundle is the basic bundle data structure, it holds the bundle name and the data.
type Bundle struct {
	Name     string
	Data     []byte
	Path     string
	Spec     *APISpec
	Manifest tykcommon.BundleManifest
}

func (b *Bundle) Verify() (err error) {
	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Verifying bundle: ", b.Spec.CustomMiddlewareBundle)

	var useSignature bool
	var bundleVerifier goverify.Verifier

	// Perform signature verification if a public key path is set:
	if config.PublicKeyPath != "" {
		if b.Manifest.Signature == "" {
			// Error: A public key is set, but the bundle isn't signed.
			err = errors.New("Bundle isn't signed")
		}
		if notificationVerifier == nil {
			bundleVerifier, err = goverify.LoadPublicKeyFromFile(config.PublicKeyPath)
		}

		if err != nil {
			return err
		}

		useSignature = true
	}

	h := md5.New()
	h.Write(b.Data)
	checksum := hex.EncodeToString(h.Sum(nil))

	var bundleData bytes.Buffer

	for _, f := range b.Manifest.FileList {
		extractedFilePath := filepath.Join(b.Path, f)

		var data []byte
		data, err = ioutil.ReadFile(extractedFilePath)
		if err != nil {
			break
		}

		bundleData.Write(data)
	}

	checksum = fmt.Sprintf("%x", md5.Sum(bundleData.Bytes()))

	if checksum != b.Manifest.Checksum {
		err = errors.New("Invalid checksum")
	}

	if useSignature {
		var signed []byte
		signed, err = b64.StdEncoding.DecodeString(b.Manifest.Signature)
		if err != nil {
			return err
		}
		err = bundleVerifier.Verify([]byte(bundleData.Bytes()), signed)
		if err != nil {
			return err
		}

	}

	return err
}

func (b *Bundle) AddToSpec() {
	b.Spec.APIDefinition.CustomMiddleware = b.Manifest.CustomMiddleware
}

// BundleGetter is used for downloading bundle data, see HttpBundleGetter for reference.
type BundleGetter interface {
	Get() ([]byte, error)
}

// HttpBundleGetter is a simple HTTP BundleGetter.
type HttpBundleGetter struct {
	Url string
}

// Get performs an HTTP GET request.
func (g *HttpBundleGetter) Get() (bundleData []byte, err error) {
	var resp *http.Response

	resp, err = http.Get(g.Url)

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, errors.New("HTTP Error")
	}

	defer resp.Body.Close()
	bundleData, err = ioutil.ReadAll(resp.Body)
	return bundleData, err
}

// BundleSaver is an interface used by bundle saver structures.
type BundleSaver interface {
	Save(*Bundle, string, *APISpec) error
}

// ZipBundleSaver is a BundleSaver for ZIP files.
type ZipBundleSaver struct {
}

// Save implements the main method of the BundleSaver interface. It makes use of archive/zip.
func (s *ZipBundleSaver) Save(bundle *Bundle, bundlePath string, spec *APISpec) (err error) {
	buf := bytes.NewReader(bundle.Data)
	reader, _ := zip.NewReader(buf, int64(len(bundle.Data)))

	for _, f := range reader.File {
		var rc io.ReadCloser
		rc, err = f.Open()

		if err != nil {
			return err
		}

		var destPath string
		destPath = filepath.Join(bundlePath, f.Name)

		isDir := f.FileHeader.Mode().IsDir()

		if isDir {
			err = os.Mkdir(destPath, 0755)
			if err != nil {
				return err
			}
		} else {
			var newFile *os.File
			newFile, err = os.Create(destPath)
			if err != nil {
				return err
			}
			_, err = io.Copy(newFile, rc)
			if err != nil {
				return err
			}
		}
	}
	return err
}

// fetchBundle will fetch a given bundle, using the right BundleGetter. The first argument is the bundle name, the base bundle URL will be used as prefix.
func fetchBundle(spec *APISpec) (thisBundle Bundle, err error) {
	var bundleUrl string

	bundleUrl = strings.Join([]string{config.BundleBaseURL, spec.CustomMiddlewareBundle}, "")

	var thisGetter BundleGetter

	var u *url.URL
	u, err = url.Parse(bundleUrl)

	switch u.Scheme {
	case "http":
		thisGetter = &HttpBundleGetter{
			Url: bundleUrl,
		}
	default:
		err = errors.New("Unknown URL scheme!")
	}

	bundleData, err := thisGetter.Get()

	thisBundle = Bundle{
		Name: spec.CustomMiddlewareBundle,
		Data: bundleData,
		Spec: spec,
	}

	return thisBundle, err
}

// saveBundle will save a bundle to the disk, see ZipBundleSaver methods for reference.
func saveBundle(bundle *Bundle, destPath string, spec *APISpec) (err error) {
	var bundleFormat = "zip"

	var bundleSaver BundleSaver

	// TODO: use enums?
	switch bundleFormat {
	case "zip":
		bundleSaver = &ZipBundleSaver{}
	}

	bundleSaver.Save(bundle, destPath, spec)

	return err
}

// loadBundleManifest will parse the manifest file and return the bundle parameters.
func loadBundleManifest(bundle *Bundle, spec *APISpec, skipVerification bool) (err error) {
	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Loading bundle: ", spec.CustomMiddlewareBundle)

	manifestPath := filepath.Join(bundle.Path, "manifest.json")
	var manifestData []byte
	manifestData, err = ioutil.ReadFile(manifestPath)

	err = json.Unmarshal(manifestData, &bundle.Manifest)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("----> Couldn't unmarshal the manifest file for bundle: ", spec.CustomMiddlewareBundle)
		return err
	}

	if skipVerification {
		return err
	}

	err = bundle.Verify()
	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("----> Bundle verification failed: ", spec.CustomMiddlewareBundle)
	}

	return err
}

// loadBundle wraps the load and save steps, it will return if an error occurs at any point.
func loadBundle(spec *APISpec) {
	var err error

	// Skip if no custom middleware bundle name is set.
	if spec.CustomMiddlewareBundle == "" {
		return
	}

	// Skip if no bundle base URL is set.
	if config.BundleBaseURL == "" {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("An API specifies a custom middleware bundle, but no bundle base URL is set in your tyk.conf! Skipping bundle: ", spec.CustomMiddlewareBundle)
		return
	}

	// Skip if the bundle destination path already exists.
	bundlePath := strings.Join([]string{spec.APIID, spec.CustomMiddlewareBundle}, "-")
	destPath := filepath.Join("/Users/matias/dev/tyk", "middleware/bundles", bundlePath)

	// The bundle exists, load and return:
	if _, err := os.Stat(destPath); err == nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("Loading existing bundle: ", spec.CustomMiddlewareBundle)

		bundle := Bundle{
			Name: spec.CustomMiddlewareBundle,
			Path: destPath,
			Spec: spec,
		}

		err = loadBundleManifest(&bundle, spec, true)
		if err != nil {
			log.WithFields(logrus.Fields{
				"prefix": "main",
			}).Info("----> Couldn't load bundle: ", spec.CustomMiddlewareBundle, " ", err)
		}

		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("----> Using bundle: ", spec.CustomMiddlewareBundle)

		bundle.AddToSpec()

		return
	}

	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Fetching Bundle: ", spec.CustomMiddlewareBundle)

	var bundle Bundle
	bundle, err = fetchBundle(spec)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Couldn't fetch bundle: ", spec.CustomMiddlewareBundle, ", ", err)
		return
	}

	err = os.Mkdir(destPath, 0755)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Error when creating bundle directory: ", spec.CustomMiddlewareBundle, ", ", err)
		// return
	}

	err = saveBundle(&bundle, destPath, spec)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Couldn't save bundle: ", spec.CustomMiddlewareBundle, ", ", err)
		return
	}

	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Debug("----> Saving Bundle: ", spec.CustomMiddlewareBundle)

	// Set the destination path:
	bundle.Path = destPath

	err = loadBundleManifest(&bundle, spec, false)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Couldn't load bundle: ", spec.CustomMiddlewareBundle, " ", err)

		removeErr := os.RemoveAll(bundle.Path)
		if removeErr != nil {
			log.WithFields(logrus.Fields{
				"prefix": "main",
			}).Error("----> Couldn't remove bundle: ", spec.CustomMiddlewareBundle, removeErr)
		}
		return
	}

	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Bundle is valid, adding to spec: ", spec.CustomMiddlewareBundle)

	bundle.AddToSpec()

}
<commit_msg>tyk-54: use Tyk directory when building bundle path.<commit_after>package main

import (
	"github.com/Sirupsen/logrus"
	"github.com/TykTechnologies/goverify"
	"github.com/TykTechnologies/tykcommon"

	"archive/zip"
	"bytes"
	"crypto/md5"
	b64 "encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
)

// Bundle is the basic bundle data structure, it holds the bundle name and the data.
type Bundle struct {
	Name     string
	Data     []byte
	Path     string
	Spec     *APISpec
	Manifest tykcommon.BundleManifest
}

func (b *Bundle) Verify() (err error) {
	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Verifying bundle: ", b.Spec.CustomMiddlewareBundle)

	var useSignature bool
	var bundleVerifier goverify.Verifier

	// Perform signature verification if a public key path is set:
	if config.PublicKeyPath != "" {
		if b.Manifest.Signature == "" {
			// Error: A public key is set, but the bundle isn't signed.
			err = errors.New("Bundle isn't signed")
		}
		if notificationVerifier == nil {
			bundleVerifier, err = goverify.LoadPublicKeyFromFile(config.PublicKeyPath)
		}

		if err != nil {
			return err
		}

		useSignature = true
	}

	h := md5.New()
	h.Write(b.Data)
	checksum := hex.EncodeToString(h.Sum(nil))

	var bundleData bytes.Buffer

	for _, f := range b.Manifest.FileList {
		extractedFilePath := filepath.Join(b.Path, f)

		var data []byte
		data, err = ioutil.ReadFile(extractedFilePath)
		if err != nil {
			break
		}

		bundleData.Write(data)
	}

	checksum = fmt.Sprintf("%x", md5.Sum(bundleData.Bytes()))

	if checksum != b.Manifest.Checksum {
		err = errors.New("Invalid checksum")
	}

	if useSignature {
		var signed []byte
		signed, err = b64.StdEncoding.DecodeString(b.Manifest.Signature)
		if err != nil {
			return err
		}
		err = bundleVerifier.Verify([]byte(bundleData.Bytes()), signed)
		if err != nil {
			return err
		}

	}

	return err
}

func (b *Bundle) AddToSpec() {
	b.Spec.APIDefinition.CustomMiddleware = b.Manifest.CustomMiddleware
}

// BundleGetter is used for downloading bundle data, see HttpBundleGetter for reference.
type BundleGetter interface {
	Get() ([]byte, error)
}

// HttpBundleGetter is a simple HTTP BundleGetter.
type HttpBundleGetter struct {
	Url string
}

// Get performs an HTTP GET request.
func (g *HttpBundleGetter) Get() (bundleData []byte, err error) {
	var resp *http.Response

	resp, err = http.Get(g.Url)

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, errors.New("HTTP Error")
	}

	defer resp.Body.Close()
	bundleData, err = ioutil.ReadAll(resp.Body)
	return bundleData, err
}

// BundleSaver is an interface used by bundle saver structures.
type BundleSaver interface {
	Save(*Bundle, string, *APISpec) error
}

// ZipBundleSaver is a BundleSaver for ZIP files.
type ZipBundleSaver struct {
}

// Save implements the main method of the BundleSaver interface. It makes use of archive/zip.
func (s *ZipBundleSaver) Save(bundle *Bundle, bundlePath string, spec *APISpec) (err error) {
	buf := bytes.NewReader(bundle.Data)
	reader, _ := zip.NewReader(buf, int64(len(bundle.Data)))

	for _, f := range reader.File {
		var rc io.ReadCloser
		rc, err = f.Open()

		if err != nil {
			return err
		}

		var destPath string
		destPath = filepath.Join(bundlePath, f.Name)

		isDir := f.FileHeader.Mode().IsDir()

		if isDir {
			err = os.Mkdir(destPath, 0755)
			if err != nil {
				return err
			}
		} else {
			var newFile *os.File
			newFile, err = os.Create(destPath)
			if err != nil {
				return err
			}
			_, err = io.Copy(newFile, rc)
			if err != nil {
				return err
			}
		}
	}
	return err
}

// fetchBundle will fetch a given bundle, using the right BundleGetter. The first argument is the bundle name, the base bundle URL will be used as prefix.
func fetchBundle(spec *APISpec) (thisBundle Bundle, err error) {
	var bundleUrl string

	bundleUrl = strings.Join([]string{config.BundleBaseURL, spec.CustomMiddlewareBundle}, "")

	var thisGetter BundleGetter

	var u *url.URL
	u, err = url.Parse(bundleUrl)

	switch u.Scheme {
	case "http":
		thisGetter = &HttpBundleGetter{
			Url: bundleUrl,
		}
	default:
		err = errors.New("Unknown URL scheme!")
	}

	bundleData, err := thisGetter.Get()

	thisBundle = Bundle{
		Name: spec.CustomMiddlewareBundle,
		Data: bundleData,
		Spec: spec,
	}

	return thisBundle, err
}

// saveBundle will save a bundle to the disk, see ZipBundleSaver methods for reference.
func saveBundle(bundle *Bundle, destPath string, spec *APISpec) (err error) {
	var bundleFormat = "zip"

	var bundleSaver BundleSaver

	// TODO: use enums?
	switch bundleFormat {
	case "zip":
		bundleSaver = &ZipBundleSaver{}
	}

	bundleSaver.Save(bundle, destPath, spec)

	return err
}

// loadBundleManifest will parse the manifest file and return the bundle parameters.
func loadBundleManifest(bundle *Bundle, spec *APISpec, skipVerification bool) (err error) {
	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Loading bundle: ", spec.CustomMiddlewareBundle)

	manifestPath := filepath.Join(bundle.Path, "manifest.json")
	var manifestData []byte
	manifestData, err = ioutil.ReadFile(manifestPath)

	err = json.Unmarshal(manifestData, &bundle.Manifest)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("----> Couldn't unmarshal the manifest file for bundle: ", spec.CustomMiddlewareBundle)
		return err
	}

	if skipVerification {
		return err
	}

	err = bundle.Verify()
	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("----> Bundle verification failed: ", spec.CustomMiddlewareBundle)
	}

	return err
}

// loadBundle wraps the load and save steps, it will return if an error occurs at any point.
func loadBundle(spec *APISpec) {
	var err error

	// Skip if no custom middleware bundle name is set.
	if spec.CustomMiddlewareBundle == "" {
		return
	}

	// Skip if no bundle base URL is set.
	if config.BundleBaseURL == "" {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("An API specifies a custom middleware bundle, but no bundle base URL is set in your tyk.conf! Skipping bundle: ", spec.CustomMiddlewareBundle)
		return
	}

	workDir, _ := os.Getwd()

	// Skip if the bundle destination path already exists.
	bundlePath := strings.Join([]string{spec.APIID, spec.CustomMiddlewareBundle}, "-")
	destPath := filepath.Join(workDir, "middleware/bundles", bundlePath)

	// The bundle exists, load and return:
	if _, err := os.Stat(destPath); err == nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("Loading existing bundle: ", spec.CustomMiddlewareBundle)

		bundle := Bundle{
			Name: spec.CustomMiddlewareBundle,
			Path: destPath,
			Spec: spec,
		}

		err = loadBundleManifest(&bundle, spec, true)
		if err != nil {
			log.WithFields(logrus.Fields{
				"prefix": "main",
			}).Info("----> Couldn't load bundle: ", spec.CustomMiddlewareBundle, " ", err)
		}

		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Info("----> Using bundle: ", spec.CustomMiddlewareBundle)

		bundle.AddToSpec()

		return
	}

	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Fetching Bundle: ", spec.CustomMiddlewareBundle)

	var bundle Bundle
	bundle, err = fetchBundle(spec)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Couldn't fetch bundle: ", spec.CustomMiddlewareBundle, ", ", err)
		return
	}

	err = os.Mkdir(destPath, 0755)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Error when creating bundle directory: ", spec.CustomMiddlewareBundle, ", ", err)
		// return
	}

	err = saveBundle(&bundle, destPath, spec)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Couldn't save bundle: ", spec.CustomMiddlewareBundle, ", ", err)
		return
	}

	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Debug("----> Saving Bundle: ", spec.CustomMiddlewareBundle)

	// Set the destination path:
	bundle.Path = destPath

	err = loadBundleManifest(&bundle, spec, false)

	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "main",
		}).Error("----> Couldn't load bundle: ", spec.CustomMiddlewareBundle, " ", err)

		removeErr := os.RemoveAll(bundle.Path)
		if removeErr != nil {
			log.WithFields(logrus.Fields{
				"prefix": "main",
			}).Error("----> Couldn't remove bundle: ", spec.CustomMiddlewareBundle, removeErr)
		}
		return
	}

	log.WithFields(logrus.Fields{
		"prefix": "main",
	}).Info("----> Bundle is valid, adding to spec: ", spec.CustomMiddlewareBundle)

	bundle.AddToSpec()

}
<|endoftext|>


<commit_before>package transactionpool

import (
	"errors"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	// The TransactionPoolSizeLimit is first checked, and then a transaction
	// set is added. The current transaction pool does not do any priority
	// ordering, so the size limit is such that the transaction pool will never
	// exceed the size of a block.
	//
	// TODO: Add a priority structure that will allow the transaction pool to
	// fill up beyond the size of a single block, without being subject to
	// manipulation.
	//
	// The first ~1/4 of the transaction pool can be filled for free. This is
	// mostly to preserve compatibility with clients that do not add fees.
	TransactionPoolSizeLimit  = 2e6 - 5e3 - modules.TransactionSetSizeLimit
	TransactionPoolSizeForFee = 500e3
)

var (
	errObjectConflict      = errors.New("transaction set conflicts with an existing transaction set")
	errFullTransactionPool = errors.New("transaction pool cannot accept more transactions")
	errLowMinerFees        = errors.New("transaction set needs more miner fees to be accepted")
	errEmptySet            = errors.New("transaction set is empty")

	TransactionMinFee = types.NewCurrency64(2).Mul(types.SiacoinPrecision)
)

// relatedObjectIDs determines all of the object ids related to a transaction.
func relatedObjectIDs(ts []types.Transaction) []ObjectID {
	oidMap := make(map[ObjectID]struct{})
	for _, t := range ts {
		for _, sci := range t.SiacoinInputs {
			oidMap[ObjectID(sci.ParentID)] = struct{}{}
		}
		for i := range t.SiacoinOutputs {
			oidMap[ObjectID(t.SiacoinOutputID(uint64(i)))] = struct{}{}
		}
		for i := range t.FileContracts {
			oidMap[ObjectID(t.FileContractID(uint64(i)))] = struct{}{}
		}
		for _, fcr := range t.FileContractRevisions {
			oidMap[ObjectID(fcr.ParentID)] = struct{}{}
		}
		for _, sp := range t.StorageProofs {
			oidMap[ObjectID(sp.ParentID)] = struct{}{}
		}
		for _, sfi := range t.SiafundInputs {
			oidMap[ObjectID(sfi.ParentID)] = struct{}{}
		}
		for i := range t.SiafundOutputs {
			oidMap[ObjectID(t.SiafundOutputID(uint64(i)))] = struct{}{}
		}
	}

	var oids []ObjectID
	for oid, _ := range oidMap {
		oids = append(oids, oid)
	}
	return oids
}

// checkMinerFees checks that the total amount of transaction fees in the
// transaction set is sufficient to earn a spot in the transaction pool.
func (tp *TransactionPool) checkMinerFees(ts []types.Transaction) error {
	// Transactions cannot be added after the TransactionPoolSizeLimit has been
	// hit.
	if tp.transactionListSize > TransactionPoolSizeLimit {
		return errFullTransactionPool
	}

	// The first TransactionPoolSizeForFee transactions do not need fees.
	if tp.transactionListSize > TransactionPoolSizeForFee {
		// Currently required fees are set on a per-transaction basis. 2 coins
		// are required per transaction if the free-fee limit has been reached,
		// adding a larger fee is not useful.
		var feeSum types.Currency
		for i := range ts {
			for _, fee := range ts[i].MinerFees {
				feeSum = feeSum.Add(fee)
			}
		}
		feeRequired := TransactionMinFee.Mul(types.NewCurrency64(uint64(len(ts))))
		if feeSum.Cmp(feeRequired) < 0 {
			return errLowMinerFees
		}
	}
	return nil
}

// checkTransactionSetComposition checks if the transaction set is valid given
// the state of the pool. It does not check that each individual transaction
// would be legal in the next block, but does check things like miner fees and
// IsStandard.
func (tp *TransactionPool) checkTransactionSetComposition(ts []types.Transaction) error {
	// Check that the transaction set is not already known.
	setID := TransactionSetID(crypto.HashObject(ts))
	_, exists := tp.transactionSets[setID]
	if exists {
		return modules.ErrDuplicateTransactionSet
	}

	// Check that the transaction set has enough fees to justify adding it to
	// the transaction list.
	err := tp.checkMinerFees(ts)
	if err != nil {
		return err
	}

	// All checks after this are expensive.
	//
	// TODO: There is no DoS prevention mechanism in place to prevent repeated
	// expensive verifications of invalid transactions that are created on the
	// fly.

	// Check that all transactions follow 'Standard.md' guidelines.
	err = tp.IsStandardTransactionSet(ts)
	if err != nil {
		return err
	}
	return nil
}

// handleConflicts detects whether the conflicts in the transaction pool are
// legal children of the new transaction pool set or not.
func (tp *TransactionPool) handleConflicts(ts []types.Transaction, conflicts []TransactionSetID) error {
	// Create a list of all the transaction ids that compose the set of
	// conflicts.
	conflictMap := make(map[types.TransactionID]TransactionSetID)
	for _, conflict := range conflicts {
		conflictSet := tp.transactionSets[conflict]
		for _, conflictTxn := range conflictSet {
			conflictMap[conflictTxn.ID()] = conflict
		}
	}

	// Discard all duplicate transactions from the input transaction set.
	var dedupSet []types.Transaction
	for _, t := range ts {
		_, exists := conflictMap[t.ID()]
		if exists {
			continue
		}
		dedupSet = append(dedupSet, t)
	}
	if len(dedupSet) == 0 {
		return modules.ErrDuplicateTransactionSet
	}
	// If transactions were pruned, it's possible that the set of
	// dependencies/conflicts has also reduced. To minimize computational load
	// on the consensus set, we want to prune out all of the conflicts that are
	// no longer relevant. As an example, consider the transaction set {A}, the
	// set {B}, and the new set {A, C}, where C is dependent on B. {A} and {B}
	// are both conflicts, but after deduplication {A} is no longer a conflict.
	// This is recursive, but it is guaranteed to run only once as the first
	// deduplication is guaranteed to be complete.
	if len(dedupSet) < len(ts) {
		oids := relatedObjectIDs(dedupSet)
		var conflicts []TransactionSetID
		for _, oid := range oids {
			conflict, exists := tp.knownObjects[oid]
			if exists {
				conflicts = append(conflicts, conflict)
			}
		}
		return tp.handleConflicts(dedupSet, conflicts)
	}

	// Merge all of the conflict sets with the input set (input set goes last
	// to preserve dependency ordering), and see if the set as a whole is both
	// small enough to be legal and valid as a set. If no, return an error. If
	// yes, add the new set to the pool, and eliminate the old set. The output
	// diff objects can be repeated, (no need to remove those). Just need to
	// remove the conflicts from tp.transactionSets.
	var superset []types.Transaction
	supersetMap := make(map[TransactionSetID]struct{})
	for _, conflict := range conflictMap {
		supersetMap[conflict] = struct{}{}
	}
	for conflict := range supersetMap {
		superset = append(superset, tp.transactionSets[conflict]...)
	}
	superset = append(superset, dedupSet...)

	// Check the composition of the transaction set, including fees and
	// IsStandard rules (this is a new set, the rules must be rechecked).
	err := tp.checkTransactionSetComposition(superset)
	if err != nil {
		return err
	}

	// Check that the transaction set is valid.
	cc, err := tp.consensusSet.TryTransactionSet(superset)
	if err != nil {
		return err
	}

	// Remove the conflicts from the transaction pool. The diffs do not need to
	// be removed, they will be overwritten later in the function.
	for _, conflict := range conflictMap {
		conflictSet := tp.transactionSets[conflict]
		tp.transactionListSize -= len(encoding.Marshal(conflictSet))
		delete(tp.transactionSets, conflict)
		delete(tp.transactionSetDiffs, conflict)
	}

	// Add the transaction set to the pool.
	setID := TransactionSetID(crypto.HashObject(superset))
	tp.transactionSets[setID] = superset
	for _, diff := range cc.SiacoinOutputDiffs {
		tp.knownObjects[ObjectID(diff.ID)] = setID
	}
	for _, diff := range cc.FileContractDiffs {
		tp.knownObjects[ObjectID(diff.ID)] = setID
	}
	for _, diff := range cc.SiafundOutputDiffs {
		tp.knownObjects[ObjectID(diff.ID)] = setID
	}
	tp.transactionSetDiffs[setID] = cc
	tp.transactionListSize += len(encoding.Marshal(superset))
	return nil
}

// acceptTransactionSet verifies that a transaction set is allowed to be in the
// transaction pool, and then adds it to the transaction pool.
func (tp *TransactionPool) acceptTransactionSet(ts []types.Transaction) error {
	if len(ts) == 0 {
		return errEmptySet
	}

	// Check the composition of the transaction set, including fees and
	// IsStandard rules.
	err := tp.checkTransactionSetComposition(ts)
	if err != nil {
		return err
	}

	// Check for conflicts with other transactions, which would indicate a
	// double-spend. Legal children of a transaction set will also trigger the
	// conflict-detector.
	oids := relatedObjectIDs(ts)
	var conflicts []TransactionSetID
	for _, oid := range oids {
		conflict, exists := tp.knownObjects[oid]
		if exists {
			conflicts = append(conflicts, conflict)
		}
	}
	if len(conflicts) > 0 {
		return tp.handleConflicts(ts, conflicts)
	}
	cc, err := tp.consensusSet.TryTransactionSet(ts)
	if err != nil {
		return err
	}

	// Add the transaction set to the pool.
	setID := TransactionSetID(crypto.HashObject(ts))
	tp.transactionSets[setID] = ts
	for _, oid := range oids {
		tp.knownObjects[oid] = setID
	}
	tp.transactionSetDiffs[setID] = cc
	tp.transactionListSize += len(encoding.Marshal(ts))
	return nil
}

// AcceptTransaction adds a transaction to the unconfirmed set of
// transactions. If the transaction is accepted, it will be relayed to
// connected peers.
func (tp *TransactionPool) AcceptTransactionSet(ts []types.Transaction) error {
	tp.mu.Lock()
	defer tp.mu.Unlock()

	err := tp.acceptTransactionSet(ts)
	if err != nil {
		return err
	}

	// Notify subscribers and broadcast the transaction set.
	go tp.gateway.Broadcast("RelayTransactionSet", ts)
	tp.updateSubscribersTransactions()

	return nil
}

// relayTransactionSet is an RPC that accepts a transaction set from a peer. If
// the accept is successful, the transaction will be relayed to the gateway's
// other peers.
func (tp *TransactionPool) relayTransactionSet(conn modules.PeerConn) error {
	var ts []types.Transaction
	err := encoding.ReadObject(conn, &ts, types.BlockSizeLimit)
	if err != nil {
		return err
	}
	return tp.AcceptTransactionSet(ts)
}
<commit_msg>restore RelayTransaction compatibility<commit_after>package transactionpool

import (
	"errors"
	"time"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	// The TransactionPoolSizeLimit is first checked, and then a transaction
	// set is added. The current transaction pool does not do any priority
	// ordering, so the size limit is such that the transaction pool will never
	// exceed the size of a block.
	//
	// TODO: Add a priority structure that will allow the transaction pool to
	// fill up beyond the size of a single block, without being subject to
	// manipulation.
	//
	// The first ~1/4 of the transaction pool can be filled for free. This is
	// mostly to preserve compatibility with clients that do not add fees.
	TransactionPoolSizeLimit  = 2e6 - 5e3 - modules.TransactionSetSizeLimit
	TransactionPoolSizeForFee = 500e3
)

var (
	errObjectConflict      = errors.New("transaction set conflicts with an existing transaction set")
	errFullTransactionPool = errors.New("transaction pool cannot accept more transactions")
	errLowMinerFees        = errors.New("transaction set needs more miner fees to be accepted")
	errEmptySet            = errors.New("transaction set is empty")

	TransactionMinFee = types.NewCurrency64(2).Mul(types.SiacoinPrecision)
)

// relatedObjectIDs determines all of the object ids related to a transaction.
func relatedObjectIDs(ts []types.Transaction) []ObjectID {
	oidMap := make(map[ObjectID]struct{})
	for _, t := range ts {
		for _, sci := range t.SiacoinInputs {
			oidMap[ObjectID(sci.ParentID)] = struct{}{}
		}
		for i := range t.SiacoinOutputs {
			oidMap[ObjectID(t.SiacoinOutputID(uint64(i)))] = struct{}{}
		}
		for i := range t.FileContracts {
			oidMap[ObjectID(t.FileContractID(uint64(i)))] = struct{}{}
		}
		for _, fcr := range t.FileContractRevisions {
			oidMap[ObjectID(fcr.ParentID)] = struct{}{}
		}
		for _, sp := range t.StorageProofs {
			oidMap[ObjectID(sp.ParentID)] = struct{}{}
		}
		for _, sfi := range t.SiafundInputs {
			oidMap[ObjectID(sfi.ParentID)] = struct{}{}
		}
		for i := range t.SiafundOutputs {
			oidMap[ObjectID(t.SiafundOutputID(uint64(i)))] = struct{}{}
		}
	}

	var oids []ObjectID
	for oid, _ := range oidMap {
		oids = append(oids, oid)
	}
	return oids
}

// checkMinerFees checks that the total amount of transaction fees in the
// transaction set is sufficient to earn a spot in the transaction pool.
func (tp *TransactionPool) checkMinerFees(ts []types.Transaction) error {
	// Transactions cannot be added after the TransactionPoolSizeLimit has been
	// hit.
	if tp.transactionListSize > TransactionPoolSizeLimit {
		return errFullTransactionPool
	}

	// The first TransactionPoolSizeForFee transactions do not need fees.
	if tp.transactionListSize > TransactionPoolSizeForFee {
		// Currently required fees are set on a per-transaction basis. 2 coins
		// are required per transaction if the free-fee limit has been reached,
		// adding a larger fee is not useful.
		var feeSum types.Currency
		for i := range ts {
			for _, fee := range ts[i].MinerFees {
				feeSum = feeSum.Add(fee)
			}
		}
		feeRequired := TransactionMinFee.Mul(types.NewCurrency64(uint64(len(ts))))
		if feeSum.Cmp(feeRequired) < 0 {
			return errLowMinerFees
		}
	}
	return nil
}

// checkTransactionSetComposition checks if the transaction set is valid given
// the state of the pool. It does not check that each individual transaction
// would be legal in the next block, but does check things like miner fees and
// IsStandard.
func (tp *TransactionPool) checkTransactionSetComposition(ts []types.Transaction) error {
	// Check that the transaction set is not already known.
	setID := TransactionSetID(crypto.HashObject(ts))
	_, exists := tp.transactionSets[setID]
	if exists {
		return modules.ErrDuplicateTransactionSet
	}

	// Check that the transaction set has enough fees to justify adding it to
	// the transaction list.
	err := tp.checkMinerFees(ts)
	if err != nil {
		return err
	}

	// All checks after this are expensive.
	//
	// TODO: There is no DoS prevention mechanism in place to prevent repeated
	// expensive verifications of invalid transactions that are created on the
	// fly.

	// Check that all transactions follow 'Standard.md' guidelines.
	err = tp.IsStandardTransactionSet(ts)
	if err != nil {
		return err
	}
	return nil
}

// handleConflicts detects whether the conflicts in the transaction pool are
// legal children of the new transaction pool set or not.
func (tp *TransactionPool) handleConflicts(ts []types.Transaction, conflicts []TransactionSetID) error {
	// Create a list of all the transaction ids that compose the set of
	// conflicts.
	conflictMap := make(map[types.TransactionID]TransactionSetID)
	for _, conflict := range conflicts {
		conflictSet := tp.transactionSets[conflict]
		for _, conflictTxn := range conflictSet {
			conflictMap[conflictTxn.ID()] = conflict
		}
	}

	// Discard all duplicate transactions from the input transaction set.
	var dedupSet []types.Transaction
	for _, t := range ts {
		_, exists := conflictMap[t.ID()]
		if exists {
			continue
		}
		dedupSet = append(dedupSet, t)
	}
	if len(dedupSet) == 0 {
		return modules.ErrDuplicateTransactionSet
	}
	// If transactions were pruned, it's possible that the set of
	// dependencies/conflicts has also reduced. To minimize computational load
	// on the consensus set, we want to prune out all of the conflicts that are
	// no longer relevant. As an example, consider the transaction set {A}, the
	// set {B}, and the new set {A, C}, where C is dependent on B. {A} and {B}
	// are both conflicts, but after deduplication {A} is no longer a conflict.
	// This is recursive, but it is guaranteed to run only once as the first
	// deduplication is guaranteed to be complete.
	if len(dedupSet) < len(ts) {
		oids := relatedObjectIDs(dedupSet)
		var conflicts []TransactionSetID
		for _, oid := range oids {
			conflict, exists := tp.knownObjects[oid]
			if exists {
				conflicts = append(conflicts, conflict)
			}
		}
		return tp.handleConflicts(dedupSet, conflicts)
	}

	// Merge all of the conflict sets with the input set (input set goes last
	// to preserve dependency ordering), and see if the set as a whole is both
	// small enough to be legal and valid as a set. If no, return an error. If
	// yes, add the new set to the pool, and eliminate the old set. The output
	// diff objects can be repeated, (no need to remove those). Just need to
	// remove the conflicts from tp.transactionSets.
	var superset []types.Transaction
	supersetMap := make(map[TransactionSetID]struct{})
	for _, conflict := range conflictMap {
		supersetMap[conflict] = struct{}{}
	}
	for conflict := range supersetMap {
		superset = append(superset, tp.transactionSets[conflict]...)
	}
	superset = append(superset, dedupSet...)

	// Check the composition of the transaction set, including fees and
	// IsStandard rules (this is a new set, the rules must be rechecked).
	err := tp.checkTransactionSetComposition(superset)
	if err != nil {
		return err
	}

	// Check that the transaction set is valid.
	cc, err := tp.consensusSet.TryTransactionSet(superset)
	if err != nil {
		return err
	}

	// Remove the conflicts from the transaction pool. The diffs do not need to
	// be removed, they will be overwritten later in the function.
	for _, conflict := range conflictMap {
		conflictSet := tp.transactionSets[conflict]
		tp.transactionListSize -= len(encoding.Marshal(conflictSet))
		delete(tp.transactionSets, conflict)
		delete(tp.transactionSetDiffs, conflict)
	}

	// Add the transaction set to the pool.
	setID := TransactionSetID(crypto.HashObject(superset))
	tp.transactionSets[setID] = superset
	for _, diff := range cc.SiacoinOutputDiffs {
		tp.knownObjects[ObjectID(diff.ID)] = setID
	}
	for _, diff := range cc.FileContractDiffs {
		tp.knownObjects[ObjectID(diff.ID)] = setID
	}
	for _, diff := range cc.SiafundOutputDiffs {
		tp.knownObjects[ObjectID(diff.ID)] = setID
	}
	tp.transactionSetDiffs[setID] = cc
	tp.transactionListSize += len(encoding.Marshal(superset))
	return nil
}

// acceptTransactionSet verifies that a transaction set is allowed to be in the
// transaction pool, and then adds it to the transaction pool.
func (tp *TransactionPool) acceptTransactionSet(ts []types.Transaction) error {
	if len(ts) == 0 {
		return errEmptySet
	}

	// Check the composition of the transaction set, including fees and
	// IsStandard rules.
	err := tp.checkTransactionSetComposition(ts)
	if err != nil {
		return err
	}

	// Check for conflicts with other transactions, which would indicate a
	// double-spend. Legal children of a transaction set will also trigger the
	// conflict-detector.
	oids := relatedObjectIDs(ts)
	var conflicts []TransactionSetID
	for _, oid := range oids {
		conflict, exists := tp.knownObjects[oid]
		if exists {
			conflicts = append(conflicts, conflict)
		}
	}
	if len(conflicts) > 0 {
		return tp.handleConflicts(ts, conflicts)
	}
	cc, err := tp.consensusSet.TryTransactionSet(ts)
	if err != nil {
		return err
	}

	// Add the transaction set to the pool.
	setID := TransactionSetID(crypto.HashObject(ts))
	tp.transactionSets[setID] = ts
	for _, oid := range oids {
		tp.knownObjects[oid] = setID
	}
	tp.transactionSetDiffs[setID] = cc
	tp.transactionListSize += len(encoding.Marshal(ts))
	return nil
}

// AcceptTransaction adds a transaction to the unconfirmed set of
// transactions. If the transaction is accepted, it will be relayed to
// connected peers.
func (tp *TransactionPool) AcceptTransactionSet(ts []types.Transaction) error {
	tp.mu.Lock()
	defer tp.mu.Unlock()

	err := tp.acceptTransactionSet(ts)
	if err != nil {
		return err
	}

	// Notify subscribers and broadcast the transaction set.
	go tp.gateway.Broadcast("RelayTransactionSet", ts)
	tp.updateSubscribersTransactions()

	// COMPAT v0.4.6
	//
	// Transactions must be broadcast individually as well so that they will be
	// seen by older nodes that don't support the "RelayTransactionSet" RPC.
	go func() {
		for _, t := range ts {
			tp.gateway.Broadcast("RelayTransaction", t)
			time.Sleep(time.Second * 15)
		}
	}()

	return nil
}

// relayTransactionSet is an RPC that accepts a transaction set from a peer. If
// the accept is successful, the transaction will be relayed to the gateway's
// other peers.
func (tp *TransactionPool) relayTransactionSet(conn modules.PeerConn) error {
	var ts []types.Transaction
	err := encoding.ReadObject(conn, &ts, types.BlockSizeLimit)
	if err != nil {
		return err
	}
	return tp.AcceptTransactionSet(ts)
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	"github.com/openshift/origin/test/extended/util"
)

var _ = g.Describe("[cli] oc adm must-gather", func() {
	defer g.GinkgoRecover()
	oc := util.NewCLI("oc-adm-must-gather", util.KubeConfigPath()).AsAdmin()
	g.It("runs successfully", func() {
		tempDir, err := ioutil.TempDir("", "test.oc-adm-must-gather.")
		o.Expect(err).ToNot(o.HaveOccurred())
		defer os.RemoveAll(tempDir)
		o.Expect(oc.Run("adm", "must-gather").Args("--dest-dir", tempDir).Execute()).To(o.Succeed())

		expectedDirectories := [][]string{
			{tempDir, "cluster-scoped-resources", "config.openshift.io"},
			{tempDir, "cluster-scoped-resources", "operator.openshift.io"},
			{tempDir, "cluster-scoped-resources", "core"},
			{tempDir, "cluster-scoped-resources", "apiregistration.k8s.io"},
			{tempDir, "namespaces", "openshift"},
			{tempDir, "namespaces", "openshift-kube-apiserver-operator"},
		}

		expectedFiles := [][]string{
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "apiservers.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "authentications.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "builds.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "clusteroperators.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "clusterversions.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "consoles.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "dnses.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "featuregates.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "images.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "infrastructures.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "ingresses.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "networks.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "oauths.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "projects.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "schedulers.yaml"},
			{tempDir, "namespaces", "openshift-kube-apiserver", "core", "configmaps.yaml"},
			{tempDir, "namespaces", "openshift-kube-apiserver", "core", "secrets.yaml"},
			{tempDir, "audit_logs", "kube-apiserver.audit_logs_listing"},
			{tempDir, "audit_logs", "openshift-apiserver.audit_logs_listing"},
			{tempDir, "host_service_logs", "masters", "crio_service.log"},
			{tempDir, "host_service_logs", "masters", "kubelet_service.log"},
		}

		for _, expectedDirectory := range expectedDirectories {
			o.Expect(path.Join(expectedDirectory...)).To(o.BeADirectory())
		}

		emptyFiles := []string{}
		for _, expectedFile := range expectedFiles {
			expectedFilePath := path.Join(expectedFile...)
			o.Expect(expectedFilePath).To(o.BeAnExistingFile())
			stat, err := os.Stat(expectedFilePath)
			o.Expect(err).ToNot(o.HaveOccurred())
			if size := stat.Size(); size < 100 {
				emptyFiles = append(emptyFiles, expectedFilePath)
			}
		}
		if len(emptyFiles) > 0 {
			o.Expect(fmt.Errorf("expected files should not be empty: %s", strings.Join(emptyFiles, ","))).NotTo(o.HaveOccurred())
		}

	})

	g.It("runs successfully with options", func() {
		tempDir, err := ioutil.TempDir("", "test.oc-adm-must-gather.")
		o.Expect(err).ToNot(o.HaveOccurred())
		defer os.RemoveAll(tempDir)
		args := []string{
			"--dest-dir", tempDir,
			"--source-dir", "/artifacts",
			"--",
			"/bin/bash", "-c",
			"ls -l > /artifacts/ls.log",
		}
		o.Expect(oc.Run("adm", "must-gather").Args(args...).Execute()).To(o.Succeed())
		expectedFilePath := path.Join(tempDir, "ls.log")
		o.Expect(expectedFilePath).To(o.BeAnExistingFile())
		stat, err := os.Stat(expectedFilePath)
		o.Expect(err).ToNot(o.HaveOccurred())
		o.Expect(stat.Size()).To(o.BeNumerically(">", 0))
	})
})
<commit_msg>test: reduce minimum size of expected files in must-gather tests<commit_after>package cli

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	"github.com/openshift/origin/test/extended/util"
)

var _ = g.Describe("[cli] oc adm must-gather", func() {
	defer g.GinkgoRecover()
	oc := util.NewCLI("oc-adm-must-gather", util.KubeConfigPath()).AsAdmin()
	g.It("runs successfully", func() {
		tempDir, err := ioutil.TempDir("", "test.oc-adm-must-gather.")
		o.Expect(err).ToNot(o.HaveOccurred())
		defer os.RemoveAll(tempDir)
		o.Expect(oc.Run("adm", "must-gather").Args("--dest-dir", tempDir).Execute()).To(o.Succeed())

		expectedDirectories := [][]string{
			{tempDir, "cluster-scoped-resources", "config.openshift.io"},
			{tempDir, "cluster-scoped-resources", "operator.openshift.io"},
			{tempDir, "cluster-scoped-resources", "core"},
			{tempDir, "cluster-scoped-resources", "apiregistration.k8s.io"},
			{tempDir, "namespaces", "openshift"},
			{tempDir, "namespaces", "openshift-kube-apiserver-operator"},
		}

		expectedFiles := [][]string{
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "apiservers.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "authentications.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "builds.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "clusteroperators.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "clusterversions.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "consoles.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "dnses.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "featuregates.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "images.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "infrastructures.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "ingresses.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "networks.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "oauths.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "projects.yaml"},
			{tempDir, "cluster-scoped-resources", "config.openshift.io", "schedulers.yaml"},
			{tempDir, "namespaces", "openshift-kube-apiserver", "core", "configmaps.yaml"},
			{tempDir, "namespaces", "openshift-kube-apiserver", "core", "secrets.yaml"},
			{tempDir, "audit_logs", "kube-apiserver.audit_logs_listing"},
			{tempDir, "audit_logs", "openshift-apiserver.audit_logs_listing"},
			{tempDir, "host_service_logs", "masters", "crio_service.log"},
			{tempDir, "host_service_logs", "masters", "kubelet_service.log"},
		}

		for _, expectedDirectory := range expectedDirectories {
			o.Expect(path.Join(expectedDirectory...)).To(o.BeADirectory())
		}

		emptyFiles := []string{}
		for _, expectedFile := range expectedFiles {
			expectedFilePath := path.Join(expectedFile...)
			o.Expect(expectedFilePath).To(o.BeAnExistingFile())
			stat, err := os.Stat(expectedFilePath)
			o.Expect(err).ToNot(o.HaveOccurred())
			if size := stat.Size(); size < 50 {
				emptyFiles = append(emptyFiles, expectedFilePath)
			}
		}
		if len(emptyFiles) > 0 {
			o.Expect(fmt.Errorf("expected files should not be empty: %s", strings.Join(emptyFiles, ","))).NotTo(o.HaveOccurred())
		}

	})

	g.It("runs successfully with options", func() {
		tempDir, err := ioutil.TempDir("", "test.oc-adm-must-gather.")
		o.Expect(err).ToNot(o.HaveOccurred())
		defer os.RemoveAll(tempDir)
		args := []string{
			"--dest-dir", tempDir,
			"--source-dir", "/artifacts",
			"--",
			"/bin/bash", "-c",
			"ls -l > /artifacts/ls.log",
		}
		o.Expect(oc.Run("adm", "must-gather").Args(args...).Execute()).To(o.Succeed())
		expectedFilePath := path.Join(tempDir, "ls.log")
		o.Expect(expectedFilePath).To(o.BeAnExistingFile())
		stat, err := os.Stat(expectedFilePath)
		o.Expect(err).ToNot(o.HaveOccurred())
		o.Expect(stat.Size()).To(o.BeNumerically(">", 0))
	})
})
<|endoftext|>


<commit_before>package namesgenerator

import (
	"fmt"
	"math/rand"
	"time"
)

type NameChecker interface {
	Exists(name string) bool
}

var (
	colors  = [...]string{"white", "silver", "gray", "black", "blue", "green", "cyan", "yellow", "gold", "orange", "brown", "red", "violet", "pink", "magenta", "purple"}
	animals = [...]string{"ant", "bird", "cat", "chicken", "cow", "dog", "fish", "fox", "horse", "lion", "monkey", "pig", "sheep", "tiger", "whale", "wolf"}
)

func GenerateRandomName(checker NameChecker) (string, error) {
	retry := 5
	rand.Seed(time.Now().UnixNano())
	name := fmt.Sprintf("%s_%s", colors[rand.Intn(len(colors))], animals[rand.Intn(len(animals))])
	for checker != nil && checker.Exists(name) && retry > 0 {
		name = fmt.Sprintf("%s%d", name, rand.Intn(10))
		retry = retry - 1
	}
	if retry == 0 {
		return name, fmt.Errorf("Error generating random name")
	}
	return name, nil
}
<commit_msg>Added some more colors and animals to increase the pool<commit_after>package namesgenerator

import (
	"fmt"
	"math/rand"
	"time"
)

type NameChecker interface {
	Exists(name string) bool
}

var (
	colors  = [...]string{"white", "silver", "gray", "black", "blue", "green", "cyan", "yellow", "gold", "orange", "brown", "red", "violet", "pink", "magenta", "purple", "maroon", "crimson", "plum", "fuchsia", "lavender", "slate", "navy", "azure", "aqua", "olive", "teal", "lime", "beige", "tan", "sienna"}
  animals = [...]string{"ant", "bear", "bird", "cat", "chicken", "cow", "deer", "dog", "donkey", "duck", "fish", "fox", "frog", "horse", "kangaroo", "koala", "lemur", "lion", "lizard", "monkey", "octopus", "pig", "shark", "sheep", "sloth", "spider", "squirrel", "tiger", "toad", "weasel", "whale", "wolf"}
)

func GenerateRandomName(checker NameChecker) (string, error) {
	retry := 5
	rand.Seed(time.Now().UnixNano())
	name := fmt.Sprintf("%s_%s", colors[rand.Intn(len(colors))], animals[rand.Intn(len(animals))])
	for checker != nil && checker.Exists(name) && retry > 0 {
		name = fmt.Sprintf("%s%d", name, rand.Intn(10))
		retry = retry - 1
	}
	if retry == 0 {
		return name, fmt.Errorf("Error generating random name")
	}
	return name, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package ghttp

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gogf/gf/encoding/gparser"
	"github.com/gogf/gf/text/gregex"
	"github.com/gogf/gf/text/gstr"
	"github.com/gogf/gf/util/gconv"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gogf/gf/os/gfile"
)

// Get send GET request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Get(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("GET", url, data...)
}

// Put send PUT request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Put(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("PUT", url, data...)
}

// Post sends request using HTTP method POST and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Post(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("POST", url, data...)
}

// Delete send DELETE request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Delete(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("DELETE", url, data...)
}

// Head send HEAD request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Head(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("HEAD", url, data...)
}

// Patch send PATCH request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Patch(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("PATCH", url, data...)
}

// Connect send CONNECT request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Connect(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("CONNECT", url, data...)
}

// Options send OPTIONS request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Options(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("OPTIONS", url, data...)
}

// Trace send TRACE request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Trace(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("TRACE", url, data...)
}

// DoRequest sends request with given HTTP method and data and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
//
// Note that it uses "multipart/form-data" as its Content-Type if it contains file uploading,
// else it uses "application/x-www-form-urlencoded". It also automatically detects the post
// content for JSON format, and for that it automatically sets the Content-Type as
// "application/json".
func (c *Client) DoRequest(method, url string, data ...interface{}) (resp *ClientResponse, err error) {
	method = strings.ToUpper(method)
	if len(c.prefix) > 0 {
		url = c.prefix + gstr.Trim(url)
	}
	param := ""
	if len(data) > 0 {
		switch c.header["Content-Type"] {
		case "application/json":
			switch data[0].(type) {
			case string, []byte:
				param = gconv.String(data[0])
			default:
				if b, err := json.Marshal(data[0]); err != nil {
					return nil, err
				} else {
					param = gconv.UnsafeBytesToStr(b)
				}
			}
		case "application/xml":
			switch data[0].(type) {
			case string, []byte:
				param = gconv.String(data[0])
			default:
				if b, err := gparser.VarToXml(data[0]); err != nil {
					return nil, err
				} else {
					param = gconv.UnsafeBytesToStr(b)
				}
			}
		default:
			param = BuildParams(data[0])
		}
	}
	req := (*http.Request)(nil)
	if strings.Contains(param, "@file:") {
		// File uploading request.
		buffer := new(bytes.Buffer)
		writer := multipart.NewWriter(buffer)
		for _, item := range strings.Split(param, "&") {
			array := strings.Split(item, "=")
			if len(array[1]) > 6 && strings.Compare(array[1][0:6], "@file:") == 0 {
				path := array[1][6:]
				if !gfile.Exists(path) {
					return nil, errors.New(fmt.Sprintf(`"%s" does not exist`, path))
				}
				if file, err := writer.CreateFormFile(array[0], path); err == nil {
					if f, err := os.Open(path); err == nil {
						if _, err = io.Copy(file, f); err != nil {
							f.Close()
							return nil, err
						}
						f.Close()
					} else {
						return nil, err
					}
				} else {
					return nil, err
				}
			} else {
				if err = writer.WriteField(array[0], array[1]); err != nil {
					return nil, err
				}
			}
		}
		// Close finishes the multipart message and writes the trailing
		// boundary end line to the output.
		if err = writer.Close(); err != nil {
			return nil, err
		}
		if req, err = http.NewRequestWithContext(c.ctx, method, url, buffer); err != nil {
			return nil, err
		} else {
			req.Header.Set("Content-Type", writer.FormDataContentType())
		}
	} else {
		// Normal request.
		paramBytes := []byte(param)
		if req, err = http.NewRequestWithContext(
			c.ctx, method, url, bytes.NewReader(paramBytes),
		); err != nil {
			return nil, err
		} else {
			if v, ok := c.header["Content-Type"]; ok {
				// Custom Content-Type.
				req.Header.Set("Content-Type", v)
			} else if len(paramBytes) > 0 {
				if (paramBytes[0] == '[' || paramBytes[0] == '{') && json.Valid(paramBytes) {
					// Auto detecting and setting the post content format: JSON.
					req.Header.Set("Content-Type", "application/json")
				} else if gregex.IsMatchString(`^[\w\[\]]+=.+`, param) {
					// If the parameters passed like "name=value", it then uses form type.
					req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
				}
			}
		}
	}
	// Custom header.
	if len(c.header) > 0 {
		for k, v := range c.header {
			req.Header.Set(k, v)
		}
	}
	// Custom Cookie.
	if len(c.cookies) > 0 {
		headerCookie := ""
		for k, v := range c.cookies {
			if len(headerCookie) > 0 {
				headerCookie += ";"
			}
			headerCookie += k + "=" + v
		}
		if len(headerCookie) > 0 {
			req.Header.Set("Cookie", headerCookie)
		}
	}
	// HTTP basic authentication.
	if len(c.authUser) > 0 {
		req.SetBasicAuth(c.authUser, c.authPass)
	}
	// Sending request.
	var r *http.Response
	for {
		if r, err = c.Do(req); err != nil {
			if c.retryCount > 0 {
				c.retryCount--
			} else {
				return nil, err
			}
		} else {
			break
		}
	}
	resp = &ClientResponse{
		Response: r,
	}
	// Auto saving cookie content.
	if c.browserMode {
		now := time.Now()
		for _, v := range r.Cookies() {
			if v.Expires.UnixNano() < now.UnixNano() {
				delete(c.cookies, v.Name)
			} else {
				c.cookies[v.Name] = v.Value
			}
		}
	}
	return resp, nil
}
<commit_msg>Get ghttp a chance to add custom host.<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package ghttp

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gogf/gf/encoding/gparser"
	"github.com/gogf/gf/text/gregex"
	"github.com/gogf/gf/text/gstr"
	"github.com/gogf/gf/util/gconv"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gogf/gf/os/gfile"
)

// Get send GET request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Get(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("GET", url, data...)
}

// Put send PUT request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Put(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("PUT", url, data...)
}

// Post sends request using HTTP method POST and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Post(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("POST", url, data...)
}

// Delete send DELETE request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Delete(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("DELETE", url, data...)
}

// Head send HEAD request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Head(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("HEAD", url, data...)
}

// Patch send PATCH request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Patch(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("PATCH", url, data...)
}

// Connect send CONNECT request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Connect(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("CONNECT", url, data...)
}

// Options send OPTIONS request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Options(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("OPTIONS", url, data...)
}

// Trace send TRACE request and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
func (c *Client) Trace(url string, data ...interface{}) (*ClientResponse, error) {
	return c.DoRequest("TRACE", url, data...)
}

// DoRequest sends request with given HTTP method and data and returns the response object.
// Note that the response object MUST be closed if it'll be never used.
//
// Note that it uses "multipart/form-data" as its Content-Type if it contains file uploading,
// else it uses "application/x-www-form-urlencoded". It also automatically detects the post
// content for JSON format, and for that it automatically sets the Content-Type as
// "application/json".
func (c *Client) DoRequest(method, url string, data ...interface{}) (resp *ClientResponse, err error) {
	method = strings.ToUpper(method)
	if len(c.prefix) > 0 {
		url = c.prefix + gstr.Trim(url)
	}
	param := ""
	if len(data) > 0 {
		switch c.header["Content-Type"] {
		case "application/json":
			switch data[0].(type) {
			case string, []byte:
				param = gconv.String(data[0])
			default:
				if b, err := json.Marshal(data[0]); err != nil {
					return nil, err
				} else {
					param = gconv.UnsafeBytesToStr(b)
				}
			}
		case "application/xml":
			switch data[0].(type) {
			case string, []byte:
				param = gconv.String(data[0])
			default:
				if b, err := gparser.VarToXml(data[0]); err != nil {
					return nil, err
				} else {
					param = gconv.UnsafeBytesToStr(b)
				}
			}
		default:
			param = BuildParams(data[0])
		}
	}
	req := (*http.Request)(nil)
	if strings.Contains(param, "@file:") {
		// File uploading request.
		buffer := new(bytes.Buffer)
		writer := multipart.NewWriter(buffer)
		for _, item := range strings.Split(param, "&") {
			array := strings.Split(item, "=")
			if len(array[1]) > 6 && strings.Compare(array[1][0:6], "@file:") == 0 {
				path := array[1][6:]
				if !gfile.Exists(path) {
					return nil, errors.New(fmt.Sprintf(`"%s" does not exist`, path))
				}
				if file, err := writer.CreateFormFile(array[0], path); err == nil {
					if f, err := os.Open(path); err == nil {
						if _, err = io.Copy(file, f); err != nil {
							f.Close()
							return nil, err
						}
						f.Close()
					} else {
						return nil, err
					}
				} else {
					return nil, err
				}
			} else {
				if err = writer.WriteField(array[0], array[1]); err != nil {
					return nil, err
				}
			}
		}
		// Close finishes the multipart message and writes the trailing
		// boundary end line to the output.
		if err = writer.Close(); err != nil {
			return nil, err
		}
		if req, err = http.NewRequestWithContext(c.ctx, method, url, buffer); err != nil {
			return nil, err
		} else {
			req.Header.Set("Content-Type", writer.FormDataContentType())
		}
	} else {
		// Normal request.
		paramBytes := []byte(param)
		if req, err = http.NewRequestWithContext(
			c.ctx, method, url, bytes.NewReader(paramBytes),
		); err != nil {
			return nil, err
		} else {
			if v, ok := c.header["Content-Type"]; ok {
				// Custom Content-Type.
				req.Header.Set("Content-Type", v)
			} else if len(paramBytes) > 0 {
				if (paramBytes[0] == '[' || paramBytes[0] == '{') && json.Valid(paramBytes) {
					// Auto detecting and setting the post content format: JSON.
					req.Header.Set("Content-Type", "application/json")
				} else if gregex.IsMatchString(`^[\w\[\]]+=.+`, param) {
					// If the parameters passed like "name=value", it then uses form type.
					req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
				}
			}
		}
	}
	// Custom header.
	if len(c.header) > 0 {
		for k, v := range c.header {
			// Set host by req.Host.
			if strings.ToLower(k) == "host" {
				req.Host = v
			}
			req.Header.Set(k, v)
		}
	}
	// Custom Cookie.
	if len(c.cookies) > 0 {
		headerCookie := ""
		for k, v := range c.cookies {
			if len(headerCookie) > 0 {
				headerCookie += ";"
			}
			headerCookie += k + "=" + v
		}
		if len(headerCookie) > 0 {
			req.Header.Set("Cookie", headerCookie)
		}
	}
	// HTTP basic authentication.
	if len(c.authUser) > 0 {
		req.SetBasicAuth(c.authUser, c.authPass)
	}
	// Sending request.
	var r *http.Response
	for {
		if r, err = c.Do(req); err != nil {
			if c.retryCount > 0 {
				c.retryCount--
			} else {
				return nil, err
			}
		} else {
			break
		}
	}
	resp = &ClientResponse{
		Response: r,
	}
	// Auto saving cookie content.
	if c.browserMode {
		now := time.Now()
		for _, v := range r.Cookies() {
			if v.Expires.UnixNano() < now.UnixNano() {
				delete(c.cookies, v.Name)
			} else {
				c.cookies[v.Name] = v.Value
			}
		}
	}
	return resp, nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package networkserver

import (
	"fmt"
	"time"

	"github.com/TheThingsNetwork/go-account-lib/rights"
	pb "github.com/TheThingsNetwork/ttn/api/networkserver"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	"github.com/TheThingsNetwork/ttn/core/networkserver/device"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/golang/protobuf/ptypes/empty"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

type networkServerManager struct {
	networkServer *networkServer
}

func (n *networkServerManager) getDevice(ctx context.Context, in *pb_lorawan.DeviceIdentifier) (*device.Device, error) {
	if !in.Validate() {
		return nil, errors.NewErrInvalidArgument("Device Identifier", "validation failed")
	}
	claims, err := n.networkServer.Component.ValidateTTNAuthContext(ctx)
	if err != nil {
		return nil, err
	}
	dev, err := n.networkServer.devices.Get(*in.AppEui, *in.DevEui)
	if err != nil {
		return nil, err
	}
	if !claims.AppRight(dev.AppID, rights.AppSettings) {
		return nil, errors.NewErrPermissionDenied(fmt.Sprintf("No access to Application %s", dev.AppID))
	}
	return dev, nil
}

func (n *networkServerManager) GetDevice(ctx context.Context, in *pb_lorawan.DeviceIdentifier) (*pb_lorawan.Device, error) {
	dev, err := n.getDevice(ctx, in)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}

	lastSeen := time.Unix(0, 0)
	if !dev.LastSeen.IsZero() {
		lastSeen = dev.LastSeen
	}

	return &pb_lorawan.Device{
		AppId:            dev.AppID,
		AppEui:           &dev.AppEUI,
		DevId:            dev.DevID,
		DevEui:           &dev.DevEUI,
		DevAddr:          &dev.DevAddr,
		NwkSKey:          &dev.NwkSKey,
		FCntUp:           dev.FCntUp,
		FCntDown:         dev.FCntDown,
		DisableFCntCheck: dev.Options.DisableFCntCheck,
		Uses32BitFCnt:    dev.Options.Uses32BitFCnt,
		LastSeen:         lastSeen.UnixNano(),
	}, nil
}

func (n *networkServerManager) SetDevice(ctx context.Context, in *pb_lorawan.Device) (*empty.Empty, error) {
	dev, err := n.getDevice(ctx, &pb_lorawan.DeviceIdentifier{AppEui: in.AppEui, DevEui: in.DevEui})
	if err != nil && errors.GetErrType(err) != errors.NotFound {
		return nil, errors.BuildGRPCError(err)
	}

	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Device Identifier")
	}

	if dev == nil {
		dev = new(device.Device)
	}
	fields := []string{}

	if dev.AppID != in.AppId {
		dev.AppID = in.AppId
		fields = append(fields, "app_id")
	}

	if dev.AppEUI != *in.AppEui {
		dev.AppEUI = *in.AppEui
		fields = append(fields, "app_eui")
	}

	if dev.DevID != in.DevId {
		dev.DevID = in.DevId
		fields = append(fields, "dev_id")
	}

	if dev.DevEUI != *in.DevEui {
		dev.DevEUI = *in.DevEui
		fields = append(fields, "dev_eui")
	}

	if dev.FCntUp != in.FCntUp {
		dev.FCntUp = in.FCntUp
		fields = append(fields, "f_cnt_up")
	}

	if dev.FCntDown != in.FCntDown {
		dev.FCntDown = in.FCntDown
		fields = append(fields, "f_cnt_down")
	}

	if dev.Options.DisableFCntCheck != in.DisableFCntCheck || dev.Options.Uses32BitFCnt != in.Uses32BitFCnt || dev.Options.ActivationConstraints != in.ActivationConstraints {
		dev.Options = device.Options{
			DisableFCntCheck:      in.DisableFCntCheck,
			Uses32BitFCnt:         in.Uses32BitFCnt,
			ActivationConstraints: in.ActivationConstraints,
		}
		fields = append(fields, "options")
	}

	if in.NwkSKey != nil && in.DevAddr != nil {
		dev.DevAddr = *in.DevAddr
		dev.NwkSKey = *in.NwkSKey
		fields = append(fields, "dev_addr", "nwk_s_key")
	}

	err = n.networkServer.devices.Set(dev, fields...)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}

	return &empty.Empty{}, nil
}

func (n *networkServerManager) DeleteDevice(ctx context.Context, in *pb_lorawan.DeviceIdentifier) (*empty.Empty, error) {
	_, err := n.getDevice(ctx, in)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}
	err = n.networkServer.devices.Delete(*in.AppEui, *in.DevEui)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}
	return &empty.Empty{}, nil
}

func (n *networkServerManager) GetPrefixes(ctx context.Context, in *pb_lorawan.PrefixesRequest) (*pb_lorawan.PrefixesResponse, error) {
	var mapping []*pb_lorawan.PrefixesResponse_PrefixMapping
	for prefix, usage := range n.networkServer.prefixes {
		mapping = append(mapping, &pb_lorawan.PrefixesResponse_PrefixMapping{
			Prefix: prefix.String(),
			Usage:  usage,
		})
	}
	return &pb_lorawan.PrefixesResponse{
		Prefixes: mapping,
	}, nil
}

func (n *networkServerManager) GetDevAddr(ctx context.Context, in *pb_lorawan.DevAddrRequest) (*pb_lorawan.DevAddrResponse, error) {
	devAddr, err := n.networkServer.getDevAddr(in.Usage...)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}
	return &pb_lorawan.DevAddrResponse{
		DevAddr: &devAddr,
	}, nil
}

func (n *networkServerManager) GetStatus(ctx context.Context, in *pb.StatusRequest) (*pb.Status, error) {
	return nil, grpcErrf(codes.Unimplemented, "Not Implemented")
}

// RegisterManager registers this networkserver as a NetworkServerManagerServer (github.com/TheThingsNetwork/ttn/api/networkserver)
func (n *networkServer) RegisterManager(s *grpc.Server) {
	server := &networkServerManager{n}
	pb.RegisterNetworkServerManagerServer(s, server)
	pb_lorawan.RegisterDeviceManagerServer(s, server)
	pb_lorawan.RegisterDevAddrManagerServer(s, server)
}
<commit_msg>Also check for rights to new device in networkserver manager<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package networkserver

import (
	"fmt"
	"time"

	"github.com/TheThingsNetwork/go-account-lib/rights"
	pb "github.com/TheThingsNetwork/ttn/api/networkserver"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	"github.com/TheThingsNetwork/ttn/core/networkserver/device"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/golang/protobuf/ptypes/empty"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

type networkServerManager struct {
	networkServer *networkServer
}

func (n *networkServerManager) getDevice(ctx context.Context, in *pb_lorawan.DeviceIdentifier) (*device.Device, error) {
	if !in.Validate() {
		return nil, errors.NewErrInvalidArgument("Device Identifier", "validation failed")
	}
	claims, err := n.networkServer.Component.ValidateTTNAuthContext(ctx)
	if err != nil {
		return nil, err
	}
	dev, err := n.networkServer.devices.Get(*in.AppEui, *in.DevEui)
	if err != nil {
		return nil, err
	}
	if !claims.AppRight(dev.AppID, rights.AppSettings) {
		return nil, errors.NewErrPermissionDenied(fmt.Sprintf("No access to Application %s", dev.AppID))
	}
	return dev, nil
}

func (n *networkServerManager) GetDevice(ctx context.Context, in *pb_lorawan.DeviceIdentifier) (*pb_lorawan.Device, error) {
	dev, err := n.getDevice(ctx, in)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}

	lastSeen := time.Unix(0, 0)
	if !dev.LastSeen.IsZero() {
		lastSeen = dev.LastSeen
	}

	return &pb_lorawan.Device{
		AppId:            dev.AppID,
		AppEui:           &dev.AppEUI,
		DevId:            dev.DevID,
		DevEui:           &dev.DevEUI,
		DevAddr:          &dev.DevAddr,
		NwkSKey:          &dev.NwkSKey,
		FCntUp:           dev.FCntUp,
		FCntDown:         dev.FCntDown,
		DisableFCntCheck: dev.Options.DisableFCntCheck,
		Uses32BitFCnt:    dev.Options.Uses32BitFCnt,
		LastSeen:         lastSeen.UnixNano(),
	}, nil
}

func (n *networkServerManager) SetDevice(ctx context.Context, in *pb_lorawan.Device) (*empty.Empty, error) {
	dev, err := n.getDevice(ctx, &pb_lorawan.DeviceIdentifier{AppEui: in.AppEui, DevEui: in.DevEui})
	if err != nil && errors.GetErrType(err) != errors.NotFound {
		return nil, errors.BuildGRPCError(err)
	}

	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Device")
	}

	claims, err := n.networkServer.Component.ValidateTTNAuthContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.AppRight(in.AppId, rights.AppSettings) {
		return nil, errors.NewErrPermissionDenied(fmt.Sprintf("No access to Application %s", dev.AppID))
	}

	if dev == nil {
		dev = new(device.Device)
	}
	fields := []string{}

	if dev.AppID != in.AppId {
		dev.AppID = in.AppId
		fields = append(fields, "app_id")
	}

	if dev.AppEUI != *in.AppEui {
		dev.AppEUI = *in.AppEui
		fields = append(fields, "app_eui")
	}

	if dev.DevID != in.DevId {
		dev.DevID = in.DevId
		fields = append(fields, "dev_id")
	}

	if dev.DevEUI != *in.DevEui {
		dev.DevEUI = *in.DevEui
		fields = append(fields, "dev_eui")
	}

	if dev.FCntUp != in.FCntUp {
		dev.FCntUp = in.FCntUp
		fields = append(fields, "f_cnt_up")
	}

	if dev.FCntDown != in.FCntDown {
		dev.FCntDown = in.FCntDown
		fields = append(fields, "f_cnt_down")
	}

	if dev.Options.DisableFCntCheck != in.DisableFCntCheck || dev.Options.Uses32BitFCnt != in.Uses32BitFCnt || dev.Options.ActivationConstraints != in.ActivationConstraints {
		dev.Options = device.Options{
			DisableFCntCheck:      in.DisableFCntCheck,
			Uses32BitFCnt:         in.Uses32BitFCnt,
			ActivationConstraints: in.ActivationConstraints,
		}
		fields = append(fields, "options")
	}

	if in.NwkSKey != nil && in.DevAddr != nil {
		dev.DevAddr = *in.DevAddr
		dev.NwkSKey = *in.NwkSKey
		fields = append(fields, "dev_addr", "nwk_s_key")
	}

	err = n.networkServer.devices.Set(dev, fields...)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}

	return &empty.Empty{}, nil
}

func (n *networkServerManager) DeleteDevice(ctx context.Context, in *pb_lorawan.DeviceIdentifier) (*empty.Empty, error) {
	_, err := n.getDevice(ctx, in)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}
	err = n.networkServer.devices.Delete(*in.AppEui, *in.DevEui)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}
	return &empty.Empty{}, nil
}

func (n *networkServerManager) GetPrefixes(ctx context.Context, in *pb_lorawan.PrefixesRequest) (*pb_lorawan.PrefixesResponse, error) {
	var mapping []*pb_lorawan.PrefixesResponse_PrefixMapping
	for prefix, usage := range n.networkServer.prefixes {
		mapping = append(mapping, &pb_lorawan.PrefixesResponse_PrefixMapping{
			Prefix: prefix.String(),
			Usage:  usage,
		})
	}
	return &pb_lorawan.PrefixesResponse{
		Prefixes: mapping,
	}, nil
}

func (n *networkServerManager) GetDevAddr(ctx context.Context, in *pb_lorawan.DevAddrRequest) (*pb_lorawan.DevAddrResponse, error) {
	devAddr, err := n.networkServer.getDevAddr(in.Usage...)
	if err != nil {
		return nil, errors.BuildGRPCError(err)
	}
	return &pb_lorawan.DevAddrResponse{
		DevAddr: &devAddr,
	}, nil
}

func (n *networkServerManager) GetStatus(ctx context.Context, in *pb.StatusRequest) (*pb.Status, error) {
	return nil, grpcErrf(codes.Unimplemented, "Not Implemented")
}

// RegisterManager registers this networkserver as a NetworkServerManagerServer (github.com/TheThingsNetwork/ttn/api/networkserver)
func (n *networkServer) RegisterManager(s *grpc.Server) {
	server := &networkServerManager{n}
	pb.RegisterNetworkServerManagerServer(s, server)
	pb_lorawan.RegisterDeviceManagerServer(s, server)
	pb_lorawan.RegisterDevAddrManagerServer(s, server)
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package local

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/stats"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	null "gopkg.in/guregu/null.v3"
)

var _ lib.Executor = &Executor{}

type vuHandle struct {
	sync.RWMutex
	vu     lib.VU
	ctx    context.Context
	cancel context.CancelFunc
}

func (h *vuHandle) run(logger *log.Logger, flow <-chan int64, out chan<- []stats.Sample) {
	h.RLock()
	ctx := h.ctx
	h.RUnlock()

	for {
		select {
		case _, ok := <-flow:
			if !ok {
				return
			}
		case <-ctx.Done():
			return
		}

		var samples []stats.Sample
		if h.vu != nil {
			s, err := h.vu.RunOnce(ctx)
			if err != nil {
				if s, ok := err.(fmt.Stringer); ok {
					logger.Error(s.String())
				} else {
					logger.Error(err.Error())
				}
			}
			samples = s
		}
		out <- samples
	}
}

type Executor struct {
	Runner lib.Runner
	Logger *log.Logger

	runLock sync.Mutex
	wg      sync.WaitGroup

	vus       []*vuHandle
	vusLock   sync.RWMutex
	numVUs    int64
	numVUsMax int64
	nextVUID  int64

	iters     int64 // Completed iterations
	partIters int64 // Partial, incomplete iterations
	endIters  int64 // End test at this many iterations

	time    int64 // Current time
	endTime int64 // End test at this timestamp

	pauseLock sync.RWMutex
	pause     chan interface{}

	// Lock for: ctx, flow, out
	lock sync.RWMutex

	// Current context, nil if a test isn't running right now.
	ctx context.Context

	// Engineward output channel for samples.
	out chan<- []stats.Sample

	// Flow control for VUs; iterations are run only after reading from this channel.
	flow chan int64
}

func New(r lib.Runner) *Executor {
	return &Executor{
		Runner:   r,
		Logger:   log.StandardLogger(),
		endIters: -1,
		endTime:  -1,
	}
}

func (e *Executor) Run(parent context.Context, out chan<- []stats.Sample) error {
	e.runLock.Lock()
	defer e.runLock.Unlock()

	ctx, cancel := context.WithCancel(parent)
	vuOut := make(chan []stats.Sample)
	vuFlow := make(chan int64)

	e.lock.Lock()
	e.ctx = ctx
	e.out = vuOut
	e.flow = vuFlow
	e.lock.Unlock()

	var cutoff time.Time
	defer func() {
		close(vuFlow)
		cancel()

		e.lock.Lock()
		e.ctx = nil
		e.out = nil
		e.flow = nil
		e.lock.Unlock()

		wait := make(chan interface{})
		go func() {
			e.wg.Wait()
			close(wait)
		}()

		var samples []stats.Sample
		for {
			select {
			case ss := <-vuOut:
				for _, s := range ss {
					if cutoff.IsZero() || s.Time.Before(cutoff) {
						samples = append(samples, s)
					}
				}
			case <-wait:
			}
			select {
			case <-wait:
				close(vuOut)
				if out != nil && len(samples) > 0 {
					out <- samples
				}
				return
			default:
			}
		}
	}()

	if err := e.scale(ctx, lib.Max(0, atomic.LoadInt64(&e.numVUs))); err != nil {
		return err
	}

	ticker := time.NewTicker(1 * time.Millisecond)
	defer ticker.Stop()

	lastTick := time.Now()
	for {
		// If the test is paused, sleep until either the pause or the test ends.
		// Also shift the last tick to omit time spent paused, but not partial ticks.
		e.pauseLock.RLock()
		pause := e.pause
		e.pauseLock.RUnlock()
		if pause != nil {
			e.Logger.Debug("Local: Pausing!")
			leftovers := time.Since(lastTick)
			select {
			case <-pause:
				e.Logger.Debug("Local: No longer paused")
				lastTick = time.Now().Add(-leftovers)
			case <-ctx.Done():
				e.Logger.Debug("Local: Terminated while in paused state")
				return nil
			}
		}

		// Dumb hack: we don't wanna start any more iterations than the max, but we can't
		// conditionally select on a channel either...so, we cheat: swap out the flow channel for a
		// nil channel (writing to nil always blocks) if we don't wanna write an iteration.
		flow := vuFlow
		end := atomic.LoadInt64(&e.endIters)
		partials := atomic.LoadInt64(&e.partIters)
		if end >= 0 && partials >= end {
			flow = nil
		}

		select {
		case flow <- partials:
			// Start an iteration if there's a VU waiting. See also: the big comment block above.
			atomic.AddInt64(&e.partIters, 1)
		case t := <-ticker.C:
			// Every tick, increment the clock and see if we passed the end point. If the test ends
			// this way, set a cutoff point; any samples collected past the cutoff point are excluded.
			d := t.Sub(lastTick)
			lastTick = t

			end := time.Duration(atomic.LoadInt64(&e.endTime))
			at := time.Duration(atomic.AddInt64(&e.time, int64(d)))
			if end >= 0 && at >= end {
				e.Logger.WithFields(log.Fields{"at": at, "end": end}).Debug("Local: Hit time limit")
				cutoff = time.Now()
				return nil
			}
		case samples := <-vuOut:
			// Every iteration ends with a write to vuOut. Check if we've hit the end point.
			if out != nil {
				out <- samples
			}

			end := atomic.LoadInt64(&e.endIters)
			at := atomic.AddInt64(&e.iters, 1)
			if end >= 0 && at >= end {
				e.Logger.WithFields(log.Fields{"at": at, "end": end}).Debug("Local: Hit iteration limit")
				return nil
			}
		case <-ctx.Done():
			// If the test is cancelled, just set the cutoff point to now and proceed down the same
			// logic as if the time limit was hit.
			e.Logger.Debug("Local: Exiting with context")
			cutoff = time.Now()
			return nil
		}
	}
}

func (e *Executor) scale(ctx context.Context, num int64) error {
	e.vusLock.Lock()
	defer e.vusLock.Unlock()

	e.lock.RLock()
	flow := e.flow
	out := e.out
	e.lock.RUnlock()

	for i, handle := range e.vus {
		handle := handle
		handle.RLock()
		cancel := handle.cancel
		handle.RUnlock()

		if i < int(num) {
			if cancel == nil {
				vuctx, cancel := context.WithCancel(ctx)
				handle.Lock()
				handle.ctx = vuctx
				handle.cancel = cancel
				handle.Unlock()

				if handle.vu != nil {
					if err := handle.vu.Reconfigure(atomic.AddInt64(&e.nextVUID, 1)); err != nil {
						return err
					}
				}

				e.wg.Add(1)
				go func() {
					handle.run(e.Logger, flow, out)
					e.wg.Done()
				}()
			}
		} else if cancel != nil {
			handle.Lock()
			handle.cancel()
			handle.cancel = nil
			handle.Unlock()
		}
	}

	atomic.StoreInt64(&e.numVUs, num)
	return nil
}

func (e *Executor) IsRunning() bool {
	e.lock.RLock()
	defer e.lock.RUnlock()
	return e.ctx != nil
}

func (e *Executor) GetRunner() lib.Runner {
	return e.Runner
}

func (e *Executor) SetLogger(l *log.Logger) {
	e.Logger = l
}

func (e *Executor) GetLogger() *log.Logger {
	return e.Logger
}

func (e *Executor) GetIterations() int64 {
	return atomic.LoadInt64(&e.iters)
}

func (e *Executor) GetEndIterations() null.Int {
	v := atomic.LoadInt64(&e.endIters)
	if v < 0 {
		return null.Int{}
	}
	return null.IntFrom(v)
}

func (e *Executor) SetEndIterations(i null.Int) {
	if !i.Valid {
		i.Int64 = -1
	}
	atomic.StoreInt64(&e.endIters, i.Int64)
}

func (e *Executor) GetTime() time.Duration {
	return time.Duration(atomic.LoadInt64(&e.time))
}

func (e *Executor) GetEndTime() lib.NullDuration {
	v := atomic.LoadInt64(&e.endTime)
	if v < 0 {
		return lib.NullDuration{}
	}
	return lib.NullDurationFrom(time.Duration(v))
}

func (e *Executor) SetEndTime(t lib.NullDuration) {
	if !t.Valid {
		t.Duration = -1
	}
	atomic.StoreInt64(&e.endTime, int64(t.Duration))
}

func (e *Executor) IsPaused() bool {
	e.pauseLock.RLock()
	defer e.pauseLock.RUnlock()
	return e.pause != nil
}

func (e *Executor) SetPaused(paused bool) {
	e.pauseLock.Lock()
	defer e.pauseLock.Unlock()

	if paused && e.pause == nil {
		e.pause = make(chan interface{})
	} else if !paused && e.pause != nil {
		close(e.pause)
		e.pause = nil
	}
}

func (e *Executor) GetVUs() int64 {
	return atomic.LoadInt64(&e.numVUs)
}

func (e *Executor) SetVUs(num int64) error {
	if num < 0 {
		return errors.New("vu count can't be negative")
	}

	if atomic.LoadInt64(&e.numVUs) == num {
		return nil
	}

	if numVUsMax := atomic.LoadInt64(&e.numVUsMax); num > numVUsMax {
		return errors.Errorf("can't raise vu count (to %d) above vu cap (%d)", num, numVUsMax)
	}

	if ctx := e.ctx; ctx != nil {
		if err := e.scale(ctx, num); err != nil {
			return err
		}
	} else {
		atomic.StoreInt64(&e.numVUs, num)
	}

	return nil
}

func (e *Executor) GetVUsMax() int64 {
	return atomic.LoadInt64(&e.numVUsMax)
}

func (e *Executor) SetVUsMax(max int64) error {
	if max < 0 {
		return errors.New("vu cap can't be negative")
	}

	numVUsMax := atomic.LoadInt64(&e.numVUsMax)

	if numVUsMax == max {
		return nil
	}

	if numVUs := atomic.LoadInt64(&e.numVUs); max < numVUs {
		return errors.Errorf("can't lower vu cap (to %d) below vu count (%d)", max, numVUs)
	}

	if max < numVUsMax {
		e.vus = e.vus[:max]
		atomic.StoreInt64(&e.numVUsMax, max)
		return nil
	}

	e.vusLock.Lock()
	defer e.vusLock.Unlock()

	vus := e.vus
	for i := numVUsMax; i < max; i++ {
		var handle vuHandle
		if e.Runner != nil {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			handle.vu = vu
		}
		vus = append(vus, &handle)
	}
	e.vus = vus

	atomic.StoreInt64(&e.numVUsMax, max)

	return nil
}
<commit_msg>More verbose logging in the local executor<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package local

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/stats"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	null "gopkg.in/guregu/null.v3"
)

var _ lib.Executor = &Executor{}

type vuHandle struct {
	sync.RWMutex
	vu     lib.VU
	ctx    context.Context
	cancel context.CancelFunc
}

func (h *vuHandle) run(logger *log.Logger, flow <-chan int64, out chan<- []stats.Sample) {
	h.RLock()
	ctx := h.ctx
	h.RUnlock()

	for {
		select {
		case _, ok := <-flow:
			if !ok {
				return
			}
		case <-ctx.Done():
			return
		}

		var samples []stats.Sample
		if h.vu != nil {
			s, err := h.vu.RunOnce(ctx)
			if err != nil {
				if s, ok := err.(fmt.Stringer); ok {
					logger.Error(s.String())
				} else {
					logger.Error(err.Error())
				}
			}
			samples = s
		}
		out <- samples
	}
}

type Executor struct {
	Runner lib.Runner
	Logger *log.Logger

	runLock sync.Mutex
	wg      sync.WaitGroup

	vus       []*vuHandle
	vusLock   sync.RWMutex
	numVUs    int64
	numVUsMax int64
	nextVUID  int64

	iters     int64 // Completed iterations
	partIters int64 // Partial, incomplete iterations
	endIters  int64 // End test at this many iterations

	time    int64 // Current time
	endTime int64 // End test at this timestamp

	pauseLock sync.RWMutex
	pause     chan interface{}

	// Lock for: ctx, flow, out
	lock sync.RWMutex

	// Current context, nil if a test isn't running right now.
	ctx context.Context

	// Engineward output channel for samples.
	out chan<- []stats.Sample

	// Flow control for VUs; iterations are run only after reading from this channel.
	flow chan int64
}

func New(r lib.Runner) *Executor {
	return &Executor{
		Runner:   r,
		Logger:   log.StandardLogger(),
		endIters: -1,
		endTime:  -1,
	}
}

func (e *Executor) Run(parent context.Context, out chan<- []stats.Sample) error {
	e.runLock.Lock()
	defer e.runLock.Unlock()

	ctx, cancel := context.WithCancel(parent)
	vuOut := make(chan []stats.Sample)
	vuFlow := make(chan int64)

	e.lock.Lock()
	e.ctx = ctx
	e.out = vuOut
	e.flow = vuFlow
	e.lock.Unlock()

	var cutoff time.Time
	defer func() {
		close(vuFlow)
		cancel()

		e.lock.Lock()
		e.ctx = nil
		e.out = nil
		e.flow = nil
		e.lock.Unlock()

		wait := make(chan interface{})
		go func() {
			e.wg.Wait()
			close(wait)
		}()

		var samples []stats.Sample
		for {
			select {
			case ss := <-vuOut:
				for _, s := range ss {
					if cutoff.IsZero() || s.Time.Before(cutoff) {
						samples = append(samples, s)
					}
				}
			case <-wait:
			}
			select {
			case <-wait:
				close(vuOut)
				if out != nil && len(samples) > 0 {
					out <- samples
				}
				return
			default:
			}
		}
	}()

	if err := e.scale(ctx, lib.Max(0, atomic.LoadInt64(&e.numVUs))); err != nil {
		return err
	}

	ticker := time.NewTicker(1 * time.Millisecond)
	defer ticker.Stop()

	lastTick := time.Now()
	for {
		// If the test is paused, sleep until either the pause or the test ends.
		// Also shift the last tick to omit time spent paused, but not partial ticks.
		e.pauseLock.RLock()
		pause := e.pause
		e.pauseLock.RUnlock()
		if pause != nil {
			e.Logger.Debug("Local: Pausing!")
			leftovers := time.Since(lastTick)
			select {
			case <-pause:
				e.Logger.Debug("Local: No longer paused")
				lastTick = time.Now().Add(-leftovers)
			case <-ctx.Done():
				e.Logger.Debug("Local: Terminated while in paused state")
				return nil
			}
		}

		// Dumb hack: we don't wanna start any more iterations than the max, but we can't
		// conditionally select on a channel either...so, we cheat: swap out the flow channel for a
		// nil channel (writing to nil always blocks) if we don't wanna write an iteration.
		flow := vuFlow
		end := atomic.LoadInt64(&e.endIters)
		partials := atomic.LoadInt64(&e.partIters)
		if end >= 0 && partials >= end {
			flow = nil
		}

		select {
		case flow <- partials:
			// Start an iteration if there's a VU waiting. See also: the big comment block above.
			atomic.AddInt64(&e.partIters, 1)
		case t := <-ticker.C:
			// Every tick, increment the clock and see if we passed the end point. If the test ends
			// this way, set a cutoff point; any samples collected past the cutoff point are excluded.
			d := t.Sub(lastTick)
			lastTick = t

			end := time.Duration(atomic.LoadInt64(&e.endTime))
			at := time.Duration(atomic.AddInt64(&e.time, int64(d)))
			if end >= 0 && at >= end {
				e.Logger.WithFields(log.Fields{"at": at, "end": end}).Debug("Local: Hit time limit")
				cutoff = time.Now()
				return nil
			}
		case samples := <-vuOut:
			// Every iteration ends with a write to vuOut. Check if we've hit the end point.
			if out != nil {
				out <- samples
			}

			end := atomic.LoadInt64(&e.endIters)
			at := atomic.AddInt64(&e.iters, 1)
			if end >= 0 && at >= end {
				e.Logger.WithFields(log.Fields{"at": at, "end": end}).Debug("Local: Hit iteration limit")
				return nil
			}
		case <-ctx.Done():
			// If the test is cancelled, just set the cutoff point to now and proceed down the same
			// logic as if the time limit was hit.
			e.Logger.Debug("Local: Exiting with context")
			cutoff = time.Now()
			return nil
		}
	}
}

func (e *Executor) scale(ctx context.Context, num int64) error {
	e.Logger.WithField("num", num).Debug("Local: Scaling...")

	e.vusLock.Lock()
	defer e.vusLock.Unlock()

	e.lock.RLock()
	flow := e.flow
	out := e.out
	e.lock.RUnlock()

	for i, handle := range e.vus {
		handle := handle
		handle.RLock()
		cancel := handle.cancel
		handle.RUnlock()

		if i < int(num) {
			if cancel == nil {
				vuctx, cancel := context.WithCancel(ctx)
				handle.Lock()
				handle.ctx = vuctx
				handle.cancel = cancel
				handle.Unlock()

				if handle.vu != nil {
					if err := handle.vu.Reconfigure(atomic.AddInt64(&e.nextVUID, 1)); err != nil {
						return err
					}
				}

				e.wg.Add(1)
				go func() {
					handle.run(e.Logger, flow, out)
					e.wg.Done()
				}()
			}
		} else if cancel != nil {
			handle.Lock()
			handle.cancel()
			handle.cancel = nil
			handle.Unlock()
		}
	}

	atomic.StoreInt64(&e.numVUs, num)
	return nil
}

func (e *Executor) IsRunning() bool {
	e.lock.RLock()
	defer e.lock.RUnlock()
	return e.ctx != nil
}

func (e *Executor) GetRunner() lib.Runner {
	return e.Runner
}

func (e *Executor) SetLogger(l *log.Logger) {
	e.Logger = l
}

func (e *Executor) GetLogger() *log.Logger {
	return e.Logger
}

func (e *Executor) GetIterations() int64 {
	return atomic.LoadInt64(&e.iters)
}

func (e *Executor) GetEndIterations() null.Int {
	v := atomic.LoadInt64(&e.endIters)
	if v < 0 {
		return null.Int{}
	}
	return null.IntFrom(v)
}

func (e *Executor) SetEndIterations(i null.Int) {
	if !i.Valid {
		i.Int64 = -1
	}
	e.Logger.WithField("i", i.Int64).Debug("Local: Setting end iterations")
	atomic.StoreInt64(&e.endIters, i.Int64)
}

func (e *Executor) GetTime() time.Duration {
	return time.Duration(atomic.LoadInt64(&e.time))
}

func (e *Executor) GetEndTime() lib.NullDuration {
	v := atomic.LoadInt64(&e.endTime)
	if v < 0 {
		return lib.NullDuration{}
	}
	return lib.NullDurationFrom(time.Duration(v))
}

func (e *Executor) SetEndTime(t lib.NullDuration) {
	if !t.Valid {
		t.Duration = -1
	}
	e.Logger.WithField("d", t.Duration).Debug("Local: Setting end time")
	atomic.StoreInt64(&e.endTime, int64(t.Duration))
}

func (e *Executor) IsPaused() bool {
	e.pauseLock.RLock()
	defer e.pauseLock.RUnlock()
	return e.pause != nil
}

func (e *Executor) SetPaused(paused bool) {
	e.Logger.WithField("paused", paused).Debug("Local: Setting paused")
	e.pauseLock.Lock()
	defer e.pauseLock.Unlock()

	if paused && e.pause == nil {
		e.pause = make(chan interface{})
	} else if !paused && e.pause != nil {
		close(e.pause)
		e.pause = nil
	}
}

func (e *Executor) GetVUs() int64 {
	return atomic.LoadInt64(&e.numVUs)
}

func (e *Executor) SetVUs(num int64) error {
	e.Logger.WithField("vus", num).Debug("Local: Setting VUs")
	if num < 0 {
		return errors.New("vu count can't be negative")
	}

	if atomic.LoadInt64(&e.numVUs) == num {
		return nil
	}

	if numVUsMax := atomic.LoadInt64(&e.numVUsMax); num > numVUsMax {
		return errors.Errorf("can't raise vu count (to %d) above vu cap (%d)", num, numVUsMax)
	}

	if ctx := e.ctx; ctx != nil {
		if err := e.scale(ctx, num); err != nil {
			return err
		}
	} else {
		atomic.StoreInt64(&e.numVUs, num)
	}

	return nil
}

func (e *Executor) GetVUsMax() int64 {
	return atomic.LoadInt64(&e.numVUsMax)
}

func (e *Executor) SetVUsMax(max int64) error {
	e.Logger.WithField("max", max).Debug("Local: Setting max VUs")
	if max < 0 {
		return errors.New("vu cap can't be negative")
	}

	numVUsMax := atomic.LoadInt64(&e.numVUsMax)

	if numVUsMax == max {
		return nil
	}

	if numVUs := atomic.LoadInt64(&e.numVUs); max < numVUs {
		return errors.Errorf("can't lower vu cap (to %d) below vu count (%d)", max, numVUs)
	}

	if max < numVUsMax {
		e.vus = e.vus[:max]
		atomic.StoreInt64(&e.numVUsMax, max)
		return nil
	}

	e.vusLock.Lock()
	defer e.vusLock.Unlock()

	vus := e.vus
	for i := numVUsMax; i < max; i++ {
		var handle vuHandle
		if e.Runner != nil {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			handle.vu = vu
		}
		vus = append(vus, &handle)
	}
	e.vus = vus

	atomic.StoreInt64(&e.numVUsMax, max)

	return nil
}
<|endoftext|>