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

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

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

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

package machine

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/docker/machine/libmachine"
	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/cert"
	"github.com/docker/machine/libmachine/check"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/drivers/plugin"
	"github.com/docker/machine/libmachine/drivers/plugin/localbinary"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/host"
	"github.com/docker/machine/libmachine/mcnutils"
	"github.com/docker/machine/libmachine/persist"
	lib_provision "github.com/docker/machine/libmachine/provision"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
	"github.com/docker/machine/libmachine/swarm"
	"github.com/docker/machine/libmachine/version"
	"github.com/pkg/errors"
	"k8s.io/minikube/pkg/minikube/command"
	"k8s.io/minikube/pkg/minikube/driver"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/localpath"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/registry"
	"k8s.io/minikube/pkg/minikube/sshutil"
	"k8s.io/minikube/pkg/provision"
)

// NewRPCClient gets a new client.
func NewRPCClient(storePath, certsDir string) libmachine.API {
	c := libmachine.NewClient(storePath, certsDir)
	c.SSHClientType = ssh.Native
	return c
}

// NewAPIClient gets a new client.
func NewAPIClient() (libmachine.API, error) {
	storePath := localpath.MiniPath()
	certsDir := localpath.MakeMiniPath("certs")

	return &LocalClient{
		certsDir:     certsDir,
		storePath:    storePath,
		Filestore:    persist.NewFilestore(storePath, certsDir, certsDir),
		legacyClient: NewRPCClient(storePath, certsDir),
	}, nil
}

// LocalClient is a non-RPC implementation
// of the libmachine API
type LocalClient struct {
	certsDir  string
	storePath string
	*persist.Filestore
	legacyClient libmachine.API
}

// NewHost creates a new Host
func (api *LocalClient) NewHost(drvName string, rawDriver []byte) (*host.Host, error) {
	def := registry.Driver(drvName)
	if def.Empty() {
		return nil, fmt.Errorf("driver %q does not exist", drvName)
	}
	if def.Init == nil {
		return api.legacyClient.NewHost(drvName, rawDriver)
	}
	d := def.Init()
	err := json.Unmarshal(rawDriver, d)
	if err != nil {
		return nil, errors.Wrapf(err, "Error getting driver %s", string(rawDriver))
	}

	return &host.Host{
		ConfigVersion: version.ConfigVersion,
		Name:          d.GetMachineName(),
		Driver:        d,
		DriverName:    d.DriverName(),
		HostOptions: &host.Options{
			AuthOptions: &auth.Options{
				CertDir:          api.certsDir,
				CaCertPath:       filepath.Join(api.certsDir, "ca.pem"),
				CaPrivateKeyPath: filepath.Join(api.certsDir, "ca-key.pem"),
				ClientCertPath:   filepath.Join(api.certsDir, "cert.pem"),
				ClientKeyPath:    filepath.Join(api.certsDir, "key.pem"),
				ServerCertPath:   filepath.Join(api.GetMachinesDir(), "server.pem"),
				ServerKeyPath:    filepath.Join(api.GetMachinesDir(), "server-key.pem"),
			},
			EngineOptions: &engine.Options{
				StorageDriver: "overlay2",
				TLSVerify:     false,
			},
			SwarmOptions: &swarm.Options{},
		},
	}, nil
}

// Load a new client, creating driver
func (api *LocalClient) Load(name string) (*host.Host, error) {
	h, err := api.Filestore.Load(name)
	if err != nil {
		return nil, errors.Wrapf(err, "filestore %q", name)
	}

	def := registry.Driver(h.DriverName)
	if def.Empty() {
		return nil, fmt.Errorf("driver %q does not exist", h.DriverName)
	}
	if def.Init == nil {
		return api.legacyClient.Load(name)
	}
	h.Driver = def.Init()
	return h, json.Unmarshal(h.RawDriver, h.Driver)
}

// Close closes the client
func (api *LocalClient) Close() error {
	if api.legacyClient != nil {
		return api.legacyClient.Close()
	}
	return nil
}

// CommandRunner returns best available command runner for this host
func CommandRunner(h *host.Host) (command.Runner, error) {
	if h.DriverName == driver.Mock {
		return &command.FakeCommandRunner{}, nil
	}
	if driver.BareMetal(h.Driver.DriverName()) {
		return command.NewExecRunner(), nil
	}

	if driver.IsKIC(h.Driver.DriverName()) {
		return command.NewKICRunner(h.Name, h.Driver.DriverName()), nil
	}
	client, err := sshutil.NewSSHClient(h.Driver)
	if err != nil {
		return nil, errors.Wrap(err, "getting ssh client for bootstrapper")
	}
	return command.NewSSHRunner(client), nil
}

// Create creates the host
func (api *LocalClient) Create(h *host.Host) error {
	def := registry.Driver(h.DriverName)
	if def.Empty() {
		return fmt.Errorf("driver %q does not exist", h.DriverName)
	}
	if def.Init == nil {
		return api.legacyClient.Create(h)
	}

	steps := []struct {
		name string
		f    func() error
	}{
		{
			"bootstrapping certificates",
			func() error { return cert.BootstrapCertificates(h.AuthOptions()) },
		},
		{
			"precreate",
			h.Driver.PreCreateCheck,
		},
		{
			"saving",
			func() error {
				return api.Save(h)
			},
		},
		{
			"creating",
			h.Driver.Create,
		},
		{
			"waiting",
			func() error {
				if driver.BareMetal(h.Driver.DriverName()) {
					return nil
				}
				return mcnutils.WaitFor(drivers.MachineInState(h.Driver, state.Running))
			},
		},
		{
			"provisioning",
			func() error {
				if driver.BareMetal(h.Driver.DriverName()) {
					return nil
				}
				var pv lib_provision.Provisioner
				if driver.IsKIC(h.Driver.DriverName()) {
					pv = provision.NewUbuntuProvisioner(h.Driver)
				} else {
					pv = provision.NewBuildrootProvisioner(h.Driver)
				}
				return pv.Provision(*h.HostOptions.SwarmOptions, *h.HostOptions.AuthOptions, *h.HostOptions.EngineOptions)
			},
		},
	}

	for _, step := range steps {
		if err := step.f(); err != nil {
			return errors.Wrap(err, step.name)
		}
	}

	return nil
}

// StartDriver starts the driver
func StartDriver() {
	cert.SetCertGenerator(&CertGenerator{})
	check.DefaultConnChecker = &ConnChecker{}
	if os.Getenv(localbinary.PluginEnvKey) == localbinary.PluginEnvVal {
		registerDriver(os.Getenv(localbinary.PluginEnvDriverName))
	}

	localbinary.CurrentBinaryIsDockerMachine = true
}

// ConnChecker can check the connection
type ConnChecker struct {
}

// Check checks the connection
func (cc *ConnChecker) Check(h *host.Host, swarm bool) (string, *auth.Options, error) {
	authOptions := h.AuthOptions()
	dockerHost, err := h.Driver.GetURL()
	if err != nil {
		return "", &auth.Options{}, err
	}
	return dockerHost, authOptions, nil
}

// CertGenerator is used to override the default machine CertGenerator with a longer timeout.
type CertGenerator struct {
	cert.X509CertGenerator
}

// ValidateCertificate is a reimplementation of the default generator with a longer timeout.
func (cg *CertGenerator) ValidateCertificate(addr string, authOptions *auth.Options) (bool, error) {
	tlsConfig, err := cg.ReadTLSConfig(addr, authOptions)
	if err != nil {
		return false, err
	}

	dialer := &net.Dialer{
		Timeout: time.Second * 40,
	}

	_, err = tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	if err != nil {
		return false, err
	}

	return true, nil
}

func registerDriver(drvName string) {
	def := registry.Driver(drvName)
	if def.Empty() {
		exit.UsageT("unsupported or missing driver: {{.name}}", out.V{"name": drvName})
	}
	plugin.RegisterDriver(def.Init())
}
<commit_msg>verify agaibn<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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

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

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

package machine

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/docker/machine/libmachine"
	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/cert"
	"github.com/docker/machine/libmachine/check"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/drivers/plugin"
	"github.com/docker/machine/libmachine/drivers/plugin/localbinary"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/host"
	"github.com/docker/machine/libmachine/mcnutils"
	"github.com/docker/machine/libmachine/persist"
	lib_provision "github.com/docker/machine/libmachine/provision"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
	"github.com/docker/machine/libmachine/swarm"
	"github.com/docker/machine/libmachine/version"
	"github.com/pkg/errors"
	"k8s.io/minikube/pkg/minikube/command"
	"k8s.io/minikube/pkg/minikube/driver"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/localpath"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/registry"
	"k8s.io/minikube/pkg/minikube/sshutil"
	"k8s.io/minikube/pkg/provision"
)

// NewRPCClient gets a new client.
func NewRPCClient(storePath, certsDir string) libmachine.API {
	c := libmachine.NewClient(storePath, certsDir)
	c.SSHClientType = ssh.Native
	return c
}

// NewAPIClient gets a new client.
func NewAPIClient() (libmachine.API, error) {
	storePath := localpath.MiniPath()
	certsDir := localpath.MakeMiniPath("certs")

	return &LocalClient{
		certsDir:     certsDir,
		storePath:    storePath,
		Filestore:    persist.NewFilestore(storePath, certsDir, certsDir),
		legacyClient: NewRPCClient(storePath, certsDir),
	}, nil
}

// LocalClient is a non-RPC implementation
// of the libmachine API
type LocalClient struct {
	certsDir  string
	storePath string
	*persist.Filestore
	legacyClient libmachine.API
}

// NewHost creates a new Host
func (api *LocalClient) NewHost(drvName string, rawDriver []byte) (*host.Host, error) {
	def := registry.Driver(drvName)
	if def.Empty() {
		return nil, fmt.Errorf("driver %q does not exist", drvName)
	}
	if def.Init == nil {
		return api.legacyClient.NewHost(drvName, rawDriver)
	}
	d := def.Init()
	err := json.Unmarshal(rawDriver, d)
	if err != nil {
		return nil, errors.Wrapf(err, "Error getting driver %s", string(rawDriver))
	}

	return &host.Host{
		ConfigVersion: version.ConfigVersion,
		Name:          d.GetMachineName(),
		Driver:        d,
		DriverName:    d.DriverName(),
		HostOptions: &host.Options{
			AuthOptions: &auth.Options{
				CertDir:          api.certsDir,
				CaCertPath:       filepath.Join(api.certsDir, "ca.pem"),
				CaPrivateKeyPath: filepath.Join(api.certsDir, "ca-key.pem"),
				ClientCertPath:   filepath.Join(api.certsDir, "cert.pem"),
				ClientKeyPath:    filepath.Join(api.certsDir, "key.pem"),
				ServerCertPath:   filepath.Join(api.GetMachinesDir(), "server.pem"),
				ServerKeyPath:    filepath.Join(api.GetMachinesDir(), "server-key.pem"),
			},
			EngineOptions: &engine.Options{
				StorageDriver: "overlay2",
				TLSVerify:     true,
			},
			SwarmOptions: &swarm.Options{},
		},
	}, nil
}

// Load a new client, creating driver
func (api *LocalClient) Load(name string) (*host.Host, error) {
	h, err := api.Filestore.Load(name)
	if err != nil {
		return nil, errors.Wrapf(err, "filestore %q", name)
	}

	def := registry.Driver(h.DriverName)
	if def.Empty() {
		return nil, fmt.Errorf("driver %q does not exist", h.DriverName)
	}
	if def.Init == nil {
		return api.legacyClient.Load(name)
	}
	h.Driver = def.Init()
	return h, json.Unmarshal(h.RawDriver, h.Driver)
}

// Close closes the client
func (api *LocalClient) Close() error {
	if api.legacyClient != nil {
		return api.legacyClient.Close()
	}
	return nil
}

// CommandRunner returns best available command runner for this host
func CommandRunner(h *host.Host) (command.Runner, error) {
	if h.DriverName == driver.Mock {
		return &command.FakeCommandRunner{}, nil
	}
	if driver.BareMetal(h.Driver.DriverName()) {
		return command.NewExecRunner(), nil
	}

	if driver.IsKIC(h.Driver.DriverName()) {
		return command.NewKICRunner(h.Name, h.Driver.DriverName()), nil
	}
	client, err := sshutil.NewSSHClient(h.Driver)
	if err != nil {
		return nil, errors.Wrap(err, "getting ssh client for bootstrapper")
	}
	return command.NewSSHRunner(client), nil
}

// Create creates the host
func (api *LocalClient) Create(h *host.Host) error {
	def := registry.Driver(h.DriverName)
	if def.Empty() {
		return fmt.Errorf("driver %q does not exist", h.DriverName)
	}
	if def.Init == nil {
		return api.legacyClient.Create(h)
	}

	steps := []struct {
		name string
		f    func() error
	}{
		{
			"bootstrapping certificates",
			func() error { return cert.BootstrapCertificates(h.AuthOptions()) },
		},
		{
			"precreate",
			h.Driver.PreCreateCheck,
		},
		{
			"saving",
			func() error {
				return api.Save(h)
			},
		},
		{
			"creating",
			h.Driver.Create,
		},
		{
			"waiting",
			func() error {
				if driver.BareMetal(h.Driver.DriverName()) {
					return nil
				}
				return mcnutils.WaitFor(drivers.MachineInState(h.Driver, state.Running))
			},
		},
		{
			"provisioning",
			func() error {
				if driver.BareMetal(h.Driver.DriverName()) {
					return nil
				}
				var pv lib_provision.Provisioner
				if driver.IsKIC(h.Driver.DriverName()) {
					pv = provision.NewUbuntuProvisioner(h.Driver)
				} else {
					pv = provision.NewBuildrootProvisioner(h.Driver)
				}
				return pv.Provision(*h.HostOptions.SwarmOptions, *h.HostOptions.AuthOptions, *h.HostOptions.EngineOptions)
			},
		},
	}

	for _, step := range steps {
		if err := step.f(); err != nil {
			return errors.Wrap(err, step.name)
		}
	}

	return nil
}

// StartDriver starts the driver
func StartDriver() {
	cert.SetCertGenerator(&CertGenerator{})
	check.DefaultConnChecker = &ConnChecker{}
	if os.Getenv(localbinary.PluginEnvKey) == localbinary.PluginEnvVal {
		registerDriver(os.Getenv(localbinary.PluginEnvDriverName))
	}

	localbinary.CurrentBinaryIsDockerMachine = true
}

// ConnChecker can check the connection
type ConnChecker struct {
}

// Check checks the connection
func (cc *ConnChecker) Check(h *host.Host, swarm bool) (string, *auth.Options, error) {
	authOptions := h.AuthOptions()
	dockerHost, err := h.Driver.GetURL()
	if err != nil {
		return "", &auth.Options{}, err
	}
	return dockerHost, authOptions, nil
}

// CertGenerator is used to override the default machine CertGenerator with a longer timeout.
type CertGenerator struct {
	cert.X509CertGenerator
}

// ValidateCertificate is a reimplementation of the default generator with a longer timeout.
func (cg *CertGenerator) ValidateCertificate(addr string, authOptions *auth.Options) (bool, error) {
	tlsConfig, err := cg.ReadTLSConfig(addr, authOptions)
	if err != nil {
		return false, err
	}

	dialer := &net.Dialer{
		Timeout: time.Second * 40,
	}

	_, err = tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	if err != nil {
		return false, err
	}

	return true, nil
}

func registerDriver(drvName string) {
	def := registry.Driver(drvName)
	if def.Empty() {
		exit.UsageT("unsupported or missing driver: {{.name}}", out.V{"name": drvName})
	}
	plugin.RegisterDriver(def.Init())
}
<|endoftext|>


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

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

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

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

// Package sysinit provides an abstraction over init systems like systemctl
package sysinit

import (
	"bytes"
	"context"
	"fmt"
	"html/template"
	"os/exec"
	"path"
	"time"

	"github.com/pkg/errors"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/minikube/vmpath"
)

var restartWrapper = `#!/bin/bash
# Wrapper script to emulate systemd restart on non-systemd systems
readonly UNIT_PATH=$1

while true; do
  if [[ -f "${UNIT_PATH}" ]]; then
	eval $(egrep "^ExecStart=" "${UNIT_PATH}" | cut -d"=" -f2-)
  fi
  sleep 1
done
`

var initScriptTmpl = template.Must(template.New("initScript").Parse(`#!/bin/bash
# OpenRC init script shim for systemd units
readonly NAME="{{.Name}}"
readonly RESTART_WRAPPER="{{.Wrapper}}"
readonly UNIT_PATH="{{.Unit}}"
readonly PID_PATH="/var/run/${NAME}.pid"

function start() {
    start-stop-daemon --oknodo --pidfile "${PID_PATH}" --background --start --make-pid --exec "${RESTART_WRAPPER}" "${UNIT_PATH}"
}

function stop() {
	if [[ -f "${PID_PATH}" ]]; then
		pkill -P "$(cat ${PID_PATH})"
	fi
	start-stop-daemon --oknodo --pidfile "${PID_PATH}" --stop
}

case "$1" in
    start)
        start
		;;
    stop)
        stop
		;;
    restart)
        stop
        start
		;;
    status)
        start-stop-daemon --pidfile "${PID_PATH}" --status
		;;
	*)
	    echo "Usage: {{.Name}} {start|stop|restart|status}"
		exit 1
		;;
esac
`))

// OpenRC is a service manager for OpenRC-like init systems
type OpenRC struct {
	r Runner
}

// Name returns the name of the init system
func (s *OpenRC) Name() string {
	return "OpenRC"
}

// Active checks if a service is running
func (s *OpenRC) Active(svc string) bool {
	_, err := s.r.RunCmd(exec.Command("sudo", "service", svc, "status"))
	return err == nil
}

// Start starts a service idempotently
func (s *OpenRC) Start(svc string) error {
	if s.Active(svc) {
		return nil
	}
	ctx, cb := context.WithTimeout(context.Background(), 5*time.Second)
	defer cb()

	rr, err := s.r.RunCmd(exec.CommandContext(ctx, "sudo", "service", svc, "start"))
	klog.Infof("start output: %s", rr.Output())
	return err
}

// Disable does nothing
func (s *OpenRC) Disable(svc string) error {
	return nil
}

// DisableNow not implemented for openRC
func (s *OpenRC) DisableNow(svc string) error {
	// supposed to do disable + stop
	// disable does nothing for OpenRC, so just Stop here
	return s.Stop(svc)
}

// Mask does nothing
func (s *OpenRC) Mask(svc string) error {
	return nil
}

// Enable does nothing
func (s *OpenRC) Enable(svc string) error {
	return nil
}

// EnableNow  not implemented for openRC
func (s *OpenRC) EnableNow(svc string) error {
	return fmt.Errorf("enable now is not implemented for OpenRC! PRs to fix are welcomed")
}

// Unmask does nothing
func (s *OpenRC) Unmask(svc string) error {
	return nil
}

// Restart restarts a service
func (s *OpenRC) Restart(svc string) error {
	rr, err := s.r.RunCmd(exec.Command("sudo", "service", svc, "restart"))
	klog.Infof("restart output: %s", rr.Output())
	return err
}

// Reload reloads a service
// currently only used by our docker-env that doesn't need openrc implementation
func (s *OpenRC) Reload(svc string) error {
	return fmt.Errorf("reload is not implemented for OpenRC yet ! Please implement if needed")
}

// Stop stops a service
func (s *OpenRC) Stop(svc string) error {
	rr, err := s.r.RunCmd(exec.Command("sudo", "service", svc, "stop"))
	klog.Infof("stop output: %s", rr.Output())
	return err
}

// ForceStop stops a service with prejuidice
func (s *OpenRC) ForceStop(svc string) error {
	return s.Stop(svc)
}

// GenerateInitShim generates any additional init files required for this service
func (s *OpenRC) GenerateInitShim(svc string, binary string, unit string) ([]assets.CopyableFile, error) {
	restartWrapperPath := path.Join(vmpath.GuestPersistentDir, "openrc-restart-wrapper.sh")

	opts := struct {
		Binary  string
		Wrapper string
		Name    string
		Unit    string
	}{
		Name:    svc,
		Binary:  binary,
		Wrapper: restartWrapperPath,
		Unit:    unit,
	}

	var b bytes.Buffer
	if err := initScriptTmpl.Execute(&b, opts); err != nil {
		return nil, errors.Wrap(err, "template execute")
	}

	files := []assets.CopyableFile{
		assets.NewMemoryAssetTarget([]byte(restartWrapper), restartWrapperPath, "0755"),
		assets.NewMemoryAssetTarget(b.Bytes(), path.Join("/etc/init.d/", svc), "0755"),
	}

	return files, nil
}
<commit_msg>Implement EnableNow for OpenRC<commit_after>/*
Copyright 2019 The Kubernetes Authors All rights reserved.

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

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

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

// Package sysinit provides an abstraction over init systems like systemctl
package sysinit

import (
	"bytes"
	"context"
	"fmt"
	"html/template"
	"os/exec"
	"path"
	"time"

	"github.com/pkg/errors"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/minikube/vmpath"
)

var restartWrapper = `#!/bin/bash
# Wrapper script to emulate systemd restart on non-systemd systems
readonly UNIT_PATH=$1

while true; do
  if [[ -f "${UNIT_PATH}" ]]; then
	eval $(egrep "^ExecStart=" "${UNIT_PATH}" | cut -d"=" -f2-)
  fi
  sleep 1
done
`

var initScriptTmpl = template.Must(template.New("initScript").Parse(`#!/bin/bash
# OpenRC init script shim for systemd units
readonly NAME="{{.Name}}"
readonly RESTART_WRAPPER="{{.Wrapper}}"
readonly UNIT_PATH="{{.Unit}}"
readonly PID_PATH="/var/run/${NAME}.pid"

function start() {
    start-stop-daemon --oknodo --pidfile "${PID_PATH}" --background --start --make-pid --exec "${RESTART_WRAPPER}" "${UNIT_PATH}"
}

function stop() {
	if [[ -f "${PID_PATH}" ]]; then
		pkill -P "$(cat ${PID_PATH})"
	fi
	start-stop-daemon --oknodo --pidfile "${PID_PATH}" --stop
}

case "$1" in
    start)
        start
		;;
    stop)
        stop
		;;
    restart)
        stop
        start
		;;
    status)
        start-stop-daemon --pidfile "${PID_PATH}" --status
		;;
	*)
	    echo "Usage: {{.Name}} {start|stop|restart|status}"
		exit 1
		;;
esac
`))

// OpenRC is a service manager for OpenRC-like init systems
type OpenRC struct {
	r Runner
}

// Name returns the name of the init system
func (s *OpenRC) Name() string {
	return "OpenRC"
}

// Active checks if a service is running
func (s *OpenRC) Active(svc string) bool {
	_, err := s.r.RunCmd(exec.Command("sudo", "service", svc, "status"))
	return err == nil
}

// Start starts a service idempotently
func (s *OpenRC) Start(svc string) error {
	if s.Active(svc) {
		return nil
	}
	ctx, cb := context.WithTimeout(context.Background(), 5*time.Second)
	defer cb()

	rr, err := s.r.RunCmd(exec.CommandContext(ctx, "sudo", "service", svc, "start"))
	klog.Infof("start output: %s", rr.Output())
	return err
}

// Disable does nothing
func (s *OpenRC) Disable(svc string) error {
	return nil
}

// DisableNow not implemented for openRC
func (s *OpenRC) DisableNow(svc string) error {
	// supposed to do disable + stop
	// disable does nothing for OpenRC, so just Stop here
	return s.Stop(svc)
}

// Mask does nothing
func (s *OpenRC) Mask(svc string) error {
	return nil
}

// Enable does nothing
func (s *OpenRC) Enable(svc string) error {
	return nil
}

// EnableNow  not implemented for openRC
func (s *OpenRC) EnableNow(svc string) error {
	// supposed to do enable + start
	// enable does nothing for OpenRC, so just Start here
	return s.Start(svc)
}

// Unmask does nothing
func (s *OpenRC) Unmask(svc string) error {
	return nil
}

// Restart restarts a service
func (s *OpenRC) Restart(svc string) error {
	rr, err := s.r.RunCmd(exec.Command("sudo", "service", svc, "restart"))
	klog.Infof("restart output: %s", rr.Output())
	return err
}

// Reload reloads a service
// currently only used by our docker-env that doesn't need openrc implementation
func (s *OpenRC) Reload(svc string) error {
	return fmt.Errorf("reload is not implemented for OpenRC yet ! Please implement if needed")
}

// Stop stops a service
func (s *OpenRC) Stop(svc string) error {
	rr, err := s.r.RunCmd(exec.Command("sudo", "service", svc, "stop"))
	klog.Infof("stop output: %s", rr.Output())
	return err
}

// ForceStop stops a service with prejuidice
func (s *OpenRC) ForceStop(svc string) error {
	return s.Stop(svc)
}

// GenerateInitShim generates any additional init files required for this service
func (s *OpenRC) GenerateInitShim(svc string, binary string, unit string) ([]assets.CopyableFile, error) {
	restartWrapperPath := path.Join(vmpath.GuestPersistentDir, "openrc-restart-wrapper.sh")

	opts := struct {
		Binary  string
		Wrapper string
		Name    string
		Unit    string
	}{
		Name:    svc,
		Binary:  binary,
		Wrapper: restartWrapperPath,
		Unit:    unit,
	}

	var b bytes.Buffer
	if err := initScriptTmpl.Execute(&b, opts); err != nil {
		return nil, errors.Wrap(err, "template execute")
	}

	files := []assets.CopyableFile{
		assets.NewMemoryAssetTarget([]byte(restartWrapper), restartWrapperPath, "0755"),
		assets.NewMemoryAssetTarget(b.Bytes(), path.Join("/etc/init.d/", svc), "0755"),
	}

	return files, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach 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.
//
// Author: Ben Darnell

package storage

import (
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/config"
	"github.com/cockroachdb/cockroach/pkg/gossip"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
)

const (
	// replicateQueueMaxSize is the max size of the replicate queue.
	replicateQueueMaxSize = 100

	// replicateQueueTimerDuration is the duration between replication of queued
	// replicas.
	replicateQueueTimerDuration = 0 // zero duration to process replication greedily
)

var (
	// minLeaseTransferInterval controls how frequently leases can be transferred
	// for rebalancing. It does not prevent transferring leases in order to allow
	// a replica to be removed from a range. The value should be some reasonable
	// fraction of the store descriptor gossip interval which currently
	minLeaseTransferInterval = gossip.GossipStoresInterval / 5
)

// replicateQueue manages a queue of replicas which may need to add an
// additional replica to their range.
type replicateQueue struct {
	*baseQueue
	allocator         Allocator
	clock             *hlc.Clock
	updateChan        chan struct{}
	lastLeaseTransfer atomic.Value // read and written by scanner & queue goroutines
}

// newReplicateQueue returns a new instance of replicateQueue.
func newReplicateQueue(
	store *Store, g *gossip.Gossip, allocator Allocator, clock *hlc.Clock, options AllocatorOptions,
) *replicateQueue {
	rq := &replicateQueue{
		allocator:  allocator,
		clock:      clock,
		updateChan: make(chan struct{}, 1),
	}
	rq.baseQueue = newBaseQueue(
		"replicate", rq, store, g,
		queueConfig{
			maxSize:              replicateQueueMaxSize,
			needsLease:           true,
			acceptsUnsplitRanges: store.TestingKnobs().ReplicateQueueAcceptsUnsplit,
			successes:            store.metrics.ReplicateQueueSuccesses,
			failures:             store.metrics.ReplicateQueueFailures,
			pending:              store.metrics.ReplicateQueuePending,
			processingNanos:      store.metrics.ReplicateQueueProcessingNanos,
			purgatory:            store.metrics.ReplicateQueuePurgatory,
		},
	)

	if g != nil { // gossip is nil for some unittests
		// Register a gossip callback to signal queue that replicas in
		// purgatory might be retried due to new store gossip.
		g.RegisterCallback(gossip.MakePrefixPattern(gossip.KeyStorePrefix), func(_ string, _ roachpb.Value) {
			select {
			case rq.updateChan <- struct{}{}:
			default:
			}
		})
	}

	return rq
}

func (rq *replicateQueue) shouldQueue(
	ctx context.Context, now hlc.Timestamp, repl *Replica, sysCfg config.SystemConfig,
) (shouldQ bool, priority float64) {
	if !repl.store.splitQueue.Disabled() && repl.needsSplitBySize() {
		// If the range exceeds the split threshold, let that finish first.
		// Ranges must fit in memory on both sender and receiver nodes while
		// being replicated. This supplements the check provided by
		// acceptsUnsplitRanges, which looks at zone config boundaries rather
		// than data size.
		//
		// This check is ignored if the split queue is disabled, since in that
		// case, the split will never come.
		return
	}

	// Find the zone config for this range.
	desc := repl.Desc()
	zone, err := sysCfg.GetZoneConfigForKey(desc.StartKey)
	if err != nil {
		log.Error(ctx, err)
		return
	}

	action, priority := rq.allocator.ComputeAction(zone, desc)
	if action != AllocatorNoop {
		if log.V(2) {
			log.Infof(ctx, "%s repair needed (%s), enqueuing", repl, action)
		}
		return true, priority
	}

	// If we hold the lease, check to see if we should transfer it.
	var leaseStoreID roachpb.StoreID
	if lease, _ := repl.getLease(); lease != nil && lease.Covers(now) {
		leaseStoreID = lease.Replica.StoreID
		if rq.canTransferLease() &&
			rq.allocator.ShouldTransferLease(zone.Constraints, leaseStoreID, desc.RangeID) {
			if log.V(2) {
				log.Infof(ctx, "%s lease transfer needed, enqueuing", repl)
			}
			return true, 0
		}
	}

	// Check for a rebalancing opportunity. Note that leaseStoreID will be 0 if
	// the range doesn't currently have a lease which will allow the current
	// replica to be considered a rebalancing source.
	target, err := rq.allocator.RebalanceTarget(
		zone.Constraints,
		desc.Replicas,
		leaseStoreID,
		desc.RangeID,
	)
	if err != nil {
		log.ErrEventf(ctx, "rebalance target failed: %s", err)
		return false, 0
	}
	if log.V(2) {
		if target != nil {
			log.Infof(ctx, "%s rebalance target found, enqueuing", repl)
		} else {
			log.Infof(ctx, "%s no rebalance target found, not enqueuing", repl)
		}
	}
	return target != nil, 0
}

func (rq *replicateQueue) process(
	ctx context.Context, now hlc.Timestamp, repl *Replica, sysCfg config.SystemConfig,
) error {
	desc := repl.Desc()
	// Find the zone config for this range.
	zone, err := sysCfg.GetZoneConfigForKey(desc.StartKey)
	if err != nil {
		return err
	}
	action, _ := rq.allocator.ComputeAction(zone, desc)

	// Avoid taking action if the range has too many dead replicas to make
	// quorum.
	deadReplicas := rq.allocator.storePool.deadReplicas(desc.RangeID, desc.Replicas)
	quorum := computeQuorum(len(desc.Replicas))
	liveReplicaCount := len(desc.Replicas) - len(deadReplicas)
	if liveReplicaCount < quorum {
		return errors.Errorf("range requires a replication change, but lacks a quorum of live nodes.")
	}

	switch action {
	case AllocatorAdd:
		log.Event(ctx, "adding a new replica")
		newStore, err := rq.allocator.AllocateTarget(
			zone.Constraints,
			desc.Replicas,
			desc.RangeID,
			true,
		)
		if err != nil {
			return err
		}
		newReplica := roachpb.ReplicaDescriptor{
			NodeID:  newStore.Node.NodeID,
			StoreID: newStore.StoreID,
		}

		log.VEventf(ctx, 1, "adding replica to %+v due to under-replication", newReplica)
		if err := rq.addReplica(ctx, repl, newReplica, desc); err != nil {
			return err
		}
	case AllocatorRemove:
		log.Event(ctx, "removing a replica")
		// If the lease holder (our local store) is an overfull store (in terms of
		// leases) allow transferring the lease away.
		leaseHolderStoreID := repl.store.StoreID()
		if rq.allocator.ShouldTransferLease(zone.Constraints, leaseHolderStoreID, desc.RangeID) {
			leaseHolderStoreID = 0
		}
		removeReplica, err := rq.allocator.RemoveTarget(
			zone.Constraints,
			desc.Replicas,
			leaseHolderStoreID,
		)
		if err != nil {
			return err
		}
		if removeReplica.StoreID == repl.store.StoreID() {
			// The local replica was selected as the removal target, but that replica
			// is the leaseholder, so transfer the lease instead. We don't check that
			// the current store has too many leases in this case under the
			// assumption that replica balance is a greater concern. Also note that
			// AllocatorRemove action takes preference over AllocatorNoop
			// (rebalancing) which is where lease transfer would otherwise occur. We
			// need to be able to transfer leases in AllocatorRemove in order to get
			// out of situations where this store is overfull and yet holds all the
			// leases.
			candidates := filterBehindReplicas(repl.RaftStatus(), desc.Replicas)
			target := rq.allocator.TransferLeaseTarget(
				zone.Constraints, candidates, repl.store.StoreID(), desc.RangeID,
				false /* checkTransferLeaseSource */)
			if target != (roachpb.ReplicaDescriptor{}) {
				log.VEventf(ctx, 1, "transferring lease to s%d", target.StoreID)
				if err := repl.AdminTransferLease(target.StoreID); err != nil {
					return errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, target.StoreID)
				}
				rq.lastLeaseTransfer.Store(timeutil.Now())
				// Do not requeue as we transferred our lease away.
				return nil
			}
		} else {
			log.VEventf(ctx, 1, "removing replica %+v due to over-replication", removeReplica)
			if err := rq.removeReplica(ctx, repl, removeReplica, desc); err != nil {
				return err
			}
		}
	case AllocatorRemoveDead:
		log.Event(ctx, "removing a dead replica")
		if len(deadReplicas) == 0 {
			if log.V(1) {
				log.Warningf(ctx, "Range of replica %s was identified as having dead replicas, but no dead replicas were found.", repl)
			}
			break
		}
		deadReplica := deadReplicas[0]
		log.VEventf(ctx, 1, "removing dead replica %+v from store", deadReplica)
		if err := repl.ChangeReplicas(ctx, roachpb.REMOVE_REPLICA, deadReplica, desc); err != nil {
			return err
		}
	case AllocatorNoop:
		// The Noop case will result if this replica was queued in order to
		// rebalance. Attempt to find a rebalancing target.
		log.Event(ctx, "considering a rebalance")

		if rq.canTransferLease() {
			// We require the lease in order to process replicas, so
			// repl.store.StoreID() corresponds to the lease-holder's store ID.
			candidates := filterBehindReplicas(repl.RaftStatus(), desc.Replicas)
			target := rq.allocator.TransferLeaseTarget(
				zone.Constraints, candidates, repl.store.StoreID(), desc.RangeID,
				true /* checkTransferLeaseSource */)
			if target != (roachpb.ReplicaDescriptor{}) {
				log.VEventf(ctx, 1, "transferring lease to s%d", target.StoreID)
				if err := repl.AdminTransferLease(target.StoreID); err != nil {
					return errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, target.StoreID)
				}
				rq.lastLeaseTransfer.Store(timeutil.Now())
				// Do not requeue as we transferred our lease away.
				return nil
			}
		}

		rebalanceStore, err := rq.allocator.RebalanceTarget(
			zone.Constraints,
			desc.Replicas,
			repl.store.StoreID(),
			desc.RangeID,
		)
		if err != nil {
			log.ErrEventf(ctx, "rebalance target failed %s", err)
			return nil
		}
		if rebalanceStore == nil {
			log.VEventf(ctx, 1, "no suitable rebalance target")
			// No action was necessary and no rebalance target was found. Return
			// without re-queuing this replica.
			return nil
		}
		rebalanceReplica := roachpb.ReplicaDescriptor{
			NodeID:  rebalanceStore.Node.NodeID,
			StoreID: rebalanceStore.StoreID,
		}
		log.VEventf(ctx, 1, "rebalancing to %+v", rebalanceReplica)
		if err := rq.addReplica(ctx, repl, rebalanceReplica, desc); err != nil {
			return err
		}
	}

	// Enqueue this replica again to see if there are more changes to be made.
	rq.MaybeAdd(repl, rq.clock.Now())
	return nil
}

func (rq *replicateQueue) addReplica(
	ctx context.Context,
	repl *Replica,
	repDesc roachpb.ReplicaDescriptor,
	desc *roachpb.RangeDescriptor,
) error {
	err := repl.ChangeReplicas(ctx, roachpb.ADD_REPLICA, repDesc, desc)
	if IsPreemptiveSnapshotError(err) {
		// If the ChangeReplicas failed because the preemptive snapshot failed, we
		// log the error but then return success indicating we should retry the
		// operation. The most likely causes of the preemptive snapshot failing are
		// a declined reservation or the remote node being unavailable. In either
		// case we don't want to wait another scanner cycle before reconsidering
		// the range.
		log.Info(ctx, err)
		return nil
	}
	return err
}

func (rq *replicateQueue) removeReplica(
	ctx context.Context,
	repl *Replica,
	repDesc roachpb.ReplicaDescriptor,
	desc *roachpb.RangeDescriptor,
) error {
	return repl.ChangeReplicas(ctx, roachpb.REMOVE_REPLICA, repDesc, desc)
}

func (rq *replicateQueue) canTransferLease() bool {
	if lastLeaseTransfer := rq.lastLeaseTransfer.Load(); lastLeaseTransfer != nil {
		return timeutil.Since(lastLeaseTransfer.(time.Time)) > minLeaseTransferInterval
	}
	return true
}

func (*replicateQueue) timer() time.Duration {
	return replicateQueueTimerDuration
}

// purgatoryChan returns the replicate queue's store update channel.
func (rq *replicateQueue) purgatoryChan() <-chan struct{} {
	return rq.updateChan
}
<commit_msg>storage: avoid busy looping when failing rebalancing reservations<commit_after>// Copyright 2015 The Cockroach 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.
//
// Author: Ben Darnell

package storage

import (
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/config"
	"github.com/cockroachdb/cockroach/pkg/gossip"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/retry"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
)

const (
	// replicateQueueMaxSize is the max size of the replicate queue.
	replicateQueueMaxSize = 100

	// replicateQueueTimerDuration is the duration between replication of queued
	// replicas.
	replicateQueueTimerDuration = 0 // zero duration to process replication greedily
)

var (
	// minLeaseTransferInterval controls how frequently leases can be transferred
	// for rebalancing. It does not prevent transferring leases in order to allow
	// a replica to be removed from a range. The value should be some reasonable
	// fraction of the store descriptor gossip interval which currently
	minLeaseTransferInterval = gossip.GossipStoresInterval / 5
)

// replicateQueue manages a queue of replicas which may need to add an
// additional replica to their range.
type replicateQueue struct {
	*baseQueue
	allocator         Allocator
	clock             *hlc.Clock
	updateChan        chan struct{}
	lastLeaseTransfer atomic.Value // read and written by scanner & queue goroutines
}

// newReplicateQueue returns a new instance of replicateQueue.
func newReplicateQueue(
	store *Store, g *gossip.Gossip, allocator Allocator, clock *hlc.Clock, options AllocatorOptions,
) *replicateQueue {
	rq := &replicateQueue{
		allocator:  allocator,
		clock:      clock,
		updateChan: make(chan struct{}, 1),
	}
	rq.baseQueue = newBaseQueue(
		"replicate", rq, store, g,
		queueConfig{
			maxSize:              replicateQueueMaxSize,
			needsLease:           true,
			acceptsUnsplitRanges: store.TestingKnobs().ReplicateQueueAcceptsUnsplit,
			successes:            store.metrics.ReplicateQueueSuccesses,
			failures:             store.metrics.ReplicateQueueFailures,
			pending:              store.metrics.ReplicateQueuePending,
			processingNanos:      store.metrics.ReplicateQueueProcessingNanos,
			purgatory:            store.metrics.ReplicateQueuePurgatory,
		},
	)

	if g != nil { // gossip is nil for some unittests
		// Register a gossip callback to signal queue that replicas in
		// purgatory might be retried due to new store gossip.
		g.RegisterCallback(gossip.MakePrefixPattern(gossip.KeyStorePrefix), func(_ string, _ roachpb.Value) {
			select {
			case rq.updateChan <- struct{}{}:
			default:
			}
		})
	}

	return rq
}

func (rq *replicateQueue) shouldQueue(
	ctx context.Context, now hlc.Timestamp, repl *Replica, sysCfg config.SystemConfig,
) (shouldQ bool, priority float64) {
	if !repl.store.splitQueue.Disabled() && repl.needsSplitBySize() {
		// If the range exceeds the split threshold, let that finish first.
		// Ranges must fit in memory on both sender and receiver nodes while
		// being replicated. This supplements the check provided by
		// acceptsUnsplitRanges, which looks at zone config boundaries rather
		// than data size.
		//
		// This check is ignored if the split queue is disabled, since in that
		// case, the split will never come.
		return
	}

	// Find the zone config for this range.
	desc := repl.Desc()
	zone, err := sysCfg.GetZoneConfigForKey(desc.StartKey)
	if err != nil {
		log.Error(ctx, err)
		return
	}

	action, priority := rq.allocator.ComputeAction(zone, desc)
	if action != AllocatorNoop {
		if log.V(2) {
			log.Infof(ctx, "%s repair needed (%s), enqueuing", repl, action)
		}
		return true, priority
	}

	// If we hold the lease, check to see if we should transfer it.
	var leaseStoreID roachpb.StoreID
	if lease, _ := repl.getLease(); lease != nil && lease.Covers(now) {
		leaseStoreID = lease.Replica.StoreID
		if rq.canTransferLease() &&
			rq.allocator.ShouldTransferLease(zone.Constraints, leaseStoreID, desc.RangeID) {
			if log.V(2) {
				log.Infof(ctx, "%s lease transfer needed, enqueuing", repl)
			}
			return true, 0
		}
	}

	// Check for a rebalancing opportunity. Note that leaseStoreID will be 0 if
	// the range doesn't currently have a lease which will allow the current
	// replica to be considered a rebalancing source.
	target, err := rq.allocator.RebalanceTarget(
		zone.Constraints,
		desc.Replicas,
		leaseStoreID,
		desc.RangeID,
	)
	if err != nil {
		log.ErrEventf(ctx, "rebalance target failed: %s", err)
		return false, 0
	}
	if log.V(2) {
		if target != nil {
			log.Infof(ctx, "%s rebalance target found, enqueuing", repl)
		} else {
			log.Infof(ctx, "%s no rebalance target found, not enqueuing", repl)
		}
	}
	return target != nil, 0
}

func (rq *replicateQueue) process(
	ctx context.Context, now hlc.Timestamp, repl *Replica, sysCfg config.SystemConfig,
) error {
	retryOpts := retry.Options{
		InitialBackoff: 50 * time.Millisecond,
		MaxBackoff:     1 * time.Second,
		Multiplier:     2,
		MaxRetries:     5,
	}

	// Use a retry loop in order to backoff in the case of preemptive
	// snapshot errors, usually signalling that a rebalancing
	// reservation could not be made with the selected target.
	for r := retry.StartWithCtx(ctx, retryOpts); r.Next(); {
		if err := rq.processOneChange(ctx, now, repl, sysCfg); err != nil {
			if IsPreemptiveSnapshotError(err) {
				// If ChangeReplicas failed because the preemptive snapshot failed, we
				// log the error but then return success indicating we should retry the
				// operation. The most likely causes of the preemptive snapshot failing are
				// a declined reservation or the remote node being unavailable. In either
				// case we don't want to wait another scanner cycle before reconsidering
				// the range.
				log.Info(ctx, err)
				continue
			}
			return err
		}
		// Enqueue this replica again to see if there are more changes to be made.
		rq.MaybeAdd(repl, rq.clock.Now())
		return nil
	}
	return errors.Errorf("failed to replicate %s after %d retries", repl, retryOpts.MaxRetries)
}

func (rq *replicateQueue) processOneChange(
	ctx context.Context, now hlc.Timestamp, repl *Replica, sysCfg config.SystemConfig,
) error {
	desc := repl.Desc()
	// Find the zone config for this range.
	zone, err := sysCfg.GetZoneConfigForKey(desc.StartKey)
	if err != nil {
		return err
	}
	action, _ := rq.allocator.ComputeAction(zone, desc)

	// Avoid taking action if the range has too many dead replicas to make
	// quorum.
	deadReplicas := rq.allocator.storePool.deadReplicas(desc.RangeID, desc.Replicas)
	quorum := computeQuorum(len(desc.Replicas))
	liveReplicaCount := len(desc.Replicas) - len(deadReplicas)
	if liveReplicaCount < quorum {
		return errors.Errorf("range requires a replication change, but lacks a quorum of live nodes.")
	}

	switch action {
	case AllocatorAdd:
		log.Event(ctx, "adding a new replica")
		newStore, err := rq.allocator.AllocateTarget(
			zone.Constraints,
			desc.Replicas,
			desc.RangeID,
			true,
		)
		if err != nil {
			return err
		}
		newReplica := roachpb.ReplicaDescriptor{
			NodeID:  newStore.Node.NodeID,
			StoreID: newStore.StoreID,
		}

		log.VEventf(ctx, 1, "adding replica to %+v due to under-replication", newReplica)
		if err := rq.addReplica(ctx, repl, newReplica, desc); err != nil {
			return err
		}
	case AllocatorRemove:
		log.Event(ctx, "removing a replica")
		// If the lease holder (our local store) is an overfull store (in terms of
		// leases) allow transferring the lease away.
		leaseHolderStoreID := repl.store.StoreID()
		if rq.allocator.ShouldTransferLease(zone.Constraints, leaseHolderStoreID, desc.RangeID) {
			leaseHolderStoreID = 0
		}
		removeReplica, err := rq.allocator.RemoveTarget(
			zone.Constraints,
			desc.Replicas,
			leaseHolderStoreID,
		)
		if err != nil {
			return err
		}
		if removeReplica.StoreID == repl.store.StoreID() {
			// The local replica was selected as the removal target, but that replica
			// is the leaseholder, so transfer the lease instead. We don't check that
			// the current store has too many leases in this case under the
			// assumption that replica balance is a greater concern. Also note that
			// AllocatorRemove action takes preference over AllocatorNoop
			// (rebalancing) which is where lease transfer would otherwise occur. We
			// need to be able to transfer leases in AllocatorRemove in order to get
			// out of situations where this store is overfull and yet holds all the
			// leases.
			candidates := filterBehindReplicas(repl.RaftStatus(), desc.Replicas)
			target := rq.allocator.TransferLeaseTarget(
				zone.Constraints, candidates, repl.store.StoreID(), desc.RangeID,
				false /* checkTransferLeaseSource */)
			if target != (roachpb.ReplicaDescriptor{}) {
				log.VEventf(ctx, 1, "transferring lease to s%d", target.StoreID)
				if err := repl.AdminTransferLease(target.StoreID); err != nil {
					return errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, target.StoreID)
				}
				rq.lastLeaseTransfer.Store(timeutil.Now())
				// Do not requeue as we transferred our lease away.
				return nil
			}
		} else {
			log.VEventf(ctx, 1, "removing replica %+v due to over-replication", removeReplica)
			if err := rq.removeReplica(ctx, repl, removeReplica, desc); err != nil {
				return err
			}
		}
	case AllocatorRemoveDead:
		log.Event(ctx, "removing a dead replica")
		if len(deadReplicas) == 0 {
			if log.V(1) {
				log.Warningf(ctx, "Range of replica %s was identified as having dead replicas, but no dead replicas were found.", repl)
			}
			break
		}
		deadReplica := deadReplicas[0]
		log.VEventf(ctx, 1, "removing dead replica %+v from store", deadReplica)
		if err := repl.ChangeReplicas(ctx, roachpb.REMOVE_REPLICA, deadReplica, desc); err != nil {
			return err
		}
	case AllocatorNoop:
		// The Noop case will result if this replica was queued in order to
		// rebalance. Attempt to find a rebalancing target.
		log.Event(ctx, "considering a rebalance")

		if rq.canTransferLease() {
			// We require the lease in order to process replicas, so
			// repl.store.StoreID() corresponds to the lease-holder's store ID.
			candidates := filterBehindReplicas(repl.RaftStatus(), desc.Replicas)
			target := rq.allocator.TransferLeaseTarget(
				zone.Constraints, candidates, repl.store.StoreID(), desc.RangeID,
				true /* checkTransferLeaseSource */)
			if target != (roachpb.ReplicaDescriptor{}) {
				log.VEventf(ctx, 1, "transferring lease to s%d", target.StoreID)
				if err := repl.AdminTransferLease(target.StoreID); err != nil {
					return errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, target.StoreID)
				}
				rq.lastLeaseTransfer.Store(timeutil.Now())
				// Do not requeue as we transferred our lease away.
				return nil
			}
		}

		rebalanceStore, err := rq.allocator.RebalanceTarget(
			zone.Constraints,
			desc.Replicas,
			repl.store.StoreID(),
			desc.RangeID,
		)
		if err != nil {
			log.ErrEventf(ctx, "rebalance target failed %s", err)
			return nil
		}
		if rebalanceStore == nil {
			log.VEventf(ctx, 1, "no suitable rebalance target")
			// No action was necessary and no rebalance target was found. Return
			// without re-queuing this replica.
			return nil
		}
		rebalanceReplica := roachpb.ReplicaDescriptor{
			NodeID:  rebalanceStore.Node.NodeID,
			StoreID: rebalanceStore.StoreID,
		}
		log.VEventf(ctx, 1, "rebalancing to %+v", rebalanceReplica)
		if err := rq.addReplica(ctx, repl, rebalanceReplica, desc); err != nil {
			return err
		}
	}

	return nil
}

func (rq *replicateQueue) addReplica(
	ctx context.Context,
	repl *Replica,
	repDesc roachpb.ReplicaDescriptor,
	desc *roachpb.RangeDescriptor,
) error {
	return repl.ChangeReplicas(ctx, roachpb.ADD_REPLICA, repDesc, desc)
}

func (rq *replicateQueue) removeReplica(
	ctx context.Context,
	repl *Replica,
	repDesc roachpb.ReplicaDescriptor,
	desc *roachpb.RangeDescriptor,
) error {
	return repl.ChangeReplicas(ctx, roachpb.REMOVE_REPLICA, repDesc, desc)
}

func (rq *replicateQueue) canTransferLease() bool {
	if lastLeaseTransfer := rq.lastLeaseTransfer.Load(); lastLeaseTransfer != nil {
		return timeutil.Since(lastLeaseTransfer.(time.Time)) > minLeaseTransferInterval
	}
	return true
}

func (*replicateQueue) timer() time.Duration {
	return replicateQueueTimerDuration
}

// purgatoryChan returns the replicate queue's store update channel.
func (rq *replicateQueue) purgatoryChan() <-chan struct{} {
	return rq.updateChan
}
<|endoftext|>


<commit_before>package app

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/go-rancher-metadata/metadata"
	"github.com/urfave/cli"
)

type StringifyError struct {
	message string
}

func (e *StringifyError) Error() string {
	return e.message
}

func IPCommand() cli.Command {
	return cli.Command{
		Name:  "ip",
		Usage: "Get IP information",
		Subcommands: []cli.Command{
			{
				Name:   "stringify",
				Usage:  "Prints a joined list of IPs",
				Action: ipStringifyAction,
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "delimiter",
						Usage: "Delimiter to use between entries",
						Value: ",",
					},
					cli.StringFlag{
						Name:  "prefix",
						Usage: "Prepend Entries with this value",
						Value: "",
					},
					cli.StringFlag{
						Name:  "suffix",
						Usage: "Add this value to the end of each entry.",
						Value: "",
					},
					cli.StringFlag{
						Name:  "source",
						Usage: "Source to lookup IPs. [metadata, dns]",
						Value: "metadata",
					},
					cli.BoolFlag{
						Name:  "use-agent-ips",
						Usage: "Use agent ips instead of rancher ips, only works with metadata source",
					},
					cli.BoolFlag{
						Name:  "use-agent-names",
						Usage: "Use agent name instead of rancher ips, only works with metadata source",
					},
				},
			}, {
				Name:   "myip",
				Usage:  "Prints the containers Rancher managed IP",
				Action: ipMyIpAction,
			},
		},
	}
}

func ipMyIpAction(c *cli.Context) error {
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	selfContainer, err := mdClient.GetSelfContainer()
	if err != nil {
		logrus.Fatalf("Failed to find IP: %v", err)
	}
	fmt.Print(selfContainer.PrimaryIp)
}

func ipStringifyAction(c *cli.Context) error {
	str := ""
	var err error

	if c.String("source") == "dns" {
		str, err = ipStringifyDNS(c)
		if err != nil {
			logrus.Fatalf("Failed to generate string: %v", err)
		}
	} else {
		str, err = ipStringifyMetadata(c)
		if err != nil {
			logrus.Fatalf("Failed to generate string: %v", err)
		}
	}

	fmt.Print(str)
}

func ipStringifyDNS(c *cli.Context) (string, error) {
	if len(c.Args()) <= 0 {
		return "", nil
	}
	ips, err := getDnsContainerIPs(c.Args().First())
	rString := joinString(
		c.String("prefix"),
		c.String("suffix"),
		c.String("delimiter"),
		ips,
	)
	return rString, err
}

func getDnsContainerIPs(host string) ([]string, error) {
	ips := []string{}
	err := &StringifyError{"Could not resolve Host: " + host}

	if testDNSResolves(host) {
		if ipBytes, errIP := net.LookupIP(host); errIP == nil {
			for _, ip := range ipBytes {
				ips = append(ips, ip.String())
			}
			return ips, errIP
		}
	}

	return ips, err
}

func testDNSResolves(host string) bool {
	retVal := false
	ticker := time.NewTicker(time.Millisecond * 500)
	timer := time.NewTimer(time.Second * 60)
	resolves := make(chan bool)

	go func() {
		for range ticker.C {
			if _, err := net.LookupIP(host); err == nil {
				retVal = true
				resolves <- true
			}
		}
	}()

	for {
		select {
		case <-timer.C:
			ticker.Stop()
			return retVal
		case <-resolves:
			ticker.Stop()
			timer.Stop()
			return retVal
		}
	}
}

func ipStringifyMetadata(c *cli.Context) (string, error) {
	split := []string{}
	rString := ""
	var err error

	if len(c.Args()) > 0 {
		split = strings.SplitN(c.Args().First(), "/", 2)
	} else {
		split, err = getSelfStackServiceName()
	}

	getMetaIPMethod := getMetadataContainerIPs
	if c.Bool("use-agent-ips") {
		getMetaIPMethod = getMetadataAgentIPs
	}

	if c.Bool("use-agent-names") {
		getMetaIPMethod = getMetadataAgentNames
	}

	if len(split) == 2 {
		ips, err := getMetaIPMethod(split[0], split[1])
		if err != nil {
			return rString, err
		}
		rString = joinString(
			c.String("prefix"),
			c.String("suffix"),
			c.String("delimiter"),
			ips,
		)
		err = nil
	} else {
		err = &StringifyError{"Not enough arguements supplied. Need stack/service or this container is not part of service"}
	}

	return rString, err
}

func getSelfStackServiceName() ([]string, error) {
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	selfContainer, err := mdClient.GetSelfContainer()
	if err != nil {
		return nil, err
	}

	return []string{selfContainer.StackName, selfContainer.ServiceName}, err
}

func getMetadataContainerIPs(stack string, service string) ([]string, error) {
	rIPs := []string{}
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	containers, err := mdClient.GetServiceContainers(service, stack)
	if err != nil {
		return rIPs, err
	}

	for _, container := range containers {
		rIPs = append(rIPs, container.PrimaryIp)
	}

	return rIPs, nil
}

func getMetadataAgentIPs(stack string, service string) ([]string, error) {
	return getMetadataAgentInfoStrings(stack, service, "AgentIP")
}

func getMetadataAgentNames(stack string, service string) ([]string, error) {
	return getMetadataAgentInfoStrings(stack, service, "Name")
}

func getMetadataAgentInfoStrings(stack, service, property string) ([]string, error) {
	rInfo := []string{}
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	containers, err := mdClient.GetServiceContainers(service, stack)
	if err != nil {
		return rInfo, err
	}

	for _, container := range containers {
		host, err := mdClient.GetHost(container.HostUUID)
		if err != nil {
			return rInfo, err
		}
		rInfo = append(rInfo, getHostInfoProperty(&host, property))
	}

	return rInfo, nil
}

func getHostInfoProperty(host *metadata.Host, property string) string {
	switch {
	case property == "Name":
		return host.Name
	case property == "AgentIP":
		return host.AgentIP
	default:
		return ""
	}
}

func joinString(pfx string, suffix string, delim string, list []string) string {
	intList := []string{}
	for _, item := range list {
		intermediate := pfx + item + suffix
		intList = append(intList, strings.Repeat(intermediate, 1))
	}
	return strings.Join(intList, delim)
}
<commit_msg>added returns<commit_after>package app

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/go-rancher-metadata/metadata"
	"github.com/urfave/cli"
)

type StringifyError struct {
	message string
}

func (e *StringifyError) Error() string {
	return e.message
}

func IPCommand() cli.Command {
	return cli.Command{
		Name:  "ip",
		Usage: "Get IP information",
		Subcommands: []cli.Command{
			{
				Name:   "stringify",
				Usage:  "Prints a joined list of IPs",
				Action: ipStringifyAction,
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "delimiter",
						Usage: "Delimiter to use between entries",
						Value: ",",
					},
					cli.StringFlag{
						Name:  "prefix",
						Usage: "Prepend Entries with this value",
						Value: "",
					},
					cli.StringFlag{
						Name:  "suffix",
						Usage: "Add this value to the end of each entry.",
						Value: "",
					},
					cli.StringFlag{
						Name:  "source",
						Usage: "Source to lookup IPs. [metadata, dns]",
						Value: "metadata",
					},
					cli.BoolFlag{
						Name:  "use-agent-ips",
						Usage: "Use agent ips instead of rancher ips, only works with metadata source",
					},
					cli.BoolFlag{
						Name:  "use-agent-names",
						Usage: "Use agent name instead of rancher ips, only works with metadata source",
					},
				},
			}, {
				Name:   "myip",
				Usage:  "Prints the containers Rancher managed IP",
				Action: ipMyIpAction,
			},
		},
	}
}

func ipMyIpAction(c *cli.Context) error {
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	selfContainer, err := mdClient.GetSelfContainer()
	if err != nil {
		logrus.Fatalf("Failed to find IP: %v", err)
	}
	fmt.Print(selfContainer.PrimaryIp)

	return nil
}

func ipStringifyAction(c *cli.Context) error {
	str := ""
	var err error

	if c.String("source") == "dns" {
		str, err = ipStringifyDNS(c)
		if err != nil {
			logrus.Fatalf("Failed to generate string: %v", err)
		}
	} else {
		str, err = ipStringifyMetadata(c)
		if err != nil {
			logrus.Fatalf("Failed to generate string: %v", err)
		}
	}

	fmt.Print(str)

	return nil
}

func ipStringifyDNS(c *cli.Context) (string, error) {
	if len(c.Args()) <= 0 {
		return "", nil
	}
	ips, err := getDnsContainerIPs(c.Args().First())
	rString := joinString(
		c.String("prefix"),
		c.String("suffix"),
		c.String("delimiter"),
		ips,
	)
	return rString, err
}

func getDnsContainerIPs(host string) ([]string, error) {
	ips := []string{}
	err := &StringifyError{"Could not resolve Host: " + host}

	if testDNSResolves(host) {
		if ipBytes, errIP := net.LookupIP(host); errIP == nil {
			for _, ip := range ipBytes {
				ips = append(ips, ip.String())
			}
			return ips, errIP
		}
	}

	return ips, err
}

func testDNSResolves(host string) bool {
	retVal := false
	ticker := time.NewTicker(time.Millisecond * 500)
	timer := time.NewTimer(time.Second * 60)
	resolves := make(chan bool)

	go func() {
		for range ticker.C {
			if _, err := net.LookupIP(host); err == nil {
				retVal = true
				resolves <- true
			}
		}
	}()

	for {
		select {
		case <-timer.C:
			ticker.Stop()
			return retVal
		case <-resolves:
			ticker.Stop()
			timer.Stop()
			return retVal
		}
	}
}

func ipStringifyMetadata(c *cli.Context) (string, error) {
	split := []string{}
	rString := ""
	var err error

	if len(c.Args()) > 0 {
		split = strings.SplitN(c.Args().First(), "/", 2)
	} else {
		split, err = getSelfStackServiceName()
	}

	getMetaIPMethod := getMetadataContainerIPs
	if c.Bool("use-agent-ips") {
		getMetaIPMethod = getMetadataAgentIPs
	}

	if c.Bool("use-agent-names") {
		getMetaIPMethod = getMetadataAgentNames
	}

	if len(split) == 2 {
		ips, err := getMetaIPMethod(split[0], split[1])
		if err != nil {
			return rString, err
		}
		rString = joinString(
			c.String("prefix"),
			c.String("suffix"),
			c.String("delimiter"),
			ips,
		)
		err = nil
	} else {
		err = &StringifyError{"Not enough arguements supplied. Need stack/service or this container is not part of service"}
	}

	return rString, err
}

func getSelfStackServiceName() ([]string, error) {
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	selfContainer, err := mdClient.GetSelfContainer()
	if err != nil {
		return nil, err
	}

	return []string{selfContainer.StackName, selfContainer.ServiceName}, err
}

func getMetadataContainerIPs(stack string, service string) ([]string, error) {
	rIPs := []string{}
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	containers, err := mdClient.GetServiceContainers(service, stack)
	if err != nil {
		return rIPs, err
	}

	for _, container := range containers {
		rIPs = append(rIPs, container.PrimaryIp)
	}

	return rIPs, nil
}

func getMetadataAgentIPs(stack string, service string) ([]string, error) {
	return getMetadataAgentInfoStrings(stack, service, "AgentIP")
}

func getMetadataAgentNames(stack string, service string) ([]string, error) {
	return getMetadataAgentInfoStrings(stack, service, "Name")
}

func getMetadataAgentInfoStrings(stack, service, property string) ([]string, error) {
	rInfo := []string{}
	mdClient, _ := metadata.NewClientAndWait(metadataURL)

	containers, err := mdClient.GetServiceContainers(service, stack)
	if err != nil {
		return rInfo, err
	}

	for _, container := range containers {
		host, err := mdClient.GetHost(container.HostUUID)
		if err != nil {
			return rInfo, err
		}
		rInfo = append(rInfo, getHostInfoProperty(&host, property))
	}

	return rInfo, nil
}

func getHostInfoProperty(host *metadata.Host, property string) string {
	switch {
	case property == "Name":
		return host.Name
	case property == "AgentIP":
		return host.AgentIP
	default:
		return ""
	}
}

func joinString(pfx string, suffix string, delim string, list []string) string {
	intList := []string{}
	for _, item := range list {
		intermediate := pfx + item + suffix
		intList = append(intList, strings.Repeat(intermediate, 1))
	}
	return strings.Join(intList, delim)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/kexec"
)

var (
	useV4              = flag.Bool("4", false, "Get a DHCPv4 lease")
	useV6              = flag.Bool("6", true, "Get a DHCPv6 lease")
	ifname             = flag.String("i", "eth0", "Interface to send packets through")
	dryRun             = flag.Bool("dryrun", false, "Do everything except assigning IP addresses, changing DNS, and kexec")
	doDebug            = flag.Bool("d", false, "Print debug output")
	skipDHCP           = flag.Bool("skip-dhcp", false, "Skip DHCP and rely on SLAAC for network configuration. This requires -netboot-url")
	overrideNetbootURL = flag.String("netboot-url", "", "Override the netboot URL normally obtained via DHCP")
	readTimeout        = flag.Int("timeout", 3, "Read timeout in seconds")
	dhcpRetries        = flag.Int("retries", 3, "Number of times a DHCP request is retried")
	userClass          = flag.String("userclass", "", "Override DHCP User Class option")
)

const (
	interfaceUpTimeout = 30 * time.Second
)

var banner = `

 _________________________________
< Net booting is so hot right now >
 ---------------------------------
        \   ^__^
         \  (oo)\_______
            (__)\       )\/\
                ||----w |
                ||     ||

`

func main() {
	flag.Parse()
	if *skipDHCP && *overrideNetbootURL == "" {
		log.Fatal("-skip-dhcp requires -netboot-url")
	}
	debug := func(string, ...interface{}) {}
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	if !*useV6 && !*useV4 {
		log.Fatal("At least one of DHCPv6 and DHCPv4 is required")
	}
	// DHCPv6
	if *useV6 {
		log.Printf("Trying to obtain a DHCPv6 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv6: IfUp failed: %v", err)
		}
		debug("DHCPv6: interface %s is up", *ifname)
		var (
			netconf  *netboot.NetConf
			bootfile string
		)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			// send a netboot request via DHCP
			modifiers := []dhcpv6.Modifier{
				dhcpv6.WithArchType(dhcpv6.EFI_X86_64),
			}
			if *userClass != "" {
				modifiers = append(modifiers, dhcpv6.WithUserClass([]byte(*userClass)))
			}
			conversation, err := netboot.RequestNetbootv6(*ifname, time.Duration(*readTimeout)*time.Second, *dhcpRetries, modifiers...)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv6: netboot request for interface %s failed: %v", *ifname, err)
			}
			// get network configuration and boot file
			netconf, bootfile, err = netboot.ConversationToNetconf(conversation)
			if err != nil {
				log.Fatalf("DHCPv6: failed to extract network configuration for %s: %v", *ifname, err)
			}
			debug("DHCPv6: network configuration: %+v", netconf)
			if !*dryRun {
				// Set up IP addresses
				log.Printf("DHCPv6: configuring network interface %s", *ifname)
				if err = netboot.ConfigureInterface(*ifname, netconf); err != nil {
					log.Fatalf("DHCPv6: cannot configure IPv6 addresses on interface %s: %v", *ifname, err)
				}
				// Set up DNS
			}
			if *overrideNetbootURL != "" {
				bootfile = *overrideNetbootURL
			}
			log.Printf("DHCPv6: boot file for interface %s is %s", *ifname, bootfile)
		}
		if *overrideNetbootURL != "" {
			bootfile = *overrideNetbootURL
		}
		debug("DHCPv6: boot file URL is %s", bootfile)
		// check for supported schemes
		if !strings.HasPrefix(bootfile, "http://") {
			log.Fatal("DHCPv6: can only handle http scheme")
		}

		log.Printf("DHCPv6: fetching boot file URL: %s", bootfile)
		resp, err := http.Get(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: http.Get of %s failed: %v", bootfile, err)
		}
		// FIXME this will not be called if something fails after this point
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			log.Fatalf("Status code is not 200 OK: %d", resp.StatusCode)
		}
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Fatalf("DHCPv6: cannot read boot file from the network: %v", err)
		}
		u, err := url.Parse(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: cannot parse URL %s: %v", bootfile, err)
		}
		// extract file name component
		if strings.HasSuffix(u.Path, "/") {
			log.Fatalf("Invalid file path, cannot end with '/': %s", u.Path)
		}
		filename := filepath.Base(u.Path)
		if filename == "." || filename == "" {
			log.Fatalf("Invalid empty file name extracted from file path %s", u.Path)
		}
		if err = ioutil.WriteFile(filename, body, 0400); err != nil {
			log.Fatalf("DHCPv6: cannot write to file %s: %v", filename, err)
		}
		debug("DHCPv6: saved boot file to %s", filename)
		if !*dryRun {
			log.Printf("DHCPv6: kexec'ing into %s", filename)
			kernel, err := os.OpenFile(filename, os.O_RDONLY, 0)
			if err != nil {
				log.Fatalf("DHCPv6: cannot open file %s: %v", filename, err)
			}
			if err = kexec.FileLoad(kernel, nil /* ramfs */, "" /* cmdline */); err != nil {
				log.Fatalf("DHCPv6: kexec.FileLoad failed: %v", err)
			}
			if err = kexec.Reboot(); err != nil {
				log.Fatalf("DHCPv6: kexec.Reboot failed: %v", err)
			}
		}
	}
	// DHCPv4
	if *useV4 {
		log.Printf("Trying to obtain a DHCPv4 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv4: IfUp failed: %v", err)
		}
		debug("DHCPv4: interface %s is up", *ifname)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			log.Print("DHCPv4: sending request")
			client := dhcpv4.NewClient()
			// TODO add options to request to netboot
			conversation, err := client.Exchange(*ifname, nil)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv4: Exchange failed: %v", err)
			}
			// TODO configure the network and DNS
			// TODO extract the next server and boot file and fetch it
			// TODO kexec into the NBP
		}
	}

}
<commit_msg>More verbose netboot output<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/kexec"
)

var (
	useV4              = flag.Bool("4", false, "Get a DHCPv4 lease")
	useV6              = flag.Bool("6", true, "Get a DHCPv6 lease")
	ifname             = flag.String("i", "eth0", "Interface to send packets through")
	dryRun             = flag.Bool("dryrun", false, "Do everything except assigning IP addresses, changing DNS, and kexec")
	doDebug            = flag.Bool("d", false, "Print debug output")
	skipDHCP           = flag.Bool("skip-dhcp", false, "Skip DHCP and rely on SLAAC for network configuration. This requires -netboot-url")
	overrideNetbootURL = flag.String("netboot-url", "", "Override the netboot URL normally obtained via DHCP")
	readTimeout        = flag.Int("timeout", 3, "Read timeout in seconds")
	dhcpRetries        = flag.Int("retries", 3, "Number of times a DHCP request is retried")
	userClass          = flag.String("userclass", "", "Override DHCP User Class option")
)

const (
	interfaceUpTimeout = 30 * time.Second
)

var banner = `

 _________________________________
< Net booting is so hot right now >
 ---------------------------------
        \   ^__^
         \  (oo)\_______
            (__)\       )\/\
                ||----w |
                ||     ||

`

func main() {
	flag.Parse()
	if *skipDHCP && *overrideNetbootURL == "" {
		log.Fatal("-skip-dhcp requires -netboot-url")
	}
	debug := func(string, ...interface{}) {}
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	if !*useV6 && !*useV4 {
		log.Fatal("At least one of DHCPv6 and DHCPv4 is required")
	}
	// DHCPv6
	if *useV6 {
		log.Printf("Trying to obtain a DHCPv6 lease on %s", *ifname)
		log.Printf("Waiting for network interface %s to come up", *ifname)
		start := time.Now()
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv6: IfUp failed: %v", err)
		}
		debug("Interface %s is up after %v", *ifname, time.Since(start))
		var (
			netconf  *netboot.NetConf
			bootfile string
		)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			// send a netboot request via DHCP
			modifiers := []dhcpv6.Modifier{
				dhcpv6.WithArchType(dhcpv6.EFI_X86_64),
			}
			if *userClass != "" {
				modifiers = append(modifiers, dhcpv6.WithUserClass([]byte(*userClass)))
			}
			conversation, err := netboot.RequestNetbootv6(*ifname, time.Duration(*readTimeout)*time.Second, *dhcpRetries, modifiers...)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv6: netboot request for interface %s failed: %v", *ifname, err)
			}
			// get network configuration and boot file
			netconf, bootfile, err = netboot.ConversationToNetconf(conversation)
			if err != nil {
				log.Fatalf("DHCPv6: failed to extract network configuration for %s: %v", *ifname, err)
			}
			debug("DHCPv6: network configuration: %+v", netconf)
			if !*dryRun {
				// Set up IP addresses
				log.Printf("DHCPv6: configuring network interface %s", *ifname)
				if err = netboot.ConfigureInterface(*ifname, netconf); err != nil {
					log.Fatalf("DHCPv6: cannot configure IPv6 addresses on interface %s: %v", *ifname, err)
				}
				// Set up DNS
			}
			if *overrideNetbootURL != "" {
				bootfile = *overrideNetbootURL
			}
			log.Printf("DHCPv6: boot file for interface %s is %s", *ifname, bootfile)
		}
		if *overrideNetbootURL != "" {
			bootfile = *overrideNetbootURL
		}
		debug("DHCPv6: boot file URL is %s", bootfile)
		// check for supported schemes
		if !strings.HasPrefix(bootfile, "http://") {
			log.Fatal("DHCPv6: can only handle http scheme")
		}

		log.Printf("DHCPv6: fetching boot file URL: %s", bootfile)
		resp, err := http.Get(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: http.Get of %s failed: %v", bootfile, err)
		}
		// FIXME this will not be called if something fails after this point
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			log.Fatalf("Status code is not 200 OK: %d", resp.StatusCode)
		}
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Fatalf("DHCPv6: cannot read boot file from the network: %v", err)
		}
		u, err := url.Parse(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: cannot parse URL %s: %v", bootfile, err)
		}
		// extract file name component
		if strings.HasSuffix(u.Path, "/") {
			log.Fatalf("Invalid file path, cannot end with '/': %s", u.Path)
		}
		filename := filepath.Base(u.Path)
		if filename == "." || filename == "" {
			log.Fatalf("Invalid empty file name extracted from file path %s", u.Path)
		}
		if err = ioutil.WriteFile(filename, body, 0400); err != nil {
			log.Fatalf("DHCPv6: cannot write to file %s: %v", filename, err)
		}
		debug("DHCPv6: saved boot file to %s", filename)
		if !*dryRun {
			log.Printf("DHCPv6: kexec'ing into %s", filename)
			kernel, err := os.OpenFile(filename, os.O_RDONLY, 0)
			if err != nil {
				log.Fatalf("DHCPv6: cannot open file %s: %v", filename, err)
			}
			if err = kexec.FileLoad(kernel, nil /* ramfs */, "" /* cmdline */); err != nil {
				log.Fatalf("DHCPv6: kexec.FileLoad failed: %v", err)
			}
			if err = kexec.Reboot(); err != nil {
				log.Fatalf("DHCPv6: kexec.Reboot failed: %v", err)
			}
		}
	}
	// DHCPv4
	if *useV4 {
		log.Printf("Trying to obtain a DHCPv4 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv4: IfUp failed: %v", err)
		}
		debug("DHCPv4: interface %s is up", *ifname)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			log.Print("DHCPv4: sending request")
			client := dhcpv4.NewClient()
			// TODO add options to request to netboot
			conversation, err := client.Exchange(*ifname, nil)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv4: Exchange failed: %v", err)
			}
			// TODO configure the network and DNS
			// TODO extract the next server and boot file and fetch it
			// TODO kexec into the NBP
		}
	}

}
<|endoftext|>


<commit_before>// The ebssurrogate package contains a packer.Builder implementation that
// builds a new EBS-backed AMI using an ephemeral instance.
package ebssurrogate

import (
	"context"
	"errors"
	"fmt"

	"github.com/aws/aws-sdk-go/service/ec2"
	awscommon "github.com/hashicorp/packer/builder/amazon/common"
	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
)

const BuilderId = "mitchellh.amazon.ebssurrogate"

type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`
	awscommon.RunConfig    `mapstructure:",squash"`
	awscommon.BlockDevices `mapstructure:",squash"`
	awscommon.AMIConfig    `mapstructure:",squash"`

	RootDevice    RootBlockDevice  `mapstructure:"ami_root_device"`
	VolumeRunTags awscommon.TagMap `mapstructure:"run_volume_tags"`
	Architecture  string           `mapstructure:"ami_architecture"`

	ctx interpolate.Context
}

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	b.config.ctx.Funcs = awscommon.TemplateFuncs
	err := config.Decode(&b.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &b.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"ami_description",
				"run_tags",
				"run_volume_tags",
				"snapshot_tags",
				"spot_tags",
				"tags",
			},
		},
	}, raws...)
	if err != nil {
		return nil, err
	}

	if b.config.PackerConfig.PackerForce {
		b.config.AMIForceDeregister = true
	}

	// Accumulate any errors
	var errs *packer.MultiError
	var warns []string
	errs = packer.MultiErrorAppend(errs, b.config.AccessConfig.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.RunConfig.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs,
		b.config.AMIConfig.Prepare(&b.config.AccessConfig, &b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.BlockDevices.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.RootDevice.Prepare(&b.config.ctx)...)

	if b.config.AMIVirtType == "" {
		errs = packer.MultiErrorAppend(errs, errors.New("ami_virtualization_type is required."))
	}

	foundRootVolume := false
	for _, launchDevice := range b.config.BlockDevices.LaunchMappings {
		if launchDevice.DeviceName == b.config.RootDevice.SourceDeviceName {
			foundRootVolume = true
			if launchDevice.OmitFromArtifact {
				errs = packer.MultiErrorAppend(errs, fmt.Errorf("You cannot set \"omit_from_artifact\": \"true\" for the root volume."))
			}
		}
	}

	if !foundRootVolume {
		errs = packer.MultiErrorAppend(errs, fmt.Errorf("no volume with name '%s' is found", b.config.RootDevice.SourceDeviceName))
	}

	if b.config.IsSpotInstance() && ((b.config.AMIENASupport != nil && *b.config.AMIENASupport) || b.config.AMISriovNetSupport) {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("Spot instances do not support modification, which is required "+
				"when either `ena_support` or `sriov_support` are set. Please ensure "+
				"you use an AMI that already has either SR-IOV or ENA enabled."))
	}

	if b.config.RunConfig.SpotPriceAutoProduct != "" {
		warns = append(warns, "spot_price_auto_product is deprecated and no "+
			"longer necessary for Packer builds. In future versions of "+
			"Packer, inclusion of spot_price_auto_product will error your "+
			"builds. Please take a look at our current documentation to "+
			"understand how Packer requests Spot instances.")
	}

	if b.config.Architecture == "" {
		b.config.Architecture = "x86_64"
	}
	valid := false
	for _, validArch := range []string{"x86_64", "arm64"} {
		if validArch == b.config.Architecture {
			valid = true
			break
		}
	}
	if !valid {
		errs = packer.MultiErrorAppend(errs, errors.New(`The only valid ami_architecture values are "x86_64" and "arm64"`))
	}
	if errs != nil && len(errs.Errors) > 0 {
		return warns, errs
	}

	packer.LogSecretFilter.Set(b.config.AccessKey, b.config.SecretKey, b.config.Token)

	return warns, nil
}

func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (packer.Artifact, error) {
	session, err := b.config.Session()
	if err != nil {
		return nil, err
	}
	ec2conn := ec2.New(session)

	// Setup the state bag and initial state for the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", &b.config)
	state.Put("access_config", &b.config.AccessConfig)
	state.Put("ami_config", &b.config.AMIConfig)
	state.Put("ec2", ec2conn)
	state.Put("awsSession", session)
	state.Put("hook", hook)
	state.Put("ui", ui)

	var instanceStep multistep.Step

	if b.config.IsSpotInstance() {
		instanceStep = &awscommon.StepRunSpotInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			BlockDevices:                      b.config.BlockDevices,
			BlockDurationMinutes:              b.config.BlockDurationMinutes,
			Ctx:                               b.config.ctx,
			Comm:                              &b.config.RunConfig.Comm,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			ExpectedRootDevice:                "ebs",
			IamInstanceProfile:                b.config.IamInstanceProfile,
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			SourceAMI:                         b.config.SourceAmi,
			SpotPrice:                         b.config.SpotPrice,
			SpotInstanceTypes:                 b.config.SpotInstanceTypes,
			SpotTags:                          b.config.SpotTags,
			Tags:                              b.config.RunTags,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
		}
	} else {
		instanceStep = &awscommon.StepRunSourceInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			BlockDevices:                      b.config.BlockDevices,
			Comm:                              &b.config.RunConfig.Comm,
			Ctx:                               b.config.ctx,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			EnableT2Unlimited:                 b.config.EnableT2Unlimited,
			ExpectedRootDevice:                "ebs",
			IamInstanceProfile:                b.config.IamInstanceProfile,
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			IsRestricted:                      b.config.IsChinaCloud() || b.config.IsGovCloud(),
			SourceAMI:                         b.config.SourceAmi,
			Tags:                              b.config.RunTags,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
		}
	}

	amiDevices := b.config.BuildAMIDevices()
	launchDevices := b.config.BuildLaunchDevices()

	// Build the steps
	steps := []multistep.Step{
		&awscommon.StepPreValidate{
			DestAmiName:     b.config.AMIName,
			ForceDeregister: b.config.AMIForceDeregister,
		},
		&awscommon.StepSourceAMIInfo{
			SourceAmi:                b.config.SourceAmi,
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
			AmiFilters:               b.config.SourceAmiFilter,
			AMIVirtType:              b.config.AMIVirtType,
		},
		&awscommon.StepNetworkInfo{
			VpcId:               b.config.VpcId,
			VpcFilter:           b.config.VpcFilter,
			SecurityGroupIds:    b.config.SecurityGroupIds,
			SecurityGroupFilter: b.config.SecurityGroupFilter,
			SubnetId:            b.config.SubnetId,
			SubnetFilter:        b.config.SubnetFilter,
			AvailabilityZone:    b.config.AvailabilityZone,
		},
		&awscommon.StepKeyPair{
			Debug:        b.config.PackerDebug,
			Comm:         &b.config.RunConfig.Comm,
			DebugKeyPath: fmt.Sprintf("ec2_%s.pem", b.config.PackerBuildName),
		},
		&awscommon.StepSecurityGroup{
			SecurityGroupFilter:    b.config.SecurityGroupFilter,
			SecurityGroupIds:       b.config.SecurityGroupIds,
			CommConfig:             &b.config.RunConfig.Comm,
			TemporarySGSourceCidrs: b.config.TemporarySGSourceCidrs,
		},
		&awscommon.StepCleanupVolumes{
			BlockDevices: b.config.BlockDevices,
		},
		instanceStep,
		&awscommon.StepGetPassword{
			Debug:     b.config.PackerDebug,
			Comm:      &b.config.RunConfig.Comm,
			Timeout:   b.config.WindowsPasswordTimeout,
			BuildName: b.config.PackerBuildName,
		},
		&communicator.StepConnect{
			Config: &b.config.RunConfig.Comm,
			Host: awscommon.SSHHost(
				ec2conn,
				b.config.SSHInterface),
			SSHConfig: b.config.RunConfig.Comm.SSHConfigFunc(),
		},
		&common.StepProvision{},
		&common.StepCleanupTempKeys{
			Comm: &b.config.RunConfig.Comm,
		},
		&awscommon.StepStopEBSBackedInstance{
			Skip:                b.config.IsSpotInstance(),
			DisableStopInstance: b.config.DisableStopInstance,
		},
		&awscommon.StepModifyEBSBackedInstance{
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
		},
		&StepSnapshotVolumes{
			LaunchDevices:   launchDevices,
			SnapshotOmitMap: b.config.GetOmissions(),
		},
		&awscommon.StepDeregisterAMI{
			AccessConfig:        &b.config.AccessConfig,
			ForceDeregister:     b.config.AMIForceDeregister,
			ForceDeleteSnapshot: b.config.AMIForceDeleteSnapshot,
			AMIName:             b.config.AMIName,
			Regions:             b.config.AMIRegions,
		},
		&StepRegisterAMI{
			RootDevice:               b.config.RootDevice,
			AMIDevices:               amiDevices,
			LaunchDevices:            launchDevices,
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
			Architecture:             b.config.Architecture,
			LaunchOmitMap:            b.config.GetOmissions(),
		},
		&awscommon.StepAMIRegionCopy{
			AccessConfig:      &b.config.AccessConfig,
			Regions:           b.config.AMIRegions,
			AMIKmsKeyId:       b.config.AMIKmsKeyId,
			RegionKeyIds:      b.config.AMIRegionKMSKeyIDs,
			EncryptBootVolume: b.config.AMIEncryptBootVolume,
			Name:              b.config.AMIName,
			OriginalRegion:    *ec2conn.Config.Region,
		},
		&awscommon.StepModifyAMIAttributes{
			Description:    b.config.AMIDescription,
			Users:          b.config.AMIUsers,
			Groups:         b.config.AMIGroups,
			ProductCodes:   b.config.AMIProductCodes,
			SnapshotUsers:  b.config.SnapshotUsers,
			SnapshotGroups: b.config.SnapshotGroups,
			Ctx:            b.config.ctx,
		},
		&awscommon.StepCreateTags{
			Tags:         b.config.AMITags,
			SnapshotTags: b.config.SnapshotTags,
			Ctx:          b.config.ctx,
		},
	}

	// Run!
	b.runner = common.NewRunner(steps, b.config.PackerConfig, ui)
	b.runner.Run(ctx, state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	if amis, ok := state.GetOk("amis"); ok {
		// Build the artifact and return it
		artifact := &awscommon.Artifact{
			Amis:           amis.(map[string]string),
			BuilderIdValue: BuilderId,
			Session:        session,
		}

		return artifact, nil
	}

	return nil, nil
}
<commit_msg>fix: change name to keypair key path file<commit_after>// The ebssurrogate package contains a packer.Builder implementation that
// builds a new EBS-backed AMI using an ephemeral instance.
package ebssurrogate

import (
	"context"
	"errors"
	"fmt"

	"github.com/aws/aws-sdk-go/service/ec2"
	awscommon "github.com/hashicorp/packer/builder/amazon/common"
	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
)

const BuilderId = "mitchellh.amazon.ebssurrogate"

type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`
	awscommon.RunConfig    `mapstructure:",squash"`
	awscommon.BlockDevices `mapstructure:",squash"`
	awscommon.AMIConfig    `mapstructure:",squash"`

	RootDevice    RootBlockDevice  `mapstructure:"ami_root_device"`
	VolumeRunTags awscommon.TagMap `mapstructure:"run_volume_tags"`
	Architecture  string           `mapstructure:"ami_architecture"`

	ctx interpolate.Context
}

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	b.config.ctx.Funcs = awscommon.TemplateFuncs
	err := config.Decode(&b.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &b.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"ami_description",
				"run_tags",
				"run_volume_tags",
				"snapshot_tags",
				"spot_tags",
				"tags",
			},
		},
	}, raws...)
	if err != nil {
		return nil, err
	}

	if b.config.PackerConfig.PackerForce {
		b.config.AMIForceDeregister = true
	}

	// Accumulate any errors
	var errs *packer.MultiError
	var warns []string
	errs = packer.MultiErrorAppend(errs, b.config.AccessConfig.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.RunConfig.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs,
		b.config.AMIConfig.Prepare(&b.config.AccessConfig, &b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.BlockDevices.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.RootDevice.Prepare(&b.config.ctx)...)

	if b.config.AMIVirtType == "" {
		errs = packer.MultiErrorAppend(errs, errors.New("ami_virtualization_type is required."))
	}

	foundRootVolume := false
	for _, launchDevice := range b.config.BlockDevices.LaunchMappings {
		if launchDevice.DeviceName == b.config.RootDevice.SourceDeviceName {
			foundRootVolume = true
			if launchDevice.OmitFromArtifact {
				errs = packer.MultiErrorAppend(errs, fmt.Errorf("You cannot set \"omit_from_artifact\": \"true\" for the root volume."))
			}
		}
	}

	if !foundRootVolume {
		errs = packer.MultiErrorAppend(errs, fmt.Errorf("no volume with name '%s' is found", b.config.RootDevice.SourceDeviceName))
	}

	if b.config.IsSpotInstance() && ((b.config.AMIENASupport != nil && *b.config.AMIENASupport) || b.config.AMISriovNetSupport) {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("Spot instances do not support modification, which is required "+
				"when either `ena_support` or `sriov_support` are set. Please ensure "+
				"you use an AMI that already has either SR-IOV or ENA enabled."))
	}

	if b.config.RunConfig.SpotPriceAutoProduct != "" {
		warns = append(warns, "spot_price_auto_product is deprecated and no "+
			"longer necessary for Packer builds. In future versions of "+
			"Packer, inclusion of spot_price_auto_product will error your "+
			"builds. Please take a look at our current documentation to "+
			"understand how Packer requests Spot instances.")
	}

	if b.config.Architecture == "" {
		b.config.Architecture = "x86_64"
	}
	valid := false
	for _, validArch := range []string{"x86_64", "arm64"} {
		if validArch == b.config.Architecture {
			valid = true
			break
		}
	}
	if !valid {
		errs = packer.MultiErrorAppend(errs, errors.New(`The only valid ami_architecture values are "x86_64" and "arm64"`))
	}
	if errs != nil && len(errs.Errors) > 0 {
		return warns, errs
	}

	packer.LogSecretFilter.Set(b.config.AccessKey, b.config.SecretKey, b.config.Token)

	return warns, nil
}

func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (packer.Artifact, error) {
	session, err := b.config.Session()
	if err != nil {
		return nil, err
	}
	ec2conn := ec2.New(session)

	// Setup the state bag and initial state for the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", &b.config)
	state.Put("access_config", &b.config.AccessConfig)
	state.Put("ami_config", &b.config.AMIConfig)
	state.Put("ec2", ec2conn)
	state.Put("awsSession", session)
	state.Put("hook", hook)
	state.Put("ui", ui)

	var instanceStep multistep.Step

	if b.config.IsSpotInstance() {
		instanceStep = &awscommon.StepRunSpotInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			BlockDevices:                      b.config.BlockDevices,
			BlockDurationMinutes:              b.config.BlockDurationMinutes,
			Ctx:                               b.config.ctx,
			Comm:                              &b.config.RunConfig.Comm,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			ExpectedRootDevice:                "ebs",
			IamInstanceProfile:                b.config.IamInstanceProfile,
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			SourceAMI:                         b.config.SourceAmi,
			SpotPrice:                         b.config.SpotPrice,
			SpotInstanceTypes:                 b.config.SpotInstanceTypes,
			SpotTags:                          b.config.SpotTags,
			Tags:                              b.config.RunTags,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
		}
	} else {
		instanceStep = &awscommon.StepRunSourceInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			BlockDevices:                      b.config.BlockDevices,
			Comm:                              &b.config.RunConfig.Comm,
			Ctx:                               b.config.ctx,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			EnableT2Unlimited:                 b.config.EnableT2Unlimited,
			ExpectedRootDevice:                "ebs",
			IamInstanceProfile:                b.config.IamInstanceProfile,
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			IsRestricted:                      b.config.IsChinaCloud() || b.config.IsGovCloud(),
			SourceAMI:                         b.config.SourceAmi,
			Tags:                              b.config.RunTags,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
		}
	}

	amiDevices := b.config.BuildAMIDevices()
	launchDevices := b.config.BuildLaunchDevices()

	// Build the steps
	steps := []multistep.Step{
		&awscommon.StepPreValidate{
			DestAmiName:     b.config.AMIName,
			ForceDeregister: b.config.AMIForceDeregister,
		},
		&awscommon.StepSourceAMIInfo{
			SourceAmi:                b.config.SourceAmi,
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
			AmiFilters:               b.config.SourceAmiFilter,
			AMIVirtType:              b.config.AMIVirtType,
		},
		&awscommon.StepNetworkInfo{
			VpcId:               b.config.VpcId,
			VpcFilter:           b.config.VpcFilter,
			SecurityGroupIds:    b.config.SecurityGroupIds,
			SecurityGroupFilter: b.config.SecurityGroupFilter,
			SubnetId:            b.config.SubnetId,
			SubnetFilter:        b.config.SubnetFilter,
			AvailabilityZone:    b.config.AvailabilityZone,
		},
		&awscommon.StepKeyPair{
			Debug:        b.config.PackerDebug,
			Comm:         &b.config.RunConfig.Comm,
			DebugKeyPath: fmt.Sprintf("ec2_%s", b.config.PackerBuildName),
		},
		&awscommon.StepSecurityGroup{
			SecurityGroupFilter:    b.config.SecurityGroupFilter,
			SecurityGroupIds:       b.config.SecurityGroupIds,
			CommConfig:             &b.config.RunConfig.Comm,
			TemporarySGSourceCidrs: b.config.TemporarySGSourceCidrs,
		},
		&awscommon.StepCleanupVolumes{
			BlockDevices: b.config.BlockDevices,
		},
		instanceStep,
		&awscommon.StepGetPassword{
			Debug:     b.config.PackerDebug,
			Comm:      &b.config.RunConfig.Comm,
			Timeout:   b.config.WindowsPasswordTimeout,
			BuildName: b.config.PackerBuildName,
		},
		&communicator.StepConnect{
			Config: &b.config.RunConfig.Comm,
			Host: awscommon.SSHHost(
				ec2conn,
				b.config.SSHInterface),
			SSHConfig: b.config.RunConfig.Comm.SSHConfigFunc(),
		},
		&common.StepProvision{},
		&common.StepCleanupTempKeys{
			Comm: &b.config.RunConfig.Comm,
		},
		&awscommon.StepStopEBSBackedInstance{
			Skip:                b.config.IsSpotInstance(),
			DisableStopInstance: b.config.DisableStopInstance,
		},
		&awscommon.StepModifyEBSBackedInstance{
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
		},
		&StepSnapshotVolumes{
			LaunchDevices:   launchDevices,
			SnapshotOmitMap: b.config.GetOmissions(),
		},
		&awscommon.StepDeregisterAMI{
			AccessConfig:        &b.config.AccessConfig,
			ForceDeregister:     b.config.AMIForceDeregister,
			ForceDeleteSnapshot: b.config.AMIForceDeleteSnapshot,
			AMIName:             b.config.AMIName,
			Regions:             b.config.AMIRegions,
		},
		&StepRegisterAMI{
			RootDevice:               b.config.RootDevice,
			AMIDevices:               amiDevices,
			LaunchDevices:            launchDevices,
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
			Architecture:             b.config.Architecture,
			LaunchOmitMap:            b.config.GetOmissions(),
		},
		&awscommon.StepAMIRegionCopy{
			AccessConfig:      &b.config.AccessConfig,
			Regions:           b.config.AMIRegions,
			AMIKmsKeyId:       b.config.AMIKmsKeyId,
			RegionKeyIds:      b.config.AMIRegionKMSKeyIDs,
			EncryptBootVolume: b.config.AMIEncryptBootVolume,
			Name:              b.config.AMIName,
			OriginalRegion:    *ec2conn.Config.Region,
		},
		&awscommon.StepModifyAMIAttributes{
			Description:    b.config.AMIDescription,
			Users:          b.config.AMIUsers,
			Groups:         b.config.AMIGroups,
			ProductCodes:   b.config.AMIProductCodes,
			SnapshotUsers:  b.config.SnapshotUsers,
			SnapshotGroups: b.config.SnapshotGroups,
			Ctx:            b.config.ctx,
		},
		&awscommon.StepCreateTags{
			Tags:         b.config.AMITags,
			SnapshotTags: b.config.SnapshotTags,
			Ctx:          b.config.ctx,
		},
	}

	// Run!
	b.runner = common.NewRunner(steps, b.config.PackerConfig, ui)
	b.runner.Run(ctx, state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	if amis, ok := state.GetOk("amis"); ok {
		// Build the artifact and return it
		artifact := &awscommon.Artifact{
			Amis:           amis.(map[string]string),
			BuilderIdValue: BuilderId,
			Session:        session,
		}

		return artifact, nil
	}

	return nil, nil
}
<|endoftext|>