<commit_before>// Copyright 2018 Project Harbor 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 oidc

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	gooidc "github.com/coreos/go-oidc"
	"github.com/goharbor/harbor/src/common"
	"github.com/goharbor/harbor/src/common/dao/group"
	"github.com/goharbor/harbor/src/common/models"
	"github.com/goharbor/harbor/src/core/config"
	"github.com/goharbor/harbor/src/lib/log"
	"golang.org/x/oauth2"
)

const (
	googleEndpoint = "https://accounts.google.com"
)

type claimsProvider interface {
	Claims(v interface{}) error
}

type providerHelper struct {
	sync.Mutex
	instance     atomic.Value
	setting      atomic.Value
	creationTime time.Time
}

func (p *providerHelper) get() (*gooidc.Provider, error) {
	if p.instance.Load() != nil {
		if time.Now().Sub(p.creationTime) > 3*time.Second {
			if err := p.create(); err != nil {
				return nil, err
			}
		}
	} else {
		p.Lock()
		defer p.Unlock()
		if p.instance.Load() == nil {
			if err := p.reloadSetting(); err != nil {
				return nil, err
			}
			if err := p.create(); err != nil {
				return nil, err
			}
			go func() {
				for {
					if err := p.reloadSetting(); err != nil {
						log.Warningf("Failed to refresh configuration, error: %v", err)
					}
					time.Sleep(3 * time.Second)
				}
			}()
		}
	}

	return p.instance.Load().(*gooidc.Provider), nil
}

func (p *providerHelper) reloadSetting() error {
	conf, err := config.OIDCSetting()
	if err != nil {
		return fmt.Errorf("failed to load OIDC setting: %v", err)
	}
	p.setting.Store(*conf)
	return nil
}

func (p *providerHelper) create() error {
	if p.setting.Load() == nil {
		return errors.New("the configuration is not loaded")
	}
	s := p.setting.Load().(models.OIDCSetting)
	ctx := clientCtx(context.Background(), s.VerifyCert)
	provider, err := gooidc.NewProvider(ctx, s.Endpoint)
	if err != nil {
		return fmt.Errorf("failed to create OIDC provider, error: %v", err)
	}
	p.instance.Store(provider)
	p.creationTime = time.Now()
	return nil
}

var provider = &providerHelper{}

var insecureTransport = &http.Transport{
	TLSClientConfig: &tls.Config{
		InsecureSkipVerify: true,
	},
}

// Token wraps the attributes of a oauth2 token plus the attribute of ID token
type Token struct {
	oauth2.Token
	RawIDToken string `json:"id_token,omitempty"`
}

// UserInfo wraps the information that is extracted via token.  It will be transformed to data object that is persisted
// in the DB
type UserInfo struct {
	Issuer           string   `json:"iss"`
	Subject          string   `json:"sub"`
	Username         string   `json:"name"`
	Email            string   `json:"email"`
	Groups           []string `json:"groups"`
	AdminGroupMember bool     `json:"admin_group_member"`
	hasGroupClaim    bool
}

func getOauthConf() (*oauth2.Config, error) {
	p, err := provider.get()
	if err != nil {
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	scopes := make([]string, 0)
	for _, sc := range setting.Scope {
		if strings.HasPrefix(p.Endpoint().AuthURL, googleEndpoint) && sc == gooidc.ScopeOfflineAccess {
			log.Warningf("Dropped unsupported scope: %s ", sc)
			continue
		}
		scopes = append(scopes, sc)
	}
	return &oauth2.Config{
		ClientID:     setting.ClientID,
		ClientSecret: setting.ClientSecret,
		Scopes:       scopes,
		RedirectURL:  setting.RedirectURL,
		Endpoint:     p.Endpoint(),
	}, nil
}

// AuthCodeURL returns the URL for OIDC provider's consent page.  The state should be verified when user is redirected
// back to Harbor.
func AuthCodeURL(state string) (string, error) {
	conf, err := getOauthConf()
	if err != nil {
		log.Errorf("Failed to get OAuth configuration, error: %v", err)
		return "", err
	}
	var options []oauth2.AuthCodeOption
	setting := provider.setting.Load().(models.OIDCSetting)
	for k, v := range setting.ExtraRedirectParms {
		options = append(options, oauth2.SetAuthURLParam(k, v))
	}
	if strings.HasPrefix(conf.Endpoint.AuthURL, googleEndpoint) { // make sure the refresh token will be returned
		options = append(options, oauth2.AccessTypeOffline)
		options = append(options, oauth2.SetAuthURLParam("prompt", "consent"))
	}
	return conf.AuthCodeURL(state, options...), nil
}

// ExchangeToken get the token from token provider via the code
func ExchangeToken(ctx context.Context, code string) (*Token, error) {
	oauth, err := getOauthConf()
	if err != nil {
		log.Errorf("Failed to get OAuth configuration, error: %v", err)
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	ctx = clientCtx(ctx, setting.VerifyCert)
	oauthToken, err := oauth.Exchange(ctx, code)
	if err != nil {
		return nil, err
	}
	return &Token{Token: *oauthToken, RawIDToken: oauthToken.Extra("id_token").(string)}, nil
}

func parseIDToken(ctx context.Context, rawIDToken string) (*gooidc.IDToken, error) {
	conf := &gooidc.Config{SkipClientIDCheck: true, SkipExpiryCheck: true}
	return verifyTokenWithConfig(ctx, rawIDToken, conf)
}

// VerifyToken verifies the ID token based on the OIDC settings
func VerifyToken(ctx context.Context, rawIDToken string) (*gooidc.IDToken, error) {
	return verifyTokenWithConfig(ctx, rawIDToken, nil)
}

func verifyTokenWithConfig(ctx context.Context, rawIDToken string, conf *gooidc.Config) (*gooidc.IDToken, error) {
	log.Debugf("Raw ID token for verification: %s", rawIDToken)
	p, err := provider.get()
	if err != nil {
		return nil, err
	}
	settings := provider.setting.Load().(models.OIDCSetting)
	if conf == nil {
		conf = &gooidc.Config{ClientID: settings.ClientID}
	}
	verifier := p.Verifier(conf)
	ctx = clientCtx(ctx, settings.VerifyCert)
	return verifier.Verify(ctx, rawIDToken)
}

func clientCtx(ctx context.Context, verifyCert bool) context.Context {
	var client *http.Client
	if !verifyCert {
		client = &http.Client{
			Transport: insecureTransport,
		}
	} else {
		client = &http.Client{}
	}
	return gooidc.ClientContext(ctx, client)
}

// refreshToken tries to refresh the token if it's expired, if it doesn't the
// original one will be returned.
func refreshToken(ctx context.Context, token *Token) (*Token, error) {
	oauthCfg, err := getOauthConf()
	if err != nil {
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	cctx := clientCtx(ctx, setting.VerifyCert)
	ts := oauthCfg.TokenSource(cctx, &token.Token)
	nt, err := ts.Token()
	if err != nil {
		return nil, err
	}
	it, ok := nt.Extra("id_token").(string)
	if !ok {
		log.Debug("id_token not exist in refresh response")
	}
	return &Token{Token: *nt, RawIDToken: it}, nil
}

// UserInfoFromToken tries to call the UserInfo endpoint of the OIDC provider, and consolidate with ID token
// to generate a UserInfo object, if the ID token is not in the input token struct, some attributes will be empty
func UserInfoFromToken(ctx context.Context, token *Token) (*UserInfo, error) {
	// #10913: preload the configuration, in case it was not previously loaded by the UI
	_, err := provider.get()
	if err != nil {
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	local, err := UserInfoFromIDToken(ctx, token, setting)
	if err != nil {
		return nil, err
	}
	remote, err := userInfoFromRemote(ctx, token, setting)
	if err != nil {
		log.Warningf("Failed to get userInfo by calling remote userinfo endpoint, error: %v ", err)
	}
	if remote != nil && local != nil {
		if remote.Subject != local.Subject {
			return nil, fmt.Errorf("the subject from userinfo: %s does not match the subject from ID token: %s, probably a security attack happened", remote.Subject, local.Subject)
		}
		return mergeUserInfo(remote, local), nil
	} else if remote != nil && local == nil {
		return remote, nil
	} else if local != nil && remote == nil {
		log.Debugf("Fall back to user data from ID token.")
		return local, nil
	}
	return nil, fmt.Errorf("failed to get userinfo from both remote and ID token")
}

func mergeUserInfo(remote, local *UserInfo) *UserInfo {
	res := &UserInfo{
		// data only contained in ID token
		Subject: local.Subject,
		Issuer:  local.Issuer,
		// Used data from userinfo
		Username: remote.Username,
		Email:    remote.Email,
	}
	if remote.hasGroupClaim {
		res.Groups = remote.Groups
		res.AdminGroupMember = remote.AdminGroupMember
		res.hasGroupClaim = true
	} else if local.hasGroupClaim {
		res.Groups = local.Groups
		res.AdminGroupMember = local.AdminGroupMember
		res.hasGroupClaim = true
	} else {
		res.Groups = []string{}
	}
	return res
}

func userInfoFromRemote(ctx context.Context, token *Token, setting models.OIDCSetting) (*UserInfo, error) {
	p, err := provider.get()
	if err != nil {
		return nil, err
	}
	cctx := clientCtx(ctx, setting.VerifyCert)
	u, err := p.UserInfo(cctx, oauth2.StaticTokenSource(&token.Token))
	if err != nil {
		return nil, err
	}
	return userInfoFromClaims(u, setting)
}

// UserInfoFromIDToken extract user info from ID token
func UserInfoFromIDToken(ctx context.Context, token *Token, setting models.OIDCSetting) (*UserInfo, error) {
	if token.RawIDToken == "" {
		return nil, nil
	}
	idt, err := parseIDToken(ctx, token.RawIDToken)
	if err != nil {
		return nil, err
	}

	return userInfoFromClaims(idt, setting)
}

func userInfoFromClaims(c claimsProvider, setting models.OIDCSetting) (*UserInfo, error) {
	res := &UserInfo{}
	if err := c.Claims(res); err != nil {
		return nil, err
	}
	if setting.UserClaim != "" {
		allClaims := make(map[string]interface{})
		if err := c.Claims(&allClaims); err != nil {
			return nil, err
		}

		username, ok := allClaims[setting.UserClaim].(string)
		if !ok {
			return nil, fmt.Errorf("OIDC. Failed to recover Username from claim. Claim '%s' is invalid or not a string", setting.UserClaim)
		}
		res.Username = username
	}
	res.Groups, res.hasGroupClaim = groupsFromClaims(c, setting.GroupsClaim)
	if len(setting.AdminGroup) > 0 {
		for _, g := range res.Groups {
			if g == setting.AdminGroup {
				res.AdminGroupMember = true
				break
			}
		}
	}
	return res, nil
}

// groupsFromClaims fetches the group name list from claimprovider, such as decoded ID token.
// If the claims does not have the claim defined as k, the second return value will be false, otherwise true
func groupsFromClaims(gp claimsProvider, k string) ([]string, bool) {
	res := make([]string, 0)
	claimMap := make(map[string]interface{})
	if err := gp.Claims(&claimMap); err != nil {
		log.Errorf("failed to fetch claims, error: %v", err)
		return res, false
	}
	g, ok := claimMap[k].([]interface{})
	if !ok {
		log.Warningf("Unable to get groups from claims, claims: %+v, groups claims key: %s", claimMap, k)
		return res, false
	}
	for _, e := range g {
		s, ok := e.(string)
		if !ok {
			log.Warningf("Element in group list is not string: %v, list: %v", e, g)
			continue
		}
		res = append(res, s)
	}
	return res, true
}

type populate func(groupNames []string) ([]int, error)

func populateGroupsDB(groupNames []string) ([]int, error) {
	return group.PopulateGroup(models.UserGroupsFromName(groupNames, common.OIDCGroupType))
}

// InjectGroupsToUser populates the group to DB and inject the group IDs to user model.
// The third optional parm is for UT only.
func InjectGroupsToUser(info *UserInfo, user *models.User, f ...populate) {
	if info == nil || user == nil {
		log.Warningf("user info or user model is nil, skip the func")
		return
	}
	var populateGroups populate
	if len(f) == 0 {
		populateGroups = populateGroupsDB
	} else {
		populateGroups = f[0]
	}
	if gids, err := populateGroups(info.Groups); err != nil {
		log.Warningf("failed to get group ID, error: %v, skip populating groups", err)
	} else {
		user.GroupIDs = gids
	}
	user.AdminRoleInAuth = info.AdminGroupMember
}

// Conn wraps connection info of an OIDC endpoint
type Conn struct {
	URL        string `json:"url"`
	VerifyCert bool   `json:"verify_cert"`
}

// TestEndpoint tests whether the endpoint is a valid OIDC endpoint.
// The nil return value indicates the success of the test
func TestEndpoint(conn Conn) error {

	// gooidc will try to call the discovery api when creating the provider and that's all we need to check
	ctx := clientCtx(context.Background(), conn.VerifyCert)
	_, err := gooidc.NewProvider(ctx, conn.URL)
	return err
}
<commit_msg>Username from /userinfo (#14038)<commit_after>// Copyright 2018 Project Harbor 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 oidc

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	gooidc "github.com/coreos/go-oidc"
	"github.com/goharbor/harbor/src/common"
	"github.com/goharbor/harbor/src/common/dao/group"
	"github.com/goharbor/harbor/src/common/models"
	"github.com/goharbor/harbor/src/core/config"
	"github.com/goharbor/harbor/src/lib/log"
	"golang.org/x/oauth2"
)

const (
	googleEndpoint = "https://accounts.google.com"
)

type claimsProvider interface {
	Claims(v interface{}) error
}

type providerHelper struct {
	sync.Mutex
	instance     atomic.Value
	setting      atomic.Value
	creationTime time.Time
}

func (p *providerHelper) get() (*gooidc.Provider, error) {
	if p.instance.Load() != nil {
		if time.Now().Sub(p.creationTime) > 3*time.Second {
			if err := p.create(); err != nil {
				return nil, err
			}
		}
	} else {
		p.Lock()
		defer p.Unlock()
		if p.instance.Load() == nil {
			if err := p.reloadSetting(); err != nil {
				return nil, err
			}
			if err := p.create(); err != nil {
				return nil, err
			}
			go func() {
				for {
					if err := p.reloadSetting(); err != nil {
						log.Warningf("Failed to refresh configuration, error: %v", err)
					}
					time.Sleep(3 * time.Second)
				}
			}()
		}
	}

	return p.instance.Load().(*gooidc.Provider), nil
}

func (p *providerHelper) reloadSetting() error {
	conf, err := config.OIDCSetting()
	if err != nil {
		return fmt.Errorf("failed to load OIDC setting: %v", err)
	}
	p.setting.Store(*conf)
	return nil
}

func (p *providerHelper) create() error {
	if p.setting.Load() == nil {
		return errors.New("the configuration is not loaded")
	}
	s := p.setting.Load().(models.OIDCSetting)
	ctx := clientCtx(context.Background(), s.VerifyCert)
	provider, err := gooidc.NewProvider(ctx, s.Endpoint)
	if err != nil {
		return fmt.Errorf("failed to create OIDC provider, error: %v", err)
	}
	p.instance.Store(provider)
	p.creationTime = time.Now()
	return nil
}

var provider = &providerHelper{}

var insecureTransport = &http.Transport{
	TLSClientConfig: &tls.Config{
		InsecureSkipVerify: true,
	},
}

// Token wraps the attributes of a oauth2 token plus the attribute of ID token
type Token struct {
	oauth2.Token
	RawIDToken string `json:"id_token,omitempty"`
}

// UserInfo wraps the information that is extracted via token.  It will be transformed to data object that is persisted
// in the DB
type UserInfo struct {
	Issuer           string   `json:"iss"`
	Subject          string   `json:"sub"`
	Username         string   `json:"name"`
	Email            string   `json:"email"`
	Groups           []string `json:"groups"`
	AdminGroupMember bool     `json:"admin_group_member"`
	hasGroupClaim    bool
}

func getOauthConf() (*oauth2.Config, error) {
	p, err := provider.get()
	if err != nil {
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	scopes := make([]string, 0)
	for _, sc := range setting.Scope {
		if strings.HasPrefix(p.Endpoint().AuthURL, googleEndpoint) && sc == gooidc.ScopeOfflineAccess {
			log.Warningf("Dropped unsupported scope: %s ", sc)
			continue
		}
		scopes = append(scopes, sc)
	}
	return &oauth2.Config{
		ClientID:     setting.ClientID,
		ClientSecret: setting.ClientSecret,
		Scopes:       scopes,
		RedirectURL:  setting.RedirectURL,
		Endpoint:     p.Endpoint(),
	}, nil
}

// AuthCodeURL returns the URL for OIDC provider's consent page.  The state should be verified when user is redirected
// back to Harbor.
func AuthCodeURL(state string) (string, error) {
	conf, err := getOauthConf()
	if err != nil {
		log.Errorf("Failed to get OAuth configuration, error: %v", err)
		return "", err
	}
	var options []oauth2.AuthCodeOption
	setting := provider.setting.Load().(models.OIDCSetting)
	for k, v := range setting.ExtraRedirectParms {
		options = append(options, oauth2.SetAuthURLParam(k, v))
	}
	if strings.HasPrefix(conf.Endpoint.AuthURL, googleEndpoint) { // make sure the refresh token will be returned
		options = append(options, oauth2.AccessTypeOffline)
		options = append(options, oauth2.SetAuthURLParam("prompt", "consent"))
	}
	return conf.AuthCodeURL(state, options...), nil
}

// ExchangeToken get the token from token provider via the code
func ExchangeToken(ctx context.Context, code string) (*Token, error) {
	oauth, err := getOauthConf()
	if err != nil {
		log.Errorf("Failed to get OAuth configuration, error: %v", err)
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	ctx = clientCtx(ctx, setting.VerifyCert)
	oauthToken, err := oauth.Exchange(ctx, code)
	if err != nil {
		return nil, err
	}
	return &Token{Token: *oauthToken, RawIDToken: oauthToken.Extra("id_token").(string)}, nil
}

func parseIDToken(ctx context.Context, rawIDToken string) (*gooidc.IDToken, error) {
	conf := &gooidc.Config{SkipClientIDCheck: true, SkipExpiryCheck: true}
	return verifyTokenWithConfig(ctx, rawIDToken, conf)
}

// VerifyToken verifies the ID token based on the OIDC settings
func VerifyToken(ctx context.Context, rawIDToken string) (*gooidc.IDToken, error) {
	return verifyTokenWithConfig(ctx, rawIDToken, nil)
}

func verifyTokenWithConfig(ctx context.Context, rawIDToken string, conf *gooidc.Config) (*gooidc.IDToken, error) {
	log.Debugf("Raw ID token for verification: %s", rawIDToken)
	p, err := provider.get()
	if err != nil {
		return nil, err
	}
	settings := provider.setting.Load().(models.OIDCSetting)
	if conf == nil {
		conf = &gooidc.Config{ClientID: settings.ClientID}
	}
	verifier := p.Verifier(conf)
	ctx = clientCtx(ctx, settings.VerifyCert)
	return verifier.Verify(ctx, rawIDToken)
}

func clientCtx(ctx context.Context, verifyCert bool) context.Context {
	var client *http.Client
	if !verifyCert {
		client = &http.Client{
			Transport: insecureTransport,
		}
	} else {
		client = &http.Client{}
	}
	return gooidc.ClientContext(ctx, client)
}

// refreshToken tries to refresh the token if it's expired, if it doesn't the
// original one will be returned.
func refreshToken(ctx context.Context, token *Token) (*Token, error) {
	oauthCfg, err := getOauthConf()
	if err != nil {
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	cctx := clientCtx(ctx, setting.VerifyCert)
	ts := oauthCfg.TokenSource(cctx, &token.Token)
	nt, err := ts.Token()
	if err != nil {
		return nil, err
	}
	it, ok := nt.Extra("id_token").(string)
	if !ok {
		log.Debug("id_token not exist in refresh response")
	}
	return &Token{Token: *nt, RawIDToken: it}, nil
}

// UserInfoFromToken tries to call the UserInfo endpoint of the OIDC provider, and consolidate with ID token
// to generate a UserInfo object, if the ID token is not in the input token struct, some attributes will be empty
func UserInfoFromToken(ctx context.Context, token *Token) (*UserInfo, error) {
	// #10913: preload the configuration, in case it was not previously loaded by the UI
	_, err := provider.get()
	if err != nil {
		return nil, err
	}
	setting := provider.setting.Load().(models.OIDCSetting)
	local, err := UserInfoFromIDToken(ctx, token, setting)
	if err != nil {
		return nil, err
	}
	remote, err := userInfoFromRemote(ctx, token, setting)
	if err != nil {
		log.Warningf("Failed to get userInfo by calling remote userinfo endpoint, error: %v ", err)
	}

	if setting.UserClaim != "" && local.Username == "" && remote.Username == "" {
		return nil, fmt.Errorf("OIDC. Failed to recover Username from claim. Claim '%s' is invalid or not a string", setting.UserClaim)
	}
	if remote != nil && local != nil {
		if remote.Subject != local.Subject {
			return nil, fmt.Errorf("the subject from userinfo: %s does not match the subject from ID token: %s, probably a security attack happened", remote.Subject, local.Subject)
		}
		return mergeUserInfo(remote, local), nil
	} else if remote != nil && local == nil {
		return remote, nil
	} else if local != nil && remote == nil {
		log.Debugf("Fall back to user data from ID token.")
		return local, nil
	}
	return nil, fmt.Errorf("failed to get userinfo from both remote and ID token")
}

func mergeUserInfo(remote, local *UserInfo) *UserInfo {
	res := &UserInfo{
		// data only contained in ID token
		Subject: local.Subject,
		Issuer:  local.Issuer,
		// Used data from userinfo
		Username: remote.Username,
		Email:    remote.Email,
	}
	if remote.hasGroupClaim {
		res.Groups = remote.Groups
		res.AdminGroupMember = remote.AdminGroupMember
		res.hasGroupClaim = true
	} else if local.hasGroupClaim {
		res.Groups = local.Groups
		res.AdminGroupMember = local.AdminGroupMember
		res.hasGroupClaim = true
	} else {
		res.Groups = []string{}
	}
	return res
}

func userInfoFromRemote(ctx context.Context, token *Token, setting models.OIDCSetting) (*UserInfo, error) {
	p, err := provider.get()
	if err != nil {
		return nil, err
	}
	cctx := clientCtx(ctx, setting.VerifyCert)
	u, err := p.UserInfo(cctx, oauth2.StaticTokenSource(&token.Token))
	if err != nil {
		return nil, err
	}
	return userInfoFromClaims(u, setting)
}

// UserInfoFromIDToken extract user info from ID token
func UserInfoFromIDToken(ctx context.Context, token *Token, setting models.OIDCSetting) (*UserInfo, error) {
	if token.RawIDToken == "" {
		return nil, nil
	}
	idt, err := parseIDToken(ctx, token.RawIDToken)
	if err != nil {
		return nil, err
	}

	return userInfoFromClaims(idt, setting)
}

func userInfoFromClaims(c claimsProvider, setting models.OIDCSetting) (*UserInfo, error) {
	res := &UserInfo{}
	if err := c.Claims(res); err != nil {
		return nil, err
	}
	if setting.UserClaim != "" {
		allClaims := make(map[string]interface{})
		if err := c.Claims(&allClaims); err != nil {
			return nil, err
		}

		if username, ok := allClaims[setting.UserClaim].(string); ok {
			res.Username = username
		} else {
			log.Warningf("OIDC. Failed to recover Username from claim. Claim '%s' is invalid or not a string", setting.UserClaim)
		}
	}
	res.Groups, res.hasGroupClaim = groupsFromClaims(c, setting.GroupsClaim)
	if len(setting.AdminGroup) > 0 {
		for _, g := range res.Groups {
			if g == setting.AdminGroup {
				res.AdminGroupMember = true
				break
			}
		}
	}
	return res, nil
}

// groupsFromClaims fetches the group name list from claimprovider, such as decoded ID token.
// If the claims does not have the claim defined as k, the second return value will be false, otherwise true
func groupsFromClaims(gp claimsProvider, k string) ([]string, bool) {
	res := make([]string, 0)
	claimMap := make(map[string]interface{})
	if err := gp.Claims(&claimMap); err != nil {
		log.Errorf("failed to fetch claims, error: %v", err)
		return res, false
	}
	g, ok := claimMap[k].([]interface{})
	if !ok {
		log.Warningf("Unable to get groups from claims, claims: %+v, groups claims key: %s", claimMap, k)
		return res, false
	}
	for _, e := range g {
		s, ok := e.(string)
		if !ok {
			log.Warningf("Element in group list is not string: %v, list: %v", e, g)
			continue
		}
		res = append(res, s)
	}
	return res, true
}

type populate func(groupNames []string) ([]int, error)

func populateGroupsDB(groupNames []string) ([]int, error) {
	return group.PopulateGroup(models.UserGroupsFromName(groupNames, common.OIDCGroupType))
}

// InjectGroupsToUser populates the group to DB and inject the group IDs to user model.
// The third optional parm is for UT only.
func InjectGroupsToUser(info *UserInfo, user *models.User, f ...populate) {
	if info == nil || user == nil {
		log.Warningf("user info or user model is nil, skip the func")
		return
	}
	var populateGroups populate
	if len(f) == 0 {
		populateGroups = populateGroupsDB
	} else {
		populateGroups = f[0]
	}
	if gids, err := populateGroups(info.Groups); err != nil {
		log.Warningf("failed to get group ID, error: %v, skip populating groups", err)
	} else {
		user.GroupIDs = gids
	}
	user.AdminRoleInAuth = info.AdminGroupMember
}

// Conn wraps connection info of an OIDC endpoint
type Conn struct {
	URL        string `json:"url"`
	VerifyCert bool   `json:"verify_cert"`
}

// TestEndpoint tests whether the endpoint is a valid OIDC endpoint.
// The nil return value indicates the success of the test
func TestEndpoint(conn Conn) error {

	// gooidc will try to call the discovery api when creating the provider and that's all we need to check
	ctx := clientCtx(context.Background(), conn.VerifyCert)
	_, err := gooidc.NewProvider(ctx, conn.URL)
	return err
}
<|endoftext|>


<commit_before><commit_msg>syscall: add comment<commit_after><|endoftext|>


<commit_before>package tdigest

import (
	"fmt"
	"log"
	"math"
	"math/rand"
)

// A TDigest is an efficient data structure for computing streaming
// approximate quantiles of a dataset. It supports two methods:
//
// Add will add a value to the TDigest, updating all quantiles. A
// weight can be specified; use weight of 1 if you don't care about
// weighting your dataset.
//
// Quantile(q) will estimate the qth quantile value of the
// dataset. The input value of q should be in the range [0.0, 1.0]; if
// it is outside that range, it will be clipped into it automatically.
//
// Calling Quantile on a TDigest with no data should return NaN.
//
// MergeInto(other) will add all of the data within a TDigest into
// other, combining them into one larger TDigest.
type TDigest interface {
	Add(val float64, weight int)
	Quantile(q float64) (val float64)
	MergeInto(other TDigest)
}

// New produces a new TDigest using the default compression level of
// 100.
func New() TDigest {
	return NewWithCompression(100)
}

// NewWithCompression produces a new TDigest with a specific
// compression level. The input compression value, which should be >=
// 1.0, will control how aggressively the TDigest compresses data
// together.
//
// The original TDigest paper suggests using a value of 100 for a good
// balance between precision and efficiency. It will land at very
// small (think like 1e-6 percentile points) errors at extreme points
// in the distribution, and compression ratios of around 500 for large
// data sets (1 millionish datapoints).
func NewWithCompression(compression float64) TDigest {
	return newCentroidSet(compression)
}

// centroid is a simple container for a mean,count pair.
type centroid struct {
	mean  float64
	count int
}

func (c *centroid) String() string {
	return fmt.Sprintf("c{%f x%d}", c.mean, c.count)
}

type centroidSet struct {
	centroids   []*centroid
	compression float64
	countTotal  int

	reclusterAt int
}

func newCentroidSet(compression float64) *centroidSet {
	return &centroidSet{
		centroids:   make([]*centroid, 0),
		compression: compression,
		countTotal:  0,
	}
}

// Find the indexes of centroids which have the minimum distance to the
// input value.
//
// TODO: Use a better data structure to avoid this loop.
func (cs *centroidSet) nearest(val float64) []int {
	var (
		nearestDist float64 = math.Inf(+1)
		thisDist    float64
		delta       float64
		result      []int = make([]int, 0)
	)
	for i, c := range cs.centroids {
		thisDist = val - c.mean
		if thisDist < 0 {
			thisDist *= -1
		}

		delta = thisDist - nearestDist
		switch {
		case delta < 0:
			// we have a new winner!
			nearestDist = thisDist
			result = result[0:0] // wipe result
			result = append(result, i)
		case delta == 0:
			// we have a tie
			result = append(result, i)
		default:
			// Since cs.centroids is sorted by mean, this means we
			// have passed the best spot, so we may as well break
			break
		}
	}
	return result
}

// returns the maximum weight that can be placed at specified index
func (cs *centroidSet) weightLimit(idx int) int {
	ptile := cs.quantileOf(idx)
	limit := int(4 * cs.compression * ptile * (1 - ptile) * float64(len(cs.centroids)))
	return limit
}

// checks whether the centroid has room for more weight
func (cs *centroidSet) centroidHasRoom(idx int) bool {
	return cs.centroids[idx].count < cs.weightLimit(idx)
}

// find which centroid to add the value to (by index)
func (cs *centroidSet) findAddTarget(val float64) int {
	var (
		nearest  []int = cs.nearest(val)
		eligible []int
	)

	for _, c := range nearest {
		// if there is room for more weight at this centroid...
		if cs.centroidHasRoom(c) {
			eligible = append(eligible, c)
		}
	}

	if len(eligible) == 0 {
		return -1
	}
	if len(eligible) == 1 {
		return eligible[0]
	}

	// Multiple eligible centroids to add to. They must be equidistant
	// from this value. Four cases are possible:
	//
	//   1. All eligible centroids' means are less than val
	//   2. Some eligible centroids' means are less than val, some are greater
	//   3. All eligible centroids' means are exactly equal to val
	//   4. All eligible centroids' means are greater than val
	//
	// If 1, then we should take the highest indexed centroid to
	// preserve ordering.  If 4, we should take the lowest for the
	// same reason. If 2 or 3, we can pick randomly.

	var anyLesser, anyGreater bool
	for _, c := range eligible {
		m := cs.centroids[c].mean
		if m < val {
			anyLesser = true
		} else if m > val {
			anyGreater = true
		}
	}

	// case 1: all are less, none are greater. Take highest one.
	if anyLesser && !anyGreater {
		greatest := eligible[0]
		for _, c := range eligible[1:] {
			if c > greatest {
				greatest = c
			}
		}
		return greatest
	}

	// case 4: all are greater, none are less. Take lowest one.
	if !anyLesser && anyGreater {
		least := eligible[0]
		for _, c := range eligible[1:] {
			if c < least {
				least = c
			}
		}
		return least
	}

	// case 2 or 3: Anything else. Pick randomly.
	return eligible[rand.Intn(len(eligible))]
}

func (cs *centroidSet) addNewCentroid(mean float64, weight int) {
	var idx int = len(cs.centroids)

	for i, c := range cs.centroids {
		// add in sorted order
		if mean < c.mean {
			idx = i
			break
		}
	}

	cs.centroids = append(cs.centroids, nil)
	copy(cs.centroids[idx+1:], cs.centroids[idx:])
	cs.centroids[idx] = &centroid{mean, weight}
}

// Add a value to the centroidSet.
func (cs *centroidSet) Add(val float64, weight int) {
	cs.countTotal += weight
	var idx = cs.findAddTarget(val)

	if idx == -1 {
		cs.addNewCentroid(val, weight)
		return
	}

	c := cs.centroids[idx]

	limit := cs.weightLimit(idx)
	// how much weight will we be adding?
	// if adding this node to this centroid would put it over the
	// weight limit, just add the most we can and recur with the remainder
	if c.count+weight > limit {
		log.Printf("splitting value %f weight=%d (adding just %d)", val, weight, limit-c.count)
		add := limit - c.count
		if add < 0 {
			// this node was already overweight
			add = 0
		}
		remainder := weight - add

		c.count += add
		c.mean = c.mean + float64(add)*(val-c.mean)/float64(c.count)

		cs.Add(val, remainder)
	} else {
		c.count += weight
		c.mean = c.mean + float64(weight)*(val-c.mean)/float64(c.count)
	}

}

// returns the approximate quantile that a particular centroid
// represents
func (cs *centroidSet) quantileOf(idx int) float64 {
	total := 0
	for _, c := range cs.centroids[:idx] {
		total += c.count
	}
	return (float64(cs.centroids[idx].count/2) + float64(total)) / float64(cs.countTotal)
}

// Quantile returns the approximate value at a quantile (eg the 99th
// percentile value would be centroidSet.quantileValue(0.99))
func (cs *centroidSet) Quantile(q float64) float64 {
	var n = len(cs.centroids)
	if n == 0 {
		return math.NaN()
	}
	if n == 1 {
		return cs.centroids[0].mean
	}

	if q < 0 {
		q = 0
	} else if q > 1 {
		q = 1
	}

	// rescale into count units instead of 0 to 1 units
	q = float64(cs.countTotal) * q
	// find the first centroid which straddles q
	var (
		qTotal float64 = 0
		i      int
	)
	for i = 0; i < n && float64(cs.centroids[i].count)/2+qTotal < q; i++ {
		qTotal += float64(cs.centroids[i].count)
	}

	if i == 0 {
		// special case 1: the targeted quantile is before the
		// left-most centroid. extrapolate from the slope from
		// centroid0 to centroid1.
		c0 := cs.centroids[0]
		c1 := cs.centroids[1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - float64(c0.count)/2 // this is negative
		return c0.mean + slope*deltaQ
	}
	if i == n {
		// special case 2: the targeted quantile is from the
		// right-most centroid. extrapolate from the slope at the
		// right edge.
		c0 := cs.centroids[n-2]
		c1 := cs.centroids[n-1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - (qTotal - float64(c1.count)/2)
		return c1.mean + slope*deltaQ
	}
	// common case: targeted quantile is between 2 centroids
	c0 := cs.centroids[i-1]
	c1 := cs.centroids[i]
	slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
	deltaQ := q - (float64(c1.count)/2 + qTotal)
	return c1.mean + slope*deltaQ
}

func (cs *centroidSet) MergeInto(other TDigest) {
	// Add each centroid in cs into other. They should be added in
	// random order.
	addOrder := rand.Perm(len(cs.centroids))
	for _, idx := range addOrder {
		c := cs.centroids[idx]
		// gradually write up the volume written so that the tdigest doesnt overload early
		added := 0
		for i := 1; i < 10; i++ {
			fmt.Printf("%v\n", other.(*centroidSet).centroids)
			toAdd := i * 2
			if added+i > c.count {
				toAdd = c.count - added
			}
			other.Add(c.mean, toAdd)
			added += toAdd
			if added >= c.count {
				break
			}
		}
		if added < c.count {
			other.Add(c.mean, c.count-added)
		}
		other.Add(c.mean, c.count)
	}
}
<commit_msg>Remove logging cruft<commit_after>package tdigest

import (
	"fmt"
	"math"
	"math/rand"
)

// A TDigest is an efficient data structure for computing streaming
// approximate quantiles of a dataset. It supports two methods:
//
// Add will add a value to the TDigest, updating all quantiles. A
// weight can be specified; use weight of 1 if you don't care about
// weighting your dataset.
//
// Quantile(q) will estimate the qth quantile value of the
// dataset. The input value of q should be in the range [0.0, 1.0]; if
// it is outside that range, it will be clipped into it automatically.
//
// Calling Quantile on a TDigest with no data should return NaN.
//
// MergeInto(other) will add all of the data within a TDigest into
// other, combining them into one larger TDigest.
type TDigest interface {
	Add(val float64, weight int)
	Quantile(q float64) (val float64)
	MergeInto(other TDigest)
}

// New produces a new TDigest using the default compression level of
// 100.
func New() TDigest {
	return NewWithCompression(100)
}

// NewWithCompression produces a new TDigest with a specific
// compression level. The input compression value, which should be >=
// 1.0, will control how aggressively the TDigest compresses data
// together.
//
// The original TDigest paper suggests using a value of 100 for a good
// balance between precision and efficiency. It will land at very
// small (think like 1e-6 percentile points) errors at extreme points
// in the distribution, and compression ratios of around 500 for large
// data sets (1 millionish datapoints).
func NewWithCompression(compression float64) TDigest {
	return newCentroidSet(compression)
}

// centroid is a simple container for a mean,count pair.
type centroid struct {
	mean  float64
	count int
}

func (c *centroid) String() string {
	return fmt.Sprintf("c{%f x%d}", c.mean, c.count)
}

type centroidSet struct {
	centroids   []*centroid
	compression float64
	countTotal  int

	reclusterAt int
}

func newCentroidSet(compression float64) *centroidSet {
	return &centroidSet{
		centroids:   make([]*centroid, 0),
		compression: compression,
		countTotal:  0,
	}
}

// Find the indexes of centroids which have the minimum distance to the
// input value.
//
// TODO: Use a better data structure to avoid this loop.
func (cs *centroidSet) nearest(val float64) []int {
	var (
		nearestDist float64 = math.Inf(+1)
		thisDist    float64
		delta       float64
		result      []int = make([]int, 0)
	)
	for i, c := range cs.centroids {
		thisDist = val - c.mean
		if thisDist < 0 {
			thisDist *= -1
		}

		delta = thisDist - nearestDist
		switch {
		case delta < 0:
			// we have a new winner!
			nearestDist = thisDist
			result = result[0:0] // wipe result
			result = append(result, i)
		case delta == 0:
			// we have a tie
			result = append(result, i)
		default:
			// Since cs.centroids is sorted by mean, this means we
			// have passed the best spot, so we may as well break
			break
		}
	}
	return result
}

// returns the maximum weight that can be placed at specified index
func (cs *centroidSet) weightLimit(idx int) int {
	ptile := cs.quantileOf(idx)
	limit := int(4 * cs.compression * ptile * (1 - ptile) * float64(len(cs.centroids)))
	return limit
}

// checks whether the centroid has room for more weight
func (cs *centroidSet) centroidHasRoom(idx int) bool {
	return cs.centroids[idx].count < cs.weightLimit(idx)
}

// find which centroid to add the value to (by index)
func (cs *centroidSet) findAddTarget(val float64) int {
	var (
		nearest  []int = cs.nearest(val)
		eligible []int
	)

	for _, c := range nearest {
		// if there is room for more weight at this centroid...
		if cs.centroidHasRoom(c) {
			eligible = append(eligible, c)
		}
	}

	if len(eligible) == 0 {
		return -1
	}
	if len(eligible) == 1 {
		return eligible[0]
	}

	// Multiple eligible centroids to add to. They must be equidistant
	// from this value. Four cases are possible:
	//
	//   1. All eligible centroids' means are less than val
	//   2. Some eligible centroids' means are less than val, some are greater
	//   3. All eligible centroids' means are exactly equal to val
	//   4. All eligible centroids' means are greater than val
	//
	// If 1, then we should take the highest indexed centroid to
	// preserve ordering.  If 4, we should take the lowest for the
	// same reason. If 2 or 3, we can pick randomly.

	var anyLesser, anyGreater bool
	for _, c := range eligible {
		m := cs.centroids[c].mean
		if m < val {
			anyLesser = true
		} else if m > val {
			anyGreater = true
		}
	}

	// case 1: all are less, none are greater. Take highest one.
	if anyLesser && !anyGreater {
		greatest := eligible[0]
		for _, c := range eligible[1:] {
			if c > greatest {
				greatest = c
			}
		}
		return greatest
	}

	// case 4: all are greater, none are less. Take lowest one.
	if !anyLesser && anyGreater {
		least := eligible[0]
		for _, c := range eligible[1:] {
			if c < least {
				least = c
			}
		}
		return least
	}

	// case 2 or 3: Anything else. Pick randomly.
	return eligible[rand.Intn(len(eligible))]
}

func (cs *centroidSet) addNewCentroid(mean float64, weight int) {
	var idx int = len(cs.centroids)

	for i, c := range cs.centroids {
		// add in sorted order
		if mean < c.mean {
			idx = i
			break
		}
	}

	cs.centroids = append(cs.centroids, nil)
	copy(cs.centroids[idx+1:], cs.centroids[idx:])
	cs.centroids[idx] = &centroid{mean, weight}
}

// Add a value to the centroidSet.
func (cs *centroidSet) Add(val float64, weight int) {
	cs.countTotal += weight
	var idx = cs.findAddTarget(val)

	if idx == -1 {
		cs.addNewCentroid(val, weight)
		return
	}

	c := cs.centroids[idx]

	limit := cs.weightLimit(idx)
	// how much weight will we be adding?
	// if adding this node to this centroid would put it over the
	// weight limit, just add the most we can and recur with the remainder
	if c.count+weight > limit {
		add := limit - c.count
		if add < 0 {
			// this node was already overweight
			add = 0
		}
		remainder := weight - add

		c.count += add
		c.mean = c.mean + float64(add)*(val-c.mean)/float64(c.count)

		cs.Add(val, remainder)
	} else {
		c.count += weight
		c.mean = c.mean + float64(weight)*(val-c.mean)/float64(c.count)
	}

}

// returns the approximate quantile that a particular centroid
// represents
func (cs *centroidSet) quantileOf(idx int) float64 {
	total := 0
	for _, c := range cs.centroids[:idx] {
		total += c.count
	}
	return (float64(cs.centroids[idx].count/2) + float64(total)) / float64(cs.countTotal)
}

// Quantile returns the approximate value at a quantile (eg the 99th
// percentile value would be centroidSet.quantileValue(0.99))
func (cs *centroidSet) Quantile(q float64) float64 {
	var n = len(cs.centroids)
	if n == 0 {
		return math.NaN()
	}
	if n == 1 {
		return cs.centroids[0].mean
	}

	if q < 0 {
		q = 0
	} else if q > 1 {
		q = 1
	}

	// rescale into count units instead of 0 to 1 units
	q = float64(cs.countTotal) * q
	// find the first centroid which straddles q
	var (
		qTotal float64 = 0
		i      int
	)
	for i = 0; i < n && float64(cs.centroids[i].count)/2+qTotal < q; i++ {
		qTotal += float64(cs.centroids[i].count)
	}

	if i == 0 {
		// special case 1: the targeted quantile is before the
		// left-most centroid. extrapolate from the slope from
		// centroid0 to centroid1.
		c0 := cs.centroids[0]
		c1 := cs.centroids[1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - float64(c0.count)/2 // this is negative
		return c0.mean + slope*deltaQ
	}
	if i == n {
		// special case 2: the targeted quantile is from the
		// right-most centroid. extrapolate from the slope at the
		// right edge.
		c0 := cs.centroids[n-2]
		c1 := cs.centroids[n-1]
		slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
		deltaQ := q - (qTotal - float64(c1.count)/2)
		return c1.mean + slope*deltaQ
	}
	// common case: targeted quantile is between 2 centroids
	c0 := cs.centroids[i-1]
	c1 := cs.centroids[i]
	slope := (c1.mean - c0.mean) / (float64(c1.count)/2 + float64(c0.count)/2)
	deltaQ := q - (float64(c1.count)/2 + qTotal)
	return c1.mean + slope*deltaQ
}

func (cs *centroidSet) MergeInto(other TDigest) {
	// Add each centroid in cs into other. They should be added in
	// random order.
	addOrder := rand.Perm(len(cs.centroids))
	for _, idx := range addOrder {
		c := cs.centroids[idx]
		// gradually write up the volume written so that the tdigest doesnt overload early
		added := 0
		for i := 1; i < 10; i++ {
			toAdd := i * 2
			if added+i > c.count {
				toAdd = c.count - added
			}
			other.Add(c.mean, toAdd)
			added += toAdd
			if added >= c.count {
				break
			}
		}
		if added < c.count {
			other.Add(c.mean, c.count-added)
		}
		other.Add(c.mean, c.count)
	}
}
<|endoftext|>


<commit_before>/*
Tensile web stress test tool

Mike Hughes 2014
intermernet AT gmail DOT com

LICENSE BSD 3 Clause
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"
)

const version = "0.1"

var (
	url string

	reqs int
	max  int

	wg sync.WaitGroup
)

func init() {
	flag.StringVar(&url, "url", "http://localhost/", "Target URL")
	flag.IntVar(&reqs, "reqs", 50, "Total requests")
	flag.IntVar(&max, "concurrent", 5, "Maximum concurrent requests")
}

type Response struct {
	*http.Response
	err error
}

// Dispatcher
func dispatcher(reqChan chan *http.Request) {
	defer close(reqChan)
	for i := 0; i < reqs; i++ {
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			log.Println(err)
		}
		reqChan <- req
	}
}

// Worker Pool
func workerPool(reqChan chan *http.Request, respChan chan Response) {
	defer close(respChan)
	t := &http.Transport{}
	defer t.CloseIdleConnections()
	for i := 0; i < max; i++ {
		wg.Add(1)
		go worker(t, reqChan, respChan)
	}
	wg.Wait()
}

// Worker
func worker(t *http.Transport, reqChan chan *http.Request, respChan chan Response) {
	defer wg.Done()
	for req := range reqChan {
		resp, err := t.RoundTrip(req)
		r := Response{resp, err}
		respChan <- r
	}
}

// Consumer
func consumer(respChan chan Response) (int64, int64) {
	var (
		conns int64
		size  int64
	)
	for r := range respChan {
		if r.err != nil {
			log.Println(r.err)
		} else {
			size += r.ContentLength
			if err := r.Body.Close(); err != nil {
				log.Println(r.err)
			}
		}
		conns++
	}
	return conns, size
}

func main() {
	flag.Parse()
	fmt.Printf("\n\tTensile web stress test tool v%s\n\n", version)
	flagErr := ""
	if reqs <= 0 {
		flagErr += "ERROR: -reqs must be greater than 0\n"
	}
	if max <= 0 {
		flagErr += "ERROR: -concurrent must be greater than 0\n"
	}
	if flagErr != "" {
		fmt.Printf("%s\n", flagErr)
		os.Exit(0)
	}
	if max > reqs {
		fmt.Println("NOTICE: Concurrent requests is greater than number of requests.")
		fmt.Println("\tChanging concurrent requests to number of requests\n")
		max = reqs
	}
	runtime.GOMAXPROCS(runtime.NumCPU())
	reqChan := make(chan *http.Request)
	respChan := make(chan Response)
	fmt.Printf("Sending %d requests to %s with %d concurrent workers.\n\n", reqs, url, max)
	start := time.Now()
	go dispatcher(reqChan)
	go workerPool(reqChan, respChan)
	fmt.Println("Waiting for replies...\n")
	conns, size := consumer(respChan)
	took := time.Since(start)
	ns := took.Nanoseconds()
	var av int64
	if conns != 0 {
		av = ns / conns
	}
	average, err := time.ParseDuration(fmt.Sprintf("%d", av) + "ns")
	if err != nil {
		log.Println(err)
	}
	fmt.Printf("Connections:\t%d\nConcurrent:\t%d\nTotal size:\t%d bytes\nTotal time:\t%s\nAverage time:\t%s\n", conns, max, size, took, average)
}
<commit_msg>Better flag value checks<commit_after>/*
Tensile web stress test tool

Mike Hughes 2014
intermernet AT gmail DOT com

LICENSE BSD 3 Clause
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"time"
)

const version = "0.1"

var (
	urlStr string

	reqs int
	max  int

	wg sync.WaitGroup

	reqsError   = "ERROR: -reqs must be greater than 0\n"
	maxError    = "ERROR: -concurrent must be greater than 0\n"
	urlError    = "ERROR: Invalid URL\n"
	schemeError = "ERROR: unsupported protocol scheme %s\n"
)

func init() {
	flag.StringVar(&urlStr, "url", "http://localhost/", "Target URL")
	flag.IntVar(&reqs, "reqs", 50, "Total requests")
	flag.IntVar(&max, "concurrent", 5, "Maximum concurrent requests")
}

type Response struct {
	*http.Response
	err error
}

// Dispatcher
func dispatcher(reqChan chan *http.Request) {
	defer close(reqChan)
	for i := 0; i < reqs; i++ {
		req, err := http.NewRequest("GET", urlStr, nil)
		if err != nil {
			log.Println(err)
		}
		reqChan <- req
	}
}

// Worker Pool
func workerPool(reqChan chan *http.Request, respChan chan Response) {
	defer close(respChan)
	t := &http.Transport{}
	defer t.CloseIdleConnections()
	for i := 0; i < max; i++ {
		wg.Add(1)
		go worker(t, reqChan, respChan)
	}
	wg.Wait()
}

// Worker
func worker(t *http.Transport, reqChan chan *http.Request, respChan chan Response) {
	defer wg.Done()
	for req := range reqChan {
		resp, err := t.RoundTrip(req)
		r := Response{resp, err}
		respChan <- r
	}
}

// Consumer
func consumer(respChan chan Response) (int64, int64) {
	var (
		conns int64
		size  int64
	)
	for r := range respChan {
		if r.err != nil {
			log.Println(r.err)
		} else {
			size += r.ContentLength
			if err := r.Body.Close(); err != nil {
				log.Println(r.err)
			}
		}
		conns++
	}
	return conns, size
}

func main() {
	// Flag checks
	flag.Parse()
	fmt.Printf("\n\tTensile web stress test tool v%s\n\n", version)
	flagErr := ""
	if reqs <= 0 {
		flagErr += reqsError
	}
	if max <= 0 {
		flagErr += maxError
	}
	if flagErr != "" {
		log.Fatal(fmt.Errorf("%s", flagErr))
	}
	if urlStr == "" {
		log.Fatal(fmt.Errorf("%s", urlError))
	}
	u, err := url.Parse(urlStr)
	if err != nil {
		log.Print(urlError)
		log.Fatal(fmt.Errorf("%s", err))
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		log.Fatal(fmt.Errorf(schemeError, u.Scheme))
	}
	if max > reqs {
		fmt.Println("NOTICE: Concurrent requests is greater than number of requests.")
		fmt.Println("\tChanging concurrent requests to number of requests\n")
		max = reqs
	}
	// Start
	runtime.GOMAXPROCS(runtime.NumCPU())
	reqChan := make(chan *http.Request)
	respChan := make(chan Response)
	fmt.Printf("Sending %d requests to %s with %d concurrent workers.\n\n", reqs, urlStr, max)
	start := time.Now()
	go dispatcher(reqChan)
	go workerPool(reqChan, respChan)
	fmt.Println("Waiting for replies...\n")
	conns, size := consumer(respChan)
	// Calculate stats
	took := time.Since(start)
	ns := took.Nanoseconds()
	var av int64
	if conns != 0 {
		av = ns / conns
	}
	average, err := time.ParseDuration(fmt.Sprintf("%d", av) + "ns")
	if err != nil {
		log.Println(err)
	}
	fmt.Printf("Connections:\t%d\nConcurrent:\t%d\nTotal size:\t%d bytes\nTotal time:\t%s\nAverage time:\t%s\n", conns, max, size, took, average)
}
<|endoftext|>


<commit_before>/*
Tensile web stress test tool

Mike Hughes 2014
intermernet AT gmail DOT com

LICENSE BSD 3 Clause
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"time"
)

const (
	app     = "Tensile web stress test tool v"
	version = "0.1"
)

var (
	reqs, max, numCPU, maxCPU, numErr, maxErr int

	urlStr, flagErr string
	reqsError       string = "ERROR: -requests (-r) must be greater than 0\n"
	maxError        string = "ERROR: -concurrent (-c) must be greater than 0\n"
	maxErrError     string = "ERROR: -maxerror (-e) must be greater than 0, or -1 for unlimited\n"
	urlError        string = "ERROR: -url (-u) cannot be blank\n"
	schemeError     string = "ERROR: unsupported protocol scheme %s\n"
	ErrLimError     string = "ERROR: maximum error limit reached:\t%d\n"
	ErrTotalError   string = "ERROR: total errors:\t%d\n"
	cpuWarn         string = "NOTICE: -cpu=%d is greater than the number of CPUs on this system\n\tChanging -cpu to %d\n\n"
	maxGTreqsWarn   string = "NOTICE: -concurrent=%d is greater than -requests\n\tChanging -concurrent to %d\n\n"

	wg sync.WaitGroup
)

func init() {
	maxCPU = runtime.NumCPU()
	flag.IntVar(&numCPU, "cpu", 1, "Number of CPUs")
	flag.IntVar(&reqs, "requests", 50, "Total requests")
	flag.IntVar(&reqs, "r", 50, "Total requests (short flag)")
	flag.IntVar(&max, "concurrent", 5, "Maximum concurrent requests")
	flag.IntVar(&max, "c", 5, "Maximum concurrent requests (short flag)")
	flag.IntVar(&maxErr, "maxerror", 1, "Maximum errors before exiting")
	flag.IntVar(&maxErr, "e", 1, "Maximum errors before exiting (short flag)")
	flag.StringVar(&urlStr, "url", "http://localhost/", "Target URL")
	flag.StringVar(&urlStr, "u", "http://localhost/", "Target URL (short flag)")
}

type response struct {
	*http.Response
	err error
}

// Close response Body
func (r *response) closeBody() {
	if err := r.Body.Close(); err != nil {
		log.Println(r.err)
	}
}

// Dispatcher
func dispatcher(reqChan chan *http.Request, quit chan bool) {
	defer close(reqChan)
	for i := 0; i < reqs; i++ {
		req, err := http.NewRequest("GET", urlStr, nil)
		if err != nil {
			log.Println(err)
		}
		select {
		case <-quit:
			return
		default:
			req.Header.Add("User-Agent", app+version)
			reqChan <- req
		}
	}
}

// Worker Pool
func workerPool(reqChan chan *http.Request, respChan chan response, quit chan bool) {
	defer close(respChan)
	t := &http.Transport{}
	defer t.CloseIdleConnections()
	defer wg.Wait()
	for i := 0; i < max; i++ {
		wg.Add(1)
		go worker(t, reqChan, respChan, quit)
	}
}

// Worker
func worker(t *http.Transport, reqChan chan *http.Request, respChan chan response, quit chan bool) {
	defer wg.Done()
	for {
		select {
		case req, ok := <-reqChan:
			if ok {
				resp, err := t.RoundTrip(req)
				respChan <- response{resp, err}
			} else {
				return
			}
		case <-quit:
			return
		}
	}
}

// Kill Workers
func killWorkers(quit chan bool) {
	for {
		select {
		case quit <- true:
		default:
			return
		}
	}
}

// Check maximum error count
func checkMaxErr(quit chan bool) bool {
	chk := false
	numErr++
	if numErr >= maxErr && maxErr != -1 {
		killWorkers(quit)
		log.Printf(ErrLimError, numErr)
		chk = true
	}
	return chk
}

// Consumer
func consumer(respChan chan response, quit chan bool) (int64, int64) {
	defer close(quit)
	var (
		conns, size int64
		prevStatus  int
	)
	for r := range respChan {
		switch {
		case r.err != nil:
			log.Println(r.err)
			if checkMaxErr(quit) {
				return conns, size
			}
		case r.StatusCode >= 400:
			if r.StatusCode != prevStatus {
				log.Printf("ERROR: %s\n", r.Status)
			}
			prevStatus = r.StatusCode
			if checkMaxErr(quit) {
				return conns, size
			}
		default:
			rSize := r.ContentLength
			if rSize >= 0 {
				size += rSize
			}
			conns++
		}
		r.closeBody()
	}
	return conns, size
}

func checkFlags() {
	flag.Parse()
	// Flag Errors
	if reqs <= 0 {
		flagErr += reqsError
	}
	if max <= 0 {
		flagErr += maxError
	}
	if maxErr == 0 || maxErr < -1 {
		flagErr += maxErrError
	}
	if urlStr == "" {
		flagErr += urlError
	}
	u, err := url.Parse(urlStr)
	if err != nil {
		flagErr += err.Error()
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		flagErr += fmt.Sprintf(schemeError, u.Scheme)
	}
	if flagErr != "" {
		log.Fatal(fmt.Errorf("\n%s", flagErr))
	}
	// Flag Warnings
	if numCPU > maxCPU {
		fmt.Printf(cpuWarn, numCPU, maxCPU)
		numCPU = maxCPU
	}
	if max > reqs {
		fmt.Printf(maxGTreqsWarn, max, reqs)
		max = reqs
	}
}

func main() {
	checkFlags()
	fmt.Printf("\n\t%s\n\n", app+version)
	runtime.GOMAXPROCS(numCPU)
	reqChan := make(chan *http.Request)
	respChan := make(chan response)
	quit := make(chan bool, max)
	fmt.Printf("Target URL:\t%s\nRequests:\t%d\nConcurrent:\t%d\nProcessors:\t%d\n\n", urlStr, reqs, max, numCPU)
	start := time.Now()
	go dispatcher(reqChan, quit)
	go workerPool(reqChan, respChan, quit)
	fmt.Printf("Waiting for replies...\n\n")
	conns, size := consumer(respChan, quit)
	if numErr > 0 {
		log.Printf(ErrTotalError, numErr)
	}
	// Calculate stats
	took := time.Since(start)
	tookNS := took.Nanoseconds()
	var averageNS int64
	if conns != 0 {
		averageNS = tookNS / conns
	}
	average, err := time.ParseDuration(fmt.Sprintf("%d", averageNS) + "ns")
	if err != nil {
		log.Println(err)
	}
	sizeHuman := ByteSize(float64(size))
	fmt.Printf("Replies:\t%d\nTotal size:\t%s\nTotal time:\t%s\nAverage time:\t%s\n", conns, sizeHuman, took, average)
}
<commit_msg>Error when -cpu < 1<commit_after>/*
Tensile web stress test tool

Mike Hughes 2014
intermernet AT gmail DOT com

LICENSE BSD 3 Clause
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"time"
)

const (
	app     = "Tensile web stress test tool v"
	version = "0.1"
)

var (
	reqs, max, numCPU, maxCPU, numErr, maxErr int

	urlStr, flagErr string
	reqsError       string = "ERROR: -requests (-r) must be greater than 0\n"
	maxError        string = "ERROR: -concurrent (-c) must be greater than 0\n"
	maxErrError     string = "ERROR: -maxerror (-e) must be greater than 0, or -1 for unlimited\n"
	urlError        string = "ERROR: -url (-u) cannot be blank\n"
	schemeError     string = "ERROR: unsupported protocol scheme %s\n"
	ErrLimError     string = "ERROR: maximum error limit reached: %d\n"
	ErrTotalError   string = "ERROR: total errors: %d\n"
	cpuWarn         string = "NOTICE: -cpu=%d is greater than the number of CPUs on this system\n\tChanging -cpu to %d\n\n"
	cpuLTE0Warn     string = "NOTICE: -cpu=%d is less than 1\n\tChanging -cpu to 1\n\n"
	maxGTreqsWarn   string = "NOTICE: -concurrent=%d is greater than -requests\n\tChanging -concurrent to %d\n\n"

	wg sync.WaitGroup
)

func init() {
	maxCPU = runtime.NumCPU()
	flag.IntVar(&numCPU, "cpu", 1, "Number of CPUs")
	flag.IntVar(&reqs, "requests", 50, "Total requests")
	flag.IntVar(&reqs, "r", 50, "Total requests (short flag)")
	flag.IntVar(&max, "concurrent", 5, "Maximum concurrent requests")
	flag.IntVar(&max, "c", 5, "Maximum concurrent requests (short flag)")
	flag.IntVar(&maxErr, "maxerror", 1, "Maximum errors before exiting")
	flag.IntVar(&maxErr, "e", 1, "Maximum errors before exiting (short flag)")
	flag.StringVar(&urlStr, "url", "http://localhost/", "Target URL")
	flag.StringVar(&urlStr, "u", "http://localhost/", "Target URL (short flag)")
}

type response struct {
	*http.Response
	err error
}

// Close response Body
func (r *response) closeBody() {
	if err := r.Body.Close(); err != nil {
		log.Println(r.err)
	}
}

// Dispatcher
func dispatcher(reqChan chan *http.Request, quit chan bool) {
	defer close(reqChan)
	for i := 0; i < reqs; i++ {
		req, err := http.NewRequest("GET", urlStr, nil)
		if err != nil {
			log.Println(err)
		}
		select {
		case <-quit:
			return
		default:
			req.Header.Add("User-Agent", app+version)
			reqChan <- req
		}
	}
}

// Worker Pool
func workerPool(reqChan chan *http.Request, respChan chan response, quit chan bool) {
	defer close(respChan)
	t := &http.Transport{}
	defer t.CloseIdleConnections()
	defer wg.Wait()
	for i := 0; i < max; i++ {
		wg.Add(1)
		go worker(t, reqChan, respChan, quit)
	}
}

// Worker
func worker(t *http.Transport, reqChan chan *http.Request, respChan chan response, quit chan bool) {
	defer wg.Done()
	for {
		select {
		case req, ok := <-reqChan:
			if ok {
				resp, err := t.RoundTrip(req)
				respChan <- response{resp, err}
			} else {
				return
			}
		case <-quit:
			return
		}
	}
}

// Kill Workers
func killWorkers(quit chan bool) {
	for {
		select {
		case quit <- true:
		default:
			return
		}
	}
}

// Check maximum error count
func checkMaxErr(quit chan bool) bool {
	chk := false
	numErr++
	if numErr >= maxErr && maxErr != -1 {
		killWorkers(quit)
		log.Printf(ErrLimError, numErr)
		chk = true
	}
	return chk
}

// Consumer
func consumer(respChan chan response, quit chan bool) (int64, int64) {
	defer close(quit)
	var (
		conns, size int64
		prevStatus  int
	)
	for r := range respChan {
		switch {
		case r.err != nil:
			log.Println(r.err)
			if checkMaxErr(quit) {
				return conns, size
			}
		case r.StatusCode >= 400:
			if r.StatusCode != prevStatus {
				log.Printf("ERROR: %s\n", r.Status)
			}
			prevStatus = r.StatusCode
			if checkMaxErr(quit) {
				return conns, size
			}
		default:
			rSize := r.ContentLength
			if rSize >= 0 {
				size += rSize
			}
			conns++
		}
		r.closeBody()
	}
	return conns, size
}

func checkFlags() {
	flag.Parse()
	// Flag Errors
	if reqs <= 0 {
		flagErr += reqsError
	}
	if max <= 0 {
		flagErr += maxError
	}
	if maxErr == 0 || maxErr < -1 {
		flagErr += maxErrError
	}
	if urlStr == "" {
		flagErr += urlError
	}
	u, err := url.Parse(urlStr)
	if err != nil {
		flagErr += err.Error()
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		flagErr += fmt.Sprintf(schemeError, u.Scheme)
	}
	if flagErr != "" {
		log.Fatal(fmt.Errorf("\n%s", flagErr))
	}
	// Flag Warnings
	if numCPU > maxCPU {
		fmt.Printf(cpuWarn, numCPU, maxCPU)
		numCPU = maxCPU
	}
	if numCPU < 1 {
		fmt.Printf(cpuLTE0Warn, numCPU)
		numCPU = 1
	}
	if max > reqs {
		fmt.Printf(maxGTreqsWarn, max, reqs)
		max = reqs
	}
}

func main() {
	checkFlags()
	fmt.Printf("\n\t%s\n\n", app+version)
	runtime.GOMAXPROCS(numCPU)
	reqChan := make(chan *http.Request)
	respChan := make(chan response)
	quit := make(chan bool, max)
	fmt.Printf("Target URL:\t%s\nRequests:\t%d\nConcurrent:\t%d\nProcessors:\t%d\n\n", urlStr, reqs, max, numCPU)
	start := time.Now()
	go dispatcher(reqChan, quit)
	go workerPool(reqChan, respChan, quit)
	fmt.Printf("Waiting for replies...\n\n")
	conns, size := consumer(respChan, quit)
	if numErr > 0 {
		log.Printf(ErrTotalError, numErr)
	}
	// Calculate stats
	took := time.Since(start)
	tookNS := took.Nanoseconds()
	var averageNS int64
	if conns != 0 {
		averageNS = tookNS / conns
	}
	average, err := time.ParseDuration(fmt.Sprintf("%d", averageNS) + "ns")
	if err != nil {
		log.Println(err)
	}
	sizeHuman := ByteSize(float64(size))
	fmt.Printf("Replies:\t%d\nTotal size:\t%s\nTotal time:\t%s\nAverage time:\t%s\n", conns, sizeHuman, took, average)
}
<|endoftext|>


<commit_before>package termbox

import "unicode/utf8"
import "errors"
import "bytes"
import "syscall"
import "unsafe"
import "strings"
import "fmt"
import "os"
import "io"

// private API

const (
	t_enter_ca = iota
	t_exit_ca
	t_show_cursor
	t_hide_cursor
	t_clear_screen
	t_sgr
	t_sgr0
	t_underline
	t_bold
	t_blink
	t_move_cursor
	t_enter_keypad
	t_exit_keypad
)

const (
	coord_invalid = -2
	attr_invalid  = Attribute(0xFFFF)
	cursor_hidden = -1
)

var (
	// term specific sequences
	keys  []string
	funcs []string

	// termbox inner state
	orig_tios      syscall.Termios
	back_buffer    cellbuf
	front_buffer   cellbuf
	termw          int
	termh          int
	input_mode     = InputEsc
	out            *os.File
	in             *os.File
	lastfg         = attr_invalid
	lastbg         = attr_invalid
	lastx          = coord_invalid
	lasty          = coord_invalid
	cursor_x       = cursor_hidden
	cursor_y       = cursor_hidden
	foreground     = ColorWhite
	background     = ColorBlack
	inbuf          = make([]byte, 0, 64)
	outbuf         bytes.Buffer
	sigwinch_input = make(chan os.Signal, 10)
	sigwinch_draw  = make(chan os.Signal, 1)
	input_comm     = make(chan []byte)
)

type cellbuf struct {
	width  int
	height int
	cells  []Cell
}

func (this *cellbuf) init(width, height int) {
	this.width = width
	this.height = height
	this.cells = make([]Cell, width*height)
}

func (this *cellbuf) resize(width, height int) {
	if this.width == width && this.height == height {
		return
	}

	oldw := this.width
	oldh := this.height
	oldcells := this.cells

	this.init(width, height)
	this.clear()

	minw, minh := oldw, oldh

	if width < minw {
		minw = width
	}
	if height < minh {
		minh = height
	}

	for i := 0; i < minh; i++ {
		srco, dsto := i*oldw, i*width
		src := oldcells[srco : srco+minw]
		dst := this.cells[dsto : dsto+minw]
		copy(dst, src)
	}
}

func (this *cellbuf) clear() {
	for i := range this.cells {
		c := &this.cells[i]
		c.Ch = ' '
		c.Fg = foreground
		c.Bg = background
	}
}

type winsize struct {
	rows    uint16
	cols    uint16
	xpixels uint16
	ypixels uint16
}

func get_term_size(fd uintptr) (int, int) {
	var sz winsize
	_, _, _ = syscall.Syscall(syscall.SYS_IOCTL,
		fd, uintptr(syscall.TIOCGWINSZ), uintptr(unsafe.Pointer(&sz)))
	return int(sz.cols), int(sz.rows)
}

func send_attr(fg, bg Attribute) {
	// TODO: implement ColorDefault
	if fg != lastfg || bg != lastbg {
		outbuf.WriteString(funcs[t_sgr0])
		// TODO: get rid of fprintf
		fmt.Fprintf(&outbuf, funcs[t_sgr], fg&0x0F, bg&0x0F)
		if fg&AttrBold != 0 {
			outbuf.WriteString(funcs[t_bold])
		}
		if bg&AttrBold != 0 {
			outbuf.WriteString(funcs[t_blink])
		}
		if fg&AttrUnderline != 0 {
			outbuf.WriteString(funcs[t_underline])
		}

		lastfg, lastbg = fg, bg
	}
}

func send_char(x, y int, ch rune) {
	var buf [8]byte
	n := utf8.EncodeRune(buf[:], ch)
	if x-1 != lastx || y != lasty {
		fmt.Fprintf(&outbuf, funcs[t_move_cursor], y+1, x+1) // TODO: get rid of fprintf
	}
	lastx, lasty = x, y
	outbuf.Write(buf[:n])
}

func is_cursor_hidden(x, y int) bool {
	return x == -1 || y == -1
}

func flush() {
	io.Copy(out, &outbuf)
	outbuf.Reset()
}

func send_clear() {
	send_attr(foreground, background)
	outbuf.WriteString(funcs[t_clear_screen])
	if !is_cursor_hidden(cursor_x, cursor_y) {
		fmt.Fprintf(&outbuf, funcs[t_move_cursor], cursor_y+1, cursor_x+1)
	}
	flush()

	// we need to invalidate cursor position too and these two vars are
	// used only for simple cursor positioning optimization, cursor
	// actually may be in the correct place, but we simply discard
	// optimization once and it gives us simple solution for the case when
	// cursor moved
	lastx = coord_invalid
	lasty = coord_invalid
}

func update_size() {
	termw, termh = get_term_size(out.Fd())
	back_buffer.resize(termw, termh)
	front_buffer.resize(termw, termh)
	front_buffer.clear()
	send_clear()
}

func tcsetattr(fd uintptr, termios *syscall.Termios) error {
	r, _, e := syscall.Syscall(syscall.SYS_IOCTL,
		fd, uintptr(syscall.TCSETS), uintptr(unsafe.Pointer(termios)))
	if r != 0 {
		return os.NewSyscallError("SYS_IOCTL", e)
	}
	return nil
}

func tcgetattr(fd uintptr, termios *syscall.Termios) error {
	r, _, e := syscall.Syscall(syscall.SYS_IOCTL,
		fd, uintptr(syscall.TCGETS), uintptr(unsafe.Pointer(termios)))
	if r != 0 {
		return os.NewSyscallError("SYS_IOCTL", e)
	}
	return nil
}

func setup_term() error {
	name := os.Getenv("TERM")
	if name == "" {
		return errors.New("termbox: TERM environment variable not set")
	}

	for _, t := range terms {
		if t.name == name {
			keys = t.keys
			funcs = t.funcs
			return nil
		}
	}

	compat_table := []struct {
		partial string
		keys    []string
		funcs   []string
	}{
		{"xterm", xterm_keys, xterm_funcs},
		{"rxvt", rxvt_unicode_keys, rxvt_unicode_funcs},
		{"linux", linux_keys, linux_funcs},
		{"Eterm", Eterm_keys, Eterm_funcs},
		{"screen", screen_keys, screen_funcs},
		// let's assume that 'cygwin' is xterm compatible
		{"cygwin", xterm_keys, xterm_funcs},
	}

	// try compatibility variants
	for _, it := range compat_table {
		if strings.Contains(name, it.partial) {
			keys = it.keys
			funcs = it.funcs
			return nil
		}
	}

	return errors.New("termbox: unsupported terminal")
}

func parse_escape_sequence(event *Event, buf []byte) int {
	bufstr := string(buf)
	for i, key := range keys {
		if strings.HasPrefix(bufstr, key) {
			event.Ch = 0
			event.Key = Key(0xFFFF - i)
			return len(key)
		}
	}
	return 0
}

func extract_event(event *Event) bool {
	if len(inbuf) == 0 {
		return false
	}

	if inbuf[0] == '\033' {
		// possible escape sequence
		n := parse_escape_sequence(event, inbuf)
		if n != 0 {
			copy(inbuf, inbuf[n:])
			inbuf = inbuf[:len(inbuf)-n]
			return true
		}

		// it's not escape sequence, then it's Alt or Esc, check input_mode
		switch input_mode {
		case InputEsc:
			// if we're in escape mode, fill Esc event, pop buffer, return success
			event.Ch = 0
			event.Key = KeyEsc
			event.Mod = 0
			copy(inbuf, inbuf[1:])
			inbuf = inbuf[:len(inbuf)-1]
			return true
		case InputAlt:
			// if we're in alt mode, set Alt modifier to event and redo parsing
			event.Mod = ModAlt
			copy(inbuf, inbuf[1:])
			inbuf = inbuf[:len(inbuf)-1]
			return extract_event(event)
		default:
			panic("unreachable")
		}
	}

	// if we're here, this is not an escape sequence and not an alt sequence
	// so, it's a FUNCTIONAL KEY or a UNICODE character

	// first of all check if it's a functional key
	if Key(inbuf[0]) <= KeySpace || Key(inbuf[0]) == KeyBackspace2 {
		// fill event, pop buffer, return success
		event.Ch = 0
		event.Key = Key(inbuf[0])
		copy(inbuf, inbuf[1:])
		inbuf = inbuf[:len(inbuf)-1]
		return true
	}

	// the only possible option is utf8 rune
	if r, n := utf8.DecodeRune(inbuf); r != utf8.RuneError {
		event.Ch = r
		event.Key = 0
		copy(inbuf, inbuf[n:])
		inbuf = inbuf[:len(inbuf)-n]
		return true
	}

	return false
}
<commit_msg>Shrink input sigwinch channel buffer to the size of 1.<commit_after>package termbox

import "unicode/utf8"
import "errors"
import "bytes"
import "syscall"
import "unsafe"
import "strings"
import "fmt"
import "os"
import "io"

// private API

const (
	t_enter_ca = iota
	t_exit_ca
	t_show_cursor
	t_hide_cursor
	t_clear_screen
	t_sgr
	t_sgr0
	t_underline
	t_bold
	t_blink
	t_move_cursor
	t_enter_keypad
	t_exit_keypad
)

const (
	coord_invalid = -2
	attr_invalid  = Attribute(0xFFFF)
	cursor_hidden = -1
)

var (
	// term specific sequences
	keys  []string
	funcs []string

	// termbox inner state
	orig_tios      syscall.Termios
	back_buffer    cellbuf
	front_buffer   cellbuf
	termw          int
	termh          int
	input_mode     = InputEsc
	out            *os.File
	in             *os.File
	lastfg         = attr_invalid
	lastbg         = attr_invalid
	lastx          = coord_invalid
	lasty          = coord_invalid
	cursor_x       = cursor_hidden
	cursor_y       = cursor_hidden
	foreground     = ColorWhite
	background     = ColorBlack
	inbuf          = make([]byte, 0, 64)
	outbuf         bytes.Buffer
	sigwinch_input = make(chan os.Signal, 1)
	sigwinch_draw  = make(chan os.Signal, 1)
	input_comm     = make(chan []byte)
)

type cellbuf struct {
	width  int
	height int
	cells  []Cell
}

func (this *cellbuf) init(width, height int) {
	this.width = width
	this.height = height
	this.cells = make([]Cell, width*height)
}

func (this *cellbuf) resize(width, height int) {
	if this.width == width && this.height == height {
		return
	}

	oldw := this.width
	oldh := this.height
	oldcells := this.cells

	this.init(width, height)
	this.clear()

	minw, minh := oldw, oldh

	if width < minw {
		minw = width
	}
	if height < minh {
		minh = height
	}

	for i := 0; i < minh; i++ {
		srco, dsto := i*oldw, i*width
		src := oldcells[srco : srco+minw]
		dst := this.cells[dsto : dsto+minw]
		copy(dst, src)
	}
}

func (this *cellbuf) clear() {
	for i := range this.cells {
		c := &this.cells[i]
		c.Ch = ' '
		c.Fg = foreground
		c.Bg = background
	}
}

type winsize struct {
	rows    uint16
	cols    uint16
	xpixels uint16
	ypixels uint16
}

func get_term_size(fd uintptr) (int, int) {
	var sz winsize
	_, _, _ = syscall.Syscall(syscall.SYS_IOCTL,
		fd, uintptr(syscall.TIOCGWINSZ), uintptr(unsafe.Pointer(&sz)))
	return int(sz.cols), int(sz.rows)
}

func send_attr(fg, bg Attribute) {
	// TODO: implement ColorDefault
	if fg != lastfg || bg != lastbg {
		outbuf.WriteString(funcs[t_sgr0])
		// TODO: get rid of fprintf
		fmt.Fprintf(&outbuf, funcs[t_sgr], fg&0x0F, bg&0x0F)
		if fg&AttrBold != 0 {
			outbuf.WriteString(funcs[t_bold])
		}
		if bg&AttrBold != 0 {
			outbuf.WriteString(funcs[t_blink])
		}
		if fg&AttrUnderline != 0 {
			outbuf.WriteString(funcs[t_underline])
		}

		lastfg, lastbg = fg, bg
	}
}

func send_char(x, y int, ch rune) {
	var buf [8]byte
	n := utf8.EncodeRune(buf[:], ch)
	if x-1 != lastx || y != lasty {
		fmt.Fprintf(&outbuf, funcs[t_move_cursor], y+1, x+1) // TODO: get rid of fprintf
	}
	lastx, lasty = x, y
	outbuf.Write(buf[:n])
}

func is_cursor_hidden(x, y int) bool {
	return x == -1 || y == -1
}

func flush() {
	io.Copy(out, &outbuf)
	outbuf.Reset()
}

func send_clear() {
	send_attr(foreground, background)
	outbuf.WriteString(funcs[t_clear_screen])
	if !is_cursor_hidden(cursor_x, cursor_y) {
		fmt.Fprintf(&outbuf, funcs[t_move_cursor], cursor_y+1, cursor_x+1)
	}
	flush()

	// we need to invalidate cursor position too and these two vars are
	// used only for simple cursor positioning optimization, cursor
	// actually may be in the correct place, but we simply discard
	// optimization once and it gives us simple solution for the case when
	// cursor moved
	lastx = coord_invalid
	lasty = coord_invalid
}

func update_size() {
	termw, termh = get_term_size(out.Fd())
	back_buffer.resize(termw, termh)
	front_buffer.resize(termw, termh)
	front_buffer.clear()
	send_clear()
}

func tcsetattr(fd uintptr, termios *syscall.Termios) error {
	r, _, e := syscall.Syscall(syscall.SYS_IOCTL,
		fd, uintptr(syscall.TCSETS), uintptr(unsafe.Pointer(termios)))
	if r != 0 {
		return os.NewSyscallError("SYS_IOCTL", e)
	}
	return nil
}

func tcgetattr(fd uintptr, termios *syscall.Termios) error {
	r, _, e := syscall.Syscall(syscall.SYS_IOCTL,
		fd, uintptr(syscall.TCGETS), uintptr(unsafe.Pointer(termios)))
	if r != 0 {
		return os.NewSyscallError("SYS_IOCTL", e)
	}
	return nil
}

func setup_term() error {
	name := os.Getenv("TERM")
	if name == "" {
		return errors.New("termbox: TERM environment variable not set")
	}

	for _, t := range terms {
		if t.name == name {
			keys = t.keys
			funcs = t.funcs
			return nil
		}
	}

	compat_table := []struct {
		partial string
		keys    []string
		funcs   []string
	}{
		{"xterm", xterm_keys, xterm_funcs},
		{"rxvt", rxvt_unicode_keys, rxvt_unicode_funcs},
		{"linux", linux_keys, linux_funcs},
		{"Eterm", Eterm_keys, Eterm_funcs},
		{"screen", screen_keys, screen_funcs},
		// let's assume that 'cygwin' is xterm compatible
		{"cygwin", xterm_keys, xterm_funcs},
	}

	// try compatibility variants
	for _, it := range compat_table {
		if strings.Contains(name, it.partial) {
			keys = it.keys
			funcs = it.funcs
			return nil
		}
	}

	return errors.New("termbox: unsupported terminal")
}

func parse_escape_sequence(event *Event, buf []byte) int {
	bufstr := string(buf)
	for i, key := range keys {
		if strings.HasPrefix(bufstr, key) {
			event.Ch = 0
			event.Key = Key(0xFFFF - i)
			return len(key)
		}
	}
	return 0
}

func extract_event(event *Event) bool {
	if len(inbuf) == 0 {
		return false
	}

	if inbuf[0] == '\033' {
		// possible escape sequence
		n := parse_escape_sequence(event, inbuf)
		if n != 0 {
			copy(inbuf, inbuf[n:])
			inbuf = inbuf[:len(inbuf)-n]
			return true
		}

		// it's not escape sequence, then it's Alt or Esc, check input_mode
		switch input_mode {
		case InputEsc:
			// if we're in escape mode, fill Esc event, pop buffer, return success
			event.Ch = 0
			event.Key = KeyEsc
			event.Mod = 0
			copy(inbuf, inbuf[1:])
			inbuf = inbuf[:len(inbuf)-1]
			return true
		case InputAlt:
			// if we're in alt mode, set Alt modifier to event and redo parsing
			event.Mod = ModAlt
			copy(inbuf, inbuf[1:])
			inbuf = inbuf[:len(inbuf)-1]
			return extract_event(event)
		default:
			panic("unreachable")
		}
	}

	// if we're here, this is not an escape sequence and not an alt sequence
	// so, it's a FUNCTIONAL KEY or a UNICODE character

	// first of all check if it's a functional key
	if Key(inbuf[0]) <= KeySpace || Key(inbuf[0]) == KeyBackspace2 {
		// fill event, pop buffer, return success
		event.Ch = 0
		event.Key = Key(inbuf[0])
		copy(inbuf, inbuf[1:])
		inbuf = inbuf[:len(inbuf)-1]
		return true
	}

	// the only possible option is utf8 rune
	if r, n := utf8.DecodeRune(inbuf); r != utf8.RuneError {
		event.Ch = r
		event.Key = 0
		copy(inbuf, inbuf[n:])
		inbuf = inbuf[:len(inbuf)-n]
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"io"
	"net/http"
	"net/http/httputil"
	"os"
	"time"

	"github.com/cleafy/promqueen/model"
	"github.com/sirupsen/logrus"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	debug      = kingpin.Flag("debug", "Enable debug mode.").Bool()
	enableGZIP = kingpin.Flag("gzip", "Enable gzip mode.").Bool()
	interval   = kingpin.Flag("interval", "Timeout waiting for ping.").Default("60s").OverrideDefaultFromEnvar("ACTION_INTERVAL").Short('i').Duration()
	umap       = kingpin.Flag("umap", "stringmap [eg. service.name=http://get.uri:port/uri].").Short('u').StringMap()
	output     = kingpin.Flag("output", "Output file.").Short('o').OverrideDefaultFromEnvar("OUTPUT_FILE").Default("metrics").String()
	Version    = "unversioned"
	filewriter io.WriteCloser
)

func closeIfNotNil(wc io.WriteCloser) {
	if wc != nil {
		if err := wc.Close(); err != nil {
			logrus.Info(err)
		}
	}
}

func writerFor() (io.Writer, error) {
	if _, err := os.Stat(*output); !os.IsNotExist(err) && filewriter != nil {
		return filewriter, nil
	}

	file, err := os.OpenFile(*output, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return nil, err
	}
	closeIfNotNil(filewriter)
	if *enableGZIP {
		filewriter = gzip.NewWriter(file)
	} else {
		filewriter = file
	}
	return filewriter, nil
}

func main() {
	kingpin.Version(Version)
	kingpin.Parse()

	if *debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	if len(*umap) <= 0 {
		kingpin.Usage()
		return
	}

	ticker := time.NewTicker(*interval)

	for range ticker.C {
		for sname, url := range *umap {
			writer, err := writerFor()
			if err != nil {
				logrus.Errorf("writeFor failed with %v", err)
				continue
			}

			resp, err := http.Get(url)
			if err != nil {
				logrus.Errorf("http.Get: %v", err)
				continue
			}
			defer resp.Body.Close()

			dump, err := httputil.DumpResponse(resp, true)
			if err != nil {
				logrus.Errorf("httputil.DumpResponse: %v", err)
				continue
			}

			frame := model.NewFrame(sname, url, dump)

			err = model.WriteFrame(writer, frame)
			if err != nil {
				logrus.Errorf("model.WriteFrame failed with %v", err)
				continue
			}
		}
	}
}
<commit_msg>Added pprof instrumentation<commit_after>package main

import (
	"compress/gzip"
	"io"
	"net/http"
	"net/http/httputil"
	"os"
	"time"
	"net/http/pprof"

	"github.com/cleafy/promqueen/model"
	"github.com/sirupsen/logrus"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	debug      = kingpin.Flag("debug", "Enable debug mode.").Bool()
	enableGZIP = kingpin.Flag("gzip", "Enable gzip mode.").Bool()
	interval   = kingpin.Flag("interval", "Timeout waiting for ping.").Default("60s").OverrideDefaultFromEnvar("ACTION_INTERVAL").Short('i').Duration()
	umap       = kingpin.Flag("umap", "stringmap [eg. service.name=http://get.uri:port/uri].").Short('u').StringMap()
	output     = kingpin.Flag("output", "Output file.").Short('o').OverrideDefaultFromEnvar("OUTPUT_FILE").Default("metrics").String()
	Version    = "unversioned"
	filewriter io.WriteCloser
)

func closeIfNotNil(wc io.WriteCloser) {
	if wc != nil {
		if err := wc.Close(); err != nil {
			logrus.Info(err)
		}
	}
}

func writerFor() (io.Writer, error) {
	if _, err := os.Stat(*output); !os.IsNotExist(err) && filewriter != nil {
		return filewriter, nil
	}

	file, err := os.OpenFile(*output, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return nil, err
	}
	closeIfNotNil(filewriter)
	if *enableGZIP {
		filewriter = gzip.NewWriter(file)
	} else {
		filewriter = file
	}
	return filewriter, nil
}

func main() {
	kingpin.Version(Version)
	kingpin.Parse()

	if *debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	if len(*umap) <= 0 {
		kingpin.Usage()
		return
	}

	ticker := time.NewTicker(*interval)

	for range ticker.C {
		for sname, url := range *umap {
			writer, err := writerFor()
			if err != nil {
				logrus.Errorf("writeFor failed with %v", err)
				continue
			}

			resp, err := http.Get(url)
			if err != nil {
				logrus.Errorf("http.Get: %v", err)
				continue
			}
			defer resp.Body.Close()

			dump, err := httputil.DumpResponse(resp, true)
			if err != nil {
				logrus.Errorf("httputil.DumpResponse: %v", err)
				continue
			}

			frame := model.NewFrame(sname, url, dump)

			err = model.WriteFrame(writer, frame)
			if err != nil {
				logrus.Errorf("model.WriteFrame failed with %v", err)
				continue
			}
		}
	}

	http.ListenAndServe("0.0.0.0:8080", nil)
}
<|endoftext|>


<commit_before>package blocks

import (
	"encoding/json"
	"encoding/xml"
	"github.com/mikedewar/aws4"
	"io/ioutil"
	"log"
	"net/url"
	"strings"
	"time"
)

var (
	AWSSQSAPIVersion    string = "2012-11-05"
	AWSSignatureVersion string = "4"
)

type fromSQSRule struct {
	SQSEndpoint  string
	AccessKey    string
	AccessSecret string
}

type Message struct {
	// this is a list in case I'm ever brave enough to up the "MaxNumberOfMessages" away from 1
	Body          []string `xml:"ReceiveMessageResult>Message>Body"`
	ReceiptHandle []string `xml:"ReceiveMessageResult>Message>ReceiptHandle"`
}

func pollSQS(endpoint string, c aws4.Client, rChan chan Message) {
	var v Message
	query := make(url.Values)
	query.Add("Action", "ReceiveMessage")
	query.Add("AttributeName", "All")
	query.Add("Version", AWSSQSAPIVersion)
	query.Add("SignatureVersion", AWSSignatureVersion)
	query.Add("WaitTimeSeconds", "0")
	query.Add("MaxNumberOfMessages", "10")

	resp, err := c.Get(endpoint + query.Encode())
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println(err.Error())
		return
	}
	err = xml.Unmarshal(body, &v)
	if err != nil {
		log.Println(err.Error())
		return
	}
	rChan <- v
}

func deleteMessage(endpoint string, c aws4.Client, ReceiptHandle string) {
	query := make(url.Values)
	query.Add("Action", "DeleteMessage")
	query.Add("ReceiptHandle", ReceiptHandle)
	query.Add("Version", AWSSQSAPIVersion)
	query.Add("SignatureVersion", AWSSignatureVersion)
	resp, err := c.Get(endpoint + query.Encode())
	if err != nil {
		log.Println(err.Error())
		return
	}
	resp.Body.Close()
}

func unpackAndDelete(body string, reciept string, endpoint string, client aws4.Client, outChan chan BMsg) {
	var SQSmsg map[string]interface{}
	var msg map[string]interface{}
	err := json.Unmarshal([]byte(body), &SQSmsg)
	if err != nil {
		log.Println(err.Error())
		return
	}
	msgString, ok := SQSmsg["Message"].(string)
	if !ok {
		log.Println("couldn't convert SQS Message to string")
		return
	}
	for _, m := range strings.Split(msgString, "\n") {
		if len(m) == 0 {
			continue
		}
		err = json.Unmarshal([]byte(m), &msg)
		if err != nil {
			log.Println(err.Error())
			continue
		}
		out := BMsg{
			Msg: msg,
		}
		outChan <- out
	}
	go deleteMessage(endpoint, client, reciept)
}

// SQSStream hooks into an Amazon SQS, and emits every message it sees into
// streamtools
func SQSStream(b *Block) {
	var rule *fromSQSRule
	var c aws4.Client
	timer := time.NewTimer(1)
	responseChan := make(chan Message)
	broadcastChan := make(chan BMsg)

	for {
		select {
		case <-timer.C:
			if rule == nil {
				timer.Reset(time.Duration(10) * time.Second)
				break
			}
			go pollSQS(rule.SQSEndpoint, c, responseChan)
			timer.Reset(time.Duration(1000) * time.Microsecond)
		case m := <-responseChan:
			if len(m.Body) == 0 {
				timer.Reset(time.Duration(10) * time.Second)
				break
			}
			log.Println("body length", len(m.Body))

			for i, body := range m.Body {
				reciept := m.ReceiptHandle[i]
				go unpackAndDelete(body, reciept, rule.SQSEndpoint, c, broadcastChan)
			}
		case outMsg := <-broadcastChan:
			broadcast(b.OutChans, outMsg)

		case msg := <-b.Routes["set_rule"]:
			if rule == nil {
				rule = &fromSQSRule{}
			}
			unmarshal(msg, rule)
			keys := &aws4.Keys{
				AccessKey: rule.AccessKey,
				SecretKey: rule.AccessSecret,
			}
			c = aws4.Client{Keys: keys}
		case msg := <-b.Routes["get_rule"]:
			if rule == nil {
				marshal(msg, &fromSQSRule{})
			} else {
				marshal(msg, rule)
			}
		case msg := <-b.AddChan:
			updateOutChans(msg, b)
		case <-b.QuitChan:
			quit(b)
			return
		}
	}
}
<commit_msg>deleteBatch<commit_after>package blocks

import (
	"encoding/json"
	"encoding/xml"
	"github.com/mikedewar/aws4"
	"io/ioutil"
	"log"
	"net/url"
	"strings"
	"time"
)

var (
	AWSSQSAPIVersion    string = "2012-11-05"
	AWSSignatureVersion string = "4"
)

type fromSQSRule struct {
	SQSEndpoint  string
	AccessKey    string
	AccessSecret string
}

type Message struct {
	// this is a list in case I'm ever brave enough to up the "MaxNumberOfMessages" away from 1
	Body          []string `xml:"ReceiveMessageResult>Message>Body"`
	ReceiptHandle []string `xml:"ReceiveMessageResult>Message>ReceiptHandle"`
}

func pollSQS(endpoint string, c aws4.Client, rChan chan Message) {
	log.Println("polling start")
	t := time.Now()
	var v Message
	query := make(url.Values)
	query.Add("Action", "ReceiveMessage")
	query.Add("AttributeName", "All")
	query.Add("Version", AWSSQSAPIVersion)
	query.Add("SignatureVersion", AWSSignatureVersion)
	query.Add("WaitTimeSeconds", "0")
	query.Add("MaxNumberOfMessages", "10")

	t1 := time.Now()
	resp, err := c.Get(endpoint + query.Encode())
	log.Println("polling Get", time.Now().Sub(t1).String())
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer resp.Body.Close()
	t1 = time.Now()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println(err.Error())
		return
	}
	log.Println("polling ReadAll", time.Now().Sub(t1).String())
	t1 = time.Now()
	err = xml.Unmarshal(body, &v)
	if err != nil {
		log.Println(err.Error())
		return
	}
	log.Println("polling Unmarshal", time.Now().Sub(t1).String())
	t1 = time.Now()
	rChan <- v
	log.Println("polling block", time.Now().Sub(t1).String())
	log.Println("polling", time.Now().Sub(t).String())
}

func deleteMessage(endpoint string, c aws4.Client, ReceiptHandle string) {
	query := make(url.Values)
	query.Add("Action", "DeleteMessage")
	query.Add("ReceiptHandle", ReceiptHandle)
	query.Add("Version", AWSSQSAPIVersion)
	query.Add("SignatureVersion", AWSSignatureVersion)
	resp, err := c.Get(endpoint + query.Encode())
	if err != nil {
		log.Println(err.Error())
		return
	}
	resp.Body.Close()
}

func deleteBatch(endpoint string, c aws4.Client, receipts []string) {
	log.Println("delete start")
	t := time.Now()
	query := make(url.Values)
	query.Add("Action", "DeleteMessage")
	query.Add("Version", AWSSQSAPIVersion)
	query.Add("SignatureVersion", AWSSignatureVersion)

	for i, r := range receipts {
		query.Add("DeleteMessageBatchRequestEntry.n.Id", "msg"+string(i))
		query.Add("DeleteMessageBatchRequestEntry.n.ReceiptHandle", r)
	}

	t1 := time.Now()
	resp, err := c.Get(endpoint + query.Encode())
	if err != nil {
		log.Println(err.Error())
	}

	log.Println("delete call", time.Now().Sub(t1).String())
	resp.Body.Close()
	log.Println("deleting", time.Now().Sub(t).String())
}

func unpack(body string, reciept string, endpoint string, client aws4.Client, outChan chan BMsg) {
	log.Println("unpack start")
	t := time.Now()
	var SQSmsg map[string]interface{}
	var msg map[string]interface{}
	err := json.Unmarshal([]byte(body), &SQSmsg)
	if err != nil {
		log.Println(err.Error())
		return
	}
	msgString, ok := SQSmsg["Message"].(string)
	if !ok {
		log.Println("couldn't convert SQS Message to string")
		return
	}
	unpackingBlock := time.Duration(0)
	for _, m := range strings.Split(msgString, "\n") {
		if len(m) == 0 {
			continue
		}
		err = json.Unmarshal([]byte(m), &msg)
		if err != nil {
			log.Println(err.Error())
			continue
		}
		out := BMsg{
			Msg: msg,
		}
		t1 := time.Now()
		outChan <- out
		unpackingBlock = unpackingBlock + time.Now().Sub(t1)
	}
	log.Println("unpacking block", unpackingBlock)
	log.Println("unpacking", time.Now().Sub(t).String())
}

// SQSStream hooks into an Amazon SQS, and emits every message it sees into
// streamtools
func SQSStream(b *Block) {
	var rule *fromSQSRule
	var c aws4.Client
	timer := time.NewTimer(1)
	responseChan := make(chan Message)
	broadcastChan := make(chan BMsg)

	for {
		select {
		case <-timer.C:
			if rule == nil {
				timer.Reset(time.Duration(1) * time.Second)
				break
			}
			go pollSQS(rule.SQSEndpoint, c, responseChan)
			timer.Reset(time.Duration(100) * time.Millisecond)
		case m := <-responseChan:
			if len(m.Body) == 0 {
				timer.Reset(time.Duration(10) * time.Second)
				break
			}
			receipts := make([]string, len(m.Body))
			for i, body := range m.Body {
				reciept := m.ReceiptHandle[i]
				go unpack(body, reciept, rule.SQSEndpoint, c, broadcastChan)
				receipts[i] = reciept
			}
			go deleteBatch(rule.SQSEndpoint, c, receipts)
		case outMsg := <-broadcastChan:
			broadcast(b.OutChans, outMsg)

		case msg := <-b.Routes["set_rule"]:
			if rule == nil {
				rule = &fromSQSRule{}
			}
			unmarshal(msg, rule)
			keys := &aws4.Keys{
				AccessKey: rule.AccessKey,
				SecretKey: rule.AccessSecret,
			}
			c = aws4.Client{Keys: keys}
		case msg := <-b.Routes["get_rule"]:
			if rule == nil {
				marshal(msg, &fromSQSRule{})
			} else {
				marshal(msg, rule)
			}
		case msg := <-b.AddChan:
			updateOutChans(msg, b)
		case <-b.QuitChan:
			quit(b)
			return
		}
	}
}
<|endoftext|>


<commit_before>// Package cmount implements a FUSE mounting system for rclone remotes.
//
// This uses the cgo based cgofuse library

// +build cmount
// +build cgo
// +build linux darwin freebsd windows

package cmount

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/billziss-gh/cgofuse/fuse"
	"github.com/pkg/errors"
	"github.com/rclone/rclone/cmd/mountlib"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/vfs"
)

func init() {
	name := "cmount"
	if runtime.GOOS == "windows" {
		name = "mount"
	}
	mountlib.NewMountCommand(name, false, mount)
	mountlib.AddRc("cmount", mount)
}

// mountOptions configures the options from the command line flags
func mountOptions(VFS *vfs.VFS, device string, mountpoint string, opt *mountlib.Options) (options []string) {
	// Options
	options = []string{
		"-o", "fsname=" + device,
		"-o", "subtype=rclone",
		"-o", fmt.Sprintf("max_readahead=%d", opt.MaxReadAhead),
		"-o", fmt.Sprintf("attr_timeout=%g", opt.AttrTimeout.Seconds()),
		// This causes FUSE to supply O_TRUNC with the Open
		// call which is more efficient for cmount.  However
		// it does not work with cgofuse on Windows with
		// WinFSP so cmount must work with or without it.
		"-o", "atomic_o_trunc",
	}
	if opt.DebugFUSE {
		options = append(options, "-o", "debug")
	}

	// OSX options
	if runtime.GOOS == "darwin" {
		if opt.NoAppleDouble {
			options = append(options, "-o", "noappledouble")
		}
		if opt.NoAppleXattr {
			options = append(options, "-o", "noapplexattr")
		}
	}

	// determine if ExtraOptions already has an opt in
	hasOption := func(optionName string) bool {
		optionName += "="
		for _, option := range opt.ExtraOptions {
			if strings.HasPrefix(option, optionName) {
				return true
			}
		}
		return false
	}

	// Windows options
	if runtime.GOOS == "windows" {
		// These cause WinFsp to mean the current user
		if !hasOption("uid") {
			options = append(options, "-o", "uid=-1")
		}
		if !hasOption("gid") {
			options = append(options, "-o", "gid=-1")
		}
		options = append(options, "--FileSystemName=rclone")
	}

	if runtime.GOOS == "darwin" || runtime.GOOS == "windows" {
		if opt.VolumeName != "" {
			options = append(options, "-o", "volname="+opt.VolumeName)
		}
	}
	if opt.AllowNonEmpty {
		options = append(options, "-o", "nonempty")
	}
	if opt.AllowOther {
		options = append(options, "-o", "allow_other")
	}
	if opt.AllowRoot {
		options = append(options, "-o", "allow_root")
	}
	if opt.DefaultPermissions {
		options = append(options, "-o", "default_permissions")
	}
	if VFS.Opt.ReadOnly {
		options = append(options, "-o", "ro")
	}
	if opt.WritebackCache {
		// FIXME? options = append(options, "-o", WritebackCache())
	}
	if opt.DaemonTimeout != 0 {
		options = append(options, "-o", fmt.Sprintf("daemon_timeout=%d", int(opt.DaemonTimeout.Seconds())))
	}
	for _, option := range opt.ExtraOptions {
		options = append(options, "-o", option)
	}
	for _, option := range opt.ExtraFlags {
		options = append(options, option)
	}
	return options
}

// waitFor runs fn() until it returns true or the timeout expires
func waitFor(fn func() bool) (ok bool) {
	const totalWait = 10 * time.Second
	const individualWait = 10 * time.Millisecond
	for i := 0; i < int(totalWait/individualWait); i++ {
		ok = fn()
		if ok {
			return ok
		}
		time.Sleep(individualWait)
	}
	return false
}

// mount the file system
//
// The mount point will be ready when this returns.
//
// returns an error, and an error channel for the serve process to
// report an error when fusermount is called.
func mount(VFS *vfs.VFS, mountpoint string, opt *mountlib.Options) (<-chan error, func() error, error) {
	f := VFS.Fs()
	fs.Debugf(f, "Mounting on %q", mountpoint)

	// Check the mountpoint - in Windows the mountpoint mustn't exist before the mount
	if runtime.GOOS != "windows" {
		fi, err := os.Stat(mountpoint)
		if err != nil {
			return nil, nil, errors.Wrap(err, "mountpoint")
		}
		if !fi.IsDir() {
			return nil, nil, errors.New("mountpoint is not a directory")
		}
	}

	// Create underlying FS
	fsys := NewFS(VFS)
	host := fuse.NewFileSystemHost(fsys)
	host.SetCapReaddirPlus(true) // only works on Windows
	host.SetCapCaseInsensitive(f.Features().CaseInsensitive)

	// Create options
	options := mountOptions(VFS, f.Name()+":"+f.Root(), mountpoint, opt)
	fs.Debugf(f, "Mounting with options: %q", options)

	// Serve the mount point in the background returning error to errChan
	errChan := make(chan error, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				errChan <- errors.Errorf("mount failed: %v", r)
			}
		}()
		var err error
		ok := host.Mount(mountpoint, options)
		if !ok {
			err = errors.New("mount failed")
			fs.Errorf(f, "Mount failed")
		}
		errChan <- err
	}()

	// unmount
	unmount := func() error {
		// Shutdown the VFS
		fsys.VFS.Shutdown()
		fs.Debugf(nil, "Calling host.Unmount")
		if host.Unmount() {
			fs.Debugf(nil, "host.Unmount succeeded")
			if runtime.GOOS == "windows" {
				if !waitFor(func() bool {
					_, err := os.Stat(mountpoint)
					return err != nil
				}) {
					fs.Errorf(nil, "mountpoint %q didn't disappear after unmount - continuing anyway", mountpoint)
				}
			}
			return nil
		}
		fs.Debugf(nil, "host.Unmount failed")
		return errors.New("host unmount failed")
	}

	// Wait for the filesystem to become ready, checking the file
	// system didn't blow up before starting
	select {
	case err := <-errChan:
		err = errors.Wrap(err, "mount stopped before calling Init")
		return nil, nil, err
	case <-fsys.ready:
	}

	// Wait for the mount point to be available on Windows
	// On Windows the Init signal comes slightly before the mount is ready
	if runtime.GOOS == "windows" {
		if !waitFor(func() bool {
			_, err := os.Stat(mountpoint)
			return err == nil
		}) {
			fs.Errorf(nil, "mountpoint %q didn't became available on mount - continuing anyway", mountpoint)
		}
	}

	return errChan, unmount, nil
}
<commit_msg>mount: make mount be cmount under macOS #4393<commit_after>// Package cmount implements a FUSE mounting system for rclone remotes.
//
// This uses the cgo based cgofuse library

// +build cmount
// +build cgo
// +build linux darwin freebsd windows

package cmount

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/billziss-gh/cgofuse/fuse"
	"github.com/pkg/errors"
	"github.com/rclone/rclone/cmd/mountlib"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/vfs"
)

func init() {
	name := "cmount"
	cmountOnly := runtime.GOOS == "windows" || runtime.GOOS == "darwin"
	if cmountOnly {
		name = "mount"
	}
	cmd := mountlib.NewMountCommand(name, false, mount)
	if cmountOnly {
		cmd.Aliases = append(cmd.Aliases, "cmount")
	}
	mountlib.AddRc("cmount", mount)
}

// mountOptions configures the options from the command line flags
func mountOptions(VFS *vfs.VFS, device string, mountpoint string, opt *mountlib.Options) (options []string) {
	// Options
	options = []string{
		"-o", "fsname=" + device,
		"-o", "subtype=rclone",
		"-o", fmt.Sprintf("max_readahead=%d", opt.MaxReadAhead),
		"-o", fmt.Sprintf("attr_timeout=%g", opt.AttrTimeout.Seconds()),
		// This causes FUSE to supply O_TRUNC with the Open
		// call which is more efficient for cmount.  However
		// it does not work with cgofuse on Windows with
		// WinFSP so cmount must work with or without it.
		"-o", "atomic_o_trunc",
	}
	if opt.DebugFUSE {
		options = append(options, "-o", "debug")
	}

	// OSX options
	if runtime.GOOS == "darwin" {
		if opt.NoAppleDouble {
			options = append(options, "-o", "noappledouble")
		}
		if opt.NoAppleXattr {
			options = append(options, "-o", "noapplexattr")
		}
	}

	// determine if ExtraOptions already has an opt in
	hasOption := func(optionName string) bool {
		optionName += "="
		for _, option := range opt.ExtraOptions {
			if strings.HasPrefix(option, optionName) {
				return true
			}
		}
		return false
	}

	// Windows options
	if runtime.GOOS == "windows" {
		// These cause WinFsp to mean the current user
		if !hasOption("uid") {
			options = append(options, "-o", "uid=-1")
		}
		if !hasOption("gid") {
			options = append(options, "-o", "gid=-1")
		}
		options = append(options, "--FileSystemName=rclone")
	}

	if runtime.GOOS == "darwin" || runtime.GOOS == "windows" {
		if opt.VolumeName != "" {
			options = append(options, "-o", "volname="+opt.VolumeName)
		}
	}
	if opt.AllowNonEmpty {
		options = append(options, "-o", "nonempty")
	}
	if opt.AllowOther {
		options = append(options, "-o", "allow_other")
	}
	if opt.AllowRoot {
		options = append(options, "-o", "allow_root")
	}
	if opt.DefaultPermissions {
		options = append(options, "-o", "default_permissions")
	}
	if VFS.Opt.ReadOnly {
		options = append(options, "-o", "ro")
	}
	if opt.WritebackCache {
		// FIXME? options = append(options, "-o", WritebackCache())
	}
	if opt.DaemonTimeout != 0 {
		options = append(options, "-o", fmt.Sprintf("daemon_timeout=%d", int(opt.DaemonTimeout.Seconds())))
	}
	for _, option := range opt.ExtraOptions {
		options = append(options, "-o", option)
	}
	for _, option := range opt.ExtraFlags {
		options = append(options, option)
	}
	return options
}

// waitFor runs fn() until it returns true or the timeout expires
func waitFor(fn func() bool) (ok bool) {
	const totalWait = 10 * time.Second
	const individualWait = 10 * time.Millisecond
	for i := 0; i < int(totalWait/individualWait); i++ {
		ok = fn()
		if ok {
			return ok
		}
		time.Sleep(individualWait)
	}
	return false
}

// mount the file system
//
// The mount point will be ready when this returns.
//
// returns an error, and an error channel for the serve process to
// report an error when fusermount is called.
func mount(VFS *vfs.VFS, mountpoint string, opt *mountlib.Options) (<-chan error, func() error, error) {
	f := VFS.Fs()
	fs.Debugf(f, "Mounting on %q", mountpoint)

	// Check the mountpoint - in Windows the mountpoint mustn't exist before the mount
	if runtime.GOOS != "windows" {
		fi, err := os.Stat(mountpoint)
		if err != nil {
			return nil, nil, errors.Wrap(err, "mountpoint")
		}
		if !fi.IsDir() {
			return nil, nil, errors.New("mountpoint is not a directory")
		}
	}

	// Create underlying FS
	fsys := NewFS(VFS)
	host := fuse.NewFileSystemHost(fsys)
	host.SetCapReaddirPlus(true) // only works on Windows
	host.SetCapCaseInsensitive(f.Features().CaseInsensitive)

	// Create options
	options := mountOptions(VFS, f.Name()+":"+f.Root(), mountpoint, opt)
	fs.Debugf(f, "Mounting with options: %q", options)

	// Serve the mount point in the background returning error to errChan
	errChan := make(chan error, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				errChan <- errors.Errorf("mount failed: %v", r)
			}
		}()
		var err error
		ok := host.Mount(mountpoint, options)
		if !ok {
			err = errors.New("mount failed")
			fs.Errorf(f, "Mount failed")
		}
		errChan <- err
	}()

	// unmount
	unmount := func() error {
		// Shutdown the VFS
		fsys.VFS.Shutdown()
		fs.Debugf(nil, "Calling host.Unmount")
		if host.Unmount() {
			fs.Debugf(nil, "host.Unmount succeeded")
			if runtime.GOOS == "windows" {
				if !waitFor(func() bool {
					_, err := os.Stat(mountpoint)
					return err != nil
				}) {
					fs.Errorf(nil, "mountpoint %q didn't disappear after unmount - continuing anyway", mountpoint)
				}
			}
			return nil
		}
		fs.Debugf(nil, "host.Unmount failed")
		return errors.New("host unmount failed")
	}

	// Wait for the filesystem to become ready, checking the file
	// system didn't blow up before starting
	select {
	case err := <-errChan:
		err = errors.Wrap(err, "mount stopped before calling Init")
		return nil, nil, err
	case <-fsys.ready:
	}

	// Wait for the mount point to be available on Windows
	// On Windows the Init signal comes slightly before the mount is ready
	if runtime.GOOS == "windows" {
		if !waitFor(func() bool {
			_, err := os.Stat(mountpoint)
			return err == nil
		}) {
			fs.Errorf(nil, "mountpoint %q didn't became available on mount - continuing anyway", mountpoint)
		}
	}

	return errChan, unmount, nil
}
<|endoftext|>