<commit_before>package s3api

import (
	"crypto/md5"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/s3api/s3err"

	"github.com/gorilla/mux"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	weed_server "github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConns:        1024,
		MaxIdleConnsPerHost: 1024,
	}}
}

func (s3a *S3ApiServer) PutObjectHandler(w http.ResponseWriter, r *http.Request) {

	// http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html

	bucket, object := getBucketAndObject(r)

	_, err := validateContentMd5(r.Header)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInvalidDigest, r.URL)
		return
	}

	dataReader := r.Body
	if s3a.iam.isEnabled() {
		rAuthType := getRequestAuthType(r)
		var s3ErrCode s3err.ErrorCode
		switch rAuthType {
		case authTypeStreamingSigned:
			dataReader, s3ErrCode = s3a.iam.newSignV4ChunkedReader(r)
		case authTypeSignedV2, authTypePresignedV2:
			_, s3ErrCode = s3a.iam.isReqAuthenticatedV2(r)
		case authTypePresigned, authTypeSigned:
			_, s3ErrCode = s3a.iam.reqSignatureV4Verify(r)
		}
		if s3ErrCode != s3err.ErrNone {
			writeErrorResponse(w, s3ErrCode, r.URL)
			return
		}
	}
	defer dataReader.Close()

	if strings.HasSuffix(object, "/") {
		if err := s3a.mkdir(s3a.option.BucketsPath, bucket+object, nil); err != nil {
			writeErrorResponse(w, s3err.ErrInternalError, r.URL)
			return
		}
	} else {
		uploadUrl := fmt.Sprintf("http://%s%s/%s%s", s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

		etag, errCode := s3a.putToFiler(r, uploadUrl, dataReader)

		if errCode != s3err.ErrNone {
			writeErrorResponse(w, errCode, r.URL)
			return
		}

		setEtag(w, etag)
	}

	writeSuccessResponseEmpty(w)
}

func urlPathEscape(object string) string {
	var escapedParts []string
	for _, part := range strings.Split(object, "/") {
		escapedParts = append(escapedParts, url.PathEscape(part))
	}
	return strings.Join(escapedParts, "/")
}

func (s3a *S3ApiServer) GetObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	if strings.HasSuffix(r.URL.Path, "/") {
		writeErrorResponse(w, s3err.ErrNotImplemented, r.URL)
		return
	}

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) HeadObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) DeleteObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s?recursive=true",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

	s3a.proxyToFiler(w, r, destUrl, func(proxyResponse *http.Response, w http.ResponseWriter) {
		for k, v := range proxyResponse.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(http.StatusNoContent)
	})
}

// / ObjectIdentifier carries key name for the object to delete.
type ObjectIdentifier struct {
	ObjectName string `xml:"Key"`
}

// DeleteObjectsRequest - xml carrying the object key names which needs to be deleted.
type DeleteObjectsRequest struct {
	// Element to enable quiet mode for the request
	Quiet bool
	// List of objects to be deleted
	Objects []ObjectIdentifier `xml:"Object"`
}

// DeleteError structure.
type DeleteError struct {
	Code    string
	Message string
	Key     string
}

// DeleteObjectsResponse container for multiple object deletes.
type DeleteObjectsResponse struct {
	XMLName xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ DeleteResult" json:"-"`

	// Collection of all deleted objects
	DeletedObjects []ObjectIdentifier `xml:"Deleted,omitempty"`

	// Collection of errors deleting certain objects.
	Errors []DeleteError `xml:"Error,omitempty"`
}

// DeleteMultipleObjectsHandler - Delete multiple objects
func (s3a *S3ApiServer) DeleteMultipleObjectsHandler(w http.ResponseWriter, r *http.Request) {

	bucket, _ := getBucketAndObject(r)

	deleteXMLBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	deleteObjects := &DeleteObjectsRequest{}
	if err := xml.Unmarshal(deleteXMLBytes, deleteObjects); err != nil {
		writeErrorResponse(w, s3err.ErrMalformedXML, r.URL)
		return
	}

	var deletedObjects []ObjectIdentifier
	var deleteErrors []DeleteError

	directoriesWithDeletion := make(map[string]int)

	s3a.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		// delete file entries
		for _, object := range deleteObjects.Objects {

			lastSeparator := strings.LastIndex(object.ObjectName, "/")
			parentDirectoryPath, entryName, isDeleteData, isRecursive := "", object.ObjectName, true, false
			if lastSeparator > 0 && lastSeparator+1 < len(object.ObjectName) {
				entryName = object.ObjectName[lastSeparator+1:]
				parentDirectoryPath = "/" + object.ObjectName[:lastSeparator]
			}
			parentDirectoryPath = fmt.Sprintf("%s/%s%s", s3a.option.BucketsPath, bucket, parentDirectoryPath)

			err := doDeleteEntry(client, parentDirectoryPath, entryName, isDeleteData, isRecursive)
			if err == nil {
				directoriesWithDeletion[parentDirectoryPath]++
				deletedObjects = append(deletedObjects, object)
			} else if strings.Contains(err.Error(), filer.MsgFailDelNonEmptyFolder) {
				deletedObjects = append(deletedObjects, object)
			} else {
				delete(directoriesWithDeletion, parentDirectoryPath)
				deleteErrors = append(deleteErrors, DeleteError{
					Code:    "",
					Message: err.Error(),
					Key:     object.ObjectName,
				})
			}
		}

		// purge empty folders, only checking folders with deletions
		for len(directoriesWithDeletion) > 0 {
			directoriesWithDeletion = s3a.doDeleteEmptyDirectories(client, directoriesWithDeletion)
		}

		return nil
	})

	deleteResp := DeleteObjectsResponse{}
	if !deleteObjects.Quiet {
		deleteResp.DeletedObjects = deletedObjects
	}
	deleteResp.Errors = deleteErrors

	writeSuccessResponseXML(w, encodeResponse(deleteResp))

}

func (s3a *S3ApiServer) doDeleteEmptyDirectories(client filer_pb.SeaweedFilerClient, directoriesWithDeletion map[string]int) (newDirectoriesWithDeletion map[string]int) {
	var allDirs []string
	for dir, _ := range directoriesWithDeletion {
		allDirs = append(allDirs, dir)
	}
	sort.Slice(allDirs, func(i, j int) bool {
		return len(allDirs[i]) > len(allDirs[j])
	})
	newDirectoriesWithDeletion = make(map[string]int)
	for _, dir := range allDirs {
		parentDir, dirName := util.FullPath(dir).DirAndName()
		if parentDir == s3a.option.BucketsPath {
			continue
		}
		if err := doDeleteEntry(client, parentDir, dirName, false, false); err != nil {
			glog.V(4).Infof("directory %s has %d deletion but still not empty: %v", dir, directoriesWithDeletion[dir], err)
		} else {
			newDirectoriesWithDeletion[parentDir]++
		}
	}
	return
}

var passThroughHeaders = []string{
	"response-cache-control",
	"response-content-disposition",
	"response-content-encoding",
	"response-content-language",
	"response-content-type",
	"response-expires",
}

func (s3a *S3ApiServer) proxyToFiler(w http.ResponseWriter, r *http.Request, destUrl string, responseFn func(proxyResponse *http.Response, w http.ResponseWriter)) {

	glog.V(2).Infof("s3 proxying %s to %s", r.Method, destUrl)

	proxyReq, err := http.NewRequest(r.Method, destUrl, r.Body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", destUrl, err)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		// handle s3 related headers
		passed := false
		for _, h := range passThroughHeaders {
			if strings.ToLower(header) == h && len(values) > 0 {
				proxyReq.Header.Add(header[len("response-"):], values[0])
				passed = true
				break
			}
		}
		if passed {
			continue
		}
		// handle other headers
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}
	defer util.CloseResponse(resp)

	if resp.ContentLength == -1 || resp.StatusCode == 404 {
		if r.Method != "DELETE" {
			writeErrorResponse(w, s3err.ErrNoSuchKey, r.URL)
			return
		}
	}

	responseFn(resp, w)

}

func passThroughResponse(proxyResponse *http.Response, w http.ResponseWriter) {
	for k, v := range proxyResponse.Header {
		w.Header()[k] = v
	}
	if proxyResponse.Header.Get("Content-Range") != "" && proxyResponse.StatusCode == 200 {
		w.WriteHeader(http.StatusPartialContent)
	} else {
		w.WriteHeader(proxyResponse.StatusCode)
	}
	io.Copy(w, proxyResponse.Body)
}

func (s3a *S3ApiServer) putToFiler(r *http.Request, uploadUrl string, dataReader io.Reader) (etag string, code s3err.ErrorCode) {

	hash := md5.New()
	var body = io.TeeReader(dataReader, hash)

	proxyReq, err := http.NewRequest("PUT", uploadUrl, body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", uploadUrl, err)
		return "", s3err.ErrInternalError
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		return "", s3err.ErrInternalError
	}
	defer resp.Body.Close()

	etag = fmt.Sprintf("%x", hash.Sum(nil))

	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		glog.Errorf("upload to filer response read %d: %v", resp.StatusCode, ra_err)
		return etag, s3err.ErrInternalError
	}
	var ret weed_server.FilerPostResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.Errorf("failing to read upload to %s : %v", uploadUrl, string(resp_body))
		return "", s3err.ErrInternalError
	}
	if ret.Error != "" {
		glog.Errorf("upload to filer error: %v", ret.Error)
		return "", filerErrorToS3Error(ret.Error)
	}

	return etag, s3err.ErrNone
}

func setEtag(w http.ResponseWriter, etag string) {
	if etag != "" {
		if strings.HasPrefix(etag, "\"") {
			w.Header().Set("ETag", etag)
		} else {
			w.Header().Set("ETag", "\""+etag+"\"")
		}
	}
}

func getBucketAndObject(r *http.Request) (bucket, object string) {
	vars := mux.Vars(r)
	bucket = vars["bucket"]
	object = vars["object"]
	if !strings.HasPrefix(object, "/") {
		object = "/" + object
	}

	return
}

func filerErrorToS3Error(errString string) s3err.ErrorCode {
	if strings.HasPrefix(errString, "existing ") && strings.HasSuffix(errString, "is a directory") {
		return s3err.ErrExistingObjectIsDirectory
	}
	return s3err.ErrInternalError
}
<commit_msg>s3api: handle 304 response code from filer<commit_after>package s3api

import (
	"crypto/md5"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/s3api/s3err"

	"github.com/gorilla/mux"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	weed_server "github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConns:        1024,
		MaxIdleConnsPerHost: 1024,
	}}
}

func (s3a *S3ApiServer) PutObjectHandler(w http.ResponseWriter, r *http.Request) {

	// http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html

	bucket, object := getBucketAndObject(r)

	_, err := validateContentMd5(r.Header)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInvalidDigest, r.URL)
		return
	}

	dataReader := r.Body
	if s3a.iam.isEnabled() {
		rAuthType := getRequestAuthType(r)
		var s3ErrCode s3err.ErrorCode
		switch rAuthType {
		case authTypeStreamingSigned:
			dataReader, s3ErrCode = s3a.iam.newSignV4ChunkedReader(r)
		case authTypeSignedV2, authTypePresignedV2:
			_, s3ErrCode = s3a.iam.isReqAuthenticatedV2(r)
		case authTypePresigned, authTypeSigned:
			_, s3ErrCode = s3a.iam.reqSignatureV4Verify(r)
		}
		if s3ErrCode != s3err.ErrNone {
			writeErrorResponse(w, s3ErrCode, r.URL)
			return
		}
	}
	defer dataReader.Close()

	if strings.HasSuffix(object, "/") {
		if err := s3a.mkdir(s3a.option.BucketsPath, bucket+object, nil); err != nil {
			writeErrorResponse(w, s3err.ErrInternalError, r.URL)
			return
		}
	} else {
		uploadUrl := fmt.Sprintf("http://%s%s/%s%s", s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

		etag, errCode := s3a.putToFiler(r, uploadUrl, dataReader)

		if errCode != s3err.ErrNone {
			writeErrorResponse(w, errCode, r.URL)
			return
		}

		setEtag(w, etag)
	}

	writeSuccessResponseEmpty(w)
}

func urlPathEscape(object string) string {
	var escapedParts []string
	for _, part := range strings.Split(object, "/") {
		escapedParts = append(escapedParts, url.PathEscape(part))
	}
	return strings.Join(escapedParts, "/")
}

func (s3a *S3ApiServer) GetObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	if strings.HasSuffix(r.URL.Path, "/") {
		writeErrorResponse(w, s3err.ErrNotImplemented, r.URL)
		return
	}

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) HeadObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) DeleteObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s?recursive=true",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, urlPathEscape(object))

	s3a.proxyToFiler(w, r, destUrl, func(proxyResponse *http.Response, w http.ResponseWriter) {
		for k, v := range proxyResponse.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(http.StatusNoContent)
	})
}

// / ObjectIdentifier carries key name for the object to delete.
type ObjectIdentifier struct {
	ObjectName string `xml:"Key"`
}

// DeleteObjectsRequest - xml carrying the object key names which needs to be deleted.
type DeleteObjectsRequest struct {
	// Element to enable quiet mode for the request
	Quiet bool
	// List of objects to be deleted
	Objects []ObjectIdentifier `xml:"Object"`
}

// DeleteError structure.
type DeleteError struct {
	Code    string
	Message string
	Key     string
}

// DeleteObjectsResponse container for multiple object deletes.
type DeleteObjectsResponse struct {
	XMLName xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ DeleteResult" json:"-"`

	// Collection of all deleted objects
	DeletedObjects []ObjectIdentifier `xml:"Deleted,omitempty"`

	// Collection of errors deleting certain objects.
	Errors []DeleteError `xml:"Error,omitempty"`
}

// DeleteMultipleObjectsHandler - Delete multiple objects
func (s3a *S3ApiServer) DeleteMultipleObjectsHandler(w http.ResponseWriter, r *http.Request) {

	bucket, _ := getBucketAndObject(r)

	deleteXMLBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	deleteObjects := &DeleteObjectsRequest{}
	if err := xml.Unmarshal(deleteXMLBytes, deleteObjects); err != nil {
		writeErrorResponse(w, s3err.ErrMalformedXML, r.URL)
		return
	}

	var deletedObjects []ObjectIdentifier
	var deleteErrors []DeleteError

	directoriesWithDeletion := make(map[string]int)

	s3a.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		// delete file entries
		for _, object := range deleteObjects.Objects {

			lastSeparator := strings.LastIndex(object.ObjectName, "/")
			parentDirectoryPath, entryName, isDeleteData, isRecursive := "", object.ObjectName, true, false
			if lastSeparator > 0 && lastSeparator+1 < len(object.ObjectName) {
				entryName = object.ObjectName[lastSeparator+1:]
				parentDirectoryPath = "/" + object.ObjectName[:lastSeparator]
			}
			parentDirectoryPath = fmt.Sprintf("%s/%s%s", s3a.option.BucketsPath, bucket, parentDirectoryPath)

			err := doDeleteEntry(client, parentDirectoryPath, entryName, isDeleteData, isRecursive)
			if err == nil {
				directoriesWithDeletion[parentDirectoryPath]++
				deletedObjects = append(deletedObjects, object)
			} else if strings.Contains(err.Error(), filer.MsgFailDelNonEmptyFolder) {
				deletedObjects = append(deletedObjects, object)
			} else {
				delete(directoriesWithDeletion, parentDirectoryPath)
				deleteErrors = append(deleteErrors, DeleteError{
					Code:    "",
					Message: err.Error(),
					Key:     object.ObjectName,
				})
			}
		}

		// purge empty folders, only checking folders with deletions
		for len(directoriesWithDeletion) > 0 {
			directoriesWithDeletion = s3a.doDeleteEmptyDirectories(client, directoriesWithDeletion)
		}

		return nil
	})

	deleteResp := DeleteObjectsResponse{}
	if !deleteObjects.Quiet {
		deleteResp.DeletedObjects = deletedObjects
	}
	deleteResp.Errors = deleteErrors

	writeSuccessResponseXML(w, encodeResponse(deleteResp))

}

func (s3a *S3ApiServer) doDeleteEmptyDirectories(client filer_pb.SeaweedFilerClient, directoriesWithDeletion map[string]int) (newDirectoriesWithDeletion map[string]int) {
	var allDirs []string
	for dir, _ := range directoriesWithDeletion {
		allDirs = append(allDirs, dir)
	}
	sort.Slice(allDirs, func(i, j int) bool {
		return len(allDirs[i]) > len(allDirs[j])
	})
	newDirectoriesWithDeletion = make(map[string]int)
	for _, dir := range allDirs {
		parentDir, dirName := util.FullPath(dir).DirAndName()
		if parentDir == s3a.option.BucketsPath {
			continue
		}
		if err := doDeleteEntry(client, parentDir, dirName, false, false); err != nil {
			glog.V(4).Infof("directory %s has %d deletion but still not empty: %v", dir, directoriesWithDeletion[dir], err)
		} else {
			newDirectoriesWithDeletion[parentDir]++
		}
	}
	return
}

var passThroughHeaders = []string{
	"response-cache-control",
	"response-content-disposition",
	"response-content-encoding",
	"response-content-language",
	"response-content-type",
	"response-expires",
}

func (s3a *S3ApiServer) proxyToFiler(w http.ResponseWriter, r *http.Request, destUrl string, responseFn func(proxyResponse *http.Response, w http.ResponseWriter)) {

	glog.V(2).Infof("s3 proxying %s to %s", r.Method, destUrl)

	proxyReq, err := http.NewRequest(r.Method, destUrl, r.Body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", destUrl, err)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		// handle s3 related headers
		passed := false
		for _, h := range passThroughHeaders {
			if strings.ToLower(header) == h && len(values) > 0 {
				proxyReq.Header.Add(header[len("response-"):], values[0])
				passed = true
				break
			}
		}
		if passed {
			continue
		}
		// handle other headers
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}
	defer util.CloseResponse(resp)

	if (resp.ContentLength == -1 || resp.StatusCode == 404) && resp.StatusCode != 304 {
		if r.Method != "DELETE" {
			writeErrorResponse(w, s3err.ErrNoSuchKey, r.URL)
			return
		}
	}

	responseFn(resp, w)

}

func passThroughResponse(proxyResponse *http.Response, w http.ResponseWriter) {
	for k, v := range proxyResponse.Header {
		w.Header()[k] = v
	}
	if proxyResponse.Header.Get("Content-Range") != "" && proxyResponse.StatusCode == 200 {
		w.WriteHeader(http.StatusPartialContent)
	} else {
		w.WriteHeader(proxyResponse.StatusCode)
	}
	io.Copy(w, proxyResponse.Body)
}

func (s3a *S3ApiServer) putToFiler(r *http.Request, uploadUrl string, dataReader io.Reader) (etag string, code s3err.ErrorCode) {

	hash := md5.New()
	var body = io.TeeReader(dataReader, hash)

	proxyReq, err := http.NewRequest("PUT", uploadUrl, body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", uploadUrl, err)
		return "", s3err.ErrInternalError
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		return "", s3err.ErrInternalError
	}
	defer resp.Body.Close()

	etag = fmt.Sprintf("%x", hash.Sum(nil))

	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		glog.Errorf("upload to filer response read %d: %v", resp.StatusCode, ra_err)
		return etag, s3err.ErrInternalError
	}
	var ret weed_server.FilerPostResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.Errorf("failing to read upload to %s : %v", uploadUrl, string(resp_body))
		return "", s3err.ErrInternalError
	}
	if ret.Error != "" {
		glog.Errorf("upload to filer error: %v", ret.Error)
		return "", filerErrorToS3Error(ret.Error)
	}

	return etag, s3err.ErrNone
}

func setEtag(w http.ResponseWriter, etag string) {
	if etag != "" {
		if strings.HasPrefix(etag, "\"") {
			w.Header().Set("ETag", etag)
		} else {
			w.Header().Set("ETag", "\""+etag+"\"")
		}
	}
}

func getBucketAndObject(r *http.Request) (bucket, object string) {
	vars := mux.Vars(r)
	bucket = vars["bucket"]
	object = vars["object"]
	if !strings.HasPrefix(object, "/") {
		object = "/" + object
	}

	return
}

func filerErrorToS3Error(errString string) s3err.ErrorCode {
	if strings.HasPrefix(errString, "existing ") && strings.HasSuffix(errString, "is a directory") {
		return s3err.ErrExistingObjectIsDirectory
	}
	return s3err.ErrInternalError
}
<|endoftext|>


<commit_before>package graph

import (
	"fmt"
	"strings"
)

type Cluster interface {
	isCluster()
	Dump(margin string)
	DumpShort() string
}

type ClusterLeaf struct {
	Nodes []NodeID
}

func (cluster *ClusterLeaf) isCluster() {
}

func (cluster *ClusterLeaf) Dump(margin string) {
	fmt.Printf("%sleaf %d\n", margin, len(cluster.Nodes))
	childMargin := margin + ".   "
	for _, n := range cluster.Nodes {
		fmt.Printf("%s%d\n", childMargin, n)
	}
}

func (cluster *ClusterLeaf) DumpShort() string {
	text := make([]string, len(cluster.Nodes))
	for i, n := range cluster.Nodes {
		text[i] = fmt.Sprintf("%d", n)
	}
	return fmt.Sprintf("(%s)", strings.Join(text, " "))
}

type ClusterLinear struct {
	Clusters []Cluster
}

func (cluster *ClusterLinear) isCluster() {
}

func (cluster *ClusterLinear) Dump(margin string) {
	fmt.Printf("%slinear %d\n", margin, len(cluster.Clusters))
	childMargin := margin + ".   "
	for _, c := range cluster.Clusters {
		c.Dump(childMargin)
	}
}

func (cluster *ClusterLinear) DumpShort() string {
	text := make([]string, len(cluster.Clusters))
	for i, c := range cluster.Clusters {
		text[i] = c.DumpShort()
	}
	return fmt.Sprintf("[%s]", strings.Join(text, " "))
}

type ClusterSwitch struct {
	Children []Cluster
}

func (cluster *ClusterSwitch) isCluster() {
}

func (cluster *ClusterSwitch) Dump(margin string) {
	fmt.Printf("%sswitch %d\n", margin, len(cluster.Children))
	childMargin := margin + ".   "
	for _, c := range cluster.Children {
		c.Dump(childMargin)
	}
}

func (cluster *ClusterSwitch) DumpShort() string {
	text := make([]string, len(cluster.Children))
	for i, c := range cluster.Children {
		text[i] = c.DumpShort()
	}
	return fmt.Sprintf("<%s>", strings.Join(text, " "))
}

type ClusterLoop struct {
	Body Cluster
}

func (cluster *ClusterLoop) isCluster() {
}

func (cluster *ClusterLoop) Dump(margin string) {
	fmt.Printf("%sloop\n", margin)
	childMargin := margin + ".   "
	cluster.Body.Dump(childMargin)
}

func (cluster *ClusterLoop) DumpShort() string {
	return fmt.Sprintf("{%s}", cluster.Body.DumpShort())
}

func makeCluster(g *Graph) Cluster {
	info, edges, postorder := analyzeStructure(g)
	return makeCluster2(g, info, edges, postorder)
}

type nodeInfo struct {
	pre           int
	post          int
	idom          NodeID
	loopHead      NodeID
	isHead        bool
	isIrreducible bool
	live          bool
}

type edgeType int

const (
	DEAD edgeType = iota
	FORWARD
	BACKWARD
	CROSS
	REENTRY
)

type loopFinder struct {
	graph     *Graph
	node      []nodeInfo
	edge      []edgeType
	postorder []NodeID
	depth     []int
	current   int
}

func (lf *loopFinder) beginTraversingNode(n NodeID, prev NodeID) {
	lf.current += 1

	lf.node[n] = nodeInfo{
		pre:      lf.current,
		idom:     prev, // A reasonable inital guess.
		loopHead: NoNode,
		live:     true,
	}

	lf.depth[n] = lf.current
}

func (lf *loopFinder) endTraversingNode(n NodeID) {
	lf.current += 1
	lf.node[n].post = lf.current

	lf.depth[n] = 0

	lf.postorder = append(lf.postorder, n)
}

func (lf *loopFinder) isUnprocessed(n NodeID) bool {
	return !lf.node[n].live
}

func (lf *loopFinder) isBeingProcessed(dst NodeID) bool {
	return lf.depth[dst] > 0
}

func (lf *loopFinder) markLoopHeader(child NodeID, head NodeID) {
	if child == head {
		// Trivial loop.
		return
	}
	childHead := lf.node[child].loopHead
	for childHead != NoNode {
		if childHead == head {
			return
		}
		if lf.depth[childHead] < lf.depth[head] {
			// Found a closer head for this child, adopt it.
			lf.node[child].loopHead = head
			child, head = head, childHead
		} else {
			child = childHead
		}
		childHead = lf.node[child].loopHead
	}
	lf.node[child].loopHead = head
}

func (lf *loopFinder) process(n NodeID, prev NodeID) {
	lf.beginTraversingNode(n, prev)
	xit := lf.graph.ExitIterator(n)
	for xit.HasNext() {
		e, next := xit.GetNext()
		if lf.isUnprocessed(next) {
			lf.edge[e] = FORWARD
			lf.process(next, n)

			// Propagage loop headers upwards.
			head := lf.node[next].loopHead
			if head != NoNode {
				lf.markLoopHeader(n, head)
			}
		} else if lf.isBeingProcessed(next) {
			lf.edge[e] = BACKWARD
			lf.node[next].isHead = true
			lf.markLoopHeader(n, next)
		} else {
			lf.edge[e] = CROSS
			if lf.node[next].loopHead != NoNode {
				// Propagate loop header from cross edge.
				otherHead := lf.node[next].loopHead
				if lf.isBeingProcessed(otherHead) {
					lf.markLoopHeader(n, otherHead)
				} else {
					lf.edge[e] = REENTRY
					lf.node[otherHead].isIrreducible = true
					// Find and mark the common loop head.
					otherHead = lf.node[otherHead].loopHead
					for otherHead != NoNode {
						if lf.isBeingProcessed(otherHead) {
							lf.markLoopHeader(n, otherHead)
							break
						}
						otherHead = lf.node[otherHead].loopHead
					}
				}
			}
		}
	}
	lf.endTraversingNode(n)
}

func (lf *loopFinder) intersect(n0 NodeID, n1 NodeID) NodeID {
	i0 := lf.node[n0].post
	i1 := lf.node[n1].post
	for i0 != i1 {
		for i0 < i1 {
			n0 = lf.node[n0].idom
			i0 = lf.node[n0].post
		}
		for i0 > i1 {
			n1 = lf.node[n1].idom
			i1 = lf.node[n1].post
		}
	}
	return n0
}

func (lf *loopFinder) cleanDeadEdges() {
	numEdges := lf.graph.NumEdges()
	for i := 0; i < numEdges; i++ {
		if lf.edge[i] == DEAD {
			lf.graph.KillEdge(EdgeID(i))
		}
	}
}

func (lf *loopFinder) findIdoms() {
	g := lf.graph
	changed := true
	for changed {
		changed = false
		for i := len(lf.postorder) - 1; i >= 0; i-- {
			n := lf.postorder[i]
			original := lf.node[n].idom
			idom := original
			eit := g.EntryIterator(n)
			for eit.HasNext() {
				src, _ := eit.GetNext()
				idom = lf.intersect(idom, src)
			}
			if idom != original {
				lf.node[n].idom = idom
				changed = true
			}
		}
	}
}

func analyzeStructure(g *Graph) ([]nodeInfo, []edgeType, []NodeID) {
	numNodes := g.NumNodes()
	lf := &loopFinder{
		graph: g.Copy(),
		node:  make([]nodeInfo, numNodes),
		edge:  make([]edgeType, g.NumEdges()),
		depth: make([]int, numNodes),
	}

	e := g.Entry()
	lf.process(e, e)
	lf.cleanDeadEdges()
	lf.findIdoms()
	return lf.node, lf.edge, lf.postorder
}

func uniqueEntry(g *Graph, src NodeID, dst NodeID) bool {
	eit := g.EntryIterator(dst)
	for eit.HasNext() {
		prev, _ := eit.GetNext()
		if prev != src {
			return false
		}
	}
	return true
}

func contract(g *Graph, src NodeID, dst NodeID, nodes []nodeInfo) {
	transferedExits := false
	eit := g.EntryIterator(dst)
	for eit.HasNext() {
		prev, e := eit.GetNext()
		if prev != src {
			panic(prev)
		}
		if !transferedExits {
			g.ReplaceEdgeWithExits(e, dst)
			transferedExits = true
		}
		g.KillEdge(e)
	}
	if !transferedExits {
		panic(dst)
	}
}

func contractLoop(g *Graph, n NodeID) {
	xit := g.ExitIterator(n)
	for xit.HasNext() {
		e, dst := xit.GetNext()
		if n == dst {
			g.KillEdge(e)
		}
	}
}

func appendCluster(src Cluster, dst Cluster, shouldFuse bool) Cluster {
	if shouldFuse {
		switch src := src.(type) {
		case *ClusterLeaf:
			switch dst := dst.(type) {
			case *ClusterLeaf:
				src.Nodes = append(src.Nodes, dst.Nodes...)
				return src
			case *ClusterLinear:
				other, ok := dst.Clusters[0].(*ClusterLeaf)
				if ok {
					other.Nodes = append(src.Nodes, other.Nodes...)
					return dst
				}
				dst.Clusters = append([]Cluster{src}, dst.Clusters...)
				return dst
			}
		case *ClusterLinear:
			switch dst := dst.(type) {
			case *ClusterLinear:
				src.Clusters = append(src.Clusters, dst.Clusters...)
				return src
			}
			src.Clusters = append(src.Clusters, dst)
			return src
		default:
			switch dst := dst.(type) {
			case *ClusterLinear:
				dst.Clusters = append([]Cluster{src}, dst.Clusters...)
				return dst
			}
		}
	} else {
		switch src := src.(type) {
		case *ClusterLinear:
			src.Clusters = append(src.Clusters, dst)
			return src
		}
	}
	return &ClusterLinear{Clusters: []Cluster{src, dst}}
}

func isClusterHead(g *Graph, src NodeID, nodes []nodeInfo, edges []edgeType) bool {
	xit := g.ExitIterator(src)
	for xit.HasNext() {
		e, dst := xit.GetNext()
		if edges[e] != BACKWARD && nodes[dst].idom != src {
			return false
		}
	}
	return true
}

func clusterRegion(g *Graph, n NodeID, currentHead NodeID, cluster Cluster, clusters []Cluster, nodes []nodeInfo, edges []edgeType) Cluster {
	for {
		ready := []NodeID{}
		readyClusters := []Cluster{}
		//readyEdges := []EdgeID{}
		//pendingEdges := []EdgeID{}
		numExits := 0
		xit := g.ExitIterator(n)
		for xit.HasNext() {
			e, dst := xit.GetNext()
			numExits += 1
			if nodes[dst].loopHead != currentHead {
				continue
			}
			if edges[e] == BACKWARD {
				continue
			}
			if uniqueEntry(g, n, dst) {
				if clusters[dst] != nil {
					ready = append(ready, dst)
					readyClusters = append(readyClusters, clusters[dst])
					clusters[dst] = nil
				}
			}
		}
		if len(ready) > 0 {
			if len(ready) > 1 {
				cluster = appendCluster(cluster, &ClusterSwitch{Children: readyClusters}, false)
			} else {
				cluster = appendCluster(cluster, readyClusters[0], numExits == 1)
			}
			for _, dst := range ready {
				contract(g, n, dst, nodes)
			}
		} else {
			return cluster
		}
	}
}

func makeCluster2(g *Graph, nodes []nodeInfo, edges []edgeType, postorder []NodeID) Cluster {
	g = g.Copy()

	clusters := make([]Cluster, g.NumNodes())

	for _, n := range postorder {
		var cluster Cluster = &ClusterLeaf{
			Nodes: []NodeID{n},
		}
		if nodes[n].isHead {
			cluster = clusterRegion(g, n, n, cluster, clusters, nodes, edges)
			contractLoop(g, n)
			cluster = &ClusterLoop{Body: cluster}
		}
		if isClusterHead(g, n, nodes, edges) {
			cluster = clusterRegion(g, n, nodes[n].loopHead, cluster, clusters, nodes, edges)
		}
		clusters[n] = cluster
	}
	return clusters[0]
}
<commit_msg>Fix graph clustering.<commit_after>package graph

import (
	"fmt"
	"strings"
)

type Cluster interface {
	isCluster()
	Dump(margin string)
	DumpShort() string
}

type ClusterLeaf struct {
	Nodes []NodeID
}

func (cluster *ClusterLeaf) isCluster() {
}

func (cluster *ClusterLeaf) Dump(margin string) {
	fmt.Printf("%sleaf %d\n", margin, len(cluster.Nodes))
	childMargin := margin + ".   "
	for _, n := range cluster.Nodes {
		fmt.Printf("%s%d\n", childMargin, n)
	}
}

func (cluster *ClusterLeaf) DumpShort() string {
	text := make([]string, len(cluster.Nodes))
	for i, n := range cluster.Nodes {
		text[i] = fmt.Sprintf("%d", n)
	}
	return fmt.Sprintf("(%s)", strings.Join(text, " "))
}

type ClusterLinear struct {
	Clusters []Cluster
}

func (cluster *ClusterLinear) isCluster() {
}

func (cluster *ClusterLinear) Dump(margin string) {
	fmt.Printf("%slinear %d\n", margin, len(cluster.Clusters))
	childMargin := margin + ".   "
	for _, c := range cluster.Clusters {
		c.Dump(childMargin)
	}
}

func (cluster *ClusterLinear) DumpShort() string {
	text := make([]string, len(cluster.Clusters))
	for i, c := range cluster.Clusters {
		text[i] = c.DumpShort()
	}
	return fmt.Sprintf("[%s]", strings.Join(text, " "))
}

type ClusterSwitch struct {
	Children []Cluster
}

func (cluster *ClusterSwitch) isCluster() {
}

func (cluster *ClusterSwitch) Dump(margin string) {
	fmt.Printf("%sswitch %d\n", margin, len(cluster.Children))
	childMargin := margin + ".   "
	for _, c := range cluster.Children {
		c.Dump(childMargin)
	}
}

func (cluster *ClusterSwitch) DumpShort() string {
	text := make([]string, len(cluster.Children))
	for i, c := range cluster.Children {
		text[i] = c.DumpShort()
	}
	return fmt.Sprintf("<%s>", strings.Join(text, " "))
}

type ClusterLoop struct {
	Body Cluster
}

func (cluster *ClusterLoop) isCluster() {
}

func (cluster *ClusterLoop) Dump(margin string) {
	fmt.Printf("%sloop\n", margin)
	childMargin := margin + ".   "
	cluster.Body.Dump(childMargin)
}

func (cluster *ClusterLoop) DumpShort() string {
	return fmt.Sprintf("{%s}", cluster.Body.DumpShort())
}

func makeCluster(g *Graph) Cluster {
	g = g.Copy()
	info, edges, postorder := analyzeStructure(g)
	return makeCluster2(g, info, edges, postorder)
}

type nodeInfo struct {
	pre           int
	post          int
	idom          NodeID
	loopHead      NodeID
	isHead        bool
	isIrreducible bool
	live          bool
}

type edgeType int

const (
	DEAD edgeType = iota
	FORWARD
	BACKWARD
	CROSS
	REENTRY
)

type loopFinder struct {
	graph     *Graph
	node      []nodeInfo
	edge      []edgeType
	postorder []NodeID
	depth     []int
	current   int
}

func (lf *loopFinder) beginTraversingNode(n NodeID, prev NodeID) {
	lf.current += 1

	lf.node[n] = nodeInfo{
		pre:      lf.current,
		idom:     prev, // A reasonable inital guess.
		loopHead: NoNode,
		live:     true,
	}

	lf.depth[n] = lf.current
}

func (lf *loopFinder) endTraversingNode(n NodeID) {
	lf.current += 1
	lf.node[n].post = lf.current

	lf.depth[n] = 0

	lf.postorder = append(lf.postorder, n)
}

func (lf *loopFinder) isUnprocessed(n NodeID) bool {
	return !lf.node[n].live
}

func (lf *loopFinder) isBeingProcessed(dst NodeID) bool {
	return lf.depth[dst] > 0
}

func (lf *loopFinder) markLoopHeader(child NodeID, head NodeID) {
	if child == head {
		// Trivial loop.
		return
	}
	childHead := lf.node[child].loopHead
	for childHead != NoNode {
		if childHead == head {
			return
		}
		if lf.depth[childHead] < lf.depth[head] {
			// Found a closer head for this child, adopt it.
			lf.node[child].loopHead = head
			child, head = head, childHead
		} else {
			child = childHead
		}
		childHead = lf.node[child].loopHead
	}
	lf.node[child].loopHead = head
}

func (lf *loopFinder) process(n NodeID, prev NodeID) {
	lf.beginTraversingNode(n, prev)
	xit := lf.graph.ExitIterator(n)
	for xit.HasNext() {
		e, next := xit.GetNext()
		if lf.isUnprocessed(next) {
			lf.edge[e] = FORWARD
			lf.process(next, n)

			// Propagage loop headers upwards.
			head := lf.node[next].loopHead
			if head != NoNode {
				lf.markLoopHeader(n, head)
			}
		} else if lf.isBeingProcessed(next) {
			lf.edge[e] = BACKWARD
			lf.node[next].isHead = true
			lf.markLoopHeader(n, next)
		} else {
			lf.edge[e] = CROSS
			if lf.node[next].loopHead != NoNode {
				// Propagate loop header from cross edge.
				otherHead := lf.node[next].loopHead
				if lf.isBeingProcessed(otherHead) {
					lf.markLoopHeader(n, otherHead)
				} else {
					lf.edge[e] = REENTRY
					lf.node[otherHead].isIrreducible = true
					// Find and mark the common loop head.
					otherHead = lf.node[otherHead].loopHead
					for otherHead != NoNode {
						if lf.isBeingProcessed(otherHead) {
							lf.markLoopHeader(n, otherHead)
							break
						}
						otherHead = lf.node[otherHead].loopHead
					}
				}
			}
		}
	}
	lf.endTraversingNode(n)
}

func (lf *loopFinder) intersect(n0 NodeID, n1 NodeID) NodeID {
	i0 := lf.node[n0].post
	i1 := lf.node[n1].post
	for i0 != i1 {
		for i0 < i1 {
			n0 = lf.node[n0].idom
			i0 = lf.node[n0].post
		}
		for i0 > i1 {
			n1 = lf.node[n1].idom
			i1 = lf.node[n1].post
		}
	}
	return n0
}

func (lf *loopFinder) cleanDeadEdges() {
	numEdges := lf.graph.NumEdges()
	for i := 0; i < numEdges; i++ {
		if lf.edge[i] == DEAD {
			lf.graph.KillEdge(EdgeID(i))
		}
	}
}

func (lf *loopFinder) findIdoms() {
	g := lf.graph
	changed := true
	for changed {
		changed = false
		for i := len(lf.postorder) - 1; i >= 0; i-- {
			n := lf.postorder[i]
			original := lf.node[n].idom
			idom := original
			eit := g.EntryIterator(n)
			for eit.HasNext() {
				src, _ := eit.GetNext()
				idom = lf.intersect(idom, src)
			}
			if idom != original {
				lf.node[n].idom = idom
				changed = true
			}
		}
	}
}

func analyzeStructure(g *Graph) ([]nodeInfo, []edgeType, []NodeID) {
	numNodes := g.NumNodes()
	lf := &loopFinder{
		graph: g,
		node:  make([]nodeInfo, numNodes),
		edge:  make([]edgeType, g.NumEdges()),
		depth: make([]int, numNodes),
	}

	e := g.Entry()
	lf.process(e, e)
	lf.cleanDeadEdges()
	lf.findIdoms()
	return lf.node, lf.edge, lf.postorder
}

func uniqueEntry(g *Graph, src NodeID, dst NodeID) bool {
	eit := g.EntryIterator(dst)
	for eit.HasNext() {
		prev, _ := eit.GetNext()
		if prev != src {
			return false
		}
	}
	return true
}

func contract(g *Graph, src NodeID, dst NodeID, nodes []nodeInfo) {
	transferedExits := false
	eit := g.EntryIterator(dst)
	for eit.HasNext() {
		prev, e := eit.GetNext()
		if prev != src {
			panic(prev)
		}
		if !transferedExits {
			g.ReplaceEdgeWithExits(e, dst)
			transferedExits = true
		}
		g.KillEdge(e)
	}
	if !transferedExits {
		panic(dst)
	}
}

func contractLoop(g *Graph, n NodeID) {
	xit := g.ExitIterator(n)
	for xit.HasNext() {
		e, dst := xit.GetNext()
		if n == dst {
			g.KillEdge(e)
		}
	}
}

func appendCluster(src Cluster, dst Cluster, shouldFuse bool) Cluster {
	if shouldFuse {
		switch src := src.(type) {
		case *ClusterLeaf:
			switch dst := dst.(type) {
			case *ClusterLeaf:
				src.Nodes = append(src.Nodes, dst.Nodes...)
				return src
			case *ClusterLinear:
				other, ok := dst.Clusters[0].(*ClusterLeaf)
				if ok {
					other.Nodes = append(src.Nodes, other.Nodes...)
					return dst
				}
				dst.Clusters = append([]Cluster{src}, dst.Clusters...)
				return dst
			}
		case *ClusterLinear:
			switch dst := dst.(type) {
			case *ClusterLinear:
				src.Clusters = append(src.Clusters, dst.Clusters...)
				return src
			}
			src.Clusters = append(src.Clusters, dst)
			return src
		default:
			switch dst := dst.(type) {
			case *ClusterLinear:
				dst.Clusters = append([]Cluster{src}, dst.Clusters...)
				return dst
			}
		}
	} else {
		switch src := src.(type) {
		case *ClusterLinear:
			src.Clusters = append(src.Clusters, dst)
			return src
		}
	}
	return &ClusterLinear{Clusters: []Cluster{src, dst}}
}

func isClusterHead(g *Graph, src NodeID, nodes []nodeInfo, edges []edgeType) bool {
	xit := g.ExitIterator(src)
	for xit.HasNext() {
		e, dst := xit.GetNext()
		if edges[e] != BACKWARD && nodes[dst].idom != src {
			return false
		}
	}
	return true
}

func clusterRegion(g *Graph, n NodeID, currentHead NodeID, cluster Cluster, clusters []Cluster, nodes []nodeInfo, edges []edgeType) Cluster {
	for {
		ready := []NodeID{}
		readyClusters := []Cluster{}
		//readyEdges := []EdgeID{}
		//pendingEdges := []EdgeID{}
		numExits := 0
		xit := g.ExitIterator(n)
		for xit.HasNext() {
			e, dst := xit.GetNext()
			numExits += 1
			if nodes[dst].loopHead != currentHead {
				continue
			}
			if edges[e] == BACKWARD {
				continue
			}
			if uniqueEntry(g, n, dst) {
				if clusters[dst] != nil {
					ready = append(ready, dst)
					readyClusters = append(readyClusters, clusters[dst])
					clusters[dst] = nil
				}
			}
		}
		if len(ready) > 0 {
			if len(ready) > 1 {
				cluster = appendCluster(cluster, &ClusterSwitch{Children: readyClusters}, false)
			} else {
				cluster = appendCluster(cluster, readyClusters[0], numExits == 1)
			}
			for _, dst := range ready {
				contract(g, n, dst, nodes)
			}
		} else {
			return cluster
		}
	}
}

func makeCluster2(g *Graph, nodes []nodeInfo, edges []edgeType, postorder []NodeID) Cluster {
	clusters := make([]Cluster, g.NumNodes())

	for _, n := range postorder {
		var cluster Cluster = &ClusterLeaf{
			Nodes: []NodeID{n},
		}
		if nodes[n].isHead {
			cluster = clusterRegion(g, n, n, cluster, clusters, nodes, edges)
			contractLoop(g, n)
			cluster = &ClusterLoop{Body: cluster}
		}
		if isClusterHead(g, n, nodes, edges) {
			cluster = clusterRegion(g, n, nodes[n].loopHead, cluster, clusters, nodes, edges)
		}
		clusters[n] = cluster
	}
	return clusters[0]
}
<|endoftext|>


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

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/googlecloudrobotics/core/src/go/pkg/gcr"
	"github.com/googlecloudrobotics/core/src/go/pkg/kubeutils"
	"github.com/googlecloudrobotics/core/src/go/pkg/robotauth"
	"github.com/googlecloudrobotics/core/src/go/pkg/setup"

	"golang.org/x/oauth2"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	apierrors "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/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

var (
	domain              = flag.String("domain", "", "Domain for the Cloud Robotics project (default: www.endpoints.<project>.cloud.goog)")
	project             = flag.String("project", "", "Project ID for the Google Cloud Platform")
	projectNumber       = flag.Uint64("project-number", 0, "Project Number for the Google Cloud Platform")
	robotName           = flag.String("robot-name", "", "Robot name")
	robotRole           = flag.String("robot-role", "", "Robot role. Optional if the robot is already registered.")
	robotType           = flag.String("robot-type", "", "Robot type. Optional if the robot is already registered.")
	robotAuthentication = flag.Bool("robot-authentication", true, "Set up robot authentication.")
	appManagement       = flag.Bool("app-management", true, "Set up app management.")
)

const (
	filesDir = "/setup-robot-files"
	helmPath = filesDir + "/helm"
)

func parseFlags() {
	flag.Parse()

	if *project == "" {
		log.Fatal("--project is required.")
	}
	if *robotName == "" {
		log.Fatal("--robot-name is required.")
	}
	if *domain == "" {
		*domain = fmt.Sprintf("www.endpoints.%s.cloud.goog", *project)
	}
}

func main() {
	parseFlags()
	envToken := os.Getenv("ACCESS_TOKEN")
	if envToken == "" {
		log.Fatal("ACCESS_TOKEN environment variable is required.")
	}

	// Set up the OAuth2 token source.
	tokenSource := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: envToken})

	if *robotRole != "" || *robotType != "" {
		if err := createOrUpdateRobot(tokenSource); err != nil {
			log.Fatalf("Failed to update robot CR %v: %v", *robotName, err)
		}
	}

	// Connect to the surrounding k8s cluster.
	localConfig, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal("Failed to load in-cluster config: ", err)
	}
	k8sLocalClientSet, err := kubernetes.NewForConfig(localConfig)
	if err != nil {
		log.Fatal("Failed to create kubernetes client set: ", err)
	}

	if *robotAuthentication {
		// Set up robot authentication.
		auth := &robotauth.RobotAuth{
			RobotName:           *robotName,
			ProjectId:           *project,
			Domain:              *domain,
			PublicKeyRegistryId: fmt.Sprintf("robot-%s", *robotName),
		}
		client := oauth2.NewClient(context.Background(), tokenSource)
		if err := setup.CreateAndPublishCredentialsToCloud(client, auth); err != nil {
			log.Fatal(err)
		}
		if err := auth.StoreInK8sSecret(k8sLocalClientSet); err != nil {
			log.Fatal(fmt.Errorf("Failed to write auth secret: %v", err))
		}
		if err := gcr.UpdateGcrCredentials(k8sLocalClientSet, auth); err != nil {
			log.Fatal(err)
		}
	}

	// Create service account and role binding for Tiller.
	// (this isn't strictly necessary until we're using auth properly, but it's
	//  one less thing to fix when RBAC is used properly)
	if _, err := k8sLocalClientSet.CoreV1().ServiceAccounts("kube-system").Create(&corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tiller",
			Namespace: "kube-system",
		},
	}); err != nil {
		log.Println("Failed to create tiller service account: ", err)
	}
	if _, err := k8sLocalClientSet.RbacV1().ClusterRoleBindings().Create(&rbacv1.ClusterRoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tiller",
			Namespace: "kube-system",
		},
		RoleRef: rbacv1.RoleRef{
			APIGroup: "rbac.authorization.k8s.io",
			Kind:     "ClusterRole",
			Name:     "cluster-admin",
		},
		Subjects: []rbacv1.Subject{{
			Kind:      "ServiceAccount",
			Name:      "tiller",
			Namespace: "kube-system",
		}},
	}); err != nil {
		log.Println("Failed to create tiller role binding: ", err)
	}

	log.Println("Initializing Helm")
	output, err := exec.Command(
		helmPath,
		"init",
		"--history-max=10",
		"--upgrade",
		"--force-upgrade",
		"--wait",
		"--service-account=tiller").CombinedOutput()
	if err != nil {
		log.Fatalf("Helm init failed: %v. Helm output:\n%s\n", err, output)
	}

	// Clean up deprecated releases.
	deleteReleaseIfPresent("robot-cluster")

	installChartOrDie("robot-base", "base-robot-0.0.1.tgz")
}

func helmValuesStringFromMap(varMap map[string]string) string {
	varList := []string{}
	for k, v := range varMap {
		varList = append(varList, fmt.Sprintf("%s=%s", k, v))
	}
	return strings.Join(varList, ",")
}

func deleteReleaseIfPresent(name string) {
	output, err := exec.Command(
		helmPath,
		"delete",
		"--purge",
		name).CombinedOutput()
	if err != nil {
		log.Printf("Helm delete of %s failed: %v. Helm output:\n%s\n", name, err, output)
	} else {
		log.Printf("%s\nSuccessfully removed %s Helm release", output, name)
	}
}

func installChartOrDie(name string, chartPath string) {
	log.Printf("Installing %s Helm chart from %s", name, chartPath)
	vars := helmValuesStringFromMap(map[string]string{
		"domain":               *domain,
		"project":              *project,
		"project_number":       strconv.FormatUint(*projectNumber, 10),
		"app_management":       strconv.FormatBool(*appManagement),
		"robot_authentication": strconv.FormatBool(*robotAuthentication),
		"robot.name":           *robotName,
	})
	output, err := exec.Command(
		helmPath,
		"upgrade",
		"--install",
		"--force",
		"--set-string",
		vars,
		name,
		filepath.Join(filesDir, chartPath)).CombinedOutput()
	if err != nil {
		log.Fatalf("Helm install of %s failed: %v. Helm output:\n%s\n", name, err, output)
	} else {
		log.Printf("%s\nSuccessfully installed %s Helm chart", output, name)
	}
}

func createOrUpdateRobot(tokenSource oauth2.TokenSource) error {
	// Set up client for cloud k8s cluster (needed only to obtain list of robots).
	k8sCloudCfg := kubeutils.BuildCloudKubernetesConfig(tokenSource, *domain)
	k8sDynamicClient, err := dynamic.NewForConfig(k8sCloudCfg)
	if err != nil {
		return err
	}
	robotGVR := schema.GroupVersionResource{
		Group:    "registry.cloudrobotics.com",
		Version:  "v1alpha1",
		Resource: "robots"}
	robotClient := k8sDynamicClient.Resource(robotGVR).Namespace("default")
	robot, err := robotClient.Get(*robotName, metav1.GetOptions{})
	if err != nil {
		if s, ok := err.(*apierrors.StatusError); ok && s.ErrStatus.Reason == metav1.StatusReasonNotFound {
			robot := &unstructured.Unstructured{}
			robot.SetKind("Robot")
			robot.SetAPIVersion("registry.cloudrobotics.com/v1alpha1")
			robot.SetName(*robotName)
			robot.SetLabels(map[string]string{"cloudrobotics.com/robot-name": *robotName})
			robot.Object["spec"] = map[string]string{
				"role":    *robotRole,
				"type":    *robotType,
				"project": *project,
			}
			robot.Object["status"] = make(map[string]interface{})
			_, err := robotClient.Create(robot, metav1.CreateOptions{})
			return err
		} else {
			return fmt.Errorf("Failed to get robot %v: %v", *robotName, err)
		}
	}
	spec, ok := robot.Object["spec"].(map[string]interface{})
	if !ok {
		return fmt.Errorf("unmarshaling robot failed: spec is not a map")
	}
	spec["role"] = *robotRole
	spec["type"] = *robotType
	spec["project"] = *project
	_, err = robotClient.Update(robot, metav1.UpdateOptions{})
	return err
}
<commit_msg>Wait for DNS before starting the robot setup.<commit_after>// Copyright 2019 The Google Cloud Robotics Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/googlecloudrobotics/core/src/go/pkg/gcr"
	"github.com/googlecloudrobotics/core/src/go/pkg/kubeutils"
	"github.com/googlecloudrobotics/core/src/go/pkg/robotauth"
	"github.com/googlecloudrobotics/core/src/go/pkg/setup"

	"golang.org/x/oauth2"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	apierrors "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/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

var (
	domain              = flag.String("domain", "", "Domain for the Cloud Robotics project (default: www.endpoints.<project>.cloud.goog)")
	project             = flag.String("project", "", "Project ID for the Google Cloud Platform")
	projectNumber       = flag.Uint64("project-number", 0, "Project Number for the Google Cloud Platform")
	robotName           = flag.String("robot-name", "", "Robot name")
	robotRole           = flag.String("robot-role", "", "Robot role. Optional if the robot is already registered.")
	robotType           = flag.String("robot-type", "", "Robot type. Optional if the robot is already registered.")
	robotAuthentication = flag.Bool("robot-authentication", true, "Set up robot authentication.")
	appManagement       = flag.Bool("app-management", true, "Set up app management.")
)

const (
	filesDir      = "/setup-robot-files"
	helmPath      = filesDir + "/helm"
	numDNSRetries = 6
)

func parseFlags() {
	flag.Parse()

	if *project == "" {
		log.Fatal("--project is required.")
	}
	if *robotName == "" {
		log.Fatal("--robot-name is required.")
	}
	if *domain == "" {
		*domain = fmt.Sprintf("www.endpoints.%s.cloud.goog", *project)
	}
}

// Since this might be the first interaction with the cluster, manually resolve the
// domain name with retries to give a better error in the case of failure.
func waitForDNS() error {
	log.Printf("DNS lookup for %q", *domain)
	delay := time.Second
	for i := 0; i < numDNSRetries; i++ {
		ips, err := net.LookupIP(*domain)
		if err == nil {
			// Check that the results contain an ipv4 addr. Initially, coredns may only
			// return ipv6 addresses in which case helm will fail.
			for _, ip := range ips {
				if ip.To4() != nil {
					return nil
				}
			}
		}
		log.Printf("... Retry dns for %q", *domain)
		time.Sleep(delay)
		delay += delay
	}
	return fmt.Errorf("DNS lookup for %q failed", *domain)
}

func main() {
	parseFlags()
	envToken := os.Getenv("ACCESS_TOKEN")
	if envToken == "" {
		log.Fatal("ACCESS_TOKEN environment variable is required.")
	}

	// Set up the OAuth2 token source.
	tokenSource := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: envToken})

	if *robotRole != "" || *robotType != "" {
		if err := createOrUpdateRobot(tokenSource); err != nil {
			log.Fatalf("Failed to update robot CR %v: %v", *robotName, err)
		}
	}

	// Connect to the surrounding k8s cluster.
	localConfig, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal("Failed to load in-cluster config: ", err)
	}
	k8sLocalClientSet, err := kubernetes.NewForConfig(localConfig)
	if err != nil {
		log.Fatal("Failed to create kubernetes client set: ", err)
	}

	if err := waitForDNS(); err != nil {
		log.Fatalf("Failed to resolve cloud cluster: %s. Please retry in 5 minutes.", err)
	}

	if *robotAuthentication {
		// Set up robot authentication.
		auth := &robotauth.RobotAuth{
			RobotName:           *robotName,
			ProjectId:           *project,
			Domain:              *domain,
			PublicKeyRegistryId: fmt.Sprintf("robot-%s", *robotName),
		}
		client := oauth2.NewClient(context.Background(), tokenSource)
		if err := setup.CreateAndPublishCredentialsToCloud(client, auth); err != nil {
			log.Fatal(err)
		}
		if err := auth.StoreInK8sSecret(k8sLocalClientSet); err != nil {
			log.Fatal(fmt.Errorf("Failed to write auth secret: %v", err))
		}
		if err := gcr.UpdateGcrCredentials(k8sLocalClientSet, auth); err != nil {
			log.Fatal(err)
		}
	}

	// Create service account and role binding for Tiller.
	// (this isn't strictly necessary until we're using auth properly, but it's
	//  one less thing to fix when RBAC is used properly)
	if _, err := k8sLocalClientSet.CoreV1().ServiceAccounts("kube-system").Create(&corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tiller",
			Namespace: "kube-system",
		},
	}); err != nil {
		log.Println("Failed to create tiller service account: ", err)
	}
	if _, err := k8sLocalClientSet.RbacV1().ClusterRoleBindings().Create(&rbacv1.ClusterRoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tiller",
			Namespace: "kube-system",
		},
		RoleRef: rbacv1.RoleRef{
			APIGroup: "rbac.authorization.k8s.io",
			Kind:     "ClusterRole",
			Name:     "cluster-admin",
		},
		Subjects: []rbacv1.Subject{{
			Kind:      "ServiceAccount",
			Name:      "tiller",
			Namespace: "kube-system",
		}},
	}); err != nil {
		log.Println("Failed to create tiller role binding: ", err)
	}

	log.Println("Initializing Helm")
	output, err := exec.Command(
		helmPath,
		"init",
		"--history-max=10",
		"--upgrade",
		"--force-upgrade",
		"--wait",
		"--service-account=tiller").CombinedOutput()
	if err != nil {
		log.Fatalf("Helm init failed: %v. Helm output:\n%s\n", err, output)
	}

	// Clean up deprecated releases.
	deleteReleaseIfPresent("robot-cluster")

	installChartOrDie("robot-base", "base-robot-0.0.1.tgz")
}

func helmValuesStringFromMap(varMap map[string]string) string {
	varList := []string{}
	for k, v := range varMap {
		varList = append(varList, fmt.Sprintf("%s=%s", k, v))
	}
	return strings.Join(varList, ",")
}

func deleteReleaseIfPresent(name string) {
	output, err := exec.Command(
		helmPath,
		"delete",
		"--purge",
		name).CombinedOutput()
	if err != nil {
		log.Printf("Helm delete of %s failed: %v. Helm output:\n%s\n", name, err, output)
	} else {
		log.Printf("%s\nSuccessfully removed %s Helm release", output, name)
	}
}

func installChartOrDie(name string, chartPath string) {
	log.Printf("Installing %s Helm chart from %s", name, chartPath)
	vars := helmValuesStringFromMap(map[string]string{
		"domain":               *domain,
		"project":              *project,
		"project_number":       strconv.FormatUint(*projectNumber, 10),
		"app_management":       strconv.FormatBool(*appManagement),
		"robot_authentication": strconv.FormatBool(*robotAuthentication),
		"robot.name":           *robotName,
	})
	output, err := exec.Command(
		helmPath,
		"upgrade",
		"--install",
		"--force",
		"--set-string",
		vars,
		name,
		filepath.Join(filesDir, chartPath)).CombinedOutput()
	if err != nil {
		log.Fatalf("Helm install of %s failed: %v. Helm output:\n%s\n", name, err, output)
	} else {
		log.Printf("%s\nSuccessfully installed %s Helm chart", output, name)
	}
}

func createOrUpdateRobot(tokenSource oauth2.TokenSource) error {
	// Set up client for cloud k8s cluster (needed only to obtain list of robots).
	k8sCloudCfg := kubeutils.BuildCloudKubernetesConfig(tokenSource, *domain)
	k8sDynamicClient, err := dynamic.NewForConfig(k8sCloudCfg)
	if err != nil {
		return err
	}
	robotGVR := schema.GroupVersionResource{
		Group:    "registry.cloudrobotics.com",
		Version:  "v1alpha1",
		Resource: "robots"}
	robotClient := k8sDynamicClient.Resource(robotGVR).Namespace("default")
	robot, err := robotClient.Get(*robotName, metav1.GetOptions{})
	if err != nil {
		if s, ok := err.(*apierrors.StatusError); ok && s.ErrStatus.Reason == metav1.StatusReasonNotFound {
			robot := &unstructured.Unstructured{}
			robot.SetKind("Robot")
			robot.SetAPIVersion("registry.cloudrobotics.com/v1alpha1")
			robot.SetName(*robotName)
			robot.SetLabels(map[string]string{"cloudrobotics.com/robot-name": *robotName})
			robot.Object["spec"] = map[string]string{
				"role":    *robotRole,
				"type":    *robotType,
				"project": *project,
			}
			robot.Object["status"] = make(map[string]interface{})
			_, err := robotClient.Create(robot, metav1.CreateOptions{})
			return err
		} else {
			return fmt.Errorf("Failed to get robot %v: %v", *robotName, err)
		}
	}
	spec, ok := robot.Object["spec"].(map[string]interface{})
	if !ok {
		return fmt.Errorf("unmarshaling robot failed: spec is not a map")
	}
	spec["role"] = *robotRole
	spec["type"] = *robotType
	spec["project"] = *project
	_, err = robotClient.Update(robot, metav1.UpdateOptions{})
	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package auth

import (
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/azure"
	"golang.org/x/crypto/pkcs12"

	"k8s.io/klog"
)

const (
	// ADFSIdentitySystem is the override value for tenantID on Azure Stack clouds.
	ADFSIdentitySystem = "adfs"
)

var (
	// ErrorNoAuth indicates that no credentials are provided.
	ErrorNoAuth = fmt.Errorf("no credentials provided for Azure cloud provider")
)

// AzureAuthConfig holds auth related part of cloud config
type AzureAuthConfig struct {
	// The cloud environment identifier. Takes values from https://github.com/Azure/go-autorest/blob/ec5f4903f77ed9927ac95b19ab8e44ada64c1356/autorest/azure/environments.go#L13
	Cloud string `json:"cloud,omitempty" yaml:"cloud,omitempty"`
	// The AAD Tenant ID for the Subscription that the cluster is deployed in
	TenantID string `json:"tenantId,omitempty" yaml:"tenantId,omitempty"`
	// The ClientID for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientID string `json:"aadClientId,omitempty" yaml:"aadClientId,omitempty"`
	// The ClientSecret for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientSecret string `json:"aadClientSecret,omitempty" yaml:"aadClientSecret,omitempty"`
	// The path of a client certificate for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientCertPath string `json:"aadClientCertPath,omitempty" yaml:"aadClientCertPath,omitempty"`
	// The password of the client certificate for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientCertPassword string `json:"aadClientCertPassword,omitempty" yaml:"aadClientCertPassword,omitempty"`
	// Use managed service identity for the virtual machine to access Azure ARM APIs
	UseManagedIdentityExtension bool `json:"useManagedIdentityExtension,omitempty" yaml:"useManagedIdentityExtension,omitempty"`
	// UserAssignedIdentityID contains the Client ID of the user assigned MSI which is assigned to the underlying VMs. If empty the user assigned identity is not used.
	// More details of the user assigned identity can be found at: https://docs.microsoft.com/en-us/azure/active-directory/managed-service-identity/overview
	// For the user assigned identity specified here to be used, the UseManagedIdentityExtension has to be set to true.
	UserAssignedIdentityID string `json:"userAssignedIdentityID,omitempty" yaml:"userAssignedIdentityID,omitempty"`
	// The ID of the Azure Subscription that the cluster is deployed in
	SubscriptionID string `json:"subscriptionId,omitempty" yaml:"subscriptionId,omitempty"`
	// IdentitySystem indicates the identity provider. Relevant only to hybrid clouds (Azure Stack).
	// Allowed values are 'azure_ad' (default), 'adfs'.
	IdentitySystem string `json:"identitySystem" yaml:"identitySystem"`
	// CloudFQDN represents the hybrid cloud's fully qualified domain name: {location}.{domain}
	// If set, cloud provider will generate its autorest.Environment instead of using one of the pre-defined ones.
	CloudFQDN string `json:"cloudFQDN" yaml:"cloudFQDN"`
}

// GetServicePrincipalToken creates a new service principal token based on the configuration
func GetServicePrincipalToken(config *AzureAuthConfig, env *azure.Environment) (*adal.ServicePrincipalToken, error) {
	var tenantID string
	if strings.EqualFold(config.IdentitySystem, ADFSIdentitySystem) {
		tenantID = ADFSIdentitySystem
	} else {
		tenantID = config.TenantID
	}

	if config.UseManagedIdentityExtension {
		klog.V(2).Infoln("azure: using managed identity extension to retrieve access token")
		msiEndpoint, err := adal.GetMSIVMEndpoint()
		if err != nil {
			return nil, fmt.Errorf("Getting the managed service identity endpoint: %v", err)
		}
		if len(config.UserAssignedIdentityID) > 0 {
			klog.V(4).Info("azure: using User Assigned MSI ID to retrieve access token")
			return adal.NewServicePrincipalTokenFromMSIWithUserAssignedID(msiEndpoint,
				env.ServiceManagementEndpoint,
				config.UserAssignedIdentityID)
		}
		klog.V(4).Info("azure: using System Assigned MSI to retrieve access token")
		return adal.NewServicePrincipalTokenFromMSI(
			msiEndpoint,
			env.ServiceManagementEndpoint)
	}

	oauthConfig, err := adal.NewOAuthConfig(env.ActiveDirectoryEndpoint, tenantID)
	if err != nil {
		return nil, fmt.Errorf("creating the OAuth config: %v", err)
	}

	if len(config.AADClientSecret) > 0 {
		klog.V(2).Infoln("azure: using client_id+client_secret to retrieve access token")
		return adal.NewServicePrincipalToken(
			*oauthConfig,
			config.AADClientID,
			config.AADClientSecret,
			env.ServiceManagementEndpoint)
	}

	if len(config.AADClientCertPath) > 0 && len(config.AADClientCertPassword) > 0 {
		klog.V(2).Infoln("azure: using jwt client_assertion (client_cert+client_private_key) to retrieve access token")
		certData, err := ioutil.ReadFile(config.AADClientCertPath)
		if err != nil {
			return nil, fmt.Errorf("reading the client certificate from file %s: %v", config.AADClientCertPath, err)
		}
		certificate, privateKey, err := decodePkcs12(certData, config.AADClientCertPassword)
		if err != nil {
			return nil, fmt.Errorf("decoding the client certificate: %v", err)
		}
		return adal.NewServicePrincipalTokenFromCertificate(
			*oauthConfig,
			config.AADClientID,
			certificate,
			privateKey,
			env.ServiceManagementEndpoint)
	}

	return nil, ErrorNoAuth
}

// ParseAzureEnvironment returns azure environment by name
func ParseAzureEnvironment(cloudName, cloudFQDN, identitySystem string) (*azure.Environment, error) {
	var env azure.Environment
	var err error
	if cloudFQDN != "" {
		resourceManagerEndpoint := fmt.Sprintf("https://management.%s/", cloudFQDN)
		nameOverride := azure.OverrideProperty{Key: azure.EnvironmentName, Value: cloudName}
		klog.V(4).Infof("Loading environment from resource manager endpoint: %s", resourceManagerEndpoint)
		env, err = azure.EnvironmentFromURL(resourceManagerEndpoint, nameOverride)
		if err == nil && strings.EqualFold(cloudName, "AzureStackCloud") {
			azureStackOverrides(env, cloudFQDN, identitySystem)
		}
	} else if cloudName == "" {
		env = azure.PublicCloud
	} else {
		env, err = azure.EnvironmentFromName(cloudName)
	}
	return &env, err
}

// decodePkcs12 decodes a PKCS#12 client certificate by extracting the public certificate and
// the private RSA key
func decodePkcs12(pkcs []byte, password string) (*x509.Certificate, *rsa.PrivateKey, error) {
	privateKey, certificate, err := pkcs12.Decode(pkcs, password)
	if err != nil {
		return nil, nil, fmt.Errorf("decoding the PKCS#12 client certificate: %v", err)
	}
	rsaPrivateKey, isRsaKey := privateKey.(*rsa.PrivateKey)
	if !isRsaKey {
		return nil, nil, fmt.Errorf("PKCS#12 certificate must contain a RSA private key")
	}

	return certificate, rsaPrivateKey, nil
}

func azureStackOverrides(env azure.Environment, cloudFQDN, identitySystem string) azure.Environment {
	// if AzureStack, make sure the generated environment matches what AKSe currently generates
	env.ManagementPortalURL = fmt.Sprintf("https://portal.%s/", cloudFQDN)
	// TODO: figure out why AKSe does this
	// why is autorest not setting ServiceManagementEndpoint?
	env.ServiceManagementEndpoint = env.TokenAudience
	// TODO: figure out why AKSe does this
	// ResourceManagerVMDNSSuffix is not referenced in k/k
	split := strings.Split(cloudFQDN, ".")
	domain := strings.Join(split[1:], ".")
	env.ResourceManagerVMDNSSuffix = fmt.Sprintf("cloudapp.%s", domain)
	// NOTE: autorest sets KeyVaultEndpoint while AKSe does not
	if strings.EqualFold(identitySystem, ADFSIdentitySystem) {
		env.ActiveDirectoryEndpoint = strings.TrimSuffix(env.ActiveDirectoryEndpoint, "/")
		env.ActiveDirectoryEndpoint = strings.TrimSuffix(env.ActiveDirectoryEndpoint, "adfs")
	}
	return env
}
<commit_msg>misc fixes<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package auth

import (
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/azure"
	"golang.org/x/crypto/pkcs12"

	"k8s.io/klog"
)

const (
	// ADFSIdentitySystem is the override value for tenantID on Azure Stack clouds.
	ADFSIdentitySystem = "adfs"
)

var (
	// ErrorNoAuth indicates that no credentials are provided.
	ErrorNoAuth = fmt.Errorf("no credentials provided for Azure cloud provider")
)

// AzureAuthConfig holds auth related part of cloud config
type AzureAuthConfig struct {
	// The cloud environment identifier. Takes values from https://github.com/Azure/go-autorest/blob/ec5f4903f77ed9927ac95b19ab8e44ada64c1356/autorest/azure/environments.go#L13
	Cloud string `json:"cloud,omitempty" yaml:"cloud,omitempty"`
	// The AAD Tenant ID for the Subscription that the cluster is deployed in
	TenantID string `json:"tenantId,omitempty" yaml:"tenantId,omitempty"`
	// The ClientID for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientID string `json:"aadClientId,omitempty" yaml:"aadClientId,omitempty"`
	// The ClientSecret for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientSecret string `json:"aadClientSecret,omitempty" yaml:"aadClientSecret,omitempty"`
	// The path of a client certificate for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientCertPath string `json:"aadClientCertPath,omitempty" yaml:"aadClientCertPath,omitempty"`
	// The password of the client certificate for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientCertPassword string `json:"aadClientCertPassword,omitempty" yaml:"aadClientCertPassword,omitempty"`
	// Use managed service identity for the virtual machine to access Azure ARM APIs
	UseManagedIdentityExtension bool `json:"useManagedIdentityExtension,omitempty" yaml:"useManagedIdentityExtension,omitempty"`
	// UserAssignedIdentityID contains the Client ID of the user assigned MSI which is assigned to the underlying VMs. If empty the user assigned identity is not used.
	// More details of the user assigned identity can be found at: https://docs.microsoft.com/en-us/azure/active-directory/managed-service-identity/overview
	// For the user assigned identity specified here to be used, the UseManagedIdentityExtension has to be set to true.
	UserAssignedIdentityID string `json:"userAssignedIdentityID,omitempty" yaml:"userAssignedIdentityID,omitempty"`
	// The ID of the Azure Subscription that the cluster is deployed in
	SubscriptionID string `json:"subscriptionId,omitempty" yaml:"subscriptionId,omitempty"`
	// IdentitySystem indicates the identity provider. Relevant only to hybrid clouds (Azure Stack).
	// Allowed values are 'azure_ad' (default), 'adfs'.
	IdentitySystem string `json:"identitySystem,omitempty" yaml:"identitySystem,omitempty"`
	// CloudFQDN represents the hybrid cloud's fully qualified domain name: {location}.{domain}
	// If set, cloud provider will generate its autorest.Environment instead of using one of the pre-defined ones.
	CloudFQDN string `json:"cloudFQDN,omitempty" yaml:"cloudFQDN,omitempty"`
}

// GetServicePrincipalToken creates a new service principal token based on the configuration
func GetServicePrincipalToken(config *AzureAuthConfig, env *azure.Environment) (*adal.ServicePrincipalToken, error) {
	var tenantID string
	if strings.EqualFold(config.IdentitySystem, ADFSIdentitySystem) {
		tenantID = ADFSIdentitySystem
	} else {
		tenantID = config.TenantID
	}

	if config.UseManagedIdentityExtension {
		klog.V(2).Infoln("azure: using managed identity extension to retrieve access token")
		msiEndpoint, err := adal.GetMSIVMEndpoint()
		if err != nil {
			return nil, fmt.Errorf("Getting the managed service identity endpoint: %v", err)
		}
		if len(config.UserAssignedIdentityID) > 0 {
			klog.V(4).Info("azure: using User Assigned MSI ID to retrieve access token")
			return adal.NewServicePrincipalTokenFromMSIWithUserAssignedID(msiEndpoint,
				env.ServiceManagementEndpoint,
				config.UserAssignedIdentityID)
		}
		klog.V(4).Info("azure: using System Assigned MSI to retrieve access token")
		return adal.NewServicePrincipalTokenFromMSI(
			msiEndpoint,
			env.ServiceManagementEndpoint)
	}

	oauthConfig, err := adal.NewOAuthConfig(env.ActiveDirectoryEndpoint, tenantID)
	if err != nil {
		return nil, fmt.Errorf("creating the OAuth config: %v", err)
	}

	if len(config.AADClientSecret) > 0 {
		klog.V(2).Infoln("azure: using client_id+client_secret to retrieve access token")
		return adal.NewServicePrincipalToken(
			*oauthConfig,
			config.AADClientID,
			config.AADClientSecret,
			env.ServiceManagementEndpoint)
	}

	if len(config.AADClientCertPath) > 0 && len(config.AADClientCertPassword) > 0 {
		klog.V(2).Infoln("azure: using jwt client_assertion (client_cert+client_private_key) to retrieve access token")
		certData, err := ioutil.ReadFile(config.AADClientCertPath)
		if err != nil {
			return nil, fmt.Errorf("reading the client certificate from file %s: %v", config.AADClientCertPath, err)
		}
		certificate, privateKey, err := decodePkcs12(certData, config.AADClientCertPassword)
		if err != nil {
			return nil, fmt.Errorf("decoding the client certificate: %v", err)
		}
		return adal.NewServicePrincipalTokenFromCertificate(
			*oauthConfig,
			config.AADClientID,
			certificate,
			privateKey,
			env.ServiceManagementEndpoint)
	}

	return nil, ErrorNoAuth
}

// ParseAzureEnvironment returns azure environment by name
func ParseAzureEnvironment(cloudName, cloudFQDN, identitySystem string) (*azure.Environment, error) {
	var env azure.Environment
	var err error
	if cloudFQDN != "" {
		resourceManagerEndpoint := fmt.Sprintf("https://management.%s/", cloudFQDN)
		nameOverride := azure.OverrideProperty{Key: azure.EnvironmentName, Value: cloudName}
		klog.V(4).Infof("Loading environment from resource manager endpoint: %s", resourceManagerEndpoint)
		env, err = azure.EnvironmentFromURL(resourceManagerEndpoint, nameOverride)
		if err == nil && strings.EqualFold(cloudName, "AzureStackCloud") {
			azureStackOverrides(&env, cloudFQDN, identitySystem)
		}
	} else if cloudName == "" {
		klog.V(4).Info("Using public cloud environment")
		env = azure.PublicCloud
	} else {
		klog.V(4).Infof("Using %s environment", cloudName)
		env, err = azure.EnvironmentFromName(cloudName)
	}
	return &env, err
}

// decodePkcs12 decodes a PKCS#12 client certificate by extracting the public certificate and
// the private RSA key
func decodePkcs12(pkcs []byte, password string) (*x509.Certificate, *rsa.PrivateKey, error) {
	privateKey, certificate, err := pkcs12.Decode(pkcs, password)
	if err != nil {
		return nil, nil, fmt.Errorf("decoding the PKCS#12 client certificate: %v", err)
	}
	rsaPrivateKey, isRsaKey := privateKey.(*rsa.PrivateKey)
	if !isRsaKey {
		return nil, nil, fmt.Errorf("PKCS#12 certificate must contain a RSA private key")
	}

	return certificate, rsaPrivateKey, nil
}

func azureStackOverrides(env *azure.Environment, cloudFQDN, identitySystem string) {
	// if AzureStack, make sure the generated environment matches what AKSe currently generates
	env.ManagementPortalURL = fmt.Sprintf("https://portal.%s/", cloudFQDN)
	// TODO: figure out why AKSe does this
	// why is autorest not setting ServiceManagementEndpoint?
	env.ServiceManagementEndpoint = env.TokenAudience
	// TODO: figure out why AKSe does this
	// ResourceManagerVMDNSSuffix is not referenced in k/k
	split := strings.Split(cloudFQDN, ".")
	domain := strings.Join(split[1:], ".")
	env.ResourceManagerVMDNSSuffix = fmt.Sprintf("cloudapp.%s", domain)
	// NOTE: autorest sets KeyVaultEndpoint while AKSe does not
	if strings.EqualFold(identitySystem, ADFSIdentitySystem) {
		env.ActiveDirectoryEndpoint = strings.TrimSuffix(env.ActiveDirectoryEndpoint, "/")
		env.ActiveDirectoryEndpoint = strings.TrimSuffix(env.ActiveDirectoryEndpoint, "adfs")
	}
}
<|endoftext|>


<commit_before>package command

import (
	"os"
	"path/filepath"
	"testing"

	"github.com/mitchellh/cli"
)

func TestInit(t *testing.T) {
	dir := tempDir(t)

	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{
		testFixturePath("init"),
		dir,
	}
	if code := c.Run(args); code != 0 {
		t.Fatalf("bad: \n%s", ui.ErrorWriter.String())
	}

	if _, err := os.Stat(filepath.Join(dir, "hello.tf")); err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestInit_multipleArgs(t *testing.T) {
	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{
		"bad",
		"bad",
	}
	if code := c.Run(args); code != 1 {
		t.Fatalf("bad: \n%s", ui.OutputWriter.String())
	}
}

func TestInit_noArgs(t *testing.T) {
	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{}
	if code := c.Run(args); code != 1 {
		t.Fatalf("bad: \n%s", ui.OutputWriter.String())
	}
}
<commit_msg>command: test init<commit_after>package command

import (
	"os"
	"path/filepath"
	"testing"

	"github.com/mitchellh/cli"
)

func TestInit(t *testing.T) {
	dir := tempDir(t)

	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{
		testFixturePath("init"),
		dir,
	}
	if code := c.Run(args); code != 0 {
		t.Fatalf("bad: \n%s", ui.ErrorWriter.String())
	}

	if _, err := os.Stat(filepath.Join(dir, "hello.tf")); err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestInit_cwd(t *testing.T) {
	dir := tempDir(t)
	if err := os.MkdirAll(dir, 0755); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Change to the temporary directory
	cwd, err := os.Getwd()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if err := os.Chdir(dir); err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.Chdir(cwd)

	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{
		testFixturePath("init"),
	}
	if code := c.Run(args); code != 0 {
		t.Fatalf("bad: \n%s", ui.ErrorWriter.String())
	}

	if _, err := os.Stat("hello.tf"); err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestInit_multipleArgs(t *testing.T) {
	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{
		"bad",
		"bad",
	}
	if code := c.Run(args); code != 1 {
		t.Fatalf("bad: \n%s", ui.OutputWriter.String())
	}
}

func TestInit_noArgs(t *testing.T) {
	ui := new(cli.MockUi)
	c := &InitCommand{
		Meta: Meta{
			ContextOpts: testCtxConfig(testProvider()),
			Ui:          ui,
		},
	}

	args := []string{}
	if code := c.Run(args); code != 1 {
		t.Fatalf("bad: \n%s", ui.OutputWriter.String())
	}
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"testing"

	"github.com/google/go-cmp/cmp"

	"github.com/hashicorp/terraform/backend"
	"github.com/hashicorp/terraform/terraform"
)

func TestMetaColorize(t *testing.T) {
	var m *Meta
	var args, args2 []string

	// Test basic, color
	m = new(Meta)
	m.Color = true
	args = []string{"foo", "bar"}
	args2 = []string{"foo", "bar"}
	args = m.process(args)
	if !reflect.DeepEqual(args, args2) {
		t.Fatalf("bad: %#v", args)
	}
	if m.Colorize().Disable {
		t.Fatal("should not be disabled")
	}

	// Test basic, no change
	m = new(Meta)
	args = []string{"foo", "bar"}
	args2 = []string{"foo", "bar"}
	args = m.process(args)
	if !reflect.DeepEqual(args, args2) {
		t.Fatalf("bad: %#v", args)
	}
	if !m.Colorize().Disable {
		t.Fatal("should be disabled")
	}

	// Test disable #1
	m = new(Meta)
	m.Color = true
	args = []string{"foo", "-no-color", "bar"}
	args2 = []string{"foo", "bar"}
	args = m.process(args)
	if !reflect.DeepEqual(args, args2) {
		t.Fatalf("bad: %#v", args)
	}
	if !m.Colorize().Disable {
		t.Fatal("should be disabled")
	}
}

func TestMetaInputMode(t *testing.T) {
	test = false
	defer func() { test = true }()

	m := new(Meta)
	args := []string{}

	fs := m.extendedFlagSet("foo")
	if err := fs.Parse(args); err != nil {
		t.Fatalf("err: %s", err)
	}

	if m.InputMode() != terraform.InputModeStd {
		t.Fatalf("bad: %#v", m.InputMode())
	}
}

func TestMetaInputMode_envVar(t *testing.T) {
	test = false
	defer func() { test = true }()
	old := os.Getenv(InputModeEnvVar)
	defer os.Setenv(InputModeEnvVar, old)

	m := new(Meta)
	args := []string{}

	fs := m.extendedFlagSet("foo")
	if err := fs.Parse(args); err != nil {
		t.Fatalf("err: %s", err)
	}

	off := terraform.InputMode(0)
	on := terraform.InputModeStd
	cases := []struct {
		EnvVar   string
		Expected terraform.InputMode
	}{
		{"false", off},
		{"0", off},
		{"true", on},
		{"1", on},
	}

	for _, tc := range cases {
		os.Setenv(InputModeEnvVar, tc.EnvVar)
		if m.InputMode() != tc.Expected {
			t.Fatalf("expected InputMode: %#v, got: %#v", tc.Expected, m.InputMode())
		}
	}
}

func TestMetaInputMode_disable(t *testing.T) {
	test = false
	defer func() { test = true }()

	m := new(Meta)
	args := []string{"-input=false"}

	fs := m.extendedFlagSet("foo")
	if err := fs.Parse(args); err != nil {
		t.Fatalf("err: %s", err)
	}

	if m.InputMode() > 0 {
		t.Fatalf("bad: %#v", m.InputMode())
	}
}

func TestMeta_initStatePaths(t *testing.T) {
	m := new(Meta)
	m.initStatePaths()

	if m.statePath != DefaultStateFilename {
		t.Fatalf("bad: %#v", m)
	}
	if m.stateOutPath != DefaultStateFilename {
		t.Fatalf("bad: %#v", m)
	}
	if m.backupPath != DefaultStateFilename+DefaultBackupExtension {
		t.Fatalf("bad: %#v", m)
	}

	m = new(Meta)
	m.statePath = "foo"
	m.initStatePaths()

	if m.stateOutPath != "foo" {
		t.Fatalf("bad: %#v", m)
	}
	if m.backupPath != "foo"+DefaultBackupExtension {
		t.Fatalf("bad: %#v", m)
	}

	m = new(Meta)
	m.stateOutPath = "foo"
	m.initStatePaths()

	if m.statePath != DefaultStateFilename {
		t.Fatalf("bad: %#v", m)
	}
	if m.backupPath != "foo"+DefaultBackupExtension {
		t.Fatalf("bad: %#v", m)
	}
}

func TestMeta_Env(t *testing.T) {
	td := tempDir(t)
	os.MkdirAll(td, 0755)
	defer os.RemoveAll(td)
	defer testChdir(t, td)()

	m := new(Meta)

	env, err := m.Workspace()
	if err != nil {
		t.Fatal(err)
	}

	if env != backend.DefaultStateName {
		t.Fatalf("expected env %q, got env %q", backend.DefaultStateName, env)
	}

	testEnv := "test_env"
	if err := m.SetWorkspace(testEnv); err != nil {
		t.Fatal("error setting env:", err)
	}

	env, _ = m.Workspace()
	if env != testEnv {
		t.Fatalf("expected env %q, got env %q", testEnv, env)
	}

	if err := m.SetWorkspace(backend.DefaultStateName); err != nil {
		t.Fatal("error setting env:", err)
	}

	env, _ = m.Workspace()
	if env != backend.DefaultStateName {
		t.Fatalf("expected env %q, got env %q", backend.DefaultStateName, env)
	}
}

func TestMeta_Workspace_override(t *testing.T) {
	defer func(value string) {
		os.Setenv(WorkspaceNameEnvVar, value)
	}(os.Getenv(WorkspaceNameEnvVar))

	m := new(Meta)

	testCases := map[string]struct {
		workspace string
		err       error
	}{
		"": {
			"default",
			nil,
		},
		"development": {
			"development",
			nil,
		},
		"invalid name": {
			"",
			invalidWorkspaceNameEnvVar,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			os.Setenv(WorkspaceNameEnvVar, name)
			workspace, err := m.Workspace()
			if workspace != tc.workspace {
				t.Errorf("Unexpected workspace\n got: %s\nwant: %s\n", workspace, tc.workspace)
			}
			if err != tc.err {
				t.Errorf("Unexpected error\n got: %s\nwant: %s\n", err, tc.err)
			}
		})
	}
}

func TestMeta_process(t *testing.T) {
	test = false
	defer func() { test = true }()

	// Create a temporary directory for our cwd
	d := tempDir(t)
	os.MkdirAll(d, 0755)
	defer os.RemoveAll(d)
	defer testChdir(t, d)()

	// At one point it was the responsibility of this process function to
	// insert fake additional -var-file options into the command line
	// if the automatic tfvars files were present. This is no longer the
	// responsibility of process (it happens in collectVariableValues instead)
	// but we're still testing with these files in place to verify that
	// they _aren't_ being interpreted by process, since that could otherwise
	// cause them to be added more than once and mess up the precedence order.
	defaultVarsfile := "terraform.tfvars"
	err := ioutil.WriteFile(
		filepath.Join(d, defaultVarsfile),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	fileFirstAlphabetical := "a-file.auto.tfvars"
	err = ioutil.WriteFile(
		filepath.Join(d, fileFirstAlphabetical),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	fileLastAlphabetical := "z-file.auto.tfvars"
	err = ioutil.WriteFile(
		filepath.Join(d, fileLastAlphabetical),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	// Regular tfvars files will not be autoloaded
	fileIgnored := "ignored.tfvars"
	err = ioutil.WriteFile(
		filepath.Join(d, fileIgnored),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	tests := []struct {
		GivenArgs    []string
		FilteredArgs []string
		ExtraCheck   func(*testing.T, *Meta)
	}{
		{
			[]string{},
			[]string{},
			func(t *testing.T, m *Meta) {
				if got, want := m.color, true; got != want {
					t.Errorf("wrong m.color value %#v; want %#v", got, want)
				}
				if got, want := m.Color, true; got != want {
					t.Errorf("wrong m.Color value %#v; want %#v", got, want)
				}
			},
		},
		{
			[]string{"-no-color"},
			[]string{},
			func(t *testing.T, m *Meta) {
				if got, want := m.color, false; got != want {
					t.Errorf("wrong m.color value %#v; want %#v", got, want)
				}
				if got, want := m.Color, false; got != want {
					t.Errorf("wrong m.Color value %#v; want %#v", got, want)
				}
			},
		},
	}

	for _, test := range tests {
		t.Run(fmt.Sprintf("%s", test.GivenArgs), func(t *testing.T) {
			m := new(Meta)
			m.Color = true // this is the default also for normal use, overridden by -no-color
			args := test.GivenArgs
			args = m.process(args)

			if !cmp.Equal(test.FilteredArgs, args) {
				t.Errorf("wrong filtered arguments\n%s", cmp.Diff(test.FilteredArgs, args))
			}

			if test.ExtraCheck != nil {
				test.ExtraCheck(t, m)
			}
		})
	}
}
<commit_msg>command: Add test for invalid selected workspace<commit_after>package command

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"testing"

	"github.com/google/go-cmp/cmp"

	"github.com/hashicorp/terraform/backend"
	"github.com/hashicorp/terraform/backend/local"
	"github.com/hashicorp/terraform/terraform"
)

func TestMetaColorize(t *testing.T) {
	var m *Meta
	var args, args2 []string

	// Test basic, color
	m = new(Meta)
	m.Color = true
	args = []string{"foo", "bar"}
	args2 = []string{"foo", "bar"}
	args = m.process(args)
	if !reflect.DeepEqual(args, args2) {
		t.Fatalf("bad: %#v", args)
	}
	if m.Colorize().Disable {
		t.Fatal("should not be disabled")
	}

	// Test basic, no change
	m = new(Meta)
	args = []string{"foo", "bar"}
	args2 = []string{"foo", "bar"}
	args = m.process(args)
	if !reflect.DeepEqual(args, args2) {
		t.Fatalf("bad: %#v", args)
	}
	if !m.Colorize().Disable {
		t.Fatal("should be disabled")
	}

	// Test disable #1
	m = new(Meta)
	m.Color = true
	args = []string{"foo", "-no-color", "bar"}
	args2 = []string{"foo", "bar"}
	args = m.process(args)
	if !reflect.DeepEqual(args, args2) {
		t.Fatalf("bad: %#v", args)
	}
	if !m.Colorize().Disable {
		t.Fatal("should be disabled")
	}
}

func TestMetaInputMode(t *testing.T) {
	test = false
	defer func() { test = true }()

	m := new(Meta)
	args := []string{}

	fs := m.extendedFlagSet("foo")
	if err := fs.Parse(args); err != nil {
		t.Fatalf("err: %s", err)
	}

	if m.InputMode() != terraform.InputModeStd {
		t.Fatalf("bad: %#v", m.InputMode())
	}
}

func TestMetaInputMode_envVar(t *testing.T) {
	test = false
	defer func() { test = true }()
	old := os.Getenv(InputModeEnvVar)
	defer os.Setenv(InputModeEnvVar, old)

	m := new(Meta)
	args := []string{}

	fs := m.extendedFlagSet("foo")
	if err := fs.Parse(args); err != nil {
		t.Fatalf("err: %s", err)
	}

	off := terraform.InputMode(0)
	on := terraform.InputModeStd
	cases := []struct {
		EnvVar   string
		Expected terraform.InputMode
	}{
		{"false", off},
		{"0", off},
		{"true", on},
		{"1", on},
	}

	for _, tc := range cases {
		os.Setenv(InputModeEnvVar, tc.EnvVar)
		if m.InputMode() != tc.Expected {
			t.Fatalf("expected InputMode: %#v, got: %#v", tc.Expected, m.InputMode())
		}
	}
}

func TestMetaInputMode_disable(t *testing.T) {
	test = false
	defer func() { test = true }()

	m := new(Meta)
	args := []string{"-input=false"}

	fs := m.extendedFlagSet("foo")
	if err := fs.Parse(args); err != nil {
		t.Fatalf("err: %s", err)
	}

	if m.InputMode() > 0 {
		t.Fatalf("bad: %#v", m.InputMode())
	}
}

func TestMeta_initStatePaths(t *testing.T) {
	m := new(Meta)
	m.initStatePaths()

	if m.statePath != DefaultStateFilename {
		t.Fatalf("bad: %#v", m)
	}
	if m.stateOutPath != DefaultStateFilename {
		t.Fatalf("bad: %#v", m)
	}
	if m.backupPath != DefaultStateFilename+DefaultBackupExtension {
		t.Fatalf("bad: %#v", m)
	}

	m = new(Meta)
	m.statePath = "foo"
	m.initStatePaths()

	if m.stateOutPath != "foo" {
		t.Fatalf("bad: %#v", m)
	}
	if m.backupPath != "foo"+DefaultBackupExtension {
		t.Fatalf("bad: %#v", m)
	}

	m = new(Meta)
	m.stateOutPath = "foo"
	m.initStatePaths()

	if m.statePath != DefaultStateFilename {
		t.Fatalf("bad: %#v", m)
	}
	if m.backupPath != "foo"+DefaultBackupExtension {
		t.Fatalf("bad: %#v", m)
	}
}

func TestMeta_Env(t *testing.T) {
	td := tempDir(t)
	os.MkdirAll(td, 0755)
	defer os.RemoveAll(td)
	defer testChdir(t, td)()

	m := new(Meta)

	env, err := m.Workspace()
	if err != nil {
		t.Fatal(err)
	}

	if env != backend.DefaultStateName {
		t.Fatalf("expected env %q, got env %q", backend.DefaultStateName, env)
	}

	testEnv := "test_env"
	if err := m.SetWorkspace(testEnv); err != nil {
		t.Fatal("error setting env:", err)
	}

	env, _ = m.Workspace()
	if env != testEnv {
		t.Fatalf("expected env %q, got env %q", testEnv, env)
	}

	if err := m.SetWorkspace(backend.DefaultStateName); err != nil {
		t.Fatal("error setting env:", err)
	}

	env, _ = m.Workspace()
	if env != backend.DefaultStateName {
		t.Fatalf("expected env %q, got env %q", backend.DefaultStateName, env)
	}
}

func TestMeta_Workspace_override(t *testing.T) {
	defer func(value string) {
		os.Setenv(WorkspaceNameEnvVar, value)
	}(os.Getenv(WorkspaceNameEnvVar))

	m := new(Meta)

	testCases := map[string]struct {
		workspace string
		err       error
	}{
		"": {
			"default",
			nil,
		},
		"development": {
			"development",
			nil,
		},
		"invalid name": {
			"",
			invalidWorkspaceNameEnvVar,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			os.Setenv(WorkspaceNameEnvVar, name)
			workspace, err := m.Workspace()
			if workspace != tc.workspace {
				t.Errorf("Unexpected workspace\n got: %s\nwant: %s\n", workspace, tc.workspace)
			}
			if err != tc.err {
				t.Errorf("Unexpected error\n got: %s\nwant: %s\n", err, tc.err)
			}
		})
	}
}

func TestMeta_Workspace_invalidSelected(t *testing.T) {
	td := tempDir(t)
	os.MkdirAll(td, 0755)
	defer os.RemoveAll(td)
	defer testChdir(t, td)()

	// this is an invalid workspace name
	workspace := "test workspace"

	// create the workspace directories
	if err := os.MkdirAll(filepath.Join(local.DefaultWorkspaceDir, workspace), 0755); err != nil {
		t.Fatal(err)
	}

	// create the workspace file to select it
	if err := os.MkdirAll(DefaultDataDir, 0755); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(filepath.Join(DefaultDataDir, local.DefaultWorkspaceFile), []byte(workspace), 0644); err != nil {
		t.Fatal(err)
	}

	m := new(Meta)

	ws, err := m.Workspace()
	if ws != workspace {
		t.Errorf("Unexpected workspace\n got: %s\nwant: %s\n", ws, workspace)
	}
	if err != nil {
		t.Errorf("Unexpected error: %s", err)
	}
}

func TestMeta_process(t *testing.T) {
	test = false
	defer func() { test = true }()

	// Create a temporary directory for our cwd
	d := tempDir(t)
	os.MkdirAll(d, 0755)
	defer os.RemoveAll(d)
	defer testChdir(t, d)()

	// At one point it was the responsibility of this process function to
	// insert fake additional -var-file options into the command line
	// if the automatic tfvars files were present. This is no longer the
	// responsibility of process (it happens in collectVariableValues instead)
	// but we're still testing with these files in place to verify that
	// they _aren't_ being interpreted by process, since that could otherwise
	// cause them to be added more than once and mess up the precedence order.
	defaultVarsfile := "terraform.tfvars"
	err := ioutil.WriteFile(
		filepath.Join(d, defaultVarsfile),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	fileFirstAlphabetical := "a-file.auto.tfvars"
	err = ioutil.WriteFile(
		filepath.Join(d, fileFirstAlphabetical),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	fileLastAlphabetical := "z-file.auto.tfvars"
	err = ioutil.WriteFile(
		filepath.Join(d, fileLastAlphabetical),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	// Regular tfvars files will not be autoloaded
	fileIgnored := "ignored.tfvars"
	err = ioutil.WriteFile(
		filepath.Join(d, fileIgnored),
		[]byte(""),
		0644)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	tests := []struct {
		GivenArgs    []string
		FilteredArgs []string
		ExtraCheck   func(*testing.T, *Meta)
	}{
		{
			[]string{},
			[]string{},
			func(t *testing.T, m *Meta) {
				if got, want := m.color, true; got != want {
					t.Errorf("wrong m.color value %#v; want %#v", got, want)
				}
				if got, want := m.Color, true; got != want {
					t.Errorf("wrong m.Color value %#v; want %#v", got, want)
				}
			},
		},
		{
			[]string{"-no-color"},
			[]string{},
			func(t *testing.T, m *Meta) {
				if got, want := m.color, false; got != want {
					t.Errorf("wrong m.color value %#v; want %#v", got, want)
				}
				if got, want := m.Color, false; got != want {
					t.Errorf("wrong m.Color value %#v; want %#v", got, want)
				}
			},
		},
	}

	for _, test := range tests {
		t.Run(fmt.Sprintf("%s", test.GivenArgs), func(t *testing.T) {
			m := new(Meta)
			m.Color = true // this is the default also for normal use, overridden by -no-color
			args := test.GivenArgs
			args = m.process(args)

			if !cmp.Equal(test.FilteredArgs, args) {
				t.Errorf("wrong filtered arguments\n%s", cmp.Diff(test.FilteredArgs, args))
			}

			if test.ExtraCheck != nil {
				test.ExtraCheck(t, m)
			}
		})
	}
}
<|endoftext|>


<commit_before>package ortstore

import (
	"bytes"
	"fmt"
	"github.com/gholt/brimtime"
	"github.com/gholt/ring"
	"github.com/gholt/valuestore"
	"github.com/pandemicsyn/ort/rediscache"
	"github.com/spaolacci/murmur3"
	"log"
	"os"
	"time"
)

type OrtStore struct {
	vs    valuestore.ValueStore
	rfile string
	r     ring.Ring
}

func getMsgRing(filename string) (ring.MsgRing, error) {
	var f *os.File
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	r, err := ring.LoadRing(f)
	return ring.NewTCPMsgRing(r), nil
}

func New(rfile string, localid int) *OrtStore {
	//mr, err := getMsgRing(rfile)

	s := &OrtStore{}
	s.rfile = rfile

	f, err := os.Open(rfile)
	if err != nil {
		panic(err)
	}
	s.r, err = ring.LoadRing(f)
	if err != nil {
		panic(err)
	}
	fmt.Println("Ring entries:")
	for k, _ := range s.r.Nodes() {
		fmt.Println(s.r.Nodes()[k].ID(), s.r.Nodes()[k].Addresses())
	}
	fmt.Println("Pretending to be:", s.r.Nodes()[localid].ID(), s.r.Nodes()[localid].Addresses())
	s.r.SetLocalNode(s.r.Nodes()[localid].ID())
	t := ring.NewTCPMsgRing(s.r)
	go func() {
		for {
			t.Listen()
		}
	}()
	l := log.New(os.Stdout, "DebugStore ", log.LstdFlags)
	s.vs = valuestore.New(valuestore.OptMsgRing(t), valuestore.OptLogDebug(l))
	s.vs.EnableAll()
	return s
}

func (vsc *OrtStore) Get(key []byte, value []byte) []byte {
	keyA, keyB := murmur3.Sum128(key)
	var err error
	_, value, err = vsc.vs.Read(keyA, keyB, value)
	if err != nil {
		fmt.Printf("Get: %#v %s\n", string(key), err)
	}
	return value
}

func (vsc *OrtStore) Set(key []byte, value []byte) {
	if bytes.Equal(key, rediscache.BYTES_SHUTDOWN) && bytes.Equal(value, rediscache.BYTES_NOW) {
		vsc.vs.DisableAll()
		vsc.vs.Flush()
		fmt.Println(vsc.vs.GatherStats(true))
		//pprof.StopCPUProfile()
		//pproffp.Close()
		os.Exit(0)
		return
	}
	keyA, keyB := murmur3.Sum128(key)
	_, err := vsc.vs.Write(keyA, keyB, brimtime.TimeToUnixMicro(time.Now()), value)
	if err != nil {
		panic(err)
	}
}
<commit_msg>Updated due to Opt/Config changes in ValueStore<commit_after>package ortstore

import (
	"bytes"
	"fmt"
	"github.com/gholt/brimtime"
	"github.com/gholt/ring"
	"github.com/gholt/valuestore"
	"github.com/pandemicsyn/ort/rediscache"
	"github.com/spaolacci/murmur3"
	"log"
	"os"
	"time"
)

type OrtStore struct {
	vs    valuestore.ValueStore
	rfile string
	r     ring.Ring
}

func getMsgRing(filename string) (ring.MsgRing, error) {
	var f *os.File
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	r, err := ring.LoadRing(f)
	return ring.NewTCPMsgRing(r), nil
}

func New(rfile string, localid int) *OrtStore {
	//mr, err := getMsgRing(rfile)

	s := &OrtStore{}
	s.rfile = rfile

	f, err := os.Open(rfile)
	if err != nil {
		panic(err)
	}
	s.r, err = ring.LoadRing(f)
	if err != nil {
		panic(err)
	}
	fmt.Println("Ring entries:")
	for k, _ := range s.r.Nodes() {
		fmt.Println(s.r.Nodes()[k].ID(), s.r.Nodes()[k].Addresses())
	}
	fmt.Println("Pretending to be:", s.r.Nodes()[localid].ID(), s.r.Nodes()[localid].Addresses())
	s.r.SetLocalNode(s.r.Nodes()[localid].ID())
	t := ring.NewTCPMsgRing(s.r)
	go func() {
		for {
			t.Listen()
		}
	}()
	l := log.New(os.Stdout, "DebugStore ", log.LstdFlags)
	s.vs = valuestore.New(&valuestore.Config{MsgRing: t, LogDebug: l})
	s.vs.EnableAll()
	return s
}

func (vsc *OrtStore) Get(key []byte, value []byte) []byte {
	keyA, keyB := murmur3.Sum128(key)
	var err error
	_, value, err = vsc.vs.Read(keyA, keyB, value)
	if err != nil {
		fmt.Printf("Get: %#v %s\n", string(key), err)
	}
	return value
}

func (vsc *OrtStore) Set(key []byte, value []byte) {
	if bytes.Equal(key, rediscache.BYTES_SHUTDOWN) && bytes.Equal(value, rediscache.BYTES_NOW) {
		vsc.vs.DisableAll()
		vsc.vs.Flush()
		fmt.Println(vsc.vs.GatherStats(true))
		//pprof.StopCPUProfile()
		//pproffp.Close()
		os.Exit(0)
		return
	}
	keyA, keyB := murmur3.Sum128(key)
	_, err := vsc.vs.Write(keyA, keyB, brimtime.TimeToUnixMicro(time.Now()), value)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>