<commit_before>// Copyright 2015 SpectoLabs. 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.

// hoverfly is an HTTP/s proxy configurable via flags/environment variables/admin HTTP API
//
// this proxy can be dynamically configured through HTTP calls when it's running, to change modes,
// export and import requests.

package main

import (
	log "github.com/Sirupsen/logrus"
	hv "github.com/SpectoLabs/hoverfly"
	"github.com/SpectoLabs/hoverfly/authentication/backends"
	"github.com/SpectoLabs/hoverfly/backends/boltdb"

	"flag"
	"fmt"
	"os"
	"strings"
)

type arrayFlags []string

func (i *arrayFlags) String() string {
	return "my string representation"
}

func (i *arrayFlags) Set(value string) error {
	*i = append(*i, value)
	return nil
}

var importFlags arrayFlags
var destinationFlags arrayFlags

func main() {
	log.SetFormatter(&log.JSONFormatter{})

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "start Hoverfly in capture mode - transparently intercepts and saves requests/response")
	synthesize := flag.Bool("synthesize", false, "start Hoverfly in synthesize mode (middleware is required)")
	modify := flag.Bool("modify", false, "start Hoverfly in modify mode - applies middleware (required) to both outgoing and incomming HTTP traffic")

	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	// database location
	database := flag.String("db", "", "database location - supply it if you want to provide specific to database (will be created there if it doesn't exist)")

	// delete current database on startup
	wipeDb := flag.Bool("wipedb", false, "supply -wipedb flag to delete all records from given database on startup")

	// metrics
	metrics := flag.Bool("metrics", false, "supply -metrics flag to enable metrics logging to stdout")

	// development
	dev := flag.Bool("dev", false, "supply -dev flag to serve directly from ./static/dist instead from statik binary")

	// import flag
	flag.Var(&importFlags, "import", "import from file or from URL (i.e. '-import my_service.json' or '-import http://mypage.com/service_x.json'")

	// destination configuration
	destination := flag.String("destination", ".", "destination URI to catch")
	flag.Var(&destinationFlags, "dest", "specify which hosts to process (i.e. '-dest fooservice.org -dest barservice.org -dest catservice.org') - other hosts will be ignored will passthrough'")

	// adding new user
	addNew := flag.Bool("add", false, "add new user '-add -username hfadmin -password hfpass'")
	addUser := flag.String("username", "", "username for new user")
	addPassword := flag.String("password", "", "password for new user")
	isAdmin := flag.Bool("admin", true, "supply '-admin false' to make this non admin user (defaults to 'true') ")

	// TODO: this should be enabled by default when UI and documentation is ready
	authEnabled := flag.Bool("auth", false, "enable authentication, currently it is disabled by default")

	flag.Parse()

	// getting settings
	cfg := hv.InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.Verbose = *verbose

	if *dev {
		// making text pretty
		log.SetFormatter(&log.TextFormatter{})
	}

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.ProxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.AdminPort = *adminPort
	}

	// development settings
	cfg.Development = *dev

	// overriding database location
	if *database != "" {
		cfg.DatabaseName = *database
	}

	if *wipeDb {
		os.Remove(cfg.DatabaseName)
	}

	// overriding default middleware setting
	cfg.Middleware = *middleware

	// setting default mode
	mode := hv.VirtualizeMode

	if *capture {
		mode = hv.CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = hv.SynthesizeMode

		if cfg.Middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = hv.ModifyMode

		if cfg.Middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// setting mode
	cfg.SetMode(mode)

	// enabling authentication if flag or env variable is set to 'true'
	if cfg.AuthEnabled || *authEnabled {
		cfg.AuthEnabled = true
	}

	if len(destinationFlags) > 0 {
		cfg.Destination = strings.Join(destinationFlags[:], "|")

	} else {
		//  setting destination regexp
		cfg.Destination = *destination
	}

	// getting boltDB
	db := boltdb.GetDB(cfg.DatabaseName)
	defer db.Close()

	cache := boltdb.NewBoltDBCache(db, []byte("requestsBucket"))

	dbClient := hv.GetNewHoverfly(cfg, cache)

	ab := backends.NewBoltDBAuthBackend(db, []byte(backends.TokenBucketName), []byte(backends.UserBucketName))

	// assigning auth backend
	dbClient.AB = ab

	// metadata backend
	metaCache := boltdb.NewBoltDBCache(db, []byte("metadataBucket"))
	md := hv.NewMetadata(metaCache)

	dbClient.MD = md

	// if add new user supplied - adding it to database
	if *addNew {
		err := ab.AddUser([]byte(*addUser), []byte(*addPassword), *isAdmin)
		if err != nil {
			log.WithFields(log.Fields{
				"error":    err.Error(),
				"username": *addUser,
			}).Fatal("failed to add new user")
		} else {
			log.WithFields(log.Fields{
				"username": *addUser,
			}).Info("user added successfuly")
		}
		return
	}

	// importing stuff
	if len(importFlags) > 0 {
		for i, v := range importFlags {
			if v != "" {
				log.WithFields(log.Fields{
					"import": v,
				}).Debug("Importing given resource")
				err := dbClient.Import(v)
				if err != nil {
					log.WithFields(log.Fields{
						"error":  err.Error(),
						"import": v,
					}).Fatal("Failed to import given resource")
				} else {
					err = dbClient.MD.Set(fmt.Sprintf("import_%d", i+1), v)
				}

			}
		}
	}

	// start metrics registry flush
	if *metrics {
		dbClient.Counter.Init()
	}

	err := dbClient.StartProxy()
	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Fatal("failed to start proxy...")
	}

	// starting admin interface, this is blocking
	dbClient.StartAdminInterface()
}
<commit_msg>added functionality to generate mitm keys<commit_after>// Copyright 2015 SpectoLabs. 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.

// hoverfly is an HTTP/s proxy configurable via flags/environment variables/admin HTTP API
//
// this proxy can be dynamically configured through HTTP calls when it's running, to change modes,
// export and import requests.

package main

import (
	log "github.com/Sirupsen/logrus"
	hv "github.com/SpectoLabs/hoverfly"
	"github.com/SpectoLabs/hoverfly/authentication/backends"
	"github.com/SpectoLabs/hoverfly/backends/boltdb"
	hvc "github.com/SpectoLabs/hoverfly/certs"
	"github.com/rusenask/goproxy"

	"flag"
	"fmt"
	"os"
	"strings"
	"crypto/tls"
	"time"
	"encoding/pem"
)

type arrayFlags []string

func (i *arrayFlags) String() string {
	return "my string representation"
}

func (i *arrayFlags) Set(value string) error {
	*i = append(*i, value)
	return nil
}

var importFlags arrayFlags
var destinationFlags arrayFlags

func main() {
	log.SetFormatter(&log.JSONFormatter{})

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "start Hoverfly in capture mode - transparently intercepts and saves requests/response")
	synthesize := flag.Bool("synthesize", false, "start Hoverfly in synthesize mode (middleware is required)")
	modify := flag.Bool("modify", false, "start Hoverfly in modify mode - applies middleware (required) to both outgoing and incomming HTTP traffic")

	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	// database location
	database := flag.String("db", "", "database location - supply it if you want to provide specific to database (will be created there if it doesn't exist)")

	// delete current database on startup
	wipeDb := flag.Bool("wipedb", false, "supply -wipedb flag to delete all records from given database on startup")

	// metrics
	metrics := flag.Bool("metrics", false, "supply -metrics flag to enable metrics logging to stdout")

	// development
	dev := flag.Bool("dev", false, "supply -dev flag to serve directly from ./static/dist instead from statik binary")

	// import flag
	flag.Var(&importFlags, "import", "import from file or from URL (i.e. '-import my_service.json' or '-import http://mypage.com/service_x.json'")

	// destination configuration
	destination := flag.String("destination", ".", "destination URI to catch")
	flag.Var(&destinationFlags, "dest", "specify which hosts to process (i.e. '-dest fooservice.org -dest barservice.org -dest catservice.org') - other hosts will be ignored will passthrough'")

	// adding new user
	addNew := flag.Bool("add", false, "add new user '-add -username hfadmin -password hfpass'")
	addUser := flag.String("username", "", "username for new user")
	addPassword := flag.String("password", "", "password for new user")
	isAdmin := flag.Bool("admin", true, "supply '-admin false' to make this non admin user (defaults to 'true') ")

	// TODO: this should be enabled by default when UI and documentation is ready
	authEnabled := flag.Bool("auth", false, "enable authentication, currently it is disabled by default")

	generateCA := flag.Bool("generate-ca-cert", false, "generate CA certificate and private key for MITM")
	certName := flag.String("cert-name", "hoverfly.proxy", "cert name")
	certOrg := flag.String("cert-org", "Hoverfly Authority", "organisation name for new cert")

	cert := flag.String("cert", "", "CA certificate used to sign MITM certificates")
	key := flag.String("key", "", "private key of the CA used to sign MITM certificates")

	flag.Parse()

	// getting settings
	cfg := hv.InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.Verbose = *verbose

	if *dev {
		// making text pretty
		log.SetFormatter(&log.TextFormatter{})
	}

	if *generateCA {
		validity := 365 * 24 * time.Hour
		x509c, priv, err := hvc.NewCertificatePair(*certName, *certOrg, validity)
		if err != nil {
			log.Fatalf("Failed to generate certificate and key pair, got error: %s", err.Error())
		}

		certOut, err := os.Create("cert.pem")
		if err != nil {
			log.Fatalf("failed to open cert.pem for writing: %s", err.Error())
		}
		pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: x509c.Raw})
		certOut.Close()
		log.Print("cert.pem created\n")

		keyOut, err := os.OpenFile("key.pem", os.O_WRONLY | os.O_CREATE | os.O_TRUNC, 0600)
		if err != nil {
			log.Fatalf("failed to open key.pem for writing: %s", err.Error())
		}
		pem.Encode(keyOut, hvc.PemBlockForKey(priv))
		keyOut.Close()
		log.Print("key.pem created.\n")

		tlsc, err := hvc.GetTlsCertificate(x509c, priv, "hoverfly.proxy", validity); if err != nil {
			log.Fatalf("failed to get tls certificate: %s", err.Error())
		}
		goproxy.GoproxyCa = *tlsc

	} else if *cert != "" && *key != "" {
		tlsc, err := tls.LoadX509KeyPair(*cert, *key)
		if err != nil {
			log.Fatalf("Failed to load certifiate and key pair, got error: %s", err.Error())
		}

		goproxy.GoproxyCa = tlsc

	}

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.ProxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.AdminPort = *adminPort
	}

	// development settings
	cfg.Development = *dev

	// overriding database location
	if *database != "" {
		cfg.DatabaseName = *database
	}

	if *wipeDb {
		os.Remove(cfg.DatabaseName)
	}

	// overriding default middleware setting
	cfg.Middleware = *middleware

	// setting default mode
	mode := hv.VirtualizeMode

	if *capture {
		mode = hv.CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = hv.SynthesizeMode

		if cfg.Middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = hv.ModifyMode

		if cfg.Middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// setting mode
	cfg.SetMode(mode)

	// enabling authentication if flag or env variable is set to 'true'
	if cfg.AuthEnabled || *authEnabled {
		cfg.AuthEnabled = true
	}

	if len(destinationFlags) > 0 {
		cfg.Destination = strings.Join(destinationFlags[:], "|")

	} else {
		//  setting destination regexp
		cfg.Destination = *destination
	}

	// getting boltDB
	db := boltdb.GetDB(cfg.DatabaseName)
	defer db.Close()

	cache := boltdb.NewBoltDBCache(db, []byte("requestsBucket"))

	dbClient := hv.GetNewHoverfly(cfg, cache)

	ab := backends.NewBoltDBAuthBackend(db, []byte(backends.TokenBucketName), []byte(backends.UserBucketName))

	// assigning auth backend
	dbClient.AB = ab

	// metadata backend
	metaCache := boltdb.NewBoltDBCache(db, []byte("metadataBucket"))
	md := hv.NewMetadata(metaCache)

	dbClient.MD = md

	// if add new user supplied - adding it to database
	if *addNew {
		err := ab.AddUser([]byte(*addUser), []byte(*addPassword), *isAdmin)
		if err != nil {
			log.WithFields(log.Fields{
				"error":    err.Error(),
				"username": *addUser,
			}).Fatal("failed to add new user")
		} else {
			log.WithFields(log.Fields{
				"username": *addUser,
			}).Info("user added successfuly")
		}
		return
	}

	// importing stuff
	if len(importFlags) > 0 {
		for i, v := range importFlags {
			if v != "" {
				log.WithFields(log.Fields{
					"import": v,
				}).Debug("Importing given resource")
				err := dbClient.Import(v)
				if err != nil {
					log.WithFields(log.Fields{
						"error":  err.Error(),
						"import": v,
					}).Fatal("Failed to import given resource")
				} else {
					err = dbClient.MD.Set(fmt.Sprintf("import_%d", i + 1), v)
				}

			}
		}
	}

	// start metrics registry flush
	if *metrics {
		dbClient.Counter.Init()
	}

	err := dbClient.StartProxy()
	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Fatal("failed to start proxy...")
	}

	// starting admin interface, this is blocking
	dbClient.StartAdminInterface()
}
<|endoftext|>


<commit_before>package install

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

func lineInFile(name string, lookFor string) bool {
	f, err := os.Open(name)
	if err != nil {
		return false
	}
	defer f.Close()
	r := bufio.NewReader(f)
	prefix := []byte{}
	for {
		line, isPrefix, err := r.ReadLine()
		if err == io.EOF {
			return false
		}
		if err != nil {
			return false
		}
		if isPrefix {
			prefix = append(prefix, line...)
			continue
		}
		line = append(prefix, line...)
		if string(line) == lookFor {
			return true
		}
		prefix = prefix[:0]
	}
}

func createFile(name string, content string) error {
	f, err := os.Create(name)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(fmt.Sprintf("%s\n", content))
	return err
}

func appendToFile(name string, content string) error {
	f, err := os.OpenFile(name, os.O_RDWR|os.O_APPEND, 0)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(fmt.Sprintf("\n%s\n", content))
	return err
}

func removeFromFile(name string, content string) error {
	backup := name + ".bck"
	err := copyFile(name, backup)
	if err != nil {
		return err
	}
	temp, err := removeContentToTempFile(name, content)
	if err != nil {
		return err
	}

	err = copyFile(temp, name)
	if err != nil {
		return err
	}

	return os.Remove(backup)
}

func removeContentToTempFile(name, content string) (string, error) {
	rf, err := os.Open(name)
	if err != nil {
		return "", err
	}
	defer rf.Close()
	wf, err := ioutil.TempFile("/tmp", "complete-")
	if err != nil {
		return "", err
	}
	defer wf.Close()

	r := bufio.NewReader(rf)
	prefix := []byte{}
	for {
		line, isPrefix, err := r.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			return "", err
		}
		if isPrefix {
			prefix = append(prefix, line...)
			continue
		}
		line = append(prefix, line...)
		str := string(line)
		if str == content {
			continue
		}
		wf.WriteString(str + "\n")
		prefix = prefix[:0]
	}
	return wf.Name(), nil
}

func copyFile(src string, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}
<commit_msg>install: create file directory before file is created<commit_after>package install

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

func lineInFile(name string, lookFor string) bool {
	f, err := os.Open(name)
	if err != nil {
		return false
	}
	defer f.Close()
	r := bufio.NewReader(f)
	prefix := []byte{}
	for {
		line, isPrefix, err := r.ReadLine()
		if err == io.EOF {
			return false
		}
		if err != nil {
			return false
		}
		if isPrefix {
			prefix = append(prefix, line...)
			continue
		}
		line = append(prefix, line...)
		if string(line) == lookFor {
			return true
		}
		prefix = prefix[:0]
	}
}

func createFile(name string, content string) error {
	// make sure file directory exists
	if err := os.MkdirAll(filepath.Dir(name), 0775); err != nil {
		return err
	}

	// create the file
	f, err := os.Create(name)
	if err != nil {
		return err
	}
	defer f.Close()

	// write file content
	_, err = f.WriteString(fmt.Sprintf("%s\n", content))
	return err
}

func appendToFile(name string, content string) error {
	f, err := os.OpenFile(name, os.O_RDWR|os.O_APPEND, 0)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(fmt.Sprintf("\n%s\n", content))
	return err
}

func removeFromFile(name string, content string) error {
	backup := name + ".bck"
	err := copyFile(name, backup)
	if err != nil {
		return err
	}
	temp, err := removeContentToTempFile(name, content)
	if err != nil {
		return err
	}

	err = copyFile(temp, name)
	if err != nil {
		return err
	}

	return os.Remove(backup)
}

func removeContentToTempFile(name, content string) (string, error) {
	rf, err := os.Open(name)
	if err != nil {
		return "", err
	}
	defer rf.Close()
	wf, err := ioutil.TempFile("/tmp", "complete-")
	if err != nil {
		return "", err
	}
	defer wf.Close()

	r := bufio.NewReader(rf)
	prefix := []byte{}
	for {
		line, isPrefix, err := r.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			return "", err
		}
		if isPrefix {
			prefix = append(prefix, line...)
			continue
		}
		line = append(prefix, line...)
		str := string(line)
		if str == content {
			continue
		}
		wf.WriteString(str + "\n")
		prefix = prefix[:0]
	}
	return wf.Name(), nil
}

func copyFile(src string, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"github.com/juju/utils"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/mongo"
	"github.com/juju/juju/state"
	"github.com/juju/juju/state/api"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/upgrades"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker"
)

type upgradingMachineAgent interface {
	ensureMongoServer(agent.Config) error
	setMachineStatus(*api.State, params.Status, string) error
	CurrentConfig() agent.Config
	ChangeConfig(AgentConfigMutator) error
}

var upgradesPerformUpgrade = upgrades.PerformUpgrade // Allow patching for tests

func NewUpgradeWorkerContext() *upgradeWorkerContext {
	return &upgradeWorkerContext{
		UpgradeComplete: make(chan struct{}),
	}
}

type upgradeWorkerContext struct {
	UpgradeComplete chan struct{}
	agent           upgradingMachineAgent
	apiState        *api.State
	jobs            []params.MachineJob
	agentConfig     agent.Config
	isStateServer   bool
	st              *state.State
}

// Initialise a upgradeWorkerContext once the agent configuration is available.
func (c *upgradeWorkerContext) InitializeFromConfig(config agent.Config) {
	if config.UpgradedToVersion() == version.Current.Number {
		logger.Infof("No need for upgrade: upgrade steps for %v have already been run.",
			version.Current.Number)
		close(c.UpgradeComplete)
	}
}

func (c *upgradeWorkerContext) Worker(
	agent upgradingMachineAgent,
	apiState *api.State,
	jobs []params.MachineJob,
) worker.Worker {
	c.agent = agent
	c.apiState = apiState
	c.jobs = jobs
	return worker.NewSimpleWorker(c.run)
}

type apiLostDuringUpgrade struct {
	err error
}

func (e *apiLostDuringUpgrade) Error() string {
	return fmt.Sprintf("API connection lost during upgrade: %v", e.err)
}

func isAPILostDuringUpgrade(err error) bool {
	_, ok := err.(*apiLostDuringUpgrade)
	return ok
}

func (c *upgradeWorkerContext) run(stop <-chan struct{}) error {
	select {
	case <-c.UpgradeComplete:
		// Our work is already done (we're probably being restarted
		// because the API connection has gone down), so do nothing.
		return nil
	default:
	}

	c.agentConfig = c.agent.CurrentConfig()

	// If the machine agent is a state server, flag that state
	// needs to be opened before running upgrade steps
	for _, job := range c.jobs {
		if job == params.JobManageEnviron {
			c.isStateServer = true
		}
	}
	// We need a *state.State for upgrades. We open it independently
	// of StateWorker, because we have no guarantees about when
	// and how often StateWorker might run.
	if c.isStateServer {
		var err error
		c.st, err = openStateForUpgrade(c.agent, c.agentConfig)
		if err != nil {
			return err
		}
		defer c.st.Close()
	}
	if err := c.runUpgrades(); err != nil {
		// Only return an error from the worker if the connection to
		// state went away (possible mongo master change). Returning
		// an error when the connection is lost will cause the agent
		// to restart.
		//
		// For other errors, the error is not returned because we want
		// the machine agent to stay running in an error state waiting
		// for user intervention.
		if isAPILostDuringUpgrade(err) {
			return err
		}
	} else {
		// Upgrade succeeded - signal that the upgrade is complete.
		close(c.UpgradeComplete)
	}
	return nil
}

// runUpgrades runs the upgrade operations for each job type and
// updates the updatedToVersion on success.
func (c *upgradeWorkerContext) runUpgrades() error {
	from := version.Current
	from.Number = c.agentConfig.UpgradedToVersion()
	if from == version.Current {
		logger.Infof("upgrade to %v already completed.", version.Current)
		return nil
	}

	a := c.agent
	tag := c.agentConfig.Tag().(names.MachineTag)

	isMaster, err := isMachineMaster(c.st, tag)
	if err != nil {
		return errors.Trace(err)
	}

	if c.isStateServer {
		// State servers need to wait for other state servers to be
		// ready to run the upgrade.
		if err := waitForOtherStateServers(c.st, isMaster); err != nil {
			logger.Errorf(`other state servers failed to come up for upgrade `+
				`to %s - aborting: %v`, version.Current, err)

			a.setMachineStatus(c.apiState, params.StatusError,
				fmt.Sprintf("upgrade to %v aborted while waiting for other "+
					"state servers: %v", version.Current, err))

			// TODO(menn0): if master, roll agent-version back to
			// previous version, once the upgrader allows downgrades
			// to previous version.

			return err
		}
	}

	err = a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
		var upgradeErr error
		a.setMachineStatus(c.apiState, params.StatusStarted,
			fmt.Sprintf("upgrading to %v", version.Current))

		context := upgrades.NewContext(agentConfig, c.apiState, c.st)
		for _, job := range c.jobs {
			target := upgradeTarget(job, isMaster)
			if target == "" {
				continue
			}
			logger.Infof("starting upgrade from %v to %v for %v %q",
				from, version.Current, target, tag)

			attempts := getUpgradeRetryStrategy()
			for attempt := attempts.Start(); attempt.Next(); {
				upgradeErr = upgradesPerformUpgrade(from.Number, target, context)
				if upgradeErr == nil {
					break
				}
				if connectionIsDead(c.apiState) {
					// API connection has gone away - abort!
					return &apiLostDuringUpgrade{upgradeErr}
				}
				retryText := "will retry"
				if !attempt.HasNext() {
					retryText = "giving up"
				}
				logger.Errorf("upgrade from %v to %v for %v %q failed (%s): %v",
					from, version.Current, target, tag, retryText, upgradeErr)
				a.setMachineStatus(c.apiState, params.StatusError,
					fmt.Sprintf("upgrade to %v failed (%s): %v",
						version.Current, retryText, upgradeErr))
			}
		}
		if upgradeErr != nil {
			return upgradeErr
		}
		agentConfig.SetUpgradedToVersion(version.Current.Number)
		return nil
	})
	if err != nil {
		logger.Errorf("upgrade to %v failed: %v", version.Current, err)
		return err
	}

	logger.Infof("upgrade to %v completed successfully.", version.Current)
	a.setMachineStatus(c.apiState, params.StatusStarted, "")
	return nil
}

var openStateForUpgrade = func(
	agent upgradingMachineAgent,
	agentConfig agent.Config,
) (*state.State, error) {
	if err := agent.ensureMongoServer(agentConfig); err != nil {
		return nil, err
	}
	var err error
	info, ok := agentConfig.MongoInfo()
	if !ok {
		return nil, fmt.Errorf("no state info available")
	}
	st, err := state.Open(info, mongo.DefaultDialOpts(), environs.NewStatePolicy())
	if err != nil {
		return nil, err
	}
	return st, nil
}

var isMachineMaster = func(st *state.State, tag names.MachineTag) (bool, error) {
	if st == nil {
		// If there is no state, we aren't a master.
		return false, nil
	}
	// Not calling the agent openState method as it does other checks
	// we really don't care about here.  All we need here is the machine
	// so we can determine if we are the master or not.
	machine, err := st.Machine(tag.Id())
	if err != nil {
		// This shouldn't happen, and if it does, the state worker will have
		// found out before us, and already errored, or is likely to error out
		// very shortly.  All we do here is return the error. The state worker
		// returns an error that will cause the agent to be terminated.
		return false, errors.Trace(err)
	}
	isMaster, err := mongo.IsMaster(st.MongoSession(), machine)
	if err != nil {
		return false, errors.Trace(err)
	}
	return isMaster, nil
}

var waitForOtherStateServers = func(st *state.State, isMaster bool) error {
	// TODO(mjs) - for now, assume that the other state servers are
	// ready. This function will be fleshed out once the UpgradeInfo
	// work is done.
	return nil
}

var getUpgradeRetryStrategy = func() utils.AttemptStrategy {
	return utils.AttemptStrategy{
		Delay: 2 * time.Minute,
		Min:   5,
	}
}

func upgradeTarget(job params.MachineJob, isMaster bool) upgrades.Target {
	switch job {
	case params.JobManageEnviron:
		if isMaster {
			return upgrades.DatabaseMaster
		}
		return upgrades.StateServer
	case params.JobHostUnits:
		return upgrades.HostMachine
	}
	return ""
}
<commit_msg>Use AreUpgradesDefined to decide agent if upgrade mode is needed<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"github.com/juju/utils"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/mongo"
	"github.com/juju/juju/state"
	"github.com/juju/juju/state/api"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/upgrades"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker"
)

type upgradingMachineAgent interface {
	ensureMongoServer(agent.Config) error
	setMachineStatus(*api.State, params.Status, string) error
	CurrentConfig() agent.Config
	ChangeConfig(AgentConfigMutator) error
}

var upgradesPerformUpgrade = upgrades.PerformUpgrade // Allow patching for tests

func NewUpgradeWorkerContext() *upgradeWorkerContext {
	return &upgradeWorkerContext{
		UpgradeComplete: make(chan struct{}),
	}
}

type upgradeWorkerContext struct {
	UpgradeComplete chan struct{}
	agent           upgradingMachineAgent
	apiState        *api.State
	jobs            []params.MachineJob
	agentConfig     agent.Config
	isStateServer   bool
	st              *state.State
}

// Initialise a upgradeWorkerContext once the agent configuration is available.
func (c *upgradeWorkerContext) InitializeFromConfig(config agent.Config) {
	if !upgrades.AreUpgradesDefined(config.UpgradedToVersion()) {
		logger.Infof("no upgrade steps required or upgrade steps for %v have already "+
			"been run.", version.Current.Number)
		close(c.UpgradeComplete)
	}
}

func (c *upgradeWorkerContext) Worker(
	agent upgradingMachineAgent,
	apiState *api.State,
	jobs []params.MachineJob,
) worker.Worker {
	c.agent = agent
	c.apiState = apiState
	c.jobs = jobs
	return worker.NewSimpleWorker(c.run)
}

type apiLostDuringUpgrade struct {
	err error
}

func (e *apiLostDuringUpgrade) Error() string {
	return fmt.Sprintf("API connection lost during upgrade: %v", e.err)
}

func isAPILostDuringUpgrade(err error) bool {
	_, ok := err.(*apiLostDuringUpgrade)
	return ok
}

func (c *upgradeWorkerContext) run(stop <-chan struct{}) error {
	select {
	case <-c.UpgradeComplete:
		// Our work is already done (we're probably being restarted
		// because the API connection has gone down), so do nothing.
		return nil
	default:
	}

	c.agentConfig = c.agent.CurrentConfig()

	// If the machine agent is a state server, flag that state
	// needs to be opened before running upgrade steps
	for _, job := range c.jobs {
		if job == params.JobManageEnviron {
			c.isStateServer = true
		}
	}
	// We need a *state.State for upgrades. We open it independently
	// of StateWorker, because we have no guarantees about when
	// and how often StateWorker might run.
	if c.isStateServer {
		var err error
		c.st, err = openStateForUpgrade(c.agent, c.agentConfig)
		if err != nil {
			return err
		}
		defer c.st.Close()
	}
	if err := c.runUpgrades(); err != nil {
		// Only return an error from the worker if the connection to
		// state went away (possible mongo master change). Returning
		// an error when the connection is lost will cause the agent
		// to restart.
		//
		// For other errors, the error is not returned because we want
		// the machine agent to stay running in an error state waiting
		// for user intervention.
		if isAPILostDuringUpgrade(err) {
			return err
		}
	} else {
		// Upgrade succeeded - signal that the upgrade is complete.
		close(c.UpgradeComplete)
	}
	return nil
}

// runUpgrades runs the upgrade operations for each job type and
// updates the updatedToVersion on success.
func (c *upgradeWorkerContext) runUpgrades() error {
	from := version.Current
	from.Number = c.agentConfig.UpgradedToVersion()
	if from == version.Current {
		logger.Infof("upgrade to %v already completed.", version.Current)
		return nil
	}

	a := c.agent
	tag := c.agentConfig.Tag().(names.MachineTag)

	isMaster, err := isMachineMaster(c.st, tag)
	if err != nil {
		return errors.Trace(err)
	}

	if c.isStateServer {
		// State servers need to wait for other state servers to be
		// ready to run the upgrade.
		if err := waitForOtherStateServers(c.st, isMaster); err != nil {
			logger.Errorf(`other state servers failed to come up for upgrade `+
				`to %s - aborting: %v`, version.Current, err)

			a.setMachineStatus(c.apiState, params.StatusError,
				fmt.Sprintf("upgrade to %v aborted while waiting for other "+
					"state servers: %v", version.Current, err))

			// TODO(menn0): if master, roll agent-version back to
			// previous version, once the upgrader allows downgrades
			// to previous version.

			return err
		}
	}

	err = a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
		var upgradeErr error
		a.setMachineStatus(c.apiState, params.StatusStarted,
			fmt.Sprintf("upgrading to %v", version.Current))

		context := upgrades.NewContext(agentConfig, c.apiState, c.st)
		for _, job := range c.jobs {
			target := upgradeTarget(job, isMaster)
			if target == "" {
				continue
			}
			logger.Infof("starting upgrade from %v to %v for %v %q",
				from, version.Current, target, tag)

			attempts := getUpgradeRetryStrategy()
			for attempt := attempts.Start(); attempt.Next(); {
				upgradeErr = upgradesPerformUpgrade(from.Number, target, context)
				if upgradeErr == nil {
					break
				}
				if connectionIsDead(c.apiState) {
					// API connection has gone away - abort!
					return &apiLostDuringUpgrade{upgradeErr}
				}
				retryText := "will retry"
				if !attempt.HasNext() {
					retryText = "giving up"
				}
				logger.Errorf("upgrade from %v to %v for %v %q failed (%s): %v",
					from, version.Current, target, tag, retryText, upgradeErr)
				a.setMachineStatus(c.apiState, params.StatusError,
					fmt.Sprintf("upgrade to %v failed (%s): %v",
						version.Current, retryText, upgradeErr))
			}
		}
		if upgradeErr != nil {
			return upgradeErr
		}
		agentConfig.SetUpgradedToVersion(version.Current.Number)
		return nil
	})
	if err != nil {
		logger.Errorf("upgrade to %v failed: %v", version.Current, err)
		return err
	}

	logger.Infof("upgrade to %v completed successfully.", version.Current)
	a.setMachineStatus(c.apiState, params.StatusStarted, "")
	return nil
}

var openStateForUpgrade = func(
	agent upgradingMachineAgent,
	agentConfig agent.Config,
) (*state.State, error) {
	if err := agent.ensureMongoServer(agentConfig); err != nil {
		return nil, err
	}
	var err error
	info, ok := agentConfig.MongoInfo()
	if !ok {
		return nil, fmt.Errorf("no state info available")
	}
	st, err := state.Open(info, mongo.DefaultDialOpts(), environs.NewStatePolicy())
	if err != nil {
		return nil, err
	}
	return st, nil
}

var isMachineMaster = func(st *state.State, tag names.MachineTag) (bool, error) {
	if st == nil {
		// If there is no state, we aren't a master.
		return false, nil
	}
	// Not calling the agent openState method as it does other checks
	// we really don't care about here.  All we need here is the machine
	// so we can determine if we are the master or not.
	machine, err := st.Machine(tag.Id())
	if err != nil {
		// This shouldn't happen, and if it does, the state worker will have
		// found out before us, and already errored, or is likely to error out
		// very shortly.  All we do here is return the error. The state worker
		// returns an error that will cause the agent to be terminated.
		return false, errors.Trace(err)
	}
	isMaster, err := mongo.IsMaster(st.MongoSession(), machine)
	if err != nil {
		return false, errors.Trace(err)
	}
	return isMaster, nil
}

var waitForOtherStateServers = func(st *state.State, isMaster bool) error {
	// TODO(mjs) - for now, assume that the other state servers are
	// ready. This function will be fleshed out once the UpgradeInfo
	// work is done.
	return nil
}

var getUpgradeRetryStrategy = func() utils.AttemptStrategy {
	return utils.AttemptStrategy{
		Delay: 2 * time.Minute,
		Min:   5,
	}
}

func upgradeTarget(job params.MachineJob, isMaster bool) upgrades.Target {
	switch job {
	case params.JobManageEnviron:
		if isMaster {
			return upgrades.DatabaseMaster
		}
		return upgrades.StateServer
	case params.JobHostUnits:
		return upgrades.HostMachine
	}
	return ""
}
<|endoftext|>


<commit_before>package ignition

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"text/template"

	"github.com/coreos/ignition/config/util"
	ignition "github.com/coreos/ignition/config/v2_2/types"

	"github.com/openshift/installer/pkg/asset"
	"github.com/openshift/installer/pkg/asset/ignition/content"
	"github.com/openshift/installer/pkg/asset/installconfig"
	"github.com/openshift/installer/pkg/asset/kubeconfig"
	"github.com/openshift/installer/pkg/asset/manifests"
	"github.com/openshift/installer/pkg/asset/tls"
	"github.com/openshift/installer/pkg/types"
)

const (
	rootDir = "/opt/tectonic"
)

// bootstrapTemplateData is the data to use to replace values in bootstrap
// template files.
type bootstrapTemplateData struct {
	ClusterDNSIP               string
	CloudProvider              string
	CloudProviderConfig        string
	DebugConfig                string
	KubeCoreRenderImage        string
	MachineConfigOperatorImage string
	EtcdCertSignerImage        string
	EtcdctlImage               string
	BootkubeImage              string
	HyperkubeImage             string
	EtcdCluster                string
}

// bootstrap is an asset that generates the ignition config for bootstrap nodes.
type bootstrap struct {
	installConfig             asset.Asset
	rootCA                    asset.Asset
	etcdCA                    asset.Asset
	ingressCertKey            asset.Asset
	kubeCA                    asset.Asset
	aggregatorCA              asset.Asset
	serviceServingCA          asset.Asset
	clusterAPIServerCertKey   asset.Asset
	etcdClientCertKey         asset.Asset
	apiServerCertKey          asset.Asset
	openshiftAPIServerCertKey asset.Asset
	apiServerProxyCertKey     asset.Asset
	adminCertKey              asset.Asset
	kubeletCertKey            asset.Asset
	mcsCertKey                asset.Asset
	serviceAccountKeyPair     asset.Asset
	kubeconfig                asset.Asset
	kubeconfigKubelet         asset.Asset
	manifests                 asset.Asset
	kubeCoreOperator          asset.Asset
}

var _ asset.Asset = (*bootstrap)(nil)

// newBootstrap creates a new bootstrap asset.
func newBootstrap(
	installConfigStock installconfig.Stock,
	tlsStock tls.Stock,
	kubeconfigStock kubeconfig.Stock,
	manifestStock manifests.Stock,
) *bootstrap {
	return &bootstrap{
		installConfig:             installConfigStock.InstallConfig(),
		rootCA:                    tlsStock.RootCA(),
		etcdCA:                    tlsStock.EtcdCA(),
		ingressCertKey:            tlsStock.IngressCertKey(),
		kubeCA:                    tlsStock.KubeCA(),
		aggregatorCA:              tlsStock.AggregatorCA(),
		serviceServingCA:          tlsStock.ServiceServingCA(),
		clusterAPIServerCertKey:   tlsStock.ClusterAPIServerCertKey(),
		etcdClientCertKey:         tlsStock.EtcdClientCertKey(),
		apiServerCertKey:          tlsStock.APIServerCertKey(),
		openshiftAPIServerCertKey: tlsStock.OpenshiftAPIServerCertKey(),
		apiServerProxyCertKey:     tlsStock.APIServerProxyCertKey(),
		adminCertKey:              tlsStock.AdminCertKey(),
		kubeletCertKey:            tlsStock.KubeletCertKey(),
		mcsCertKey:                tlsStock.MCSCertKey(),
		serviceAccountKeyPair:     tlsStock.ServiceAccountKeyPair(),
		kubeconfig:                kubeconfigStock.KubeconfigAdmin(),
		kubeconfigKubelet:         kubeconfigStock.KubeconfigKubelet(),
		manifests:                 manifestStock.Manifests(),
		kubeCoreOperator:          manifestStock.KubeCoreOperator(),
	}
}

// Dependencies returns the assets on which the bootstrap asset depends.
func (a *bootstrap) Dependencies() []asset.Asset {
	return []asset.Asset{
		a.installConfig,
		a.rootCA,
		a.etcdCA,
		a.ingressCertKey,
		a.kubeCA,
		a.aggregatorCA,
		a.serviceServingCA,
		a.clusterAPIServerCertKey,
		a.etcdClientCertKey,
		a.apiServerCertKey,
		a.openshiftAPIServerCertKey,
		a.apiServerProxyCertKey,
		a.adminCertKey,
		a.kubeletCertKey,
		a.mcsCertKey,
		a.serviceAccountKeyPair,
		a.kubeconfig,
		a.kubeconfigKubelet,
		a.manifests,
		a.kubeCoreOperator,
	}
}

// Generate generates the ignition config for the bootstrap asset.
func (a *bootstrap) Generate(dependencies map[asset.Asset]*asset.State) (*asset.State, error) {
	installConfig, err := installconfig.GetInstallConfig(a.installConfig, dependencies)
	if err != nil {
		return nil, err
	}

	templateData, err := a.getTemplateData(installConfig)
	if err != nil {
		return nil, err
	}

	config := ignition.Config{}

	a.addBootstrapFiles(&config, dependencies)
	a.addBootkubeFiles(&config, dependencies, templateData)
	a.addTectonicFiles(&config, dependencies, templateData)
	a.addTLSCertFiles(&config, dependencies)

	config.Systemd.Units = append(
		config.Systemd.Units,
		ignition.Unit{Name: "bootkube.service", Contents: content.BootkubeSystemdContents},
		ignition.Unit{Name: "tectonic.service", Contents: content.TectonicSystemdContents, Enabled: util.BoolToPtr(true)},
		ignition.Unit{Name: "kubelet.service", Contents: applyTemplateData(content.KubeletSystemdTemplate, templateData), Enabled: util.BoolToPtr(true)},
	)

	config.Passwd.Users = append(
		config.Passwd.Users,
		ignition.PasswdUser{Name: "core", SSHAuthorizedKeys: []ignition.SSHAuthorizedKey{ignition.SSHAuthorizedKey(installConfig.Admin.SSHKey)}},
	)

	data, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

	return &asset.State{
		Contents: []asset.Content{{
			Name: "bootstrap.ign",
			Data: data,
		}},
	}, nil
}

// Name returns the human-friendly name of the asset.
func (a *bootstrap) Name() string {
	return "Bootstrap Ignition Config"
}

// getTemplateData returns the data to use to execute bootstrap templates.
func (a *bootstrap) getTemplateData(installConfig *types.InstallConfig) (*bootstrapTemplateData, error) {
	clusterDNSIP, err := installconfig.ClusterDNSIP(installConfig)
	if err != nil {
		return nil, err
	}
	etcdEndpoints := make([]string, masterCount(installConfig))
	for i := range etcdEndpoints {
		etcdEndpoints[i] = fmt.Sprintf("https://%s-etcd-%d.%s:2379", installConfig.Name, i, installConfig.BaseDomain)
	}
	return &bootstrapTemplateData{
		ClusterDNSIP:               clusterDNSIP,
		CloudProvider:              getCloudProvider(installConfig),
		CloudProviderConfig:        getCloudProviderConfig(installConfig),
		DebugConfig:                "",
		KubeCoreRenderImage:        "quay.io/coreos/kube-core-renderer-dev:436b1b4395ae54d866edc88864c9b01797cebac1",
		MachineConfigOperatorImage: "docker.io/openshift/origin-machine-config-operator:v4.0.0",
		EtcdCertSignerImage:        "quay.io/coreos/kube-etcd-signer-server:678cc8e6841e2121ebfdb6e2db568fce290b67d6",
		EtcdctlImage:               "quay.io/coreos/etcd:v3.2.14",
		BootkubeImage:              "quay.io/coreos/bootkube:v0.10.0",
		HyperkubeImage:             "openshift/origin-node:latest",
		EtcdCluster:                strings.Join(etcdEndpoints, ","),
	}, nil
}

func (a *bootstrap) addBootstrapFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State) {
	config.Storage.Files = append(
		config.Storage.Files,
		fileFromBytes("/etc/kubernetes/kubeconfig", 0600, dependencies[a.kubeconfigKubelet].Contents[0].Data),
		fileFromBytes("/var/lib/kubeconfig", 0600, dependencies[a.kubeconfigKubelet].Contents[0].Data),
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0644, dependencies[a.kubeCoreOperator].Contents)...,
	)
}

func (a *bootstrap) addBootkubeFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State, templateData *bootstrapTemplateData) {
	config.Storage.Files = append(
		config.Storage.Files,
		fileFromString("/opt/tectonic/bootkube.sh", 0555, applyTemplateData(content.BootkubeShFileTemplate, templateData)),
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0600, dependencies[a.kubeconfig].Contents)...,
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0644, dependencies[a.manifests].Contents)...,
	)
}

func (a *bootstrap) addTectonicFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State, templateData *bootstrapTemplateData) {
	// TODO (staebler) - missing manifests from tectonic module
	config.Storage.Files = append(
		config.Storage.Files,
		fileFromString("/opt/tectonic/tectonic.sh", 0555, content.TectonicShFileContents),
	)
}

func (a *bootstrap) addTLSCertFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State) {
	for _, asset := range []asset.Asset{
		a.rootCA,
		a.kubeCA,
		a.aggregatorCA,
		a.serviceServingCA,
		a.etcdCA,
		a.clusterAPIServerCertKey,
		a.etcdClientCertKey,
		a.apiServerCertKey,
		a.openshiftAPIServerCertKey,
		a.apiServerProxyCertKey,
		a.adminCertKey,
		a.kubeletCertKey,
		a.mcsCertKey,
		a.serviceAccountKeyPair,
	} {
		config.Storage.Files = append(config.Storage.Files, filesFromContents(rootDir, 0600, dependencies[asset].Contents)...)
	}
}

func getCloudProvider(installConfig *types.InstallConfig) string {
	if installConfig.AWS != nil {
		return "aws"
	}
	return ""
}

func getCloudProviderConfig(installConfig *types.InstallConfig) string {
	return ""
}

func applyTemplateData(template *template.Template, templateData interface{}) string {
	buf := &bytes.Buffer{}
	if err := template.Execute(buf, templateData); err != nil {
		panic(err)
	}
	return buf.String()
}
<commit_msg>asset/ignition: include tectonic manifests<commit_after>package ignition

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"text/template"

	"github.com/coreos/ignition/config/util"
	ignition "github.com/coreos/ignition/config/v2_2/types"

	"github.com/openshift/installer/pkg/asset"
	"github.com/openshift/installer/pkg/asset/ignition/content"
	"github.com/openshift/installer/pkg/asset/installconfig"
	"github.com/openshift/installer/pkg/asset/kubeconfig"
	"github.com/openshift/installer/pkg/asset/manifests"
	"github.com/openshift/installer/pkg/asset/tls"
	"github.com/openshift/installer/pkg/types"
)

const (
	rootDir = "/opt/tectonic"
)

// bootstrapTemplateData is the data to use to replace values in bootstrap
// template files.
type bootstrapTemplateData struct {
	ClusterDNSIP               string
	CloudProvider              string
	CloudProviderConfig        string
	DebugConfig                string
	KubeCoreRenderImage        string
	MachineConfigOperatorImage string
	EtcdCertSignerImage        string
	EtcdctlImage               string
	BootkubeImage              string
	HyperkubeImage             string
	EtcdCluster                string
}

// bootstrap is an asset that generates the ignition config for bootstrap nodes.
type bootstrap struct {
	installConfig             asset.Asset
	rootCA                    asset.Asset
	etcdCA                    asset.Asset
	ingressCertKey            asset.Asset
	kubeCA                    asset.Asset
	aggregatorCA              asset.Asset
	serviceServingCA          asset.Asset
	clusterAPIServerCertKey   asset.Asset
	etcdClientCertKey         asset.Asset
	apiServerCertKey          asset.Asset
	openshiftAPIServerCertKey asset.Asset
	apiServerProxyCertKey     asset.Asset
	adminCertKey              asset.Asset
	kubeletCertKey            asset.Asset
	mcsCertKey                asset.Asset
	serviceAccountKeyPair     asset.Asset
	kubeconfig                asset.Asset
	kubeconfigKubelet         asset.Asset
	manifests                 asset.Asset
	tectonic                  asset.Asset
	kubeCoreOperator          asset.Asset
}

var _ asset.Asset = (*bootstrap)(nil)

// newBootstrap creates a new bootstrap asset.
func newBootstrap(
	installConfigStock installconfig.Stock,
	tlsStock tls.Stock,
	kubeconfigStock kubeconfig.Stock,
	manifestStock manifests.Stock,
) *bootstrap {
	return &bootstrap{
		installConfig:             installConfigStock.InstallConfig(),
		rootCA:                    tlsStock.RootCA(),
		etcdCA:                    tlsStock.EtcdCA(),
		ingressCertKey:            tlsStock.IngressCertKey(),
		kubeCA:                    tlsStock.KubeCA(),
		aggregatorCA:              tlsStock.AggregatorCA(),
		serviceServingCA:          tlsStock.ServiceServingCA(),
		clusterAPIServerCertKey:   tlsStock.ClusterAPIServerCertKey(),
		etcdClientCertKey:         tlsStock.EtcdClientCertKey(),
		apiServerCertKey:          tlsStock.APIServerCertKey(),
		openshiftAPIServerCertKey: tlsStock.OpenshiftAPIServerCertKey(),
		apiServerProxyCertKey:     tlsStock.APIServerProxyCertKey(),
		adminCertKey:              tlsStock.AdminCertKey(),
		kubeletCertKey:            tlsStock.KubeletCertKey(),
		mcsCertKey:                tlsStock.MCSCertKey(),
		serviceAccountKeyPair:     tlsStock.ServiceAccountKeyPair(),
		kubeconfig:                kubeconfigStock.KubeconfigAdmin(),
		kubeconfigKubelet:         kubeconfigStock.KubeconfigKubelet(),
		manifests:                 manifestStock.Manifests(),
		tectonic:                  manifestStock.Tectonic(),
		kubeCoreOperator:          manifestStock.KubeCoreOperator(),
	}
}

// Dependencies returns the assets on which the bootstrap asset depends.
func (a *bootstrap) Dependencies() []asset.Asset {
	return []asset.Asset{
		a.installConfig,
		a.rootCA,
		a.etcdCA,
		a.ingressCertKey,
		a.kubeCA,
		a.aggregatorCA,
		a.serviceServingCA,
		a.clusterAPIServerCertKey,
		a.etcdClientCertKey,
		a.apiServerCertKey,
		a.openshiftAPIServerCertKey,
		a.apiServerProxyCertKey,
		a.adminCertKey,
		a.kubeletCertKey,
		a.mcsCertKey,
		a.serviceAccountKeyPair,
		a.kubeconfig,
		a.kubeconfigKubelet,
		a.manifests,
		a.tectonic,
		a.kubeCoreOperator,
	}
}

// Generate generates the ignition config for the bootstrap asset.
func (a *bootstrap) Generate(dependencies map[asset.Asset]*asset.State) (*asset.State, error) {
	installConfig, err := installconfig.GetInstallConfig(a.installConfig, dependencies)
	if err != nil {
		return nil, err
	}

	templateData, err := a.getTemplateData(installConfig)
	if err != nil {
		return nil, err
	}

	config := ignition.Config{}

	a.addBootstrapFiles(&config, dependencies)
	a.addBootkubeFiles(&config, dependencies, templateData)
	a.addTectonicFiles(&config, dependencies, templateData)
	a.addTLSCertFiles(&config, dependencies)

	config.Systemd.Units = append(
		config.Systemd.Units,
		ignition.Unit{Name: "bootkube.service", Contents: content.BootkubeSystemdContents},
		ignition.Unit{Name: "tectonic.service", Contents: content.TectonicSystemdContents, Enabled: util.BoolToPtr(true)},
		ignition.Unit{Name: "kubelet.service", Contents: applyTemplateData(content.KubeletSystemdTemplate, templateData), Enabled: util.BoolToPtr(true)},
	)

	config.Passwd.Users = append(
		config.Passwd.Users,
		ignition.PasswdUser{Name: "core", SSHAuthorizedKeys: []ignition.SSHAuthorizedKey{ignition.SSHAuthorizedKey(installConfig.Admin.SSHKey)}},
	)

	data, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

	return &asset.State{
		Contents: []asset.Content{{
			Name: "bootstrap.ign",
			Data: data,
		}},
	}, nil
}

// Name returns the human-friendly name of the asset.
func (a *bootstrap) Name() string {
	return "Bootstrap Ignition Config"
}

// getTemplateData returns the data to use to execute bootstrap templates.
func (a *bootstrap) getTemplateData(installConfig *types.InstallConfig) (*bootstrapTemplateData, error) {
	clusterDNSIP, err := installconfig.ClusterDNSIP(installConfig)
	if err != nil {
		return nil, err
	}
	etcdEndpoints := make([]string, masterCount(installConfig))
	for i := range etcdEndpoints {
		etcdEndpoints[i] = fmt.Sprintf("https://%s-etcd-%d.%s:2379", installConfig.Name, i, installConfig.BaseDomain)
	}
	return &bootstrapTemplateData{
		ClusterDNSIP:               clusterDNSIP,
		CloudProvider:              getCloudProvider(installConfig),
		CloudProviderConfig:        getCloudProviderConfig(installConfig),
		DebugConfig:                "",
		KubeCoreRenderImage:        "quay.io/coreos/kube-core-renderer-dev:436b1b4395ae54d866edc88864c9b01797cebac1",
		MachineConfigOperatorImage: "docker.io/openshift/origin-machine-config-operator:v4.0.0",
		EtcdCertSignerImage:        "quay.io/coreos/kube-etcd-signer-server:678cc8e6841e2121ebfdb6e2db568fce290b67d6",
		EtcdctlImage:               "quay.io/coreos/etcd:v3.2.14",
		BootkubeImage:              "quay.io/coreos/bootkube:v0.10.0",
		HyperkubeImage:             "openshift/origin-node:latest",
		EtcdCluster:                strings.Join(etcdEndpoints, ","),
	}, nil
}

func (a *bootstrap) addBootstrapFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State) {
	config.Storage.Files = append(
		config.Storage.Files,
		fileFromBytes("/etc/kubernetes/kubeconfig", 0600, dependencies[a.kubeconfigKubelet].Contents[0].Data),
		fileFromBytes("/var/lib/kubeconfig", 0600, dependencies[a.kubeconfigKubelet].Contents[0].Data),
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0644, dependencies[a.kubeCoreOperator].Contents)...,
	)
}

func (a *bootstrap) addBootkubeFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State, templateData *bootstrapTemplateData) {
	config.Storage.Files = append(
		config.Storage.Files,
		fileFromString("/opt/tectonic/bootkube.sh", 0555, applyTemplateData(content.BootkubeShFileTemplate, templateData)),
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0600, dependencies[a.kubeconfig].Contents)...,
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0644, dependencies[a.manifests].Contents)...,
	)
}

func (a *bootstrap) addTectonicFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State, templateData *bootstrapTemplateData) {
	config.Storage.Files = append(
		config.Storage.Files,
		fileFromString("/opt/tectonic/tectonic.sh", 0555, content.TectonicShFileContents),
	)
	config.Storage.Files = append(
		config.Storage.Files,
		filesFromContents(rootDir, 0644, dependencies[a.tectonic].Contents)...,
	)
}

func (a *bootstrap) addTLSCertFiles(config *ignition.Config, dependencies map[asset.Asset]*asset.State) {
	for _, asset := range []asset.Asset{
		a.rootCA,
		a.kubeCA,
		a.aggregatorCA,
		a.serviceServingCA,
		a.etcdCA,
		a.clusterAPIServerCertKey,
		a.etcdClientCertKey,
		a.apiServerCertKey,
		a.openshiftAPIServerCertKey,
		a.apiServerProxyCertKey,
		a.adminCertKey,
		a.kubeletCertKey,
		a.mcsCertKey,
		a.serviceAccountKeyPair,
	} {
		config.Storage.Files = append(config.Storage.Files, filesFromContents(rootDir, 0600, dependencies[asset].Contents)...)
	}
}

func getCloudProvider(installConfig *types.InstallConfig) string {
	if installConfig.AWS != nil {
		return "aws"
	}
	return ""
}

func getCloudProviderConfig(installConfig *types.InstallConfig) string {
	return ""
}

func applyTemplateData(template *template.Template, templateData interface{}) string {
	buf := &bytes.Buffer{}
	if err := template.Execute(buf, templateData); err != nil {
		panic(err)
	}
	return buf.String()
}
<|endoftext|>


<commit_before>package fs

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

	"github.com/gopasspw/gopass/pkg/fsutil"
	"github.com/gopasspw/gopass/pkg/out"

	"github.com/blang/semver"
)

// Store is a fs based store
type Store struct {
	path string
}

// New creates a new store
func New(dir string) *Store {
	return &Store{
		path: filepath.Clean(dir),
	}
}

// Get retrieves the named content
func (s *Store) Get(ctx context.Context, name string) ([]byte, error) {
	path := filepath.Join(s.path, filepath.Clean(name))
	out.Debug(ctx, "fs.Get(%s) - %s", name, path)
	return ioutil.ReadFile(path)
}

// Set writes the given content
func (s *Store) Set(ctx context.Context, name string, value []byte) error {
	filename := filepath.Join(s.path, filepath.Clean(name))
	filedir := filepath.Dir(filename)
	if !fsutil.IsDir(filedir) {
		if err := os.MkdirAll(filedir, 0700); err != nil {
			return err
		}
	}
	out.Debug(ctx, "fs.Set(%s) - %s", name, filepath.Join(s.path, name))
	return ioutil.WriteFile(filepath.Join(s.path, name), value, 0644)
}

// Delete removes the named entity
func (s *Store) Delete(ctx context.Context, name string) error {
	path := filepath.Join(s.path, filepath.Clean(name))
	out.Debug(ctx, "fs.Delete(%s) - %s", name, path)
	return os.Remove(path)
}

// Exists checks if the named entity exists
func (s *Store) Exists(ctx context.Context, name string) bool {
	path := filepath.Join(s.path, filepath.Clean(name))
	out.Debug(ctx, "fs.Exists(%s) - %s", name, path)
	return fsutil.IsFile(path)
}

// List returns a list of all entities
func (s *Store) List(ctx context.Context, prefix string) ([]string, error) {
	prefix = strings.TrimPrefix(prefix, "/")
	out.Debug(ctx, "fs.List(%s)", prefix)
	files := make([]string, 0, 100)
	if err := filepath.Walk(s.path, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() && strings.HasPrefix(info.Name(), ".") && path != s.path {
			return filepath.SkipDir
		}
		if info.IsDir() {
			return nil
		}
		if path == s.path {
			return nil
		}
		name := strings.TrimPrefix(path, s.path+string(filepath.Separator))
		if !strings.HasPrefix(name, prefix) {
			return nil
		}
		files = append(files, name)
		return nil
	}); err != nil {
		return nil, err
	}
	sort.Strings(files)
	return files, nil
}

// IsDir returns true if the named entity is a directory
func (s *Store) IsDir(ctx context.Context, name string) bool {
	path := filepath.Join(s.path, filepath.Clean(name))
	isDir := fsutil.IsDir(path)
	out.Debug(ctx, "fs.Isdir(%s) - %s -> %t", name, path, isDir)
	return isDir
}

// Prune removes a named directory
func (s *Store) Prune(ctx context.Context, prefix string) error {
	path := filepath.Join(s.path, filepath.Clean(prefix))
	out.Debug(ctx, "fs.Prune(%s) - %s", prefix, path)
	return os.RemoveAll(path)
}

// Name returns the name of this backend
func (s *Store) Name() string {
	return "fs"
}

// Version returns the version of this backend
func (s *Store) Version(context.Context) semver.Version {
	return semver.Version{Minor: 1}
}

// String implements fmt.Stringer
func (s *Store) String() string {
	return fmt.Sprintf("fs(v0.1.0,path:%s)", s.path)
}

// Available will check if this backend is useable
func (s *Store) Available(ctx context.Context) error {
	return s.Set(ctx, ".test", []byte("test"))
}
<commit_msg>Try to resolve symlinks (#886)<commit_after>package fs

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

	"github.com/gopasspw/gopass/pkg/fsutil"
	"github.com/gopasspw/gopass/pkg/out"

	"github.com/blang/semver"
)

// Store is a fs based store
type Store struct {
	path string
}

// New creates a new store
func New(dir string) *Store {
	if d, err := filepath.EvalSymlinks(dir); err == nil {
		dir = d
	}
	return &Store{
		path: dir,
	}
}

// Get retrieves the named content
func (s *Store) Get(ctx context.Context, name string) ([]byte, error) {
	path := filepath.Join(s.path, filepath.Clean(name))
	out.Debug(ctx, "fs.Get(%s) - %s", name, path)
	return ioutil.ReadFile(path)
}

// Set writes the given content
func (s *Store) Set(ctx context.Context, name string, value []byte) error {
	filename := filepath.Join(s.path, filepath.Clean(name))
	filedir := filepath.Dir(filename)
	if !fsutil.IsDir(filedir) {
		if err := os.MkdirAll(filedir, 0700); err != nil {
			return err
		}
	}
	out.Debug(ctx, "fs.Set(%s) - %s", name, filepath.Join(s.path, name))
	return ioutil.WriteFile(filepath.Join(s.path, name), value, 0644)
}

// Delete removes the named entity
func (s *Store) Delete(ctx context.Context, name string) error {
	path := filepath.Join(s.path, filepath.Clean(name))
	out.Debug(ctx, "fs.Delete(%s) - %s", name, path)
	return os.Remove(path)
}

// Exists checks if the named entity exists
func (s *Store) Exists(ctx context.Context, name string) bool {
	path := filepath.Join(s.path, filepath.Clean(name))
	out.Debug(ctx, "fs.Exists(%s) - %s", name, path)
	return fsutil.IsFile(path)
}

// List returns a list of all entities
func (s *Store) List(ctx context.Context, prefix string) ([]string, error) {
	prefix = strings.TrimPrefix(prefix, "/")
	out.Debug(ctx, "fs.List(%s)", prefix)
	files := make([]string, 0, 100)
	if err := filepath.Walk(s.path, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() && strings.HasPrefix(info.Name(), ".") && path != s.path {
			return filepath.SkipDir
		}
		if info.IsDir() {
			return nil
		}
		if path == s.path {
			return nil
		}
		name := strings.TrimPrefix(path, s.path+string(filepath.Separator))
		if !strings.HasPrefix(name, prefix) {
			return nil
		}
		files = append(files, name)
		return nil
	}); err != nil {
		return nil, err
	}
	sort.Strings(files)
	return files, nil
}

// IsDir returns true if the named entity is a directory
func (s *Store) IsDir(ctx context.Context, name string) bool {
	path := filepath.Join(s.path, filepath.Clean(name))
	isDir := fsutil.IsDir(path)
	out.Debug(ctx, "fs.Isdir(%s) - %s -> %t", name, path, isDir)
	return isDir
}

// Prune removes a named directory
func (s *Store) Prune(ctx context.Context, prefix string) error {
	path := filepath.Join(s.path, filepath.Clean(prefix))
	out.Debug(ctx, "fs.Prune(%s) - %s", prefix, path)
	return os.RemoveAll(path)
}

// Name returns the name of this backend
func (s *Store) Name() string {
	return "fs"
}

// Version returns the version of this backend
func (s *Store) Version(context.Context) semver.Version {
	return semver.Version{Minor: 1}
}

// String implements fmt.Stringer
func (s *Store) String() string {
	return fmt.Sprintf("fs(v0.1.0,path:%s)", s.path)
}

// Available will check if this backend is useable
func (s *Store) Available(ctx context.Context) error {
	return s.Set(ctx, ".test", []byte("test"))
}
<|endoftext|>


<commit_before>package snickers_test

import (
	"github.com/flavioribeiro/snickers/db"
	"github.com/flavioribeiro/snickers/db/memory"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Database", func() {
	Context("in memory", func() {
		var (
			dbInstance *memory.Database
		)

		BeforeEach(func() {
			// TODO solve this by transforming database in interface/object (#6)
			dbInstance, _ = memory.GetDatabase()
			memory.ClearDatabase()
		})

		It("should be able to store a preset", func() {
			examplePreset := db.Preset{
				Name:         "examplePreset",
				Description:  "This is an example of preset",
				Container:    "mp4",
				Profile:      "high",
				ProfileLevel: "3.1",
				RateControl:  "VBR",
				Video: db.VideoPreset{
					Width:         "720",
					Height:        "1080",
					Codec:         "h264",
					Bitrate:       "10000",
					GopSize:       "90",
					GopMode:       "fixed",
					InterlaceMode: "progressive",
				},
				Audio: db.AudioPreset{
					Codec:   "aac",
					Bitrate: "64000",
				},
			}
			expected := map[string]db.Preset{"examplePreset": examplePreset}
			Expect(dbInstance.StorePreset(examplePreset)).To(Equal(expected))
		})

		It("should be able to retrieve a preset by its name", func() {
			preset1 := db.Preset{
				Name:        "presetOne",
				Description: "This is preset one",
			}

			preset2 := db.Preset{
				Name:        "presetTwo",
				Description: "This is preset two",
			}

			dbInstance.StorePreset(preset1)
			dbInstance.StorePreset(preset2)

			Expect(dbInstance.RetrievePreset("presetOne")).To(Equal(preset1))
		})

		It("should be able to list presets", func() {
			preset1 := db.Preset{
				Name:        "presetOne",
				Description: "This is preset one",
			}

			preset2 := db.Preset{
				Name:        "presetTwo",
				Description: "This is preset two",
			}

			dbInstance.StorePreset(preset1)
			dbInstance.StorePreset(preset2)

			expected := []db.Preset{preset1, preset2}

			Expect(dbInstance.GetPresets()).To(Equal(expected))
		})
	})
})
<commit_msg>db: fix tests<commit_after>package snickers_test

import (
	"github.com/flavioribeiro/snickers/db"
	"github.com/flavioribeiro/snickers/db/memory"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Database", func() {
	Context("in memory", func() {
		var (
			dbInstance *memory.Database
		)

		BeforeEach(func() {
			// TODO solve this by transforming database in interface/object (#6)
			dbInstance, _ = memory.GetDatabase()
			memory.ClearDatabase()
		})

		It("should be able to store a preset", func() {
			examplePreset := db.Preset{
				Name:         "examplePreset",
				Description:  "This is an example of preset",
				Container:    "mp4",
				Profile:      "high",
				ProfileLevel: "3.1",
				RateControl:  "VBR",
				Video: db.VideoPreset{
					Width:         "720",
					Height:        "1080",
					Codec:         "h264",
					Bitrate:       "10000",
					GopSize:       "90",
					GopMode:       "fixed",
					InterlaceMode: "progressive",
				},
				Audio: db.AudioPreset{
					Codec:   "aac",
					Bitrate: "64000",
				},
			}
			expected := map[string]db.Preset{"examplePreset": examplePreset}
			Expect(dbInstance.StorePreset(examplePreset)).To(Equal(expected))
		})

		It("should be able to retrieve a preset by its name", func() {
			preset1 := db.Preset{
				Name:        "presetOne",
				Description: "This is preset one",
			}

			preset2 := db.Preset{
				Name:        "presetTwo",
				Description: "This is preset two",
			}

			dbInstance.StorePreset(preset1)
			dbInstance.StorePreset(preset2)

			Expect(dbInstance.RetrievePreset("presetOne")).To(Equal(preset1))
		})

		It("should be able to list presets", func() {
			preset1 := db.Preset{
				Name:        "presetOne",
				Description: "This is preset one",
			}

			preset2 := db.Preset{
				Name:        "presetTwo",
				Description: "This is preset two",
			}

			dbInstance.StorePreset(preset1)
			Expect(len(dbInstance.GetPresets())).To(Equal(1))

			dbInstance.StorePreset(preset2)
			Expect(len(dbInstance.GetPresets())).To(Equal(2))

		})
	})
})
<|endoftext|>


<commit_before>/*
 * MinIO Cloud Storage, (C) 2020 MinIO, Inc.
 *
 * 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 bandwidth

import (
	"context"
	"sync"
	"time"

	"github.com/minio/minio/pkg/bandwidth"
	"github.com/minio/minio/pkg/pubsub"
)

// throttleBandwidth gets the throttle for bucket with the configured value
func (m *Monitor) throttleBandwidth(ctx context.Context, bucket string, bandwidthBytesPerSecond int64, clusterBandwidth int64) *throttle {
	m.lock.Lock()
	defer m.lock.Unlock()
	throttle, ok := m.bucketThrottle[bucket]
	if !ok {
		throttle = newThrottle(ctx, bandwidthBytesPerSecond, clusterBandwidth)
		m.bucketThrottle[bucket] = throttle
		return throttle
	}
	throttle.SetBandwidth(bandwidthBytesPerSecond, clusterBandwidth)
	return throttle
}

// SubscribeToBuckets subscribes to buckets. Empty array for monitoring all buckets.
func (m *Monitor) SubscribeToBuckets(subCh chan interface{}, doneCh <-chan struct{}, buckets []string) {
	m.pubsub.Subscribe(subCh, doneCh, func(f interface{}) bool {
		if buckets != nil || len(buckets) == 0 {
			return true
		}
		report, ok := f.(*bandwidth.Report)
		if !ok {
			return false
		}
		for _, b := range buckets {
			_, ok := report.BucketStats[b]
			if ok {
				return true
			}
		}
		return false
	})
}

// Monitor implements the monitoring for bandwidth measurements.
type Monitor struct {
	lock sync.Mutex // lock for all updates

	activeBuckets map[string]*bucketMeasurement // Buckets with objects in flight

	bucketMovingAvgTicker *time.Ticker // Ticker for calculating moving averages

	pubsub *pubsub.PubSub // PubSub for reporting bandwidths.

	bucketThrottle map[string]*throttle

	startProcessing sync.Once

	doneCh <-chan struct{}
}

// NewMonitor returns a monitor with defaults.
func NewMonitor(doneCh <-chan struct{}) *Monitor {
	m := &Monitor{
		activeBuckets:         make(map[string]*bucketMeasurement),
		bucketMovingAvgTicker: time.NewTicker(2 * time.Second),
		pubsub:                pubsub.New(),
		bucketThrottle:        make(map[string]*throttle),
		doneCh:                doneCh,
	}
	return m
}

// SelectionFunction for buckets
type SelectionFunction func(bucket string) bool

// SelectBuckets will select all the buckets passed in.
func SelectBuckets(buckets ...string) SelectionFunction {
	if len(buckets) == 0 {
		return func(bucket string) bool {
			return true
		}
	}
	return func(bucket string) bool {
		for _, b := range buckets {
			if b == "" || b == bucket {
				return true
			}
		}
		return false
	}
}

// GetReport gets the report for all bucket bandwidth details.
func (m *Monitor) GetReport(selectBucket SelectionFunction) *bandwidth.Report {
	m.lock.Lock()
	defer m.lock.Unlock()
	return m.getReport(selectBucket)
}

func (m *Monitor) getReport(selectBucket SelectionFunction) *bandwidth.Report {
	report := &bandwidth.Report{
		BucketStats: make(map[string]bandwidth.Details),
	}
	for bucket, bucketMeasurement := range m.activeBuckets {
		if !selectBucket(bucket) {
			continue
		}
		report.BucketStats[bucket] = bandwidth.Details{
			LimitInBytesPerSecond:            m.bucketThrottle[bucket].clusterBandwidth,
			CurrentBandwidthInBytesPerSecond: bucketMeasurement.getExpMovingAvgBytesPerSecond(),
		}
	}
	return report
}

func (m *Monitor) process(doneCh <-chan struct{}) {
	for {
		select {
		case <-m.bucketMovingAvgTicker.C:
			m.processAvg()
		case <-doneCh:
			return
		}
	}
}

func (m *Monitor) getBucketMeasurement(bucket string, initTime time.Time) *bucketMeasurement {
	bucketTracker, ok := m.activeBuckets[bucket]
	if !ok {
		bucketTracker = newBucketMeasurement(initTime)
		m.activeBuckets[bucket] = bucketTracker
	}
	return bucketTracker
}

func (m *Monitor) processAvg() {
	m.lock.Lock()
	defer m.lock.Unlock()
	for _, bucketMeasurement := range m.activeBuckets {
		bucketMeasurement.updateExponentialMovingAverage(time.Now())
	}
	m.pubsub.Publish(m.getReport(SelectBuckets()))
}

// track returns the measurement object for bucket and object
func (m *Monitor) track(bucket string, object string, timeNow time.Time) *bucketMeasurement {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.startProcessing.Do(func() {
		go m.process(m.doneCh)
	})
	b := m.getBucketMeasurement(bucket, timeNow)
	return b
}

// DeleteBucket deletes monitoring the 'bucket'
func (m *Monitor) DeleteBucket(bucket string) {
	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.activeBuckets, bucket)
	delete(m.bucketThrottle, bucket)
}
<commit_msg>fix possible crash in bucket bandwidth monitor (#11986)<commit_after>/*
 * MinIO Cloud Storage, (C) 2020 MinIO, Inc.
 *
 * 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 bandwidth

import (
	"context"
	"sync"
	"time"

	"github.com/minio/minio/pkg/bandwidth"
	"github.com/minio/minio/pkg/pubsub"
)

// throttleBandwidth gets the throttle for bucket with the configured value
func (m *Monitor) throttleBandwidth(ctx context.Context, bucket string, bandwidthBytesPerSecond int64, clusterBandwidth int64) *throttle {
	m.lock.Lock()
	defer m.lock.Unlock()
	throttle, ok := m.bucketThrottle[bucket]
	if !ok {
		throttle = newThrottle(ctx, bandwidthBytesPerSecond, clusterBandwidth)
		m.bucketThrottle[bucket] = throttle
		return throttle
	}
	throttle.SetBandwidth(bandwidthBytesPerSecond, clusterBandwidth)
	return throttle
}

// SubscribeToBuckets subscribes to buckets. Empty array for monitoring all buckets.
func (m *Monitor) SubscribeToBuckets(subCh chan interface{}, doneCh <-chan struct{}, buckets []string) {
	m.pubsub.Subscribe(subCh, doneCh, func(f interface{}) bool {
		if buckets != nil || len(buckets) == 0 {
			return true
		}
		report, ok := f.(*bandwidth.Report)
		if !ok {
			return false
		}
		for _, b := range buckets {
			_, ok := report.BucketStats[b]
			if ok {
				return true
			}
		}
		return false
	})
}

// Monitor implements the monitoring for bandwidth measurements.
type Monitor struct {
	lock sync.Mutex // lock for all updates

	activeBuckets map[string]*bucketMeasurement // Buckets with objects in flight

	bucketMovingAvgTicker *time.Ticker // Ticker for calculating moving averages

	pubsub *pubsub.PubSub // PubSub for reporting bandwidths.

	bucketThrottle map[string]*throttle

	startProcessing sync.Once

	doneCh <-chan struct{}
}

// NewMonitor returns a monitor with defaults.
func NewMonitor(doneCh <-chan struct{}) *Monitor {
	m := &Monitor{
		activeBuckets:         make(map[string]*bucketMeasurement),
		bucketMovingAvgTicker: time.NewTicker(2 * time.Second),
		pubsub:                pubsub.New(),
		bucketThrottle:        make(map[string]*throttle),
		doneCh:                doneCh,
	}
	return m
}

// SelectionFunction for buckets
type SelectionFunction func(bucket string) bool

// SelectBuckets will select all the buckets passed in.
func SelectBuckets(buckets ...string) SelectionFunction {
	if len(buckets) == 0 {
		return func(bucket string) bool {
			return true
		}
	}
	return func(bucket string) bool {
		for _, b := range buckets {
			if b == "" || b == bucket {
				return true
			}
		}
		return false
	}
}

// GetReport gets the report for all bucket bandwidth details.
func (m *Monitor) GetReport(selectBucket SelectionFunction) *bandwidth.Report {
	m.lock.Lock()
	defer m.lock.Unlock()
	return m.getReport(selectBucket)
}

func (m *Monitor) getReport(selectBucket SelectionFunction) *bandwidth.Report {
	report := &bandwidth.Report{
		BucketStats: make(map[string]bandwidth.Details),
	}
	for bucket, bucketMeasurement := range m.activeBuckets {
		if !selectBucket(bucket) {
			continue
		}
		bucketThrottle, ok := m.bucketThrottle[bucket]
		if !ok {
			continue
		}
		report.BucketStats[bucket] = bandwidth.Details{
			LimitInBytesPerSecond:            bucketThrottle.clusterBandwidth,
			CurrentBandwidthInBytesPerSecond: bucketMeasurement.getExpMovingAvgBytesPerSecond(),
		}
	}
	return report
}

func (m *Monitor) process(doneCh <-chan struct{}) {
	for {
		select {
		case <-m.bucketMovingAvgTicker.C:
			m.processAvg()
		case <-doneCh:
			return
		}
	}
}

func (m *Monitor) getBucketMeasurement(bucket string, initTime time.Time) *bucketMeasurement {
	bucketTracker, ok := m.activeBuckets[bucket]
	if !ok {
		bucketTracker = newBucketMeasurement(initTime)
		m.activeBuckets[bucket] = bucketTracker
	}
	return bucketTracker
}

func (m *Monitor) processAvg() {
	m.lock.Lock()
	defer m.lock.Unlock()
	for _, bucketMeasurement := range m.activeBuckets {
		bucketMeasurement.updateExponentialMovingAverage(time.Now())
	}
	m.pubsub.Publish(m.getReport(SelectBuckets()))
}

// track returns the measurement object for bucket and object
func (m *Monitor) track(bucket string, object string, timeNow time.Time) *bucketMeasurement {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.startProcessing.Do(func() {
		go m.process(m.doneCh)
	})
	b := m.getBucketMeasurement(bucket, timeNow)
	return b
}

// DeleteBucket deletes monitoring the 'bucket'
func (m *Monitor) DeleteBucket(bucket string) {
	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.activeBuckets, bucket)
	delete(m.bucketThrottle, bucket)
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus 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 main

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

	kingpin "gopkg.in/alecthomas/kingpin.v2"
	yaml "gopkg.in/yaml.v2"

	"github.com/prometheus/common/model"
	"github.com/prometheus/common/version"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/pkg/rulefmt"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/util/promlint"
)

func main() {
	app := kingpin.New(filepath.Base(os.Args[0]), "Tooling for the Prometheus monitoring system.")
	app.Version(version.Print("promtool"))
	app.HelpFlag.Short('h')

	checkCmd := app.Command("check", "Check the resources for validity.")

	checkConfigCmd := checkCmd.Command("config", "Check if the config files are valid or not.")
	configFiles := checkConfigCmd.Arg(
		"config-files",
		"The config files to check.",
	).Required().ExistingFiles()

	checkRulesCmd := checkCmd.Command("rules", "Check if the rule files are valid or not.")
	ruleFiles := checkRulesCmd.Arg(
		"rule-files",
		"The rule files to check.",
	).Required().ExistingFiles()

	checkMetricsCmd := checkCmd.Command("metrics", checkMetricsUsage)

	updateCmd := app.Command("update", "Update the resources to newer formats.")
	updateRulesCmd := updateCmd.Command("rules", "Update rules from the 1.x to 2.x format.")
	ruleFilesUp := updateRulesCmd.Arg("rule-files", "The rule files to update.").Required().ExistingFiles()

	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
	case checkConfigCmd.FullCommand():
		os.Exit(CheckConfig(*configFiles...))

	case checkRulesCmd.FullCommand():
		os.Exit(CheckRules(*ruleFiles...))

	case checkMetricsCmd.FullCommand():
		os.Exit(CheckMetrics())

	case updateRulesCmd.FullCommand():
		os.Exit(UpdateRules(*ruleFilesUp...))

	}

}

// CheckConfig validates configuration files.
func CheckConfig(files ...string) int {
	failed := false

	for _, f := range files {
		ruleFiles, err := checkConfig(f)
		if err != nil {
			fmt.Fprintln(os.Stderr, "  FAILED:", err)
			failed = true
		} else {
			fmt.Printf("  SUCCESS: %d rule files found\n", len(ruleFiles))
		}
		fmt.Println()

		for _, rf := range ruleFiles {
			if n, err := checkRules(rf); err != nil {
				fmt.Fprintln(os.Stderr, "  FAILED:", err)
				failed = true
			} else {
				fmt.Printf("  SUCCESS: %d rules found\n", n)
			}
			fmt.Println()
		}
	}
	if failed {
		return 1
	}
	return 0
}

func checkFileExists(fn string) error {
	// Nothing set, nothing to error on.
	if fn == "" {
		return nil
	}
	_, err := os.Stat(fn)
	return err
}

func checkConfig(filename string) ([]string, error) {
	fmt.Println("Checking", filename)

	cfg, err := config.LoadFile(filename)
	if err != nil {
		return nil, err
	}

	var ruleFiles []string
	for _, rf := range cfg.RuleFiles {
		rfs, err := filepath.Glob(rf)
		if err != nil {
			return nil, err
		}
		// If an explicit file was given, error if it is not accessible.
		if !strings.Contains(rf, "*") {
			if len(rfs) == 0 {
				return nil, fmt.Errorf("%q does not point to an existing file", rf)
			}
			if err := checkFileExists(rfs[0]); err != nil {
				return nil, fmt.Errorf("error checking rule file %q: %s", rfs[0], err)
			}
		}
		ruleFiles = append(ruleFiles, rfs...)
	}

	for _, scfg := range cfg.ScrapeConfigs {
		if err := checkFileExists(scfg.HTTPClientConfig.BearerTokenFile); err != nil {
			return nil, fmt.Errorf("error checking bearer token file %q: %s", scfg.HTTPClientConfig.BearerTokenFile, err)
		}

		if err := checkTLSConfig(scfg.HTTPClientConfig.TLSConfig); err != nil {
			return nil, err
		}

		for _, kd := range scfg.ServiceDiscoveryConfig.KubernetesSDConfigs {
			if err := checkTLSConfig(kd.TLSConfig); err != nil {
				return nil, err
			}
		}
	}

	return ruleFiles, nil
}

func checkTLSConfig(tlsConfig config.TLSConfig) error {
	if err := checkFileExists(tlsConfig.CertFile); err != nil {
		return fmt.Errorf("error checking client cert file %q: %s", tlsConfig.CertFile, err)
	}
	if err := checkFileExists(tlsConfig.KeyFile); err != nil {
		return fmt.Errorf("error checking client key file %q: %s", tlsConfig.KeyFile, err)
	}

	if len(tlsConfig.CertFile) > 0 && len(tlsConfig.KeyFile) == 0 {
		return fmt.Errorf("client cert file %q specified without client key file", tlsConfig.CertFile)
	}
	if len(tlsConfig.KeyFile) > 0 && len(tlsConfig.CertFile) == 0 {
		return fmt.Errorf("client key file %q specified without client cert file", tlsConfig.KeyFile)
	}

	return nil
}

// CheckRules validates rule files.
func CheckRules(files ...string) int {
	failed := false

	for _, f := range files {
		if n, errs := checkRules(f); errs != nil {
			fmt.Fprintln(os.Stderr, "  FAILED:")
			for _, e := range errs {
				fmt.Fprintln(os.Stderr, e.Error())
			}
			failed = true
		} else {
			fmt.Printf("  SUCCESS: %d rules found\n", n)
		}
		fmt.Println()
	}
	if failed {
		return 1
	}
	return 0
}

func checkRules(filename string) (int, []error) {
	fmt.Println("Checking", filename)

	rgs, errs := rulefmt.ParseFile(filename)
	if errs != nil {
		return 0, errs
	}

	numRules := 0
	for _, rg := range rgs.Groups {
		numRules += len(rg.Rules)
	}

	return numRules, nil
}

// UpdateRules updates the rule files.
func UpdateRules(files ...string) int {
	failed := false

	for _, f := range files {
		if err := updateRules(f); err != nil {
			fmt.Fprintln(os.Stderr, "  FAILED:", err)
			failed = true
		}
	}

	if failed {
		return 1
	}
	return 0
}

func updateRules(filename string) error {
	fmt.Println("Updating", filename)

	content, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}

	rules, err := promql.ParseStmts(string(content))
	if err != nil {
		return err
	}

	yamlRG := &rulefmt.RuleGroups{
		Groups: []rulefmt.RuleGroup{{
			Name: filename,
		}},
	}

	yamlRules := make([]rulefmt.Rule, 0, len(rules))

	for _, rule := range rules {
		switch r := rule.(type) {
		case *promql.AlertStmt:
			yamlRules = append(yamlRules, rulefmt.Rule{
				Alert:       r.Name,
				Expr:        r.Expr.String(),
				For:         model.Duration(r.Duration),
				Labels:      r.Labels.Map(),
				Annotations: r.Annotations.Map(),
			})
		case *promql.RecordStmt:
			yamlRules = append(yamlRules, rulefmt.Rule{
				Record: r.Name,
				Expr:   r.Expr.String(),
				Labels: r.Labels.Map(),
			})
		default:
			panic("unknown statement type")
		}
	}

	yamlRG.Groups[0].Rules = yamlRules
	y, err := yaml.Marshal(yamlRG)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filename+".yml", y, 0666)
}

var checkMetricsUsage = strings.TrimSpace(`
Pass Prometheus metrics over stdin to lint them for consistency and correctness.

examples:

$ cat metrics.prom | promtool check metrics

$ curl -s http://localhost:9090/metrics | promtool check metrics
`)

// CheckMetrics performs a linting pass on input metrics.
func CheckMetrics() int {
	l := promlint.New(os.Stdin)
	problems, err := l.Lint()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error while linting:", err)
		return 1
	}

	for _, p := range problems {
		fmt.Fprintln(os.Stderr, p.Metric, p.Text)
	}

	if len(problems) > 0 {
		return 3
	}

	return 0
}
<commit_msg>check if file_sd files exist in checkConfig<commit_after>// Copyright 2015 The Prometheus 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 main

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

	kingpin "gopkg.in/alecthomas/kingpin.v2"
	yaml "gopkg.in/yaml.v2"

	"github.com/prometheus/common/model"
	"github.com/prometheus/common/version"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/pkg/rulefmt"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/util/promlint"
)

func main() {
	app := kingpin.New(filepath.Base(os.Args[0]), "Tooling for the Prometheus monitoring system.")
	app.Version(version.Print("promtool"))
	app.HelpFlag.Short('h')

	checkCmd := app.Command("check", "Check the resources for validity.")

	checkConfigCmd := checkCmd.Command("config", "Check if the config files are valid or not.")
	configFiles := checkConfigCmd.Arg(
		"config-files",
		"The config files to check.",
	).Required().ExistingFiles()

	checkRulesCmd := checkCmd.Command("rules", "Check if the rule files are valid or not.")
	ruleFiles := checkRulesCmd.Arg(
		"rule-files",
		"The rule files to check.",
	).Required().ExistingFiles()

	checkMetricsCmd := checkCmd.Command("metrics", checkMetricsUsage)

	updateCmd := app.Command("update", "Update the resources to newer formats.")
	updateRulesCmd := updateCmd.Command("rules", "Update rules from the 1.x to 2.x format.")
	ruleFilesUp := updateRulesCmd.Arg("rule-files", "The rule files to update.").Required().ExistingFiles()

	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
	case checkConfigCmd.FullCommand():
		os.Exit(CheckConfig(*configFiles...))

	case checkRulesCmd.FullCommand():
		os.Exit(CheckRules(*ruleFiles...))

	case checkMetricsCmd.FullCommand():
		os.Exit(CheckMetrics())

	case updateRulesCmd.FullCommand():
		os.Exit(UpdateRules(*ruleFilesUp...))

	}

}

// CheckConfig validates configuration files.
func CheckConfig(files ...string) int {
	failed := false

	for _, f := range files {
		ruleFiles, err := checkConfig(f)
		if err != nil {
			fmt.Fprintln(os.Stderr, "  FAILED:", err)
			failed = true
		} else {
			fmt.Printf("  SUCCESS: %d rule files found\n", len(ruleFiles))
		}
		fmt.Println()

		for _, rf := range ruleFiles {
			if n, err := checkRules(rf); err != nil {
				fmt.Fprintln(os.Stderr, "  FAILED:", err)
				failed = true
			} else {
				fmt.Printf("  SUCCESS: %d rules found\n", n)
			}
			fmt.Println()
		}
	}
	if failed {
		return 1
	}
	return 0
}

func checkFileExists(fn string) error {
	// Nothing set, nothing to error on.
	if fn == "" {
		return nil
	}
	_, err := os.Stat(fn)
	return err
}

func checkConfig(filename string) ([]string, error) {
	fmt.Println("Checking", filename)

	cfg, err := config.LoadFile(filename)
	if err != nil {
		return nil, err
	}

	var ruleFiles []string
	for _, rf := range cfg.RuleFiles {
		rfs, err := filepath.Glob(rf)
		if err != nil {
			return nil, err
		}
		// If an explicit file was given, error if it is not accessible.
		if !strings.Contains(rf, "*") {
			if len(rfs) == 0 {
				return nil, fmt.Errorf("%q does not point to an existing file", rf)
			}
			if err := checkFileExists(rfs[0]); err != nil {
				return nil, fmt.Errorf("error checking rule file %q: %s", rfs[0], err)
			}
		}
		ruleFiles = append(ruleFiles, rfs...)
	}

	for _, scfg := range cfg.ScrapeConfigs {
		if err := checkFileExists(scfg.HTTPClientConfig.BearerTokenFile); err != nil {
			return nil, fmt.Errorf("error checking bearer token file %q: %s", scfg.HTTPClientConfig.BearerTokenFile, err)
		}

		if err := checkTLSConfig(scfg.HTTPClientConfig.TLSConfig); err != nil {
			return nil, err
		}

		for _, kd := range scfg.ServiceDiscoveryConfig.KubernetesSDConfigs {
			if err := checkTLSConfig(kd.TLSConfig); err != nil {
				return nil, err
			}
		}

		for _, filesd := range scfg.ServiceDiscoveryConfig.FileSDConfigs {
			for _, file := range filesd.Files {
				files, err := filepath.Glob(file)
				if err != nil {
					return nil, err
				}
				if len(files) != 0 {
					// There was at least one match for the glob and we can assume checkFileExists
					// for all matches would pass, we can continue the loop.
					continue
				}
				fmt.Printf("  WARNING: file %q for file_sd in scrape job %q does not exist\n", file, scfg.JobName)
			}
		}
	}

	return ruleFiles, nil
}

func checkTLSConfig(tlsConfig config.TLSConfig) error {
	if err := checkFileExists(tlsConfig.CertFile); err != nil {
		return fmt.Errorf("error checking client cert file %q: %s", tlsConfig.CertFile, err)
	}
	if err := checkFileExists(tlsConfig.KeyFile); err != nil {
		return fmt.Errorf("error checking client key file %q: %s", tlsConfig.KeyFile, err)
	}

	if len(tlsConfig.CertFile) > 0 && len(tlsConfig.KeyFile) == 0 {
		return fmt.Errorf("client cert file %q specified without client key file", tlsConfig.CertFile)
	}
	if len(tlsConfig.KeyFile) > 0 && len(tlsConfig.CertFile) == 0 {
		return fmt.Errorf("client key file %q specified without client cert file", tlsConfig.KeyFile)
	}

	return nil
}

// CheckRules validates rule files.
func CheckRules(files ...string) int {
	failed := false

	for _, f := range files {
		if n, errs := checkRules(f); errs != nil {
			fmt.Fprintln(os.Stderr, "  FAILED:")
			for _, e := range errs {
				fmt.Fprintln(os.Stderr, e.Error())
			}
			failed = true
		} else {
			fmt.Printf("  SUCCESS: %d rules found\n", n)
		}
		fmt.Println()
	}
	if failed {
		return 1
	}
	return 0
}

func checkRules(filename string) (int, []error) {
	fmt.Println("Checking", filename)

	rgs, errs := rulefmt.ParseFile(filename)
	if errs != nil {
		return 0, errs
	}

	numRules := 0
	for _, rg := range rgs.Groups {
		numRules += len(rg.Rules)
	}

	return numRules, nil
}

// UpdateRules updates the rule files.
func UpdateRules(files ...string) int {
	failed := false

	for _, f := range files {
		if err := updateRules(f); err != nil {
			fmt.Fprintln(os.Stderr, "  FAILED:", err)
			failed = true
		}
	}

	if failed {
		return 1
	}
	return 0
}

func updateRules(filename string) error {
	fmt.Println("Updating", filename)

	content, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}

	rules, err := promql.ParseStmts(string(content))
	if err != nil {
		return err
	}

	yamlRG := &rulefmt.RuleGroups{
		Groups: []rulefmt.RuleGroup{{
			Name: filename,
		}},
	}

	yamlRules := make([]rulefmt.Rule, 0, len(rules))

	for _, rule := range rules {
		switch r := rule.(type) {
		case *promql.AlertStmt:
			yamlRules = append(yamlRules, rulefmt.Rule{
				Alert:       r.Name,
				Expr:        r.Expr.String(),
				For:         model.Duration(r.Duration),
				Labels:      r.Labels.Map(),
				Annotations: r.Annotations.Map(),
			})
		case *promql.RecordStmt:
			yamlRules = append(yamlRules, rulefmt.Rule{
				Record: r.Name,
				Expr:   r.Expr.String(),
				Labels: r.Labels.Map(),
			})
		default:
			panic("unknown statement type")
		}
	}

	yamlRG.Groups[0].Rules = yamlRules
	y, err := yaml.Marshal(yamlRG)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filename+".yml", y, 0666)
}

var checkMetricsUsage = strings.TrimSpace(`
Pass Prometheus metrics over stdin to lint them for consistency and correctness.

examples:

$ cat metrics.prom | promtool check metrics

$ curl -s http://localhost:9090/metrics | promtool check metrics
`)

// CheckMetrics performs a linting pass on input metrics.
func CheckMetrics() int {
	l := promlint.New(os.Stdin)
	problems, err := l.Lint()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error while linting:", err)
		return 1
	}

	for _, p := range problems {
		fmt.Fprintln(os.Stderr, p.Metric, p.Text)
	}

	if len(problems) > 0 {
		return 3
	}

	return 0
}
<|endoftext|>