<commit_before>// Package microsoft provides authentication strategies using Microsoft.
package microsoft

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"golang.org/x/oauth2"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/pkg/log"
)

const (
	apiURL = "https://graph.microsoft.com"
	// Microsoft requires this scope to access user's profile
	scopeUser = "user.read"
	// Microsoft requires this scope to list groups the user is a member of
	// and resolve their UUIDs to groups names.
	scopeGroups = "directory.read.all"
)

// Config holds configuration options for microsoft logins.
type Config struct {
	ClientID           string   `json:"clientID"`
	ClientSecret       string   `json:"clientSecret"`
	RedirectURI        string   `json:"redirectURI"`
	Tenant             string   `json:"tenant"`
	OnlySecurityGroups bool     `json:"onlySecurityGroups"`
	Groups             []string `json:"groups"`
}

// Open returns a strategy for logging in through Microsoft.
func (c *Config) Open(id string, logger log.Logger) (connector.Connector, error) {
	m := microsoftConnector{
		redirectURI:        c.RedirectURI,
		clientID:           c.ClientID,
		clientSecret:       c.ClientSecret,
		tenant:             c.Tenant,
		onlySecurityGroups: c.OnlySecurityGroups,
		groups:             c.Groups,
		logger:             logger,
	}
	// By default allow logins from both personal and business/school
	// accounts.
	if m.tenant == "" {
		m.tenant = "common"
	}

	return &m, nil
}

type connectorData struct {
	AccessToken  string    `json:"accessToken"`
	RefreshToken string    `json:"refreshToken"`
	Expiry       time.Time `json:"expiry"`
}

var (
	_ connector.CallbackConnector = (*microsoftConnector)(nil)
	_ connector.RefreshConnector  = (*microsoftConnector)(nil)
)

type microsoftConnector struct {
	redirectURI        string
	clientID           string
	clientSecret       string
	tenant             string
	onlySecurityGroups bool
	groups             []string
	logger             log.Logger
}

func (c *microsoftConnector) isOrgTenant() bool {
	return c.tenant != "common" && c.tenant != "consumers" && c.tenant != "organizations"
}

func (c *microsoftConnector) groupsRequired(groupScope bool) bool {
	return (len(c.groups) > 0 || groupScope) && c.isOrgTenant()
}

func (c *microsoftConnector) oauth2Config(scopes connector.Scopes) *oauth2.Config {
	microsoftScopes := []string{scopeUser}
	if c.groupsRequired(scopes.Groups) {
		microsoftScopes = append(microsoftScopes, scopeGroups)
	}

	return &oauth2.Config{
		ClientID:     c.clientID,
		ClientSecret: c.clientSecret,
		Endpoint: oauth2.Endpoint{
			AuthURL:  "https://login.microsoftonline.com/" + c.tenant + "/oauth2/v2.0/authorize",
			TokenURL: "https://login.microsoftonline.com/" + c.tenant + "/oauth2/v2.0/token",
		},
		Scopes:      microsoftScopes,
		RedirectURL: c.redirectURI,
	}
}

func (c *microsoftConnector) LoginURL(scopes connector.Scopes, callbackURL, state string) (string, error) {
	if c.redirectURI != callbackURL {
		return "", fmt.Errorf("expected callback URL %q did not match the URL in the config %q", callbackURL, c.redirectURI)
	}

	return c.oauth2Config(scopes).AuthCodeURL(state), nil
}

func (c *microsoftConnector) HandleCallback(s connector.Scopes, r *http.Request) (identity connector.Identity, err error) {
	q := r.URL.Query()
	if errType := q.Get("error"); errType != "" {
		return identity, &oauth2Error{errType, q.Get("error_description")}
	}

	oauth2Config := c.oauth2Config(s)

	ctx := r.Context()

	token, err := oauth2Config.Exchange(ctx, q.Get("code"))
	if err != nil {
		return identity, fmt.Errorf("microsoft: failed to get token: %v", err)
	}

	client := oauth2Config.Client(ctx, token)

	user, err := c.user(ctx, client)
	if err != nil {
		return identity, fmt.Errorf("microsoft: get user: %v", err)
	}

	identity = connector.Identity{
		UserID:        user.ID,
		Username:      user.Name,
		Email:         user.Email,
		EmailVerified: true,
	}

	if c.groupsRequired(s.Groups) {
		groups, err := c.getGroups(ctx, client, user.ID)
		if err != nil {
			return identity, fmt.Errorf("microsoft: get groups: %v", err)
		}
		identity.Groups = groups
	}

	if s.OfflineAccess {
		data := connectorData{
			AccessToken:  token.AccessToken,
			RefreshToken: token.RefreshToken,
			Expiry:       token.Expiry,
		}
		connData, err := json.Marshal(data)
		if err != nil {
			return identity, fmt.Errorf("microsoft: marshal connector data: %v", err)
		}
		identity.ConnectorData = connData
	}

	return identity, nil
}

type tokenNotifyFunc func(*oauth2.Token) error

// notifyRefreshTokenSource is essentially `oauth2.ResuseTokenSource` with `TokenNotifyFunc` added.
type notifyRefreshTokenSource struct {
	new oauth2.TokenSource
	mu  sync.Mutex // guards t
	t   *oauth2.Token
	f   tokenNotifyFunc // called when token refreshed so new refresh token can be persisted
}

// Token returns the current token if it's still valid, else will
// refresh the current token (using r.Context for HTTP client
// information) and return the new one.
func (s *notifyRefreshTokenSource) Token() (*oauth2.Token, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.t.Valid() {
		return s.t, nil
	}
	t, err := s.new.Token()
	if err != nil {
		return nil, err
	}
	s.t = t
	return t, s.f(t)
}

func (c *microsoftConnector) Refresh(ctx context.Context, s connector.Scopes, identity connector.Identity) (connector.Identity, error) {
	if len(identity.ConnectorData) == 0 {
		return identity, errors.New("microsoft: no upstream access token found")
	}

	var data connectorData
	if err := json.Unmarshal(identity.ConnectorData, &data); err != nil {
		return identity, fmt.Errorf("microsoft: unmarshal access token: %v", err)
	}
	tok := &oauth2.Token{
		AccessToken:  data.AccessToken,
		RefreshToken: data.RefreshToken,
		Expiry:       data.Expiry,
	}

	client := oauth2.NewClient(ctx, &notifyRefreshTokenSource{
		new: c.oauth2Config(s).TokenSource(ctx, tok),
		t:   tok,
		f: func(tok *oauth2.Token) error {
			data := connectorData{
				AccessToken:  tok.AccessToken,
				RefreshToken: tok.RefreshToken,
				Expiry:       tok.Expiry,
			}
			connData, err := json.Marshal(data)
			if err != nil {
				return fmt.Errorf("microsoft: marshal connector data: %v", err)
			}
			identity.ConnectorData = connData
			return nil
		},
	})
	user, err := c.user(ctx, client)
	if err != nil {
		return identity, fmt.Errorf("microsoft: get user: %v", err)
	}

	identity.Username = user.Name
	identity.Email = user.Email

	if c.groupsRequired(s.Groups) {
		groups, err := c.getGroups(ctx, client, user.ID)
		if err != nil {
			return identity, fmt.Errorf("microsoft: get groups: %v", err)
		}
		identity.Groups = groups
	}

	return identity, nil
}

// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/resources/user
// id                - The unique identifier for the user. Inherited from
//                     directoryObject. Key. Not nullable. Read-only.
// displayName       - The name displayed in the address book for the user.
//                     This is usually the combination of the user's first name,
//                     middle initial and last name. This property is required
//                     when a user is created and it cannot be cleared during
//                     updates. Supports $filter and $orderby.
// userPrincipalName - The user principal name (UPN) of the user.
//                     The UPN is an Internet-style login name for the user
//                     based on the Internet standard RFC 822. By convention,
//                     this should map to the user's email name. The general
//                     format is alias@domain, where domain must be present in
//                     the tenant’s collection of verified domains. This
//                     property is required when a user is created. The
//                     verified domains for the tenant can be accessed from the
//                     verifiedDomains property of organization. Supports
//                     $filter and $orderby.
type user struct {
	ID    string `json:"id"`
	Name  string `json:"displayName"`
	Email string `json:"userPrincipalName"`
}

func (c *microsoftConnector) user(ctx context.Context, client *http.Client) (u user, err error) {
	// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/api/user_get
	req, err := http.NewRequest("GET", apiURL+"/v1.0/me?$select=id,displayName,userPrincipalName", nil)
	if err != nil {
		return u, fmt.Errorf("new req: %v", err)
	}

	resp, err := client.Do(req.WithContext(ctx))
	if err != nil {
		return u, fmt.Errorf("get URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return u, newGraphError(resp.Body)
	}

	if err := json.NewDecoder(resp.Body).Decode(&u); err != nil {
		return u, fmt.Errorf("JSON decode: %v", err)
	}

	return u, err
}

// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/resources/group
// displayName - The display name for the group. This property is required when
//               a group is created and it cannot be cleared during updates.
//               Supports $filter and $orderby.
type group struct {
	Name string `json:"displayName"`
}

func (c *microsoftConnector) getGroups(ctx context.Context, client *http.Client, userID string) (groups []string, err error) {
	ids, err := c.getGroupIDs(ctx, client)
	if err != nil {
		return groups, err
	}

	groups, err = c.getGroupNames(ctx, client, ids)
	if err != nil {
		return
	}

	// ensure that the user is in at least one required group
	isInGroups := false
	if len(c.groups) > 0 {
		gs := make(map[string]struct{})
		for _, g := range c.groups {
			gs[g] = struct{}{}
		}

		for _, g := range groups {
			if _, ok := gs[g]; ok {
				isInGroups = true
				break
			}
		}
	}
	if len(c.groups) > 0 && !isInGroups {
		return nil, fmt.Errorf("microsoft: user %v not in required groups", userID)
	}

	return
}

func (c *microsoftConnector) getGroupIDs(ctx context.Context, client *http.Client) (ids []string, err error) {
	// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/api/user_getmembergroups
	in := &struct {
		SecurityEnabledOnly bool `json:"securityEnabledOnly"`
	}{c.onlySecurityGroups}
	reqURL := apiURL + "/v1.0/me/getMemberGroups"
	for {
		var out []string
		var next string

		next, err = c.post(ctx, client, reqURL, in, &out)
		if err != nil {
			return ids, err
		}

		ids = append(ids, out...)
		if next == "" {
			return
		}
		reqURL = next
	}
}

func (c *microsoftConnector) getGroupNames(ctx context.Context, client *http.Client, ids []string) (groups []string, err error) {
	if len(ids) == 0 {
		return
	}

	// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/api/directoryobject_getbyids
	in := &struct {
		IDs   []string `json:"ids"`
		Types []string `json:"types"`
	}{ids, []string{"group"}}
	reqURL := apiURL + "/v1.0/directoryObjects/getByIds"
	for {
		var out []group
		var next string

		next, err = c.post(ctx, client, reqURL, in, &out)
		if err != nil {
			return groups, err
		}

		for _, g := range out {
			groups = append(groups, g.Name)
		}
		if next == "" {
			return
		}
		reqURL = next
	}
}

func (c *microsoftConnector) post(ctx context.Context, client *http.Client, reqURL string, in interface{}, out interface{}) (string, error) {
	var payload bytes.Buffer

	err := json.NewEncoder(&payload).Encode(in)
	if err != nil {
		return "", fmt.Errorf("microsoft: JSON encode: %v", err)
	}

	req, err := http.NewRequest("POST", reqURL, &payload)
	if err != nil {
		return "", fmt.Errorf("new req: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req.WithContext(ctx))
	if err != nil {
		return "", fmt.Errorf("post URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", newGraphError(resp.Body)
	}

	var next string
	if err = json.NewDecoder(resp.Body).Decode(&struct {
		NextLink *string     `json:"@odata.nextLink"`
		Value    interface{} `json:"value"`
	}{&next, out}); err != nil {
		return "", fmt.Errorf("JSON decode: %v", err)
	}

	return next, nil
}

type graphError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

func (e *graphError) Error() string {
	return e.Code + ": " + e.Message
}

func newGraphError(r io.Reader) error {
	// https://developer.microsoft.com/en-us/graph/docs/concepts/errors
	var ge graphError
	if err := json.NewDecoder(r).Decode(&struct {
		Error *graphError `json:"error"`
	}{&ge}); err != nil {
		return fmt.Errorf("JSON error decode: %v", err)
	}
	return &ge
}

type oauth2Error struct {
	error            string
	errorDescription string
}

func (e *oauth2Error) Error() string {
	if e.errorDescription == "" {
		return e.error
	}
	return e.error + ": " + e.errorDescription
}
<commit_msg>dexidp#1440 Add offline_access scope, if required<commit_after>// Package microsoft provides authentication strategies using Microsoft.
package microsoft

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"golang.org/x/oauth2"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/pkg/log"
)

const (
	apiURL = "https://graph.microsoft.com"
	// Microsoft requires this scope to access user's profile
	scopeUser = "user.read"
	// Microsoft requires this scope to list groups the user is a member of
	// and resolve their UUIDs to groups names.
	scopeGroups = "directory.read.all"
	// Microsoft requires this scope to return a refresh token
	// see https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-permissions-and-consent#offline_access
	scopeOfflineAccess = "offline_access"
)

// Config holds configuration options for microsoft logins.
type Config struct {
	ClientID           string   `json:"clientID"`
	ClientSecret       string   `json:"clientSecret"`
	RedirectURI        string   `json:"redirectURI"`
	Tenant             string   `json:"tenant"`
	OnlySecurityGroups bool     `json:"onlySecurityGroups"`
	Groups             []string `json:"groups"`
}

// Open returns a strategy for logging in through Microsoft.
func (c *Config) Open(id string, logger log.Logger) (connector.Connector, error) {
	m := microsoftConnector{
		redirectURI:        c.RedirectURI,
		clientID:           c.ClientID,
		clientSecret:       c.ClientSecret,
		tenant:             c.Tenant,
		onlySecurityGroups: c.OnlySecurityGroups,
		groups:             c.Groups,
		logger:             logger,
	}
	// By default allow logins from both personal and business/school
	// accounts.
	if m.tenant == "" {
		m.tenant = "common"
	}

	return &m, nil
}

type connectorData struct {
	AccessToken  string    `json:"accessToken"`
	RefreshToken string    `json:"refreshToken"`
	Expiry       time.Time `json:"expiry"`
}

var (
	_ connector.CallbackConnector = (*microsoftConnector)(nil)
	_ connector.RefreshConnector  = (*microsoftConnector)(nil)
)

type microsoftConnector struct {
	redirectURI        string
	clientID           string
	clientSecret       string
	tenant             string
	onlySecurityGroups bool
	groups             []string
	logger             log.Logger
}

func (c *microsoftConnector) isOrgTenant() bool {
	return c.tenant != "common" && c.tenant != "consumers" && c.tenant != "organizations"
}

func (c *microsoftConnector) groupsRequired(groupScope bool) bool {
	return (len(c.groups) > 0 || groupScope) && c.isOrgTenant()
}

func (c *microsoftConnector) oauth2Config(scopes connector.Scopes) *oauth2.Config {
	microsoftScopes := []string{scopeUser}
	if c.groupsRequired(scopes.Groups) {
		microsoftScopes = append(microsoftScopes, scopeGroups)
	}

	if scopes.OfflineAccess {
		microsoftScopes = append(microsoftScopes, scopeOfflineAccess)
	}

	return &oauth2.Config{
		ClientID:     c.clientID,
		ClientSecret: c.clientSecret,
		Endpoint: oauth2.Endpoint{
			AuthURL:  "https://login.microsoftonline.com/" + c.tenant + "/oauth2/v2.0/authorize",
			TokenURL: "https://login.microsoftonline.com/" + c.tenant + "/oauth2/v2.0/token",
		},
		Scopes:      microsoftScopes,
		RedirectURL: c.redirectURI,
	}
}

func (c *microsoftConnector) LoginURL(scopes connector.Scopes, callbackURL, state string) (string, error) {
	if c.redirectURI != callbackURL {
		return "", fmt.Errorf("expected callback URL %q did not match the URL in the config %q", callbackURL, c.redirectURI)
	}

	return c.oauth2Config(scopes).AuthCodeURL(state), nil
}

func (c *microsoftConnector) HandleCallback(s connector.Scopes, r *http.Request) (identity connector.Identity, err error) {
	q := r.URL.Query()
	if errType := q.Get("error"); errType != "" {
		return identity, &oauth2Error{errType, q.Get("error_description")}
	}

	oauth2Config := c.oauth2Config(s)

	ctx := r.Context()

	token, err := oauth2Config.Exchange(ctx, q.Get("code"))
	if err != nil {
		return identity, fmt.Errorf("microsoft: failed to get token: %v", err)
	}

	client := oauth2Config.Client(ctx, token)

	user, err := c.user(ctx, client)
	if err != nil {
		return identity, fmt.Errorf("microsoft: get user: %v", err)
	}

	identity = connector.Identity{
		UserID:        user.ID,
		Username:      user.Name,
		Email:         user.Email,
		EmailVerified: true,
	}

	if c.groupsRequired(s.Groups) {
		groups, err := c.getGroups(ctx, client, user.ID)
		if err != nil {
			return identity, fmt.Errorf("microsoft: get groups: %v", err)
		}
		identity.Groups = groups
	}

	if s.OfflineAccess {
		data := connectorData{
			AccessToken:  token.AccessToken,
			RefreshToken: token.RefreshToken,
			Expiry:       token.Expiry,
		}
		connData, err := json.Marshal(data)
		if err != nil {
			return identity, fmt.Errorf("microsoft: marshal connector data: %v", err)
		}
		identity.ConnectorData = connData
	}

	return identity, nil
}

type tokenNotifyFunc func(*oauth2.Token) error

// notifyRefreshTokenSource is essentially `oauth2.ResuseTokenSource` with `TokenNotifyFunc` added.
type notifyRefreshTokenSource struct {
	new oauth2.TokenSource
	mu  sync.Mutex // guards t
	t   *oauth2.Token
	f   tokenNotifyFunc // called when token refreshed so new refresh token can be persisted
}

// Token returns the current token if it's still valid, else will
// refresh the current token (using r.Context for HTTP client
// information) and return the new one.
func (s *notifyRefreshTokenSource) Token() (*oauth2.Token, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.t.Valid() {
		return s.t, nil
	}
	t, err := s.new.Token()
	if err != nil {
		return nil, err
	}
	s.t = t
	return t, s.f(t)
}

func (c *microsoftConnector) Refresh(ctx context.Context, s connector.Scopes, identity connector.Identity) (connector.Identity, error) {
	if len(identity.ConnectorData) == 0 {
		return identity, errors.New("microsoft: no upstream access token found")
	}

	var data connectorData
	if err := json.Unmarshal(identity.ConnectorData, &data); err != nil {
		return identity, fmt.Errorf("microsoft: unmarshal access token: %v", err)
	}
	tok := &oauth2.Token{
		AccessToken:  data.AccessToken,
		RefreshToken: data.RefreshToken,
		Expiry:       data.Expiry,
	}

	client := oauth2.NewClient(ctx, &notifyRefreshTokenSource{
		new: c.oauth2Config(s).TokenSource(ctx, tok),
		t:   tok,
		f: func(tok *oauth2.Token) error {
			data := connectorData{
				AccessToken:  tok.AccessToken,
				RefreshToken: tok.RefreshToken,
				Expiry:       tok.Expiry,
			}
			connData, err := json.Marshal(data)
			if err != nil {
				return fmt.Errorf("microsoft: marshal connector data: %v", err)
			}
			identity.ConnectorData = connData
			return nil
		},
	})
	user, err := c.user(ctx, client)
	if err != nil {
		return identity, fmt.Errorf("microsoft: get user: %v", err)
	}

	identity.Username = user.Name
	identity.Email = user.Email

	if c.groupsRequired(s.Groups) {
		groups, err := c.getGroups(ctx, client, user.ID)
		if err != nil {
			return identity, fmt.Errorf("microsoft: get groups: %v", err)
		}
		identity.Groups = groups
	}

	return identity, nil
}

// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/resources/user
// id                - The unique identifier for the user. Inherited from
//                     directoryObject. Key. Not nullable. Read-only.
// displayName       - The name displayed in the address book for the user.
//                     This is usually the combination of the user's first name,
//                     middle initial and last name. This property is required
//                     when a user is created and it cannot be cleared during
//                     updates. Supports $filter and $orderby.
// userPrincipalName - The user principal name (UPN) of the user.
//                     The UPN is an Internet-style login name for the user
//                     based on the Internet standard RFC 822. By convention,
//                     this should map to the user's email name. The general
//                     format is alias@domain, where domain must be present in
//                     the tenant’s collection of verified domains. This
//                     property is required when a user is created. The
//                     verified domains for the tenant can be accessed from the
//                     verifiedDomains property of organization. Supports
//                     $filter and $orderby.
type user struct {
	ID    string `json:"id"`
	Name  string `json:"displayName"`
	Email string `json:"userPrincipalName"`
}

func (c *microsoftConnector) user(ctx context.Context, client *http.Client) (u user, err error) {
	// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/api/user_get
	req, err := http.NewRequest("GET", apiURL+"/v1.0/me?$select=id,displayName,userPrincipalName", nil)
	if err != nil {
		return u, fmt.Errorf("new req: %v", err)
	}

	resp, err := client.Do(req.WithContext(ctx))
	if err != nil {
		return u, fmt.Errorf("get URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return u, newGraphError(resp.Body)
	}

	if err := json.NewDecoder(resp.Body).Decode(&u); err != nil {
		return u, fmt.Errorf("JSON decode: %v", err)
	}

	return u, err
}

// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/resources/group
// displayName - The display name for the group. This property is required when
//               a group is created and it cannot be cleared during updates.
//               Supports $filter and $orderby.
type group struct {
	Name string `json:"displayName"`
}

func (c *microsoftConnector) getGroups(ctx context.Context, client *http.Client, userID string) (groups []string, err error) {
	ids, err := c.getGroupIDs(ctx, client)
	if err != nil {
		return groups, err
	}

	groups, err = c.getGroupNames(ctx, client, ids)
	if err != nil {
		return
	}

	// ensure that the user is in at least one required group
	isInGroups := false
	if len(c.groups) > 0 {
		gs := make(map[string]struct{})
		for _, g := range c.groups {
			gs[g] = struct{}{}
		}

		for _, g := range groups {
			if _, ok := gs[g]; ok {
				isInGroups = true
				break
			}
		}
	}
	if len(c.groups) > 0 && !isInGroups {
		return nil, fmt.Errorf("microsoft: user %v not in required groups", userID)
	}

	return
}

func (c *microsoftConnector) getGroupIDs(ctx context.Context, client *http.Client) (ids []string, err error) {
	// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/api/user_getmembergroups
	in := &struct {
		SecurityEnabledOnly bool `json:"securityEnabledOnly"`
	}{c.onlySecurityGroups}
	reqURL := apiURL + "/v1.0/me/getMemberGroups"
	for {
		var out []string
		var next string

		next, err = c.post(ctx, client, reqURL, in, &out)
		if err != nil {
			return ids, err
		}

		ids = append(ids, out...)
		if next == "" {
			return
		}
		reqURL = next
	}
}

func (c *microsoftConnector) getGroupNames(ctx context.Context, client *http.Client, ids []string) (groups []string, err error) {
	if len(ids) == 0 {
		return
	}

	// https://developer.microsoft.com/en-us/graph/docs/api-reference/v1.0/api/directoryobject_getbyids
	in := &struct {
		IDs   []string `json:"ids"`
		Types []string `json:"types"`
	}{ids, []string{"group"}}
	reqURL := apiURL + "/v1.0/directoryObjects/getByIds"
	for {
		var out []group
		var next string

		next, err = c.post(ctx, client, reqURL, in, &out)
		if err != nil {
			return groups, err
		}

		for _, g := range out {
			groups = append(groups, g.Name)
		}
		if next == "" {
			return
		}
		reqURL = next
	}
}

func (c *microsoftConnector) post(ctx context.Context, client *http.Client, reqURL string, in interface{}, out interface{}) (string, error) {
	var payload bytes.Buffer

	err := json.NewEncoder(&payload).Encode(in)
	if err != nil {
		return "", fmt.Errorf("microsoft: JSON encode: %v", err)
	}

	req, err := http.NewRequest("POST", reqURL, &payload)
	if err != nil {
		return "", fmt.Errorf("new req: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req.WithContext(ctx))
	if err != nil {
		return "", fmt.Errorf("post URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", newGraphError(resp.Body)
	}

	var next string
	if err = json.NewDecoder(resp.Body).Decode(&struct {
		NextLink *string     `json:"@odata.nextLink"`
		Value    interface{} `json:"value"`
	}{&next, out}); err != nil {
		return "", fmt.Errorf("JSON decode: %v", err)
	}

	return next, nil
}

type graphError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

func (e *graphError) Error() string {
	return e.Code + ": " + e.Message
}

func newGraphError(r io.Reader) error {
	// https://developer.microsoft.com/en-us/graph/docs/concepts/errors
	var ge graphError
	if err := json.NewDecoder(r).Decode(&struct {
		Error *graphError `json:"error"`
	}{&ge}); err != nil {
		return fmt.Errorf("JSON error decode: %v", err)
	}
	return &ge
}

type oauth2Error struct {
	error            string
	errorDescription string
}

func (e *oauth2Error) Error() string {
	if e.errorDescription == "" {
		return e.error
	}
	return e.error + ": " + e.errorDescription
}
<|endoftext|>


<commit_before>package connectors

import (
	"encoding/json"
	"fmt"
	"github.com/soprasteria/intools-engine/common/logs"
	"github.com/soprasteria/intools-engine/common/utils"
	"github.com/soprasteria/intools-engine/common/websocket"
	"github.com/soprasteria/intools-engine/executors"
	"github.com/soprasteria/intools-engine/intools"
	"github.com/samalba/dockerclient"
	"gopkg.in/robfig/cron.v2"
	"sync"
	"time"
)

func InitSchedule(c *Connector) cron.EntryID {
	if intools.Engine.GetCron() != nil {
		crontab := fmt.Sprintf("@every %ds", c.Refresh)
		logs.Debug.Printf("Schedule %s:%s %s", c.Group, c.Name, crontab)
		entryId, _ := intools.Engine.GetCron().AddJob(crontab, c)
		return entryId
	}
	return 0
}

func RemoveScheduleJob(entryId cron.EntryID) {
	if intools.Engine.GetCron() != nil {
		logs.Debug.Printf("Remove schedule job with cronId: %s", entryId)
		intools.Engine.GetCron().Remove(entryId)
	}
}

func Exec(connector *Connector) (*executors.Executor, error) {
	executor := &executors.Executor{}

	//Saving connector to redis
	go SaveConnector(connector)

	//Get all containers
	containers, err := intools.Engine.GetDockerClient().ListContainers(true, false, "")
	if err != nil {
		logs.Error.Println(err)
		return nil, err
	}

	//Searching for the container with the same name
	containerExists := false
	previousContainerId := "-1"
	for _, c := range containers {
		for _, n := range c.Names {
			if n == fmt.Sprintf("/%s", connector.GetContainerName()) {
				containerExists = true
				previousContainerId = c.Id
			}
		}
	}

	//If it exists, remove it
	if containerExists {
		logs.Trace.Printf("Removing container %s [/%s]", previousContainerId[:11], connector.GetContainerName())
		err := intools.Engine.GetDockerClient().RemoveContainer(previousContainerId, true, true)
		if err != nil {
			logs.Error.Println("Cannot remove container " + previousContainerId[:11])
			logs.Error.Println(err)
			return nil, err
		}
	}

	//Create container
	logs.Debug.Println("ContainerConfig ", connector.ContainerConfig)
	ContainerId, err := intools.Engine.GetDockerClient().CreateContainer(connector.ContainerConfig, connector.GetContainerName(), intools.Engine.GetDockerAuth())
	if err != nil {
		logs.Error.Println("Cannot create container " + connector.GetContainerName())
		logs.Error.Println(err)
		return nil, err
	}
	//Save the short ContainerId
	executor.ContainerId = ContainerId[:11]
	executor.Host = intools.Engine.GetDockerHost()

	logs.Trace.Printf("%s [/%s] successfully created", executor.ContainerId, connector.GetContainerName())
	hostConfig := &dockerclient.HostConfig{}

	//Prepare the waiting group to sync execution of the container
	var wg sync.WaitGroup
	wg.Add(1)

	//Start the container
	err = intools.Engine.GetDockerClient().StartContainer(ContainerId, hostConfig)
	if err != nil {
		logs.Error.Println("Cannot start container " + executor.ContainerId)
		logs.Error.Println(err)
		return nil, err
	}

	logs.Trace.Printf("%s [/%s] successfully started", executor.ContainerId, connector.GetContainerName())
	logs.Debug.Println(executor.ContainerId + " will be stopped after " + fmt.Sprint(connector.Timeout) + " seconds")
	//Trigger stop of the container after the timeout
	intools.Engine.GetDockerClient().StopContainer(ContainerId, connector.Timeout)

	//Wait for the end of the execution of the container
	for {
		//Each time inspect the container
		inspect, err := intools.Engine.GetDockerClient().InspectContainer(ContainerId)
		if err != nil {
			logs.Error.Println("Cannot inpect container " + executor.ContainerId)
			logs.Error.Println(err)
			return executor, err
		}
		if !inspect.State.Running {
			//When the container is not running
			logs.Debug.Println(executor.ContainerId + " is stopped")
			executor.Running = false
			executor.Terminated = true
			executor.ExitCode = inspect.State.ExitCode
			executor.StartedAt = inspect.State.StartedAt
			executor.FinishedAt = inspect.State.FinishedAt
			//Trigger next part of the waiting group
			wg.Done()
			//Exit from the waiting loop
			break
		} else {
			//Wait
			logs.Debug.Println(executor.ContainerId + " is running...")
			time.Sleep(5 * time.Second)
		}
	}

	//Next part : after the container has been executed
	wg.Wait()

	logStdOutOptions := &dockerclient.LogOptions{
		Follow:     true,
		Stdout:     true,
		Stderr:     false,
		Timestamps: false,
		Tail:       0,
	}

	logStdErrOptions := &dockerclient.LogOptions{
		Follow:     true,
		Stdout:     false,
		Stderr:     true,
		Timestamps: false,
		Tail:       0,
	}

	//Get the stdout and stderr
	logsStdOutReader, err := intools.Engine.GetDockerClient().ContainerLogs(ContainerId, logStdOutOptions)
	logsStdErrReader, err := intools.Engine.GetDockerClient().ContainerLogs(ContainerId, logStdErrOptions)

	if err != nil {
		logs.Error.Println("-cannot read logs from server")
	} else {
		containerLogs, err := utils.ReadLogs(logsStdOutReader)
		if err != nil {
			return executor, err
		} else {
			executor.Stdout = containerLogs
			executor.JsonStdout = new(map[string]interface{})
			errJsonStdOut := json.Unmarshal([]byte(executor.Stdout), executor.JsonStdout)
			executor.Valid = true
			if errJsonStdOut != nil {
				logs.Warning.Printf("Unable to parse stdout from container %s", executor.ContainerId)
				logs.Warning.Println(errJsonStdOut)
			}
		}
		containerLogs, err = utils.ReadLogs(logsStdErrReader)
		if err != nil {
			return executor, err
		} else {
			executor.Stderr = containerLogs
		}
	}

	// Broadcast result to registered clients
	lightConnector := &websocket.LightConnector{
		GroupId:     connector.Group,
		ConnectorId: connector.Name,
		Value:       executor.JsonStdout,
	}
	websocket.ConnectorBuffer <- lightConnector

	//Save result to redis
	defer SaveExecutor(connector, executor)

	return executor, nil
}
<commit_msg>Pulling image before starting container<commit_after>package connectors

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/samalba/dockerclient"
	"github.com/soprasteria/intools-engine/common/logs"
	"github.com/soprasteria/intools-engine/common/utils"
	"github.com/soprasteria/intools-engine/common/websocket"
	"github.com/soprasteria/intools-engine/executors"
	"github.com/soprasteria/intools-engine/intools"
	"gopkg.in/robfig/cron.v2"
)

func InitSchedule(c *Connector) cron.EntryID {
	if intools.Engine.GetCron() != nil {
		crontab := fmt.Sprintf("@every %ds", c.Refresh)
		logs.Debug.Printf("Schedule %s:%s %s", c.Group, c.Name, crontab)
		entryId, _ := intools.Engine.GetCron().AddJob(crontab, c)
		return entryId
	}
	return 0
}

func RemoveScheduleJob(entryId cron.EntryID) {
	if intools.Engine.GetCron() != nil {
		logs.Debug.Printf("Remove schedule job with cronId: %s", entryId)
		intools.Engine.GetCron().Remove(entryId)
	}
}

func Exec(connector *Connector) (*executors.Executor, error) {
	executor := &executors.Executor{}

	//Saving connector to redis
	go SaveConnector(connector)

	//Get all containers
	containers, err := intools.Engine.GetDockerClient().ListContainers(true, false, "")
	if err != nil {
		logs.Error.Println(err)
		return nil, err
	}

	//Searching for the container with the same name
	containerExists := false
	previousContainerId := "-1"
	for _, c := range containers {
		for _, n := range c.Names {
			if n == fmt.Sprintf("/%s", connector.GetContainerName()) {
				containerExists = true
				previousContainerId = c.Id
			}
		}
	}

	//If it exists, remove it
	if containerExists {
		logs.Trace.Printf("Removing container %s [/%s]", previousContainerId[:11], connector.GetContainerName())
		err := intools.Engine.GetDockerClient().RemoveContainer(previousContainerId, true, true)
		if err != nil {
			logs.Error.Println("Cannot remove container " + previousContainerId[:11])
			logs.Error.Println(err)
			return nil, err
		}
	}

	// Pulls the image
	logs.Debug.Println("Pulling image ", connector.ContainerConfig.Image)
	err = intools.Engine.GetDockerClient().PullImage(connector.ContainerConfig.Image, intools.Engine.GetDockerAuth())
	if err != nil {
		logs.Error.Println("Cannot pull image " + connector.ContainerConfig.Image + " for container " + connector.GetContainerName())
		logs.Error.Println(err)
		return nil, err
	}

	//Create container
	logs.Debug.Println("ContainerConfig ", connector.ContainerConfig)
	ContainerId, err := intools.Engine.GetDockerClient().CreateContainer(connector.ContainerConfig, connector.GetContainerName(), intools.Engine.GetDockerAuth())
	if err != nil {
		logs.Error.Println("Cannot create container " + connector.GetContainerName())
		logs.Error.Println(err)
		return nil, err
	}
	//Save the short ContainerId
	executor.ContainerId = ContainerId[:11]
	executor.Host = intools.Engine.GetDockerHost()

	logs.Trace.Printf("%s [/%s] successfully created", executor.ContainerId, connector.GetContainerName())
	hostConfig := &dockerclient.HostConfig{}

	//Prepare the waiting group to sync execution of the container
	var wg sync.WaitGroup
	wg.Add(1)

	//Start the container
	err = intools.Engine.GetDockerClient().StartContainer(ContainerId, hostConfig)
	if err != nil {
		logs.Error.Println("Cannot start container " + executor.ContainerId)
		logs.Error.Println(err)
		return nil, err
	}

	logs.Trace.Printf("%s [/%s] successfully started", executor.ContainerId, connector.GetContainerName())
	logs.Debug.Println(executor.ContainerId + " will be stopped after " + fmt.Sprint(connector.Timeout) + " seconds")
	//Trigger stop of the container after the timeout
	intools.Engine.GetDockerClient().StopContainer(ContainerId, connector.Timeout)

	//Wait for the end of the execution of the container
	for {
		//Each time inspect the container
		inspect, err := intools.Engine.GetDockerClient().InspectContainer(ContainerId)
		if err != nil {
			logs.Error.Println("Cannot inpect container " + executor.ContainerId)
			logs.Error.Println(err)
			return executor, err
		}
		if !inspect.State.Running {
			//When the container is not running
			logs.Debug.Println(executor.ContainerId + " is stopped")
			executor.Running = false
			executor.Terminated = true
			executor.ExitCode = inspect.State.ExitCode
			executor.StartedAt = inspect.State.StartedAt
			executor.FinishedAt = inspect.State.FinishedAt
			//Trigger next part of the waiting group
			wg.Done()
			//Exit from the waiting loop
			break
		} else {
			//Wait
			logs.Debug.Println(executor.ContainerId + " is running...")
			time.Sleep(5 * time.Second)
		}
	}

	//Next part : after the container has been executed
	wg.Wait()

	logStdOutOptions := &dockerclient.LogOptions{
		Follow:     true,
		Stdout:     true,
		Stderr:     false,
		Timestamps: false,
		Tail:       0,
	}

	logStdErrOptions := &dockerclient.LogOptions{
		Follow:     true,
		Stdout:     false,
		Stderr:     true,
		Timestamps: false,
		Tail:       0,
	}

	//Get the stdout and stderr
	logsStdOutReader, err := intools.Engine.GetDockerClient().ContainerLogs(ContainerId, logStdOutOptions)
	logsStdErrReader, err := intools.Engine.GetDockerClient().ContainerLogs(ContainerId, logStdErrOptions)

	if err != nil {
		logs.Error.Println("-cannot read logs from server")
	} else {
		containerLogs, err := utils.ReadLogs(logsStdOutReader)
		if err != nil {
			return executor, err
		} else {
			logs.Debug.Printf("container logs %s", containerLogs)
			executor.Stdout = containerLogs
			executor.JsonStdout = new(map[string]interface{})
			errJsonStdOut := json.Unmarshal([]byte(executor.Stdout), executor.JsonStdout)
			executor.Valid = true
			if errJsonStdOut != nil {
				logs.Warning.Printf("Unable to parse stdout from container %s", executor.ContainerId)
				logs.Warning.Println(errJsonStdOut)
			}
		}
		containerLogs, err = utils.ReadLogs(logsStdErrReader)
		if err != nil {
			return executor, err
		} else {
			executor.Stderr = containerLogs
		}
	}

	// Broadcast result to registered clients
	lightConnector := &websocket.LightConnector{
		GroupId:     connector.Group,
		ConnectorId: connector.Name,
		Value:       executor.JsonStdout,
	}
	websocket.ConnectorBuffer <- lightConnector

	//Save result to redis
	defer SaveExecutor(connector, executor)

	return executor, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package validation

import (
	"fmt"

	gm "github.com/getgauge/gauge/gauge_messages"
)

var message = map[gm.StepValidateResponse_ErrorType]string{
	gm.StepValidateResponse_STEP_IMPLEMENTATION_NOT_FOUND: "Add the following missing implementations to fix `Step implementation not found` errors.",
}

func showSuggestion(validationErrors validationErrors) {
	if !HideSuggestion {
		for t, errs := range groupErrors(validationErrors) {
			fmt.Println(getSuggestionMessage(t))
			for _, err := range errs {
				fmt.Println(err.Suggestion())
			}
		}
	}
}

func getSuggestionMessage(t gm.StepValidateResponse_ErrorType) string {
	if msg, ok := message[t]; ok {
		return msg
	}
	return fmt.Sprintf("Suggestions for fixing `%s` errors.\n", getMessage(t.String()))
}

func groupErrors(validationErrors validationErrors) map[gm.StepValidateResponse_ErrorType][]StepValidationError {
	errMap := make(map[gm.StepValidateResponse_ErrorType][]StepValidationError)
	for _, errs := range validationErrors {
		for _, v := range errs {
			if e, ok := v.(StepValidationError); ok && e.suggestion != "" {
				errType := *(v.(StepValidationError).errorType)
				errMap[errType] = append(errMap[errType], e)
			}
		}
	}
	return errMap
}
<commit_msg>Add new line before printing validation error suggestion<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package validation

import (
	"fmt"

	gm "github.com/getgauge/gauge/gauge_messages"
)

var message = map[gm.StepValidateResponse_ErrorType]string{
	gm.StepValidateResponse_STEP_IMPLEMENTATION_NOT_FOUND: "Add the following missing implementations to fix `Step implementation not found` errors.\n",
}

func showSuggestion(validationErrors validationErrors) {
	if !HideSuggestion {
		for t, errs := range groupErrors(validationErrors) {
			fmt.Println(getSuggestionMessage(t))
			for _, err := range errs {
				fmt.Println(err.Suggestion())
			}
		}
	}
}

func getSuggestionMessage(t gm.StepValidateResponse_ErrorType) string {
	if msg, ok := message[t]; ok {
		return msg
	}
	return fmt.Sprintf("Suggestions for fixing `%s` errors.\n", getMessage(t.String()))
}

func groupErrors(validationErrors validationErrors) map[gm.StepValidateResponse_ErrorType][]StepValidationError {
	errMap := make(map[gm.StepValidateResponse_ErrorType][]StepValidationError)
	for _, errs := range validationErrors {
		for _, v := range errs {
			if e, ok := v.(StepValidationError); ok && e.suggestion != "" {
				errType := *(v.(StepValidationError).errorType)
				errMap[errType] = append(errMap[errType], e)
			}
		}
	}
	return errMap
}
<|endoftext|>


<commit_before>package db

import (
	"container/list"
	"errors"
	"expvar"
	"sync"
	"time"

	"github.com/couchbase/sync_gateway/base"
)

// Number of recently-accessed doc revisions to cache in RAM
var KDefaultRevisionCacheCapacity uint32 = 5000

// An LRU cache of document revision bodies, together with their channel access.
type RevisionCache struct {
	cache      map[IDAndRev]*list.Element // Fast lookup of list element by doc/rev ID
	lruList    *list.List                 // List ordered by most recent access (Front is newest)
	capacity   uint32                     // Max number of revisions to cache
	loaderFunc RevisionCacheLoaderFunc    // Function which does actual loading of something from rev cache
	lock       sync.Mutex                 // For thread-safety
	statsCache *expvar.Map                // Per-db stats related to cache
}

// Revision information as returned by the rev cache
type DocumentRevision struct {
	RevID       string
	Body        Body
	History     Revisions
	Channels    base.Set
	Expiry      *time.Time
	Attachments AttachmentsMeta
	Delta       *RevCacheDelta
}

// Callback function signature for loading something from the rev cache
type RevisionCacheLoaderFunc func(id IDAndRev) (body Body, history Revisions, channels base.Set, attachments AttachmentsMeta, expiry *time.Time, err error)

// The cache payload data. Stored as the Value of a list Element.
type revCacheValue struct {
	key         IDAndRev        // doc/rev IDs
	body        Body            // Revision body (a pristine shallow copy)
	history     Revisions       // Rev history encoded like a "_revisions" property
	channels    base.Set        // Set of channels that have access
	expiry      *time.Time      // Document expiry
	attachments AttachmentsMeta // Document _attachments property
	delta       *RevCacheDelta  // Available delta *from* this revision
	err         error           // Error from loaderFunc if it failed
	lock        sync.RWMutex    // Synchronizes access to this struct
}

type RevCacheDelta struct {
	ToRevID    string
	DeltaBytes []byte
}

// Creates a revision cache with the given capacity and an optional loader function.
func NewRevisionCache(capacity uint32, loaderFunc RevisionCacheLoaderFunc, statsCache *expvar.Map) *RevisionCache {

	if capacity == 0 {
		capacity = KDefaultRevisionCacheCapacity
	}

	return &RevisionCache{
		cache:      map[IDAndRev]*list.Element{},
		lruList:    list.New(),
		capacity:   capacity,
		loaderFunc: loaderFunc,
		statsCache: statsCache,
	}
}

// Looks up a revision from the cache.
// Returns the body of the revision, its history, and the set of channels it's in.
// If the cache has a loaderFunction, it will be called if the revision isn't in the cache;
// any error returned by the loaderFunction will be returned from Get.
func (rc *RevisionCache) Get(docid, revid string) (DocumentRevision, error) {
	return rc.getFromCache(docid, revid, BodyShallowCopy, rc.loaderFunc != nil)
}

// Looks up a revision from the cache only.  Will not fall back to loader function if not
// present in the cache.
func (rc *RevisionCache) GetCached(docid, revid string) (DocumentRevision, error) {
	return rc.getFromCache(docid, revid, BodyShallowCopy, false)
}

// Returns the body of the revision based on the specified body copy policy (deep/shallow/none)
func (rc *RevisionCache) GetWithCopy(docid, revid string, copyType BodyCopyType) (DocumentRevision, error) {
	return rc.getFromCache(docid, revid, copyType, rc.loaderFunc != nil)
}

// Attempt to update the delta on a revision cache entry.  If the entry is no longer resident in the cache,
// fails silently
func (rc *RevisionCache) UpdateDelta(docID, revID string, toRevID string, delta []byte) {
	value := rc.getValue(docID, revID, false)
	if value != nil {
		value.updateDelta(toRevID, delta)
	}
}

func (rc *RevisionCache) getFromCache(docid, revid string, copyType BodyCopyType, loadOnCacheMiss bool) (DocumentRevision, error) {
	value := rc.getValue(docid, revid, loadOnCacheMiss)
	if value == nil {
		return DocumentRevision{}, nil
	}
	docRev, statEvent, err := value.load(rc.loaderFunc, copyType)
	rc.statsRecorderFunc(statEvent)

	if err != nil {
		rc.removeValue(value) // don't keep failed loads in the cache
	}
	return docRev, err
}

// Attempts to retrieve the active revision for a document from the cache.  Requires retrieval
// of the document from the bucket to guarantee the current active revision, but does minimal unmarshalling
// of the retrieved document to get the current rev from _sync metadata.  If active rev is already in the
// rev cache, will use it.  Otherwise will add to the rev cache using the raw document obtained in the
// initial retrieval.
func (rc *RevisionCache) GetActive(docid string, context *DatabaseContext) (docRev DocumentRevision, err error) {

	// Look up active rev for doc
	bucketDoc, getErr := context.GetDocument(docid, DocUnmarshalSync)
	if getErr != nil {
		return DocumentRevision{}, getErr
	}
	if bucketDoc == nil {
		return DocumentRevision{}, nil
	}

	// Retrieve from or add to rev cache
	value := rc.getValue(docid, bucketDoc.CurrentRev, true)
	docRev, statEvent, err := value.loadForDoc(bucketDoc, context, BodyShallowCopy)
	rc.statsRecorderFunc(statEvent)

	if err != nil {
		rc.removeValue(value) // don't keep failed loads in the cache
	}
	return docRev, err
}

func (rc *RevisionCache) statsRecorderFunc(cacheHit bool) {
	if rc.statsCache == nil {
		return
	}
	if cacheHit {
		rc.statsCache.Add(base.StatKeyRevisionCacheHits, 1)
	} else {
		rc.statsCache.Add(base.StatKeyRevisionCacheMisses, 1)
	}
}

// Adds a revision to the cache.
func (rc *RevisionCache) Put(docid string, docRev DocumentRevision) {
	if docRev.History == nil {
		panic("Missing history for RevisionCache.Put")
	}
	value := rc.getValue(docid, docRev.RevID, true)
	value.store(docRev)
}

func (rc *RevisionCache) getValue(docid, revid string, create bool) (value *revCacheValue) {
	if docid == "" || revid == "" {
		panic("RevisionCache: invalid empty doc/rev id")
	}
	key := IDAndRev{DocID: docid, RevID: revid}
	rc.lock.Lock()
	defer rc.lock.Unlock()
	if elem := rc.cache[key]; elem != nil {
		rc.lruList.MoveToFront(elem)
		value = elem.Value.(*revCacheValue)
	} else if create {
		value = &revCacheValue{key: key}
		rc.cache[key] = rc.lruList.PushFront(value)
		for len(rc.cache) > int(rc.capacity) {
			rc.purgeOldest_()
		}
	}
	return
}

func (rc *RevisionCache) removeValue(value *revCacheValue) {
	rc.lock.Lock()
	if element := rc.cache[value.key]; element != nil && element.Value == value {
		rc.lruList.Remove(element)
		delete(rc.cache, value.key)
	}
	rc.lock.Unlock()
}

func (rc *RevisionCache) purgeOldest_() {
	value := rc.lruList.Remove(rc.lruList.Back()).(*revCacheValue)
	delete(rc.cache, value.key)
}

// Gets the body etc. out of a revCacheValue. If they aren't present already, the loader func
// will be called. This is synchronized so that the loader will only be called once even if
// multiple goroutines try to load at the same time.
func (value *revCacheValue) load(loaderFunc RevisionCacheLoaderFunc, copyType BodyCopyType) (docRev DocumentRevision, cacheHit bool, err error) {

	// Attempt to read cached value
	value.lock.RLock()
	if value.body != nil || value.err != nil {
		docRev, err = value._asDocumentRevision(copyType)
		value.lock.RUnlock()
		return docRev, true, err
	}
	value.lock.RUnlock()

	// Cached value not found - attempt to load if loaderFunc provided
	if loaderFunc == nil {
		return DocumentRevision{}, false, errors.New("No loader function defined for revision cache")
	}
	value.lock.Lock()
	// Check if the value was loaded while we waited for the lock - if so, return.
	if value.body != nil || value.err != nil {
		cacheHit = true
	} else {
		cacheHit = false
		value.body, value.history, value.channels, value.attachments, value.expiry, value.err = loaderFunc(value.key)
	}

	docRev, err = value._asDocumentRevision(copyType)
	value.lock.Unlock()
	return docRev, cacheHit, err
}

// _asDocumentRevision copies the rev cache value into a DocumentRevision.  Requires callers hold at least the read lock on value.lock
func (value *revCacheValue) _asDocumentRevision(copyType BodyCopyType) (DocumentRevision, error) {
	return DocumentRevision{
		RevID:       value.key.RevID,
		Body:        value.body.Copy(copyType), // Never let the caller mutate the stored body
		History:     value.history,
		Channels:    value.channels,
		Expiry:      value.expiry,
		Attachments: value.attachments.ShallowCopy(), // Avoid caller mutating the stored attachments
		Delta:       value.delta,
	}, value.err
}

// Retrieves the body etc. out of a revCacheValue.  If they aren't already present, loads into the cache value using
// the provided document.
func (value *revCacheValue) loadForDoc(doc *document, context *DatabaseContext, copyType BodyCopyType) (docRev DocumentRevision, cacheHit bool, err error) {

	value.lock.RLock()
	if value.body != nil || value.err != nil {
		docRev, err = value._asDocumentRevision(copyType)
		value.lock.RUnlock()
		return docRev, true, err
	}
	value.lock.RUnlock()

	value.lock.Lock()
	// Check if the value was loaded while we waited for the lock - if so, return.
	if value.body != nil || value.err != nil {
		cacheHit = true
	} else {
		cacheHit = false
		value.body, value.history, value.channels, value.attachments, value.expiry, value.err = context.revCacheLoaderForDocument(doc, value.key.RevID)
	}

	docRev, err = value._asDocumentRevision(copyType)
	value.lock.Unlock()
	return docRev, cacheHit, err
}

// Stores a body etc. into a revCacheValue if there isn't one already.
func (value *revCacheValue) store(docRev DocumentRevision) {
	value.lock.Lock()
	defer value.lock.Unlock()
	if value.body == nil {
		value.body = docRev.Body.ShallowCopy() // Don't store a body the caller might later mutate
		value.body[BodyId] = value.key.DocID   // Rev cache includes id and rev in the body.  Ensure they are set in case callers aren't passing
		value.body[BodyRev] = value.key.RevID
		value.history = docRev.History
		value.channels = docRev.Channels
		value.expiry = docRev.Expiry
		value.attachments = docRev.Attachments.ShallowCopy() // Don't store attachments the caller might later mutate
		value.err = nil
	}
}

func (value *revCacheValue) updateDelta(toRevID string, deltaBytes []byte) {
	value.lock.Lock()
	defer value.lock.Unlock()
	value.delta = &RevCacheDelta{
		ToRevID:    toRevID,
		DeltaBytes: deltaBytes,
	}
}
<commit_msg>Avoid defer when unlocking rev cache (#3956)<commit_after>package db

import (
	"container/list"
	"errors"
	"expvar"
	"sync"
	"time"

	"github.com/couchbase/sync_gateway/base"
)

// Number of recently-accessed doc revisions to cache in RAM
var KDefaultRevisionCacheCapacity uint32 = 5000

// An LRU cache of document revision bodies, together with their channel access.
type RevisionCache struct {
	cache      map[IDAndRev]*list.Element // Fast lookup of list element by doc/rev ID
	lruList    *list.List                 // List ordered by most recent access (Front is newest)
	capacity   uint32                     // Max number of revisions to cache
	loaderFunc RevisionCacheLoaderFunc    // Function which does actual loading of something from rev cache
	lock       sync.Mutex                 // For thread-safety
	statsCache *expvar.Map                // Per-db stats related to cache
}

// Revision information as returned by the rev cache
type DocumentRevision struct {
	RevID       string
	Body        Body
	History     Revisions
	Channels    base.Set
	Expiry      *time.Time
	Attachments AttachmentsMeta
	Delta       *RevCacheDelta
}

// Callback function signature for loading something from the rev cache
type RevisionCacheLoaderFunc func(id IDAndRev) (body Body, history Revisions, channels base.Set, attachments AttachmentsMeta, expiry *time.Time, err error)

// The cache payload data. Stored as the Value of a list Element.
type revCacheValue struct {
	key         IDAndRev        // doc/rev IDs
	body        Body            // Revision body (a pristine shallow copy)
	history     Revisions       // Rev history encoded like a "_revisions" property
	channels    base.Set        // Set of channels that have access
	expiry      *time.Time      // Document expiry
	attachments AttachmentsMeta // Document _attachments property
	delta       *RevCacheDelta  // Available delta *from* this revision
	err         error           // Error from loaderFunc if it failed
	lock        sync.RWMutex    // Synchronizes access to this struct
}

type RevCacheDelta struct {
	ToRevID    string
	DeltaBytes []byte
}

// Creates a revision cache with the given capacity and an optional loader function.
func NewRevisionCache(capacity uint32, loaderFunc RevisionCacheLoaderFunc, statsCache *expvar.Map) *RevisionCache {

	if capacity == 0 {
		capacity = KDefaultRevisionCacheCapacity
	}

	return &RevisionCache{
		cache:      map[IDAndRev]*list.Element{},
		lruList:    list.New(),
		capacity:   capacity,
		loaderFunc: loaderFunc,
		statsCache: statsCache,
	}
}

// Looks up a revision from the cache.
// Returns the body of the revision, its history, and the set of channels it's in.
// If the cache has a loaderFunction, it will be called if the revision isn't in the cache;
// any error returned by the loaderFunction will be returned from Get.
func (rc *RevisionCache) Get(docid, revid string) (DocumentRevision, error) {
	return rc.getFromCache(docid, revid, BodyShallowCopy, rc.loaderFunc != nil)
}

// Looks up a revision from the cache only.  Will not fall back to loader function if not
// present in the cache.
func (rc *RevisionCache) GetCached(docid, revid string) (DocumentRevision, error) {
	return rc.getFromCache(docid, revid, BodyShallowCopy, false)
}

// Returns the body of the revision based on the specified body copy policy (deep/shallow/none)
func (rc *RevisionCache) GetWithCopy(docid, revid string, copyType BodyCopyType) (DocumentRevision, error) {
	return rc.getFromCache(docid, revid, copyType, rc.loaderFunc != nil)
}

// Attempt to update the delta on a revision cache entry.  If the entry is no longer resident in the cache,
// fails silently
func (rc *RevisionCache) UpdateDelta(docID, revID string, toRevID string, delta []byte) {
	value := rc.getValue(docID, revID, false)
	if value != nil {
		value.updateDelta(toRevID, delta)
	}
}

func (rc *RevisionCache) getFromCache(docid, revid string, copyType BodyCopyType, loadOnCacheMiss bool) (DocumentRevision, error) {
	value := rc.getValue(docid, revid, loadOnCacheMiss)
	if value == nil {
		return DocumentRevision{}, nil
	}
	docRev, statEvent, err := value.load(rc.loaderFunc, copyType)
	rc.statsRecorderFunc(statEvent)

	if err != nil {
		rc.removeValue(value) // don't keep failed loads in the cache
	}
	return docRev, err
}

// Attempts to retrieve the active revision for a document from the cache.  Requires retrieval
// of the document from the bucket to guarantee the current active revision, but does minimal unmarshalling
// of the retrieved document to get the current rev from _sync metadata.  If active rev is already in the
// rev cache, will use it.  Otherwise will add to the rev cache using the raw document obtained in the
// initial retrieval.
func (rc *RevisionCache) GetActive(docid string, context *DatabaseContext) (docRev DocumentRevision, err error) {

	// Look up active rev for doc
	bucketDoc, getErr := context.GetDocument(docid, DocUnmarshalSync)
	if getErr != nil {
		return DocumentRevision{}, getErr
	}
	if bucketDoc == nil {
		return DocumentRevision{}, nil
	}

	// Retrieve from or add to rev cache
	value := rc.getValue(docid, bucketDoc.CurrentRev, true)
	docRev, statEvent, err := value.loadForDoc(bucketDoc, context, BodyShallowCopy)
	rc.statsRecorderFunc(statEvent)

	if err != nil {
		rc.removeValue(value) // don't keep failed loads in the cache
	}
	return docRev, err
}

func (rc *RevisionCache) statsRecorderFunc(cacheHit bool) {
	if rc.statsCache == nil {
		return
	}
	if cacheHit {
		rc.statsCache.Add(base.StatKeyRevisionCacheHits, 1)
	} else {
		rc.statsCache.Add(base.StatKeyRevisionCacheMisses, 1)
	}
}

// Adds a revision to the cache.
func (rc *RevisionCache) Put(docid string, docRev DocumentRevision) {
	if docRev.History == nil {
		panic("Missing history for RevisionCache.Put")
	}
	value := rc.getValue(docid, docRev.RevID, true)
	value.store(docRev)
}

func (rc *RevisionCache) getValue(docid, revid string, create bool) (value *revCacheValue) {
	if docid == "" || revid == "" {
		panic("RevisionCache: invalid empty doc/rev id")
	}
	key := IDAndRev{DocID: docid, RevID: revid}
	rc.lock.Lock()
	if elem := rc.cache[key]; elem != nil {
		rc.lruList.MoveToFront(elem)
		value = elem.Value.(*revCacheValue)
	} else if create {
		value = &revCacheValue{key: key}
		rc.cache[key] = rc.lruList.PushFront(value)
		for len(rc.cache) > int(rc.capacity) {
			rc.purgeOldest_()
		}
	}
	rc.lock.Unlock()
	return
}

func (rc *RevisionCache) removeValue(value *revCacheValue) {
	rc.lock.Lock()
	if element := rc.cache[value.key]; element != nil && element.Value == value {
		rc.lruList.Remove(element)
		delete(rc.cache, value.key)
	}
	rc.lock.Unlock()
}

func (rc *RevisionCache) purgeOldest_() {
	value := rc.lruList.Remove(rc.lruList.Back()).(*revCacheValue)
	delete(rc.cache, value.key)
}

// Gets the body etc. out of a revCacheValue. If they aren't present already, the loader func
// will be called. This is synchronized so that the loader will only be called once even if
// multiple goroutines try to load at the same time.
func (value *revCacheValue) load(loaderFunc RevisionCacheLoaderFunc, copyType BodyCopyType) (docRev DocumentRevision, cacheHit bool, err error) {

	// Attempt to read cached value
	value.lock.RLock()
	if value.body != nil || value.err != nil {
		docRev, err = value._asDocumentRevision(copyType)
		value.lock.RUnlock()
		return docRev, true, err
	}
	value.lock.RUnlock()

	// Cached value not found - attempt to load if loaderFunc provided
	if loaderFunc == nil {
		return DocumentRevision{}, false, errors.New("No loader function defined for revision cache")
	}
	value.lock.Lock()
	// Check if the value was loaded while we waited for the lock - if so, return.
	if value.body != nil || value.err != nil {
		cacheHit = true
	} else {
		cacheHit = false
		value.body, value.history, value.channels, value.attachments, value.expiry, value.err = loaderFunc(value.key)
	}

	docRev, err = value._asDocumentRevision(copyType)
	value.lock.Unlock()
	return docRev, cacheHit, err
}

// _asDocumentRevision copies the rev cache value into a DocumentRevision.  Requires callers hold at least the read lock on value.lock
func (value *revCacheValue) _asDocumentRevision(copyType BodyCopyType) (DocumentRevision, error) {
	return DocumentRevision{
		RevID:       value.key.RevID,
		Body:        value.body.Copy(copyType), // Never let the caller mutate the stored body
		History:     value.history,
		Channels:    value.channels,
		Expiry:      value.expiry,
		Attachments: value.attachments.ShallowCopy(), // Avoid caller mutating the stored attachments
		Delta:       value.delta,
	}, value.err
}

// Retrieves the body etc. out of a revCacheValue.  If they aren't already present, loads into the cache value using
// the provided document.
func (value *revCacheValue) loadForDoc(doc *document, context *DatabaseContext, copyType BodyCopyType) (docRev DocumentRevision, cacheHit bool, err error) {

	value.lock.RLock()
	if value.body != nil || value.err != nil {
		docRev, err = value._asDocumentRevision(copyType)
		value.lock.RUnlock()
		return docRev, true, err
	}
	value.lock.RUnlock()

	value.lock.Lock()
	// Check if the value was loaded while we waited for the lock - if so, return.
	if value.body != nil || value.err != nil {
		cacheHit = true
	} else {
		cacheHit = false
		value.body, value.history, value.channels, value.attachments, value.expiry, value.err = context.revCacheLoaderForDocument(doc, value.key.RevID)
	}

	docRev, err = value._asDocumentRevision(copyType)
	value.lock.Unlock()
	return docRev, cacheHit, err
}

// Stores a body etc. into a revCacheValue if there isn't one already.
func (value *revCacheValue) store(docRev DocumentRevision) {
	value.lock.Lock()
	defer value.lock.Unlock()
	if value.body == nil {
		value.body = docRev.Body.ShallowCopy() // Don't store a body the caller might later mutate
		value.body[BodyId] = value.key.DocID   // Rev cache includes id and rev in the body.  Ensure they are set in case callers aren't passing
		value.body[BodyRev] = value.key.RevID
		value.history = docRev.History
		value.channels = docRev.Channels
		value.expiry = docRev.Expiry
		value.attachments = docRev.Attachments.ShallowCopy() // Don't store attachments the caller might later mutate
		value.err = nil
	}
}

func (value *revCacheValue) updateDelta(toRevID string, deltaBytes []byte) {
	value.lock.Lock()
	defer value.lock.Unlock()
	value.delta = &RevCacheDelta{
		ToRevID:    toRevID,
		DeltaBytes: deltaBytes,
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * 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.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @Copyright 	2017-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 */

package jwk

import (
	"crypto/rsa"
	"strings"

	jwt2 "github.com/dgrijalva/jwt-go"
	"github.com/ory/fosite/token/jwt"
	"github.com/pkg/errors"
)

func NewRS256JWTStrategy(m Manager, set string) (*RS256JWTStrategy, error) {
	j := &RS256JWTStrategy{
		Manager:          m,
		RS256JWTStrategy: &jwt.RS256JWTStrategy{},
		Set:              set,
	}
	if err := j.refresh(); err != nil {
		return nil, err
	}
	return j, nil
}

type RS256JWTStrategy struct {
	RS256JWTStrategy *jwt.RS256JWTStrategy
	Manager          Manager
	Set              string

	publicKey    *rsa.PublicKey
	privateKey   *rsa.PrivateKey
	publicKeyID  string
	privateKeyID string
}

func (j *RS256JWTStrategy) Hash(in []byte) ([]byte, error) {
	return j.RS256JWTStrategy.Hash(in)
}

// GetSigningMethodLength will return the length of the signing method
func (j *RS256JWTStrategy) GetSigningMethodLength() int {
	return j.RS256JWTStrategy.GetSigningMethodLength()
}

func (j *RS256JWTStrategy) GetSignature(token string) (string, error) {
	return j.RS256JWTStrategy.GetSignature(token)
}

func (j *RS256JWTStrategy) Generate(claims jwt2.Claims, header jwt.Mapper) (string, string, error) {
	if err := j.refresh(); err != nil {
		return "", "", err
	}

	return j.RS256JWTStrategy.Generate(claims, header)
}

func (j *RS256JWTStrategy) Validate(token string) (string, error) {
	if err := j.refresh(); err != nil {
		return "", err
	}

	return j.RS256JWTStrategy.Validate(token)
}

func (j *RS256JWTStrategy) Decode(token string) (*jwt2.Token, error) {
	if err := j.refresh(); err != nil {
		return nil, err
	}

	return j.RS256JWTStrategy.Decode(token)
}

func (j *RS256JWTStrategy) GetPublicKeyID() (string, error) {
	if err := j.refresh(); err != nil {
		return "", err
	}

	return j.publicKeyID, nil
}

func (j *RS256JWTStrategy) refresh() error {
	keys, err := j.Manager.GetKeySet(j.Set)
	if err != nil {
		return err
	}

	public, err := FindKeyByPrefix(keys, "public")
	if err != nil {
		return err
	}

	private, err := FindKeyByPrefix(keys, "private")
	if err != nil {
		return err
	}

	if strings.Replace(public.KeyID, "public:", "", 1) != strings.Replace(private.KeyID, "private:", "", 1) {
		return errors.New("public and private key pair kids do not match")
	}

	if k, ok := private.Key.(*rsa.PrivateKey); !ok {
		return errors.New("unable to type assert key to *rsa.PublicKey")
	} else {
		j.privateKey = k
		j.RS256JWTStrategy.PrivateKey = k
	}

	if k, ok := public.Key.(*rsa.PublicKey); !ok {
		return errors.New("unable to type assert key to *rsa.PublicKey")
	} else {
		j.publicKey = k
		j.publicKeyID = public.KeyID
	}

	if j.privateKey.PublicKey != *j.publicKey {
		return errors.New("public and private key pair fetched from store does not match")
	}

	return nil
}
<commit_msg>jwk: Tests for simple equality in JWT strategy<commit_after>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * 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.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @Copyright 	2017-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 */

package jwk

import (
	"crypto/rsa"
	"strings"

	jwt2 "github.com/dgrijalva/jwt-go"
	"github.com/ory/fosite/token/jwt"
	"github.com/pkg/errors"
)

func NewRS256JWTStrategy(m Manager, set string) (*RS256JWTStrategy, error) {
	j := &RS256JWTStrategy{
		Manager:          m,
		RS256JWTStrategy: &jwt.RS256JWTStrategy{},
		Set:              set,
	}
	if err := j.refresh(); err != nil {
		return nil, err
	}
	return j, nil
}

type RS256JWTStrategy struct {
	RS256JWTStrategy *jwt.RS256JWTStrategy
	Manager          Manager
	Set              string

	publicKey    *rsa.PublicKey
	privateKey   *rsa.PrivateKey
	publicKeyID  string
	privateKeyID string
}

func (j *RS256JWTStrategy) Hash(in []byte) ([]byte, error) {
	return j.RS256JWTStrategy.Hash(in)
}

// GetSigningMethodLength will return the length of the signing method
func (j *RS256JWTStrategy) GetSigningMethodLength() int {
	return j.RS256JWTStrategy.GetSigningMethodLength()
}

func (j *RS256JWTStrategy) GetSignature(token string) (string, error) {
	return j.RS256JWTStrategy.GetSignature(token)
}

func (j *RS256JWTStrategy) Generate(claims jwt2.Claims, header jwt.Mapper) (string, string, error) {
	if err := j.refresh(); err != nil {
		return "", "", err
	}

	return j.RS256JWTStrategy.Generate(claims, header)
}

func (j *RS256JWTStrategy) Validate(token string) (string, error) {
	if err := j.refresh(); err != nil {
		return "", err
	}

	return j.RS256JWTStrategy.Validate(token)
}

func (j *RS256JWTStrategy) Decode(token string) (*jwt2.Token, error) {
	if err := j.refresh(); err != nil {
		return nil, err
	}

	return j.RS256JWTStrategy.Decode(token)
}

func (j *RS256JWTStrategy) GetPublicKeyID() (string, error) {
	if err := j.refresh(); err != nil {
		return "", err
	}

	return j.publicKeyID, nil
}

func (j *RS256JWTStrategy) refresh() error {
	keys, err := j.Manager.GetKeySet(j.Set)
	if err != nil {
		return err
	}

	public, err := FindKeyByPrefix(keys, "public")
	if err != nil {
		return err
	}

	private, err := FindKeyByPrefix(keys, "private")
	if err != nil {
		return err
	}

	if strings.Replace(public.KeyID, "public:", "", 1) != strings.Replace(private.KeyID, "private:", "", 1) {
		return errors.New("public and private key pair kids do not match")
	}

	if k, ok := private.Key.(*rsa.PrivateKey); !ok {
		return errors.New("unable to type assert key to *rsa.PublicKey")
	} else {
		j.privateKey = k
		j.RS256JWTStrategy.PrivateKey = k
	}

	if k, ok := public.Key.(*rsa.PublicKey); !ok {
		return errors.New("unable to type assert key to *rsa.PublicKey")
	} else {
		j.publicKey = k
		j.publicKeyID = public.KeyID
	}

	if j.privateKey.PublicKey.E != j.publicKey.E ||
		j.privateKey.PublicKey.N.String() != j.publicKey.N.String() {
		return errors.New("public and private key pair fetched from store does not match")
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"strings"
	"testing"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/datastore"
	"cloud.google.com/go/storage"
	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
	"golang.org/x/net/context"
	dlppb "google.golang.org/genproto/googleapis/privacy/dlp/v2"
)

func TestInspectString(t *testing.T) {
	testutil.SystemTest(t)
	tests := []struct {
		s    string
		want bool
	}{
		{
			s:    "My SSN is 111222333",
			want: true,
		},
		{
			s: "Does not match",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectString(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, test.s)
		if got := buf.String(); test.want != strings.Contains(got, "US_SOCIAL_SECURITY_NUMBER") {
			if test.want {
				t.Errorf("inspectString(%s) = %q, want 'US_SOCIAL_SECURITY_NUMBER' substring", test.s, got)
			} else {
				t.Errorf("inspectString(%s) = %q, want to not contain 'US_SOCIAL_SECURITY_NUMBER'", test.s, got)
			}
		}
	}
}

func TestInspectFile(t *testing.T) {
	testutil.SystemTest(t)
	tests := []struct {
		s    string
		want bool
	}{
		{
			s:    "My SSN is 111222333",
			want: true,
		},
		{
			s: "Does not match",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectFile(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, dlppb.ByteContentItem_TEXT_UTF8, strings.NewReader(test.s))
		if got := buf.String(); test.want != strings.Contains(got, "US_SOCIAL_SECURITY_NUMBER") {
			if test.want {
				t.Errorf("inspectString(%s) = %q, want 'US_SOCIAL_SECURITY_NUMBER' substring", test.s, got)
			} else {
				t.Errorf("inspectString(%s) = %q, want to not contain 'US_SOCIAL_SECURITY_NUMBER'", test.s, got)
			}
		}
	}
}

const (
	ssnFileName             = "fake_ssn.txt"
	nothingEventfulFileName = "nothing_eventful.txt"
	bucketName              = "golang-samples-dlp-test"
)

func writeTestGCSFiles(t *testing.T, projectID string) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		t.Fatalf("Failed to create client: %v", err)
	}
	bucket := client.Bucket(bucketName)
	_, err = bucket.Attrs(ctx)
	if err != nil {
		switch err {
		case storage.ErrObjectNotExist:
			if err := bucket.Create(ctx, projectID, nil); err != nil {
				t.Fatalf("Failed to create bucket: %v", err)
			}
		default:
			t.Fatalf("error getting bucket attrs: %v", err)
		}
	}
	if err := writeObject(ctx, bucket, ssnFileName, "My SSN is 111222333"); err != nil {
		t.Fatalf("error writing object: %v", err)
	}
	if err := writeObject(ctx, bucket, nothingEventfulFileName, "Nothing eventful"); err != nil {
		t.Fatalf("error writing object: %v", err)
	}
}

func writeObject(ctx context.Context, bucket *storage.BucketHandle, fileName, content string) error {
	obj := bucket.Object(fileName)
	_, err := obj.Attrs(ctx)
	if err != nil {
		switch err {
		case storage.ErrObjectNotExist:
			w := obj.NewWriter(ctx)
			w.Write([]byte(content))
			if err := w.Close(); err != nil {
				return err
			}
		default:
			return err
		}
	}
	return nil
}

func TestInspectGCS(t *testing.T) {
	t.Skip("flaky due to timeout")
	testutil.SystemTest(t)
	writeTestGCSFiles(t, projectID)
	tests := []struct {
		fileName string
		want     string
	}{
		{
			fileName: ssnFileName,
			want:     "US_SOCIAL_SECURITY_NUMBER",
		},
		{
			fileName: nothingEventfulFileName,
			want:     "No results",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectGCSFile(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, "test-topic", "test-sub", bucketName, test.fileName)
		if got := buf.String(); !strings.Contains(got, test.want) {
			t.Errorf("inspectString(%s) = %q, want %q substring", test.fileName, got, test.want)
		}
	}
}

type SSNTask struct {
	Description string
}

type BoringTask struct {
	Description string
}

func writeTestDatastoreFiles(t *testing.T, projectID string) {
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, projectID)
	if err != nil {
		t.Fatalf("Failed to create client: %v", err)
	}
	kind := "SSNTask"
	name := "ssntask1"
	ssnKey := datastore.NameKey(kind, name, nil)
	task := SSNTask{
		Description: "My SSN is 111222333",
	}
	if _, err := client.Put(ctx, ssnKey, &task); err != nil {
		t.Fatalf("Failed to save task: %v", err)
	}

	kind = "BoringTask"
	name = "boringtask1"
	boringKey := datastore.NameKey(kind, name, nil)
	boringTask := BoringTask{
		Description: "Nothing meaningful",
	}
	if _, err := client.Put(ctx, boringKey, &boringTask); err != nil {
		t.Fatalf("Failed to save task: %v", err)
	}
}

func TestInspectDatastore(t *testing.T) {
	t.Skip("flaky due to timeout")
	testutil.SystemTest(t)
	writeTestDatastoreFiles(t, projectID)
	tests := []struct {
		kind string
		want string
	}{
		{
			kind: "SSNTask",
			want: "US_SOCIAL_SECURITY_NUMBER",
		},
		{
			kind: "BoringTask",
			want: "No results",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectDatastore(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, "test-topic", "test-sub", projectID, "", test.kind)
		if got := buf.String(); !strings.Contains(got, test.want) {
			t.Errorf("inspectDatastore(%s) = %q, want %q substring", test.kind, got, test.want)
		}
	}
}

type Item struct {
	Description string
}

const (
	harmlessTable = "harmless"
	harmfulTable  = "harmful"
	bqDatasetID   = "golang_samples_dlp"
)

func createBigqueryTestFiles(projectID, datasetID string) error {
	ctx := context.Background()
	client, err := bigquery.NewClient(ctx, projectID)
	if err != nil {
		return err
	}
	defer client.Close()
	d := client.Dataset(datasetID)
	if _, err := d.Metadata(ctx); err != nil {
		if err := d.Create(ctx, &bigquery.DatasetMetadata{}); err != nil {
			return err
		}
	}
	schema, err := bigquery.InferSchema(Item{})
	if err != nil {
		return err
	}
	if err := uploadBigQuery(ctx, d, schema, harmlessTable, "Nothing meaningful"); err != nil {
		return err
	}
	return uploadBigQuery(ctx, d, schema, harmfulTable, "My SSN is 111222333")
}

func uploadBigQuery(ctx context.Context, d *bigquery.Dataset, schema bigquery.Schema, table, content string) error {
	t := d.Table(table)
	if _, err := t.Metadata(ctx); err == nil {
		return nil
	}
	if err := t.Create(ctx, &bigquery.TableMetadata{Schema: schema}); err != nil {
		return err
	}
	source := bigquery.NewReaderSource(strings.NewReader(content))
	l := t.LoaderFrom(source)
	job, err := l.Run(ctx)
	if err != nil {
		return err
	}
	status, err := job.Wait(ctx)
	if err != nil {
		return err
	}
	return status.Err()
}

func TestInspectBigquery(t *testing.T) {
	t.Skip("flaky due to timeout")
	testutil.SystemTest(t)
	if err := createBigqueryTestFiles(projectID, bqDatasetID); err != nil {
		t.Fatalf("error creating test BigQuery files: %v", err)
	}
	tests := []struct {
		table string
		want  string
	}{
		{
			table: harmfulTable,
			want:  "US_SOCIAL_SECURITY_NUMBER",
		},
		{
			table: harmlessTable,
			want:  "No results",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectBigquery(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, "test-topic", "test-sub", projectID, bqDatasetID, test.table)
		if got := buf.String(); !strings.Contains(got, test.want) {
			t.Errorf("inspectBigquery(%s) = %q, want %q substring", test.table, got, test.want)
		}
	}
}
<commit_msg>dlp: re-enable tests (#509)<commit_after>// Copyright 2018 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"strings"
	"testing"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/datastore"
	"cloud.google.com/go/storage"
	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
	"golang.org/x/net/context"
	dlppb "google.golang.org/genproto/googleapis/privacy/dlp/v2"
)

func TestInspectString(t *testing.T) {
	testutil.SystemTest(t)
	tests := []struct {
		s    string
		want bool
	}{
		{
			s:    "My SSN is 111222333",
			want: true,
		},
		{
			s: "Does not match",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectString(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, test.s)
		if got := buf.String(); test.want != strings.Contains(got, "US_SOCIAL_SECURITY_NUMBER") {
			if test.want {
				t.Errorf("inspectString(%s) = %q, want 'US_SOCIAL_SECURITY_NUMBER' substring", test.s, got)
			} else {
				t.Errorf("inspectString(%s) = %q, want to not contain 'US_SOCIAL_SECURITY_NUMBER'", test.s, got)
			}
		}
	}
}

func TestInspectFile(t *testing.T) {
	testutil.SystemTest(t)
	tests := []struct {
		s    string
		want bool
	}{
		{
			s:    "My SSN is 111222333",
			want: true,
		},
		{
			s: "Does not match",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectFile(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, dlppb.ByteContentItem_TEXT_UTF8, strings.NewReader(test.s))
		if got := buf.String(); test.want != strings.Contains(got, "US_SOCIAL_SECURITY_NUMBER") {
			if test.want {
				t.Errorf("inspectString(%s) = %q, want 'US_SOCIAL_SECURITY_NUMBER' substring", test.s, got)
			} else {
				t.Errorf("inspectString(%s) = %q, want to not contain 'US_SOCIAL_SECURITY_NUMBER'", test.s, got)
			}
		}
	}
}

const (
	ssnFileName             = "fake_ssn.txt"
	nothingEventfulFileName = "nothing_eventful.txt"
	bucketName              = "golang-samples-dlp-test"
)

func writeTestGCSFiles(t *testing.T, projectID string) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		t.Fatalf("Failed to create client: %v", err)
	}
	bucket := client.Bucket(bucketName)
	_, err = bucket.Attrs(ctx)
	if err != nil {
		switch err {
		case storage.ErrObjectNotExist:
			if err := bucket.Create(ctx, projectID, nil); err != nil {
				t.Fatalf("Failed to create bucket: %v", err)
			}
		default:
			t.Fatalf("error getting bucket attrs: %v", err)
		}
	}
	if err := writeObject(ctx, bucket, ssnFileName, "My SSN is 111222333"); err != nil {
		t.Fatalf("error writing object: %v", err)
	}
	if err := writeObject(ctx, bucket, nothingEventfulFileName, "Nothing eventful"); err != nil {
		t.Fatalf("error writing object: %v", err)
	}
}

func writeObject(ctx context.Context, bucket *storage.BucketHandle, fileName, content string) error {
	obj := bucket.Object(fileName)
	_, err := obj.Attrs(ctx)
	if err != nil {
		switch err {
		case storage.ErrObjectNotExist:
			w := obj.NewWriter(ctx)
			w.Write([]byte(content))
			if err := w.Close(); err != nil {
				return err
			}
		default:
			return err
		}
	}
	return nil
}

func TestInspectGCS(t *testing.T) {
	testutil.SystemTest(t)
	writeTestGCSFiles(t, projectID)
	tests := []struct {
		fileName string
		want     string
	}{
		{
			fileName: ssnFileName,
			want:     "US_SOCIAL_SECURITY_NUMBER",
		},
		{
			fileName: nothingEventfulFileName,
			want:     "No results",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectGCSFile(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, "test-topic", "test-sub", bucketName, test.fileName)
		if got := buf.String(); !strings.Contains(got, test.want) {
			t.Errorf("inspectString(%s) = %q, want %q substring", test.fileName, got, test.want)
		}
	}
}

type SSNTask struct {
	Description string
}

type BoringTask struct {
	Description string
}

func writeTestDatastoreFiles(t *testing.T, projectID string) {
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, projectID)
	if err != nil {
		t.Fatalf("Failed to create client: %v", err)
	}
	kind := "SSNTask"
	name := "ssntask1"
	ssnKey := datastore.NameKey(kind, name, nil)
	task := SSNTask{
		Description: "My SSN is 111222333",
	}
	if _, err := client.Put(ctx, ssnKey, &task); err != nil {
		t.Fatalf("Failed to save task: %v", err)
	}

	kind = "BoringTask"
	name = "boringtask1"
	boringKey := datastore.NameKey(kind, name, nil)
	boringTask := BoringTask{
		Description: "Nothing meaningful",
	}
	if _, err := client.Put(ctx, boringKey, &boringTask); err != nil {
		t.Fatalf("Failed to save task: %v", err)
	}
}

func TestInspectDatastore(t *testing.T) {
	testutil.SystemTest(t)
	writeTestDatastoreFiles(t, projectID)
	tests := []struct {
		kind string
		want string
	}{
		{
			kind: "SSNTask",
			want: "US_SOCIAL_SECURITY_NUMBER",
		},
		{
			kind: "BoringTask",
			want: "No results",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectDatastore(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, "test-topic", "test-sub", projectID, "", test.kind)
		if got := buf.String(); !strings.Contains(got, test.want) {
			t.Errorf("inspectDatastore(%s) = %q, want %q substring", test.kind, got, test.want)
		}
	}
}

type Item struct {
	Description string
}

const (
	harmlessTable = "harmless"
	harmfulTable  = "harmful"
	bqDatasetID   = "golang_samples_dlp"
)

func createBigqueryTestFiles(projectID, datasetID string) error {
	ctx := context.Background()
	client, err := bigquery.NewClient(ctx, projectID)
	if err != nil {
		return err
	}
	defer client.Close()
	d := client.Dataset(datasetID)
	if _, err := d.Metadata(ctx); err != nil {
		if err := d.Create(ctx, &bigquery.DatasetMetadata{}); err != nil {
			return err
		}
	}
	schema, err := bigquery.InferSchema(Item{})
	if err != nil {
		return err
	}
	if err := uploadBigQuery(ctx, d, schema, harmlessTable, "Nothing meaningful"); err != nil {
		return err
	}
	return uploadBigQuery(ctx, d, schema, harmfulTable, "My SSN is 111222333")
}

func uploadBigQuery(ctx context.Context, d *bigquery.Dataset, schema bigquery.Schema, table, content string) error {
	t := d.Table(table)
	if _, err := t.Metadata(ctx); err == nil {
		return nil
	}
	if err := t.Create(ctx, &bigquery.TableMetadata{Schema: schema}); err != nil {
		return err
	}
	source := bigquery.NewReaderSource(strings.NewReader(content))
	l := t.LoaderFrom(source)
	job, err := l.Run(ctx)
	if err != nil {
		return err
	}
	status, err := job.Wait(ctx)
	if err != nil {
		return err
	}
	return status.Err()
}

func TestInspectBigquery(t *testing.T) {
	testutil.SystemTest(t)
	if err := createBigqueryTestFiles(projectID, bqDatasetID); err != nil {
		t.Fatalf("error creating test BigQuery files: %v", err)
	}
	tests := []struct {
		table string
		want  string
	}{
		{
			table: harmfulTable,
			want:  "US_SOCIAL_SECURITY_NUMBER",
		},
		{
			table: harmlessTable,
			want:  "No results",
		},
	}
	for _, test := range tests {
		buf := new(bytes.Buffer)
		inspectBigquery(buf, client, projectID, dlppb.Likelihood_POSSIBLE, 0, true, []string{"US_SOCIAL_SECURITY_NUMBER"}, "test-topic", "test-sub", projectID, bqDatasetID, test.table)
		if got := buf.String(); !strings.Contains(got, test.want) {
			t.Errorf("inspectBigquery(%s) = %q, want %q substring", test.table, got, test.want)
		}
	}
}
<|endoftext|>


<commit_before>package vault

import (
	"context"
	"sort"

	"github.com/hashicorp/vault/helper/namespace"
	"github.com/hashicorp/vault/logical"
)

// Capabilities is used to fetch the capabilities of the given token on the
// given path
func (c *Core) Capabilities(ctx context.Context, token, path string) ([]string, error) {
	if path == "" {
		return nil, &logical.StatusBadRequest{Err: "missing path"}
	}

	if token == "" {
		return nil, &logical.StatusBadRequest{Err: "missing token"}
	}

	te, err := c.tokenStore.Lookup(ctx, token)
	if err != nil {
		return nil, err
	}
	if te == nil {
		return nil, &logical.StatusBadRequest{Err: "invalid token"}
	}

	tokenNS, err := NamespaceByID(ctx, te.NamespaceID, c)
	if err != nil {
		return nil, err
	}
	if tokenNS == nil {
		return nil, namespace.ErrNoNamespace
	}

	var policyCount int
	policyNames := make(map[string][]string)
	policyNames[tokenNS.ID] = te.Policies
	policyCount += len(te.Policies)

	// Attach token's namespace information to the context
	ctx = namespace.ContextWithNamespace(ctx, tokenNS)
	entity, identityPolicies, err := c.fetchEntityAndDerivedPolicies(ctx, tokenNS, te.EntityID)
	if err != nil {
		return nil, err
	}
	if entity != nil && entity.Disabled {
		c.logger.Warn("permission denied as the entity on the token is disabled")
		return nil, logical.ErrPermissionDenied
	}
	if te.EntityID != "" && entity == nil {
		c.logger.Warn("permission denied as the entity on the token is invalid")
		return nil, logical.ErrPermissionDenied
	}

	for nsID, nsPolicies := range identityPolicies {
		policyNames[nsID] = append(policyNames[nsID], nsPolicies...)
		policyCount += len(nsPolicies)
	}

	if policyCount == 0 {
		return []string{DenyCapability}, nil
	}

	acl, err := c.policyStore.ACL(ctx, entity, policyNames)
	if err != nil {
		return nil, err
	}

	capabilities := acl.Capabilities(ctx, path)
	sort.Strings(capabilities)
	return capabilities, nil
}
<commit_msg>Fix Capabilities check when in a child namespace (#5406)<commit_after>package vault

import (
	"context"
	"sort"

	"github.com/hashicorp/vault/helper/namespace"
	"github.com/hashicorp/vault/logical"
)

// Capabilities is used to fetch the capabilities of the given token on the
// given path
func (c *Core) Capabilities(ctx context.Context, token, path string) ([]string, error) {
	if path == "" {
		return nil, &logical.StatusBadRequest{Err: "missing path"}
	}

	if token == "" {
		return nil, &logical.StatusBadRequest{Err: "missing token"}
	}

	te, err := c.tokenStore.Lookup(ctx, token)
	if err != nil {
		return nil, err
	}
	if te == nil {
		return nil, &logical.StatusBadRequest{Err: "invalid token"}
	}

	tokenNS, err := NamespaceByID(ctx, te.NamespaceID, c)
	if err != nil {
		return nil, err
	}
	if tokenNS == nil {
		return nil, namespace.ErrNoNamespace
	}

	var policyCount int
	policyNames := make(map[string][]string)
	policyNames[tokenNS.ID] = te.Policies
	policyCount += len(te.Policies)

	entity, identityPolicies, err := c.fetchEntityAndDerivedPolicies(ctx, tokenNS, te.EntityID)
	if err != nil {
		return nil, err
	}
	if entity != nil && entity.Disabled {
		c.logger.Warn("permission denied as the entity on the token is disabled")
		return nil, logical.ErrPermissionDenied
	}
	if te.EntityID != "" && entity == nil {
		c.logger.Warn("permission denied as the entity on the token is invalid")
		return nil, logical.ErrPermissionDenied
	}

	for nsID, nsPolicies := range identityPolicies {
		policyNames[nsID] = append(policyNames[nsID], nsPolicies...)
		policyCount += len(nsPolicies)
	}

	if policyCount == 0 {
		return []string{DenyCapability}, nil
	}

	// Construct the corresponding ACL object. ACL construction should be
	// performed on the token's namespace.
	tokenCtx := namespace.ContextWithNamespace(ctx, tokenNS)
	acl, err := c.policyStore.ACL(tokenCtx, entity, policyNames)
	if err != nil {
		return nil, err
	}

	capabilities := acl.Capabilities(ctx, path)
	sort.Strings(capabilities)
	return capabilities, nil
}
<|endoftext|>