<commit_before>// The supervisor acts as the glue between the cli and the storm topology
// - The storm topology communicates with the supervisor in order to determine settings, etc.
// - The cli communicates with the supervisor to modify settings, get results, etc.
// @author Robin Verlangen

package main

import (
	"bufio"
	"encoding/base64"
	"flag"
	"fmt"
	"github.com/RobinUS2/golang-jresp"
	"github.com/julienschmidt/httprouter"
	"log"
	"net/http"
	"strings"
)

var serverPort int
var basicAuthUsr string
var basicAuthPwd string
var dbFile string
var filterManager *FilterManager

func init() {
	flag.IntVar(&serverPort, "port", 1525, "Server port")
	flag.StringVar(&basicAuthUsr, "auth-user", "cloud", "Username")
	flag.StringVar(&basicAuthPwd, "auth-password", "pelican", "Password")
	flag.StringVar(&dbFile, "db-file", "cloudpelican_lsd_supervisor.db", "Database file")
	flag.Parse()
}

func main() {
	// Filter manager
	filterManager = NewFilterManager()

	// Routing
	router := httprouter.New()

	// Docs
	router.GET("/", GetHome)

	// Filters
	router.POST("/filter", PostFilter)                // Create new filter
	router.GET("/filter/:id/result", GetFilterResult) // Get results of a single filter
	router.PUT("/filter/:id/result", PutFilterResult) // Store new results into a filter
	router.GET("/filter", GetFilter)                  // Get all filters
	router.DELETE("/filter/:id", DeleteFilter)        // Delete a filter

	// Start webserver
	log.Println(fmt.Sprintf("Starting supervisor service at port %d", serverPort))
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", serverPort), router))
}

func GetHome(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	jresp.Set("hello", "This is the CloudPelican supervisor")
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func PostFilter(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()

	// Validate
	regex := strings.TrimSpace(r.URL.Query().Get("regex"))
	if len(regex) < 1 {
		jresp.Error("Please provide a regex")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	name := strings.TrimSpace(r.URL.Query().Get("name"))
	if len(name) < 1 {
		jresp.Error("Please provide a name")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}

	// Create filter
	id, err := filterManager.CreateFilter(name, r.RemoteAddr, regex)
	if err != nil {
		jresp.Error(fmt.Sprintf("Failed to create filter: %s", err))
		fmt.Fprint(w, jresp.ToString(false))
		return
	}

	// OK :)
	jresp.Set("filter_id", id)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func GetFilterResult(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	id := strings.TrimSpace(ps.ByName("id"))
	if len(id) < 1 {
		jresp.Error("Please provide an ID")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	filter := filterManager.GetFilter(id)
	if filter == nil {
		jresp.Error(fmt.Sprintf("Filter %s not found", id))
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	jresp.Set("results", filter.Results)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func PutFilterResult(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	id := strings.TrimSpace(ps.ByName("id"))
	if len(id) < 1 {
		jresp.Error("Please provide an ID")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	filter := filterManager.GetFilter(id)
	if filter == nil {
		jresp.Error(fmt.Sprintf("Filter %s not found", id))
		fmt.Fprint(w, jresp.ToString(false))
		return
	}

	// Read body
	scanner := bufio.NewScanner(r.Body)
	scanner.Split(bufio.ScanLines)
	var lines []string = make([]string, 0)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	// Add results
	res := filter.AddResults(lines)
	jresp.Set("ack", res)
	jresp.Set("lines", len(lines))
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func GetFilter(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	filters := filterManager.GetFilters()
	var filtersNoRes []*Filter = make([]*Filter, 0)
	for _, filter := range filters {
		filter.Results = nil
		filtersNoRes = append(filtersNoRes, filter)
	}
	jresp.Set("filters", filtersNoRes)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func DeleteFilter(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	id := strings.TrimSpace(ps.ByName("id"))
	if len(id) < 1 {
		jresp.Error("Please provide an ID")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	res := filterManager.DeleteFilter(id)
	jresp.Set("deleted", res)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func basicAuth(w http.ResponseWriter, r *http.Request) bool {
	if r.Header["Authorization"] == nil || len(r.Header["Authorization"]) < 1 {
		log.Printf("%s", r.Header)
		http.Error(w, "bad syntax a", http.StatusBadRequest)
		return false
	}
	auth := strings.SplitN(r.Header["Authorization"][0], " ", 2)

	if len(auth) != 2 || auth[0] != "Basic" {
		log.Printf("%s", r.Header)
		http.Error(w, "bad syntax b", http.StatusBadRequest)
		return false
	}

	payload, _ := base64.StdEncoding.DecodeString(auth[1])
	pair := strings.SplitN(string(payload), ":", 2)

	if len(pair) != 2 || !validateAuth(pair[0], pair[1]) {
		http.Error(w, "authorization failed", http.StatusUnauthorized)
		return false
	}
	return true
}

func validateAuth(username, password string) bool {
	if username == basicAuthUsr && password == basicAuthPwd {
		return true
	}
	return false
}
<commit_msg>Clear res after retrieve<commit_after>// The supervisor acts as the glue between the cli and the storm topology
// - The storm topology communicates with the supervisor in order to determine settings, etc.
// - The cli communicates with the supervisor to modify settings, get results, etc.
// @author Robin Verlangen

package main

import (
	"bufio"
	"encoding/base64"
	"flag"
	"fmt"
	"github.com/RobinUS2/golang-jresp"
	"github.com/julienschmidt/httprouter"
	"log"
	"net/http"
	"strings"
)

var serverPort int
var basicAuthUsr string
var basicAuthPwd string
var dbFile string
var filterManager *FilterManager

func init() {
	flag.IntVar(&serverPort, "port", 1525, "Server port")
	flag.StringVar(&basicAuthUsr, "auth-user", "cloud", "Username")
	flag.StringVar(&basicAuthPwd, "auth-password", "pelican", "Password")
	flag.StringVar(&dbFile, "db-file", "cloudpelican_lsd_supervisor.db", "Database file")
	flag.Parse()
}

func main() {
	// Filter manager
	filterManager = NewFilterManager()

	// Routing
	router := httprouter.New()

	// Docs
	router.GET("/", GetHome)

	// Filters
	router.POST("/filter", PostFilter)                // Create new filter
	router.GET("/filter/:id/result", GetFilterResult) // Get results of a single filter
	router.PUT("/filter/:id/result", PutFilterResult) // Store new results into a filter
	router.GET("/filter", GetFilter)                  // Get all filters
	router.DELETE("/filter/:id", DeleteFilter)        // Delete a filter

	// Start webserver
	log.Println(fmt.Sprintf("Starting supervisor service at port %d", serverPort))
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", serverPort), router))
}

func GetHome(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	jresp.Set("hello", "This is the CloudPelican supervisor")
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func PostFilter(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()

	// Validate
	regex := strings.TrimSpace(r.URL.Query().Get("regex"))
	if len(regex) < 1 {
		jresp.Error("Please provide a regex")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	name := strings.TrimSpace(r.URL.Query().Get("name"))
	if len(name) < 1 {
		jresp.Error("Please provide a name")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}

	// Create filter
	id, err := filterManager.CreateFilter(name, r.RemoteAddr, regex)
	if err != nil {
		jresp.Error(fmt.Sprintf("Failed to create filter: %s", err))
		fmt.Fprint(w, jresp.ToString(false))
		return
	}

	// OK :)
	jresp.Set("filter_id", id)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func GetFilterResult(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	id := strings.TrimSpace(ps.ByName("id"))
	if len(id) < 1 {
		jresp.Error("Please provide an ID")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	filter := filterManager.GetFilter(id)
	if filter == nil {
		jresp.Error(fmt.Sprintf("Filter %s not found", id))
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	filter.resultsMux.RLock()
	jresp.Set("results", filter.Results)
	filter.resultsMux.RUnlock()
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))

	// Clear results
	filter.resultsMux.Lock()
	filter.Results = make([]string, 0)
	filter.resultsMux.Unlock()
	filter.Save()
}

func PutFilterResult(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	id := strings.TrimSpace(ps.ByName("id"))
	if len(id) < 1 {
		jresp.Error("Please provide an ID")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	filter := filterManager.GetFilter(id)
	if filter == nil {
		jresp.Error(fmt.Sprintf("Filter %s not found", id))
		fmt.Fprint(w, jresp.ToString(false))
		return
	}

	// Read body
	scanner := bufio.NewScanner(r.Body)
	scanner.Split(bufio.ScanLines)
	var lines []string = make([]string, 0)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	// Add results
	res := filter.AddResults(lines)
	jresp.Set("ack", res)
	jresp.Set("lines", len(lines))
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func GetFilter(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	filters := filterManager.GetFilters()
	var filtersNoRes []*Filter = make([]*Filter, 0)
	for _, filter := range filters {
		filter.Results = nil
		filtersNoRes = append(filtersNoRes, filter)
	}
	jresp.Set("filters", filtersNoRes)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func DeleteFilter(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	if !basicAuth(w, r) {
		return
	}
	jresp := jresp.NewJsonResp()
	id := strings.TrimSpace(ps.ByName("id"))
	if len(id) < 1 {
		jresp.Error("Please provide an ID")
		fmt.Fprint(w, jresp.ToString(false))
		return
	}
	res := filterManager.DeleteFilter(id)
	jresp.Set("deleted", res)
	jresp.OK()
	fmt.Fprint(w, jresp.ToString(false))
}

func basicAuth(w http.ResponseWriter, r *http.Request) bool {
	if r.Header["Authorization"] == nil || len(r.Header["Authorization"]) < 1 {
		log.Printf("%s", r.Header)
		http.Error(w, "bad syntax a", http.StatusBadRequest)
		return false
	}
	auth := strings.SplitN(r.Header["Authorization"][0], " ", 2)

	if len(auth) != 2 || auth[0] != "Basic" {
		log.Printf("%s", r.Header)
		http.Error(w, "bad syntax b", http.StatusBadRequest)
		return false
	}

	payload, _ := base64.StdEncoding.DecodeString(auth[1])
	pair := strings.SplitN(string(payload), ":", 2)

	if len(pair) != 2 || !validateAuth(pair[0], pair[1]) {
		http.Error(w, "authorization failed", http.StatusUnauthorized)
		return false
	}
	return true
}

func validateAuth(username, password string) bool {
	if username == basicAuthUsr && password == basicAuthPwd {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package trust

import (
	"fmt"
	"os"

	"github.com/docker/cli/cli"
	"github.com/docker/cli/cli/command"
	"github.com/docker/cli/cli/trust"
	"github.com/docker/notary/client"
	"github.com/docker/notary/tuf/data"
	"github.com/spf13/cobra"
)

type signerRemoveOptions struct {
	forceYes bool
}

func newSignerRemoveCommand(dockerCli command.Cli) *cobra.Command {
	options := signerRemoveOptions{}
	cmd := &cobra.Command{
		Use:   "signer-remove NAME IMAGE [IMAGE...]",
		Short: "Remove a signer",
		Args:  cli.RequiresMinArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			return removeSigner(dockerCli, args[0], args[1:], &options)
		},
	}
	flags := cmd.Flags()
	flags.BoolVarP(&options.forceYes, "yes", "y", false, "Answer yes to removing most recent signer (no confirmation)")
	return cmd
}

func removeSigner(cli command.Cli, signer string, images []string, options *signerRemoveOptions) error {
	for _, image := range images {
		if err := removeSingleSigner(cli, image, signer, options.forceYes); err != nil {
			return err
		}
	}

	return nil
}

func isLastSignerForReleases(roleWithSig data.Role, allRoles []client.RoleWithSignatures) (bool, error) {
	var releasesRoleWithSigs client.RoleWithSignatures
	for _, role := range allRoles {
		if role.Name == releasesRoleTUFName {
			releasesRoleWithSigs = role
			break
		}
	}
	counter := len(releasesRoleWithSigs.Signatures)
	if counter == 0 {
		return false, fmt.Errorf("All signed tags are currently revoked, use docker trust sign to fix")
	}
	for _, signature := range releasesRoleWithSigs.Signatures {
		for _, key := range roleWithSig.KeyIDs {
			if signature.KeyID == key {
				counter--
			}
		}
	}
	return counter < releasesRoleWithSigs.Threshold, nil
}

func removeSingleSigner(cli command.Cli, image, signerName string, forceYes bool) error {
	_, ref, repoInfo, authConfig, err := getImageReferencesAndAuth(cli, image)
	if err != nil {
		return err
	}
	signerDelegation := data.RoleName("targets/" + signerName)
	if signerDelegation == releasesRoleTUFName {
		return fmt.Errorf("releases is a reserved keyword and cannot be removed")
	}
	notaryRepo, err := trust.GetNotaryRepository(cli, repoInfo, *authConfig, "push", "pull")
	if err != nil {
		return trust.NotaryError(ref.Name(), err)
	}
	delegationRoles, err := notaryRepo.GetDelegationRoles()
	if err != nil {
		return fmt.Errorf("Error retrieving signers for %s", image)
	}
	var role data.Role
	for _, delRole := range delegationRoles {
		if delRole.Name == signerDelegation {
			role = delRole
			break
		}
	}
	if role.Name == "" {
		return fmt.Errorf("No signer %s for image %s", signerName, image)
	}
	allRoles, err := notaryRepo.ListRoles()
	if err != nil {
		return err
	}
	if ok, err := isLastSignerForReleases(role, allRoles); ok && !forceYes {
		removeSigner := command.PromptForConfirmation(os.Stdin, cli.Out(), fmt.Sprintf("The signer %s signed the last released version of %s. "+
			"Removing them will lead to signed tags from %s being unpullable."+
			"Are you sure you want to continue?: ",
			signerName, image, image,
		))

		if !removeSigner {
			fmt.Fprintf(cli.Out(), "\nAborting action.\n")
			return nil
		}
	} else if err != nil {
		fmt.Fprintf(cli.Out(), err.Error())
	}
	if err = notaryRepo.RemoveDelegationKeys(releasesRoleTUFName, role.KeyIDs); err != nil {
		return err
	}
	if err = notaryRepo.RemoveDelegationRole(signerDelegation); err != nil {
		return err
	}
	if err = notaryRepo.Publish(); err != nil {
		return err
	}
	fmt.Fprintf(cli.Out(), "Successfully removed %s from %s\n", signerName, image)
	return nil
}
<commit_msg>docker trust remove: add newline to error print<commit_after>package trust

import (
	"fmt"
	"os"

	"github.com/docker/cli/cli"
	"github.com/docker/cli/cli/command"
	"github.com/docker/cli/cli/trust"
	"github.com/docker/notary/client"
	"github.com/docker/notary/tuf/data"
	"github.com/spf13/cobra"
)

type signerRemoveOptions struct {
	forceYes bool
}

func newSignerRemoveCommand(dockerCli command.Cli) *cobra.Command {
	options := signerRemoveOptions{}
	cmd := &cobra.Command{
		Use:   "signer-remove NAME IMAGE [IMAGE...]",
		Short: "Remove a signer",
		Args:  cli.RequiresMinArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			return removeSigner(dockerCli, args[0], args[1:], &options)
		},
	}
	flags := cmd.Flags()
	flags.BoolVarP(&options.forceYes, "yes", "y", false, "Answer yes to removing most recent signer (no confirmation)")
	return cmd
}

func removeSigner(cli command.Cli, signer string, images []string, options *signerRemoveOptions) error {
	for _, image := range images {
		if err := removeSingleSigner(cli, image, signer, options.forceYes); err != nil {
			return err
		}
	}

	return nil
}

func isLastSignerForReleases(roleWithSig data.Role, allRoles []client.RoleWithSignatures) (bool, error) {
	var releasesRoleWithSigs client.RoleWithSignatures
	for _, role := range allRoles {
		if role.Name == releasesRoleTUFName {
			releasesRoleWithSigs = role
			break
		}
	}
	counter := len(releasesRoleWithSigs.Signatures)
	if counter == 0 {
		return false, fmt.Errorf("All signed tags are currently revoked, use docker trust sign to fix")
	}
	for _, signature := range releasesRoleWithSigs.Signatures {
		for _, key := range roleWithSig.KeyIDs {
			if signature.KeyID == key {
				counter--
			}
		}
	}
	return counter < releasesRoleWithSigs.Threshold, nil
}

func removeSingleSigner(cli command.Cli, image, signerName string, forceYes bool) error {
	_, ref, repoInfo, authConfig, err := getImageReferencesAndAuth(cli, image)
	if err != nil {
		return err
	}
	signerDelegation := data.RoleName("targets/" + signerName)
	if signerDelegation == releasesRoleTUFName {
		return fmt.Errorf("releases is a reserved keyword and cannot be removed")
	}
	notaryRepo, err := trust.GetNotaryRepository(cli, repoInfo, *authConfig, "push", "pull")
	if err != nil {
		return trust.NotaryError(ref.Name(), err)
	}
	delegationRoles, err := notaryRepo.GetDelegationRoles()
	if err != nil {
		return fmt.Errorf("Error retrieving signers for %s", image)
	}
	var role data.Role
	for _, delRole := range delegationRoles {
		if delRole.Name == signerDelegation {
			role = delRole
			break
		}
	}
	if role.Name == "" {
		return fmt.Errorf("No signer %s for image %s", signerName, image)
	}
	allRoles, err := notaryRepo.ListRoles()
	if err != nil {
		return err
	}
	if ok, err := isLastSignerForReleases(role, allRoles); ok && !forceYes {
		removeSigner := command.PromptForConfirmation(os.Stdin, cli.Out(), fmt.Sprintf("The signer %s signed the last released version of %s. "+
			"Removing them will lead to signed tags from %s being unpullable."+
			"Are you sure you want to continue?: ",
			signerName, image, image,
		))

		if !removeSigner {
			fmt.Fprintf(cli.Out(), "\nAborting action.\n")
			return nil
		}
	} else if err != nil {
		fmt.Fprintln(cli.Out(), err.Error())
	}
	if err = notaryRepo.RemoveDelegationKeys(releasesRoleTUFName, role.KeyIDs); err != nil {
		return err
	}
	if err = notaryRepo.RemoveDelegationRole(signerDelegation); err != nil {
		return err
	}
	if err = notaryRepo.Publish(); err != nil {
		return err
	}
	fmt.Fprintf(cli.Out(), "Successfully removed %s from %s\n", signerName, image)
	return nil
}
<|endoftext|>


<commit_before>package cli_test

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/kopia/kopia/internal/testutil"
	"github.com/kopia/kopia/tests/testenv"
)

func TestServerControl(t *testing.T) {
	env := testenv.NewCLITest(t, testenv.RepoFormatNotImportant, testenv.NewInProcRunner(t))

	dir0 := testutil.TempDirectory(t)
	dir1 := testutil.TempDirectory(t)
	dir2 := testutil.TempDirectory(t)
	dir3 := testutil.TempDirectory(t)

	env.RunAndExpectSuccess(t, "repo", "create", "filesystem", "--path", env.RepoDir, "--override-username=another-user", "--override-hostname=another-host")
	env.RunAndExpectSuccess(t, "snap", "create", dir0)

	env.RunAndExpectSuccess(t, "repo", "connect", "filesystem", "--path", env.RepoDir, "--override-username=test-user", "--override-hostname=test-host")
	env.RunAndExpectSuccess(t, "snap", "create", dir1)
	env.RunAndExpectSuccess(t, "snap", "create", dir2)

	serverStarted := make(chan struct{})
	serverStopped := make(chan struct{})

	var sp testutil.ServerParameters

	go func() {
		kill := env.RunAndProcessStderr(t, sp.ProcessOutput,
			"server", "start", "--insecure", "--random-server-control-password", "--address=127.0.0.1:0")

		close(serverStarted)

		defer kill()

		close(serverStopped)
	}()

	select {
	case <-serverStarted:
		t.Logf("server started on %v", sp.BaseURL)

	case <-time.After(5 * time.Second):
		t.Fatalf("server did not start in time")
	}

	const (
		pollFrequency = 100 * time.Millisecond
		waitTimeout   = 5 * time.Second
	)

	require.Eventually(t, func() bool {
		lines := env.RunAndExpectSuccess(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
		return hasLine(lines, "IDLE: test-user@test-host:"+dir1) && hasLine(lines, "IDLE: test-user@test-host:"+dir2)
	}, waitTimeout, pollFrequency)

	lines := env.RunAndExpectSuccess(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--remote")
	require.Len(t, lines, 3)
	require.Contains(t, lines, "IDLE: test-user@test-host:"+dir1)
	require.Contains(t, lines, "IDLE: test-user@test-host:"+dir2)
	require.Contains(t, lines, "REMOTE: another-user@another-host:"+dir0)

	// create snapshot outside of the server
	env.RunAndExpectSuccess(t, "snap", "create", dir3)
	env.RunAndExpectSuccess(t, "server", "refresh", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	require.Eventually(t, func() bool {
		return hasLine(
			env.RunAndExpectSuccess(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--remote"),
			"IDLE: test-user@test-host:"+dir3)
	}, waitTimeout, pollFrequency)

	env.RunAndExpectSuccess(t, "server", "flush", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	// trigger server snapshot
	env.RunAndExpectSuccess(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--all")
	env.RunAndExpectSuccess(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, dir1)
	env.RunAndExpectFailure(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "no-such-dir")

	// neither dir nor --all specified
	env.RunAndExpectFailure(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	// cancel snapshot
	env.RunAndExpectSuccess(t, "server", "cancel", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--all")

	env.RunAndExpectSuccess(t, "server", "pause", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, dir1)
	env.RunAndExpectSuccess(t, "server", "resume", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, dir1)

	env.RunAndExpectSuccess(t, "server", "shutdown", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	select {
	case <-serverStopped:
		t.Logf("server shut down")

	case <-time.After(5 * time.Second):
		t.Fatalf("server did not shutdown in time")
	}

	// this will fail since the server is down
	env.RunAndExpectFailure(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
	env.RunAndExpectFailure(t, "server", "flush", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
	env.RunAndExpectFailure(t, "server", "refresh", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
	env.RunAndExpectFailure(t, "server", "shutdown", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
}

func hasLine(lines []string, lookFor string) bool {
	for _, l := range lines {
		if l == lookFor {
			return true
		}
	}

	return false
}
<commit_msg>fix(ci): fixed flaky TestServerControl (#1840)<commit_after>package cli_test

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/kopia/kopia/internal/testutil"
	"github.com/kopia/kopia/tests/testenv"
)

func TestServerControl(t *testing.T) {
	env := testenv.NewCLITest(t, testenv.RepoFormatNotImportant, testenv.NewInProcRunner(t))

	dir0 := testutil.TempDirectory(t)
	dir1 := testutil.TempDirectory(t)
	dir2 := testutil.TempDirectory(t)
	dir3 := testutil.TempDirectory(t)

	env.RunAndExpectSuccess(t, "repo", "create", "filesystem", "--path", env.RepoDir, "--override-username=another-user", "--override-hostname=another-host")
	env.RunAndExpectSuccess(t, "snap", "create", dir0)

	env.RunAndExpectSuccess(t, "repo", "connect", "filesystem", "--path", env.RepoDir, "--override-username=test-user", "--override-hostname=test-host")
	env.RunAndExpectSuccess(t, "snap", "create", dir1)
	env.RunAndExpectSuccess(t, "snap", "create", dir2)

	serverStarted := make(chan struct{})
	serverStopped := make(chan struct{})

	var sp testutil.ServerParameters

	go func() {
		kill := env.RunAndProcessStderr(t, sp.ProcessOutput,
			"server", "start", "--insecure", "--random-server-control-password", "--address=127.0.0.1:0")

		close(serverStarted)

		defer kill()

		close(serverStopped)
	}()

	select {
	case <-serverStarted:
		t.Logf("server started on %v", sp.BaseURL)

	case <-time.After(5 * time.Second):
		t.Fatalf("server did not start in time")
	}

	const (
		pollFrequency = 100 * time.Millisecond
		waitTimeout   = 15 * time.Second
	)

	require.Eventually(t, func() bool {
		lines := env.RunAndExpectSuccess(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
		t.Logf("lines: %v", lines)
		return hasLine(lines, "IDLE: test-user@test-host:"+dir1) && hasLine(lines, "IDLE: test-user@test-host:"+dir2)
	}, waitTimeout, pollFrequency)

	lines := env.RunAndExpectSuccess(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--remote")
	require.Len(t, lines, 3)
	require.Contains(t, lines, "IDLE: test-user@test-host:"+dir1)
	require.Contains(t, lines, "IDLE: test-user@test-host:"+dir2)
	require.Contains(t, lines, "REMOTE: another-user@another-host:"+dir0)

	// create snapshot outside of the server
	env.RunAndExpectSuccess(t, "snap", "create", dir3)
	env.RunAndExpectSuccess(t, "server", "refresh", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	require.Eventually(t, func() bool {
		lines := env.RunAndExpectSuccess(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--remote")
		t.Logf("lines: %v", lines)
		return hasLine(lines, "IDLE: test-user@test-host:"+dir3)
	}, waitTimeout, pollFrequency)

	env.RunAndExpectSuccess(t, "server", "flush", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	// trigger server snapshot
	env.RunAndExpectSuccess(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--all")
	env.RunAndExpectSuccess(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, dir1)
	env.RunAndExpectFailure(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "no-such-dir")

	// neither dir nor --all specified
	env.RunAndExpectFailure(t, "server", "snapshot", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	// cancel snapshot
	env.RunAndExpectSuccess(t, "server", "cancel", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, "--all")

	env.RunAndExpectSuccess(t, "server", "pause", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, dir1)
	env.RunAndExpectSuccess(t, "server", "resume", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword, dir1)

	env.RunAndExpectSuccess(t, "server", "shutdown", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)

	select {
	case <-serverStopped:
		t.Logf("server shut down")

	case <-time.After(15 * time.Second):
		t.Fatalf("server did not shutdown in time")
	}

	// this will fail since the server is down
	env.RunAndExpectFailure(t, "server", "status", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
	env.RunAndExpectFailure(t, "server", "flush", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
	env.RunAndExpectFailure(t, "server", "refresh", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
	env.RunAndExpectFailure(t, "server", "shutdown", "--address", sp.BaseURL, "--server-control-password", sp.ServerControlPassword)
}

func hasLine(lines []string, lookFor string) bool {
	for _, l := range lines {
		if l == lookFor {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2015 The LUCI 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 lib

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/tree"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/client/archiver"
	"go.chromium.org/luci/client/isolated"
	"go.chromium.org/luci/common/data/text/units"
	"go.chromium.org/luci/common/errors"
	isol "go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/system/signals"
)

// CmdArchive returns an object for the `archive` subcommand.
func CmdArchive(options CommandOptions) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "archive <options>...",
		ShortDesc: "creates a .isolated file and uploads the tree to an isolate server",
		LongDesc: `Given a list of files and directories, creates a .isolated file and uploads the
tree to to an isolate server.

When specifying directories and files, you must also specify a current working
directory for that file or directory. The current working directory will not
be included in the archived path. For example, to isolate './usr/foo/bar' and
have it appear as 'foo/bar' in the .isolated, specify '-files ./usr:foo/bar' or
'-files usr:foo/bar'. When the .isolated is then downloaded, it will then appear
under 'foo/bar' in the desired directory.

Note that '.' may be omitted in general, so to upload 'foo' from the current
working directory, '-files :foo' is sufficient.`,
		CommandRun: func() subcommands.CommandRun {
			c := archiveRun{
				CommandOptions: options,
			}
			c.commonFlags.Init(options.DefaultAuthOpts)
			c.Flags.Var(&c.dirs, "dirs", "Directory(ies) to archive. Specify as <working directory>:<relative path to dir>")
			c.Flags.Var(&c.files, "files", "Individual file(s) to archive. Specify as <working directory>:<relative path to file>")
			c.Flags.StringVar(&c.dumpHash, "dump-hash", "",
				"Write the composite isolated hash to a file")
			c.Flags.StringVar(&c.isolated, "isolated", "",
				"Write the composite isolated to a file")
			c.Flags.StringVar(&c.dumpStatsJSON, "dump-stats-json", "",
				"Write the upload stats to this file as JSON")
			return &c
		},
	}
}

type archiveRun struct {
	commonFlags
	CommandOptions
	dirs          isolated.ScatterGather
	files         isolated.ScatterGather
	dumpHash      string
	isolated      string
	dumpStatsJSON string
}

func (c *archiveRun) Parse(a subcommands.Application, args []string) error {
	if err := c.commonFlags.Parse(); err != nil {
		return err
	}
	if len(args) != 0 {
		return errors.Reason("position arguments not expected").Err()
	}
	return nil
}

// getRoot returns root directory if there is only one working directory.
func getRoot(dirs, files isolated.ScatterGather) (string, error) {
	var rel0, wd0 string
	pickedOne := false
	for rel, wd := range dirs {
		if !pickedOne {
			rel0 = rel
			wd0 = wd
			pickedOne = true
			continue
		}

		if wd0 != wd {
			return "", errors.Reason("different root (working) directory is not supported: %s:%s vs %s:%s", wd0, rel0, wd, rel).Err()
		}
	}

	for rel, wd := range files {
		if !pickedOne {
			rel0 = rel
			wd0 = wd
			pickedOne = true
			continue
		}

		if wd0 != wd {
			return "", errors.Reason("different root (working) directory is not supported: %s:%s vs %s:%s", wd0, rel0, wd, rel).Err()
		}
	}

	if !pickedOne {
		return "", errors.Reason("-dirs or -files should be specified at least once").Err()
	}

	return wd0, nil
}

func (c *archiveRun) doCASAarchive(ctx context.Context) error {
	root, err := getRoot(c.dirs, c.files)
	if err != nil {
		return err
	}

	is := command.InputSpec{}
	for dir := range c.dirs {
		is.Inputs = append(is.Inputs, dir)
	}
	for file := range c.files {
		is.Inputs = append(is.Inputs, file)
	}

	rootDg, chunkers, _, err := tree.ComputeMerkleTree(root, &is, chunker.DefaultChunkSize, filemetadata.NewNoopCache())
	if err != nil {
		return errors.Annotate(err, "failed to call ComputeMerkleTree").Err()
	}

	client, err := c.casFlags.NewClient(ctx)
	if err != nil {
		return errors.Annotate(err, "failed to create cas client").Err()
	}
	defer client.Close()

	if err := client.UploadIfMissing(ctx, chunkers...); err != nil {
		return errors.Annotate(err, "failed to call UploadIfMissing").Err()
	}

	if c.dumpHash != "" {
		if err := ioutil.WriteFile(c.dumpHash, []byte(rootDg.String()), 0644); err != nil {
			return err
		}
	}

	return nil
}

func (c *archiveRun) doIsolatedArchive(ctx context.Context) (stats *archiver.Stats, err error) {
	isolatedClient, isolErr := c.createIsolatedClient(ctx, c.CommandOptions)
	if isolErr != nil {
		err = errors.Annotate(isolErr, "failed to create isolated client").Err()
		return
	}
	out := os.Stdout
	arch := archiver.New(ctx, isolatedClient, out)
	defer func() {
		// This waits for all uploads.
		if cerr := arch.Close(); err == nil {
			err = cerr
		}
		// We must take the stats until after all the uploads have finished
		if err == nil {
			stats = arch.Stats()
		}
	}()

	opts := isolated.ArchiveOptions{
		Files:    c.files,
		Dirs:     c.dirs,
		Isolated: c.isolated,
	}
	if len(c.isolated) != 0 {
		var dumpIsolated *os.File
		dumpIsolated, err = os.Create(c.isolated)
		if err != nil {
			return
		}
		// This is OK to close before arch because isolated.Archive
		// does the writing (it's not handed off elsewhere).
		defer dumpIsolated.Close()
		opts.LeakIsolated = dumpIsolated
	}
	item := isolated.Archive(ctx, arch, &opts)
	if err = item.Error(); err != nil {
		return
	}

	item.WaitForHashed()
	if len(c.dumpHash) != 0 {
		if err = ioutil.WriteFile(c.dumpHash, []byte(item.Digest()), 0644); err != nil {
			return
		}
	}
	return
}

// Does the archive by uploading to isolate-server, then return the archive stats and error.
func (c *archiveRun) doArchive(a subcommands.Application, args []string) (stats *archiver.Stats, err error) {
	ctx, cancel := context.WithCancel(c.defaultFlags.MakeLoggingContext(os.Stderr))
	signals.HandleInterrupt(cancel)

	if c.casFlags.Instance != "" {
		// TODO(crbug.com/1110569): get stats
		return &archiver.Stats{}, c.doCASAarchive(ctx)
	}

	return c.doIsolatedArchive(ctx)
}

func (c *archiveRun) postprocessStats(stats *archiver.Stats, start time.Time) error {
	if !c.defaultFlags.Quiet {
		duration := time.Since(start)
		fmt.Fprintf(os.Stderr, "Hits    : %5d (%s)\n", stats.TotalHits(), stats.TotalBytesHits())
		fmt.Fprintf(os.Stderr, "Misses  : %5d (%s)\n", stats.TotalMisses(), stats.TotalBytesPushed())
		fmt.Fprintf(os.Stderr, "Duration: %s\n", units.Round(duration, time.Millisecond))
	}
	if c.dumpStatsJSON != "" {
		return dumpStatsJSON(c.dumpStatsJSON, stats)
	}
	return nil
}

func (c *archiveRun) Run(a subcommands.Application, args []string, _ subcommands.Env) int {
	if err := c.Parse(a, args); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	cl, err := c.defaultFlags.StartTracing()
	if err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	defer cl.Close()
	defer c.profilerFlags.Stop()
	start := time.Now()
	stats, err := c.doArchive(a, args)
	if err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	if err := c.postprocessStats(stats, start); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

func dumpStatsJSON(jsonPath string, stats *archiver.Stats) error {
	hits := make([]int64, len(stats.Hits))
	for i, h := range stats.Hits {
		hits[i] = int64(h)
	}
	sort.Slice(hits, func(i, j int) bool { return hits[i] < hits[j] })
	itemsHot, err := isol.Pack(hits)
	if err != nil {
		return errors.Annotate(err, "failed to pack itemsHot").Err()
	}

	pushed := make([]int64, len(stats.Pushed))
	for i, p := range stats.Pushed {
		pushed[i] = int64(p.Size)
	}
	sort.Slice(pushed, func(i, j int) bool { return pushed[i] < pushed[j] })
	itemsCold, err := isol.Pack(pushed)
	if err != nil {
		return errors.Annotate(err, "failed to pack itemsCold").Err()
	}

	statsJSON, err := json.Marshal(struct {
		ItemsCold []byte `json:"items_cold"`
		ItemsHot  []byte `json:"items_hot"`
	}{
		ItemsCold: itemsCold,
		ItemsHot:  itemsHot,
	})
	if err != nil {
		return errors.Annotate(err, "failed to marshal result json").Err()
	}
	if err := ioutil.WriteFile(jsonPath, statsJSON, 0664); err != nil {
		return errors.Annotate(err, "failed to write stats json to %s", jsonPath).Err()
	}
	return nil
}
<commit_msg>[isolated] Silence output when -quiet is true<commit_after>// Copyright 2015 The LUCI 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 lib

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/tree"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/client/archiver"
	"go.chromium.org/luci/client/isolated"
	"go.chromium.org/luci/common/data/text/units"
	"go.chromium.org/luci/common/errors"
	isol "go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/system/signals"
)

// CmdArchive returns an object for the `archive` subcommand.
func CmdArchive(options CommandOptions) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "archive <options>...",
		ShortDesc: "creates a .isolated file and uploads the tree to an isolate server",
		LongDesc: `Given a list of files and directories, creates a .isolated file and uploads the
tree to to an isolate server.

When specifying directories and files, you must also specify a current working
directory for that file or directory. The current working directory will not
be included in the archived path. For example, to isolate './usr/foo/bar' and
have it appear as 'foo/bar' in the .isolated, specify '-files ./usr:foo/bar' or
'-files usr:foo/bar'. When the .isolated is then downloaded, it will then appear
under 'foo/bar' in the desired directory.

Note that '.' may be omitted in general, so to upload 'foo' from the current
working directory, '-files :foo' is sufficient.`,
		CommandRun: func() subcommands.CommandRun {
			c := archiveRun{
				CommandOptions: options,
			}
			c.commonFlags.Init(options.DefaultAuthOpts)
			c.Flags.Var(&c.dirs, "dirs", "Directory(ies) to archive. Specify as <working directory>:<relative path to dir>")
			c.Flags.Var(&c.files, "files", "Individual file(s) to archive. Specify as <working directory>:<relative path to file>")
			c.Flags.StringVar(&c.dumpHash, "dump-hash", "",
				"Write the composite isolated hash to a file")
			c.Flags.StringVar(&c.isolated, "isolated", "",
				"Write the composite isolated to a file")
			c.Flags.StringVar(&c.dumpStatsJSON, "dump-stats-json", "",
				"Write the upload stats to this file as JSON")
			return &c
		},
	}
}

type archiveRun struct {
	commonFlags
	CommandOptions
	dirs          isolated.ScatterGather
	files         isolated.ScatterGather
	dumpHash      string
	isolated      string
	dumpStatsJSON string
}

func (c *archiveRun) Parse(a subcommands.Application, args []string) error {
	if err := c.commonFlags.Parse(); err != nil {
		return err
	}
	if len(args) != 0 {
		return errors.Reason("position arguments not expected").Err()
	}
	return nil
}

// getRoot returns root directory if there is only one working directory.
func getRoot(dirs, files isolated.ScatterGather) (string, error) {
	var rel0, wd0 string
	pickedOne := false
	for rel, wd := range dirs {
		if !pickedOne {
			rel0 = rel
			wd0 = wd
			pickedOne = true
			continue
		}

		if wd0 != wd {
			return "", errors.Reason("different root (working) directory is not supported: %s:%s vs %s:%s", wd0, rel0, wd, rel).Err()
		}
	}

	for rel, wd := range files {
		if !pickedOne {
			rel0 = rel
			wd0 = wd
			pickedOne = true
			continue
		}

		if wd0 != wd {
			return "", errors.Reason("different root (working) directory is not supported: %s:%s vs %s:%s", wd0, rel0, wd, rel).Err()
		}
	}

	if !pickedOne {
		return "", errors.Reason("-dirs or -files should be specified at least once").Err()
	}

	return wd0, nil
}

func (c *archiveRun) doCASAarchive(ctx context.Context) error {
	root, err := getRoot(c.dirs, c.files)
	if err != nil {
		return err
	}

	is := command.InputSpec{}
	for dir := range c.dirs {
		is.Inputs = append(is.Inputs, dir)
	}
	for file := range c.files {
		is.Inputs = append(is.Inputs, file)
	}

	rootDg, chunkers, _, err := tree.ComputeMerkleTree(root, &is, chunker.DefaultChunkSize, filemetadata.NewNoopCache())
	if err != nil {
		return errors.Annotate(err, "failed to call ComputeMerkleTree").Err()
	}

	client, err := c.casFlags.NewClient(ctx)
	if err != nil {
		return errors.Annotate(err, "failed to create cas client").Err()
	}
	defer client.Close()

	if err := client.UploadIfMissing(ctx, chunkers...); err != nil {
		return errors.Annotate(err, "failed to call UploadIfMissing").Err()
	}

	if c.dumpHash != "" {
		if err := ioutil.WriteFile(c.dumpHash, []byte(rootDg.String()), 0644); err != nil {
			return err
		}
	}

	return nil
}

func (c *archiveRun) doIsolatedArchive(ctx context.Context) (stats *archiver.Stats, err error) {
	isolatedClient, isolErr := c.createIsolatedClient(ctx, c.CommandOptions)
	if isolErr != nil {
		err = errors.Annotate(isolErr, "failed to create isolated client").Err()
		return
	}
	var out io.Writer = os.Stdout
	if c.defaultFlags.Quiet {
		out = ioutil.Discard
	}
	arch := archiver.New(ctx, isolatedClient, out)
	defer func() {
		// This waits for all uploads.
		if cerr := arch.Close(); err == nil {
			err = cerr
		}
		// We must take the stats until after all the uploads have finished
		if err == nil {
			stats = arch.Stats()
		}
	}()

	opts := isolated.ArchiveOptions{
		Files:    c.files,
		Dirs:     c.dirs,
		Isolated: c.isolated,
	}
	if len(c.isolated) != 0 {
		var dumpIsolated *os.File
		dumpIsolated, err = os.Create(c.isolated)
		if err != nil {
			return
		}
		// This is OK to close before arch because isolated.Archive
		// does the writing (it's not handed off elsewhere).
		defer dumpIsolated.Close()
		opts.LeakIsolated = dumpIsolated
	}
	item := isolated.Archive(ctx, arch, &opts)
	if err = item.Error(); err != nil {
		return
	}

	item.WaitForHashed()
	if len(c.dumpHash) != 0 {
		if err = ioutil.WriteFile(c.dumpHash, []byte(item.Digest()), 0644); err != nil {
			return
		}
	}
	return
}

// Does the archive by uploading to isolate-server, then return the archive stats and error.
func (c *archiveRun) doArchive(a subcommands.Application, args []string) (stats *archiver.Stats, err error) {
	ctx, cancel := context.WithCancel(c.defaultFlags.MakeLoggingContext(os.Stderr))
	signals.HandleInterrupt(cancel)

	if c.casFlags.Instance != "" {
		// TODO(crbug.com/1110569): get stats
		return &archiver.Stats{}, c.doCASAarchive(ctx)
	}

	return c.doIsolatedArchive(ctx)
}

func (c *archiveRun) postprocessStats(stats *archiver.Stats, start time.Time) error {
	if !c.defaultFlags.Quiet {
		duration := time.Since(start)
		fmt.Fprintf(os.Stderr, "Hits    : %5d (%s)\n", stats.TotalHits(), stats.TotalBytesHits())
		fmt.Fprintf(os.Stderr, "Misses  : %5d (%s)\n", stats.TotalMisses(), stats.TotalBytesPushed())
		fmt.Fprintf(os.Stderr, "Duration: %s\n", units.Round(duration, time.Millisecond))
	}
	if c.dumpStatsJSON != "" {
		return dumpStatsJSON(c.dumpStatsJSON, stats)
	}
	return nil
}

func (c *archiveRun) Run(a subcommands.Application, args []string, _ subcommands.Env) int {
	if err := c.Parse(a, args); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	cl, err := c.defaultFlags.StartTracing()
	if err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	defer cl.Close()
	defer c.profilerFlags.Stop()
	start := time.Now()
	stats, err := c.doArchive(a, args)
	if err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	if err := c.postprocessStats(stats, start); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

func dumpStatsJSON(jsonPath string, stats *archiver.Stats) error {
	hits := make([]int64, len(stats.Hits))
	for i, h := range stats.Hits {
		hits[i] = int64(h)
	}
	sort.Slice(hits, func(i, j int) bool { return hits[i] < hits[j] })
	itemsHot, err := isol.Pack(hits)
	if err != nil {
		return errors.Annotate(err, "failed to pack itemsHot").Err()
	}

	pushed := make([]int64, len(stats.Pushed))
	for i, p := range stats.Pushed {
		pushed[i] = int64(p.Size)
	}
	sort.Slice(pushed, func(i, j int) bool { return pushed[i] < pushed[j] })
	itemsCold, err := isol.Pack(pushed)
	if err != nil {
		return errors.Annotate(err, "failed to pack itemsCold").Err()
	}

	statsJSON, err := json.Marshal(struct {
		ItemsCold []byte `json:"items_cold"`
		ItemsHot  []byte `json:"items_hot"`
	}{
		ItemsCold: itemsCold,
		ItemsHot:  itemsHot,
	})
	if err != nil {
		return errors.Annotate(err, "failed to marshal result json").Err()
	}
	if err := ioutil.WriteFile(jsonPath, statsJSON, 0664); err != nil {
		return errors.Annotate(err, "failed to write stats json to %s", jsonPath).Err()
	}
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"time"

	cloudTypes "github.com/atlassian/gostatsd/cloudprovider/types"
	"github.com/atlassian/gostatsd/types"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/spf13/viper"
)

const (
	// ProviderName is the name of AWS cloud provider.
	ProviderName = "aws"
)

const sampleConfig = `
[aws]
	# maximum number of retries in case of retriable errors
	max_retries = 5 # optional, default to 3
`

// Provider represents an AWS provider.
type Provider struct {
	Metadata *ec2metadata.EC2Metadata
	Ec2      *ec2.EC2
}

// DescribeInstances is an implementation of EC2.Instances.
func (p *Provider) describeInstances(ctx context.Context, request *ec2.DescribeInstancesInput) ([]*ec2.Instance, error) {
	// Instances are paged
	results := []*ec2.Instance{}

	for {
		req, response := p.Ec2.DescribeInstancesRequest(request)
		req.HTTPRequest = req.HTTPRequest.WithContext(ctx)
		err := req.Send()
		if err != nil {
			return nil, fmt.Errorf("error listing AWS instances: %v", err)
		}

		for _, reservation := range response.Reservations {
			results = append(results, reservation.Instances...)
		}

		if response.NextToken == nil {
			break
		}
		request.NextToken = response.NextToken
	}

	return results, nil
}

func newEc2Filter(name string, value string) *ec2.Filter {
	return &ec2.Filter{
		Name: aws.String(name),
		Values: []*string{
			aws.String(value),
		},
	}
}

// Instance returns the instance details from aws.
func (p *Provider) Instance(ctx context.Context, IP types.IP) (*cloudTypes.Instance, error) {
	filters := []*ec2.Filter{newEc2Filter("private-ip-address", string(IP))}
	request := &ec2.DescribeInstancesInput{
		Filters: filters,
	}

	instances, err := p.describeInstances(ctx, request)
	if err != nil {
		return nil, err
	}
	if len(instances) == 0 {
		return nil, errors.New("no instances returned")
	}

	i := instances[0]
	region, err := azToRegion(aws.StringValue(i.Placement.AvailabilityZone))
	if err != nil {
		log.Errorf("Error getting instance region: %v", err)
	}
	tags := make(types.Tags, len(i.Tags))
	for idx, tag := range i.Tags {
		tags[idx] = fmt.Sprintf("%s:%s",
			types.NormalizeTagKey(aws.StringValue(tag.Key)),
			aws.StringValue(tag.Value))
	}
	instance := &cloudTypes.Instance{
		ID:     aws.StringValue(i.InstanceId),
		Region: region,
		Tags:   tags,
	}
	return instance, nil
}

// ProviderName returns the name of the provider.
func (p *Provider) ProviderName() string {
	return ProviderName
}

// SampleConfig returns the sample config for the datadog backend.
func (p *Provider) SampleConfig() string {
	return sampleConfig
}

// SelfIP returns host's IPv4 address.
func (p *Provider) SelfIP() (types.IP, error) {
	ip, err := p.Metadata.GetMetadata("local-ipv4")
	return types.IP(ip), err
}

// Derives the region from a valid az name.
// Returns an error if the az is known invalid (empty).
func azToRegion(az string) (string, error) {
	if az == "" {
		return "", errors.New("invalid (empty) AZ")
	}
	region := az[:len(az)-1]
	return region, nil
}

// NewProviderFromViper returns a new aws provider.
func NewProviderFromViper(v *viper.Viper) (cloudTypes.Interface, error) {
	a := getSubViper(v, "aws")
	a.SetDefault("max_retries", 3)
	a.SetDefault("http_timeout", 3*time.Second)
	httpTimeout := a.GetDuration("http_timeout")
	if httpTimeout <= 0 {
		return nil, errors.New("http client timeout must be positive")
	}

	// This is the main config without credentials.
	config := &aws.Config{
		MaxRetries: aws.Int(a.GetInt("max_retries")),
		HTTPClient: &http.Client{
			Timeout: httpTimeout,
		},
	}
	metadata := ec2metadata.New(session.New(config))
	az, err := metadata.GetMetadata("placement/availability-zone")
	if err != nil {
		return nil, fmt.Errorf("error getting availability zone: %v", err)
	}
	region, err := azToRegion(az)
	if err != nil {
		return nil, fmt.Errorf("error getting aws region: %v", err)
	}
	return &Provider{
		Metadata: metadata,
		Ec2: ec2.New(session.New(config.Copy(&aws.Config{
			Credentials: credentials.NewChainCredentials(
				[]credentials.Provider{
					&credentials.EnvProvider{},
					&ec2rolecreds.EC2RoleProvider{
						Client: metadata,
					},
					&credentials.SharedCredentialsProvider{},
				}),
			Region: aws.String(region),
		}))),
	}, nil
}

func getSubViper(v *viper.Viper, key string) *viper.Viper {
	n := v.Sub(key)
	if n == nil {
		n = viper.New()
	}
	return n
}
<commit_msg>Use built-in pagination<commit_after>package aws

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"time"

	cloudTypes "github.com/atlassian/gostatsd/cloudprovider/types"
	"github.com/atlassian/gostatsd/types"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/spf13/viper"
)

const (
	// ProviderName is the name of AWS cloud provider.
	ProviderName = "aws"
)

const sampleConfig = `
[aws]
	# maximum number of retries in case of retriable errors
	max_retries = 5 # optional, default to 3
`

// Provider represents an AWS provider.
type Provider struct {
	Metadata *ec2metadata.EC2Metadata
	Ec2      *ec2.EC2
}

func newEc2Filter(name string, value string) *ec2.Filter {
	return &ec2.Filter{
		Name: aws.String(name),
		Values: []*string{
			aws.String(value),
		},
	}
}

// Instance returns the instance details from aws.
func (p *Provider) Instance(ctx context.Context, IP types.IP) (*cloudTypes.Instance, error) {
	req, _ := p.Ec2.DescribeInstancesRequest(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			newEc2Filter("private-ip-address", string(IP)),
		},
	})
	req.HTTPRequest = req.HTTPRequest.WithContext(ctx)
	var inst *ec2.Instance
	err := req.EachPage(func(data interface{}, isLastPage bool) bool {
		for _, reservation := range data.(*ec2.DescribeInstancesOutput).Reservations {
			for _, instance := range reservation.Instances {
				inst = instance
				return false
			}
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing AWS instances: %v", err)
	}
	if inst == nil {
		return nil, errors.New("no instances found")
	}
	region, err := azToRegion(aws.StringValue(inst.Placement.AvailabilityZone))
	if err != nil {
		log.Errorf("Error getting instance region: %v", err)
	}
	tags := make(types.Tags, len(inst.Tags))
	for idx, tag := range inst.Tags {
		tags[idx] = fmt.Sprintf("%s:%s",
			types.NormalizeTagKey(aws.StringValue(tag.Key)),
			aws.StringValue(tag.Value))
	}
	instance := &cloudTypes.Instance{
		ID:     aws.StringValue(inst.InstanceId),
		Region: region,
		Tags:   tags,
	}
	return instance, nil
}

// ProviderName returns the name of the provider.
func (p *Provider) ProviderName() string {
	return ProviderName
}

// SampleConfig returns the sample config for the datadog backend.
func (p *Provider) SampleConfig() string {
	return sampleConfig
}

// SelfIP returns host's IPv4 address.
func (p *Provider) SelfIP() (types.IP, error) {
	ip, err := p.Metadata.GetMetadata("local-ipv4")
	return types.IP(ip), err
}

// Derives the region from a valid az name.
// Returns an error if the az is known invalid (empty).
func azToRegion(az string) (string, error) {
	if az == "" {
		return "", errors.New("invalid (empty) AZ")
	}
	region := az[:len(az)-1]
	return region, nil
}

// NewProviderFromViper returns a new aws provider.
func NewProviderFromViper(v *viper.Viper) (cloudTypes.Interface, error) {
	a := getSubViper(v, "aws")
	a.SetDefault("max_retries", 3)
	a.SetDefault("http_timeout", 3*time.Second)
	httpTimeout := a.GetDuration("http_timeout")
	if httpTimeout <= 0 {
		return nil, errors.New("http client timeout must be positive")
	}

	// This is the main config without credentials.
	config := &aws.Config{
		MaxRetries: aws.Int(a.GetInt("max_retries")),
		HTTPClient: &http.Client{
			Timeout: httpTimeout,
		},
	}
	metadata := ec2metadata.New(session.New(config))
	az, err := metadata.GetMetadata("placement/availability-zone")
	if err != nil {
		return nil, fmt.Errorf("error getting availability zone: %v", err)
	}
	region, err := azToRegion(az)
	if err != nil {
		return nil, fmt.Errorf("error getting aws region: %v", err)
	}
	return &Provider{
		Metadata: metadata,
		Ec2: ec2.New(session.New(config.Copy(&aws.Config{
			Credentials: credentials.NewChainCredentials(
				[]credentials.Provider{
					&credentials.EnvProvider{},
					&ec2rolecreds.EC2RoleProvider{
						Client: metadata,
					},
					&credentials.SharedCredentialsProvider{},
				}),
			Region: aws.String(region),
		}))),
	}, nil
}

func getSubViper(v *viper.Viper, key string) *viper.Viper {
	n := v.Sub(key)
	if n == nil {
		n = viper.New()
	}
	return n
}
<|endoftext|>


<commit_before>package rabbus

import (
	"context"
	"errors"
	"sync"
	"time"

	amqpwrap "github.com/rafaeljesus/rabbus/internal/amqp"
	"github.com/rafaeljesus/retry-go"
	"github.com/sony/gobreaker"
	"github.com/streadway/amqp"
)

const (
	// Transient means higher throughput but messages will not be restored on broker restart.
	Transient uint8 = 1
	// Persistent messages will be restored to durable queues and lost on non-durable queues during server restart.
	Persistent uint8 = 2
	// ContentTypeJSON define json content type
	ContentTypeJSON = "application/json"
	// ContentTypePlain define plain text content type
	ContentTypePlain = "plain/text"

	contentEncoding = "UTF-8"
)

type (
	// Option represents an option you can pass to New.
	// See the documentation for the individual options.
	Option func(*Rabbus) error
	// OnStateChangeFunc is the callback function when circuit breaker state changes.
	OnStateChangeFunc func(name, from, to string)

	// Message carries fields for sending messages.
	Message struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// Payload the message payload.
		Payload []byte
		// DeliveryMode indicates if the is Persistent or Transient.
		DeliveryMode uint8
		// ContentType the message content-type.
		ContentType string
		// Headers the message application headers
		Headers map[string]interface{}
	}

	// ListenConfig carries fields for listening messages.
	ListenConfig struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// PassiveExchange determines a passive exchange connection it uses
		// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
		PassiveExchange bool
		// Queue the queue name
		Queue string
	}

	// Delivery wraps amqp.Delivery struct
	Delivery struct {
		amqp.Delivery
	}

	// Rabbus interpret (implement) Rabbus interface definition
	Rabbus struct {
		Amqp
		mu         sync.RWMutex
		breaker    *gobreaker.CircuitBreaker
		emit       chan Message
		emitErr    chan error
		emitOk     chan struct{}
		reconn     chan struct{}
		exDeclared map[string]struct{}
		config
		conDeclared int // conDeclared is a counter for the declared consumers
	}

	// Amqp expose a interface for interacting with amqp broker
	Amqp interface {
		// Publish wraps amqp.Publish method
		Publish(exchange, key string, opts amqp.Publishing) error
		// CreateConsumer creates a amqp consumer
		CreateConsumer(exchange, key, kind, queue string, durable bool) (<-chan amqp.Delivery, error)
		// WithExchange creates a amqp exchange
		WithExchange(exchange, kind string, durable bool) error
		// WithQos wrapper over amqp.Qos method
		WithQos(count, size int, global bool) error
		// NotifyClose wrapper over notifyClose method
		NotifyClose(c chan *amqp.Error) chan *amqp.Error
		// Close closes the running amqp connection and channel
		Close() error
	}

	config struct {
		dsn                string
		durable, passiveex bool
		retrycfg
		breaker
		qos
	}

	retrycfg struct {
		attempts              int
		sleep, reconnectSleep time.Duration
	}

	breaker struct {
		interval, timeout time.Duration
		threshold         uint32
		onStateChange     OnStateChangeFunc
	}

	qos struct {
		prefetchCount, prefetchSize int
		global                      bool
	}
)

func (lc ListenConfig) validate() error {
	if lc.Exchange == "" {
		return ErrMissingExchange
	}

	if lc.Kind == "" {
		return ErrMissingKind
	}

	if lc.Queue == "" {
		return ErrMissingQueue
	}

	return nil
}

// New returns a new Rabbus configured with the
// variables from the config parameter, or returning an non-nil err
// if an error occurred while creating connection and channel.
func New(dsn string, options ...Option) (*Rabbus, error) {
	r := &Rabbus{
		emit:       make(chan Message),
		emitErr:    make(chan error),
		emitOk:     make(chan struct{}),
		reconn:     make(chan struct{}, 10),
		exDeclared: make(map[string]struct{}),
	}

	for _, o := range options {
		if err := o(r); err != nil {
			return nil, err
		}
	}

	if r.Amqp == nil {
		amqpWrapper, err := amqpwrap.New(dsn, r.config.passiveex)
		if err != nil {
			return nil, err
		}
		r.Amqp = amqpWrapper
	}

	if err := r.WithQos(
		r.config.qos.prefetchCount,
		r.config.qos.prefetchSize,
		r.config.qos.global,
	); err != nil {
		return nil, err
	}

	r.config.dsn = dsn
	r.breaker = gobreaker.NewCircuitBreaker(newBreakerSettings(r.config))

	return r, nil
}

// Run starts rabbus channels for emitting and listening for amqp connection close
// returns ctx error in case of any.
func (r *Rabbus) Run(ctx context.Context) error {
	for {
		select {
		case m, ok := <-r.emit:
			if ok {
				r.produce(m)
			}
		case err, ok := <-r.NotifyClose(make(chan *amqp.Error)):
			if ok {
				r.handleAmqpClose(err)
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// EmitAsync emits a message to RabbitMQ, but does not wait for the response from broker.
func (r *Rabbus) EmitAsync() chan<- Message { return r.emit }

// EmitErr returns an error if encoding payload fails, or if after circuit breaker is open or retries attempts exceed.
func (r *Rabbus) EmitErr() <-chan error { return r.emitErr }

// EmitOk returns true when the message was sent.
func (r *Rabbus) EmitOk() <-chan struct{} { return r.emitOk }

// Listen to a message from RabbitMQ, returns
// an error if exchange, queue name and function handler not passed or if an error occurred while creating
// amqp consumer.
func (r *Rabbus) Listen(c ListenConfig) (chan ConsumerMessage, error) {
	if err := c.validate(); err != nil {
		return nil, err
	}

	msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
	if err != nil {
		return nil, err
	}

	r.conDeclared++ // increase the declared consumer's counter
	r.mu.Lock()
	r.exDeclared[c.Exchange] = struct{}{}
	r.mu.Unlock()

	messages := make(chan ConsumerMessage, 256)
	go r.wrapMessage(c, msgs, messages)
	go r.listenReconn(c, messages)

	return messages, nil
}

// Close channels and attempt to close channel and connection.
func (r *Rabbus) Close() error {
	close(r.emit)
	close(r.emitOk)
	close(r.emitErr)
	close(r.reconn)

	return r.Amqp.Close()
}

func (r *Rabbus) produce(m Message) {
	if _, ok := r.exDeclared[m.Exchange]; !ok {
		if err := r.WithExchange(m.Exchange, m.Kind, r.config.durable); err != nil {
			r.emitErr <- err
			return
		}
		r.exDeclared[m.Exchange] = struct{}{}
	}

	if m.ContentType == "" {
		m.ContentType = ContentTypeJSON
	}

	if m.DeliveryMode == 0 {
		m.DeliveryMode = Persistent
	}

	opts := amqp.Publishing{
		Headers:         amqp.Table(m.Headers),
		ContentType:     m.ContentType,
		ContentEncoding: contentEncoding,
		DeliveryMode:    m.DeliveryMode,
		Timestamp:       time.Now(),
		Body:            m.Payload,
	}

	if _, err := r.breaker.Execute(func() (interface{}, error) {
		return nil, retry.Do(func() error {
			return r.Publish(m.Exchange, m.Key, opts)
		}, r.config.retrycfg.attempts, r.config.retrycfg.sleep)
	}); err != nil {
		r.emitErr <- err
		return
	}

	r.emitOk <- struct{}{}
}

// Durable indicates of the queue will survive broker restarts. Default to true.
func Durable(durable bool) Option {
	return func(r *Rabbus) error {
		r.config.durable = durable
		return nil
	}
}

// PassiveExchange forces passive connection with all exchanges using
// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
func PassiveExchange(passiveex bool) Option {
	return func(r *Rabbus) error {
		r.config.passiveex = passiveex
		return nil
	}
}

// PrefetchCount limit the number of unacknowledged messages.
func PrefetchCount(count int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchCount = count
		return nil
	}
}

// PrefetchSize when greater than zero, the server will try to keep at least
// that many bytes of deliveries flushed to the network before receiving
// acknowledgments from the consumers.
func PrefetchSize(size int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchSize = size
		return nil
	}
}

// QosGlobal when global is true, these Qos settings apply to all existing and future
// consumers on all channels on the same connection. When false, the Channel.Qos
// settings will apply to all existing and future consumers on this channel.
// RabbitMQ does not implement the global flag.
func QosGlobal(global bool) Option {
	return func(r *Rabbus) error {
		r.config.qos.global = global
		return nil
	}
}

// Attempts is the max number of retries on broker outages.
func Attempts(attempts int) Option {
	return func(r *Rabbus) error {
		r.config.retrycfg.attempts = attempts
		return nil
	}
}

// Sleep is the sleep time of the retry mechanism.
func Sleep(sleep time.Duration) Option {
	return func(r *Rabbus) error {
		if sleep == 0 {
			r.config.retrycfg.reconnectSleep = time.Second * 10
		}
		r.config.retrycfg.sleep = sleep
		return nil
	}
}

// BreakerInterval is the cyclic period of the closed state for CircuitBreaker to clear the internal counts,
// If Interval is 0, CircuitBreaker doesn't clear the internal counts during the closed state.
func BreakerInterval(interval time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.interval = interval
		return nil
	}
}

// BreakerTimeout is the period of the open state, after which the state of CircuitBreaker becomes half-open.
// If Timeout is 0, the timeout value of CircuitBreaker is set to 60 seconds.
func BreakerTimeout(timeout time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.timeout = timeout
		return nil
	}
}

// Threshold when a threshold of failures has been reached, future calls to the broker will not run.
// During this state, the circuit breaker will periodically allow the calls to run and, if it is successful,
// will start running the function again. Default value is 5.
func Threshold(threshold uint32) Option {
	return func(r *Rabbus) error {
		if threshold == 0 {
			threshold = 5
		}
		r.config.breaker.threshold = threshold
		return nil
	}
}

// OnStateChange is called whenever the state of CircuitBreaker changes.
func OnStateChange(fn OnStateChangeFunc) Option {
	return func(r *Rabbus) error {
		r.config.breaker.onStateChange = fn
		return nil
	}
}

// AmqpProvider expose a interface for interacting with amqp broker
func AmqpProvider(provider Amqp) Option {
	return func(r *Rabbus) error {
		if provider != nil {
			r.Amqp = provider
			return nil
		}
		return errors.New("unexpected amqp provider")
	}
}

func (r *Rabbus) wrapMessage(c ListenConfig, sourceChan <-chan amqp.Delivery, targetChan chan ConsumerMessage) {
	for m := range sourceChan {
		targetChan <- newConsumerMessage(m)
	}
}

func (r *Rabbus) handleAmqpClose(err error) {
	for {
		time.Sleep(time.Second)
		aw, err := amqpwrap.New(r.config.dsn, r.config.passiveex)
		if err != nil {
			continue
		}

		r.mu.Lock()
		r.Amqp = aw
		r.mu.Unlock()
		for i := 1; i <= r.conDeclared; i++ {
			r.reconn <- struct{}{}
		}
		break
	}
}

func (r *Rabbus) listenReconn(c ListenConfig, messages chan ConsumerMessage) {
	for range r.reconn {
		msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
		if err != nil {
			r.Close()
			continue
		}

		go r.wrapMessage(c, msgs, messages)
		go r.listenReconn(c, messages)
		break
	}
}

func newBreakerSettings(c config) gobreaker.Settings {
	s := gobreaker.Settings{}
	s.Name = "rabbus-circuit-breaker"
	s.Interval = c.breaker.interval
	s.Timeout = c.breaker.timeout
	s.ReadyToTrip = func(counts gobreaker.Counts) bool {
		return counts.ConsecutiveFailures > c.breaker.threshold
	}
	if c.breaker.onStateChange != nil {
		s.OnStateChange = func(name string, from gobreaker.State, to gobreaker.State) {
			c.breaker.onStateChange(name, from.String(), to.String())
		}
	}
	return s
}
<commit_msg>the counter is locked by the mutex<commit_after>package rabbus

import (
	"context"
	"errors"
	"sync"
	"time"

	amqpwrap "github.com/rafaeljesus/rabbus/internal/amqp"
	"github.com/rafaeljesus/retry-go"
	"github.com/sony/gobreaker"
	"github.com/streadway/amqp"
)

const (
	// Transient means higher throughput but messages will not be restored on broker restart.
	Transient uint8 = 1
	// Persistent messages will be restored to durable queues and lost on non-durable queues during server restart.
	Persistent uint8 = 2
	// ContentTypeJSON define json content type
	ContentTypeJSON = "application/json"
	// ContentTypePlain define plain text content type
	ContentTypePlain = "plain/text"

	contentEncoding = "UTF-8"
)

type (
	// Option represents an option you can pass to New.
	// See the documentation for the individual options.
	Option func(*Rabbus) error
	// OnStateChangeFunc is the callback function when circuit breaker state changes.
	OnStateChangeFunc func(name, from, to string)

	// Message carries fields for sending messages.
	Message struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// Payload the message payload.
		Payload []byte
		// DeliveryMode indicates if the is Persistent or Transient.
		DeliveryMode uint8
		// ContentType the message content-type.
		ContentType string
		// Headers the message application headers
		Headers map[string]interface{}
	}

	// ListenConfig carries fields for listening messages.
	ListenConfig struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// PassiveExchange determines a passive exchange connection it uses
		// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
		PassiveExchange bool
		// Queue the queue name
		Queue string
	}

	// Delivery wraps amqp.Delivery struct
	Delivery struct {
		amqp.Delivery
	}

	// Rabbus interpret (implement) Rabbus interface definition
	Rabbus struct {
		Amqp
		mu         sync.RWMutex
		breaker    *gobreaker.CircuitBreaker
		emit       chan Message
		emitErr    chan error
		emitOk     chan struct{}
		reconn     chan struct{}
		exDeclared map[string]struct{}
		config
		conDeclared int // conDeclared is a counter for the declared consumers
	}

	// Amqp expose a interface for interacting with amqp broker
	Amqp interface {
		// Publish wraps amqp.Publish method
		Publish(exchange, key string, opts amqp.Publishing) error
		// CreateConsumer creates a amqp consumer
		CreateConsumer(exchange, key, kind, queue string, durable bool) (<-chan amqp.Delivery, error)
		// WithExchange creates a amqp exchange
		WithExchange(exchange, kind string, durable bool) error
		// WithQos wrapper over amqp.Qos method
		WithQos(count, size int, global bool) error
		// NotifyClose wrapper over notifyClose method
		NotifyClose(c chan *amqp.Error) chan *amqp.Error
		// Close closes the running amqp connection and channel
		Close() error
	}

	config struct {
		dsn                string
		durable, passiveex bool
		retrycfg
		breaker
		qos
	}

	retrycfg struct {
		attempts              int
		sleep, reconnectSleep time.Duration
	}

	breaker struct {
		interval, timeout time.Duration
		threshold         uint32
		onStateChange     OnStateChangeFunc
	}

	qos struct {
		prefetchCount, prefetchSize int
		global                      bool
	}
)

func (lc ListenConfig) validate() error {
	if lc.Exchange == "" {
		return ErrMissingExchange
	}

	if lc.Kind == "" {
		return ErrMissingKind
	}

	if lc.Queue == "" {
		return ErrMissingQueue
	}

	return nil
}

// New returns a new Rabbus configured with the
// variables from the config parameter, or returning an non-nil err
// if an error occurred while creating connection and channel.
func New(dsn string, options ...Option) (*Rabbus, error) {
	r := &Rabbus{
		emit:       make(chan Message),
		emitErr:    make(chan error),
		emitOk:     make(chan struct{}),
		reconn:     make(chan struct{}, 10),
		exDeclared: make(map[string]struct{}),
	}

	for _, o := range options {
		if err := o(r); err != nil {
			return nil, err
		}
	}

	if r.Amqp == nil {
		amqpWrapper, err := amqpwrap.New(dsn, r.config.passiveex)
		if err != nil {
			return nil, err
		}
		r.Amqp = amqpWrapper
	}

	if err := r.WithQos(
		r.config.qos.prefetchCount,
		r.config.qos.prefetchSize,
		r.config.qos.global,
	); err != nil {
		return nil, err
	}

	r.config.dsn = dsn
	r.breaker = gobreaker.NewCircuitBreaker(newBreakerSettings(r.config))

	return r, nil
}

// Run starts rabbus channels for emitting and listening for amqp connection close
// returns ctx error in case of any.
func (r *Rabbus) Run(ctx context.Context) error {
	for {
		select {
		case m, ok := <-r.emit:
			if ok {
				r.produce(m)
			}
		case err, ok := <-r.NotifyClose(make(chan *amqp.Error)):
			if ok {
				r.handleAmqpClose(err)
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// EmitAsync emits a message to RabbitMQ, but does not wait for the response from broker.
func (r *Rabbus) EmitAsync() chan<- Message { return r.emit }

// EmitErr returns an error if encoding payload fails, or if after circuit breaker is open or retries attempts exceed.
func (r *Rabbus) EmitErr() <-chan error { return r.emitErr }

// EmitOk returns true when the message was sent.
func (r *Rabbus) EmitOk() <-chan struct{} { return r.emitOk }

// Listen to a message from RabbitMQ, returns
// an error if exchange, queue name and function handler not passed or if an error occurred while creating
// amqp consumer.
func (r *Rabbus) Listen(c ListenConfig) (chan ConsumerMessage, error) {
	if err := c.validate(); err != nil {
		return nil, err
	}

	msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
	if err != nil {
		return nil, err
	}

	r.mu.Lock()
	r.conDeclared++ // increase the declared consumers counter
	r.exDeclared[c.Exchange] = struct{}{}
	r.mu.Unlock()

	messages := make(chan ConsumerMessage, 256)
	go r.wrapMessage(c, msgs, messages)
	go r.listenReconn(c, messages)

	return messages, nil
}

// Close channels and attempt to close channel and connection.
func (r *Rabbus) Close() error {
	close(r.emit)
	close(r.emitOk)
	close(r.emitErr)
	close(r.reconn)

	return r.Amqp.Close()
}

func (r *Rabbus) produce(m Message) {
	if _, ok := r.exDeclared[m.Exchange]; !ok {
		if err := r.WithExchange(m.Exchange, m.Kind, r.config.durable); err != nil {
			r.emitErr <- err
			return
		}
		r.exDeclared[m.Exchange] = struct{}{}
	}

	if m.ContentType == "" {
		m.ContentType = ContentTypeJSON
	}

	if m.DeliveryMode == 0 {
		m.DeliveryMode = Persistent
	}

	opts := amqp.Publishing{
		Headers:         amqp.Table(m.Headers),
		ContentType:     m.ContentType,
		ContentEncoding: contentEncoding,
		DeliveryMode:    m.DeliveryMode,
		Timestamp:       time.Now(),
		Body:            m.Payload,
	}

	if _, err := r.breaker.Execute(func() (interface{}, error) {
		return nil, retry.Do(func() error {
			return r.Publish(m.Exchange, m.Key, opts)
		}, r.config.retrycfg.attempts, r.config.retrycfg.sleep)
	}); err != nil {
		r.emitErr <- err
		return
	}

	r.emitOk <- struct{}{}
}

// Durable indicates of the queue will survive broker restarts. Default to true.
func Durable(durable bool) Option {
	return func(r *Rabbus) error {
		r.config.durable = durable
		return nil
	}
}

// PassiveExchange forces passive connection with all exchanges using
// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
func PassiveExchange(passiveex bool) Option {
	return func(r *Rabbus) error {
		r.config.passiveex = passiveex
		return nil
	}
}

// PrefetchCount limit the number of unacknowledged messages.
func PrefetchCount(count int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchCount = count
		return nil
	}
}

// PrefetchSize when greater than zero, the server will try to keep at least
// that many bytes of deliveries flushed to the network before receiving
// acknowledgments from the consumers.
func PrefetchSize(size int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchSize = size
		return nil
	}
}

// QosGlobal when global is true, these Qos settings apply to all existing and future
// consumers on all channels on the same connection. When false, the Channel.Qos
// settings will apply to all existing and future consumers on this channel.
// RabbitMQ does not implement the global flag.
func QosGlobal(global bool) Option {
	return func(r *Rabbus) error {
		r.config.qos.global = global
		return nil
	}
}

// Attempts is the max number of retries on broker outages.
func Attempts(attempts int) Option {
	return func(r *Rabbus) error {
		r.config.retrycfg.attempts = attempts
		return nil
	}
}

// Sleep is the sleep time of the retry mechanism.
func Sleep(sleep time.Duration) Option {
	return func(r *Rabbus) error {
		if sleep == 0 {
			r.config.retrycfg.reconnectSleep = time.Second * 10
		}
		r.config.retrycfg.sleep = sleep
		return nil
	}
}

// BreakerInterval is the cyclic period of the closed state for CircuitBreaker to clear the internal counts,
// If Interval is 0, CircuitBreaker doesn't clear the internal counts during the closed state.
func BreakerInterval(interval time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.interval = interval
		return nil
	}
}

// BreakerTimeout is the period of the open state, after which the state of CircuitBreaker becomes half-open.
// If Timeout is 0, the timeout value of CircuitBreaker is set to 60 seconds.
func BreakerTimeout(timeout time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.timeout = timeout
		return nil
	}
}

// Threshold when a threshold of failures has been reached, future calls to the broker will not run.
// During this state, the circuit breaker will periodically allow the calls to run and, if it is successful,
// will start running the function again. Default value is 5.
func Threshold(threshold uint32) Option {
	return func(r *Rabbus) error {
		if threshold == 0 {
			threshold = 5
		}
		r.config.breaker.threshold = threshold
		return nil
	}
}

// OnStateChange is called whenever the state of CircuitBreaker changes.
func OnStateChange(fn OnStateChangeFunc) Option {
	return func(r *Rabbus) error {
		r.config.breaker.onStateChange = fn
		return nil
	}
}

// AmqpProvider expose a interface for interacting with amqp broker
func AmqpProvider(provider Amqp) Option {
	return func(r *Rabbus) error {
		if provider != nil {
			r.Amqp = provider
			return nil
		}
		return errors.New("unexpected amqp provider")
	}
}

func (r *Rabbus) wrapMessage(c ListenConfig, sourceChan <-chan amqp.Delivery, targetChan chan ConsumerMessage) {
	for m := range sourceChan {
		targetChan <- newConsumerMessage(m)
	}
}

func (r *Rabbus) handleAmqpClose(err error) {
	for {
		time.Sleep(time.Second)
		aw, err := amqpwrap.New(r.config.dsn, r.config.passiveex)
		if err != nil {
			continue
		}

		r.mu.Lock()
		r.Amqp = aw
		r.mu.Unlock()
		for i := 1; i <= r.conDeclared; i++ {
			r.reconn <- struct{}{}
		}
		break
	}
}

func (r *Rabbus) listenReconn(c ListenConfig, messages chan ConsumerMessage) {
	for range r.reconn {
		msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
		if err != nil {
			r.Close()
			continue
		}

		go r.wrapMessage(c, msgs, messages)
		go r.listenReconn(c, messages)
		break
	}
}

func newBreakerSettings(c config) gobreaker.Settings {
	s := gobreaker.Settings{}
	s.Name = "rabbus-circuit-breaker"
	s.Interval = c.breaker.interval
	s.Timeout = c.breaker.timeout
	s.ReadyToTrip = func(counts gobreaker.Counts) bool {
		return counts.ConsecutiveFailures > c.breaker.threshold
	}
	if c.breaker.onStateChange != nil {
		s.OnStateChange = func(name string, from gobreaker.State, to gobreaker.State) {
			c.breaker.onStateChange(name, from.String(), to.String())
		}
	}
	return s
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 cmd

import (
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"k8s.io/minikube/pkg/minikube/cluster"
	"k8s.io/minikube/pkg/minikube/config"
	"k8s.io/minikube/pkg/minikube/constants"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/kubeconfig"
	"k8s.io/minikube/pkg/minikube/machine"
	"k8s.io/minikube/pkg/minikube/out"
)

// updateContextCmd represents the update-context command
var updateContextCmd = &cobra.Command{
	Use:   "update-context",
	Short: "Verify the IP address of the running cluster in kubeconfig.",
	Long: `Retrieves the IP address of the running cluster, checks it
			with IP in kubeconfig, and corrects kubeconfig if incorrect.`,
	Run: func(cmd *cobra.Command, args []string) {
		api, err := machine.NewAPIClient()
		if err != nil {
			exit.WithError("Error getting client", err)
		}
		defer api.Close()
		machineName := viper.GetString(config.MachineProfile)
		ip, err := cluster.GetHostDriverIP(api, machineName)
		if err != nil {
			exit.WithError("Error host driver ip status", err)
		}
		updated := false
		kubeConfigPath := os.Getenv("KUBECONFIG")
		if kubeConfigPath == "" {
			updated, err = kubeconfig.UpdateIP(ip, machineName, constants.KubeconfigPath)
		} else {
			updated, err = kubeconfig.UpdateIP(ip, machineName, kubeConfigPath)
		}
		if err != nil {
			exit.WithError("update config", err)
		}
		if updated {
			out.T(out.Celebrate, "{{.machine}} IP has been updated to point at {{.ip}}", out.V{"machine": machineName, "ip": ip})
		} else {
			out.T(out.Meh, "{{.machine}} IP was already correctly configured for {{.ip}}", out.V{"machine": machineName, "ip": ip})
		}

	},
}
<commit_msg>Minor: change assignment to declaration<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 cmd

import (
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"k8s.io/minikube/pkg/minikube/cluster"
	"k8s.io/minikube/pkg/minikube/config"
	"k8s.io/minikube/pkg/minikube/constants"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/kubeconfig"
	"k8s.io/minikube/pkg/minikube/machine"
	"k8s.io/minikube/pkg/minikube/out"
)

// updateContextCmd represents the update-context command
var updateContextCmd = &cobra.Command{
	Use:   "update-context",
	Short: "Verify the IP address of the running cluster in kubeconfig.",
	Long: `Retrieves the IP address of the running cluster, checks it
			with IP in kubeconfig, and corrects kubeconfig if incorrect.`,
	Run: func(cmd *cobra.Command, args []string) {
		api, err := machine.NewAPIClient()
		if err != nil {
			exit.WithError("Error getting client", err)
		}
		defer api.Close()
		machineName := viper.GetString(config.MachineProfile)
		ip, err := cluster.GetHostDriverIP(api, machineName)
		if err != nil {
			exit.WithError("Error host driver ip status", err)
		}
		var updated bool
		kubeConfigPath := os.Getenv("KUBECONFIG")
		if kubeConfigPath == "" {
			updated, err = kubeconfig.UpdateIP(ip, machineName, constants.KubeconfigPath)
		} else {
			updated, err = kubeconfig.UpdateIP(ip, machineName, kubeConfigPath)
		}
		if err != nil {
			exit.WithError("update config", err)
		}
		if updated {
			out.T(out.Celebrate, "{{.machine}} IP has been updated to point at {{.ip}}", out.V{"machine": machineName, "ip": ip})
		} else {
			out.T(out.Meh, "{{.machine}} IP was already correctly configured for {{.ip}}", out.V{"machine": machineName, "ip": ip})
		}

	},
}
<|endoftext|>


<commit_before>package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	crand "crypto/rand"
	"encoding/hex"
	"io"
	"sync"

	. "github.com/tendermint/go-common"
)

var gRandInfo *randInfo

func init() {
	gRandInfo = &randInfo{}
	gRandInfo.AddSeed(randBytes(32)) // Init
}

// Add additional bytes of randomness, e.g. from hardware, user-input, etc.
// It is OK to call it multiple times.  It does not deminish security.
func Seed(seedBytes []byte) {
	gRandInfo.AddSeed(seedBytes)
}

// This only uses the OS's randomness
func randBytes(numBytes int) []byte {
	b := make([]byte, numBytes)
	_, err := crand.Read(b)
	if err != nil {
		PanicCrisis(err)
	}
	return b
}

// This uses the OS and the Seed(s).
func CRandBytes(numBytes int) []byte {
	b := make([]byte, numBytes)
	_, err := gRandInfo.Read(b)
	if err != nil {
		PanicCrisis(err)
	}
	return b
}

// RandHex(24) gives 96 bits of randomness, strong enough for most purposes.
func CRandHex(numDigits int) string {
	return hex.EncodeToString(CRandBytes(numDigits / 2))
}

// Returns a crand.Reader mixed with user-supplied entropy
func CReader() io.Reader {
	return gRandInfo
}

//--------------------------------------------------------------------------------

type randInfo struct {
	mtx          sync.Mutex
	seedBytes    [32]byte
	cipherAES256 cipher.Block
	streamAES256 cipher.Stream
	reader       io.Reader
}

// You can call this as many times as you'd like.
// XXX TODO review
func (ri *randInfo) AddSeed(seedBytes []byte) {
	ri.mtx.Lock()
	defer ri.mtx.Unlock()
	// Make new ri.seedBytes
	hashBytes := Sha256(seedBytes)
	hashBytes32 := [32]byte{}
	copy(hashBytes32[:], hashBytes)
	ri.seedBytes = xorBytes32(ri.seedBytes, hashBytes32)
	// Create new cipher.Block
	var err error
	ri.cipherAES256, err = aes.NewCipher(ri.seedBytes[:])
	if err != nil {
		PanicSanity("Error creating AES256 cipher: " + err.Error())
	}
	// Create new stream
	ri.streamAES256 = cipher.NewCTR(ri.cipherAES256, randBytes(aes.BlockSize))
	// Create new reader
	ri.reader = &cipher.StreamReader{S: ri.streamAES256, R: crand.Reader}
}

func (ri *randInfo) Read(b []byte) (n int, err error) {
	ri.mtx.Lock()
	defer ri.mtx.Unlock()
	return ri.reader.Read(b)
}

func xorBytes32(bytesA [32]byte, bytesB [32]byte) (res [32]byte) {
	for i, b := range bytesA {
		res[i] = b ^ bytesB[i]
	}
	return res
}
<commit_msg>s/Seed/MixEntropy/g<commit_after>package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	crand "crypto/rand"
	"encoding/hex"
	"io"
	"sync"

	. "github.com/tendermint/go-common"
)

var gRandInfo *randInfo

func init() {
	gRandInfo = &randInfo{}
	gRandInfo.MixEntropy(randBytes(32)) // Init
}

// Mix additional bytes of randomness, e.g. from hardware, user-input, etc.
// It is OK to call it multiple times.  It does not diminish security.
func MixEntropy(seedBytes []byte) {
	gRandInfo.MixEntropy(seedBytes)
}

// This only uses the OS's randomness
func randBytes(numBytes int) []byte {
	b := make([]byte, numBytes)
	_, err := crand.Read(b)
	if err != nil {
		PanicCrisis(err)
	}
	return b
}

// This uses the OS and the Seed(s).
func CRandBytes(numBytes int) []byte {
	b := make([]byte, numBytes)
	_, err := gRandInfo.Read(b)
	if err != nil {
		PanicCrisis(err)
	}
	return b
}

// RandHex(24) gives 96 bits of randomness, strong enough for most purposes.
func CRandHex(numDigits int) string {
	return hex.EncodeToString(CRandBytes(numDigits / 2))
}

// Returns a crand.Reader mixed with user-supplied entropy
func CReader() io.Reader {
	return gRandInfo
}

//--------------------------------------------------------------------------------

type randInfo struct {
	mtx          sync.Mutex
	seedBytes    [32]byte
	cipherAES256 cipher.Block
	streamAES256 cipher.Stream
	reader       io.Reader
}

// You can call this as many times as you'd like.
// XXX TODO review
func (ri *randInfo) MixEntropy(seedBytes []byte) {
	ri.mtx.Lock()
	defer ri.mtx.Unlock()
	// Make new ri.seedBytes
	hashBytes := Sha256(seedBytes)
	hashBytes32 := [32]byte{}
	copy(hashBytes32[:], hashBytes)
	ri.seedBytes = xorBytes32(ri.seedBytes, hashBytes32)
	// Create new cipher.Block
	var err error
	ri.cipherAES256, err = aes.NewCipher(ri.seedBytes[:])
	if err != nil {
		PanicSanity("Error creating AES256 cipher: " + err.Error())
	}
	// Create new stream
	ri.streamAES256 = cipher.NewCTR(ri.cipherAES256, randBytes(aes.BlockSize))
	// Create new reader
	ri.reader = &cipher.StreamReader{S: ri.streamAES256, R: crand.Reader}
}

func (ri *randInfo) Read(b []byte) (n int, err error) {
	ri.mtx.Lock()
	defer ri.mtx.Unlock()
	return ri.reader.Read(b)
}

func xorBytes32(bytesA [32]byte, bytesB [32]byte) (res [32]byte) {
	for i, b := range bytesA {
		res[i] = b ^ bytesB[i]
	}
	return res
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"syscall"

	"bufio"
	"github.com/rgbkrk/libvirt-go"
	"html"
	"strings"
)

type virtlauncher struct {
	domainXML  string
	domainUUID string
	connURI    string
	user       string
	pass       string
}

func main() {
	xmlPath := flag.String("domain-path", "/var/run/virt-launcher/dom.xml", "Where to look for the domain xml.")
	downwardAPIPath := flag.String("downward-api-path", "", "Load domain from this downward API file")
	conUri := flag.String("libvirt-uri", "qemu:///system", "Libvirt connection string.")
	user := flag.String("user", "vdsm@ovirt", "Libvirt user")
	pass := flag.String("pass", "shibboleth", "Libvirt password")
	receiveOnly := flag.Bool("receive-only", false, "Do not create the domain")
	flag.Parse()

	launcher := virtlauncher{
		connURI: *conUri,
		user:    *user,
		pass:    *pass,
	}

	if !*receiveOnly {
		launcher.ReadDomainXML(*xmlPath, *downwardAPIPath)
		launcher.CreateDomain()
	}

	waitUntilSignal()
}

func (vl *virtlauncher) CreateDomain() {
	conn, err := libvirt.NewVirConnectionWithAuth(vl.connURI, vl.user, vl.pass)
	if err != nil {
		panic(fmt.Sprintf("Could not connect to libvirt using %s: %s", vl.connURI, err))
	}
	defer func() {
		if _, closeErr := conn.CloseConnection(); closeErr != nil {
			log.Fatalf("CloseConnection() failed: %s", closeErr)
		}
		log.Print("Connection closed")
	}()

	log.Print("Libvirt connection established")

	// Launch VM
	_, err = conn.DomainCreateXML(vl.domainXML, 0)
	if err != nil {
		panic(fmt.Sprintf("Could not create the libvirt domain: %s", err))
	}

	log.Print("Domain started")
}

func waitUntilSignal() {
	// Wait for termination
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT,
	)

	log.Printf("Waiting forever...")
	s := <-c
	log.Print("Got signal: ", s)
}

func (vl *virtlauncher) ReadDomainXML(xmlPath string, downwardAPIPath string) {
	if downwardAPIPath == "" {
		log.Print("Loading Domain from XML file")
		rawXML, err := ioutil.ReadFile(xmlPath)
		if err != nil {
			log.Fatal(err)
		}
		vl.domainXML = string(rawXML)
	} else {
		log.Print("Loading Domain from downward API file")
		f, err := os.Open(downwardAPIPath)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, `domainXML="`) {
				vl.domainXML = DecodeDomainXML(strings.Trim(strings.TrimPrefix(line, "domainXML="), `"`))
			}
		}

	}

	if vl.domainXML == "" {
		panic("Could not load domain XML. The resulting XML is empty")
	}
	log.Print("Domain description loaded.")
}

func DecodeDomainXML(domainXML string) string {
	decodedXML := html.UnescapeString(string(domainXML))
	decodedXML = strings.Replace(decodedXML, "\\\\", "\\", -1)
	decodedXML = strings.Replace(decodedXML, "\\n", "\n", -1)
	return decodedXML
}
<commit_msg>virt-launcher: monitor qemu process<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"bufio"
	"github.com/rgbkrk/libvirt-go"
	"html"
	"strconv"
	"strings"
)

type virtlauncher struct {
	domainXML  string
	domainUUID string
	connURI    string
	user       string
	pass       string
}

type monitor struct {
	timeout   time.Duration
	pid       int
	exename   string
	start     time.Time
	isDone    bool
	debugMode bool
}

func (mon *monitor) refresh() {
	if mon.isDone {
		log.Print("Called refresh after done!")
		return
	}

	if mon.debugMode {
		log.Printf("Refreshing executable %s pid %d", mon.exename, mon.pid)
	}

	// is the procecess there?
	if mon.pid == 0 {
		var err error
		mon.pid, err = pidOf(mon.exename)
		if err == nil {
			log.Printf("Found PID for %s: %d", mon.exename, mon.pid)
		} else {
			if mon.debugMode {
				log.Printf("Missing PID for %s", mon.exename)
			}
			// if the proces is not there yet, is it too late?
			elapsed := time.Since(mon.start)
			if mon.timeout > 0 && elapsed >= mon.timeout {
				log.Printf("%s not found after timeout", mon.exename)
				mon.isDone = true
			}
		}
		return
	}

	// is the process gone? mon.pid != 0 -> mon.pid == 0
	// note libvirt deliver one event for this, but since we need
	// to poll procfs anyway to detect incoming QEMUs after migrations,
	// we choose to not use this. Bonus: we can close the connection
	// and open it only when needed, which is a tiny part of the
	// virt-launcher lifetime.
	if !pidExists(mon.pid) {
		log.Printf("Process %s is gone!", mon.exename)
		mon.pid = 0
		mon.isDone = true
		return
	}

	return
}

func (mon *monitor) RunForever(startTimeout time.Duration) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT,
	)

	// random value, no real rationale
	rate := 500 * time.Millisecond

	if mon.debugMode {
		timeoutRepr := fmt.Sprintf("%v", startTimeout)
		if startTimeout == 0 {
			timeoutRepr = "disabled"
		}
		log.Printf("Monitoring loop: rate %v start timeout %s", rate, timeoutRepr)
	}

	ticker := time.NewTicker(rate)

	gotSignal := false
	mon.isDone = false
	mon.timeout = startTimeout
	mon.start = time.Now()

	log.Printf("Waiting forever...")
	for !gotSignal && !mon.isDone {
		select {
		case <-ticker.C:
			mon.refresh()
		case s := <-c:
			log.Print("Got signal: ", s)
			gotSignal = true
		}
	}

	ticker.Stop()
	log.Printf("Exiting...")
}

func main() {
	startTimeout := 0 * time.Second

	xmlPath := flag.String("domain-path", "/var/run/virt-launcher/dom.xml", "Where to look for the domain xml.")
	downwardAPIPath := flag.String("downward-api-path", "", "Load domain from this downward API file")
	conUri := flag.String("libvirt-uri", "qemu:///system", "Libvirt connection string.")
	user := flag.String("user", "vdsm@ovirt", "Libvirt user")
	pass := flag.String("pass", "shibboleth", "Libvirt password")
	receiveOnly := flag.Bool("receive-only", false, "Do not create the domain")
	qemuTimeout := flag.Duration("qemu-timeout", startTimeout, "Amount of time to wait for qemu")
	debugMode := flag.Bool("debug", false, "Enable debug messages")
	flag.Parse()

	mon := monitor{
		exename:   "qemu",
		debugMode: *debugMode,
	}

	launcher := virtlauncher{
		connURI: *conUri,
		user:    *user,
		pass:    *pass,
	}

	if !*receiveOnly {
		launcher.ReadDomainXML(*xmlPath, *downwardAPIPath)
		launcher.CreateDomain()
	}

	mon.RunForever(*qemuTimeout)
}

func (vl *virtlauncher) CreateDomain() {
	conn, err := libvirt.NewVirConnectionWithAuth(vl.connURI, vl.user, vl.pass)
	if err != nil {
		panic(fmt.Sprintf("Could not connect to libvirt using %s: %s", vl.connURI, err))
	}
	defer func() {
		if _, closeErr := conn.CloseConnection(); closeErr != nil {
			log.Fatalf("CloseConnection() failed: %s", closeErr)
		}
		log.Print("Connection closed")
	}()

	log.Print("Libvirt connection established")

	// Launch VM
	_, err = conn.DomainCreateXML(vl.domainXML, 0)
	if err != nil {
		panic(fmt.Sprintf("Could not create the libvirt domain: %s", err))
	}

	log.Print("Domain started")
}

func (vl *virtlauncher) ReadDomainXML(xmlPath string, downwardAPIPath string) {
	if downwardAPIPath == "" {
		log.Print("Loading Domain from XML file")
		rawXML, err := ioutil.ReadFile(xmlPath)
		if err != nil {
			log.Fatal(err)
		}
		vl.domainXML = string(rawXML)
	} else {
		log.Print("Loading Domain from downward API file")
		f, err := os.Open(downwardAPIPath)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, `domainXML="`) {
				vl.domainXML = DecodeDomainXML(strings.Trim(strings.TrimPrefix(line, "domainXML="), `"`))
			}
		}

	}

	if vl.domainXML == "" {
		panic("Could not load domain XML. The resulting XML is empty")
	}
	log.Print("Domain description loaded.")
}

func DecodeDomainXML(domainXML string) string {
	decodedXML := html.UnescapeString(string(domainXML))
	decodedXML = strings.Replace(decodedXML, "\\\\", "\\", -1)
	decodedXML = strings.Replace(decodedXML, "\\n", "\n", -1)
	return decodedXML
}

func readProcCmdline(pathname string) ([]string, error) {
	content, err := ioutil.ReadFile(pathname)
	if err != nil {
		return nil, err
	}

	return strings.Split(string(content), "\x00"), nil
}

func pidOf(exename string) (int, error) {
	entries, err := filepath.Glob("/proc/*/cmdline")
	if err != nil {
		return 0, err
	}
	for _, entry := range entries {
		argv, err := readProcCmdline(entry)
		if err != nil {
			return 0, err
		}

		// we need to support both
		// - /usr/bin/qemu-system-$ARCH (fedora)
		// - /usr/libexec/qemu-kvm (*EL, CentOS)
		match, _ := filepath.Match(fmt.Sprintf("%s*", exename), filepath.Base(argv[0]))

		if match {
			//   <empty> /    proc     /    $PID   /   cmdline
			// items[0] sep items[1] sep items[2] sep  items[3]
			items := strings.Split(entry, string(os.PathSeparator))
			pid, err := strconv.Atoi(items[2])
			if err != nil {
				return 0, err
			}

			return pid, nil
		}
	}
	return 0, fmt.Errorf("Process %s not found in /proc", exename)
}

func pidExists(pid int) bool {
	path := fmt.Sprintf("/proc/%d/cmdline", pid)
	if _, err := os.Stat(path); err == nil {
		return true
	}
	return false
}
<|endoftext|>