<commit_before>package stepman

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-pathutil/pathutil"
	"github.com/bitrise-io/stepman/models"
	"gopkg.in/yaml.v2"
)

// DebugMode ...
var DebugMode bool

func parseStepYml(collectionURI, pth, id, version string) (models.StepModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepModel{}, err
	}

	var stepModel models.StepModel
	if err := yaml.Unmarshal(bytes, &stepModel); err != nil {
		return models.StepModel{}, err
	}

	if err := stepModel.Normalize(); err != nil {
		return models.StepModel{}, err
	}

	if err := stepModel.Validate(); err != nil {
		return models.StepModel{}, err
	}

	if err := stepModel.FillMissingDeafults(); err != nil {
		return models.StepModel{}, err
	}

	return stepModel, nil
}

// ParseStepCollection ...
func ParseStepCollection(pth string) (models.StepCollectionModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	if err := yaml.Unmarshal(bytes, &stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, nil
}

// DownloadStep ...
func DownloadStep(collection models.StepCollectionModel, id, version string) error {
	log.Debugf("Download Step: %#v (%#v)\n", id, version)
	downloadLocations, err := collection.GetDownloadLocations(id, version)
	if err != nil {
		return err
	}

	stepPth := GetStepCacheDirPath(collection.SteplibSource, id, version)
	if exist, err := pathutil.IsPathExists(stepPth); err != nil {
		return err
	} else if exist {
		log.Info("[STEPMAN] - Step already downloaded")
		return nil
	}

	success := false
	for _, downloadLocation := range downloadLocations {
		switch downloadLocation.Type {
		case "zip":
			log.Info("[STEPMAN] - Downloading step from:", downloadLocation.Src)
			if err := DownloadAndUnZIP(downloadLocation.Src, stepPth); err != nil {
				log.Error("[STEPMAN] - Failed to download step.zip:", err)
			} else {
				success = true
				return nil
			}
		case "git":
			log.Info("[STEPMAN] - Git clone step from:", downloadLocation.Src)
			if err := DoGitCloneWithVersion(downloadLocation.Src, stepPth, version); err != nil {
				log.Errorf("[STEPMAN] - Failed to clone step (%s): %v", downloadLocation.Src, err)
			} else {
				success = true
				return nil
			}
		default:
			return fmt.Errorf("[STEPMAN] - Failed to download: Invalid download location (%#v) for step %#v (%#v)", downloadLocation, id, version)
		}
	}

	if !success {
		return errors.New("Failed to download step")
	}
	return nil
}

// GetStepCacheDirPath ...
// Step's Cache dir path, where it's code lives.
func GetStepCacheDirPath(collectionURI string, id, version string) string {
	return GetCacheBaseDir(collectionURI) + "/" + id + "/" + version
}

// GetStepCollectionDirPath ...
// Step's Collection dir path, where it's spec (step.yml) lives.
func GetStepCollectionDirPath(collectionURI string, id, version string) string {
	return GetCollectionBaseDirPath(collectionURI) + "/steps/" + id + "/" + version
}

func addStepVersionToStepGroup(step models.StepModel, version string, stepGroup models.StepGroupModel) (models.StepGroupModel, error) {
	if stepGroup.LatestVersionNumber != "" {
		r, err := models.CompareVersions(stepGroup.LatestVersionNumber, version)
		if err != nil {
			return models.StepGroupModel{}, err
		}
		if r == 1 {
			stepGroup.LatestVersionNumber = version
		}
	} else {
		stepGroup.LatestVersionNumber = version
	}
	log.Debugf("SetGroup: %#v, versionParam: %#v, stepParam: %#v", stepGroup, version, step)
	stepGroup.Versions[version] = step
	return stepGroup, nil
}

func generateFormattedJSONForStepsSpec(collectionURI string, templateCollection models.StepCollectionModel) ([]byte, error) {
	collection := models.StepCollectionModel{
		FormatVersion:        templateCollection.FormatVersion,
		GeneratedAtTimeStamp: time.Now().Unix(),
		SteplibSource:        collectionURI,
		DownloadLocations:    templateCollection.DownloadLocations,
	}

	stepHash := models.StepHash{}

	stepsSpecDir := GetCollectionBaseDirPath(collectionURI)
	log.Debugln("  stepsSpecDir: ", stepsSpecDir)
	err := filepath.Walk(stepsSpecDir, func(path string, f os.FileInfo, err error) error {
		truncatedPath := strings.Replace(path, stepsSpecDir+"/", "", -1)
		match, matchErr := regexp.MatchString("([a-z]+).yml", truncatedPath)
		if matchErr != nil {
			return matchErr
		}

		if match {
			components := strings.Split(truncatedPath, "/")
			if len(components) == 4 {
				id := components[1]
				version := components[2]

				log.Debugf("Start parsing (StepId:%s) (Version:%s)", id, version)
				step, parseErr := parseStepYml(collectionURI, path, id, version)
				if parseErr != nil {
					log.Debugf("  Failed to parse StepId: %v Version: %v", id, version)
					return parseErr
				}
				stepGroup, found := stepHash[id]
				if !found {
					stepGroup = models.StepGroupModel{
						Versions: map[string]models.StepModel{},
					}
				}
				stepGroup, err = addStepVersionToStepGroup(step, version, stepGroup)
				if err != nil {
					log.Debugf("  Failed to add step to step-group. (StepId:%v) (Version: %v) | Error: %v", id, version, err)
				}

				stepHash[id] = stepGroup
			} else {
				log.Debug("  * Path:", truncatedPath)
				log.Debug("  * Legth:", len(components))
			}
		}

		return err
	})
	if err != nil {
		log.Error("[STEPMAN] - Failed to walk through path:", err)
	}

	// log.Debugf("  collected steps: %#v\n", stepHash)
	collection.Steps = stepHash

	var bytes []byte
	// if DebugMode == true {
	bytes, err = json.MarshalIndent(collection, "", "\t")
	// } else {
	// 	bytes, err = json.Marshal(collection)
	// }
	if err != nil {
		log.Error("[STEPMAN] - Failed to parse json:", err)
		return []byte{}, err
	}

	return bytes, nil
}

// WriteStepSpecToFile ...
func WriteStepSpecToFile(collectionURI string, templateCollection models.StepCollectionModel) error {
	pth := GetStepSpecPath(collectionURI)

	if exist, err := pathutil.IsPathExists(pth); err != nil {
		log.Error("[STEPMAN] - Failed to check path:", err)
		return err
	} else if !exist {
		dir, _ := path.Split(pth)
		err := os.MkdirAll(dir, 0777)
		if err != nil {
			return err
		}
	} else {
		err := os.Remove(pth)
		if err != nil {
			return err
		}
	}

	file, err := os.OpenFile(pth, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		return err
	}
	defer func() {
		err := file.Close()
		if err != nil {
			log.Error("[STEPMAN] - Failed to close file:", err)
		}
	}()

	jsonContBytes, err := generateFormattedJSONForStepsSpec(collectionURI, templateCollection)
	if err != nil {
		return err
	}

	_, err = file.Write(jsonContBytes)
	if err != nil {
		return err
	}

	return nil
}

// ReadStepSpec ...
func ReadStepSpec(collectionURI string) (models.StepCollectionModel, error) {
	log.Debugln("-> ReadStepSpec: ", collectionURI)

	pth := GetStepSpecPath(collectionURI)
	file, err := os.Open(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	parser := json.NewDecoder(file)
	if err = parser.Decode(&stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, err
}

// DownloadAndUnZIP ...
func DownloadAndUnZIP(url, pth string) error {
	filePath := os.TempDir() + "step.zip"
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil {
			log.Fatal("Failed to close file:", err)
		}
		if err := os.Remove(filePath); err != nil {
			log.Fatal("Failed to remove file:", err)
		}
	}()

	response, err := http.Get(url)
	if err != nil {
		return err
	}
	defer func() {
		if err := response.Body.Close(); err != nil {
			log.Fatal("Failed to close response body:", err)
		}
	}()

	if response.StatusCode != http.StatusOK {
		errorMsg := "Failed to download step.zip from: " + url
		return errors.New(errorMsg)
	}

	log.Info("Successfully downloaded step.zip")
	if _, err := io.Copy(file, response.Body); err != nil {
		return err
	}

	return unzip(filePath, pth)
}

func unzip(src, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer func() {
		if err := r.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	if err := os.MkdirAll(dest, 0755); err != nil {
		return err
	}

	// Closure to address file descriptors issue with all the deferred .Close() methods
	extractAndWriteFile := func(f *zip.File) error {
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer func() {
			if err := rc.Close(); err != nil {
				log.Fatal(err)
			}
		}()

		path := filepath.Join(dest, f.Name)

		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(path, f.Mode()); err != nil {
				return err
			}
		} else {
			f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
			if err != nil {
				return err
			}
			defer func() {
				if err := f.Close(); err != nil {
					log.Fatal(err)
				}
			}()

			if _, err = io.Copy(f, rc); err != nil {
				return err
			}
		}
		return nil
	}

	for _, f := range r.File {
		if err := extractAndWriteFile(f); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>added missing return err<commit_after>package stepman

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-pathutil/pathutil"
	"github.com/bitrise-io/stepman/models"
	"gopkg.in/yaml.v2"
)

// DebugMode ...
var DebugMode bool

func parseStepYml(collectionURI, pth, id, version string) (models.StepModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepModel{}, err
	}

	var stepModel models.StepModel
	if err := yaml.Unmarshal(bytes, &stepModel); err != nil {
		return models.StepModel{}, err
	}

	if err := stepModel.Normalize(); err != nil {
		return models.StepModel{}, err
	}

	if err := stepModel.Validate(); err != nil {
		return models.StepModel{}, err
	}

	if err := stepModel.FillMissingDeafults(); err != nil {
		return models.StepModel{}, err
	}

	return stepModel, nil
}

// ParseStepCollection ...
func ParseStepCollection(pth string) (models.StepCollectionModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	if err := yaml.Unmarshal(bytes, &stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, nil
}

// DownloadStep ...
func DownloadStep(collection models.StepCollectionModel, id, version string) error {
	log.Debugf("Download Step: %#v (%#v)\n", id, version)
	downloadLocations, err := collection.GetDownloadLocations(id, version)
	if err != nil {
		return err
	}

	stepPth := GetStepCacheDirPath(collection.SteplibSource, id, version)
	if exist, err := pathutil.IsPathExists(stepPth); err != nil {
		return err
	} else if exist {
		log.Info("[STEPMAN] - Step already downloaded")
		return nil
	}

	success := false
	for _, downloadLocation := range downloadLocations {
		switch downloadLocation.Type {
		case "zip":
			log.Info("[STEPMAN] - Downloading step from:", downloadLocation.Src)
			if err := DownloadAndUnZIP(downloadLocation.Src, stepPth); err != nil {
				log.Error("[STEPMAN] - Failed to download step.zip:", err)
			} else {
				success = true
				return nil
			}
		case "git":
			log.Info("[STEPMAN] - Git clone step from:", downloadLocation.Src)
			if err := DoGitCloneWithVersion(downloadLocation.Src, stepPth, version); err != nil {
				log.Errorf("[STEPMAN] - Failed to clone step (%s): %v", downloadLocation.Src, err)
			} else {
				success = true
				return nil
			}
		default:
			return fmt.Errorf("[STEPMAN] - Failed to download: Invalid download location (%#v) for step %#v (%#v)", downloadLocation, id, version)
		}
	}

	if !success {
		return errors.New("Failed to download step")
	}
	return nil
}

// GetStepCacheDirPath ...
// Step's Cache dir path, where it's code lives.
func GetStepCacheDirPath(collectionURI string, id, version string) string {
	return GetCacheBaseDir(collectionURI) + "/" + id + "/" + version
}

// GetStepCollectionDirPath ...
// Step's Collection dir path, where it's spec (step.yml) lives.
func GetStepCollectionDirPath(collectionURI string, id, version string) string {
	return GetCollectionBaseDirPath(collectionURI) + "/steps/" + id + "/" + version
}

func addStepVersionToStepGroup(step models.StepModel, version string, stepGroup models.StepGroupModel) (models.StepGroupModel, error) {
	if stepGroup.LatestVersionNumber != "" {
		r, err := models.CompareVersions(stepGroup.LatestVersionNumber, version)
		if err != nil {
			return models.StepGroupModel{}, err
		}
		if r == 1 {
			stepGroup.LatestVersionNumber = version
		}
	} else {
		stepGroup.LatestVersionNumber = version
	}
	log.Debugf("SetGroup: %#v, versionParam: %#v, stepParam: %#v", stepGroup, version, step)
	stepGroup.Versions[version] = step
	return stepGroup, nil
}

func generateFormattedJSONForStepsSpec(collectionURI string, templateCollection models.StepCollectionModel) ([]byte, error) {
	collection := models.StepCollectionModel{
		FormatVersion:        templateCollection.FormatVersion,
		GeneratedAtTimeStamp: time.Now().Unix(),
		SteplibSource:        collectionURI,
		DownloadLocations:    templateCollection.DownloadLocations,
	}

	stepHash := models.StepHash{}

	stepsSpecDir := GetCollectionBaseDirPath(collectionURI)
	log.Debugln("  stepsSpecDir: ", stepsSpecDir)
	err := filepath.Walk(stepsSpecDir, func(path string, f os.FileInfo, err error) error {
		truncatedPath := strings.Replace(path, stepsSpecDir+"/", "", -1)
		match, matchErr := regexp.MatchString("([a-z]+).yml", truncatedPath)
		if matchErr != nil {
			return matchErr
		}

		if match {
			components := strings.Split(truncatedPath, "/")
			if len(components) == 4 {
				id := components[1]
				version := components[2]

				log.Debugf("Start parsing (StepId:%s) (Version:%s)", id, version)
				step, parseErr := parseStepYml(collectionURI, path, id, version)
				if parseErr != nil {
					log.Debugf("  Failed to parse StepId: %v Version: %v", id, version)
					return parseErr
				}
				stepGroup, found := stepHash[id]
				if !found {
					stepGroup = models.StepGroupModel{
						Versions: map[string]models.StepModel{},
					}
				}
				stepGroup, err = addStepVersionToStepGroup(step, version, stepGroup)
				if err != nil {
					log.Debugf("  Failed to add step to step-group. (StepId:%v) (Version: %v) | Error: %v", id, version, err)
					return err
				}

				stepHash[id] = stepGroup
			} else {
				log.Debug("  * Path:", truncatedPath)
				log.Debug("  * Legth:", len(components))
			}
		}

		return err
	})
	if err != nil {
		log.Error("[STEPMAN] - Failed to walk through path:", err)
	}

	// log.Debugf("  collected steps: %#v\n", stepHash)
	collection.Steps = stepHash

	var bytes []byte
	// if DebugMode == true {
	bytes, err = json.MarshalIndent(collection, "", "\t")
	// } else {
	// 	bytes, err = json.Marshal(collection)
	// }
	if err != nil {
		log.Error("[STEPMAN] - Failed to parse json:", err)
		return []byte{}, err
	}

	return bytes, nil
}

// WriteStepSpecToFile ...
func WriteStepSpecToFile(collectionURI string, templateCollection models.StepCollectionModel) error {
	pth := GetStepSpecPath(collectionURI)

	if exist, err := pathutil.IsPathExists(pth); err != nil {
		log.Error("[STEPMAN] - Failed to check path:", err)
		return err
	} else if !exist {
		dir, _ := path.Split(pth)
		err := os.MkdirAll(dir, 0777)
		if err != nil {
			return err
		}
	} else {
		err := os.Remove(pth)
		if err != nil {
			return err
		}
	}

	file, err := os.OpenFile(pth, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		return err
	}
	defer func() {
		err := file.Close()
		if err != nil {
			log.Error("[STEPMAN] - Failed to close file:", err)
		}
	}()

	jsonContBytes, err := generateFormattedJSONForStepsSpec(collectionURI, templateCollection)
	if err != nil {
		return err
	}

	_, err = file.Write(jsonContBytes)
	if err != nil {
		return err
	}

	return nil
}

// ReadStepSpec ...
func ReadStepSpec(collectionURI string) (models.StepCollectionModel, error) {
	log.Debugln("-> ReadStepSpec: ", collectionURI)

	pth := GetStepSpecPath(collectionURI)
	file, err := os.Open(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	parser := json.NewDecoder(file)
	if err = parser.Decode(&stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, err
}

// DownloadAndUnZIP ...
func DownloadAndUnZIP(url, pth string) error {
	filePath := os.TempDir() + "step.zip"
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil {
			log.Fatal("Failed to close file:", err)
		}
		if err := os.Remove(filePath); err != nil {
			log.Fatal("Failed to remove file:", err)
		}
	}()

	response, err := http.Get(url)
	if err != nil {
		return err
	}
	defer func() {
		if err := response.Body.Close(); err != nil {
			log.Fatal("Failed to close response body:", err)
		}
	}()

	if response.StatusCode != http.StatusOK {
		errorMsg := "Failed to download step.zip from: " + url
		return errors.New(errorMsg)
	}

	log.Info("Successfully downloaded step.zip")
	if _, err := io.Copy(file, response.Body); err != nil {
		return err
	}

	return unzip(filePath, pth)
}

func unzip(src, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer func() {
		if err := r.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	if err := os.MkdirAll(dest, 0755); err != nil {
		return err
	}

	// Closure to address file descriptors issue with all the deferred .Close() methods
	extractAndWriteFile := func(f *zip.File) error {
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer func() {
			if err := rc.Close(); err != nil {
				log.Fatal(err)
			}
		}()

		path := filepath.Join(dest, f.Name)

		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(path, f.Mode()); err != nil {
				return err
			}
		} else {
			f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
			if err != nil {
				return err
			}
			defer func() {
				if err := f.Close(); err != nil {
					log.Fatal(err)
				}
			}()

			if _, err = io.Copy(f, rc); err != nil {
				return err
			}
		}
		return nil
	}

	for _, f := range r.File {
		if err := extractAndWriteFile(f); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package nakadi

import (
	"context"
	"testing"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestStreamAPI_startStream(t *testing.T) {
	retryCh := make(chan error, 1)
	okCh := make(chan struct{}, 1)

	stream := &mockStreamer{}
	_, opener, _ := setupMockStream(retryCh, okCh)

	opener.On("openStream").Once().Return(nil, assert.AnError)
	opener.unblock <- struct{}{}
	time.Sleep(50 * time.Millisecond)

	select {
	case err := <-retryCh:
		assert.EqualError(t, assert.AnError, err.Error())
	default:
		assert.Fail(t, "no retry error detected")
	}

	opener.On("openStream").Once().Return(stream, nil)
	opener.unblock <- struct{}{}
	time.Sleep(50 * time.Millisecond)

	select {
	case <-okCh:
		// nothing
	default:
		assert.Fail(t, "no retry ok detected")
	}
}

func TestStreamAPI_NextEvents(t *testing.T) {
	expectedCursor := Cursor{NakadiStreamID: "stream-id"}
	expectedEvents := []byte(`"events":[{"metadata":{"eid":"74450ab6-5461-11e7-9dd2-87c3afa8811f"})]`)

	t.Run("fail with error", func(t *testing.T) {
		stream := &mockStreamer{unblock: make(chan struct{}, 1)}
		streamAPI, opener, _ := setupMockStream(nil, nil)

		opener.On("openStream").Once().Return(stream, nil)
		opener.unblock <- struct{}{}
		stream.On("nextEvents").Once().Return(Cursor{}, nil, assert.AnError)
		stream.On("closeStream").Once().Return(nil)

		stream.unblock <- struct{}{}
		_, _, err := streamAPI.NextEvents()

		assert.EqualError(t, assert.AnError, err.Error())
	})

	t.Run("successful read events", func(t *testing.T) {
		stream := &mockStreamer{unblock: make(chan struct{}, 1)}
		streamAPI, opener, _ := setupMockStream(nil, nil)

		opener.On("openStream").Once().Return(stream, nil)
		opener.unblock <- struct{}{}
		stream.On("nextEvents").Twice().Return(expectedCursor, expectedEvents, nil)

		stream.unblock <- struct{}{}
		cursor, events, err := streamAPI.NextEvents()

		require.NoError(t, err)
		assert.Equal(t, expectedCursor, cursor)
		assert.Equal(t, expectedEvents, events)

		stream.unblock <- struct{}{}
		cursor, events, err = streamAPI.NextEvents()

		require.NoError(t, err)
		assert.Equal(t, expectedCursor, cursor)
		assert.Equal(t, expectedEvents, events)

		opener.unblock <- struct{}{} // provoke an error when openStream is called
	})

	t.Run("fail with canceled", func(t *testing.T) {
		stream := &mockStreamer{unblock: make(chan struct{}, 1)}
		streamAPI, opener, _ := setupMockStream(nil, nil)

		opener.On("openStream").Once().Return(stream, nil)
		opener.unblock <- struct{}{}
		stream.On("nextEvents").Once().Return(Cursor{}, nil, context.Canceled)
		stream.On("closeStream").Once().Return(nil)

		stream.unblock <- struct{}{}
		_, _, err := streamAPI.NextEvents()

		assert.EqualError(t, context.Canceled, err.Error())

		opener.unblock <- struct{}{} // provoke an error when openStream is called
	})
}

func TestStreamAPI_CommitCursor(t *testing.T) {
	retryCh := make(chan error, 1)
	okCh := make(chan struct{}, 1)
	errorCh := make(chan error, 1)
	expectedCursor := Cursor{NakadiStreamID: "stream-id"}

	t.Run("fail with time out", func(t *testing.T) {
		streamAPI, _, committer := setupMockStream(retryCh, okCh)
		committer.On("commitCursor", expectedCursor).Twice().Return(assert.AnError)

		go func() {
			err := streamAPI.CommitCursor(expectedCursor)
			errorCh <- err
		}()

		committer.unblock <- struct{}{}

		err := <-retryCh
		assert.EqualError(t, assert.AnError, err.Error())

		time.Sleep(600 * time.Millisecond)
		committer.unblock <- struct{}{}

		err = <-errorCh
		assert.EqualError(t, assert.AnError, err.Error())
	})

	t.Run("fail retry succeed", func(t *testing.T) {
		streamAPI, _, committer := setupMockStream(retryCh, okCh)
		committer.On("commitCursor", expectedCursor).Once().Return(assert.AnError)

		go func() {
			err := streamAPI.CommitCursor(expectedCursor)
			errorCh <- err
		}()

		committer.unblock <- struct{}{}

		err := <-retryCh
		assert.EqualError(t, assert.AnError, err.Error())

		committer.On("commitCursor", expectedCursor).Once().Return(nil)
		committer.unblock <- struct{}{}

		err = <-errorCh
		assert.NoError(t, err)

		select {
		case <-okCh:
			// nothing
		default:
			assert.Fail(t, "no retry ok detected")
		}
	})

	t.Run("success", func(t *testing.T) {
		streamAPI, _, committer := setupMockStream(nil, nil)
		committer.On("commitCursor", expectedCursor).Once().Return(nil)

		committer.unblock <- struct{}{}
		err := streamAPI.CommitCursor(expectedCursor)

		assert.NoError(t, err)
	})
}

func TestStreamAPI_Close(t *testing.T) {
	errorCh := make(chan error, 1)
	streamAPI, opener, _ := setupMockStream(nil, nil)

	stream := &mockStreamer{unblock: make(chan struct{}, 1)}
	opener.On("openStream").Once().Return(stream, nil)
	stream.On("nextEvents").Once().Return(Cursor{}, nil, assert.AnError)
	stream.On("closeStream").Once().Return(nil)

	go func() {
		err := streamAPI.Close()
		errorCh <- err
	}()

	opener.unblock <- struct{}{}
	time.Sleep(50 * time.Millisecond)
	stream.unblock <- struct{}{}

	err := <-errorCh
	assert.NoError(t, err)

	<-streamAPI.eventCh
	stream.AssertCalled(t, "closeStream")
}

func setupMockStream(errCh chan error, okCh chan struct{}) (*StreamAPI, *mockStreamOpener, *mockCommitter) {
	ctx, cancel := context.WithCancel(context.Background())

	streamBackOff := backoff.NewExponentialBackOff()
	streamBackOff.InitialInterval = 1 * time.Millisecond
	streamBackOff.MaxInterval = 100 * time.Millisecond
	streamBackOff.MaxElapsedTime = 500 * time.Millisecond

	commitBackOff := backoff.NewExponentialBackOff()
	commitBackOff.InitialInterval = 1 * time.Millisecond
	commitBackOff.MaxInterval = 100 * time.Millisecond
	commitBackOff.MaxElapsedTime = 500 * time.Millisecond

	opener := &mockStreamOpener{unblock: make(chan struct{}, 1)}
	committer := &mockCommitter{unblock: make(chan struct{}, 1)}

	stream := &StreamAPI{
		opener:        opener,
		committer:     committer,
		eventCh:       make(chan eventsOrError, 10),
		ctx:           ctx,
		cancel:        cancel,
		streamBackOff: backoff.WithContext(streamBackOff, ctx),
		commitBackOff: backoff.WithContext(commitBackOff, ctx),
		notifyErr: func(err error, _ time.Duration) {
			if errCh != nil {
				errCh <- err
			}
		},
		notifyOK: func() {
			if okCh != nil {
				okCh <- struct{}{}
			}
		},
	}

	go stream.startStream()

	return stream, opener, committer
}

type mockStreamOpener struct {
	mock.Mock
	unblock chan struct{}
}

func (so *mockStreamOpener) openStream() (streamer, error) {
	<-so.unblock
	args := so.Called()
	if args.Error(1) != nil {
		return nil, args.Error(1)
	}
	return args.Get(0).(streamer), nil
}

type mockStreamer struct {
	mock.Mock
	unblock chan struct{}
}

func (s *mockStreamer) nextEvents() (Cursor, []byte, error) {
	<-s.unblock
	args := s.Called()
	if args.Error(2) != nil {
		return Cursor{}, nil, args.Error(2)
	}
	return args.Get(0).(Cursor), args.Get(1).([]byte), nil
}

func (s *mockStreamer) closeStream() error {
	return s.Called().Error(0)
}

type mockCommitter struct {
	mock.Mock
	unblock chan struct{}
}

func (c *mockCommitter) commitCursor(cursor Cursor) error {
	<-c.unblock
	return c.Called(cursor).Error(0)
}
<commit_msg>Replaced unblocked channels with WaitUntil method from mock library<commit_after>package nakadi

import (
	"context"
	"testing"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestStreamAPI_startStream(t *testing.T) {
	retryCh := make(chan error, 1)
	okCh := make(chan struct{}, 1)
	blockCh := make(chan time.Time, 1)
	stream := &mockStreamer{}

	_, opener, _ := setupMockStream(retryCh, okCh)

	opener.On("openStream").Once().Return(nil, assert.AnError)
	opener.On("openStream").Once().Return(stream, nil).WaitUntil(blockCh)
	stream.On("nextEvents").Return(Cursor{}, nil, assert.AnError).WaitUntil(blockCh)

	select {
	case err := <-retryCh:
		assert.EqualError(t, assert.AnError, err.Error())
	case <-time.After(50 * time.Millisecond):
		assert.Fail(t, "no retry error detected")
	}

	blockCh <- time.Now()
	select {
	case <-okCh:
		// nothing
	case <-time.After(50 * time.Millisecond):
		assert.Fail(t, "no retry ok detected")
	}
}

func TestStreamAPI_NextEvents(t *testing.T) {
	expectedCursor := Cursor{NakadiStreamID: "stream-id"}
	expectedEvents := []byte(`"events":[{"metadata":{"eid":"74450ab6-5461-11e7-9dd2-87c3afa8811f"})]`)

	t.Run("fail with error", func(t *testing.T) {
		stream := &mockStreamer{}
		blockCh := make(chan time.Time, 1)
		streamAPI, opener, _ := setupMockStream(nil, nil)

		opener.On("openStream").Return(stream, nil).WaitUntil(blockCh)
		blockCh <- time.Now()

		stream.On("nextEvents").Return(Cursor{}, nil, assert.AnError).WaitUntil(blockCh)
		blockCh <- time.Now()
		stream.On("closeStream").Return(nil)

		_, _, err := streamAPI.NextEvents()

		assert.EqualError(t, assert.AnError, err.Error())
	})

	t.Run("successful read events", func(t *testing.T) {
		stream := &mockStreamer{}
		blockCh := make(chan time.Time, 1)
		streamAPI, opener, _ := setupMockStream(nil, nil)

		opener.On("openStream").Return(stream, nil).WaitUntil(blockCh)
		blockCh <- time.Now()
		stream.On("nextEvents").Return(expectedCursor, expectedEvents, nil).WaitUntil(blockCh)

		blockCh <- time.Now()
		cursor, events, err := streamAPI.NextEvents()

		require.NoError(t, err)
		assert.Equal(t, expectedCursor, cursor)
		assert.Equal(t, expectedEvents, events)

		blockCh <- time.Now()
		cursor, events, err = streamAPI.NextEvents()

		require.NoError(t, err)
		assert.Equal(t, expectedCursor, cursor)
		assert.Equal(t, expectedEvents, events)
	})

	t.Run("fail with canceled", func(t *testing.T) {
		stream := &mockStreamer{}
		blockCh := make(chan time.Time, 1)
		streamAPI, opener, _ := setupMockStream(nil, nil)

		opener.On("openStream").Return(stream, nil).WaitUntil(blockCh)
		blockCh <- time.Now()

		stream.On("nextEvents").Return(Cursor{}, nil, context.Canceled).WaitUntil(blockCh)
		stream.On("closeStream").Return(nil)

		blockCh <- time.Now()
		_, _, err := streamAPI.NextEvents()

		assert.EqualError(t, context.Canceled, err.Error())
	})
}

func TestStreamAPI_CommitCursor(t *testing.T) {
	retryCh := make(chan error, 1)
	okCh := make(chan struct{}, 1)
	errorCh := make(chan error, 1)
	blockCh := make(chan time.Time, 1)
	blockStreamer := make(chan time.Time, 1)
	expectedCursor := Cursor{NakadiStreamID: "stream-id"}

	t.Run("fail with time out", func(t *testing.T) {
		streamAPI, opener, committer := setupMockStream(retryCh, okCh)
		opener.On("openStream").WaitUntil(blockStreamer)
		committer.On("commitCursor", expectedCursor).WaitUntil(blockCh).Twice().Return(assert.AnError)

		go func() {
			err := streamAPI.CommitCursor(expectedCursor)
			errorCh <- err
		}()

		blockCh <- time.Now()

		err := <-retryCh
		assert.EqualError(t, assert.AnError, err.Error())

		time.Sleep(600 * time.Millisecond)
		blockCh <- time.Now()

		err = <-errorCh
		assert.EqualError(t, assert.AnError, err.Error())
	})

	t.Run("fail retry succeed", func(t *testing.T) {
		streamAPI, opener, committer := setupMockStream(retryCh, okCh)
		opener.On("openStream").WaitUntil(blockStreamer)
		committer.On("commitCursor", expectedCursor).WaitUntil(blockCh).Once().Return(assert.AnError)

		go func() {
			err := streamAPI.CommitCursor(expectedCursor)
			errorCh <- err
		}()

		blockCh <- time.Now()

		err := <-retryCh
		assert.EqualError(t, assert.AnError, err.Error())

		committer.On("commitCursor", expectedCursor).WaitUntil(blockCh).Once().Return(nil)
		blockCh <- time.Now()

		err = <-errorCh
		assert.NoError(t, err)

		select {
		case <-okCh:
			// nothing
		default:
			assert.Fail(t, "no retry ok detected")
		}
	})

	t.Run("success", func(t *testing.T) {
		streamAPI, opener, committer := setupMockStream(nil, nil)
		opener.On("openStream").WaitUntil(blockStreamer)
		committer.On("commitCursor", expectedCursor).Once().Return(nil).WaitUntil(blockCh)
		blockCh <- time.Now()
		err := streamAPI.CommitCursor(expectedCursor)

		assert.NoError(t, err)
	})
}

func TestStreamAPI_Close(t *testing.T) {
	errorCh := make(chan error, 1)
	blockCh := make(chan time.Time, 1)
	streamAPI, opener, _ := setupMockStream(nil, nil)

	stream := &mockStreamer{}
	opener.On("openStream").WaitUntil(blockCh).Return(stream, nil)
	stream.On("nextEvents").WaitUntil(blockCh).Once().Return(Cursor{}, nil, assert.AnError)
	stream.On("closeStream").Return(nil)

	go func() {
		err := streamAPI.Close()
		errorCh <- err
	}()

	blockCh <- time.Now()
	time.Sleep(50 * time.Millisecond)
	blockCh <- time.Now()

	err := <-errorCh
	assert.NoError(t, err)

	<-streamAPI.eventCh
	stream.AssertCalled(t, "closeStream")
}

func setupMockStream(errCh chan error, okCh chan struct{}) (*StreamAPI, *mockStreamOpener, *mockCommitter) {
	ctx, cancel := context.WithCancel(context.Background())

	streamBackOff := backoff.NewExponentialBackOff()
	streamBackOff.InitialInterval = 1 * time.Millisecond
	streamBackOff.MaxInterval = 100 * time.Millisecond
	streamBackOff.MaxElapsedTime = 500 * time.Millisecond

	commitBackOff := backoff.NewExponentialBackOff()
	commitBackOff.InitialInterval = 1 * time.Millisecond
	commitBackOff.MaxInterval = 100 * time.Millisecond
	commitBackOff.MaxElapsedTime = 500 * time.Millisecond

	opener := &mockStreamOpener{}
	committer := &mockCommitter{}

	stream := &StreamAPI{
		opener:        opener,
		committer:     committer,
		eventCh:       make(chan eventsOrError, 10),
		ctx:           ctx,
		cancel:        cancel,
		streamBackOff: backoff.WithContext(streamBackOff, ctx),
		commitBackOff: backoff.WithContext(commitBackOff, ctx),
		notifyErr: func(err error, _ time.Duration) {
			if errCh != nil {
				errCh <- err
			}
		},
		notifyOK: func() {
			if okCh != nil {
				okCh <- struct{}{}
			}
		},
	}

	go stream.startStream()

	return stream, opener, committer
}

type mockStreamOpener struct {
	mock.Mock
}

func (so *mockStreamOpener) openStream() (streamer, error) {
	args := so.Called()
	if args.Error(1) != nil {
		return nil, args.Error(1)
	}
	return args.Get(0).(streamer), nil
}

type mockStreamer struct {
	mock.Mock
}

func (s *mockStreamer) nextEvents() (Cursor, []byte, error) {
	args := s.Called()
	if args.Error(2) != nil {
		return Cursor{}, nil, args.Error(2)
	}
	return args.Get(0).(Cursor), args.Get(1).([]byte), nil
}

func (s *mockStreamer) closeStream() error {
	return s.Called().Error(0)
}

type mockCommitter struct {
	mock.Mock
}

func (c *mockCommitter) commitCursor(cursor Cursor) error {
	return c.Called(cursor).Error(0)
}
<|endoftext|>


<commit_before>package strutil_test

import (
	. "github.com/101loops/bdd"
	"github.com/101loops/strutil"
)

var _ = Describe("String Utility", func() {

	It("deletes all spaces", func() {
		Check(strutil.DeleteWhitespace(""), Equals, "")
		Check(strutil.DeleteWhitespace(" a  b "), Equals, "ab")
		Check(strutil.DeleteWhitespace("a \n \t b"), Equals, "ab")
	})

	It("trims all spaces", func() {
		Check(strutil.ShrinkWhitespace(""), Equals, "")
		Check(strutil.ShrinkWhitespace(" "), Equals, "")
		Check(strutil.ShrinkWhitespace(" test "), Equals, "test")
		Check(strutil.ShrinkWhitespace(" te  st "), Equals, "te st")
		Check(strutil.ShrinkWhitespace(" te\tst "), Equals, "te st")
		Check(strutil.ShrinkWhitespace(" te\nst "), Equals, "te st")
	})

	It("reverses string", func() {
		Check(strutil.Reverse(""), Equals, "")
		Check(strutil.Reverse("abc"), Equals, "cba")
	})

	It("checks if empty", func() {
		Check(strutil.IsEmpty(""), IsTrue)

		Check(strutil.IsEmpty(" "), IsFalse)
		Check(strutil.IsEmpty("abc"), IsFalse)
	})

	It("checks if blank", func() {
		Check(strutil.IsBlank(""), IsTrue)
		Check(strutil.IsBlank("\n"), IsTrue)
		Check(strutil.IsBlank("\t"), IsTrue)
		Check(strutil.IsBlank(" "), IsTrue)

		Check(strutil.IsBlank("abc"), IsFalse)
		Check(strutil.IsBlank(" abc "), IsFalse)
	})

	It("returns index after substring", func() {
		Check(strutil.IndexAfter("", ""), Equals, 0)
		Check(strutil.IndexAfter("", "a"), Equals, -1)
		Check(strutil.IndexAfter("abc", "ab"), Equals, 2)
		Check(strutil.IndexAfter("abc", "bc"), Equals, 3)
		Check(strutil.IndexAfter("abc", "xyz"), Equals, -1)
	})

	It("returns lines", func() {
		Check(strutil.Lines(""), Equals, []string{""})
		Check(strutil.Lines("a"), Equals, []string{"a"})
		Check(strutil.Lines("a\nb\nc"), Equals, []string{"a", "b", "c"})
	})

	It("wraps text to fixed width", func() {
		Check(strutil.WordWrap(``, 5), Equals, ``)
		//Check(WordWrap(`abcdef`, 5), Equals, `abcdef`)
		//Check(WordWrap(`abcdefghijklmno`, 5), Equals, `abcdef`)
	})

	It("surrounds a string", func() {
		Check(strutil.Surround("Go", "'"), Equals, "'Go'")
	})
})
<commit_msg>add missing test case<commit_after>package strutil_test

import (
	. "github.com/101loops/bdd"
	"github.com/101loops/strutil"
)

var _ = Describe("String Utility", func() {

	It("deletes all spaces", func() {
		Check(strutil.DeleteWhitespace(""), Equals, "")
		Check(strutil.DeleteWhitespace("abc"), Equals, "abc")
		Check(strutil.DeleteWhitespace(" a  b "), Equals, "ab")
		Check(strutil.DeleteWhitespace("a \n \t b"), Equals, "ab")
	})

	It("trims all spaces", func() {
		Check(strutil.ShrinkWhitespace(""), Equals, "")
		Check(strutil.ShrinkWhitespace(" "), Equals, "")
		Check(strutil.ShrinkWhitespace(" test "), Equals, "test")
		Check(strutil.ShrinkWhitespace(" te  st "), Equals, "te st")
		Check(strutil.ShrinkWhitespace(" te\tst "), Equals, "te st")
		Check(strutil.ShrinkWhitespace(" te\nst "), Equals, "te st")
	})

	It("reverses string", func() {
		Check(strutil.Reverse(""), Equals, "")
		Check(strutil.Reverse("abc"), Equals, "cba")
	})

	It("checks if empty", func() {
		Check(strutil.IsEmpty(""), IsTrue)

		Check(strutil.IsEmpty(" "), IsFalse)
		Check(strutil.IsEmpty("abc"), IsFalse)
	})

	It("checks if blank", func() {
		Check(strutil.IsBlank(""), IsTrue)
		Check(strutil.IsBlank("\n"), IsTrue)
		Check(strutil.IsBlank("\t"), IsTrue)
		Check(strutil.IsBlank(" "), IsTrue)

		Check(strutil.IsBlank("abc"), IsFalse)
		Check(strutil.IsBlank(" abc "), IsFalse)
	})

	It("returns index after substring", func() {
		Check(strutil.IndexAfter("", ""), Equals, 0)
		Check(strutil.IndexAfter("", "a"), Equals, -1)
		Check(strutil.IndexAfter("abc", "ab"), Equals, 2)
		Check(strutil.IndexAfter("abc", "bc"), Equals, 3)
		Check(strutil.IndexAfter("abc", "xyz"), Equals, -1)
	})

	It("returns lines", func() {
		Check(strutil.Lines(""), Equals, []string{""})
		Check(strutil.Lines("a"), Equals, []string{"a"})
		Check(strutil.Lines("a\nb\nc"), Equals, []string{"a", "b", "c"})
	})

	It("wraps text to fixed width", func() {
		Check(strutil.WordWrap(``, 5), Equals, ``)
		//Check(WordWrap(`abcdef`, 5), Equals, `abcdef`)
		//Check(WordWrap(`abcdefghijklmno`, 5), Equals, `abcdef`)
	})

	It("surrounds a string", func() {
		Check(strutil.Surround("Go", "'"), Equals, "'Go'")
	})
})
<|endoftext|>


<commit_before>package appregistration

import (
	stork_api "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

const (
	// IBMApp registration name
	IBMApp = "ibm"
	// CouchBaseApp registration name
	CouchBaseApp = "couchbase"
	// RedisClusterApp registration name
	RedisClusterApp = "redis"
	// CassandraApp registration name
	CassandraApp = "cassandra"
	// WeblogicDomainApp registration name
	WeblogicDomainApp = "weblogic"
	// RabbitMQApp registration name
	RabbitMQApp = "rabbitmq"
	// PerconaDBApp registration name
	PerconaDBApp = "perconadb"
	// PrometheusApp registration name
	PrometheusApp = "prometheus"
	// KafkaApp registration name
	KafkaApp = "kafka"
	// PostgressApp registration name
	PostgressApp = "postgress"
	// MongoDBCommunityApp registration name
	MongoDBCommunityApp = "mongodbcommunity"
)

// GetSupportedCRD returns the list of supported CRDs.
// Note: If you add support for a new CRD add it to the supported
// GroupVersionResource list as well.
func GetSupportedCRD() map[string][]stork_api.ApplicationResource {
	// supported CRD registration
	defCRD := make(map[string][]stork_api.ApplicationResource)
	// IBM CRD's
	defCRD[IBMApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPCA",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPConsole",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPOrderer",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPPeer",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
	}
	//CouchBase CRD's
	defCRD[CouchBaseApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseBucket",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseCluster",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.paused",
				Type: "bool",
			},
			PodsPath: "status.members.ready",
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseEphemeralBucket",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseMemcachedBucket",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseReplication",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseUser",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseGroup",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseRoleBinding",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseBackup",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseBackupRestore",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
	}
	// datastax/Cassandra CRD's
	defCRD[CassandraApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CassandraDatacenter",
				Group:   "cassandra.datastax.com",
				Version: "v1beta1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.stopped",
				Type: "bool",
			},
		},
	}
	// redis cluster CRD's
	defCRD[RedisClusterApp] =
		[]stork_api.ApplicationResource{
			{
				GroupVersionKind: metav1.GroupVersionKind{
					Kind:    "RedisEnterpriseCluster",
					Group:   "app.redislabs.com",
					Version: "v1",
				},
				KeepStatus: false,
			},
			{
				GroupVersionKind: metav1.GroupVersionKind{
					Kind:    "RedisEnterpriseDatabase",
					Group:   "app.redislabs.com",
					Version: "v1",
				},
				KeepStatus: false,
			},
		}
		// weblogic domain crds
	defCRD[WeblogicDomainApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "Domain",
				Group:   "weblogic.oracle",
				Version: "v8",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path:  "spec.serverStartPolicy",
				Type:  "string",
				Value: "NEVER",
			},
		},
	}
	// rabbitmq app crds
	defCRD[RabbitMQApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "RabbitmqCluster",
				Group:   "rabbitmq.com",
				Version: "v1beta1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
	}

	// perconadb app crds
	defCRD[PerconaDBApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "PerconaXtraDBCluster",
				Group:   "pxc.percona.com",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.pxc.size",
				Type: "int",
			},
		},
	}
	// prometheus app crds
	defCRD[PrometheusApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "Prometheus",
				Group:   "monitoring.coreos.com",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
	}

	// kafka app crds
	defCRD[KafkaApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "Kafka",
				Group:   "kafka.strimzi.io",
				Version: "v1beta2",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.kafka.replicas",
				Type: "int",
			},
		},
	}
	// postgress app crds
	defCRD[PostgressApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "postgresql",
				Group:   "acid.zalan.do",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.numberOfInstances",
				Type: "int",
			},
		},
	}
	// mongdb app crds
	defCRD[MongoDBCommunityApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "MongoDBCommunity",
				Group:   "mongodbcommunity.mongodb.com",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.members",
				Type: "int",
			},
		},
	}

	return defCRD
}

// GetSupportedGVR returns a list of GroupVersionResource schema for the supported CRDs
func GetSupportedGVR() map[schema.GroupVersionResource]string {

	return map[schema.GroupVersionResource]string{
		// IBM CRD's
		{
			Resource: "ibpcas",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IBPCAsList",
		{
			Resource: "ibpconsoles",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IBPConsolesList",
		{
			Resource: "ibporderers",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IPBOrderersList",
		{
			Resource: "ibppeers",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IBPeersList",
		{
			Resource: "couchbasebuckets",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseBucketsList",
		{
			Resource: "couchbaseclusters",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseClustersList",
		{
			Resource: "couchbaseephemeralbuckets",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseEphemeralBucketsList",
		{
			Resource: "couchbasememcachedbuckets",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseMemcachedBucketsList",
		{
			Resource: "couchbasereplications",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseReplicationsList",
		{
			Resource: "couchbaseusers",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseUsersList",
		{
			Resource: "couchbasegroups",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseGroupsList",
		{
			Resource: "couchbaserolebindings",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseRoleBindingsList",
		{
			Resource: "couchbasebackups",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseBackupsList",
		{
			Resource: "couchbasebackuprestores",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseBackupRestoresList",
		// datastax/Cassandra CRD's
		{
			Resource: "cassandradatacenters",
			Group:    "cassandra.datastax.com",
			Version:  "v1beta1",
		}: "CassandraDatacentersList",
		// redis cluster CRD's
		{
			Resource: "redisenterpriseclusters",
			Group:    "app.redislabs.com",
			Version:  "v1",
		}: "RedisEnterpriseClustersList",
		{
			Resource: "redisenterprisedatabases",
			Group:    "app.redislabs.com",
			Version:  "v1",
		}: "RedisEnterpriseDatabasesList",
		// weblogic domain crds
		{
			Resource: "domains",
			Group:    "weblogic.oracle",
			Version:  "v8",
		}: "DomainsList",
		// mongodb crds
		{
			Resource: "MongoDBCommunity",
			Group:    "mongodbcommunity.mongodb.com",
			Version:  "v1",
		}: "MongoDBCommunityList",
	}
}
<commit_msg>Update PerconaXtraDBCluster suspend option to gracefully shutdown cluster<commit_after>package appregistration

import (
	stork_api "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

const (
	// IBMApp registration name
	IBMApp = "ibm"
	// CouchBaseApp registration name
	CouchBaseApp = "couchbase"
	// RedisClusterApp registration name
	RedisClusterApp = "redis"
	// CassandraApp registration name
	CassandraApp = "cassandra"
	// WeblogicDomainApp registration name
	WeblogicDomainApp = "weblogic"
	// RabbitMQApp registration name
	RabbitMQApp = "rabbitmq"
	// PerconaDBApp registration name
	PerconaDBApp = "perconadb"
	// PrometheusApp registration name
	PrometheusApp = "prometheus"
	// KafkaApp registration name
	KafkaApp = "kafka"
	// PostgressApp registration name
	PostgressApp = "postgress"
	// MongoDBCommunityApp registration name
	MongoDBCommunityApp = "mongodbcommunity"
)

// GetSupportedCRD returns the list of supported CRDs.
// Note: If you add support for a new CRD add it to the supported
// GroupVersionResource list as well.
func GetSupportedCRD() map[string][]stork_api.ApplicationResource {
	// supported CRD registration
	defCRD := make(map[string][]stork_api.ApplicationResource)
	// IBM CRD's
	defCRD[IBMApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPCA",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPConsole",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPOrderer",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "IBPPeer",
				Group:   "ibp.com",
				Version: "v1alpha1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
	}
	//CouchBase CRD's
	defCRD[CouchBaseApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseBucket",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseCluster",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.paused",
				Type: "bool",
			},
			PodsPath: "status.members.ready",
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseEphemeralBucket",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseMemcachedBucket",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseReplication",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseUser",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseGroup",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseRoleBinding",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseBackup",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CouchbaseBackupRestore",
				Group:   "couchbase.com",
				Version: "v2",
			},
			KeepStatus: false,
		},
	}
	// datastax/Cassandra CRD's
	defCRD[CassandraApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "CassandraDatacenter",
				Group:   "cassandra.datastax.com",
				Version: "v1beta1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.stopped",
				Type: "bool",
			},
		},
	}
	// redis cluster CRD's
	defCRD[RedisClusterApp] =
		[]stork_api.ApplicationResource{
			{
				GroupVersionKind: metav1.GroupVersionKind{
					Kind:    "RedisEnterpriseCluster",
					Group:   "app.redislabs.com",
					Version: "v1",
				},
				KeepStatus: false,
			},
			{
				GroupVersionKind: metav1.GroupVersionKind{
					Kind:    "RedisEnterpriseDatabase",
					Group:   "app.redislabs.com",
					Version: "v1",
				},
				KeepStatus: false,
			},
		}
		// weblogic domain crds
	defCRD[WeblogicDomainApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "Domain",
				Group:   "weblogic.oracle",
				Version: "v8",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path:  "spec.serverStartPolicy",
				Type:  "string",
				Value: "NEVER",
			},
		},
	}
	// rabbitmq app crds
	defCRD[RabbitMQApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "RabbitmqCluster",
				Group:   "rabbitmq.com",
				Version: "v1beta1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
	}

	// perconadb app crds
	defCRD[PerconaDBApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "PerconaXtraDBCluster",
				Group:   "pxc.percona.com",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.pause",
				Type: "bool",
			},
		},
	}
	// prometheus app crds
	defCRD[PrometheusApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "Prometheus",
				Group:   "monitoring.coreos.com",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.replicas",
				Type: "int",
			},
		},
	}

	// kafka app crds
	defCRD[KafkaApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "Kafka",
				Group:   "kafka.strimzi.io",
				Version: "v1beta2",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.kafka.replicas",
				Type: "int",
			},
		},
	}
	// postgress app crds
	defCRD[PostgressApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "postgresql",
				Group:   "acid.zalan.do",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.numberOfInstances",
				Type: "int",
			},
		},
	}
	// mongdb app crds
	defCRD[MongoDBCommunityApp] = []stork_api.ApplicationResource{
		{
			GroupVersionKind: metav1.GroupVersionKind{
				Kind:    "MongoDBCommunity",
				Group:   "mongodbcommunity.mongodb.com",
				Version: "v1",
			},
			KeepStatus: false,
			SuspendOptions: stork_api.SuspendOptions{
				Path: "spec.members",
				Type: "int",
			},
		},
	}

	return defCRD
}

// GetSupportedGVR returns a list of GroupVersionResource schema for the supported CRDs
func GetSupportedGVR() map[schema.GroupVersionResource]string {

	return map[schema.GroupVersionResource]string{
		// IBM CRD's
		{
			Resource: "ibpcas",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IBPCAsList",
		{
			Resource: "ibpconsoles",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IBPConsolesList",
		{
			Resource: "ibporderers",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IPBOrderersList",
		{
			Resource: "ibppeers",
			Group:    "ibp.com",
			Version:  "v1alpha1",
		}: "IBPeersList",
		{
			Resource: "couchbasebuckets",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseBucketsList",
		{
			Resource: "couchbaseclusters",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseClustersList",
		{
			Resource: "couchbaseephemeralbuckets",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseEphemeralBucketsList",
		{
			Resource: "couchbasememcachedbuckets",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseMemcachedBucketsList",
		{
			Resource: "couchbasereplications",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseReplicationsList",
		{
			Resource: "couchbaseusers",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseUsersList",
		{
			Resource: "couchbasegroups",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseGroupsList",
		{
			Resource: "couchbaserolebindings",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseRoleBindingsList",
		{
			Resource: "couchbasebackups",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseBackupsList",
		{
			Resource: "couchbasebackuprestores",
			Group:    "couchbase.com",
			Version:  "v2",
		}: "CouchbaseBackupRestoresList",
		// datastax/Cassandra CRD's
		{
			Resource: "cassandradatacenters",
			Group:    "cassandra.datastax.com",
			Version:  "v1beta1",
		}: "CassandraDatacentersList",
		// redis cluster CRD's
		{
			Resource: "redisenterpriseclusters",
			Group:    "app.redislabs.com",
			Version:  "v1",
		}: "RedisEnterpriseClustersList",
		{
			Resource: "redisenterprisedatabases",
			Group:    "app.redislabs.com",
			Version:  "v1",
		}: "RedisEnterpriseDatabasesList",
		// weblogic domain crds
		{
			Resource: "domains",
			Group:    "weblogic.oracle",
			Version:  "v8",
		}: "DomainsList",
		// mongodb crds
		{
			Resource: "MongoDBCommunity",
			Group:    "mongodbcommunity.mongodb.com",
			Version:  "v1",
		}: "MongoDBCommunityList",
	}
}
<|endoftext|>


<commit_before>package saml

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"encoding/xml"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/crewjam/saml"
	"github.com/crewjam/saml/samlsp"
	"github.com/gorilla/mux"
	"github.com/rancher/rancher/pkg/auth/tokens"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	log "github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

type IDPMetadata struct {
	XMLName           xml.Name                `xml:"urn:oasis:names:tc:SAML:2.0:metadata EntityDescriptor"`
	ValidUntil        time.Time               `xml:"validUntil,attr"`
	EntityID          string                  `xml:"entityID,attr"`
	IDPSSODescriptors []saml.IDPSSODescriptor `xml:"IDPSSODescriptor"`
	SPSSODescriptors  []saml.SPSSODescriptor  `xml:"SPSSODescriptor"`
}

var root *mux.Router
var appliedVersion string
var initMu sync.Mutex

func InitializeSamlServiceProvider(configToSet *v3.SamlConfig, name string) error {

	initMu.Lock()
	defer initMu.Unlock()

	if configToSet.ResourceVersion == appliedVersion {
		return nil
	}

	var privKey *rsa.PrivateKey
	var cert *x509.Certificate
	var err error
	var ok bool

	if configToSet.IDPMetadataContent == "" {
		return fmt.Errorf("SAML: Cannot initialize saml SP properly, missing IDP URL/metadata in the config %v", configToSet)
	}

	if configToSet.SpCert == "" {
		return fmt.Errorf("SAML: Cannot initialize saml SP properly, missing SpCert in the config %v", configToSet)
	}

	if configToSet.SpKey == "" {
		return fmt.Errorf("SAML: Cannot initialize saml SP properly, missing SpKey in the config %v", configToSet)
	}

	if configToSet.SpKey != "" {
		// used from ssh.ParseRawPrivateKey

		block, _ := pem.Decode([]byte(configToSet.SpKey))
		if block == nil {
			return fmt.Errorf("SAML: no key found")
		}

		if strings.Contains(block.Headers["Proc-Type"], "ENCRYPTED") {
			return fmt.Errorf("SAML: cannot decode encrypted private keys")
		}

		switch block.Type {
		case "RSA PRIVATE KEY":
			privKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
			if err != nil {
				return fmt.Errorf("SAML: error parsing PKCS1 RSA key: %v", err)
			}
		case "PRIVATE KEY":
			pk, err := x509.ParsePKCS8PrivateKey(block.Bytes)
			if err != nil {
				return fmt.Errorf("SAML: error parsing PKCS8 RSA key: %v", err)
			}
			privKey, ok = pk.(*rsa.PrivateKey)
			if !ok {
				return fmt.Errorf("SAML: unable to get rsa key")
			}
		default:
			return fmt.Errorf("SAML: unsupported key type %q", block.Type)
		}
	}

	if configToSet.SpCert != "" {
		block, _ := pem.Decode([]byte(configToSet.SpCert))
		if block == nil {
			return fmt.Errorf("SAML: failed to parse PEM block containing the private key")
		}

		cert, err = x509.ParseCertificate(block.Bytes)
		if err != nil {
			return fmt.Errorf("SAML: failed to parse DER encoded public key: " + err.Error())
		}
	}

	provider := SamlProviders[name]

	rancherAPIHost := strings.TrimRight(configToSet.RancherAPIHost, "/")
	samlURL := rancherAPIHost + "/v1-saml/"
	samlURL += name
	actURL, err := url.Parse(samlURL)
	if err != nil {
		return fmt.Errorf("SAML: error in parsing URL")
	}

	metadataURL := *actURL
	metadataURL.Path = metadataURL.Path + "/saml/metadata"
	acsURL := *actURL
	acsURL.Path = acsURL.Path + "/saml/acs"

	sp := saml.ServiceProvider{
		Key:         privKey,
		Certificate: cert,
		MetadataURL: metadataURL,
		AcsURL:      acsURL,
	}

	// XML unmarshal throws an error for IdP Metadata cacheDuration field, as it's of type xml Duration. Using a separate struct for unmarshaling for now
	idm := &IDPMetadata{}
	if configToSet.IDPMetadataContent != "" {
		sp.IDPMetadata = &saml.EntityDescriptor{}
		if err := xml.NewDecoder(strings.NewReader(configToSet.IDPMetadataContent)).Decode(idm); err != nil {
			return fmt.Errorf("SAML: cannot initialize saml SP, cannot decode IDP Metadata content from the config %v, error %v", configToSet, err)
		}
	}

	sp.IDPMetadata.XMLName = idm.XMLName
	sp.IDPMetadata.ValidUntil = idm.ValidUntil
	sp.IDPMetadata.EntityID = idm.EntityID
	sp.IDPMetadata.SPSSODescriptors = idm.SPSSODescriptors
	sp.IDPMetadata.IDPSSODescriptors = idm.IDPSSODescriptors
	if name == ADFSName || name == OKTAName {
		sp.AuthnNameIDFormat = saml.UnspecifiedNameIDFormat
	}

	provider.serviceProvider = &sp

	cookieStore := samlsp.ClientCookies{
		ServiceProvider: &sp,
		Name:            "token",
		Domain:          actURL.Host,
	}
	provider.clientState = &cookieStore

	SamlProviders[name] = provider

	switch name {
	case PingName:
		root.Get("PingACS").HandlerFunc(provider.ServeHTTP)
		root.Get("PingMetadata").HandlerFunc(provider.ServeHTTP)
	case ADFSName:
		root.Get("AdfsACS").HandlerFunc(provider.ServeHTTP)
		root.Get("AdfsMetadata").HandlerFunc(provider.ServeHTTP)
	case KeyCloakName:
		root.Get("KeyCloakACS").HandlerFunc(provider.ServeHTTP)
		root.Get("KeyCloakMetadata").HandlerFunc(provider.ServeHTTP)
	case OKTAName:
		root.Get("OktaACS").HandlerFunc(provider.ServeHTTP)
		root.Get("OktaMetadata").HandlerFunc(provider.ServeHTTP)
	}

	appliedVersion = configToSet.ResourceVersion

	return nil
}

func AuthHandler() http.Handler {
	root = mux.NewRouter()
	root.Methods("POST").Path("/v1-saml/ping/saml/acs").Name("PingACS")
	root.Methods("GET").Path("/v1-saml/ping/saml/metadata").Name("PingMetadata")

	root.Methods("POST").Path("/v1-saml/adfs/saml/acs").Name("AdfsACS")
	root.Methods("GET").Path("/v1-saml/adfs/saml/metadata").Name("AdfsMetadata")

	root.Methods("POST").Path("/v1-saml/keycloak/saml/acs").Name("KeyCloakACS")
	root.Methods("GET").Path("/v1-saml/keycloak/saml/metadata").Name("KeyCloakMetadata")

	root.Methods("POST").Path("/v1-saml/okta/saml/acs").Name("OktaACS")
	root.Methods("GET").Path("/v1-saml/okta/saml/metadata").Name("OktaMetadata")

	return root
}

func (s *Provider) getSamlPrincipals(config *v3.SamlConfig, samlData map[string][]string) (v3.Principal, []v3.Principal, error) {
	var userPrincipal v3.Principal
	var groupPrincipals []v3.Principal
	uid, ok := samlData[config.UIDField]
	if !ok {
		// UID field provided by user is actually not there in SAMLResponse, without this we cannot differentiate between users and create separate principals
		return userPrincipal, groupPrincipals, fmt.Errorf("SAML: Unique ID field is not provided in SAML Response")
	}
	userPrincipal = v3.Principal{
		ObjectMeta:    metav1.ObjectMeta{Name: s.userType + "://" + uid[0]},
		Provider:      s.name,
		PrincipalType: "user",
		Me:            true,
	}

	displayName, ok := samlData[config.DisplayNameField]
	if ok {
		userPrincipal.DisplayName = displayName[0]
	}

	userName, ok := samlData[config.UserNameField]
	if ok {
		userPrincipal.LoginName = userName[0]
	}

	groups, ok := samlData[config.GroupsField]
	if ok {
		for _, group := range groups {
			group := v3.Principal{
				ObjectMeta:    metav1.ObjectMeta{Name: s.groupType + "://" + group},
				DisplayName:   group,
				Provider:      s.name,
				PrincipalType: "group",
				MemberOf:      true,
			}
			groupPrincipals = append(groupPrincipals, group)
		}
	}

	return userPrincipal, groupPrincipals, nil
}

// HandleSamlAssertion processes/handles the assertion obtained by the POST to /saml/acs from IdP
func (s *Provider) HandleSamlAssertion(w http.ResponseWriter, r *http.Request, assertion *saml.Assertion) {
	var groupPrincipals []v3.Principal
	var userPrincipal v3.Principal

	if relayState := r.Form.Get("RelayState"); relayState != "" {
		// delete the cookie
		s.clientState.DeleteState(w, r, relayState)
	}

	redirectURL := s.clientState.GetState(r, "Rancher_FinalRedirectURL")
	rancherAction := s.clientState.GetState(r, "Rancher_Action")
	if rancherAction == loginAction {
		redirectURL += "/login?"
	} else if rancherAction == testAndEnableAction {
		// the first query param is config=saml_provider_name set by UI
		redirectURL += "&"
	}

	samlData := make(map[string][]string)

	for _, attributeStatement := range assertion.AttributeStatements {
		for _, attr := range attributeStatement.Attributes {
			attrName := attr.FriendlyName
			if attrName == "" {
				attrName = attr.Name
			}
			for _, value := range attr.Values {
				samlData[attrName] = append(samlData[attrName], value.Value)
			}
		}
	}

	config, err := s.getSamlConfig()
	if err != nil {
		log.Errorf("SAML: Error getting saml config %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		return
	}

	userPrincipal, groupPrincipals, err = s.getSamlPrincipals(config, samlData)
	if err != nil {
		log.Error(err)
		http.Redirect(w, r, redirectURL+"errorCode=422&errorMsg=Invalid saml attributes", http.StatusFound)
		return
	}
	allowedPrincipals := config.AllowedPrincipalIDs

	allowed, err := s.userMGR.CheckAccess(config.AccessMode, allowedPrincipals, userPrincipal.Name, groupPrincipals)
	if err != nil {
		log.Errorf("SAML: Error during login while checking access %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		return
	}
	if !allowed {
		log.Errorf("SAML: User does not have access %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=403", http.StatusFound)
		return
	}

	userID := s.clientState.GetState(r, "Rancher_UserID")
	if userID != "" && rancherAction == testAndEnableAction {
		user, err := s.userMGR.SetPrincipalOnCurrentUserByUserID(userID, userPrincipal)
		if err != nil {
			log.Errorf("SAML: Error setting principal on current user %v", err)
			http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
			return
		}

		config.Enabled = true
		err = s.saveSamlConfig(config)
		if err != nil {
			log.Errorf("SAML: Error saving saml config %v", err)
			http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
			return
		}

		isSecure := false
		if r.URL.Scheme == "https" {
			isSecure = true
		}
		err = setRancherToken(w, r, s.tokenMGR, user.Name, userPrincipal, groupPrincipals, isSecure)
		if err != nil {
			log.Errorf("SAML: Failed creating token with error: %v", err)
			http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		}
		// delete the cookies
		s.clientState.DeleteState(w, r, "Rancher_UserID")
		s.clientState.DeleteState(w, r, "Rancher_Action")
		redirectURL := s.clientState.GetState(r, "Rancher_FinalRedirectURL")
		s.clientState.DeleteState(w, r, "Rancher_FinalRedirectURL")
		if redirectURL != "" {
			// delete the cookie
			s.clientState.DeleteState(w, r, "Rancher_FinalRedirectURL")
			http.Redirect(w, r, redirectURL, http.StatusFound)
		}
		return
	}

	displayName := userPrincipal.DisplayName
	if displayName == "" {
		displayName = userPrincipal.LoginName
	}
	user, err := s.userMGR.EnsureUser(userPrincipal.Name, displayName)
	if err != nil {
		log.Errorf("SAML: Failed getting user with error: %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		return
	}

	if user.Enabled != nil && !*user.Enabled {
		log.Errorf("SAML: User %v permission denied", user.Name)
		http.Redirect(w, r, redirectURL+"errorCode=403", http.StatusFound)
		return
	}

	err = setRancherToken(w, r, s.tokenMGR, user.Name, userPrincipal, groupPrincipals, true)
	if err != nil {
		log.Errorf("SAML: Failed creating token with error: %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
	}
	redirectURL = s.clientState.GetState(r, "Rancher_FinalRedirectURL")
	if redirectURL != "" {
		// delete the cookie
		s.clientState.DeleteState(w, r, "Rancher_FinalRedirectURL")
		http.Redirect(w, r, redirectURL, http.StatusFound)
	}
	return
}

func setRancherToken(w http.ResponseWriter, r *http.Request, tokenMGR *tokens.Manager, userID string, userPrincipal v3.Principal,
	groupPrincipals []v3.Principal, isSecure bool) error {
	rToken, err := tokenMGR.NewLoginToken(userID, userPrincipal, groupPrincipals, "", 0, "")
	if err != nil {
		return err
	}
	tokenCookie := &http.Cookie{
		Name:     "R_SESS",
		Value:    rToken.ObjectMeta.Name + ":" + rToken.Token,
		Secure:   isSecure,
		Path:     "/",
		HttpOnly: true,
	}
	http.SetCookie(w, tokenCookie)
	return nil
}
<commit_msg>Send translation key as error message<commit_after>package saml

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"encoding/xml"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/crewjam/saml"
	"github.com/crewjam/saml/samlsp"
	"github.com/gorilla/mux"
	"github.com/rancher/rancher/pkg/auth/tokens"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	log "github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

type IDPMetadata struct {
	XMLName           xml.Name                `xml:"urn:oasis:names:tc:SAML:2.0:metadata EntityDescriptor"`
	ValidUntil        time.Time               `xml:"validUntil,attr"`
	EntityID          string                  `xml:"entityID,attr"`
	IDPSSODescriptors []saml.IDPSSODescriptor `xml:"IDPSSODescriptor"`
	SPSSODescriptors  []saml.SPSSODescriptor  `xml:"SPSSODescriptor"`
}

var root *mux.Router
var appliedVersion string
var initMu sync.Mutex

const UITranslationKeyForErrorMessage = "invalidSamlAttrs"

func InitializeSamlServiceProvider(configToSet *v3.SamlConfig, name string) error {

	initMu.Lock()
	defer initMu.Unlock()

	if configToSet.ResourceVersion == appliedVersion {
		return nil
	}

	var privKey *rsa.PrivateKey
	var cert *x509.Certificate
	var err error
	var ok bool

	if configToSet.IDPMetadataContent == "" {
		return fmt.Errorf("SAML: Cannot initialize saml SP properly, missing IDP URL/metadata in the config %v", configToSet)
	}

	if configToSet.SpCert == "" {
		return fmt.Errorf("SAML: Cannot initialize saml SP properly, missing SpCert in the config %v", configToSet)
	}

	if configToSet.SpKey == "" {
		return fmt.Errorf("SAML: Cannot initialize saml SP properly, missing SpKey in the config %v", configToSet)
	}

	if configToSet.SpKey != "" {
		// used from ssh.ParseRawPrivateKey

		block, _ := pem.Decode([]byte(configToSet.SpKey))
		if block == nil {
			return fmt.Errorf("SAML: no key found")
		}

		if strings.Contains(block.Headers["Proc-Type"], "ENCRYPTED") {
			return fmt.Errorf("SAML: cannot decode encrypted private keys")
		}

		switch block.Type {
		case "RSA PRIVATE KEY":
			privKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
			if err != nil {
				return fmt.Errorf("SAML: error parsing PKCS1 RSA key: %v", err)
			}
		case "PRIVATE KEY":
			pk, err := x509.ParsePKCS8PrivateKey(block.Bytes)
			if err != nil {
				return fmt.Errorf("SAML: error parsing PKCS8 RSA key: %v", err)
			}
			privKey, ok = pk.(*rsa.PrivateKey)
			if !ok {
				return fmt.Errorf("SAML: unable to get rsa key")
			}
		default:
			return fmt.Errorf("SAML: unsupported key type %q", block.Type)
		}
	}

	if configToSet.SpCert != "" {
		block, _ := pem.Decode([]byte(configToSet.SpCert))
		if block == nil {
			return fmt.Errorf("SAML: failed to parse PEM block containing the private key")
		}

		cert, err = x509.ParseCertificate(block.Bytes)
		if err != nil {
			return fmt.Errorf("SAML: failed to parse DER encoded public key: " + err.Error())
		}
	}

	provider := SamlProviders[name]

	rancherAPIHost := strings.TrimRight(configToSet.RancherAPIHost, "/")
	samlURL := rancherAPIHost + "/v1-saml/"
	samlURL += name
	actURL, err := url.Parse(samlURL)
	if err != nil {
		return fmt.Errorf("SAML: error in parsing URL")
	}

	metadataURL := *actURL
	metadataURL.Path = metadataURL.Path + "/saml/metadata"
	acsURL := *actURL
	acsURL.Path = acsURL.Path + "/saml/acs"

	sp := saml.ServiceProvider{
		Key:         privKey,
		Certificate: cert,
		MetadataURL: metadataURL,
		AcsURL:      acsURL,
	}

	// XML unmarshal throws an error for IdP Metadata cacheDuration field, as it's of type xml Duration. Using a separate struct for unmarshaling for now
	idm := &IDPMetadata{}
	if configToSet.IDPMetadataContent != "" {
		sp.IDPMetadata = &saml.EntityDescriptor{}
		if err := xml.NewDecoder(strings.NewReader(configToSet.IDPMetadataContent)).Decode(idm); err != nil {
			return fmt.Errorf("SAML: cannot initialize saml SP, cannot decode IDP Metadata content from the config %v, error %v", configToSet, err)
		}
	}

	sp.IDPMetadata.XMLName = idm.XMLName
	sp.IDPMetadata.ValidUntil = idm.ValidUntil
	sp.IDPMetadata.EntityID = idm.EntityID
	sp.IDPMetadata.SPSSODescriptors = idm.SPSSODescriptors
	sp.IDPMetadata.IDPSSODescriptors = idm.IDPSSODescriptors
	if name == ADFSName || name == OKTAName {
		sp.AuthnNameIDFormat = saml.UnspecifiedNameIDFormat
	}

	provider.serviceProvider = &sp

	cookieStore := samlsp.ClientCookies{
		ServiceProvider: &sp,
		Name:            "token",
		Domain:          actURL.Host,
	}
	provider.clientState = &cookieStore

	SamlProviders[name] = provider

	switch name {
	case PingName:
		root.Get("PingACS").HandlerFunc(provider.ServeHTTP)
		root.Get("PingMetadata").HandlerFunc(provider.ServeHTTP)
	case ADFSName:
		root.Get("AdfsACS").HandlerFunc(provider.ServeHTTP)
		root.Get("AdfsMetadata").HandlerFunc(provider.ServeHTTP)
	case KeyCloakName:
		root.Get("KeyCloakACS").HandlerFunc(provider.ServeHTTP)
		root.Get("KeyCloakMetadata").HandlerFunc(provider.ServeHTTP)
	case OKTAName:
		root.Get("OktaACS").HandlerFunc(provider.ServeHTTP)
		root.Get("OktaMetadata").HandlerFunc(provider.ServeHTTP)
	}

	appliedVersion = configToSet.ResourceVersion

	return nil
}

func AuthHandler() http.Handler {
	root = mux.NewRouter()
	root.Methods("POST").Path("/v1-saml/ping/saml/acs").Name("PingACS")
	root.Methods("GET").Path("/v1-saml/ping/saml/metadata").Name("PingMetadata")

	root.Methods("POST").Path("/v1-saml/adfs/saml/acs").Name("AdfsACS")
	root.Methods("GET").Path("/v1-saml/adfs/saml/metadata").Name("AdfsMetadata")

	root.Methods("POST").Path("/v1-saml/keycloak/saml/acs").Name("KeyCloakACS")
	root.Methods("GET").Path("/v1-saml/keycloak/saml/metadata").Name("KeyCloakMetadata")

	root.Methods("POST").Path("/v1-saml/okta/saml/acs").Name("OktaACS")
	root.Methods("GET").Path("/v1-saml/okta/saml/metadata").Name("OktaMetadata")

	return root
}

func (s *Provider) getSamlPrincipals(config *v3.SamlConfig, samlData map[string][]string) (v3.Principal, []v3.Principal, error) {
	var userPrincipal v3.Principal
	var groupPrincipals []v3.Principal
	uid, ok := samlData[config.UIDField]
	if !ok {
		// UID field provided by user is actually not there in SAMLResponse, without this we cannot differentiate between users and create separate principals
		return userPrincipal, groupPrincipals, fmt.Errorf("SAML: Unique ID field is not provided in SAML Response")
	}
	userPrincipal = v3.Principal{
		ObjectMeta:    metav1.ObjectMeta{Name: s.userType + "://" + uid[0]},
		Provider:      s.name,
		PrincipalType: "user",
		Me:            true,
	}

	displayName, ok := samlData[config.DisplayNameField]
	if ok {
		userPrincipal.DisplayName = displayName[0]
	}

	userName, ok := samlData[config.UserNameField]
	if ok {
		userPrincipal.LoginName = userName[0]
	}

	groups, ok := samlData[config.GroupsField]
	if ok {
		for _, group := range groups {
			group := v3.Principal{
				ObjectMeta:    metav1.ObjectMeta{Name: s.groupType + "://" + group},
				DisplayName:   group,
				Provider:      s.name,
				PrincipalType: "group",
				MemberOf:      true,
			}
			groupPrincipals = append(groupPrincipals, group)
		}
	}

	return userPrincipal, groupPrincipals, nil
}

// HandleSamlAssertion processes/handles the assertion obtained by the POST to /saml/acs from IdP
func (s *Provider) HandleSamlAssertion(w http.ResponseWriter, r *http.Request, assertion *saml.Assertion) {
	var groupPrincipals []v3.Principal
	var userPrincipal v3.Principal

	if relayState := r.Form.Get("RelayState"); relayState != "" {
		// delete the cookie
		s.clientState.DeleteState(w, r, relayState)
	}

	redirectURL := s.clientState.GetState(r, "Rancher_FinalRedirectURL")
	rancherAction := s.clientState.GetState(r, "Rancher_Action")
	if rancherAction == loginAction {
		redirectURL += "/login?"
	} else if rancherAction == testAndEnableAction {
		// the first query param is config=saml_provider_name set by UI
		redirectURL += "&"
	}

	samlData := make(map[string][]string)

	for _, attributeStatement := range assertion.AttributeStatements {
		for _, attr := range attributeStatement.Attributes {
			attrName := attr.FriendlyName
			if attrName == "" {
				attrName = attr.Name
			}
			for _, value := range attr.Values {
				samlData[attrName] = append(samlData[attrName], value.Value)
			}
		}
	}

	config, err := s.getSamlConfig()
	if err != nil {
		log.Errorf("SAML: Error getting saml config %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		return
	}

	userPrincipal, groupPrincipals, err = s.getSamlPrincipals(config, samlData)
	if err != nil {
		log.Error(err)
		// UI uses this translation key to get the error message
		http.Redirect(w, r, redirectURL+"errorCode=422&errorMsg="+UITranslationKeyForErrorMessage, http.StatusFound)
		return
	}
	allowedPrincipals := config.AllowedPrincipalIDs

	allowed, err := s.userMGR.CheckAccess(config.AccessMode, allowedPrincipals, userPrincipal.Name, groupPrincipals)
	if err != nil {
		log.Errorf("SAML: Error during login while checking access %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		return
	}
	if !allowed {
		log.Errorf("SAML: User does not have access %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=403", http.StatusFound)
		return
	}

	userID := s.clientState.GetState(r, "Rancher_UserID")
	if userID != "" && rancherAction == testAndEnableAction {
		user, err := s.userMGR.SetPrincipalOnCurrentUserByUserID(userID, userPrincipal)
		if err != nil {
			log.Errorf("SAML: Error setting principal on current user %v", err)
			http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
			return
		}

		config.Enabled = true
		err = s.saveSamlConfig(config)
		if err != nil {
			log.Errorf("SAML: Error saving saml config %v", err)
			http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
			return
		}

		isSecure := false
		if r.URL.Scheme == "https" {
			isSecure = true
		}
		err = setRancherToken(w, r, s.tokenMGR, user.Name, userPrincipal, groupPrincipals, isSecure)
		if err != nil {
			log.Errorf("SAML: Failed creating token with error: %v", err)
			http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		}
		// delete the cookies
		s.clientState.DeleteState(w, r, "Rancher_UserID")
		s.clientState.DeleteState(w, r, "Rancher_Action")
		redirectURL := s.clientState.GetState(r, "Rancher_FinalRedirectURL")
		s.clientState.DeleteState(w, r, "Rancher_FinalRedirectURL")
		if redirectURL != "" {
			// delete the cookie
			s.clientState.DeleteState(w, r, "Rancher_FinalRedirectURL")
			http.Redirect(w, r, redirectURL, http.StatusFound)
		}
		return
	}

	displayName := userPrincipal.DisplayName
	if displayName == "" {
		displayName = userPrincipal.LoginName
	}
	user, err := s.userMGR.EnsureUser(userPrincipal.Name, displayName)
	if err != nil {
		log.Errorf("SAML: Failed getting user with error: %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
		return
	}

	if user.Enabled != nil && !*user.Enabled {
		log.Errorf("SAML: User %v permission denied", user.Name)
		http.Redirect(w, r, redirectURL+"errorCode=403", http.StatusFound)
		return
	}

	err = setRancherToken(w, r, s.tokenMGR, user.Name, userPrincipal, groupPrincipals, true)
	if err != nil {
		log.Errorf("SAML: Failed creating token with error: %v", err)
		http.Redirect(w, r, redirectURL+"errorCode=500", http.StatusFound)
	}
	redirectURL = s.clientState.GetState(r, "Rancher_FinalRedirectURL")
	if redirectURL != "" {
		// delete the cookie
		s.clientState.DeleteState(w, r, "Rancher_FinalRedirectURL")
		http.Redirect(w, r, redirectURL, http.StatusFound)
	}
	return
}

func setRancherToken(w http.ResponseWriter, r *http.Request, tokenMGR *tokens.Manager, userID string, userPrincipal v3.Principal,
	groupPrincipals []v3.Principal, isSecure bool) error {
	rToken, err := tokenMGR.NewLoginToken(userID, userPrincipal, groupPrincipals, "", 0, "")
	if err != nil {
		return err
	}
	tokenCookie := &http.Cookie{
		Name:     "R_SESS",
		Value:    rToken.ObjectMeta.Name + ":" + rToken.Token,
		Secure:   isSecure,
		Path:     "/",
		HttpOnly: true,
	}
	http.SetCookie(w, tokenCookie)
	return nil
}
<|endoftext|>