<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/jumoel/bitbucket-enforcer/gobucket"
	"github.com/jumoel/bitbucket-enforcer/log"
	dotenv "github.com/jumoel/bitbucket-enforcer/vendor/godotenv"
)

type branchManagement struct {
	PreventDelete []string
	PreventRebase []string
	AllowPushes   map[string]struct {
		Groups []string
		Users  []string
	}
}

type accessManagement struct {
	Users  map[string]string // usernames => permissions
	Groups map[string]string // groupnames => permissions
}

type repositorySettings struct {
	LandingPage      string
	Private          interface{}
	Forks            string
	IssueTracker     string
	DeployKeys       publicKeyList
	PostHooks        []string
	BranchManagement branchManagement
	AccessManagement accessManagement
}

type publicKey struct {
	Name string
	Key  string
}

type publicKeyList []publicKey

type bbServices []gobucket.Service
type matchType int

const (
	matchNone matchType = iota
	matchContent
	matchExact
)

const sleepTime = 1 * time.Second

var configDir = flag.String("configdir", "configs", "the folder containing repository configrations")
var verbose = flag.Bool("v", false, "print more output")
var bbAPI *gobucket.APIClient

func main() {
	log.SetPrefix("bitbucket-enforcer")

	flag.Parse()

	err := dotenv.Load()
	if err != nil {
		log.Notice(".env error", err)
	}

	bbUsername := os.Getenv("BITBUCKET_ENFORCER_USERNAME")
	bbKey := os.Getenv("BITBUCKET_ENFORCER_API_KEY")

	bbAPI = gobucket.New(bbUsername, bbKey)

	scanRepositories(bbUsername)
}

func scanRepositories(bbUsername string) {
	var enforcementMatcher = regexp.MustCompile(`-enforce(?:=([a-zA-Z0-9]+))?`)

	var lastEtag string
	var changed bool

	for _ = range time.Tick(sleepTime) {
		var err error
		if changed, lastEtag, err = bbAPI.RepositoriesChanged(bbUsername, lastEtag); err != nil {
			log.Error(fmt.Sprintf("Error determining if repository list has changed (%s)", err))
			continue
		}

		if !changed {
			if *verbose {
				log.Info("No repository changes, sleeping.")
			}
			continue
		}

		log.Info("Repository list changed")

		repos, err := bbAPI.GetRepositories(bbUsername)

		if err != nil {
			log.Error("Error getting repository list", err)
			continue
		}

		for _, repo := range repos {
			if strings.Contains(repo.Description, "-noenforce") {
				if *verbose {
					log.Info(fmt.Sprintf("Skipping <%s> because of '-noenforce'\n", repo.FullName))
				}
				continue
			}

			if strings.Contains(repo.Description, "-enforced") {
				if *verbose {
					log.Info(fmt.Sprintf("Skipping <%s> because of '-enforced'\n", repo.FullName))
				}
				continue
			}

			matches := enforcementMatcher.FindStringSubmatch(repo.Description)

			enforcementPolicy := "default"
			if len(matches) > 0 {
				enforcementPolicy = matches[1]
			}

			log.Info(fmt.Sprintf("Enforcing repo '%s' with policy '%s'", repo.FullName, enforcementPolicy))

			parts := strings.Split(repo.FullName, "/")
			err := enforcePolicy(parts[0], parts[1], enforcementPolicy)

			if err != nil {
				log.Warning(fmt.Sprintf("Could not enforce policy '%s' on repo '%s'. Will be processed again next cycle. (%s)", enforcementPolicy, repo.FullName, err))
			} else {
				newDescription := strings.TrimSpace(fmt.Sprintf("%s\n\n-enforced", repo.Description))

				if err := bbAPI.SetDescription(parts[0], parts[1], newDescription); err != nil {
					log.Warning("Could not set description on repo '%s'. Will be processed again next cycle.", repo.FullName)
				}
			}
		}
	}
}

func enforcePolicy(owner string, repo string, policyname string) error {
	policy, err := parseConfig(policyname)

	if err != nil {
		log.Error(fmt.Sprintf("Error parsing parsing policy '%s': ", policyname), err)
		return err
	}

	if policy.Private != nil {
		if err := bbAPI.SetPrivacy(owner, repo, policy.Private.(bool)); err != nil {
			log.Warning("Error setting privacy: ", err)
			return err
		}
	}

	if policy.Forks != "" {
		if err := bbAPI.SetForks(owner, repo, policy.Forks); err != nil {
			log.Warning("Error fork policy: ", err)
			return err
		}
	}

	if policy.LandingPage != "" {
		if err := bbAPI.SetLandingPage(owner, repo, policy.LandingPage); err != nil {
			log.Warning("Error setting landing page: ", err)
			return err
		}
	}

	if len(policy.DeployKeys) > 0 {
		if err := enforceDeployKeys(owner, repo, policy.DeployKeys); err != nil {
			log.Warning("Error setting deploy keys: ", err)
			return err
		}
	}

	if len(policy.PostHooks) > 0 {
		if err := enforcePOSTHooks(owner, repo, policy.PostHooks); err != nil {
			log.Warning("Error setting POST hooks: ", err)
			return err
		}
	}

	if policy.IssueTracker != "" {
		if err := bbAPI.SetIssueTracker(owner, repo, policy.IssueTracker); err != nil {
			log.Warning("Error setting issue tracker: ", err)
			return err
		}
	}

	if err := enforceBranchManagement(owner, repo, policy.BranchManagement); err != nil {
		log.Warning("Error setting branch policies: ", err)
		return err
	}

	if err := enforceAccessManagement(owner, repo, policy.AccessManagement); err != nil {
		log.Warning("Error setting access policies: ", err)
		return err
	}

	return nil
}

func enforceAccessManagement(owner string, repo string, policies accessManagement) error {
	for username, privilege := range policies.Users {
		if err := bbAPI.AddUserPrivilege(owner, repo, username, privilege); err != nil {
			return err
		}
	}

	for groupname, privilege := range policies.Groups {
		if err := bbAPI.AddGroupPrivilege(owner, repo, groupname, privilege); err != nil {
			return err
		}
	}

	return nil
}

func enforceBranchManagement(owner string, repo string, policies branchManagement) error {
	for _, branch := range policies.PreventDelete {
		if err := bbAPI.AddBranchRestriction(owner, repo, "delete", branch, nil, nil); err != nil {
			return err
		}
	}

	for _, branch := range policies.PreventRebase {
		if err := bbAPI.AddBranchRestriction(owner, repo, "force", branch, nil, nil); err != nil {
			return err
		}
	}

	for branch, permissions := range policies.AllowPushes {
		if err := bbAPI.AddBranchRestriction(owner, repo, "push", branch, permissions.Users, permissions.Groups); err != nil {
			return err
		}
	}

	return nil
}

func (hooks *bbServices) hasPOSTHook(URL string) bool {
	for _, hook := range *hooks {
		if hook.Service.Type == "POST" {
			for _, field := range hook.Service.Fields {
				if field.Name == "URL" && field.Value == URL {
					return true
				}
			}
		}
	}

	return false
}

func enforcePOSTHooks(owner string, repo string, hookURLs []string) error {
	hookList, err := bbAPI.GetServices(owner, repo)

	if err != nil {
		return err
	}

	var currentHooks bbServices = hookList

	for _, url := range hookURLs {
		if !currentHooks.hasPOSTHook(url) {
			if err := bbAPI.AddService(owner, repo, "POST", map[string]string{"URL": url}); err != nil {
				return err
			}
		}
	}

	return nil
}

func (keys *publicKeyList) hasKey(needle gobucket.DeployKey) (matchType, int) {
	for index, key := range *keys {
		if key.Key == needle.Key && key.Name == needle.Label {
			return matchExact, index
		} else if key.Key == needle.Key {
			return matchContent, index
		}
	}

	return matchNone, -1
}

/*
This method ensures the presence of all required keys.
- It removes keys with matching content but mismatching names. Afterwards they
  are added again, this time with the correct name.
- It adds keys that are not present.
- It doesn't remove keys that are present in Bitbucket but not in the policy
  file.
*/
func enforceDeployKeys(owner string, repo string, keys publicKeyList) error {
	currkeys, _ := bbAPI.GetDeployKeys(owner, repo)

	newkeys := make(publicKeyList, len(keys))
	copy(newkeys, keys)

	for _, key := range currkeys {
		match, matchIndex := newkeys.hasKey(key)

		if match == matchContent {
			// Delete the key from BB so it can be reuploaded with proper name
			if err := bbAPI.DeleteDeployKey(owner, repo, key.ID); err != nil {
				return err
			}
		} else if match == matchExact {
			// Don't waste time reuploading key as it is an exact match
			newkeys = append(newkeys[:matchIndex], newkeys[(matchIndex+1):]...)
		}
	}

	for _, key := range newkeys {
		if err := bbAPI.AddDeployKey(owner, repo, key.Name, key.Key); err != nil {
			return err
		}
	}

	return nil
}

func parseConfig(configFile string) (repositorySettings, error) {
	rawConfig, err := ioutil.ReadFile(fmt.Sprintf("%s/%s.json", *configDir, configFile))
	if err != nil {
		return repositorySettings{}, err
	}

	var config repositorySettings
	if err := json.Unmarshal(rawConfig, &config); err != nil {
		return repositorySettings{}, err
	}

	if *verbose {
		log.Info("Loaded config: ", config)
	}

	return config, nil
}
<commit_msg>Reduce polling rate<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/jumoel/bitbucket-enforcer/gobucket"
	"github.com/jumoel/bitbucket-enforcer/log"
	dotenv "github.com/jumoel/bitbucket-enforcer/vendor/godotenv"
)

type branchManagement struct {
	PreventDelete []string
	PreventRebase []string
	AllowPushes   map[string]struct {
		Groups []string
		Users  []string
	}
}

type accessManagement struct {
	Users  map[string]string // usernames => permissions
	Groups map[string]string // groupnames => permissions
}

type repositorySettings struct {
	LandingPage      string
	Private          interface{}
	Forks            string
	IssueTracker     string
	DeployKeys       publicKeyList
	PostHooks        []string
	BranchManagement branchManagement
	AccessManagement accessManagement
}

type publicKey struct {
	Name string
	Key  string
}

type publicKeyList []publicKey

type bbServices []gobucket.Service
type matchType int

const (
	matchNone matchType = iota
	matchContent
	matchExact
)

const sleepTime = 5 * time.Second

var configDir = flag.String("configdir", "configs", "the folder containing repository configrations")
var verbose = flag.Bool("v", false, "print more output")
var bbAPI *gobucket.APIClient

func main() {
	log.SetPrefix("bitbucket-enforcer")

	flag.Parse()

	err := dotenv.Load()
	if err != nil {
		log.Notice(".env error", err)
	}

	bbUsername := os.Getenv("BITBUCKET_ENFORCER_USERNAME")
	bbKey := os.Getenv("BITBUCKET_ENFORCER_API_KEY")

	bbAPI = gobucket.New(bbUsername, bbKey)

	scanRepositories(bbUsername)
}

func scanRepositories(bbUsername string) {
	var enforcementMatcher = regexp.MustCompile(`-enforce(?:=([a-zA-Z0-9]+))?`)

	var lastEtag string
	var changed bool

	for _ = range time.Tick(sleepTime) {
		var err error
		if changed, lastEtag, err = bbAPI.RepositoriesChanged(bbUsername, lastEtag); err != nil {
			log.Error(fmt.Sprintf("Error determining if repository list has changed (%s)", err))
			continue
		}

		if !changed {
			if *verbose {
				log.Info("No repository changes, sleeping.")
			}
			continue
		}

		log.Info("Repository list changed")

		repos, err := bbAPI.GetRepositories(bbUsername)

		if err != nil {
			log.Error("Error getting repository list", err)
			continue
		}

		for _, repo := range repos {
			if strings.Contains(repo.Description, "-noenforce") {
				if *verbose {
					log.Info(fmt.Sprintf("Skipping <%s> because of '-noenforce'\n", repo.FullName))
				}
				continue
			}

			if strings.Contains(repo.Description, "-enforced") {
				if *verbose {
					log.Info(fmt.Sprintf("Skipping <%s> because of '-enforced'\n", repo.FullName))
				}
				continue
			}

			matches := enforcementMatcher.FindStringSubmatch(repo.Description)

			enforcementPolicy := "default"
			if len(matches) > 0 {
				enforcementPolicy = matches[1]
			}

			log.Info(fmt.Sprintf("Enforcing repo '%s' with policy '%s'", repo.FullName, enforcementPolicy))

			parts := strings.Split(repo.FullName, "/")
			err := enforcePolicy(parts[0], parts[1], enforcementPolicy)

			if err != nil {
				log.Warning(fmt.Sprintf("Could not enforce policy '%s' on repo '%s'. Will be processed again next cycle. (%s)", enforcementPolicy, repo.FullName, err))
			} else {
				newDescription := strings.TrimSpace(fmt.Sprintf("%s\n\n-enforced", repo.Description))

				if err := bbAPI.SetDescription(parts[0], parts[1], newDescription); err != nil {
					log.Warning("Could not set description on repo '%s'. Will be processed again next cycle.", repo.FullName)
				}
			}
		}
	}
}

func enforcePolicy(owner string, repo string, policyname string) error {
	policy, err := parseConfig(policyname)

	if err != nil {
		log.Error(fmt.Sprintf("Error parsing parsing policy '%s': ", policyname), err)
		return err
	}

	if policy.Private != nil {
		if err := bbAPI.SetPrivacy(owner, repo, policy.Private.(bool)); err != nil {
			log.Warning("Error setting privacy: ", err)
			return err
		}
	}

	if policy.Forks != "" {
		if err := bbAPI.SetForks(owner, repo, policy.Forks); err != nil {
			log.Warning("Error fork policy: ", err)
			return err
		}
	}

	if policy.LandingPage != "" {
		if err := bbAPI.SetLandingPage(owner, repo, policy.LandingPage); err != nil {
			log.Warning("Error setting landing page: ", err)
			return err
		}
	}

	if len(policy.DeployKeys) > 0 {
		if err := enforceDeployKeys(owner, repo, policy.DeployKeys); err != nil {
			log.Warning("Error setting deploy keys: ", err)
			return err
		}
	}

	if len(policy.PostHooks) > 0 {
		if err := enforcePOSTHooks(owner, repo, policy.PostHooks); err != nil {
			log.Warning("Error setting POST hooks: ", err)
			return err
		}
	}

	if policy.IssueTracker != "" {
		if err := bbAPI.SetIssueTracker(owner, repo, policy.IssueTracker); err != nil {
			log.Warning("Error setting issue tracker: ", err)
			return err
		}
	}

	if err := enforceBranchManagement(owner, repo, policy.BranchManagement); err != nil {
		log.Warning("Error setting branch policies: ", err)
		return err
	}

	if err := enforceAccessManagement(owner, repo, policy.AccessManagement); err != nil {
		log.Warning("Error setting access policies: ", err)
		return err
	}

	return nil
}

func enforceAccessManagement(owner string, repo string, policies accessManagement) error {
	for username, privilege := range policies.Users {
		if err := bbAPI.AddUserPrivilege(owner, repo, username, privilege); err != nil {
			return err
		}
	}

	for groupname, privilege := range policies.Groups {
		if err := bbAPI.AddGroupPrivilege(owner, repo, groupname, privilege); err != nil {
			return err
		}
	}

	return nil
}

func enforceBranchManagement(owner string, repo string, policies branchManagement) error {
	for _, branch := range policies.PreventDelete {
		if err := bbAPI.AddBranchRestriction(owner, repo, "delete", branch, nil, nil); err != nil {
			return err
		}
	}

	for _, branch := range policies.PreventRebase {
		if err := bbAPI.AddBranchRestriction(owner, repo, "force", branch, nil, nil); err != nil {
			return err
		}
	}

	for branch, permissions := range policies.AllowPushes {
		if err := bbAPI.AddBranchRestriction(owner, repo, "push", branch, permissions.Users, permissions.Groups); err != nil {
			return err
		}
	}

	return nil
}

func (hooks *bbServices) hasPOSTHook(URL string) bool {
	for _, hook := range *hooks {
		if hook.Service.Type == "POST" {
			for _, field := range hook.Service.Fields {
				if field.Name == "URL" && field.Value == URL {
					return true
				}
			}
		}
	}

	return false
}

func enforcePOSTHooks(owner string, repo string, hookURLs []string) error {
	hookList, err := bbAPI.GetServices(owner, repo)

	if err != nil {
		return err
	}

	var currentHooks bbServices = hookList

	for _, url := range hookURLs {
		if !currentHooks.hasPOSTHook(url) {
			if err := bbAPI.AddService(owner, repo, "POST", map[string]string{"URL": url}); err != nil {
				return err
			}
		}
	}

	return nil
}

func (keys *publicKeyList) hasKey(needle gobucket.DeployKey) (matchType, int) {
	for index, key := range *keys {
		if key.Key == needle.Key && key.Name == needle.Label {
			return matchExact, index
		} else if key.Key == needle.Key {
			return matchContent, index
		}
	}

	return matchNone, -1
}

/*
This method ensures the presence of all required keys.
- It removes keys with matching content but mismatching names. Afterwards they
  are added again, this time with the correct name.
- It adds keys that are not present.
- It doesn't remove keys that are present in Bitbucket but not in the policy
  file.
*/
func enforceDeployKeys(owner string, repo string, keys publicKeyList) error {
	currkeys, _ := bbAPI.GetDeployKeys(owner, repo)

	newkeys := make(publicKeyList, len(keys))
	copy(newkeys, keys)

	for _, key := range currkeys {
		match, matchIndex := newkeys.hasKey(key)

		if match == matchContent {
			// Delete the key from BB so it can be reuploaded with proper name
			if err := bbAPI.DeleteDeployKey(owner, repo, key.ID); err != nil {
				return err
			}
		} else if match == matchExact {
			// Don't waste time reuploading key as it is an exact match
			newkeys = append(newkeys[:matchIndex], newkeys[(matchIndex+1):]...)
		}
	}

	for _, key := range newkeys {
		if err := bbAPI.AddDeployKey(owner, repo, key.Name, key.Key); err != nil {
			return err
		}
	}

	return nil
}

func parseConfig(configFile string) (repositorySettings, error) {
	rawConfig, err := ioutil.ReadFile(fmt.Sprintf("%s/%s.json", *configDir, configFile))
	if err != nil {
		return repositorySettings{}, err
	}

	var config repositorySettings
	if err := json.Unmarshal(rawConfig, &config); err != nil {
		return repositorySettings{}, err
	}

	if *verbose {
		log.Info("Loaded config: ", config)
	}

	return config, nil
}
<|endoftext|>


<commit_before>package env

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"

	"github.com/zouyx/agollo/v3/component/log"
	jsonConfig "github.com/zouyx/agollo/v3/env/config/json"
)

const suffix = ".json"

var (
	configFileMap  = make(map[string]string, 1)
	jsonFileConfig = &jsonConfig.ConfigFile{}
)

//WriteWithRaw
func WriteWithRaw(f func(config *ApolloConfig, configPath string) error) func(config *ApolloConfig, configPath string) error {
	return func(config *ApolloConfig, configPath string) error {
		filePath := fmt.Sprintf("%s/%s", configPath, config.NamespaceName)
		file, e := os.Create(filePath)
		if e != nil {
			return e
		}
		defer file.Close()
		_, e = file.WriteString(config.Configurations["content"])
		if e != nil {
			return e
		}

		return f(config, configPath)
	}
}

//WriteConfigFile write config to file
func WriteConfigFile(config *ApolloConfig, configPath string) error {
	return jsonFileConfig.Write(config, GetConfigFile(configPath, config.NamespaceName))
}

//GetConfigFile get real config file
func GetConfigFile(configDir string, namespace string) string {
	fullPath := configFileMap[namespace]
	if fullPath == "" {
		filePath := fmt.Sprintf("%s%s", namespace, suffix)
		if configDir != "" {
			configFileMap[namespace] = fmt.Sprintf("%s/%s", configDir, filePath)
		} else {
			configFileMap[namespace] = filePath
		}
	}
	return configFileMap[namespace]
}

//LoadConfigFile load config from file
func LoadConfigFile(configDir string, namespace string) (*ApolloConfig, error) {
	configFilePath := GetConfigFile(configDir, namespace)
	log.Info("load config file from :", configFilePath)
	c, e := jsonFileConfig.Load(configFilePath, func(b []byte) (interface{}, error) {
		config := &ApolloConfig{}
		e := json.NewDecoder(bytes.NewBuffer(b)).Decode(config)
		return config, e
	})

	if c == nil || e != nil {
		log.Errorf("loadConfigFile fail,error:", e)
		return nil, e
	}

	return c.(*ApolloConfig), e
}
<commit_msg>修正代码注释<commit_after>package env

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"

	"github.com/zouyx/agollo/v3/component/log"
	jsonConfig "github.com/zouyx/agollo/v3/env/config/json"
)

const suffix = ".json"

var (
	configFileMap  = make(map[string]string, 1)
	jsonFileConfig = &jsonConfig.ConfigFile{}
)

//WriteWithRaw decorator for WriteConfigFile
func WriteWithRaw(f func(config *ApolloConfig, configPath string) error) func(config *ApolloConfig, configPath string) error {
	return func(config *ApolloConfig, configPath string) error {
		filePath := fmt.Sprintf("%s/%s", configPath, config.NamespaceName)
		file, e := os.Create(filePath)
		if e != nil {
			return e
		}
		defer file.Close()
		_, e = file.WriteString(config.Configurations["content"])
		if e != nil {
			return e
		}

		return f(config, configPath)
	}
}

//WriteConfigFile write config to file
func WriteConfigFile(config *ApolloConfig, configPath string) error {
	return jsonFileConfig.Write(config, GetConfigFile(configPath, config.NamespaceName))
}

//GetConfigFile get real config file
func GetConfigFile(configDir string, namespace string) string {
	fullPath := configFileMap[namespace]
	if fullPath == "" {
		filePath := fmt.Sprintf("%s%s", namespace, suffix)
		if configDir != "" {
			configFileMap[namespace] = fmt.Sprintf("%s/%s", configDir, filePath)
		} else {
			configFileMap[namespace] = filePath
		}
	}
	return configFileMap[namespace]
}

//LoadConfigFile load config from file
func LoadConfigFile(configDir string, namespace string) (*ApolloConfig, error) {
	configFilePath := GetConfigFile(configDir, namespace)
	log.Info("load config file from :", configFilePath)
	c, e := jsonFileConfig.Load(configFilePath, func(b []byte) (interface{}, error) {
		config := &ApolloConfig{}
		e := json.NewDecoder(bytes.NewBuffer(b)).Decode(config)
		return config, e
	})

	if c == nil || e != nil {
		log.Errorf("loadConfigFile fail,error:", e)
		return nil, e
	}

	return c.(*ApolloConfig), e
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package badges

import (
	"sync"

	"encoding/json"

	"github.com/keybase/client/go/gregor"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
)

// BadgeState represents the number of badges on the app. It's threadsafe.
// Useable from both the client service and gregor server.
// See service:Badger for the service part that owns this.
type BadgeState struct {
	sync.Mutex

	log   logger.Logger
	state keybase1.BadgeState

	inboxVers chat1.InboxVers
	// Map from ConversationID.String to BadgeConversationInfo.
	chatUnreadMap map[string]keybase1.BadgeConversationInfo
}

// NewBadgeState creates a new empty BadgeState.
func NewBadgeState(log logger.Logger) *BadgeState {
	return &BadgeState{
		log:           log,
		inboxVers:     chat1.InboxVers(0),
		chatUnreadMap: make(map[string]keybase1.BadgeConversationInfo),
	}
}

// Exports the state summary
func (b *BadgeState) Export() (keybase1.BadgeState, error) {
	b.Lock()
	defer b.Unlock()

	b.state.Conversations = []keybase1.BadgeConversationInfo{}
	for _, info := range b.chatUnreadMap {
		b.state.Conversations = append(b.state.Conversations, info)
	}
	b.state.InboxVers = int(b.inboxVers)

	return b.state, nil
}

type problemSetBody struct {
	Count int `json:"count"`
}

// UpdateWithGregor updates the badge state from a gregor state.
func (b *BadgeState) UpdateWithGregor(gstate gregor.State) error {
	b.Lock()
	defer b.Unlock()

	b.state.NewTlfs = 0
	b.state.NewFollowers = 0
	b.state.RekeysNeeded = 0

	items, err := gstate.Items()
	if err != nil {
		return err
	}
	for _, item := range items {
		categoryObj := item.Category()
		if categoryObj == nil {
			continue
		}
		category := categoryObj.String()
		switch category {
		case "tlf":
			jsw, err := jsonw.Unmarshal(item.Body().Bytes())
			if err != nil {
				b.log.Warning("BadgeState encountered non-json 'tlf' item: %v", err)
				continue
			}
			itemType, err := jsw.AtKey("type").GetString()
			if err != nil {
				b.log.Warning("BadgeState encountered gregor 'tlf' item without 'type': %v", err)
				continue
			}
			if itemType != "created" {
				continue
			}
			b.state.NewTlfs++
		case "kbfs_tlf_problem_set_count", "kbfs_tlf_sbs_problem_set_count":
			var body problemSetBody
			if err := json.Unmarshal(item.Body().Bytes(), &body); err != nil {
				b.log.Warning("BadgeState encountered non-json 'problem set' item: %v", err)
				continue
			}
			b.state.RekeysNeeded += body.Count
		case "follow":
			b.state.NewFollowers++
		case "new_git_repo":
			jsw, err := jsonw.Unmarshal(item.Body().Bytes())
			if err != nil {
				b.log.Warning("BadgeState encountered non-json 'new_git_repo' item: %v", err)
				continue
			}
			globalUniqueID, err := jsw.AtKey("global_unique_id").GetString()
			if err != nil {
				b.log.Warning("BadgeState encountered gregor 'new_git_repo' item without 'global_unique_id': %v", err)
				continue
			}
			b.state.NewGitRepoGlobalUniqueIDs = append(b.state.NewGitRepoGlobalUniqueIDs, globalUniqueID)
		}
	}

	return nil
}

func (b *BadgeState) UpdateWithChat(update chat1.UnreadUpdate, inboxVers chat1.InboxVers) {
	b.Lock()
	defer b.Unlock()

	// Skip stale updates
	if inboxVers < b.inboxVers {
		return
	}

	b.inboxVers = inboxVers
	b.updateWithChat(update)
}

func (b *BadgeState) UpdateWithChatFull(update chat1.UnreadUpdateFull) {
	b.Lock()
	defer b.Unlock()

	if update.Ignore {
		return
	}

	// Skip stale updates
	if update.InboxVers < b.inboxVers {
		return
	}

	b.chatUnreadMap = make(map[string]keybase1.BadgeConversationInfo)

	for _, upd := range update.Updates {
		b.updateWithChat(upd)
	}

	b.inboxVers = update.InboxVers
}

func (b *BadgeState) Clear() {
	b.Lock()
	defer b.Unlock()

	b.state = keybase1.BadgeState{}
	b.inboxVers = chat1.InboxVers(0)
	b.chatUnreadMap = make(map[string]keybase1.BadgeConversationInfo)
}

func (b *BadgeState) updateWithChat(update chat1.UnreadUpdate) {
	b.chatUnreadMap[update.ConvID.String()] = keybase1.BadgeConversationInfo{
		ConvID:         keybase1.ChatConversationID(update.ConvID),
		UnreadMessages: update.UnreadMessages,
		BadgeCounts:    update.UnreadNotifyingMessages,
	}
}
<commit_msg>fix forgetting to clear the repos list in badge state<commit_after>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package badges

import (
	"sync"

	"encoding/json"

	"github.com/keybase/client/go/gregor"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
)

// BadgeState represents the number of badges on the app. It's threadsafe.
// Useable from both the client service and gregor server.
// See service:Badger for the service part that owns this.
type BadgeState struct {
	sync.Mutex

	log   logger.Logger
	state keybase1.BadgeState

	inboxVers chat1.InboxVers
	// Map from ConversationID.String to BadgeConversationInfo.
	chatUnreadMap map[string]keybase1.BadgeConversationInfo
}

// NewBadgeState creates a new empty BadgeState.
func NewBadgeState(log logger.Logger) *BadgeState {
	return &BadgeState{
		log:           log,
		inboxVers:     chat1.InboxVers(0),
		chatUnreadMap: make(map[string]keybase1.BadgeConversationInfo),
	}
}

// Exports the state summary
func (b *BadgeState) Export() (keybase1.BadgeState, error) {
	b.Lock()
	defer b.Unlock()

	b.state.Conversations = []keybase1.BadgeConversationInfo{}
	for _, info := range b.chatUnreadMap {
		b.state.Conversations = append(b.state.Conversations, info)
	}
	b.state.InboxVers = int(b.inboxVers)

	return b.state, nil
}

type problemSetBody struct {
	Count int `json:"count"`
}

// UpdateWithGregor updates the badge state from a gregor state.
func (b *BadgeState) UpdateWithGregor(gstate gregor.State) error {
	b.Lock()
	defer b.Unlock()

	b.state.NewTlfs = 0
	b.state.NewFollowers = 0
	b.state.RekeysNeeded = 0
	b.state.NewGitRepoGlobalUniqueIDs = []string{}

	items, err := gstate.Items()
	if err != nil {
		return err
	}
	for _, item := range items {
		categoryObj := item.Category()
		if categoryObj == nil {
			continue
		}
		category := categoryObj.String()
		switch category {
		case "tlf":
			jsw, err := jsonw.Unmarshal(item.Body().Bytes())
			if err != nil {
				b.log.Warning("BadgeState encountered non-json 'tlf' item: %v", err)
				continue
			}
			itemType, err := jsw.AtKey("type").GetString()
			if err != nil {
				b.log.Warning("BadgeState encountered gregor 'tlf' item without 'type': %v", err)
				continue
			}
			if itemType != "created" {
				continue
			}
			b.state.NewTlfs++
		case "kbfs_tlf_problem_set_count", "kbfs_tlf_sbs_problem_set_count":
			var body problemSetBody
			if err := json.Unmarshal(item.Body().Bytes(), &body); err != nil {
				b.log.Warning("BadgeState encountered non-json 'problem set' item: %v", err)
				continue
			}
			b.state.RekeysNeeded += body.Count
		case "follow":
			b.state.NewFollowers++
		case "new_git_repo":
			jsw, err := jsonw.Unmarshal(item.Body().Bytes())
			if err != nil {
				b.log.Warning("BadgeState encountered non-json 'new_git_repo' item: %v", err)
				continue
			}
			globalUniqueID, err := jsw.AtKey("global_unique_id").GetString()
			if err != nil {
				b.log.Warning("BadgeState encountered gregor 'new_git_repo' item without 'global_unique_id': %v", err)
				continue
			}
			b.state.NewGitRepoGlobalUniqueIDs = append(b.state.NewGitRepoGlobalUniqueIDs, globalUniqueID)
		}
	}

	return nil
}

func (b *BadgeState) UpdateWithChat(update chat1.UnreadUpdate, inboxVers chat1.InboxVers) {
	b.Lock()
	defer b.Unlock()

	// Skip stale updates
	if inboxVers < b.inboxVers {
		return
	}

	b.inboxVers = inboxVers
	b.updateWithChat(update)
}

func (b *BadgeState) UpdateWithChatFull(update chat1.UnreadUpdateFull) {
	b.Lock()
	defer b.Unlock()

	if update.Ignore {
		return
	}

	// Skip stale updates
	if update.InboxVers < b.inboxVers {
		return
	}

	b.chatUnreadMap = make(map[string]keybase1.BadgeConversationInfo)

	for _, upd := range update.Updates {
		b.updateWithChat(upd)
	}

	b.inboxVers = update.InboxVers
}

func (b *BadgeState) Clear() {
	b.Lock()
	defer b.Unlock()

	b.state = keybase1.BadgeState{}
	b.inboxVers = chat1.InboxVers(0)
	b.chatUnreadMap = make(map[string]keybase1.BadgeConversationInfo)
}

func (b *BadgeState) updateWithChat(update chat1.UnreadUpdate) {
	b.chatUnreadMap[update.ConvID.String()] = keybase1.BadgeConversationInfo{
		ConvID:         keybase1.ChatConversationID(update.ConvID),
		UnreadMessages: update.UnreadMessages,
		BadgeCounts:    update.UnreadNotifyingMessages,
	}
}
<|endoftext|>


<commit_before>package shinylog

import (
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
)

type ShinyLogger struct {
	mu             sync.Mutex
	happyLogger    log.Logger
	errorLogger    log.Logger
	sadLogger      log.Logger
	suppressOutput bool
	disableColor   bool
}

type loggerOptions struct {
	isError, printNewline, includeLocation bool
}

var errorOptions loggerOptions
var outOptions loggerOptions

func init() {
	errorOptions = loggerOptions{isError: true, printNewline: true, includeLocation: true}
	outOptions = loggerOptions{isError: false, printNewline: true, includeLocation: false}
}

func NewShinyLogger(out, err interface {
	io.Writer
}) *ShinyLogger {
	happyLogger := log.New(out, "", 0)
	errorLogger := log.New(err, "", 0)
	sadLogger := log.New(err, "", log.Lshortfile)
	var mu sync.Mutex
	return &ShinyLogger{mu, *happyLogger, *errorLogger, *sadLogger, false, false}
}

func NewTraceLogger(out interface {
	io.Writer
}) *log.Logger {
	return log.New(out, "", log.Ldate|log.Ltime|log.Lmicroseconds)
}

const (
	red         = "\x1b[31m"
	green       = "\x1b[32m"
	brightgreen = "\x1b[1;32m"
	yellow      = "\x1b[33m"
	blue        = "\x1b[34m"
	magenta     = "\x1b[35m"
	reset       = "\x1b[0m"
)

var dlm sync.RWMutex
var defaultLogger *ShinyLogger = NewShinyLogger(os.Stdout, os.Stderr)
var traceLogger *log.Logger = nil

func DefaultLogger() *ShinyLogger {
	dlm.RLock()
	defer dlm.RUnlock()
	return defaultLogger
}

func SetDefaultLogger(sl *ShinyLogger) {
	dlm.Lock()
	defaultLogger = sl
	dlm.Unlock()
}

func TraceLogger() *log.Logger {
	dlm.RLock()
	defer dlm.RUnlock()
	return traceLogger
}

func SetTraceLogger(sl *log.Logger) {
	dlm.Lock()
	traceLogger = sl
	dlm.Unlock()
}

func Suppress()                           { DefaultLogger().Suppress() }
func DisableColor()                       { DefaultLogger().DisableColor() }
func Colorized(msg string) (printed bool) { return DefaultLogger().Colorized(msg) }
func Error(err error) bool                { return DefaultLogger().Error(err) }
func FatalError(err error)                { DefaultLogger().FatalError(err) }
func FatalErrorString(msg string)         { DefaultLogger().FatalErrorString(msg) }
func ErrorString(msg string) bool         { return DefaultLogger().ErrorString(msg) }
func StdErrorString(msg string) bool      { return DefaultLogger().StdErrorString(msg) }
func Red(msg string) bool                 { return DefaultLogger().Red(msg) }
func Green(msg string) bool               { return DefaultLogger().Green(msg) }
func Brightgreen(msg string) bool         { return DefaultLogger().Brightgreen(msg) }
func Yellow(msg string) bool              { return DefaultLogger().Yellow(msg) }
func Blue(msg string) bool                { return DefaultLogger().Blue(msg) }
func Magenta(msg string) bool             { return DefaultLogger().Magenta(msg) }

func TraceEnabled() bool {
	return TraceLogger() != nil
}

func Trace(format string, v ...interface{}) bool {
	if TraceEnabled() {
		TraceLogger().Printf(format, v...)
		return true
	}
	return false
}

func (l *ShinyLogger) Suppress() {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.suppressOutput = true
}

func (l *ShinyLogger) DisableColor() {
	l.disableColor = true
}

func (l *ShinyLogger) Colorized(msg string) (printed bool) {
	return l.colorized(3, msg, outOptions)
}

func (l *ShinyLogger) ColorizedSansNl(msg string) (printed bool) {
	return l.colorized(3, msg, loggerOptions{isError: false, printNewline: false, includeLocation: false})
}

// If we send SIGTERM rather than explicitly exiting,
// the signal can be handled and the master can clean up.
// This is a workaround for Go not having `atexit` :(.
func terminate() {
	proc, _ := os.FindProcess(os.Getpid())
	proc.Signal(syscall.SIGTERM)
}

func (l *ShinyLogger) FatalErrorString(msg string) {
	l.colorized(3, "{red}"+msg, errorOptions)
	terminate()
}

func (l *ShinyLogger) FatalError(err error) {
	l.colorized(3, "{red}"+err.Error(), errorOptions)
	terminate()
}

func (l *ShinyLogger) Error(err error) bool {
	return l.colorized(3, "{red}"+err.Error(), errorOptions)
}

func (l *ShinyLogger) ErrorString(msg string) bool {
	return l.colorized(3, "{red}"+msg, errorOptions)
}

func (l *ShinyLogger) StdErrorString(msg string) bool {
	return l.colorized(3, "{red}"+msg, loggerOptions{isError: true, printNewline: true, includeLocation: false})
}

func (l *ShinyLogger) Red(msg string) bool {
	return l.colorized(3, "{red}"+msg, outOptions)
}

func (l *ShinyLogger) Green(msg string) bool {
	return l.colorized(3, "{green}"+msg, outOptions)
}

func (l *ShinyLogger) Brightgreen(msg string) bool {
	return l.colorized(3, "{brightgreen}"+msg, outOptions)
}

func (l *ShinyLogger) Yellow(msg string) bool {
	return l.colorized(3, "{yellow}"+msg, outOptions)
}

func (l *ShinyLogger) Blue(msg string) bool {
	return l.colorized(3, "{blue}"+msg, outOptions)
}

func (l *ShinyLogger) Magenta(msg string) bool {
	return l.colorized(3, "{magenta}"+msg, outOptions)
}

func (l *ShinyLogger) formatColors(msg string) string {
	if l.disableColor {
		msg = strings.Replace(msg, "{red}", "", -1)
		msg = strings.Replace(msg, "{green}", "", -1)
		msg = strings.Replace(msg, "{brightgreen}", "", -1)
		msg = strings.Replace(msg, "{yellow}", "", -1)
		msg = strings.Replace(msg, "{blue}", "", -1)
		msg = strings.Replace(msg, "{magenta}", "", -1)
		msg = strings.Replace(msg, "{reset}", "", -1)
	} else {
		msg = strings.Replace(msg, "{red}", red, -1)
		msg = strings.Replace(msg, "{green}", green, -1)
		msg = strings.Replace(msg, "{brightgreen}", brightgreen, -1)
		msg = strings.Replace(msg, "{yellow}", yellow, -1)
		msg = strings.Replace(msg, "{blue}", blue, -1)
		msg = strings.Replace(msg, "{magenta}", magenta, -1)
		msg = strings.Replace(msg, "{reset}", reset, -1)
	}
	return msg
}

func (l *ShinyLogger) colorized(callDepth int, msg string, options loggerOptions) (printed bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	if !l.suppressOutput {
		msg = l.formatColors(msg)

		if l == DefaultLogger() {
			callDepth += 1 // this was called through a proxy method
		}
		if options.isError {
			if options.includeLocation {
				l.sadLogger.Output(callDepth, msg+reset)
			} else {
				l.errorLogger.Output(callDepth, msg+reset)
			}
		} else {
			if options.printNewline {
				fmt.Println(msg + reset)
			} else {
				fmt.Print(msg + reset)
			}
		}
	}
	return !l.suppressOutput
}
<commit_msg>improve naming<commit_after>package shinylog

import (
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
)

type ShinyLogger struct {
	mu             sync.Mutex
	happyLogger    log.Logger
	errorLogger    log.Logger
	sadLogger      log.Logger
	suppressOutput bool
	disableColor   bool
}

type loggerOptions struct {
	isError, printNewline, includeLocation bool
}

var errorOptions loggerOptions
var stdoutOptions loggerOptions
var stderrOptions loggerOptions

func init() {
	errorOptions = loggerOptions{isError: true, printNewline: true, includeLocation: true}
	stdoutOptions = loggerOptions{isError: false, printNewline: true, includeLocation: false}
	stderrOptions = loggerOptions{isError: false, printNewline: true, includeLocation: false}
}

func NewShinyLogger(out, err interface {
	io.Writer
}) *ShinyLogger {
	happyLogger := log.New(out, "", 0)
	errorLogger := log.New(err, "", 0)
	sadLogger := log.New(err, "", log.Lshortfile)
	var mu sync.Mutex
	return &ShinyLogger{mu, *happyLogger, *errorLogger, *sadLogger, false, false}
}

func NewTraceLogger(out interface {
	io.Writer
}) *log.Logger {
	return log.New(out, "", log.Ldate|log.Ltime|log.Lmicroseconds)
}

const (
	red         = "\x1b[31m"
	green       = "\x1b[32m"
	brightgreen = "\x1b[1;32m"
	yellow      = "\x1b[33m"
	blue        = "\x1b[34m"
	magenta     = "\x1b[35m"
	reset       = "\x1b[0m"
)

var dlm sync.RWMutex
var defaultLogger *ShinyLogger = NewShinyLogger(os.Stdout, os.Stderr)
var traceLogger *log.Logger = nil

func DefaultLogger() *ShinyLogger {
	dlm.RLock()
	defer dlm.RUnlock()
	return defaultLogger
}

func SetDefaultLogger(sl *ShinyLogger) {
	dlm.Lock()
	defaultLogger = sl
	dlm.Unlock()
}

func TraceLogger() *log.Logger {
	dlm.RLock()
	defer dlm.RUnlock()
	return traceLogger
}

func SetTraceLogger(sl *log.Logger) {
	dlm.Lock()
	traceLogger = sl
	dlm.Unlock()
}

func Suppress()                           { DefaultLogger().Suppress() }
func DisableColor()                       { DefaultLogger().DisableColor() }
func Colorized(msg string) (printed bool) { return DefaultLogger().Colorized(msg) }
func Error(err error) bool                { return DefaultLogger().Error(err) }
func FatalError(err error)                { DefaultLogger().FatalError(err) }
func FatalErrorString(msg string)         { DefaultLogger().FatalErrorString(msg) }
func ErrorString(msg string) bool         { return DefaultLogger().ErrorString(msg) }
func StdErrorString(msg string) bool      { return DefaultLogger().StdErrorString(msg) }
func Red(msg string) bool                 { return DefaultLogger().Red(msg) }
func Green(msg string) bool               { return DefaultLogger().Green(msg) }
func Brightgreen(msg string) bool         { return DefaultLogger().Brightgreen(msg) }
func Yellow(msg string) bool              { return DefaultLogger().Yellow(msg) }
func Blue(msg string) bool                { return DefaultLogger().Blue(msg) }
func Magenta(msg string) bool             { return DefaultLogger().Magenta(msg) }

func TraceEnabled() bool {
	return TraceLogger() != nil
}

func Trace(format string, v ...interface{}) bool {
	if TraceEnabled() {
		TraceLogger().Printf(format, v...)
		return true
	}
	return false
}

func (l *ShinyLogger) Suppress() {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.suppressOutput = true
}

func (l *ShinyLogger) DisableColor() {
	l.disableColor = true
}

func (l *ShinyLogger) Colorized(msg string) (printed bool) {
	return l.colorized(3, msg, stdoutOptions)
}

func (l *ShinyLogger) ColorizedSansNl(msg string) (printed bool) {
	return l.colorized(3, msg, loggerOptions{isError: false, printNewline: false, includeLocation: false})
}

// If we send SIGTERM rather than explicitly exiting,
// the signal can be handled and the master can clean up.
// This is a workaround for Go not having `atexit` :(.
func terminate() {
	proc, _ := os.FindProcess(os.Getpid())
	proc.Signal(syscall.SIGTERM)
}

func (l *ShinyLogger) FatalErrorString(msg string) {
	l.colorized(3, "{red}"+msg, errorOptions)
	terminate()
}

func (l *ShinyLogger) FatalError(err error) {
	l.colorized(3, "{red}"+err.Error(), errorOptions)
	terminate()
}

func (l *ShinyLogger) Error(err error) bool {
	return l.colorized(3, "{red}"+err.Error(), errorOptions)
}

func (l *ShinyLogger) ErrorString(msg string) bool {
	return l.colorized(3, "{red}"+msg, errorOptions)
}

func (l *ShinyLogger) StdErrorString(msg string) bool {
	return l.colorized(3, "{red}"+msg, stderrOptions)
}

func (l *ShinyLogger) Red(msg string) bool {
	return l.colorized(3, "{red}"+msg, stdoutOptions)
}

func (l *ShinyLogger) Green(msg string) bool {
	return l.colorized(3, "{green}"+msg, stdoutOptions)
}

func (l *ShinyLogger) Brightgreen(msg string) bool {
	return l.colorized(3, "{brightgreen}"+msg, stdoutOptions)
}

func (l *ShinyLogger) Yellow(msg string) bool {
	return l.colorized(3, "{yellow}"+msg, stdoutOptions)
}

func (l *ShinyLogger) Blue(msg string) bool {
	return l.colorized(3, "{blue}"+msg, stdoutOptions)
}

func (l *ShinyLogger) Magenta(msg string) bool {
	return l.colorized(3, "{magenta}"+msg, stdoutOptions)
}

func (l *ShinyLogger) formatColors(msg string) string {
	if l.disableColor {
		msg = strings.Replace(msg, "{red}", "", -1)
		msg = strings.Replace(msg, "{green}", "", -1)
		msg = strings.Replace(msg, "{brightgreen}", "", -1)
		msg = strings.Replace(msg, "{yellow}", "", -1)
		msg = strings.Replace(msg, "{blue}", "", -1)
		msg = strings.Replace(msg, "{magenta}", "", -1)
		msg = strings.Replace(msg, "{reset}", "", -1)
	} else {
		msg = strings.Replace(msg, "{red}", red, -1)
		msg = strings.Replace(msg, "{green}", green, -1)
		msg = strings.Replace(msg, "{brightgreen}", brightgreen, -1)
		msg = strings.Replace(msg, "{yellow}", yellow, -1)
		msg = strings.Replace(msg, "{blue}", blue, -1)
		msg = strings.Replace(msg, "{magenta}", magenta, -1)
		msg = strings.Replace(msg, "{reset}", reset, -1)
	}
	return msg
}

func (l *ShinyLogger) colorized(callDepth int, msg string, options loggerOptions) (printed bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	if !l.suppressOutput {
		msg = l.formatColors(msg)

		if l == DefaultLogger() {
			callDepth += 1 // this was called through a proxy method
		}
		if options.isError {
			if options.includeLocation {
				l.sadLogger.Output(callDepth, msg+reset)
			} else {
				l.errorLogger.Output(callDepth, msg+reset)
			}
		} else {
			if options.printNewline {
				fmt.Println(msg + reset)
			} else {
				fmt.Print(msg + reset)
			}
		}
	}
	return !l.suppressOutput
}
<|endoftext|>


<commit_before>// Copyright 2017 The OpenSDS 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.

/*
This module implements a entry into the OpenSDS northbound service.
*/

package controller

import (
	"encoding/json"
	"errors"
	"fmt"

	log "github.com/golang/glog"

	"github.com/opensds/opensds/pkg/controller/policy"
	"github.com/opensds/opensds/pkg/controller/selector"
	"github.com/opensds/opensds/pkg/controller/volume"
	"github.com/opensds/opensds/pkg/db"
	pb "github.com/opensds/opensds/pkg/dock/proto"
	"github.com/opensds/opensds/pkg/model"
)

const (
	CREATE_LIFECIRCLE_FLAG = iota + 1
	GET_LIFECIRCLE_FLAG
	LIST_LIFECIRCLE_FLAG
	DELETE_LIFECIRCLE_FLAG
)

var Brain *Controller

func NewController() *Controller {
	return &Controller{
		selector:         selector.NewSelector(),
		volumeController: volume.NewController(),
	}
}

type Controller struct {
	selector         selector.Selector
	volumeController volume.Controller
	policyController policy.Controller
}

func (c *Controller) CreateVolume(in *model.VolumeSpec) (*model.VolumeSpec, error) {
	var profile *model.ProfileSpec
	var err error

	if in.ProfileId == "" {
		log.Warning("Use default profile when user doesn't specify profile.")
		profile, err = db.C.GetDefaultProfile()
	} else {
		profile, err = db.C.GetProfile(in.ProfileId)
	}
	if err != nil {
		log.Error("Get profile failed: ", err)
		return nil, err
	}

	if in.Size <= 0 {
		errMsg := fmt.Sprintf("Invalid volume size: %d", in.Size)
		log.Error(errMsg)
		return nil, errors.New(errMsg)
	}

	if in.AvailabilityZone == "" {
		log.Warning("Use default availability zone when user doesn't specify availabilityZone.")
		in.AvailabilityZone = "default"
	}

	var filterRequest map[string]interface{}
	if profile.Extra != nil {
		filterRequest = profile.Extra
	} else {
		filterRequest = make(map[string]interface{})
	}
	filterRequest["size"] = in.Size
	filterRequest["availabilityZone"] = in.AvailabilityZone

	polInfo, err := c.selector.SelectSupportedPool(filterRequest)
	if err != nil {
		log.Error("When search supported pool resource:", err)
		return nil, err
	}
	dockInfo, err := db.C.GetDock(polInfo.DockId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return nil, err
	}

	c.volumeController.SetDock(dockInfo)
	opt := &pb.CreateVolumeOpts{
		Id:               in.GetId(),
		Name:             in.GetName(),
		Description:      in.GetDescription(),
		Size:             in.GetSize(),
		AvailabilityZone: in.GetAvailabilityZone(),
		ProfileId:        profile.GetId(),
		PoolId:           polInfo.GetId(),
		PoolName:         polInfo.GetName(),
		DockId:           dockInfo.GetId(),
		DriverName:       dockInfo.GetDriverName(),
	}
	result, err := c.volumeController.CreateVolume(opt)
	if err != nil {
		return nil, err
	}

	// Select the storage tag according to the lifecycle flag.
	c.policyController = policy.NewController(profile)
	c.policyController.Setup(CREATE_LIFECIRCLE_FLAG)
	c.policyController.SetDock(dockInfo)

	var errChan = make(chan error, 1)
	volBody, _ := json.Marshal(result)
	go c.policyController.ExecuteAsyncPolicy(opt, string(volBody), errChan)

	return result, nil
}

func (c *Controller) DeleteVolume(in *model.VolumeSpec) error {
	prf, err := db.C.GetProfile(in.ProfileId)
	if err != nil {
		log.Error("when search profile in db:", err)
		return err
	}

	// Select the storage tag according to the lifecycle flag.
	c.policyController = policy.NewController(prf)
	c.policyController.Setup(DELETE_LIFECIRCLE_FLAG)

	dockInfo, err := db.C.GetDockByPoolId(in.PoolId)
	if err != nil {
		log.Error("When search dock in db by pool id: ", err)
		return err
	}
	c.policyController.SetDock(dockInfo)
	c.volumeController.SetDock(dockInfo)

	opt := &pb.DeleteVolumeOpts{
		Id:         in.GetId(),
		Metadata:   in.GetMetadata(),
		DockId:     dockInfo.GetId(),
		DriverName: dockInfo.GetDriverName(),
	}

	var errChan = make(chan error, 1)
	go c.policyController.ExecuteAsyncPolicy(opt, "", errChan)

	if err := <-errChan; err != nil {
		log.Error("When execute async policy:", err)
		return err
	}

	return c.volumeController.DeleteVolume(opt)
}

func (c *Controller) CreateVolumeAttachment(in *model.VolumeAttachmentSpec) (*model.VolumeAttachmentSpec, error) {
	volume, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in create volume attachment method: ", err)
		return nil, err
	}
	dockInfo, err := db.C.GetDockByPoolId(volume.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return nil, err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.CreateVolumeAttachment(
		&pb.CreateAttachmentOpts{
			Id:       in.GetId(),
			VolumeId: in.GetVolumeId(),
			HostInfo: &pb.HostInfo{
				Platform:  in.GetPlatform(),
				OsType:    in.GetOsType(),
				Ip:        in.GetIp(),
				Host:      in.GetHost(),
				Initiator: in.GetInitiator(),
			},
			Metadata:   in.GetMetadata(),
			DockId:     dockInfo.GetId(),
			DriverName: dockInfo.GetDriverName(),
		},
	)
}

func (c *Controller) UpdateVolumeAttachment(in *model.VolumeAttachmentSpec) (*model.VolumeAttachmentSpec, error) {
	return nil, errors.New("Not implemented!")
}

func (c *Controller) DeleteVolumeAttachment(in *model.VolumeAttachmentSpec) error {
	volume, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in delete volume attachment method: ", err)
		return err
	}
	dockInfo, err := db.C.GetDockByPoolId(volume.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.DeleteVolumeAttachment(
		&pb.DeleteAttachmentOpts{
			Id:       in.GetId(),
			VolumeId: in.GetVolumeId(),
			HostInfo: &pb.HostInfo{
				Platform:  in.GetPlatform(),
				OsType:    in.GetOsType(),
				Ip:        in.GetIp(),
				Host:      in.GetHost(),
				Initiator: in.GetInitiator(),
			},
			Metadata:   in.GetMetadata(),
			DockId:     dockInfo.GetId(),
			DriverName: dockInfo.GetDriverName(),
		},
	)
}

func (c *Controller) CreateVolumeSnapshot(in *model.VolumeSnapshotSpec) (*model.VolumeSnapshotSpec, error) {
	volume, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in create volume snapshot method: ", err)
		return nil, err
	}

	dockInfo, err := db.C.GetDockByPoolId(volume.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return nil, err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.CreateVolumeSnapshot(
		&pb.CreateVolumeSnapshotOpts{
			Id:          in.GetId(),
			Name:        in.GetName(),
			Description: in.GetDescription(),
			Size:        in.GetSize(),
			VolumeId:    in.GetVolumeId(),
			Metadata:    in.GetMetadata(),
		},
	)
}

func (c *Controller) DeleteVolumeSnapshot(in *model.VolumeSnapshotSpec) error {
	volume, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in delete volume snapshot method: ", err)
		return err
	}
	dockInfo, err := db.C.GetDockByPoolId(volume.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.DeleteVolumeSnapshot(
		&pb.DeleteVolumeSnapshotOpts{
			Id:       in.GetId(),
			VolumeId: in.GetVolumeId(),
			Metadata: in.GetMetadata(),
		},
	)
}
<commit_msg>Fixed bugs that creating snapshot failed<commit_after>// Copyright 2017 The OpenSDS 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.

/*
This module implements a entry into the OpenSDS northbound service.
*/

package controller

import (
	"encoding/json"
	"errors"
	"fmt"

	log "github.com/golang/glog"

	"github.com/opensds/opensds/pkg/controller/policy"
	"github.com/opensds/opensds/pkg/controller/selector"
	"github.com/opensds/opensds/pkg/controller/volume"
	"github.com/opensds/opensds/pkg/db"
	pb "github.com/opensds/opensds/pkg/dock/proto"
	"github.com/opensds/opensds/pkg/model"
)

const (
	CREATE_LIFECIRCLE_FLAG = iota + 1
	GET_LIFECIRCLE_FLAG
	LIST_LIFECIRCLE_FLAG
	DELETE_LIFECIRCLE_FLAG
)

var Brain *Controller

func NewController() *Controller {
	return &Controller{
		selector:         selector.NewSelector(),
		volumeController: volume.NewController(),
	}
}

type Controller struct {
	selector         selector.Selector
	volumeController volume.Controller
	policyController policy.Controller
}

func (c *Controller) CreateVolume(in *model.VolumeSpec) (*model.VolumeSpec, error) {
	var profile *model.ProfileSpec
	var err error

	if in.ProfileId == "" {
		log.Warning("Use default profile when user doesn't specify profile.")
		profile, err = db.C.GetDefaultProfile()
	} else {
		profile, err = db.C.GetProfile(in.ProfileId)
	}
	if err != nil {
		log.Error("Get profile failed: ", err)
		return nil, err
	}

	if in.Size <= 0 {
		errMsg := fmt.Sprintf("Invalid volume size: %d", in.Size)
		log.Error(errMsg)
		return nil, errors.New(errMsg)
	}

	if in.AvailabilityZone == "" {
		log.Warning("Use default availability zone when user doesn't specify availabilityZone.")
		in.AvailabilityZone = "default"
	}

	var filterRequest map[string]interface{}
	if profile.Extra != nil {
		filterRequest = profile.Extra
	} else {
		filterRequest = make(map[string]interface{})
	}
	filterRequest["size"] = in.Size
	filterRequest["availabilityZone"] = in.AvailabilityZone

	polInfo, err := c.selector.SelectSupportedPool(filterRequest)
	if err != nil {
		log.Error("When search supported pool resource:", err)
		return nil, err
	}
	dockInfo, err := db.C.GetDock(polInfo.DockId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return nil, err
	}

	c.volumeController.SetDock(dockInfo)
	opt := &pb.CreateVolumeOpts{
		Id:               in.GetId(),
		Name:             in.GetName(),
		Description:      in.GetDescription(),
		Size:             in.GetSize(),
		AvailabilityZone: in.GetAvailabilityZone(),
		ProfileId:        profile.GetId(),
		PoolId:           polInfo.GetId(),
		PoolName:         polInfo.GetName(),
		DockId:           dockInfo.GetId(),
		DriverName:       dockInfo.GetDriverName(),
	}
	result, err := c.volumeController.CreateVolume(opt)
	if err != nil {
		return nil, err
	}

	// Select the storage tag according to the lifecycle flag.
	c.policyController = policy.NewController(profile)
	c.policyController.Setup(CREATE_LIFECIRCLE_FLAG)
	c.policyController.SetDock(dockInfo)

	var errChan = make(chan error, 1)
	volBody, _ := json.Marshal(result)
	go c.policyController.ExecuteAsyncPolicy(opt, string(volBody), errChan)

	return result, nil
}

func (c *Controller) DeleteVolume(in *model.VolumeSpec) error {
	prf, err := db.C.GetProfile(in.ProfileId)
	if err != nil {
		log.Error("when search profile in db:", err)
		return err
	}

	// Select the storage tag according to the lifecycle flag.
	c.policyController = policy.NewController(prf)
	c.policyController.Setup(DELETE_LIFECIRCLE_FLAG)

	dockInfo, err := db.C.GetDockByPoolId(in.PoolId)
	if err != nil {
		log.Error("When search dock in db by pool id: ", err)
		return err
	}
	c.policyController.SetDock(dockInfo)
	c.volumeController.SetDock(dockInfo)

	opt := &pb.DeleteVolumeOpts{
		Id:         in.GetId(),
		Metadata:   in.GetMetadata(),
		DockId:     dockInfo.GetId(),
		DriverName: dockInfo.GetDriverName(),
	}

	var errChan = make(chan error, 1)
	go c.policyController.ExecuteAsyncPolicy(opt, "", errChan)

	if err := <-errChan; err != nil {
		log.Error("When execute async policy:", err)
		return err
	}

	return c.volumeController.DeleteVolume(opt)
}

func (c *Controller) CreateVolumeAttachment(in *model.VolumeAttachmentSpec) (*model.VolumeAttachmentSpec, error) {
	vol, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in create volume attachment method: ", err)
		return nil, err
	}
	dockInfo, err := db.C.GetDockByPoolId(vol.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return nil, err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.CreateVolumeAttachment(
		&pb.CreateAttachmentOpts{
			Id:       in.GetId(),
			VolumeId: in.GetVolumeId(),
			HostInfo: &pb.HostInfo{
				Platform:  in.GetPlatform(),
				OsType:    in.GetOsType(),
				Ip:        in.GetIp(),
				Host:      in.GetHost(),
				Initiator: in.GetInitiator(),
			},
			Metadata:   in.GetMetadata(),
			DockId:     dockInfo.GetId(),
			DriverName: dockInfo.GetDriverName(),
		},
	)
}

func (c *Controller) UpdateVolumeAttachment(in *model.VolumeAttachmentSpec) (*model.VolumeAttachmentSpec, error) {
	return nil, errors.New("Not implemented!")
}

func (c *Controller) DeleteVolumeAttachment(in *model.VolumeAttachmentSpec) error {
	vol, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in delete volume attachment method: ", err)
		return err
	}
	dockInfo, err := db.C.GetDockByPoolId(vol.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.DeleteVolumeAttachment(
		&pb.DeleteAttachmentOpts{
			Id:       in.GetId(),
			VolumeId: in.GetVolumeId(),
			HostInfo: &pb.HostInfo{
				Platform:  in.GetPlatform(),
				OsType:    in.GetOsType(),
				Ip:        in.GetIp(),
				Host:      in.GetHost(),
				Initiator: in.GetInitiator(),
			},
			Metadata:   in.GetMetadata(),
			DockId:     dockInfo.GetId(),
			DriverName: dockInfo.GetDriverName(),
		},
	)
}

func (c *Controller) CreateVolumeSnapshot(in *model.VolumeSnapshotSpec) (*model.VolumeSnapshotSpec, error) {
	vol, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in create volume snapshot method: ", err)
		return nil, err
	}

	dockInfo, err := db.C.GetDockByPoolId(vol.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return nil, err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.CreateVolumeSnapshot(
		&pb.CreateVolumeSnapshotOpts{
			Id:          in.GetId(),
			Name:        in.GetName(),
			Description: in.GetDescription(),
			Size:        in.GetSize(),
			VolumeId:    in.GetVolumeId(),
			Metadata:    in.GetMetadata(),
			DockId:      dockInfo.GetId(),
			DriverName:  dockInfo.GetDriverName(),
		},
	)
}

func (c *Controller) DeleteVolumeSnapshot(in *model.VolumeSnapshotSpec) error {
	vol, err := db.C.GetVolume(in.VolumeId)
	if err != nil {
		log.Error("Get volume failed in delete volume snapshot method: ", err)
		return err
	}
	dockInfo, err := db.C.GetDockByPoolId(vol.PoolId)
	if err != nil {
		log.Error("When search supported dock resource:", err)
		return err
	}
	c.volumeController.SetDock(dockInfo)

	return c.volumeController.DeleteVolumeSnapshot(
		&pb.DeleteVolumeSnapshotOpts{
			Id:         in.GetId(),
			VolumeId:   in.GetVolumeId(),
			Metadata:   in.GetMetadata(),
			DockId:     dockInfo.GetId(),
			DriverName: dockInfo.GetDriverName(),
		},
	)
}
<|endoftext|>


<commit_before>package ephemera

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/samalba/dockerclient"
	"github.com/satori/go.uuid"
)

// Callback used to listen to Docker's events
func eventCallback(event *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Printf("Received event: %#v\n", *event)
}

type Container struct {
	Name      string
	ID        string
	Image     string
	IP        string
	Proxy     http.Handler
	Started   bool
	StartedAt time.Time
	TTL       time.Duration
	Config    *dockerclient.ContainerConfig
	e         *Ephemera
}

var (
	containerPrefix = "ephemera"
	dockerDebug     = false
)

func UUID() string {
	return uuid.NewV4().String()
}

func (c *Container) WaitKill() {
	<-time.After(c.TTL)
	c.Kill()
	return
}
func (c *Container) String() string {
	return fmt.Sprintf("<Container %v [img=%v,started=%v,ttl=%v]>", c.Name, c.Config.Image, c.Started, c.TTL)
}

func (c *Container) Start() {
	if c.Started {
		return
	}
	containerId, err := c.e.docker.CreateContainer(c.Config, fmt.Sprintf("%v-%v", containerPrefix, c.Name))
	if err != nil {
		log.Fatal(err)
	}
	// Start the container
	hostConfig := &dockerclient.HostConfig{}
	err = c.e.docker.StartContainer(containerId, hostConfig)
	if err != nil {
		log.Fatal(err)
	}
	time.Sleep(250 * time.Millisecond)
	info, _ := c.e.docker.InspectContainer(containerId)
	c.IP = info.NetworkSettings.IPAddress
	c.ID = containerId
	c.Started = true
	c.StartedAt = time.Now()
}

func (c *Container) Kill() {
	c.e.Lock()
	defer c.e.Unlock()
	c.e.docker.StopContainer(c.ID, 5)
	c.e.docker.RemoveContainer(c.ID, true, true)
	delete(c.e.containers, c.Name)
}

type Ephemera struct {
	sync.Mutex
	ttl        time.Duration
	image      string
	containers map[string]*Container
	docker     *dockerclient.DockerClient
	handler    http.Handler
}

func (e *Ephemera) KillAll() {
	for _, c := range e.containers {
		log.Printf("kill %v", c)
		c.Kill()
	}
}
func (e *Ephemera) RegisterHandler(r *mux.Router) {
	r.HandleFunc("/demo/new", e.newHandler)
	r.PathPrefix("/demo/{id}").Handler(http.HandlerFunc(e.proxyHandler))
}

func (e *Ephemera) NewContainer(img string, ttl time.Duration) *Container {
	e.Lock()
	defer e.Unlock()
	container := &Container{
		e:       e,
		Name:    UUID(),
		Image:   img,
		TTL:     ttl,
		Started: false,
		Config: &dockerclient.ContainerConfig{
			Image: img,
		},
	}
	e.containers[container.Name] = container
	return container
}

func (e *Ephemera) proxyHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	log.Println("/demo/%v requested", id)
	if c, ok := e.containers[id]; ok {
		c.Proxy.ServeHTTP(w, r)
		return
	}
	log.Printf("unknown id %v", id)
}

func (e *Ephemera) newHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("New container request")
	c := e.NewContainer(e.image, e.ttl)
	c.Start()
	log.Printf("container started: %v", c)
	go c.WaitKill()
	u, _ := url.Parse(fmt.Sprintf("http://%v:8080", c.IP))
	c.Proxy = http.StripPrefix("/demo/"+c.Name, httputil.NewSingleHostReverseProxy(u))
	log.Printf("container proxy setup /demo/%v => %v", c.Name, c.IP)
	if r.URL.Query().Get("redirect") != "0" {
		http.Redirect(w, r, "/demo/"+c.Name, http.StatusTemporaryRedirect)
		return
	}
	w.Write([]byte(c.Name))
	return
}

func New(dockerURI, image string, ttl time.Duration) (*Ephemera, error) {
	if dockerURI == "" {
		dockerURI = "unix:///var/run/docker.sock"
	}
	// Init the Docker client
	docker, err := dockerclient.NewDockerClient(dockerURI, nil)
	if err != nil {
		return nil, err
	}
	if dockerDebug {
		docker.StartMonitorEvents(eventCallback, nil)
	}
	return &Ephemera{
		containers: map[string]*Container{},
		docker:     docker,
		ttl:        ttl,
		image:      image,
	}, nil
}
<commit_msg>added docs<commit_after>package ephemera

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/samalba/dockerclient"
	"github.com/satori/go.uuid"
)

var (
	containerPrefix = "ephemera"
	dockerDebug     = false
)

// Callback used to listen to Docker's events
func eventCallback(event *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Printf("Received event: %#v\n", *event)
}

// Container represents a ephemeral container.
type Container struct {
	Name      string
	ID        string
	Image     string
	IP        string
	Proxy     http.Handler
	Started   bool
	StartedAt time.Time
	TTL       time.Duration
	Config    *dockerclient.ContainerConfig
	e         *Ephemera
}

// WaitKill blocks till the TTL is elapsed and kill the container.
func (c *Container) WaitKill() {
	<-time.After(c.TTL)
	c.Kill()
	return
}

// String implements fmt.Stringer
func (c *Container) String() string {
	return fmt.Sprintf("<Container %v [img=%v,started=%v,ttl=%v]>", c.Name, c.Config.Image, c.Started, c.TTL)
}

// Start actually start the container
func (c *Container) Start() {
	if c.Started {
		return
	}
	containerId, err := c.e.docker.CreateContainer(c.Config, fmt.Sprintf("%v-%v", containerPrefix, c.Name))
	if err != nil {
		log.Fatal(err)
	}
	// Start the container
	hostConfig := &dockerclient.HostConfig{}
	err = c.e.docker.StartContainer(containerId, hostConfig)
	if err != nil {
		log.Fatal(err)
	}
	time.Sleep(250 * time.Millisecond)
	info, _ := c.e.docker.InspectContainer(containerId)
	c.IP = info.NetworkSettings.IPAddress
	c.ID = containerId
	c.Started = true
	c.StartedAt = time.Now()
}

// Kill stops and removes the container.
func (c *Container) Kill() {
	c.e.Lock()
	defer c.e.Unlock()
	c.e.docker.StopContainer(c.ID, 5)
	c.e.docker.RemoveContainer(c.ID, true, true)
	delete(c.e.containers, c.Name)
}

type Ephemera struct {
	sync.Mutex
	ttl        time.Duration
	image      string
	containers map[string]*Container
	docker     *dockerclient.DockerClient
	handler    http.Handler
}

// KillAll kills all the spawned containers still alive.
func (e *Ephemera) KillAll() {
	for _, c := range e.containers {
		log.Printf("kill %v", c)
		c.Kill()
	}
}

// RegisterHandler registers /demo/new and /demo/{id} routes.
func (e *Ephemera) RegisterHandler(r *mux.Router) {
	r.HandleFunc("/demo/new", e.newHandler)
	r.PathPrefix("/demo/{id}").Handler(http.HandlerFunc(e.proxyHandler))
}

// Spawn a new container with the given Docker image and TTL.
// The container will be killed only if WaitKill/Kill is called manually.
func (e *Ephemera) NewContainer(img string, ttl time.Duration) *Container {
	e.Lock()
	defer e.Unlock()
	container := &Container{
		e:       e,
		Name:    uuid.NewV4().String(),
		Image:   img,
		TTL:     ttl,
		Started: false,
		Config: &dockerclient.ContainerConfig{
			Image: img,
		},
	}
	e.containers[container.Name] = container
	return container
}

func (e *Ephemera) proxyHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	log.Println("/demo/%v requested", id)
	if c, ok := e.containers[id]; ok {
		c.Proxy.ServeHTTP(w, r)
		return
	}
	log.Printf("unknown id %v", id)
}

func (e *Ephemera) newHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("New container request")
	c := e.NewContainer(e.image, e.ttl)
	c.Start()
	log.Printf("container started: %v", c)
	go c.WaitKill()
	u, _ := url.Parse(fmt.Sprintf("http://%v:8080", c.IP))
	c.Proxy = http.StripPrefix("/demo/"+c.Name, httputil.NewSingleHostReverseProxy(u))
	log.Printf("container proxy setup /demo/%v => %v", c.Name, c.IP)
	if r.URL.Query().Get("redirect") != "0" {
		http.Redirect(w, r, "/demo/"+c.Name, http.StatusTemporaryRedirect)
		return
	}
	w.Write([]byte(c.Name))
	return
}

// New initializes a new Ephemera instance.
func New(dockerURI, image string, ttl time.Duration) (*Ephemera, error) {
	if dockerURI == "" {
		dockerURI = "unix:///var/run/docker.sock"
	}
	// Init the Docker client
	docker, err := dockerclient.NewDockerClient(dockerURI, nil)
	if err != nil {
		return nil, err
	}
	if dockerDebug {
		docker.StartMonitorEvents(eventCallback, nil)
	}
	return &Ephemera{
		containers: map[string]*Container{},
		docker:     docker,
		ttl:        ttl,
		image:      image,
	}, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google, Inc. All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package v1alpha1

import (
	"reflect"
	"testing"

	corev1 "k8s.io/api/core/v1"
)

const (
	generation int64 = 5
)

func testGeneration(t *testing.T) {
	r := Revision{}
	if a := r.GetGeneration(); a != 0 {
		t.Errorf("empty revision generation should be 0 was: %d", a)
	}

	r.SetGeneration(5)
	if e, a := generation, r.GetGeneration(); e != a {
		t.Errorf("getgeneration mismatch expected: %d got: %d", e, a)
	}

}

func TestIsReady(t *testing.T) {
	cases := []struct {
		name    string
		status  RevisionStatus
		isReady bool
	}{
		{
			name:    "empty status should not be ready",
			status:  RevisionStatus{},
			isReady: false,
		},
		{
			name: "Different condition type should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionBuildComplete,
						Status: corev1.ConditionTrue,
					},
				},
			},
			isReady: false,
		},
		{
			name: "False condition status should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionFalse,
					},
				},
			},
			isReady: false,
		},
		{
			name: "Unknown condition status should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionUnknown,
					},
				},
			},
			isReady: false,
		},
		{
			name: "Missing condition status should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type: RevisionConditionReady,
					},
				},
			},
			isReady: false,
		},
		{
			name: "True condition status should be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionTrue,
					},
				},
			},
			isReady: true,
		},
		{
			name: "Multiple conditions with ready status should be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionBuildComplete,
						Status: corev1.ConditionTrue,
					},
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionTrue,
					},
				},
			},
			isReady: true,
		},
		{
			name: "Multiple conditions with ready status false should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionBuildComplete,
						Status: corev1.ConditionTrue,
					},
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionFalse,
					},
				},
			},
			isReady: false,
		},
	}

	for _, tc := range cases {
		if e, a := tc.isReady, tc.status.IsReady(); e != a {
			t.Errorf("%q expected: %v got: %v", tc.name, e, a)
		}
	}
}

/*
func TestMismatchedConditions(t *testing.T) {
	rs := RevisionStatus{
		Conditions: []RevisionCondition{
			{
				Type:    RevisionConditionBuildComplete,
				Status:  corev1.ConditionTrue,
				Message: "existing",
			},
		},
	}

	rc := &RevisionCondition{
		Type:    RevisionConditionReady,
		Status:  corev1.ConditionTrue,
		Message: "new one",
	}
	// Trying to add ConditionReady or update BuildComplete but screw up
	// the type here, or in the above definition.
	// End result is is that only the new one is added
	rs.SetCondition(RevisionConditionBuildComplete, rc)
	if len(rs.Conditions) != 2 {
		t.Errorf("BUG?: %+v", rs.Conditions)
	}
}
*/

func TestGetSetCondition(t *testing.T) {
	rs := RevisionStatus{}
	if a := rs.GetCondition(RevisionConditionReady); a != nil {
		t.Errorf("empty RevisionStatus returned %v when expected nil", a)
	}

	rc := &RevisionCondition{
		Type:   RevisionConditionBuildComplete,
		Status: corev1.ConditionTrue,
	}
	// Set Condition and make sure it's the only thing returned
	rs.SetCondition(RevisionConditionBuildComplete, rc)
	if e, a := rc, rs.GetCondition(RevisionConditionBuildComplete); !reflect.DeepEqual(e, a) {
		t.Errorf("GetCondition expected %v got: %v", e, a)
	}
	if a := rs.GetCondition(RevisionConditionReady); a != nil {
		t.Errorf("GetCondition expected nil got: %v", a)
	}
	// Remove and make sure it's no longer there
	rs.RemoveCondition(RevisionConditionBuildComplete)
	if a := rs.GetCondition(RevisionConditionBuildComplete); a != nil {
		t.Errorf("empty RevisionStatus returned %v when expected nil", a)
	}

}
<commit_msg>oops, did not see request to remove this (#111)<commit_after>/*
Copyright 2018 Google, Inc. All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package v1alpha1

import (
	"reflect"
	"testing"

	corev1 "k8s.io/api/core/v1"
)

const (
	generation int64 = 5
)

func testGeneration(t *testing.T) {
	r := Revision{}
	if a := r.GetGeneration(); a != 0 {
		t.Errorf("empty revision generation should be 0 was: %d", a)
	}

	r.SetGeneration(5)
	if e, a := generation, r.GetGeneration(); e != a {
		t.Errorf("getgeneration mismatch expected: %d got: %d", e, a)
	}

}

func TestIsReady(t *testing.T) {
	cases := []struct {
		name    string
		status  RevisionStatus
		isReady bool
	}{
		{
			name:    "empty status should not be ready",
			status:  RevisionStatus{},
			isReady: false,
		},
		{
			name: "Different condition type should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionBuildComplete,
						Status: corev1.ConditionTrue,
					},
				},
			},
			isReady: false,
		},
		{
			name: "False condition status should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionFalse,
					},
				},
			},
			isReady: false,
		},
		{
			name: "Unknown condition status should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionUnknown,
					},
				},
			},
			isReady: false,
		},
		{
			name: "Missing condition status should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type: RevisionConditionReady,
					},
				},
			},
			isReady: false,
		},
		{
			name: "True condition status should be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionTrue,
					},
				},
			},
			isReady: true,
		},
		{
			name: "Multiple conditions with ready status should be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionBuildComplete,
						Status: corev1.ConditionTrue,
					},
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionTrue,
					},
				},
			},
			isReady: true,
		},
		{
			name: "Multiple conditions with ready status false should not be ready",
			status: RevisionStatus{
				Conditions: []RevisionCondition{
					{
						Type:   RevisionConditionBuildComplete,
						Status: corev1.ConditionTrue,
					},
					{
						Type:   RevisionConditionReady,
						Status: corev1.ConditionFalse,
					},
				},
			},
			isReady: false,
		},
	}

	for _, tc := range cases {
		if e, a := tc.isReady, tc.status.IsReady(); e != a {
			t.Errorf("%q expected: %v got: %v", tc.name, e, a)
		}
	}
}

func TestGetSetCondition(t *testing.T) {
	rs := RevisionStatus{}
	if a := rs.GetCondition(RevisionConditionReady); a != nil {
		t.Errorf("empty RevisionStatus returned %v when expected nil", a)
	}

	rc := &RevisionCondition{
		Type:   RevisionConditionBuildComplete,
		Status: corev1.ConditionTrue,
	}
	// Set Condition and make sure it's the only thing returned
	rs.SetCondition(RevisionConditionBuildComplete, rc)
	if e, a := rc, rs.GetCondition(RevisionConditionBuildComplete); !reflect.DeepEqual(e, a) {
		t.Errorf("GetCondition expected %v got: %v", e, a)
	}
	if a := rs.GetCondition(RevisionConditionReady); a != nil {
		t.Errorf("GetCondition expected nil got: %v", a)
	}
	// Remove and make sure it's no longer there
	rs.RemoveCondition(RevisionConditionBuildComplete)
	if a := rs.GetCondition(RevisionConditionBuildComplete); a != nil {
		t.Errorf("empty RevisionStatus returned %v when expected nil", a)
	}

}
<|endoftext|>


<commit_before>package install

import (
	"context"
	"flag"
	"fmt"
	"strings"

	mf "github.com/jcrossley3/manifestival"
	servingv1alpha1 "github.com/openshift-knative/knative-serving-operator/pkg/apis/serving/v1alpha1"
	"github.com/openshift-knative/knative-serving-operator/version"
	configv1 "github.com/openshift/api/config/v1"

	"github.com/operator-framework/operator-sdk/pkg/k8sutil"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
	logf "sigs.k8s.io/controller-runtime/pkg/runtime/log"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

var (
	filename = flag.String("filename", "deploy/resources",
		"The filename containing the YAML resources to apply")
	recursive = flag.Bool("recursive", false,
		"If filename is a directory, process all manifests recursively")
	autoinstall = flag.Bool("install", false,
		"Automatically creates an Install resource if none exist")
	olm = flag.Bool("olm", false,
		"Ignores resources managed by the Operator Lifecycle Manager")
	namespace = flag.String("namespace", "",
		"Overrides the hard-coded namespace references in the manifest")
	log = logf.Log.WithName("controller_install")
)

// Add creates a new Install Controller and adds it to the Manager. The Manager will set fields on the Controller
// and Start it when the Manager is Started.
func Add(mgr manager.Manager) error {
	manifest, err := mf.NewYamlManifest(*filename, *recursive, mgr.GetClient())
	if err != nil {
		return err
	}
	return add(mgr, newReconciler(mgr, manifest))
}

// newReconciler returns a new reconcile.Reconciler
func newReconciler(mgr manager.Manager, man mf.Manifest) reconcile.Reconciler {
	return &ReconcileInstall{client: mgr.GetClient(), scheme: mgr.GetScheme(), config: man}
}

// add adds a new Controller to mgr with r as the reconcile.Reconciler
func add(mgr manager.Manager, r reconcile.Reconciler) error {
	// Register scheme
	if err := configv1.Install(mgr.GetScheme()); err != nil {
		log.Error(err, "Unable to register scheme")
	}

	// Create a new controller
	c, err := controller.New("install-controller", mgr, controller.Options{Reconciler: r})
	if err != nil {
		return err
	}

	// Watch for changes to primary resource Install
	err = c.Watch(&source.Kind{Type: &servingv1alpha1.Install{}}, &handler.EnqueueRequestForObject{})
	if err != nil {
		return err
	}

	// Auto-create Install
	if *autoinstall {
		ns, _ := k8sutil.GetWatchNamespace()
		go autoInstall(mgr.GetClient(), ns)
	}
	return nil
}

var _ reconcile.Reconciler = &ReconcileInstall{}

// ReconcileInstall reconciles a Install object
type ReconcileInstall struct {
	// This client, initialized using mgr.Client() above, is a split client
	// that reads objects from the cache and writes to the apiserver
	client client.Client
	scheme *runtime.Scheme
	config mf.Manifest
}

// Reconcile reads that state of the cluster for a Install object and makes changes based on the state read
// and what is in the Install.Spec
// Note:
// The Controller will requeue the Request to be processed again if the returned error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileInstall) Reconcile(request reconcile.Request) (reconcile.Result, error) {
	reqLogger := log.WithValues("Request.Namespace", request.Namespace, "Request.Name", request.Name)
	reqLogger.Info("Reconciling Install")

	// Fetch the Install instance
	instance := &servingv1alpha1.Install{}
	err := r.client.Get(context.TODO(), request.NamespacedName, instance)
	if err != nil {
		if errors.IsNotFound(err) {
			r.config.DeleteAll()
			return reconcile.Result{}, nil
		}
		// Error reading the object - requeue the request.
		return reconcile.Result{}, err
	}

	stages := []func(*servingv1alpha1.Install) error{
		r.install,
		r.deleteObsoleteResources,
		r.checkForMinikube,
		r.updateServiceNetwork,
		r.updateDomain,
		r.configure,
	}

	for _, stage := range stages {
		if err := stage(instance); err != nil {
			return reconcile.Result{}, err
		}
	}

	return reconcile.Result{}, nil
}

// Apply the embedded resources
func (r *ReconcileInstall) install(instance *servingv1alpha1.Install) error {
	// Filter resources as appropriate
	filters := []mf.FilterFn{mf.ByOwner(instance)}
	switch {
	case *olm:
		sa, err := k8sutil.GetOperatorName()
		if err != nil {
			return err
		}
		filters = append(filters,
			mf.ByOLM,
			mf.ByNamespace(instance.GetNamespace()),
			mf.ByServiceAccount(sa))
	case len(*namespace) > 0:
		filters = append(filters, mf.ByNamespace(*namespace))
	}
	r.config.Filter(filters...)

	if instance.Status.Version == version.Version {
		// we've already successfully applied our YAML
		return nil
	}
	// Apply the resources in the YAML file
	if err := r.config.ApplyAll(); err != nil {
		return err
	}

	// Update status
	instance.Status.Resources = r.config.ResourceNames()
	instance.Status.Version = version.Version
	if err := r.client.Status().Update(context.TODO(), instance); err != nil {
		return err
	}
	return nil
}

// Set ConfigMap values from Install spec
func (r *ReconcileInstall) configure(instance *servingv1alpha1.Install) error {
	for suffix, config := range instance.Spec.Config {
		name := "config-" + suffix
		cm := r.config.Find("v1", "ConfigMap", name)
		if cm == nil {
			log.Error(fmt.Errorf("ConfigMap '%s' not found", name), "Invalid Install spec")
			continue
		}
		for k, v := range config {
			unstructured.SetNestedField(cm.Object, v, "data", k)
		}
		if err := r.config.Apply(cm); err != nil {
			return err
		}
	}
	return nil
}

// Delete obsolete istio-system resources, if any
func (r *ReconcileInstall) deleteObsoleteResources(instance *servingv1alpha1.Install) error {
	resource := &unstructured.Unstructured{}
	resource.SetNamespace("istio-system")
	resource.SetName("knative-ingressgateway")
	resource.SetAPIVersion("v1")
	resource.SetKind("Service")
	if err := r.config.Delete(resource); err != nil {
		return err
	}
	resource.SetAPIVersion("apps/v1")
	resource.SetKind("Deployment")
	if err := r.config.Delete(resource); err != nil {
		return err
	}
	resource.SetAPIVersion("autoscaling/v1")
	resource.SetKind("HorizontalPodAutoscaler")
	return r.config.Delete(resource)
}

// Configure minikube if we're soaking in it
func (r *ReconcileInstall) checkForMinikube(instance *servingv1alpha1.Install) error {
	node := &v1.Node{}
	err := r.client.Get(context.TODO(), types.NamespacedName{Name: "minikube"}, node)
	if err != nil {
		if errors.IsNotFound(err) {
			return nil // not running on minikube!
		}
		return err
	}

	cm, err := r.config.Get(r.config.Find("v1", "ConfigMap", "config-network"))
	if err != nil {
		return err
	}
	if cm == nil {
		log.Error(err, "Missing ConfigMap", "name", "config-network")
		return nil // no sense in trying if the CM is gone
	}
	const k, v = "istio.sidecar.includeOutboundIPRanges", "10.0.0.1/24"
	if _, found, _ := unstructured.NestedString(cm.Object, "data", k); found {
		log.V(1).Info("Detected minikube; egress already configured", k, v)
		return nil // already set
	}
	log.Info("Detected minikube; configuring egress", k, v)
	unstructured.SetNestedField(cm.Object, v, "data", k)
	return r.client.Update(context.TODO(), cm)

}

// Get Service Network from cluster resource
func (r *ReconcileInstall) getServiceNetwork() string {
	networkConfig := &configv1.Network{}
	serviceNetwork := ""
	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: "cluster"}, networkConfig); err != nil {
		log.V(1).Info("OpenShift Network Config is not available.")
	} else if len(networkConfig.Spec.ServiceNetwork) > 0 {
		serviceNetwork = strings.Join(networkConfig.Spec.ServiceNetwork, ",")
		log.Info("OpenShift Network Config is available", "Service Network", serviceNetwork)
	}
	return serviceNetwork
}

func (r *ReconcileInstall) getDomain() string {
	ingressConfig := &configv1.Ingress{}
	domain := ""
	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: "cluster"}, ingressConfig); err != nil {
		log.V(1).Info("OpenShift Ingress Config is not available.")
	} else {
		domain = ingressConfig.Spec.Domain
		log.Info("OpenShift Ingress Config is available", "Domain", domain)
	}

	return domain
}

// Set domain in the Config Map
func (r *ReconcileInstall) updateDomain(instance *servingv1alpha1.Install) error {

	// retrieve domain for configuring for ingress traffic
	domain := r.getDomain()

	// If domain is available, update config-domain config map
	if len(domain) > 0 {

		cm := &v1.ConfigMap{}
		u := r.config.Find("v1", "ConfigMap", "config-domain")
		key := types.NamespacedName{Namespace: u.GetNamespace(), Name: u.GetName()}
		if err := r.client.Get(context.TODO(), key, cm); err != nil {
			return err
		}
		cm.Data[domain] = ""
		return r.client.Update(context.TODO(), cm)
	}

	return nil
}

// Set istio.sidecar.includeOutboundIPRanges property with service network
func (r *ReconcileInstall) updateServiceNetwork(instance *servingv1alpha1.Install) error {

	// retrieve service networks for configuring egress traffic
	serviceNetwork := r.getServiceNetwork()

	// If service network is available, update config-network config map
	if len(serviceNetwork) > 0 {

		cm := &v1.ConfigMap{}
		u := r.config.Find("v1", "ConfigMap", "config-network")
		key := types.NamespacedName{Namespace: u.GetNamespace(), Name: u.GetName()}
		if err := r.client.Get(context.TODO(), key, cm); err != nil {
			return err
		}
		cm.Data["istio.sidecar.includeOutboundIPRanges"] = serviceNetwork
		return r.client.Update(context.TODO(), cm)

	}

	return nil
}

func autoInstall(c client.Client, ns string) error {
	installList := &servingv1alpha1.InstallList{}
	err := c.List(context.TODO(), &client.ListOptions{Namespace: ns}, installList)
	if err != nil {
		log.Error(err, "Unable to list Installs")
		return err
	}
	if len(installList.Items) == 0 {
		install := &servingv1alpha1.Install{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "auto-install",
				Namespace: ns,
			},
		}
		err = c.Create(context.TODO(), install)
		if err != nil {
			log.Error(err, "Unable to create Install")
		}
	}
	return nil
}
<commit_msg>Apparently, we've never fetched Installs correctly<commit_after>package install

import (
	"context"
	"flag"
	"fmt"
	"strings"

	mf "github.com/jcrossley3/manifestival"
	servingv1alpha1 "github.com/openshift-knative/knative-serving-operator/pkg/apis/serving/v1alpha1"
	"github.com/openshift-knative/knative-serving-operator/version"
	configv1 "github.com/openshift/api/config/v1"

	"github.com/operator-framework/operator-sdk/pkg/k8sutil"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
	logf "sigs.k8s.io/controller-runtime/pkg/runtime/log"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

var (
	filename = flag.String("filename", "deploy/resources",
		"The filename containing the YAML resources to apply")
	recursive = flag.Bool("recursive", false,
		"If filename is a directory, process all manifests recursively")
	autoinstall = flag.Bool("install", false,
		"Automatically creates an Install resource if none exist")
	olm = flag.Bool("olm", false,
		"Ignores resources managed by the Operator Lifecycle Manager")
	namespace = flag.String("namespace", "",
		"Overrides the hard-coded namespace references in the manifest")
	log = logf.Log.WithName("controller_install")
)

// Add creates a new Install Controller and adds it to the Manager. The Manager will set fields on the Controller
// and Start it when the Manager is Started.
func Add(mgr manager.Manager) error {
	manifest, err := mf.NewYamlManifest(*filename, *recursive, mgr.GetClient())
	if err != nil {
		return err
	}
	return add(mgr, newReconciler(mgr, manifest))
}

// newReconciler returns a new reconcile.Reconciler
func newReconciler(mgr manager.Manager, man mf.Manifest) reconcile.Reconciler {
	return &ReconcileInstall{client: mgr.GetClient(), scheme: mgr.GetScheme(), config: man}
}

// add adds a new Controller to mgr with r as the reconcile.Reconciler
func add(mgr manager.Manager, r reconcile.Reconciler) error {
	// Register scheme
	if err := configv1.Install(mgr.GetScheme()); err != nil {
		log.Error(err, "Unable to register scheme")
	}

	// Create a new controller
	c, err := controller.New("install-controller", mgr, controller.Options{Reconciler: r})
	if err != nil {
		return err
	}

	// Watch for changes to primary resource Install
	err = c.Watch(&source.Kind{Type: &servingv1alpha1.Install{}}, &handler.EnqueueRequestForObject{})
	if err != nil {
		return err
	}

	// Make an attempt to auto-create an Install CR
	if *autoinstall {
		ns, _ := k8sutil.GetWatchNamespace()
		c, _ := client.New(mgr.GetConfig(), client.Options{})
		go autoInstall(c, ns)
	}
	return nil
}

var _ reconcile.Reconciler = &ReconcileInstall{}

// ReconcileInstall reconciles a Install object
type ReconcileInstall struct {
	// This client, initialized using mgr.Client() above, is a split client
	// that reads objects from the cache and writes to the apiserver
	client client.Client
	scheme *runtime.Scheme
	config mf.Manifest
}

// Reconcile reads that state of the cluster for a Install object and makes changes based on the state read
// and what is in the Install.Spec
// Note:
// The Controller will requeue the Request to be processed again if the returned error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileInstall) Reconcile(request reconcile.Request) (reconcile.Result, error) {
	reqLogger := log.WithValues("Request.Namespace", request.Namespace, "Request.Name", request.Name)
	reqLogger.Info("Reconciling Install")

	// Fetch the Install instance
	instance := &servingv1alpha1.Install{}
	err := r.client.Get(context.TODO(), request.NamespacedName, instance)
	if err != nil {
		if errors.IsNotFound(err) {
			r.config.DeleteAll()
			return reconcile.Result{}, nil
		}
		// Error reading the object - requeue the request.
		return reconcile.Result{}, err
	}

	stages := []func(*servingv1alpha1.Install) error{
		r.install,
		r.deleteObsoleteResources,
		r.checkForMinikube,
		r.updateServiceNetwork,
		r.updateDomain,
		r.configure,
	}

	for _, stage := range stages {
		if err := stage(instance); err != nil {
			return reconcile.Result{}, err
		}
	}

	return reconcile.Result{}, nil
}

// Apply the embedded resources
func (r *ReconcileInstall) install(instance *servingv1alpha1.Install) error {
	// Filter resources as appropriate
	filters := []mf.FilterFn{mf.ByOwner(instance)}
	switch {
	case *olm:
		sa, err := k8sutil.GetOperatorName()
		if err != nil {
			return err
		}
		filters = append(filters,
			mf.ByOLM,
			mf.ByNamespace(instance.GetNamespace()),
			mf.ByServiceAccount(sa))
	case len(*namespace) > 0:
		filters = append(filters, mf.ByNamespace(*namespace))
	}
	r.config.Filter(filters...)

	if instance.Status.Version == version.Version {
		// we've already successfully applied our YAML
		return nil
	}
	// Apply the resources in the YAML file
	if err := r.config.ApplyAll(); err != nil {
		return err
	}

	// Update status
	instance.Status.Resources = r.config.ResourceNames()
	instance.Status.Version = version.Version
	if err := r.client.Status().Update(context.TODO(), instance); err != nil {
		return err
	}
	return nil
}

// Set ConfigMap values from Install spec
func (r *ReconcileInstall) configure(instance *servingv1alpha1.Install) error {
	for suffix, config := range instance.Spec.Config {
		name := "config-" + suffix
		cm := r.config.Find("v1", "ConfigMap", name)
		if cm == nil {
			log.Error(fmt.Errorf("ConfigMap '%s' not found", name), "Invalid Install spec")
			continue
		}
		for k, v := range config {
			unstructured.SetNestedField(cm.Object, v, "data", k)
		}
		if err := r.config.Apply(cm); err != nil {
			return err
		}
	}
	return nil
}

// Delete obsolete istio-system resources, if any
func (r *ReconcileInstall) deleteObsoleteResources(instance *servingv1alpha1.Install) error {
	resource := &unstructured.Unstructured{}
	resource.SetNamespace("istio-system")
	resource.SetName("knative-ingressgateway")
	resource.SetAPIVersion("v1")
	resource.SetKind("Service")
	if err := r.config.Delete(resource); err != nil {
		return err
	}
	resource.SetAPIVersion("apps/v1")
	resource.SetKind("Deployment")
	if err := r.config.Delete(resource); err != nil {
		return err
	}
	resource.SetAPIVersion("autoscaling/v1")
	resource.SetKind("HorizontalPodAutoscaler")
	return r.config.Delete(resource)
}

// Configure minikube if we're soaking in it
func (r *ReconcileInstall) checkForMinikube(instance *servingv1alpha1.Install) error {
	node := &v1.Node{}
	err := r.client.Get(context.TODO(), types.NamespacedName{Name: "minikube"}, node)
	if err != nil {
		if errors.IsNotFound(err) {
			return nil // not running on minikube!
		}
		return err
	}

	cm, err := r.config.Get(r.config.Find("v1", "ConfigMap", "config-network"))
	if err != nil {
		return err
	}
	if cm == nil {
		log.Error(err, "Missing ConfigMap", "name", "config-network")
		return nil // no sense in trying if the CM is gone
	}
	const k, v = "istio.sidecar.includeOutboundIPRanges", "10.0.0.1/24"
	if _, found, _ := unstructured.NestedString(cm.Object, "data", k); found {
		log.V(1).Info("Detected minikube; egress already configured", k, v)
		return nil // already set
	}
	log.Info("Detected minikube; configuring egress", k, v)
	unstructured.SetNestedField(cm.Object, v, "data", k)
	return r.client.Update(context.TODO(), cm)

}

// Get Service Network from cluster resource
func (r *ReconcileInstall) getServiceNetwork() string {
	networkConfig := &configv1.Network{}
	serviceNetwork := ""
	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: "cluster"}, networkConfig); err != nil {
		log.V(1).Info("OpenShift Network Config is not available.")
	} else if len(networkConfig.Spec.ServiceNetwork) > 0 {
		serviceNetwork = strings.Join(networkConfig.Spec.ServiceNetwork, ",")
		log.Info("OpenShift Network Config is available", "Service Network", serviceNetwork)
	}
	return serviceNetwork
}

func (r *ReconcileInstall) getDomain() string {
	ingressConfig := &configv1.Ingress{}
	domain := ""
	if err := r.client.Get(context.TODO(), types.NamespacedName{Name: "cluster"}, ingressConfig); err != nil {
		log.V(1).Info("OpenShift Ingress Config is not available.")
	} else {
		domain = ingressConfig.Spec.Domain
		log.Info("OpenShift Ingress Config is available", "Domain", domain)
	}

	return domain
}

// Set domain in the Config Map
func (r *ReconcileInstall) updateDomain(instance *servingv1alpha1.Install) error {

	// retrieve domain for configuring for ingress traffic
	domain := r.getDomain()

	// If domain is available, update config-domain config map
	if len(domain) > 0 {

		cm := &v1.ConfigMap{}
		u := r.config.Find("v1", "ConfigMap", "config-domain")
		key := types.NamespacedName{Namespace: u.GetNamespace(), Name: u.GetName()}
		if err := r.client.Get(context.TODO(), key, cm); err != nil {
			return err
		}
		cm.Data[domain] = ""
		return r.client.Update(context.TODO(), cm)
	}

	return nil
}

// Set istio.sidecar.includeOutboundIPRanges property with service network
func (r *ReconcileInstall) updateServiceNetwork(instance *servingv1alpha1.Install) error {

	// retrieve service networks for configuring egress traffic
	serviceNetwork := r.getServiceNetwork()

	// If service network is available, update config-network config map
	if len(serviceNetwork) > 0 {

		cm := &v1.ConfigMap{}
		u := r.config.Find("v1", "ConfigMap", "config-network")
		key := types.NamespacedName{Namespace: u.GetNamespace(), Name: u.GetName()}
		if err := r.client.Get(context.TODO(), key, cm); err != nil {
			return err
		}
		cm.Data["istio.sidecar.includeOutboundIPRanges"] = serviceNetwork
		return r.client.Update(context.TODO(), cm)

	}

	return nil
}

func autoInstall(c client.Client, ns string) error {
	installList := &servingv1alpha1.InstallList{}
	err := c.List(context.TODO(), &client.ListOptions{Namespace: ns}, installList)
	if err != nil {
		log.Error(err, "Unable to list Installs")
		return err
	}
	if len(installList.Items) == 0 {
		install := &servingv1alpha1.Install{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "auto-install",
				Namespace: ns,
			},
		}
		err = c.Create(context.TODO(), install)
		if err != nil {
			log.Error(err, "Unable to create Install")
		}
	}
	return nil
}
<|endoftext|>