<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"time"

	gclient "code.cloudfoundry.org/garden/client"
	gconn "code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/lager"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	bclient "github.com/concourse/baggageclaim/client"
	"github.com/concourse/concourse"
	"github.com/concourse/concourse/worker"
	"github.com/concourse/flag"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
)

type WorkerCommand struct {
	Worker WorkerConfig

	TSA worker.TSAConfig `group:"TSA Configuration" namespace:"tsa"`

	Certs Certs

	WorkDir flag.Dir `long:"work-dir" required:"true" description:"Directory in which to place container data."`

	BindIP   flag.IP `long:"bind-ip"   default:"127.0.0.1" description:"IP address on which to listen for the Garden server."`
	BindPort uint16  `long:"bind-port" default:"7777"      description:"Port on which to listen for the Garden server."`

	SweepInterval time.Duration `long:"sweep-interval" default:"30s" description:"Interval on which containers and volumes will be garbage collected from the worker."`

	RebalanceInterval time.Duration `long:"rebalance-interval" description:"Duration after which the registration should be swapped to another random SSH gateway."`

	DrainTimeout time.Duration `long:"drain-timeout" default:"1h" description:"Duration after which a worker should give up draining forwarded connections on shutdown."`

	Garden GardenBackend `group:"Garden Configuration" namespace:"garden"`

	Baggageclaim baggageclaimcmd.BaggageclaimCommand `group:"Baggageclaim Configuration" namespace:"baggageclaim"`

	ResourceTypes flag.Dir `long:"resource-types" description:"Path to directory containing resource types the worker should advertise."`

	Logger flag.Lager
}

func (cmd *WorkerCommand) Execute(args []string) error {
	runner, err := cmd.Runner(args)
	if err != nil {
		return err
	}

	return <-ifrit.Invoke(sigmon.New(runner)).Wait()
}

func (cmd *WorkerCommand) Runner(args []string) (ifrit.Runner, error) {
	if cmd.ResourceTypes == "" {
		cmd.ResourceTypes = flag.Dir(discoverAsset("resource-types"))
	}

	logger, _ := cmd.Logger.Logger("worker")

	atcWorker, gardenRunner, err := cmd.gardenRunner(logger.Session("garden"))
	if err != nil {
		return nil, err
	}

	atcWorker.Version = concourse.WorkerVersion

	baggageclaimRunner, err := cmd.baggageclaimRunner(logger.Session("baggageclaim"))
	if err != nil {
		return nil, err
	}

	members := grouper.Members{
		{
			Name:   "garden",
			Runner: NewLoggingRunner(logger.Session("garden-runner"), gardenRunner),
		},
		{
			Name:   "baggageclaim",
			Runner: NewLoggingRunner(logger.Session("baggageclaim-runner"), baggageclaimRunner),
		},
	}

	if cmd.TSA.WorkerPrivateKey != nil {
		tsaClient := cmd.TSA.Client(atcWorker)
		gardenClient := gclient.New(gconn.NewWithLogger("tcp", cmd.gardenAddr(), logger.Session("garden-connection")))
		baggageclaimClient := bclient.New(cmd.baggageclaimURL(), http.DefaultTransport)

		beacon := &worker.Beacon{
			Logger: logger.Session("beacon"),

			Client: tsaClient,

			RebalanceInterval: cmd.RebalanceInterval,
			DrainTimeout:      cmd.DrainTimeout,

			LocalGardenNetwork: "tcp",
			LocalGardenAddr:    cmd.gardenAddr(),

			LocalBaggageclaimNetwork: "tcp",
			LocalBaggageclaimAddr:    cmd.baggageclaimAddr(),
		}

		members = append(members, grouper.Member{
			Name: "beacon",
			Runner: NewLoggingRunner(
				logger.Session("beacon-runner"),
				worker.NewBeaconRunner(
					logger.Session("beacon-runner"),
					beacon,
					tsaClient,
				),
			),
		})

		members = append(members, grouper.Member{
			Name: "sweeper",
			Runner: NewLoggingRunner(
				logger.Session("sweeper"),
				&worker.SweepRunner{
					Logger:             logger.Session("sweeper-runner"),
					Interval:           cmd.SweepInterval,
					TSAClient:          tsaClient,
					GardenClient:       gardenClient,
					BaggageclaimClient: baggageclaimClient,
				},
			),
		})
	}

	return grouper.NewParallel(os.Interrupt, members), nil
}

func (cmd *WorkerCommand) gardenAddr() string {
	return fmt.Sprintf("%s:%d", cmd.BindIP, cmd.BindPort)
}

func (cmd *WorkerCommand) baggageclaimAddr() string {
	return fmt.Sprintf("%s:%d", cmd.Baggageclaim.BindIP, cmd.Baggageclaim.BindPort)
}

func (cmd *WorkerCommand) baggageclaimURL() string {
	return fmt.Sprintf("http://%s", cmd.baggageclaimAddr())
}

func (cmd *WorkerCommand) workerName() (string, error) {
	if cmd.Worker.Name != "" {
		return cmd.Worker.Name, nil
	}

	return os.Hostname()
}

func (cmd *WorkerCommand) baggageclaimRunner(logger lager.Logger) (ifrit.Runner, error) {
	volumesDir := filepath.Join(cmd.WorkDir.Path(), "volumes")

	err := os.MkdirAll(volumesDir, 0755)
	if err != nil {
		return nil, err
	}

	cmd.Baggageclaim.VolumesDir = flag.Dir(volumesDir)

	cmd.Baggageclaim.OverlaysDir = filepath.Join(cmd.WorkDir.Path(), "overlays")

	return cmd.Baggageclaim.Runner(nil)
}
<commit_msg>bin: prevent sweeper from hanging during shutdown<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"time"

	gclient "code.cloudfoundry.org/garden/client"
	gconn "code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/lager"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	bclient "github.com/concourse/baggageclaim/client"
	"github.com/concourse/concourse"
	"github.com/concourse/concourse/worker"
	"github.com/concourse/flag"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
)

type WorkerCommand struct {
	Worker WorkerConfig

	TSA worker.TSAConfig `group:"TSA Configuration" namespace:"tsa"`

	Certs Certs

	WorkDir flag.Dir `long:"work-dir" required:"true" description:"Directory in which to place container data."`

	BindIP   flag.IP `long:"bind-ip"   default:"127.0.0.1" description:"IP address on which to listen for the Garden server."`
	BindPort uint16  `long:"bind-port" default:"7777"      description:"Port on which to listen for the Garden server."`

	SweepInterval time.Duration `long:"sweep-interval" default:"30s" description:"Interval on which containers and volumes will be garbage collected from the worker."`

	RebalanceInterval time.Duration `long:"rebalance-interval" description:"Duration after which the registration should be swapped to another random SSH gateway."`

	DrainTimeout time.Duration `long:"drain-timeout" default:"1h" description:"Duration after which a worker should give up draining forwarded connections on shutdown."`

	Garden GardenBackend `group:"Garden Configuration" namespace:"garden"`

	Baggageclaim baggageclaimcmd.BaggageclaimCommand `group:"Baggageclaim Configuration" namespace:"baggageclaim"`

	ResourceTypes flag.Dir `long:"resource-types" description:"Path to directory containing resource types the worker should advertise."`

	Logger flag.Lager
}

func (cmd *WorkerCommand) Execute(args []string) error {
	runner, err := cmd.Runner(args)
	if err != nil {
		return err
	}

	return <-ifrit.Invoke(sigmon.New(runner)).Wait()
}

func (cmd *WorkerCommand) Runner(args []string) (ifrit.Runner, error) {
	if cmd.ResourceTypes == "" {
		cmd.ResourceTypes = flag.Dir(discoverAsset("resource-types"))
	}

	logger, _ := cmd.Logger.Logger("worker")

	atcWorker, gardenRunner, err := cmd.gardenRunner(logger.Session("garden"))
	if err != nil {
		return nil, err
	}

	atcWorker.Version = concourse.WorkerVersion

	baggageclaimRunner, err := cmd.baggageclaimRunner(logger.Session("baggageclaim"))
	if err != nil {
		return nil, err
	}

	members := grouper.Members{
		{
			Name:   "garden",
			Runner: NewLoggingRunner(logger.Session("garden-runner"), gardenRunner),
		},
		{
			Name:   "baggageclaim",
			Runner: NewLoggingRunner(logger.Session("baggageclaim-runner"), baggageclaimRunner),
		},
	}

	if cmd.TSA.WorkerPrivateKey != nil {
		tsaClient := cmd.TSA.Client(atcWorker)

		beacon := &worker.Beacon{
			Logger: logger.Session("beacon"),

			Client: tsaClient,

			RebalanceInterval: cmd.RebalanceInterval,
			DrainTimeout:      cmd.DrainTimeout,

			LocalGardenNetwork: "tcp",
			LocalGardenAddr:    cmd.gardenAddr(),

			LocalBaggageclaimNetwork: "tcp",
			LocalBaggageclaimAddr:    cmd.baggageclaimAddr(),
		}

		members = append(members, grouper.Member{
			Name: "beacon",
			Runner: NewLoggingRunner(
				logger.Session("beacon-runner"),
				worker.NewBeaconRunner(
					logger.Session("beacon-runner"),
					beacon,
					tsaClient,
				),
			),
		})

		gardenClient := gclient.New(
			gconn.NewWithLogger(
				"tcp",
				cmd.gardenAddr(),
				logger.Session("garden-connection"),
			),
		)

		baggageclaimClient := bclient.NewWithHTTPClient(
			cmd.baggageclaimURL(),

			// ensure we don't use baggageclaim's default retryhttp client; all
			// traffic should be local, so any failures are unlikely to be transient.
			// we don't want a retry loop to block up sweeping and prevent the worker
			// from existing.
			&http.Client{
				Transport: &http.Transport{
					// don't let a slow (possibly stuck) baggageclaim server slow down
					// sweeping too much
					ResponseHeaderTimeout: 1 * time.Minute,
				},
			},
		)

		members = append(members, grouper.Member{
			Name: "sweeper",
			Runner: NewLoggingRunner(
				logger.Session("sweeper"),
				&worker.SweepRunner{
					Logger:             logger.Session("sweeper-runner"),
					Interval:           cmd.SweepInterval,
					TSAClient:          tsaClient,
					GardenClient:       gardenClient,
					BaggageclaimClient: baggageclaimClient,
				},
			),
		})
	}

	return grouper.NewParallel(os.Interrupt, members), nil
}

func (cmd *WorkerCommand) gardenAddr() string {
	return fmt.Sprintf("%s:%d", cmd.BindIP, cmd.BindPort)
}

func (cmd *WorkerCommand) baggageclaimAddr() string {
	return fmt.Sprintf("%s:%d", cmd.Baggageclaim.BindIP, cmd.Baggageclaim.BindPort)
}

func (cmd *WorkerCommand) baggageclaimURL() string {
	return fmt.Sprintf("http://%s", cmd.baggageclaimAddr())
}

func (cmd *WorkerCommand) workerName() (string, error) {
	if cmd.Worker.Name != "" {
		return cmd.Worker.Name, nil
	}

	return os.Hostname()
}

func (cmd *WorkerCommand) baggageclaimRunner(logger lager.Logger) (ifrit.Runner, error) {
	volumesDir := filepath.Join(cmd.WorkDir.Path(), "volumes")

	err := os.MkdirAll(volumesDir, 0755)
	if err != nil {
		return nil, err
	}

	cmd.Baggageclaim.VolumesDir = flag.Dir(volumesDir)

	cmd.Baggageclaim.OverlaysDir = filepath.Join(cmd.WorkDir.Path(), "overlays")

	return cmd.Baggageclaim.Runner(nil)
}
<|endoftext|>


<commit_before><commit_msg>added explicit casting the interface to int64 + convert int64 to int<commit_after><|endoftext|>


<commit_before>package postgresql

import (
	"database/sql"

	"github.com/DMarby/picsum-photos/database"

	// Import the postgresql driver
	_ "github.com/jackc/pgx/stdlib"
	"github.com/jmoiron/sqlx"
)

// Provider implements a file-based image storage
type Provider struct {
	db *sqlx.DB
}

// New returns a new Provider instance
func New(address string) (*Provider, error) {
	db, err := sqlx.Connect("pgx", address)
	if err != nil {
		return nil, err
	}

	// Use Unsafe so that the app doesn't fail if we add new columns to the database
	return &Provider{
		db: db.Unsafe(),
	}, nil
}

// Get returns the image data for an image id
func (p *Provider) Get(id string) (*database.Image, error) {
	i := &database.Image{}
	err := p.db.Get(i, "select * from image where id = $1", id)

	if err != nil {
		if err == sql.ErrNoRows {
			return nil, database.ErrNotFound
		}

		return nil, err
	}

	return i, nil
}

// GetRandom returns a random image ID
func (p *Provider) GetRandom() (id string, err error) {
	// This will be slow on large tables
	err = p.db.Get(&id, "select id from image order by random() limit 1")
	return
}

// List returns a list of all the images
func (p *Provider) List() ([]database.Image, error) {
	i := []database.Image{}
	err := p.db.Select(&i, "select * from image")

	if err != nil {
		return nil, err
	}

	return i, nil
}

// Shutdown shuts down the database client
func (p *Provider) Shutdown() {
	p.db.Close()
}
<commit_msg>Correct comment in postgresql database<commit_after>package postgresql

import (
	"database/sql"

	"github.com/DMarby/picsum-photos/database"

	// Import the postgresql driver
	_ "github.com/jackc/pgx/stdlib"
	"github.com/jmoiron/sqlx"
)

// Provider implements a postgresql based storage
type Provider struct {
	db *sqlx.DB
}

// New returns a new Provider instance
func New(address string) (*Provider, error) {
	db, err := sqlx.Connect("pgx", address)
	if err != nil {
		return nil, err
	}

	// Use Unsafe so that the app doesn't fail if we add new columns to the database
	return &Provider{
		db: db.Unsafe(),
	}, nil
}

// Get returns the image data for an image id
func (p *Provider) Get(id string) (*database.Image, error) {
	i := &database.Image{}
	err := p.db.Get(i, "select * from image where id = $1", id)

	if err != nil {
		if err == sql.ErrNoRows {
			return nil, database.ErrNotFound
		}

		return nil, err
	}

	return i, nil
}

// GetRandom returns a random image ID
func (p *Provider) GetRandom() (id string, err error) {
	// This will be slow on large tables
	err = p.db.Get(&id, "select id from image order by random() limit 1")
	return
}

// List returns a list of all the images
func (p *Provider) List() ([]database.Image, error) {
	i := []database.Image{}
	err := p.db.Select(&i, "select * from image")

	if err != nil {
		return nil, err
	}

	return i, nil
}

// Shutdown shuts down the database client
func (p *Provider) Shutdown() {
	p.db.Close()
}
<|endoftext|>


<commit_before>package providers

import (
	"fmt"
	"strings"

	"github.com/coreos/go-semver/semver"
	"github.com/juju/errgo"
	logging "github.com/op/go-logging"
)

type InstanceConnection interface {
	SSHClient

	// Sync the filesystems on the instance
	Sync(log *logging.Logger) error

	// Exec executes a command on the instance
	Exec(log *logging.Logger, command string) (string, error)

	// EnableService calls `systemctl enable <name>`
	EnableService(log *logging.Logger, name string) error
	// RunScript uploads a script with given content and executes it
	RunScript(log *logging.Logger, scriptContent, scriptPath string) error

	GetClusterID(log *logging.Logger) (string, error)

	GetGluonEnv(log *logging.Logger) (string, error)

	GetMachineID(log *logging.Logger) (string, error)

	GetVaultCrt(log *logging.Logger) (string, error)

	GetVaultAddr(log *logging.Logger) (string, error)

	GetWeaveEnv(log *logging.Logger) (string, error)

	GetWeaveSeed(log *logging.Logger) (string, error)

	GetOSRelease(log *logging.Logger) (semver.Version, error)

	// IsEtcdProxyFromService queries the ETCD2 service on the instance to look for an ETCD_PROXY variable.
	IsEtcdProxyFromService(log *logging.Logger) (bool, error)

	// AddEtcdMember calls etcdctl to add a member to ETCD
	AddEtcdMember(log *logging.Logger, name, clusterIP string) error

	// RemoveEtcdMember calls etcdctl to remove a member from ETCD
	RemoveEtcdMember(log *logging.Logger, name, clusterIP string) error
}

type instanceConnection struct {
	client SSHClient
	host   string
}

// SSHConnect creates a new SSH connection to the given user on the given host.
func SSHConnect(userName, host string) (InstanceConnection, error) {
	client, err := DialSSH(userName, host)
	if err != nil {
		return nil, maskAny(err)
	}
	return &instanceConnection{client, host}, nil
}

func (s *instanceConnection) Close() error {
	return maskAny(s.client.Close())
}

func (s *instanceConnection) Run(log *logging.Logger, command, stdin string, quiet bool) (string, error) {
	out, err := s.client.Run(log, command, stdin, quiet)
	return out, maskAny(err)
}

// Sync the filesystems on the instance
func (s *instanceConnection) Sync(log *logging.Logger) error {
	if _, err := s.Run(log, "sudo sync", "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

// Exec executes a command on the instance
func (s *instanceConnection) Exec(log *logging.Logger, command string) (string, error) {
	stdout, err := s.Run(log, command, "", false)
	if err != nil {
		return stdout, maskAny(err)
	}
	return stdout, nil
}

// EnableService calls `systemctl enable <name>`
func (s *instanceConnection) EnableService(log *logging.Logger, name string) error {
	if _, err := s.Run(log, "sudo systemctl enable "+name, "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

// RunScript uploads a script with given content and executes it
func (s *instanceConnection) RunScript(log *logging.Logger, scriptContent, scriptPath string) error {
	if _, err := s.Run(log, fmt.Sprintf("sudo tee %s", scriptPath), scriptContent, false); err != nil {
		return maskAny(err)
	}
	if _, err := s.Run(log, fmt.Sprintf("sudo chmod +x %s", scriptPath), "", false); err != nil {
		return maskAny(err)
	}
	if _, err := s.Run(log, fmt.Sprintf("sudo %s", scriptPath), "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

func (s *instanceConnection) GetClusterID(log *logging.Logger) (string, error) {
	log.Debugf("Fetching cluster-id on %s", s.host)
	id, err := s.Run(log, "sudo cat /etc/pulcy/cluster-id", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetGluonEnv(log *logging.Logger) (string, error) {
	log.Debugf("Fetching gluon.env on %s", s.host)
	// gluon.env does not have to exists, so ignore errors by the `|| echo ""` parts.
	id, err := s.Run(log, "sh -c 'sudo cat /etc/pulcy/gluon.env || echo \"\"'", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetMachineID(log *logging.Logger) (string, error) {
	log.Debugf("Fetching machine-id on %s", s.host)
	id, err := s.Run(log, "cat /etc/machine-id", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetVaultCrt(log *logging.Logger) (string, error) {
	log.Debugf("Fetching vault.crt on %s", s.host)
	id, err := s.Run(log, "sudo cat /etc/pulcy/vault.crt", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetVaultAddr(log *logging.Logger) (string, error) {
	const prefix = "VAULT_ADDR="
	log.Debugf("Fetching vault-addr on %s", s.host)
	env, err := s.Run(log, "sudo cat /etc/pulcy/vault.env", "", false)
	if err != nil {
		return "", maskAny(err)
	}
	for _, line := range strings.Split(env, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, prefix) {
			return strings.TrimSpace(line[len(prefix):]), nil
		}
	}
	return "", maskAny(errgo.New("VAULT_ADDR not found in /etc/pulcy/vault.env"))
}

func (s *instanceConnection) GetWeaveEnv(log *logging.Logger) (string, error) {
	log.Debugf("Fetching weave.env on %s", s.host)
	id, err := s.Run(log, "sudo cat /etc/pulcy/weave.env", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetWeaveSeed(log *logging.Logger) (string, error) {
	log.Debugf("Fetching weave-seed on %s", s.host)
	// weave-seed does not have to exists, so ignore errors by the `|| echo ""` parts.
	id, err := s.Run(log, "sh -c 'sudo cat /etc/pulcy/weave-seed || echo \"\"'", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetOSRelease(log *logging.Logger) (semver.Version, error) {
	const prefix = "DISTRIB_RELEASE="
	log.Debugf("Fetching OS release on %s", s.host)
	env, err := s.Run(log, "cat /etc/lsb-release", "", false)
	if err != nil {
		return semver.Version{}, maskAny(err)
	}
	for _, line := range strings.Split(env, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, prefix) {
			v, err := semver.NewVersion(strings.TrimSpace(line[len(prefix):]))
			if err != nil {
				return semver.Version{}, maskAny(err)
			}
			return *v, nil
		}
	}
	return semver.Version{}, maskAny(errgo.Newf("%s not found in /etc/lsb-release", prefix))
}

// IsEtcdProxyFromService queries the ETCD2 service on the instance to look for an ETCD_PROXY variable.
func (s *instanceConnection) IsEtcdProxyFromService(log *logging.Logger) (bool, error) {
	log.Debugf("Fetching etcd proxy status on %s", s.host)
	cat, err := s.Run(log, "sudo systemctl cat etcd2.service", "", false)
	return strings.Contains(cat, "ETCD_PROXY"), maskAny(err)
}

// AddEtcdMember calls etcdctl to add a member to ETCD
func (s *instanceConnection) AddEtcdMember(log *logging.Logger, name, clusterIP string) error {
	log.Infof("Adding %s(%s) to etcd on %s", name, clusterIP, s.host)
	cmd := []string{
		"etcdctl",
		"member",
		"add",
		name,
		fmt.Sprintf("http://%s:2380", clusterIP),
	}
	if _, err := s.Run(log, strings.Join(cmd, " "), "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

// RemoveEtcdMember calls etcdctl to remove a member from ETCD
func (s *instanceConnection) RemoveEtcdMember(log *logging.Logger, name, clusterIP string) error {
	log.Infof("Removing %s(%s) from etcd on %s", name, clusterIP, s.host)
	id, err := s.Run(log, fmt.Sprintf("sh -c 'etcdctl member list | grep %s | cut -d: -f1 | cut -d[ -f1'", clusterIP), "", false)
	if err != nil {
		return maskAny(err)
	}
	cmd := []string{
		"etcdctl",
		"member",
		"remove",
		id,
	}
	if _, err := s.Run(log, strings.Join(cmd, " "), "", false); err != nil {
		return maskAny(err)
	}
	return nil
}
<commit_msg>Fixed ETCD proxy detection in case etcd2.service does not exist<commit_after>package providers

import (
	"fmt"
	"strings"

	"github.com/coreos/go-semver/semver"
	"github.com/juju/errgo"
	logging "github.com/op/go-logging"
)

type InstanceConnection interface {
	SSHClient

	// Sync the filesystems on the instance
	Sync(log *logging.Logger) error

	// Exec executes a command on the instance
	Exec(log *logging.Logger, command string) (string, error)

	// EnableService calls `systemctl enable <name>`
	EnableService(log *logging.Logger, name string) error
	// RunScript uploads a script with given content and executes it
	RunScript(log *logging.Logger, scriptContent, scriptPath string) error

	GetClusterID(log *logging.Logger) (string, error)

	GetGluonEnv(log *logging.Logger) (string, error)

	GetMachineID(log *logging.Logger) (string, error)

	GetVaultCrt(log *logging.Logger) (string, error)

	GetVaultAddr(log *logging.Logger) (string, error)

	GetWeaveEnv(log *logging.Logger) (string, error)

	GetWeaveSeed(log *logging.Logger) (string, error)

	GetOSRelease(log *logging.Logger) (semver.Version, error)

	// IsEtcdProxyFromService queries the ETCD2 service on the instance to look for an ETCD_PROXY variable.
	IsEtcdProxyFromService(log *logging.Logger) (bool, error)

	// AddEtcdMember calls etcdctl to add a member to ETCD
	AddEtcdMember(log *logging.Logger, name, clusterIP string) error

	// RemoveEtcdMember calls etcdctl to remove a member from ETCD
	RemoveEtcdMember(log *logging.Logger, name, clusterIP string) error
}

type instanceConnection struct {
	client SSHClient
	host   string
}

// SSHConnect creates a new SSH connection to the given user on the given host.
func SSHConnect(userName, host string) (InstanceConnection, error) {
	client, err := DialSSH(userName, host)
	if err != nil {
		return nil, maskAny(err)
	}
	return &instanceConnection{client, host}, nil
}

func (s *instanceConnection) Close() error {
	return maskAny(s.client.Close())
}

func (s *instanceConnection) Run(log *logging.Logger, command, stdin string, quiet bool) (string, error) {
	out, err := s.client.Run(log, command, stdin, quiet)
	return out, maskAny(err)
}

// Sync the filesystems on the instance
func (s *instanceConnection) Sync(log *logging.Logger) error {
	if _, err := s.Run(log, "sudo sync", "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

// Exec executes a command on the instance
func (s *instanceConnection) Exec(log *logging.Logger, command string) (string, error) {
	stdout, err := s.Run(log, command, "", false)
	if err != nil {
		return stdout, maskAny(err)
	}
	return stdout, nil
}

// EnableService calls `systemctl enable <name>`
func (s *instanceConnection) EnableService(log *logging.Logger, name string) error {
	if _, err := s.Run(log, "sudo systemctl enable "+name, "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

// RunScript uploads a script with given content and executes it
func (s *instanceConnection) RunScript(log *logging.Logger, scriptContent, scriptPath string) error {
	if _, err := s.Run(log, fmt.Sprintf("sudo tee %s", scriptPath), scriptContent, false); err != nil {
		return maskAny(err)
	}
	if _, err := s.Run(log, fmt.Sprintf("sudo chmod +x %s", scriptPath), "", false); err != nil {
		return maskAny(err)
	}
	if _, err := s.Run(log, fmt.Sprintf("sudo %s", scriptPath), "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

func (s *instanceConnection) GetClusterID(log *logging.Logger) (string, error) {
	log.Debugf("Fetching cluster-id on %s", s.host)
	id, err := s.Run(log, "sudo cat /etc/pulcy/cluster-id", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetGluonEnv(log *logging.Logger) (string, error) {
	log.Debugf("Fetching gluon.env on %s", s.host)
	// gluon.env does not have to exists, so ignore errors by the `|| echo ""` parts.
	id, err := s.Run(log, "sh -c 'sudo cat /etc/pulcy/gluon.env || echo \"\"'", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetMachineID(log *logging.Logger) (string, error) {
	log.Debugf("Fetching machine-id on %s", s.host)
	id, err := s.Run(log, "cat /etc/machine-id", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetVaultCrt(log *logging.Logger) (string, error) {
	log.Debugf("Fetching vault.crt on %s", s.host)
	id, err := s.Run(log, "sudo cat /etc/pulcy/vault.crt", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetVaultAddr(log *logging.Logger) (string, error) {
	const prefix = "VAULT_ADDR="
	log.Debugf("Fetching vault-addr on %s", s.host)
	env, err := s.Run(log, "sudo cat /etc/pulcy/vault.env", "", false)
	if err != nil {
		return "", maskAny(err)
	}
	for _, line := range strings.Split(env, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, prefix) {
			return strings.TrimSpace(line[len(prefix):]), nil
		}
	}
	return "", maskAny(errgo.New("VAULT_ADDR not found in /etc/pulcy/vault.env"))
}

func (s *instanceConnection) GetWeaveEnv(log *logging.Logger) (string, error) {
	log.Debugf("Fetching weave.env on %s", s.host)
	id, err := s.Run(log, "sudo cat /etc/pulcy/weave.env", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetWeaveSeed(log *logging.Logger) (string, error) {
	log.Debugf("Fetching weave-seed on %s", s.host)
	// weave-seed does not have to exists, so ignore errors by the `|| echo ""` parts.
	id, err := s.Run(log, "sh -c 'sudo cat /etc/pulcy/weave-seed || echo \"\"'", "", false)
	return id, maskAny(err)
}

func (s *instanceConnection) GetOSRelease(log *logging.Logger) (semver.Version, error) {
	const prefix = "DISTRIB_RELEASE="
	log.Debugf("Fetching OS release on %s", s.host)
	env, err := s.Run(log, "cat /etc/lsb-release", "", false)
	if err != nil {
		return semver.Version{}, maskAny(err)
	}
	for _, line := range strings.Split(env, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, prefix) {
			v, err := semver.NewVersion(strings.TrimSpace(line[len(prefix):]))
			if err != nil {
				return semver.Version{}, maskAny(err)
			}
			return *v, nil
		}
	}
	return semver.Version{}, maskAny(errgo.Newf("%s not found in /etc/lsb-release", prefix))
}

// IsEtcdProxyFromService queries the ETCD2 service on the instance to look for an ETCD_PROXY variable.
func (s *instanceConnection) IsEtcdProxyFromService(log *logging.Logger) (bool, error) {
	log.Debugf("Fetching etcd proxy status on %s", s.host)
	// weave-seed does not have to exists, so ignore errors by the `|| echo ""` parts.
	cat, err := s.Run(log, "sh -c 'sudo systemctl cat etcd2.service || echo \"\"'", "", false)
	return cat == "" || strings.Contains(cat, "ETCD_PROXY"), maskAny(err)
}

// AddEtcdMember calls etcdctl to add a member to ETCD
func (s *instanceConnection) AddEtcdMember(log *logging.Logger, name, clusterIP string) error {
	log.Infof("Adding %s(%s) to etcd on %s", name, clusterIP, s.host)
	cmd := []string{
		"etcdctl",
		"member",
		"add",
		name,
		fmt.Sprintf("http://%s:2380", clusterIP),
	}
	if _, err := s.Run(log, strings.Join(cmd, " "), "", false); err != nil {
		return maskAny(err)
	}
	return nil
}

// RemoveEtcdMember calls etcdctl to remove a member from ETCD
func (s *instanceConnection) RemoveEtcdMember(log *logging.Logger, name, clusterIP string) error {
	log.Infof("Removing %s(%s) from etcd on %s", name, clusterIP, s.host)
	id, err := s.Run(log, fmt.Sprintf("sh -c 'etcdctl member list | grep %s | cut -d: -f1 | cut -d[ -f1'", clusterIP), "", false)
	if err != nil {
		return maskAny(err)
	}
	cmd := []string{
		"etcdctl",
		"member",
		"remove",
		id,
	}
	if _, err := s.Run(log, strings.Join(cmd, " "), "", false); err != nil {
		return maskAny(err)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"github.com/dotcloud/docker"
	"github.com/globocom/config"
	"github.com/globocom/docker-cluster/cluster"
	"github.com/globocom/tsuru/app"
	"github.com/globocom/tsuru/db"
	"labix.org/v2/mgo/bson"
	"launchpad.net/gocheck"
	"net/http/httptest"
	"strings"
)

type FlattenSuite struct {
	apps      []app.App
	conn      *db.Storage
	server    *httptest.Server
	cleanup   func()
	calls     int
	scheduler *fakeScheduler
}

var _ = gocheck.Suite(&FlattenSuite{})

func (s *FlattenSuite) SetUpSuite(c *gocheck.C) {
	var err error
	config.Set("docker:collection", "docker")
	config.Set("docker:repository-namespace", "tsuru")
	config.Set("database:url", "127.0.0.1:27017")
	config.Set("database:name", "docker_flatten")
	s.conn, err = db.Conn()
	units := []app.Unit{{Name: "4fa6e0f0c678"}, {Name: "e90e34656806"}}
	app1 := app.App{Name: "app1", Platform: "python", Deploys: 40, Units: units}
	err = s.conn.Apps().Insert(app1)
	c.Assert(err, gocheck.IsNil)
	app2 := app.App{Name: "app2", Platform: "python", Deploys: 20, Units: units}
	err = s.conn.Apps().Insert(app2)
	c.Assert(err, gocheck.IsNil)
	app3 := app.App{Name: "app3", Platform: "python", Deploys: 3, Units: units}
	err = s.conn.Apps().Insert(app3)
	c.Assert(err, gocheck.IsNil)
	app4 := app.App{Name: "app4", Platform: "python", Deploys: 19, Units: units}
	err = s.conn.Apps().Insert(app4)
	c.Assert(err, gocheck.IsNil)
	s.apps = append(s.apps, []app.App{app1, app2, app3, app4}...)
	s.cleanup, s.server = startDockerTestServer("4567", &s.calls)
	c.Assert(err, gocheck.IsNil)
	config.Set("docker:registry", strings.Replace(s.server.URL, "http://", "", 1))
	node := cluster.Node{ID: "server", Address: s.server.URL}
	s.scheduler = &fakeScheduler{}
	dCluster, _ = cluster.New(s.scheduler, node)
	dCluster.SetStorage(&mapStorage{})
	err = newImage(assembleImageName("app1"), s.server.URL)
	c.Assert(err, gocheck.IsNil)
	err = newImage(assembleImageName("app2"), s.server.URL)
	c.Assert(err, gocheck.IsNil)
	err = collection().Insert(&container{ID: "app1id", AppName: "app1", Image: assembleImageName("app1")})
	c.Assert(err, gocheck.IsNil)
	err = collection().Insert(&container{ID: "app2id", AppName: "app2", Image: assembleImageName("app2")})
	c.Assert(err, gocheck.IsNil)
}

func (s *FlattenSuite) TearDownSuite(c *gocheck.C) {
	collection().RemoveAll(nil)
	names := make([]string, len(s.apps))
	for i, a := range s.apps {
		names[i] = a.GetName()
	}
	_, err := s.conn.Apps().RemoveAll(bson.M{"name": bson.M{"$in": names}})
	c.Assert(err, gocheck.IsNil)
	collection().RemoveAll(nil)
	s.cleanup()
	defer config.Set("docker:registry", "")
}

func (s *FlattenSuite) TestImagesToFlattenRetrievesOnlyUnitsWith20DeploysOrMore(c *gocheck.C) {
	images := imagesToFlatten()
	c.Assert(len(images), gocheck.Equals, 2)
	expected := []string{assembleImageName("app1"), assembleImageName("app2")}
	c.Assert(images, gocheck.DeepEquals, expected)
}

func (s *FlattenSuite) TestFlatten(c *gocheck.C) {
	s.scheduler.container = &docker.Container{ID: "containerid"}
	Flatten()
	c.Assert(s.calls, gocheck.Equals, 10) //create, export, import, remove old img, remove container twice
}
<commit_msg>provision/docker/flatten_test: spliting SetUpSuite method by responsibility<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"github.com/dotcloud/docker"
	"github.com/globocom/config"
	"github.com/globocom/docker-cluster/cluster"
	"github.com/globocom/tsuru/app"
	"github.com/globocom/tsuru/db"
	"labix.org/v2/mgo/bson"
	"launchpad.net/gocheck"
	"net/http/httptest"
	"strings"
)

type FlattenSuite struct {
	apps      []app.App
	conn      *db.Storage
	server    *httptest.Server
	cleanup   func()
	calls     int
	scheduler *fakeScheduler
}

var _ = gocheck.Suite(&FlattenSuite{})

func (s *FlattenSuite) createApps(c *gocheck.C) {
	units := []app.Unit{{Name: "4fa6e0f0c678"}, {Name: "e90e34656806"}}
	app1 := app.App{Name: "app1", Platform: "python", Deploys: 40, Units: units}
	err := s.conn.Apps().Insert(app1)
	c.Assert(err, gocheck.IsNil)
	app2 := app.App{Name: "app2", Platform: "python", Deploys: 20, Units: units}
	err = s.conn.Apps().Insert(app2)
	c.Assert(err, gocheck.IsNil)
	app3 := app.App{Name: "app3", Platform: "python", Deploys: 3, Units: units}
	err = s.conn.Apps().Insert(app3)
	c.Assert(err, gocheck.IsNil)
	app4 := app.App{Name: "app4", Platform: "python", Deploys: 19, Units: units}
	err = s.conn.Apps().Insert(app4)
	c.Assert(err, gocheck.IsNil)
	s.apps = append(s.apps, []app.App{app1, app2, app3, app4}...)
}

func (s *FlattenSuite) setConfig() {
	config.Set("docker:collection", "docker")
	config.Set("docker:repository-namespace", "tsuru")
	config.Set("database:url", "127.0.0.1:27017")
	config.Set("database:name", "docker_flatten")
}

func (s *FlattenSuite) createContainers(c *gocheck.C) {
	err := newImage(assembleImageName("app1"), s.server.URL)
	c.Assert(err, gocheck.IsNil)
	err = newImage(assembleImageName("app2"), s.server.URL)
	c.Assert(err, gocheck.IsNil)
	err = collection().Insert(&container{ID: "app1id", AppName: "app1", Image: assembleImageName("app1")})
	c.Assert(err, gocheck.IsNil)
	err = collection().Insert(&container{ID: "app2id", AppName: "app2", Image: assembleImageName("app2")})
	c.Assert(err, gocheck.IsNil)
}

func (s *FlattenSuite) setupDockerCluster() {
	node := cluster.Node{ID: "server", Address: s.server.URL}
	s.scheduler = &fakeScheduler{}
	dCluster, _ = cluster.New(s.scheduler, node)
	dCluster.SetStorage(&mapStorage{})
}

func (s *FlattenSuite) SetUpSuite(c *gocheck.C) {
	var err error
	s.setConfig()
	s.conn, err = db.Conn()
	s.createApps(c)
	s.cleanup, s.server = startDockerTestServer("4567", &s.calls)
	c.Assert(err, gocheck.IsNil)
	config.Set("docker:registry", strings.Replace(s.server.URL, "http://", "", 1))
	s.setupDockerCluster()
	s.createContainers(c)
}

func (s *FlattenSuite) TearDownSuite(c *gocheck.C) {
	collection().RemoveAll(nil)
	names := make([]string, len(s.apps))
	for i, a := range s.apps {
		names[i] = a.GetName()
	}
	_, err := s.conn.Apps().RemoveAll(bson.M{"name": bson.M{"$in": names}})
	c.Assert(err, gocheck.IsNil)
	collection().RemoveAll(nil)
	s.cleanup()
	defer config.Set("docker:registry", "")
}

func (s *FlattenSuite) TestImagesToFlattenRetrievesOnlyUnitsWith20DeploysOrMore(c *gocheck.C) {
	images := imagesToFlatten()
	c.Assert(len(images), gocheck.Equals, 2)
	expected := []string{assembleImageName("app1"), assembleImageName("app2")}
	c.Assert(images, gocheck.DeepEquals, expected)
}

func (s *FlattenSuite) TestFlatten(c *gocheck.C) {
	s.scheduler.container = &docker.Container{ID: "containerid"}
	Flatten()
	c.Assert(s.calls, gocheck.Equals, 10) //create, export, import, remove old img, remove container twice
}
<|endoftext|>


<commit_before>// R.I.P Shadowsocks
package shadowsocks

import (
	"crypto/rand"
	"io"
	"sync"

	"github.com/v2ray/v2ray-core/app"
	"github.com/v2ray/v2ray-core/app/dispatcher"
	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/crypto"
	v2io "github.com/v2ray/v2ray-core/common/io"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/common/protocol"
	"github.com/v2ray/v2ray-core/common/serial"
	"github.com/v2ray/v2ray-core/proxy"
	"github.com/v2ray/v2ray-core/proxy/internal"
	"github.com/v2ray/v2ray-core/transport/hub"
)

type Shadowsocks struct {
	packetDispatcher dispatcher.PacketDispatcher
	config           *Config
	port             v2net.Port
	accepting        bool
	tcpHub           *hub.TCPHub
	udpHub           *hub.UDPHub
	udpServer        *hub.UDPServer
}

func NewShadowsocks(config *Config, packetDispatcher dispatcher.PacketDispatcher) *Shadowsocks {
	return &Shadowsocks{
		config:           config,
		packetDispatcher: packetDispatcher,
	}
}

func (this *Shadowsocks) Port() v2net.Port {
	return this.port
}

func (this *Shadowsocks) Close() {
	this.accepting = false
	// TODO: synchronization
	if this.tcpHub != nil {
		this.tcpHub.Close()
		this.tcpHub = nil
	}

	if this.udpHub != nil {
		this.udpHub.Close()
		this.udpHub = nil
	}

}

func (this *Shadowsocks) Listen(port v2net.Port) error {
	if this.accepting {
		if this.port == port {
			return nil
		} else {
			return proxy.ErrorAlreadyListening
		}
	}

	tcpHub, err := hub.ListenTCP(port, this.handleConnection)
	if err != nil {
		log.Error("Shadowsocks: Failed to listen TCP on port ", port, ": ", err)
		return err
	}
	this.tcpHub = tcpHub

	if this.config.UDP {
		this.udpServer = hub.NewUDPServer(this.packetDispatcher)
		udpHub, err := hub.ListenUDP(port, this.handlerUDPPayload)
		if err != nil {
			log.Error("Shadowsocks: Failed to listen UDP on port ", port, ": ", err)
			return err
		}
		this.udpHub = udpHub
	}

	this.port = port
	this.accepting = true

	return nil
}

func (this *Shadowsocks) handlerUDPPayload(payload *alloc.Buffer, source v2net.Destination) {
	defer payload.Release()

	ivLen := this.config.Cipher.IVSize()
	iv := payload.Value[:ivLen]
	key := this.config.Key
	payload.SliceFrom(ivLen)

	stream, err := this.config.Cipher.NewDecodingStream(key, iv)
	if err != nil {
		log.Error("Shadowsocks: Failed to create decoding stream: ", err)
		return
	}

	reader := crypto.NewCryptionReader(stream, payload)

	request, err := ReadRequest(reader, NewAuthenticator(HeaderKeyGenerator(key, iv)), true)
	if err != nil {
		log.Access(source, serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
		log.Warning("Shadowsocks: Invalid request from ", source, ": ", err)
		return
	}
	defer request.Release()

	dest := v2net.UDPDestination(request.Address, request.Port)
	log.Access(source, dest, log.AccessAccepted, serial.StringLiteral(""))
	log.Info("Shadowsocks: Tunnelling request to ", dest)

	this.udpServer.Dispatch(source, dest, request.DetachUDPPayload(), func(destination v2net.Destination, payload *alloc.Buffer) {
		defer payload.Release()

		response := alloc.NewBuffer().Slice(0, ivLen)
		defer response.Release()

		rand.Read(response.Value)
		respIv := response.Value

		stream, err := this.config.Cipher.NewEncodingStream(key, respIv)
		if err != nil {
			log.Error("Shadowsocks: Failed to create encoding stream: ", err)
			return
		}

		writer := crypto.NewCryptionWriter(stream, response)

		switch {
		case request.Address.IsIPv4():
			writer.Write([]byte{AddrTypeIPv4})
			writer.Write(request.Address.IP())
		case request.Address.IsIPv6():
			writer.Write([]byte{AddrTypeIPv6})
			writer.Write(request.Address.IP())
		case request.Address.IsDomain():
			writer.Write([]byte{AddrTypeDomain, byte(len(request.Address.Domain()))})
			writer.Write([]byte(request.Address.Domain()))
		}

		writer.Write(request.Port.Bytes())
		writer.Write(payload.Value)

		if request.OTA {
			respAuth := NewAuthenticator(HeaderKeyGenerator(key, respIv))
			respAuth.Authenticate(response.Value, response.Value[ivLen:])
		}

		this.udpHub.WriteTo(response.Value, source)
	})
}

func (this *Shadowsocks) handleConnection(conn hub.Connection) {
	defer conn.Close()

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()

	timedReader := v2net.NewTimeOutReader(16, conn)
	defer timedReader.Release()

	bufferedReader := v2io.NewBufferedReader(timedReader)
	defer bufferedReader.Release()

	ivLen := this.config.Cipher.IVSize()
	_, err := io.ReadFull(bufferedReader, buffer.Value[:ivLen])
	if err != nil {
		log.Access(conn.RemoteAddr(), serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
		log.Error("Shadowsocks: Failed to read IV: ", err)
		return
	}

	iv := buffer.Value[:ivLen]
	key := this.config.Key

	stream, err := this.config.Cipher.NewDecodingStream(key, iv)
	if err != nil {
		log.Error("Shadowsocks: Failed to create decoding stream: ", err)
		return
	}

	reader := crypto.NewCryptionReader(stream, bufferedReader)

	request, err := ReadRequest(reader, NewAuthenticator(HeaderKeyGenerator(key, iv)), false)
	if err != nil {
		log.Access(conn.RemoteAddr(), serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
		log.Warning("Shadowsocks: Invalid request from ", conn.RemoteAddr(), ": ", err)
		return
	}
	defer request.Release()
	bufferedReader.SetCached(false)

	userSettings := protocol.GetUserSettings(this.config.Level)
	timedReader.SetTimeOut(userSettings.PayloadReadTimeout)

	dest := v2net.TCPDestination(request.Address, request.Port)
	log.Access(conn.RemoteAddr(), dest, log.AccessAccepted, serial.StringLiteral(""))
	log.Info("Shadowsocks: Tunnelling request to ", dest)

	ray := this.packetDispatcher.DispatchToOutbound(dest)

	var writeFinish sync.Mutex
	writeFinish.Lock()
	go func() {
		if payload, err := ray.InboundOutput().Read(); err == nil {
			payload.SliceBack(ivLen)
			rand.Read(payload.Value[:ivLen])

			stream, err := this.config.Cipher.NewEncodingStream(key, payload.Value[:ivLen])
			if err != nil {
				log.Error("Shadowsocks: Failed to create encoding stream: ", err)
				return
			}
			stream.XORKeyStream(payload.Value[ivLen:], payload.Value[ivLen:])

			conn.Write(payload.Value)
			payload.Release()

			writer := crypto.NewCryptionWriter(stream, conn)
			v2writer := v2io.NewAdaptiveWriter(writer)

			v2io.Pipe(ray.InboundOutput(), v2writer)
			ray.InboundOutput().Release()
			writer.Release()
			v2writer.Release()
		}
		writeFinish.Unlock()
	}()

	var payloadReader v2io.Reader
	if request.OTA {
		payloadAuth := NewAuthenticator(ChunkKeyGenerator(iv))
		payloadReader = NewChunkReader(reader, payloadAuth)
	} else {
		payloadReader = v2io.NewAdaptiveReader(reader)
	}

	v2io.Pipe(payloadReader, ray.InboundInput())
	ray.InboundInput().Close()
	payloadReader.Release()

	writeFinish.Lock()
}

func init() {
	internal.MustRegisterInboundHandlerCreator("shadowsocks",
		func(space app.Space, rawConfig interface{}) (proxy.InboundHandler, error) {
			if !space.HasApp(dispatcher.APP_ID) {
				return nil, internal.ErrorBadConfiguration
			}
			return NewShadowsocks(
				rawConfig.(*Config),
				space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher)), nil
		})
}
<commit_msg>fix nil reference issue in shadowsocks<commit_after>// R.I.P Shadowsocks
package shadowsocks

import (
	"crypto/rand"
	"io"
	"sync"

	"github.com/v2ray/v2ray-core/app"
	"github.com/v2ray/v2ray-core/app/dispatcher"
	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/crypto"
	v2io "github.com/v2ray/v2ray-core/common/io"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/common/protocol"
	"github.com/v2ray/v2ray-core/common/serial"
	"github.com/v2ray/v2ray-core/proxy"
	"github.com/v2ray/v2ray-core/proxy/internal"
	"github.com/v2ray/v2ray-core/transport/hub"
)

type Shadowsocks struct {
	packetDispatcher dispatcher.PacketDispatcher
	config           *Config
	port             v2net.Port
	accepting        bool
	tcpHub           *hub.TCPHub
	udpHub           *hub.UDPHub
	udpServer        *hub.UDPServer
}

func NewShadowsocks(config *Config, packetDispatcher dispatcher.PacketDispatcher) *Shadowsocks {
	return &Shadowsocks{
		config:           config,
		packetDispatcher: packetDispatcher,
	}
}

func (this *Shadowsocks) Port() v2net.Port {
	return this.port
}

func (this *Shadowsocks) Close() {
	this.accepting = false
	// TODO: synchronization
	if this.tcpHub != nil {
		this.tcpHub.Close()
		this.tcpHub = nil
	}

	if this.udpHub != nil {
		this.udpHub.Close()
		this.udpHub = nil
	}

}

func (this *Shadowsocks) Listen(port v2net.Port) error {
	if this.accepting {
		if this.port == port {
			return nil
		} else {
			return proxy.ErrorAlreadyListening
		}
	}

	tcpHub, err := hub.ListenTCP(port, this.handleConnection)
	if err != nil {
		log.Error("Shadowsocks: Failed to listen TCP on port ", port, ": ", err)
		return err
	}
	this.tcpHub = tcpHub

	if this.config.UDP {
		this.udpServer = hub.NewUDPServer(this.packetDispatcher)
		udpHub, err := hub.ListenUDP(port, this.handlerUDPPayload)
		if err != nil {
			log.Error("Shadowsocks: Failed to listen UDP on port ", port, ": ", err)
			return err
		}
		this.udpHub = udpHub
	}

	this.port = port
	this.accepting = true

	return nil
}

func (this *Shadowsocks) handlerUDPPayload(payload *alloc.Buffer, source v2net.Destination) {
	defer payload.Release()

	ivLen := this.config.Cipher.IVSize()
	iv := payload.Value[:ivLen]
	key := this.config.Key
	payload.SliceFrom(ivLen)

	stream, err := this.config.Cipher.NewDecodingStream(key, iv)
	if err != nil {
		log.Error("Shadowsocks: Failed to create decoding stream: ", err)
		return
	}

	reader := crypto.NewCryptionReader(stream, payload)

	request, err := ReadRequest(reader, NewAuthenticator(HeaderKeyGenerator(key, iv)), true)
	if err != nil {
		log.Access(source, serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
		log.Warning("Shadowsocks: Invalid request from ", source, ": ", err)
		return
	}
	//defer request.Release()

	dest := v2net.UDPDestination(request.Address, request.Port)
	log.Access(source, dest, log.AccessAccepted, serial.StringLiteral(""))
	log.Info("Shadowsocks: Tunnelling request to ", dest)

	this.udpServer.Dispatch(source, dest, request.DetachUDPPayload(), func(destination v2net.Destination, payload *alloc.Buffer) {
		defer payload.Release()

		response := alloc.NewBuffer().Slice(0, ivLen)
		defer response.Release()

		rand.Read(response.Value)
		respIv := response.Value

		stream, err := this.config.Cipher.NewEncodingStream(key, respIv)
		if err != nil {
			log.Error("Shadowsocks: Failed to create encoding stream: ", err)
			return
		}

		writer := crypto.NewCryptionWriter(stream, response)

		switch {
		case request.Address.IsIPv4():
			writer.Write([]byte{AddrTypeIPv4})
			writer.Write(request.Address.IP())
		case request.Address.IsIPv6():
			writer.Write([]byte{AddrTypeIPv6})
			writer.Write(request.Address.IP())
		case request.Address.IsDomain():
			writer.Write([]byte{AddrTypeDomain, byte(len(request.Address.Domain()))})
			writer.Write([]byte(request.Address.Domain()))
		}

		writer.Write(request.Port.Bytes())
		writer.Write(payload.Value)

		if request.OTA {
			respAuth := NewAuthenticator(HeaderKeyGenerator(key, respIv))
			respAuth.Authenticate(response.Value, response.Value[ivLen:])
		}

		this.udpHub.WriteTo(response.Value, source)
	})
}

func (this *Shadowsocks) handleConnection(conn hub.Connection) {
	defer conn.Close()

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()

	timedReader := v2net.NewTimeOutReader(16, conn)
	defer timedReader.Release()

	bufferedReader := v2io.NewBufferedReader(timedReader)
	defer bufferedReader.Release()

	ivLen := this.config.Cipher.IVSize()
	_, err := io.ReadFull(bufferedReader, buffer.Value[:ivLen])
	if err != nil {
		log.Access(conn.RemoteAddr(), serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
		log.Error("Shadowsocks: Failed to read IV: ", err)
		return
	}

	iv := buffer.Value[:ivLen]
	key := this.config.Key

	stream, err := this.config.Cipher.NewDecodingStream(key, iv)
	if err != nil {
		log.Error("Shadowsocks: Failed to create decoding stream: ", err)
		return
	}

	reader := crypto.NewCryptionReader(stream, bufferedReader)

	request, err := ReadRequest(reader, NewAuthenticator(HeaderKeyGenerator(key, iv)), false)
	if err != nil {
		log.Access(conn.RemoteAddr(), serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
		log.Warning("Shadowsocks: Invalid request from ", conn.RemoteAddr(), ": ", err)
		return
	}
	defer request.Release()
	bufferedReader.SetCached(false)

	userSettings := protocol.GetUserSettings(this.config.Level)
	timedReader.SetTimeOut(userSettings.PayloadReadTimeout)

	dest := v2net.TCPDestination(request.Address, request.Port)
	log.Access(conn.RemoteAddr(), dest, log.AccessAccepted, serial.StringLiteral(""))
	log.Info("Shadowsocks: Tunnelling request to ", dest)

	ray := this.packetDispatcher.DispatchToOutbound(dest)

	var writeFinish sync.Mutex
	writeFinish.Lock()
	go func() {
		if payload, err := ray.InboundOutput().Read(); err == nil {
			payload.SliceBack(ivLen)
			rand.Read(payload.Value[:ivLen])

			stream, err := this.config.Cipher.NewEncodingStream(key, payload.Value[:ivLen])
			if err != nil {
				log.Error("Shadowsocks: Failed to create encoding stream: ", err)
				return
			}
			stream.XORKeyStream(payload.Value[ivLen:], payload.Value[ivLen:])

			conn.Write(payload.Value)
			payload.Release()

			writer := crypto.NewCryptionWriter(stream, conn)
			v2writer := v2io.NewAdaptiveWriter(writer)

			v2io.Pipe(ray.InboundOutput(), v2writer)
			ray.InboundOutput().Release()
			writer.Release()
			v2writer.Release()
		}
		writeFinish.Unlock()
	}()

	var payloadReader v2io.Reader
	if request.OTA {
		payloadAuth := NewAuthenticator(ChunkKeyGenerator(iv))
		payloadReader = NewChunkReader(reader, payloadAuth)
	} else {
		payloadReader = v2io.NewAdaptiveReader(reader)
	}

	v2io.Pipe(payloadReader, ray.InboundInput())
	ray.InboundInput().Close()
	payloadReader.Release()

	writeFinish.Lock()
}

func init() {
	internal.MustRegisterInboundHandlerCreator("shadowsocks",
		func(space app.Space, rawConfig interface{}) (proxy.InboundHandler, error) {
			if !space.HasApp(dispatcher.APP_ID) {
				return nil, internal.ErrorBadConfiguration
			}
			return NewShadowsocks(
				rawConfig.(*Config),
				space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher)), nil
		})
}
<|endoftext|>


<commit_before>// +build functional

package cri_containerd

import (
	"context"
	"strconv"
	"strings"
	"testing"

	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
)

//TODO:
// Create POD, Container interfaces and make
// Runp, create, exec, start, stop as methods instead of funcs
func execContainer(ctx context.Context, t *testing.T, client runtime.RuntimeServiceClient, request *runtime.ExecSyncRequest) *runtime.ExecSyncResponse {
	response, err := client.ExecSync(ctx, request)
	if err != nil {
		t.Fatalf("failed ExecSync in container: %s, with: %v", request.ContainerId, err)
	}
	return response
}

func runexecContainerTestWithSandbox(t *testing.T, sandboxRequest *runtime.RunPodSandboxRequest, request *runtime.CreateContainerRequest, execReq *runtime.ExecSyncRequest) *runtime.ExecSyncResponse {
	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, sandboxRequest)
	defer func() {
		stopAndRemovePodSandbox(t, client, ctx, podID)
	}()
	request.PodSandboxId = podID
	request.SandboxConfig = sandboxRequest.Config
	containerID := createContainer(t, client, ctx, request)
	defer func() {
		stopAndRemoveContainer(t, client, ctx, containerID)
	}()
	_, err := client.StartContainer(ctx, &runtime.StartContainerRequest{
		ContainerId: containerID,
	})
	if err != nil {
		t.Fatalf("failed StartContainer request for container: %s, with: %v", containerID, err)
	}

	execReq.ContainerId = containerID
	return execContainer(ctx, t, client, execReq)
}

func execContainerLCOW(t *testing.T, uid int64, cmd []string) *runtime.ExecSyncResponse {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause, imageLcowCosmos})

	//run podsandbox request
	sandboxRequest := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name() + "-Sandbox2",
				Uid:       "0",
				Namespace: testNamespace,
			},
			Linux: &runtime.LinuxPodSandboxConfig{
				SecurityContext: &runtime.LinuxSandboxSecurityContext{
					Privileged: true,
				},
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	//create container request
	request := &runtime.CreateContainerRequest{
		Config: &runtime.ContainerConfig{
			Metadata: &runtime.ContainerMetadata{
				Name: t.Name() + "-Container2",
			},
			Image: &runtime.ImageSpec{
				Image: imageLcowCosmos,
			},
			// Hold this command open until killed
			Command: []string{
				"sleep",
				"100000",
			},
			Linux: &runtime.LinuxContainerConfig{
				SecurityContext: &runtime.LinuxContainerSecurityContext{
					RunAsUser: &runtime.Int64Value{
						Value: uid,
					},
				},
			},
		},
	}

	//exec request
	execRequest := &runtime.ExecSyncRequest{
		ContainerId: "",
		Cmd:         cmd,
		Timeout:     20,
	}

	return runexecContainerTestWithSandbox(t, sandboxRequest, request, execRequest)
}

func Test_ExecContainer_RunAs_LCOW(t *testing.T) {
	// this is just saying 'give me the UID of the process with pid = 1; ignore headers'
	r := execContainerLCOW(t,
		9000, // cosmosarno user
		[]string{
			"ps",
			"-o",
			"uid",
			"-p",
			"1",
			"--no-headers",
		})

	output := strings.TrimSpace(string(r.Stdout))
	errorMsg := string(r.Stderr)
	exitCode := int(r.ExitCode)

	t.Logf("exec request exited with code: %d", exitCode)
	t.Logf("exec request output: %v", output)

	if exitCode != 0 {
		t.Fatalf("Test %v exited with exit code: %d, Test_CreateContainer_RunAs_LCOW", errorMsg, exitCode)
	}

	if output != "9000" {
		t.Fatalf("failed to start container with runas option: error: %v, exitcode: %d", errorMsg, exitCode)
	}
}

func Test_ExecContainer_LCOW_HasEntropy(t *testing.T) {
	r := execContainerLCOW(t, 9000, []string{"cat", "/proc/sys/kernel/random/entropy_avail"})
	if r.ExitCode != 0 {
		t.Fatalf("failed with exit code %d to cat entropy_avail: %s", r.ExitCode, r.Stderr)
	}
	output := strings.TrimSpace(string(r.Stdout))
	bits, err := strconv.ParseInt(output, 10, 0)
	if err != nil {
		t.Fatalf("could not parse entropy output %s: %s", output, err)
	}
	if bits < 2000 {
		t.Fatalf("%d is fewer than 2000 bits entropy", bits)
	}
	t.Logf("got %d bits entropy", bits)
}
<commit_msg>Fix LCOW run-as-user test to use own PID namespace<commit_after>// +build functional

package cri_containerd

import (
	"context"
	"strconv"
	"strings"
	"testing"

	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
)

//TODO:
// Create POD, Container interfaces and make
// Runp, create, exec, start, stop as methods instead of funcs
func execContainer(ctx context.Context, t *testing.T, client runtime.RuntimeServiceClient, request *runtime.ExecSyncRequest) *runtime.ExecSyncResponse {
	response, err := client.ExecSync(ctx, request)
	if err != nil {
		t.Fatalf("failed ExecSync in container: %s, with: %v", request.ContainerId, err)
	}
	return response
}

func runexecContainerTestWithSandbox(t *testing.T, sandboxRequest *runtime.RunPodSandboxRequest, request *runtime.CreateContainerRequest, execReq *runtime.ExecSyncRequest) *runtime.ExecSyncResponse {
	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, sandboxRequest)
	defer func() {
		stopAndRemovePodSandbox(t, client, ctx, podID)
	}()
	request.PodSandboxId = podID
	request.SandboxConfig = sandboxRequest.Config
	containerID := createContainer(t, client, ctx, request)
	defer func() {
		stopAndRemoveContainer(t, client, ctx, containerID)
	}()
	_, err := client.StartContainer(ctx, &runtime.StartContainerRequest{
		ContainerId: containerID,
	})
	if err != nil {
		t.Fatalf("failed StartContainer request for container: %s, with: %v", containerID, err)
	}

	execReq.ContainerId = containerID
	return execContainer(ctx, t, client, execReq)
}

func execContainerLCOW(t *testing.T, uid int64, cmd []string) *runtime.ExecSyncResponse {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause, imageLcowCosmos})

	//run podsandbox request
	sandboxRequest := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name() + "-Sandbox2",
				Uid:       "0",
				Namespace: testNamespace,
			},
			Linux: &runtime.LinuxPodSandboxConfig{
				SecurityContext: &runtime.LinuxSandboxSecurityContext{
					Privileged: true,
				},
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	//create container request
	request := &runtime.CreateContainerRequest{
		Config: &runtime.ContainerConfig{
			Metadata: &runtime.ContainerMetadata{
				Name: t.Name() + "-Container2",
			},
			Image: &runtime.ImageSpec{
				Image: imageLcowCosmos,
			},
			// Hold this command open until killed
			Command: []string{
				"sleep",
				"100000",
			},
			Linux: &runtime.LinuxContainerConfig{
				SecurityContext: &runtime.LinuxContainerSecurityContext{
					// Our tests rely on the init process for the workload
					// container being pid 1, but the CRI default is to use the
					// pod's pid namespace.
					NamespaceOptions: &runtime.NamespaceOption{
						Pid: runtime.NamespaceMode_CONTAINER,
					},
					RunAsUser: &runtime.Int64Value{
						Value: uid,
					},
				},
			},
		},
	}

	//exec request
	execRequest := &runtime.ExecSyncRequest{
		ContainerId: "",
		Cmd:         cmd,
		Timeout:     20,
	}

	return runexecContainerTestWithSandbox(t, sandboxRequest, request, execRequest)
}

func Test_ExecContainer_RunAs_LCOW(t *testing.T) {
	// this is just saying 'give me the UID of the process with pid = 1; ignore headers'
	r := execContainerLCOW(t,
		9000, // cosmosarno user
		[]string{
			"ps",
			"-o",
			"uid",
			"-p",
			"1",
			"--no-headers",
		})

	output := strings.TrimSpace(string(r.Stdout))
	errorMsg := string(r.Stderr)
	exitCode := int(r.ExitCode)

	t.Logf("exec request exited with code: %d", exitCode)
	t.Logf("exec request output: %v", output)

	if exitCode != 0 {
		t.Fatalf("Test %v exited with exit code: %d, Test_CreateContainer_RunAs_LCOW", errorMsg, exitCode)
	}

	if output != "9000" {
		t.Fatalf("failed to start container with runas option: error: %v, exitcode: %d", errorMsg, exitCode)
	}
}

func Test_ExecContainer_LCOW_HasEntropy(t *testing.T) {
	r := execContainerLCOW(t, 9000, []string{"cat", "/proc/sys/kernel/random/entropy_avail"})
	if r.ExitCode != 0 {
		t.Fatalf("failed with exit code %d to cat entropy_avail: %s", r.ExitCode, r.Stderr)
	}
	output := strings.TrimSpace(string(r.Stdout))
	bits, err := strconv.ParseInt(output, 10, 0)
	if err != nil {
		t.Fatalf("could not parse entropy output %s: %s", output, err)
	}
	if bits < 2000 {
		t.Fatalf("%d is fewer than 2000 bits entropy", bits)
	}
	t.Logf("got %d bits entropy", bits)
}
<|endoftext|>


<commit_before>package infodownloader

import (
	"bytes"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/peerconn"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/semaphore"
)

const blockSize = 16 * 1024

const maxQueuedBlocks = 10

// InfoDownloader downloads all blocks of a piece from a peer.
type InfoDownloader struct {
	extID     uint8
	totalSize uint32
	Peer      *peerconn.Conn
	blocks    []block
	semaphore *semaphore.Semaphore
	DataC     chan Data
	// RejectC chan *piece.Block
	resultC chan Result
	closeC  chan struct{}
}

type Data struct {
	Index uint32
	Data  []byte
}

type block struct {
	index     uint32
	size      uint32
	requested bool
	data      []byte
}

type Result struct {
	Peer  *peerconn.Conn
	Bytes []byte
	Error error
}

func New(pe *peerconn.Conn, extID uint8, totalSize uint32, resultC chan Result) *InfoDownloader {
	numBlocks := totalSize / blockSize
	mod := totalSize % blockSize
	if mod != 0 {
		numBlocks++
	}
	blocks := make([]block, numBlocks)
	for i := range blocks {
		blocks[i] = block{
			index: uint32(i),
			size:  blockSize,
		}
	}
	if mod != 0 && len(blocks) > 0 {
		blocks[len(blocks)-1].size = mod
	}
	return &InfoDownloader{
		extID:     extID,
		totalSize: totalSize,
		Peer:      pe,
		blocks:    blocks,
		semaphore: semaphore.New(maxQueuedBlocks),
		DataC:     make(chan Data),
		// RejectC: make(chan *piece.Block),
		resultC: resultC,
		closeC:  make(chan struct{}),
	}
}

func (d *InfoDownloader) Close() {
	close(d.closeC)
}

func (d *InfoDownloader) Run() {
	result := Result{
		Peer: d.Peer,
	}
	defer func() {
		select {
		case d.resultC <- result:
		case <-d.closeC:
		}
	}()
	d.semaphore.Start()
	for {
		select {
		case <-d.semaphore.Ready:
			b := d.nextBlock()
			if b == nil {
				d.semaphore.Stop()
				break
			}
			msg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: d.extID,
				Payload: peerprotocol.ExtensionMetadataMessage{
					Type:  peerprotocol.ExtensionMetadataMessageTypeRequest,
					Piece: b.index,
				},
			}
			d.Peer.SendMessage(msg)
		case msg := <-d.DataC:
			if msg.Index >= uint32(len(d.blocks)) {
				result.Error = fmt.Errorf("peer sent invalid index for metadata message: %q", msg.Index)
				return
			}
			b := &d.blocks[msg.Index]
			if uint32(len(msg.Data)) != b.size {
				result.Error = fmt.Errorf("peer sent invalid size for metadata message: %q", len(msg.Data))
				return
			}
			if b.requested && b.data == nil {
				d.semaphore.Signal(1)
			}
			b.data = msg.Data
			if d.allDone() {
				result.Bytes = d.assembleBlocks().Bytes()
				return
			}
		// TODO handle rejects in info downloader
		// case blk := <-d.RejectC:
		// 	b := d.blocks[blk.Index]
		// 	if !b.requested {
		// 		d.Peer.Close()
		// 		d.ErrC <- errors.New("received invalid reject message")
		// 		return
		// 	}
		// 	d.blocks[blk.Index].requested = false
		case <-d.closeC:
			return
		}
	}
}

func (d *InfoDownloader) nextBlock() *block {
	for i := range d.blocks {
		if !d.blocks[i].requested {
			d.blocks[i].requested = true
			return &d.blocks[i]
		}
	}
	return nil
}

func (d *InfoDownloader) allDone() bool {
	for i := range d.blocks {
		if d.blocks[i].data == nil {
			return false
		}
	}
	return true
}

func (d *InfoDownloader) assembleBlocks() *bytes.Buffer {
	buf := bytes.NewBuffer(make([]byte, 0, d.totalSize))
	for i := range d.blocks {
		buf.Write(d.blocks[i].data)
	}
	return buf
}
<commit_msg>remove semaphore usage<commit_after>package infodownloader

import (
	"bytes"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/peerconn"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
)

const blockSize = 16 * 1024

const maxQueuedBlocks = 10

// InfoDownloader downloads all blocks of a piece from a peer.
type InfoDownloader struct {
	extID     uint8
	totalSize uint32
	limiter   chan struct{}
	Peer      *peerconn.Conn
	blocks    []block
	DataC     chan Data
	// RejectC chan *piece.Block
	resultC chan Result
	closeC  chan struct{}
}

type Data struct {
	Index uint32
	Data  []byte
}

type block struct {
	index     uint32
	size      uint32
	requested bool
	data      []byte
}

type Result struct {
	Peer  *peerconn.Conn
	Bytes []byte
	Error error
}

func New(pe *peerconn.Conn, extID uint8, totalSize uint32, resultC chan Result) *InfoDownloader {
	numBlocks := totalSize / blockSize
	mod := totalSize % blockSize
	if mod != 0 {
		numBlocks++
	}
	blocks := make([]block, numBlocks)
	for i := range blocks {
		blocks[i] = block{
			index: uint32(i),
			size:  blockSize,
		}
	}
	if mod != 0 && len(blocks) > 0 {
		blocks[len(blocks)-1].size = mod
	}
	return &InfoDownloader{
		extID:     extID,
		totalSize: totalSize,
		Peer:      pe,
		blocks:    blocks,
		limiter:   make(chan struct{}, maxQueuedBlocks),
		DataC:     make(chan Data),
		// RejectC: make(chan *piece.Block),
		resultC: resultC,
		closeC:  make(chan struct{}),
	}
}

func (d *InfoDownloader) Close() {
	close(d.closeC)
}

func (d *InfoDownloader) Run() {
	result := Result{
		Peer: d.Peer,
	}
	defer func() {
		select {
		case d.resultC <- result:
		case <-d.closeC:
		}
	}()
	for {
		select {
		case d.limiter <- struct{}{}:
			b := d.nextBlock()
			if b == nil {
				d.limiter = nil
				break
			}
			msg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: d.extID,
				Payload: peerprotocol.ExtensionMetadataMessage{
					Type:  peerprotocol.ExtensionMetadataMessageTypeRequest,
					Piece: b.index,
				},
			}
			d.Peer.SendMessage(msg)
		case msg := <-d.DataC:
			if msg.Index >= uint32(len(d.blocks)) {
				result.Error = fmt.Errorf("peer sent invalid index for metadata message: %q", msg.Index)
				return
			}
			b := &d.blocks[msg.Index]
			if uint32(len(msg.Data)) != b.size {
				result.Error = fmt.Errorf("peer sent invalid size for metadata message: %q", len(msg.Data))
				return
			}
			if b.requested && b.data == nil {
				if d.limiter == nil {
					d.limiter = make(chan struct{}, maxQueuedBlocks)
				}
			}
			b.data = msg.Data
			if d.allDone() {
				result.Bytes = d.assembleBlocks().Bytes()
				return
			}
		// TODO handle rejects in info downloader
		// case blk := <-d.RejectC:
		// 	b := d.blocks[blk.Index]
		// 	if !b.requested {
		// 		d.Peer.Close()
		// 		d.ErrC <- errors.New("received invalid reject message")
		// 		return
		// 	}
		// 	d.blocks[blk.Index].requested = false
		case <-d.closeC:
			return
		}
	}
}

func (d *InfoDownloader) nextBlock() *block {
	for i := range d.blocks {
		if !d.blocks[i].requested {
			d.blocks[i].requested = true
			return &d.blocks[i]
		}
	}
	return nil
}

func (d *InfoDownloader) allDone() bool {
	for i := range d.blocks {
		if d.blocks[i].data == nil {
			return false
		}
	}
	return true
}

func (d *InfoDownloader) assembleBlocks() *bytes.Buffer {
	buf := bytes.NewBuffer(make([]byte, 0, d.totalSize))
	for i := range d.blocks {
		buf.Write(d.blocks[i].data)
	}
	return buf
}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 fieldmanager

import (
	"fmt"
	"reflect"
	"time"

	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apiserver/pkg/endpoints/handlers/fieldmanager/internal"
	"k8s.io/klog/v2"
	"sigs.k8s.io/structured-merge-diff/v4/fieldpath"
	"sigs.k8s.io/structured-merge-diff/v4/merge"
)

// DefaultMaxUpdateManagers defines the default maximum retained number of managedFields entries from updates
// if the number of update managers exceeds this, the oldest entries will be merged until the number is below the maximum.
// TODO(jennybuckley): Determine if this is really the best value. Ideally we wouldn't unnecessarily merge too many entries.
const DefaultMaxUpdateManagers int = 10

// DefaultTrackOnCreateProbability defines the default probability that the field management of an object
// starts being tracked from the object's creation, instead of from the first time the object is applied to.
const DefaultTrackOnCreateProbability float32 = 1

var atMostEverySecond = internal.NewAtMostEvery(time.Second)

// Managed groups a fieldpath.ManagedFields together with the timestamps associated with each operation.
type Managed interface {
	// Fields gets the fieldpath.ManagedFields.
	Fields() fieldpath.ManagedFields

	// Times gets the timestamps associated with each operation.
	Times() map[string]*metav1.Time
}

// Manager updates the managed fields and merges applied configurations.
type Manager interface {
	// Update is used when the object has already been merged (non-apply
	// use-case), and simply updates the managed fields in the output
	// object.
	//  * `liveObj` is not mutated by this function
	//  * `newObj` may be mutated by this function
	// Returns the new object with managedFields removed, and the object's new
	// proposed managedFields separately.
	Update(liveObj, newObj runtime.Object, managed Managed, manager string) (runtime.Object, Managed, error)

	// Apply is used when server-side apply is called, as it merges the
	// object and updates the managed fields.
	//  * `liveObj` is not mutated by this function
	//  * `newObj` may be mutated by this function
	// Returns the new object with managedFields removed, and the object's new
	// proposed managedFields separately.
	Apply(liveObj, appliedObj runtime.Object, managed Managed, fieldManager string, force bool) (runtime.Object, Managed, error)
}

// FieldManager updates the managed fields and merge applied
// configurations.
type FieldManager struct {
	fieldManager Manager
	subresource  string
}

// NewFieldManager creates a new FieldManager that decodes, manages, then re-encodes managedFields
// on update and apply requests.
func NewFieldManager(f Manager, subresource string) *FieldManager {
	return &FieldManager{fieldManager: f, subresource: subresource}
}

// NewDefaultFieldManager creates a new FieldManager that merges apply requests
// and update managed fields for other types of requests.
func NewDefaultFieldManager(typeConverter TypeConverter, objectConverter runtime.ObjectConvertor, objectDefaulter runtime.ObjectDefaulter, objectCreater runtime.ObjectCreater, kind schema.GroupVersionKind, hub schema.GroupVersion, subresource string, resetFields map[fieldpath.APIVersion]*fieldpath.Set) (*FieldManager, error) {
	f, err := NewStructuredMergeManager(typeConverter, objectConverter, objectDefaulter, kind.GroupVersion(), hub, resetFields)
	if err != nil {
		return nil, fmt.Errorf("failed to create field manager: %v", err)
	}
	return newDefaultFieldManager(f, typeConverter, objectConverter, objectCreater, kind, subresource), nil
}

// NewDefaultCRDFieldManager creates a new FieldManager specifically for
// CRDs. This allows for the possibility of fields which are not defined
// in models, as well as having no models defined at all.
func NewDefaultCRDFieldManager(typeConverter TypeConverter, objectConverter runtime.ObjectConvertor, objectDefaulter runtime.ObjectDefaulter, objectCreater runtime.ObjectCreater, kind schema.GroupVersionKind, hub schema.GroupVersion, subresource string, resetFields map[fieldpath.APIVersion]*fieldpath.Set) (_ *FieldManager, err error) {
	f, err := NewCRDStructuredMergeManager(typeConverter, objectConverter, objectDefaulter, kind.GroupVersion(), hub, resetFields)
	if err != nil {
		return nil, fmt.Errorf("failed to create field manager: %v", err)
	}
	return newDefaultFieldManager(f, typeConverter, objectConverter, objectCreater, kind, subresource), nil
}

// newDefaultFieldManager is a helper function which wraps a Manager with certain default logic.
func newDefaultFieldManager(f Manager, typeConverter TypeConverter, objectConverter runtime.ObjectConvertor, objectCreater runtime.ObjectCreater, kind schema.GroupVersionKind, subresource string) *FieldManager {
	return NewFieldManager(
		NewLastAppliedUpdater(
			NewLastAppliedManager(
				NewProbabilisticSkipNonAppliedManager(
					NewCapManagersManager(
						NewBuildManagerInfoManager(
							NewManagedFieldsUpdater(
								NewStripMetaManager(f),
							), kind.GroupVersion(), subresource,
						), DefaultMaxUpdateManagers,
					), objectCreater, kind, DefaultTrackOnCreateProbability,
				), typeConverter, objectConverter, kind.GroupVersion()),
		), subresource,
	)
}

// DecodeManagedFields converts ManagedFields from the wire format (api format)
// to the format used by sigs.k8s.io/structured-merge-diff
func DecodeManagedFields(encodedManagedFields []metav1.ManagedFieldsEntry) (Managed, error) {
	return internal.DecodeManagedFields(encodedManagedFields)
}

func decodeLiveOrNew(liveObj, newObj runtime.Object, ignoreManagedFieldsFromRequestObject bool) (Managed, error) {
	liveAccessor, err := meta.Accessor(liveObj)
	if err != nil {
		return nil, err
	}

	// We take the managedFields of the live object in case the request tries to
	// manually set managedFields via a subresource.
	if ignoreManagedFieldsFromRequestObject {
		return emptyManagedFieldsOnErr(DecodeManagedFields(liveAccessor.GetManagedFields()))
	}

	// If the object doesn't have metadata, we should just return without trying to
	// set the managedFields at all, so creates/updates/patches will work normally.
	newAccessor, err := meta.Accessor(newObj)
	if err != nil {
		return nil, err
	}

	if isResetManagedFields(newAccessor.GetManagedFields()) {
		return internal.NewEmptyManaged(), nil
	}

	// If the managed field is empty or we failed to decode it,
	// let's try the live object. This is to prevent clients who
	// don't understand managedFields from deleting it accidentally.
	managed, err := DecodeManagedFields(newAccessor.GetManagedFields())
	if err != nil || len(managed.Fields()) == 0 {
		return emptyManagedFieldsOnErr(DecodeManagedFields(liveAccessor.GetManagedFields()))
	}
	return managed, nil
}

func emptyManagedFieldsOnErr(managed Managed, err error) (Managed, error) {
	if err != nil {
		return internal.NewEmptyManaged(), nil
	}
	return managed, nil
}

// Update is used when the object has already been merged (non-apply
// use-case), and simply updates the managed fields in the output
// object.
func (f *FieldManager) Update(liveObj, newObj runtime.Object, manager string) (object runtime.Object, err error) {
	// First try to decode the managed fields provided in the update,
	// This is necessary to allow directly updating managed fields.
	isSubresource := f.subresource != ""
	managed, err := decodeLiveOrNew(liveObj, newObj, isSubresource)
	if err != nil {
		return newObj, nil
	}

	internal.RemoveObjectManagedFields(newObj)

	if object, managed, err = f.fieldManager.Update(liveObj, newObj, managed, manager); err != nil {
		return nil, err
	}

	if err = internal.EncodeObjectManagedFields(object, managed); err != nil {
		return nil, fmt.Errorf("failed to encode managed fields: %v", err)
	}

	return object, nil
}

// UpdateNoErrors is the same as Update, but it will not return
// errors. If an error happens, the object is returned with
// managedFields cleared.
func (f *FieldManager) UpdateNoErrors(liveObj, newObj runtime.Object, manager string) runtime.Object {
	obj, err := f.Update(liveObj, newObj, manager)
	if err != nil {
		atMostEverySecond.Do(func() {
			ns, name := "unknown", "unknown"
			accessor, err := meta.Accessor(newObj)
			if err == nil {
				ns = accessor.GetNamespace()
				name = accessor.GetName()
			}

			klog.ErrorS(err, "[SHOULD NOT HAPPEN] failed to update managedFields", "VersionKind",
				newObj.GetObjectKind().GroupVersionKind(), "namespace", ns, "name", name)
		})
		// Explicitly remove managedFields on failure, so that
		// we can't have garbage in it.
		internal.RemoveObjectManagedFields(newObj)
		return newObj
	}
	return obj
}

// Returns true if the managedFields indicate that the user is trying to
// reset the managedFields, i.e. if the list is non-nil but empty, or if
// the list has one empty item.
func isResetManagedFields(managedFields []metav1.ManagedFieldsEntry) bool {
	if len(managedFields) == 0 {
		return managedFields != nil
	}

	if len(managedFields) == 1 {
		return reflect.DeepEqual(managedFields[0], metav1.ManagedFieldsEntry{})
	}

	return false
}

// Apply is used when server-side apply is called, as it merges the
// object and updates the managed fields.
func (f *FieldManager) Apply(liveObj, appliedObj runtime.Object, manager string, force bool) (object runtime.Object, err error) {
	// If the object doesn't have metadata, apply isn't allowed.
	accessor, err := meta.Accessor(liveObj)
	if err != nil {
		return nil, fmt.Errorf("couldn't get accessor: %v", err)
	}

	// Decode the managed fields in the live object, since it isn't allowed in the patch.
	managed, err := DecodeManagedFields(accessor.GetManagedFields())
	if err != nil {
		return nil, fmt.Errorf("failed to decode managed fields: %v", err)
	}

	object, managed, err = f.fieldManager.Apply(liveObj, appliedObj, managed, manager, force)
	if err != nil {
		if conflicts, ok := err.(merge.Conflicts); ok {
			return nil, internal.NewConflictError(conflicts)
		}
		return nil, err
	}

	if err = internal.EncodeObjectManagedFields(object, managed); err != nil {
		return nil, fmt.Errorf("failed to encode managed fields: %v", err)
	}

	return object, nil
}
<commit_msg>expose FieldsToSet and SetToFields<commit_after>/*
Copyright 2018 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 fieldmanager

import (
	"fmt"
	"reflect"
	"time"

	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apiserver/pkg/endpoints/handlers/fieldmanager/internal"
	"k8s.io/klog/v2"
	"sigs.k8s.io/structured-merge-diff/v4/fieldpath"
	"sigs.k8s.io/structured-merge-diff/v4/merge"
)

// DefaultMaxUpdateManagers defines the default maximum retained number of managedFields entries from updates
// if the number of update managers exceeds this, the oldest entries will be merged until the number is below the maximum.
// TODO(jennybuckley): Determine if this is really the best value. Ideally we wouldn't unnecessarily merge too many entries.
const DefaultMaxUpdateManagers int = 10

// DefaultTrackOnCreateProbability defines the default probability that the field management of an object
// starts being tracked from the object's creation, instead of from the first time the object is applied to.
const DefaultTrackOnCreateProbability float32 = 1

var atMostEverySecond = internal.NewAtMostEvery(time.Second)

// Managed groups a fieldpath.ManagedFields together with the timestamps associated with each operation.
type Managed interface {
	// Fields gets the fieldpath.ManagedFields.
	Fields() fieldpath.ManagedFields

	// Times gets the timestamps associated with each operation.
	Times() map[string]*metav1.Time
}

// Manager updates the managed fields and merges applied configurations.
type Manager interface {
	// Update is used when the object has already been merged (non-apply
	// use-case), and simply updates the managed fields in the output
	// object.
	//  * `liveObj` is not mutated by this function
	//  * `newObj` may be mutated by this function
	// Returns the new object with managedFields removed, and the object's new
	// proposed managedFields separately.
	Update(liveObj, newObj runtime.Object, managed Managed, manager string) (runtime.Object, Managed, error)

	// Apply is used when server-side apply is called, as it merges the
	// object and updates the managed fields.
	//  * `liveObj` is not mutated by this function
	//  * `newObj` may be mutated by this function
	// Returns the new object with managedFields removed, and the object's new
	// proposed managedFields separately.
	Apply(liveObj, appliedObj runtime.Object, managed Managed, fieldManager string, force bool) (runtime.Object, Managed, error)
}

// FieldManager updates the managed fields and merge applied
// configurations.
type FieldManager struct {
	fieldManager Manager
	subresource  string
}

// NewFieldManager creates a new FieldManager that decodes, manages, then re-encodes managedFields
// on update and apply requests.
func NewFieldManager(f Manager, subresource string) *FieldManager {
	return &FieldManager{fieldManager: f, subresource: subresource}
}

// NewDefaultFieldManager creates a new FieldManager that merges apply requests
// and update managed fields for other types of requests.
func NewDefaultFieldManager(typeConverter TypeConverter, objectConverter runtime.ObjectConvertor, objectDefaulter runtime.ObjectDefaulter, objectCreater runtime.ObjectCreater, kind schema.GroupVersionKind, hub schema.GroupVersion, subresource string, resetFields map[fieldpath.APIVersion]*fieldpath.Set) (*FieldManager, error) {
	f, err := NewStructuredMergeManager(typeConverter, objectConverter, objectDefaulter, kind.GroupVersion(), hub, resetFields)
	if err != nil {
		return nil, fmt.Errorf("failed to create field manager: %v", err)
	}
	return newDefaultFieldManager(f, typeConverter, objectConverter, objectCreater, kind, subresource), nil
}

// NewDefaultCRDFieldManager creates a new FieldManager specifically for
// CRDs. This allows for the possibility of fields which are not defined
// in models, as well as having no models defined at all.
func NewDefaultCRDFieldManager(typeConverter TypeConverter, objectConverter runtime.ObjectConvertor, objectDefaulter runtime.ObjectDefaulter, objectCreater runtime.ObjectCreater, kind schema.GroupVersionKind, hub schema.GroupVersion, subresource string, resetFields map[fieldpath.APIVersion]*fieldpath.Set) (_ *FieldManager, err error) {
	f, err := NewCRDStructuredMergeManager(typeConverter, objectConverter, objectDefaulter, kind.GroupVersion(), hub, resetFields)
	if err != nil {
		return nil, fmt.Errorf("failed to create field manager: %v", err)
	}
	return newDefaultFieldManager(f, typeConverter, objectConverter, objectCreater, kind, subresource), nil
}

// newDefaultFieldManager is a helper function which wraps a Manager with certain default logic.
func newDefaultFieldManager(f Manager, typeConverter TypeConverter, objectConverter runtime.ObjectConvertor, objectCreater runtime.ObjectCreater, kind schema.GroupVersionKind, subresource string) *FieldManager {
	return NewFieldManager(
		NewLastAppliedUpdater(
			NewLastAppliedManager(
				NewProbabilisticSkipNonAppliedManager(
					NewCapManagersManager(
						NewBuildManagerInfoManager(
							NewManagedFieldsUpdater(
								NewStripMetaManager(f),
							), kind.GroupVersion(), subresource,
						), DefaultMaxUpdateManagers,
					), objectCreater, kind, DefaultTrackOnCreateProbability,
				), typeConverter, objectConverter, kind.GroupVersion()),
		), subresource,
	)
}

func FieldsToSet(f metav1.FieldsV1) (s fieldpath.Set, err error) {
	return internal.FieldsToSet(f)
}

// SetToFields creates a trie of fields from an input set of paths
func SetToFields(s fieldpath.Set) (f metav1.FieldsV1, err error) {
	return internal.SetToFields(s)
}

// DecodeManagedFields converts ManagedFields from the wire format (api format)
// to the format used by sigs.k8s.io/structured-merge-diff
func DecodeManagedFields(encodedManagedFields []metav1.ManagedFieldsEntry) (Managed, error) {
	return internal.DecodeManagedFields(encodedManagedFields)
}

func decodeLiveOrNew(liveObj, newObj runtime.Object, ignoreManagedFieldsFromRequestObject bool) (Managed, error) {
	liveAccessor, err := meta.Accessor(liveObj)
	if err != nil {
		return nil, err
	}

	// We take the managedFields of the live object in case the request tries to
	// manually set managedFields via a subresource.
	if ignoreManagedFieldsFromRequestObject {
		return emptyManagedFieldsOnErr(DecodeManagedFields(liveAccessor.GetManagedFields()))
	}

	// If the object doesn't have metadata, we should just return without trying to
	// set the managedFields at all, so creates/updates/patches will work normally.
	newAccessor, err := meta.Accessor(newObj)
	if err != nil {
		return nil, err
	}

	if isResetManagedFields(newAccessor.GetManagedFields()) {
		return internal.NewEmptyManaged(), nil
	}

	// If the managed field is empty or we failed to decode it,
	// let's try the live object. This is to prevent clients who
	// don't understand managedFields from deleting it accidentally.
	managed, err := DecodeManagedFields(newAccessor.GetManagedFields())
	if err != nil || len(managed.Fields()) == 0 {
		return emptyManagedFieldsOnErr(DecodeManagedFields(liveAccessor.GetManagedFields()))
	}
	return managed, nil
}

func emptyManagedFieldsOnErr(managed Managed, err error) (Managed, error) {
	if err != nil {
		return internal.NewEmptyManaged(), nil
	}
	return managed, nil
}

// Update is used when the object has already been merged (non-apply
// use-case), and simply updates the managed fields in the output
// object.
func (f *FieldManager) Update(liveObj, newObj runtime.Object, manager string) (object runtime.Object, err error) {
	// First try to decode the managed fields provided in the update,
	// This is necessary to allow directly updating managed fields.
	isSubresource := f.subresource != ""
	managed, err := decodeLiveOrNew(liveObj, newObj, isSubresource)
	if err != nil {
		return newObj, nil
	}

	internal.RemoveObjectManagedFields(newObj)

	if object, managed, err = f.fieldManager.Update(liveObj, newObj, managed, manager); err != nil {
		return nil, err
	}

	if err = internal.EncodeObjectManagedFields(object, managed); err != nil {
		return nil, fmt.Errorf("failed to encode managed fields: %v", err)
	}

	return object, nil
}

// UpdateNoErrors is the same as Update, but it will not return
// errors. If an error happens, the object is returned with
// managedFields cleared.
func (f *FieldManager) UpdateNoErrors(liveObj, newObj runtime.Object, manager string) runtime.Object {
	obj, err := f.Update(liveObj, newObj, manager)
	if err != nil {
		atMostEverySecond.Do(func() {
			ns, name := "unknown", "unknown"
			accessor, err := meta.Accessor(newObj)
			if err == nil {
				ns = accessor.GetNamespace()
				name = accessor.GetName()
			}

			klog.ErrorS(err, "[SHOULD NOT HAPPEN] failed to update managedFields", "VersionKind",
				newObj.GetObjectKind().GroupVersionKind(), "namespace", ns, "name", name)
		})
		// Explicitly remove managedFields on failure, so that
		// we can't have garbage in it.
		internal.RemoveObjectManagedFields(newObj)
		return newObj
	}
	return obj
}

// Returns true if the managedFields indicate that the user is trying to
// reset the managedFields, i.e. if the list is non-nil but empty, or if
// the list has one empty item.
func isResetManagedFields(managedFields []metav1.ManagedFieldsEntry) bool {
	if len(managedFields) == 0 {
		return managedFields != nil
	}

	if len(managedFields) == 1 {
		return reflect.DeepEqual(managedFields[0], metav1.ManagedFieldsEntry{})
	}

	return false
}

// Apply is used when server-side apply is called, as it merges the
// object and updates the managed fields.
func (f *FieldManager) Apply(liveObj, appliedObj runtime.Object, manager string, force bool) (object runtime.Object, err error) {
	// If the object doesn't have metadata, apply isn't allowed.
	accessor, err := meta.Accessor(liveObj)
	if err != nil {
		return nil, fmt.Errorf("couldn't get accessor: %v", err)
	}

	// Decode the managed fields in the live object, since it isn't allowed in the patch.
	managed, err := DecodeManagedFields(accessor.GetManagedFields())
	if err != nil {
		return nil, fmt.Errorf("failed to decode managed fields: %v", err)
	}

	object, managed, err = f.fieldManager.Apply(liveObj, appliedObj, managed, manager, force)
	if err != nil {
		if conflicts, ok := err.(merge.Conflicts); ok {
			return nil, internal.NewConflictError(conflicts)
		}
		return nil, err
	}

	if err = internal.EncodeObjectManagedFields(object, managed); err != nil {
		return nil, fmt.Errorf("failed to encode managed fields: %v", err)
	}

	return object, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 util

import (
	"fmt"
	"strings"
	"time"

	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	pkgruntime "k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/record"
	"k8s.io/kubernetes/pkg/api"
	kubeclientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
)

// Type of the operation that can be executed in Federated.
type FederatedOperationType string

const (
	OperationTypeAdd    = "add"
	OperationTypeUpdate = "update"
	OperationTypeDelete = "delete"
)

// FederatedOperation definition contains type (add/update/delete) and the object itself.
type FederatedOperation struct {
	Type        FederatedOperationType
	ClusterName string
	Obj         pkgruntime.Object
	Key         string
}

// A helper that executes the given set of updates on federation, in parallel.
type FederatedUpdater interface {
	// Executes the given set of operations.
	Update([]FederatedOperation) error
}

// A function that executes some operation using the passed client and object.
type FederatedOperationHandler func(kubeclientset.Interface, pkgruntime.Object) error

type federatedUpdaterImpl struct {
	federation FederationView

	kind string

	timeout time.Duration

	eventRecorder record.EventRecorder

	addFunction    FederatedOperationHandler
	updateFunction FederatedOperationHandler
	deleteFunction FederatedOperationHandler
}

func NewFederatedUpdater(federation FederationView, kind string, timeout time.Duration, recorder record.EventRecorder, add, update, del FederatedOperationHandler) FederatedUpdater {
	return &federatedUpdaterImpl{
		federation:     federation,
		kind:           kind,
		timeout:        timeout,
		eventRecorder:  recorder,
		addFunction:    add,
		updateFunction: update,
		deleteFunction: del,
	}
}

func (fu *federatedUpdaterImpl) recordEvent(obj runtime.Object, eventType, eventVerb string, args ...interface{}) {
	messageFmt := eventVerb + " %s %q in cluster %s"
	fu.eventRecorder.Eventf(obj, api.EventTypeNormal, eventType, messageFmt, args...)
}

// Update executes the given set of operations within the timeout specified for
// the instance. Timeout is best-effort. There is no guarantee that the
// underlying operations are stopped when it is reached. However the function
// will return after the timeout with a non-nil error.
func (fu *federatedUpdaterImpl) Update(ops []FederatedOperation) error {
	done := make(chan error, len(ops))
	for _, op := range ops {
		go func(op FederatedOperation) {
			clusterName := op.ClusterName

			// TODO: Ensure that the clientset has reasonable timeout.
			clientset, err := fu.federation.GetClientsetForCluster(clusterName)
			if err != nil {
				done <- err
				return
			}

			eventArgs := []interface{}{fu.kind, op.Key, clusterName}
			baseEventType := fmt.Sprintf("%s", op.Type)
			eventType := fmt.Sprintf("%sInCluster", strings.Title(baseEventType))

			switch op.Type {
			case OperationTypeAdd:
				// TODO s+OperationTypeAdd+OperationTypeCreate+
				baseEventType = "create"
				eventType := "CreateInCluster"

				fu.recordEvent(op.Obj, eventType, "Creating", eventArgs...)
				err = fu.addFunction(clientset, op.Obj)
			case OperationTypeUpdate:
				fu.recordEvent(op.Obj, eventType, "Updating", eventArgs...)
				err = fu.updateFunction(clientset, op.Obj)
			case OperationTypeDelete:
				fu.recordEvent(op.Obj, eventType, "Deleting", eventArgs...)
				err = fu.deleteFunction(clientset, op.Obj)
				// IsNotFound error is fine since that means the object is deleted already.
				if err != nil && !errors.IsNotFound(err) {
					err = nil
				}
			}

			if err != nil {
				eventType := eventType + "Failed"
				messageFmt := "Failed to " + baseEventType + " %s %q in cluster %s: %v"
				eventArgs = append(eventArgs, err)
				fu.eventRecorder.Eventf(op.Obj, api.EventTypeWarning, eventType, messageFmt, eventArgs...)
			}

			done <- err
		}(op)
	}
	start := time.Now()
	for i := 0; i < len(ops); i++ {
		now := time.Now()
		if !now.Before(start.Add(fu.timeout)) {
			return fmt.Errorf("failed to finish all operations in %v", fu.timeout)
		}
		select {
		case err := <-done:
			if err != nil {
				return err
			}
		case <-time.After(start.Add(fu.timeout).Sub(now)):
			return fmt.Errorf("failed to finish all operations in %v", fu.timeout)
		}
	}
	// All operations finished in time.
	return nil
}
<commit_msg>fed: Fix bad logic of deletion error handling for federated updater<commit_after>/*
Copyright 2016 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 util

import (
	"fmt"
	"strings"
	"time"

	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	pkgruntime "k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/record"
	"k8s.io/kubernetes/pkg/api"
	kubeclientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
)

// Type of the operation that can be executed in Federated.
type FederatedOperationType string

const (
	OperationTypeAdd    = "add"
	OperationTypeUpdate = "update"
	OperationTypeDelete = "delete"
)

// FederatedOperation definition contains type (add/update/delete) and the object itself.
type FederatedOperation struct {
	Type        FederatedOperationType
	ClusterName string
	Obj         pkgruntime.Object
	Key         string
}

// A helper that executes the given set of updates on federation, in parallel.
type FederatedUpdater interface {
	// Executes the given set of operations.
	Update([]FederatedOperation) error
}

// A function that executes some operation using the passed client and object.
type FederatedOperationHandler func(kubeclientset.Interface, pkgruntime.Object) error

type federatedUpdaterImpl struct {
	federation FederationView

	kind string

	timeout time.Duration

	eventRecorder record.EventRecorder

	addFunction    FederatedOperationHandler
	updateFunction FederatedOperationHandler
	deleteFunction FederatedOperationHandler
}

func NewFederatedUpdater(federation FederationView, kind string, timeout time.Duration, recorder record.EventRecorder, add, update, del FederatedOperationHandler) FederatedUpdater {
	return &federatedUpdaterImpl{
		federation:     federation,
		kind:           kind,
		timeout:        timeout,
		eventRecorder:  recorder,
		addFunction:    add,
		updateFunction: update,
		deleteFunction: del,
	}
}

func (fu *federatedUpdaterImpl) recordEvent(obj runtime.Object, eventType, eventVerb string, args ...interface{}) {
	messageFmt := eventVerb + " %s %q in cluster %s"
	fu.eventRecorder.Eventf(obj, api.EventTypeNormal, eventType, messageFmt, args...)
}

// Update executes the given set of operations within the timeout specified for
// the instance. Timeout is best-effort. There is no guarantee that the
// underlying operations are stopped when it is reached. However the function
// will return after the timeout with a non-nil error.
func (fu *federatedUpdaterImpl) Update(ops []FederatedOperation) error {
	done := make(chan error, len(ops))
	for _, op := range ops {
		go func(op FederatedOperation) {
			clusterName := op.ClusterName

			// TODO: Ensure that the clientset has reasonable timeout.
			clientset, err := fu.federation.GetClientsetForCluster(clusterName)
			if err != nil {
				done <- err
				return
			}

			eventArgs := []interface{}{fu.kind, op.Key, clusterName}
			baseEventType := fmt.Sprintf("%s", op.Type)
			eventType := fmt.Sprintf("%sInCluster", strings.Title(baseEventType))

			switch op.Type {
			case OperationTypeAdd:
				// TODO s+OperationTypeAdd+OperationTypeCreate+
				baseEventType = "create"
				eventType := "CreateInCluster"

				fu.recordEvent(op.Obj, eventType, "Creating", eventArgs...)
				err = fu.addFunction(clientset, op.Obj)
			case OperationTypeUpdate:
				fu.recordEvent(op.Obj, eventType, "Updating", eventArgs...)
				err = fu.updateFunction(clientset, op.Obj)
			case OperationTypeDelete:
				fu.recordEvent(op.Obj, eventType, "Deleting", eventArgs...)
				err = fu.deleteFunction(clientset, op.Obj)
				// IsNotFound error is fine since that means the object is deleted already.
				if errors.IsNotFound(err) {
					err = nil
				}
			}

			if err != nil {
				eventType := eventType + "Failed"
				messageFmt := "Failed to " + baseEventType + " %s %q in cluster %s: %v"
				eventArgs = append(eventArgs, err)
				fu.eventRecorder.Eventf(op.Obj, api.EventTypeWarning, eventType, messageFmt, eventArgs...)
			}

			done <- err
		}(op)
	}
	start := time.Now()
	for i := 0; i < len(ops); i++ {
		now := time.Now()
		if !now.Before(start.Add(fu.timeout)) {
			return fmt.Errorf("failed to finish all operations in %v", fu.timeout)
		}
		select {
		case err := <-done:
			if err != nil {
				return err
			}
		case <-time.After(start.Add(fu.timeout).Sub(now)):
			return fmt.Errorf("failed to finish all operations in %v", fu.timeout)
		}
	}
	// All operations finished in time.
	return nil
}
<|endoftext|>