<commit_before>package readservice

import (
	"context"
	"errors"

	"github.com/gogo/protobuf/proto"
	"github.com/influxdata/influxdb/v2/kit/tracing"
	"github.com/influxdata/influxdb/v2/models"
	"github.com/influxdata/influxdb/v2/storage/reads"
	"github.com/influxdata/influxdb/v2/storage/reads/datatypes"
	"github.com/influxdata/influxdb/v2/tsdb/cursors"
	"github.com/influxdata/influxql"
)

type store struct {
	viewer    reads.Viewer
	groupCap  GroupCapability
	windowCap WindowAggregateCapability
}

// NewStore creates a store used to query time-series data.
func NewStore(viewer reads.Viewer) reads.Store {
	return &store{
		viewer: viewer,
		groupCap: GroupCapability{
			Count: true,
			Sum:   true,
			First: true,
			Last:  true,
		},
		windowCap: WindowAggregateCapability{
			Count: true,
			Sum:   true,
			First: true,
		},
	}
}

func (s *store) ReadFilter(ctx context.Context, req *datatypes.ReadFilterRequest) (reads.ResultSet, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.ReadSource == nil {
		return nil, tracing.LogError(span, errors.New("missing read source"))
	}

	source, err := getReadSource(*req.ReadSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}

	var cur reads.SeriesCursor
	if cur, err = reads.NewIndexSeriesCursor(ctx, source.GetOrgID(), source.GetBucketID(), req.Predicate, s.viewer); err != nil {
		return nil, tracing.LogError(span, err)
	} else if cur == nil {
		return nil, nil
	}

	return reads.NewFilteredResultSet(ctx, req, cur), nil
}

func (s *store) GetGroupCapability(ctx context.Context) reads.GroupCapability {
	return s.groupCap
}

func (s *store) ReadGroup(ctx context.Context, req *datatypes.ReadGroupRequest) (reads.GroupResultSet, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.ReadSource == nil {
		return nil, tracing.LogError(span, errors.New("missing read source"))
	}

	source, err := getReadSource(*req.ReadSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}

	newCursor := func() (reads.SeriesCursor, error) {
		return reads.NewIndexSeriesCursor(ctx, source.GetOrgID(), source.GetBucketID(), req.Predicate, s.viewer)
	}

	return reads.NewGroupResultSet(ctx, req, newCursor), nil
}

func (s *store) TagKeys(ctx context.Context, req *datatypes.TagKeysRequest) (cursors.StringIterator, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.TagsSource == nil {
		return nil, tracing.LogError(span, errors.New("missing tags source"))
	}

	if req.Range.Start == 0 {
		req.Range.Start = models.MinNanoTime
	}
	if req.Range.End == 0 {
		req.Range.End = models.MaxNanoTime
	}

	var expr influxql.Expr
	var err error
	if root := req.Predicate.GetRoot(); root != nil {
		expr, err = reads.NodeToExpr(root, nil)
		if err != nil {
			return nil, tracing.LogError(span, err)
		}

		if found := reads.HasFieldValueKey(expr); found {
			return nil, tracing.LogError(span, errors.New("field values unsupported"))
		}
		expr = influxql.Reduce(influxql.CloneExpr(expr), nil)
		if reads.IsTrueBooleanLiteral(expr) {
			expr = nil
		}
	}

	readSource, err := getReadSource(*req.TagsSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}
	return s.viewer.TagKeys(ctx, readSource.GetOrgID(), readSource.GetBucketID(), req.Range.Start, req.Range.End, expr)
}

func (s *store) TagValues(ctx context.Context, req *datatypes.TagValuesRequest) (cursors.StringIterator, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.TagsSource == nil {
		return nil, tracing.LogError(span, errors.New("missing tags source"))
	}

	if req.Range.Start == 0 {
		req.Range.Start = models.MinNanoTime
	}
	if req.Range.End == 0 {
		req.Range.End = models.MaxNanoTime
	}

	if req.TagKey == "" {
		return nil, tracing.LogError(span, errors.New("missing tag key"))
	}

	var expr influxql.Expr
	var err error
	if root := req.Predicate.GetRoot(); root != nil {
		expr, err = reads.NodeToExpr(root, nil)
		if err != nil {
			return nil, tracing.LogError(span, err)
		}

		if found := reads.HasFieldValueKey(expr); found {
			return nil, tracing.LogError(span, errors.New("field values unsupported"))
		}
		expr = influxql.Reduce(influxql.CloneExpr(expr), nil)
		if reads.IsTrueBooleanLiteral(expr) {
			expr = nil
		}
	}

	readSource, err := getReadSource(*req.TagsSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}
	return s.viewer.TagValues(ctx, readSource.GetOrgID(), readSource.GetBucketID(), req.TagKey, req.Range.Start, req.Range.End, expr)
}

func (s *store) GetSource(orgID, bucketID uint64) proto.Message {
	return &readSource{
		BucketID:       bucketID,
		OrganizationID: orgID,
	}
}

func (s *store) GetWindowAggregateCapability(ctx context.Context) reads.WindowAggregateCapability {
	return s.windowCap
}

// WindowAggregate will invoke a ReadWindowAggregateRequest against the Store.
func (s *store) WindowAggregate(ctx context.Context, req *datatypes.ReadWindowAggregateRequest) (reads.ResultSet, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.ReadSource == nil {
		return nil, tracing.LogError(span, errors.New("missing read source"))
	}

	source, err := getReadSource(*req.ReadSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}

	var cur reads.SeriesCursor
	if cur, err = reads.NewIndexSeriesCursor(ctx, source.GetOrgID(), source.GetBucketID(), req.Predicate, s.viewer); err != nil {
		return nil, tracing.LogError(span, err)
	} else if cur == nil {
		return nil, nil
	}

	return reads.NewWindowAggregateResultSet(ctx, req, cur)
}

type GroupCapability struct {
	Count bool
	Sum   bool
	First bool
	Last  bool
}

func (c GroupCapability) HaveCount() bool { return c.Count }
func (c GroupCapability) HaveSum() bool   { return c.Sum }
func (c GroupCapability) HaveFirst() bool { return c.First }
func (c GroupCapability) HaveLast() bool  { return c.Last }

type WindowAggregateCapability struct {
	Min   bool
	Max   bool
	Mean  bool
	Count bool
	Sum   bool
	First bool
	Last  bool
}

func (w WindowAggregateCapability) HaveMin() bool   { return w.Min }
func (w WindowAggregateCapability) HaveMax() bool   { return w.Max }
func (w WindowAggregateCapability) HaveMean() bool  { return w.Mean }
func (w WindowAggregateCapability) HaveCount() bool { return w.Count }
func (w WindowAggregateCapability) HaveSum() bool   { return w.Sum }
func (w WindowAggregateCapability) HaveFirst() bool { return w.First }
func (w WindowAggregateCapability) HaveLast() bool  { return w.Last }
<commit_msg>feat(storage): turn on capability for last selection<commit_after>package readservice

import (
	"context"
	"errors"

	"github.com/gogo/protobuf/proto"
	"github.com/influxdata/influxdb/v2/kit/tracing"
	"github.com/influxdata/influxdb/v2/models"
	"github.com/influxdata/influxdb/v2/storage/reads"
	"github.com/influxdata/influxdb/v2/storage/reads/datatypes"
	"github.com/influxdata/influxdb/v2/tsdb/cursors"
	"github.com/influxdata/influxql"
)

type store struct {
	viewer    reads.Viewer
	groupCap  GroupCapability
	windowCap WindowAggregateCapability
}

// NewStore creates a store used to query time-series data.
func NewStore(viewer reads.Viewer) reads.Store {
	return &store{
		viewer: viewer,
		groupCap: GroupCapability{
			Count: true,
			Sum:   true,
			First: true,
			Last:  true,
		},
		windowCap: WindowAggregateCapability{
			Count: true,
			Sum:   true,
			First: true,
			Last:  true,
		},
	}
}

func (s *store) ReadFilter(ctx context.Context, req *datatypes.ReadFilterRequest) (reads.ResultSet, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.ReadSource == nil {
		return nil, tracing.LogError(span, errors.New("missing read source"))
	}

	source, err := getReadSource(*req.ReadSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}

	var cur reads.SeriesCursor
	if cur, err = reads.NewIndexSeriesCursor(ctx, source.GetOrgID(), source.GetBucketID(), req.Predicate, s.viewer); err != nil {
		return nil, tracing.LogError(span, err)
	} else if cur == nil {
		return nil, nil
	}

	return reads.NewFilteredResultSet(ctx, req, cur), nil
}

func (s *store) GetGroupCapability(ctx context.Context) reads.GroupCapability {
	return s.groupCap
}

func (s *store) ReadGroup(ctx context.Context, req *datatypes.ReadGroupRequest) (reads.GroupResultSet, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.ReadSource == nil {
		return nil, tracing.LogError(span, errors.New("missing read source"))
	}

	source, err := getReadSource(*req.ReadSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}

	newCursor := func() (reads.SeriesCursor, error) {
		return reads.NewIndexSeriesCursor(ctx, source.GetOrgID(), source.GetBucketID(), req.Predicate, s.viewer)
	}

	return reads.NewGroupResultSet(ctx, req, newCursor), nil
}

func (s *store) TagKeys(ctx context.Context, req *datatypes.TagKeysRequest) (cursors.StringIterator, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.TagsSource == nil {
		return nil, tracing.LogError(span, errors.New("missing tags source"))
	}

	if req.Range.Start == 0 {
		req.Range.Start = models.MinNanoTime
	}
	if req.Range.End == 0 {
		req.Range.End = models.MaxNanoTime
	}

	var expr influxql.Expr
	var err error
	if root := req.Predicate.GetRoot(); root != nil {
		expr, err = reads.NodeToExpr(root, nil)
		if err != nil {
			return nil, tracing.LogError(span, err)
		}

		if found := reads.HasFieldValueKey(expr); found {
			return nil, tracing.LogError(span, errors.New("field values unsupported"))
		}
		expr = influxql.Reduce(influxql.CloneExpr(expr), nil)
		if reads.IsTrueBooleanLiteral(expr) {
			expr = nil
		}
	}

	readSource, err := getReadSource(*req.TagsSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}
	return s.viewer.TagKeys(ctx, readSource.GetOrgID(), readSource.GetBucketID(), req.Range.Start, req.Range.End, expr)
}

func (s *store) TagValues(ctx context.Context, req *datatypes.TagValuesRequest) (cursors.StringIterator, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.TagsSource == nil {
		return nil, tracing.LogError(span, errors.New("missing tags source"))
	}

	if req.Range.Start == 0 {
		req.Range.Start = models.MinNanoTime
	}
	if req.Range.End == 0 {
		req.Range.End = models.MaxNanoTime
	}

	if req.TagKey == "" {
		return nil, tracing.LogError(span, errors.New("missing tag key"))
	}

	var expr influxql.Expr
	var err error
	if root := req.Predicate.GetRoot(); root != nil {
		expr, err = reads.NodeToExpr(root, nil)
		if err != nil {
			return nil, tracing.LogError(span, err)
		}

		if found := reads.HasFieldValueKey(expr); found {
			return nil, tracing.LogError(span, errors.New("field values unsupported"))
		}
		expr = influxql.Reduce(influxql.CloneExpr(expr), nil)
		if reads.IsTrueBooleanLiteral(expr) {
			expr = nil
		}
	}

	readSource, err := getReadSource(*req.TagsSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}
	return s.viewer.TagValues(ctx, readSource.GetOrgID(), readSource.GetBucketID(), req.TagKey, req.Range.Start, req.Range.End, expr)
}

func (s *store) GetSource(orgID, bucketID uint64) proto.Message {
	return &readSource{
		BucketID:       bucketID,
		OrganizationID: orgID,
	}
}

func (s *store) GetWindowAggregateCapability(ctx context.Context) reads.WindowAggregateCapability {
	return s.windowCap
}

// WindowAggregate will invoke a ReadWindowAggregateRequest against the Store.
func (s *store) WindowAggregate(ctx context.Context, req *datatypes.ReadWindowAggregateRequest) (reads.ResultSet, error) {
	span, ctx := tracing.StartSpanFromContext(ctx)
	defer span.Finish()

	if req.ReadSource == nil {
		return nil, tracing.LogError(span, errors.New("missing read source"))
	}

	source, err := getReadSource(*req.ReadSource)
	if err != nil {
		return nil, tracing.LogError(span, err)
	}

	var cur reads.SeriesCursor
	if cur, err = reads.NewIndexSeriesCursor(ctx, source.GetOrgID(), source.GetBucketID(), req.Predicate, s.viewer); err != nil {
		return nil, tracing.LogError(span, err)
	} else if cur == nil {
		return nil, nil
	}

	return reads.NewWindowAggregateResultSet(ctx, req, cur)
}

type GroupCapability struct {
	Count bool
	Sum   bool
	First bool
	Last  bool
}

func (c GroupCapability) HaveCount() bool { return c.Count }
func (c GroupCapability) HaveSum() bool   { return c.Sum }
func (c GroupCapability) HaveFirst() bool { return c.First }
func (c GroupCapability) HaveLast() bool  { return c.Last }

type WindowAggregateCapability struct {
	Min   bool
	Max   bool
	Mean  bool
	Count bool
	Sum   bool
	First bool
	Last  bool
}

func (w WindowAggregateCapability) HaveMin() bool   { return w.Min }
func (w WindowAggregateCapability) HaveMax() bool   { return w.Max }
func (w WindowAggregateCapability) HaveMean() bool  { return w.Mean }
func (w WindowAggregateCapability) HaveCount() bool { return w.Count }
func (w WindowAggregateCapability) HaveSum() bool   { return w.Sum }
func (w WindowAggregateCapability) HaveFirst() bool { return w.First }
func (w WindowAggregateCapability) HaveLast() bool  { return w.Last }
<|endoftext|>


<commit_before>// Package slack implements the OAuth2 protocol for authenticating users through slack.
// This package can be used as a reference implementation of an OAuth2 provider for Goth.
package slack

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"fmt"
	"github.com/markbates/goth"
	"golang.org/x/oauth2"
)

const (
	authURL         string = "https://slack.com/oauth/authorize"
	tokenURL        string = "https://slack.com/api/oauth.access"
	endpointUser    string = "https://slack.com/api/auth.test"
	endpointProfile string = "https://slack.com/api/users.info"
)

// Provider is the implementation of `goth.Provider` for accessing Slack.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// New creates a new Slack provider and sets up important connection details.
// You should always call `slack.New` to get a new provider.  Never try to
// create one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "slack",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the slack package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Slack for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	return &Session{
		AuthURL: p.config.AuthCodeURL(state),
	}, nil
}

// FetchUser will go to Slack and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	sess := session.(*Session)
	user := goth.User{
		AccessToken:  sess.AccessToken,
		Provider:     p.Name(),
		RefreshToken: sess.RefreshToken,
		ExpiresAt:    sess.ExpiresAt,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	// Get the userID, slack needs userID in order to get user profile info
	response, err := p.Client().Get(endpointUser + "?token=" + url.QueryEscape(sess.AccessToken))
	if err != nil {
		if response != nil {
			response.Body.Close()
		}
		return user, err
	}

	if response.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, response.StatusCode)
	}

	bits, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return user, err
	}

	u := struct {
		UserID string `json:"user_id"`
	}{}

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&u)

	// Get user profile info
	response, err = p.Client().Get(endpointProfile + "?token=" + url.QueryEscape(sess.AccessToken) + "&user=" + u.UserID)
	if err != nil {
		if response != nil {
			response.Body.Close()
		}
		return user, err
	}
	defer response.Body.Close()

	bits, err = ioutil.ReadAll(response.Body)
	if err != nil {
		return user, err
	}

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
	if err != nil {
		return user, err
	}

	err = userFromReader(bytes.NewReader(bits), &user)
	return user, err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	} else {
		c.Scopes = append(c.Scopes, "users:read")
	}
	return c
}

func userFromReader(r io.Reader, user *goth.User) error {
	u := struct {
		User struct {
			NickName string `json:"name"`
			ID       string `json:"id"`
			Profile  struct {
				Email     string `json:"email"`
				Name      string `json:"real_name"`
				AvatarURL string `json:"image_32"`
				FirstName string `json:"first_name"`
				LastName string `json:"last_name"`
			} `json:"profile"`
		} `json:"user"`
	}{}
	err := json.NewDecoder(r).Decode(&u)
	if err != nil {
		return err
	}
	user.Email = u.User.Profile.Email
	user.Name = u.User.Profile.Name
	user.NickName = u.User.NickName
	user.UserID = u.User.ID
	user.AvatarURL = u.User.Profile.AvatarURL
	user.FirstName = u.User.Profile.FirstName
	user.LastName = u.User.Profile.LastName
	return nil
}

//RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return false
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	return nil, nil
}
<commit_msg>Adding comments and fixing formatting, to appease golint<commit_after>// Package slack implements the OAuth2 protocol for authenticating users through slack.
// This package can be used as a reference implementation of an OAuth2 provider for Goth.
package slack

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"fmt"

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
)

// URLs and endpoints
const (
	authURL         string = "https://slack.com/oauth/authorize"
	tokenURL        string = "https://slack.com/api/oauth.access"
	endpointUser    string = "https://slack.com/api/auth.test"
	endpointProfile string = "https://slack.com/api/users.info"
)

// Provider is the implementation of `goth.Provider` for accessing Slack.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// New creates a new Slack provider and sets up important connection details.
// You should always call `slack.New` to get a new provider.  Never try to
// create one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "slack",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

// Client returns the http.Client used in the provider.
func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the slack package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Slack for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	return &Session{
		AuthURL: p.config.AuthCodeURL(state),
	}, nil
}

// FetchUser will go to Slack and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	sess := session.(*Session)
	user := goth.User{
		AccessToken:  sess.AccessToken,
		Provider:     p.Name(),
		RefreshToken: sess.RefreshToken,
		ExpiresAt:    sess.ExpiresAt,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	// Get the userID, slack needs userID in order to get user profile info
	response, err := p.Client().Get(endpointUser + "?token=" + url.QueryEscape(sess.AccessToken))
	if err != nil {
		if response != nil {
			response.Body.Close()
		}
		return user, err
	}

	if response.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, response.StatusCode)
	}

	bits, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return user, err
	}

	u := struct {
		UserID string `json:"user_id"`
	}{}

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&u)

	// Get user profile info
	response, err = p.Client().Get(endpointProfile + "?token=" + url.QueryEscape(sess.AccessToken) + "&user=" + u.UserID)
	if err != nil {
		if response != nil {
			response.Body.Close()
		}
		return user, err
	}
	defer response.Body.Close()

	bits, err = ioutil.ReadAll(response.Body)
	if err != nil {
		return user, err
	}

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
	if err != nil {
		return user, err
	}

	err = userFromReader(bytes.NewReader(bits), &user)
	return user, err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	} else {
		c.Scopes = append(c.Scopes, "users:read")
	}
	return c
}

func userFromReader(r io.Reader, user *goth.User) error {
	u := struct {
		User struct {
			NickName string `json:"name"`
			ID       string `json:"id"`
			Profile  struct {
				Email     string `json:"email"`
				Name      string `json:"real_name"`
				AvatarURL string `json:"image_32"`
				FirstName string `json:"first_name"`
				LastName  string `json:"last_name"`
			} `json:"profile"`
		} `json:"user"`
	}{}
	err := json.NewDecoder(r).Decode(&u)
	if err != nil {
		return err
	}
	user.Email = u.User.Profile.Email
	user.Name = u.User.Profile.Name
	user.NickName = u.User.NickName
	user.UserID = u.User.ID
	user.AvatarURL = u.User.Profile.AvatarURL
	user.FirstName = u.User.Profile.FirstName
	user.LastName = u.User.Profile.LastName
	return nil
}

//RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return false
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>package acceptance_test

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"

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

var _ = Describe("external connectivity", func() {
	var (
		appA                          string
		orgName                       string
		spaceName                     string
		appRoute                      string
		originalRunningSecurityGroups []string
	)

	BeforeEach(func() {
		appA = fmt.Sprintf("appA-%d", rand.Int31())

		Auth(testConfig.TestUser, testConfig.TestUserPassword)

		orgName = "test-org"
		Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

		spaceName = "test-space"
		Expect(cf.Cf("create-space", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName, "-s", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))

		pushApp(appA)
		appRoute = fmt.Sprintf("http://%s.%s/", appA, config.AppsDomain)

		allSecurityGroups := getAllSecurityGroups()
		for _, sg := range allSecurityGroups {
			Expect(cf.Cf("bind-running-security-group", sg).Wait(Timeout_Short)).To(gexec.Exit(0))
		}
		originalRunningSecurityGroups = getRunningSecurityGroups()
	})

	AfterEach(func() {
		appReport(appA, Timeout_Short)

		// clean up everything
		Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))

		By("adding back all the original security groups")
		for _, sg := range originalRunningSecurityGroups {
			Expect(cf.Cf("bind-running-security-group", sg).Wait(Timeout_Short)).To(gexec.Exit(0))
		}
	})

	Describe("accessing the app via the router", func() {
		It("allows external traffic to reach the app", func(done Done) {
			Consistently(func() bool {
				resp, err := http.Get(appRoute)
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(200))
				respBytes, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(respBytes).To(ContainSubstring(`{"ListenAddresses":[`))
				return true
			}, "10s", "1s").Should(BeTrue())

			close(done)
		}, 60 /* <-- overall spec timeout in seconds */)
	})

	Describe("reaching the internet", func() {
		It("can reach the internet only if the application security groups allow it", func() {
			By("checking that the internet is reachable")
			Consistently(func() bool {
				resp, err := http.Get(appRoute + "proxy/example.com")
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(200))
				respBytes, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(respBytes).To(ContainSubstring("Example Domain"))
				return true
			}, "10s", "1s").Should(BeTrue())

			By("removing all the original security groups")
			for _, sg := range originalRunningSecurityGroups {
				Expect(cf.Cf("unbind-running-security-group", sg).Wait(Timeout_Short)).To(gexec.Exit(0))
			}

			By("restarting the app")
			Expect(cf.Cf("restart", appA).Wait(Timeout_Push)).To(gexec.Exit(0))

			By("checking that the connection fails")
			Consistently(func() bool {
				resp, err := http.Get(appRoute + "proxy/example.com")
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(500))
				respBytes, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(respBytes).To(ContainSubstring("example.com"))
				return true
			}, "10s", "1s").Should(BeTrue())
		})
	})
})

func getRunningSecurityGroups() []string {
	session := cf.Cf("running-security-groups")
	Expect(session.Wait(Timeout_Short)).To(gexec.Exit(0))

	candidateGroups := strings.Split(string(session.Out.Contents()), "\n")[3:]
	actualGroups := []string{}
	for _, l := range candidateGroups {
		trimmed := strings.TrimSpace(l)
		if trimmed != "" {
			actualGroups = append(actualGroups, trimmed)
		}
	}
	return actualGroups
}

func getAllSecurityGroups() []string {
	session := cf.Cf("security-groups")
	Expect(session.Wait(Timeout_Short)).To(gexec.Exit(0))

	candidateGroups := strings.Split(string(session.Out.Contents()), "\n")[4:]
	actualGroups := []string{}
	for _, l := range candidateGroups {
		fields := strings.Fields(l)
		if len(fields) < 2 {
			continue
		}
		trimmed := strings.TrimSpace(fields[1])
		if trimmed != "" {
			actualGroups = append(actualGroups, trimmed)
		}
	}
	return actualGroups
}
<commit_msg>Compress external connectivity test from two specs to just one<commit_after>package acceptance_test

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"

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

var _ = Describe("external connectivity", func() {
	var (
		appA                          string
		orgName                       string
		spaceName                     string
		appRoute                      string
		originalRunningSecurityGroups []string
	)

	BeforeEach(func() {
		appA = fmt.Sprintf("appA-%d", rand.Int31())

		Auth(testConfig.TestUser, testConfig.TestUserPassword)

		orgName = "test-org"
		Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

		spaceName = "test-space"
		Expect(cf.Cf("create-space", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName, "-s", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))

		pushApp(appA)
		appRoute = fmt.Sprintf("http://%s.%s/", appA, config.AppsDomain)

		allSecurityGroups := getAllSecurityGroups()
		for _, sg := range allSecurityGroups {
			Expect(cf.Cf("bind-running-security-group", sg).Wait(Timeout_Short)).To(gexec.Exit(0))
		}
		originalRunningSecurityGroups = getRunningSecurityGroups()
	})

	AfterEach(func() {
		appReport(appA, Timeout_Short)

		// clean up everything
		Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))

		By("adding back all the original security groups")
		for _, sg := range originalRunningSecurityGroups {
			Expect(cf.Cf("bind-running-security-group", sg).Wait(Timeout_Short)).To(gexec.Exit(0))
		}
	})

	Describe("basic (legacy) network behavior for an app", func() {
		It("is reachable from the router, and can reach the internet only if allowed", func(done Done) {
			By("checking that the app is reachable via the router")
			Consistently(func() bool {
				resp, err := http.Get(appRoute)
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(200))
				respBytes, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(respBytes).To(ContainSubstring(`{"ListenAddresses":[`))
				return true
			}, "10s", "1s").Should(BeTrue())

			By("checking that it can reach the internet")
			Consistently(func() bool {
				resp, err := http.Get(appRoute + "proxy/example.com")
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(200))
				respBytes, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(respBytes).To(ContainSubstring("Example Domain"))
				return true
			}, "10s", "1s").Should(BeTrue())

			By("removing all the original security groups")
			for _, sg := range originalRunningSecurityGroups {
				Expect(cf.Cf("unbind-running-security-group", sg).Wait(Timeout_Short)).To(gexec.Exit(0))
			}

			By("restarting the app")
			Expect(cf.Cf("restart", appA).Wait(Timeout_Push)).To(gexec.Exit(0))

			By("checking that the app cannot reach the internet")
			Consistently(func() bool {
				resp, err := http.Get(appRoute + "proxy/example.com")
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(500))
				respBytes, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(respBytes).To(ContainSubstring("example.com"))
				return true
			}, "10s", "1s").Should(BeTrue())

			close(done)
		}, 60 /* <-- overall spec timeout in seconds */)
	})
})

func getRunningSecurityGroups() []string {
	session := cf.Cf("running-security-groups")
	Expect(session.Wait(Timeout_Short)).To(gexec.Exit(0))

	candidateGroups := strings.Split(string(session.Out.Contents()), "\n")[3:]
	actualGroups := []string{}
	for _, l := range candidateGroups {
		trimmed := strings.TrimSpace(l)
		if trimmed != "" {
			actualGroups = append(actualGroups, trimmed)
		}
	}
	return actualGroups
}

func getAllSecurityGroups() []string {
	session := cf.Cf("security-groups")
	Expect(session.Wait(Timeout_Short)).To(gexec.Exit(0))

	candidateGroups := strings.Split(string(session.Out.Contents()), "\n")[4:]
	actualGroups := []string{}
	for _, l := range candidateGroups {
		fields := strings.Fields(l)
		if len(fields) < 2 {
			continue
		}
		trimmed := strings.TrimSpace(fields[1])
		if trimmed != "" {
			actualGroups = append(actualGroups, trimmed)
		}
	}
	return actualGroups
}
<|endoftext|>


<commit_before>package configloader

import (
	"bytes"
	"os"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/iron-io/titan/runner/agent"
	drivercommon "github.com/iron-io/titan/runner/drivers/common"
	"github.com/pivotal-golang/bytefmt"
	"github.com/spf13/viper"
)

// This package provides a function to read a complete runner configuration
// from various sources (currently JSON file and environment variables). Just
// importing the configloader package will cause it to read in configuration
// from the sources. Callers may then use RunnerConfiguration() to get an
// initialized runner configuration. Other methods provide access to some common functionality.
// A basic configuration looks like:
//
//     {
//       "runner": {
//         ...
//       },
//       "stats": {
//         ...
//       },
//       "log_level": "..."
//     }
//
// Callers can use LogLevel() to obtain the log level.
// configloader will also set the runner's log package's log level to the value obtained here.
// StatsConfig() returns a stats configuration which can be passed to titan
// stats.New(). The runner's stats are not initialized using this
// configuration. Use the environment passed to agent.NewRunner() to handle
// this for now.
//
// For more configuration sub-trees related to the app specific behaviour that
// wraps a runner, use ExtendedConfiguration(key) to obtain a configuration
// sub-tree.
//
// This package is currently shared by the Titan runner and Swapi runner. Be
// careful when making changes to preserve compatibility.
//
// The following top-level keys are reserved for runner configuration, since
// they reflect the names of the environment variables captured in the process:
// - log_level
// - driver
// - concurrency
// - memory_per_job
//
// configloader will terminate if any configuration does not match valid values.
// For this reason, it should only be used from a binary's main().

var runnerConfig agent.Config
var logLevel string
var apiURL string

func init() {
	// Set up defaults.
	viper.SetConfigName("config")
	viper.AddConfigPath(".")
	replacer := strings.NewReplacer(".", "_")
	viper.SetEnvKeyReplacer(replacer)
	viper.AutomaticEnv()
	viper.SetDefault("log_level", "info")
	viper.SetDefault("driver", "docker")
	viper.SetDefault("api_url", "http://localhost:8080")
	viper.SetDefault("memory_per_job", "256M")
	// If we do not set it here, viper will not load the relevant environment var. Set it to int's default value.
	viper.SetDefault("concurrency", "0")

	// Initialize viper from config.json.
	err := viper.ReadInConfig()
	if err != nil {
		if _, ok := err.(viper.UnsupportedConfigError); ok {
			logrus.WithError(err).Error("Couldn't read config file, this is fine, it's not required")
			// ignore
		} else {
			logrus.WithError(err).Fatal("Error reading config file")
		}
	}

	// Initialize viper from CONFIG environment variable.
	if os.Getenv("CONFIG") != "" {
		viper.SetConfigType("json")
		err = viper.ReadConfig(bytes.NewBufferString(os.Getenv("CONFIG")))
		if err != nil {
			logrus.WithError(err).Fatal("Error reading CONFIG from env")
		}
	}

	// Set up defaults.
	dconfig := &drivercommon.Config{}
	dconfig.Defaults()
	memPerJob, err := bytefmt.ToBytes(viper.GetString("memory_per_job"))
	if err != nil {
		logrus.WithError(err).WithFields(logrus.Fields{"memory_per_job": viper.GetString("memory_per_job")}).Fatal("Invalid MEMORY_PER_JOB variable")
	}
	// todo: we should pass the entire config to the driver so they can use the same variable. Or make a config interface and pass it along.
	dconfig.Memory = memPerJob
	runnerConfig.DriverConfig = dconfig

	// Let loaded configuration overwrite defaults.
	if viper.InConfig("runner") {
		err = viper.Sub("runner").Unmarshal(&runnerConfig)
		if err != nil {
			logrus.WithError(err).Fatal("could not unmarshal runner configuration from config")
		}
	}

	// Finally, let top-level reserved keys overwrite any existing settings.
	runnerConfig.Driver = viper.GetString("driver")

	if concurrency := viper.GetInt("concurrency"); concurrency > 0 {
		runnerConfig.Concurrency = concurrency
	}

	apiURL = viper.GetString("api_url")

	logrus.WithFields(logrus.Fields{"runner_configuration": runnerConfig}).Info("configloader runner configuration")
}

func RunnerConfiguration() *agent.Config {
	return &runnerConfig
}

func LogLevel() string {
	return logLevel
}

func ApiURL() string {
	return apiURL
}

func ExtendedConfiguration(key string) *viper.Viper {
	if !viper.InConfig(key) {
		return nil
	}

	return viper.Sub(key)
}
<commit_msg>are you not entertained?<commit_after>package configloader

import (
	"bytes"
	"os"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/iron-io/titan/runner/agent"
	drivercommon "github.com/iron-io/titan/runner/drivers/common"
	"github.com/pivotal-golang/bytefmt"
	"github.com/spf13/viper"
)

// This package provides a function to read a complete runner configuration
// from various sources (currently JSON file and environment variables). Just
// importing the configloader package will cause it to read in configuration
// from the sources. Callers may then use RunnerConfiguration() to get an
// initialized runner configuration. Other methods provide access to some common functionality.
// A basic configuration looks like:
//
//     {
//       "runner": {
//         ...
//       },
//       "stats": {
//         ...
//       },
//       "log_level": "..."
//     }
//
// Callers can use LogLevel() to obtain the log level.
// configloader will also set the runner's log package's log level to the value obtained here.
// StatsConfig() returns a stats configuration which can be passed to titan
// stats.New(). The runner's stats are not initialized using this
// configuration. Use the environment passed to agent.NewRunner() to handle
// this for now.
//
// For more configuration sub-trees related to the app specific behaviour that
// wraps a runner, use ExtendedConfiguration(key) to obtain a configuration
// sub-tree.
//
// This package is currently shared by the Titan runner and Swapi runner. Be
// careful when making changes to preserve compatibility.
//
// The following top-level keys are reserved for runner configuration, since
// they reflect the names of the environment variables captured in the process:
// - log_level
// - driver
// - concurrency
// - memory_per_job
//
// configloader will terminate if any configuration does not match valid values.
// For this reason, it should only be used from a binary's main().

var runnerConfig agent.Config
var logLevel string
var apiURL string

func init() {
	// Set up defaults.
	viper.SetConfigName("config")
	viper.AddConfigPath(".")
	replacer := strings.NewReplacer(".", "_")
	viper.SetEnvKeyReplacer(replacer)
	viper.AutomaticEnv()
	viper.SetDefault("log_level", "info")
	viper.SetDefault("driver", "docker")
	viper.SetDefault("api_url", "http://localhost:8080")
	viper.SetDefault("memory_per_job", "256M")
	// If we do not set it here, viper will not load the relevant environment var. Set it to int's default value.
	viper.SetDefault("concurrency", "0")

	// Initialize viper from config.json.
	err := viper.ReadInConfig()
	if err != nil {
		if _, ok := err.(viper.UnsupportedConfigError); ok {
			logrus.WithError(err).Error("Couldn't read config file, this is fine, it's not required")
			// ignore
		} else {
			logrus.WithError(err).Fatal("Error reading config file")
		}
	}

	// Initialize viper from CONFIG environment variable.
	if os.Getenv("CONFIG") != "" {
		viper.SetConfigType("json")
		err = viper.ReadConfig(bytes.NewBufferString(os.Getenv("CONFIG")))
		if err != nil {
			logrus.WithError(err).Fatal("Error reading CONFIG from env")
		}
	}

	// Set up defaults.
	dconfig := &drivercommon.Config{}
	dconfig.Defaults()
	memPerJob, err := bytefmt.ToBytes(viper.GetString("memory_per_job"))
	if err != nil {
		logrus.WithError(err).WithFields(logrus.Fields{"memory_per_job": viper.GetString("memory_per_job")}).Fatal("Invalid MEMORY_PER_JOB variable")
	}
	// todo: we should pass the entire config to the driver so they can use the same variable. Or make a config interface and pass it along.
	dconfig.Memory = memPerJob
	runnerConfig.DriverConfig = dconfig

	// Let loaded configuration overwrite defaults.
	if viper.InConfig("runner") {
		err = viper.Sub("runner").Unmarshal(&runnerConfig)
		if err != nil {
			logrus.WithError(err).Fatal("could not unmarshal runner configuration from config")
		}
	}

	// Finally, let top-level reserved keys overwrite any existing settings.
	runnerConfig.Driver = viper.GetString("driver")

	if concurrency := viper.GetInt("concurrency"); concurrency > 0 {
		runnerConfig.Concurrency = concurrency
	}

	apiURL = viper.GetString("api_url")
	logLevel = viper.GetString("log_level")

	{
		logrus.WithFields(logrus.Fields{"level": logLevel}).Info("configloader setting log level")
		level, err := logrus.ParseLevel(logLevel)
		if err != nil {
			logrus.WithFields(logrus.Fields{"level": logLevel}).Warn("Could not parse log level, setting to info")
			level = logrus.InfoLevel
		}
	}

	logrus.WithFields(logrus.Fields{"runner_configuration": runnerConfig}).Info("configloader runner configuration")
}

func RunnerConfiguration() *agent.Config {
	return &runnerConfig
}

func LogLevel() string {
	return logLevel
}

func ApiURL() string {
	return apiURL
}

func ExtendedConfiguration(key string) *viper.Viper {
	if !viper.InConfig(key) {
		return nil
	}

	return viper.Sub(key)
}
<|endoftext|>


<commit_before>package forwarded

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/herb-go/herb/middleware"
)

const xForwardedForHeader = "X-Forwarded-For"
const xForwardedHostHeader = "X-Forwarded-Host"
const xForwardedProtoHeader = "X-Forwarded-Proto"

var s *httptest.Server

func doRequest(t *testing.T, header http.Header) (map[string]string, int) {
	var url = "http://" + s.Listener.Addr().String() + "/"
	var req, err = http.NewRequest("GET", url, nil)
	if err != nil {
		t.Fatal(err)
	}
	for k, v := range header {
		for _, value := range v {
			req.Header.Add(k, value)
		}
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, resp.StatusCode
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var data = map[string]string{}
	err = json.Unmarshal(body, &data)
	if err != nil {
		t.Fatal(err)
	}
	return data, resp.StatusCode
}
func TestMiddleware(t *testing.T) {
	var app = middleware.New()
	var data map[string]string
	var statusCode int
	var middleware = Middleware{}
	app.
		Use(middleware.ServeMiddleware).
		HandleFunc(func(w http.ResponseWriter, r *http.Request) {
			var result = map[string]string{}
			result["RemoteAddr"] = r.RemoteAddr
			result["Host"] = r.Host
			result["Scheme"] = r.URL.Scheme
			var data, err = json.Marshal(result)
			if err != nil {
				panic(err)
			}
			_, err = w.Write(data)
			if err != nil {
				panic(err)
			}
		})
	s = httptest.NewServer(app)
	defer s.Close()
	var host = s.Listener.Addr().String()
	var addr = "127.0.0.1"
	var scheme = ""
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
	}
	if strings.Split(data["RemoteAddr"], ":")[0] != addr {
		t.Error(data["RemoteAddr"])
	}
	if data["Host"] != host {
		t.Error(data["Host"])
	}
	if data["scheme"] != scheme {
		t.Error(data["scheme"])
	}

	middleware.Status = 1
	data, statusCode = doRequest(t, nil)

	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, nil)
	if strings.Split(data["RemoteAddr"], ":")[0] != addr {
		t.Error(data["RemoteAddr"])
	}
	if data["Host"] != host {
		t.Error(data["Host"])
	}
	if data["scheme"] != scheme {
		t.Error(data["scheme"])
	}

	middleware.ForwardedForHeader = xForwardedForHeader
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if strings.Split(data["RemoteAddr"], ":")[0] != addr {
			t.Error(data["RemoteAddr"])
		}
	}
	var forwardedIP = "1234567"
	data, statusCode = doRequest(t, map[string][]string{
		xForwardedForHeader: []string{forwardedIP},
	})
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if strings.Split(data["RemoteAddr"], ":")[0] != forwardedIP {
			t.Error(data["RemoteAddr"])
		}
	}

	middleware.ForwardedHostHeader = xForwardedHostHeader
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Host"] != host {
			t.Error(data["Host"])
		}
	}
	var forwardedHost = "github.com"
	data, statusCode = doRequest(t, map[string][]string{
		xForwardedHostHeader: []string{forwardedHost},
	})
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Host"] != forwardedHost {
			t.Error(data["Host"])
		}
	}

	middleware.ForwardedProtoHeader = xForwardedProtoHeader
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Scheme"] != "" {
			t.Error(data["Scheme"])
		}
	}
	var forwardedScheme = "https"
	data, statusCode = doRequest(t, map[string][]string{
		xForwardedHostHeader: []string{forwardedHost},
	})
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Scheme"] != forwardedScheme {
			t.Error(data["Scheme"])
		}
	}
	var forwardedToken = "X-Forwarded-Token"
	var forwardedTokenWrong = "wrongtoken"
	var forwardedValue = "value"
	var forwardedValueWrong = "wrongvalue"

	middleware.ForwardedTokenHeader = forwardedToken
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusBadRequest {
		t.Fatal(statusCode)
	}

	data, statusCode = doRequest(t, map[string][]string{
		forwardedTokenWrong: []string{forwardedValue},
	})
	middleware.ForwardedTokenValue = forwardedValue

	if statusCode != http.StatusBadRequest {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, map[string][]string{
		forwardedToken: []string{forwardedValueWrong},
	})
	if statusCode != http.StatusBadRequest {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, map[string][]string{
		forwardedToken: []string{forwardedValue},
	})
	if statusCode != http.StatusOK {
		t.Error(middleware)
		t.Fatal(statusCode)
	}

	middleware.Status = 0
	data, statusCode = doRequest(t, nil)

	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, nil)
	if strings.Split(data["RemoteAddr"], ":")[0] != addr {
		t.Error(data["RemoteAddr"])
	}
	if data["Host"] != host {
		t.Error(data["Host"])
	}
	if data["scheme"] != scheme {
		t.Error(data["scheme"])
	}
}
<commit_msg>add set token fail action<commit_after>package forwarded

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/herb-go/herb/middleware"
)

const xForwardedForHeader = "X-Forwarded-For"
const xForwardedHostHeader = "X-Forwarded-Host"
const xForwardedProtoHeader = "X-Forwarded-Proto"

var s *httptest.Server

func doRequest(t *testing.T, header http.Header) (map[string]string, int) {
	var url = "http://" + s.Listener.Addr().String() + "/"
	var req, err = http.NewRequest("GET", url, nil)
	if err != nil {
		t.Fatal(err)
	}
	for k, v := range header {
		for _, value := range v {
			req.Header.Add(k, value)
		}
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, resp.StatusCode
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var data = map[string]string{}
	err = json.Unmarshal(body, &data)
	if err != nil {
		t.Fatal(err)
	}
	return data, resp.StatusCode
}
func TestMiddleware(t *testing.T) {
	var app = middleware.New()
	var data map[string]string
	var statusCode int
	var middleware = Middleware{}
	app.
		Use(middleware.ServeMiddleware).
		HandleFunc(func(w http.ResponseWriter, r *http.Request) {
			var result = map[string]string{}
			result["RemoteAddr"] = r.RemoteAddr
			result["Host"] = r.Host
			result["Scheme"] = r.URL.Scheme
			var data, err = json.Marshal(result)
			if err != nil {
				panic(err)
			}
			_, err = w.Write(data)
			if err != nil {
				panic(err)
			}
		})
	s = httptest.NewServer(app)
	defer s.Close()
	var host = s.Listener.Addr().String()
	var addr = "127.0.0.1"
	var scheme = ""
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
	}
	if strings.Split(data["RemoteAddr"], ":")[0] != addr {
		t.Error(data["RemoteAddr"])
	}
	if data["Host"] != host {
		t.Error(data["Host"])
	}
	if data["scheme"] != scheme {
		t.Error(data["scheme"])
	}

	middleware.Status = 1
	data, statusCode = doRequest(t, nil)

	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, nil)
	if strings.Split(data["RemoteAddr"], ":")[0] != addr {
		t.Error(data["RemoteAddr"])
	}
	if data["Host"] != host {
		t.Error(data["Host"])
	}
	if data["scheme"] != scheme {
		t.Error(data["scheme"])
	}

	middleware.ForwardedForHeader = xForwardedForHeader
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if strings.Split(data["RemoteAddr"], ":")[0] != addr {
			t.Error(data["RemoteAddr"])
		}
	}
	var forwardedIP = "1234567"
	data, statusCode = doRequest(t, map[string][]string{
		xForwardedForHeader: []string{forwardedIP},
	})
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if strings.Split(data["RemoteAddr"], ":")[0] != forwardedIP {
			t.Error(data["RemoteAddr"])
		}
	}

	middleware.ForwardedHostHeader = xForwardedHostHeader
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Host"] != host {
			t.Error(data["Host"])
		}
	}
	var forwardedHost = "github.com"
	data, statusCode = doRequest(t, map[string][]string{
		xForwardedHostHeader: []string{forwardedHost},
	})
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Host"] != forwardedHost {
			t.Error(data["Host"])
		}
	}

	middleware.ForwardedProtoHeader = xForwardedProtoHeader
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Scheme"] != "" {
			t.Error(data["Scheme"])
		}
	}
	var forwardedScheme = "https"
	data, statusCode = doRequest(t, map[string][]string{
		xForwardedHostHeader: []string{forwardedHost},
	})
	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
		if data["Scheme"] != forwardedScheme {
			t.Error(data["Scheme"])
		}
	}
	var forwardedToken = "X-Forwarded-Token"
	var forwardedTokenWrong = "wrongtoken"
	var forwardedValue = "value"
	var forwardedValueWrong = "wrongvalue"

	middleware.ForwardedTokenHeader = forwardedToken
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusBadRequest {
		t.Fatal(statusCode)
	}

	data, statusCode = doRequest(t, map[string][]string{
		forwardedTokenWrong: []string{forwardedValue},
	})
	middleware.ForwardedTokenValue = forwardedValue

	if statusCode != http.StatusBadRequest {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, map[string][]string{
		forwardedToken: []string{forwardedValueWrong},
	})
	if statusCode != http.StatusBadRequest {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, map[string][]string{
		forwardedToken: []string{forwardedValue},
	})
	if statusCode != http.StatusOK {
		t.Error(middleware)
		t.Fatal(statusCode)
	}

	middleware.SetTokenFailedAction(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, http.StatusText(http.StatusForbidden), http.StatusForbidden)
	})
	data, statusCode = doRequest(t, nil)
	if statusCode != http.StatusForbidden {
		t.Fatal(statusCode)
	}
	middleware.Status = 0
	data, statusCode = doRequest(t, nil)

	if statusCode != http.StatusOK {
		t.Fatal(statusCode)
	}
	data, statusCode = doRequest(t, nil)
	if strings.Split(data["RemoteAddr"], ":")[0] != addr {
		t.Error(data["RemoteAddr"])
	}
	if data["Host"] != host {
		t.Error(data["Host"])
	}
	if data["scheme"] != scheme {
		t.Error(data["scheme"])
	}
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"
	"sync"
	"time"
)

// TombstoneGC is used to track creation of tombstones
// so that they can be garbage collected after their TTL
// expires. The tombstones allow queries to provide monotonic
// index values within the TTL window. The GC is used to
// prevent monotonic growth in storage usage. This is a trade off
// between the length of the TTL and the storage overhead.
//
// In practice, this is required to fix the issue of delete
// visibility. When data is deleted from the KV store, the
// "latest" row can go backwards if the newest row is removed.
// The tombstones provide a way to ensure time doesn't move
// backwards within some interval.
//
type TombstoneGC struct {
	ttl         time.Duration
	granularity time.Duration

	// enabled controls if we actually setup any timers.
	enabled bool

	// expires maps the time of expiration to the highest
	// tombstone value that should be expired.
	expires map[time.Time]*expireInterval

	// expireCh is used to stream expiration
	expireCh chan uint64

	// lock is used to ensure safe access to all the fields
	lock sync.Mutex
}

// expireInterval is used to track the maximum index
// to expire in a given interval with a timer
type expireInterval struct {
	maxIndex uint64
	timer    *time.Timer
}

// NewTombstoneGC is used to construct a new TombstoneGC given
// a TTL for tombstones and a tracking granularity. Longer TTLs
// ensure correct behavior for more time, but use more storage.
// A shorter granularity increases the number of Raft transactions
// and reduce how far past the TTL we perform GC.
func NewTombstoneGC(ttl, granularity time.Duration) (*TombstoneGC, error) {
	// Sanity check the inputs
	if ttl <= 0 || granularity <= 0 {
		return nil, fmt.Errorf("Tombstone TTL and granularity must be positive")
	}

	t := &TombstoneGC{
		ttl:         ttl,
		granularity: granularity,
		enabled:     false,
		expires:     make(map[time.Time]*expireInterval),
		expireCh:    make(chan uint64, 1),
	}
	return t, nil
}

// ExpireCh is used to return a channel that streams the next index
// that should be expired
func (t *TombstoneGC) ExpireCh() <-chan uint64 {
	return t.expireCh
}

// SetEnabled is used to control if the tombstone GC is
// enabled. Should only be enabled by the leader node.
func (t *TombstoneGC) SetEnabled(enabled bool) {
	t.lock.Lock()
	defer t.lock.Unlock()
	if enabled == t.enabled {
		return
	}

	// Stop all the timers and clear
	if !enabled {
		for _, exp := range t.expires {
			exp.timer.Stop()
		}
		t.expires = make(map[time.Time]*expireInterval)
	}

	// Update the status
	t.enabled = enabled
}

// Hint is used to indicate that keys at the given index have been
// deleted, and that their GC should be scheduled.
func (t *TombstoneGC) Hint(index uint64) {
	expires := t.nextExpires()

	t.lock.Lock()
	defer t.lock.Unlock()
	if !t.enabled {
		return
	}

	// Check for an existing expiration timer
	exp, ok := t.expires[expires]
	if ok {
		// Increment the highest index to be expired at that time
		if index > exp.maxIndex {
			exp.maxIndex = index
		}
		return
	}

	// Create new expiration time
	t.expires[expires] = &expireInterval{
		maxIndex: index,
		timer: time.AfterFunc(expires.Sub(time.Now()), func() {
			t.expireTime(expires)
		}),
	}
}

// PendingExpiration is used to check if any expirations are pending
func (t *TombstoneGC) PendingExpiration() bool {
	t.lock.Lock()
	defer t.lock.Unlock()
	return len(t.expires) > 0
}

// nextExpires is used to calculate the next expiration time
func (t *TombstoneGC) nextExpires() time.Time {
	expires := time.Now().Add(t.ttl)
	remain := expires.UnixNano() % int64(t.granularity)
	adj := expires.Add(t.granularity - time.Duration(remain))
	return adj
}

// expireTime is used to expire the entries at the given time
func (t *TombstoneGC) expireTime(expires time.Time) {
	// Get the maximum index and clear the entry
	t.lock.Lock()
	exp := t.expires[expires]
	delete(t.expires, expires)
	t.lock.Unlock()

	// Notify the expires channel
	t.expireCh <- exp.maxIndex
}
<commit_msg>Cleans up comments.<commit_after>package state

import (
	"fmt"
	"sync"
	"time"
)

// TombstoneGC is used to track creation of tombstones so that they can be
// garbage collected after their TTL expires. The tombstones allow queries to
// provide monotonic index values within the TTL window. The GC is used to
// prevent monotonic growth in storage usage. This is a trade off between the
// length of the TTL and the storage overhead.
//
// In practice, this is required to fix the issue of delete visibility. When
// data is deleted from the KV store, the "latest" row can go backwards if the
// newest row is removed. The tombstones provide a way to ensure time doesn't
// move backwards within some interval.
//
type TombstoneGC struct {
	// ttl sets the TTL for tombstones.
	ttl time.Duration

	// granularity determines how we bin TTLs into timers.
	granularity time.Duration

	// enabled controls if we actually setup any timers.
	enabled bool

	// expires maps the time of expiration to the highest tombstone value
	// that should be expired.
	expires map[time.Time]*expireInterval

	// expireCh is used to stream expiration to the leader for processing.
	expireCh chan uint64

	// lock is used to ensure safe access to all the fields.
	lock sync.Mutex
}

// expireInterval is used to track the maximum index to expire in a given
// interval with a timer.
type expireInterval struct {
	// maxIndex has the highest tombstone index that should be GC-d.
	maxIndex uint64

	// timer is the timer tracking this bin.
	timer *time.Timer
}

// NewTombstoneGC is used to construct a new TombstoneGC given a TTL for
// tombstones and a tracking granularity. Longer TTLs ensure correct behavior
// for more time, but use more storage. A shorter granularity increases the
// number of Raft transactions and reduce how far past the TTL we perform GC.
func NewTombstoneGC(ttl, granularity time.Duration) (*TombstoneGC, error) {
	// Sanity check the inputs
	if ttl <= 0 || granularity <= 0 {
		return nil, fmt.Errorf("Tombstone TTL and granularity must be positive")
	}

	t := &TombstoneGC{
		ttl:         ttl,
		granularity: granularity,
		enabled:     false,
		expires:     make(map[time.Time]*expireInterval),
		expireCh:    make(chan uint64, 1),
	}
	return t, nil
}

// ExpireCh is used to return a channel that streams the next index that should
// be expired.
func (t *TombstoneGC) ExpireCh() <-chan uint64 {
	return t.expireCh
}

// SetEnabled is used to control if the tombstone GC is
// enabled. Should only be enabled by the leader node.
func (t *TombstoneGC) SetEnabled(enabled bool) {
	t.lock.Lock()
	defer t.lock.Unlock()
	if enabled == t.enabled {
		return
	}

	// Stop all the timers and clear
	if !enabled {
		for _, exp := range t.expires {
			exp.timer.Stop()
		}
		t.expires = make(map[time.Time]*expireInterval)
	}

	// Update the status
	t.enabled = enabled
}

// Hint is used to indicate that keys at the given index have been
// deleted, and that their GC should be scheduled.
func (t *TombstoneGC) Hint(index uint64) {
	expires := t.nextExpires()

	t.lock.Lock()
	defer t.lock.Unlock()
	if !t.enabled {
		return
	}

	// Check for an existing expiration timer and bump its index if we
	// find one.
	exp, ok := t.expires[expires]
	if ok {
		if index > exp.maxIndex {
			exp.maxIndex = index
		}
		return
	}

	// Create a new expiration timer.
	t.expires[expires] = &expireInterval{
		maxIndex: index,
		timer: time.AfterFunc(expires.Sub(time.Now()), func() {
			t.expireTime(expires)
		}),
	}
}

// PendingExpiration is used to check if any expirations are pending.
func (t *TombstoneGC) PendingExpiration() bool {
	t.lock.Lock()
	defer t.lock.Unlock()

	return len(t.expires) > 0
}

// nextExpires is used to calculate the next expiration time, based on the
// granularity that is set. This allows us to bin expirations and avoid a ton
// of timers.
func (t *TombstoneGC) nextExpires() time.Time {
	expires := time.Now().Add(t.ttl)
	remain := expires.UnixNano() % int64(t.granularity)
	adj := expires.Add(t.granularity - time.Duration(remain))
	return adj
}

// expireTime is used to expire the entries at the given time.
func (t *TombstoneGC) expireTime(expires time.Time) {
	// Get the maximum index and clear the entry
	t.lock.Lock()
	exp := t.expires[expires]
	delete(t.expires, expires)
	t.lock.Unlock()

	// Notify the expires channel
	t.expireCh <- exp.maxIndex
}
<|endoftext|>


<commit_before>package tabular

import (
	"bufio"
	"html"
	"io"
	"strings"
)

// HTMLOpts represents options passed to the HTML writer.
type HTMLOpts struct {
	Caption string
	Indent  int

	TableClass string
	RowClass   string
	HeadClass  string
	DataClass  string
}

// NewHTMLWriter creates a new HTML dataset writer.
func NewHTMLWriter(opts *HTMLOpts) *HTMLWriter {
	w := &HTMLWriter{opts}
	return w
}

// HTMLWriter represents a HTML dataset writer.
type HTMLWriter struct {
	opts *HTMLOpts
}

// Name returns name of the writer.
func (wh *HTMLWriter) Name() string {
	return "html"
}

// NeedsHeaders returns true if headers are required.
func (wh *HTMLWriter) NeedsHeaders() bool {
	return false
}

// Write writes dataset to writer.
func (wh *HTMLWriter) Write(d *Dataset, w io.Writer) error {
	tw := newHTMLTableWriter(d, w, wh.opts)
	return tw.write()
}

func newHTMLTableWriter(d *Dataset, w io.Writer, opts *HTMLOpts) *htmlTableWriter {
	return &htmlTableWriter{
		d:    d,
		w:    bufio.NewWriter(w),
		opts: opts,
	}
}

type htmlTableWriter struct {
	d    *Dataset
	w    *bufio.Writer
	opts *HTMLOpts
	err  error
}

func (h *htmlTableWriter) write() error {
	level := 0
	h.writeStartElem("table", level, h.opts.TableClass, true)

	if h.opts.Caption != "" {
		h.writeInlineElem("caption", h.opts.Caption, "", level+1)
	}

	if h.d.HasHeaders() {
		h.writeStartElem("thead", level+1, "", true)
		h.writeHeaders(level + 2)
		h.writeEndElem("thead", level+1, true)
	}

	h.writeStartElem("tbody", level+1, "", true)
	h.writeRows(level + 2)
	h.writeEndElem("tbody", level+1, true)

	h.writeEndElem("table", level+0, true)

	return h.flush()
}

func (h *htmlTableWriter) writeHeaders(level int) {
	h.writeStartElem("tr", level, h.opts.RowClass, true)
	for _, hdr := range h.d.Headers() {
		h.writeHeader(hdr, level+1)
	}
	h.writeEndElem("tr", level, true)
}

func (h *htmlTableWriter) writeHeader(hdr *Header, level int) {
	h.writeInlineElem("th", hdr.Title, h.opts.HeadClass, level)
}

func (h *htmlTableWriter) writeRows(level int) {
	for _, row := range h.d.Rows() {
		h.writeRow(row, level)
	}
}

func (h *htmlTableWriter) writeRow(row *Row, level int) {
	h.writeStartElem("tr", level, h.opts.RowClass, true)
	for _, item := range row.Items() {
		h.writeRowItem(item, level+1)
	}
	h.writeEndElem("tr", level, true)
}

func (h *htmlTableWriter) writeRowItem(item string, level int) {
	h.writeInlineElem("td", item, h.opts.DataClass, level)
}

func (h *htmlTableWriter) writeElem(name string, val string, class string, level int) {
	h.writeStartElem(name, level, class, false)
	h.writeEscaped(val)
	h.writeEndElem(name, level, true)
}

func (h *htmlTableWriter) writeInlineElem(name string, val string, class string, level int) {
	h.writeStartElem(name, level, class, false)
	h.writeEscaped(val)
	h.writeInlineEndElem(name, level, true)
}

func (h *htmlTableWriter) writeStartElem(name string, level int, class string, newline bool) {
	if class != "" {
		h.writeIndent(`<`+name+` class="`+h.escapeAttr(class)+`">`, level, true, false, newline)
	} else {
		h.writeIndent("<"+name+">", level, true, false, newline)
	}
}

func (h *htmlTableWriter) writeEndElem(name string, level int, newline bool) {
	h.writeIndent("</"+name+">", level, false, false, newline)
}

func (h *htmlTableWriter) writeInlineEndElem(name string, level int, newline bool) {
	h.writeIndent("</"+name+">", level, false, true, newline)
}

func (h *htmlTableWriter) writeIndent(val string, level int, start bool, inline bool, newline bool) {
	hasIndent := h.opts.Indent > 0
	if hasIndent && (start || !inline) {
		h.writeString(strings.Repeat(" ", h.opts.Indent*level))
	}

	h.writeString(val)
	if hasIndent && newline {
		h.writeString("\n")
	}
}

func (h *htmlTableWriter) writeEscaped(val string) {
	h.writeString(h.escapeVal(val))
}

func (h *htmlTableWriter) writeString(val string) {
	if h.err != nil {
		return
	}
	_, err := h.w.WriteString(val)
	h.err = err
}

func (h *htmlTableWriter) flush() error {
	if h.err != nil {
		return h.err
	}
	return h.w.Flush()
}

func (h *htmlTableWriter) escapeVal(val string) string {
	return html.EscapeString(val)
}

func (h *htmlTableWriter) escapeAttr(val string) string {
	return html.EscapeString(val)
}
<commit_msg>Drop unused writeElem from htmlTableWriter.<commit_after>package tabular

import (
	"bufio"
	"html"
	"io"
	"strings"
)

// HTMLOpts represents options passed to the HTML writer.
type HTMLOpts struct {
	Caption string
	Indent  int

	TableClass string
	RowClass   string
	HeadClass  string
	DataClass  string
}

// NewHTMLWriter creates a new HTML dataset writer.
func NewHTMLWriter(opts *HTMLOpts) *HTMLWriter {
	w := &HTMLWriter{opts}
	return w
}

// HTMLWriter represents a HTML dataset writer.
type HTMLWriter struct {
	opts *HTMLOpts
}

// Name returns name of the writer.
func (wh *HTMLWriter) Name() string {
	return "html"
}

// NeedsHeaders returns true if headers are required.
func (wh *HTMLWriter) NeedsHeaders() bool {
	return false
}

// Write writes dataset to writer.
func (wh *HTMLWriter) Write(d *Dataset, w io.Writer) error {
	tw := newHTMLTableWriter(d, w, wh.opts)
	return tw.write()
}

func newHTMLTableWriter(d *Dataset, w io.Writer, opts *HTMLOpts) *htmlTableWriter {
	return &htmlTableWriter{
		d:    d,
		w:    bufio.NewWriter(w),
		opts: opts,
	}
}

type htmlTableWriter struct {
	d    *Dataset
	w    *bufio.Writer
	opts *HTMLOpts
	err  error
}

func (h *htmlTableWriter) write() error {
	level := 0
	h.writeStartElem("table", level, h.opts.TableClass, true)

	if h.opts.Caption != "" {
		h.writeInlineElem("caption", h.opts.Caption, "", level+1)
	}

	if h.d.HasHeaders() {
		h.writeStartElem("thead", level+1, "", true)
		h.writeHeaders(level + 2)
		h.writeEndElem("thead", level+1, true)
	}

	h.writeStartElem("tbody", level+1, "", true)
	h.writeRows(level + 2)
	h.writeEndElem("tbody", level+1, true)

	h.writeEndElem("table", level+0, true)

	return h.flush()
}

func (h *htmlTableWriter) writeHeaders(level int) {
	h.writeStartElem("tr", level, h.opts.RowClass, true)
	for _, hdr := range h.d.Headers() {
		h.writeHeader(hdr, level+1)
	}
	h.writeEndElem("tr", level, true)
}

func (h *htmlTableWriter) writeHeader(hdr *Header, level int) {
	h.writeInlineElem("th", hdr.Title, h.opts.HeadClass, level)
}

func (h *htmlTableWriter) writeRows(level int) {
	for _, row := range h.d.Rows() {
		h.writeRow(row, level)
	}
}

func (h *htmlTableWriter) writeRow(row *Row, level int) {
	h.writeStartElem("tr", level, h.opts.RowClass, true)
	for _, item := range row.Items() {
		h.writeRowItem(item, level+1)
	}
	h.writeEndElem("tr", level, true)
}

func (h *htmlTableWriter) writeRowItem(item string, level int) {
	h.writeInlineElem("td", item, h.opts.DataClass, level)
}

func (h *htmlTableWriter) writeInlineElem(name string, val string, class string, level int) {
	h.writeStartElem(name, level, class, false)
	h.writeEscaped(val)
	h.writeInlineEndElem(name, level, true)
}

func (h *htmlTableWriter) writeStartElem(name string, level int, class string, newline bool) {
	if class != "" {
		h.writeIndent(`<`+name+` class="`+h.escapeAttr(class)+`">`, level, true, false, newline)
	} else {
		h.writeIndent("<"+name+">", level, true, false, newline)
	}
}

func (h *htmlTableWriter) writeEndElem(name string, level int, newline bool) {
	h.writeIndent("</"+name+">", level, false, false, newline)
}

func (h *htmlTableWriter) writeInlineEndElem(name string, level int, newline bool) {
	h.writeIndent("</"+name+">", level, false, true, newline)
}

func (h *htmlTableWriter) writeIndent(val string, level int, start bool, inline bool, newline bool) {
	hasIndent := h.opts.Indent > 0
	if hasIndent && (start || !inline) {
		h.writeString(strings.Repeat(" ", h.opts.Indent*level))
	}

	h.writeString(val)
	if hasIndent && newline {
		h.writeString("\n")
	}
}

func (h *htmlTableWriter) writeEscaped(val string) {
	h.writeString(h.escapeVal(val))
}

func (h *htmlTableWriter) writeString(val string) {
	if h.err != nil {
		return
	}
	_, err := h.w.WriteString(val)
	h.err = err
}

func (h *htmlTableWriter) flush() error {
	if h.err != nil {
		return h.err
	}
	return h.w.Flush()
}

func (h *htmlTableWriter) escapeVal(val string) string {
	return html.EscapeString(val)
}

func (h *htmlTableWriter) escapeAttr(val string) string {
	return html.EscapeString(val)
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Rook Authors. All rights reserved.

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

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

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

package v1

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

/*
 * Liveness probes
 */

// GetMonLivenessProbe returns the liveness probe for the MON service
func GetMonLivenessProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.LivenessProbe[ResourcesKeyMon].Probe
}

// GetMgrLivenessProbe returns the liveness probe for the MGR service
func GetMgrLivenessProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.LivenessProbe[ResourcesKeyMgr].Probe
}

// GetOSDLivenessProbe returns the liveness probe for the OSD service
func GetOSDLivenessProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.LivenessProbe[ResourcesKeyOSD].Probe
}

// GetMdsLivenessProbe returns the liveness probe for the MDS service
func GetMdsLivenessProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.LivenessProbe[ResourcesKeyMDS].Probe
}

/*
 * Startup probes
 */

// GetMonStartupProbe returns the startup probe for the MON service
func GetMonStartupProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.StartupProbe[ResourcesKeyMon].Probe
}

// GetMgrStartupProbe returns the startup probe for the MGR service
func GetMgrStartupProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.StartupProbe[ResourcesKeyMgr].Probe
}

// GetOSDStartupProbe returns the startup probe for the OSD service
func GetOSDStartupProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.StartupProbe[ResourcesKeyOSD].Probe
}

// GetMdsStartupProbe returns the startup probe for the MDS service
func GetMdsStartupProbe(l CephClusterHealthCheckSpec) *corev1.Probe {
	return l.StartupProbe[ResourcesKeyMDS].Probe
}
<commit_msg>core: remove obsolete probe helper<commit_after><|endoftext|>


<commit_before>package metrics

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/sapcc/kubernikus/pkg/api/models"
)

const (
	metricNamespace = "kubernikus"
)

var klusterPhases = []models.KlusterPhase{
	models.KlusterPhasePending,
	models.KlusterPhaseCreating,
	models.KlusterPhaseRunning,
	models.KlusterPhaseTerminating,
}

var klusterBootDurationSummary = prometheus.NewSummaryVec(
	prometheus.SummaryOpts{
		Namespace: metricNamespace,
		Name:      "kluster_boot_duration",
		Help:      "Duration until kluster got from phase pending to running",
	},
	[]string{},
)

var klusterInfo = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "kluster_info",
		Help:      "detailed information on a kluster",
	},
	[]string{"kluster_namespace", "kluster_name", "kluster_version", "creator", "account", "project_id"},
)

var klusterStatusPhase = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "kluster_status_phase",
		Help:      "the phase the kluster is currently in",
	},
	[]string{"kluster_id", "phase"},
)

var nodePoolSize = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "node_pool_size",
		Help:      "size of a node pool",
	},
	[]string{"kluster_id", "node_pool", "image_name", "flavor_name"},
)

var nodePoolStatus = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "node_pool_status",
		Help:      "status of the node pool and the number of nodes nodes in that status",
	},
	[]string{"kluster_id", "node_pool", "status"},
)

func SetMetricKlusterInfo(namespace, name, version, projectID string, annotations, labels map[string]string) {
	promLabels := prometheus.Labels{
		"kluster_namespace": namespace,
		"kluster_name":      name,
		"kluster_version":   version,
		"creator":           getCreatorFromAnnotations(annotations),
		"account":           getAccountFromLabels(labels),
		"project_id":        projectID,
	}
	klusterInfo.With(promLabels).Set(1)
}

/*
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="<phase>"} 			< 1|0 >
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="creating"} 		1
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="running"} 			0
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="pending"} 			0
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="terminating"} 	0
*/
func SetMetricKlusterStatusPhase(klusterName string, klusterPhase models.KlusterPhase) {
	// Set current phase to 1, others to 0
	for _, phase := range klusterPhases {
		labels := prometheus.Labels{
			"kluster_id": klusterName,
			"phase":      string(phase),
		}
		klusterStatusPhase.With(labels).Set(boolToFloat64(klusterPhase == phase))
	}
}

func SetMetricKlusterTerminated(klusterName string) {
	for _, phase := range klusterPhases {
		labels := prometheus.Labels{
			"kluster_id": klusterName,
			"phase":      string(phase),
		}
		klusterStatusPhase.With(labels).Set(0)
	}
}

/*
kubernikus_node_pool_size{"kluster_id"="<id", "node_pool"="<name>", "image_name"="<name>", "flavor_name"="<name>"} <node_pool_size>
*/
func setMetricNodePoolSize(klusterID, nodePoolName, imageName, flavorName string, size int64) {
	nodePoolSize.With(prometheus.Labels{
		"kluster_id":  klusterID,
		"node_pool":   nodePoolName,
		"image_name":  imageName,
		"flavor_name": flavorName,
	}).Set(float64(size))
}

/*
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="<status>"} < number of nodes in that status >
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="ready"} 1
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="running"} 1
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="healthy"} 1
*/
func setMetricNodePoolStatus(klusterID, nodePoolName string, status map[string]int64) {
	if status != nil {
		for s, v := range status {
			nodePoolStatus.With(prometheus.Labels{
				"kluster_id": klusterID,
				"node_pool":  nodePoolName,
				"status":     s,
			}).Set(float64(v))
		}
	}
}

func SetMetricBootDurationSummary(creationTimestamp, now time.Time) {
	klusterBootDurationSummary.With(prometheus.Labels{}).Observe(now.Sub(creationTimestamp).Seconds())
}

func boolToFloat64(b bool) float64 {
	if b {
		return 1
	}
	return 0
}

func getCreatorFromAnnotations(annotations map[string]string) string {
	creator, ok := annotations["creator"]
	if !ok {
		return "NA"
	}
	return creator
}

func getAccountFromLabels(labels map[string]string) string {
	account, ok := labels["account"]
	if !ok {
		return "NA"
	}
	return account
}

func init() {
	prometheus.MustRegister(
		klusterInfo,
		klusterStatusPhase,
		klusterBootDurationSummary,
		nodePoolSize,
		nodePoolStatus,
		LaunchOperationsLatency,
		LaunchOperationsTotal,
		LaunchSuccessfulOperationsTotal,
		LaunchFailedOperationsTotal,
	)
}

func ExposeMetrics(host string, metricPort int, stopCh <-chan struct{}, wg *sync.WaitGroup, logger log.Logger) error {
	logger.Log(
		"msg", "Exposing metrics",
		"host", host,
		"port", metricPort,
		"path", "/metrics")
	defer wg.Done()
	wg.Add(1)
	for {
		select {
		case <-stopCh:
			return nil
		default:
			http.Handle("/metrics", promhttp.Handler())
			return http.ListenAndServe(
				fmt.Sprintf("%v:%v", host, metricPort),
				nil,
			)
		}
	}
}
<commit_msg>Fix shutdown of operator metrics exporter<commit_after>package metrics

import (
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/sapcc/kubernikus/pkg/api/models"
)

const (
	metricNamespace = "kubernikus"
)

var klusterPhases = []models.KlusterPhase{
	models.KlusterPhasePending,
	models.KlusterPhaseCreating,
	models.KlusterPhaseRunning,
	models.KlusterPhaseTerminating,
}

var klusterBootDurationSummary = prometheus.NewSummaryVec(
	prometheus.SummaryOpts{
		Namespace: metricNamespace,
		Name:      "kluster_boot_duration",
		Help:      "Duration until kluster got from phase pending to running",
	},
	[]string{},
)

var klusterInfo = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "kluster_info",
		Help:      "detailed information on a kluster",
	},
	[]string{"kluster_namespace", "kluster_name", "kluster_version", "creator", "account", "project_id"},
)

var klusterStatusPhase = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "kluster_status_phase",
		Help:      "the phase the kluster is currently in",
	},
	[]string{"kluster_id", "phase"},
)

var nodePoolSize = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "node_pool_size",
		Help:      "size of a node pool",
	},
	[]string{"kluster_id", "node_pool", "image_name", "flavor_name"},
)

var nodePoolStatus = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Namespace: metricNamespace,
		Name:      "node_pool_status",
		Help:      "status of the node pool and the number of nodes nodes in that status",
	},
	[]string{"kluster_id", "node_pool", "status"},
)

func SetMetricKlusterInfo(namespace, name, version, projectID string, annotations, labels map[string]string) {
	promLabels := prometheus.Labels{
		"kluster_namespace": namespace,
		"kluster_name":      name,
		"kluster_version":   version,
		"creator":           getCreatorFromAnnotations(annotations),
		"account":           getAccountFromLabels(labels),
		"project_id":        projectID,
	}
	klusterInfo.With(promLabels).Set(1)
}

/*
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="<phase>"} 			< 1|0 >
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="creating"} 		1
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="running"} 			0
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="pending"} 			0
kubernikus_kluster_status_phase{"kluster_id"="<id>","phase"="terminating"} 	0
*/
func SetMetricKlusterStatusPhase(klusterName string, klusterPhase models.KlusterPhase) {
	// Set current phase to 1, others to 0
	for _, phase := range klusterPhases {
		labels := prometheus.Labels{
			"kluster_id": klusterName,
			"phase":      string(phase),
		}
		klusterStatusPhase.With(labels).Set(boolToFloat64(klusterPhase == phase))
	}
}

func SetMetricKlusterTerminated(klusterName string) {
	for _, phase := range klusterPhases {
		labels := prometheus.Labels{
			"kluster_id": klusterName,
			"phase":      string(phase),
		}
		klusterStatusPhase.With(labels).Set(0)
	}
}

/*
kubernikus_node_pool_size{"kluster_id"="<id", "node_pool"="<name>", "image_name"="<name>", "flavor_name"="<name>"} <node_pool_size>
*/
func setMetricNodePoolSize(klusterID, nodePoolName, imageName, flavorName string, size int64) {
	nodePoolSize.With(prometheus.Labels{
		"kluster_id":  klusterID,
		"node_pool":   nodePoolName,
		"image_name":  imageName,
		"flavor_name": flavorName,
	}).Set(float64(size))
}

/*
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="<status>"} < number of nodes in that status >
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="ready"} 1
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="running"} 1
kubernikus_node_pool_status{"kluster_id"="<id", "node_pool"="<name>", "status"="healthy"} 1
*/
func setMetricNodePoolStatus(klusterID, nodePoolName string, status map[string]int64) {
	if status != nil {
		for s, v := range status {
			nodePoolStatus.With(prometheus.Labels{
				"kluster_id": klusterID,
				"node_pool":  nodePoolName,
				"status":     s,
			}).Set(float64(v))
		}
	}
}

func SetMetricBootDurationSummary(creationTimestamp, now time.Time) {
	klusterBootDurationSummary.With(prometheus.Labels{}).Observe(now.Sub(creationTimestamp).Seconds())
}

func boolToFloat64(b bool) float64 {
	if b {
		return 1
	}
	return 0
}

func getCreatorFromAnnotations(annotations map[string]string) string {
	creator, ok := annotations["creator"]
	if !ok {
		return "NA"
	}
	return creator
}

func getAccountFromLabels(labels map[string]string) string {
	account, ok := labels["account"]
	if !ok {
		return "NA"
	}
	return account
}

func init() {
	prometheus.MustRegister(
		klusterInfo,
		klusterStatusPhase,
		klusterBootDurationSummary,
		nodePoolSize,
		nodePoolStatus,
		LaunchOperationsLatency,
		LaunchOperationsTotal,
		LaunchSuccessfulOperationsTotal,
		LaunchFailedOperationsTotal,
	)
}

func ExposeMetrics(host string, metricPort int, stopCh <-chan struct{}, wg *sync.WaitGroup, logger log.Logger) {
	wg.Add(1)
	defer wg.Done()
	ln, err := net.Listen("tcp", fmt.Sprintf("%v:%v", host, metricPort))
	logger.Log(
		"msg", "Exposing metrics",
		"host", host,
		"port", metricPort,
		"err", err)
	if err != nil {
		return
	}
	go http.Serve(ln, promhttp.Handler())
	<-stopCh
	ln.Close()
}
<|endoftext|>


<commit_before>package alerting

import (
	"errors"
	"fmt"
	"time"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/components/imguploader"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/metrics"
	"github.com/grafana/grafana/pkg/services/rendering"
	"github.com/grafana/grafana/pkg/setting"

	m "github.com/grafana/grafana/pkg/models"
)

type NotifierPlugin struct {
	Type            string          `json:"type"`
	Name            string          `json:"name"`
	Description     string          `json:"description"`
	OptionsTemplate string          `json:"optionsTemplate"`
	Factory         NotifierFactory `json:"-"`
}

type NotificationService interface {
	SendIfNeeded(context *EvalContext) error
}

func NewNotificationService(renderService rendering.Service) NotificationService {
	return &notificationService{
		log:           log.New("alerting.notifier"),
		renderService: renderService,
	}
}

type notificationService struct {
	log           log.Logger
	renderService rendering.Service
}

func (n *notificationService) SendIfNeeded(context *EvalContext) error {
	notifierStates, err := n.getNeededNotifiers(context.Rule.OrgId, context.Rule.Notifications, context)
	if err != nil {
		return err
	}

	if len(notifierStates) == 0 {
		return nil
	}

	if notifierStates.ShouldUploadImage() {
		if err = n.uploadImage(context); err != nil {
			n.log.Error("Failed to upload alert panel image.", "error", err)
		}
	}

	return n.sendNotifications(context, notifierStates)
}

func (n *notificationService) sendAndMarkAsComplete(evalContext *EvalContext, notifierState *notifierState) error {
	not := notifierState.notifier
	n.log.Debug("Sending notification", "type", not.GetType(), "id", not.GetNotifierId(), "isDefault", not.GetIsDefault())
	metrics.M_Alerting_Notification_Sent.WithLabelValues(not.GetType()).Inc()

	err := not.Notify(evalContext)

	if err != nil {
		n.log.Error("failed to send notification", "id", not.GetNotifierId())
	} else {
		notifierState.state.UpdatedAt = time.Now().UTC().Unix()
	}

	if evalContext.IsTestRun {
		return nil
	}

	cmd := &m.SetAlertNotificationStateToCompleteCommand{
		Id:      not.GetNotifierId(),
		Version: notifierState.state.Version,
	}

	if err = bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		if err == m.ErrAlertNotificationStateVersionConflict {
			n.log.Error("notification state out of sync", "id", not.GetNotifierId())
			return nil
		}

		return err
	}

	return nil
}

func (n *notificationService) sendNotification(evalContext *EvalContext, notifierState *notifierState) error {
	if !evalContext.IsTestRun {
		setPendingCmd := &m.SetAlertNotificationStateToPendingCommand{
			Id:                           notifierState.state.NotifierId,
			Version:                      notifierState.state.Version,
			AlertRuleStateUpdatedVersion: evalContext.Rule.StateChanges,
		}

		err := bus.DispatchCtx(evalContext.Ctx, setPendingCmd)
		if err == m.ErrAlertNotificationStateVersionConflict {
			return nil
		}

		if err != nil {
			return err
		}

		// We need to update state version to be able to log
		// unexpected version conflicts when marking notifications as ok
		notifierState.state.Version = setPendingCmd.ResultVersion
	}

	return n.sendAndMarkAsComplete(evalContext, notifierState)
}

func (n *notificationService) sendNotifications(evalContext *EvalContext, notifierStates notifierStateSlice) error {
	for _, notifierState := range notifierStates {
		err := n.sendNotification(evalContext, notifierState)
		if err != nil {
			n.log.Error("failed to send notification", "id", notifierState.notifier.GetNotifierId())
		}
	}

	return nil
}

func (n *notificationService) uploadImage(context *EvalContext) (err error) {
	uploader, err := imguploader.NewImageUploader()
	if err != nil {
		return err
	}

	renderOpts := rendering.Opts{
		Width:           1000,
		Height:          500,
		Timeout:         alertTimeout / 2,
		OrgId:           context.Rule.OrgId,
		OrgRole:         m.ROLE_ADMIN,
		ConcurrentLimit: setting.AlertingRenderLimit,
	}

	ref, err := context.GetDashboardUID()
	if err != nil {
		return err
	}

	renderOpts.Path = fmt.Sprintf("d-solo/%s/%s?panelId=%d", ref.Uid, ref.Slug, context.Rule.PanelId)

	result, err := n.renderService.Render(context.Ctx, renderOpts)
	if err != nil {
		return err
	}

	context.ImageOnDiskPath = result.FilePath
	context.ImagePublicUrl, err = uploader.Upload(context.Ctx, context.ImageOnDiskPath)
	if err != nil {
		return err
	}

	if context.ImagePublicUrl != "" {
		n.log.Info("uploaded screenshot of alert to external image store", "url", context.ImagePublicUrl)
	}

	return nil
}

func (n *notificationService) getNeededNotifiers(orgId int64, notificationIds []int64, evalContext *EvalContext) (notifierStateSlice, error) {
	query := &m.GetAlertNotificationsToSendQuery{OrgId: orgId, Ids: notificationIds}

	if err := bus.Dispatch(query); err != nil {
		return nil, err
	}

	var result notifierStateSlice
	for _, notification := range query.Result {
		not, err := n.createNotifierFor(notification)
		if err != nil {
			n.log.Error("Could not create notifier", "notifier", notification.Id, "error", err)
			continue
		}

		query := &m.GetOrCreateNotificationStateQuery{
			NotifierId: notification.Id,
			AlertId:    evalContext.Rule.Id,
			OrgId:      evalContext.Rule.OrgId,
		}

		err = bus.DispatchCtx(evalContext.Ctx, query)
		if err != nil {
			n.log.Error("Could not get notification state.", "notifier", notification.Id, "error", err)
			continue
		}

		if not.ShouldNotify(evalContext.Ctx, evalContext, query.Result) {
			result = append(result, &notifierState{
				notifier: not,
				state:    query.Result,
			})
		}
	}

	return result, nil
}

func (n *notificationService) createNotifierFor(model *m.AlertNotification) (Notifier, error) {
	notifierPlugin, found := notifierFactories[model.Type]
	if !found {
		return nil, errors.New("Unsupported notification type")
	}

	return notifierPlugin.Factory(model)
}

type NotifierFactory func(notification *m.AlertNotification) (Notifier, error)

var notifierFactories = make(map[string]*NotifierPlugin)

func RegisterNotifier(plugin *NotifierPlugin) {
	notifierFactories[plugin.Type] = plugin
}

func GetNotifiers() []*NotifierPlugin {
	list := make([]*NotifierPlugin, 0)

	for _, value := range notifierFactories {
		list = append(list, value)
	}

	return list
}
<commit_msg>improve local variable name<commit_after>package alerting

import (
	"errors"
	"fmt"
	"time"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/components/imguploader"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/metrics"
	"github.com/grafana/grafana/pkg/services/rendering"
	"github.com/grafana/grafana/pkg/setting"

	m "github.com/grafana/grafana/pkg/models"
)

type NotifierPlugin struct {
	Type            string          `json:"type"`
	Name            string          `json:"name"`
	Description     string          `json:"description"`
	OptionsTemplate string          `json:"optionsTemplate"`
	Factory         NotifierFactory `json:"-"`
}

type NotificationService interface {
	SendIfNeeded(context *EvalContext) error
}

func NewNotificationService(renderService rendering.Service) NotificationService {
	return &notificationService{
		log:           log.New("alerting.notifier"),
		renderService: renderService,
	}
}

type notificationService struct {
	log           log.Logger
	renderService rendering.Service
}

func (n *notificationService) SendIfNeeded(context *EvalContext) error {
	notifierStates, err := n.getNeededNotifiers(context.Rule.OrgId, context.Rule.Notifications, context)
	if err != nil {
		return err
	}

	if len(notifierStates) == 0 {
		return nil
	}

	if notifierStates.ShouldUploadImage() {
		if err = n.uploadImage(context); err != nil {
			n.log.Error("Failed to upload alert panel image.", "error", err)
		}
	}

	return n.sendNotifications(context, notifierStates)
}

func (n *notificationService) sendAndMarkAsComplete(evalContext *EvalContext, notifierState *notifierState) error {
	notifier := notifierState.notifier
	n.log.Debug("Sending notification", "type", notifier.GetType(), "id", notifier.GetNotifierId(), "isDefault", notifier.GetIsDefault())
	metrics.M_Alerting_Notification_Sent.WithLabelValues(notifier.GetType()).Inc()

	err := notifier.Notify(evalContext)

	if err != nil {
		n.log.Error("failed to send notification", "id", notifier.GetNotifierId())
	} else {
		notifierState.state.UpdatedAt = time.Now().UTC().Unix()
	}

	if evalContext.IsTestRun {
		return nil
	}

	cmd := &m.SetAlertNotificationStateToCompleteCommand{
		Id:      notifier.GetNotifierId(),
		Version: notifierState.state.Version,
	}

	if err = bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		if err == m.ErrAlertNotificationStateVersionConflict {
			n.log.Error("notification state out of sync", "id", notifier.GetNotifierId())
			return nil
		}

		return err
	}

	return nil
}

func (n *notificationService) sendNotification(evalContext *EvalContext, notifierState *notifierState) error {
	if !evalContext.IsTestRun {
		setPendingCmd := &m.SetAlertNotificationStateToPendingCommand{
			Id:                           notifierState.state.NotifierId,
			Version:                      notifierState.state.Version,
			AlertRuleStateUpdatedVersion: evalContext.Rule.StateChanges,
		}

		err := bus.DispatchCtx(evalContext.Ctx, setPendingCmd)
		if err == m.ErrAlertNotificationStateVersionConflict {
			return nil
		}

		if err != nil {
			return err
		}

		// We need to update state version to be able to log
		// unexpected version conflicts when marking notifications as ok
		notifierState.state.Version = setPendingCmd.ResultVersion
	}

	return n.sendAndMarkAsComplete(evalContext, notifierState)
}

func (n *notificationService) sendNotifications(evalContext *EvalContext, notifierStates notifierStateSlice) error {
	for _, notifierState := range notifierStates {
		err := n.sendNotification(evalContext, notifierState)
		if err != nil {
			n.log.Error("failed to send notification", "id", notifierState.notifier.GetNotifierId())
		}
	}

	return nil
}

func (n *notificationService) uploadImage(context *EvalContext) (err error) {
	uploader, err := imguploader.NewImageUploader()
	if err != nil {
		return err
	}

	renderOpts := rendering.Opts{
		Width:           1000,
		Height:          500,
		Timeout:         alertTimeout / 2,
		OrgId:           context.Rule.OrgId,
		OrgRole:         m.ROLE_ADMIN,
		ConcurrentLimit: setting.AlertingRenderLimit,
	}

	ref, err := context.GetDashboardUID()
	if err != nil {
		return err
	}

	renderOpts.Path = fmt.Sprintf("d-solo/%s/%s?panelId=%d", ref.Uid, ref.Slug, context.Rule.PanelId)

	result, err := n.renderService.Render(context.Ctx, renderOpts)
	if err != nil {
		return err
	}

	context.ImageOnDiskPath = result.FilePath
	context.ImagePublicUrl, err = uploader.Upload(context.Ctx, context.ImageOnDiskPath)
	if err != nil {
		return err
	}

	if context.ImagePublicUrl != "" {
		n.log.Info("uploaded screenshot of alert to external image store", "url", context.ImagePublicUrl)
	}

	return nil
}

func (n *notificationService) getNeededNotifiers(orgId int64, notificationIds []int64, evalContext *EvalContext) (notifierStateSlice, error) {
	query := &m.GetAlertNotificationsToSendQuery{OrgId: orgId, Ids: notificationIds}

	if err := bus.Dispatch(query); err != nil {
		return nil, err
	}

	var result notifierStateSlice
	for _, notification := range query.Result {
		not, err := n.createNotifierFor(notification)
		if err != nil {
			n.log.Error("Could not create notifier", "notifier", notification.Id, "error", err)
			continue
		}

		query := &m.GetOrCreateNotificationStateQuery{
			NotifierId: notification.Id,
			AlertId:    evalContext.Rule.Id,
			OrgId:      evalContext.Rule.OrgId,
		}

		err = bus.DispatchCtx(evalContext.Ctx, query)
		if err != nil {
			n.log.Error("Could not get notification state.", "notifier", notification.Id, "error", err)
			continue
		}

		if not.ShouldNotify(evalContext.Ctx, evalContext, query.Result) {
			result = append(result, &notifierState{
				notifier: not,
				state:    query.Result,
			})
		}
	}

	return result, nil
}

func (n *notificationService) createNotifierFor(model *m.AlertNotification) (Notifier, error) {
	notifierPlugin, found := notifierFactories[model.Type]
	if !found {
		return nil, errors.New("Unsupported notification type")
	}

	return notifierPlugin.Factory(model)
}

type NotifierFactory func(notification *m.AlertNotification) (Notifier, error)

var notifierFactories = make(map[string]*NotifierPlugin)

func RegisterNotifier(plugin *NotifierPlugin) {
	notifierFactories[plugin.Type] = plugin
}

func GetNotifiers() []*NotifierPlugin {
	list := make([]*NotifierPlugin, 0)

	for _, value := range notifierFactories {
		list = append(list, value)
	}

	return list
}
<|endoftext|>