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

package signer

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/minio/minio-go/v7/pkg/s3utils"
)

// Signature and API related constants.
const (
	signV2Algorithm = "AWS"
)

// Encode input URL path to URL encoded path.
func encodeURL2Path(req *http.Request, virtualHost bool) (path string) {
	if virtualHost {
		reqHost := getHostAddr(req)
		dotPos := strings.Index(reqHost, ".")
		if dotPos > -1 {
			bucketName := reqHost[:dotPos]
			path = "/" + bucketName
			path += req.URL.Path
			path = s3utils.EncodePath(path)
			return
		}
	}
	path = s3utils.EncodePath(req.URL.Path)
	return
}

// PreSignV2 - presign the request in following style.
// https://${S3_BUCKET}.s3.amazonaws.com/${S3_OBJECT}?AWSAccessKeyId=${S3_ACCESS_KEY}&Expires=${TIMESTAMP}&Signature=${SIGNATURE}.
func PreSignV2(req http.Request, accessKeyID, secretAccessKey string, expires int64, virtualHost bool) *http.Request {
	// Presign is not needed for anonymous credentials.
	if accessKeyID == "" || secretAccessKey == "" {
		return &req
	}

	d := time.Now().UTC()
	// Find epoch expires when the request will expire.
	epochExpires := d.Unix() + expires

	// Add expires header if not present.
	if expiresStr := req.Header.Get("Expires"); expiresStr == "" {
		req.Header.Set("Expires", strconv.FormatInt(epochExpires, 10))
	}

	// Get presigned string to sign.
	stringToSign := preStringToSignV2(req, virtualHost)
	hm := hmac.New(sha1.New, []byte(secretAccessKey))
	hm.Write([]byte(stringToSign))

	// Calculate signature.
	signature := base64.StdEncoding.EncodeToString(hm.Sum(nil))

	query := req.URL.Query()
	// Handle specially for Google Cloud Storage.
	if strings.Contains(getHostAddr(&req), ".storage.googleapis.com") {
		query.Set("GoogleAccessId", accessKeyID)
	} else {
		query.Set("AWSAccessKeyId", accessKeyID)
	}

	// Fill in Expires for presigned query.
	query.Set("Expires", strconv.FormatInt(epochExpires, 10))

	// Encode query and save.
	req.URL.RawQuery = s3utils.QueryEncode(query)

	// Save signature finally.
	req.URL.RawQuery += "&Signature=" + s3utils.EncodePath(signature)

	// Return.
	return &req
}

// PostPresignSignatureV2 - presigned signature for PostPolicy
// request.
func PostPresignSignatureV2(policyBase64, secretAccessKey string) string {
	hm := hmac.New(sha1.New, []byte(secretAccessKey))
	hm.Write([]byte(policyBase64))
	signature := base64.StdEncoding.EncodeToString(hm.Sum(nil))
	return signature
}

// Authorization = "AWS" + " " + AWSAccessKeyId + ":" + Signature;
// Signature = Base64( HMAC-SHA1( YourSecretAccessKeyID, UTF-8-Encoding-Of( StringToSign ) ) );
//
// StringToSign = HTTP-Verb + "\n" +
//  	Content-Md5 + "\n" +
//  	Content-Type + "\n" +
//  	Date + "\n" +
//  	CanonicalizedProtocolHeaders +
//  	CanonicalizedResource;
//
// CanonicalizedResource = [ "/" + Bucket ] +
//  	<HTTP-Request-URI, from the protocol name up to the query string> +
//  	[ subresource, if present. For example "?acl", "?location", "?logging", or "?torrent"];
//
// CanonicalizedProtocolHeaders = <described below>

// SignV2 sign the request before Do() (AWS Signature Version 2).
func SignV2(req http.Request, accessKeyID, secretAccessKey string, virtualHost bool) *http.Request {
	// Signature calculation is not needed for anonymous credentials.
	if accessKeyID == "" || secretAccessKey == "" {
		return &req
	}

	// Initial time.
	d := time.Now().UTC()

	// Add date if not present.
	if date := req.Header.Get("Date"); date == "" {
		req.Header.Set("Date", d.Format(http.TimeFormat))
	}

	// Calculate HMAC for secretAccessKey.
	stringToSign := stringToSignV2(req, virtualHost)
	hm := hmac.New(sha1.New, []byte(secretAccessKey))
	hm.Write([]byte(stringToSign))

	// Prepare auth header.
	authHeader := new(bytes.Buffer)
	authHeader.WriteString(fmt.Sprintf("%s %s:", signV2Algorithm, accessKeyID))
	encoder := base64.NewEncoder(base64.StdEncoding, authHeader)
	encoder.Write(hm.Sum(nil))
	encoder.Close()

	// Set Authorization header.
	req.Header.Set("Authorization", authHeader.String())

	return &req
}

// From the Amazon docs:
//
// StringToSign = HTTP-Verb + "\n" +
// 	 Content-Md5 + "\n" +
//	 Content-Type + "\n" +
//	 Expires + "\n" +
//	 CanonicalizedProtocolHeaders +
//	 CanonicalizedResource;
func preStringToSignV2(req http.Request, virtualHost bool) string {
	buf := new(bytes.Buffer)
	// Write standard headers.
	writePreSignV2Headers(buf, req)
	// Write canonicalized protocol headers if any.
	writeCanonicalizedHeaders(buf, req)
	// Write canonicalized Query resources if any.
	writeCanonicalizedResource(buf, req, virtualHost)
	return buf.String()
}

// writePreSignV2Headers - write preSign v2 required headers.
func writePreSignV2Headers(buf *bytes.Buffer, req http.Request) {
	buf.WriteString(req.Method + "\n")
	buf.WriteString(req.Header.Get("Content-Md5") + "\n")
	buf.WriteString(req.Header.Get("Content-Type") + "\n")
	buf.WriteString(req.Header.Get("Expires") + "\n")
}

// From the Amazon docs:
//
// StringToSign = HTTP-Verb + "\n" +
// 	 Content-Md5 + "\n" +
//	 Content-Type + "\n" +
//	 Date + "\n" +
//	 CanonicalizedProtocolHeaders +
//	 CanonicalizedResource;
func stringToSignV2(req http.Request, virtualHost bool) string {
	buf := new(bytes.Buffer)
	// Write standard headers.
	writeSignV2Headers(buf, req)
	// Write canonicalized protocol headers if any.
	writeCanonicalizedHeaders(buf, req)
	// Write canonicalized Query resources if any.
	writeCanonicalizedResource(buf, req, virtualHost)
	return buf.String()
}

// writeSignV2Headers - write signV2 required headers.
func writeSignV2Headers(buf *bytes.Buffer, req http.Request) {
	buf.WriteString(req.Method + "\n")
	buf.WriteString(req.Header.Get("Content-Md5") + "\n")
	buf.WriteString(req.Header.Get("Content-Type") + "\n")
	buf.WriteString(req.Header.Get("Date") + "\n")
}

// writeCanonicalizedHeaders - write canonicalized headers.
func writeCanonicalizedHeaders(buf *bytes.Buffer, req http.Request) {
	var protoHeaders []string
	vals := make(map[string][]string)
	for k, vv := range req.Header {
		// All the AMZ headers should be lowercase
		lk := strings.ToLower(k)
		if strings.HasPrefix(lk, "x-amz") {
			protoHeaders = append(protoHeaders, lk)
			vals[lk] = vv
		}
	}
	sort.Strings(protoHeaders)
	for _, k := range protoHeaders {
		buf.WriteString(k)
		buf.WriteByte(':')
		for idx, v := range vals[k] {
			if idx > 0 {
				buf.WriteByte(',')
			}
			buf.WriteString(v)
		}
		buf.WriteByte('\n')
	}
}

// AWS S3 Signature V2 calculation rule is give here:
// http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html#RESTAuthenticationStringToSign

// Whitelist resource list that will be used in query string for signature-V2 calculation.
// The list should be alphabetically sorted
var resourceList = []string{
	"acl",
	"delete",
	"lifecycle",
	"location",
	"logging",
	"notification",
	"partNumber",
	"policy",
	"replication",
	"requestPayment",
	"response-cache-control",
	"response-content-disposition",
	"response-content-encoding",
	"response-content-language",
	"response-content-type",
	"response-expires",
	"torrent",
	"uploadId",
	"uploads",
	"versionId",
	"versioning",
	"versions",
	"website",
}

// From the Amazon docs:
//
// CanonicalizedResource = [ "/" + Bucket ] +
// 	  <HTTP-Request-URI, from the protocol name up to the query string> +
// 	  [ sub-resource, if present. For example "?acl", "?location", "?logging", or "?torrent"];
func writeCanonicalizedResource(buf *bytes.Buffer, req http.Request, virtualHost bool) {
	// Save request URL.
	requestURL := req.URL
	// Get encoded URL path.
	buf.WriteString(encodeURL2Path(&req, virtualHost))
	if requestURL.RawQuery != "" {
		var n int
		vals, _ := url.ParseQuery(requestURL.RawQuery)
		// Verify if any sub resource queries are present, if yes
		// canonicallize them.
		for _, resource := range resourceList {
			if vv, ok := vals[resource]; ok && len(vv) > 0 {
				n++
				// First element
				switch n {
				case 1:
					buf.WriteByte('?')
				// The rest
				default:
					buf.WriteByte('&')
				}
				buf.WriteString(resource)
				// Request parameters
				if len(vv[0]) > 0 {
					buf.WriteByte('=')
					buf.WriteString(vv[0])
				}
			}
		}
	}
}
<commit_msg>add missing query params for signature v2<commit_after>/*
 * MinIO Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2015-2017 MinIO, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package signer

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/minio/minio-go/v7/pkg/s3utils"
)

// Signature and API related constants.
const (
	signV2Algorithm = "AWS"
)

// Encode input URL path to URL encoded path.
func encodeURL2Path(req *http.Request, virtualHost bool) (path string) {
	if virtualHost {
		reqHost := getHostAddr(req)
		dotPos := strings.Index(reqHost, ".")
		if dotPos > -1 {
			bucketName := reqHost[:dotPos]
			path = "/" + bucketName
			path += req.URL.Path
			path = s3utils.EncodePath(path)
			return
		}
	}
	path = s3utils.EncodePath(req.URL.Path)
	return
}

// PreSignV2 - presign the request in following style.
// https://${S3_BUCKET}.s3.amazonaws.com/${S3_OBJECT}?AWSAccessKeyId=${S3_ACCESS_KEY}&Expires=${TIMESTAMP}&Signature=${SIGNATURE}.
func PreSignV2(req http.Request, accessKeyID, secretAccessKey string, expires int64, virtualHost bool) *http.Request {
	// Presign is not needed for anonymous credentials.
	if accessKeyID == "" || secretAccessKey == "" {
		return &req
	}

	d := time.Now().UTC()
	// Find epoch expires when the request will expire.
	epochExpires := d.Unix() + expires

	// Add expires header if not present.
	if expiresStr := req.Header.Get("Expires"); expiresStr == "" {
		req.Header.Set("Expires", strconv.FormatInt(epochExpires, 10))
	}

	// Get presigned string to sign.
	stringToSign := preStringToSignV2(req, virtualHost)
	hm := hmac.New(sha1.New, []byte(secretAccessKey))
	hm.Write([]byte(stringToSign))

	// Calculate signature.
	signature := base64.StdEncoding.EncodeToString(hm.Sum(nil))

	query := req.URL.Query()
	// Handle specially for Google Cloud Storage.
	if strings.Contains(getHostAddr(&req), ".storage.googleapis.com") {
		query.Set("GoogleAccessId", accessKeyID)
	} else {
		query.Set("AWSAccessKeyId", accessKeyID)
	}

	// Fill in Expires for presigned query.
	query.Set("Expires", strconv.FormatInt(epochExpires, 10))

	// Encode query and save.
	req.URL.RawQuery = s3utils.QueryEncode(query)

	// Save signature finally.
	req.URL.RawQuery += "&Signature=" + s3utils.EncodePath(signature)

	// Return.
	return &req
}

// PostPresignSignatureV2 - presigned signature for PostPolicy
// request.
func PostPresignSignatureV2(policyBase64, secretAccessKey string) string {
	hm := hmac.New(sha1.New, []byte(secretAccessKey))
	hm.Write([]byte(policyBase64))
	signature := base64.StdEncoding.EncodeToString(hm.Sum(nil))
	return signature
}

// Authorization = "AWS" + " " + AWSAccessKeyId + ":" + Signature;
// Signature = Base64( HMAC-SHA1( YourSecretAccessKeyID, UTF-8-Encoding-Of( StringToSign ) ) );
//
// StringToSign = HTTP-Verb + "\n" +
//  	Content-Md5 + "\n" +
//  	Content-Type + "\n" +
//  	Date + "\n" +
//  	CanonicalizedProtocolHeaders +
//  	CanonicalizedResource;
//
// CanonicalizedResource = [ "/" + Bucket ] +
//  	<HTTP-Request-URI, from the protocol name up to the query string> +
//  	[ subresource, if present. For example "?acl", "?location", "?logging", or "?torrent"];
//
// CanonicalizedProtocolHeaders = <described below>

// SignV2 sign the request before Do() (AWS Signature Version 2).
func SignV2(req http.Request, accessKeyID, secretAccessKey string, virtualHost bool) *http.Request {
	// Signature calculation is not needed for anonymous credentials.
	if accessKeyID == "" || secretAccessKey == "" {
		return &req
	}

	// Initial time.
	d := time.Now().UTC()

	// Add date if not present.
	if date := req.Header.Get("Date"); date == "" {
		req.Header.Set("Date", d.Format(http.TimeFormat))
	}

	// Calculate HMAC for secretAccessKey.
	stringToSign := stringToSignV2(req, virtualHost)
	hm := hmac.New(sha1.New, []byte(secretAccessKey))
	hm.Write([]byte(stringToSign))

	// Prepare auth header.
	authHeader := new(bytes.Buffer)
	authHeader.WriteString(fmt.Sprintf("%s %s:", signV2Algorithm, accessKeyID))
	encoder := base64.NewEncoder(base64.StdEncoding, authHeader)
	encoder.Write(hm.Sum(nil))
	encoder.Close()

	// Set Authorization header.
	req.Header.Set("Authorization", authHeader.String())

	return &req
}

// From the Amazon docs:
//
// StringToSign = HTTP-Verb + "\n" +
// 	 Content-Md5 + "\n" +
//	 Content-Type + "\n" +
//	 Expires + "\n" +
//	 CanonicalizedProtocolHeaders +
//	 CanonicalizedResource;
func preStringToSignV2(req http.Request, virtualHost bool) string {
	buf := new(bytes.Buffer)
	// Write standard headers.
	writePreSignV2Headers(buf, req)
	// Write canonicalized protocol headers if any.
	writeCanonicalizedHeaders(buf, req)
	// Write canonicalized Query resources if any.
	writeCanonicalizedResource(buf, req, virtualHost)
	return buf.String()
}

// writePreSignV2Headers - write preSign v2 required headers.
func writePreSignV2Headers(buf *bytes.Buffer, req http.Request) {
	buf.WriteString(req.Method + "\n")
	buf.WriteString(req.Header.Get("Content-Md5") + "\n")
	buf.WriteString(req.Header.Get("Content-Type") + "\n")
	buf.WriteString(req.Header.Get("Expires") + "\n")
}

// From the Amazon docs:
//
// StringToSign = HTTP-Verb + "\n" +
// 	 Content-Md5 + "\n" +
//	 Content-Type + "\n" +
//	 Date + "\n" +
//	 CanonicalizedProtocolHeaders +
//	 CanonicalizedResource;
func stringToSignV2(req http.Request, virtualHost bool) string {
	buf := new(bytes.Buffer)
	// Write standard headers.
	writeSignV2Headers(buf, req)
	// Write canonicalized protocol headers if any.
	writeCanonicalizedHeaders(buf, req)
	// Write canonicalized Query resources if any.
	writeCanonicalizedResource(buf, req, virtualHost)
	return buf.String()
}

// writeSignV2Headers - write signV2 required headers.
func writeSignV2Headers(buf *bytes.Buffer, req http.Request) {
	buf.WriteString(req.Method + "\n")
	buf.WriteString(req.Header.Get("Content-Md5") + "\n")
	buf.WriteString(req.Header.Get("Content-Type") + "\n")
	buf.WriteString(req.Header.Get("Date") + "\n")
}

// writeCanonicalizedHeaders - write canonicalized headers.
func writeCanonicalizedHeaders(buf *bytes.Buffer, req http.Request) {
	var protoHeaders []string
	vals := make(map[string][]string)
	for k, vv := range req.Header {
		// All the AMZ headers should be lowercase
		lk := strings.ToLower(k)
		if strings.HasPrefix(lk, "x-amz") {
			protoHeaders = append(protoHeaders, lk)
			vals[lk] = vv
		}
	}
	sort.Strings(protoHeaders)
	for _, k := range protoHeaders {
		buf.WriteString(k)
		buf.WriteByte(':')
		for idx, v := range vals[k] {
			if idx > 0 {
				buf.WriteByte(',')
			}
			buf.WriteString(v)
		}
		buf.WriteByte('\n')
	}
}

// AWS S3 Signature V2 calculation rule is give here:
// http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html#RESTAuthenticationStringToSign

// Whitelist resource list that will be used in query string for signature-V2 calculation.
//
// This list should be kept alphabetically sorted, do not hastily edit.
var resourceList = []string{
	"acl",
	"cors",
	"delete",
	"encryption",
	"legal-hold",
	"lifecycle",
	"location",
	"logging",
	"notification",
	"partNumber",
	"policy",
	"replication",
	"requestPayment",
	"response-cache-control",
	"response-content-disposition",
	"response-content-encoding",
	"response-content-language",
	"response-content-type",
	"response-expires",
	"retention",
	"select",
	"select-type",
	"tagging",
	"torrent",
	"uploadId",
	"uploads",
	"versionId",
	"versioning",
	"versions",
	"website",
}

// From the Amazon docs:
//
// CanonicalizedResource = [ "/" + Bucket ] +
// 	  <HTTP-Request-URI, from the protocol name up to the query string> +
// 	  [ sub-resource, if present. For example "?acl", "?location", "?logging", or "?torrent"];
func writeCanonicalizedResource(buf *bytes.Buffer, req http.Request, virtualHost bool) {
	// Save request URL.
	requestURL := req.URL
	// Get encoded URL path.
	buf.WriteString(encodeURL2Path(&req, virtualHost))
	if requestURL.RawQuery != "" {
		var n int
		vals, _ := url.ParseQuery(requestURL.RawQuery)
		// Verify if any sub resource queries are present, if yes
		// canonicallize them.
		for _, resource := range resourceList {
			if vv, ok := vals[resource]; ok && len(vv) > 0 {
				n++
				// First element
				switch n {
				case 1:
					buf.WriteByte('?')
				// The rest
				default:
					buf.WriteByte('&')
				}
				buf.WriteString(resource)
				// Request parameters
				if len(vv[0]) > 0 {
					buf.WriteByte('=')
					buf.WriteString(vv[0])
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package cloudwatch

import (
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/credentials/endpointcreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/sts"
)

type cache struct {
	credential *credentials.Credentials
	expiration *time.Time
}

var awsCredentialCache map[string]cache = make(map[string]cache)
var credentialCacheLock sync.RWMutex

func GetCredentials(dsInfo *DatasourceInfo) (*credentials.Credentials, error) {
	cacheKey := dsInfo.AccessKey + ":" + dsInfo.Profile + ":" + dsInfo.AssumeRoleArn
	credentialCacheLock.RLock()
	if _, ok := awsCredentialCache[cacheKey]; ok {
		if awsCredentialCache[cacheKey].expiration != nil &&
			(*awsCredentialCache[cacheKey].expiration).After(time.Now().UTC()) {
			result := awsCredentialCache[cacheKey].credential
			credentialCacheLock.RUnlock()
			return result, nil
		}
	}
	credentialCacheLock.RUnlock()

	accessKeyId := ""
	secretAccessKey := ""
	sessionToken := ""
	var expiration *time.Time
	expiration = nil
	if dsInfo.AuthType == "arn" && strings.Index(dsInfo.AssumeRoleArn, "arn:aws:iam:") == 0 {
		params := &sts.AssumeRoleInput{
			RoleArn:         aws.String(dsInfo.AssumeRoleArn),
			RoleSessionName: aws.String("GrafanaSession"),
			DurationSeconds: aws.Int64(900),
		}

		stsSess, err := session.NewSession()
		if err != nil {
			return nil, err
		}
		stsCreds := credentials.NewChainCredentials(
			[]credentials.Provider{
				&credentials.EnvProvider{},
				&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
				remoteCredProvider(stsSess),
			})
		stsConfig := &aws.Config{
			Region:      aws.String(dsInfo.Region),
			Credentials: stsCreds,
		}

		sess, err := session.NewSession(stsConfig)
		if err != nil {
			return nil, err
		}
		svc := sts.New(sess, stsConfig)
		resp, err := svc.AssumeRole(params)
		if err != nil {
			return nil, err
		}
		if resp.Credentials != nil {
			accessKeyId = *resp.Credentials.AccessKeyId
			secretAccessKey = *resp.Credentials.SecretAccessKey
			sessionToken = *resp.Credentials.SessionToken
			expiration = resp.Credentials.Expiration
		}
	} else {
		now := time.Now()
		e := now.Add(5 * time.Minute)
		expiration = &e
	}

	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     accessKeyId,
				SecretAccessKey: secretAccessKey,
				SessionToken:    sessionToken,
			}},
			&credentials.EnvProvider{},
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     dsInfo.AccessKey,
				SecretAccessKey: dsInfo.SecretKey,
			}},
			&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
			remoteCredProvider(sess),
		})

	credentialCacheLock.Lock()
	awsCredentialCache[cacheKey] = cache{
		credential: creds,
		expiration: expiration,
	}
	credentialCacheLock.Unlock()

	return creds, nil
}

func remoteCredProvider(sess *session.Session) credentials.Provider {
	ecsCredURI := os.Getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")

	if len(ecsCredURI) > 0 {
		return ecsCredProvider(sess, ecsCredURI)
	}
	return ec2RoleProvider(sess)
}

func ecsCredProvider(sess *session.Session, uri string) credentials.Provider {
	const host = `169.254.170.2`

	d := defaults.Get()
	return endpointcreds.NewProviderClient(
		*d.Config,
		d.Handlers,
		fmt.Sprintf("http://%s%s", host, uri),
		func(p *endpointcreds.Provider) { p.ExpiryWindow = 5 * time.Minute })
}

func ec2RoleProvider(sess *session.Session) credentials.Provider {
	return &ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: 5 * time.Minute}
}

func (e *CloudWatchExecutor) getDsInfo(region string) *DatasourceInfo {
	authType := e.DataSource.JsonData.Get("authType").MustString()
	assumeRoleArn := e.DataSource.JsonData.Get("assumeRoleArn").MustString()
	accessKey := ""
	secretKey := ""
	for key, value := range e.DataSource.SecureJsonData.Decrypt() {
		if key == "accessKey" {
			accessKey = value
		}
		if key == "secretKey" {
			secretKey = value
		}
	}

	datasourceInfo := &DatasourceInfo{
		Region:        region,
		Profile:       e.DataSource.Database,
		AuthType:      authType,
		AssumeRoleArn: assumeRoleArn,
		AccessKey:     accessKey,
		SecretKey:     secretKey,
	}

	return datasourceInfo
}

func (e *CloudWatchExecutor) getAwsConfig(dsInfo *DatasourceInfo) (*aws.Config, error) {
	creds, err := GetCredentials(dsInfo)
	if err != nil {
		return nil, err
	}

	cfg := &aws.Config{
		Region:      aws.String(dsInfo.Region),
		Credentials: creds,
	}
	return cfg, nil
}

func (e *CloudWatchExecutor) getClient(region string) (*cloudwatch.CloudWatch, error) {
	datasourceInfo := e.getDsInfo(region)
	cfg, err := e.getAwsConfig(datasourceInfo)
	if err != nil {
		return nil, err
	}

	sess, err := session.NewSession(cfg)
	if err != nil {
		return nil, err
	}

	client := cloudwatch.New(sess, cfg)
	return client, nil
}
<commit_msg>Fix go fmt<commit_after>package cloudwatch

import (
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/credentials/endpointcreds"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/sts"
)

type cache struct {
	credential *credentials.Credentials
	expiration *time.Time
}

var awsCredentialCache map[string]cache = make(map[string]cache)
var credentialCacheLock sync.RWMutex

func GetCredentials(dsInfo *DatasourceInfo) (*credentials.Credentials, error) {
	cacheKey := dsInfo.AccessKey + ":" + dsInfo.Profile + ":" + dsInfo.AssumeRoleArn
	credentialCacheLock.RLock()
	if _, ok := awsCredentialCache[cacheKey]; ok {
		if awsCredentialCache[cacheKey].expiration != nil &&
			(*awsCredentialCache[cacheKey].expiration).After(time.Now().UTC()) {
			result := awsCredentialCache[cacheKey].credential
			credentialCacheLock.RUnlock()
			return result, nil
		}
	}
	credentialCacheLock.RUnlock()

	accessKeyId := ""
	secretAccessKey := ""
	sessionToken := ""
	var expiration *time.Time
	expiration = nil
	if dsInfo.AuthType == "arn" && strings.Index(dsInfo.AssumeRoleArn, "arn:aws:iam:") == 0 {
		params := &sts.AssumeRoleInput{
			RoleArn:         aws.String(dsInfo.AssumeRoleArn),
			RoleSessionName: aws.String("GrafanaSession"),
			DurationSeconds: aws.Int64(900),
		}

		stsSess, err := session.NewSession()
		if err != nil {
			return nil, err
		}
		stsCreds := credentials.NewChainCredentials(
			[]credentials.Provider{
				&credentials.EnvProvider{},
				&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
				remoteCredProvider(stsSess),
			})
		stsConfig := &aws.Config{
			Region:      aws.String(dsInfo.Region),
			Credentials: stsCreds,
		}

		sess, err := session.NewSession(stsConfig)
		if err != nil {
			return nil, err
		}
		svc := sts.New(sess, stsConfig)
		resp, err := svc.AssumeRole(params)
		if err != nil {
			return nil, err
		}
		if resp.Credentials != nil {
			accessKeyId = *resp.Credentials.AccessKeyId
			secretAccessKey = *resp.Credentials.SecretAccessKey
			sessionToken = *resp.Credentials.SessionToken
			expiration = resp.Credentials.Expiration
		}
	} else {
		now := time.Now()
		e := now.Add(5 * time.Minute)
		expiration = &e
	}

	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     accessKeyId,
				SecretAccessKey: secretAccessKey,
				SessionToken:    sessionToken,
			}},
			&credentials.EnvProvider{},
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     dsInfo.AccessKey,
				SecretAccessKey: dsInfo.SecretKey,
			}},
			&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
			remoteCredProvider(sess),
		})

	credentialCacheLock.Lock()
	awsCredentialCache[cacheKey] = cache{
		credential: creds,
		expiration: expiration,
	}
	credentialCacheLock.Unlock()

	return creds, nil
}

func remoteCredProvider(sess *session.Session) credentials.Provider {
	ecsCredURI := os.Getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")

	if len(ecsCredURI) > 0 {
		return ecsCredProvider(sess, ecsCredURI)
	}
	return ec2RoleProvider(sess)
}

func ecsCredProvider(sess *session.Session, uri string) credentials.Provider {
	const host = `169.254.170.2`

	d := defaults.Get()
	return endpointcreds.NewProviderClient(
		*d.Config,
		d.Handlers,
		fmt.Sprintf("http://%s%s", host, uri),
		func(p *endpointcreds.Provider) { p.ExpiryWindow = 5 * time.Minute })
}

func ec2RoleProvider(sess *session.Session) credentials.Provider {
	return &ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: 5 * time.Minute}
}

func (e *CloudWatchExecutor) getDsInfo(region string) *DatasourceInfo {
	authType := e.DataSource.JsonData.Get("authType").MustString()
	assumeRoleArn := e.DataSource.JsonData.Get("assumeRoleArn").MustString()
	accessKey := ""
	secretKey := ""
	for key, value := range e.DataSource.SecureJsonData.Decrypt() {
		if key == "accessKey" {
			accessKey = value
		}
		if key == "secretKey" {
			secretKey = value
		}
	}

	datasourceInfo := &DatasourceInfo{
		Region:        region,
		Profile:       e.DataSource.Database,
		AuthType:      authType,
		AssumeRoleArn: assumeRoleArn,
		AccessKey:     accessKey,
		SecretKey:     secretKey,
	}

	return datasourceInfo
}

func (e *CloudWatchExecutor) getAwsConfig(dsInfo *DatasourceInfo) (*aws.Config, error) {
	creds, err := GetCredentials(dsInfo)
	if err != nil {
		return nil, err
	}

	cfg := &aws.Config{
		Region:      aws.String(dsInfo.Region),
		Credentials: creds,
	}
	return cfg, nil
}

func (e *CloudWatchExecutor) getClient(region string) (*cloudwatch.CloudWatch, error) {
	datasourceInfo := e.getDsInfo(region)
	cfg, err := e.getAwsConfig(datasourceInfo)
	if err != nil {
		return nil, err
	}

	sess, err := session.NewSession(cfg)
	if err != nil {
		return nil, err
	}

	client := cloudwatch.New(sess, cfg)
	return client, nil
}
<|endoftext|>


<commit_before>package pixel

import (
	"time"
	"log"
	"fmt"
	"math"
	"sync"

	"github.com/tarm/serial"
)

type Pixel struct{
	PortName      string
	PortSpeed     int
	lastPixelData PixelData
	serial        *serial.Port
	mutex         sync.Mutex
}

type PixelData struct {
	Value      int
	Message    string
	Blink      int
	Brightness int
}

func (p *Pixel) Connect(){
	p.mutex = sync.Mutex{}
	p.lastPixelData = PixelData{ -1, "", 0, 100 }

	c := &serial.Config{Name: p.PortName, Baud: p.PortSpeed}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Fatalf("Could not open port %s, %s", c.Name, err)
	}
	p.serial = s

	// port not opened before 1500 milliseconds pause
	time.Sleep(1500 * time.Millisecond)
}

func (p Pixel) GetStatus() PixelData{
	return p.lastPixelData
}

func (p *Pixel) SetStatus(pd PixelData){
	p.mutex.Lock()
	animationDuration := 3000 // ms

	switch pd.Blink {
	case 1:
		pd.Value += 100
	case 2:
		pd.Value += 200
	}

	delta := pd.Value - p.lastPixelData.Value
	stepTime := float64(animationDuration) / math.Abs(float64(delta))

	var step int
	if delta > 0{
		step = 1
	} else {
		step = -1
	}

	if delta > 0{
		// smooth switch color
		if(pd.Blink == 0 && p.lastPixelData.Value > 0 && pd.Value > 0) {
			for i := p.lastPixelData.Value; i != pd.Value; i += step {
				p.sendSerial(PixelData{i, "", 0, pd.Brightness })
				time.Sleep(time.Millisecond * time.Duration(stepTime))
			}
		}
	} else {
		// sharp switch color
		if(pd.Blink == 0 && p.lastPixelData.Value > 0 && pd.Value > 0) {
			for i := 0; i < 3; i++ {
				p.sendSerial(PixelData{pd.Value, "", 0, pd.Brightness })
				time.Sleep(time.Millisecond * 250)
				p.sendSerial(PixelData{p.lastPixelData.Value, "", 0, pd.Brightness })
				time.Sleep(time.Millisecond * 250)
			}
			p.sendSerial(PixelData{pd.Value, "", 0, pd.Brightness })
		}
	}

	time.Sleep(100 * time.Duration(stepTime))
	log.Printf("Pixel.setStatus: %v\n", pd)
	p.sendSerial(pd)

	p.lastPixelData = pd

	// if success value, turn off led
	if pd.Value == 100{
		time.Sleep(time.Millisecond * 5000)
		p.sendSerial(PixelData{ -1, "", 0, 100 })
		p.lastPixelData = PixelData{ -1, "", 0, 100 }
		log.Printf("lastPixelData: %v", p.lastPixelData)
	}

	time.Sleep(1000 * time.Millisecond)
	p.mutex.Unlock()
}

func (p Pixel) sendSerial (pd PixelData) (int, error){
	command := fmt.Sprintf("%d|%s|%d\n",pd.Value, pd.Message, pd.Brightness)
	//log.Println(command)
	n, err := p.serial.Write([]byte(command))
	if err != nil {
		log.Fatalf("Could not write to port, %s", err)
	}
	return n, err
}
<commit_msg>refactor(pixel): set Pixel embedded lock<commit_after>package pixel

import (
	"time"
	"log"
	"fmt"
	"math"
	"sync"

	"github.com/tarm/serial"
)

type Pixel struct{
	sync.Mutex
	PortName      string
	PortSpeed     int
	lastPixelData PixelData
	serial        *serial.Port
}

type PixelData struct {
	Value      int
	Message    string
	Blink      int
	Brightness int
}

func (p *Pixel) Connect(){
	p.lastPixelData = PixelData{ -1, "", 0, 100 }

	c := &serial.Config{Name: p.PortName, Baud: p.PortSpeed}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Fatalf("Could not open port %s, %s", c.Name, err)
	}
	p.serial = s

	// port not opened before 1500 milliseconds pause
	time.Sleep(1500 * time.Millisecond)
}

func (p Pixel) GetStatus() PixelData{
	return p.lastPixelData
}

func (p *Pixel) SetStatus(pd PixelData){
	p.Lock()
	animationDuration := 3000 // ms

	switch pd.Blink {
	case 1:
		pd.Value += 100
	case 2:
		pd.Value += 200
	}

	delta := pd.Value - p.lastPixelData.Value
	stepTime := float64(animationDuration) / math.Abs(float64(delta))

	var step int
	if delta > 0{
		step = 1
	} else {
		step = -1
	}

	if delta > 0{
		// smooth switch color
		if(pd.Blink == 0 && p.lastPixelData.Value > 0 && pd.Value > 0) {
			for i := p.lastPixelData.Value; i != pd.Value; i += step {
				p.sendSerial(PixelData{i, "", 0, pd.Brightness })
				time.Sleep(time.Millisecond * time.Duration(stepTime))
			}
		}
	} else {
		// sharp switch color
		if(pd.Blink == 0 && p.lastPixelData.Value > 0 && pd.Value > 0) {
			for i := 0; i < 3; i++ {
				p.sendSerial(PixelData{pd.Value, "", 0, pd.Brightness })
				time.Sleep(time.Millisecond * 250)
				p.sendSerial(PixelData{p.lastPixelData.Value, "", 0, pd.Brightness })
				time.Sleep(time.Millisecond * 250)
			}
			p.sendSerial(PixelData{pd.Value, "", 0, pd.Brightness })
		}
	}

	time.Sleep(100 * time.Duration(stepTime))
	log.Printf("Pixel.setStatus: %v\n", pd)
	p.sendSerial(pd)

	p.lastPixelData = pd

	// if success value, turn off led
	if pd.Value == 100{
		time.Sleep(time.Millisecond * 5000)
		p.sendSerial(PixelData{ -1, "", 0, 100 })
		p.lastPixelData = PixelData{ -1, "", 0, 100 }
		log.Printf("lastPixelData: %v", p.lastPixelData)
	}

	time.Sleep(1000 * time.Millisecond)
	p.Unlock()
}

func (p Pixel) sendSerial (pd PixelData) (int, error){
	command := fmt.Sprintf("%d|%s|%d\n",pd.Value, pd.Message, pd.Brightness)
	//log.Println(command)
	n, err := p.serial.Write([]byte(command))
	if err != nil {
		log.Fatalf("Could not write to port, %s", err)
	}
	return n, err
}
<|endoftext|>


<commit_before>package wats

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
)

func unbindSecurityGroups() []string {
	var securityGroups []string

	AsUser(environment.AdminUserContext(), time.Minute, func() {
		out, err := runCfWithOutput("curl", "/v2/config/running_security_groups")
		Expect(err).NotTo(HaveOccurred())
		var result map[string]interface{}
		err = json.Unmarshal(out.Contents(), &result)
		Expect(err).NotTo(HaveOccurred())

		resources := result["resources"].([]interface{})
		for _, group := range resources {
			foo := group.(map[string]interface{})
			entity := foo["entity"].(map[string]interface{})
			name := entity["name"].(string)
			securityGroups = append(securityGroups, name)
			_, err = runCfWithOutput("unbind-running-security-group", name)
			Expect(err).NotTo(HaveOccurred())
		}
	})
	return securityGroups
}

func bindSecurityGroups(groups []string) {
	AsUser(environment.AdminUserContext(), time.Minute, func() {
		for _, group := range groups {
			_, err := runCfWithOutput("bind-running-security-group", group)
			Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to recreate running-security-group %s", group))
		}
	})
}

var _ = Describe("Security Groups", func() {
	type NoraCurlResponse struct {
		Stdout     string
		Stderr     string
		ReturnCode int `json:"return_code"`
	}

	It("allows traffic and then blocks traffic", func() {
		groups := unbindSecurityGroups()
		defer bindSecurityGroups(groups)

		By("pushing it")
		Expect(pushNora(appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))

		By("staging and running it on Diego")
		enableDiego(appName)
		Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))

		By("verifying it's up")
		Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))

		secureAddress := config.GetSecureAddress()
		secureHost, securePort, err := net.SplitHostPort(secureAddress)
		Expect(err).NotTo(HaveOccurred())

		// test app egress rules
		curlResponse := func() int {
			var noraCurlResponse NoraCurlResponse
			resp := helpers.CurlApp(config, appName, fmt.Sprintf("/curl/%s/%s", secureHost, securePort))
			json.Unmarshal([]byte(resp), &noraCurlResponse)
			return noraCurlResponse.ReturnCode
		}
		firstCurlError := curlResponse()
		Expect(firstCurlError).ShouldNot(Equal(0))

		// apply security group
		rules := fmt.Sprintf(`[{"destination":"%s","ports":"%s","protocol":"tcp"}]`, secureHost, securePort)

		file, _ := ioutil.TempFile(os.TempDir(), "DATS-sg-rules")
		defer os.Remove(file.Name())
		file.WriteString(rules)
		file.Close()

		rulesPath := file.Name()
		securityGroupName := fmt.Sprintf("DATS-SG-%s", generator.PrefixedRandomName(config.GetNamePrefix(), "SECURITY-GROUP"))

		AsUser(environment.AdminUserContext(), time.Minute, func() {
			Expect(cf.Cf("create-security-group", securityGroupName, rulesPath).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			Expect(cf.Cf("bind-security-group",
				securityGroupName,
				environment.RegularUserContext().Org,
				environment.RegularUserContext().Space).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		})
		defer func() {
			AsUser(environment.AdminUserContext(), time.Minute, func() {
				Expect(cf.Cf("delete-security-group", securityGroupName, "-f").Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			})
		}()

		Expect(cf.Cf("restart", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))

		// test app egress rules
		Eventually(curlResponse).Should(Equal(0))

		// unapply security group
		AsUser(environment.AdminUserContext(), time.Minute, func() {
			Expect(cf.Cf("unbind-security-group",
				securityGroupName, environment.RegularUserContext().Org,
				environment.RegularUserContext().Space).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		})

		By("restarting it - without security group")
		Expect(cf.Cf("restart", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))

		// test app egress rules
		Eventually(curlResponse).Should(Equal(firstCurlError))
	})
})
<commit_msg>add ICMP ignore test<commit_after>package wats

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
)

func unbindSecurityGroups() []string {
	var securityGroups []string

	AsUser(environment.AdminUserContext(), time.Minute, func() {
		out, err := runCfWithOutput("curl", "/v2/config/running_security_groups")
		Expect(err).NotTo(HaveOccurred())
		var result map[string]interface{}
		err = json.Unmarshal(out.Contents(), &result)
		Expect(err).NotTo(HaveOccurred())

		resources := result["resources"].([]interface{})
		for _, group := range resources {
			foo := group.(map[string]interface{})
			entity := foo["entity"].(map[string]interface{})
			name := entity["name"].(string)
			securityGroups = append(securityGroups, name)
			_, err = runCfWithOutput("unbind-running-security-group", name)
			Expect(err).NotTo(HaveOccurred())
		}
	})
	return securityGroups
}

func bindSecurityGroups(groups []string) {
	AsUser(environment.AdminUserContext(), time.Minute, func() {
		for _, group := range groups {
			_, err := runCfWithOutput("bind-running-security-group", group)
			Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to recreate running-security-group %s", group))
		}
	})
}

var _ = Describe("Security Groups", func() {
	type NoraCurlResponse struct {
		Stdout     string
		Stderr     string
		ReturnCode int `json:"return_code"`
	}

	It("allows traffic and then blocks traffic", func() {
		groups := unbindSecurityGroups()
		defer bindSecurityGroups(groups)

		By("pushing it")
		Expect(pushNora(appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))

		By("staging and running it on Diego")
		enableDiego(appName)
		Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))

		By("verifying it's up")
		Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))

		secureAddress := config.GetSecureAddress()
		secureHost, securePort, err := net.SplitHostPort(secureAddress)
		Expect(err).NotTo(HaveOccurred())

		// test app egress rules
		curlResponse := func() int {
			var noraCurlResponse NoraCurlResponse
			resp := helpers.CurlApp(config, appName, fmt.Sprintf("/curl/%s/%s", secureHost, securePort))
			json.Unmarshal([]byte(resp), &noraCurlResponse)
			return noraCurlResponse.ReturnCode
		}
		firstCurlError := curlResponse()
		Expect(firstCurlError).ShouldNot(Equal(0))

		// apply security group
		rules := fmt.Sprintf(`[{"destination":"%s","ports":"%s","protocol":"tcp"}]`, secureHost, securePort)

		file, _ := ioutil.TempFile(os.TempDir(), "DATS-sg-rules")
		defer os.Remove(file.Name())
		file.WriteString(rules)
		file.Close()

		rulesPath := file.Name()
		securityGroupName := fmt.Sprintf("DATS-SG-%s", generator.PrefixedRandomName(config.GetNamePrefix(), "SECURITY-GROUP"))

		AsUser(environment.AdminUserContext(), time.Minute, func() {
			Expect(cf.Cf("create-security-group", securityGroupName, rulesPath).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			Expect(cf.Cf("bind-security-group",
				securityGroupName,
				environment.RegularUserContext().Org,
				environment.RegularUserContext().Space).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		})
		defer func() {
			AsUser(environment.AdminUserContext(), time.Minute, func() {
				Expect(cf.Cf("delete-security-group", securityGroupName, "-f").Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			})
		}()

		Expect(cf.Cf("restart", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))

		// test app egress rules
		Eventually(curlResponse).Should(Equal(0))

		// unapply security group
		AsUser(environment.AdminUserContext(), time.Minute, func() {
			Expect(cf.Cf("unbind-security-group",
				securityGroupName, environment.RegularUserContext().Org,
				environment.RegularUserContext().Space).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		})

		By("restarting it - without security group")
		Expect(cf.Cf("restart", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
		Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))

		// test app egress rules
		Eventually(curlResponse).Should(Equal(firstCurlError))
	})

	Context("when an icmp rule is applied", func() {
		var (
			icmpRuleFile      string
			securityGroupName string
		)

		BeforeEach(func() {
			icmpRule := `[{"code": 1,"destination":"0.0.0.0/0","protocol":"icmp","type":0}]`
			securityGroupName = fmt.Sprintf("DATS-SG-%s", generator.PrefixedRandomName(config.GetNamePrefix(), "SECURITY-GROUP"))

			file, err := ioutil.TempFile("", securityGroupName)
			Expect(err).ToNot(HaveOccurred())
			_, err = file.WriteString(icmpRule)
			Expect(err).ToNot(HaveOccurred())
			Expect(file.Close()).To(Succeed())

			icmpRuleFile = file.Name()

			AsUser(environment.AdminUserContext(), 2*time.Minute, func() {
				Expect(cf.Cf("create-security-group", securityGroupName, icmpRuleFile).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
				Expect(cf.Cf("bind-security-group",
					securityGroupName,
					environment.RegularUserContext().Org,
					environment.RegularUserContext().Space).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			})
		})

		AfterEach(func() {
			Expect(os.Remove(icmpRuleFile)).To(Succeed())
			AsUser(environment.AdminUserContext(), 2*time.Minute, func() {
				Expect(cf.Cf("unbind-security-group",
					securityGroupName, environment.RegularUserContext().Org,
					environment.RegularUserContext().Space).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
				Expect(cf.Cf("delete-security-group", securityGroupName, "-f").Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			})
		})

		It("ignores the rule and can push an app", func() {
			By("pushing it", func() {
				Expect(pushNora(appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			})

			By("staging and running it on Diego", func() {
				enableDiego(appName)
				Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(gexec.Exit(0))
			})

			By("verifying it's up", func() {
				Eventually(helpers.CurlingAppRoot(config, appName)).Should(ContainSubstring("hello i am nora"))
			})
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"encoding/json"
	"os"
	"io/ioutil"
	"reflect"
)

type Named interface {
	Name() string
}

type GameObject struct {
	name string
	components []*Componenter
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: %s <game JSON file>", os.Args[0])
		os.Exit(1)
	}

	filename := os.Args[1]
	var filecontents []byte
	filecontents, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	var jsonData map[string]interface{}
	if err := json.Unmarshal(filecontents, &jsonData); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	var objects = []GameObject{}
	objs := jsonData["objects"].([]interface{})
	for _, jobj := range objs {
		objmap := jobj.(map[string]interface{})
		var obj GameObject
		obj.name = objmap["name"].(string)
		components := objmap["components"].([]interface{})
		for _, jcomp := range components {
			comp := jcomp.(map[string]interface{})
			typeName := comp["type"]
			compInst := ComponentNameMap[typeName.(string)]()
			obj.components = append(obj.components, &compInst)

			for jvaluename, jvaluedata := range comp["values"].(map[string]interface{}) {
				switch typ := jvaluedata.(type) {
				case string:
					reflect.ValueOf(compInst).Elem().FieldByName(jvaluename).SetString(jvaluedata.(string))
				case int:
					reflect.ValueOf(compInst).Elem().FieldByName(jvaluename).SetInt(jvaluedata.(int64))
				case float64:
					reflect.ValueOf(compInst).Elem().FieldByName(jvaluename).SetInt(int64(jvaluedata.(float64)))
				default:
					fmt.Println("Unknown type", typ)
				}
			}
		}
		objects = append(objects, obj)
	}

	runGame(objects)
}

func runGame(objects []GameObject) {
	for _, obj := range objects {
		for _, comp := range obj.components {
			(*comp).Start()
		}
	}
}
<commit_msg>do not use float64 when decoding JSON<commit_after>package main

import (
	"fmt"
	"encoding/json"
	"os"
	"reflect"
)

type Named interface {
	Name() string
}

type GameObject struct {
	name string
	components []*Componenter
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: %s <game JSON file>", os.Args[0])
		os.Exit(1)
	}

	filename := os.Args[1]
	f, err := os.Open(filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	var jsonData map[string]interface{}
	dec := json.NewDecoder(f)
	dec.UseNumber()
	if err := dec.Decode(&jsonData); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	var objects = []GameObject{}
	objs := jsonData["objects"].([]interface{})
	for _, jobj := range objs {
		objmap := jobj.(map[string]interface{})
		var obj GameObject
		obj.name = objmap["name"].(string)
		components := objmap["components"].([]interface{})
		for _, jcomp := range components {
			comp := jcomp.(map[string]interface{})
			typeName := comp["type"]
			compInst := ComponentNameMap[typeName.(string)]()
			obj.components = append(obj.components, &compInst)

			for jvaluename, jvaluedata := range comp["values"].(map[string]interface{}) {
				compValue := reflect.ValueOf(compInst).Elem()
				fieldValue := compValue.FieldByName(jvaluename)
				typ := fieldValue.Kind()
				switch typ {
				case reflect.Bool:
					fieldValue.SetBool(jvaluedata.(bool))
				case reflect.Int:
					v, err := jvaluedata.(json.Number).Int64()
					if err != nil {
						fmt.Println("Error on field %s: %s", jvaluename, err)
					}
					fieldValue.SetInt(v)
				case reflect.String:
					fieldValue.SetString(jvaluedata.(string))
				default:
					fmt.Println("Unknown type", typ)
				}
			}
		}
		objects = append(objects, obj)
	}

	runGame(objects)
}

func runGame(objects []GameObject) {
	for _, obj := range objects {
		for _, comp := range obj.components {
			(*comp).Start()
		}
	}
}
<|endoftext|>


<commit_before>package etcd2topo

import (
	"fmt"
	"path"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/vt/topo"
)

// Watch is part of the topo.Backend interface
func (s *Server) Watch(ctx context.Context, cell, filePath string) (*topo.WatchData, <-chan *topo.WatchData, topo.CancelFunc) {
	c, err := s.clientForCell(ctx, cell)
	if err != nil {
		return &topo.WatchData{Err: fmt.Errorf("Watch cannot get cell: %v", err)}, nil, nil
	}
	nodePath := path.Join(c.root, filePath)

	// Get the initial version of the file
	initial, err := s.global.cli.Get(ctx, nodePath)
	if err != nil {
		// Generic error.
		return &topo.WatchData{Err: convertError(err)}, nil, nil
	}
	if len(initial.Kvs) != 1 {
		// Node doesn't exist.
		return &topo.WatchData{Err: topo.ErrNoNode}, nil, nil
	}
	wd := &topo.WatchData{
		Contents: initial.Kvs[0].Value,
		Version:  EtcdVersion(initial.Kvs[0].ModRevision),
	}

	// Create a context, will be used to cancel the watch.
	watchCtx, watchCancel := context.WithCancel(context.Background())

	// Create the Watcher.
	watcher := s.global.cli.Watch(watchCtx, nodePath, clientv3.WithRev(initial.Kvs[0].ModRevision))
	if watcher == nil {
		return &topo.WatchData{Err: fmt.Errorf("Watch failed")}, nil, nil
	}

	// Create the notifications channel, send updates to it.
	notifications := make(chan *topo.WatchData, 10)
	go func() {
		defer close(notifications)

		for {
			select {
			case <-watchCtx.Done():
				// This includes context cancelation errors.
				notifications <- &topo.WatchData{
					Err: convertError(watchCtx.Err()),
				}
				return
			case wresp := <-watcher:
				if wresp.Canceled {
					// Final notification.
					notifications <- &topo.WatchData{
						Err: convertError(wresp.Err()),
					}
					return
				}

				for _, ev := range wresp.Events {
					switch ev.Type {
					case mvccpb.PUT:
						notifications <- &topo.WatchData{
							Contents: ev.Kv.Value,
							Version:  EtcdVersion(ev.Kv.Version),
						}
					case mvccpb.DELETE:
						// Node is gone, send a final notice.
						notifications <- &topo.WatchData{
							Err: topo.ErrNoNode,
						}
						return
					default:
						notifications <- &topo.WatchData{
							Err: fmt.Errorf("unexpected event received: %v", ev),
						}
						return
					}
				}
			}
		}
	}()

	return wd, notifications, topo.CancelFunc(watchCancel)
}
<commit_msg>etcd2topo watch fix.<commit_after>package etcd2topo

import (
	"fmt"
	"path"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/vt/topo"
)

// Watch is part of the topo.Backend interface
func (s *Server) Watch(ctx context.Context, cell, filePath string) (*topo.WatchData, <-chan *topo.WatchData, topo.CancelFunc) {
	c, err := s.clientForCell(ctx, cell)
	if err != nil {
		return &topo.WatchData{Err: fmt.Errorf("Watch cannot get cell: %v", err)}, nil, nil
	}
	nodePath := path.Join(c.root, filePath)

	// Get the initial version of the file
	initial, err := s.global.cli.Get(ctx, nodePath)
	if err != nil {
		// Generic error.
		return &topo.WatchData{Err: convertError(err)}, nil, nil
	}
	if len(initial.Kvs) != 1 {
		// Node doesn't exist.
		return &topo.WatchData{Err: topo.ErrNoNode}, nil, nil
	}
	wd := &topo.WatchData{
		Contents: initial.Kvs[0].Value,
		Version:  EtcdVersion(initial.Kvs[0].ModRevision),
	}

	// Create a context, will be used to cancel the watch.
	watchCtx, watchCancel := context.WithCancel(context.Background())

	// Create the Watcher.  We start watching from the response we
	// got, not from the file original version, as the server may
	// not have that much history.
	watcher := s.global.cli.Watch(watchCtx, nodePath, clientv3.WithRev(initial.Header.Revision))
	if watcher == nil {
		return &topo.WatchData{Err: fmt.Errorf("Watch failed")}, nil, nil
	}

	// Create the notifications channel, send updates to it.
	notifications := make(chan *topo.WatchData, 10)
	go func() {
		defer close(notifications)

		for {
			select {
			case <-watchCtx.Done():
				// This includes context cancelation errors.
				notifications <- &topo.WatchData{
					Err: convertError(watchCtx.Err()),
				}
				return
			case wresp := <-watcher:
				if wresp.Canceled {
					// Final notification.
					notifications <- &topo.WatchData{
						Err: convertError(wresp.Err()),
					}
					return
				}

				for _, ev := range wresp.Events {
					switch ev.Type {
					case mvccpb.PUT:
						notifications <- &topo.WatchData{
							Contents: ev.Kv.Value,
							Version:  EtcdVersion(ev.Kv.Version),
						}
					case mvccpb.DELETE:
						// Node is gone, send a final notice.
						notifications <- &topo.WatchData{
							Err: topo.ErrNoNode,
						}
						return
					default:
						notifications <- &topo.WatchData{
							Err: fmt.Errorf("unexpected event received: %v", ev),
						}
						return
					}
				}
			}
		}
	}()

	return wd, notifications, topo.CancelFunc(watchCancel)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

package zk2topo

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/log"
)

const (
	// maxAttempts is how many times we retry queries.  At 2 for
	// now, so if a query fails because the session expired, we
	// just try to reconnect once and go on.
	maxAttempts = 2

	// PermDirectory are default permissions for a node.
	PermDirectory = zk.PermAdmin | zk.PermCreate | zk.PermDelete | zk.PermRead | zk.PermWrite

	// PermFile allows a zk node to emulate file behavior by
	// disallowing child nodes.
	PermFile = zk.PermAdmin | zk.PermRead | zk.PermWrite
)

var (
	maxConcurrency = flag.Int("topo_zk_max_concurrency", 64, "maximum number of pending requests to send to a Zookeeper server.")

	baseTimeout = flag.Duration("topo_zk_base_timeout", 30*time.Second, "zk base timeout (see zk.Connect)")

	certPath = flag.String("topo_zk_tls_cert", "", "the cert to use to connect to the zk topo server, requires topo_zk_tls_key, enables TLS")
	keyPath  = flag.String("topo_zk_tls_key", "", "the key to use to connect to the zk topo server, enables TLS")
	caPath   = flag.String("topo_zk_tls_ca", "", "the server ca to use to validate servers when connecting to the zk topo server")
	authFile = flag.String("topo_zk_auth_file", "", "auth to use when connecting to the zk topo server, file contents should be <scheme>:<auth>, e.g., digest:user:pass")
)

// Time returns a time.Time from a ZK int64 milliseconds since Epoch time.
func Time(i int64) time.Time {
	return time.Unix(i/1000, i%1000*1000000)
}

// ZkTime returns a ZK time (int64) from a time.Time
func ZkTime(t time.Time) int64 {
	return t.Unix()*1000 + int64(t.Nanosecond()/1000000)
}

// ZkConn is a wrapper class on top of a zk.Conn.
// It will do a few things for us:
// - add the context parameter. However, we do not enforce its deadlines
//   necessarily.
// - enforce a max concurrency of access to Zookeeper. We just don't
//   want to make too many calls concurrently, to not take too many resources.
// - retry some calls to Zookeeper. If we were disconnected from the
//   server, we want to try connecting again before failing.
type ZkConn struct {
	// addr is set at construction time, and immutable.
	addr string

	// sem protects concurrent calls to Zookeeper.
	sem *sync2.Semaphore

	// mu protects the following fields.
	mu   sync.Mutex
	conn *zk.Conn
}

// Connect to the Zookeeper servers specified in addr
// addr can be a comma separated list of servers and each server can be a DNS entry with multiple values.
// Connects to the endpoints in a randomized order to avoid hot spots.
func Connect(addr string) *ZkConn {
	return &ZkConn{
		addr: addr,
		sem:  sync2.NewSemaphore(*maxConcurrency, 0),
	}
}

// Get is part of the Conn interface.
func (c *ZkConn) Get(ctx context.Context, path string) (data []byte, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		data, stat, err = conn.Get(path)
		return err
	})
	return
}

func (c *ZkConn) GetW(ctx context.Context, path string) (data []byte, stat *zk.Stat, watch <-chan zk.Event, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		data, stat, watch, err = conn.GetW(path)
		return err
	})
	return
}

func (c *ZkConn) Children(ctx context.Context, path string) (children []string, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		children, stat, err = conn.Children(path)
		return err
	})
	return
}

func (c *ZkConn) ChildrenW(ctx context.Context, path string) (children []string, stat *zk.Stat, watch <-chan zk.Event, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		children, stat, watch, err = conn.ChildrenW(path)
		return err
	})
	return
}

func (c *ZkConn) Exists(ctx context.Context, path string) (exists bool, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		exists, stat, err = conn.Exists(path)
		return err
	})
	return
}

func (c *ZkConn) ExistsW(ctx context.Context, path string) (exists bool, stat *zk.Stat, watch <-chan zk.Event, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		exists, stat, watch, err = conn.ExistsW(path)
		return err
	})
	return
}

// Create is part of the Conn interface.
func (c *ZkConn) Create(ctx context.Context, path string, value []byte, flags int32, aclv []zk.ACL) (pathCreated string, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		pathCreated, err = conn.Create(path, value, flags, aclv)
		return err
	})
	return
}

// Set is part of the Conn interface.
func (c *ZkConn) Set(ctx context.Context, path string, value []byte, version int32) (stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		stat, err = conn.Set(path, value, version)
		return err
	})
	return
}

// Delete is part of the Conn interface.
func (c *ZkConn) Delete(ctx context.Context, path string, version int32) error {
	return c.withRetry(ctx, func(conn *zk.Conn) error {
		return conn.Delete(path, version)
	})
}

func (c *ZkConn) GetACL(ctx context.Context, path string) (aclv []zk.ACL, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		aclv, stat, err = conn.GetACL(path)
		return err
	})
	return
}

func (c *ZkConn) SetACL(ctx context.Context, path string, aclv []zk.ACL, version int32) error {
	return c.withRetry(ctx, func(conn *zk.Conn) error {
		_, err := conn.SetACL(path, aclv, version)
		return err
	})
}

func (c *ZkConn) AddAuth(ctx context.Context, scheme string, auth []byte) error {
	return c.withRetry(ctx, func(conn *zk.Conn) error {
		err := conn.AddAuth(scheme, auth)
		return err
	})
}

// Close is part of the Conn interface.
func (c *ZkConn) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.conn != nil {
		c.conn.Close()
	}
	return nil
}

// withRetry encapsulates the retry logic and concurrent access to
// Zookeeper.
//
// Some errors are not handled gracefully by the Zookeeper client. This is
// sort of odd, but in general it doesn't affect the kind of code you
// need to have a truly reliable client.
//
// However, it can manifest itself as an annoying transient error that
// is likely avoidable when trying simple operations like Get.
// To that end, we retry when possible to minimize annoyance at
// higher levels.
//
// https://issues.apache.org/jira/browse/ZOOKEEPER-22
func (c *ZkConn) withRetry(ctx context.Context, action func(conn *zk.Conn) error) (err error) {

	// Handle concurrent access to a Zookeeper server here.
	c.sem.Acquire()
	defer c.sem.Release()

	for i := 0; i < maxAttempts; i++ {
		if i > 0 {
			// Add a bit of backoff time before retrying:
			// 1 second base + up to 5 seconds.
			time.Sleep(1*time.Second + time.Duration(rand.Int63n(5e9)))
		}

		// Get the current connection, or connect.
		var conn *zk.Conn
		conn, err = c.getConn(ctx)
		if err != nil {
			// We can't connect, try again.
			continue
		}

		// Execute the action.
		err = action(conn)
		if err != zk.ErrConnectionClosed {
			// It worked, or it failed for another reason
			// than connection related.
			return
		}

		// We got an error, because the connection was closed.
		// Let's clear up our errored connection and try again.
		c.mu.Lock()
		if c.conn == conn {
			c.conn = nil
		}
		c.mu.Unlock()
	}
	return
}

// getConn returns the connection in a thread safe way. It will try to connect
// if not connected yet.
func (c *ZkConn) getConn(ctx context.Context) (*zk.Conn, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		conn, events, err := dialZk(ctx, c.addr)
		if err != nil {
			return nil, err
		}
		c.conn = conn
		go c.handleSessionEvents(conn, events)
		c.maybeAddAuth(ctx)
	}
	return c.conn, nil
}

// maybeAddAuth calls AddAuth if the `-topo_zk_auth_file` flag was specified
func (c *ZkConn) maybeAddAuth(ctx context.Context) {
	if *authFile == "" {
		return
	}
	authInfoBytes, err := ioutil.ReadFile(*authFile)
	if err != nil {
		log.Errorf("failed to read topo_zk_auth_file: %v", err)
		return
	}
	authInfo := string(authInfoBytes)
	authInfoParts := strings.SplitN(authInfo, ":", 2)
	if len(authInfoParts) != 2 {
		log.Errorf("failed to parse topo_zk_auth_file contents, expected format <scheme>:<auth> but saw: %s", authInfo)
		return
	}
	err = c.conn.AddAuth(authInfoParts[0], []byte(authInfoParts[1]))
	if err != nil {
		log.Errorf("failed to add auth from topo_zk_auth_file: %v", err)
		return
	}
}

// handleSessionEvents is processing events from the session channel.
// When it detects that the connection is not working any more, it
// clears out the connection record.
func (c *ZkConn) handleSessionEvents(conn *zk.Conn, session <-chan zk.Event) {
	for event := range session {
		closeRequired := false

		switch event.State {
		case zk.StateExpired, zk.StateConnecting:
			closeRequired = true
			fallthrough
		case zk.StateDisconnected:
			c.mu.Lock()
			if c.conn == conn {
				// The ZkConn still references this
				// connection, let's nil it.
				c.conn = nil
			}
			c.mu.Unlock()
			if closeRequired {
				conn.Close()
			}
			log.Infof("zk conn: session for addr %v ended: %v", c.addr, event)
			return
		}
		log.Infof("zk conn: session for addr %v event: %v", c.addr, event)
	}
}

// dialZk dials the server, and waits until connection.
func dialZk(ctx context.Context, addr string) (*zk.Conn, <-chan zk.Event, error) {
	servers := strings.Split(addr, ",")
	options := zk.WithDialer(net.DialTimeout)
	// If TLS is enabled use a TLS enabled dialer option
	if *certPath != "" && *keyPath != "" {
		if strings.Contains(addr, ",") {
			log.Fatalf("This TLS zk code requires that the all the zk servers validate to a single server name.")
		}

		serverName := strings.Split(addr, ":")[0]

		log.Infof("Using TLS ZK, connecting to %v server name %v", addr, serverName)
		cert, err := tls.LoadX509KeyPair(*certPath, *keyPath)
		if err != nil {
			log.Fatalf("Unable to load cert %v and key %v, err %v", *certPath, *keyPath, err)
		}

		clientCACert, err := ioutil.ReadFile(*caPath)
		if err != nil {
			log.Fatalf("Unable to open ca cert %v, err %v", *caPath, err)
		}

		clientCertPool := x509.NewCertPool()
		clientCertPool.AppendCertsFromPEM(clientCACert)

		tlsConfig := &tls.Config{
			Certificates: []tls.Certificate{cert},
			RootCAs:      clientCertPool,
			ServerName:   serverName,
		}

		tlsConfig.BuildNameToCertificate()

		options = zk.WithDialer(func(network, address string, timeout time.Duration) (net.Conn, error) {
			d := net.Dialer{Timeout: timeout}

			return tls.DialWithDialer(&d, network, address, tlsConfig)
		})
	}

	// zk.Connect automatically shuffles the servers
	zconn, session, err := zk.Connect(servers, *baseTimeout, options)
	if err != nil {
		return nil, nil, err
	}

	// Wait for connection, skipping transition states.
	for {
		select {
		case <-ctx.Done():
			zconn.Close()
			return nil, nil, ctx.Err()
		case event := <-session:
			switch event.State {
			case zk.StateConnected:
				// success
				return zconn, session, nil

			case zk.StateAuthFailed:
				// fast fail this one
				zconn.Close()
				return nil, nil, fmt.Errorf("zk connect failed: StateAuthFailed")
			}
		}
	}
}
<commit_msg>Add doc comments back in<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

package zk2topo

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/log"
)

const (
	// maxAttempts is how many times we retry queries.  At 2 for
	// now, so if a query fails because the session expired, we
	// just try to reconnect once and go on.
	maxAttempts = 2

	// PermDirectory are default permissions for a node.
	PermDirectory = zk.PermAdmin | zk.PermCreate | zk.PermDelete | zk.PermRead | zk.PermWrite

	// PermFile allows a zk node to emulate file behavior by
	// disallowing child nodes.
	PermFile = zk.PermAdmin | zk.PermRead | zk.PermWrite
)

var (
	maxConcurrency = flag.Int("topo_zk_max_concurrency", 64, "maximum number of pending requests to send to a Zookeeper server.")

	baseTimeout = flag.Duration("topo_zk_base_timeout", 30*time.Second, "zk base timeout (see zk.Connect)")

	certPath = flag.String("topo_zk_tls_cert", "", "the cert to use to connect to the zk topo server, requires topo_zk_tls_key, enables TLS")
	keyPath  = flag.String("topo_zk_tls_key", "", "the key to use to connect to the zk topo server, enables TLS")
	caPath   = flag.String("topo_zk_tls_ca", "", "the server ca to use to validate servers when connecting to the zk topo server")
	authFile = flag.String("topo_zk_auth_file", "", "auth to use when connecting to the zk topo server, file contents should be <scheme>:<auth>, e.g., digest:user:pass")
)

// Time returns a time.Time from a ZK int64 milliseconds since Epoch time.
func Time(i int64) time.Time {
	return time.Unix(i/1000, i%1000*1000000)
}

// ZkTime returns a ZK time (int64) from a time.Time
func ZkTime(t time.Time) int64 {
	return t.Unix()*1000 + int64(t.Nanosecond()/1000000)
}

// ZkConn is a wrapper class on top of a zk.Conn.
// It will do a few things for us:
// - add the context parameter. However, we do not enforce its deadlines
//   necessarily.
// - enforce a max concurrency of access to Zookeeper. We just don't
//   want to make too many calls concurrently, to not take too many resources.
// - retry some calls to Zookeeper. If we were disconnected from the
//   server, we want to try connecting again before failing.
type ZkConn struct {
	// addr is set at construction time, and immutable.
	addr string

	// sem protects concurrent calls to Zookeeper.
	sem *sync2.Semaphore

	// mu protects the following fields.
	mu   sync.Mutex
	conn *zk.Conn
}

// Connect to the Zookeeper servers specified in addr
// addr can be a comma separated list of servers and each server can be a DNS entry with multiple values.
// Connects to the endpoints in a randomized order to avoid hot spots.
func Connect(addr string) *ZkConn {
	return &ZkConn{
		addr: addr,
		sem:  sync2.NewSemaphore(*maxConcurrency, 0),
	}
}

// Get is part of the Conn interface.
func (c *ZkConn) Get(ctx context.Context, path string) (data []byte, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		data, stat, err = conn.Get(path)
		return err
	})
	return
}

// GetW is part of the Conn interface.
func (c *ZkConn) GetW(ctx context.Context, path string) (data []byte, stat *zk.Stat, watch <-chan zk.Event, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		data, stat, watch, err = conn.GetW(path)
		return err
	})
	return
}

// Children is part of the Conn interface.
func (c *ZkConn) Children(ctx context.Context, path string) (children []string, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		children, stat, err = conn.Children(path)
		return err
	})
	return
}

// ChildrenW is part of the Conn interface.
func (c *ZkConn) ChildrenW(ctx context.Context, path string) (children []string, stat *zk.Stat, watch <-chan zk.Event, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		children, stat, watch, err = conn.ChildrenW(path)
		return err
	})
	return
}

// Exists is part of the Conn interface.
func (c *ZkConn) Exists(ctx context.Context, path string) (exists bool, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		exists, stat, err = conn.Exists(path)
		return err
	})
	return
}

// ExistsW is part of the Conn interface.
func (c *ZkConn) ExistsW(ctx context.Context, path string) (exists bool, stat *zk.Stat, watch <-chan zk.Event, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		exists, stat, watch, err = conn.ExistsW(path)
		return err
	})
	return
}

// Create is part of the Conn interface.
func (c *ZkConn) Create(ctx context.Context, path string, value []byte, flags int32, aclv []zk.ACL) (pathCreated string, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		pathCreated, err = conn.Create(path, value, flags, aclv)
		return err
	})
	return
}

// Set is part of the Conn interface.
func (c *ZkConn) Set(ctx context.Context, path string, value []byte, version int32) (stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		stat, err = conn.Set(path, value, version)
		return err
	})
	return
}

// Delete is part of the Conn interface.
func (c *ZkConn) Delete(ctx context.Context, path string, version int32) error {
	return c.withRetry(ctx, func(conn *zk.Conn) error {
		return conn.Delete(path, version)
	})
}

// GetACL is part of the Conn interface.
func (c *ZkConn) GetACL(ctx context.Context, path string) (aclv []zk.ACL, stat *zk.Stat, err error) {
	err = c.withRetry(ctx, func(conn *zk.Conn) error {
		aclv, stat, err = conn.GetACL(path)
		return err
	})
	return
}

// SetACL is part of the Conn interface.
func (c *ZkConn) SetACL(ctx context.Context, path string, aclv []zk.ACL, version int32) error {
	return c.withRetry(ctx, func(conn *zk.Conn) error {
		_, err := conn.SetACL(path, aclv, version)
		return err
	})
}

// AddAuth is part of the Conn interface.
func (c *ZkConn) AddAuth(ctx context.Context, scheme string, auth []byte) error {
	return c.withRetry(ctx, func(conn *zk.Conn) error {
		err := conn.AddAuth(scheme, auth)
		return err
	})
}

// Close is part of the Conn interface.
func (c *ZkConn) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.conn != nil {
		c.conn.Close()
	}
	return nil
}

// withRetry encapsulates the retry logic and concurrent access to
// Zookeeper.
//
// Some errors are not handled gracefully by the Zookeeper client. This is
// sort of odd, but in general it doesn't affect the kind of code you
// need to have a truly reliable client.
//
// However, it can manifest itself as an annoying transient error that
// is likely avoidable when trying simple operations like Get.
// To that end, we retry when possible to minimize annoyance at
// higher levels.
//
// https://issues.apache.org/jira/browse/ZOOKEEPER-22
func (c *ZkConn) withRetry(ctx context.Context, action func(conn *zk.Conn) error) (err error) {

	// Handle concurrent access to a Zookeeper server here.
	c.sem.Acquire()
	defer c.sem.Release()

	for i := 0; i < maxAttempts; i++ {
		if i > 0 {
			// Add a bit of backoff time before retrying:
			// 1 second base + up to 5 seconds.
			time.Sleep(1*time.Second + time.Duration(rand.Int63n(5e9)))
		}

		// Get the current connection, or connect.
		var conn *zk.Conn
		conn, err = c.getConn(ctx)
		if err != nil {
			// We can't connect, try again.
			continue
		}

		// Execute the action.
		err = action(conn)
		if err != zk.ErrConnectionClosed {
			// It worked, or it failed for another reason
			// than connection related.
			return
		}

		// We got an error, because the connection was closed.
		// Let's clear up our errored connection and try again.
		c.mu.Lock()
		if c.conn == conn {
			c.conn = nil
		}
		c.mu.Unlock()
	}
	return
}

// getConn returns the connection in a thread safe way. It will try to connect
// if not connected yet.
func (c *ZkConn) getConn(ctx context.Context) (*zk.Conn, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		conn, events, err := dialZk(ctx, c.addr)
		if err != nil {
			return nil, err
		}
		c.conn = conn
		go c.handleSessionEvents(conn, events)
		c.maybeAddAuth(ctx)
	}
	return c.conn, nil
}

// maybeAddAuth calls AddAuth if the `-topo_zk_auth_file` flag was specified
func (c *ZkConn) maybeAddAuth(ctx context.Context) {
	if *authFile == "" {
		return
	}
	authInfoBytes, err := ioutil.ReadFile(*authFile)
	if err != nil {
		log.Errorf("failed to read topo_zk_auth_file: %v", err)
		return
	}
	authInfo := string(authInfoBytes)
	authInfoParts := strings.SplitN(authInfo, ":", 2)
	if len(authInfoParts) != 2 {
		log.Errorf("failed to parse topo_zk_auth_file contents, expected format <scheme>:<auth> but saw: %s", authInfo)
		return
	}
	err = c.conn.AddAuth(authInfoParts[0], []byte(authInfoParts[1]))
	if err != nil {
		log.Errorf("failed to add auth from topo_zk_auth_file: %v", err)
		return
	}
}

// handleSessionEvents is processing events from the session channel.
// When it detects that the connection is not working any more, it
// clears out the connection record.
func (c *ZkConn) handleSessionEvents(conn *zk.Conn, session <-chan zk.Event) {
	for event := range session {
		closeRequired := false

		switch event.State {
		case zk.StateExpired, zk.StateConnecting:
			closeRequired = true
			fallthrough
		case zk.StateDisconnected:
			c.mu.Lock()
			if c.conn == conn {
				// The ZkConn still references this
				// connection, let's nil it.
				c.conn = nil
			}
			c.mu.Unlock()
			if closeRequired {
				conn.Close()
			}
			log.Infof("zk conn: session for addr %v ended: %v", c.addr, event)
			return
		}
		log.Infof("zk conn: session for addr %v event: %v", c.addr, event)
	}
}

// dialZk dials the server, and waits until connection.
func dialZk(ctx context.Context, addr string) (*zk.Conn, <-chan zk.Event, error) {
	servers := strings.Split(addr, ",")
	options := zk.WithDialer(net.DialTimeout)
	// If TLS is enabled use a TLS enabled dialer option
	if *certPath != "" && *keyPath != "" {
		if strings.Contains(addr, ",") {
			log.Fatalf("This TLS zk code requires that the all the zk servers validate to a single server name.")
		}

		serverName := strings.Split(addr, ":")[0]

		log.Infof("Using TLS ZK, connecting to %v server name %v", addr, serverName)
		cert, err := tls.LoadX509KeyPair(*certPath, *keyPath)
		if err != nil {
			log.Fatalf("Unable to load cert %v and key %v, err %v", *certPath, *keyPath, err)
		}

		clientCACert, err := ioutil.ReadFile(*caPath)
		if err != nil {
			log.Fatalf("Unable to open ca cert %v, err %v", *caPath, err)
		}

		clientCertPool := x509.NewCertPool()
		clientCertPool.AppendCertsFromPEM(clientCACert)

		tlsConfig := &tls.Config{
			Certificates: []tls.Certificate{cert},
			RootCAs:      clientCertPool,
			ServerName:   serverName,
		}

		tlsConfig.BuildNameToCertificate()

		options = zk.WithDialer(func(network, address string, timeout time.Duration) (net.Conn, error) {
			d := net.Dialer{Timeout: timeout}

			return tls.DialWithDialer(&d, network, address, tlsConfig)
		})
	}

	// zk.Connect automatically shuffles the servers
	zconn, session, err := zk.Connect(servers, *baseTimeout, options)
	if err != nil {
		return nil, nil, err
	}

	// Wait for connection, skipping transition states.
	for {
		select {
		case <-ctx.Done():
			zconn.Close()
			return nil, nil, ctx.Err()
		case event := <-session:
			switch event.State {
			case zk.StateConnected:
				// success
				return zconn, session, nil

			case zk.StateAuthFailed:
				// fast fail this one
				zconn.Close()
				return nil, nil, fmt.Errorf("zk connect failed: StateAuthFailed")
			}
		}
	}
}
<|endoftext|>


<commit_before>/**
Sharder allows to balance requests among configured nodes
according to some request ID. Each request with the same ID
will be routed to the same node.

Author: Aleksey Morarash <aleksey.morarash@gmail.com>
Since: 13 Oct 2017
Copyright: 2017, Aleksey Morarash <aleksey.morarash@gmail.com>
*/

package tcpcall

import (
	"crypto/md5"
	"encoding/binary"
	"errors"
	"sync"
	"time"
)

type Sharder struct {
	// Configuration used to create Sharder instance
	config SharderConfig
	// List of servers actually used
	nodes []string
	// Map server hostname to connection pool
	conns map[string]*Pool
	// Controls access to the map of connection pools
	mu *sync.RWMutex
}

// Sharder configuration
type SharderConfig struct {
	// Callback function used to get list of servers
	NodesGetter func() ([]string, error)
	// How often NodesGetter will be called
	ReconfigPeriod time.Duration
	// How many connections will be established to each server
	ConnsPerNode int
	// Request send max retry count. Negative value means count of
	// currently established connections to one server,
	// 0 means no retries will performed at all.
	MaxRequestRetries int
}

// Create new Sharder instance with given configuration.
func NewSharder(config SharderConfig) *Sharder {
	sharder := &Sharder{
		config: config,
		nodes:  []string{},
		conns:  map[string]*Pool{},
		mu:     &sync.RWMutex{},
	}
	go sharder.reconfigLoop()
	return sharder
}

// Create default configuration for Sharder.
func NewSharderConfig() SharderConfig {
	return SharderConfig{
		NodesGetter: func() ([]string, error) {
			return []string{}, nil
		},
		ReconfigPeriod:    10 * time.Second,
		ConnsPerNode:      5,
		MaxRequestRetries: -1,
	}
}

// Send request to one of configured servers. ID will not
// be sent, but used only to balance request between servers.
func (s *Sharder) Req(id, body []byte, timeout time.Duration) (reply []byte, err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	nodesCount := len(s.nodes)
	if nodesCount == 0 {
		return nil, errors.New("not configured")
	}
	i := shard(id, nodesCount)
	return s.conns[s.nodes[i]].Req(body, timeout)
}

// Sharding function.
func shard(id []byte, n int) int {
	hash := md5.Sum(id)
	return int(binary.BigEndian.Uint64(hash[:8]) % uint64(n))
}

// Goroutine.
// Reconfigure clients pool on the fly.
func (s *Sharder) reconfigLoop() {
	for {
		if nodes, err := s.config.NodesGetter(); err == nil {
			s.setNodes(nodes)
			time.Sleep(s.config.ReconfigPeriod)
		} else {
			time.Sleep(s.config.ReconfigPeriod / 10)
		}
	}
}

// Apply new node list to the clients pool.
func (s *Sharder) setNodes(newNodes []string) {
	newNodes = uniq(newNodes)
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, n := range newNodes {
		if _, ok := s.conns[n]; !ok {
			poolCfg := NewPoolConf()
			peers := make([]string, s.config.ConnsPerNode)
			for i := 0; i < s.config.ConnsPerNode; i++ {
				peers[i] = n
			}
			poolCfg.Peers = peers
			poolCfg.MaxRequestRetries = s.config.MaxRequestRetries
			s.conns[n] = NewPool(poolCfg)
		}
	}
	for k, p := range s.conns {
		if !inList(k, newNodes) {
			p.Close()
			delete(s.conns, k)
		}
	}
	s.nodes = newNodes
}

// Leave only unique elements from a string array.
func uniq(a []string) []string {
	r := []string{}
	for _, v := range a {
		if !inList(v, r) {
			r = append(r, v)
		}
	}
	return r
}

// Return truth if array a contains string s at least once.
func inList(s string, a []string) bool {
	for _, e := range a {
		if s == e {
			return true
		}
	}
	return false
}
<commit_msg>golang: refactor SharderConfig to SharderConf<commit_after>/**
Sharder allows to balance requests among configured nodes
according to some request ID. Each request with the same ID
will be routed to the same node.

Author: Aleksey Morarash <aleksey.morarash@gmail.com>
Since: 13 Oct 2017
Copyright: 2017, Aleksey Morarash <aleksey.morarash@gmail.com>
*/

package tcpcall

import (
	"crypto/md5"
	"encoding/binary"
	"errors"
	"sync"
	"time"
)

type Sharder struct {
	// Configuration used to create Sharder instance
	config SharderConf
	// List of servers actually used
	nodes []string
	// Map server hostname to connection pool
	conns map[string]*Pool
	// Controls access to the map of connection pools
	mu *sync.RWMutex
}

// Sharder configuration
type SharderConf struct {
	// Callback function used to get list of servers
	NodesGetter func() ([]string, error)
	// How often NodesGetter will be called
	ReconfigPeriod time.Duration
	// How many connections will be established to each server
	ConnsPerNode int
	// Request send max retry count. Negative value means count of
	// currently established connections to one server,
	// 0 means no retries will performed at all.
	MaxRequestRetries int
}

// Create new Sharder instance with given configuration.
func NewSharder(config SharderConf) *Sharder {
	sharder := &Sharder{
		config: config,
		nodes:  []string{},
		conns:  map[string]*Pool{},
		mu:     &sync.RWMutex{},
	}
	go sharder.reconfigLoop()
	return sharder
}

// Create default configuration for Sharder.
func NewSharderConf() SharderConf {
	return SharderConf{
		NodesGetter: func() ([]string, error) {
			return []string{}, nil
		},
		ReconfigPeriod:    10 * time.Second,
		ConnsPerNode:      5,
		MaxRequestRetries: -1,
	}
}

// Send request to one of configured servers. ID will not
// be sent, but used only to balance request between servers.
func (s *Sharder) Req(id, body []byte, timeout time.Duration) (reply []byte, err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	nodesCount := len(s.nodes)
	if nodesCount == 0 {
		return nil, errors.New("not configured")
	}
	i := shard(id, nodesCount)
	return s.conns[s.nodes[i]].Req(body, timeout)
}

// Sharding function.
func shard(id []byte, n int) int {
	hash := md5.Sum(id)
	return int(binary.BigEndian.Uint64(hash[:8]) % uint64(n))
}

// Goroutine.
// Reconfigure clients pool on the fly.
func (s *Sharder) reconfigLoop() {
	for {
		if nodes, err := s.config.NodesGetter(); err == nil {
			s.setNodes(nodes)
			time.Sleep(s.config.ReconfigPeriod)
		} else {
			time.Sleep(s.config.ReconfigPeriod / 10)
		}
	}
}

// Apply new node list to the clients pool.
func (s *Sharder) setNodes(newNodes []string) {
	newNodes = uniq(newNodes)
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, n := range newNodes {
		if _, ok := s.conns[n]; !ok {
			poolCfg := NewPoolConf()
			peers := make([]string, s.config.ConnsPerNode)
			for i := 0; i < s.config.ConnsPerNode; i++ {
				peers[i] = n
			}
			poolCfg.Peers = peers
			poolCfg.MaxRequestRetries = s.config.MaxRequestRetries
			s.conns[n] = NewPool(poolCfg)
		}
	}
	for k, p := range s.conns {
		if !inList(k, newNodes) {
			p.Close()
			delete(s.conns, k)
		}
	}
	s.nodes = newNodes
}

// Leave only unique elements from a string array.
func uniq(a []string) []string {
	r := []string{}
	for _, v := range a {
		if !inList(v, r) {
			r = append(r, v)
		}
	}
	return r
}

// Return truth if array a contains string s at least once.
func inList(s string, a []string) bool {
	for _, e := range a {
		if s == e {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package stream

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/ericxtang/m3u8"
	"github.com/golang/glog"
	"github.com/livepeer/go-livepeer/common"
)

const DefaultHLSStreamCap = uint(500)
const DefaultHLSStreamWin = uint(3)

// const DefaultMediaWinLen = uint(5)
const DefaultSegWaitTime = time.Second * 10
const SegWaitInterval = time.Second

var ErrAddHLSSegment = errors.New("ErrAddHLSSegment")

//BasicHLSVideoStream is a basic implementation of HLSVideoStream
type BasicHLSVideoStream struct {
	plCache    *m3u8.MediaPlaylist //StrmID -> MediaPlaylist
	sqMap      map[string]*HLSSegment
	lock       sync.Locker
	strmID     string
	subscriber func(*HLSSegment, bool)
	winSize    uint
}

func NewBasicHLSVideoStream(strmID string, wSize uint) *BasicHLSVideoStream {
	pl, err := m3u8.NewMediaPlaylist(wSize, DefaultHLSStreamCap)
	if err != nil {
		return nil
	}

	return &BasicHLSVideoStream{
		plCache: pl,
		// variant: variant,
		sqMap:   make(map[string]*HLSSegment),
		lock:    &sync.Mutex{},
		strmID:  strmID,
		winSize: wSize,
	}
}

//SetSubscriber sets the callback function that will be called when a new hls segment is inserted
func (s *BasicHLSVideoStream) SetSubscriber(f func(seg *HLSSegment, eof bool)) {
	s.subscriber = f
}

//GetStreamID returns the streamID
func (s *BasicHLSVideoStream) GetStreamID() string { return s.strmID }

//GetStreamFormat always returns HLS
func (s *BasicHLSVideoStream) GetStreamFormat() VideoFormat { return HLS }

//GetStreamPlaylist returns the media playlist represented by the streamID
func (s *BasicHLSVideoStream) GetStreamPlaylist() (*m3u8.MediaPlaylist, error) {
	if s.plCache.Count() < s.winSize {
		return nil, nil
	}

	return s.plCache, nil
}

// func (s *BasicHLSVideoStream) GetStreamVariant() *m3u8.Variant {
// 	return s.variant
// }

//GetHLSSegment gets the HLS segment.  It blocks until something is found, or timeout happens.
func (s *BasicHLSVideoStream) GetHLSSegment(segName string) (*HLSSegment, error) {
	seg, ok := s.sqMap[segName]
	if !ok {
		return nil, ErrNotFound
	}
	return seg, nil
}

//AddHLSSegment adds the hls segment to the right stream
func (s *BasicHLSVideoStream) AddHLSSegment(seg *HLSSegment) error {
	if _, ok := s.sqMap[seg.Name]; ok {
		return nil //Already have the seg.
	}
	glog.V(common.VERBOSE).Infof("Adding segment: %v", seg.Name)

	s.lock.Lock()
	defer s.lock.Unlock()

	//Add segment to media playlist
	s.plCache.AppendSegment(&m3u8.MediaSegment{SeqId: seg.SeqNo, Duration: seg.Duration, URI: seg.Name})
	if s.plCache.Count() > s.winSize {
		s.plCache.Remove()
	}

	//Add to buffer
	s.sqMap[seg.Name] = seg

	//Call subscriber
	if s.subscriber != nil {
		s.subscriber(seg, false)
	}

	return nil
}

func (s *BasicHLSVideoStream) End() {
	if s.subscriber != nil {
		s.subscriber(nil, true)
	}
}

func (s BasicHLSVideoStream) String() string {
	return fmt.Sprintf("StreamID: %v, Type: %v, len: %v", s.GetStreamID(), s.GetStreamFormat(), len(s.sqMap))
}

// //AddVariant adds a new variant playlist (and therefore, a new HLS video stream) to the master playlist.
// func (s *BasicHLSVideoStream) AddVariant(strmID string, variant *m3u8.Variant) error {
// 	if variant == nil {
// 		glog.Errorf("Cannot add nil variant")
// 		return ErrAddVariant
// 	}

// 	_, ok := s.variantMediaPlCache[strmID]
// 	if ok {
// 		glog.Errorf("Variant %v already exists", strmID)
// 		return ErrAddVariant
// 	}

// 	for _, v := range s.masterPlCache.Variants {
// 		if v.Bandwidth == variant.Bandwidth && v.Resolution == variant.Resolution {
// 			glog.Errorf("Variant with Bandwidth %v and Resolution %v already exists", v.Bandwidth, v.Resolution)
// 			return ErrAddVariant
// 		}
// 	}

// 	//Append to master playlist
// 	s.masterPlCache.Append(variant.URI, variant.Chunklist, variant.VariantParams)

// 	//Add to mediaPLCache
// 	s.variantMediaPlCache[strmID] = variant.Chunklist

// 	//Create the "media playlist specific" lock
// 	s.lockMap[strmID] = &sync.Mutex{}

// 	return nil
// }
<commit_msg>clean up<commit_after>package stream

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/ericxtang/m3u8"
	"github.com/golang/glog"
	"github.com/livepeer/go-livepeer/common"
)

const DefaultHLSStreamCap = uint(500)
const DefaultHLSStreamWin = uint(3)

// const DefaultMediaWinLen = uint(5)
const DefaultSegWaitTime = time.Second * 10
const SegWaitInterval = time.Second

var ErrAddHLSSegment = errors.New("ErrAddHLSSegment")

//BasicHLSVideoStream is a basic implementation of HLSVideoStream
type BasicHLSVideoStream struct {
	plCache    *m3u8.MediaPlaylist //StrmID -> MediaPlaylist
	sqMap      map[string]*HLSSegment
	lock       sync.Locker
	strmID     string
	subscriber func(*HLSSegment, bool)
	winSize    uint
}

func NewBasicHLSVideoStream(strmID string, wSize uint) *BasicHLSVideoStream {
	pl, err := m3u8.NewMediaPlaylist(wSize, DefaultHLSStreamCap)
	if err != nil {
		return nil
	}

	return &BasicHLSVideoStream{
		plCache: pl,
		// variant: variant,
		sqMap:   make(map[string]*HLSSegment),
		lock:    &sync.Mutex{},
		strmID:  strmID,
		winSize: wSize,
	}
}

//SetSubscriber sets the callback function that will be called when a new hls segment is inserted
func (s *BasicHLSVideoStream) SetSubscriber(f func(seg *HLSSegment, eof bool)) {
	s.subscriber = f
}

//GetStreamID returns the streamID
func (s *BasicHLSVideoStream) GetStreamID() string { return s.strmID }

//GetStreamFormat always returns HLS
func (s *BasicHLSVideoStream) GetStreamFormat() VideoFormat { return HLS }

//GetStreamPlaylist returns the media playlist represented by the streamID
func (s *BasicHLSVideoStream) GetStreamPlaylist() (*m3u8.MediaPlaylist, error) {
	if s.plCache.Count() < s.winSize {
		return nil, nil
	}

	return s.plCache, nil
}

//GetHLSSegment gets the HLS segment.  It blocks until something is found, or timeout happens.
func (s *BasicHLSVideoStream) GetHLSSegment(segName string) (*HLSSegment, error) {
	seg, ok := s.sqMap[segName]
	if !ok {
		return nil, ErrNotFound
	}
	return seg, nil
}

//AddHLSSegment adds the hls segment to the right stream
func (s *BasicHLSVideoStream) AddHLSSegment(seg *HLSSegment) error {
	if _, ok := s.sqMap[seg.Name]; ok {
		return nil //Already have the seg.
	}
	glog.V(common.VERBOSE).Infof("Adding segment: %v", seg.Name)

	s.lock.Lock()
	defer s.lock.Unlock()

	//Add segment to media playlist
	s.plCache.AppendSegment(&m3u8.MediaSegment{SeqId: seg.SeqNo, Duration: seg.Duration, URI: seg.Name})
	if s.plCache.Count() > s.winSize {
		s.plCache.Remove()
	}

	//Add to buffer
	s.sqMap[seg.Name] = seg

	//Call subscriber
	if s.subscriber != nil {
		s.subscriber(seg, false)
	}

	return nil
}

func (s *BasicHLSVideoStream) End() {
	if s.subscriber != nil {
		s.subscriber(nil, true)
	}
}

func (s BasicHLSVideoStream) String() string {
	return fmt.Sprintf("StreamID: %v, Type: %v, len: %v", s.GetStreamID(), s.GetStreamFormat(), len(s.sqMap))
}
<|endoftext|>