<commit_before>package graphite

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
	"time"

	"go.opentelemetry.io/otel/attribute"

	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
	"github.com/grafana/grafana-plugin-sdk-go/data"

	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/httpclient"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/infra/tracing"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tsdb/legacydata"
)

type Service struct {
	logger log.Logger
	im     instancemgmt.InstanceManager
	tracer tracing.Tracer
}

const (
	TargetFullModelField = "targetFull"
	TargetModelField     = "target"
)

func ProvideService(httpClientProvider httpclient.Provider, tracer tracing.Tracer) *Service {
	return &Service{
		logger: log.New("tsdb.graphite"),
		im:     datasource.NewInstanceManager(newInstanceSettings(httpClientProvider)),
		tracer: tracer,
	}
}

type datasourceInfo struct {
	HTTPClient *http.Client
	URL        string
	Id         int64
}

func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.InstanceFactoryFunc {
	return func(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
		opts, err := settings.HTTPClientOptions()
		if err != nil {
			return nil, err
		}

		client, err := httpClientProvider.New(opts)
		if err != nil {
			return nil, err
		}

		model := datasourceInfo{
			HTTPClient: client,
			URL:        settings.URL,
			Id:         settings.ID,
		}

		return model, nil
	}
}

func (s *Service) getDSInfo(pluginCtx backend.PluginContext) (*datasourceInfo, error) {
	i, err := s.im.Get(pluginCtx)
	if err != nil {
		return nil, err
	}
	instance := i.(datasourceInfo)
	return &instance, nil
}

func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
	if len(req.Queries) == 0 {
		return nil, fmt.Errorf("query contains no queries")
	}

	// get datasource info from context
	dsInfo, err := s.getDSInfo(req.PluginContext)
	if err != nil {
		return nil, err
	}

	// take the first query in the request list, since all query should share the same timerange
	q := req.Queries[0]

	/*
		graphite doc about from and until, with sdk we are getting absolute instead of relative time
		https://graphite-api.readthedocs.io/en/latest/api.html#from-until
	*/
	from, until := epochMStoGraphiteTime(q.TimeRange)
	formData := url.Values{
		"from":          []string{from},
		"until":         []string{until},
		"format":        []string{"json"},
		"maxDataPoints": []string{"500"},
	}

	// Calculate and get the last target of Graphite Request
	var target string
	emptyQueries := make([]string, 0)
	for _, query := range req.Queries {
		model, err := simplejson.NewJson(query.JSON)
		if err != nil {
			return nil, err
		}
		s.logger.Debug("graphite", "query", model)
		currTarget := ""
		if fullTarget, err := model.Get(TargetFullModelField).String(); err == nil {
			currTarget = fullTarget
		} else {
			currTarget = model.Get(TargetModelField).MustString()
		}
		if currTarget == "" {
			s.logger.Debug("graphite", "empty query target", model)
			emptyQueries = append(emptyQueries, fmt.Sprintf("Query: %v has no target", model))
			continue
		}
		target = fixIntervalFormat(currTarget)
	}

	var result = backend.QueryDataResponse{}

	if target == "" {
		s.logger.Error("No targets in query model", "models without targets", strings.Join(emptyQueries, "\n"))
		return &result, errors.New("no query target found for the alert rule")
	}

	formData["target"] = []string{target}

	if setting.Env == setting.Dev {
		s.logger.Debug("Graphite request", "params", formData)
	}

	graphiteReq, err := s.createRequest(ctx, dsInfo, formData)
	if err != nil {
		return &result, err
	}

	ctx, span := s.tracer.Start(ctx, "graphite query")
	span.SetAttributes("target", target, attribute.Key("target").String(target))
	span.SetAttributes("from", from, attribute.Key("from").String(from))
	span.SetAttributes("until", until, attribute.Key("until").String(until))
	span.SetAttributes("datasource_id", dsInfo.Id, attribute.Key("datasource_id").Int64(dsInfo.Id))
	span.SetAttributes("org_id", req.PluginContext.OrgID, attribute.Key("org_id").Int64(req.PluginContext.OrgID))

	defer span.End()
	s.tracer.Inject(ctx, graphiteReq.Header, span)

	res, err := dsInfo.HTTPClient.Do(graphiteReq)
	if err != nil {
		return &result, err
	}

	frames, err := s.toDataFrames(res)
	if err != nil {
		return &result, err
	}

	result = backend.QueryDataResponse{
		Responses: make(backend.Responses),
	}

	result.Responses["A"] = backend.DataResponse{
		Frames: frames,
	}

	return &result, nil
}

func (s *Service) parseResponse(res *http.Response) ([]TargetResponseDTO, error) {
	body, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := res.Body.Close(); err != nil {
			s.logger.Warn("Failed to close response body", "err", err)
		}
	}()

	if res.StatusCode/100 != 2 {
		s.logger.Info("Request failed", "status", res.Status, "body", string(body))
		return nil, fmt.Errorf("request failed, status: %s", res.Status)
	}

	var data []TargetResponseDTO
	err = json.Unmarshal(body, &data)
	if err != nil {
		s.logger.Info("Failed to unmarshal graphite response", "error", err, "status", res.Status, "body", string(body))
		return nil, err
	}

	return data, nil
}

func (s *Service) toDataFrames(response *http.Response) (frames data.Frames, error error) {
	responseData, err := s.parseResponse(response)
	if err != nil {
		return nil, err
	}

	frames = data.Frames{}
	for _, series := range responseData {
		timeVector := make([]time.Time, 0, len(series.DataPoints))
		values := make([]*float64, 0, len(series.DataPoints))
		name := series.Target

		for _, dataPoint := range series.DataPoints {
			var timestamp, value, err = parseDataTimePoint(dataPoint)
			if err != nil {
				return nil, err
			}
			timeVector = append(timeVector, timestamp)
			values = append(values, value)
		}

		tags := make(map[string]string)
		for name, value := range series.Tags {
			switch value := value.(type) {
			case string:
				tags[name] = value
			case float64:
				tags[name] = strconv.FormatFloat(value, 'f', -1, 64)
			}
		}

		frames = append(frames, data.NewFrame(name,
			data.NewField("time", nil, timeVector),
			data.NewField("value", tags, values).SetConfig(&data.FieldConfig{DisplayNameFromDS: name})))

		if setting.Env == setting.Dev {
			s.logger.Debug("Graphite response", "target", series.Target, "datapoints", len(series.DataPoints))
		}
	}
	return frames, nil
}

func (s *Service) createRequest(ctx context.Context, dsInfo *datasourceInfo, data url.Values) (*http.Request, error) {
	u, err := url.Parse(dsInfo.URL)
	if err != nil {
		return nil, err
	}
	u.Path = path.Join(u.Path, "render")

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), strings.NewReader(data.Encode()))
	if err != nil {
		s.logger.Info("Failed to create request", "error", err)
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return req, err
}

func fixIntervalFormat(target string) string {
	rMinute := regexp.MustCompile(`'(\d+)m'`)
	target = rMinute.ReplaceAllStringFunc(target, func(m string) string {
		return strings.ReplaceAll(m, "m", "min")
	})
	rMonth := regexp.MustCompile(`'(\d+)M'`)
	target = rMonth.ReplaceAllStringFunc(target, func(M string) string {
		return strings.ReplaceAll(M, "M", "mon")
	})
	return target
}

func epochMStoGraphiteTime(tr backend.TimeRange) (string, string) {
	return fmt.Sprintf("%d", tr.From.UTC().Unix()), fmt.Sprintf("%d", tr.To.UTC().Unix())
}

/**
 * Graphite should always return timestamp as a number but values might be nil when data is missing
 */
func parseDataTimePoint(dataTimePoint legacydata.DataTimePoint) (time.Time, *float64, error) {
	if !dataTimePoint[1].Valid {
		return time.Time{}, nil, errors.New("failed to parse data point timestamp")
	}

	timestamp := time.Unix(int64(dataTimePoint[1].Float64), 0).UTC()

	if dataTimePoint[0].Valid {
		var value = new(float64)
		*value = dataTimePoint[0].Float64
		return timestamp, value, nil
	} else {
		return timestamp, nil, nil
	}
}
<commit_msg>Add error information to graphite queries tracing (#55249)<commit_after>package graphite

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
	"github.com/grafana/grafana-plugin-sdk-go/data"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"

	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/httpclient"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/infra/tracing"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tsdb/legacydata"
)

type Service struct {
	logger log.Logger
	im     instancemgmt.InstanceManager
	tracer tracing.Tracer
}

const (
	TargetFullModelField = "targetFull"
	TargetModelField     = "target"
)

func ProvideService(httpClientProvider httpclient.Provider, tracer tracing.Tracer) *Service {
	return &Service{
		logger: log.New("tsdb.graphite"),
		im:     datasource.NewInstanceManager(newInstanceSettings(httpClientProvider)),
		tracer: tracer,
	}
}

type datasourceInfo struct {
	HTTPClient *http.Client
	URL        string
	Id         int64
}

func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.InstanceFactoryFunc {
	return func(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
		opts, err := settings.HTTPClientOptions()
		if err != nil {
			return nil, err
		}

		client, err := httpClientProvider.New(opts)
		if err != nil {
			return nil, err
		}

		model := datasourceInfo{
			HTTPClient: client,
			URL:        settings.URL,
			Id:         settings.ID,
		}

		return model, nil
	}
}

func (s *Service) getDSInfo(pluginCtx backend.PluginContext) (*datasourceInfo, error) {
	i, err := s.im.Get(pluginCtx)
	if err != nil {
		return nil, err
	}
	instance := i.(datasourceInfo)
	return &instance, nil
}

func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
	if len(req.Queries) == 0 {
		return nil, fmt.Errorf("query contains no queries")
	}

	// get datasource info from context
	dsInfo, err := s.getDSInfo(req.PluginContext)
	if err != nil {
		return nil, err
	}

	// take the first query in the request list, since all query should share the same timerange
	q := req.Queries[0]

	/*
		graphite doc about from and until, with sdk we are getting absolute instead of relative time
		https://graphite-api.readthedocs.io/en/latest/api.html#from-until
	*/
	from, until := epochMStoGraphiteTime(q.TimeRange)
	formData := url.Values{
		"from":          []string{from},
		"until":         []string{until},
		"format":        []string{"json"},
		"maxDataPoints": []string{"500"},
	}

	// Calculate and get the last target of Graphite Request
	var target string
	emptyQueries := make([]string, 0)
	for _, query := range req.Queries {
		model, err := simplejson.NewJson(query.JSON)
		if err != nil {
			return nil, err
		}
		s.logger.Debug("graphite", "query", model)
		currTarget := ""
		if fullTarget, err := model.Get(TargetFullModelField).String(); err == nil {
			currTarget = fullTarget
		} else {
			currTarget = model.Get(TargetModelField).MustString()
		}
		if currTarget == "" {
			s.logger.Debug("graphite", "empty query target", model)
			emptyQueries = append(emptyQueries, fmt.Sprintf("Query: %v has no target", model))
			continue
		}
		target = fixIntervalFormat(currTarget)
	}

	var result = backend.QueryDataResponse{}

	if target == "" {
		s.logger.Error("No targets in query model", "models without targets", strings.Join(emptyQueries, "\n"))
		return &result, errors.New("no query target found for the alert rule")
	}

	formData["target"] = []string{target}

	if setting.Env == setting.Dev {
		s.logger.Debug("Graphite request", "params", formData)
	}

	graphiteReq, err := s.createRequest(ctx, dsInfo, formData)
	if err != nil {
		return &result, err
	}

	ctx, span := s.tracer.Start(ctx, "graphite query")
	defer span.End()

	span.SetAttributes("target", target, attribute.Key("target").String(target))
	span.SetAttributes("from", from, attribute.Key("from").String(from))
	span.SetAttributes("until", until, attribute.Key("until").String(until))
	span.SetAttributes("datasource_id", dsInfo.Id, attribute.Key("datasource_id").Int64(dsInfo.Id))
	span.SetAttributes("org_id", req.PluginContext.OrgID, attribute.Key("org_id").Int64(req.PluginContext.OrgID))

	s.tracer.Inject(ctx, graphiteReq.Header, span)

	res, err := dsInfo.HTTPClient.Do(graphiteReq)
	if res != nil {
		span.SetAttributes("graphite.response.code", res.StatusCode, attribute.Key("graphite.response.code").Int(res.StatusCode))
	}
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		return &result, err
	}

	frames, err := s.toDataFrames(res)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		return &result, err
	}

	result = backend.QueryDataResponse{
		Responses: make(backend.Responses),
	}

	result.Responses["A"] = backend.DataResponse{
		Frames: frames,
	}

	return &result, nil
}

func (s *Service) parseResponse(res *http.Response) ([]TargetResponseDTO, error) {
	body, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := res.Body.Close(); err != nil {
			s.logger.Warn("Failed to close response body", "err", err)
		}
	}()

	if res.StatusCode/100 != 2 {
		s.logger.Info("Request failed", "status", res.Status, "body", string(body))
		return nil, fmt.Errorf("request failed, status: %s", res.Status)
	}

	var data []TargetResponseDTO
	err = json.Unmarshal(body, &data)
	if err != nil {
		s.logger.Info("Failed to unmarshal graphite response", "error", err, "status", res.Status, "body", string(body))
		return nil, err
	}

	return data, nil
}

func (s *Service) toDataFrames(response *http.Response) (frames data.Frames, error error) {
	responseData, err := s.parseResponse(response)
	if err != nil {
		return nil, err
	}

	frames = data.Frames{}
	for _, series := range responseData {
		timeVector := make([]time.Time, 0, len(series.DataPoints))
		values := make([]*float64, 0, len(series.DataPoints))
		name := series.Target

		for _, dataPoint := range series.DataPoints {
			var timestamp, value, err = parseDataTimePoint(dataPoint)
			if err != nil {
				return nil, err
			}
			timeVector = append(timeVector, timestamp)
			values = append(values, value)
		}

		tags := make(map[string]string)
		for name, value := range series.Tags {
			switch value := value.(type) {
			case string:
				tags[name] = value
			case float64:
				tags[name] = strconv.FormatFloat(value, 'f', -1, 64)
			}
		}

		frames = append(frames, data.NewFrame(name,
			data.NewField("time", nil, timeVector),
			data.NewField("value", tags, values).SetConfig(&data.FieldConfig{DisplayNameFromDS: name})))

		if setting.Env == setting.Dev {
			s.logger.Debug("Graphite response", "target", series.Target, "datapoints", len(series.DataPoints))
		}
	}
	return frames, nil
}

func (s *Service) createRequest(ctx context.Context, dsInfo *datasourceInfo, data url.Values) (*http.Request, error) {
	u, err := url.Parse(dsInfo.URL)
	if err != nil {
		return nil, err
	}
	u.Path = path.Join(u.Path, "render")

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), strings.NewReader(data.Encode()))
	if err != nil {
		s.logger.Info("Failed to create request", "error", err)
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return req, err
}

func fixIntervalFormat(target string) string {
	rMinute := regexp.MustCompile(`'(\d+)m'`)
	target = rMinute.ReplaceAllStringFunc(target, func(m string) string {
		return strings.ReplaceAll(m, "m", "min")
	})
	rMonth := regexp.MustCompile(`'(\d+)M'`)
	target = rMonth.ReplaceAllStringFunc(target, func(M string) string {
		return strings.ReplaceAll(M, "M", "mon")
	})
	return target
}

func epochMStoGraphiteTime(tr backend.TimeRange) (string, string) {
	return fmt.Sprintf("%d", tr.From.UTC().Unix()), fmt.Sprintf("%d", tr.To.UTC().Unix())
}

/**
 * Graphite should always return timestamp as a number but values might be nil when data is missing
 */
func parseDataTimePoint(dataTimePoint legacydata.DataTimePoint) (time.Time, *float64, error) {
	if !dataTimePoint[1].Valid {
		return time.Time{}, nil, errors.New("failed to parse data point timestamp")
	}

	timestamp := time.Unix(int64(dataTimePoint[1].Float64), 0).UTC()

	if dataTimePoint[0].Valid {
		var value = new(float64)
		*value = dataTimePoint[0].Float64
		return timestamp, value, nil
	} else {
		return timestamp, nil, nil
	}
}
<|endoftext|>


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

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

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

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

package keymutex

import (
	"fmt"
	"github.com/golang/glog"
	"sync"
)

// KeyMutex is a thread-safe interface for aquiring locks on arbitrary strings.
type KeyMutex interface {
	// Aquires a lock associated with the specified ID, creates the lock if one doesn't already exist.
	LockKey(id string)

	// Releases the lock associated with the specified ID.
	// Returns an error if the specified ID doesn't exist.
	UnlockKey(id string) error
}

// Returns a new instance of a key mutex.
func NewKeyMutex() KeyMutex {
	return &keyMutex{
		mutexMap: make(map[string]*sync.Mutex),
	}
}

type keyMutex struct {
	sync.RWMutex
	mutexMap map[string]*sync.Mutex
}

// Aquires a lock associated with the specified ID (creates the lock if one doesn't already exist).
func (km *keyMutex) LockKey(id string) {
	glog.V(5).Infof("LockKey(...) called for id %q\r\n", id)
	mutex := km.getOrCreateLock(id)
	mutex.Lock()
	glog.V(5).Infof("LockKey(...) for id %q completed.\r\n", id)
}

// Releases the lock associated with the specified ID.
// Returns an error if the specified ID doesn't exist.
func (km *keyMutex) UnlockKey(id string) error {
	glog.V(5).Infof("UnlockKey(...) called for id %q\r\n", id)
	km.RLock()
	defer km.RUnlock()
	mutex, exists := km.mutexMap[id]
	if !exists {
		return fmt.Errorf("id %q not found", id)
	}
	glog.V(5).Infof("UnlockKey(...) for id. Mutex found, trying to unlock it. %q\r\n", id)

	mutex.Unlock()
	glog.V(5).Infof("UnlockKey(...) for id %q completed.\r\n", id)
	return nil
}

// Returns lock associated with the specified ID, or creates the lock if one doesn't already exist.
func (km *keyMutex) getOrCreateLock(id string) *sync.Mutex {
	km.Lock()
	defer km.Unlock()

	if _, exists := km.mutexMap[id]; !exists {
		km.mutexMap[id] = &sync.Mutex{}
	}

	return km.mutexMap[id]
}
<commit_msg>fixed spelling typo<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package keymutex

import (
	"fmt"
	"github.com/golang/glog"
	"sync"
)

// KeyMutex is a thread-safe interface for acquiring locks on arbitrary strings.
type KeyMutex interface {
	// Aquires a lock associated with the specified ID, creates the lock if one doesn't already exist.
	LockKey(id string)

	// Releases the lock associated with the specified ID.
	// Returns an error if the specified ID doesn't exist.
	UnlockKey(id string) error
}

// Returns a new instance of a key mutex.
func NewKeyMutex() KeyMutex {
	return &keyMutex{
		mutexMap: make(map[string]*sync.Mutex),
	}
}

type keyMutex struct {
	sync.RWMutex
	mutexMap map[string]*sync.Mutex
}

// Aquires a lock associated with the specified ID (creates the lock if one doesn't already exist).
func (km *keyMutex) LockKey(id string) {
	glog.V(5).Infof("LockKey(...) called for id %q\r\n", id)
	mutex := km.getOrCreateLock(id)
	mutex.Lock()
	glog.V(5).Infof("LockKey(...) for id %q completed.\r\n", id)
}

// Releases the lock associated with the specified ID.
// Returns an error if the specified ID doesn't exist.
func (km *keyMutex) UnlockKey(id string) error {
	glog.V(5).Infof("UnlockKey(...) called for id %q\r\n", id)
	km.RLock()
	defer km.RUnlock()
	mutex, exists := km.mutexMap[id]
	if !exists {
		return fmt.Errorf("id %q not found", id)
	}
	glog.V(5).Infof("UnlockKey(...) for id. Mutex found, trying to unlock it. %q\r\n", id)

	mutex.Unlock()
	glog.V(5).Infof("UnlockKey(...) for id %q completed.\r\n", id)
	return nil
}

// Returns lock associated with the specified ID, or creates the lock if one doesn't already exist.
func (km *keyMutex) getOrCreateLock(id string) *sync.Mutex {
	km.Lock()
	defer km.Unlock()

	if _, exists := km.mutexMap[id]; !exists {
		km.mutexMap[id] = &sync.Mutex{}
	}

	return km.mutexMap[id]
}
<|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 common

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	"k8s.io/api/networking/v1beta1"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	client "k8s.io/client-go/kubernetes/typed/networking/v1beta1"
	"k8s.io/klog"
	"k8s.io/kubernetes/pkg/util/slice"
)

const (
	// FinalizerKey is the string representing the Ingress finalizer.
	FinalizerKey = "networking.gke.io/ingress-finalizer"
	// FinalizerKeyV2 is the string representing the Ingress finalizer version.
	// Ingress with V2 finalizer uses V2 frontend naming scheme.
	FinalizerKeyV2 = "networking.gke.io/ingress-finalizer-V2"
	// TODO remove the 2 definitions once they are added in legacy-cloud-providers/gce
	// LegacyILBFinalizer key is used to identify ILB services whose resources are managed by service controller.
	LegacyILBFinalizer = "gke.networking.io/l4-ilb-v1"
	// ILBFinalizerV2 is the finalizer used by newer controllers that implement Internal LoadBalancer services.
	ILBFinalizerV2 = "gke.networking.io/l4-ilb-v2"
)

// IsDeletionCandidate is true if the passed in meta contains an ingress finalizer.
func IsDeletionCandidate(m meta_v1.ObjectMeta) bool {
	return IsDeletionCandidateForGivenFinalizer(m, FinalizerKey) || IsDeletionCandidateForGivenFinalizer(m, FinalizerKeyV2)
}

// IsDeletionCandidateForGivenFinalizer is true if the passed in meta contains the specified finalizer.
func IsDeletionCandidateForGivenFinalizer(m meta_v1.ObjectMeta, key string) bool {
	return m.DeletionTimestamp != nil && HasGivenFinalizer(m, key)
}

// HasFinalizer is true if the passed in meta has an ingress finalizer.
func HasFinalizer(m meta_v1.ObjectMeta) bool {
	return HasGivenFinalizer(m, FinalizerKey) || HasGivenFinalizer(m, FinalizerKeyV2)
}

// HasGivenFinalizer is true if the passed in meta has the specified finalizer.
func HasGivenFinalizer(m meta_v1.ObjectMeta, key string) bool {
	return slice.ContainsString(m.Finalizers, key, nil)
}

// EnsureFinalizer ensures that the specified finalizer exists on given Ingress.
func EnsureFinalizer(ing *v1beta1.Ingress, ingClient client.IngressInterface, finalizerKey string) (*v1beta1.Ingress, error) {
	updated := ing.DeepCopy()
	if needToAddFinalizer(ing.ObjectMeta, finalizerKey) {
		updated.ObjectMeta.Finalizers = append(updated.ObjectMeta.Finalizers, finalizerKey)
		// TODO(smatti): Make this optimistic concurrency control compliant on write.
		// Refer: https://github.com/eBay/Kubernetes/blob/master/docs/devel/api-conventions.md#concurrency-control-and-consistency
		if _, err := ingClient.Update(context.TODO(), updated, meta_v1.UpdateOptions{}); err != nil {
			return nil, fmt.Errorf("error updating Ingress %s/%s: %v", ing.Namespace, ing.Name, err)
		}
		klog.V(2).Infof("Added finalizer %q for Ingress %s/%s", finalizerKey, ing.Namespace, ing.Name)
	}
	return updated, nil
}

// needToAddFinalizer is true if the passed in meta does not contain the specified finalizer.
func needToAddFinalizer(m meta_v1.ObjectMeta, key string) bool {
	return m.DeletionTimestamp == nil && !HasGivenFinalizer(m, key)
}

// EnsureDeleteFinalizer ensures that the specified finalizer is deleted from given Ingress.
func EnsureDeleteFinalizer(ing *v1beta1.Ingress, ingClient client.IngressInterface, finalizerKey string) error {
	if HasGivenFinalizer(ing.ObjectMeta, finalizerKey) {
		updated := ing.DeepCopy()
		updated.ObjectMeta.Finalizers = slice.RemoveString(updated.ObjectMeta.Finalizers, finalizerKey, nil)
		if _, err := ingClient.Update(context.TODO(), updated, meta_v1.UpdateOptions{}); err != nil {
			return fmt.Errorf("error updating Ingress %s/%s: %v", ing.Namespace, ing.Name, err)
		}
		klog.V(2).Infof("Removed finalizer %q for Ingress %s/%s", finalizerKey, ing.Namespace, ing.Name)
	}
	return nil
}

// EnsureServiceFinalizer patches the service to add finalizer.
func EnsureServiceFinalizer(service *v1.Service, key string, kubeClient kubernetes.Interface) error {
	if HasGivenFinalizer(service.ObjectMeta, key) {
		return nil
	}

	// Make a copy so we don't mutate the shared informer cache.
	updated := service.DeepCopy()
	updated.ObjectMeta.Finalizers = append(updated.ObjectMeta.Finalizers, key)

	klog.V(2).Infof("Adding finalizer %s to service %s/%s", key, updated.Namespace, updated.Name)
	_, err := kubeClient.CoreV1().Services(updated.Namespace).Update(context.TODO(), updated, meta_v1.UpdateOptions{})
	return err
}

// removeFinalizer patches the service to remove finalizer.
func EnsureDeleteServiceFinalizer(service *v1.Service, key string, kubeClient kubernetes.Interface) error {
	if !HasGivenFinalizer(service.ObjectMeta, key) {
		return nil
	}

	// Make a copy so we don't mutate the shared informer cache.
	updated := service.DeepCopy()
	updated.ObjectMeta.Finalizers = slice.RemoveString(updated.ObjectMeta.Finalizers, key, nil)

	klog.V(2).Infof("Removing finalizer from service %s/%s", updated.Namespace, updated.Name)
	_, err := kubeClient.CoreV1().Services(updated.Namespace).Update(context.TODO(), updated, meta_v1.UpdateOptions{})
	return err
}
<commit_msg>Use patch for updating ingress finalizers<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 common

import (
	"context"
	"encoding/json"
	"fmt"

	v1 "k8s.io/api/core/v1"
	"k8s.io/api/networking/v1beta1"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/client-go/kubernetes"
	client "k8s.io/client-go/kubernetes/typed/networking/v1beta1"
	"k8s.io/klog"
	"k8s.io/kubernetes/pkg/util/slice"
)

const (
	// FinalizerKey is the string representing the Ingress finalizer.
	FinalizerKey = "networking.gke.io/ingress-finalizer"
	// FinalizerKeyV2 is the string representing the Ingress finalizer version.
	// Ingress with V2 finalizer uses V2 frontend naming scheme.
	FinalizerKeyV2 = "networking.gke.io/ingress-finalizer-V2"
	// TODO remove the 2 definitions once they are added in legacy-cloud-providers/gce
	// LegacyILBFinalizer key is used to identify ILB services whose resources are managed by service controller.
	LegacyILBFinalizer = "gke.networking.io/l4-ilb-v1"
	// ILBFinalizerV2 is the finalizer used by newer controllers that implement Internal LoadBalancer services.
	ILBFinalizerV2 = "gke.networking.io/l4-ilb-v2"
)

// IsDeletionCandidate is true if the passed in meta contains an ingress finalizer.
func IsDeletionCandidate(m meta_v1.ObjectMeta) bool {
	return IsDeletionCandidateForGivenFinalizer(m, FinalizerKey) || IsDeletionCandidateForGivenFinalizer(m, FinalizerKeyV2)
}

// IsDeletionCandidateForGivenFinalizer is true if the passed in meta contains the specified finalizer.
func IsDeletionCandidateForGivenFinalizer(m meta_v1.ObjectMeta, key string) bool {
	return m.DeletionTimestamp != nil && HasGivenFinalizer(m, key)
}

// HasFinalizer is true if the passed in meta has an ingress finalizer.
func HasFinalizer(m meta_v1.ObjectMeta) bool {
	return HasGivenFinalizer(m, FinalizerKey) || HasGivenFinalizer(m, FinalizerKeyV2)
}

// HasGivenFinalizer is true if the passed in meta has the specified finalizer.
func HasGivenFinalizer(m meta_v1.ObjectMeta, key string) bool {
	return slice.ContainsString(m.Finalizers, key, nil)
}

// EnsureFinalizer ensures that the specified finalizer exists on given Ingress.
func EnsureFinalizer(ing *v1beta1.Ingress, ingClient client.IngressInterface, finalizerKey string) (*v1beta1.Ingress, error) {
	updated := ing.DeepCopy()
	if needToAddFinalizer(ing.ObjectMeta, finalizerKey) {
		updated.ObjectMeta.Finalizers = append(updated.ObjectMeta.Finalizers, finalizerKey)
		if _, err := patchIngressFinalizer(ingClient, ing, updated); err != nil {
			return nil, fmt.Errorf("error patching Ingress %s/%s: %v", ing.Namespace, ing.Name, err)
		}
		klog.V(2).Infof("Added finalizer %q for Ingress %s/%s", finalizerKey, ing.Namespace, ing.Name)
	}
	return updated, nil
}

// needToAddFinalizer is true if the passed in meta does not contain the specified finalizer.
func needToAddFinalizer(m meta_v1.ObjectMeta, key string) bool {
	return m.DeletionTimestamp == nil && !HasGivenFinalizer(m, key)
}

// EnsureDeleteFinalizer ensures that the specified finalizer is deleted from given Ingress.
func EnsureDeleteFinalizer(ing *v1beta1.Ingress, ingClient client.IngressInterface, finalizerKey string) error {
	if HasGivenFinalizer(ing.ObjectMeta, finalizerKey) {
		updated := ing.DeepCopy()
		updated.ObjectMeta.Finalizers = slice.RemoveString(updated.ObjectMeta.Finalizers, finalizerKey, nil)
		if _, err := patchIngressFinalizer(ingClient, ing, updated); err != nil {
			return fmt.Errorf("error patching Ingress %s/%s: %v", ing.Namespace, ing.Name, err)
		}
		klog.V(2).Infof("Removed finalizer %q for Ingress %s/%s", finalizerKey, ing.Namespace, ing.Name)
	}
	return nil
}

func patchIngressFinalizer(ic client.IngressInterface, oldIngress, newIngress *v1beta1.Ingress) (*v1beta1.Ingress, error) {
	ingKey := fmt.Sprintf("%s/%s", oldIngress.Namespace, oldIngress.Name)
	oldData, err := json.Marshal(oldIngress)
	if err != nil {
		return nil, fmt.Errorf("failed to Marshal oldData for ingress %s: %v", ingKey, err)
	}

	newData, err := json.Marshal(newIngress)
	if err != nil {
		return nil, fmt.Errorf("failed to Marshal newData for ingress %s: %v", ingKey, err)
	}

	patchBytes, err := strategicpatch.CreateTwoWayMergePatch(oldData, newData, v1beta1.Ingress{})
	if err != nil {
		return nil, fmt.Errorf("failed to create TwoWayMergePatch for ingress %s: %v", ingKey, err)
	}

	klog.V(2).Infof("Patch bytes for ingress %s: %s", ingKey, patchBytes)
	return ic.Patch(context.TODO(), oldIngress.Name, types.StrategicMergePatchType, patchBytes, meta_v1.PatchOptions{}, "status")
}

// EnsureServiceFinalizer patches the service to add finalizer.
func EnsureServiceFinalizer(service *v1.Service, key string, kubeClient kubernetes.Interface) error {
	if HasGivenFinalizer(service.ObjectMeta, key) {
		return nil
	}

	// Make a copy so we don't mutate the shared informer cache.
	updated := service.DeepCopy()
	updated.ObjectMeta.Finalizers = append(updated.ObjectMeta.Finalizers, key)

	klog.V(2).Infof("Adding finalizer %s to service %s/%s", key, updated.Namespace, updated.Name)
	_, err := kubeClient.CoreV1().Services(updated.Namespace).Update(context.TODO(), updated, meta_v1.UpdateOptions{})
	return err
}

// removeFinalizer patches the service to remove finalizer.
func EnsureDeleteServiceFinalizer(service *v1.Service, key string, kubeClient kubernetes.Interface) error {
	if !HasGivenFinalizer(service.ObjectMeta, key) {
		return nil
	}

	// Make a copy so we don't mutate the shared informer cache.
	updated := service.DeepCopy()
	updated.ObjectMeta.Finalizers = slice.RemoveString(updated.ObjectMeta.Finalizers, key, nil)

	klog.V(2).Infof("Removing finalizer from service %s/%s", updated.Namespace, updated.Name)
	_, err := kubeClient.CoreV1().Services(updated.Namespace).Update(context.TODO(), updated, meta_v1.UpdateOptions{})
	return err
}
<|endoftext|>


<commit_before>//go:build !providerless
// +build !providerless

/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package aws

import (
	"encoding/json"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// IPPermissionSet maps IP strings of strings to EC2 IpPermissions
type IPPermissionSet map[string]*ec2.IpPermission

// IPPermissionPredicate is an predicate to test whether IPPermission matches some condition.
type IPPermissionPredicate interface {
	// Test checks whether specified IPPermission matches condition.
	Test(perm *ec2.IpPermission) bool
}

// NewIPPermissionSet creates a new IPPermissionSet
func NewIPPermissionSet(items ...*ec2.IpPermission) IPPermissionSet {
	s := make(IPPermissionSet, len(items))
	s.Insert(items...)
	return s
}

// Ungroup splits permissions out into individual permissions
// EC2 will combine permissions with the same port but different SourceRanges together, for example
// We ungroup them so we can process them
func (s IPPermissionSet) Ungroup() IPPermissionSet {
	l := []*ec2.IpPermission{}
	for _, p := range s.List() {
		if len(p.IpRanges) <= 1 {
			l = append(l, p)
			continue
		}
		for _, ipRange := range p.IpRanges {
			c := &ec2.IpPermission{}
			*c = *p
			c.IpRanges = []*ec2.IpRange{ipRange}
			l = append(l, c)
		}
	}

	l2 := []*ec2.IpPermission{}
	for _, p := range l {
		if len(p.UserIdGroupPairs) <= 1 {
			l2 = append(l2, p)
			continue
		}
		for _, u := range p.UserIdGroupPairs {
			c := &ec2.IpPermission{}
			*c = *p
			c.UserIdGroupPairs = []*ec2.UserIdGroupPair{u}
			l2 = append(l, c)
		}
	}

	l3 := []*ec2.IpPermission{}
	for _, p := range l2 {
		if len(p.PrefixListIds) <= 1 {
			l3 = append(l3, p)
			continue
		}
		for _, v := range p.PrefixListIds {
			c := &ec2.IpPermission{}
			*c = *p
			c.PrefixListIds = []*ec2.PrefixListId{v}
			l3 = append(l3, c)
		}
	}

	return NewIPPermissionSet(l3...)
}

// Insert adds items to the set.
func (s IPPermissionSet) Insert(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		s[k] = p
	}
}

// Delete delete permission from the set.
func (s IPPermissionSet) Delete(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		delete(s, k)
	}
}

// DeleteIf delete permission from the set if permission matches predicate.
func (s IPPermissionSet) DeleteIf(predicate IPPermissionPredicate) {
	for k, p := range s {
		if predicate.Test(p) {
			delete(s, k)
		}
	}
}

// List returns the contents as a slice.  Order is not defined.
func (s IPPermissionSet) List() []*ec2.IpPermission {
	res := make([]*ec2.IpPermission, 0, len(s))
	for _, v := range s {
		res = append(res, v)
	}
	return res
}

// IsSuperset returns true if and only if s is a superset of s2.
func (s IPPermissionSet) IsSuperset(s2 IPPermissionSet) bool {
	for k := range s2 {
		_, found := s[k]
		if !found {
			return false
		}
	}
	return true
}

// Equal returns true if and only if s is equal (as a set) to s2.
// Two sets are equal if their membership is identical.
// (In practice, this means same elements, order doesn't matter)
func (s IPPermissionSet) Equal(s2 IPPermissionSet) bool {
	return len(s) == len(s2) && s.IsSuperset(s2)
}

// Difference returns a set of objects that are not in s2
// For example:
// s1 = {a1, a2, a3}
// s2 = {a1, a2, a4, a5}
// s1.Difference(s2) = {a3}
// s2.Difference(s1) = {a4, a5}
func (s IPPermissionSet) Difference(s2 IPPermissionSet) IPPermissionSet {
	result := NewIPPermissionSet()
	for k, v := range s {
		_, found := s2[k]
		if !found {
			result[k] = v
		}
	}
	return result
}

// Len returns the size of the set.
func (s IPPermissionSet) Len() int {
	return len(s)
}

func keyForIPPermission(p *ec2.IpPermission) string {
	v, err := json.Marshal(p)
	if err != nil {
		panic(fmt.Sprintf("error building JSON representation of ec2.IpPermission: %v", err))
	}
	return string(v)
}

var _ IPPermissionPredicate = IPPermissionMatchDesc{}

// IPPermissionMatchDesc checks whether specific IPPermission contains description.
type IPPermissionMatchDesc struct {
	Description string
}

// Test whether specific IPPermission contains description.
func (p IPPermissionMatchDesc) Test(perm *ec2.IpPermission) bool {
	for _, v4Range := range perm.IpRanges {
		if aws.StringValue(v4Range.Description) == p.Description {
			return true
		}
	}
	for _, v6Range := range perm.Ipv6Ranges {
		if aws.StringValue(v6Range.Description) == p.Description {
			return true
		}
	}
	for _, prefixListID := range perm.PrefixListIds {
		if aws.StringValue(prefixListID.Description) == p.Description {
			return true
		}
	}
	for _, group := range perm.UserIdGroupPairs {
		if aws.StringValue(group.Description) == p.Description {
			return true
		}
	}
	return false
}

var _ IPPermissionPredicate = IPPermissionNotMatch{}

// IPPermissionNotMatch is the *not* operator for Predicate
type IPPermissionNotMatch struct {
	Predicate IPPermissionPredicate
}

// Test whether specific IPPermission not match the embed predicate.
func (p IPPermissionNotMatch) Test(perm *ec2.IpPermission) bool {
	return !p.Predicate.Test(perm)
}
<commit_msg>reset legacy-cloud-providers modify<commit_after>//go:build !providerless
// +build !providerless

/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package aws

import (
	"encoding/json"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// IPPermissionSet maps IP strings of strings to EC2 IpPermissions
type IPPermissionSet map[string]*ec2.IpPermission

// IPPermissionPredicate is an predicate to test whether IPPermission matches some condition.
type IPPermissionPredicate interface {
	// Test checks whether specified IPPermission matches condition.
	Test(perm *ec2.IpPermission) bool
}

// NewIPPermissionSet creates a new IPPermissionSet
func NewIPPermissionSet(items ...*ec2.IpPermission) IPPermissionSet {
	s := make(IPPermissionSet)
	s.Insert(items...)
	return s
}

// Ungroup splits permissions out into individual permissions
// EC2 will combine permissions with the same port but different SourceRanges together, for example
// We ungroup them so we can process them
func (s IPPermissionSet) Ungroup() IPPermissionSet {
	l := []*ec2.IpPermission{}
	for _, p := range s.List() {
		if len(p.IpRanges) <= 1 {
			l = append(l, p)
			continue
		}
		for _, ipRange := range p.IpRanges {
			c := &ec2.IpPermission{}
			*c = *p
			c.IpRanges = []*ec2.IpRange{ipRange}
			l = append(l, c)
		}
	}

	l2 := []*ec2.IpPermission{}
	for _, p := range l {
		if len(p.UserIdGroupPairs) <= 1 {
			l2 = append(l2, p)
			continue
		}
		for _, u := range p.UserIdGroupPairs {
			c := &ec2.IpPermission{}
			*c = *p
			c.UserIdGroupPairs = []*ec2.UserIdGroupPair{u}
			l2 = append(l, c)
		}
	}

	l3 := []*ec2.IpPermission{}
	for _, p := range l2 {
		if len(p.PrefixListIds) <= 1 {
			l3 = append(l3, p)
			continue
		}
		for _, v := range p.PrefixListIds {
			c := &ec2.IpPermission{}
			*c = *p
			c.PrefixListIds = []*ec2.PrefixListId{v}
			l3 = append(l3, c)
		}
	}

	return NewIPPermissionSet(l3...)
}

// Insert adds items to the set.
func (s IPPermissionSet) Insert(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		s[k] = p
	}
}

// Delete delete permission from the set.
func (s IPPermissionSet) Delete(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		delete(s, k)
	}
}

// DeleteIf delete permission from the set if permission matches predicate.
func (s IPPermissionSet) DeleteIf(predicate IPPermissionPredicate) {
	for k, p := range s {
		if predicate.Test(p) {
			delete(s, k)
		}
	}
}

// List returns the contents as a slice.  Order is not defined.
func (s IPPermissionSet) List() []*ec2.IpPermission {
	res := make([]*ec2.IpPermission, 0, len(s))
	for _, v := range s {
		res = append(res, v)
	}
	return res
}

// IsSuperset returns true if and only if s is a superset of s2.
func (s IPPermissionSet) IsSuperset(s2 IPPermissionSet) bool {
	for k := range s2 {
		_, found := s[k]
		if !found {
			return false
		}
	}
	return true
}

// Equal returns true if and only if s is equal (as a set) to s2.
// Two sets are equal if their membership is identical.
// (In practice, this means same elements, order doesn't matter)
func (s IPPermissionSet) Equal(s2 IPPermissionSet) bool {
	return len(s) == len(s2) && s.IsSuperset(s2)
}

// Difference returns a set of objects that are not in s2
// For example:
// s1 = {a1, a2, a3}
// s2 = {a1, a2, a4, a5}
// s1.Difference(s2) = {a3}
// s2.Difference(s1) = {a4, a5}
func (s IPPermissionSet) Difference(s2 IPPermissionSet) IPPermissionSet {
	result := NewIPPermissionSet()
	for k, v := range s {
		_, found := s2[k]
		if !found {
			result[k] = v
		}
	}
	return result
}

// Len returns the size of the set.
func (s IPPermissionSet) Len() int {
	return len(s)
}

func keyForIPPermission(p *ec2.IpPermission) string {
	v, err := json.Marshal(p)
	if err != nil {
		panic(fmt.Sprintf("error building JSON representation of ec2.IpPermission: %v", err))
	}
	return string(v)
}

var _ IPPermissionPredicate = IPPermissionMatchDesc{}

// IPPermissionMatchDesc checks whether specific IPPermission contains description.
type IPPermissionMatchDesc struct {
	Description string
}

// Test whether specific IPPermission contains description.
func (p IPPermissionMatchDesc) Test(perm *ec2.IpPermission) bool {
	for _, v4Range := range perm.IpRanges {
		if aws.StringValue(v4Range.Description) == p.Description {
			return true
		}
	}
	for _, v6Range := range perm.Ipv6Ranges {
		if aws.StringValue(v6Range.Description) == p.Description {
			return true
		}
	}
	for _, prefixListID := range perm.PrefixListIds {
		if aws.StringValue(prefixListID.Description) == p.Description {
			return true
		}
	}
	for _, group := range perm.UserIdGroupPairs {
		if aws.StringValue(group.Description) == p.Description {
			return true
		}
	}
	return false
}

var _ IPPermissionPredicate = IPPermissionNotMatch{}

// IPPermissionNotMatch is the *not* operator for Predicate
type IPPermissionNotMatch struct {
	Predicate IPPermissionPredicate
}

// Test whether specific IPPermission not match the embed predicate.
func (p IPPermissionNotMatch) Test(perm *ec2.IpPermission) bool {
	return !p.Predicate.Test(perm)
}
<|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.
*/

// register package is introduced in order to avoid circle imports between openstack and cloudprovider packages.
package register

import (
	"fmt"

	"github.com/gophercloud/gophercloud"
	gopenstack "github.com/gophercloud/gophercloud/openstack"
	"k8s.io/client-go/kubernetes"

	"k8s.io/cloud-provider-openstack/pkg/autohealing/cloudprovider"
	"k8s.io/cloud-provider-openstack/pkg/autohealing/cloudprovider/openstack"
	"k8s.io/cloud-provider-openstack/pkg/autohealing/config"
	"k8s.io/cloud-provider-openstack/pkg/client"
)

func registerOpenStack(cfg config.Config, kubeClient kubernetes.Interface) (cloudprovider.CloudProvider, error) {
	client, err := client.NewOpenStackClient(&cfg.OpenStack, "magnum-auto-healer")
	if err != nil {
		return nil, err
	}

	eoOpts := gophercloud.EndpointOpts{
		Region:       cfg.OpenStack.Region,
		Availability: cfg.OpenStack.EndpointType,
	}

	// get nova service client
	var novaClient *gophercloud.ServiceClient
	novaClient, err = gopenstack.NewComputeV2(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Nova service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}

	// get heat service client
	var heatClient *gophercloud.ServiceClient
	heatClient, err = gopenstack.NewOrchestrationV1(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Heat service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}

	// get magnum service client
	var magnumClient *gophercloud.ServiceClient
	magnumClient, err = gopenstack.NewContainerInfraV1(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Magnum service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}
	magnumClient.Microversion = "latest"

	// get cinder service client
	var cinderClient *gophercloud.ServiceClient
	cinderClient, err = gopenstack.NewBlockStorageV2(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Cinder service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}

	var p cloudprovider.CloudProvider
	p = openstack.OpenStackCloudProvider{
		KubeClient: kubeClient,
		Nova:       novaClient,
		Heat:       heatClient,
		Magnum:     magnumClient,
		Cinder:     cinderClient,
		Config:     cfg,
	}

	return p, nil
}

func init() {
	cloudprovider.RegisterCloudProvider(openstack.ProviderName, registerOpenStack)
}
<commit_msg>magnum-auto-healer: use Cinder v3 (#1769)<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.
*/

// register package is introduced in order to avoid circle imports between openstack and cloudprovider packages.
package register

import (
	"fmt"

	"github.com/gophercloud/gophercloud"
	gopenstack "github.com/gophercloud/gophercloud/openstack"
	"k8s.io/client-go/kubernetes"

	"k8s.io/cloud-provider-openstack/pkg/autohealing/cloudprovider"
	"k8s.io/cloud-provider-openstack/pkg/autohealing/cloudprovider/openstack"
	"k8s.io/cloud-provider-openstack/pkg/autohealing/config"
	"k8s.io/cloud-provider-openstack/pkg/client"
)

func registerOpenStack(cfg config.Config, kubeClient kubernetes.Interface) (cloudprovider.CloudProvider, error) {
	client, err := client.NewOpenStackClient(&cfg.OpenStack, "magnum-auto-healer")
	if err != nil {
		return nil, err
	}

	eoOpts := gophercloud.EndpointOpts{
		Region:       cfg.OpenStack.Region,
		Availability: cfg.OpenStack.EndpointType,
	}

	// get nova service client
	var novaClient *gophercloud.ServiceClient
	novaClient, err = gopenstack.NewComputeV2(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Nova service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}

	// get heat service client
	var heatClient *gophercloud.ServiceClient
	heatClient, err = gopenstack.NewOrchestrationV1(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Heat service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}

	// get magnum service client
	var magnumClient *gophercloud.ServiceClient
	magnumClient, err = gopenstack.NewContainerInfraV1(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Magnum service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}
	magnumClient.Microversion = "latest"

	// get cinder service client
	var cinderClient *gophercloud.ServiceClient
	cinderClient, err = gopenstack.NewBlockStorageV3(client, eoOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to find Cinder service endpoint in the region %s: %v", cfg.OpenStack.Region, err)
	}

	var p cloudprovider.CloudProvider
	p = openstack.OpenStackCloudProvider{
		KubeClient: kubeClient,
		Nova:       novaClient,
		Heat:       heatClient,
		Magnum:     magnumClient,
		Cinder:     cinderClient,
		Config:     cfg,
	}

	return p, nil
}

func init() {
	cloudprovider.RegisterCloudProvider(openstack.ProviderName, registerOpenStack)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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 openstack

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strings"
	"time"

	"gopkg.in/gcfg.v1"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/openstack"
	"github.com/rackspace/gophercloud/openstack/compute/v2/servers"
	"github.com/rackspace/gophercloud/pagination"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/cloudprovider"
	"k8s.io/kubernetes/pkg/types"
)

const ProviderName = "openstack"

var ErrNotFound = errors.New("Failed to find object")
var ErrMultipleResults = errors.New("Multiple results where only one expected")
var ErrNoAddressFound = errors.New("No address found for host")
var ErrAttrNotFound = errors.New("Expected attribute not found")

const (
	MiB = 1024 * 1024
	GB  = 1000 * 1000 * 1000
)

// encoding.TextUnmarshaler interface for time.Duration
type MyDuration struct {
	time.Duration
}

func (d *MyDuration) UnmarshalText(text []byte) error {
	res, err := time.ParseDuration(string(text))
	if err != nil {
		return err
	}
	d.Duration = res
	return nil
}

type LoadBalancer struct {
	network *gophercloud.ServiceClient
	compute *gophercloud.ServiceClient
	opts    LoadBalancerOpts
}

type LoadBalancerOpts struct {
	LBVersion            string     `gcfg:"lb-version"` // overrides autodetection. v1 or v2
	SubnetId             string     `gcfg:"subnet-id"`  // required
	FloatingNetworkId    string     `gcfg:"floating-network-id"`
	LBMethod             string     `gcfg:"lb-method"`
	CreateMonitor        bool       `gcfg:"create-monitor"`
	MonitorDelay         MyDuration `gcfg:"monitor-delay"`
	MonitorTimeout       MyDuration `gcfg:"monitor-timeout"`
	MonitorMaxRetries    uint       `gcfg:"monitor-max-retries"`
	ManageSecurityGroups bool       `gcfg:"manage-security-groups"`
	NodeSecurityGroupID  string     `gcfg:"node-security-group"`
}

// OpenStack is an implementation of cloud provider Interface for OpenStack.
type OpenStack struct {
	provider *gophercloud.ProviderClient
	region   string
	lbOpts   LoadBalancerOpts
	// InstanceID of the server where this OpenStack object is instantiated.
	localInstanceID string
}

type Config struct {
	Global struct {
		AuthUrl    string `gcfg:"auth-url"`
		Username   string
		UserId     string `gcfg:"user-id"`
		Password   string
		ApiKey     string `gcfg:"api-key"`
		TenantId   string `gcfg:"tenant-id"`
		TenantName string `gcfg:"tenant-name"`
		DomainId   string `gcfg:"domain-id"`
		DomainName string `gcfg:"domain-name"`
		Region     string
	}
	LoadBalancer LoadBalancerOpts
}

func init() {
	cloudprovider.RegisterCloudProvider(ProviderName, func(config io.Reader) (cloudprovider.Interface, error) {
		cfg, err := readConfig(config)
		if err != nil {
			return nil, err
		}
		return newOpenStack(cfg)
	})
}

func (cfg Config) toAuthOptions() gophercloud.AuthOptions {
	return gophercloud.AuthOptions{
		IdentityEndpoint: cfg.Global.AuthUrl,
		Username:         cfg.Global.Username,
		UserID:           cfg.Global.UserId,
		Password:         cfg.Global.Password,
		APIKey:           cfg.Global.ApiKey,
		TenantID:         cfg.Global.TenantId,
		TenantName:       cfg.Global.TenantName,
		DomainID:         cfg.Global.DomainId,
		DomainName:       cfg.Global.DomainName,

		// Persistent service, so we need to be able to renew tokens.
		AllowReauth: true,
	}
}

func readConfig(config io.Reader) (Config, error) {
	if config == nil {
		err := fmt.Errorf("no OpenStack cloud provider config file given")
		return Config{}, err
	}

	var cfg Config
	err := gcfg.ReadInto(&cfg, config)
	return cfg, err
}

func readInstanceID() (string, error) {
	// Try to find instance ID on the local filesystem (created by cloud-init)
	const instanceIDFile = "/var/lib/cloud/data/instance-id"
	idBytes, err := ioutil.ReadFile(instanceIDFile)
	if err == nil {
		instanceID := string(idBytes)
		instanceID = strings.TrimSpace(instanceID)
		glog.V(3).Infof("Got instance id from %s: %s", instanceIDFile, instanceID)
		if instanceID != "" {
			return instanceID, nil
		}
		// Fall through to metadata server lookup
	}

	md, err := getMetadata()
	if err != nil {
		return "", err
	}

	return md.Uuid, nil
}

func newOpenStack(cfg Config) (*OpenStack, error) {
	provider, err := openstack.AuthenticatedClient(cfg.toAuthOptions())
	if err != nil {
		return nil, err
	}

	id, err := readInstanceID()
	if err != nil {
		return nil, err
	}

	os := OpenStack{
		provider:        provider,
		region:          cfg.Global.Region,
		lbOpts:          cfg.LoadBalancer,
		localInstanceID: id,
	}

	return &os, nil
}

// mapNodeNameToServerName maps a k8s NodeName to an OpenStack Server Name
// This is a simple string cast.
func mapNodeNameToServerName(nodeName types.NodeName) string {
	return string(nodeName)
}

// mapServerToNodeName maps an OpenStack Server to a k8s NodeName
func mapServerToNodeName(server *servers.Server) types.NodeName {
	return types.NodeName(server.Name)
}

func getServerByName(client *gophercloud.ServiceClient, name types.NodeName) (*servers.Server, error) {
	opts := servers.ListOpts{
		Name:   fmt.Sprintf("^%s$", regexp.QuoteMeta(mapNodeNameToServerName(name))),
		Status: "ACTIVE",
	}
	pager := servers.List(client, opts)

	serverList := make([]servers.Server, 0, 1)

	err := pager.EachPage(func(page pagination.Page) (bool, error) {
		s, err := servers.ExtractServers(page)
		if err != nil {
			return false, err
		}
		serverList = append(serverList, s...)
		if len(serverList) > 1 {
			return false, ErrMultipleResults
		}
		return true, nil
	})
	if err != nil {
		return nil, err
	}

	if len(serverList) == 0 {
		return nil, ErrNotFound
	} else if len(serverList) > 1 {
		return nil, ErrMultipleResults
	}

	return &serverList[0], nil
}

func getAddressesByName(client *gophercloud.ServiceClient, name types.NodeName) ([]api.NodeAddress, error) {
	srv, err := getServerByName(client, name)
	if err != nil {
		return nil, err
	}

	addrs := []api.NodeAddress{}

	for network, netblob := range srv.Addresses {
		list, ok := netblob.([]interface{})
		if !ok {
			continue
		}

		for _, item := range list {
			var addressType api.NodeAddressType

			props, ok := item.(map[string]interface{})
			if !ok {
				continue
			}

			extIPType, ok := props["OS-EXT-IPS:type"]
			if (ok && extIPType == "floating") || (!ok && network == "public") {
				addressType = api.NodeExternalIP
			} else {
				addressType = api.NodeInternalIP
			}

			tmp, ok := props["addr"]
			if !ok {
				continue
			}
			addr, ok := tmp.(string)
			if !ok {
				continue
			}

			api.AddToNodeAddresses(&addrs,
				api.NodeAddress{
					Type:    addressType,
					Address: addr,
				},
			)
		}
	}

	// AccessIPs are usually duplicates of "public" addresses.
	if srv.AccessIPv4 != "" {
		api.AddToNodeAddresses(&addrs,
			api.NodeAddress{
				Type:    api.NodeExternalIP,
				Address: srv.AccessIPv4,
			},
		)
	}

	if srv.AccessIPv6 != "" {
		api.AddToNodeAddresses(&addrs,
			api.NodeAddress{
				Type:    api.NodeExternalIP,
				Address: srv.AccessIPv6,
			},
		)
	}

	return addrs, nil
}

func getAddressByName(client *gophercloud.ServiceClient, name types.NodeName) (string, error) {
	addrs, err := getAddressesByName(client, name)
	if err != nil {
		return "", err
	} else if len(addrs) == 0 {
		return "", ErrNoAddressFound
	}

	for _, addr := range addrs {
		if addr.Type == api.NodeInternalIP {
			return addr.Address, nil
		}
	}

	return addrs[0].Address, nil
}

func (os *OpenStack) Clusters() (cloudprovider.Clusters, bool) {
	return nil, false
}

// ProviderName returns the cloud provider ID.
func (os *OpenStack) ProviderName() string {
	return ProviderName
}

// ScrubDNS filters DNS settings for pods.
func (os *OpenStack) ScrubDNS(nameservers, searches []string) (nsOut, srchOut []string) {
	return nameservers, searches
}

func (os *OpenStack) LoadBalancer() (cloudprovider.LoadBalancer, bool) {
	glog.V(4).Info("openstack.LoadBalancer() called")

	// TODO: Search for and support Rackspace loadbalancer API, and others.
	network, err := openstack.NewNetworkV2(os.provider, gophercloud.EndpointOpts{
		Region: os.region,
	})
	if err != nil {
		glog.Warningf("Failed to find neutron endpoint: %v", err)
		return nil, false
	}

	compute, err := openstack.NewComputeV2(os.provider, gophercloud.EndpointOpts{
		Region: os.region,
	})
	if err != nil {
		glog.Warningf("Failed to find compute endpoint: %v", err)
		return nil, false
	}

	lbversion := os.lbOpts.LBVersion
	if lbversion == "" {
		// No version specified, try newest supported by server
		netExts, err := networkExtensions(network)
		if err != nil {
			glog.Warningf("Failed to list neutron extensions: %v", err)
			return nil, false
		}

		if netExts["lbaasv2"] {
			lbversion = "v2"
		} else if netExts["lbaas"] {
			lbversion = "v1"
		} else {
			glog.Warningf("Failed to find neutron LBaaS extension (v1 or v2)")
			return nil, false
		}
		glog.V(3).Infof("Using LBaaS extension %v", lbversion)
	}

	glog.V(1).Info("Claiming to support LoadBalancer")

	if os.lbOpts.LBVersion == "v2" {
		return &LbaasV2{LoadBalancer{network, compute, os.lbOpts}}, true
	} else if lbversion == "v1" {
		return &LbaasV1{LoadBalancer{network, compute, os.lbOpts}}, true
	} else {
		glog.Warningf("Config error: unrecognised lb-version \"%v\"", lbversion)
		return nil, false
	}
}

func isNotFound(err error) bool {
	e, ok := err.(*gophercloud.UnexpectedResponseCodeError)
	return ok && e.Actual == http.StatusNotFound
}

func (os *OpenStack) Zones() (cloudprovider.Zones, bool) {
	glog.V(1).Info("Claiming to support Zones")

	return os, true
}
func (os *OpenStack) GetZone() (cloudprovider.Zone, error) {
	md, err := getMetadata()
	if err != nil {
		return cloudprovider.Zone{}, err
	}

	zone := cloudprovider.Zone{
		FailureDomain: md.AvailabilityZone,
		Region:        os.region,
	}
	glog.V(1).Infof("Current zone is %v", zone)

	return zone, nil
}

func (os *OpenStack) Routes() (cloudprovider.Routes, bool) {
	return nil, false
}
<commit_msg>Support trust id as a scope in the OpenStack authentication logic.<commit_after>/*
Copyright 2014 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 openstack

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strings"
	"time"

	"gopkg.in/gcfg.v1"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/openstack"
	"github.com/rackspace/gophercloud/openstack/compute/v2/servers"
	"github.com/rackspace/gophercloud/openstack/identity/v3/extensions/trust"
	token3 "github.com/rackspace/gophercloud/openstack/identity/v3/tokens"
	"github.com/rackspace/gophercloud/pagination"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/cloudprovider"
	"k8s.io/kubernetes/pkg/types"
)

const ProviderName = "openstack"

var ErrNotFound = errors.New("Failed to find object")
var ErrMultipleResults = errors.New("Multiple results where only one expected")
var ErrNoAddressFound = errors.New("No address found for host")
var ErrAttrNotFound = errors.New("Expected attribute not found")

const (
	MiB = 1024 * 1024
	GB  = 1000 * 1000 * 1000
)

// encoding.TextUnmarshaler interface for time.Duration
type MyDuration struct {
	time.Duration
}

func (d *MyDuration) UnmarshalText(text []byte) error {
	res, err := time.ParseDuration(string(text))
	if err != nil {
		return err
	}
	d.Duration = res
	return nil
}

type LoadBalancer struct {
	network *gophercloud.ServiceClient
	compute *gophercloud.ServiceClient
	opts    LoadBalancerOpts
}

type LoadBalancerOpts struct {
	LBVersion            string     `gcfg:"lb-version"` // overrides autodetection. v1 or v2
	SubnetId             string     `gcfg:"subnet-id"`  // required
	FloatingNetworkId    string     `gcfg:"floating-network-id"`
	LBMethod             string     `gcfg:"lb-method"`
	CreateMonitor        bool       `gcfg:"create-monitor"`
	MonitorDelay         MyDuration `gcfg:"monitor-delay"`
	MonitorTimeout       MyDuration `gcfg:"monitor-timeout"`
	MonitorMaxRetries    uint       `gcfg:"monitor-max-retries"`
	ManageSecurityGroups bool       `gcfg:"manage-security-groups"`
	NodeSecurityGroupID  string     `gcfg:"node-security-group"`
}

// OpenStack is an implementation of cloud provider Interface for OpenStack.
type OpenStack struct {
	provider *gophercloud.ProviderClient
	region   string
	lbOpts   LoadBalancerOpts
	// InstanceID of the server where this OpenStack object is instantiated.
	localInstanceID string
}

type Config struct {
	Global struct {
		AuthUrl    string `gcfg:"auth-url"`
		Username   string
		UserId     string `gcfg:"user-id"`
		Password   string
		ApiKey     string `gcfg:"api-key"`
		TenantId   string `gcfg:"tenant-id"`
		TenantName string `gcfg:"tenant-name"`
		TrustId    string `gcfg:"trust-id"`
		DomainId   string `gcfg:"domain-id"`
		DomainName string `gcfg:"domain-name"`
		Region     string
	}
	LoadBalancer LoadBalancerOpts
}

func init() {
	cloudprovider.RegisterCloudProvider(ProviderName, func(config io.Reader) (cloudprovider.Interface, error) {
		cfg, err := readConfig(config)
		if err != nil {
			return nil, err
		}
		return newOpenStack(cfg)
	})
}

func (cfg Config) toAuthOptions() gophercloud.AuthOptions {
	return gophercloud.AuthOptions{
		IdentityEndpoint: cfg.Global.AuthUrl,
		Username:         cfg.Global.Username,
		UserID:           cfg.Global.UserId,
		Password:         cfg.Global.Password,
		APIKey:           cfg.Global.ApiKey,
		TenantID:         cfg.Global.TenantId,
		TenantName:       cfg.Global.TenantName,
		DomainID:         cfg.Global.DomainId,
		DomainName:       cfg.Global.DomainName,

		// Persistent service, so we need to be able to renew tokens.
		AllowReauth: true,
	}
}

func readConfig(config io.Reader) (Config, error) {
	if config == nil {
		err := fmt.Errorf("no OpenStack cloud provider config file given")
		return Config{}, err
	}

	var cfg Config
	err := gcfg.ReadInto(&cfg, config)
	return cfg, err
}

func readInstanceID() (string, error) {
	// Try to find instance ID on the local filesystem (created by cloud-init)
	const instanceIDFile = "/var/lib/cloud/data/instance-id"
	idBytes, err := ioutil.ReadFile(instanceIDFile)
	if err == nil {
		instanceID := string(idBytes)
		instanceID = strings.TrimSpace(instanceID)
		glog.V(3).Infof("Got instance id from %s: %s", instanceIDFile, instanceID)
		if instanceID != "" {
			return instanceID, nil
		}
		// Fall through to metadata server lookup
	}

	md, err := getMetadata()
	if err != nil {
		return "", err
	}

	return md.Uuid, nil
}

func newOpenStack(cfg Config) (*OpenStack, error) {
	provider, err := openstack.NewClient(cfg.Global.AuthUrl)
	if err != nil {
		return nil, err
	}
	if cfg.Global.TrustId != "" {
		authOptionsExt := trust.AuthOptionsExt{
			TrustID:     cfg.Global.TrustId,
			AuthOptions: token3.AuthOptions{AuthOptions: cfg.toAuthOptions()},
		}
		err = trust.AuthenticateV3Trust(provider, authOptionsExt)
	} else {
		err = openstack.Authenticate(provider, cfg.toAuthOptions())
	}

	if err != nil {
		return nil, err
	}

	id, err := readInstanceID()
	if err != nil {
		return nil, err
	}

	os := OpenStack{
		provider:        provider,
		region:          cfg.Global.Region,
		lbOpts:          cfg.LoadBalancer,
		localInstanceID: id,
	}

	return &os, nil
}

// mapNodeNameToServerName maps a k8s NodeName to an OpenStack Server Name
// This is a simple string cast.
func mapNodeNameToServerName(nodeName types.NodeName) string {
	return string(nodeName)
}

// mapServerToNodeName maps an OpenStack Server to a k8s NodeName
func mapServerToNodeName(server *servers.Server) types.NodeName {
	return types.NodeName(server.Name)
}

func getServerByName(client *gophercloud.ServiceClient, name types.NodeName) (*servers.Server, error) {
	opts := servers.ListOpts{
		Name:   fmt.Sprintf("^%s$", regexp.QuoteMeta(mapNodeNameToServerName(name))),
		Status: "ACTIVE",
	}
	pager := servers.List(client, opts)

	serverList := make([]servers.Server, 0, 1)

	err := pager.EachPage(func(page pagination.Page) (bool, error) {
		s, err := servers.ExtractServers(page)
		if err != nil {
			return false, err
		}
		serverList = append(serverList, s...)
		if len(serverList) > 1 {
			return false, ErrMultipleResults
		}
		return true, nil
	})
	if err != nil {
		return nil, err
	}

	if len(serverList) == 0 {
		return nil, ErrNotFound
	} else if len(serverList) > 1 {
		return nil, ErrMultipleResults
	}

	return &serverList[0], nil
}

func getAddressesByName(client *gophercloud.ServiceClient, name types.NodeName) ([]api.NodeAddress, error) {
	srv, err := getServerByName(client, name)
	if err != nil {
		return nil, err
	}

	addrs := []api.NodeAddress{}

	for network, netblob := range srv.Addresses {
		list, ok := netblob.([]interface{})
		if !ok {
			continue
		}

		for _, item := range list {
			var addressType api.NodeAddressType

			props, ok := item.(map[string]interface{})
			if !ok {
				continue
			}

			extIPType, ok := props["OS-EXT-IPS:type"]
			if (ok && extIPType == "floating") || (!ok && network == "public") {
				addressType = api.NodeExternalIP
			} else {
				addressType = api.NodeInternalIP
			}

			tmp, ok := props["addr"]
			if !ok {
				continue
			}
			addr, ok := tmp.(string)
			if !ok {
				continue
			}

			api.AddToNodeAddresses(&addrs,
				api.NodeAddress{
					Type:    addressType,
					Address: addr,
				},
			)
		}
	}

	// AccessIPs are usually duplicates of "public" addresses.
	if srv.AccessIPv4 != "" {
		api.AddToNodeAddresses(&addrs,
			api.NodeAddress{
				Type:    api.NodeExternalIP,
				Address: srv.AccessIPv4,
			},
		)
	}

	if srv.AccessIPv6 != "" {
		api.AddToNodeAddresses(&addrs,
			api.NodeAddress{
				Type:    api.NodeExternalIP,
				Address: srv.AccessIPv6,
			},
		)
	}

	return addrs, nil
}

func getAddressByName(client *gophercloud.ServiceClient, name types.NodeName) (string, error) {
	addrs, err := getAddressesByName(client, name)
	if err != nil {
		return "", err
	} else if len(addrs) == 0 {
		return "", ErrNoAddressFound
	}

	for _, addr := range addrs {
		if addr.Type == api.NodeInternalIP {
			return addr.Address, nil
		}
	}

	return addrs[0].Address, nil
}

func (os *OpenStack) Clusters() (cloudprovider.Clusters, bool) {
	return nil, false
}

// ProviderName returns the cloud provider ID.
func (os *OpenStack) ProviderName() string {
	return ProviderName
}

// ScrubDNS filters DNS settings for pods.
func (os *OpenStack) ScrubDNS(nameservers, searches []string) (nsOut, srchOut []string) {
	return nameservers, searches
}

func (os *OpenStack) LoadBalancer() (cloudprovider.LoadBalancer, bool) {
	glog.V(4).Info("openstack.LoadBalancer() called")

	// TODO: Search for and support Rackspace loadbalancer API, and others.
	network, err := openstack.NewNetworkV2(os.provider, gophercloud.EndpointOpts{
		Region: os.region,
	})
	if err != nil {
		glog.Warningf("Failed to find neutron endpoint: %v", err)
		return nil, false
	}

	compute, err := openstack.NewComputeV2(os.provider, gophercloud.EndpointOpts{
		Region: os.region,
	})
	if err != nil {
		glog.Warningf("Failed to find compute endpoint: %v", err)
		return nil, false
	}

	lbversion := os.lbOpts.LBVersion
	if lbversion == "" {
		// No version specified, try newest supported by server
		netExts, err := networkExtensions(network)
		if err != nil {
			glog.Warningf("Failed to list neutron extensions: %v", err)
			return nil, false
		}

		if netExts["lbaasv2"] {
			lbversion = "v2"
		} else if netExts["lbaas"] {
			lbversion = "v1"
		} else {
			glog.Warningf("Failed to find neutron LBaaS extension (v1 or v2)")
			return nil, false
		}
		glog.V(3).Infof("Using LBaaS extension %v", lbversion)
	}

	glog.V(1).Info("Claiming to support LoadBalancer")

	if os.lbOpts.LBVersion == "v2" {
		return &LbaasV2{LoadBalancer{network, compute, os.lbOpts}}, true
	} else if lbversion == "v1" {
		return &LbaasV1{LoadBalancer{network, compute, os.lbOpts}}, true
	} else {
		glog.Warningf("Config error: unrecognised lb-version \"%v\"", lbversion)
		return nil, false
	}
}

func isNotFound(err error) bool {
	e, ok := err.(*gophercloud.UnexpectedResponseCodeError)
	return ok && e.Actual == http.StatusNotFound
}

func (os *OpenStack) Zones() (cloudprovider.Zones, bool) {
	glog.V(1).Info("Claiming to support Zones")

	return os, true
}
func (os *OpenStack) GetZone() (cloudprovider.Zone, error) {
	md, err := getMetadata()
	if err != nil {
		return cloudprovider.Zone{}, err
	}

	zone := cloudprovider.Zone{
		FailureDomain: md.AvailabilityZone,
		Region:        os.region,
	}
	glog.V(1).Infof("Current zone is %v", zone)

	return zone, nil
}

func (os *OpenStack) Routes() (cloudprovider.Routes, bool) {
	return nil, false
}
<|endoftext|>


<commit_before>package integration

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"time"

	"github.com/cloudfoundry/bosh-agent/agentclient"
	"github.com/cloudfoundry/bosh-agent/agentclient/http"
	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	"github.com/cloudfoundry/bosh-utils/httpclient"
	"github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type TestEnvironment struct {
	cmdRunner        boshsys.CmdRunner
	currentDeviceNum int
	sshTunnelProc    boshsys.Process
	logger           logger.Logger
	agentClient      agentclient.AgentClient
}

func NewTestEnvironment(
	cmdRunner boshsys.CmdRunner,
) *TestEnvironment {
	return &TestEnvironment{
		cmdRunner:        cmdRunner,
		currentDeviceNum: 2,
		logger:           logger.NewLogger(logger.LevelDebug),
	}
}

func (t *TestEnvironment) SetupConfigDrive() error {
	loopDeviceNum, err := t.AttachLoopDevice(10)
	if err != nil {
		return err
	}

	setupConfigDriveTemplate := `
sudo mkfs -t ext3 -m 1 -v /dev/loop%d
sudo e2label /dev/loop%d config-2
sudo rm -rf /tmp/config-drive
sudo mkdir /tmp/config-drive
sudo mount /dev/disk/by-label/config-2 /tmp/config-drive
sudo chown vagrant:vagrant /tmp/config-drive
sudo mkdir -p /tmp/config-drive/ec2/latest
sudo cp %s/meta-data.json /tmp/config-drive/ec2/latest/meta-data.json
sudo cp %s/user-data.json /tmp/config-drive/ec2/latest/user-data.json
sudo umount /tmp/config-drive
`
	setupConfigDriveScript := fmt.Sprintf(setupConfigDriveTemplate, loopDeviceNum, loopDeviceNum, t.assetsDir(), t.assetsDir())

	_, err = t.RunCommand(setupConfigDriveScript)
	return err
}

func (t *TestEnvironment) unmountDir(dir string) error {
	mountPoints, err := t.RunCommand(fmt.Sprintf(`sudo mount | grep "on %s" | cut -d ' ' -f 3`, dir))
	if err != nil {
		return err
	}

	for _, mountPoint := range strings.Split(mountPoints, "\n") {
		if mountPoint != "" {
			t.RunCommand(fmt.Sprintf("sudo umount -l %s", mountPoint))
		}
	}

	_, err = t.RunCommand(fmt.Sprintf("sudo rm -rf %s", dir))
	return err
}

func (t *TestEnvironment) CleanupDataDir() error {
	t.RunCommand(`sudo /var/vcap/bosh/bin/monit stop all`)

	err := t.unmountDir("/var/log")
	if err != nil {
		return err
	}

	return t.unmountDir("/var/vcap/data")
}

// ConfigureAgentForGenericInfrastructure executes the agent_runit.sh asset.
// Required for reverse-compatibility with older bosh-lite
// (remove once a new warden stemcell is built).
func (t *TestEnvironment) ConfigureAgentForGenericInfrastructure() error {
	_, err := t.RunCommand(
		fmt.Sprintf(
			"sudo cp %s/agent_runit.sh /etc/service/agent/run",
			t.assetsDir(),
		),
	)
	return err
}

func (t *TestEnvironment) CleanupLogFile() error {
	_, err := t.RunCommand("sudo truncate -s 0 /var/vcap/bosh/log/current")
	return err
}

func (t *TestEnvironment) LogFileContains(content string) bool {
	_, err := t.RunCommand(fmt.Sprintf(`sudo grep "%s" /var/vcap/bosh/log/current`, content))
	return err == nil
}

func (t *TestEnvironment) AttachDevice(devicePath string, partitionSize, numPartitions int) error {
	partitionPath := devicePath
	for i := 0; i <= numPartitions; i++ {
		if i > 0 {
			partitionPath = fmt.Sprintf("%s%d", devicePath, i)
		}

		loopDeviceNum, err := t.AttachLoopDevice(partitionSize)
		if err != nil {
			return err
		}

		t.RemoveDevice(partitionPath)

		_, err = t.RunCommand(fmt.Sprintf("sudo mknod %s b 7 %d", partitionPath, loopDeviceNum))
		if err != nil {
			return err
		}
	}
	return nil
}

func (t *TestEnvironment) AttachPartitionedRootDevice(devicePath string, sizeInMB, rootPartitionSizeInMB int) (string, string, error) {
	// Partitioner requires fs backed device
	_, err := t.RunCommand(fmt.Sprintf("sudo mknod %s b 7 99", devicePath))
	if err != nil {
		return "", "", err
	}

	attachDeviceTemplate := `
sudo rm -rf /virtual-root-fs
sudo dd if=/dev/zero of=/virtual-root-fs bs=1M count=%d
sudo losetup %s /virtual-root-fs
`
	attachDeviceScript := fmt.Sprintf(attachDeviceTemplate, sizeInMB, devicePath)
	_, err = t.RunCommand(attachDeviceScript)
	if err != nil {
		return "", "", err
	}

	err = t.AttachDevice(devicePath, sizeInMB, 3)
	if err != nil {
		return "", "", err
	}

	// Create only first partition, agent will partition the rest for ephemeral disk
	partitionTemplate := `
echo ',%d,L,' | sudo sfdisk -uM %s
`
	partitionScript := fmt.Sprintf(partitionTemplate, rootPartitionSizeInMB, devicePath)
	_, err = t.RunCommand(partitionScript)
	if err != nil {
		return "", "", err
	}

	rootLink, err := t.RunCommand("df / | grep /dev/ | cut -d' ' -f1")
	if err != nil {
		return "", "", err
	}

	oldRootDevice, err := t.RunCommand(fmt.Sprintf("readlink -f %s", rootLink))
	if err != nil {
		return "", "", err
	}

	// Agent reads the symlink to get root device
	// Replace the symlink with our fake device
	err = t.SwitchRootDevice(devicePath, rootLink)
	if err != nil {
		return "", "", err
	}

	return strings.TrimSpace(oldRootDevice), strings.TrimSpace(rootLink), nil
}

func (t *TestEnvironment) SwitchRootDevice(devicePath, rootLink string) error {
	_, err := t.RunCommand(fmt.Sprintf("sudo rm -f %s", rootLink))
	if err != nil {
		return err
	}

	_, err = t.RunCommand(fmt.Sprintf("sudo ln -s %s1 %s", devicePath, rootLink))
	if err != nil {
		return err
	}

	return nil
}

func (t *TestEnvironment) DetachDevice(devicePath string) error {
	mountPoint, err := t.RunCommand(fmt.Sprintf("sudo mount | grep %s | cut -d ' ' -f 3", devicePath))
	if err != nil {
		return err
	}

	if mountPoint != "" {
		_, err = t.RunCommand(fmt.Sprintf("sudo umount -l %s", mountPoint))
		if err != nil {
			return err
		}
	}
	_, err = t.RunCommand(fmt.Sprintf("sudo rm -f %s*", devicePath))
	return err
}

func (t *TestEnvironment) RemoveDevice(devicePath string) error {
	_, err := t.RunCommand(fmt.Sprintf("sudo rm -f %s", devicePath))
	return err
}

func (t *TestEnvironment) AttachLoopDevice(size int) (int, error) {
	loopDeviceNum := t.currentDeviceNum
	devicePath := fmt.Sprintf("/dev/loop%d", t.currentDeviceNum)

	t.DetachLoopDevice(devicePath)

	attachDeviceTemplate := `
sudo rm -rf /virtualfs-%d
sudo dd if=/dev/zero of=/virtualfs-%d bs=1M count=%d
sudo losetup %s /virtualfs-%d
`
	attachDeviceScript := fmt.Sprintf(attachDeviceTemplate, loopDeviceNum, loopDeviceNum, size, devicePath, loopDeviceNum)
	_, err := t.RunCommand(attachDeviceScript)
	if err != nil {
		return 0, err
	}

	t.currentDeviceNum++

	return loopDeviceNum, nil
}

func (t *TestEnvironment) DetachLoopDevice(devicePath string) error {
	_, err := t.RunCommand(fmt.Sprintf("sudo losetup -d %s", devicePath))
	return err
}

func (t *TestEnvironment) UpdateAgentConfig(configFile string) error {
	_, err := t.RunCommand(
		fmt.Sprintf(
			"sudo cp %s/%s /var/vcap/bosh/agent.json",
			t.assetsDir(),
			configFile,
		),
	)
	return err
}

func (t *TestEnvironment) RestartAgent() error {
	err := t.StopAgent()
	if err != nil {
		return err
	}

	return t.StartAgent()
}

func (t *TestEnvironment) StopAgent() error {
	_, err := t.RunCommand("nohup sudo sv stop agent &")
	return err
}

func (t *TestEnvironment) StartAgent() error {
	_, err := t.RunCommand("nohup sudo sv start agent &")
	return err
}

type emptyReader struct{}

func (er emptyReader) Read(p []byte) (int, error) {
	time.Sleep(1 * time.Second)
	return 0, nil
}

func (t *TestEnvironment) StartAgentTunnel(mbusUser, mbusPass string, mbusPort int) (agentclient.AgentClient, error) {
	if t.sshTunnelProc != nil {
		return nil, fmt.Errorf("Already running")
	}

	sshCmd := boshsys.Command{
		Name: "vagrant",
		Args: []string{
			"ssh",
			"--",
			fmt.Sprintf("-L16868:127.0.0.1:%d", mbusPort),
		},
		Stdin: emptyReader{},
	}
	newTunnelProc, err := t.cmdRunner.RunComplexCommandAsync(sshCmd)
	if err != nil {
		return nil, err
	}
	t.sshTunnelProc = newTunnelProc

	httpClient := httpclient.NewHTTPClient(httpclient.DefaultClient, t.logger)
	mbusURL := fmt.Sprintf("https://%s:%s@localhost:16868", mbusUser, mbusPass)
	client := http.NewAgentClient(mbusURL, "fake-director-uuid", 1*time.Second, 10, httpClient, t.logger)

	for i := 1; i < 1000000; i++ {
		t.logger.Debug("test environment", "Trying to contact agent via ssh tunnel...")
		time.Sleep(1 * time.Second)
		_, err := client.Ping()
		if err == nil {
			break
		}
		t.logger.Debug("test environment", err.Error())
	}
	return client, nil
}

func (t *TestEnvironment) StopAgentTunnel() error {
	if t.sshTunnelProc == nil {
		return fmt.Errorf("Not running")
	}
	t.sshTunnelProc.Wait()
	t.sshTunnelProc = nil
	return nil
}

func (t *TestEnvironment) StartRegistry(settings boshsettings.Settings) error {
	settingsJSON, err := json.Marshal(settings)
	if err != nil {
		return err
	}

	_, err = t.RunCommand("sudo rm -f /var/vcap/bosh/settings.json")
	if err != nil {
		return err
	}

	t.RunCommand("sudo killall -9 fake-registry")

	_, err = t.RunCommand(
		fmt.Sprintf(
			`nohup %s/tmp/fake-registry -user user -password pass -host 127.0.0.1 -port 9090 -instance instance-id -settings %s &> /dev/null &`,
			t.agentDir(),
			strconv.Quote(string(settingsJSON)),
		),
	)
	return err
}

func (t *TestEnvironment) GetVMNetworks() (boshsettings.Networks, error) {
	stdout, _, _, err := t.cmdRunner.RunCommand("vagrant", "status")
	if err != nil {
		return boshsettings.Networks{}, err
	}

	if strings.Contains(stdout, "virtualbox") {
		return boshsettings.Networks{
			"eth0": {
				Type: "dynamic",
			},
			"eth1": {
				Type:    "manual",
				IP:      "192.168.50.4",
				Netmask: "255.255.255.0",
			},
		}, nil
	}

	if strings.Contains(stdout, "aws") {
		return boshsettings.Networks{
			"eth0": {
				Type: "dynamic",
			},
		}, nil
	}

	return boshsettings.Networks{}, nil
}

func (t *TestEnvironment) GetFileContents(filePath string) (string, error) {
	return t.RunCommand(
		fmt.Sprintf(
			`cat %s`,
			filePath,
		),
	)
}

func (t *TestEnvironment) RunCommand(command string) (string, error) {
	stdout, _, _, err := t.cmdRunner.RunCommand(
		"vagrant",
		"ssh",
		"-c",
		command,
	)

	return stdout, err
}

func (t *TestEnvironment) agentDir() string {
	return "/home/vagrant/go/src/github.com/cloudfoundry/bosh-agent"
}

func (t *TestEnvironment) assetsDir() string {
	return fmt.Sprintf("%s/integration/assets", t.agentDir())
}
<commit_msg>Refactor integration test<commit_after>package integration

import (
	"encoding/json"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"time"

	"github.com/cloudfoundry/bosh-agent/agentclient"
	"github.com/cloudfoundry/bosh-agent/agentclient/http"
	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	"github.com/cloudfoundry/bosh-utils/httpclient"
	"github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type TestEnvironment struct {
	cmdRunner        boshsys.CmdRunner
	currentDeviceNum int
	sshTunnelProc    boshsys.Process
	logger           logger.Logger
	agentClient      agentclient.AgentClient
	deviceMap        map[int]string
}

func NewTestEnvironment(
	cmdRunner boshsys.CmdRunner,
) *TestEnvironment {
	return &TestEnvironment{
		cmdRunner:        cmdRunner,
		currentDeviceNum: 2,
		logger:           logger.NewLogger(logger.LevelDebug),
		deviceMap:        make(map[int]string),
	}
}

func (t *TestEnvironment) SetupConfigDrive() error {
	deviceNum, err := t.AttachLoopDevice(10)
	if err != nil {
		return err
	}

	setupConfigDriveTemplate := `
sudo mkfs -t ext3 -m 1 -v %s
sudo e2label %s config-2
sudo rm -rf /tmp/config-drive
sudo mkdir /tmp/config-drive
sudo mount /dev/disk/by-label/config-2 /tmp/config-drive
sudo chown vagrant:vagrant /tmp/config-drive
sudo mkdir -p /tmp/config-drive/ec2/latest
sudo cp %s/meta-data.json /tmp/config-drive/ec2/latest/meta-data.json
sudo cp %s/user-data.json /tmp/config-drive/ec2/latest/user-data.json
sudo umount /tmp/config-drive
`
	setupConfigDriveScript := fmt.Sprintf(setupConfigDriveTemplate, t.deviceMap[deviceNum], t.deviceMap[deviceNum], t.assetsDir(), t.assetsDir())

	_, err = t.RunCommand(setupConfigDriveScript)
	return err
}

type byLen []string

func (a byLen) Len() int           { return len(a) }
func (a byLen) Less(i, j int) bool { return len(a[i]) > len(a[j]) }
func (a byLen) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func (t *TestEnvironment) DetachDevice(dir string) error {
	mountPoints, err := t.RunCommand(fmt.Sprintf(`sudo mount | grep "on %s" | cut -d ' ' -f 3`, dir))
	if err != nil {
		return err
	}

	mountPointsSlice := strings.Split(mountPoints, "\n")
	sort.Sort(byLen(mountPointsSlice))
	for _, mountPoint := range mountPointsSlice {
		if mountPoint != "" {
			t.RunCommand(fmt.Sprintf("sudo fuser -km %s", mountPoint))
			t.RunCommand(fmt.Sprintf("sudo umount %s", mountPoint))
		}
	}

	_, err = t.RunCommand(fmt.Sprintf("sudo rm -rf %s", dir))
	return err
}

func (t *TestEnvironment) CleanupDataDir() error {
	t.RunCommand(`sudo /var/vcap/bosh/bin/monit stop all`)

	err := t.DetachDevice("/var/log")
	if err != nil {
		return err
	}

	return t.DetachDevice("/var/vcap/data")
}

// ConfigureAgentForGenericInfrastructure executes the agent_runit.sh asset.
// Required for reverse-compatibility with older bosh-lite
// (remove once a new warden stemcell is built).
func (t *TestEnvironment) ConfigureAgentForGenericInfrastructure() error {
	_, err := t.RunCommand(
		fmt.Sprintf(
			"sudo cp %s/agent_runit.sh /etc/service/agent/run",
			t.assetsDir(),
		),
	)
	return err
}

func (t *TestEnvironment) CleanupLogFile() error {
	_, err := t.RunCommand("sudo truncate -s 0 /var/vcap/bosh/log/current")
	return err
}

func (t *TestEnvironment) LogFileContains(content string) bool {
	_, err := t.RunCommand(fmt.Sprintf(`sudo grep "%s" /var/vcap/bosh/log/current`, content))
	return err == nil
}

func (t *TestEnvironment) AttachDevice(devicePath string, partitionSize, numPartitions int) error {
	partitionPath := devicePath
	for i := 0; i <= numPartitions; i++ {
		if i > 0 {
			partitionPath = fmt.Sprintf("%s%d", devicePath, i)
		}

		deviceNum, err := t.AttachLoopDevice(partitionSize)
		if err != nil {
			return err
		}

		output, err := t.RunCommand(fmt.Sprintf("ls -al %s | cut -d' ' -f 6", t.deviceMap[deviceNum]))
		minorNum := strings.TrimSpace(output)
		if err != nil {
			return err
		}

		t.RemoveDevice(partitionPath)

		_, err = t.RunCommand(fmt.Sprintf("sudo mknod %s b 7 %s", partitionPath, minorNum))
		if err != nil {
			return err
		}
	}
	return nil
}

func (t *TestEnvironment) AttachPartitionedRootDevice(devicePath string, sizeInMB, rootPartitionSizeInMB int) (string, string, error) {
	// Partitioner requires fs backed device
	_, err := t.RunCommand(fmt.Sprintf("sudo mknod %s b 7 99", devicePath))
	if err != nil {
		return "", "", err
	}

	attachDeviceTemplate := `
sudo rm -rf /virtual-root-fs
sudo dd if=/dev/zero of=/virtual-root-fs bs=1M count=%d
sudo losetup %s /virtual-root-fs
`
	attachDeviceScript := fmt.Sprintf(attachDeviceTemplate, sizeInMB, devicePath)
	_, err = t.RunCommand(attachDeviceScript)
	if err != nil {
		return "", "", err
	}

	err = t.AttachDevice(devicePath, sizeInMB, 3)
	if err != nil {
		return "", "", err
	}

	// Create only first partition, agent will partition the rest for ephemeral disk
	partitionTemplate := `
echo ',%d,L,' | sudo sfdisk -uM %s
`
	partitionScript := fmt.Sprintf(partitionTemplate, rootPartitionSizeInMB, devicePath)
	_, err = t.RunCommand(partitionScript)
	if err != nil {
		return "", "", err
	}

	rootLink, err := t.RunCommand("df / | grep /dev/ | cut -d' ' -f1")
	if err != nil {
		return "", "", err
	}

	oldRootDevice, err := t.RunCommand(fmt.Sprintf("readlink -f %s", rootLink))
	if err != nil {
		return "", "", err
	}

	// Agent reads the symlink to get root device
	// Replace the symlink with our fake device
	err = t.SwitchRootDevice(devicePath, rootLink)
	if err != nil {
		return "", "", err
	}

	return strings.TrimSpace(oldRootDevice), strings.TrimSpace(rootLink), nil
}

func (t *TestEnvironment) SwitchRootDevice(devicePath, rootLink string) error {
	_, err := t.RunCommand(fmt.Sprintf("sudo rm -f %s", rootLink))
	if err != nil {
		return err
	}

	_, err = t.RunCommand(fmt.Sprintf("sudo ln -s %s1 %s", devicePath, rootLink))
	if err != nil {
		return err
	}

	return nil
}

func (t *TestEnvironment) RemoveDevice(devicePath string) error {
	_, err := t.RunCommand(fmt.Sprintf("sudo rm -f %s", devicePath))
	return err
}

func (t *TestEnvironment) AttachLoopDevice(size int) (int, error) {
	deviceNum := t.currentDeviceNum

	output, err := t.RunCommand("sudo losetup -f")
	devicePath := strings.TrimSpace(output)
	if err != nil {
		return 0, err
	}

	if oldDevicePath, ok := t.deviceMap[deviceNum]; ok {
		t.DetachLoopDevice(oldDevicePath)
	}

	attachDeviceTemplate := `
sudo rm -rf /virtualfs-%d
sudo dd if=/dev/zero of=/virtualfs-%d bs=1M count=%d
sudo losetup %s /virtualfs-%d
`
	attachDeviceScript := fmt.Sprintf(attachDeviceTemplate, deviceNum, deviceNum, size, devicePath, deviceNum)
	_, err = t.RunCommand(attachDeviceScript)
	if err != nil {
		return 0, err
	}

	t.deviceMap[deviceNum] = devicePath
	t.currentDeviceNum++

	return deviceNum, nil
}

func (t *TestEnvironment) DetachLoopDevice(devicePath string) error {
	_, err := t.RunCommand(fmt.Sprintf("sudo losetup -d %s", devicePath))
	return err
}

func (t *TestEnvironment) UpdateAgentConfig(configFile string) error {
	_, err := t.RunCommand(
		fmt.Sprintf(
			"sudo cp %s/%s /var/vcap/bosh/agent.json",
			t.assetsDir(),
			configFile,
		),
	)
	return err
}

func (t *TestEnvironment) RestartAgent() error {
	err := t.StopAgent()
	if err != nil {
		return err
	}

	return t.StartAgent()
}

func (t *TestEnvironment) StopAgent() error {
	_, err := t.RunCommand("nohup sudo sv stop agent &")
	return err
}

func (t *TestEnvironment) StartAgent() error {
	_, err := t.RunCommand("nohup sudo sv start agent &")
	return err
}

type emptyReader struct{}

func (er emptyReader) Read(p []byte) (int, error) {
	time.Sleep(1 * time.Second)
	return 0, nil
}

func (t *TestEnvironment) StartAgentTunnel(mbusUser, mbusPass string, mbusPort int) (agentclient.AgentClient, error) {
	if t.sshTunnelProc != nil {
		return nil, fmt.Errorf("Already running")
	}

	sshCmd := boshsys.Command{
		Name: "vagrant",
		Args: []string{
			"ssh",
			"--",
			fmt.Sprintf("-L16868:127.0.0.1:%d", mbusPort),
		},
		Stdin: emptyReader{},
	}
	newTunnelProc, err := t.cmdRunner.RunComplexCommandAsync(sshCmd)
	if err != nil {
		return nil, err
	}
	t.sshTunnelProc = newTunnelProc

	httpClient := httpclient.NewHTTPClient(httpclient.DefaultClient, t.logger)
	mbusURL := fmt.Sprintf("https://%s:%s@localhost:16868", mbusUser, mbusPass)
	client := http.NewAgentClient(mbusURL, "fake-director-uuid", 1*time.Second, 10, httpClient, t.logger)

	for i := 1; i < 1000000; i++ {
		t.logger.Debug("test environment", "Trying to contact agent via ssh tunnel...")
		time.Sleep(1 * time.Second)
		_, err := client.Ping()
		if err == nil {
			break
		}
		t.logger.Debug("test environment", err.Error())
	}
	return client, nil
}

func (t *TestEnvironment) StopAgentTunnel() error {
	if t.sshTunnelProc == nil {
		return fmt.Errorf("Not running")
	}
	t.sshTunnelProc.Wait()
	t.sshTunnelProc = nil
	return nil
}

func (t *TestEnvironment) StartRegistry(settings boshsettings.Settings) error {
	settingsJSON, err := json.Marshal(settings)
	if err != nil {
		return err
	}

	_, err = t.RunCommand("sudo rm -f /var/vcap/bosh/settings.json")
	if err != nil {
		return err
	}

	t.RunCommand("sudo killall -9 fake-registry")

	_, err = t.RunCommand(
		fmt.Sprintf(
			`nohup %s/tmp/fake-registry -user user -password pass -host 127.0.0.1 -port 9090 -instance instance-id -settings %s &> /dev/null &`,
			t.agentDir(),
			strconv.Quote(string(settingsJSON)),
		),
	)
	return err
}

func (t *TestEnvironment) GetVMNetworks() (boshsettings.Networks, error) {
	stdout, _, _, err := t.cmdRunner.RunCommand("vagrant", "status")
	if err != nil {
		return boshsettings.Networks{}, err
	}

	if strings.Contains(stdout, "virtualbox") {
		return boshsettings.Networks{
			"eth0": {
				Type: "dynamic",
			},
			"eth1": {
				Type:    "manual",
				IP:      "192.168.50.4",
				Netmask: "255.255.255.0",
			},
		}, nil
	}

	if strings.Contains(stdout, "aws") {
		return boshsettings.Networks{
			"eth0": {
				Type: "dynamic",
			},
		}, nil
	}

	return boshsettings.Networks{}, nil
}

func (t *TestEnvironment) GetFileContents(filePath string) (string, error) {
	return t.RunCommand(
		fmt.Sprintf(
			`cat %s`,
			filePath,
		),
	)
}

func (t *TestEnvironment) RunCommand(command string) (string, error) {
	stdout, _, _, err := t.cmdRunner.RunCommand(
		"vagrant",
		"ssh",
		"-c",
		command,
	)

	return stdout, err
}

func (t *TestEnvironment) agentDir() string {
	return "/home/vagrant/go/src/github.com/cloudfoundry/bosh-agent"
}

func (t *TestEnvironment) assetsDir() string {
	return fmt.Sprintf("%s/integration/assets", t.agentDir())
}
<|endoftext|>