<commit_before>package agent

import (
	"encoding/json"
	"fmt"
	"github.com/golang/glog"
	. "github.com/infradash/dash/pkg/dash"
	"github.com/qorio/maestro/pkg/docker"
	_ "github.com/qorio/maestro/pkg/mqtt"
	"github.com/qorio/maestro/pkg/pubsub"
	"github.com/qorio/maestro/pkg/registry"
	"github.com/qorio/maestro/pkg/zk"
	"github.com/qorio/omni/runtime"
	"github.com/qorio/omni/version"
	"net/http"
	"path"
	"strings"
	"sync"
	"time"
)

type Agent struct {
	QualifyByTags

	ZkSettings
	DockerSettings

	RegistryContainerEntry

	ListenPort   int `json:"listen_port"`
	DockerUIPort int `json:"dockerui_port"`

	UiDocRoot string `json:"ui_doc_root,omitempty"`
	EnableUI  bool   `json:"enable_ui,omitempty"`

	Initializer *ConfigLoader `json:"config_loader"`

	selfRegister bool `json:"-"`

	// json skips these fields
	endpoint       http.Handler      `json:"-"`
	zk             zk.ZK             `json:"-"`
	docker         *docker.Docker    `json:"-"`
	self_container *docker.Container `json:"-"`

	lock          sync.Mutex
	domains       map[string]*Domain
	domainConfigs map[string]DomainConfig

	StatusPubsubTopic string `json:"status_topic,omitempty"`
}

// Checks that all the information required for agent start up is met.
func (this *Agent) checkPreconditions() {
	if this.Host == "" {
		panic(ErrNoHost)
	}
	if this.Name == "" {
		panic(ErrNoName)
	}
}

func (this *Agent) is_running_dockerui() bool {
	return this.EnableUI && this.UiDocRoot != ""
}

// Block until SIGTERM
func (this *Agent) Run() {

	this.checkPreconditions()

	glog.Infoln("Agent", this.GetIdentity())

	this.domains = make(map[string]*Domain, 0)
	this.domainConfigs = make(map[string]DomainConfig, 0)

	err := this.ConnectServices()
	if err != nil {
		panic(err)
	}

	endpoint, err := NewApiEndPoint(this)
	if err != nil {
		panic(err)
	}

	if this.is_running_dockerui() {
		mux := http.NewServeMux()
		glog.Infoln("Starting UI with docroot=", this.UiDocRoot, "DockerPort=", this.DockerPort)
		fileHandler := http.FileServer(http.Dir(this.UiDocRoot))

		dockerApiHandler := this.createDockerApiHandler(this.DockerPort)
		mux.Handle("/dockerapi/", http.StripPrefix("/dockerapi", dockerApiHandler))
		mux.Handle("/", fileHandler)
		go func() {
			p := fmt.Sprintf(":%d", this.DockerUIPort)
			glog.Infoln("UI Listening on", p)
			if err := http.ListenAndServe(p, mux); err != nil {
				panic(err)
			}
		}()
	}

	if this.selfRegister {
		err = this.DiscoverSelfInDocker()
		if err != nil {
			panic(err)
		}

		err = this.Register()
		if err != nil {
			panic(err)
		}
	}

	this.endpoint = endpoint

	if this.Initializer != nil {
		glog.Infoln("Loading configuration from", this.Initializer.ConfigUrl)

		var list []DomainConfig
		loaded, err := this.Initializer.Load(&list, this.AuthToken, this.zk)
		if err != nil {
			panic(err)
		}

		if loaded {

			glog.Infoln("Loaded and applied configuration. Processing.")

			for _, per_domain := range list {

				applied := new(DomainConfig)
				err := ApplyVarSubs(per_domain, applied,
					MergeMaps(map[string]interface{}{
						"Domain": per_domain.Domain,
					}, EscapeVars(ConfigVariables[1:]...)))
				if err != nil {
					panic(err)
				}
				per_domain = *applied

				if _, config_err := this.ConfigureDomain(&per_domain); config_err != nil {
					panic(config_err)
				}
			}

			glog.Infoln("Start running discovery / container monitors")
			matcher := new(DiscoveryContainerMatcher).Init()
			for _, domain := range this.domains {
				watches, err := domain.GetContainerWatcherSpecs()
				if err != nil {
					panic(err)
				}
				for svc, watch := range watches {

					glog.Infoln(domain.Domain, svc, "Container matcher for discovery:", *watch)
					matcher.C(domain.Domain, svc, watch)

					glog.Infoln(domain.Domain, svc, "Set up container monitor:", *watch)
					domain.WatchContainer(svc, watch)
				}
			}

			err = this.DiscoverRunningContainers(
				matcher.Match,
				func(c *docker.Container, match_rule *ContainerMatchRule) {

					// Need to increment a counter regardless of the container's state
					cc := get_sequence_by_image(c.Image)
					glog.V(100).Infoln("Service=", match_rule.Service, "Image=", c.Image, "SequenceCounter=", cc)

					if d, has := this.domains[match_rule.Domain]; has {

						glog.V(100).Infoln("Service=", match_rule.Service, "Id=", c.Id[0:12],
							"FinishedAt=", c.DockerData.State.FinishedAt)

						// Match the name but we need to take into account if it's running or not.
						switch {

						case c.DockerData.State.Restarting:
							d.tracker.Starting(match_rule.Service, c)

						case c.DockerData.State.Running, c.DockerData.State.Restarting:
							d.tracker.Running(match_rule.Service, c)

							glog.Infoln("Registering container Id=", c.Id, "Image=", c.Image)
							if entry, err := BuildRegistryEntry(c, match_rule.GetMatchContainerPort()); entry != nil {
								entry.Domain = match_rule.Domain
								entry.Service = string(match_rule.Service)
								entry.Host = this.Host

								err = entry.Register(this.zk)

								if err != nil {
									glog.Warningln("Error during registration:", err)
								}
								k, v, _ := entry.KeyValue()
								glog.Infoln("Registered", k, v)
							} else {
								glog.Warning("Error building registry", err, "for", *c)
							}

						case c.DockerData.State.FinishedAt.Before(time.Now()):
							glog.V(100).Infoln("Container", "Id=", c.Id[0:12], "Name=", c.Name, "stopped.")
							d.tracker.Stopped(match_rule.Service, c)
						}
					}
				})
			if err != nil {
				glog.Infoln("Error discovering containers:", err)
			}

			// Configure and start up services
			glog.Infoln("Configure domains")
			for _, domain := range this.domains {
				glog.Infoln("Starting services: Domain=", domain.Domain)
				_, config_err := domain.StartServices(this.QualifyByTags)
				if config_err != nil {
					panic(config_err)
				}
			}

			glog.Infoln("Synchronize local states with scheduler")
			for _, domain := range this.domains {
				err := domain.SynchronizeSchedule()
				if err != nil {
					glog.Warningln("Failed to synchronize scheduling for Domain=", domain.Identity, "Err=", err)
				}
			}
		}
	}

	runtime.MinimalContainer(this.ListenPort,
		func() http.Handler {
			return endpoint
		},
		func() error {
			err := endpoint.Stop()
			glog.Infoln("Stopped endpoint", err)
			err = this.zk.Close()
			glog.Infoln("Stopped zk", err)
			return err
		})
}

func (this *Agent) GetIdentity() string {
	return fmt.Sprintf("%s:%s", this.Name, this.Host)
}

func (this *Agent) ConnectServices() error {
	glog.Infoln("Connecting to zookeeper:", this.Hosts)
	zk, err := zk.Connect(strings.Split(this.Hosts, ","), this.Timeout)
	if err != nil {
		return err
	}
	this.zk = zk
	glog.Infoln("Connected to zookeeper:", this.Hosts)

	glog.Infoln("Connecting to docker:", this.DockerSettings)
	if this.Cert != "" {

		docker, err := docker.NewTLSClient(this.DockerPort, this.Cert, this.Key, this.Ca)
		if err != nil {
			return err
		}
		this.docker = docker
		glog.Infoln("Connected to docker:", this.DockerPort)

	} else {
		docker, err := docker.NewClient(this.DockerPort)
		if err != nil {
			return err
		}
		this.docker = docker
		glog.Infoln("Connected to docker:", this.DockerPort)

	}

	// Set up callbacks
	if this.docker != nil {
		// This is where we get callbacks for the containers that the agent initiates.
		// The event api can include containers that are started manually.
	}

	glog.Infoln("Start Zookeeper events channel")
	go func() {

		status := func(map[string]interface{}) {
			// no-op
		}
		if this.StatusPubsubTopic != "" {
			id := path.Join(this.Domain, this.Name, this.Host)
			topic := pubsub.Topic(this.StatusPubsubTopic + "/" + id)
			if topic.Valid() {
				if pb, err := topic.Broker().PubSub(id); err == nil {
					status = func(evt map[string]interface{}) {
						msg, err := json.Marshal(evt)
						if err == nil {
							pb.Publish(topic, msg)
						}
					}
				}
			} else {
				panic(topic)
			}
		}

		events := zk.Events()
		for {
			evt := <-events
			glog.Infoln("ZKEvent:", evt.JSON())

			// send as pubsub
			// TODO - Redpill compatible:
			/*
				type Event struct {
					Status      string `json:"status"`
					Title       string `json:"title,omitempty"`
					Description string `json:"description,omitempty"`
					Note        string `json:"note,omitempty"`
					User        string `json:"user,omitempty"`
					Type        string `json:"type,omitempty"`
					Url         string `json:"url,omitempty"`
					Timestamp   int64  `json:"timestamp,omitempty"`
					ObjectId    string `json:"object_id"`
					ObjectType  string `json:"object_type"`
				}
			*/
			m := evt.AsMap()
			m["object_id"] = path.Join(this.Domain, this.Name, this.Host)
			m["object_type"] = "agent"
			m["timestamp"] = time.Now().Unix()
			m["description"] = fmt.Sprint(m["type"], ":", m["state"], "@", "server=", m["server"])
			m["title"] = "zookeeper event from agent " + path.Join(this.Domain, this.Name, this.Host)
			m["user"] = "dash"
			switch m["state"] {
			case "state-disconnected", "state-auth-failed", "state-expired":
				m["status"] = "fatal"
			case "state-connected", "state-has-session":
				m["status"] = "ok"
			}
			status(m)
		}
	}()

	return nil
}

func (this *Agent) Register() error {
	attempts := 0
	for {
		err := this._register()
		if err != nil {
			if attempts == 12 {
				return err
			} else {
				time.Sleep(5 * time.Second)
				attempts += 1
			}
		} else {
			return nil
		}
	}
}

func (this *Agent) info() interface{} {
	info := map[string]interface{}{
		"api":       fmt.Sprintf("%s:%d", this.Host, this.ListenPort),
		"dockerapi": fmt.Sprintf("http://%s:%d/dockerapi", this.Host, this.ListenPort),
		"version":   *version.BuildInfo(),
	}
	if this.is_running_dockerui() {
		info["dockerui"] = fmt.Sprintf("http://%s:%d/dockerapi", this.Host, this.DockerUIPort)
	}

	return info
}

func (this *Agent) _register() error {
	if this.zk == nil {
		return ErrNotConnectedToRegistry
	}

	key := registry.NewPath("dash", this.Host)
	err := zk.CreateOrSet(this.zk, key, this.info(), true)
	glog.Infoln("Register key=", key, "err=", err)
	if err == nil {
		// Update this only on successful registration
		this.Registration = key.Path()
	}
	return err
}

// Containers in this domain
func (this *Agent) ListContainers(domain, service string) ([]*docker.Container, error) {
	d, has := this.domains[domain]
	if !has {
		return nil, ErrNoDomain
	}
	return d.ListContainers(ServiceKey(service))
}

func (this *Agent) WatchContainer(domain, service string, spec *WatchContainerSpec) error {
	d, has := this.domains[domain]
	if !has {
		return ErrNoDomain
	}

	this.lock.Lock()
	defer this.lock.Unlock()
	return d.WatchContainer(ServiceKey(service), spec)
}

func (this *Agent) ConfigureDomain(config *DomainConfig) (*Domain, error) {
	if config.Domain == "" {
		return nil, ErrNoDomain
	}

	this.lock.Lock()
	defer this.lock.Unlock()

	domain, has := this.domains[config.Domain]
	if !has {

		domain = &Domain{
			Domain:                 config.Domain,
			RegistryContainerEntry: config.RegistryContainerEntry,
			zk:                 this.zk,
			docker:             this.docker,
			container_watchers: make(map[string]chan<- bool, 0),
			triggers:           NewZkWatcher(this.zk),
			agent:              this,
			tracker:            NewContainerTracker(config.Domain),
			schedulers:         make(map[ServiceKey]*Scheduler),
		}

		_, err := domain.StartScheduleExecutor()
		if err != nil {
			glog.Warningln("Cannot start schedule executor")
			return nil, err
		}

		domain.Host = this.Host
		domain.Port = this.Port

		err = domain.Register()
		if err != nil {
			return nil, err
		}

		this.domains[config.Domain] = domain
		this.domainConfigs[config.Domain] = *config
		domain.Config = config
	}
	return this.domains[config.Domain], nil
}
<commit_msg>Change the endpoint reporting for DockerUI<commit_after>package agent

import (
	"encoding/json"
	"fmt"
	"github.com/golang/glog"
	. "github.com/infradash/dash/pkg/dash"
	"github.com/qorio/maestro/pkg/docker"
	_ "github.com/qorio/maestro/pkg/mqtt"
	"github.com/qorio/maestro/pkg/pubsub"
	"github.com/qorio/maestro/pkg/registry"
	"github.com/qorio/maestro/pkg/zk"
	"github.com/qorio/omni/runtime"
	"github.com/qorio/omni/version"
	"net/http"
	"path"
	"strings"
	"sync"
	"time"
)

type Agent struct {
	QualifyByTags

	ZkSettings
	DockerSettings

	RegistryContainerEntry

	ListenPort   int `json:"listen_port"`
	DockerUIPort int `json:"dockerui_port"`

	UiDocRoot string `json:"ui_doc_root,omitempty"`
	EnableUI  bool   `json:"enable_ui,omitempty"`

	Initializer *ConfigLoader `json:"config_loader"`

	selfRegister bool `json:"-"`

	// json skips these fields
	endpoint       http.Handler      `json:"-"`
	zk             zk.ZK             `json:"-"`
	docker         *docker.Docker    `json:"-"`
	self_container *docker.Container `json:"-"`

	lock          sync.Mutex
	domains       map[string]*Domain
	domainConfigs map[string]DomainConfig

	StatusPubsubTopic string `json:"status_topic,omitempty"`
}

// Checks that all the information required for agent start up is met.
func (this *Agent) checkPreconditions() {
	if this.Host == "" {
		panic(ErrNoHost)
	}
	if this.Name == "" {
		panic(ErrNoName)
	}
}

func (this *Agent) is_running_dockerui() bool {
	return this.EnableUI && this.UiDocRoot != ""
}

// Block until SIGTERM
func (this *Agent) Run() {

	this.checkPreconditions()

	glog.Infoln("Agent", this.GetIdentity())

	this.domains = make(map[string]*Domain, 0)
	this.domainConfigs = make(map[string]DomainConfig, 0)

	err := this.ConnectServices()
	if err != nil {
		panic(err)
	}

	endpoint, err := NewApiEndPoint(this)
	if err != nil {
		panic(err)
	}

	if this.is_running_dockerui() {
		mux := http.NewServeMux()
		glog.Infoln("Starting UI with docroot=", this.UiDocRoot, "DockerPort=", this.DockerPort)
		fileHandler := http.FileServer(http.Dir(this.UiDocRoot))

		dockerApiHandler := this.createDockerApiHandler(this.DockerPort)
		mux.Handle("/dockerapi/", http.StripPrefix("/dockerapi", dockerApiHandler))
		mux.Handle("/", fileHandler)
		go func() {
			p := fmt.Sprintf(":%d", this.DockerUIPort)
			glog.Infoln("UI Listening on", p)
			if err := http.ListenAndServe(p, mux); err != nil {
				panic(err)
			}
		}()
	}

	if this.selfRegister {
		err = this.DiscoverSelfInDocker()
		if err != nil {
			panic(err)
		}

		err = this.Register()
		if err != nil {
			panic(err)
		}
	}

	this.endpoint = endpoint

	if this.Initializer != nil {
		glog.Infoln("Loading configuration from", this.Initializer.ConfigUrl)

		var list []DomainConfig
		loaded, err := this.Initializer.Load(&list, this.AuthToken, this.zk)
		if err != nil {
			panic(err)
		}

		if loaded {

			glog.Infoln("Loaded and applied configuration. Processing.")

			for _, per_domain := range list {

				applied := new(DomainConfig)
				err := ApplyVarSubs(per_domain, applied,
					MergeMaps(map[string]interface{}{
						"Domain": per_domain.Domain,
					}, EscapeVars(ConfigVariables[1:]...)))
				if err != nil {
					panic(err)
				}
				per_domain = *applied

				if _, config_err := this.ConfigureDomain(&per_domain); config_err != nil {
					panic(config_err)
				}
			}

			glog.Infoln("Start running discovery / container monitors")
			matcher := new(DiscoveryContainerMatcher).Init()
			for _, domain := range this.domains {
				watches, err := domain.GetContainerWatcherSpecs()
				if err != nil {
					panic(err)
				}
				for svc, watch := range watches {

					glog.Infoln(domain.Domain, svc, "Container matcher for discovery:", *watch)
					matcher.C(domain.Domain, svc, watch)

					glog.Infoln(domain.Domain, svc, "Set up container monitor:", *watch)
					domain.WatchContainer(svc, watch)
				}
			}

			err = this.DiscoverRunningContainers(
				matcher.Match,
				func(c *docker.Container, match_rule *ContainerMatchRule) {

					// Need to increment a counter regardless of the container's state
					cc := get_sequence_by_image(c.Image)
					glog.V(100).Infoln("Service=", match_rule.Service, "Image=", c.Image, "SequenceCounter=", cc)

					if d, has := this.domains[match_rule.Domain]; has {

						glog.V(100).Infoln("Service=", match_rule.Service, "Id=", c.Id[0:12],
							"FinishedAt=", c.DockerData.State.FinishedAt)

						// Match the name but we need to take into account if it's running or not.
						switch {

						case c.DockerData.State.Restarting:
							d.tracker.Starting(match_rule.Service, c)

						case c.DockerData.State.Running, c.DockerData.State.Restarting:
							d.tracker.Running(match_rule.Service, c)

							glog.Infoln("Registering container Id=", c.Id, "Image=", c.Image)
							if entry, err := BuildRegistryEntry(c, match_rule.GetMatchContainerPort()); entry != nil {
								entry.Domain = match_rule.Domain
								entry.Service = string(match_rule.Service)
								entry.Host = this.Host

								err = entry.Register(this.zk)

								if err != nil {
									glog.Warningln("Error during registration:", err)
								}
								k, v, _ := entry.KeyValue()
								glog.Infoln("Registered", k, v)
							} else {
								glog.Warning("Error building registry", err, "for", *c)
							}

						case c.DockerData.State.FinishedAt.Before(time.Now()):
							glog.V(100).Infoln("Container", "Id=", c.Id[0:12], "Name=", c.Name, "stopped.")
							d.tracker.Stopped(match_rule.Service, c)
						}
					}
				})
			if err != nil {
				glog.Infoln("Error discovering containers:", err)
			}

			// Configure and start up services
			glog.Infoln("Configure domains")
			for _, domain := range this.domains {
				glog.Infoln("Starting services: Domain=", domain.Domain)
				_, config_err := domain.StartServices(this.QualifyByTags)
				if config_err != nil {
					panic(config_err)
				}
			}

			glog.Infoln("Synchronize local states with scheduler")
			for _, domain := range this.domains {
				err := domain.SynchronizeSchedule()
				if err != nil {
					glog.Warningln("Failed to synchronize scheduling for Domain=", domain.Identity, "Err=", err)
				}
			}
		}
	}

	runtime.MinimalContainer(this.ListenPort,
		func() http.Handler {
			return endpoint
		},
		func() error {
			err := endpoint.Stop()
			glog.Infoln("Stopped endpoint", err)
			err = this.zk.Close()
			glog.Infoln("Stopped zk", err)
			return err
		})
}

func (this *Agent) GetIdentity() string {
	return fmt.Sprintf("%s:%s", this.Name, this.Host)
}

func (this *Agent) ConnectServices() error {
	glog.Infoln("Connecting to zookeeper:", this.Hosts)
	zk, err := zk.Connect(strings.Split(this.Hosts, ","), this.Timeout)
	if err != nil {
		return err
	}
	this.zk = zk
	glog.Infoln("Connected to zookeeper:", this.Hosts)

	glog.Infoln("Connecting to docker:", this.DockerSettings)
	if this.Cert != "" {

		docker, err := docker.NewTLSClient(this.DockerPort, this.Cert, this.Key, this.Ca)
		if err != nil {
			return err
		}
		this.docker = docker
		glog.Infoln("Connected to docker:", this.DockerPort)

	} else {
		docker, err := docker.NewClient(this.DockerPort)
		if err != nil {
			return err
		}
		this.docker = docker
		glog.Infoln("Connected to docker:", this.DockerPort)

	}

	// Set up callbacks
	if this.docker != nil {
		// This is where we get callbacks for the containers that the agent initiates.
		// The event api can include containers that are started manually.
	}

	glog.Infoln("Start Zookeeper events channel")
	go func() {

		status := func(map[string]interface{}) {
			// no-op
		}
		if this.StatusPubsubTopic != "" {
			id := path.Join(this.Domain, this.Name, this.Host)
			topic := pubsub.Topic(this.StatusPubsubTopic + "/" + id)
			if topic.Valid() {
				if pb, err := topic.Broker().PubSub(id); err == nil {
					status = func(evt map[string]interface{}) {
						msg, err := json.Marshal(evt)
						if err == nil {
							pb.Publish(topic, msg)
						}
					}
				}
			} else {
				panic(topic)
			}
		}

		events := zk.Events()
		for {
			evt := <-events
			glog.Infoln("ZKEvent:", evt.JSON())

			// send as pubsub
			// TODO - Redpill compatible:
			/*
				type Event struct {
					Status      string `json:"status"`
					Title       string `json:"title,omitempty"`
					Description string `json:"description,omitempty"`
					Note        string `json:"note,omitempty"`
					User        string `json:"user,omitempty"`
					Type        string `json:"type,omitempty"`
					Url         string `json:"url,omitempty"`
					Timestamp   int64  `json:"timestamp,omitempty"`
					ObjectId    string `json:"object_id"`
					ObjectType  string `json:"object_type"`
				}
			*/
			m := evt.AsMap()
			m["object_id"] = path.Join(this.Domain, this.Name, this.Host)
			m["object_type"] = "agent"
			m["timestamp"] = time.Now().Unix()
			m["description"] = fmt.Sprint(m["type"], ":", m["state"], "@", "server=", m["server"])
			m["title"] = "zookeeper event from agent " + path.Join(this.Domain, this.Name, this.Host)
			m["user"] = "dash"
			switch m["state"] {
			case "state-disconnected", "state-auth-failed", "state-expired":
				m["status"] = "fatal"
			case "state-connected", "state-has-session":
				m["status"] = "ok"
			}
			status(m)
		}
	}()

	return nil
}

func (this *Agent) Register() error {
	attempts := 0
	for {
		err := this._register()
		if err != nil {
			if attempts == 12 {
				return err
			} else {
				time.Sleep(5 * time.Second)
				attempts += 1
			}
		} else {
			return nil
		}
	}
}

func (this *Agent) info() interface{} {
	info := map[string]interface{}{
		"api":       fmt.Sprintf("%s:%d", this.Host, this.ListenPort),
		"dockerapi": fmt.Sprintf("http://%s:%d/dockerapi", this.Host, this.ListenPort),
		"version":   *version.BuildInfo(),
	}
	if this.is_running_dockerui() {
		info["dockerui"] = fmt.Sprintf("http://%s:%d/", this.Host, this.DockerUIPort)
	}

	return info
}

func (this *Agent) _register() error {
	if this.zk == nil {
		return ErrNotConnectedToRegistry
	}

	key := registry.NewPath("dash", this.Host)
	err := zk.CreateOrSet(this.zk, key, this.info(), true)
	glog.Infoln("Register key=", key, "err=", err)
	if err == nil {
		// Update this only on successful registration
		this.Registration = key.Path()
	}
	return err
}

// Containers in this domain
func (this *Agent) ListContainers(domain, service string) ([]*docker.Container, error) {
	d, has := this.domains[domain]
	if !has {
		return nil, ErrNoDomain
	}
	return d.ListContainers(ServiceKey(service))
}

func (this *Agent) WatchContainer(domain, service string, spec *WatchContainerSpec) error {
	d, has := this.domains[domain]
	if !has {
		return ErrNoDomain
	}

	this.lock.Lock()
	defer this.lock.Unlock()
	return d.WatchContainer(ServiceKey(service), spec)
}

func (this *Agent) ConfigureDomain(config *DomainConfig) (*Domain, error) {
	if config.Domain == "" {
		return nil, ErrNoDomain
	}

	this.lock.Lock()
	defer this.lock.Unlock()

	domain, has := this.domains[config.Domain]
	if !has {

		domain = &Domain{
			Domain:                 config.Domain,
			RegistryContainerEntry: config.RegistryContainerEntry,
			zk:                 this.zk,
			docker:             this.docker,
			container_watchers: make(map[string]chan<- bool, 0),
			triggers:           NewZkWatcher(this.zk),
			agent:              this,
			tracker:            NewContainerTracker(config.Domain),
			schedulers:         make(map[ServiceKey]*Scheduler),
		}

		_, err := domain.StartScheduleExecutor()
		if err != nil {
			glog.Warningln("Cannot start schedule executor")
			return nil, err
		}

		domain.Host = this.Host
		domain.Port = this.Port

		err = domain.Register()
		if err != nil {
			return nil, err
		}

		this.domains[config.Domain] = domain
		this.domainConfigs[config.Domain] = *config
		domain.Config = config
	}
	return this.domains[config.Domain], nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package cephfs

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os/exec"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"k8s.io/klog"

	"github.com/ceph/ceph-csi/pkg/util"
	"github.com/container-storage-interface/spec/lib/go/csi"
	"k8s.io/kubernetes/pkg/util/mount"
)

type volumeID string

func makeVolumeID(volName string) volumeID {
	return volumeID("csi-cephfs-" + volName)
}

func closePipeOnError(pipe io.Closer, err error) {
	if err != nil {
		if err = pipe.Close(); err != nil {
			klog.Warningf("failed to close pipe: %v", err)
		}
	}
}

func execCommand(program string, args ...string) (stdout, stderr []byte, err error) {
	cmd := exec.Command(program, args...) // nolint: gosec
	klog.V(4).Infof("cephfs: EXEC %s %s", program, util.StripSecretInArgs(args))

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, fmt.Errorf("cannot open stdout pipe for %s %v: %v", program, args, err)
	}

	defer closePipeOnError(stdoutPipe, err)

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, fmt.Errorf("cannot open stdout pipe for %s %v: %v", program, args, err)
	}

	defer closePipeOnError(stderrPipe, err)

	if err = cmd.Start(); err != nil {
		return nil, nil, fmt.Errorf("failed to run %s %v: %v", program, args, err)
	}

	stdout, err = ioutil.ReadAll(stdoutPipe)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read from stdout for %s %v: %v", program, args, err)
	}

	stderr, err = ioutil.ReadAll(stderrPipe)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read from stderr for %s %v: %v", program, args, err)
	}

	if waitErr := cmd.Wait(); waitErr != nil {
		return nil, nil, fmt.Errorf("an error occurred while running %s %v: %v: %s", program, args, waitErr, stderr)
	}

	return
}

func execCommandErr(program string, args ...string) error {
	if _, _, err := execCommand(program, args...); err != nil {
		return err
	}

	return nil
}

func execCommandJSON(v interface{}, program string, args ...string) error {
	stdout, _, err := execCommand(program, args...)
	if err != nil {
		return err
	}

	if err = json.Unmarshal(stdout, v); err != nil {
		return fmt.Errorf("failed to unmarshal JSON for %s %v: %s: %v", program, args, stdout, err)
	}

	return nil
}

// Used in isMountPoint()
var dummyMount = mount.New("")

func isMountPoint(p string) (bool, error) {
	notMnt, err := dummyMount.IsLikelyNotMountPoint(p)
	if err != nil {
		return false, status.Error(codes.Internal, err.Error())
	}

	return !notMnt, nil
}

//
// Controller service request validation
//

func (cs *ControllerServer) validateCreateVolumeRequest(req *csi.CreateVolumeRequest) error {
	if err := cs.Driver.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME); err != nil {
		return fmt.Errorf("invalid CreateVolumeRequest: %v", err)
	}

	if req.GetName() == "" {
		return status.Error(codes.InvalidArgument, "Volume Name cannot be empty")
	}

	reqCaps := req.GetVolumeCapabilities()
	if reqCaps == nil {
		return status.Error(codes.InvalidArgument, "Volume Capabilities cannot be empty")
	}

	for _, cap := range reqCaps {
		if cap.GetBlock() != nil {
			return status.Error(codes.Unimplemented, "block volume not supported")
		}
	}

	return nil
}

func (cs *ControllerServer) validateDeleteVolumeRequest() error {
	if err := cs.Driver.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME); err != nil {
		return fmt.Errorf("invalid DeleteVolumeRequest: %v", err)
	}

	return nil
}

//
// Node service request validation
//

func validateNodeStageVolumeRequest(req *csi.NodeStageVolumeRequest) error {
	if req.GetVolumeCapability() == nil {
		return errors.New("volume capability missing in request")
	}

	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetStagingTargetPath() == "" {
		return errors.New("staging target path missing in request")
	}

	if req.GetSecrets() == nil || len(req.GetSecrets()) == 0 {
		return errors.New("stage secrets cannot be nil or empty")
	}

	return nil
}

func validateNodeUnstageVolumeRequest(req *csi.NodeUnstageVolumeRequest) error {
	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetStagingTargetPath() == "" {
		return errors.New("staging target path missing in request")
	}

	return nil
}

func validateNodePublishVolumeRequest(req *csi.NodePublishVolumeRequest) error {
	if req.GetVolumeCapability() == nil {
		return errors.New("volume capability missing in request")
	}

	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetTargetPath() == "" {
		return errors.New("varget path missing in request")
	}

	return nil
}

func validateNodeUnpublishVolumeRequest(req *csi.NodeUnpublishVolumeRequest) error {
	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetTargetPath() == "" {
		return errors.New("target path missing in request")
	}

	return nil
}
<commit_msg>skip redundant error check<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package cephfs

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os/exec"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"k8s.io/klog"

	"github.com/ceph/ceph-csi/pkg/util"
	"github.com/container-storage-interface/spec/lib/go/csi"
	"k8s.io/kubernetes/pkg/util/mount"
)

type volumeID string

func makeVolumeID(volName string) volumeID {
	return volumeID("csi-cephfs-" + volName)
}

func closePipeOnError(pipe io.Closer, err error) {
	if err != nil {
		if err = pipe.Close(); err != nil {
			klog.Warningf("failed to close pipe: %v", err)
		}
	}
}

func execCommand(program string, args ...string) (stdout, stderr []byte, err error) {
	cmd := exec.Command(program, args...) // nolint: gosec
	klog.V(4).Infof("cephfs: EXEC %s %s", program, util.StripSecretInArgs(args))

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, fmt.Errorf("cannot open stdout pipe for %s %v: %v", program, args, err)
	}

	defer closePipeOnError(stdoutPipe, err)

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, fmt.Errorf("cannot open stdout pipe for %s %v: %v", program, args, err)
	}

	defer closePipeOnError(stderrPipe, err)

	if err = cmd.Start(); err != nil {
		return nil, nil, fmt.Errorf("failed to run %s %v: %v", program, args, err)
	}

	stdout, err = ioutil.ReadAll(stdoutPipe)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read from stdout for %s %v: %v", program, args, err)
	}

	stderr, err = ioutil.ReadAll(stderrPipe)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read from stderr for %s %v: %v", program, args, err)
	}

	if waitErr := cmd.Wait(); waitErr != nil {
		return nil, nil, fmt.Errorf("an error occurred while running %s %v: %v: %s", program, args, waitErr, stderr)
	}

	return
}

func execCommandErr(program string, args ...string) error {
	_, _, err := execCommand(program, args...)
	return err
}

func execCommandJSON(v interface{}, program string, args ...string) error {
	stdout, _, err := execCommand(program, args...)
	if err != nil {
		return err
	}

	if err = json.Unmarshal(stdout, v); err != nil {
		return fmt.Errorf("failed to unmarshal JSON for %s %v: %s: %v", program, args, stdout, err)
	}

	return nil
}

// Used in isMountPoint()
var dummyMount = mount.New("")

func isMountPoint(p string) (bool, error) {
	notMnt, err := dummyMount.IsLikelyNotMountPoint(p)
	if err != nil {
		return false, status.Error(codes.Internal, err.Error())
	}

	return !notMnt, nil
}

// Controller service request validation
func (cs *ControllerServer) validateCreateVolumeRequest(req *csi.CreateVolumeRequest) error {
	if err := cs.Driver.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME); err != nil {
		return fmt.Errorf("invalid CreateVolumeRequest: %v", err)
	}

	if req.GetName() == "" {
		return status.Error(codes.InvalidArgument, "Volume Name cannot be empty")
	}

	reqCaps := req.GetVolumeCapabilities()
	if reqCaps == nil {
		return status.Error(codes.InvalidArgument, "Volume Capabilities cannot be empty")
	}

	for _, cap := range reqCaps {
		if cap.GetBlock() != nil {
			return status.Error(codes.Unimplemented, "block volume not supported")
		}
	}

	return nil
}

func (cs *ControllerServer) validateDeleteVolumeRequest() error {
	if err := cs.Driver.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME); err != nil {
		return fmt.Errorf("invalid DeleteVolumeRequest: %v", err)
	}

	return nil
}

// Node service request validation
func validateNodeStageVolumeRequest(req *csi.NodeStageVolumeRequest) error {
	if req.GetVolumeCapability() == nil {
		return errors.New("volume capability missing in request")
	}

	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetStagingTargetPath() == "" {
		return errors.New("staging target path missing in request")
	}

	if req.GetSecrets() == nil || len(req.GetSecrets()) == 0 {
		return errors.New("stage secrets cannot be nil or empty")
	}

	return nil
}

func validateNodeUnstageVolumeRequest(req *csi.NodeUnstageVolumeRequest) error {
	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetStagingTargetPath() == "" {
		return errors.New("staging target path missing in request")
	}

	return nil
}

func validateNodePublishVolumeRequest(req *csi.NodePublishVolumeRequest) error {
	if req.GetVolumeCapability() == nil {
		return errors.New("volume capability missing in request")
	}

	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetTargetPath() == "" {
		return errors.New("target path missing in request")
	}

	return nil
}

func validateNodeUnpublishVolumeRequest(req *csi.NodeUnpublishVolumeRequest) error {
	if req.GetVolumeId() == "" {
		return errors.New("volume ID missing in request")
	}

	if req.GetTargetPath() == "" {
		return errors.New("target path missing in request")
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Kohei YOSHIDA <https://yosida95.com/>.
//
// 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 chame

import (
	"fmt"
	"mime"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/golang/glog"
)

type Chame struct {
	store  Store
	proxy  Proxy
	ctypes []string
}

func New(cfg *Config) http.Handler {
	chame := &Chame{
		store: cfg.Store,
		proxy: cfg.Proxy,
	}
	if chame.proxy == nil {
		chame.proxy = &HTTPProxy{
			httpCFactory: cfg.NewHTTPClient,
		}
	}

	ctypes := cfg.ProxyContentType
	if ctypes == nil {
		ctypes = defaultProxyContentType
	}
	chame.ctypes = make([]string, len(ctypes))
	for i, ctype := range ctypes {
		chame.ctypes[i] = strings.ToLower(ctype)
	}

	mux := http.NewServeMux()
	mux.Handle("/", cfg.applyInterceptors(http.HandlerFunc(chame.ServeHome)))
	mux.Handle(proxyPrefix, cfg.applyInterceptors(http.HandlerFunc(chame.ServeProxy)))
	return mux
}

func (chame *Chame) ServeHome(w http.ResponseWriter, req *http.Request) {
	emitCommonHeaders(w)
	if req.URL.Path != "/" {
		httpError(w, http.StatusNotFound)
		return
	}
	if !httpErrorIfMethodNotAllowed(w, req, http.MethodGet) {
		return
	}
	w.Header().Set(headerKeyContentType, "text/plain;charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "Hello, this is chame!\nVisit https://github.com/yosida95/chame.")
}

const proxyPrefix = "/i/"

func (chame *Chame) ServeProxy(w http.ResponseWriter, userReq *http.Request) {
	emitCommonHeaders(w)
	if !httpErrorIfMethodNotAllowed(w, userReq, http.MethodGet) {
		return
	}
	ctx, cancel := newProxyContext(userReq.Context())
	defer cancel()

	signedURL := userReq.URL.Path[len(proxyPrefix):]
	decoded, err := DecodeToken(ctx, chame.store, signedURL)
	if err != nil {
		glog.Errorf("chame: failed to decode signed token %q: %v", signedURL, err)
		httpError(w, http.StatusBadRequest)
		return
	}
	reqUrl, err := url.Parse(decoded)
	if err != nil {
		glog.Errorf("chame: malformed URL: %v", err)
		httpError(w, http.StatusBadRequest)
		return
	}

	filtered := make(http.Header)
	copyHeadersOnlyIn(filtered, userReq.Header, passThroughReqHeaders)

	w = chame.newResponseWriter(w)
	chame.proxy.Do(w, &ProxyRequest{
		Context: ctx,
		URL:     reqUrl,
		Header:  filtered,
	})
}

func (chame *Chame) isAcceptableContentType(ctype string) bool {
	// As https://tools.ietf.org/html/rfc2045#section-5.1 said,
	// it is case-insensitive.
	ctype = strings.ToLower(ctype)
	for _, allowed := range chame.ctypes {
		if ctype == allowed {
			return true
		}
	}
	return false
}

type responsewriter struct {
	http.ResponseWriter
	once    sync.Once
	headers http.Header
	discard bool

	isAcceptableContentType func(string) bool
}

var _ http.ResponseWriter = (*responsewriter)(nil)

func (chame *Chame) newResponseWriter(w http.ResponseWriter) *responsewriter {
	return &responsewriter{
		ResponseWriter: w,
		headers:        make(http.Header),

		isAcceptableContentType: chame.isAcceptableContentType,
	}
}

func (w *responsewriter) Header() http.Header { return w.headers }

func (w *responsewriter) WriteHeader(code int) {
	w.once.Do(func() {
		dest := w.ResponseWriter.Header()
		src := w.headers
		copyHeadersOnlyIn(dest, src, passThroughRespHeaders)

		ctype, _, err := mime.ParseMediaType(dest.Get(headerKeyContentType))
		if err != nil || !w.isAcceptableContentType(ctype) {
			glog.Infof("chame: unacceptable Content-Type")
			dest.Del(headerKeyContentType)
			dest.Del(headerKeyContentLength)
			httpError(w.ResponseWriter, http.StatusBadRequest)
			w.discard = true
			return
		}
		w.ResponseWriter.WriteHeader(code)
	})
}

func (w *responsewriter) Write(p []byte) (int, error) {
	w.WriteHeader(http.StatusOK)
	if w.discard {
		return len(p), nil
	}
	return w.ResponseWriter.Write(p)
}
<commit_msg>add ending at the end of `/` response<commit_after>// Copyright 2017 Kohei YOSHIDA <https://yosida95.com/>.
//
// 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 chame

import (
	"fmt"
	"mime"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/golang/glog"
)

type Chame struct {
	store  Store
	proxy  Proxy
	ctypes []string
}

func New(cfg *Config) http.Handler {
	chame := &Chame{
		store: cfg.Store,
		proxy: cfg.Proxy,
	}
	if chame.proxy == nil {
		chame.proxy = &HTTPProxy{
			httpCFactory: cfg.NewHTTPClient,
		}
	}

	ctypes := cfg.ProxyContentType
	if ctypes == nil {
		ctypes = defaultProxyContentType
	}
	chame.ctypes = make([]string, len(ctypes))
	for i, ctype := range ctypes {
		chame.ctypes[i] = strings.ToLower(ctype)
	}

	mux := http.NewServeMux()
	mux.Handle("/", cfg.applyInterceptors(http.HandlerFunc(chame.ServeHome)))
	mux.Handle(proxyPrefix, cfg.applyInterceptors(http.HandlerFunc(chame.ServeProxy)))
	return mux
}

func (chame *Chame) ServeHome(w http.ResponseWriter, req *http.Request) {
	emitCommonHeaders(w)
	if req.URL.Path != "/" {
		httpError(w, http.StatusNotFound)
		return
	}
	if !httpErrorIfMethodNotAllowed(w, req, http.MethodGet) {
		return
	}
	w.Header().Set(headerKeyContentType, "text/plain;charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "Hello, this is chame!\nVisit https://github.com/yosida95/chame.\n")
}

const proxyPrefix = "/i/"

func (chame *Chame) ServeProxy(w http.ResponseWriter, userReq *http.Request) {
	emitCommonHeaders(w)
	if !httpErrorIfMethodNotAllowed(w, userReq, http.MethodGet) {
		return
	}
	ctx, cancel := newProxyContext(userReq.Context())
	defer cancel()

	signedURL := userReq.URL.Path[len(proxyPrefix):]
	decoded, err := DecodeToken(ctx, chame.store, signedURL)
	if err != nil {
		glog.Errorf("chame: failed to decode signed token %q: %v", signedURL, err)
		httpError(w, http.StatusBadRequest)
		return
	}
	reqUrl, err := url.Parse(decoded)
	if err != nil {
		glog.Errorf("chame: malformed URL: %v", err)
		httpError(w, http.StatusBadRequest)
		return
	}

	filtered := make(http.Header)
	copyHeadersOnlyIn(filtered, userReq.Header, passThroughReqHeaders)

	w = chame.newResponseWriter(w)
	chame.proxy.Do(w, &ProxyRequest{
		Context: ctx,
		URL:     reqUrl,
		Header:  filtered,
	})
}

func (chame *Chame) isAcceptableContentType(ctype string) bool {
	// As https://tools.ietf.org/html/rfc2045#section-5.1 said,
	// it is case-insensitive.
	ctype = strings.ToLower(ctype)
	for _, allowed := range chame.ctypes {
		if ctype == allowed {
			return true
		}
	}
	return false
}

type responsewriter struct {
	http.ResponseWriter
	once    sync.Once
	headers http.Header
	discard bool

	isAcceptableContentType func(string) bool
}

var _ http.ResponseWriter = (*responsewriter)(nil)

func (chame *Chame) newResponseWriter(w http.ResponseWriter) *responsewriter {
	return &responsewriter{
		ResponseWriter: w,
		headers:        make(http.Header),

		isAcceptableContentType: chame.isAcceptableContentType,
	}
}

func (w *responsewriter) Header() http.Header { return w.headers }

func (w *responsewriter) WriteHeader(code int) {
	w.once.Do(func() {
		dest := w.ResponseWriter.Header()
		src := w.headers
		copyHeadersOnlyIn(dest, src, passThroughRespHeaders)

		ctype, _, err := mime.ParseMediaType(dest.Get(headerKeyContentType))
		if err != nil || !w.isAcceptableContentType(ctype) {
			glog.Infof("chame: unacceptable Content-Type")
			dest.Del(headerKeyContentType)
			dest.Del(headerKeyContentLength)
			httpError(w.ResponseWriter, http.StatusBadRequest)
			w.discard = true
			return
		}
		w.ResponseWriter.WriteHeader(code)
	})
}

func (w *responsewriter) Write(p []byte) (int, error) {
	w.WriteHeader(http.StatusOK)
	if w.discard {
		return len(p), nil
	}
	return w.ResponseWriter.Write(p)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package drain

import (
	"context"
	"fmt"
	"io"
	"math"
	"time"

	corev1 "k8s.io/api/core/v1"
	policyv1beta1 "k8s.io/api/policy/v1beta1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/cli-runtime/pkg/resource"
	"k8s.io/client-go/kubernetes"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
)

const (
	// EvictionKind represents the kind of evictions object
	EvictionKind = "Eviction"
	// EvictionSubresource represents the kind of evictions object as pod's subresource
	EvictionSubresource = "pods/eviction"
	podSkipMsgTemplate  = "pod %q has DeletionTimestamp older than %v seconds, skipping\n"
)

// Helper contains the parameters to control the behaviour of drainer
type Helper struct {
	Ctx                 context.Context
	Client              kubernetes.Interface
	Force               bool
	GracePeriodSeconds  int
	IgnoreAllDaemonSets bool
	Timeout             time.Duration
	DeleteLocalData     bool
	Selector            string
	PodSelector         string

	// DisableEviction forces drain to use delete rather than evict
	DisableEviction bool

	// SkipWaitForDeleteTimeoutSeconds ignores pods that have a
	// DeletionTimeStamp > N seconds. It's up to the user to decide when this
	// option is appropriate; examples include the Node is unready and the pods
	// won't drain otherwise
	SkipWaitForDeleteTimeoutSeconds int

	Out    io.Writer
	ErrOut io.Writer

	DryRunStrategy cmdutil.DryRunStrategy
	DryRunVerifier *resource.DryRunVerifier

	// OnPodDeletedOrEvicted is called when a pod is evicted/deleted; for printing progress output
	OnPodDeletedOrEvicted func(pod *corev1.Pod, usingEviction bool)
}

type waitForDeleteParams struct {
	ctx                             context.Context
	pods                            []corev1.Pod
	interval                        time.Duration
	timeout                         time.Duration
	usingEviction                   bool
	getPodFn                        func(string, string) (*corev1.Pod, error)
	onDoneFn                        func(pod *corev1.Pod, usingEviction bool)
	globalTimeout                   time.Duration
	skipWaitForDeleteTimeoutSeconds int
	out                             io.Writer
}

// CheckEvictionSupport uses Discovery API to find out if the server support
// eviction subresource If support, it will return its groupVersion; Otherwise,
// it will return an empty string
func CheckEvictionSupport(clientset kubernetes.Interface) (string, error) {
	discoveryClient := clientset.Discovery()
	groupList, err := discoveryClient.ServerGroups()
	if err != nil {
		return "", err
	}
	foundPolicyGroup := false
	var policyGroupVersion string
	for _, group := range groupList.Groups {
		if group.Name == "policy" {
			foundPolicyGroup = true
			policyGroupVersion = group.PreferredVersion.GroupVersion
			break
		}
	}
	if !foundPolicyGroup {
		return "", nil
	}
	resourceList, err := discoveryClient.ServerResourcesForGroupVersion("v1")
	if err != nil {
		return "", err
	}
	for _, resource := range resourceList.APIResources {
		if resource.Name == EvictionSubresource && resource.Kind == EvictionKind {
			return policyGroupVersion, nil
		}
	}
	return "", nil
}

func (d *Helper) makeDeleteOptions() *metav1.DeleteOptions {
	deleteOptions := &metav1.DeleteOptions{}
	if d.GracePeriodSeconds >= 0 {
		gracePeriodSeconds := int64(d.GracePeriodSeconds)
		deleteOptions.GracePeriodSeconds = &gracePeriodSeconds
	}
	if d.DryRunStrategy == cmdutil.DryRunServer {
		deleteOptions.DryRun = []string{metav1.DryRunAll}
	}
	return deleteOptions
}

// DeletePod will delete the given pod, or return an error if it couldn't
func (d *Helper) DeletePod(pod corev1.Pod) error {
	if d.DryRunStrategy == cmdutil.DryRunServer {
		if err := d.DryRunVerifier.HasSupport(pod.GroupVersionKind()); err != nil {
			return err
		}
	}
	return d.Client.CoreV1().Pods(pod.Namespace).Delete(context.TODO(), pod.Name, d.makeDeleteOptions())
}

// EvictPod will evict the give pod, or return an error if it couldn't
func (d *Helper) EvictPod(pod corev1.Pod, policyGroupVersion string) error {
	if d.DryRunStrategy == cmdutil.DryRunServer {
		if err := d.DryRunVerifier.HasSupport(pod.GroupVersionKind()); err != nil {
			return err
		}
	}
	eviction := &policyv1beta1.Eviction{
		TypeMeta: metav1.TypeMeta{
			APIVersion: policyGroupVersion,
			Kind:       EvictionKind,
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      pod.Name,
			Namespace: pod.Namespace,
		},
		DeleteOptions: d.makeDeleteOptions(),
	}
	// Remember to change change the URL manipulation func when Eviction's version change
	return d.Client.PolicyV1beta1().Evictions(eviction.Namespace).Evict(eviction)
}

// GetPodsForDeletion receives resource info for a node, and returns those pods as PodDeleteList,
// or error if it cannot list pods. All pods that are ready to be deleted can be obtained with .Pods(),
// and string with all warning can be obtained with .Warnings(), and .Errors() for all errors that
// occurred during deletion.
func (d *Helper) GetPodsForDeletion(nodeName string) (*podDeleteList, []error) {
	labelSelector, err := labels.Parse(d.PodSelector)
	if err != nil {
		return nil, []error{err}
	}

	podList, err := d.Client.CoreV1().Pods(metav1.NamespaceAll).List(context.TODO(), metav1.ListOptions{
		LabelSelector: labelSelector.String(),
		FieldSelector: fields.SelectorFromSet(fields.Set{"spec.nodeName": nodeName}).String()})
	if err != nil {
		return nil, []error{err}
	}

	pods := []podDelete{}

	for _, pod := range podList.Items {
		var status podDeleteStatus
		for _, filter := range d.makeFilters() {
			status = filter(pod)
			if !status.delete {
				// short-circuit as soon as pod is filtered out
				// at that point, there is no reason to run pod
				// through any additional filters
				break
			}
		}
		// Add the pod to podDeleteList no matter what podDeleteStatus is,
		// those pods whose podDeleteStatus is false like DaemonSet will
		// be catched by list.errors()
		pods = append(pods, podDelete{
			pod:    pod,
			status: status,
		})
	}

	list := &podDeleteList{items: pods}

	if errs := list.errors(); len(errs) > 0 {
		return list, errs
	}

	return list, nil
}

// DeleteOrEvictPods deletes or evicts the pods on the api server
func (d *Helper) DeleteOrEvictPods(pods []corev1.Pod) error {
	if len(pods) == 0 {
		return nil
	}

	// TODO(justinsb): unnecessary?
	getPodFn := func(namespace, name string) (*corev1.Pod, error) {
		return d.Client.CoreV1().Pods(namespace).Get(context.TODO(), name, metav1.GetOptions{})
	}

	if !d.DisableEviction {
		policyGroupVersion, err := CheckEvictionSupport(d.Client)
		if err != nil {
			return err
		}

		if len(policyGroupVersion) > 0 {
			return d.evictPods(pods, policyGroupVersion, getPodFn)
		}
	}

	return d.deletePods(pods, getPodFn)
}

func (d *Helper) evictPods(pods []corev1.Pod, policyGroupVersion string, getPodFn func(namespace, name string) (*corev1.Pod, error)) error {
	returnCh := make(chan error, 1)
	// 0 timeout means infinite, we use MaxInt64 to represent it.
	var globalTimeout time.Duration
	if d.Timeout == 0 {
		globalTimeout = time.Duration(math.MaxInt64)
	} else {
		globalTimeout = d.Timeout
	}
	ctx, cancel := context.WithTimeout(d.getContext(), globalTimeout)
	defer cancel()
	for _, pod := range pods {
		go func(pod corev1.Pod, returnCh chan error) {
			for {
				switch d.DryRunStrategy {
				case cmdutil.DryRunServer:
					fmt.Fprintf(d.Out, "evicting pod %s/%s (server dry run)\n", pod.Namespace, pod.Name)
				default:
					fmt.Fprintf(d.Out, "evicting pod %s/%s\n", pod.Namespace, pod.Name)
				}
				select {
				case <-ctx.Done():
					// return here or we'll leak a goroutine.
					returnCh <- fmt.Errorf("error when evicting pod %q: global timeout reached: %v", pod.Name, globalTimeout)
					return
				default:
				}
				err := d.EvictPod(pod, policyGroupVersion)
				if err == nil {
					break
				} else if apierrors.IsNotFound(err) {
					returnCh <- nil
					return
				} else if apierrors.IsTooManyRequests(err) {
					fmt.Fprintf(d.ErrOut, "error when evicting pod %q (will retry after 5s): %v\n", pod.Name, err)
					time.Sleep(5 * time.Second)
				} else {
					returnCh <- fmt.Errorf("error when evicting pod %q: %v", pod.Name, err)
					return
				}
			}
			if d.DryRunStrategy == cmdutil.DryRunServer {
				returnCh <- nil
				return
			}
			params := waitForDeleteParams{
				ctx:                             ctx,
				pods:                            []corev1.Pod{pod},
				interval:                        1 * time.Second,
				timeout:                         time.Duration(math.MaxInt64),
				usingEviction:                   true,
				getPodFn:                        getPodFn,
				onDoneFn:                        d.OnPodDeletedOrEvicted,
				globalTimeout:                   globalTimeout,
				skipWaitForDeleteTimeoutSeconds: d.SkipWaitForDeleteTimeoutSeconds,
				out:                             d.Out,
			}
			_, err := waitForDelete(params)
			if err == nil {
				returnCh <- nil
			} else {
				returnCh <- fmt.Errorf("error when waiting for pod %q terminating: %v", pod.Name, err)
			}
		}(pod, returnCh)
	}

	doneCount := 0
	var errors []error

	numPods := len(pods)
	for doneCount < numPods {
		select {
		case err := <-returnCh:
			doneCount++
			if err != nil {
				errors = append(errors, err)
			}
		default:
		}
	}

	return utilerrors.NewAggregate(errors)
}

func (d *Helper) deletePods(pods []corev1.Pod, getPodFn func(namespace, name string) (*corev1.Pod, error)) error {
	// 0 timeout means infinite, we use MaxInt64 to represent it.
	var globalTimeout time.Duration
	if d.Timeout == 0 {
		globalTimeout = time.Duration(math.MaxInt64)
	} else {
		globalTimeout = d.Timeout
	}
	for _, pod := range pods {
		err := d.DeletePod(pod)
		if err != nil && !apierrors.IsNotFound(err) {
			return err
		}
	}
	ctx := d.getContext()
	params := waitForDeleteParams{
		ctx:                             ctx,
		pods:                            pods,
		interval:                        1 * time.Second,
		timeout:                         globalTimeout,
		usingEviction:                   false,
		getPodFn:                        getPodFn,
		onDoneFn:                        d.OnPodDeletedOrEvicted,
		globalTimeout:                   globalTimeout,
		skipWaitForDeleteTimeoutSeconds: d.SkipWaitForDeleteTimeoutSeconds,
		out:                             d.Out,
	}
	_, err := waitForDelete(params)
	return err
}

func waitForDelete(params waitForDeleteParams) ([]corev1.Pod, error) {
	pods := params.pods
	err := wait.PollImmediate(params.interval, params.timeout, func() (bool, error) {
		pendingPods := []corev1.Pod{}
		for i, pod := range pods {
			p, err := params.getPodFn(pod.Namespace, pod.Name)
			if apierrors.IsNotFound(err) || (p != nil && p.ObjectMeta.UID != pod.ObjectMeta.UID) {
				if params.onDoneFn != nil {
					params.onDoneFn(&pod, params.usingEviction)
				}
				continue
			} else if err != nil {
				return false, err
			} else {
				if shouldSkipPod(*p, params.skipWaitForDeleteTimeoutSeconds) {
					fmt.Fprintf(params.out, podSkipMsgTemplate, pod.Name, params.skipWaitForDeleteTimeoutSeconds)
					continue
				}
				pendingPods = append(pendingPods, pods[i])
			}
		}
		pods = pendingPods
		if len(pendingPods) > 0 {
			select {
			case <-params.ctx.Done():
				return false, fmt.Errorf("global timeout reached: %v", params.globalTimeout)
			default:
				return false, nil
			}
		}
		return true, nil
	})
	return pods, err
}

// Since Helper does not have a constructor, we can't enforce Helper.Ctx != nil
// Multiple public methods prevent us from initializing the context in a single
// place as well.
func (d *Helper) getContext() context.Context {
	if d.Ctx != nil {
		return d.Ctx
	}
	return context.Background()
}
<commit_msg>more artisanal fixes<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package drain

import (
	"context"
	"fmt"
	"io"
	"math"
	"time"

	corev1 "k8s.io/api/core/v1"
	policyv1beta1 "k8s.io/api/policy/v1beta1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/cli-runtime/pkg/resource"
	"k8s.io/client-go/kubernetes"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
)

const (
	// EvictionKind represents the kind of evictions object
	EvictionKind = "Eviction"
	// EvictionSubresource represents the kind of evictions object as pod's subresource
	EvictionSubresource = "pods/eviction"
	podSkipMsgTemplate  = "pod %q has DeletionTimestamp older than %v seconds, skipping\n"
)

// Helper contains the parameters to control the behaviour of drainer
type Helper struct {
	Ctx                 context.Context
	Client              kubernetes.Interface
	Force               bool
	GracePeriodSeconds  int
	IgnoreAllDaemonSets bool
	Timeout             time.Duration
	DeleteLocalData     bool
	Selector            string
	PodSelector         string

	// DisableEviction forces drain to use delete rather than evict
	DisableEviction bool

	// SkipWaitForDeleteTimeoutSeconds ignores pods that have a
	// DeletionTimeStamp > N seconds. It's up to the user to decide when this
	// option is appropriate; examples include the Node is unready and the pods
	// won't drain otherwise
	SkipWaitForDeleteTimeoutSeconds int

	Out    io.Writer
	ErrOut io.Writer

	DryRunStrategy cmdutil.DryRunStrategy
	DryRunVerifier *resource.DryRunVerifier

	// OnPodDeletedOrEvicted is called when a pod is evicted/deleted; for printing progress output
	OnPodDeletedOrEvicted func(pod *corev1.Pod, usingEviction bool)
}

type waitForDeleteParams struct {
	ctx                             context.Context
	pods                            []corev1.Pod
	interval                        time.Duration
	timeout                         time.Duration
	usingEviction                   bool
	getPodFn                        func(string, string) (*corev1.Pod, error)
	onDoneFn                        func(pod *corev1.Pod, usingEviction bool)
	globalTimeout                   time.Duration
	skipWaitForDeleteTimeoutSeconds int
	out                             io.Writer
}

// CheckEvictionSupport uses Discovery API to find out if the server support
// eviction subresource If support, it will return its groupVersion; Otherwise,
// it will return an empty string
func CheckEvictionSupport(clientset kubernetes.Interface) (string, error) {
	discoveryClient := clientset.Discovery()
	groupList, err := discoveryClient.ServerGroups()
	if err != nil {
		return "", err
	}
	foundPolicyGroup := false
	var policyGroupVersion string
	for _, group := range groupList.Groups {
		if group.Name == "policy" {
			foundPolicyGroup = true
			policyGroupVersion = group.PreferredVersion.GroupVersion
			break
		}
	}
	if !foundPolicyGroup {
		return "", nil
	}
	resourceList, err := discoveryClient.ServerResourcesForGroupVersion("v1")
	if err != nil {
		return "", err
	}
	for _, resource := range resourceList.APIResources {
		if resource.Name == EvictionSubresource && resource.Kind == EvictionKind {
			return policyGroupVersion, nil
		}
	}
	return "", nil
}

func (d *Helper) makeDeleteOptions() metav1.DeleteOptions {
	deleteOptions := metav1.DeleteOptions{}
	if d.GracePeriodSeconds >= 0 {
		gracePeriodSeconds := int64(d.GracePeriodSeconds)
		deleteOptions.GracePeriodSeconds = &gracePeriodSeconds
	}
	if d.DryRunStrategy == cmdutil.DryRunServer {
		deleteOptions.DryRun = []string{metav1.DryRunAll}
	}
	return deleteOptions
}

// DeletePod will delete the given pod, or return an error if it couldn't
func (d *Helper) DeletePod(pod corev1.Pod) error {
	if d.DryRunStrategy == cmdutil.DryRunServer {
		if err := d.DryRunVerifier.HasSupport(pod.GroupVersionKind()); err != nil {
			return err
		}
	}
	return d.Client.CoreV1().Pods(pod.Namespace).Delete(context.TODO(), pod.Name, d.makeDeleteOptions())
}

// EvictPod will evict the give pod, or return an error if it couldn't
func (d *Helper) EvictPod(pod corev1.Pod, policyGroupVersion string) error {
	if d.DryRunStrategy == cmdutil.DryRunServer {
		if err := d.DryRunVerifier.HasSupport(pod.GroupVersionKind()); err != nil {
			return err
		}
	}

	delOpts := d.makeDeleteOptions()
	eviction := &policyv1beta1.Eviction{
		TypeMeta: metav1.TypeMeta{
			APIVersion: policyGroupVersion,
			Kind:       EvictionKind,
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      pod.Name,
			Namespace: pod.Namespace,
		},
		DeleteOptions: &delOpts,
	}

	// Remember to change change the URL manipulation func when Eviction's version change
	return d.Client.PolicyV1beta1().Evictions(eviction.Namespace).Evict(eviction)
}

// GetPodsForDeletion receives resource info for a node, and returns those pods as PodDeleteList,
// or error if it cannot list pods. All pods that are ready to be deleted can be obtained with .Pods(),
// and string with all warning can be obtained with .Warnings(), and .Errors() for all errors that
// occurred during deletion.
func (d *Helper) GetPodsForDeletion(nodeName string) (*podDeleteList, []error) {
	labelSelector, err := labels.Parse(d.PodSelector)
	if err != nil {
		return nil, []error{err}
	}

	podList, err := d.Client.CoreV1().Pods(metav1.NamespaceAll).List(context.TODO(), metav1.ListOptions{
		LabelSelector: labelSelector.String(),
		FieldSelector: fields.SelectorFromSet(fields.Set{"spec.nodeName": nodeName}).String()})
	if err != nil {
		return nil, []error{err}
	}

	pods := []podDelete{}

	for _, pod := range podList.Items {
		var status podDeleteStatus
		for _, filter := range d.makeFilters() {
			status = filter(pod)
			if !status.delete {
				// short-circuit as soon as pod is filtered out
				// at that point, there is no reason to run pod
				// through any additional filters
				break
			}
		}
		// Add the pod to podDeleteList no matter what podDeleteStatus is,
		// those pods whose podDeleteStatus is false like DaemonSet will
		// be catched by list.errors()
		pods = append(pods, podDelete{
			pod:    pod,
			status: status,
		})
	}

	list := &podDeleteList{items: pods}

	if errs := list.errors(); len(errs) > 0 {
		return list, errs
	}

	return list, nil
}

// DeleteOrEvictPods deletes or evicts the pods on the api server
func (d *Helper) DeleteOrEvictPods(pods []corev1.Pod) error {
	if len(pods) == 0 {
		return nil
	}

	// TODO(justinsb): unnecessary?
	getPodFn := func(namespace, name string) (*corev1.Pod, error) {
		return d.Client.CoreV1().Pods(namespace).Get(context.TODO(), name, metav1.GetOptions{})
	}

	if !d.DisableEviction {
		policyGroupVersion, err := CheckEvictionSupport(d.Client)
		if err != nil {
			return err
		}

		if len(policyGroupVersion) > 0 {
			return d.evictPods(pods, policyGroupVersion, getPodFn)
		}
	}

	return d.deletePods(pods, getPodFn)
}

func (d *Helper) evictPods(pods []corev1.Pod, policyGroupVersion string, getPodFn func(namespace, name string) (*corev1.Pod, error)) error {
	returnCh := make(chan error, 1)
	// 0 timeout means infinite, we use MaxInt64 to represent it.
	var globalTimeout time.Duration
	if d.Timeout == 0 {
		globalTimeout = time.Duration(math.MaxInt64)
	} else {
		globalTimeout = d.Timeout
	}
	ctx, cancel := context.WithTimeout(d.getContext(), globalTimeout)
	defer cancel()
	for _, pod := range pods {
		go func(pod corev1.Pod, returnCh chan error) {
			for {
				switch d.DryRunStrategy {
				case cmdutil.DryRunServer:
					fmt.Fprintf(d.Out, "evicting pod %s/%s (server dry run)\n", pod.Namespace, pod.Name)
				default:
					fmt.Fprintf(d.Out, "evicting pod %s/%s\n", pod.Namespace, pod.Name)
				}
				select {
				case <-ctx.Done():
					// return here or we'll leak a goroutine.
					returnCh <- fmt.Errorf("error when evicting pod %q: global timeout reached: %v", pod.Name, globalTimeout)
					return
				default:
				}
				err := d.EvictPod(pod, policyGroupVersion)
				if err == nil {
					break
				} else if apierrors.IsNotFound(err) {
					returnCh <- nil
					return
				} else if apierrors.IsTooManyRequests(err) {
					fmt.Fprintf(d.ErrOut, "error when evicting pod %q (will retry after 5s): %v\n", pod.Name, err)
					time.Sleep(5 * time.Second)
				} else {
					returnCh <- fmt.Errorf("error when evicting pod %q: %v", pod.Name, err)
					return
				}
			}
			if d.DryRunStrategy == cmdutil.DryRunServer {
				returnCh <- nil
				return
			}
			params := waitForDeleteParams{
				ctx:                             ctx,
				pods:                            []corev1.Pod{pod},
				interval:                        1 * time.Second,
				timeout:                         time.Duration(math.MaxInt64),
				usingEviction:                   true,
				getPodFn:                        getPodFn,
				onDoneFn:                        d.OnPodDeletedOrEvicted,
				globalTimeout:                   globalTimeout,
				skipWaitForDeleteTimeoutSeconds: d.SkipWaitForDeleteTimeoutSeconds,
				out:                             d.Out,
			}
			_, err := waitForDelete(params)
			if err == nil {
				returnCh <- nil
			} else {
				returnCh <- fmt.Errorf("error when waiting for pod %q terminating: %v", pod.Name, err)
			}
		}(pod, returnCh)
	}

	doneCount := 0
	var errors []error

	numPods := len(pods)
	for doneCount < numPods {
		select {
		case err := <-returnCh:
			doneCount++
			if err != nil {
				errors = append(errors, err)
			}
		default:
		}
	}

	return utilerrors.NewAggregate(errors)
}

func (d *Helper) deletePods(pods []corev1.Pod, getPodFn func(namespace, name string) (*corev1.Pod, error)) error {
	// 0 timeout means infinite, we use MaxInt64 to represent it.
	var globalTimeout time.Duration
	if d.Timeout == 0 {
		globalTimeout = time.Duration(math.MaxInt64)
	} else {
		globalTimeout = d.Timeout
	}
	for _, pod := range pods {
		err := d.DeletePod(pod)
		if err != nil && !apierrors.IsNotFound(err) {
			return err
		}
	}
	ctx := d.getContext()
	params := waitForDeleteParams{
		ctx:                             ctx,
		pods:                            pods,
		interval:                        1 * time.Second,
		timeout:                         globalTimeout,
		usingEviction:                   false,
		getPodFn:                        getPodFn,
		onDoneFn:                        d.OnPodDeletedOrEvicted,
		globalTimeout:                   globalTimeout,
		skipWaitForDeleteTimeoutSeconds: d.SkipWaitForDeleteTimeoutSeconds,
		out:                             d.Out,
	}
	_, err := waitForDelete(params)
	return err
}

func waitForDelete(params waitForDeleteParams) ([]corev1.Pod, error) {
	pods := params.pods
	err := wait.PollImmediate(params.interval, params.timeout, func() (bool, error) {
		pendingPods := []corev1.Pod{}
		for i, pod := range pods {
			p, err := params.getPodFn(pod.Namespace, pod.Name)
			if apierrors.IsNotFound(err) || (p != nil && p.ObjectMeta.UID != pod.ObjectMeta.UID) {
				if params.onDoneFn != nil {
					params.onDoneFn(&pod, params.usingEviction)
				}
				continue
			} else if err != nil {
				return false, err
			} else {
				if shouldSkipPod(*p, params.skipWaitForDeleteTimeoutSeconds) {
					fmt.Fprintf(params.out, podSkipMsgTemplate, pod.Name, params.skipWaitForDeleteTimeoutSeconds)
					continue
				}
				pendingPods = append(pendingPods, pods[i])
			}
		}
		pods = pendingPods
		if len(pendingPods) > 0 {
			select {
			case <-params.ctx.Done():
				return false, fmt.Errorf("global timeout reached: %v", params.globalTimeout)
			default:
				return false, nil
			}
		}
		return true, nil
	})
	return pods, err
}

// Since Helper does not have a constructor, we can't enforce Helper.Ctx != nil
// Multiple public methods prevent us from initializing the context in a single
// place as well.
func (d *Helper) getContext() context.Context {
	if d.Ctx != nil {
		return d.Ctx
	}
	return context.Background()
}
<|endoftext|>


<commit_before>package ec2

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strconv"
	"strings"

	// Community:
	log "github.com/Sirupsen/logrus"
	"github.com/katosys/kato/pkg/tools"
)

//-----------------------------------------------------------------------------
// func: Add
//-----------------------------------------------------------------------------

// Add a new instance to the cluster.
func (d *Data) Add() {

	// Set current command:
	d.command = "add"

	// Load state from state file:
	if err := d.loadState(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Retrieve the CoreOS AMI ID:
	var err error
	if d.AmiID, err = d.retrieveCoreOSAmiID(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Execute the udata|run pipeline:
	if _, err := tools.ExecutePipeline(
		d.forgeUdataCommand(), d.forgeRunCommand()); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Publish DNS records:
	if err := publishDNSRecords("instanceID", d.Roles); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Warning(err)
	}
}

//-----------------------------------------------------------------------------
// func: retrieveCoreOSAmiID
//-----------------------------------------------------------------------------

func (d *Data) retrieveCoreOSAmiID() (string, error) {

	// Send the request:
	res, err := http.Get("https://coreos.com/dist/aws/aws-" +
		d.CoreOSChannel + ".json")
	if err != nil {
		return "", err
	}

	// Retrieve the data:
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}

	// Close the handler:
	if err = res.Body.Close(); err != nil {
		return "", err
	}

	// Decode JSON into Go values:
	var jsonData map[string]interface{}
	if err := json.Unmarshal(data, &jsonData); err != nil {
		return "", err
	}

	// Store the AMI ID:
	amis := jsonData[d.Region].(map[string]interface{})
	amiID := amis["hvm"].(string)

	// Log this action:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": amiID}).
		Info("Latest CoreOS " + d.CoreOSChannel + " AMI located")

	// Return the AMI ID:
	return amiID, nil
}

//-----------------------------------------------------------------------------
// func: forgeUdataCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeUdataCommand() *exec.Cmd {

	// Udata arguments bundle:
	args := []string{"udata",
		"--roles", d.Roles,
		"--cluster-id", d.ClusterID,
		"--cluster-state", d.ClusterState,
		"--quorum-count", strconv.Itoa(d.QuorumCount),
		"--master-count", strconv.Itoa(d.MasterCount),
		"--host-name", d.HostName,
		"--host-id", d.HostID,
		"--domain", d.Domain,
		"--ec2-region", d.Region,
		"--dns-provider", d.DNSProvider,
		"--dns-api-key", d.DNSApiKey,
		"--etcd-token", d.EtcdToken,
		"--calico-ip-pool", d.CalicoIPPool,
		"--rexray-storage-driver", "ebs",
		"--iaas-provider", "ec2",
		"--prometheus",
		"--gzip-udata",
	}

	// Append flags if present:
	if d.SysdigAccessKey != "" {
		args = append(args, "--sysdig-access-key", d.SysdigAccessKey)
	}
	if d.DatadogAPIKey != "" {
		args = append(args, "--datadog-api-key", d.DatadogAPIKey)
	}
	if d.SlackWebhook != "" {
		args = append(args, "--slack-webhook", d.SlackWebhook)
	}
	if d.CaCertPath != "" {
		args = append(args, "--ca-cert-path", d.CaCertPath)
	}
	for _, z := range d.StubZones {
		args = append(args, "--stub-zone", z)
	}
	if d.SMTPURL != "" {
		args = append(args, "--smtp-url", d.SMTPURL)
	}
	if d.AdminEmail != "" {
		args = append(args, "--admin-email", d.AdminEmail)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: forgeRunCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeRunCommand() *exec.Cmd {

	// Ec2 run arguments bundle:
	args := []string{"ec2", "run",
		"--tag-name", d.HostName + "-" + d.HostID + "." + d.Domain,
		"--region", d.Region,
		"--zone", d.Zone,
		"--ami-id", d.AmiID,
		"--instance-type", d.InstanceType,
		"--key-pair", d.KeyPair,
		"--subnet-id", d.ExtSubnetID,
		"--security-group-ids", strings.Join(d.securityGroupIDs(d.Roles), ","),
		"--iam-role", "kato",
		"--source-dest-check", "false",
		"--public-ip", "true",
	}

	// Append flags if present:
	if strings.Contains(d.Roles, "master") {
		i, _ := strconv.Atoi(d.HostID)
		args = append(args, "--private-ip", tools.OffsetIP(d.ExtSubnetCidr, 10+i))
	}
	if strings.Contains(d.Roles, "worker") {
		args = append(args, "--elb-name", d.ClusterID)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: securityGroupIDs
//-----------------------------------------------------------------------------

func (d *Data) securityGroupIDs(roles string) (list []string) {
	for _, role := range strings.Split(roles, ",") {
		switch role {
		case "quorum":
			list = append(list, d.QuorumSecGrp)
		case "master":
			list = append(list, d.MasterSecGrp)
		case "worker":
			list = append(list, d.WorkerSecGrp)
		case "border":
			list = append(list, d.BorderSecGrp)
		}
	}
	return
}

//-----------------------------------------------------------------------------
// func: publishDNSRecords
//-----------------------------------------------------------------------------

func publishDNSRecords(instance, roles string) error {

	// Retrieve the instance IPs:

	// For every role in this instance:
	for _, role := range strings.Split(roles, ",") {
		fmt.Println(role)
	}

	// Success:
	return nil
}
<commit_msg>Function publishDNSRecords (WiP)<commit_after>package ec2

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strconv"
	"strings"

	// Community:
	log "github.com/Sirupsen/logrus"
	"github.com/katosys/kato/pkg/tools"
)

//-----------------------------------------------------------------------------
// func: Add
//-----------------------------------------------------------------------------

// Add a new instance to the cluster.
func (d *Data) Add() {

	// Set current command:
	d.command = "add"

	// Load state from state file:
	if err := d.loadState(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Retrieve the CoreOS AMI ID:
	var err error
	if d.AmiID, err = d.retrieveCoreOSAmiID(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Execute the udata|run pipeline:
	out, err := tools.ExecutePipeline(d.forgeUdataCommand(), d.forgeRunCommand())
	if err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Publish DNS records:
	if err := d.publishDNSRecords(d.Roles, out); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Warning(err)
	}
}

//-----------------------------------------------------------------------------
// func: retrieveCoreOSAmiID
//-----------------------------------------------------------------------------

func (d *Data) retrieveCoreOSAmiID() (string, error) {

	// Send the request:
	res, err := http.Get("https://coreos.com/dist/aws/aws-" +
		d.CoreOSChannel + ".json")
	if err != nil {
		return "", err
	}

	// Retrieve the data:
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}

	// Close the handler:
	if err = res.Body.Close(); err != nil {
		return "", err
	}

	// Decode JSON into Go values:
	var jsonData map[string]interface{}
	if err := json.Unmarshal(data, &jsonData); err != nil {
		return "", err
	}

	// Store the AMI ID:
	amis := jsonData[d.Region].(map[string]interface{})
	amiID := amis["hvm"].(string)

	// Log this action:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": amiID}).
		Info("Latest CoreOS " + d.CoreOSChannel + " AMI located")

	// Return the AMI ID:
	return amiID, nil
}

//-----------------------------------------------------------------------------
// func: forgeUdataCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeUdataCommand() *exec.Cmd {

	// Udata arguments bundle:
	args := []string{"udata",
		"--roles", d.Roles,
		"--cluster-id", d.ClusterID,
		"--cluster-state", d.ClusterState,
		"--quorum-count", strconv.Itoa(d.QuorumCount),
		"--master-count", strconv.Itoa(d.MasterCount),
		"--host-name", d.HostName,
		"--host-id", d.HostID,
		"--domain", d.Domain,
		"--ec2-region", d.Region,
		"--dns-provider", d.DNSProvider,
		"--dns-api-key", d.DNSApiKey,
		"--etcd-token", d.EtcdToken,
		"--calico-ip-pool", d.CalicoIPPool,
		"--rexray-storage-driver", "ebs",
		"--iaas-provider", "ec2",
		"--prometheus",
		"--gzip-udata",
	}

	// Append flags if present:
	if d.SysdigAccessKey != "" {
		args = append(args, "--sysdig-access-key", d.SysdigAccessKey)
	}
	if d.DatadogAPIKey != "" {
		args = append(args, "--datadog-api-key", d.DatadogAPIKey)
	}
	if d.SlackWebhook != "" {
		args = append(args, "--slack-webhook", d.SlackWebhook)
	}
	if d.CaCertPath != "" {
		args = append(args, "--ca-cert-path", d.CaCertPath)
	}
	for _, z := range d.StubZones {
		args = append(args, "--stub-zone", z)
	}
	if d.SMTPURL != "" {
		args = append(args, "--smtp-url", d.SMTPURL)
	}
	if d.AdminEmail != "" {
		args = append(args, "--admin-email", d.AdminEmail)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: forgeRunCommand
//-----------------------------------------------------------------------------

func (d *Data) forgeRunCommand() *exec.Cmd {

	// Ec2 run arguments bundle:
	args := []string{"ec2", "run",
		"--tag-name", d.HostName + "-" + d.HostID + "." + d.Domain,
		"--region", d.Region,
		"--zone", d.Zone,
		"--ami-id", d.AmiID,
		"--instance-type", d.InstanceType,
		"--key-pair", d.KeyPair,
		"--subnet-id", d.ExtSubnetID,
		"--security-group-ids", strings.Join(d.securityGroupIDs(d.Roles), ","),
		"--iam-role", "kato",
		"--source-dest-check", "false",
		"--public-ip", "true",
	}

	// Append flags if present:
	if strings.Contains(d.Roles, "master") {
		i, _ := strconv.Atoi(d.HostID)
		args = append(args, "--private-ip", tools.OffsetIP(d.ExtSubnetCidr, 10+i))
	}
	if strings.Contains(d.Roles, "worker") {
		args = append(args, "--elb-name", d.ClusterID)
	}

	// Forge the command and return:
	return exec.Command("katoctl", args...)
}

//-----------------------------------------------------------------------------
// func: securityGroupIDs
//-----------------------------------------------------------------------------

func (d *Data) securityGroupIDs(roles string) (list []string) {
	for _, role := range strings.Split(roles, ",") {
		switch role {
		case "quorum":
			list = append(list, d.QuorumSecGrp)
		case "master":
			list = append(list, d.MasterSecGrp)
		case "worker":
			list = append(list, d.WorkerSecGrp)
		case "border":
			list = append(list, d.BorderSecGrp)
		}
	}
	return
}

//-----------------------------------------------------------------------------
// func: publishDNSRecords
//-----------------------------------------------------------------------------

func (d *Data) publishDNSRecords(roles string, out []byte) error {

	// Retrieve the instance IPs:
	var dat map[string]string
	if err := json.Unmarshal(out, &dat); err != nil {
		return err
	}

	// For every role in this instance:
	for _, role := range strings.Split(roles, ",") {

		log.Info("katoctl " + d.DNSProvider +
			"--api-key " + d.DNSApiKey +
			"record add --zone " + "int." +
			d.Domain + " " + dat["internal"] +
			":A:" + role + "-" + d.HostID)

		log.Info("katoctl " + d.DNSProvider +
			"--api-key " + d.DNSApiKey +
			"record add --zone " + "ext." +
			d.Domain + " " + dat["external"] +
			":A:" + role + "-" + d.HostID)
	}

	return nil
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package elf

import (
	"bufio"
	"debug/elf"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
)

const (
	dataSection   = ".data"
	mapSection    = "maps"
	btfMapSection = ".maps"

	nullTerminator     = byte(0)
	relocSectionPrefix = ".rel"

	// nMapRelocations is an approximation of the number of offsets in an
	// ELF relating to map names that need to be adjusted. It's used for
	// map initialization within 'Symbols'.
	//
	// $ readelf -a bpf/bpf_lxc.o | grep "cilium_.*_" | grep "^0000" | wc -l
	// 51
	nMapRelocations = 64
)

type symbolKind uint32

const (
	symbolData   = symbolKind(1)
	symbolString = symbolKind(2)
)

func (k symbolKind) String() string {
	switch k {
	case symbolData:
		return "data"
	case symbolString:
		return "string"
	}
	return "unknown"
}

// symbol stores the location and type of a symbol within the ELF file.
type symbol struct {
	name      string
	kind      symbolKind
	offset    uint64
	offsetBTF uint64
	size      uint64
}

func newSymbol(name string, kind symbolKind, offset, size uint64) symbol {
	return symbol{
		name:   name,
		kind:   kind,
		offset: offset,
		size:   size,
	}
}

func newVariable(name string, offset, size uint64) symbol {
	return newSymbol(name, symbolData, offset, size)
}

func newString(name string, offset uint64) symbol {
	return newSymbol(name, symbolString, offset, uint64(len(name)))
}

type symbolSlice []symbol

// sort a slice of symbols by offset.
func (c symbolSlice) sort() symbolSlice {
	sort.Slice(c, func(i, j int) bool { return c[i].offset < c[j].offset })
	return c
}

type symbols struct {
	// data caches static 32-bit variables by name.
	data map[string]symbol
	// strings caches string symbols by name.
	strings map[string]symbol
}

func (s *symbols) sort() symbolSlice {
	result := make(symbolSlice, 0)
	for _, c := range s.data {
		result = append(result, c)
	}
	for _, c := range s.strings {
		result = append(result, c)
	}
	return result.sort()
}

// isGlobalData returns true if the symbol meets all of the following criteria:
// - symbol is of type STT_NOTYPE or STT_OBJECT
// - symbol has a global binding (STB_GLOBAL)
// - symbol has default visibility (STV_DEFAULT)
func isGlobalData(sym elf.Symbol) bool {
	return (elf.ST_TYPE(sym.Info) == elf.STT_NOTYPE ||
		elf.ST_TYPE(sym.Info) == elf.STT_OBJECT) &&
		elf.ST_BIND(sym.Info) == elf.STB_GLOBAL &&
		elf.ST_VISIBILITY(sym.Other) == elf.STV_DEFAULT
}

func readStringOffset(e *elf.File, r io.ReadSeeker, symbolOffset int64) (uint64, error) {
	if _, err := r.Seek(symbolOffset, io.SeekStart); err != nil {
		return 0, err
	}

	switch e.Class {
	case elf.ELFCLASS32:
		var sym32 elf.Sym32
		if err := binary.Read(r, e.ByteOrder, &sym32); err != nil {
			return 0, err
		}
		return uint64(sym32.Name), nil
	case elf.ELFCLASS64:
		var sym64 elf.Sym64
		if err := binary.Read(r, e.ByteOrder, &sym64); err != nil {
			return 0, err
		}
		return uint64(sym64.Name), nil
	}
	return 0, fmt.Errorf("unsupported ELF type %d", e.Class)
}

// extractFrom processes the specified ELF and populates the received symbols
// object with data and string offsets in the file, for later substitution.
func (s *symbols) extractFrom(e *elf.File) error {
	dataOffsets := make(map[string]symbol)
	stringOffsets := make(map[string]symbol, nMapRelocations)

	symbols, err := e.Symbols()
	if err != nil {
		return err
	}
	symtab := e.SectionByType(elf.SHT_SYMTAB)
	strtab := e.Sections[symtab.Link]

	// Scan symbol table for offsets of static data and symbol names.
	symbolReader := symtab.Open()
	for i, sym := range symbols {
		// BTF extensions like line info not recognized by normal ELF parsers
		if elf.ST_TYPE(sym.Info) == elf.STT_FILE {
			continue
		}

		// Get ELF section reference by its index encoded in the symbol.
		section := e.Sections[sym.Section]

		switch {
		case section.Flags&elf.SHF_COMPRESSED > 0:
			return fmt.Errorf("compressed %s section not supported", section.Name)

		// Skip local symbols that are objects or untyped. These are usually
		// program segments referred to using long jumps in bytecode.
		case !isGlobalData(sym):
			// LBB is a common llvm symbol prefix (basic block);
			// Don't flood the logs with messages about it.
			if !strings.HasPrefix(sym.Name, "LBB") {
				log.Debugf("Skipping %s", sym.Name)
			}
			continue

		// Find the absolute offset to the variable's value this symbol represents
		// in the .data section.
		// This implements substitution of static data.
		case section.Name == dataSection:
			// Offset from start of binary to variable inside .data
			offset := section.Offset + sym.Value
			dataOffsets[sym.Name] = newVariable(sym.Name, offset, sym.Size)
			log.WithField(fieldSymbol, sym.Name).Debugf("Found variable with offset %d", offset)

		// Find the absolute offset to the map's entry in .strtab.
		// This implements renaming maps.
		case section.Name == mapSection || section.Name == btfMapSection:
			// From the Golang Documentation:
			//   "For compatibility with Go 1.0, Symbols omits the
			//   the null symbol at index 0."
			// We must reverse this when reading directly.
			symbolOffset := int64(i+1) * int64(symtab.Entsize)
			symOffsetInStrtab, err := readStringOffset(e, symbolReader, symbolOffset)
			if err != nil {
				return err
			}

			// Offset from start of binary to name inside .strtab
			symOffset := strtab.Offset + symOffsetInStrtab
			stringOffsets[sym.Name] = newString(sym.Name, symOffset)
			log.WithField(fieldSymbol, sym.Name).Debugf("Found symbol with offset %d", symOffset)

		default:
			log.WithField(fieldSymbol, sym.Name).Debugf("Found symbol referring to unknown section id %d", sym.Section)
		}
	}

	// Scan string table for offsets of section names.
	stringReader := bufio.NewReader(strtab.Open())
	var elfString string
	for off := uint64(0); off < strtab.Size; off += uint64(len(elfString)) {
		// off is the offset within the string table.
		elfString, err = stringReader.ReadString(nullTerminator)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return err
		}

		// We only need to worry about sections with relocations.
		if !strings.HasPrefix(elfString, relocSectionPrefix) {
			continue
		}

		elfEnd := len(elfString)
		if elfString[elfEnd-1] == nullTerminator {
			elfEnd--
		}
		relocOffset := uint64(len(relocSectionPrefix))
		secName := elfString[relocOffset:elfEnd]
		if sec := e.Section(secName); sec != nil {
			// Offset from start of binary to name inside .strtab
			globalOffset := strtab.Offset + off + relocOffset
			stringOffsets[secName] = newString(secName, globalOffset)
			log.WithField(fieldSymbol, secName).Debugf("Found section with offset %d", globalOffset)
		}
	}

	// If .BTF section is present in the ELF, get the offset of each symbol
	// into the BTF string table to be able to overwrite them later.
	btfSec := e.Section(".BTF")
	if btfSec != nil {
		// Read BTF header from the .BTF section.
		h, err := readBTFHeader(btfSec, e.ByteOrder)
		if err != nil {
			return fmt.Errorf("reading BTF section header: %w", err)
		}

		// Populate offsetBTF fields of all symbols.
		if err := findBTFSymbols(stringOffsets, btfSec, h); err != nil {
			return fmt.Errorf("finding BTF string offsets: %w", err)
		}
	}

	s.data = dataOffsets
	s.strings = stringOffsets
	return nil
}
<commit_msg>elf: avoid flooding debug log with empty symbol names<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package elf

import (
	"bufio"
	"debug/elf"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
)

const (
	dataSection   = ".data"
	mapSection    = "maps"
	btfMapSection = ".maps"

	nullTerminator     = byte(0)
	relocSectionPrefix = ".rel"

	// nMapRelocations is an approximation of the number of offsets in an
	// ELF relating to map names that need to be adjusted. It's used for
	// map initialization within 'Symbols'.
	//
	// $ readelf -a bpf/bpf_lxc.o | grep "cilium_.*_" | grep "^0000" | wc -l
	// 51
	nMapRelocations = 64
)

type symbolKind uint32

const (
	symbolData   = symbolKind(1)
	symbolString = symbolKind(2)
)

func (k symbolKind) String() string {
	switch k {
	case symbolData:
		return "data"
	case symbolString:
		return "string"
	}
	return "unknown"
}

// symbol stores the location and type of a symbol within the ELF file.
type symbol struct {
	name      string
	kind      symbolKind
	offset    uint64
	offsetBTF uint64
	size      uint64
}

func newSymbol(name string, kind symbolKind, offset, size uint64) symbol {
	return symbol{
		name:   name,
		kind:   kind,
		offset: offset,
		size:   size,
	}
}

func newVariable(name string, offset, size uint64) symbol {
	return newSymbol(name, symbolData, offset, size)
}

func newString(name string, offset uint64) symbol {
	return newSymbol(name, symbolString, offset, uint64(len(name)))
}

type symbolSlice []symbol

// sort a slice of symbols by offset.
func (c symbolSlice) sort() symbolSlice {
	sort.Slice(c, func(i, j int) bool { return c[i].offset < c[j].offset })
	return c
}

type symbols struct {
	// data caches static 32-bit variables by name.
	data map[string]symbol
	// strings caches string symbols by name.
	strings map[string]symbol
}

func (s *symbols) sort() symbolSlice {
	result := make(symbolSlice, 0)
	for _, c := range s.data {
		result = append(result, c)
	}
	for _, c := range s.strings {
		result = append(result, c)
	}
	return result.sort()
}

// isGlobalData returns true if the symbol meets all of the following criteria:
// - symbol is of type STT_NOTYPE or STT_OBJECT
// - symbol has a global binding (STB_GLOBAL)
// - symbol has default visibility (STV_DEFAULT)
func isGlobalData(sym elf.Symbol) bool {
	return (elf.ST_TYPE(sym.Info) == elf.STT_NOTYPE ||
		elf.ST_TYPE(sym.Info) == elf.STT_OBJECT) &&
		elf.ST_BIND(sym.Info) == elf.STB_GLOBAL &&
		elf.ST_VISIBILITY(sym.Other) == elf.STV_DEFAULT
}

func readStringOffset(e *elf.File, r io.ReadSeeker, symbolOffset int64) (uint64, error) {
	if _, err := r.Seek(symbolOffset, io.SeekStart); err != nil {
		return 0, err
	}

	switch e.Class {
	case elf.ELFCLASS32:
		var sym32 elf.Sym32
		if err := binary.Read(r, e.ByteOrder, &sym32); err != nil {
			return 0, err
		}
		return uint64(sym32.Name), nil
	case elf.ELFCLASS64:
		var sym64 elf.Sym64
		if err := binary.Read(r, e.ByteOrder, &sym64); err != nil {
			return 0, err
		}
		return uint64(sym64.Name), nil
	}
	return 0, fmt.Errorf("unsupported ELF type %d", e.Class)
}

// extractFrom processes the specified ELF and populates the received symbols
// object with data and string offsets in the file, for later substitution.
func (s *symbols) extractFrom(e *elf.File) error {
	dataOffsets := make(map[string]symbol)
	stringOffsets := make(map[string]symbol, nMapRelocations)

	symbols, err := e.Symbols()
	if err != nil {
		return err
	}
	symtab := e.SectionByType(elf.SHT_SYMTAB)
	strtab := e.Sections[symtab.Link]

	// Scan symbol table for offsets of static data and symbol names.
	symbolReader := symtab.Open()
	for i, sym := range symbols {
		// BTF extensions like line info not recognized by normal ELF parsers
		if elf.ST_TYPE(sym.Info) == elf.STT_FILE {
			continue
		}

		// Get ELF section reference by its index encoded in the symbol.
		section := e.Sections[sym.Section]

		switch {
		case section.Flags&elf.SHF_COMPRESSED > 0:
			return fmt.Errorf("compressed %s section not supported", section.Name)

		// Skip local symbols that are objects or untyped. These are usually
		// program segments referred to using long jumps in bytecode.
		case !isGlobalData(sym):
			// Some symbols don't have names and 'LBB0 is a common llvm symbol prefix
			// (basic block). Don't flood the logs with messages about them.
			if sym.Name != "" && !strings.HasPrefix(sym.Name, "LBB") {
				log.Debugf("Skipping %s", sym.Name)
			}
			continue

		// Find the absolute offset to the variable's value this symbol represents
		// in the .data section.
		// This implements substitution of static data.
		case section.Name == dataSection:
			// Offset from start of binary to variable inside .data
			offset := section.Offset + sym.Value
			dataOffsets[sym.Name] = newVariable(sym.Name, offset, sym.Size)
			log.WithField(fieldSymbol, sym.Name).Debugf("Found variable with offset %d", offset)

		// Find the absolute offset to the map's entry in .strtab.
		// This implements renaming maps.
		case section.Name == mapSection || section.Name == btfMapSection:
			// From the Golang Documentation:
			//   "For compatibility with Go 1.0, Symbols omits the
			//   the null symbol at index 0."
			// We must reverse this when reading directly.
			symbolOffset := int64(i+1) * int64(symtab.Entsize)
			symOffsetInStrtab, err := readStringOffset(e, symbolReader, symbolOffset)
			if err != nil {
				return err
			}

			// Offset from start of binary to name inside .strtab
			symOffset := strtab.Offset + symOffsetInStrtab
			stringOffsets[sym.Name] = newString(sym.Name, symOffset)
			log.WithField(fieldSymbol, sym.Name).Debugf("Found symbol with offset %d", symOffset)

		default:
			log.WithField(fieldSymbol, sym.Name).Debugf("Found symbol referring to unknown section id %d", sym.Section)
		}
	}

	// Scan string table for offsets of section names.
	stringReader := bufio.NewReader(strtab.Open())
	var elfString string
	for off := uint64(0); off < strtab.Size; off += uint64(len(elfString)) {
		// off is the offset within the string table.
		elfString, err = stringReader.ReadString(nullTerminator)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return err
		}

		// We only need to worry about sections with relocations.
		if !strings.HasPrefix(elfString, relocSectionPrefix) {
			continue
		}

		elfEnd := len(elfString)
		if elfString[elfEnd-1] == nullTerminator {
			elfEnd--
		}
		relocOffset := uint64(len(relocSectionPrefix))
		secName := elfString[relocOffset:elfEnd]
		if sec := e.Section(secName); sec != nil {
			// Offset from start of binary to name inside .strtab
			globalOffset := strtab.Offset + off + relocOffset
			stringOffsets[secName] = newString(secName, globalOffset)
			log.WithField(fieldSymbol, secName).Debugf("Found section with offset %d", globalOffset)
		}
	}

	// If .BTF section is present in the ELF, get the offset of each symbol
	// into the BTF string table to be able to overwrite them later.
	btfSec := e.Section(".BTF")
	if btfSec != nil {
		// Read BTF header from the .BTF section.
		h, err := readBTFHeader(btfSec, e.ByteOrder)
		if err != nil {
			return fmt.Errorf("reading BTF section header: %w", err)
		}

		// Populate offsetBTF fields of all symbols.
		if err := findBTFSymbols(stringOffsets, btfSec, h); err != nil {
			return fmt.Errorf("finding BTF string offsets: %w", err)
		}
	}

	s.data = dataOffsets
	s.strings = stringOffsets
	return nil
}
<|endoftext|>