<commit_before>package client

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/camptocamp/bivac/pkg/volume"
)

// Client contains informations needed to connect to a Bivac API
type Client struct {
	remoteAddress string
	psk           string
}

// NewClient returns a Bivac client
func NewClient(remoteAddress string, psk string) (c *Client, err error) {
	c = &Client{
		remoteAddress: remoteAddress,
		psk:           psk,
	}

	var pingResponse map[string]string
	err = c.newRequest(&pingResponse, "GET", "/ping", "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	if pingResponse["type"] != "pong" {
		err = fmt.Errorf("wrong response from the Bivac instance: %v", pingResponse)
		return
	}
	return
}

// GetVolumes returns the list of the volumes managed by Bivac
func (c *Client) GetVolumes() (volumes []volume.Volume, err error) {
	err = c.newRequest(&volumes, "GET", "/volumes", "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	return
}

// BackupVolume requests a backup of a volume
func (c *Client) BackupVolume(volumeName string, force bool) (err error) {
	err = c.newRequest(nil, "POST", fmt.Sprintf("/backup/%s?force=%s", volumeName, strconv.FormatBool(force)), "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	return
}

// RestoreVolume requests a restore of a volume
func (c *Client) RestoreVolume(
	volumeName string,
	force bool,
	snapshotName string,
) (err error) {
	err = c.newRequest(
		nil,
		"POST",
		fmt.Sprintf(
			"/restore/%s/%s?force=%s",
			volumeName,
			snapshotName,
			strconv.FormatBool(force),
		),
		"",
	)
	if err != nil {
		err = fmt.Errorf(
			"failed to connect to the remote Bivac instance: %s",
			err,
		)
		return
	}
	return
}

// RunRawCommand runs a custom Restic command on a volume's repository and returns the output
func (c *Client) RunRawCommand(volumeID string, cmd []string) (output string, err error) {
	var response map[string]interface{}

	postValue := make(map[string][]string)
	postValue["cmd"] = cmd

	postValueEncoded, _ := json.Marshal(postValue)

	err = c.newRequest(&response, "POST", fmt.Sprintf("/restic/%s", volumeID), string(postValueEncoded))
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}

	output = response["data"].(string)
	return
}

// GetInformations returns informations about the Bivac manager
func (c *Client) GetInformations() (informations map[string]string, err error) {
	var data struct {
		Type string `json:"type"`
		Data map[string]string
	}
	err = c.newRequest(&data, "GET", "/info", "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	informations = data.Data
	return
}

func (c *Client) newRequest(data interface{}, method, endpoint, value string) (err error) {
	client := &http.Client{}
	req, err := http.NewRequest(method, c.remoteAddress+endpoint, bytes.NewBuffer([]byte(value)))
	if err != nil {
		err = fmt.Errorf("failed to build request: %s", err)
		return
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.psk))

	res, err := client.Do(req)
	if err != nil {
		err = fmt.Errorf("failed to send request: %s", err)
		return
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		err = fmt.Errorf("failed to read body: %s", err)
		return
	}

	if res.StatusCode == http.StatusOK {
		if err := json.Unmarshal(body, &data); err != nil {
			err = fmt.Errorf("failed to unmarshal response from the Bivac instance: %s", err)
			return err
		}
	} else {
		err = fmt.Errorf("received wrong status code from the Bivac instance: [%d] %s", res.StatusCode, string(body))
		return
	}
	return
}
<commit_msg>fix: decode base64 output from raw command<commit_after>package client

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	"github.com/camptocamp/bivac/pkg/volume"
)

// Client contains informations needed to connect to a Bivac API
type Client struct {
	remoteAddress string
	psk           string
}

// NewClient returns a Bivac client
func NewClient(remoteAddress string, psk string) (c *Client, err error) {
	c = &Client{
		remoteAddress: remoteAddress,
		psk:           psk,
	}

	var pingResponse map[string]string
	err = c.newRequest(&pingResponse, "GET", "/ping", "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	if pingResponse["type"] != "pong" {
		err = fmt.Errorf("wrong response from the Bivac instance: %v", pingResponse)
		return
	}
	return
}

// GetVolumes returns the list of the volumes managed by Bivac
func (c *Client) GetVolumes() (volumes []volume.Volume, err error) {
	err = c.newRequest(&volumes, "GET", "/volumes", "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	return
}

// BackupVolume requests a backup of a volume
func (c *Client) BackupVolume(volumeName string, force bool) (err error) {
	err = c.newRequest(nil, "POST", fmt.Sprintf("/backup/%s?force=%s", volumeName, strconv.FormatBool(force)), "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	return
}

// RestoreVolume requests a restore of a volume
func (c *Client) RestoreVolume(
	volumeName string,
	force bool,
	snapshotName string,
) (err error) {
	err = c.newRequest(
		nil,
		"POST",
		fmt.Sprintf(
			"/restore/%s/%s?force=%s",
			volumeName,
			snapshotName,
			strconv.FormatBool(force),
		),
		"",
	)
	if err != nil {
		err = fmt.Errorf(
			"failed to connect to the remote Bivac instance: %s",
			err,
		)
		return
	}
	return
}

// RunRawCommand runs a custom Restic command on a volume's repository and returns the output
func (c *Client) RunRawCommand(volumeID string, cmd []string) (output string, err error) {
	var response map[string]interface{}

	postValue := make(map[string][]string)
	postValue["cmd"] = cmd

	postValueEncoded, _ := json.Marshal(postValue)

	err = c.newRequest(&response, "POST", fmt.Sprintf("/restic/%s", volumeID), string(postValueEncoded))
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	decodedOutput, err := base64.StdEncoding.DecodeString(strings.Replace(response["data"].(string), " ", "", -1))
	if err != nil {
		err = fmt.Errorf("failed to decode output: %s -> `%s`", err, strings.Replace(output, " ", "", -1))
	}
	output = string(decodedOutput)
	return
}

// GetInformations returns informations about the Bivac manager
func (c *Client) GetInformations() (informations map[string]string, err error) {
	var data struct {
		Type string `json:"type"`
		Data map[string]string
	}
	err = c.newRequest(&data, "GET", "/info", "")
	if err != nil {
		err = fmt.Errorf("failed to connect to the remote Bivac instance: %s", err)
		return
	}
	informations = data.Data
	return
}

func (c *Client) newRequest(data interface{}, method, endpoint, value string) (err error) {
	client := &http.Client{}
	req, err := http.NewRequest(method, c.remoteAddress+endpoint, bytes.NewBuffer([]byte(value)))
	if err != nil {
		err = fmt.Errorf("failed to build request: %s", err)
		return
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.psk))

	res, err := client.Do(req)
	if err != nil {
		err = fmt.Errorf("failed to send request: %s", err)
		return
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		err = fmt.Errorf("failed to read body: %s", err)
		return
	}

	if res.StatusCode == http.StatusOK {
		if err := json.Unmarshal(body, &data); err != nil {
			err = fmt.Errorf("failed to unmarshal response from the Bivac instance: %s", err)
			return err
		}
	} else {
		err = fmt.Errorf("received wrong status code from the Bivac instance: [%d] %s", res.StatusCode, string(body))
		return
	}
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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

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

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

package commands

import (
	"fmt"
	"strings"

	"github.com/GoogleContainerTools/kaniko/pkg/constants"
	"github.com/GoogleContainerTools/kaniko/pkg/dockerfile"
	"github.com/GoogleContainerTools/kaniko/pkg/util"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/moby/buildkit/frontend/dockerfile/instructions"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// for testing
var (
	Lookup = util.Lookup
)

type UserCommand struct {
	BaseCommand
	cmd *instructions.UserCommand
}

func (r *UserCommand) ExecuteCommand(config *v1.Config, buildArgs *dockerfile.BuildArgs) error {
	logrus.Info("cmd: USER")
	u := r.cmd.User
	userAndGroup := strings.Split(u, ":")
	replacementEnvs := buildArgs.ReplacementEnvs(config.Env)
	userStr, err := util.ResolveEnvironmentReplacement(userAndGroup[0], replacementEnvs, false)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("resolving user %s", userAndGroup[0]))
	}

	if len(userAndGroup) > 1 {
		groupStr, err := util.ResolveEnvironmentReplacement(userAndGroup[1], replacementEnvs, false)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("resolving group %s", userAndGroup[1]))
		}
		userStr = userStr + ":" + groupStr
	}

	config.User = userStr

	// Remove any HOME environment variable previously set.
	env := make([]string, 0, len(config.Env))
	for _, v := range config.Env {
		if strings.HasPrefix(v, constants.HOME+"=") {
			continue
		}
		env = append(env, v)
	}
	config.Env = env

	return nil
}

func (r *UserCommand) String() string {
	return r.cmd.String()
}
<commit_msg>user: keep existing HOME env<commit_after>/*
Copyright 2018 Google LLC

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

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

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

package commands

import (
	"fmt"
	"strings"

	"github.com/GoogleContainerTools/kaniko/pkg/dockerfile"
	"github.com/GoogleContainerTools/kaniko/pkg/util"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/moby/buildkit/frontend/dockerfile/instructions"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// for testing
var (
	Lookup = util.Lookup
)

type UserCommand struct {
	BaseCommand
	cmd *instructions.UserCommand
}

func (r *UserCommand) ExecuteCommand(config *v1.Config, buildArgs *dockerfile.BuildArgs) error {
	logrus.Info("cmd: USER")
	u := r.cmd.User
	userAndGroup := strings.Split(u, ":")
	replacementEnvs := buildArgs.ReplacementEnvs(config.Env)
	userStr, err := util.ResolveEnvironmentReplacement(userAndGroup[0], replacementEnvs, false)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("resolving user %s", userAndGroup[0]))
	}

	if len(userAndGroup) > 1 {
		groupStr, err := util.ResolveEnvironmentReplacement(userAndGroup[1], replacementEnvs, false)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("resolving group %s", userAndGroup[1]))
		}
		userStr = userStr + ":" + groupStr
	}

	config.User = userStr

	return nil
}

func (r *UserCommand) String() string {
	return r.cmd.String()
}
<|endoftext|>


<commit_before>package daemon

import (
	"database/sql"
	"fmt"
	"path/filepath"

	"nimona.io/internal/net"
	"nimona.io/pkg/config"
	"nimona.io/pkg/context"
	"nimona.io/pkg/feedmanager"
	"nimona.io/pkg/hyperspace/resolver"
	"nimona.io/pkg/network"
	"nimona.io/pkg/objectmanager"
	"nimona.io/pkg/objectstore"
	"nimona.io/pkg/peer"
	"nimona.io/pkg/preferences"
	"nimona.io/pkg/sqlobjectstore"
)

type (
	Daemon interface {
		Config() config.Config
		Preferences() preferences.Preferences
		Network() network.Network
		Resolver() resolver.Resolver
		ObjectStore() objectstore.Store
		ObjectManager() objectmanager.ObjectManager
		FeedManager() feedmanager.FeedManager
		// daemon specific methods
		Close()
	}
	daemon struct {
		config        config.Config
		preferences   preferences.Preferences
		configOptions []config.Option
		network       network.Network
		resolver      resolver.Resolver
		objectstore   objectstore.Store
		objectmanager objectmanager.ObjectManager
		feedmanager   feedmanager.FeedManager
		// internal
		listener net.Listener
	}
	Option func(d *daemon) error
)

func New(ctx context.Context, opts ...Option) (Daemon, error) {
	d := &daemon{}

	// apply options
	for _, o := range opts {
		if err := o(d); err != nil {
			return nil, err
		}
	}

	// load config with given options
	cfg, err := config.New(d.configOptions...)
	if err != nil {
		return nil, fmt.Errorf("loading config: %w", err)
	}

	// construct new network
	ntw := network.New(
		ctx,
		network.WithPeerKey(cfg.Peer.PrivateKey),
	)

	if cfg.Peer.BindAddress != "" {
		// start listening
		lis, err := ntw.Listen(
			ctx,
			cfg.Peer.BindAddress,
			network.ListenOnLocalIPs,
			// network.ListenOnExternalPort,
		)
		if err != nil {
			return nil, fmt.Errorf("listening: %w", err)
		}
		d.listener = lis
	}

	// convert shorthands into connection infos
	bootstrapPeers := []*peer.ConnectionInfo{}
	for _, s := range cfg.Peer.Bootstraps {
		bootstrapPeer, err := s.GetConnectionInfo()
		if err != nil {
			return nil, fmt.Errorf("parsing bootstraps: %w", err)
		}
		bootstrapPeers = append(bootstrapPeers, bootstrapPeer)
	}

	// add bootstrap peers as relays
	ntw.RegisterRelays(bootstrapPeers...)

	// construct preferences db
	pdb, err := sql.Open("sqlite3", filepath.Join(cfg.Path, "preferences.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file for preferences: %w", err)
	}

	// construct preferences
	prf, err := preferences.NewSQLProvider(pdb)
	if err != nil {
		return nil, fmt.Errorf("constructing preferences provider: %w", err)
	}

	// construct object store
	db, err := sql.Open("sqlite3", filepath.Join(cfg.Path, "nimona.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file: %w", err)
	}

	str, err := sqlobjectstore.New(db)
	if err != nil {
		return nil, fmt.Errorf("starting sql store: %w", err)
	}

	// construct new resolver
	res := resolver.New(
		ctx,
		ntw,
		str,
		resolver.WithBoostrapPeers(bootstrapPeers...),
	)

	// construct manager
	man := objectmanager.New(
		ctx,
		ntw,
		res,
		str,
	)

	// construct feed manager
	fdm, err := feedmanager.New(
		ctx,
		ntw,
		res,
		str,
		man,
	)
	if err != nil {
		return nil, fmt.Errorf("constructing feed manager, %w", err)
	}

	d.config = *cfg
	d.preferences = prf
	d.network = ntw
	d.resolver = res
	d.objectstore = str
	d.objectmanager = man
	d.feedmanager = fdm

	return d, nil
}

func (d *daemon) Config() config.Config {
	return d.config
}

func (d *daemon) Preferences() preferences.Preferences {
	return d.preferences
}

func (d *daemon) Network() network.Network {
	return d.network
}

func (d *daemon) Resolver() resolver.Resolver {
	return d.resolver
}

func (d *daemon) ObjectStore() objectstore.Store {
	return d.objectstore
}

func (d *daemon) ObjectManager() objectmanager.ObjectManager {
	return d.objectmanager
}

func (d *daemon) FeedManager() feedmanager.FeedManager {
	return d.feedmanager
}

func (d *daemon) Close() {
	if d.listener != nil {
		d.listener.Close() // nolint: errcheck
	}
}
<commit_msg>feat(daemon): add keystreammanager<commit_after>package daemon

import (
	"database/sql"
	"fmt"
	"path/filepath"

	"nimona.io/internal/net"
	"nimona.io/pkg/config"
	"nimona.io/pkg/context"
	"nimona.io/pkg/feedmanager"
	"nimona.io/pkg/hyperspace/resolver"
	"nimona.io/pkg/keystream"
	"nimona.io/pkg/network"
	"nimona.io/pkg/objectmanager"
	"nimona.io/pkg/objectstore"
	"nimona.io/pkg/peer"
	"nimona.io/pkg/preferences"
	"nimona.io/pkg/sqlobjectstore"
)

type (
	Daemon interface {
		Config() config.Config
		Preferences() preferences.Preferences
		Network() network.Network
		Resolver() resolver.Resolver
		ObjectStore() objectstore.Store
		ObjectManager() objectmanager.ObjectManager
		FeedManager() feedmanager.FeedManager
		KeyStreamManager() keystream.KeyStreamManager
		// daemon specific methods
		Close()
	}
	daemon struct {
		config          config.Config
		preferences     preferences.Preferences
		configOptions   []config.Option
		network         network.Network
		resolver        resolver.Resolver
		objectstore     objectstore.Store
		objectmanager   objectmanager.ObjectManager
		feedmanager     feedmanager.FeedManager
		keystreamanager keystream.KeyStreamManager
		// internal
		listener net.Listener
	}
	Option func(d *daemon) error
)

func New(ctx context.Context, opts ...Option) (Daemon, error) {
	d := &daemon{}

	// apply options
	for _, o := range opts {
		if err := o(d); err != nil {
			return nil, err
		}
	}

	// load config with given options
	cfg, err := config.New(d.configOptions...)
	if err != nil {
		return nil, fmt.Errorf("loading config: %w", err)
	}

	// construct new network
	ntw := network.New(
		ctx,
		network.WithPeerKey(cfg.Peer.PrivateKey),
	)

	if cfg.Peer.BindAddress != "" {
		// start listening
		lis, err := ntw.Listen(
			ctx,
			cfg.Peer.BindAddress,
			network.ListenOnLocalIPs,
			// network.ListenOnExternalPort,
		)
		if err != nil {
			return nil, fmt.Errorf("listening: %w", err)
		}
		d.listener = lis
	}

	// convert shorthands into connection infos
	bootstrapPeers := []*peer.ConnectionInfo{}
	for _, s := range cfg.Peer.Bootstraps {
		bootstrapPeer, err := s.GetConnectionInfo()
		if err != nil {
			return nil, fmt.Errorf("parsing bootstraps: %w", err)
		}
		bootstrapPeers = append(bootstrapPeers, bootstrapPeer)
	}

	// add bootstrap peers as relays
	ntw.RegisterRelays(bootstrapPeers...)

	// construct preferences db
	pdb, err := sql.Open("sqlite3", filepath.Join(cfg.Path, "preferences.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file for preferences: %w", err)
	}

	// construct preferences
	prf, err := preferences.NewSQLProvider(pdb)
	if err != nil {
		return nil, fmt.Errorf("constructing preferences provider: %w", err)
	}

	// construct object store
	db, err := sql.Open("sqlite3", filepath.Join(cfg.Path, "nimona.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file: %w", err)
	}

	str, err := sqlobjectstore.New(db)
	if err != nil {
		return nil, fmt.Errorf("starting sql store: %w", err)
	}

	// construct new resolver
	res := resolver.New(
		ctx,
		ntw,
		str,
		resolver.WithBoostrapPeers(bootstrapPeers...),
	)

	// construct manager
	man := objectmanager.New(
		ctx,
		ntw,
		res,
		str,
	)

	// construct feed manager
	fdm, err := feedmanager.New(
		ctx,
		ntw,
		res,
		str,
		man,
	)
	if err != nil {
		return nil, fmt.Errorf("constructing feed manager, %w", err)
	}

	// construct key stream manager
	ksm := keystream.NewKeyManager()

	d.config = *cfg
	d.preferences = prf
	d.network = ntw
	d.resolver = res
	d.objectstore = str
	d.objectmanager = man
	d.feedmanager = fdm
	d.keystreamanager = ksm

	return d, nil
}

func (d *daemon) Config() config.Config {
	return d.config
}

func (d *daemon) Preferences() preferences.Preferences {
	return d.preferences
}

func (d *daemon) Network() network.Network {
	return d.network
}

func (d *daemon) Resolver() resolver.Resolver {
	return d.resolver
}

func (d *daemon) ObjectStore() objectstore.Store {
	return d.objectstore
}

func (d *daemon) ObjectManager() objectmanager.ObjectManager {
	return d.objectmanager
}

func (d *daemon) FeedManager() feedmanager.FeedManager {
	return d.feedmanager
}

func (d *daemon) KeyStreamManager() keystream.KeyStreamManager {
	return d.keystreamanager
}

func (d *daemon) Close() {
	if d.listener != nil {
		d.listener.Close() // nolint: errcheck
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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

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

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

package executor

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/GoogleContainerTools/kaniko/pkg/cache"
	"github.com/GoogleContainerTools/kaniko/pkg/config"
	"github.com/GoogleContainerTools/kaniko/pkg/constants"
	"github.com/GoogleContainerTools/kaniko/pkg/creds"
	"github.com/GoogleContainerTools/kaniko/pkg/timing"
	"github.com/GoogleContainerTools/kaniko/pkg/util"
	"github.com/GoogleContainerTools/kaniko/pkg/version"
	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/layout"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/afero"
)

type withUserAgent struct {
	t http.RoundTripper
}

// for testing
var (
	newRetry = transport.NewRetry
)

const (
	UpstreamClientUaKey = "UPSTREAM_CLIENT_TYPE"
)

// DockerConfLocation returns the file system location of the Docker
// configuration file under the directory set in the DOCKER_CONFIG environment
// variable.  If that variable is not set, it returns the OS-equivalent of
// "/kaniko/.docker/config.json".
func DockerConfLocation() string {
	configFile := "config.json"
	if dockerConfDir := os.Getenv("DOCKER_CONFIG"); dockerConfDir != "" {
		return filepath.Join(dockerConfDir, configFile)
	}
	return string(os.PathSeparator) + filepath.Join("kaniko", ".docker", configFile)
}

func (w *withUserAgent) RoundTrip(r *http.Request) (*http.Response, error) {
	ua := []string{fmt.Sprintf("kaniko/%s", version.Version())}
	if upstream := os.Getenv(UpstreamClientUaKey); upstream != "" {
		ua = append(ua, upstream)
	}
	r.Header.Set("User-Agent", strings.Join(ua, ","))
	return w.t.RoundTrip(r)
}

// for testing
var (
	fs                        = afero.NewOsFs()
	execCommand               = exec.Command
	checkRemotePushPermission = remote.CheckPushPermission
)

// CheckPushPermissions checks that the configured credentials can be used to
// push to every specified destination.
func CheckPushPermissions(opts *config.KanikoOptions) error {
	if opts.NoPush {
		return nil
	}

	checked := map[string]bool{}
	for _, destination := range opts.Destinations {
		destRef, err := name.NewTag(destination, name.WeakValidation)
		if err != nil {
			return errors.Wrap(err, "getting tag for destination")
		}
		if checked[destRef.Context().RepositoryStr()] {
			continue
		}

		registryName := destRef.Repository.Registry.Name()
		// Historically kaniko was pre-configured by default with gcr credential helper,
		// in here we keep the backwards compatibility by enabling the GCR helper only
		// when gcr.io (or pkg.dev) is in one of the destinations.
		if registryName == "gcr.io" || strings.HasSuffix(registryName, ".gcr.io") || strings.HasSuffix(registryName, ".pkg.dev") {
			// Checking for existence of docker.config as it's normally required for
			// authenticated registries and prevent overwriting user provided docker conf
			if _, err := fs.Stat(DockerConfLocation()); os.IsNotExist(err) {
				flags := fmt.Sprintf("--registries=%s", registryName)
				if err := execCommand("docker-credential-gcr", "configure-docker", flags).Run(); err != nil {
					return errors.Wrap(err, "error while configuring docker-credential-gcr helper")
				}
			}
		}
		if opts.Insecure || opts.InsecureRegistries.Contains(registryName) {
			newReg, err := name.NewRegistry(registryName, name.WeakValidation, name.Insecure)
			if err != nil {
				return errors.Wrap(err, "getting new insecure registry")
			}
			destRef.Repository.Registry = newReg
		}
		tr := newRetry(util.MakeTransport(opts, registryName))
		if err := checkRemotePushPermission(destRef, creds.GetKeychain(), tr); err != nil {
			return errors.Wrapf(err, "checking push permission for %q", destRef)
		}
		checked[destRef.Context().RepositoryStr()] = true
	}
	return nil
}

func getDigest(image v1.Image) ([]byte, error) {
	digest, err := image.Digest()
	if err != nil {
		return nil, err
	}
	return []byte(digest.String()), nil
}

// DoPush is responsible for pushing image to the destinations specified in opts
func DoPush(image v1.Image, opts *config.KanikoOptions) error {
	t := timing.Start("Total Push Time")
	var digestByteArray []byte
	var builder strings.Builder
	if opts.DigestFile != "" || opts.ImageNameDigestFile != "" {
		var err error
		digestByteArray, err = getDigest(image)
		if err != nil {
			return errors.Wrap(err, "error fetching digest")
		}
	}

	if opts.DigestFile != "" {
		err := ioutil.WriteFile(opts.DigestFile, digestByteArray, 0644)
		if err != nil {
			return errors.Wrap(err, "writing digest to file failed")
		}
	}

	if opts.OCILayoutPath != "" {
		path, err := layout.Write(opts.OCILayoutPath, empty.Index)
		if err != nil {
			return errors.Wrap(err, "writing empty layout")
		}
		if err := path.AppendImage(image); err != nil {
			return errors.Wrap(err, "appending image")
		}
	}

	destRefs := []name.Tag{}
	for _, destination := range opts.Destinations {
		destRef, err := name.NewTag(destination, name.WeakValidation)
		if err != nil {
			return errors.Wrap(err, "getting tag for destination")
		}
		if opts.ImageNameDigestFile != "" {
			imageName := []byte(destRef.Repository.Name() + "@")
			builder.Write(append(imageName, digestByteArray...))
			builder.WriteString("\n")
		}
		destRefs = append(destRefs, destRef)
	}

	if opts.ImageNameDigestFile != "" {
		err := ioutil.WriteFile(opts.ImageNameDigestFile, []byte(builder.String()), 0644)
		if err != nil {
			return errors.Wrap(err, "writing image name with digest to file failed")
		}
	}

	if opts.TarPath != "" {
		tagToImage := map[name.Tag]v1.Image{}
		for _, destRef := range destRefs {
			tagToImage[destRef] = image
		}
		return tarball.MultiWriteToFile(opts.TarPath, tagToImage)
	}

	if opts.NoPush {
		logrus.Info("Skipping push to container registry due to --no-push flag")
		return nil
	}

	// continue pushing unless an error occurs
	for _, destRef := range destRefs {
		registryName := destRef.Repository.Registry.Name()
		if opts.Insecure || opts.InsecureRegistries.Contains(registryName) {
			newReg, err := name.NewRegistry(registryName, name.WeakValidation, name.Insecure)
			if err != nil {
				return errors.Wrap(err, "getting new insecure registry")
			}
			destRef.Repository.Registry = newReg
		}

		pushAuth, err := creds.GetKeychain().Resolve(destRef.Context().Registry)
		if err != nil {
			return errors.Wrap(err, "resolving pushAuth")
		}

		tr := newRetry(util.MakeTransport(opts, registryName))
		rt := &withUserAgent{t: tr}

		if err := remote.Write(destRef, image, remote.WithAuth(pushAuth), remote.WithTransport(rt)); err != nil {
			return errors.Wrap(err, fmt.Sprintf("failed to push to destination %s", destRef))
		}
	}
	timing.DefaultRun.Stop(t)
	return writeImageOutputs(image, destRefs)
}

func writeImageOutputs(image v1.Image, destRefs []name.Tag) error {
	dir := os.Getenv("BUILDER_OUTPUT")
	if dir == "" {
		return nil
	}
	f, err := fs.Create(filepath.Join(dir, "images"))
	if err != nil {
		return err
	}
	defer f.Close()

	d, err := image.Digest()
	if err != nil {
		return err
	}

	type imageOutput struct {
		Name   string `json:"name"`
		Digest string `json:"digest"`
	}
	for _, r := range destRefs {
		if err := json.NewEncoder(f).Encode(imageOutput{
			Name:   r.String(),
			Digest: d.String(),
		}); err != nil {
			return err
		}
	}
	return nil
}

// pushLayerToCache pushes layer (tagged with cacheKey) to opts.Cache
// if opts.Cache doesn't exist, infer the cache from the given destination
func pushLayerToCache(opts *config.KanikoOptions, cacheKey string, tarPath string, createdBy string) error {
	layer, err := tarball.LayerFromFile(tarPath)
	if err != nil {
		return err
	}
	cache, err := cache.Destination(opts, cacheKey)
	if err != nil {
		return errors.Wrap(err, "getting cache destination")
	}
	logrus.Infof("Pushing layer %s to cache now", cache)
	empty := empty.Image
	empty, err = mutate.CreatedAt(empty, v1.Time{Time: time.Now()})
	if err != nil {
		return errors.Wrap(err, "setting empty image created time")
	}

	empty, err = mutate.Append(empty,
		mutate.Addendum{
			Layer: layer,
			History: v1.History{
				Author:    constants.Author,
				CreatedBy: createdBy,
			},
		},
	)
	if err != nil {
		return errors.Wrap(err, "appending layer onto empty image")
	}
	cacheOpts := *opts
	cacheOpts.TarPath = ""   // tarPath doesn't make sense for Docker layers
	cacheOpts.NoPush = false // we want to push cached layers
	cacheOpts.Destinations = []string{cache}
	cacheOpts.InsecureRegistries = opts.InsecureRegistries
	cacheOpts.SkipTLSVerifyRegistries = opts.SkipTLSVerifyRegistries
	return DoPush(empty, &cacheOpts)
}
<commit_msg>change repo string to just string<commit_after>/*
Copyright 2018 Google LLC

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

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

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

package executor

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/GoogleContainerTools/kaniko/pkg/cache"
	"github.com/GoogleContainerTools/kaniko/pkg/config"
	"github.com/GoogleContainerTools/kaniko/pkg/constants"
	"github.com/GoogleContainerTools/kaniko/pkg/creds"
	"github.com/GoogleContainerTools/kaniko/pkg/timing"
	"github.com/GoogleContainerTools/kaniko/pkg/util"
	"github.com/GoogleContainerTools/kaniko/pkg/version"
	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/layout"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/afero"
)

type withUserAgent struct {
	t http.RoundTripper
}

// for testing
var (
	newRetry = transport.NewRetry
)

const (
	UpstreamClientUaKey = "UPSTREAM_CLIENT_TYPE"
)

// DockerConfLocation returns the file system location of the Docker
// configuration file under the directory set in the DOCKER_CONFIG environment
// variable.  If that variable is not set, it returns the OS-equivalent of
// "/kaniko/.docker/config.json".
func DockerConfLocation() string {
	configFile := "config.json"
	if dockerConfDir := os.Getenv("DOCKER_CONFIG"); dockerConfDir != "" {
		return filepath.Join(dockerConfDir, configFile)
	}
	return string(os.PathSeparator) + filepath.Join("kaniko", ".docker", configFile)
}

func (w *withUserAgent) RoundTrip(r *http.Request) (*http.Response, error) {
	ua := []string{fmt.Sprintf("kaniko/%s", version.Version())}
	if upstream := os.Getenv(UpstreamClientUaKey); upstream != "" {
		ua = append(ua, upstream)
	}
	r.Header.Set("User-Agent", strings.Join(ua, ","))
	return w.t.RoundTrip(r)
}

// for testing
var (
	fs                        = afero.NewOsFs()
	execCommand               = exec.Command
	checkRemotePushPermission = remote.CheckPushPermission
)

// CheckPushPermissions checks that the configured credentials can be used to
// push to every specified destination.
func CheckPushPermissions(opts *config.KanikoOptions) error {
	if opts.NoPush {
		return nil
	}

	checked := map[string]bool{}
	for _, destination := range opts.Destinations {
		destRef, err := name.NewTag(destination, name.WeakValidation)
		if err != nil {
			return errors.Wrap(err, "getting tag for destination")
		}
		if checked[destRef.Context().String()] {
			continue
		}

		registryName := destRef.Repository.Registry.Name()
		// Historically kaniko was pre-configured by default with gcr credential helper,
		// in here we keep the backwards compatibility by enabling the GCR helper only
		// when gcr.io (or pkg.dev) is in one of the destinations.
		if registryName == "gcr.io" || strings.HasSuffix(registryName, ".gcr.io") || strings.HasSuffix(registryName, ".pkg.dev") {
			// Checking for existence of docker.config as it's normally required for
			// authenticated registries and prevent overwriting user provided docker conf
			if _, err := fs.Stat(DockerConfLocation()); os.IsNotExist(err) {
				flags := fmt.Sprintf("--registries=%s", registryName)
				if err := execCommand("docker-credential-gcr", "configure-docker", flags).Run(); err != nil {
					return errors.Wrap(err, "error while configuring docker-credential-gcr helper")
				}
			}
		}
		if opts.Insecure || opts.InsecureRegistries.Contains(registryName) {
			newReg, err := name.NewRegistry(registryName, name.WeakValidation, name.Insecure)
			if err != nil {
				return errors.Wrap(err, "getting new insecure registry")
			}
			destRef.Repository.Registry = newReg
		}
		tr := newRetry(util.MakeTransport(opts, registryName))
		if err := checkRemotePushPermission(destRef, creds.GetKeychain(), tr); err != nil {
			return errors.Wrapf(err, "checking push permission for %q", destRef)
		}
		checked[destRef.Context().String()] = true
	}
	return nil
}

func getDigest(image v1.Image) ([]byte, error) {
	digest, err := image.Digest()
	if err != nil {
		return nil, err
	}
	return []byte(digest.String()), nil
}

// DoPush is responsible for pushing image to the destinations specified in opts
func DoPush(image v1.Image, opts *config.KanikoOptions) error {
	t := timing.Start("Total Push Time")
	var digestByteArray []byte
	var builder strings.Builder
	if opts.DigestFile != "" || opts.ImageNameDigestFile != "" {
		var err error
		digestByteArray, err = getDigest(image)
		if err != nil {
			return errors.Wrap(err, "error fetching digest")
		}
	}

	if opts.DigestFile != "" {
		err := ioutil.WriteFile(opts.DigestFile, digestByteArray, 0644)
		if err != nil {
			return errors.Wrap(err, "writing digest to file failed")
		}
	}

	if opts.OCILayoutPath != "" {
		path, err := layout.Write(opts.OCILayoutPath, empty.Index)
		if err != nil {
			return errors.Wrap(err, "writing empty layout")
		}
		if err := path.AppendImage(image); err != nil {
			return errors.Wrap(err, "appending image")
		}
	}

	destRefs := []name.Tag{}
	for _, destination := range opts.Destinations {
		destRef, err := name.NewTag(destination, name.WeakValidation)
		if err != nil {
			return errors.Wrap(err, "getting tag for destination")
		}
		if opts.ImageNameDigestFile != "" {
			imageName := []byte(destRef.Repository.Name() + "@")
			builder.Write(append(imageName, digestByteArray...))
			builder.WriteString("\n")
		}
		destRefs = append(destRefs, destRef)
	}

	if opts.ImageNameDigestFile != "" {
		err := ioutil.WriteFile(opts.ImageNameDigestFile, []byte(builder.String()), 0644)
		if err != nil {
			return errors.Wrap(err, "writing image name with digest to file failed")
		}
	}

	if opts.TarPath != "" {
		tagToImage := map[name.Tag]v1.Image{}
		for _, destRef := range destRefs {
			tagToImage[destRef] = image
		}
		return tarball.MultiWriteToFile(opts.TarPath, tagToImage)
	}

	if opts.NoPush {
		logrus.Info("Skipping push to container registry due to --no-push flag")
		return nil
	}

	// continue pushing unless an error occurs
	for _, destRef := range destRefs {
		registryName := destRef.Repository.Registry.Name()
		if opts.Insecure || opts.InsecureRegistries.Contains(registryName) {
			newReg, err := name.NewRegistry(registryName, name.WeakValidation, name.Insecure)
			if err != nil {
				return errors.Wrap(err, "getting new insecure registry")
			}
			destRef.Repository.Registry = newReg
		}

		pushAuth, err := creds.GetKeychain().Resolve(destRef.Context().Registry)
		if err != nil {
			return errors.Wrap(err, "resolving pushAuth")
		}

		tr := newRetry(util.MakeTransport(opts, registryName))
		rt := &withUserAgent{t: tr}

		if err := remote.Write(destRef, image, remote.WithAuth(pushAuth), remote.WithTransport(rt)); err != nil {
			return errors.Wrap(err, fmt.Sprintf("failed to push to destination %s", destRef))
		}
	}
	timing.DefaultRun.Stop(t)
	return writeImageOutputs(image, destRefs)
}

func writeImageOutputs(image v1.Image, destRefs []name.Tag) error {
	dir := os.Getenv("BUILDER_OUTPUT")
	if dir == "" {
		return nil
	}
	f, err := fs.Create(filepath.Join(dir, "images"))
	if err != nil {
		return err
	}
	defer f.Close()

	d, err := image.Digest()
	if err != nil {
		return err
	}

	type imageOutput struct {
		Name   string `json:"name"`
		Digest string `json:"digest"`
	}
	for _, r := range destRefs {
		if err := json.NewEncoder(f).Encode(imageOutput{
			Name:   r.String(),
			Digest: d.String(),
		}); err != nil {
			return err
		}
	}
	return nil
}

// pushLayerToCache pushes layer (tagged with cacheKey) to opts.Cache
// if opts.Cache doesn't exist, infer the cache from the given destination
func pushLayerToCache(opts *config.KanikoOptions, cacheKey string, tarPath string, createdBy string) error {
	layer, err := tarball.LayerFromFile(tarPath)
	if err != nil {
		return err
	}
	cache, err := cache.Destination(opts, cacheKey)
	if err != nil {
		return errors.Wrap(err, "getting cache destination")
	}
	logrus.Infof("Pushing layer %s to cache now", cache)
	empty := empty.Image
	empty, err = mutate.CreatedAt(empty, v1.Time{Time: time.Now()})
	if err != nil {
		return errors.Wrap(err, "setting empty image created time")
	}

	empty, err = mutate.Append(empty,
		mutate.Addendum{
			Layer: layer,
			History: v1.History{
				Author:    constants.Author,
				CreatedBy: createdBy,
			},
		},
	)
	if err != nil {
		return errors.Wrap(err, "appending layer onto empty image")
	}
	cacheOpts := *opts
	cacheOpts.TarPath = ""   // tarPath doesn't make sense for Docker layers
	cacheOpts.NoPush = false // we want to push cached layers
	cacheOpts.Destinations = []string{cache}
	cacheOpts.InsecureRegistries = opts.InsecureRegistries
	cacheOpts.SkipTLSVerifyRegistries = opts.SkipTLSVerifyRegistries
	return DoPush(empty, &cacheOpts)
}
<|endoftext|>


<commit_before>package image

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

	"github.com/pkg/errors"
	"github.com/rancher/norman/types/convert"
	v32 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	libhelm "github.com/rancher/rancher/pkg/catalog/helm"
	util "github.com/rancher/rancher/pkg/cluster"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/settings"
	rketypes "github.com/rancher/rke/types"
	img "github.com/rancher/rke/types/image"
	"gopkg.in/yaml.v2"
)

type OSType int

type chart struct {
	dir     string
	version string
}

const (
	Linux OSType = iota
	Windows
)

func Resolve(image string) string {
	return ResolveWithCluster(image, nil)
}

func ResolveWithCluster(image string, cluster *v3.Cluster) string {
	reg := util.GetPrivateRepoURL(cluster)
	if reg != "" && !strings.HasPrefix(image, reg) {
		//Images from Dockerhub Library repo, we add rancher prefix when using private registry
		if !strings.Contains(image, "/") {
			image = "rancher/" + image
		}
		return path.Join(reg, image)
	}

	return image
}

func getChartAndVersion(path string) (map[string]chart, error) {
	rtn := map[string]chart{}
	helm := libhelm.Helm{
		LocalPath: path,
		IconPath:  path,
		Hash:      "",
	}
	index, err := helm.LoadIndex()
	if err != nil {
		return nil, err
	}
	for k, versions := range index.IndexFile.Entries {
		// because versions is sorted in reverse order, the first one will be the latest version
		if len(versions) > 0 {
			newestVersionedChart := versions[0]
			rtn[k] = chart{
				dir:     newestVersionedChart.Dir,
				version: newestVersionedChart.Version}
		}
	}

	return rtn, nil
}

func pickImagesFromValuesYAML(imagesSet map[string]map[string]bool, charts map[string]chart, basePath, path string, info os.FileInfo, osType OSType) error {
	if info.Name() != "values.yaml" {
		return nil
	}
	relPath, err := filepath.Rel(basePath, path)
	if err != nil {
		return err
	}
	var chartNameAndVersion string
	for name, chart := range charts {
		if strings.HasPrefix(relPath, chart.dir) {
			chartNameAndVersion = fmt.Sprintf("%s:%s", name, chart.version)
			break
		}
	}
	if chartNameAndVersion == "" {
		// path does not belong to a given chart
		return nil
	}
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	dataInterface := map[interface{}]interface{}{}
	if err := yaml.Unmarshal(data, &dataInterface); err != nil {
		return err
	}

	walkthroughMap(dataInterface, func(inputMap map[interface{}]interface{}) {
		generateImages(chartNameAndVersion, inputMap, imagesSet, osType)
	})
	return nil
}

func generateImages(chartNameAndVersion string, inputMap map[interface{}]interface{}, output map[string]map[string]bool, osType OSType) {
	r, repoOk := inputMap["repository"]
	t, tagOk := inputMap["tag"]
	if !repoOk || !tagOk {
		return
	}
	repo, repoOk := r.(string)
	if !repoOk {
		return
	}

	// distinguish images by os
	os := inputMap["os"]
	switch os {
	case "windows": // must have indicate `os: windows` if the image is using in Windows cluster
		if osType != Windows {
			return
		}
	default:
		if osType != Linux {
			return
		}
	}

	imageName := fmt.Sprintf("%s:%v", repo, t)
	addSourceToImage(output, imageName, chartNameAndVersion)
}

func addSourceToImage(imagesSet map[string]map[string]bool, image string, sources ...string) {
	if imagesSet[image] == nil {
		imagesSet[image] = make(map[string]bool)
	}
	for _, source := range sources {
		imagesSet[image][source] = true
	}
}

func walkthroughMap(inputMap map[interface{}]interface{}, walkFunc func(map[interface{}]interface{})) {
	walkFunc(inputMap)
	for _, value := range inputMap {
		if v, ok := value.(map[interface{}]interface{}); ok {
			walkthroughMap(v, walkFunc)
		}
	}
}

func GetImages(systemChartPath, chartPath string, k3sUpgradeImages, imagesFromArgs []string, rkeSystemImages map[string]rketypes.RKESystemImages, osType OSType) ([]string, []string, error) {
	// fetch images from system charts
	imagesSet := make(map[string]map[string]bool)
	if systemChartPath != "" {
		if err := fetchImagesFromCharts(systemChartPath, osType, imagesSet); err != nil {
			return nil, nil, errors.Wrap(err, "failed to fetch images from system charts")
		}
	}

	// fetch images from charts
	if chartPath != "" {
		if err := fetchImagesFromCharts(chartPath, osType, imagesSet); err != nil {
			return nil, nil, errors.Wrap(err, "failed to fetch images from charts")
		}
	}

	// fetch images from system images
	if len(rkeSystemImages) > 0 {
		if err := fetchImagesFromSystem(rkeSystemImages, osType, imagesSet); err != nil {
			return nil, nil, errors.Wrap(err, "failed to fetch images from system images")
		}
	}

	setRequirementImages(osType, imagesSet)

	// set rancher images from args
	setImages("rancher", imagesFromArgs, imagesSet)

	// set images for k3s-upgrade
	setImages("k3sUpgrade", k3sUpgradeImages, imagesSet)

	convertMirroredImages(imagesSet)

	imagesList, imagesAndSourcesList := generateImageAndSourceLists(imagesSet)

	return imagesList, imagesAndSourcesList, nil
}

func setImages(source string, imagesFromArgs []string, imagesSet map[string]map[string]bool) {
	for _, image := range imagesFromArgs {
		addSourceToImage(imagesSet, image, source)
	}
}

func convertMirroredImages(imagesSet map[string]map[string]bool) {
	for image := range imagesSet {
		convertedImage := img.Mirror(image)
		if image == convertedImage {
			continue
		}
		for source, val := range imagesSet[image] {
			if !val {
				continue
			}
			addSourceToImage(imagesSet, convertedImage, source)
		}
		delete(imagesSet, image)
	}
}

func fetchImagesFromCharts(path string, osType OSType, imagesSet map[string]map[string]bool) error {
	chartVersion, err := getChartAndVersion(path)
	if err != nil {
		return errors.Wrapf(err, "failed to get chart and version from %q", path)
	}

	err = filepath.Walk(path, func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		return pickImagesFromValuesYAML(imagesSet, chartVersion, path, p, info, osType)
	})
	if err != nil {
		return errors.Wrap(err, "failed to pick images from values.yaml")
	}

	return nil
}

func fetchImagesFromSystem(rkeSystemImages map[string]rketypes.RKESystemImages, osType OSType, imagesSet map[string]map[string]bool) error {
	collectionImagesList := []interface{}{
		rkeSystemImages,
	}
	switch osType {
	case Linux:
		collectionImagesList = append(collectionImagesList, v32.ToolsSystemImages)
	}

	images, err := flatImagesFromCollections(collectionImagesList...)
	if err != nil {
		return err
	}

	for _, image := range images {
		addSourceToImage(imagesSet, image, "system")

	}
	return nil
}

func flatImagesFromCollections(cols ...interface{}) (images []string, err error) {
	for _, col := range cols {
		colObj := map[string]interface{}{}
		if err := convert.ToObj(col, &colObj); err != nil {
			return []string{}, err
		}

		images = append(images, fetchImagesFromCollection(colObj)...)
	}
	return images, nil
}

func fetchImagesFromCollection(obj map[string]interface{}) (images []string) {
	for _, v := range obj {
		switch t := v.(type) {
		case string:
			images = append(images, t)
		case map[string]interface{}:
			images = append(images, fetchImagesFromCollection(t)...)
		}
	}
	return images
}

func setRequirementImages(osType OSType, imagesSet map[string]map[string]bool) {
	coreLabel := "core"
	switch osType {
	case Linux:
		addSourceToImage(imagesSet, settings.ShellImage.Get(), coreLabel)
		addSourceToImage(imagesSet, "busybox", coreLabel)
	}
}

func generateImageAndSourceLists(imagesSet map[string]map[string]bool) ([]string, []string) {
	var images, imagesAndSources []string
	// unique
	for image := range imagesSet {
		images = append(images, image)
	}

	// sort
	sort.Slice(images, func(i, j int) bool {
		return images[i] < images[j]
	})

	for _, image := range images {
		imagesAndSources = append(imagesAndSources, fmt.Sprintf("%s %s", image, getSourcesList(imagesSet[image])))
	}

	return images, imagesAndSources
}

func getSourcesList(imageSources map[string]bool) string {
	var sources []string

	for source, val := range imageSources {
		if !val {
			continue
		}
		sources = append(sources, source)
	}
	sort.Strings(sources)
	return strings.Join(sources, ",")
}
<commit_msg>Allow walkthroughMap to walk through list entries<commit_after>package image

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

	"github.com/pkg/errors"
	"github.com/rancher/norman/types/convert"
	v32 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	libhelm "github.com/rancher/rancher/pkg/catalog/helm"
	util "github.com/rancher/rancher/pkg/cluster"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/settings"
	rketypes "github.com/rancher/rke/types"
	img "github.com/rancher/rke/types/image"
	"gopkg.in/yaml.v2"
)

type OSType int

type chart struct {
	dir     string
	version string
}

const (
	Linux OSType = iota
	Windows
)

func Resolve(image string) string {
	return ResolveWithCluster(image, nil)
}

func ResolveWithCluster(image string, cluster *v3.Cluster) string {
	reg := util.GetPrivateRepoURL(cluster)
	if reg != "" && !strings.HasPrefix(image, reg) {
		//Images from Dockerhub Library repo, we add rancher prefix when using private registry
		if !strings.Contains(image, "/") {
			image = "rancher/" + image
		}
		return path.Join(reg, image)
	}

	return image
}

func getChartAndVersion(path string) (map[string]chart, error) {
	rtn := map[string]chart{}
	helm := libhelm.Helm{
		LocalPath: path,
		IconPath:  path,
		Hash:      "",
	}
	index, err := helm.LoadIndex()
	if err != nil {
		return nil, err
	}
	for k, versions := range index.IndexFile.Entries {
		// because versions is sorted in reverse order, the first one will be the latest version
		if len(versions) > 0 {
			newestVersionedChart := versions[0]
			rtn[k] = chart{
				dir:     newestVersionedChart.Dir,
				version: newestVersionedChart.Version}
		}
	}

	return rtn, nil
}

func pickImagesFromValuesYAML(imagesSet map[string]map[string]bool, charts map[string]chart, basePath, path string, info os.FileInfo, osType OSType) error {
	if info.Name() != "values.yaml" {
		return nil
	}
	relPath, err := filepath.Rel(basePath, path)
	if err != nil {
		return err
	}
	var chartNameAndVersion string
	for name, chart := range charts {
		if strings.HasPrefix(relPath, chart.dir) {
			chartNameAndVersion = fmt.Sprintf("%s:%s", name, chart.version)
			break
		}
	}
	if chartNameAndVersion == "" {
		// path does not belong to a given chart
		return nil
	}
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	dataInterface := map[interface{}]interface{}{}
	if err := yaml.Unmarshal(data, &dataInterface); err != nil {
		return err
	}

	walkthroughMap(dataInterface, func(inputMap map[interface{}]interface{}) {
		generateImages(chartNameAndVersion, inputMap, imagesSet, osType)
	})
	return nil
}

func generateImages(chartNameAndVersion string, inputMap map[interface{}]interface{}, output map[string]map[string]bool, osType OSType) {
	r, repoOk := inputMap["repository"]
	t, tagOk := inputMap["tag"]
	if !repoOk || !tagOk {
		return
	}
	repo, repoOk := r.(string)
	if !repoOk {
		return
	}

	// distinguish images by os
	os := inputMap["os"]
	switch os {
	case "windows": // must have indicate `os: windows` if the image is using in Windows cluster
		if osType != Windows {
			return
		}
	default:
		if osType != Linux {
			return
		}
	}

	imageName := fmt.Sprintf("%s:%v", repo, t)
	addSourceToImage(output, imageName, chartNameAndVersion)
}

func addSourceToImage(imagesSet map[string]map[string]bool, image string, sources ...string) {
	if imagesSet[image] == nil {
		imagesSet[image] = make(map[string]bool)
	}
	for _, source := range sources {
		imagesSet[image][source] = true
	}
}

func walkthroughMap(data interface{}, walkFunc func(map[interface{}]interface{})) {
	inputMap, isMap := data.(map[interface{}]interface{})
	if isMap {
		// Run the walkFunc on the root node and each child node
		walkFunc(inputMap)
		for _, value := range inputMap {
			walkthroughMap(value, walkFunc)
		}
		return
	}
	// Check if data is a list
	inputList, isList := data.([]interface{})
	if isList {
		// Run the walkFunc on each element in the root node, ignoring the root itself
		for _, elem := range inputList {
			walkthroughMap(elem, walkFunc)
		}
		return
	}
	// Element is neither a map nor a list
	return
}

func GetImages(systemChartPath, chartPath string, k3sUpgradeImages, imagesFromArgs []string, rkeSystemImages map[string]rketypes.RKESystemImages, osType OSType) ([]string, []string, error) {
	// fetch images from system charts
	imagesSet := make(map[string]map[string]bool)
	if systemChartPath != "" {
		if err := fetchImagesFromCharts(systemChartPath, osType, imagesSet); err != nil {
			return nil, nil, errors.Wrap(err, "failed to fetch images from system charts")
		}
	}

	// fetch images from charts
	if chartPath != "" {
		if err := fetchImagesFromCharts(chartPath, osType, imagesSet); err != nil {
			return nil, nil, errors.Wrap(err, "failed to fetch images from charts")
		}
	}

	// fetch images from system images
	if len(rkeSystemImages) > 0 {
		if err := fetchImagesFromSystem(rkeSystemImages, osType, imagesSet); err != nil {
			return nil, nil, errors.Wrap(err, "failed to fetch images from system images")
		}
	}

	setRequirementImages(osType, imagesSet)

	// set rancher images from args
	setImages("rancher", imagesFromArgs, imagesSet)

	// set images for k3s-upgrade
	setImages("k3sUpgrade", k3sUpgradeImages, imagesSet)

	convertMirroredImages(imagesSet)

	imagesList, imagesAndSourcesList := generateImageAndSourceLists(imagesSet)

	return imagesList, imagesAndSourcesList, nil
}

func setImages(source string, imagesFromArgs []string, imagesSet map[string]map[string]bool) {
	for _, image := range imagesFromArgs {
		addSourceToImage(imagesSet, image, source)
	}
}

func convertMirroredImages(imagesSet map[string]map[string]bool) {
	for image := range imagesSet {
		convertedImage := img.Mirror(image)
		if image == convertedImage {
			continue
		}
		for source, val := range imagesSet[image] {
			if !val {
				continue
			}
			addSourceToImage(imagesSet, convertedImage, source)
		}
		delete(imagesSet, image)
	}
}

func fetchImagesFromCharts(path string, osType OSType, imagesSet map[string]map[string]bool) error {
	chartVersion, err := getChartAndVersion(path)
	if err != nil {
		return errors.Wrapf(err, "failed to get chart and version from %q", path)
	}

	err = filepath.Walk(path, func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		return pickImagesFromValuesYAML(imagesSet, chartVersion, path, p, info, osType)
	})
	if err != nil {
		return errors.Wrap(err, "failed to pick images from values.yaml")
	}

	return nil
}

func fetchImagesFromSystem(rkeSystemImages map[string]rketypes.RKESystemImages, osType OSType, imagesSet map[string]map[string]bool) error {
	collectionImagesList := []interface{}{
		rkeSystemImages,
	}
	switch osType {
	case Linux:
		collectionImagesList = append(collectionImagesList, v32.ToolsSystemImages)
	}

	images, err := flatImagesFromCollections(collectionImagesList...)
	if err != nil {
		return err
	}

	for _, image := range images {
		addSourceToImage(imagesSet, image, "system")

	}
	return nil
}

func flatImagesFromCollections(cols ...interface{}) (images []string, err error) {
	for _, col := range cols {
		colObj := map[string]interface{}{}
		if err := convert.ToObj(col, &colObj); err != nil {
			return []string{}, err
		}

		images = append(images, fetchImagesFromCollection(colObj)...)
	}
	return images, nil
}

func fetchImagesFromCollection(obj map[string]interface{}) (images []string) {
	for _, v := range obj {
		switch t := v.(type) {
		case string:
			images = append(images, t)
		case map[string]interface{}:
			images = append(images, fetchImagesFromCollection(t)...)
		}
	}
	return images
}

func setRequirementImages(osType OSType, imagesSet map[string]map[string]bool) {
	coreLabel := "core"
	switch osType {
	case Linux:
		addSourceToImage(imagesSet, settings.ShellImage.Get(), coreLabel)
		addSourceToImage(imagesSet, "busybox", coreLabel)
	}
}

func generateImageAndSourceLists(imagesSet map[string]map[string]bool) ([]string, []string) {
	var images, imagesAndSources []string
	// unique
	for image := range imagesSet {
		images = append(images, image)
	}

	// sort
	sort.Slice(images, func(i, j int) bool {
		return images[i] < images[j]
	})

	for _, image := range images {
		imagesAndSources = append(imagesAndSources, fmt.Sprintf("%s %s", image, getSourcesList(imagesSet[image])))
	}

	return images, imagesAndSources
}

func getSourcesList(imageSources map[string]bool) string {
	var sources []string

	for source, val := range imageSources {
		if !val {
			continue
		}
		sources = append(sources, source)
	}
	sort.Strings(sources)
	return strings.Join(sources, ",")
}
<|endoftext|>


<commit_before>package ledger

import (
	"bufio"
	"fmt"
	"io"
	"math/big"
	"sort"
	"strings"
	"time"
)

var TransactionDateFormat string
var DisplayPrecision int

func init() {
	TransactionDateFormat = "2006/01/02"
	DisplayPrecision = 2
}

type Account struct {
	Name    string
	Balance *big.Rat
}

type sortAccounts []*Account

func (s sortAccounts) Len() int      { return len(s) }
func (s sortAccounts) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type sortAccountsByName struct{ sortAccounts }

func (s sortAccountsByName) Less(i, j int) bool {
	return s.sortAccounts[i].Name < s.sortAccounts[j].Name
}

type Transaction struct {
	Payee          string
	Date           time.Time
	AccountChanges []Account
}

type sortTransactions []*Transaction

func (s sortTransactions) Len() int      { return len(s) }
func (s sortTransactions) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type sortTransactionsByDate struct{ sortTransactions }

func (s sortTransactionsByDate) Less(i, j int) bool {
	return s.sortTransactions[i].Date.Before(s.sortTransactions[j].Date)
}

// Parses a ledger file and returns a list of Transactions.
//
// Transactions are sorted by date.
func ParseLedger(ledgerReader io.Reader) (generalLedger []*Transaction, err error) {
	var trans *Transaction
	scanner := bufio.NewScanner(ledgerReader)
	var line string
	var lineCount int
	for scanner.Scan() {
		line = scanner.Text()
		lineCount++
		if strings.HasPrefix(line, ";") {
			// nop
		} else if len(line) == 0 {
			if trans != nil {
				transErr := balanceTransaction(trans)
				if transErr != nil {
					return generalLedger, fmt.Errorf("%d: Unable to balance transaction, %s", lineCount, transErr)
				}
				generalLedger = append(generalLedger, trans)
				trans = nil
			}
		} else if trans == nil {
			lineSplit := strings.SplitN(line, " ", 2)
			if len(lineSplit) != 2 {
				return generalLedger, fmt.Errorf("%d: Unable to parse payee line: %s", lineCount, line)
			}
			dateString := lineSplit[0]
			transDate, dateErr := time.Parse(TransactionDateFormat, dateString)
			if dateErr != nil {
				return generalLedger, fmt.Errorf("%d: Unable to parse date: %s", lineCount, dateString)
			}
			payeeString := lineSplit[1]
			trans = &Transaction{Payee: payeeString, Date: transDate}
		} else {
			var accChange Account
			lineSplit := strings.Split(line, " ")
			nonEmptyWords := []string{}
			for _, word := range lineSplit {
				if len(word) > 0 {
					nonEmptyWords = append(nonEmptyWords, word)
				}
			}
			lastIndex := len(nonEmptyWords) - 1
			rationalNum := new(big.Rat)
			_, balErr := rationalNum.SetString(nonEmptyWords[lastIndex])
			if balErr == false {
				// Assuming no balance and whole line is account name
				accChange.Name = strings.Join(nonEmptyWords, " ")
			} else {
				accChange.Name = strings.Join(nonEmptyWords[:lastIndex], " ")
				accChange.Balance = rationalNum
			}
			trans.AccountChanges = append(trans.AccountChanges, accChange)
		}
	}
	sort.Sort(sortTransactionsByDate{generalLedger})
	return generalLedger, scanner.Err()
}

// Given a list of transactions and filter strings, returns account balances of
// all accounts that have any filter as a substring of the account name. Also
// returns balances for each account level depth as a separate record.
//
// Accounts are sorted by name.
func GetBalances(generalLedger []*Transaction, filterArr []string) []*Account {
	balances := make(map[string]*big.Rat)
	for _, trans := range generalLedger {
		for _, accChange := range trans.AccountChanges {
			inFilter := len(filterArr) == 0
			for _, filter := range filterArr {
				if strings.Contains(accChange.Name, filter) {
					inFilter = true
				}
			}
			if inFilter {
				accHier := strings.Split(accChange.Name, ":")
				accDepth := len(accHier)
				for currDepth := accDepth; currDepth > 0; currDepth-- {
					currAccName := strings.Join(accHier[:currDepth], ":")
					if ratNum, ok := balances[currAccName]; !ok {
						ratNum = new(big.Rat)
						ratNum.SetString(accChange.Balance.RatString())
						balances[currAccName] = ratNum
					} else {
						ratNum.Add(ratNum, accChange.Balance)
					}
				}
			}
		}
	}

	accList := make([]*Account, len(balances))
	count := 0
	for accName, accBalance := range balances {
		account := &Account{Name: accName, Balance: accBalance}
		accList[count] = account
		count++
	}

	sort.Sort(sortAccountsByName{accList})
	return accList
}

// Prints out account balances formated to a windows of a width of columns.
// Only shows accounts with names less than or equal to the given depth.
func PrintBalances(accountList []*Account, printZeroBalances bool, depth, columns int) {
	overallBalance := new(big.Rat)
	for _, account := range accountList {
		accDepth := len(strings.Split(account.Name, ":"))
		if accDepth == 1 {
			overallBalance.Add(overallBalance, account.Balance)
		}
		if (printZeroBalances || account.Balance.Sign() != 0) && (depth < 0 || accDepth <= depth) {
			outBalanceString := account.Balance.FloatString(DisplayPrecision)
			spaceCount := columns - len(account.Name) - len(outBalanceString)
			fmt.Printf("%s%s%s\n", account.Name, strings.Repeat(" ", spaceCount), outBalanceString)
		}
	}
	fmt.Println(strings.Repeat("-", columns))
	outBalanceString := overallBalance.FloatString(DisplayPrecision)
	spaceCount := columns - len(outBalanceString)
	fmt.Printf("%s%s\n", strings.Repeat(" ", spaceCount), outBalanceString)
}

// Prints a transaction formatted to fit in specified column width.
func PrintTransaction(w io.Writer, trans *Transaction, columns int) {
	fmt.Fprintf(w, "%s %s\n", trans.Date.Format(TransactionDateFormat), trans.Payee)
	for _, accChange := range trans.AccountChanges {
		outBalanceString := accChange.Balance.FloatString(DisplayPrecision)
		spaceCount := columns - 4 - len(accChange.Name) - len(outBalanceString)
		fmt.Fprintf(w, "    %s%s%s\n", accChange.Name, strings.Repeat(" ", spaceCount), outBalanceString)
	}
	fmt.Fprintln(w, "")
}

// Prints all transactions as a formatted ledger file.
func PrintLedger(w io.Writer, generalLedger []*Transaction, columns int) {
	for _, trans := range generalLedger {
		PrintTransaction(w, trans, columns)
	}
}

// Prints each transaction that matches the given filters.
func PrintRegister(generalLedger []*Transaction, filterArr []string, columns int) {
	runningBalance := new(big.Rat)
	for _, trans := range generalLedger {
		for _, accChange := range trans.AccountChanges {
			inFilter := len(filterArr) == 0
			for _, filter := range filterArr {
				if strings.Contains(accChange.Name, filter) {
					inFilter = true
				}
			}
			if inFilter {
				runningBalance.Add(runningBalance, accChange.Balance)
				writtenBytes, _ := fmt.Printf("%s %s", trans.Date.Format(TransactionDateFormat), trans.Payee)
				outBalanceString := accChange.Balance.FloatString(DisplayPrecision)
				outRunningBalanceString := runningBalance.FloatString(DisplayPrecision)
				spaceCount := columns - writtenBytes - 2 - len(outBalanceString) - len(outRunningBalanceString)
				if spaceCount < 0 {
					spaceCount = 0
				}
				fmt.Printf("%s%s %s", strings.Repeat(" ", spaceCount), outBalanceString, outRunningBalanceString)
				fmt.Println("")
			}
		}
	}
}

// Takes a transaction and balances it. This is mainly to fill in the empty part
// with the remaining balance.
func balanceTransaction(input *Transaction) error {
	balance := new(big.Rat)
	var emptyAccPtr *Account
	var emptyAccIndex int
	for accIndex, accChange := range input.AccountChanges {
		if accChange.Balance == nil {
			if emptyAccPtr != nil {
				return fmt.Errorf("More than one account change empty!")
			}
			emptyAccPtr = &accChange
			emptyAccIndex = accIndex
		} else {
			balance = balance.Add(balance, accChange.Balance)
		}
	}
	if balance.Sign() != 0 {
		if emptyAccPtr == nil {
			return fmt.Errorf("No empty account change to place extra balance!")
		}
		input.AccountChanges[emptyAccIndex].Balance = balance.Neg(balance)
	}
	return nil
}
<commit_msg>Parse last transaction of the ledger file.<commit_after>package ledger

import (
	"bufio"
	"fmt"
	"io"
	"math/big"
	"sort"
	"strings"
	"time"
)

var TransactionDateFormat string
var DisplayPrecision int

func init() {
	TransactionDateFormat = "2006/01/02"
	DisplayPrecision = 2
}

type Account struct {
	Name    string
	Balance *big.Rat
}

type sortAccounts []*Account

func (s sortAccounts) Len() int      { return len(s) }
func (s sortAccounts) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type sortAccountsByName struct{ sortAccounts }

func (s sortAccountsByName) Less(i, j int) bool {
	return s.sortAccounts[i].Name < s.sortAccounts[j].Name
}

type Transaction struct {
	Payee          string
	Date           time.Time
	AccountChanges []Account
}

type sortTransactions []*Transaction

func (s sortTransactions) Len() int      { return len(s) }
func (s sortTransactions) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type sortTransactionsByDate struct{ sortTransactions }

func (s sortTransactionsByDate) Less(i, j int) bool {
	return s.sortTransactions[i].Date.Before(s.sortTransactions[j].Date)
}

// Parses a ledger file and returns a list of Transactions.
//
// Transactions are sorted by date.
func ParseLedger(ledgerReader io.Reader) (generalLedger []*Transaction, err error) {
	var trans *Transaction
	scanner := bufio.NewScanner(ledgerReader)
	var line string
	var lineCount int
	for scanner.Scan() {
		line = scanner.Text()
		lineCount++
		if strings.HasPrefix(line, ";") {
			// nop
		} else if len(line) == 0 {
			if trans != nil {
				transErr := balanceTransaction(trans)
				if transErr != nil {
					return generalLedger, fmt.Errorf("%d: Unable to balance transaction, %s", lineCount, transErr)
				}
				generalLedger = append(generalLedger, trans)
				trans = nil
			}
		} else if trans == nil {
			lineSplit := strings.SplitN(line, " ", 2)
			if len(lineSplit) != 2 {
				return generalLedger, fmt.Errorf("%d: Unable to parse payee line: %s", lineCount, line)
			}
			dateString := lineSplit[0]
			transDate, dateErr := time.Parse(TransactionDateFormat, dateString)
			if dateErr != nil {
				return generalLedger, fmt.Errorf("%d: Unable to parse date: %s", lineCount, dateString)
			}
			payeeString := lineSplit[1]
			trans = &Transaction{Payee: payeeString, Date: transDate}
		} else {
			var accChange Account
			lineSplit := strings.Split(line, " ")
			nonEmptyWords := []string{}
			for _, word := range lineSplit {
				if len(word) > 0 {
					nonEmptyWords = append(nonEmptyWords, word)
				}
			}
			lastIndex := len(nonEmptyWords) - 1
			rationalNum := new(big.Rat)
			_, balErr := rationalNum.SetString(nonEmptyWords[lastIndex])
			if balErr == false {
				// Assuming no balance and whole line is account name
				accChange.Name = strings.Join(nonEmptyWords, " ")
			} else {
				accChange.Name = strings.Join(nonEmptyWords[:lastIndex], " ")
				accChange.Balance = rationalNum
			}
			trans.AccountChanges = append(trans.AccountChanges, accChange)
		}
	}
	// If the file does not end on empty line, we must attempt to balance last
	// transaction of the file.
	if trans != nil {
		transErr := balanceTransaction(trans)
		if transErr != nil {
			return generalLedger, fmt.Errorf("%d: Unable to balance transaction, %s", lineCount, transErr)
		}
		generalLedger = append(generalLedger, trans)
		trans = nil
	}
	sort.Sort(sortTransactionsByDate{generalLedger})
	return generalLedger, scanner.Err()
}

// Given a list of transactions and filter strings, returns account balances of
// all accounts that have any filter as a substring of the account name. Also
// returns balances for each account level depth as a separate record.
//
// Accounts are sorted by name.
func GetBalances(generalLedger []*Transaction, filterArr []string) []*Account {
	balances := make(map[string]*big.Rat)
	for _, trans := range generalLedger {
		for _, accChange := range trans.AccountChanges {
			inFilter := len(filterArr) == 0
			for _, filter := range filterArr {
				if strings.Contains(accChange.Name, filter) {
					inFilter = true
				}
			}
			if inFilter {
				accHier := strings.Split(accChange.Name, ":")
				accDepth := len(accHier)
				for currDepth := accDepth; currDepth > 0; currDepth-- {
					currAccName := strings.Join(accHier[:currDepth], ":")
					if ratNum, ok := balances[currAccName]; !ok {
						ratNum = new(big.Rat)
						ratNum.SetString(accChange.Balance.RatString())
						balances[currAccName] = ratNum
					} else {
						ratNum.Add(ratNum, accChange.Balance)
					}
				}
			}
		}
	}

	accList := make([]*Account, len(balances))
	count := 0
	for accName, accBalance := range balances {
		account := &Account{Name: accName, Balance: accBalance}
		accList[count] = account
		count++
	}

	sort.Sort(sortAccountsByName{accList})
	return accList
}

// Prints out account balances formated to a windows of a width of columns.
// Only shows accounts with names less than or equal to the given depth.
func PrintBalances(accountList []*Account, printZeroBalances bool, depth, columns int) {
	overallBalance := new(big.Rat)
	for _, account := range accountList {
		accDepth := len(strings.Split(account.Name, ":"))
		if accDepth == 1 {
			overallBalance.Add(overallBalance, account.Balance)
		}
		if (printZeroBalances || account.Balance.Sign() != 0) && (depth < 0 || accDepth <= depth) {
			outBalanceString := account.Balance.FloatString(DisplayPrecision)
			spaceCount := columns - len(account.Name) - len(outBalanceString)
			fmt.Printf("%s%s%s\n", account.Name, strings.Repeat(" ", spaceCount), outBalanceString)
		}
	}
	fmt.Println(strings.Repeat("-", columns))
	outBalanceString := overallBalance.FloatString(DisplayPrecision)
	spaceCount := columns - len(outBalanceString)
	fmt.Printf("%s%s\n", strings.Repeat(" ", spaceCount), outBalanceString)
}

// Prints a transaction formatted to fit in specified column width.
func PrintTransaction(w io.Writer, trans *Transaction, columns int) {
	fmt.Fprintf(w, "%s %s\n", trans.Date.Format(TransactionDateFormat), trans.Payee)
	for _, accChange := range trans.AccountChanges {
		outBalanceString := accChange.Balance.FloatString(DisplayPrecision)
		spaceCount := columns - 4 - len(accChange.Name) - len(outBalanceString)
		fmt.Fprintf(w, "    %s%s%s\n", accChange.Name, strings.Repeat(" ", spaceCount), outBalanceString)
	}
	fmt.Fprintln(w, "")
}

// Prints all transactions as a formatted ledger file.
func PrintLedger(w io.Writer, generalLedger []*Transaction, columns int) {
	for _, trans := range generalLedger {
		PrintTransaction(w, trans, columns)
	}
}

// Prints each transaction that matches the given filters.
func PrintRegister(generalLedger []*Transaction, filterArr []string, columns int) {
	runningBalance := new(big.Rat)
	for _, trans := range generalLedger {
		for _, accChange := range trans.AccountChanges {
			inFilter := len(filterArr) == 0
			for _, filter := range filterArr {
				if strings.Contains(accChange.Name, filter) {
					inFilter = true
				}
			}
			if inFilter {
				runningBalance.Add(runningBalance, accChange.Balance)
				writtenBytes, _ := fmt.Printf("%s %s", trans.Date.Format(TransactionDateFormat), trans.Payee)
				outBalanceString := accChange.Balance.FloatString(DisplayPrecision)
				outRunningBalanceString := runningBalance.FloatString(DisplayPrecision)
				spaceCount := columns - writtenBytes - 2 - len(outBalanceString) - len(outRunningBalanceString)
				if spaceCount < 0 {
					spaceCount = 0
				}
				fmt.Printf("%s%s %s", strings.Repeat(" ", spaceCount), outBalanceString, outRunningBalanceString)
				fmt.Println("")
			}
		}
	}
}

// Takes a transaction and balances it. This is mainly to fill in the empty part
// with the remaining balance.
func balanceTransaction(input *Transaction) error {
	balance := new(big.Rat)
	var emptyAccPtr *Account
	var emptyAccIndex int
	for accIndex, accChange := range input.AccountChanges {
		if accChange.Balance == nil {
			if emptyAccPtr != nil {
				return fmt.Errorf("More than one account change empty!")
			}
			emptyAccPtr = &accChange
			emptyAccIndex = accIndex
		} else {
			balance = balance.Add(balance, accChange.Balance)
		}
	}
	if balance.Sign() != 0 {
		if emptyAccPtr == nil {
			return fmt.Errorf("No empty account change to place extra balance!")
		}
		input.AccountChanges[emptyAccIndex].Balance = balance.Neg(balance)
	}
	return nil
}
<|endoftext|>


<commit_before>package router

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/matt-deboer/mpp/pkg/locator"
	"github.com/matt-deboer/mpp/pkg/selector"
	"github.com/matt-deboer/mpp/pkg/version"
	"github.com/vulcand/oxy/buffer"
	"github.com/vulcand/oxy/forward"
)

// Router provides dynamic routing of http requests based on a configurable strategy
type Router struct {
	locators        []locator.Locator
	selector        *selector.Selector
	selection       *selector.Result
	forward         http.Handler
	buffer          *buffer.Buffer
	rewriter        urlRewriter
	affinityOptions []AffinityOption
	interval        time.Duration
	metrics         *metrics
	// used to mark control of the selection process
	theConch            chan struct{}
	selectionInProgress sync.RWMutex
}

// Status contains a snapshot status summary of the router state
type Status struct {
	Endpoints           []*locator.PrometheusEndpoint
	Strategy            string
	StrategyDescription string
	AffinityOptions     string
	ComparisonMetric    string
	Interval            time.Duration
}

type urlRewriter func(u *url.URL)

var noOpRewriter = func(u *url.URL) {}

// NewRouter constructs a new router based on the provided stategy and locators
func NewRouter(interval time.Duration, affinityOptions []AffinityOption,
	locators []locator.Locator, strategyArgs ...string) (*Router, error) {

	sel, err := selector.NewSelector(locators, strategyArgs...)
	if err != nil {
		return nil, err
	}

	r := &Router{
		locators:        locators,
		selector:        sel,
		affinityOptions: affinityOptions,
		interval:        interval,
		rewriter:        noOpRewriter,
		metrics:         newMetrics(version.Name),
		selection:       &selector.Result{},
		theConch:        make(chan struct{}, 1),
	}

	// Set up the lock
	r.theConch <- struct{}{}
	r.doSelection()
	go func() {
		// TODO: create shutdown channel for this
		for {
			if log.GetLevel() >= log.DebugLevel {
				log.Debugf("Backend selection is sleeping for %s", interval)
			}
			time.Sleep(r.interval)
			r.doSelection()
		}
	}()

	r.forward, _ = forward.New()
	r.buffer, _ = buffer.New(&internalRouter{
		router:   r,
		affinity: newAffinityProvider(affinityOptions),
	},
		buffer.Retry(`IsNetworkError() && Attempts() < 2`))
	return r, nil
}

func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	r.buffer.ServeHTTP(w, retryableRequest(req))
}

func (r *Router) doSelection() {
	select {
	case _ = <-r.theConch:
		r.selectionInProgress.Lock()
		defer r.selectionInProgress.Unlock()
		if log.GetLevel() >= log.DebugLevel {
			log.Debugf("Got selection lock; performing selection")
		}

		result, err := r.selector.Select()

		if len(result.Selection) == 0 {
			if err != nil {
				log.Errorf("Selector returned no valid selection, and error: %v", err)
				if r.selection == nil {
					r.selection = result
					r.rewriter = noOpRewriter
				}
			} else {
				r.selection = result
				r.rewriter = noOpRewriter
				log.Warnf("Selector returned no valid selection")
			}
		} else {
			if log.GetLevel() >= log.DebugLevel {
				log.Debugf("Selected targets: %v", result.Selection)
			}
			if r.selection == nil || !equal(r.selection.Selection, result.Selection) {
				log.Infof("New targets differ from current selection %v; updating rewriter => %v", r.selection, result)
				r.rewriter = func(u *url.URL) {
					selection := result.Selection
					i := r.selector.Strategy.NextIndex(selection)
					target := selection[i]
					u.Host = target.Host
					u.Scheme = target.Scheme
				}
			} else if log.GetLevel() >= log.DebugLevel {
				log.Debugf("Selection is unchanged: %v", r.selection)
			}
			r.selection = result
		}

		r.metrics.selectedBackends.Set(float64(len(result.Selection)))
		r.metrics.selectionEvents.Inc()
		if log.GetLevel() >= log.DebugLevel {
			log.Debugf("Returning selection lock")
		}
		r.theConch <- struct{}{}
	default:
		if log.GetLevel() >= log.DebugLevel {
			log.Debugf("Selection is already in-progress; awaiting result")
		}
		r.selectionInProgress.RLock()
		defer r.selectionInProgress.RUnlock()
	}
}

func equal(a, b []*url.URL) bool {
	for i, v := range a {
		if *v != *b[i] {
			return false
		}
	}
	return len(a) == len(b)
}

func contains(a []*url.URL, u *url.URL) bool {
	for _, v := range a {
		if *u == *v {
			return true
		}
	}
	return false
}

func backend(u *url.URL) string {
	return fmt.Sprintf("%s://%s", u.Scheme, u.Host)
}

// Status returns a summary of the router's current state
func (r *Router) Status() *Status {
	return &Status{
		Endpoints:           r.selection.Candidates,
		Strategy:            r.selector.Strategy.Name(),
		StrategyDescription: r.selector.Strategy.Description(),
		ComparisonMetric:    r.selector.Strategy.ComparisonMetricName(),
		AffinityOptions:     strings.Trim(fmt.Sprintf("%v", r.affinityOptions), "[]"),
		Interval:            r.interval,
	}
}
<commit_msg>increased logging level for some selection based scenarios<commit_after>package router

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/matt-deboer/mpp/pkg/locator"
	"github.com/matt-deboer/mpp/pkg/selector"
	"github.com/matt-deboer/mpp/pkg/version"
	"github.com/vulcand/oxy/buffer"
	"github.com/vulcand/oxy/forward"
)

// Router provides dynamic routing of http requests based on a configurable strategy
type Router struct {
	locators        []locator.Locator
	selector        *selector.Selector
	selection       *selector.Result
	forward         http.Handler
	buffer          *buffer.Buffer
	rewriter        urlRewriter
	affinityOptions []AffinityOption
	interval        time.Duration
	metrics         *metrics
	// used to mark control of the selection process
	theConch            chan struct{}
	selectionInProgress sync.RWMutex
}

// Status contains a snapshot status summary of the router state
type Status struct {
	Endpoints           []*locator.PrometheusEndpoint
	Strategy            string
	StrategyDescription string
	AffinityOptions     string
	ComparisonMetric    string
	Interval            time.Duration
}

type urlRewriter func(u *url.URL)

var noOpRewriter = func(u *url.URL) {}

// NewRouter constructs a new router based on the provided stategy and locators
func NewRouter(interval time.Duration, affinityOptions []AffinityOption,
	locators []locator.Locator, strategyArgs ...string) (*Router, error) {

	sel, err := selector.NewSelector(locators, strategyArgs...)
	if err != nil {
		return nil, err
	}

	r := &Router{
		locators:        locators,
		selector:        sel,
		affinityOptions: affinityOptions,
		interval:        interval,
		rewriter:        noOpRewriter,
		metrics:         newMetrics(version.Name),
		selection:       &selector.Result{},
		theConch:        make(chan struct{}, 1),
	}

	// Set up the lock
	r.theConch <- struct{}{}
	r.doSelection()
	go func() {
		// TODO: create shutdown channel for this
		for {
			if log.GetLevel() >= log.DebugLevel {
				log.Debugf("Backend selection is sleeping for %s", interval)
			}
			time.Sleep(r.interval)
			r.doSelection()
		}
	}()

	r.forward, _ = forward.New()
	r.buffer, _ = buffer.New(&internalRouter{
		router:   r,
		affinity: newAffinityProvider(affinityOptions),
	},
		buffer.Retry(`IsNetworkError() && Attempts() < 2`))
	return r, nil
}

func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	r.buffer.ServeHTTP(w, retryableRequest(req))
}

func (r *Router) doSelection() {
	select {
	case _ = <-r.theConch:
		r.selectionInProgress.Lock()
		defer r.selectionInProgress.Unlock()
		if log.GetLevel() >= log.DebugLevel {
			log.Debugf("Got selection lock; performing selection")
		}

		result, err := r.selector.Select()

		if len(result.Selection) == 0 {
			if err != nil {
				log.Errorf("Selector returned no valid selection, and error: %v", err)
				if r.selection == nil {
					r.selection = result
					r.rewriter = noOpRewriter
				}
			} else {
				r.selection = result
				r.rewriter = noOpRewriter
				log.Warnf("Selector returned no valid selection")
			}
		} else {
			if log.GetLevel() >= log.DebugLevel {
				log.Debugf("Selected targets: %v", result.Selection)
			}
			if r.selection == nil || !equal(r.selection.Selection, result.Selection) {
				log.Infof("New targets differ from current selection %v; updating rewriter => %v", r.selection, result)
				r.rewriter = func(u *url.URL) {
					selection := result.Selection
					i := r.selector.Strategy.NextIndex(selection)
					target := selection[i]
					u.Host = target.Host
					u.Scheme = target.Scheme
				}
			} else {
				log.Infof("Selection is unchanged: %v, out of candidates: %v", r.selection.Selection, r.selection.Candidates)
			}
			r.selection = result
		}

		r.metrics.selectedBackends.Set(float64(len(result.Selection)))
		r.metrics.selectionEvents.Inc()
		if log.GetLevel() >= log.DebugLevel {
			log.Debugf("Returning selection lock")
		}
		r.theConch <- struct{}{}
	default:
		log.Warnf("Selection is already in-progress; awaiting result")
		r.selectionInProgress.RLock()
		defer r.selectionInProgress.RUnlock()
	}
}

func equal(a, b []*url.URL) bool {
	for i, v := range a {
		if *v != *b[i] {
			return false
		}
	}
	return len(a) == len(b)
}

func contains(a []*url.URL, u *url.URL) bool {
	for _, v := range a {
		if *u == *v {
			return true
		}
	}
	return false
}

func backend(u *url.URL) string {
	return fmt.Sprintf("%s://%s", u.Scheme, u.Host)
}

// Status returns a summary of the router's current state
func (r *Router) Status() *Status {
	return &Status{
		Endpoints:           r.selection.Candidates,
		Strategy:            r.selector.Strategy.Name(),
		StrategyDescription: r.selector.Strategy.Description(),
		ComparisonMetric:    r.selector.Strategy.ComparisonMetricName(),
		AffinityOptions:     strings.Trim(fmt.Sprintf("%v", r.affinityOptions), "[]"),
		Interval:            r.interval,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes 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 server

import (
	"os"
	"os/signal"
)

var onlyOneSignalHandler = make(chan struct{})
var shutdownHandler = make(chan os.Signal, 2)

// SetupSignalHandler registered for SIGTERM and SIGINT. A stop channel is returned
// which is closed on one of these signals. If a second signal is caught, the program
// is terminated with exit code 1.
func SetupSignalHandler() <-chan struct{} {
	close(onlyOneSignalHandler) // panics when called twice

	stop := make(chan struct{})
	signal.Notify(shutdownHandler, shutdownSignals...)
	go func() {
		<-shutdownHandler
		close(stop)
		<-shutdownHandler
		os.Exit(1) // second signal. Exit directly.
	}()

	return stop
}

// RequestShutdown emulates a received event that is considered as shutdown signal (SIGTERM/SIGINT)
func RequestShutdown() {
	select {
	case shutdownHandler <- shutdownSignals[0]:
	default:
	}
}
<commit_msg>windows/svc: workaround-exit mechanism that works for signal-less binaries<commit_after>/*
Copyright 2017 The Kubernetes 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 server

import (
	"os"
	"os/signal"
)

var onlyOneSignalHandler = make(chan struct{})
var shutdownHandler chan os.Signal

// SetupSignalHandler registered for SIGTERM and SIGINT. A stop channel is returned
// which is closed on one of these signals. If a second signal is caught, the program
// is terminated with exit code 1.
func SetupSignalHandler() <-chan struct{} {
	close(onlyOneSignalHandler) // panics when called twice

	shutdownHandler = make(chan os.Signal, 2)

	stop := make(chan struct{})
	signal.Notify(shutdownHandler, shutdownSignals...)
	go func() {
		<-shutdownHandler
		close(stop)
		<-shutdownHandler
		os.Exit(1) // second signal. Exit directly.
	}()

	return stop
}

// RequestShutdown emulates a received event that is considered as shutdown signal (SIGTERM/SIGINT)
// This returns whether a handler was notified
func RequestShutdown() bool {
	if shutdownHandler != nil {
		select {
		case shutdownHandler <- shutdownSignals[0]:
			return true
		default:
		}
	}

	return false
}
<|endoftext|>