<commit_before>// Command juggler-callee implements a testing callee that provides
// simple URI functions.
//
//     - test.echo (string) : returns the received string
//     - test.reverse (string) : reverses each rune in the received string
//     - test.delay (string) : sleeps for the duration received as string, converted to number (in ms)
//
package main

import (
	"encoding/json"
	"expvar"
	"flag"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/PuerkitoBio/juggler/broker"
	"github.com/PuerkitoBio/juggler/broker/redisbroker"
	"github.com/PuerkitoBio/juggler/callee"
	"github.com/PuerkitoBio/juggler/msg"
	"github.com/garyburd/redigo/redis"
)

var (
	redisAddrFlag             = flag.String("redis", ":6379", "Redis `address`.")
	redisPoolMaxActiveFlag    = flag.Int("redis-max-active", 100, "Maximum active redis `connections`.")
	redisPoolMaxIdleFlag      = flag.Int("redis-max-idle", 10, "Maximum idle redis `connections`.")
	redisPoolIdleTimeoutFlag  = flag.Duration("redis-idle-timeout", time.Minute, "Redis idle connection `timeout`.")
	brokerResultCapFlag       = flag.Int("broker-result-cap", 100, "Capacity of the `results` queue.")
	brokerBlockingTimeoutFlag = flag.Duration("broker-blocking-timeout", 0, "Blocking `timeout` when polling for call requests.")
	workersFlag               = flag.Int("workers", 1, "Number of concurrent `workers` processing call requests.")
	httpServerPortFlag        = flag.Int("port", 9001, "HTTP server `port` to serve debug endpoints.")
	helpFlag                  = flag.Bool("help", false, "Show help.")
)

var uris = map[string]callee.Thunk{
	"test.echo":    echoThunk,
	"test.reverse": reverseThunk,
	"test.delay":   delayThunk,
}

func main() {
	flag.Parse()
	if *helpFlag {
		flag.Usage()
		return
	}
	if *workersFlag <= 0 {
		*workersFlag = 1
	}

	pool := newRedisPool(*redisAddrFlag)
	c := &callee.Callee{Broker: newBroker(pool)}

	vars := expvar.NewMap("callee")

	// start a web server to serve pprof and expvar data
	log.Printf("serving debug endpoints on %d", *httpServerPortFlag)
	go func() {
		log.Println(http.ListenAndServe(":"+strconv.Itoa(*httpServerPortFlag), nil))
	}()

	log.Printf("listening for call requests on %s with %d workers", *redisAddrFlag, *workersFlag)
	keys := make([]string, 0, len(uris))
	for k := range uris {
		keys = append(keys, k)
	}

	cc, err := c.Broker.Calls(keys...)
	if err != nil {
		log.Fatalf("Calls failed: %v", err)
	}
	defer cc.Close()

	wg := sync.WaitGroup{}
	wg.Add(*workersFlag)
	for i := 0; i < *workersFlag; i++ {
		go func() {
			defer wg.Done()

			ch := cc.Calls()
			for cp := range ch {
				log.Printf("received request %v %s", cp.MsgUUID, cp.URI)
				vars.Add("Requests", 1)
				vars.Add("Requests."+cp.URI, 1)

				if err := c.InvokeAndStoreResult(cp, uris[cp.URI]); err != nil {
					if err != callee.ErrCallExpired {
						log.Printf("InvokeAndStoreResult failed: %v", err)
						vars.Add("Failed", 1)
						vars.Add("Failed."+cp.URI, 1)
						continue
					}
					log.Printf("expired request %v %s", cp.MsgUUID, cp.URI)
					vars.Add("Expired", 1)
					vars.Add("Expired."+cp.URI, 1)
					continue
				}
				log.Printf("sent result %v %s", cp.MsgUUID, cp.URI)
				vars.Add("Succeeded", 1)
				vars.Add("Succeded."+cp.URI, 1)
			}
		}()
	}
	wg.Wait()
}

func logWrapThunk(t callee.Thunk) callee.Thunk {
	return func(cp *msg.CallPayload) (interface{}, error) {
		log.Printf("received call for %s from %v", cp.URI, cp.MsgUUID)
		v, err := t(cp)
		log.Printf("sending result for %s from %v", cp.URI, cp.MsgUUID)
		return v, err
	}
}

func delayThunk(cp *msg.CallPayload) (interface{}, error) {
	var s string
	if err := json.Unmarshal(cp.Args, &s); err != nil {
		return nil, err
	}
	i, err := strconv.Atoi(s)
	if err != nil {
		return nil, err
	}
	return delay(i), nil
}

func delay(i int) int {
	time.Sleep(time.Duration(i) * time.Millisecond)
	return i
}

func reverseThunk(cp *msg.CallPayload) (interface{}, error) {
	var s string
	if err := json.Unmarshal(cp.Args, &s); err != nil {
		return nil, err
	}
	return reverse(s), nil
}

func reverse(s string) string {
	chars := []rune(s)
	for i, j := 0, len(chars)-1; i < j; i, j = i+1, j-1 {
		chars[i], chars[j] = chars[j], chars[i]
	}
	return string(chars)
}

func echoThunk(cp *msg.CallPayload) (interface{}, error) {
	var s string
	if err := json.Unmarshal(cp.Args, &s); err != nil {
		return nil, err
	}
	return echo(s), nil
}

func echo(s string) string {
	return s
}

func newBroker(pool *redis.Pool) broker.CalleeBroker {
	return &redisbroker.Broker{
		Pool:            pool,
		Dial:            pool.Dial,
		BlockingTimeout: *brokerBlockingTimeoutFlag,
		ResultCap:       *brokerResultCapFlag,
	}
}

func newRedisPool(addr string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     *redisPoolMaxIdleFlag,
		MaxActive:   *redisPoolMaxActiveFlag,
		IdleTimeout: *redisPoolIdleTimeoutFlag,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", addr)
			if err != nil {
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}
<commit_msg>juggler/cmd: support multiple URIs for test.delay in callee<commit_after>// Command juggler-callee implements a testing callee that provides
// simple URI functions.
//
//     - test.echo (string) : returns the received string
//     - test.reverse (string) : reverses each rune in the received string
//     - test.delay (string) : sleeps for the duration received as string, converted to number (in ms)
//
package main

import (
	"encoding/json"
	"expvar"
	"flag"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/PuerkitoBio/juggler/broker"
	"github.com/PuerkitoBio/juggler/broker/redisbroker"
	"github.com/PuerkitoBio/juggler/callee"
	"github.com/PuerkitoBio/juggler/msg"
	"github.com/garyburd/redigo/redis"
)

var (
	brokerBlockingTimeoutFlag = flag.Duration("broker-blocking-timeout", 0, "Blocking `timeout` when polling for call requests.")
	brokerResultCapFlag       = flag.Int("broker-result-cap", 100, "Capacity of the `results` queue.")
	helpFlag                  = flag.Bool("help", false, "Show help.")
	numDelayURIsFlag          = flag.Int("n", 100, "Number of test.delay `URIs`.")
	redisAddrFlag             = flag.String("redis", ":6379", "Redis `address`.")
	httpServerPortFlag        = flag.Int("port", 9001, "HTTP server `port` to serve debug endpoints.")
	redisPoolIdleTimeoutFlag  = flag.Duration("redis-idle-timeout", time.Minute, "Redis idle connection `timeout`.")
	redisPoolMaxActiveFlag    = flag.Int("redis-max-active", 100, "Maximum active redis `connections`.")
	redisPoolMaxIdleFlag      = flag.Int("redis-max-idle", 10, "Maximum idle redis `connections`.")
	workersFlag               = flag.Int("workers", 1, "Number of concurrent `workers` processing call requests.")
)

var uris = map[string]callee.Thunk{
	"test.echo":    echoThunk,
	"test.reverse": reverseThunk,
	"test.delay":   delayThunk,
}

func main() {
	flag.Parse()
	if *helpFlag {
		flag.Usage()
		return
	}
	if *workersFlag <= 0 {
		*workersFlag = 1
	}

	for i := 0; i < *numDelayURIsFlag; i++ {
		uris["test.delay."+strconv.Itoa(i)] = delayThunk
	}

	pool := newRedisPool(*redisAddrFlag)
	c := &callee.Callee{Broker: newBroker(pool)}

	vars := expvar.NewMap("callee")

	// start a web server to serve pprof and expvar data
	log.Printf("serving debug endpoints on %d", *httpServerPortFlag)
	go func() {
		log.Println(http.ListenAndServe(":"+strconv.Itoa(*httpServerPortFlag), nil))
	}()

	log.Printf("listening for call requests on %s with %d workers", *redisAddrFlag, *workersFlag)
	keys := make([]string, 0, len(uris))
	for k := range uris {
		keys = append(keys, k)
	}

	cc, err := c.Broker.Calls(keys...)
	if err != nil {
		log.Fatalf("Calls failed: %v", err)
	}
	defer cc.Close()

	wg := sync.WaitGroup{}
	wg.Add(*workersFlag)
	for i := 0; i < *workersFlag; i++ {
		go func() {
			defer wg.Done()

			ch := cc.Calls()
			for cp := range ch {
				log.Printf("received request %v %s", cp.MsgUUID, cp.URI)
				vars.Add("Requests", 1)
				vars.Add("Requests."+cp.URI, 1)

				if err := c.InvokeAndStoreResult(cp, uris[cp.URI]); err != nil {
					if err != callee.ErrCallExpired {
						log.Printf("InvokeAndStoreResult failed: %v", err)
						vars.Add("Failed", 1)
						vars.Add("Failed."+cp.URI, 1)
						continue
					}
					log.Printf("expired request %v %s", cp.MsgUUID, cp.URI)
					vars.Add("Expired", 1)
					vars.Add("Expired."+cp.URI, 1)
					continue
				}
				log.Printf("sent result %v %s", cp.MsgUUID, cp.URI)
				vars.Add("Succeeded", 1)
				vars.Add("Succeded."+cp.URI, 1)
			}
		}()
	}
	wg.Wait()
}

func logWrapThunk(t callee.Thunk) callee.Thunk {
	return func(cp *msg.CallPayload) (interface{}, error) {
		log.Printf("received call for %s from %v", cp.URI, cp.MsgUUID)
		v, err := t(cp)
		log.Printf("sending result for %s from %v", cp.URI, cp.MsgUUID)
		return v, err
	}
}

func delayThunk(cp *msg.CallPayload) (interface{}, error) {
	var s string
	if err := json.Unmarshal(cp.Args, &s); err != nil {
		return nil, err
	}
	i, err := strconv.Atoi(s)
	if err != nil {
		return nil, err
	}
	return delay(i), nil
}

func delay(i int) int {
	time.Sleep(time.Duration(i) * time.Millisecond)
	return i
}

func reverseThunk(cp *msg.CallPayload) (interface{}, error) {
	var s string
	if err := json.Unmarshal(cp.Args, &s); err != nil {
		return nil, err
	}
	return reverse(s), nil
}

func reverse(s string) string {
	chars := []rune(s)
	for i, j := 0, len(chars)-1; i < j; i, j = i+1, j-1 {
		chars[i], chars[j] = chars[j], chars[i]
	}
	return string(chars)
}

func echoThunk(cp *msg.CallPayload) (interface{}, error) {
	var s string
	if err := json.Unmarshal(cp.Args, &s); err != nil {
		return nil, err
	}
	return echo(s), nil
}

func echo(s string) string {
	return s
}

func newBroker(pool *redis.Pool) broker.CalleeBroker {
	return &redisbroker.Broker{
		Pool:            pool,
		Dial:            pool.Dial,
		BlockingTimeout: *brokerBlockingTimeoutFlag,
		ResultCap:       *brokerResultCapFlag,
	}
}

func newRedisPool(addr string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     *redisPoolMaxIdleFlag,
		MaxActive:   *redisPoolMaxActiveFlag,
		IdleTimeout: *redisPoolIdleTimeoutFlag,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", addr)
			if err != nil {
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Mirantis

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 (
	"flag"

	"github.com/Mirantis/virtlet/tests/e2e/framework"
	"github.com/Mirantis/virtlet/tests/longevity"
	"github.com/golang/glog"
)

func main() {
	baseTest := flag.Bool("base", false, "Run base longevity tests")
	stressTest := flag.Bool("stress", false, "Run longevity stress tests")

	flag.Parse()

	glog.Infof("Starting Virtlet longevity tests...")
	controller, err := framework.NewController("")
	if err != nil {
		glog.Fatal(err)
	}

	instances := []*longevity.VMInstance{}

	if *baseTest {
		instances = append(instances, longevity.GetBaseTests(controller)...)
	}
	if *stressTest {
		instances = append(instances, longevity.GetStressTests(controller)...)
	}

	err = longevity.Run(controller, instances)
	if err != nil {
		glog.Fatal(err)
	}
	controller.Finalize()
}
<commit_msg>When there are no tests selected then just return<commit_after>/*
Copyright 2018 Mirantis

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 (
	"flag"

	"github.com/Mirantis/virtlet/tests/e2e/framework"
	"github.com/Mirantis/virtlet/tests/longevity"
	"github.com/golang/glog"
)

func main() {
	baseTest := flag.Bool("base", false, "Run base longevity tests")
	stressTest := flag.Bool("stress", false, "Run longevity stress tests")

	flag.Parse()

	glog.Infof("Starting Virtlet longevity tests...")
	controller, err := framework.NewController("")
	if err != nil {
		glog.Fatal(err)
	}

	instances := []*longevity.VMInstance{}

	if *baseTest {
		instances = append(instances, longevity.GetBaseTests(controller)...)
	}
	if *stressTest {
		instances = append(instances, longevity.GetStressTests(controller)...)
	}

	if len(instances) == 0 {
		glog.Fatal("No tests selected. Quitting...")
	}

	err = longevity.Run(controller, instances)
	if err != nil {
		glog.Fatal(err)
	}
	controller.Finalize()
}
<|endoftext|>


<commit_before>// mailbox_sender - sender worker.
/*
Usage:

	mailbox_sender [flags]

The flags are:

	`-cid`: Campaign id
	`-g`: User groups
	`-p`: HTML file path
	`-t`: Mail Subject
	`-d`: Dry run all but not to send mail

Example:

	mailbox_sender -cid cbc6eb46 -g testuser -p ./email_1.html -t "#1 New Paper!" -d

*/
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"github.com/toomore/mailbox/campaign"
	"github.com/toomore/mailbox/mails"
	"github.com/toomore/mailbox/utils"
)

var (
	cid     = flag.String("cid", "", "campaign ID")
	dryRun  = flag.Bool("d", false, "Dry run")
	groups  = flag.String("g", "", "User groups")
	path    = flag.String("p", "", "HTML file path")
	subject = flag.String("t", "", "mail subject")
)

func replaceReader(html *[]byte, cid string, seed string, uid string) {
	data := url.Values{}
	data.Set("c", cid)
	data.Set("u", uid)
	hm := campaign.MakeMacSeed(seed, data)
	*html = bytes.Replace(
		*html,
		[]byte("{{READER}}"),
		[]byte(fmt.Sprintf("https://%s/read/%x?%s", os.Getenv("mailbox_web_site"), hm, data.Encode())),
		1)
}

func main() {
	flag.Parse()
	file, err := os.Open(*path)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatal(err)
	}
	seed := campaign.GetSeed(*cid)
	conn := utils.GetConn()
	rows, err := conn.Query(`SELECT id,email,f_name,l_name from user where alive=1 and groups=?`, *groups)
	defer rows.Close()
	if err != nil {
		log.Fatal(err)
	}

	var count int
	for rows.Next() {
		var (
			email string
			fname string
			lname string
			msg   []byte
			no    string
		)
		rows.Scan(&no, &email, &fname, &lname)

		msg = body
		replaceReader(&msg, *cid, seed, no)
		params := mails.GenParams(
			fmt.Sprintf("%s %s <%s>", fname, lname, email),
			string(msg),
			*subject)
		if !*dryRun {
			mails.Send(params)
		}
		count++
	}
	log.Printf("\n  cid: %s, groups: %s, count: %d\n  Subject: `%s`\n", *cid, *groups, count, *subject)
}
<commit_msg>Add uid into mailbox_sender<commit_after>// mailbox_sender - sender worker.
/*
Usage:

	mailbox_sender [flags]

The flags are:

	`-cid`: Campaign id
	`-g`: User groups
	`-p`: HTML file path
	`-t`: Mail Subject
	`-d`: Dry run all but not to send mail
	`-uid`: User ID

`-uid`, `-g` can't use together

Example:

	mailbox_sender -cid cbc6eb46 -g testuser -p ./email_1.html -t "#1 New Paper!" -d

*/
package main

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"github.com/toomore/mailbox/campaign"
	"github.com/toomore/mailbox/mails"
	"github.com/toomore/mailbox/utils"
)

var (
	cid     = flag.String("cid", "", "campaign ID")
	dryRun  = flag.Bool("d", false, "Dry run")
	groups  = flag.String("g", "", "User groups")
	path    = flag.String("p", "", "HTML file path")
	subject = flag.String("t", "", "mail subject")
	uid     = flag.String("uid", "", "User ID")
)

func replaceReader(html *[]byte, cid string, seed string, uid string) {
	data := url.Values{}
	data.Set("c", cid)
	data.Set("u", uid)
	hm := campaign.MakeMacSeed(seed, data)
	*html = bytes.Replace(
		*html,
		[]byte("{{READER}}"),
		[]byte(fmt.Sprintf("https://%s/read/%x?%s", os.Getenv("mailbox_web_site"), hm, data.Encode())),
		1)
}

func main() {
	flag.Parse()
	file, err := os.Open(*path)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatal(err)
	}
	seed := campaign.GetSeed(*cid)
	var rows *sql.Rows
	if *uid != "" {
		rows, err = utils.GetConn().Query(`SELECT id,email,f_name,l_name from user where alive=1 and id=?`, *uid)
	} else {
		rows, err = utils.GetConn().Query(`SELECT id,email,f_name,l_name from user where alive=1 and groups=?`, *groups)
	}
	defer rows.Close()
	if err != nil {
		log.Fatal(err)
	}

	var count int
	for rows.Next() {
		var (
			email string
			fname string
			lname string
			msg   []byte
			no    string
		)
		rows.Scan(&no, &email, &fname, &lname)

		msg = body
		replaceReader(&msg, *cid, seed, no)
		params := mails.GenParams(
			fmt.Sprintf("%s %s <%s>", fname, lname, email),
			string(msg),
			*subject)
		if !*dryRun {
			mails.Send(params)
		}
		count++
	}
	if *uid != "" {
		log.Printf("\n  cid: %s, uid: %s, count: %d\n  Subject: `%s`\n", *cid, *uid, count, *subject)
	} else {
		log.Printf("\n  cid: %s, groups: %s, count: %d\n  Subject: `%s`\n", *cid, *groups, count, *subject)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package modelcmd

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/api"
	"github.com/juju/juju/api/controller"
	"github.com/juju/juju/api/modelmanager"
	"github.com/juju/juju/api/usermanager"
	"github.com/juju/juju/jujuclient"
)

var (
	// ErrNoControllerSpecified is returned by commands that operate on
	// a controller if there is no current controller, no controller has been
	// explicitly specified, and there is no default controller.
	ErrNoControllerSpecified = errors.New(`no controller

Please create a new controller using "juju bootstrap" or
use "juju switch" to set the current controller/model.
`)
	// ErrNotLoggedInToController is returned by commands that operate on
	// a controller if there is no current controller, no controller has been
	// explicitly specified, and there is no default controller but there are
	// controllers that client knows about, i.e. the user needs to log in to one of them.
	ErrNotLoggedInToController = errors.New(`not logged in

Please use "juju controllers" to view all controllers available to you.
You can login into an existing controller using "juju login -c <controller>".
`)

	// ErrNoAccountSpecified is returned by commands that operate on a
	// controller if there is no current account associated with the
	// controller.
	ErrNoAccountSpecified = errors.New("no account specified")
)

// ControllerCommand is intended to be a base for all commands
// that need to operate on controllers as opposed to models.
type ControllerCommand interface {
	CommandBase

	// SetClientStore is called prior to the wrapped command's Init method
	// with the default controller store. It may also be called to override the
	// default controller store for testing.
	SetClientStore(jujuclient.ClientStore)

	// ClientStore returns the controller store that the command is
	// associated with.
	ClientStore() jujuclient.ClientStore

	// SetControllerName is called prior to the wrapped command's Init method with
	// the active controller name. The controller name is guaranteed to be non-empty
	// at entry of Init. It records the current model name in the
	// ControllerCommandBase.
	SetControllerName(controllerName string) error

	// ControllerName returns the name of the controller or model used to
	// determine that API end point.
	ControllerName() string

	// SetAPIOpener allows the replacement of the default API opener,
	// which ends up calling NewAPIRoot
	SetAPIOpener(opener APIOpener)
}

// ControllerCommandBase is a convenience type for embedding in commands
// that wish to implement ControllerCommand.
type ControllerCommandBase struct {
	JujuCommandBase

	store          jujuclient.ClientStore
	controllerName string
	accountName    string

	// opener is the strategy used to open the API connection.
	opener APIOpener
}

// SetClientStore implements the ControllerCommand interface.
func (c *ControllerCommandBase) SetClientStore(store jujuclient.ClientStore) {
	c.store = store
}

// ClientStore implements the ControllerCommand interface.
func (c *ControllerCommandBase) ClientStore() jujuclient.ClientStore {
	return c.store
}

// SetControllerName implements the ControllerCommand interface.
func (c *ControllerCommandBase) SetControllerName(controllerName string) error {
	actualControllerName, err := ResolveControllerName(c.ClientStore(), controllerName)
	if err != nil {
		return errors.Trace(err)
	}
	controllerName = actualControllerName

	accountName, err := c.store.CurrentAccount(controllerName)
	if err != nil && !errors.IsNotFound(err) {
		return errors.Trace(err)
	}
	c.controllerName = controllerName
	c.accountName = accountName
	return nil
}

// ControllerName implements the ControllerCommand interface.
func (c *ControllerCommandBase) ControllerName() string {
	return c.controllerName
}

// AccountName implements the ControllerCommand interface.
func (c *ControllerCommandBase) AccountName() string {
	return c.accountName
}

// SetAPIOpener specifies the strategy used by the command to open
// the API connection.
func (c *ControllerCommandBase) SetAPIOpener(opener APIOpener) {
	c.opener = opener
}

// NewModelManagerAPIClient returns an API client for the
// ModelManager on the current controller using the current credentials.
func (c *ControllerCommandBase) NewModelManagerAPIClient() (*modelmanager.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return modelmanager.NewClient(root), nil
}

// NewControllerAPIClient returns an API client for the Controller on
// the current controller using the current credentials.
func (c *ControllerCommandBase) NewControllerAPIClient() (*controller.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return controller.NewClient(root), nil
}

// NewUserManagerAPIClient returns an API client for the UserManager on the
// current controller using the current credentials.
func (c *ControllerCommandBase) NewUserManagerAPIClient() (*usermanager.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return usermanager.NewClient(root), nil
}

// NewAPIRoot returns a restricted API for the current controller using the current
// credentials.  Only the UserManager and ModelManager may be accessed
// through this API connection.
func (c *ControllerCommandBase) NewAPIRoot() (api.Connection, error) {
	if c.controllerName == "" {
		return nil, errors.Trace(ErrNoControllerSpecified)
	}
	if c.accountName == "" {
		return nil, errors.Trace(ErrNoAccountSpecified)
	}
	opener := c.opener
	if opener == nil {
		opener = OpenFunc(c.JujuCommandBase.NewAPIRoot)
	}
	return opener.Open(c.store, c.controllerName, c.accountName, "")
}

// ModelUUIDs returns the model UUIDs for the given model names.
func (c *ControllerCommandBase) ModelUUIDs(modelNames []string) ([]string, error) {
	var result []string
	store := c.ClientStore()
	controllerName := c.ControllerName()
	accountName := c.AccountName()
	for _, modelName := range modelNames {
		model, err := store.ModelByName(controllerName, accountName, modelName)
		if errors.IsNotFound(err) {
			// The model isn't known locally, so query the models available in the controller.
			logger.Infof("model %q not cached locally, refreshing models from controller", modelName)
			if err := c.RefreshModels(store, controllerName, accountName); err != nil {
				return nil, errors.Annotatef(err, "refreshing model %q", modelName)
			}
			model, err = store.ModelByName(controllerName, accountName, modelName)
		}
		if err != nil {
			return nil, errors.Annotatef(err, "model %q not found", modelName)
		}
		result = append(result, model.ModelUUID)
	}
	return result, nil
}

// WrapControllerOption sets various parameters of the
// ControllerCommand wrapper.
type WrapControllerOption func(*sysCommandWrapper)

// ControllerSkipFlags instructs the wrapper to skip -c
// and --controller flag definition.
func ControllerSkipFlags(w *sysCommandWrapper) {
	w.setFlags = false
}

// ControllerSkipDefault instructs the wrapper not to
// use the default controller name.
func ControllerSkipDefault(w *sysCommandWrapper) {
	w.useDefaultControllerName = false
}

// ControllerAPIOpener instructs the underlying controller command to use a
// different APIOpener strategy.
func ControllerAPIOpener(opener APIOpener) WrapControllerOption {
	return func(w *sysCommandWrapper) {
		w.ControllerCommand.SetAPIOpener(opener)
	}
}

// WrapController wraps the specified ControllerCommand, returning a Command
// that proxies to each of the ControllerCommand methods.
func WrapController(c ControllerCommand, options ...WrapControllerOption) cmd.Command {
	wrapper := &sysCommandWrapper{
		ControllerCommand:        c,
		setFlags:                 true,
		useDefaultControllerName: true,
	}
	for _, option := range options {
		option(wrapper)
	}
	return WrapBase(wrapper)
}

type sysCommandWrapper struct {
	ControllerCommand
	setFlags                 bool
	useDefaultControllerName bool
	controllerName           string
}

// SetFlags implements Command.SetFlags, then calls the wrapped command's SetFlags.
func (w *sysCommandWrapper) SetFlags(f *gnuflag.FlagSet) {
	if w.setFlags {
		f.StringVar(&w.controllerName, "c", "", "Controller to operate in")
		f.StringVar(&w.controllerName, "controller", "", "")
	}
	w.ControllerCommand.SetFlags(f)
}

// Init implements Command.Init, then calls the wrapped command's Init.
func (w *sysCommandWrapper) Init(args []string) error {
	store := w.ClientStore()
	if store == nil {
		store = jujuclient.NewFileClientStore()
		w.SetClientStore(store)
	}
	if w.setFlags {
		if w.controllerName == "" && w.useDefaultControllerName {
			currentController, err := store.CurrentController()
			if errors.IsNotFound(err) {
				return ErrNoControllerSpecified
			}
			if err != nil {
				return errors.Trace(err)
			}
			w.controllerName = currentController
		}
		if w.controllerName == "" && !w.useDefaultControllerName {
			return ErrNoControllerSpecified
		}
	}
	if w.controllerName != "" {
		if err := w.SetControllerName(w.controllerName); err != nil {
			return errors.Trace(err)
		}
	}
	return w.ControllerCommand.Init(args)
}

// ResolveControllerName returns the canonical name of a controller given
// an unambiguous identifier for that controller.
// Locally created controllers (i.e. those whose names begin with "local.")
// may be identified with or without the "local." prefix if there exists no
// other controller in the store with the same unprefixed name.
func ResolveControllerName(store jujuclient.ControllerStore, controllerName string) (string, error) {
	_, err := store.ControllerByName(controllerName)
	if err == nil {
		return controllerName, nil
	}
	if !errors.IsNotFound(err) {
		return "", err
	}
	var secondErr error
	localName := "local." + controllerName
	_, secondErr = store.ControllerByName(localName)
	// If fallback name not found, return the original error.
	if errors.IsNotFound(secondErr) {
		return "", err
	}
	return localName, secondErr
}
<commit_msg>Changed 'no controller' message.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package modelcmd

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/api"
	"github.com/juju/juju/api/controller"
	"github.com/juju/juju/api/modelmanager"
	"github.com/juju/juju/api/usermanager"
	"github.com/juju/juju/jujuclient"
)

var (
	// ErrNoControllerSpecified is returned by commands that operate on
	// a controller if there is no current controller, no controller has been
	// explicitly specified, and there is no default controller.
	ErrNoControllerSpecified = errors.New(`no controller

Please either create your own new controller using "juju bootstrap" or
connect to another controller that you have been given access to using "juju register".
`)
	// ErrNotLoggedInToController is returned by commands that operate on
	// a controller if there is no current controller, no controller has been
	// explicitly specified, and there is no default controller but there are
	// controllers that client knows about, i.e. the user needs to log in to one of them.
	ErrNotLoggedInToController = errors.New(`not logged in

Please use "juju controllers" to view all controllers available to you.
You can login into an existing controller using "juju login -c <controller>".
`)

	// ErrNoAccountSpecified is returned by commands that operate on a
	// controller if there is no current account associated with the
	// controller.
	ErrNoAccountSpecified = errors.New("no account specified")
)

// ControllerCommand is intended to be a base for all commands
// that need to operate on controllers as opposed to models.
type ControllerCommand interface {
	CommandBase

	// SetClientStore is called prior to the wrapped command's Init method
	// with the default controller store. It may also be called to override the
	// default controller store for testing.
	SetClientStore(jujuclient.ClientStore)

	// ClientStore returns the controller store that the command is
	// associated with.
	ClientStore() jujuclient.ClientStore

	// SetControllerName is called prior to the wrapped command's Init method with
	// the active controller name. The controller name is guaranteed to be non-empty
	// at entry of Init. It records the current model name in the
	// ControllerCommandBase.
	SetControllerName(controllerName string) error

	// ControllerName returns the name of the controller or model used to
	// determine that API end point.
	ControllerName() string

	// SetAPIOpener allows the replacement of the default API opener,
	// which ends up calling NewAPIRoot
	SetAPIOpener(opener APIOpener)
}

// ControllerCommandBase is a convenience type for embedding in commands
// that wish to implement ControllerCommand.
type ControllerCommandBase struct {
	JujuCommandBase

	store          jujuclient.ClientStore
	controllerName string
	accountName    string

	// opener is the strategy used to open the API connection.
	opener APIOpener
}

// SetClientStore implements the ControllerCommand interface.
func (c *ControllerCommandBase) SetClientStore(store jujuclient.ClientStore) {
	c.store = store
}

// ClientStore implements the ControllerCommand interface.
func (c *ControllerCommandBase) ClientStore() jujuclient.ClientStore {
	return c.store
}

// SetControllerName implements the ControllerCommand interface.
func (c *ControllerCommandBase) SetControllerName(controllerName string) error {
	actualControllerName, err := ResolveControllerName(c.ClientStore(), controllerName)
	if err != nil {
		return errors.Trace(err)
	}
	controllerName = actualControllerName

	accountName, err := c.store.CurrentAccount(controllerName)
	if err != nil && !errors.IsNotFound(err) {
		return errors.Trace(err)
	}
	c.controllerName = controllerName
	c.accountName = accountName
	return nil
}

// ControllerName implements the ControllerCommand interface.
func (c *ControllerCommandBase) ControllerName() string {
	return c.controllerName
}

// AccountName implements the ControllerCommand interface.
func (c *ControllerCommandBase) AccountName() string {
	return c.accountName
}

// SetAPIOpener specifies the strategy used by the command to open
// the API connection.
func (c *ControllerCommandBase) SetAPIOpener(opener APIOpener) {
	c.opener = opener
}

// NewModelManagerAPIClient returns an API client for the
// ModelManager on the current controller using the current credentials.
func (c *ControllerCommandBase) NewModelManagerAPIClient() (*modelmanager.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return modelmanager.NewClient(root), nil
}

// NewControllerAPIClient returns an API client for the Controller on
// the current controller using the current credentials.
func (c *ControllerCommandBase) NewControllerAPIClient() (*controller.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return controller.NewClient(root), nil
}

// NewUserManagerAPIClient returns an API client for the UserManager on the
// current controller using the current credentials.
func (c *ControllerCommandBase) NewUserManagerAPIClient() (*usermanager.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return usermanager.NewClient(root), nil
}

// NewAPIRoot returns a restricted API for the current controller using the current
// credentials.  Only the UserManager and ModelManager may be accessed
// through this API connection.
func (c *ControllerCommandBase) NewAPIRoot() (api.Connection, error) {
	if c.controllerName == "" {
		return nil, errors.Trace(ErrNoControllerSpecified)
	}
	if c.accountName == "" {
		return nil, errors.Trace(ErrNoAccountSpecified)
	}
	opener := c.opener
	if opener == nil {
		opener = OpenFunc(c.JujuCommandBase.NewAPIRoot)
	}
	return opener.Open(c.store, c.controllerName, c.accountName, "")
}

// ModelUUIDs returns the model UUIDs for the given model names.
func (c *ControllerCommandBase) ModelUUIDs(modelNames []string) ([]string, error) {
	var result []string
	store := c.ClientStore()
	controllerName := c.ControllerName()
	accountName := c.AccountName()
	for _, modelName := range modelNames {
		model, err := store.ModelByName(controllerName, accountName, modelName)
		if errors.IsNotFound(err) {
			// The model isn't known locally, so query the models available in the controller.
			logger.Infof("model %q not cached locally, refreshing models from controller", modelName)
			if err := c.RefreshModels(store, controllerName, accountName); err != nil {
				return nil, errors.Annotatef(err, "refreshing model %q", modelName)
			}
			model, err = store.ModelByName(controllerName, accountName, modelName)
		}
		if err != nil {
			return nil, errors.Annotatef(err, "model %q not found", modelName)
		}
		result = append(result, model.ModelUUID)
	}
	return result, nil
}

// WrapControllerOption sets various parameters of the
// ControllerCommand wrapper.
type WrapControllerOption func(*sysCommandWrapper)

// ControllerSkipFlags instructs the wrapper to skip -c
// and --controller flag definition.
func ControllerSkipFlags(w *sysCommandWrapper) {
	w.setFlags = false
}

// ControllerSkipDefault instructs the wrapper not to
// use the default controller name.
func ControllerSkipDefault(w *sysCommandWrapper) {
	w.useDefaultControllerName = false
}

// ControllerAPIOpener instructs the underlying controller command to use a
// different APIOpener strategy.
func ControllerAPIOpener(opener APIOpener) WrapControllerOption {
	return func(w *sysCommandWrapper) {
		w.ControllerCommand.SetAPIOpener(opener)
	}
}

// WrapController wraps the specified ControllerCommand, returning a Command
// that proxies to each of the ControllerCommand methods.
func WrapController(c ControllerCommand, options ...WrapControllerOption) cmd.Command {
	wrapper := &sysCommandWrapper{
		ControllerCommand:        c,
		setFlags:                 true,
		useDefaultControllerName: true,
	}
	for _, option := range options {
		option(wrapper)
	}
	return WrapBase(wrapper)
}

type sysCommandWrapper struct {
	ControllerCommand
	setFlags                 bool
	useDefaultControllerName bool
	controllerName           string
}

// SetFlags implements Command.SetFlags, then calls the wrapped command's SetFlags.
func (w *sysCommandWrapper) SetFlags(f *gnuflag.FlagSet) {
	if w.setFlags {
		f.StringVar(&w.controllerName, "c", "", "Controller to operate in")
		f.StringVar(&w.controllerName, "controller", "", "")
	}
	w.ControllerCommand.SetFlags(f)
}

// Init implements Command.Init, then calls the wrapped command's Init.
func (w *sysCommandWrapper) Init(args []string) error {
	store := w.ClientStore()
	if store == nil {
		store = jujuclient.NewFileClientStore()
		w.SetClientStore(store)
	}
	if w.setFlags {
		if w.controllerName == "" && w.useDefaultControllerName {
			currentController, err := store.CurrentController()
			if errors.IsNotFound(err) {
				return ErrNoControllerSpecified
			}
			if err != nil {
				return errors.Trace(err)
			}
			w.controllerName = currentController
		}
		if w.controllerName == "" && !w.useDefaultControllerName {
			return ErrNoControllerSpecified
		}
	}
	if w.controllerName != "" {
		if err := w.SetControllerName(w.controllerName); err != nil {
			return errors.Trace(err)
		}
	}
	return w.ControllerCommand.Init(args)
}

// ResolveControllerName returns the canonical name of a controller given
// an unambiguous identifier for that controller.
// Locally created controllers (i.e. those whose names begin with "local.")
// may be identified with or without the "local." prefix if there exists no
// other controller in the store with the same unprefixed name.
func ResolveControllerName(store jujuclient.ControllerStore, controllerName string) (string, error) {
	_, err := store.ControllerByName(controllerName)
	if err == nil {
		return controllerName, nil
	}
	if !errors.IsNotFound(err) {
		return "", err
	}
	var secondErr error
	localName := "local." + controllerName
	_, secondErr = store.ControllerByName(localName)
	// If fallback name not found, return the original error.
	if errors.IsNotFound(secondErr) {
		return "", err
	}
	return localName, secondErr
}
<|endoftext|>


<commit_before>package main_test

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"

	"google.golang.org/grpc/grpclog"

	"code.cloudfoundry.org/bbs"
	bbsconfig "code.cloudfoundry.org/bbs/cmd/bbs/config"
	bbstestrunner "code.cloudfoundry.org/bbs/cmd/bbs/testrunner"
	"code.cloudfoundry.org/bbs/encryption"
	"code.cloudfoundry.org/bbs/test_helpers"
	"code.cloudfoundry.org/bbs/test_helpers/sqlrunner"
	"code.cloudfoundry.org/consuladapter/consulrunner"
	"code.cloudfoundry.org/diego-logging-client/testhelpers"
	"code.cloudfoundry.org/durationjson"
	"code.cloudfoundry.org/inigo/helpers/portauthority"
	"code.cloudfoundry.org/lager/lagerflags"
	"code.cloudfoundry.org/locket"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/config"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

var (
	cellID              string
	representativePath  string
	natsPort            int
	serverPort          uint16
	serverPortSecurable uint16
	consulRunner        *consulrunner.ClusterRunner

	bbsConfig        bbsconfig.BBSConfig
	bbsBinPath       string
	bbsURL           *url.URL
	bbsRunner        *ginkgomon.Runner
	bbsProcess       ifrit.Process
	bbsClient        bbs.InternalClient
	auctioneerServer *ghttp.Server
	locketBinPath    string
	node             int

	testIngressServer *testhelpers.TestIngressServer

	sqlProcess    ifrit.Process
	sqlRunner     sqlrunner.SQLRunner
	portAllocator portauthority.PortAllocator

	fixturesPath = path.Join(os.Getenv("GOPATH"), "src/code.cloudfoundry.org/rep/cmd/rep/fixtures")
)

func TestRep(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Rep Integration Suite")
}

var _ = SynchronizedBeforeSuite(func() []byte {
	bbsConfig, err := gexec.Build("code.cloudfoundry.org/bbs/cmd/bbs", "-race")
	Expect(err).NotTo(HaveOccurred())

	locketPath, err := gexec.Build("code.cloudfoundry.org/locket/cmd/locket", "-race")
	Expect(err).NotTo(HaveOccurred())

	representative, err := gexec.Build("code.cloudfoundry.org/rep/cmd/rep", "-race")
	Expect(err).NotTo(HaveOccurred())

	return []byte(strings.Join([]string{representative, locketPath, bbsConfig}, ","))
}, func(pathsByte []byte) {

	node = GinkgoParallelNode()
	startPort := 1050 * node
	portRange := 1000
	endPort := startPort + portRange

	var err error
	portAllocator, err = portauthority.New(startPort, endPort)
	Expect(err).NotTo(HaveOccurred())

	grpclog.SetLogger(log.New(ioutil.Discard, "", 0))

	// tests here are fairly Eventually driven which tends to flake out under
	// load (for insignificant reasons); bump the default a bit higher than the
	// default (1 second)
	SetDefaultEventuallyTimeout(5 * time.Second)

	strPath := string(pathsByte)
	representativePath = strings.Split(strPath, ",")[0]
	locketBinPath = strings.Split(strPath, ",")[1]
	bbsBinPath = strings.Split(strPath, ",")[2]

	cellID = "the_rep_id-" + strconv.Itoa(GinkgoParallelNode())

	serverPort, err = portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	serverPortSecurable, err = portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())

	dbName := fmt.Sprintf("diego_%d", GinkgoParallelNode())

	sqlRunner = test_helpers.NewSQLRunner(dbName)
	sqlProcess = ginkgomon.Invoke(sqlRunner)

	consulRunner = consulrunner.NewClusterRunner(
		consulrunner.ClusterRunnerConfig{
			StartingPort: 9001 + config.GinkgoConfig.ParallelNode*consulrunner.PortOffsetLength,
			NumNodes:     1,
			Scheme:       "http",
		},
	)

	consulRunner.Start()

	bbsPort, err := portAllocator.ClaimPorts(2)
	Expect(err).NotTo(HaveOccurred())
	healthPort := bbsPort + 1
	bbsAddress := fmt.Sprintf("127.0.0.1:%d", bbsPort)
	healthAddress := fmt.Sprintf("127.0.0.1:%d", healthPort)

	bbsURL = &url.URL{
		Scheme: "https",
		Host:   bbsAddress,
	}

	fixturesPath := path.Join(os.Getenv("GOPATH"), "src/code.cloudfoundry.org/rep/cmd/rep/fixtures")

	auctioneerServer = ghttp.NewServer()
	auctioneerServer.UnhandledRequestStatusCode = http.StatusAccepted
	auctioneerServer.AllowUnhandledRequests = true

	bbsConfig = bbsconfig.BBSConfig{
		ListenAddress:                  bbsAddress,
		AdvertiseURL:                   bbsURL.String(),
		AuctioneerAddress:              auctioneerServer.URL(),
		DatabaseDriver:                 sqlRunner.DriverName(),
		DatabaseConnectionString:       sqlRunner.ConnectionString(),
		DetectConsulCellRegistrations:  true,
		ConsulCluster:                  consulRunner.ConsulCluster(),
		HealthAddress:                  healthAddress,
		LocksLocketEnabled:             false,
		CellRegistrationsLocketEnabled: false,
		EncryptionConfig: encryption.EncryptionConfig{
			EncryptionKeys: map[string]string{"label": "key"},
			ActiveKeyLabel: "label",
		},

		CaFile:   path.Join(fixturesPath, "green-certs", "server-ca.crt"),
		CertFile: path.Join(fixturesPath, "green-certs", "server.crt"),
		KeyFile:  path.Join(fixturesPath, "green-certs", "server.key"),

		SessionName:                     "bbs",
		CommunicationTimeout:            durationjson.Duration(10 * time.Second),
		RequireSSL:                      true,
		DesiredLRPCreationTimeout:       durationjson.Duration(1 * time.Minute),
		ExpireCompletedTaskDuration:     durationjson.Duration(2 * time.Minute),
		ExpirePendingTaskDuration:       durationjson.Duration(30 * time.Minute),
		EnableConsulServiceRegistration: false,
		ConvergeRepeatInterval:          durationjson.Duration(30 * time.Second),
		KickTaskDuration:                durationjson.Duration(30 * time.Second),
		LockTTL:                         durationjson.Duration(locket.DefaultSessionTTL),
		LockRetryInterval:               durationjson.Duration(locket.RetryInterval),
		ReportInterval:                  durationjson.Duration(1 * time.Minute),
		ConvergenceWorkers:              20,
		UpdateWorkers:                   1000,
		TaskCallbackWorkers:             1000,
		MaxOpenDatabaseConnections:      200,
		MaxIdleDatabaseConnections:      200,
		AuctioneerRequireTLS:            false,
		RepClientSessionCacheSize:       0,
		RepRequireTLS:                   false,
		LagerConfig: lagerflags.LagerConfig{
			LogLevel: "info",
		},
		GenerateSuspectActualLRPs: false,
	}
})

var _ = BeforeEach(func() {
	consulRunner.WaitUntilReady()
	consulRunner.Reset()

	bbsRunner = bbstestrunner.New(bbsBinPath, bbsConfig)
	bbsProcess = ginkgomon.Invoke(bbsRunner)

	metronCAFile := path.Join(fixturesPath, "metron", "CA.crt")
	metronServerCertFile := path.Join(fixturesPath, "metron", "metron.crt")
	metronServerKeyFile := path.Join(fixturesPath, "metron", "metron.key")

	var err error
	testIngressServer, err = testhelpers.NewTestIngressServer(metronServerCertFile, metronServerKeyFile, metronCAFile)
	Expect(err).NotTo(HaveOccurred())
	Expect(testIngressServer.Start()).To(Succeed())
})

var _ = AfterEach(func() {
	sqlRunner.Reset()

	ginkgomon.Kill(bbsProcess)
})

var _ = SynchronizedAfterSuite(func() {
	ginkgomon.Kill(sqlProcess)
	if consulRunner != nil {
		consulRunner.Stop()
	}
	if runner != nil {
		runner.KillWithFire()
	}
	if auctioneerServer != nil {
		auctioneerServer.Close()
	}
}, func() {
	gexec.CleanupBuildArtifacts()
})
<commit_msg>remove GenerateSuspectActualLRPs<commit_after>package main_test

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"

	"google.golang.org/grpc/grpclog"

	"code.cloudfoundry.org/bbs"
	bbsconfig "code.cloudfoundry.org/bbs/cmd/bbs/config"
	bbstestrunner "code.cloudfoundry.org/bbs/cmd/bbs/testrunner"
	"code.cloudfoundry.org/bbs/encryption"
	"code.cloudfoundry.org/bbs/test_helpers"
	"code.cloudfoundry.org/bbs/test_helpers/sqlrunner"
	"code.cloudfoundry.org/consuladapter/consulrunner"
	"code.cloudfoundry.org/diego-logging-client/testhelpers"
	"code.cloudfoundry.org/durationjson"
	"code.cloudfoundry.org/inigo/helpers/portauthority"
	"code.cloudfoundry.org/lager/lagerflags"
	"code.cloudfoundry.org/locket"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/config"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

var (
	cellID              string
	representativePath  string
	natsPort            int
	serverPort          uint16
	serverPortSecurable uint16
	consulRunner        *consulrunner.ClusterRunner

	bbsConfig        bbsconfig.BBSConfig
	bbsBinPath       string
	bbsURL           *url.URL
	bbsRunner        *ginkgomon.Runner
	bbsProcess       ifrit.Process
	bbsClient        bbs.InternalClient
	auctioneerServer *ghttp.Server
	locketBinPath    string
	node             int

	testIngressServer *testhelpers.TestIngressServer

	sqlProcess    ifrit.Process
	sqlRunner     sqlrunner.SQLRunner
	portAllocator portauthority.PortAllocator

	fixturesPath = path.Join(os.Getenv("GOPATH"), "src/code.cloudfoundry.org/rep/cmd/rep/fixtures")
)

func TestRep(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Rep Integration Suite")
}

var _ = SynchronizedBeforeSuite(func() []byte {
	bbsConfig, err := gexec.Build("code.cloudfoundry.org/bbs/cmd/bbs", "-race")
	Expect(err).NotTo(HaveOccurred())

	locketPath, err := gexec.Build("code.cloudfoundry.org/locket/cmd/locket", "-race")
	Expect(err).NotTo(HaveOccurred())

	representative, err := gexec.Build("code.cloudfoundry.org/rep/cmd/rep", "-race")
	Expect(err).NotTo(HaveOccurred())

	return []byte(strings.Join([]string{representative, locketPath, bbsConfig}, ","))
}, func(pathsByte []byte) {

	node = GinkgoParallelNode()
	startPort := 1050 * node
	portRange := 1000
	endPort := startPort + portRange

	var err error
	portAllocator, err = portauthority.New(startPort, endPort)
	Expect(err).NotTo(HaveOccurred())

	grpclog.SetLogger(log.New(ioutil.Discard, "", 0))

	// tests here are fairly Eventually driven which tends to flake out under
	// load (for insignificant reasons); bump the default a bit higher than the
	// default (1 second)
	SetDefaultEventuallyTimeout(5 * time.Second)

	strPath := string(pathsByte)
	representativePath = strings.Split(strPath, ",")[0]
	locketBinPath = strings.Split(strPath, ",")[1]
	bbsBinPath = strings.Split(strPath, ",")[2]

	cellID = "the_rep_id-" + strconv.Itoa(GinkgoParallelNode())

	serverPort, err = portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	serverPortSecurable, err = portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())

	dbName := fmt.Sprintf("diego_%d", GinkgoParallelNode())

	sqlRunner = test_helpers.NewSQLRunner(dbName)
	sqlProcess = ginkgomon.Invoke(sqlRunner)

	consulRunner = consulrunner.NewClusterRunner(
		consulrunner.ClusterRunnerConfig{
			StartingPort: 9001 + config.GinkgoConfig.ParallelNode*consulrunner.PortOffsetLength,
			NumNodes:     1,
			Scheme:       "http",
		},
	)

	consulRunner.Start()

	bbsPort, err := portAllocator.ClaimPorts(2)
	Expect(err).NotTo(HaveOccurred())
	healthPort := bbsPort + 1
	bbsAddress := fmt.Sprintf("127.0.0.1:%d", bbsPort)
	healthAddress := fmt.Sprintf("127.0.0.1:%d", healthPort)

	bbsURL = &url.URL{
		Scheme: "https",
		Host:   bbsAddress,
	}

	fixturesPath := path.Join(os.Getenv("GOPATH"), "src/code.cloudfoundry.org/rep/cmd/rep/fixtures")

	auctioneerServer = ghttp.NewServer()
	auctioneerServer.UnhandledRequestStatusCode = http.StatusAccepted
	auctioneerServer.AllowUnhandledRequests = true

	bbsConfig = bbsconfig.BBSConfig{
		ListenAddress:                  bbsAddress,
		AdvertiseURL:                   bbsURL.String(),
		AuctioneerAddress:              auctioneerServer.URL(),
		DatabaseDriver:                 sqlRunner.DriverName(),
		DatabaseConnectionString:       sqlRunner.ConnectionString(),
		DetectConsulCellRegistrations:  true,
		ConsulCluster:                  consulRunner.ConsulCluster(),
		HealthAddress:                  healthAddress,
		LocksLocketEnabled:             false,
		CellRegistrationsLocketEnabled: false,
		EncryptionConfig: encryption.EncryptionConfig{
			EncryptionKeys: map[string]string{"label": "key"},
			ActiveKeyLabel: "label",
		},

		CaFile:   path.Join(fixturesPath, "green-certs", "server-ca.crt"),
		CertFile: path.Join(fixturesPath, "green-certs", "server.crt"),
		KeyFile:  path.Join(fixturesPath, "green-certs", "server.key"),

		SessionName:                     "bbs",
		CommunicationTimeout:            durationjson.Duration(10 * time.Second),
		RequireSSL:                      true,
		DesiredLRPCreationTimeout:       durationjson.Duration(1 * time.Minute),
		ExpireCompletedTaskDuration:     durationjson.Duration(2 * time.Minute),
		ExpirePendingTaskDuration:       durationjson.Duration(30 * time.Minute),
		EnableConsulServiceRegistration: false,
		ConvergeRepeatInterval:          durationjson.Duration(30 * time.Second),
		KickTaskDuration:                durationjson.Duration(30 * time.Second),
		LockTTL:                         durationjson.Duration(locket.DefaultSessionTTL),
		LockRetryInterval:               durationjson.Duration(locket.RetryInterval),
		ReportInterval:                  durationjson.Duration(1 * time.Minute),
		ConvergenceWorkers:              20,
		UpdateWorkers:                   1000,
		TaskCallbackWorkers:             1000,
		MaxOpenDatabaseConnections:      200,
		MaxIdleDatabaseConnections:      200,
		AuctioneerRequireTLS:            false,
		RepClientSessionCacheSize:       0,
		RepRequireTLS:                   false,
		LagerConfig: lagerflags.LagerConfig{
			LogLevel: "info",
		},
	}
})

var _ = BeforeEach(func() {
	consulRunner.WaitUntilReady()
	consulRunner.Reset()

	bbsRunner = bbstestrunner.New(bbsBinPath, bbsConfig)
	bbsProcess = ginkgomon.Invoke(bbsRunner)

	metronCAFile := path.Join(fixturesPath, "metron", "CA.crt")
	metronServerCertFile := path.Join(fixturesPath, "metron", "metron.crt")
	metronServerKeyFile := path.Join(fixturesPath, "metron", "metron.key")

	var err error
	testIngressServer, err = testhelpers.NewTestIngressServer(metronServerCertFile, metronServerKeyFile, metronCAFile)
	Expect(err).NotTo(HaveOccurred())
	Expect(testIngressServer.Start()).To(Succeed())
})

var _ = AfterEach(func() {
	sqlRunner.Reset()

	ginkgomon.Kill(bbsProcess)
})

var _ = SynchronizedAfterSuite(func() {
	ginkgomon.Kill(sqlProcess)
	if consulRunner != nil {
		consulRunner.Stop()
	}
	if runner != nil {
		runner.KillWithFire()
	}
	if auctioneerServer != nil {
		auctioneerServer.Close()
	}
}, func() {
	gexec.CleanupBuildArtifacts()
})
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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.

// Implements ruby/appengine buildpack.
// The appengine buildpack sets the image entrypoint.
package main

import (
	"errors"
	"path/filepath"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/appengine"
	gcp "github.com/GoogleCloudPlatform/buildpacks/pkg/gcpbuildpack"
)

const (
	bundleIndicator  = "Gemfile.lock"
	bundle2Indicator = "gems.locked"
	railsIndicator   = "bin/rails"
	railsCommand     = "bin/rails server"
	rackIndicator    = "config.ru"
	rackCommand      = "rackup --port $PORT"
)

func main() {
	gcp.Main(detectFn, buildFn)
}

func detectFn(ctx *gcp.Context) error {
	// Always opt-in.
	return nil
}

func buildFn(ctx *gcp.Context) error {
	// Ruby sometimes writes to local directories tmp/ and log/, so we link these to writable areas.
	localTemp := filepath.Join(ctx.ApplicationRoot(), "tmp")
	localLog := filepath.Join(ctx.ApplicationRoot(), "log")
	ctx.RemoveAll(localTemp)
	ctx.RemoveAll(localLog)
	ctx.Symlink("/tmp", localTemp)
	ctx.Symlink("/var/log", localLog)

	return appengine.Build(ctx, "ruby",
		func(ctx *gcp.Context) (*appengine.Entrypoint, error) {
			return entrypoint(ctx, ctx.ApplicationRoot())
		})
}

func entrypoint(ctx *gcp.Context, srcDir string) (*appengine.Entrypoint, error) {
	var ep string
	ctx.Logf("WARNING: No entrypoint specified. Attempting to infer entrypoint, but it is recommended to set an explicit `entrypoint` in app.yaml.")
	if ctx.FileExists(srcDir, railsIndicator) {
		ep = maybeBundle(ctx, srcDir, railsCommand)
	} else if ctx.FileExists(srcDir, rackIndicator) {
		ep = maybeBundle(ctx, srcDir, rackCommand)
	} else {
		return nil, errors.New("unable to infer entrypoint, please set the `entrypoint` field in app.yaml: https://cloud.google.com/appengine/docs/standard/ruby/runtime#application_startup")
	}
	ctx.Logf("Using inferred entrypoint: %q", ep)
	return &appengine.Entrypoint{
		Type:    appengine.EntrypointGenerated.String(),
		Command: ep,
	}, nil
}

func maybeBundle(ctx *gcp.Context, srcDir, cmd string) string {
	if ctx.FileExists(srcDir, bundleIndicator) || ctx.FileExists(srcDir, bundle2Indicator) {
		return "bundle exec " + cmd
	}
	return cmd
}
<commit_msg>Change failure attribution when failing to infer entrypoint in GAE Ruby to user<commit_after>// Copyright 2020 Google LLC
//
// 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.

// Implements ruby/appengine buildpack.
// The appengine buildpack sets the image entrypoint.
package main

import (
	"path/filepath"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/appengine"
	gcp "github.com/GoogleCloudPlatform/buildpacks/pkg/gcpbuildpack"
)

const (
	bundleIndicator  = "Gemfile.lock"
	bundle2Indicator = "gems.locked"
	railsIndicator   = "bin/rails"
	railsCommand     = "bin/rails server"
	rackIndicator    = "config.ru"
	rackCommand      = "rackup --port $PORT"
)

func main() {
	gcp.Main(detectFn, buildFn)
}

func detectFn(ctx *gcp.Context) error {
	// Always opt-in.
	return nil
}

func buildFn(ctx *gcp.Context) error {
	// Ruby sometimes writes to local directories tmp/ and log/, so we link these to writable areas.
	localTemp := filepath.Join(ctx.ApplicationRoot(), "tmp")
	localLog := filepath.Join(ctx.ApplicationRoot(), "log")
	ctx.RemoveAll(localTemp)
	ctx.RemoveAll(localLog)
	ctx.Symlink("/tmp", localTemp)
	ctx.Symlink("/var/log", localLog)

	return appengine.Build(ctx, "ruby",
		func(ctx *gcp.Context) (*appengine.Entrypoint, error) {
			return entrypoint(ctx, ctx.ApplicationRoot())
		})
}

func entrypoint(ctx *gcp.Context, srcDir string) (*appengine.Entrypoint, error) {
	var ep string
	ctx.Logf("WARNING: No entrypoint specified. Attempting to infer entrypoint, but it is recommended to set an explicit `entrypoint` in app.yaml.")
	if ctx.FileExists(srcDir, railsIndicator) {
		ep = maybeBundle(ctx, srcDir, railsCommand)
	} else if ctx.FileExists(srcDir, rackIndicator) {
		ep = maybeBundle(ctx, srcDir, rackCommand)
	} else {
		return nil, gcp.UserErrorf("unable to infer entrypoint, please set the `entrypoint` field in app.yaml: https://cloud.google.com/appengine/docs/standard/ruby/runtime#application_startup")
	}
	ctx.Logf("Using inferred entrypoint: %q", ep)
	return &appengine.Entrypoint{
		Type:    appengine.EntrypointGenerated.String(),
		Command: ep,
	}, nil
}

func maybeBundle(ctx *gcp.Context, srcDir, cmd string) string {
	if ctx.FileExists(srcDir, bundleIndicator) || ctx.FileExists(srcDir, bundle2Indicator) {
		return "bundle exec " + cmd
	}
	return cmd
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).

package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/kardianos/osext"
	"github.com/syncthing/relaysrv/client"
	"github.com/syncthing/syncthing/lib/sync"
	"github.com/syncthing/syncthing/lib/tlsutil"
)

type relay struct {
	URL string `json:"url"`
}

func (r relay) String() string {
	return r.URL
}

type request struct {
	relay  relay
	uri    *url.URL
	result chan result
}

type result struct {
	err      error
	eviction time.Duration
}

var (
	binDir       string
	testCert     []tls.Certificate
	listen       string
	dir          string
	evictionTime time.Duration
	debug        bool

	requests = make(chan request, 10)

	mut             sync.RWMutex           = sync.NewRWMutex()
	knownRelays     []relay                = make([]relay, 0)
	permanentRelays []relay                = make([]relay, 0)
	evictionTimers  map[string]*time.Timer = make(map[string]*time.Timer)
)

func main() {
	flag.StringVar(&listen, "listen", ":80", "Listen address")
	flag.StringVar(&dir, "keys", "", "Directory where http-cert.pem and http-key.pem is stored for TLS listening")
	flag.BoolVar(&debug, "debug", false, "Enable debug output")
	flag.DurationVar(&evictionTime, "eviction", time.Hour, "After how long the relay is evicted")

	flag.Parse()

	var listener net.Listener
	var err error

	binDir, err = osext.ExecutableFolder()
	if err != nil {
		log.Fatalln("Failed to locate executable directory")
	}

	loadPermanentRelays()
	loadOrCreateTestCertificate()
	go requestProcessor()

	if dir != "" {
		if debug {
			log.Println("Starting TLS listener")
		}
		certFile, keyFile := filepath.Join(dir, "http-cert.pem"), filepath.Join(dir, "http-key.pem")
		cert, err := tls.LoadX509KeyPair(certFile, keyFile)
		if err != nil {
			log.Fatalln("Failed to load HTTP X509 key pair:", err)
		}

		tlsCfg := &tls.Config{
			Certificates: []tls.Certificate{cert},
			MinVersion:   tls.VersionTLS10, // No SSLv3
			CipherSuites: []uint16{
				// No RC4
				tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
				tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
				tls.TLS_RSA_WITH_AES_128_CBC_SHA,
				tls.TLS_RSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA,
				tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA,
			},
		}

		listener, err = tls.Listen("tcp", listen, tlsCfg)
	} else {
		if debug {
			log.Println("Starting plain listener")
		}
		listener, err = net.Listen("tcp", listen)
	}

	if err != nil {
		log.Fatalln("listen:", err)
	}

	handler := http.NewServeMux()
	handler.HandleFunc("/", handleRequest)

	srv := http.Server{
		Handler:     handler,
		ReadTimeout: 10 * time.Second,
	}

	err = srv.Serve(listener)
	if err != nil {
		log.Fatalln("serve:", err)
	}
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		handleGetRequest(w, r)
	case "POST":
		handlePostRequest(w, r)
	default:
		if debug {
			log.Println("Unhandled HTTP method", r.Method)
		}
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

func handleGetRequest(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	mut.RLock()
	relays := append(permanentRelays, knownRelays...)
	mut.RUnlock()

	// Shuffle
	for i := range relays {
		j := rand.Intn(i + 1)
		relays[i], relays[j] = relays[j], relays[i]
	}

	json.NewEncoder(w).Encode(map[string][]relay{
		"relays": relays,
	})
}

func handlePostRequest(w http.ResponseWriter, r *http.Request) {
	var newRelay relay
	err := json.NewDecoder(r.Body).Decode(&newRelay)
	r.Body.Close()

	if err != nil {
		if debug {
			log.Println("Failed to parse payload")
		}
		http.Error(w, err.Error(), 500)
		return
	}

	for _, current := range permanentRelays {
		if current == newRelay {
			if debug {
				log.Println("Asked to add a relay", newRelay, "which exists in permanent list")
			}
			http.Error(w, "Invalid request", 500)
			return
		}
	}

	uri, err := url.Parse(newRelay.URL)
	if err != nil {
		if debug {
			log.Println("Failed to parse URI", newRelay.URL)
		}
		http.Error(w, err.Error(), 500)
		return
	}

	reschan := make(chan result)

	select {
	case requests <- request{newRelay, uri, reschan}:
		result := <-reschan
		if result.err != nil {
			http.Error(w, result.err.Error(), 500)
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		json.NewEncoder(w).Encode(map[string]time.Duration{
			"evictionIn": result.eviction,
		})

	default:
		if debug {
			log.Println("Dropping request")
		}
		w.WriteHeader(429)
	}
}

func loadPermanentRelays() {
	path, err := osext.ExecutableFolder()
	if err != nil {
		log.Println("Failed to locate executable directory")
		return
	}

	content, err := ioutil.ReadFile(filepath.Join(path, "relays"))
	if err != nil {
		return
	}

	for _, line := range strings.Split(string(content), "\n") {
		if len(line) == 0 {
			continue
		}
		permanentRelays = append(permanentRelays, relay{
			URL: line,
		})
		if debug {
			log.Println("Adding permanent relay", line)
		}
	}
}

func loadOrCreateTestCertificate() {
	certFile, keyFile := filepath.Join(binDir, "cert.pem"), filepath.Join(binDir, "key.pem")
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err == nil {
		testCert = []tls.Certificate{cert}
		return
	}

	cert, err = tlsutil.NewCertificate(certFile, keyFile, "relaypoolsrv", 3072)
	if err != nil {
		log.Fatalln("Failed to create test X509 key pair:", err)
	}
	testCert = []tls.Certificate{cert}
}

func requestProcessor() {
	for request := range requests {
		if debug {
			log.Println("Request for", request.relay)
		}
		if !client.TestRelay(request.uri, testCert, 250*time.Millisecond, 4) {
			if debug {
				log.Println("Test for relay", request.relay, "failed")
			}
			request.result <- result{fmt.Errorf("test failed"), 0}
			continue
		}

		mut.Lock()
		timer, ok := evictionTimers[request.relay.URL]
		if ok {
			if debug {
				log.Println("Stopping existing timer for", request.relay)
			}
			timer.Stop()
		}

		for _, current := range knownRelays {
			if current == request.relay {
				if debug {
					log.Println("Relay", request.relay, "already exists")
				}
				goto found
			}
		}

		if debug {
			log.Println("Adding new relay", request.relay)
		}
		knownRelays = append(knownRelays, request.relay)

	found:
		evictionTimers[request.relay.URL] = time.AfterFunc(evictionTime, evict(request.relay))
		mut.Unlock()
		request.result <- result{nil, evictionTime}
	}

}

func evict(relay relay) func() {
	return func() {
		mut.Lock()
		defer mut.Unlock()
		if debug {
			log.Println("Evicting", relay)
		}
		for i, current := range knownRelays {
			if current == relay {
				if debug {
					log.Println("Evicted", relay)
				}
				last := len(knownRelays) - 1
				knownRelays[i] = knownRelays[last]
				knownRelays = knownRelays[:last]
			}
		}
		delete(evictionTimers, relay.URL)
	}
}
<commit_msg>Relays are matched by ip:port pairs<commit_after>// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).

package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/kardianos/osext"
	"github.com/syncthing/relaysrv/client"
	"github.com/syncthing/syncthing/lib/sync"
	"github.com/syncthing/syncthing/lib/tlsutil"
)

type relay struct {
	URL string `json:"url"`
	uri *url.URL
}

func (r relay) String() string {
	return r.URL
}

type request struct {
	relay  relay
	uri    *url.URL
	result chan result
}

type result struct {
	err      error
	eviction time.Duration
}

var (
	binDir       string
	testCert     []tls.Certificate
	listen       string
	dir          string
	evictionTime time.Duration
	debug        bool

	requests = make(chan request, 10)

	mut             sync.RWMutex           = sync.NewRWMutex()
	knownRelays     []relay                = make([]relay, 0)
	permanentRelays []relay                = make([]relay, 0)
	evictionTimers  map[string]*time.Timer = make(map[string]*time.Timer)
)

func main() {
	flag.StringVar(&listen, "listen", ":80", "Listen address")
	flag.StringVar(&dir, "keys", "", "Directory where http-cert.pem and http-key.pem is stored for TLS listening")
	flag.BoolVar(&debug, "debug", false, "Enable debug output")
	flag.DurationVar(&evictionTime, "eviction", time.Hour, "After how long the relay is evicted")

	flag.Parse()

	var listener net.Listener
	var err error

	binDir, err = osext.ExecutableFolder()
	if err != nil {
		log.Fatalln("Failed to locate executable directory")
	}

	loadPermanentRelays()
	loadOrCreateTestCertificate()
	go requestProcessor()

	if dir != "" {
		if debug {
			log.Println("Starting TLS listener on", listen)
		}
		certFile, keyFile := filepath.Join(dir, "http-cert.pem"), filepath.Join(dir, "http-key.pem")
		cert, err := tls.LoadX509KeyPair(certFile, keyFile)
		if err != nil {
			log.Fatalln("Failed to load HTTP X509 key pair:", err)
		}

		tlsCfg := &tls.Config{
			Certificates: []tls.Certificate{cert},
			MinVersion:   tls.VersionTLS10, // No SSLv3
			CipherSuites: []uint16{
				// No RC4
				tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
				tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
				tls.TLS_RSA_WITH_AES_128_CBC_SHA,
				tls.TLS_RSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA,
				tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA,
			},
		}

		listener, err = tls.Listen("tcp", listen, tlsCfg)
	} else {
		if debug {
			log.Println("Starting plain listener on", listen)
		}
		listener, err = net.Listen("tcp", listen)
	}

	if err != nil {
		log.Fatalln("listen:", err)
	}

	handler := http.NewServeMux()
	handler.HandleFunc("/", handleRequest)

	srv := http.Server{
		Handler:     handler,
		ReadTimeout: 10 * time.Second,
	}

	err = srv.Serve(listener)
	if err != nil {
		log.Fatalln("serve:", err)
	}
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		handleGetRequest(w, r)
	case "POST":
		handlePostRequest(w, r)
	default:
		if debug {
			log.Println("Unhandled HTTP method", r.Method)
		}
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

func handleGetRequest(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	mut.RLock()
	relays := append(permanentRelays, knownRelays...)
	mut.RUnlock()

	// Shuffle
	for i := range relays {
		j := rand.Intn(i + 1)
		relays[i], relays[j] = relays[j], relays[i]
	}

	json.NewEncoder(w).Encode(map[string][]relay{
		"relays": relays,
	})
}

func handlePostRequest(w http.ResponseWriter, r *http.Request) {
	var newRelay relay
	err := json.NewDecoder(r.Body).Decode(&newRelay)
	r.Body.Close()

	if err != nil {
		if debug {
			log.Println("Failed to parse payload")
		}
		http.Error(w, err.Error(), 500)
		return
	}

	uri, err := url.Parse(newRelay.URL)
	if err != nil {
		if debug {
			log.Println("Failed to parse URI", newRelay.URL)
		}
		http.Error(w, err.Error(), 500)
		return
	}

	host, port, err := net.SplitHostPort(uri.Host)
	if err != nil {
		if debug {
			log.Println("Failed to split URI", newRelay.URL)
		}
		http.Error(w, err.Error(), 500)
		return
	}

	// The client did not provide an IP address, work it out.
	if host == "" {
		rhost, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			if debug {
				log.Println("Failed to split remote address", r.RemoteAddr)
			}
			http.Error(w, err.Error(), 500)
			return
		}
		uri.Host = net.JoinHostPort(rhost, port)
		newRelay.URL = uri.String()
	}
	newRelay.uri = uri

	for _, current := range permanentRelays {
		if current.uri.Host == newRelay.uri.Host {
			if debug {
				log.Println("Asked to add a relay", newRelay, "which exists in permanent list")
			}
			http.Error(w, "Invalid request", 500)
			return
		}
	}

	reschan := make(chan result)

	select {
	case requests <- request{newRelay, uri, reschan}:
		result := <-reschan
		if result.err != nil {
			http.Error(w, result.err.Error(), 500)
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		json.NewEncoder(w).Encode(map[string]time.Duration{
			"evictionIn": result.eviction,
		})

	default:
		if debug {
			log.Println("Dropping request")
		}
		w.WriteHeader(429)
	}
}

func loadPermanentRelays() {
	path, err := osext.ExecutableFolder()
	if err != nil {
		log.Println("Failed to locate executable directory")
		return
	}

	content, err := ioutil.ReadFile(filepath.Join(path, "relays"))
	if err != nil {
		return
	}

	for _, line := range strings.Split(string(content), "\n") {
		if len(line) == 0 {
			continue
		}

		uri, err := url.Parse(line)
		if err != nil {
			if debug {
				log.Println("Skipping permanent relay", line, "due to parse error", err)
			}
			continue
		}

		permanentRelays = append(permanentRelays, relay{
			URL: line,
			uri: uri,
		})
		if debug {
			log.Println("Adding permanent relay", line)
		}
	}
}

func loadOrCreateTestCertificate() {
	certFile, keyFile := filepath.Join(binDir, "cert.pem"), filepath.Join(binDir, "key.pem")
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err == nil {
		testCert = []tls.Certificate{cert}
		return
	}

	cert, err = tlsutil.NewCertificate(certFile, keyFile, "relaypoolsrv", 3072)
	if err != nil {
		log.Fatalln("Failed to create test X509 key pair:", err)
	}
	testCert = []tls.Certificate{cert}
}

func requestProcessor() {
	for request := range requests {
		if debug {
			log.Println("Request for", request.relay)
		}
		if !client.TestRelay(request.uri, testCert, 250*time.Millisecond, 4) {
			if debug {
				log.Println("Test for relay", request.relay, "failed")
			}
			request.result <- result{fmt.Errorf("test failed"), 0}
			continue
		}

		mut.Lock()
		timer, ok := evictionTimers[request.relay.uri.Host]
		if ok {
			if debug {
				log.Println("Stopping existing timer for", request.relay)
			}
			timer.Stop()
		}

		for _, current := range knownRelays {
			if current.uri.Host == request.relay.uri.Host {
				if debug {
					log.Println("Relay", request.relay, "already exists")
				}
				goto found
			}
		}

		if debug {
			log.Println("Adding new relay", request.relay)
		}
		knownRelays = append(knownRelays, request.relay)

	found:
		evictionTimers[request.relay.uri.Host] = time.AfterFunc(evictionTime, evict(request.relay))
		mut.Unlock()
		request.result <- result{nil, evictionTime}
	}

}

func evict(relay relay) func() {
	return func() {
		mut.Lock()
		defer mut.Unlock()
		if debug {
			log.Println("Evicting", relay)
		}
		for i, current := range knownRelays {
			if current.uri.Host == relay.uri.Host {
				if debug {
					log.Println("Evicted", relay)
				}
				last := len(knownRelays) - 1
				knownRelays[i] = knownRelays[last]
				knownRelays = knownRelays[:last]
			}
		}
		delete(evictionTimers, relay.uri.Host)
	}
}
<|endoftext|>


<commit_before>package main

import "encoding/json"

func testSanitizers(s string) {
	jsonData, _ := json.Marshal(s)
	ignore(make([]byte, len(jsonData)+1)) // NOT OK: data might be big

	ignore(make([]byte, int64(len(jsonData))+1)) // OK: sanitized by widening to 64 bits

	if len(jsonData) < 1000 {
		ignore(make([]byte, len(jsonData)+1)) // OK: there is an upper bound check on len(jsonData)
	}

	{
		newlength := len(jsonData) + 2 // OK: there is an upper bound check which dominates `make`
		_ := newlength - 1
		if newlength < 1000 {
			ignore(make([]byte, newlength))
		}
	}

	{
		newlength := len(jsonData) + 3 // NOT OK: newlength is changed after the upper bound check (even though it's made smaller)
		if newlength < 1000 {
			newlength = newlength - 1
			ignore(make([]byte, newlength))
		}
	}

	{
		newlength := len(jsonData) + 4 // NOT OK: there is an upper bound check but it doesn't dominate `make`
		if newlength < 1000 {
			_ := newlength + 2
		}
		ignore(make([]byte, newlength))
	}

	{
		newlength := len(jsonData) + 5 // OK: there is an upper bound check which dominates `make`
		if newlength > 1000 {
			return
		}
		ignore(make([]byte, newlength))
	}
}
<commit_msg>Fix frontend errors in test.<commit_after>package main

import "encoding/json"

func testSanitizers(s string) {
	jsonData, _ := json.Marshal(s)
	ignore(make([]byte, len(jsonData)+1)) // NOT OK: data might be big

	ignore(make([]byte, int64(len(jsonData))+1)) // OK: sanitized by widening to 64 bits

	if len(jsonData) < 1000 {
		ignore(make([]byte, len(jsonData)+1)) // OK: there is an upper bound check on len(jsonData)
	}

	{
		newlength := len(jsonData) + 2 // OK: there is an upper bound check which dominates `make`
		ignore(newlength - 1)
		if newlength < 1000 {
			ignore(make([]byte, newlength))
		}
	}

	{
		newlength := len(jsonData) + 3 // NOT OK: newlength is changed after the upper bound check (even though it's made smaller)
		if newlength < 1000 {
			newlength = newlength - 1
			ignore(make([]byte, newlength))
		}
	}

	{
		newlength := len(jsonData) + 4 // NOT OK: there is an upper bound check but it doesn't dominate `make`
		if newlength < 1000 {
			ignore(newlength + 2)
		}
		ignore(make([]byte, newlength))
	}

	{
		newlength := len(jsonData) + 5 // OK: there is an upper bound check which dominates `make`
		if newlength > 1000 {
			return
		}
		ignore(make([]byte, newlength))
	}
}
<|endoftext|>


<commit_before>// Copyright The OpenTelemetry Authors
//
// 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 fluentforwardreceiver // import "github.com/open-telemetry/opentelemetry-collector-contrib/receiver/fluentforwardreceiver"

import (
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"time"

	"github.com/tinylib/msgp/msgp"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/plog"
)

const tagAttributeKey = "fluent.tag"

// Most of this logic is derived directly from
// https://github.com/fluent/fluentd/wiki/Forward-Protocol-Specification-v1,
// which describes the fields in much greater detail.

type Event interface {
	DecodeMsg(dc *msgp.Reader) error
	LogRecords() plog.LogRecordSlice
	Chunk() string
	Compressed() string
}

type OptionsMap map[string]interface{}

// Chunk returns the `chunk` option or blank string if it was not set.
func (om OptionsMap) Chunk() string {
	c, _ := om["chunk"].(string)
	return c
}

func (om OptionsMap) Compressed() string {
	compressed, _ := om["compressed"].(string)
	return compressed
}

type EventMode int

type Peeker interface {
	Peek(n int) ([]byte, error)
}

// Values for enum EventMode.
const (
	UnknownMode EventMode = iota
	MessageMode
	ForwardMode
	PackedForwardMode
)

func (em EventMode) String() string {
	switch em {
	case UnknownMode:
		return "unknown"
	case MessageMode:
		return "message"
	case ForwardMode:
		return "forward"
	case PackedForwardMode:
		return "packedforward"
	default:
		panic("programmer bug")
	}
}

// parseInterfaceToMap takes map of interface objects and returns
// AttributeValueMap
func parseInterfaceToMap(msi map[string]interface{}, dest pcommon.Value) {
	am := dest.SetEmptyMap()
	am.EnsureCapacity(len(msi))
	for k, value := range msi {
		parseToAttributeValue(value, am.PutEmpty(k))
	}
}

// parseInterfaceToArray takes array of interface objects and returns
// AttributeValueArray
func parseInterfaceToArray(ai []interface{}, dest pcommon.Value) {
	av := dest.SetEmptySlice()
	av.EnsureCapacity(len(ai))
	for _, value := range ai {
		parseToAttributeValue(value, av.AppendEmpty())
	}
}

// parseToAttributeValue converts interface object to AttributeValue
func parseToAttributeValue(val interface{}, dest pcommon.Value) {
	// See https://github.com/tinylib/msgp/wiki/Type-Mapping-Rules
	switch r := val.(type) {
	case bool:
		dest.SetBool(r)
	case string:
		dest.SetStr(r)
	case uint64:
		dest.SetInt(int64(r))
	case int64:
		dest.SetInt(r)
	// Sometimes strings come in as bytes array
	case []byte:
		dest.SetStr(string(r))
	case map[string]interface{}:
		parseInterfaceToMap(r, dest)
	case []interface{}:
		parseInterfaceToArray(r, dest)
	case float32:
		dest.SetDouble(float64(r))
	case float64:
		dest.SetDouble(r)
	case nil:
	default:
		dest.SetStr(fmt.Sprintf("%v", val))
	}
}

func timeFromTimestamp(ts interface{}) (time.Time, error) {
	switch v := ts.(type) {
	case uint64:
		return time.Unix(int64(v), 0), nil
	case int64:
		return time.Unix(v, 0), nil
	case *eventTimeExt:
		return time.Time(*v), nil
	default:
		return time.Time{}, fmt.Errorf("unknown type of value: %v", ts)
	}
}

func decodeTimestampToLogRecord(dc *msgp.Reader, lr plog.LogRecord) error {
	tsIntf, err := dc.ReadIntf()
	if err != nil {
		return msgp.WrapError(err, "Time")
	}

	ts, err := timeFromTimestamp(tsIntf)
	if err != nil {
		return msgp.WrapError(err, "Time")
	}

	lr.SetTimestamp(pcommon.NewTimestampFromTime(ts))
	return nil
}

func parseRecordToLogRecord(dc *msgp.Reader, lr plog.LogRecord) error {
	attrs := lr.Attributes()

	recordLen, err := dc.ReadMapHeader()
	if err != nil {
		return msgp.WrapError(err, "Record")
	}

	for recordLen > 0 {
		recordLen--
		key, err := dc.ReadString()
		if err != nil {
			// The protocol doesn't specify this but apparently some map keys
			// can be binary type instead of string
			keyBytes, keyBytesErr := dc.ReadBytes(nil)
			if keyBytesErr != nil {
				return msgp.WrapError(keyBytesErr, "Record")
			}
			key = string(keyBytes)
		}
		val, err := dc.ReadIntf()
		if err != nil {
			return msgp.WrapError(err, "Record", key)
		}

		// fluentd uses message, fluentbit log.
		if key == "message" || key == "log" {
			parseToAttributeValue(val, lr.Body())
		} else {
			parseToAttributeValue(val, attrs.PutEmpty(key))
		}
	}

	return nil
}

type MessageEventLogRecord struct {
	plog.LogRecordSlice
	OptionsMap
}

func (melr *MessageEventLogRecord) LogRecords() plog.LogRecordSlice {
	return melr.LogRecordSlice
}

func (melr *MessageEventLogRecord) DecodeMsg(dc *msgp.Reader) error {
	melr.LogRecordSlice = plog.NewLogRecordSlice()
	log := melr.LogRecordSlice.AppendEmpty()

	var arrLen uint32
	var err error

	arrLen, err = dc.ReadArrayHeader()
	if err != nil {
		return msgp.WrapError(err)
	}
	if arrLen > 4 || arrLen < 3 {
		return msgp.ArrayError{Wanted: 3, Got: arrLen}
	}

	tag, err := dc.ReadString()
	if err != nil {
		return msgp.WrapError(err, "Tag")
	}

	attrs := log.Attributes()
	attrs.PutString(tagAttributeKey, tag)

	err = decodeTimestampToLogRecord(dc, log)
	if err != nil {
		return msgp.WrapError(err, "Time")
	}

	err = parseRecordToLogRecord(dc, log)
	if err != nil {
		return err
	}

	if arrLen == 4 {
		melr.OptionsMap, err = parseOptions(dc)
		if err != nil {
			return err
		}
	}
	return nil
}

func parseOptions(dc *msgp.Reader) (OptionsMap, error) {
	var optionLen uint32
	optionLen, err := dc.ReadMapHeader()
	if err != nil {
		return nil, msgp.WrapError(err, "Option")
	}
	out := make(OptionsMap, optionLen)

	for optionLen > 0 {
		optionLen--
		key, err := dc.ReadString()
		if err != nil {
			return nil, msgp.WrapError(err, "Option")
		}
		val, err := dc.ReadIntf()
		if err != nil {
			return nil, msgp.WrapError(err, "Option", key)
		}
		out[key] = val
	}
	return out, nil
}

type ForwardEventLogRecords struct {
	plog.LogRecordSlice
	OptionsMap
}

func (fe *ForwardEventLogRecords) LogRecords() plog.LogRecordSlice {
	return fe.LogRecordSlice
}

func (fe *ForwardEventLogRecords) DecodeMsg(dc *msgp.Reader) (err error) {
	fe.LogRecordSlice = plog.NewLogRecordSlice()

	var arrLen uint32
	arrLen, err = dc.ReadArrayHeader()
	if err != nil {
		err = msgp.WrapError(err)
		return
	}
	if arrLen < 2 || arrLen > 3 {
		err = msgp.ArrayError{Wanted: 2, Got: arrLen}
		return
	}

	tag, err := dc.ReadString()
	if err != nil {
		return msgp.WrapError(err, "Tag")
	}

	entryLen, err := dc.ReadArrayHeader()
	if err != nil {
		err = msgp.WrapError(err, "Record")
		return
	}

	fe.LogRecordSlice.EnsureCapacity(int(entryLen))
	for i := 0; i < int(entryLen); i++ {
		lr := fe.LogRecordSlice.AppendEmpty()

		err = parseEntryToLogRecord(dc, lr)
		if err != nil {
			return msgp.WrapError(err, "Entries", i)
		}
		fe.LogRecordSlice.At(i).Attributes().PutString(tagAttributeKey, tag)
	}

	if arrLen == 3 {
		fe.OptionsMap, err = parseOptions(dc)
		if err != nil {
			return err
		}
	}

	return
}

func parseEntryToLogRecord(dc *msgp.Reader, lr plog.LogRecord) error {
	arrLen, err := dc.ReadArrayHeader()
	if err != nil {
		return msgp.WrapError(err)
	}
	if arrLen != 2 {
		return msgp.ArrayError{Wanted: 2, Got: arrLen}
	}

	err = decodeTimestampToLogRecord(dc, lr)
	if err != nil {
		return msgp.WrapError(err, "Time")
	}

	return parseRecordToLogRecord(dc, lr)
}

type PackedForwardEventLogRecords struct {
	plog.LogRecordSlice
	OptionsMap
}

func (pfe *PackedForwardEventLogRecords) LogRecords() plog.LogRecordSlice {
	return pfe.LogRecordSlice
}

// DecodeMsg implements msgp.Decodable.  This was originally code generated but
// then manually copied here in order to handle the optional Options field.
func (pfe *PackedForwardEventLogRecords) DecodeMsg(dc *msgp.Reader) error {
	pfe.LogRecordSlice = plog.NewLogRecordSlice()

	arrLen, err := dc.ReadArrayHeader()
	if err != nil {
		return msgp.WrapError(err)
	}
	if arrLen < 2 || arrLen > 3 {
		return msgp.ArrayError{Wanted: 2, Got: arrLen}
	}

	tag, err := dc.ReadString()
	if err != nil {
		return msgp.WrapError(err, "Tag")
	}

	entriesFirstByte, err := dc.R.Peek(1)
	if err != nil {
		return msgp.WrapError(err, "EntriesRaw")
	}

	entriesType := msgp.NextType(entriesFirstByte)
	// We have to read out the entries raw all the way first because we don't
	// know whether it is compressed or not until we read the options map which
	// comes after.  I guess we could use some kind of detection logic to
	// determine if it is gzipped by peeking and just ignoring options, but
	// this seems simpler for now.
	var entriesRaw []byte
	switch entriesType {
	case msgp.StrType:
		var entriesStr string
		entriesStr, err = dc.ReadString()
		if err != nil {
			return msgp.WrapError(err, "EntriesRaw")
		}
		entriesRaw = []byte(entriesStr)
	case msgp.BinType:
		entriesRaw, err = dc.ReadBytes(nil)
		if err != nil {
			return msgp.WrapError(err, "EntriesRaw")
		}
	default:
		return msgp.WrapError(fmt.Errorf("invalid type %d", entriesType), "EntriesRaw")
	}

	if arrLen == 3 {
		pfe.OptionsMap, err = parseOptions(dc)
		if err != nil {
			return err
		}
	}

	err = pfe.parseEntries(entriesRaw, pfe.Compressed() == "gzip", tag)
	if err != nil {
		return err
	}

	return nil
}

func (pfe *PackedForwardEventLogRecords) parseEntries(entriesRaw []byte, isGzipped bool, tag string) error {
	var reader io.Reader
	reader = bytes.NewReader(entriesRaw)

	if isGzipped {
		var err error
		reader, err = gzip.NewReader(reader)
		if err != nil {
			return err
		}
		defer reader.(*gzip.Reader).Close()
	}

	msgpReader := msgp.NewReader(reader)
	for {
		lr := plog.NewLogRecord()
		err := parseEntryToLogRecord(msgpReader, lr)
		if err != nil {
			if errors.Is(msgp.Cause(err), io.EOF) {
				return nil
			}
			return err
		}

		lr.Attributes().PutString(tagAttributeKey, tag)

		lr.MoveTo(pfe.LogRecordSlice.AppendEmpty())
	}
}
<commit_msg>[chore] fix small nits in fluentforwardreceiver (#14685)<commit_after>// Copyright The OpenTelemetry Authors
//
// 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 fluentforwardreceiver // import "github.com/open-telemetry/opentelemetry-collector-contrib/receiver/fluentforwardreceiver"

import (
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"time"

	"github.com/tinylib/msgp/msgp"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/plog"
)

const tagAttributeKey = "fluent.tag"

// Most of this logic is derived directly from
// https://github.com/fluent/fluentd/wiki/Forward-Protocol-Specification-v1,
// which describes the fields in much greater detail.

type Event interface {
	DecodeMsg(dc *msgp.Reader) error
	LogRecords() plog.LogRecordSlice
	Chunk() string
	Compressed() string
}

type OptionsMap map[string]interface{}

// Chunk returns the `chunk` option or blank string if it was not set.
func (om OptionsMap) Chunk() string {
	c, _ := om["chunk"].(string)
	return c
}

func (om OptionsMap) Compressed() string {
	compressed, _ := om["compressed"].(string)
	return compressed
}

type EventMode int

type Peeker interface {
	Peek(n int) ([]byte, error)
}

// Values for enum EventMode.
const (
	UnknownMode EventMode = iota
	MessageMode
	ForwardMode
	PackedForwardMode
)

func (em EventMode) String() string {
	switch em {
	case UnknownMode:
		return "unknown"
	case MessageMode:
		return "message"
	case ForwardMode:
		return "forward"
	case PackedForwardMode:
		return "packedforward"
	default:
		panic("programmer bug")
	}
}

// parseInterfaceToMap takes map of interface objects and returns
// AttributeValueMap
func parseInterfaceToMap(msi map[string]interface{}, dest pcommon.Value) {
	am := dest.SetEmptyMap()
	am.EnsureCapacity(len(msi))
	for k, value := range msi {
		parseToAttributeValue(value, am.PutEmpty(k))
	}
}

// parseInterfaceToArray takes array of interface objects and returns
// AttributeValueArray
func parseInterfaceToArray(ai []interface{}, dest pcommon.Value) {
	av := dest.SetEmptySlice()
	av.EnsureCapacity(len(ai))
	for _, value := range ai {
		parseToAttributeValue(value, av.AppendEmpty())
	}
}

// parseToAttributeValue converts interface object to AttributeValue
func parseToAttributeValue(val interface{}, dest pcommon.Value) {
	// See https://github.com/tinylib/msgp/wiki/Type-Mapping-Rules
	switch r := val.(type) {
	case bool:
		dest.SetBool(r)
	case string:
		dest.SetStr(r)
	case uint64:
		dest.SetInt(int64(r))
	case int64:
		dest.SetInt(r)
	// Sometimes strings come in as bytes array
	case []byte:
		dest.SetStr(string(r))
	case map[string]interface{}:
		parseInterfaceToMap(r, dest)
	case []interface{}:
		parseInterfaceToArray(r, dest)
	case float32:
		dest.SetDouble(float64(r))
	case float64:
		dest.SetDouble(r)
	case nil:
	default:
		dest.SetStr(fmt.Sprintf("%v", val))
	}
}

func timeFromTimestamp(ts interface{}) (time.Time, error) {
	switch v := ts.(type) {
	case uint64:
		return time.Unix(int64(v), 0), nil
	case int64:
		return time.Unix(v, 0), nil
	case *eventTimeExt:
		return time.Time(*v), nil
	default:
		return time.Time{}, fmt.Errorf("unknown type of value: %v", ts)
	}
}

func parseRecordToLogRecord(dc *msgp.Reader, lr plog.LogRecord) error {
	tsIntf, err := dc.ReadIntf()
	if err != nil {
		return msgp.WrapError(err, "Time")
	}

	ts, err := timeFromTimestamp(tsIntf)
	if err != nil {
		return msgp.WrapError(err, "Time")
	}

	lr.SetTimestamp(pcommon.NewTimestampFromTime(ts))

	recordLen, err := dc.ReadMapHeader()
	if err != nil {
		return msgp.WrapError(err, "Record")
	}

	for recordLen > 0 {
		recordLen--
		key, err := dc.ReadString()
		if err != nil {
			// The protocol doesn't specify this but apparently some map keys
			// can be binary type instead of string
			keyBytes, keyBytesErr := dc.ReadBytes(nil)
			if keyBytesErr != nil {
				return msgp.WrapError(keyBytesErr, "Record")
			}
			key = string(keyBytes)
		}
		val, err := dc.ReadIntf()
		if err != nil {
			return msgp.WrapError(err, "Record", key)
		}

		// fluentd uses message, fluentbit log.
		if key == "message" || key == "log" {
			parseToAttributeValue(val, lr.Body())
		} else {
			parseToAttributeValue(val, lr.Attributes().PutEmpty(key))
		}
	}

	return nil
}

type MessageEventLogRecord struct {
	plog.LogRecordSlice
	OptionsMap
}

func (melr *MessageEventLogRecord) LogRecords() plog.LogRecordSlice {
	return melr.LogRecordSlice
}

func (melr *MessageEventLogRecord) DecodeMsg(dc *msgp.Reader) error {
	melr.LogRecordSlice = plog.NewLogRecordSlice()
	var arrLen uint32
	var err error

	arrLen, err = dc.ReadArrayHeader()
	if err != nil {
		return err
	}
	if arrLen > 4 || arrLen < 3 {
		return msgp.ArrayError{Wanted: 3, Got: arrLen}
	}

	tag, err := dc.ReadString()
	if err != nil {
		return msgp.WrapError(err, "Tag")
	}

	log := melr.LogRecordSlice.AppendEmpty()
	attrs := log.Attributes()
	attrs.PutString(tagAttributeKey, tag)
	err = parseRecordToLogRecord(dc, log)
	if err != nil {
		return err
	}

	if arrLen == 4 {
		melr.OptionsMap, err = parseOptions(dc)
		return err
	}
	return nil
}

func parseOptions(dc *msgp.Reader) (OptionsMap, error) {
	var optionLen uint32
	optionLen, err := dc.ReadMapHeader()
	if err != nil {
		return nil, msgp.WrapError(err, "Option")
	}
	out := make(OptionsMap, optionLen)

	for optionLen > 0 {
		optionLen--
		key, err := dc.ReadString()
		if err != nil {
			return nil, msgp.WrapError(err, "Option")
		}
		val, err := dc.ReadIntf()
		if err != nil {
			return nil, msgp.WrapError(err, "Option", key)
		}
		out[key] = val
	}
	return out, nil
}

type ForwardEventLogRecords struct {
	plog.LogRecordSlice
	OptionsMap
}

func (fe *ForwardEventLogRecords) LogRecords() plog.LogRecordSlice {
	return fe.LogRecordSlice
}

func (fe *ForwardEventLogRecords) DecodeMsg(dc *msgp.Reader) error {
	fe.LogRecordSlice = plog.NewLogRecordSlice()

	arrLen, err := dc.ReadArrayHeader()
	if err != nil {
		return err
	}
	if arrLen < 2 || arrLen > 3 {
		return msgp.ArrayError{Wanted: 2, Got: arrLen}
	}

	tag, err := dc.ReadString()
	if err != nil {
		return msgp.WrapError(err, "Tag")
	}

	entryLen, err := dc.ReadArrayHeader()
	if err != nil {
		return msgp.WrapError(err, "Record")
	}

	fe.LogRecordSlice.EnsureCapacity(int(entryLen))
	for i := 0; i < int(entryLen); i++ {
		lr := fe.LogRecordSlice.AppendEmpty()

		err = parseEntryToLogRecord(dc, lr)
		if err != nil {
			return msgp.WrapError(err, "Entries", i)
		}
		lr.Attributes().PutString(tagAttributeKey, tag)
	}

	if arrLen == 3 {
		fe.OptionsMap, err = parseOptions(dc)
		return err
	}

	return nil
}

func parseEntryToLogRecord(dc *msgp.Reader, lr plog.LogRecord) error {
	arrLen, err := dc.ReadArrayHeader()
	if err != nil {
		return err
	}
	if arrLen != 2 {
		return msgp.ArrayError{Wanted: 2, Got: arrLen}
	}
	return parseRecordToLogRecord(dc, lr)
}

type PackedForwardEventLogRecords struct {
	plog.LogRecordSlice
	OptionsMap
}

func (pfe *PackedForwardEventLogRecords) LogRecords() plog.LogRecordSlice {
	return pfe.LogRecordSlice
}

// DecodeMsg implements msgp.Decodable.  This was originally code generated but
// then manually copied here in order to handle the optional Options field.
func (pfe *PackedForwardEventLogRecords) DecodeMsg(dc *msgp.Reader) error {
	pfe.LogRecordSlice = plog.NewLogRecordSlice()

	arrLen, err := dc.ReadArrayHeader()
	if err != nil {
		return err
	}
	if arrLen < 2 || arrLen > 3 {
		return msgp.ArrayError{Wanted: 2, Got: arrLen}
	}

	tag, err := dc.ReadString()
	if err != nil {
		return msgp.WrapError(err, "Tag")
	}

	entriesFirstByte, err := dc.R.Peek(1)
	if err != nil {
		return msgp.WrapError(err, "EntriesRaw")
	}

	entriesType := msgp.NextType(entriesFirstByte)
	// We have to read out the entries raw all the way first because we don't
	// know whether it is compressed or not until we read the options map which
	// comes after.  I guess we could use some kind of detection logic to
	// determine if it is gzipped by peeking and just ignoring options, but
	// this seems simpler for now.
	var entriesRaw []byte
	switch entriesType {
	case msgp.StrType:
		var entriesStr string
		entriesStr, err = dc.ReadString()
		if err != nil {
			return msgp.WrapError(err, "EntriesRaw")
		}
		entriesRaw = []byte(entriesStr)
	case msgp.BinType:
		entriesRaw, err = dc.ReadBytes(nil)
		if err != nil {
			return msgp.WrapError(err, "EntriesRaw")
		}
	default:
		return msgp.WrapError(fmt.Errorf("invalid type %d", entriesType), "EntriesRaw")
	}

	if arrLen == 3 {
		pfe.OptionsMap, err = parseOptions(dc)
		if err != nil {
			return err
		}
	}

	err = pfe.parseEntries(entriesRaw, pfe.Compressed() == "gzip", tag)
	if err != nil {
		return err
	}

	return nil
}

func (pfe *PackedForwardEventLogRecords) parseEntries(entriesRaw []byte, isGzipped bool, tag string) error {
	var reader io.Reader
	reader = bytes.NewReader(entriesRaw)

	if isGzipped {
		var err error
		reader, err = gzip.NewReader(reader)
		if err != nil {
			return err
		}
		defer reader.(*gzip.Reader).Close()
	}

	msgpReader := msgp.NewReader(reader)
	// Allocate only once, since the MoveTo cleans the lr, so we can reuse.
	lr := plog.NewLogRecord()
	for {
		err := parseEntryToLogRecord(msgpReader, lr)
		if err != nil {
			if errors.Is(msgp.Cause(err), io.EOF) {
				return nil
			}
			return err
		}
		lr.Attributes().PutString(tagAttributeKey, tag)
		lr.MoveTo(pfe.LogRecordSlice.AppendEmpty())
	}
}
<|endoftext|>