<commit_before>// Copyright 2016-2018 Authors of Cilium
//
// 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 proxy

import (
	"fmt"
	"sync"

	"github.com/cilium/cilium/pkg/completion"
	"github.com/cilium/cilium/pkg/envoy"

	"github.com/spf13/viper"
)

// the global Envoy instance
var envoyProxy *envoy.Envoy

// envoyRedirect implements the Redirect interface for an l7 proxy.
type envoyRedirect struct {
	redirectID string
	xdsServer  *envoy.XDSServer
}

var envoyOnce sync.Once

// createEnvoyRedirect creates a redirect with corresponding proxy
// configuration. This will launch a proxy instance.
func createEnvoyRedirect(r *Redirect, stateDir string, xdsServer *envoy.XDSServer, wg *completion.WaitGroup) (RedirectImplementation, error) {
	envoyOnce.Do(func() {
		// Start Envoy on first invocation
		envoyProxy = envoy.StartEnvoy(9901, stateDir, viper.GetString("envoy-log"), 0)
	})

	if envoyProxy != nil {
		redirectID := r.id
		redir := &envoyRedirect{
			redirectID: redirectID,
			xdsServer:  xdsServer,
		}

		ip := r.localEndpoint.GetIPv4Address()
		if ip == "" {
			ip = r.localEndpoint.GetIPv6Address()
		}
		if ip == "" {
			return nil, fmt.Errorf("%s: Cannot create redirect, proxy local endpoint has no IP address", r.id)
		}
		xdsServer.AddListener(redirectID, ip, r.ProxyPort, r.ingress, wg)

		return redir, nil
	}

	return nil, fmt.Errorf("%s: Envoy proxy process failed to start, cannot add redirect", r.id)
}

// UpdateRules replaces old l7 rules of a redirect with new ones.
func (r *envoyRedirect) UpdateRules(wg *completion.WaitGroup) error {
	return nil
}

// Close the redirect.
func (r *envoyRedirect) Close(wg *completion.WaitGroup) {
	if envoyProxy != nil {
		r.xdsServer.RemoveListener(r.redirectID, wg)
	}
}
<commit_msg>envoy: Include the redirect port into each Envoy listener name<commit_after>// Copyright 2016-2018 Authors of Cilium
//
// 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 proxy

import (
	"fmt"
	"sync"

	"github.com/cilium/cilium/pkg/completion"
	"github.com/cilium/cilium/pkg/envoy"

	"github.com/spf13/viper"
)

// the global Envoy instance
var envoyProxy *envoy.Envoy

// envoyRedirect implements the Redirect interface for an l7 proxy.
type envoyRedirect struct {
	listenerName string
	xdsServer    *envoy.XDSServer
}

var envoyOnce sync.Once

// createEnvoyRedirect creates a redirect with corresponding proxy
// configuration. This will launch a proxy instance.
func createEnvoyRedirect(r *Redirect, stateDir string, xdsServer *envoy.XDSServer, wg *completion.WaitGroup) (RedirectImplementation, error) {
	envoyOnce.Do(func() {
		// Start Envoy on first invocation
		envoyProxy = envoy.StartEnvoy(9901, stateDir, viper.GetString("envoy-log"), 0)
	})

	if envoyProxy != nil {
		redir := &envoyRedirect{
			listenerName: fmt.Sprintf("%s:%d", r.id, r.ProxyPort),
			xdsServer:    xdsServer,
		}

		ip := r.localEndpoint.GetIPv4Address()
		if ip == "" {
			ip = r.localEndpoint.GetIPv6Address()
		}
		if ip == "" {
			return nil, fmt.Errorf("%s: Cannot create redirect, proxy local endpoint has no IP address", r.id)
		}
		xdsServer.AddListener(redir.listenerName, ip, r.ProxyPort, r.ingress, wg)

		return redir, nil
	}

	return nil, fmt.Errorf("%s: Envoy proxy process failed to start, cannot add redirect", r.id)
}

// UpdateRules replaces old l7 rules of a redirect with new ones.
func (r *envoyRedirect) UpdateRules(wg *completion.WaitGroup) error {
	return nil
}

// Close the redirect.
func (r *envoyRedirect) Close(wg *completion.WaitGroup) {
	if envoyProxy != nil {
		r.xdsServer.RemoveListener(r.listenerName, wg)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// topotools package contains a few utility classes to handle topo.Server
// objects, and transitions.
package topotools

import (
	"sync"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/topo"
)

// CopyKeyspaces will create the keyspaces in the destination topo
func CopyKeyspaces(fromTS, toTS topo.Server) {
	keyspaces, err := fromTS.GetKeyspaces()
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			if err := toTS.CreateKeyspace(keyspace); err != nil {
				if err == topo.ErrNodeExists {
					log.Warningf("keyspace %v already exists", keyspace)
				} else {
					rec.RecordError(err)
				}
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyKeyspaces failed: %v", rec.Error())
	}
}

// CopyShards will create the shards in the destination topo
func CopyShards(fromTS, toTS topo.Server, deleteKeyspaceShards bool) {
	keyspaces, err := fromTS.GetKeyspaces()
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(keyspace)
			if err != nil {
				rec.RecordError(err)
				return
			}

			if deleteKeyspaceShards {
				if err := toTS.DeleteKeyspaceShards(keyspace); err != nil {
					rec.RecordError(err)
					return
				}
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()
					if err := topo.CreateShard(toTS, keyspace, shard); err != nil {
						if err == topo.ErrNodeExists {
							log.Warningf("shard %v/%v already exists", keyspace, shard)
						} else {
							rec.RecordError(err)
							return
						}
					}

					si, err := fromTS.GetShard(keyspace, shard)
					if err != nil {
						rec.RecordError(err)
						return
					}

					if err := toTS.UpdateShard(si); err != nil {
						rec.RecordError(err)
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}

// CopyTablets will create the tablets in the destination topo
func CopyTablets(fromTS, toTS topo.Server) {
	cells, err := fromTS.GetKnownCells()
	if err != nil {
		log.Fatalf("fromTS.GetKnownCells failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, cell := range cells {
		wg.Add(1)
		go func(cell string) {
			defer wg.Done()
			tabletAliases, err := fromTS.GetTabletsByCell(cell)
			if err != nil {
				rec.RecordError(err)
			} else {
				for _, tabletAlias := range tabletAliases {
					wg.Add(1)
					go func(tabletAlias topo.TabletAlias) {
						defer wg.Done()

						// read the source tablet
						ti, err := fromTS.GetTablet(tabletAlias)
						if err != nil {
							rec.RecordError(err)
							return
						}

						// try to create the destination
						err = toTS.CreateTablet(ti.Tablet)
						if err == topo.ErrNodeExists {
							// update the destination tablet
							log.Warningf("tablet %v already exists, updating it", tabletAlias)
							err = toTS.UpdateTabletFields(ti.GetAlias(), func(t *topo.Tablet) error {
								*t = *ti.Tablet
								return nil
							})
						}
						if err != nil {
							rec.RecordError(err)
							return
						}
					}(tabletAlias)
				}
			}
		}(cell)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyTablets failed: %v", rec.Error())
	}
}

// CopyShardReplications will create the ShardReplication objects in
// the destination topo
func CopyShardReplications(fromTS, toTS topo.Server) {
	keyspaces, err := fromTS.GetKeyspaces()
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(keyspace)
			if err != nil {
				rec.RecordError(err)
				return
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()

					// read the source shard to get the cells
					si, err := fromTS.GetShard(keyspace, shard)
					if err != nil {
						rec.RecordError(err)
						return
					}

					for _, cell := range si.Cells {
						sri, err := fromTS.GetShardReplication(cell, keyspace, shard)
						if err != nil {
							rec.RecordError(err)
							continue
						}

						err = toTS.CreateShardReplication(cell, keyspace, shard, sri.ShardReplication)
						switch err {
						case nil:
							// good
						case topo.ErrNodeExists:
							if err := toTS.UpdateShardReplicationFields(cell, keyspace, shard, func(oldSR *topo.ShardReplication) error {
								*oldSR = *sri.ShardReplication
								return nil
							}); err != nil {
								rec.RecordError(err)
							}
						default:
							rec.RecordError(err)
						}
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}
<commit_msg>More descriptive error messages.<commit_after>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// topotools package contains a few utility classes to handle topo.Server
// objects, and transitions.
package topotools

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/topo"
)

// CopyKeyspaces will create the keyspaces in the destination topo
func CopyKeyspaces(fromTS, toTS topo.Server) {
	keyspaces, err := fromTS.GetKeyspaces()
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			if err := toTS.CreateKeyspace(keyspace); err != nil {
				if err == topo.ErrNodeExists {
					log.Warningf("keyspace %v already exists", keyspace)
				} else {
					rec.RecordError(fmt.Errorf("Failed to create keyspace %v: %v", keyspace, err))
				}
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyKeyspaces failed: %v", rec.Error())
	}
}

// CopyShards will create the shards in the destination topo
func CopyShards(fromTS, toTS topo.Server, deleteKeyspaceShards bool) {
	keyspaces, err := fromTS.GetKeyspaces()
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetShardNames(%v) failed: %v", keyspace, err))
				return
			}

			if deleteKeyspaceShards {
				if err := toTS.DeleteKeyspaceShards(keyspace); err != nil {
					rec.RecordError(fmt.Errorf("DeleteKeyspaceShards(%v) failed: %v", keyspace, err))
					return
				}
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()
					if err := topo.CreateShard(toTS, keyspace, shard); err != nil {
						if err == topo.ErrNodeExists {
							log.Warningf("shard %v/%v already exists", keyspace, shard)
						} else {
							rec.RecordError(fmt.Errorf("CreateShard(%v, %v) failed: %v", keyspace, shard, err))
							return
						}
					}

					si, err := fromTS.GetShard(keyspace, shard)
					if err != nil {
						rec.RecordError(fmt.Errorf("GetShard(%v, %v) failed: %v", keyspace, shard, err))
						return
					}

					if err := toTS.UpdateShard(si); err != nil {
						rec.RecordError(fmt.Errorf("UpdateShard(%v, %v) failed: %v", keyspace, shard, err))
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}

// CopyTablets will create the tablets in the destination topo
func CopyTablets(fromTS, toTS topo.Server) {
	cells, err := fromTS.GetKnownCells()
	if err != nil {
		log.Fatalf("fromTS.GetKnownCells failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, cell := range cells {
		wg.Add(1)
		go func(cell string) {
			defer wg.Done()
			tabletAliases, err := fromTS.GetTabletsByCell(cell)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetTabletsByCell(%v) failed: %v", cell, err))
			} else {
				for _, tabletAlias := range tabletAliases {
					wg.Add(1)
					go func(tabletAlias topo.TabletAlias) {
						defer wg.Done()

						// read the source tablet
						ti, err := fromTS.GetTablet(tabletAlias)
						if err != nil {
							rec.RecordError(fmt.Errorf("GetTablet(%v) failed: %v", tabletAlias, err))
							return
						}

						// try to create the destination
						err = toTS.CreateTablet(ti.Tablet)
						if err == topo.ErrNodeExists {
							// update the destination tablet
							log.Warningf("tablet %v already exists, updating it", tabletAlias)
							err = toTS.UpdateTabletFields(ti.GetAlias(), func(t *topo.Tablet) error {
								*t = *ti.Tablet
								return nil
							})
						}
						if err != nil {
							rec.RecordError(fmt.Errorf("CreateTablet(%v) failed: %v", tabletAlias, err))
							return
						}
					}(tabletAlias)
				}
			}
		}(cell)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyTablets failed: %v", rec.Error())
	}
}

// CopyShardReplications will create the ShardReplication objects in
// the destination topo
func CopyShardReplications(fromTS, toTS topo.Server) {
	keyspaces, err := fromTS.GetKeyspaces()
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces failed: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetShardNames(%v) failed: %v", keyspace, err))
				return
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()

					// read the source shard to get the cells
					si, err := fromTS.GetShard(keyspace, shard)
					if err != nil {
						rec.RecordError(fmt.Errorf("GetShard(%v, %v) failed: %v", keyspace, shard, err))
						return
					}

					for _, cell := range si.Cells {
						sri, err := fromTS.GetShardReplication(cell, keyspace, shard)
						if err != nil {
							rec.RecordError(fmt.Errorf("GetShardReplication(%v, %v, %v) failed: %v", cell, keyspace, shard, err))
							continue
						}

						err = toTS.CreateShardReplication(cell, keyspace, shard, sri.ShardReplication)
						switch err {
						case nil:
							// good
						case topo.ErrNodeExists:
							if err := toTS.UpdateShardReplicationFields(cell, keyspace, shard, func(oldSR *topo.ShardReplication) error {
								*oldSR = *sri.ShardReplication
								return nil
							}); err != nil {
								rec.RecordError(fmt.Errorf("UpdateShardReplicationFields(%v, %v, %v) failed: %v", cell, keyspace, shard, err))
							}
						default:
							rec.RecordError(fmt.Errorf("CreateShardReplication(%v, %v, %v) failed: %v", cell, keyspace, shard, err))
						}
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}
<|endoftext|>


<commit_before>package auth

import (
	"bytes"
	"context"
	"encoding/base64"
	"net/http"
	"strings"
)

const (
	schemeBasic         = "basic"
	schemeBearer        = "bearer"
	authorizationHeader = "Authorization"
)

// RequireUser creates auth middleware with given configuration.
func RequireUser(config *Config) func(http.Handler) http.Handler {
	config = config.setDefaults()
	return func(next http.Handler) http.Handler {
		return &middleware{
			config: config,
			next:   next,
		}
	}
}

// RequireAdmin creates auth middleware that authenticates only admin users.
func RequireAdmin(config *Config) func(http.Handler) http.Handler {
	config = config.setDefaults()
	return func(next http.Handler) http.Handler {
		return &middleware{
			config:       config,
			next:         next,
			requireAdmin: true,
		}
	}
}

type middleware struct {
	config       *Config
	next         http.Handler
	requireAdmin bool
}

// ServeHTTP implementation.
func (m *middleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx, err := m.authenticate(r)
	if err == nil {
		if ctx != nil && ctx != r.Context() {
			r = r.WithContext(ctx)
		}
		m.next.ServeHTTP(w, r)
	} else {
		sendError(w, err)
	}
}

// Validates auth header or auth_token.
func (m *middleware) authenticate(r *http.Request) (context.Context, *Error) {
	var h = r.Header.Get(authorizationHeader)
	if len(h) > 0 {
		return m.validateHeader(r, h)
	}

	// auth token as part of url
	if r.Method == "GET" {
		var token = r.URL.Query().Get(m.config.TokenKey)
		if len(token) > 0 {
			return m.jwtHandler(r, token)
		}
	}

	return nil, errBadAuthorizationHeader
}

// Validates authorization header.
func (m *middleware) validateHeader(r *http.Request, auth string) (context.Context, *Error) {
	if len(auth) == 0 {
		return nil, errBadAuthorizationHeader
	}

	var f = strings.Fields(auth)
	if len(f) != 2 {
		return nil, errBadAuthorizationHeader
	}

	var scheme = strings.ToLower(f[0])
	var token = f[1]

	switch scheme {
	case schemeBasic:
		return m.basicHandler(r, token)
	case schemeBearer:
		return m.jwtHandler(r, token)
	default:
		return nil, errUnsupportedAuthScheme
	}
}

func (m *middleware) basicHandler(r *http.Request, tokenString string) (context.Context, *Error) {
	var str, err = base64.StdEncoding.DecodeString(tokenString)
	if err != nil {
		return nil, errBadAuthorizationHeader.cause(err)
	}

	creds := bytes.SplitN(str, []byte(":"), 2)
	userName := string(creds[0])
	user, err := m.config.UserStore.ValidateCredentials(userName, string(creds[1]))
	if err != nil {
		return nil, errBadCredentials.cause(err)
	}

	return m.validateUser(r, user)
}

func (m *middleware) jwtHandler(r *http.Request, tokenString string) (context.Context, *Error) {
	token, err := parseToken(m.config, tokenString, getClientIP(r), false)
	if err != nil {
		return nil, err
	}

	user, error := m.config.UserStore.FindUserByID(token.UserID)
	if error != nil {
		return nil, errUserNotFound.cause(error)
	}

	return m.validateUser(r, user)
}

func (m *middleware) validateUser(r *http.Request, user User) (context.Context, *Error) {
	err := m.checkUser(user)
	if err != nil {
		return nil, err
	}
	return withUser(r.Context(), user), nil
}

func (m *middleware) checkUser(user User) *Error {
	if m.requireAdmin && !user.IsAdmin() {
		return errNotAdmin
	}
	return nil
}
<commit_msg>reuse BasicAuth<commit_after>package auth

import (
	"context"
	"net/http"
	"strings"
)

const (
	schemeBasic         = "basic"
	schemeBearer        = "bearer"
	authorizationHeader = "Authorization"
)

// RequireUser creates auth middleware with given configuration.
func RequireUser(config *Config) func(http.Handler) http.Handler {
	config = config.setDefaults()
	return func(next http.Handler) http.Handler {
		return &middleware{
			config: config,
			next:   next,
		}
	}
}

// RequireAdmin creates auth middleware that authenticates only admin users.
func RequireAdmin(config *Config) func(http.Handler) http.Handler {
	config = config.setDefaults()
	return func(next http.Handler) http.Handler {
		return &middleware{
			config:       config,
			next:         next,
			requireAdmin: true,
		}
	}
}

type middleware struct {
	config       *Config
	next         http.Handler
	requireAdmin bool
}

// ServeHTTP implementation.
func (m *middleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx, err := m.authenticate(r)
	if err == nil {
		if ctx != nil && ctx != r.Context() {
			r = r.WithContext(ctx)
		}
		m.next.ServeHTTP(w, r)
	} else {
		sendError(w, err)
	}
}

// Validates auth header or auth_token.
func (m *middleware) authenticate(r *http.Request) (context.Context, *Error) {
	var h = r.Header.Get(authorizationHeader)
	if len(h) > 0 {
		return m.validateHeader(r, h)
	}

	// auth token as part of url
	if r.Method == "GET" {
		var token = r.URL.Query().Get(m.config.TokenKey)
		if len(token) > 0 {
			return m.validateJWT(r, token)
		}
	}

	return nil, errBadAuthorizationHeader
}

// Validates authorization header.
func (m *middleware) validateHeader(r *http.Request, auth string) (context.Context, *Error) {
	if len(auth) == 0 {
		return nil, errBadAuthorizationHeader
	}

	var f = strings.Fields(auth)
	if len(f) != 2 {
		return nil, errBadAuthorizationHeader
	}

	var scheme = strings.ToLower(f[0])
	var token = f[1]

	switch scheme {
	case schemeBasic:
		return m.validateBasicAuth(r)
	case schemeBearer:
		return m.validateJWT(r, token)
	default:
		return nil, errUnsupportedAuthScheme
	}
}

func (m *middleware) validateBasicAuth(r *http.Request) (context.Context, *Error) {
	username, password, ok := r.BasicAuth()
	if !ok {
		return nil, errBadAuthorizationHeader
	}

	user, err := m.config.UserStore.ValidateCredentials(username, password)
	if err != nil {
		return nil, errBadCredentials.cause(err)
	}

	return m.validateUser(r, user)
}

func (m *middleware) validateJWT(r *http.Request, tokenString string) (context.Context, *Error) {
	token, err := parseToken(m.config, tokenString, getClientIP(r), false)
	if err != nil {
		return nil, err
	}

	user, error := m.config.UserStore.FindUserByID(token.UserID)
	if error != nil {
		return nil, errUserNotFound.cause(error)
	}

	return m.validateUser(r, user)
}

func (m *middleware) validateUser(r *http.Request, user User) (context.Context, *Error) {
	err := m.checkUser(user)
	if err != nil {
		return nil, err
	}
	return withUser(r.Context(), user), nil
}

func (m *middleware) checkUser(user User) *Error {
	if m.requireAdmin && !user.IsAdmin() {
		return errNotAdmin
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Google LLC
//
// 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 storage

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"net/http"
	"strings"
	"testing"

	"cloud.google.com/go/internal/testutil"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"
)

var testEncryptionKey = []byte("secret-key-that-is-32-bytes-long")

func TestErrorOnObjectsInsertCall(t *testing.T) {
	t.Skip("Skipping until https://github.com/googleapis/google-api-go-client/issues/392 is resolved")

	t.Parallel()
	ctx := context.Background()
	const contents = "hello world"

	doWrite := func(mt *mockTransport) *Writer {
		client := mockClient(t, mt)
		wc := client.Bucket("bucketname").Object("filename1").NewWriter(ctx)
		wc.ContentType = "text/plain"

		// We can't check that the Write fails, since it depends on the write to the
		// underling mockTransport failing which is racy.
		wc.Write([]byte(contents))
		return wc
	}

	wc := doWrite(&mockTransport{})
	// Close must always return an error though since it waits for the transport to
	// have closed.
	if err := wc.Close(); err == nil {
		t.Errorf("expected error on close, got nil")
	}

	// Retry on 5xx
	mt := &mockTransport{}
	mt.addResult(&http.Response{StatusCode: 503, Body: bodyReader("")}, nil)
	mt.addResult(&http.Response{StatusCode: 200, Body: bodyReader("{}")}, nil)

	wc = doWrite(mt)
	if err := wc.Close(); err != nil {
		t.Errorf("got %v, want nil", err)
	}
	got := string(mt.gotBody)
	if !strings.Contains(got, contents) {
		t.Errorf("got body %q, which does not contain %q", got, contents)
	}
}

func TestEncryption(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	mt := &mockTransport{}
	mt.addResult(&http.Response{StatusCode: 200, Body: bodyReader("{}")}, nil)
	client := mockClient(t, mt)
	obj := client.Bucket("bucketname").Object("filename1")
	wc := obj.Key(testEncryptionKey).NewWriter(ctx)
	if _, err := wc.Write([]byte("hello world")); err != nil {
		t.Fatal(err)
	}
	if err := wc.Close(); err != nil {
		t.Fatal(err)
	}
	if got, want := mt.gotReq.Header.Get("x-goog-encryption-algorithm"), "AES256"; got != want {
		t.Errorf("algorithm: got %q, want %q", got, want)
	}
	gotKey, err := base64.StdEncoding.DecodeString(mt.gotReq.Header.Get("x-goog-encryption-key"))
	if err != nil {
		t.Fatalf("decoding key: %v", err)
	}
	if !testutil.Equal(gotKey, testEncryptionKey) {
		t.Errorf("key: got %v, want %v", gotKey, testEncryptionKey)
	}
	wantHash := sha256.Sum256(testEncryptionKey)
	gotHash, err := base64.StdEncoding.DecodeString(mt.gotReq.Header.Get("x-goog-encryption-key-sha256"))
	if err != nil {
		t.Fatalf("decoding hash: %v", err)
	}
	if !testutil.Equal(gotHash, wantHash[:]) { // wantHash is an array
		t.Errorf("hash: got\n%v, want\n%v", gotHash, wantHash)
	}

	// Using a customer-supplied encryption key and a KMS key together is an error.
	checkKMSError := func(msg string, err error) {
		if err == nil {
			t.Errorf("%s: got nil, want error", msg)
		} else if !strings.Contains(err.Error(), "KMS") {
			t.Errorf(`%s: got %q, want it to contain "KMS"`, msg, err)
		}
	}

	wc = obj.Key(testEncryptionKey).NewWriter(ctx)
	wc.KMSKeyName = "key"
	_, err = wc.Write([]byte{})
	checkKMSError("Write", err)
	checkKMSError("Close", wc.Close())
}

// This test demonstrates the data race on Writer.err that can happen when the
// Writer's context is cancelled. To see the race, comment out the w.mu.Lock/Unlock
// lines in writer.go and run this test with -race.
func TestRaceOnCancel(t *testing.T) {
	client := mockClient(t, &mockTransport{})
	cctx, cancel := context.WithCancel(context.Background())
	w := client.Bucket("b").Object("o").NewWriter(cctx)
	w.ChunkSize = googleapi.MinUploadChunkSize
	buf := make([]byte, w.ChunkSize)
	// This Write starts the goroutine in Writer.open. That reads the first chunk in its entirety
	// before sending the request (see google.golang.org/api/gensupport.PrepareUpload),
	// so to exhibit the race we must provide ChunkSize bytes.  The goroutine then makes the RPC (L137).
	w.Write(buf)
	// Canceling the context causes the call to return context.Canceled, which makes the open goroutine
	// write to w.err (L151).
	cancel()
	// This call to Write concurrently reads w.err (L169).
	w.Write([]byte(nil))
}

func TestCancelDoesNotLeak(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	const contents = "hello world"
	mt := mockTransport{}

	client, err := NewClient(ctx, option.WithHTTPClient(&http.Client{Transport: &mt}))
	if err != nil {
		t.Fatal(err)
	}

	wc := client.Bucket("bucketname").Object("filename1").NewWriter(ctx)
	wc.ContentType = "text/plain"

	// We can't check that the Write fails, since it depends on the write to the
	// underling mockTransport failing which is racy.
	wc.Write([]byte(contents))

	cancel()
}

func TestCloseDoesNotLeak(t *testing.T) {
	ctx := context.Background()
	const contents = "hello world"
	mt := mockTransport{}

	client, err := NewClient(ctx, option.WithHTTPClient(&http.Client{Transport: &mt}))
	if err != nil {
		t.Fatal(err)
	}

	wc := client.Bucket("bucketname").Object("filename1").NewWriter(ctx)
	wc.ContentType = "text/plain"

	// We can't check that the Write fails, since it depends on the write to the
	// underling mockTransport failing which is racy.
	wc.Write([]byte(contents))

	wc.Close()
}
<commit_msg>test(storage): re-enable skipped write retry test (#2676)<commit_after>// Copyright 2014 Google LLC
//
// 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 storage

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"net/http"
	"strings"
	"testing"

	"cloud.google.com/go/internal/testutil"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"
)

var testEncryptionKey = []byte("secret-key-that-is-32-bytes-long")

func TestErrorOnObjectsInsertCall(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	const contents = "hello world"

	doWrite := func(mt *mockTransport) *Writer {
		client := mockClient(t, mt)
		wc := client.Bucket("bucketname").Object("filename1").NewWriter(ctx)
		wc.ContentType = "text/plain"

		// We can't check that the Write fails, since it depends on the write to the
		// underling mockTransport failing which is racy.
		wc.Write([]byte(contents))
		return wc
	}

	wc := doWrite(&mockTransport{})
	// Close must always return an error though since it waits for the transport to
	// have closed.
	if err := wc.Close(); err == nil {
		t.Errorf("expected error on close, got nil")
	}

	// Retry on 5xx
	mt := &mockTransport{}
	mt.addResult(&http.Response{StatusCode: 503, Body: bodyReader("")}, nil)
	mt.addResult(&http.Response{StatusCode: 200, Body: bodyReader("{}")}, nil)

	wc = doWrite(mt)
	if err := wc.Close(); err != nil {
		t.Errorf("got %v, want nil", err)
	}
	got := string(mt.gotBody)
	if !strings.Contains(got, contents) {
		t.Errorf("got body %q, which does not contain %q", got, contents)
	}
}

func TestEncryption(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	mt := &mockTransport{}
	mt.addResult(&http.Response{StatusCode: 200, Body: bodyReader("{}")}, nil)
	client := mockClient(t, mt)
	obj := client.Bucket("bucketname").Object("filename1")
	wc := obj.Key(testEncryptionKey).NewWriter(ctx)
	if _, err := wc.Write([]byte("hello world")); err != nil {
		t.Fatal(err)
	}
	if err := wc.Close(); err != nil {
		t.Fatal(err)
	}
	if got, want := mt.gotReq.Header.Get("x-goog-encryption-algorithm"), "AES256"; got != want {
		t.Errorf("algorithm: got %q, want %q", got, want)
	}
	gotKey, err := base64.StdEncoding.DecodeString(mt.gotReq.Header.Get("x-goog-encryption-key"))
	if err != nil {
		t.Fatalf("decoding key: %v", err)
	}
	if !testutil.Equal(gotKey, testEncryptionKey) {
		t.Errorf("key: got %v, want %v", gotKey, testEncryptionKey)
	}
	wantHash := sha256.Sum256(testEncryptionKey)
	gotHash, err := base64.StdEncoding.DecodeString(mt.gotReq.Header.Get("x-goog-encryption-key-sha256"))
	if err != nil {
		t.Fatalf("decoding hash: %v", err)
	}
	if !testutil.Equal(gotHash, wantHash[:]) { // wantHash is an array
		t.Errorf("hash: got\n%v, want\n%v", gotHash, wantHash)
	}

	// Using a customer-supplied encryption key and a KMS key together is an error.
	checkKMSError := func(msg string, err error) {
		if err == nil {
			t.Errorf("%s: got nil, want error", msg)
		} else if !strings.Contains(err.Error(), "KMS") {
			t.Errorf(`%s: got %q, want it to contain "KMS"`, msg, err)
		}
	}

	wc = obj.Key(testEncryptionKey).NewWriter(ctx)
	wc.KMSKeyName = "key"
	_, err = wc.Write([]byte{})
	checkKMSError("Write", err)
	checkKMSError("Close", wc.Close())
}

// This test demonstrates the data race on Writer.err that can happen when the
// Writer's context is cancelled. To see the race, comment out the w.mu.Lock/Unlock
// lines in writer.go and run this test with -race.
func TestRaceOnCancel(t *testing.T) {
	client := mockClient(t, &mockTransport{})
	cctx, cancel := context.WithCancel(context.Background())
	w := client.Bucket("b").Object("o").NewWriter(cctx)
	w.ChunkSize = googleapi.MinUploadChunkSize
	buf := make([]byte, w.ChunkSize)
	// This Write starts the goroutine in Writer.open. That reads the first chunk in its entirety
	// before sending the request (see google.golang.org/api/gensupport.PrepareUpload),
	// so to exhibit the race we must provide ChunkSize bytes.  The goroutine then makes the RPC (L137).
	w.Write(buf)
	// Canceling the context causes the call to return context.Canceled, which makes the open goroutine
	// write to w.err (L151).
	cancel()
	// This call to Write concurrently reads w.err (L169).
	w.Write([]byte(nil))
}

func TestCancelDoesNotLeak(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	const contents = "hello world"
	mt := mockTransport{}

	client, err := NewClient(ctx, option.WithHTTPClient(&http.Client{Transport: &mt}))
	if err != nil {
		t.Fatal(err)
	}

	wc := client.Bucket("bucketname").Object("filename1").NewWriter(ctx)
	wc.ContentType = "text/plain"

	// We can't check that the Write fails, since it depends on the write to the
	// underling mockTransport failing which is racy.
	wc.Write([]byte(contents))

	cancel()
}

func TestCloseDoesNotLeak(t *testing.T) {
	ctx := context.Background()
	const contents = "hello world"
	mt := mockTransport{}

	client, err := NewClient(ctx, option.WithHTTPClient(&http.Client{Transport: &mt}))
	if err != nil {
		t.Fatal(err)
	}

	wc := client.Bucket("bucketname").Object("filename1").NewWriter(ctx)
	wc.ContentType = "text/plain"

	// We can't check that the Write fails, since it depends on the write to the
	// underling mockTransport failing which is racy.
	wc.Write([]byte(contents))

	wc.Close()
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ecr"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSEcrLifecyclePolicy_basic(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSEcrLifecyclePolicyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccEcrLifecyclePolicyConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSEcrLifecyclePolicyExists("aws_ecr_lifecycle_policy.foo"),
				),
			},
		},
	})
}

func testAccCheckAWSEcrLifecyclePolicyDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).ecrconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_ecr_lifecycle_policy" {
			continue
		}

		input := &ecr.GetLifecyclePolicyInput{
			RegistryId:     aws.String(rs.Primary.Attributes["registry_id"]),
			RepositoryName: aws.String(rs.Primary.Attributes["repository"]),
		}

		_, err := conn.GetLifecyclePolicy(input)
		if err != nil {
			if aerr, ok := err.(awserr.Error); ok {
				switch aerr.Code() {
				case ecr.ErrCodeRepositoryNotFoundException:
					return nil
				default:
					return err
				}
			}
			return err
		}
	}

	return nil
}

func testAccCheckAWSEcrLifecyclePolicyExists(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		_, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		return nil
	}
}

const testAccEcrLifecyclePolicyConfig = `
resource "aws_ecr_repository" "foo" {
  name = "bar"
}
resource "aws_ecr_lifecycle_policy" "foo" {
	repository = "${aws_ecr_repository.foo.name}"
  policy = <<EOF
{
  "rules": [
    {
      "rulePriority": 1,
      "description": "Expire images older than 14 days",
      "selection": {
        "tagStatus": "untagged",
        "countType": "sinceImagePushed",
        "countUnit": "days",
        "countNumber": 14
      },
      "action": {
        "type": "expire"
      }
    }
  ]
}
EOF
}
`
<commit_msg>t/ecr_lifecycle_policy: Randomize repo name<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ecr"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSEcrLifecyclePolicy_basic(t *testing.T) {
	randString := acctest.RandString(10)
	rName := fmt.Sprintf("tf-acc-test-lifecycle-%s", randString)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSEcrLifecyclePolicyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccEcrLifecyclePolicyConfig(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSEcrLifecyclePolicyExists("aws_ecr_lifecycle_policy.foo"),
				),
			},
		},
	})
}

func testAccCheckAWSEcrLifecyclePolicyDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).ecrconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_ecr_lifecycle_policy" {
			continue
		}

		input := &ecr.GetLifecyclePolicyInput{
			RegistryId:     aws.String(rs.Primary.Attributes["registry_id"]),
			RepositoryName: aws.String(rs.Primary.Attributes["repository"]),
		}

		_, err := conn.GetLifecyclePolicy(input)
		if err != nil {
			if aerr, ok := err.(awserr.Error); ok {
				switch aerr.Code() {
				case ecr.ErrCodeRepositoryNotFoundException:
					return nil
				default:
					return err
				}
			}
			return err
		}
	}

	return nil
}

func testAccCheckAWSEcrLifecyclePolicyExists(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		_, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		return nil
	}
}

func testAccEcrLifecyclePolicyConfig(rName string) string {
	return fmt.Sprintf(`
resource "aws_ecr_repository" "foo" {
  name = "%s"
}
resource "aws_ecr_lifecycle_policy" "foo" {
	repository = "${aws_ecr_repository.foo.name}"
  policy = <<EOF
{
  "rules": [
    {
      "rulePriority": 1,
      "description": "Expire images older than 14 days",
      "selection": {
        "tagStatus": "untagged",
        "countType": "sinceImagePushed",
        "countUnit": "days",
        "countNumber": 14
      },
      "action": {
        "type": "expire"
      }
    }
  ]
}
EOF
}
`, rName)
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

package revision

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"testing"
	"time"

	"go.uber.org/atomic"
	logtesting "knative.dev/pkg/logging/testing"
	"knative.dev/pkg/ptr"

	"github.com/google/go-containerregistry/pkg/authn/k8schain"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/util/workqueue"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
)

var (
	errDigest    = errors.New("digest error")
	fakeRevision = rev("rev", "first-image", "second-image")
)

func TestResolveInBackground(t *testing.T) {
	tests := []struct {
		name         string
		resolver     resolveFunc
		timeout      *time.Duration
		wantStatuses []v1.ContainerStatus
		wantError    error
	}{{
		name: "success",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			return img + "-digest", nil
		},
		wantStatuses: []v1.ContainerStatus{{
			Name:        "first",
			ImageDigest: "first-image-digest",
		}, {
			Name:        "second",
			ImageDigest: "second-image-digest",
		}},
	}, {
		name: "passing params",
		resolver: func(_ context.Context, img string, opt k8schain.Options, skip sets.String) (string, error) {
			return fmt.Sprintf("%s-%s-%s", img, opt.ServiceAccountName, skip.List()[0]), nil
		},
		wantStatuses: []v1.ContainerStatus{{
			Name:        "first",
			ImageDigest: "first-image-san-skip",
		}, {
			Name:        "second",
			ImageDigest: "second-image-san-skip",
		}},
	}, {
		name: "one slow resolve",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			if img == "first-image" {
				// make the first resolve arrive after the second.
				time.Sleep(50 * time.Millisecond)
			}
			return img + "-digest", nil
		},
		wantStatuses: []v1.ContainerStatus{{
			Name:        "first",
			ImageDigest: "first-image-digest",
		}, {
			Name:        "second",
			ImageDigest: "second-image-digest",
		}},
	}, {
		name: "resolver entirely fails",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			return img + "-digest", errDigest
		},
		wantError: errDigest,
	}, {
		name: "resolver fails one image",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			if img == "second-image" {
				return "", errDigest
			}

			return img + "-digest", nil
		},
		wantError: errDigest,
	}, {
		name:    "timeout",
		timeout: ptr.Duration(10 * time.Millisecond),
		resolver: func(ctx context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			if img == "second-image" {
				select {
				case <-time.After(10 * time.Second):
				case <-ctx.Done():
				}
			}

			return img + "-digest", ctx.Err()
		},
		wantError: context.DeadlineExceeded,
	}}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			timeout := 5 * time.Second
			if tt.timeout != nil {
				timeout = *tt.timeout
			}

			ready := make(chan types.NamespacedName)
			cb := func(rev types.NamespacedName) {
				ready <- rev
				close(ready)
			}

			logger := logtesting.TestLogger(t)
			subject := newBackgroundResolver(logger, tt.resolver, workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()), cb)

			stop := make(chan struct{})
			done := subject.Start(stop, 10)

			defer func() {
				close(stop)
				<-done
			}()

			for i := 0; i < 2; i++ {
				t.Run(fmt.Sprint("iteration", i), func(t *testing.T) {
					statuses, err := subject.Resolve(fakeRevision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), timeout)
					if err != nil || statuses != nil {
						// Initial result should be nil, nil since we have nothing in cache.
						t.Errorf("Resolve() = %v, %v, wanted nil, nil", statuses, err)
					}

					select {
					case <-ready:
						// got the result.
					case <-time.After(2 * time.Second):
						// This shouldn't happen in these tests. The callback is always
						// eventually called.
						t.Fatalf("Resolver did not report ready")
					}

					statuses, err = subject.Resolve(fakeRevision, k8schain.Options{}, nil, timeout)
					if got, want := err, tt.wantError; !errors.Is(got, want) {
						t.Errorf("Resolve() = _, %q, wanted %q", got, want)
					}
					if got, want := statuses, tt.wantStatuses; !reflect.DeepEqual(got, want) {
						t.Errorf("Resolve() = %v, wanted %v", got, want)
					}

					// Clear, then we'll loop and make sure that we look everything up from scratch.
					subject.Clear(types.NamespacedName{Namespace: fakeRevision.Namespace, Name: fakeRevision.Name})
					ready = make(chan types.NamespacedName)
				})
			}
		})
	}
}

func TestRateLimitGlobal(t *testing.T) {
	logger := logtesting.TestLogger(t)

	var resolves atomic.Int32
	var resolver resolveFunc = func(_ context.Context, _ string, _ k8schain.Options, _ sets.String) (string, error) {
		resolves.Inc()
		return "", errors.New("failed")
	}

	queue := workqueue.NewRateLimitingQueue(workqueue.NewItemExponentialFailureRateLimiter(1*time.Second, 5*time.Second))
	subject := newBackgroundResolver(logger, resolver, queue, func(types.NamespacedName) {})

	stop := make(chan struct{})
	done := subject.Start(stop, 1)

	defer func() {
		close(stop)
		<-done
	}()

	for i := 0; i < 5; i++ {
		name := fmt.Sprint("rev", i)
		subject.Resolve(rev(name, name+"img1", name+"img2"), k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
	}

	// Rate limit base rate in this test is 1 second, so this should give time for at most 1 resolve.
	time.Sleep(500 * time.Millisecond)

	if r := resolves.Load(); r > 1 {
		t.Fatalf("Expected resolves to be rate limited, but was called %d times", r)
	}
}

func TestRateLimitPerItem(t *testing.T) {
	logger := logtesting.TestLogger(t)

	var resolver resolveFunc = func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
		if img == "img1" {
			return "", nil
		}

		return "", errors.New("failed")
	}

	queue := workqueue.NewRateLimitingQueue(workqueue.NewItemExponentialFailureRateLimiter(50*time.Millisecond, 5*time.Second))

	enqueue := make(chan struct{})
	subject := newBackgroundResolver(logger, resolver, queue, func(types.NamespacedName) {
		enqueue <- struct{}{}
	})

	stop := make(chan struct{})
	done := subject.Start(stop, 1)

	defer func() {
		close(stop)
		<-done
	}()

	start := time.Now()
	revision := rev("rev", "img1", "img2")
	for i := 0; i < 5; i++ {
		subject.Clear(types.NamespacedName{Name: revision.Name, Namespace: revision.Namespace})
		resolution, err := subject.Resolve(revision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err != nil || resolution != nil {
			t.Fatalf("Expected Resolve to be nil, nil but got %v, %v", resolution, err)
		}

		<-enqueue
		_, err = subject.Resolve(revision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err == nil {
			t.Fatalf("Expected Resolve to fail")
		}
	}

	if took := time.Since(start); took < 600*time.Millisecond {
		// Per-item time is 50ms, so after 5 cycles of back-off should take at least 800ms.
		// (Otherwise will take only ~250ms)
		t.Fatal("Expected resolves to take longer than 600ms, but took", took)
	}

	t.Run("Does not affect other revisions", func(t *testing.T) {
		start := time.Now()

		resolution, err := subject.Resolve(rev("another-revision", "img1", "img2"), k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err != nil || resolution != nil {
			t.Fatalf("Expected Resolve to be nil, nil but got %v, %v", resolution, err)
		}

		<-enqueue
		if took := time.Since(start); took > 500*time.Millisecond {
			t.Fatal("Expected per-item limit not to affect other revisions, but took", took)
		}
	})

	t.Run("Forget clears per-item rate limit", func(t *testing.T) {
		subject.Forget(types.NamespacedName{Name: revision.Name, Namespace: revision.Namespace})

		start = time.Now()
		resolution, err := subject.Resolve(revision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err != nil || resolution != nil {
			t.Fatalf("Expected Resolve to be nil, nil but got %v, %v", resolution, err)
		}

		<-enqueue
		if took := time.Since(start); took > 500*time.Millisecond {
			t.Fatal("Expected Forget to remove revision from rate limiter, but took", took)
		}
	})
}

type resolveFunc func(context.Context, string, k8schain.Options, sets.String) (string, error)

func (r resolveFunc) Resolve(c context.Context, s string, o k8schain.Options, t sets.String) (string, error) {
	return r(c, s, o, t)
}

func rev(name, firstImage, secondImage string) *v1.Revision {
	return &v1.Revision{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: "ns",
		},
		Spec: v1.RevisionSpec{
			PodSpec: corev1.PodSpec{
				Containers: []corev1.Container{{
					Name:  "first",
					Image: firstImage,
				}, {
					Name:  "second",
					Image: secondImage,
				}},
			},
		},
	}
}
<commit_msg>Speed up TestRateLimitPerItem (#11908)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

package revision

import (
	"context"
	"errors"
	"fmt"
	"math"
	"reflect"
	"testing"
	"time"

	"go.uber.org/atomic"
	logtesting "knative.dev/pkg/logging/testing"
	"knative.dev/pkg/ptr"

	"github.com/google/go-containerregistry/pkg/authn/k8schain"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/util/workqueue"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
)

var (
	errDigest    = errors.New("digest error")
	fakeRevision = rev("rev", "first-image", "second-image")
)

func TestResolveInBackground(t *testing.T) {
	tests := []struct {
		name         string
		resolver     resolveFunc
		timeout      *time.Duration
		wantStatuses []v1.ContainerStatus
		wantError    error
	}{{
		name: "success",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			return img + "-digest", nil
		},
		wantStatuses: []v1.ContainerStatus{{
			Name:        "first",
			ImageDigest: "first-image-digest",
		}, {
			Name:        "second",
			ImageDigest: "second-image-digest",
		}},
	}, {
		name: "passing params",
		resolver: func(_ context.Context, img string, opt k8schain.Options, skip sets.String) (string, error) {
			return fmt.Sprintf("%s-%s-%s", img, opt.ServiceAccountName, skip.List()[0]), nil
		},
		wantStatuses: []v1.ContainerStatus{{
			Name:        "first",
			ImageDigest: "first-image-san-skip",
		}, {
			Name:        "second",
			ImageDigest: "second-image-san-skip",
		}},
	}, {
		name: "one slow resolve",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			if img == "first-image" {
				// make the first resolve arrive after the second.
				time.Sleep(50 * time.Millisecond)
			}
			return img + "-digest", nil
		},
		wantStatuses: []v1.ContainerStatus{{
			Name:        "first",
			ImageDigest: "first-image-digest",
		}, {
			Name:        "second",
			ImageDigest: "second-image-digest",
		}},
	}, {
		name: "resolver entirely fails",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			return img + "-digest", errDigest
		},
		wantError: errDigest,
	}, {
		name: "resolver fails one image",
		resolver: func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			if img == "second-image" {
				return "", errDigest
			}

			return img + "-digest", nil
		},
		wantError: errDigest,
	}, {
		name:    "timeout",
		timeout: ptr.Duration(10 * time.Millisecond),
		resolver: func(ctx context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
			if img == "second-image" {
				select {
				case <-time.After(10 * time.Second):
				case <-ctx.Done():
				}
			}

			return img + "-digest", ctx.Err()
		},
		wantError: context.DeadlineExceeded,
	}}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			timeout := 5 * time.Second
			if tt.timeout != nil {
				timeout = *tt.timeout
			}

			ready := make(chan types.NamespacedName)
			cb := func(rev types.NamespacedName) {
				ready <- rev
				close(ready)
			}

			logger := logtesting.TestLogger(t)
			subject := newBackgroundResolver(logger, tt.resolver, workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()), cb)

			stop := make(chan struct{})
			done := subject.Start(stop, 10)

			defer func() {
				close(stop)
				<-done
			}()

			for i := 0; i < 2; i++ {
				t.Run(fmt.Sprint("iteration", i), func(t *testing.T) {
					statuses, err := subject.Resolve(fakeRevision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), timeout)
					if err != nil || statuses != nil {
						// Initial result should be nil, nil since we have nothing in cache.
						t.Errorf("Resolve() = %v, %v, wanted nil, nil", statuses, err)
					}

					select {
					case <-ready:
						// got the result.
					case <-time.After(2 * time.Second):
						// This shouldn't happen in these tests. The callback is always
						// eventually called.
						t.Fatalf("Resolver did not report ready")
					}

					statuses, err = subject.Resolve(fakeRevision, k8schain.Options{}, nil, timeout)
					if got, want := err, tt.wantError; !errors.Is(got, want) {
						t.Errorf("Resolve() = _, %q, wanted %q", got, want)
					}
					if got, want := statuses, tt.wantStatuses; !reflect.DeepEqual(got, want) {
						t.Errorf("Resolve() = %v, wanted %v", got, want)
					}

					// Clear, then we'll loop and make sure that we look everything up from scratch.
					subject.Clear(types.NamespacedName{Namespace: fakeRevision.Namespace, Name: fakeRevision.Name})
					ready = make(chan types.NamespacedName)
				})
			}
		})
	}
}

func TestRateLimitGlobal(t *testing.T) {
	logger := logtesting.TestLogger(t)

	var resolves atomic.Int32
	var resolver resolveFunc = func(_ context.Context, _ string, _ k8schain.Options, _ sets.String) (string, error) {
		resolves.Inc()
		return "", errors.New("failed")
	}

	queue := workqueue.NewRateLimitingQueue(workqueue.NewItemExponentialFailureRateLimiter(1*time.Second, 5*time.Second))
	subject := newBackgroundResolver(logger, resolver, queue, func(types.NamespacedName) {})

	stop := make(chan struct{})
	done := subject.Start(stop, 1)

	defer func() {
		close(stop)
		<-done
	}()

	for i := 0; i < 5; i++ {
		name := fmt.Sprint("rev", i)
		subject.Resolve(rev(name, name+"img1", name+"img2"), k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
	}

	// Rate limit base rate in this test is 1 second, so this should give time for at most 1 resolve.
	time.Sleep(500 * time.Millisecond)

	if r := resolves.Load(); r > 1 {
		t.Fatalf("Expected resolves to be rate limited, but was called %d times", r)
	}
}

func TestRateLimitPerItem(t *testing.T) {
	logger := logtesting.TestLogger(t)

	var resolver resolveFunc = func(_ context.Context, img string, _ k8schain.Options, _ sets.String) (string, error) {
		if img == "img1" {
			return "", nil
		}

		return "", errors.New("failed")
	}

	baseDelay := 10 * time.Millisecond
	queue := workqueue.NewRateLimitingQueue(workqueue.NewItemExponentialFailureRateLimiter(baseDelay, 5*time.Second))

	enqueue := make(chan struct{})
	subject := newBackgroundResolver(logger, resolver, queue, func(types.NamespacedName) {
		enqueue <- struct{}{}
	})

	stop := make(chan struct{})
	done := subject.Start(stop, 1)

	defer func() {
		close(stop)
		<-done
	}()

	revision := rev("rev", "img1", "img2")
	for i := 0; i < 4; i++ {
		subject.Clear(types.NamespacedName{Name: revision.Name, Namespace: revision.Namespace})
		start := time.Now()
		resolution, err := subject.Resolve(revision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err != nil || resolution != nil {
			t.Fatalf("Expected Resolve to be nil, nil but got %v, %v", resolution, err)
		}

		<-enqueue

		_, err = subject.Resolve(revision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err == nil {
			t.Fatalf("Expected Resolve to fail")
		}

		latency := time.Since(start)
		expected := time.Duration(math.Pow(2, float64(i))) * baseDelay
		if latency < expected {
			t.Fatalf("latency = %s, want at least %s", latency, expected)
		}
	}

	t.Run("Does not affect other revisions", func(t *testing.T) {
		start := time.Now()

		resolution, err := subject.Resolve(rev("another-revision", "img1", "img2"), k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err != nil || resolution != nil {
			t.Fatalf("Expected Resolve to be nil, nil but got %v, %v", resolution, err)
		}

		<-enqueue
		if took := time.Since(start); took > baseDelay*2 {
			t.Fatal("Expected per-item limit not to affect other revisions, but took", took)
		}
	})

	t.Run("Forget clears per-item rate limit", func(t *testing.T) {
		subject.Forget(types.NamespacedName{Name: revision.Name, Namespace: revision.Namespace})

		start := time.Now()
		resolution, err := subject.Resolve(revision, k8schain.Options{ServiceAccountName: "san"}, sets.NewString("skip"), 0)
		if err != nil || resolution != nil {
			t.Fatalf("Expected Resolve to be nil, nil but got %v, %v", resolution, err)
		}

		<-enqueue
		if took := time.Since(start); took > baseDelay*2 {
			t.Fatal("Expected Forget to remove revision from rate limiter, but took", took)
		}
	})
}

type resolveFunc func(context.Context, string, k8schain.Options, sets.String) (string, error)

func (r resolveFunc) Resolve(c context.Context, s string, o k8schain.Options, t sets.String) (string, error) {
	return r(c, s, o, t)
}

func rev(name, firstImage, secondImage string) *v1.Revision {
	return &v1.Revision{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: "ns",
		},
		Spec: v1.RevisionSpec{
			PodSpec: corev1.PodSpec{
				Containers: []corev1.Container{{
					Name:  "first",
					Image: firstImage,
				}, {
					Name:  "second",
					Image: secondImage,
				}},
			},
		},
	}
}
<|endoftext|>


<commit_before>// Package recovery provides tooling to help with control plane disaster recovery. Recover() uses a
// Backend to extract the control plane from a store, such as etcd, and use those to write assets
// that can be used by `bootkube start` to reboot the control plane.
//
// The recovery tool assumes that the component names for the control plane elements are the same as
// what is output by `bootkube render`. The `bootkube start` command also makes this assumption.
// It also assumes that kubeconfig on the kubelet is located at /etc/kubernetes/kubeconfig, though
// that can be changed in the bootstrap manifests that are rendered.
package recovery

import (
	"context"
	"fmt"
	"io/ioutil"
	"path"
	"path/filepath"
	"reflect"

	"github.com/ghodss/yaml"
	"k8s.io/api/core/v1"
	"k8s.io/api/extensions/v1beta1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"

	"github.com/kubernetes-incubator/bootkube/pkg/asset"
)

const (
	k8sAppLabel            = "k8s-app"   // The label used in versions > v0.4.2
	componentAppLabel      = "component" // The label used in versions <= v0.4.2
	kubeletKubeConfigPath  = "/etc/kubernetes"
	apiServerContainerName = "kube-apiserver"
)

var (
	// bootstrapK8sApps contains the components (as identified by the the label "k8s-app") that we
	// will extract to construct the temporary bootstrap control plane.
	bootstrapK8sApps = map[string]struct{}{
		apiServerContainerName:    {},
		"kube-controller-manager": {},
		"kube-scheduler":          {},
	}
	// kubeConfigK8sContainers contains the names of the bootstrap container specs that need to add a
	// --kubeconfig flag to run in non-self-hosted mode.
	kubeConfigK8sContainers = map[string]struct{}{
		"kube-controller-manager": {},
		"kube-scheduler":          {},
	}
	// typeMetas contains a mapping from API object types to the TypeMeta struct that should be
	// populated for them when they are serialized.
	typeMetas    = make(map[reflect.Type]metav1.TypeMeta)
	metaAccessor = meta.NewAccessor()
)

func init() {
	addTypeMeta := func(obj runtime.Object, gv schema.GroupVersion) {
		t := reflect.TypeOf(obj)
		typeMetas[t] = metav1.TypeMeta{
			APIVersion: gv.String(),
			Kind:       t.Elem().Name(),
		}
	}
	addTypeMeta(&v1.ConfigMap{}, v1.SchemeGroupVersion)
	addTypeMeta(&v1beta1.DaemonSet{}, v1beta1.SchemeGroupVersion)
	addTypeMeta(&v1beta1.Deployment{}, v1beta1.SchemeGroupVersion)
	addTypeMeta(&v1.Pod{}, v1.SchemeGroupVersion)
	addTypeMeta(&v1.Secret{}, v1.SchemeGroupVersion)
}

// Backend defines an interface for any backend that can populate a controlPlane struct.
type Backend interface {
	read(context.Context) (*controlPlane, error)
}

// controlPlane holds the control plane objects that are recovered from a backend.
type controlPlane struct {
	configMaps  v1.ConfigMapList
	daemonSets  v1beta1.DaemonSetList
	deployments v1beta1.DeploymentList
	secrets     v1.SecretList
}

// Recover recovers a control plane using the provided backend and kubeConfigPath, returning assets
// for the existing control plane and a bootstrap control plane that can be used with `bootkube
// start` to re-bootstrap the control plane.
func Recover(ctx context.Context, backend Backend, kubeConfigPath string) (asset.Assets, error) {
	cp, err := backend.read(ctx)
	if err != nil {
		return nil, err
	}

	as, err := cp.renderBootstrap()
	if err != nil {
		return nil, err
	}

	kc, err := renderKubeConfig(kubeConfigPath)
	if err != nil {
		return nil, err
	}
	as = append(as, kc)

	return as, nil
}

// renderBootstrap returns assets for a bootstrap control plane that can be used with `bootkube
// start` to re-bootstrap a control plane. These assets are derived from the self-hosted control
// plane that was recovered by the backend, but modified for direct injection into a kubelet.
func (cp *controlPlane) renderBootstrap() (asset.Assets, error) {
	pods, err := extractBootstrapPods(cp.daemonSets.Items, cp.deployments.Items)
	if err != nil {
		return nil, err
	}
	requiredConfigMaps, requiredSecrets := fixUpBootstrapPods(pods)
	as, err := outputBootstrapPods(pods)
	if err != nil {
		return nil, err
	}
	configMaps, err := outputBootstrapConfigMaps(cp.configMaps, requiredConfigMaps)
	if err != nil {
		return nil, err
	}
	as = append(as, configMaps...)

	secrets, err := outputBootstrapSecrets(cp.secrets, requiredSecrets)
	if err != nil {
		return nil, err
	}
	as = append(as, secrets...)
	return as, nil
}

// extractBootstrapPods extracts bootstrap pod specs from daemonsets and deployments.
func extractBootstrapPods(daemonSets []v1beta1.DaemonSet, deployments []v1beta1.Deployment) ([]v1.Pod, error) {
	var pods []v1.Pod
	for _, ds := range daemonSets {
		if isBootstrapApp(ds.Labels) {
			pod := v1.Pod{Spec: ds.Spec.Template.Spec}
			if err := setBootstrapPodMetadata(&pod, ds.ObjectMeta); err != nil {
				return nil, err
			}
			pods = append(pods, pod)
		}
	}
	for _, ds := range deployments {
		if isBootstrapApp(ds.Labels) {
			pod := v1.Pod{Spec: ds.Spec.Template.Spec}
			if err := setBootstrapPodMetadata(&pod, ds.ObjectMeta); err != nil {
				return nil, err
			}
			pods = append(pods, pod)
		}
	}
	return pods, nil
}

// isBootstrapApp returns true if this app belongs to the bootstrap control plane, based on its
// labels.
func isBootstrapApp(labels map[string]string) bool {
	k8sApp := labels[k8sAppLabel]
	if k8sApp == "" {
		k8sApp = labels[componentAppLabel]
	}
	_, ok := bootstrapK8sApps[k8sApp]
	return ok
}

// setBootstrapPodMetadata creates valid metadata for a bootstrap pod. Currently it sets the
// TypeMeta and Name, Namespace, and Annotations on the ObjectMeta.
func setBootstrapPodMetadata(pod *v1.Pod, parent metav1.ObjectMeta) error {
	if err := setTypeMeta(pod); err != nil {
		return err
	}
	pod.ObjectMeta = metav1.ObjectMeta{
		Annotations: parent.Annotations,
		Name:        "bootstrap-" + parent.Name,
		Namespace:   parent.Namespace,
	}
	return nil
}

// fixUpBootstrapPods modifies extracted bootstrap pod specs to have correct metadata and point to
// filesystem-mount-based secrets, and removes any security contexts that might prevent the pods
// from accessing those secrets. It returns mappings from configMap and secret names to output
// paths that must also be rendered in order for the bootstrap pods to be functional.
func fixUpBootstrapPods(pods []v1.Pod) (requiredConfigMaps, requiredSecrets map[string]string) {
	requiredConfigMaps, requiredSecrets = make(map[string]string), make(map[string]string)
	for i := range pods {
		pod := &pods[i]

		// Fix SecurityContext to ensure the pod runs as root.
		if pod.Spec.SecurityContext != nil {
			pod.Spec.SecurityContext.RunAsNonRoot = nil
			pod.Spec.SecurityContext.RunAsUser = nil
		}
		// Fix hostNetwork: true because bootstrap assets can not rely on overlay networks.
		if pod.Spec.HostNetwork == false {
			pod.Spec.HostNetwork = true
		}
		// Change secret volumes to point to file mounts.
		for i := range pod.Spec.Volumes {
			vol := &pod.Spec.Volumes[i]
			if vol.Secret != nil {
				pathSuffix := filepath.Join("secrets", vol.Secret.SecretName)
				requiredSecrets[vol.Secret.SecretName] = filepath.Join(asset.AssetPathSecrets, pathSuffix)
				vol.HostPath = &v1.HostPathVolumeSource{Path: filepath.Join(asset.BootstrapSecretsDir, pathSuffix)}
				vol.Secret = nil
			} else if vol.ConfigMap != nil {
				pathSuffix := filepath.Join("config-maps", vol.ConfigMap.Name)
				requiredConfigMaps[vol.ConfigMap.Name] = filepath.Join(asset.AssetPathSecrets, pathSuffix)
				vol.HostPath = &v1.HostPathVolumeSource{Path: path.Join(asset.BootstrapSecretsDir, pathSuffix)}
				vol.ConfigMap = nil
			}
		}

		// Make sure the kubeconfig is in the commandline.
		for i := range pod.Spec.Containers {
			cn := &pod.Spec.Containers[i]

			// Fix SecurityContext to ensure the container runs as root.
			if cn.SecurityContext != nil {
				cn.SecurityContext.RunAsNonRoot = nil
				cn.SecurityContext.RunAsUser = nil
			}

			// Assumes the bootkube naming convention is used. Could also just make sure the image uses hyperkube.
			if _, ok := kubeConfigK8sContainers[cn.Name]; ok {
				cn.Command = append(cn.Command, "--kubeconfig=/kubeconfig/kubeconfig")
				cn.VolumeMounts = append(cn.VolumeMounts, v1.VolumeMount{
					MountPath: "/kubeconfig",
					Name:      "kubeconfig",
					ReadOnly:  true,
				})
			}
		}

		// Add a mount for the kubeconfig.
		pod.Spec.Volumes = append(pod.Spec.Volumes, v1.Volume{
			VolumeSource: v1.VolumeSource{HostPath: &v1.HostPathVolumeSource{Path: kubeletKubeConfigPath}},
			Name:         "kubeconfig",
		})
	}
	return
}

// outputBootstrapPods outputs the bootstrap pod definitions.
func outputBootstrapPods(pods []v1.Pod) (asset.Assets, error) {
	var as asset.Assets
	for _, pod := range pods {
		a, err := serializeObjToYAML(path.Join(asset.AssetPathBootstrapManifests, pod.Name+".yaml"), &pod)
		if err != nil {
			return nil, err
		}
		as = append(as, a)
	}
	return as, nil
}

// outputBootstrapConfigMaps creates assets for all the configMap names in the requiredConfigMaps
// set. It returns an error if any configMap cannot be found in the provided configMaps list.
func outputBootstrapConfigMaps(configMaps v1.ConfigMapList, requiredConfigMaps map[string]string) (asset.Assets, error) {
	return outputKeyValueData(&configMaps, requiredConfigMaps, func(obj runtime.Object) map[string][]byte {
		configMap, ok := obj.(*v1.ConfigMap)
		if !ok || configMap == nil {
			return nil
		}
		output := make(map[string][]byte)
		for k, v := range configMap.Data {
			output[k] = []byte(v)
		}
		return output
	})
}

// outputBootstrapSecrets creates assets for all the secret names in the requiredSecrets set. It
// returns an error if any secret cannot be found in the provided secrets list.
func outputBootstrapSecrets(secrets v1.SecretList, requiredSecrets map[string]string) (asset.Assets, error) {
	return outputKeyValueData(&secrets, requiredSecrets, func(obj runtime.Object) map[string][]byte {
		if secret, ok := obj.(*v1.Secret); ok && secret != nil {
			return secret.Data
		}
		return nil
	})
}

// outputKeyValueData takes a key-value object (such as a Secret or ConfigMap) and outputs assets
// for each key-value pair. See outputBootstrapConfigMaps or outputBootstrapSecrets for usage.
func outputKeyValueData(objList runtime.Object, requiredObjs map[string]string, extractData func(runtime.Object) map[string][]byte) (asset.Assets, error) {
	var as asset.Assets
	objs, err := meta.ExtractList(objList)
	if err != nil {
		return nil, err
	}
	for _, obj := range objs {
		name, err := metaAccessor.Name(obj)
		if err != nil {
			return nil, err
		}
		if namePrefix, ok := requiredObjs[name]; ok {
			for key, data := range extractData(obj) {
				as = append(as, asset.Asset{
					Name: path.Join(namePrefix, key),
					Data: data,
				})
			}
			delete(requiredObjs, name)
		}
	}
	if len(requiredObjs) > 0 {
		var missingObjs []string
		for obj := range requiredObjs {
			missingObjs = append(missingObjs, obj)
		}
		return nil, fmt.Errorf("failed to extract some required objects: %v", missingObjs)
	}
	return as, nil
}

// renderKubeConfig outputs kubeconfig assets to ensure that the kubeconfig will be rendered to the
// assetDir for use by `bootkube start`.
func renderKubeConfig(kubeConfigPath string) (asset.Asset, error) {
	kubeConfig, err := ioutil.ReadFile(kubeConfigPath)
	if err != nil {
		return asset.Asset{}, err
	}
	return asset.Asset{
		Name: asset.AssetPathAdminKubeConfig, // used by `bootkube start`.
		Data: kubeConfig,
	}, nil
}

// setTypeMeta sets the TypeMeta for a runtime.Object.
// TODO(diegs): find the apimachinery code that does this, and use that instead.
func setTypeMeta(obj runtime.Object) error {
	typeMeta, ok := typeMetas[reflect.TypeOf(obj)]
	if !ok {
		return fmt.Errorf("don't know about type: %T", obj)
	}
	metaAccessor.SetAPIVersion(obj, typeMeta.APIVersion)
	metaAccessor.SetKind(obj, typeMeta.Kind)
	return nil
}

// serializeObjToYAML serializes a runtime.Object into a YAML asset.
func serializeObjToYAML(assetName string, obj runtime.Object) (asset.Asset, error) {
	data, err := yaml.Marshal(obj)
	if err != nil {
		return asset.Asset{}, err
	}
	return asset.Asset{
		Name: assetName,
		Data: data,
	}, nil
}
<commit_msg>Delete repeated 'the'<commit_after>// Package recovery provides tooling to help with control plane disaster recovery. Recover() uses a
// Backend to extract the control plane from a store, such as etcd, and use those to write assets
// that can be used by `bootkube start` to reboot the control plane.
//
// The recovery tool assumes that the component names for the control plane elements are the same as
// what is output by `bootkube render`. The `bootkube start` command also makes this assumption.
// It also assumes that kubeconfig on the kubelet is located at /etc/kubernetes/kubeconfig, though
// that can be changed in the bootstrap manifests that are rendered.
package recovery

import (
	"context"
	"fmt"
	"io/ioutil"
	"path"
	"path/filepath"
	"reflect"

	"github.com/ghodss/yaml"
	"k8s.io/api/core/v1"
	"k8s.io/api/extensions/v1beta1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"

	"github.com/kubernetes-incubator/bootkube/pkg/asset"
)

const (
	k8sAppLabel            = "k8s-app"   // The label used in versions > v0.4.2
	componentAppLabel      = "component" // The label used in versions <= v0.4.2
	kubeletKubeConfigPath  = "/etc/kubernetes"
	apiServerContainerName = "kube-apiserver"
)

var (
	// bootstrapK8sApps contains the components (as identified by the label "k8s-app") that we
	// will extract to construct the temporary bootstrap control plane.
	bootstrapK8sApps = map[string]struct{}{
		apiServerContainerName:    {},
		"kube-controller-manager": {},
		"kube-scheduler":          {},
	}
	// kubeConfigK8sContainers contains the names of the bootstrap container specs that need to add a
	// --kubeconfig flag to run in non-self-hosted mode.
	kubeConfigK8sContainers = map[string]struct{}{
		"kube-controller-manager": {},
		"kube-scheduler":          {},
	}
	// typeMetas contains a mapping from API object types to the TypeMeta struct that should be
	// populated for them when they are serialized.
	typeMetas    = make(map[reflect.Type]metav1.TypeMeta)
	metaAccessor = meta.NewAccessor()
)

func init() {
	addTypeMeta := func(obj runtime.Object, gv schema.GroupVersion) {
		t := reflect.TypeOf(obj)
		typeMetas[t] = metav1.TypeMeta{
			APIVersion: gv.String(),
			Kind:       t.Elem().Name(),
		}
	}
	addTypeMeta(&v1.ConfigMap{}, v1.SchemeGroupVersion)
	addTypeMeta(&v1beta1.DaemonSet{}, v1beta1.SchemeGroupVersion)
	addTypeMeta(&v1beta1.Deployment{}, v1beta1.SchemeGroupVersion)
	addTypeMeta(&v1.Pod{}, v1.SchemeGroupVersion)
	addTypeMeta(&v1.Secret{}, v1.SchemeGroupVersion)
}

// Backend defines an interface for any backend that can populate a controlPlane struct.
type Backend interface {
	read(context.Context) (*controlPlane, error)
}

// controlPlane holds the control plane objects that are recovered from a backend.
type controlPlane struct {
	configMaps  v1.ConfigMapList
	daemonSets  v1beta1.DaemonSetList
	deployments v1beta1.DeploymentList
	secrets     v1.SecretList
}

// Recover recovers a control plane using the provided backend and kubeConfigPath, returning assets
// for the existing control plane and a bootstrap control plane that can be used with `bootkube
// start` to re-bootstrap the control plane.
func Recover(ctx context.Context, backend Backend, kubeConfigPath string) (asset.Assets, error) {
	cp, err := backend.read(ctx)
	if err != nil {
		return nil, err
	}

	as, err := cp.renderBootstrap()
	if err != nil {
		return nil, err
	}

	kc, err := renderKubeConfig(kubeConfigPath)
	if err != nil {
		return nil, err
	}
	as = append(as, kc)

	return as, nil
}

// renderBootstrap returns assets for a bootstrap control plane that can be used with `bootkube
// start` to re-bootstrap a control plane. These assets are derived from the self-hosted control
// plane that was recovered by the backend, but modified for direct injection into a kubelet.
func (cp *controlPlane) renderBootstrap() (asset.Assets, error) {
	pods, err := extractBootstrapPods(cp.daemonSets.Items, cp.deployments.Items)
	if err != nil {
		return nil, err
	}
	requiredConfigMaps, requiredSecrets := fixUpBootstrapPods(pods)
	as, err := outputBootstrapPods(pods)
	if err != nil {
		return nil, err
	}
	configMaps, err := outputBootstrapConfigMaps(cp.configMaps, requiredConfigMaps)
	if err != nil {
		return nil, err
	}
	as = append(as, configMaps...)

	secrets, err := outputBootstrapSecrets(cp.secrets, requiredSecrets)
	if err != nil {
		return nil, err
	}
	as = append(as, secrets...)
	return as, nil
}

// extractBootstrapPods extracts bootstrap pod specs from daemonsets and deployments.
func extractBootstrapPods(daemonSets []v1beta1.DaemonSet, deployments []v1beta1.Deployment) ([]v1.Pod, error) {
	var pods []v1.Pod
	for _, ds := range daemonSets {
		if isBootstrapApp(ds.Labels) {
			pod := v1.Pod{Spec: ds.Spec.Template.Spec}
			if err := setBootstrapPodMetadata(&pod, ds.ObjectMeta); err != nil {
				return nil, err
			}
			pods = append(pods, pod)
		}
	}
	for _, ds := range deployments {
		if isBootstrapApp(ds.Labels) {
			pod := v1.Pod{Spec: ds.Spec.Template.Spec}
			if err := setBootstrapPodMetadata(&pod, ds.ObjectMeta); err != nil {
				return nil, err
			}
			pods = append(pods, pod)
		}
	}
	return pods, nil
}

// isBootstrapApp returns true if this app belongs to the bootstrap control plane, based on its
// labels.
func isBootstrapApp(labels map[string]string) bool {
	k8sApp := labels[k8sAppLabel]
	if k8sApp == "" {
		k8sApp = labels[componentAppLabel]
	}
	_, ok := bootstrapK8sApps[k8sApp]
	return ok
}

// setBootstrapPodMetadata creates valid metadata for a bootstrap pod. Currently it sets the
// TypeMeta and Name, Namespace, and Annotations on the ObjectMeta.
func setBootstrapPodMetadata(pod *v1.Pod, parent metav1.ObjectMeta) error {
	if err := setTypeMeta(pod); err != nil {
		return err
	}
	pod.ObjectMeta = metav1.ObjectMeta{
		Annotations: parent.Annotations,
		Name:        "bootstrap-" + parent.Name,
		Namespace:   parent.Namespace,
	}
	return nil
}

// fixUpBootstrapPods modifies extracted bootstrap pod specs to have correct metadata and point to
// filesystem-mount-based secrets, and removes any security contexts that might prevent the pods
// from accessing those secrets. It returns mappings from configMap and secret names to output
// paths that must also be rendered in order for the bootstrap pods to be functional.
func fixUpBootstrapPods(pods []v1.Pod) (requiredConfigMaps, requiredSecrets map[string]string) {
	requiredConfigMaps, requiredSecrets = make(map[string]string), make(map[string]string)
	for i := range pods {
		pod := &pods[i]

		// Fix SecurityContext to ensure the pod runs as root.
		if pod.Spec.SecurityContext != nil {
			pod.Spec.SecurityContext.RunAsNonRoot = nil
			pod.Spec.SecurityContext.RunAsUser = nil
		}
		// Fix hostNetwork: true because bootstrap assets can not rely on overlay networks.
		if pod.Spec.HostNetwork == false {
			pod.Spec.HostNetwork = true
		}
		// Change secret volumes to point to file mounts.
		for i := range pod.Spec.Volumes {
			vol := &pod.Spec.Volumes[i]
			if vol.Secret != nil {
				pathSuffix := filepath.Join("secrets", vol.Secret.SecretName)
				requiredSecrets[vol.Secret.SecretName] = filepath.Join(asset.AssetPathSecrets, pathSuffix)
				vol.HostPath = &v1.HostPathVolumeSource{Path: filepath.Join(asset.BootstrapSecretsDir, pathSuffix)}
				vol.Secret = nil
			} else if vol.ConfigMap != nil {
				pathSuffix := filepath.Join("config-maps", vol.ConfigMap.Name)
				requiredConfigMaps[vol.ConfigMap.Name] = filepath.Join(asset.AssetPathSecrets, pathSuffix)
				vol.HostPath = &v1.HostPathVolumeSource{Path: path.Join(asset.BootstrapSecretsDir, pathSuffix)}
				vol.ConfigMap = nil
			}
		}

		// Make sure the kubeconfig is in the commandline.
		for i := range pod.Spec.Containers {
			cn := &pod.Spec.Containers[i]

			// Fix SecurityContext to ensure the container runs as root.
			if cn.SecurityContext != nil {
				cn.SecurityContext.RunAsNonRoot = nil
				cn.SecurityContext.RunAsUser = nil
			}

			// Assumes the bootkube naming convention is used. Could also just make sure the image uses hyperkube.
			if _, ok := kubeConfigK8sContainers[cn.Name]; ok {
				cn.Command = append(cn.Command, "--kubeconfig=/kubeconfig/kubeconfig")
				cn.VolumeMounts = append(cn.VolumeMounts, v1.VolumeMount{
					MountPath: "/kubeconfig",
					Name:      "kubeconfig",
					ReadOnly:  true,
				})
			}
		}

		// Add a mount for the kubeconfig.
		pod.Spec.Volumes = append(pod.Spec.Volumes, v1.Volume{
			VolumeSource: v1.VolumeSource{HostPath: &v1.HostPathVolumeSource{Path: kubeletKubeConfigPath}},
			Name:         "kubeconfig",
		})
	}
	return
}

// outputBootstrapPods outputs the bootstrap pod definitions.
func outputBootstrapPods(pods []v1.Pod) (asset.Assets, error) {
	var as asset.Assets
	for _, pod := range pods {
		a, err := serializeObjToYAML(path.Join(asset.AssetPathBootstrapManifests, pod.Name+".yaml"), &pod)
		if err != nil {
			return nil, err
		}
		as = append(as, a)
	}
	return as, nil
}

// outputBootstrapConfigMaps creates assets for all the configMap names in the requiredConfigMaps
// set. It returns an error if any configMap cannot be found in the provided configMaps list.
func outputBootstrapConfigMaps(configMaps v1.ConfigMapList, requiredConfigMaps map[string]string) (asset.Assets, error) {
	return outputKeyValueData(&configMaps, requiredConfigMaps, func(obj runtime.Object) map[string][]byte {
		configMap, ok := obj.(*v1.ConfigMap)
		if !ok || configMap == nil {
			return nil
		}
		output := make(map[string][]byte)
		for k, v := range configMap.Data {
			output[k] = []byte(v)
		}
		return output
	})
}

// outputBootstrapSecrets creates assets for all the secret names in the requiredSecrets set. It
// returns an error if any secret cannot be found in the provided secrets list.
func outputBootstrapSecrets(secrets v1.SecretList, requiredSecrets map[string]string) (asset.Assets, error) {
	return outputKeyValueData(&secrets, requiredSecrets, func(obj runtime.Object) map[string][]byte {
		if secret, ok := obj.(*v1.Secret); ok && secret != nil {
			return secret.Data
		}
		return nil
	})
}

// outputKeyValueData takes a key-value object (such as a Secret or ConfigMap) and outputs assets
// for each key-value pair. See outputBootstrapConfigMaps or outputBootstrapSecrets for usage.
func outputKeyValueData(objList runtime.Object, requiredObjs map[string]string, extractData func(runtime.Object) map[string][]byte) (asset.Assets, error) {
	var as asset.Assets
	objs, err := meta.ExtractList(objList)
	if err != nil {
		return nil, err
	}
	for _, obj := range objs {
		name, err := metaAccessor.Name(obj)
		if err != nil {
			return nil, err
		}
		if namePrefix, ok := requiredObjs[name]; ok {
			for key, data := range extractData(obj) {
				as = append(as, asset.Asset{
					Name: path.Join(namePrefix, key),
					Data: data,
				})
			}
			delete(requiredObjs, name)
		}
	}
	if len(requiredObjs) > 0 {
		var missingObjs []string
		for obj := range requiredObjs {
			missingObjs = append(missingObjs, obj)
		}
		return nil, fmt.Errorf("failed to extract some required objects: %v", missingObjs)
	}
	return as, nil
}

// renderKubeConfig outputs kubeconfig assets to ensure that the kubeconfig will be rendered to the
// assetDir for use by `bootkube start`.
func renderKubeConfig(kubeConfigPath string) (asset.Asset, error) {
	kubeConfig, err := ioutil.ReadFile(kubeConfigPath)
	if err != nil {
		return asset.Asset{}, err
	}
	return asset.Asset{
		Name: asset.AssetPathAdminKubeConfig, // used by `bootkube start`.
		Data: kubeConfig,
	}, nil
}

// setTypeMeta sets the TypeMeta for a runtime.Object.
// TODO(diegs): find the apimachinery code that does this, and use that instead.
func setTypeMeta(obj runtime.Object) error {
	typeMeta, ok := typeMetas[reflect.TypeOf(obj)]
	if !ok {
		return fmt.Errorf("don't know about type: %T", obj)
	}
	metaAccessor.SetAPIVersion(obj, typeMeta.APIVersion)
	metaAccessor.SetKind(obj, typeMeta.Kind)
	return nil
}

// serializeObjToYAML serializes a runtime.Object into a YAML asset.
func serializeObjToYAML(assetName string, obj runtime.Object) (asset.Asset, error) {
	data, err := yaml.Marshal(obj)
	if err != nil {
		return asset.Asset{}, err
	}
	return asset.Asset{
		Name: assetName,
		Data: data,
	}, nil
}
<|endoftext|>


<commit_before>package folderimpl

import (
	"context"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/infra/db"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/featuremgmt"
	"github.com/grafana/grafana/pkg/services/folder"
	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/util"
)

type sqlStore struct {
	db  db.DB
	log log.Logger
	cfg *setting.Cfg
	fm  featuremgmt.FeatureManager
}

// sqlStore implements the store interface.
var _ store = (*sqlStore)(nil)

func ProvideStore(db db.DB, cfg *setting.Cfg, features featuremgmt.FeatureManager) *sqlStore {
	return &sqlStore{db: db, log: log.New("folder-store"), cfg: cfg, fm: features}
}

func (ss *sqlStore) Create(ctx context.Context, cmd folder.CreateFolderCommand) (*folder.Folder, error) {
	if cmd.UID == "" {
		return nil, folder.ErrBadRequest.Errorf("missing UID")
	}

	var foldr *folder.Folder
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		var sqlOrArgs []interface{}
		if cmd.ParentUID == "" {
			sql := "INSERT INTO folder(org_id, uid, title, description, created, updated) VALUES(?, ?, ?, ?, ?, ?)"
			sqlOrArgs = []interface{}{sql, cmd.OrgID, cmd.UID, cmd.Title, cmd.Description, time.Now(), time.Now()}
		} else {
			if cmd.ParentUID != folder.GeneralFolderUID {
				if _, err := ss.Get(ctx, folder.GetFolderQuery{
					UID:   &cmd.ParentUID,
					OrgID: cmd.OrgID,
				}); err != nil {
					return err
				}
			}
			sql := "INSERT INTO folder(org_id, uid, parent_uid, title, description, created, updated) VALUES(?, ?, ?, ?, ?, ?, ?)"
			sqlOrArgs = []interface{}{sql, cmd.OrgID, cmd.UID, cmd.ParentUID, cmd.Title, cmd.Description, time.Now(), time.Now()}
		}
		res, err := sess.Exec(sqlOrArgs...)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to insert folder: %w", err)
		}
		id, err := res.LastInsertId()
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get last inserted id: %w", err)
		}

		foldr, err = ss.Get(ctx, folder.GetFolderQuery{
			ID: &id,
		})
		if err != nil {
			return err
		}
		return nil
	})
	return foldr, err
}

func (ss *sqlStore) Delete(ctx context.Context, uid string, orgID int64) error {
	return ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		_, err := sess.Exec("DELETE FROM folder WHERE uid=? AND org_id=?", uid, orgID)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to delete folder: %w", err)
		}
		/*
			affected, err := res.RowsAffected()
			if err != nil {
				return folder.ErrDatabaseError.Errorf("failed to get affected rows: %w", err)
			}
				if affected == 0 {
					return folder.ErrFolderNotFound.Errorf("folder not found uid:%s org_id:%d", uid, orgID)
				}
		*/
		return nil
	})
}

func (ss *sqlStore) Update(ctx context.Context, cmd folder.UpdateFolderCommand) (*folder.Folder, error) {
	if cmd.Folder == nil {
		return nil, folder.ErrBadRequest.Errorf("invalid update command: missing folder")
	}

	cmd.Folder.Updated = time.Now()
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		description := cmd.Folder.Description
		if cmd.NewDescription != nil {
			description = *cmd.NewDescription
		}

		title := cmd.Folder.Title
		if cmd.NewTitle != nil {
			title = *cmd.NewTitle
		}

		uid := cmd.Folder.UID
		if cmd.NewUID != nil {
			uid = *cmd.NewUID
		}

		res, err := sess.Exec("UPDATE folder SET description = ?, title = ?, uid = ?, updated = ? WHERE uid = ? AND org_id = ?", description, title, uid, cmd.Folder.Updated, cmd.Folder.UID, cmd.Folder.OrgID)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to update folder: %w", err)
		}

		affected, err := res.RowsAffected()
		if err != nil {
			return folder.ErrInternal.Errorf("failed to get affected row: %w", err)
		}
		if affected == 0 {
			return folder.ErrInternal.Errorf("no folders are updated")
		}

		cmd.Folder.Description = description
		cmd.Folder.Title = title
		cmd.Folder.UID = uid
		return nil
	})

	return cmd.Folder, err
}

func (ss *sqlStore) Get(ctx context.Context, q folder.GetFolderQuery) (*folder.Folder, error) {
	foldr := &folder.Folder{}
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		exists := false
		var err error

		switch {
		case q.ID != nil:
			exists, err = sess.SQL("SELECT * FROM folder WHERE id = ?", q.ID).Get(foldr)
		case q.Title != nil:
			exists, err = sess.SQL("SELECT * FROM folder WHERE title = ? AND org_id = ?", q.Title, q.OrgID).Get(foldr)
		case q.UID != nil:
			exists, err = sess.SQL("SELECT * FROM folder WHERE uid = ? AND org_id = ?", q.UID, q.OrgID).Get(foldr)
		default:
			return folder.ErrBadRequest.Errorf("one of ID, UID, or Title must be included in the command")
		}
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get folder: %w", err)
		}
		if !exists {
			return folder.ErrFolderNotFound.Errorf("folder not found")
		}
		return nil
	})
	return foldr, err
}

func (ss *sqlStore) GetParents(ctx context.Context, q folder.GetParentsQuery) ([]*folder.Folder, error) {
	var folders []*folder.Folder
	if ss.db.GetDBType() == migrator.MySQL {
		return ss.getParentsMySQL(ctx, q)
	}

	recQuery := `
		WITH RECURSIVE RecQry AS (
			SELECT * FROM folder WHERE uid = ? AND org_id = ?
			UNION ALL SELECT f.* FROM folder f INNER JOIN RecQry r ON f.uid = r.parent_uid and f.org_id = r.org_id
		)
		SELECT * FROM RecQry;
	`

	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		err := sess.SQL(recQuery, q.UID, q.OrgID).Find(&folders)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get folder parents: %w", err)
		}
		return nil
	})
	return util.Reverse(folders[1:]), err
}

func (ss *sqlStore) GetChildren(ctx context.Context, q folder.GetTreeQuery) ([]*folder.Folder, error) {
	var folders []*folder.Folder

	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		sql := strings.Builder{}
		sql.Write([]byte("SELECT * FROM folder WHERE parent_uid=? AND org_id=?"))

		if q.Limit != 0 {
			var offset int64 = 1
			if q.Page != 0 {
				offset = q.Page
			}
			sql.Write([]byte(ss.db.GetDialect().LimitOffset(q.Limit, offset)))
		}
		err := sess.SQL(sql.String(), q.UID, q.OrgID).Find(&folders)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get folder children: %w", err)
		}
		return nil
	})
	return folders, err
}

func (ss *sqlStore) getParentsMySQL(ctx context.Context, cmd folder.GetParentsQuery) ([]*folder.Folder, error) {
	var foldrs []*folder.Folder
	var foldr *folder.Folder
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		uid := cmd.UID
		for uid != folder.GeneralFolderUID && len(foldrs) < 8 {
			err := sess.Where("uid=? AND org_id=>", uid, cmd.OrgID).Find(foldr)
			if err != nil {
				return folder.ErrDatabaseError.Errorf("failed to get folder parents: %w", err)
			}
			foldrs = append(foldrs, foldr)
			uid = foldr.ParentUID
		}
		return nil
	})
	return foldrs, err
}
<commit_msg>Nested Folder: Modify store Update() (#58183)<commit_after>package folderimpl

import (
	"context"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/infra/db"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/featuremgmt"
	"github.com/grafana/grafana/pkg/services/folder"
	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/util"
)

type sqlStore struct {
	db  db.DB
	log log.Logger
	cfg *setting.Cfg
	fm  featuremgmt.FeatureManager
}

// sqlStore implements the store interface.
var _ store = (*sqlStore)(nil)

func ProvideStore(db db.DB, cfg *setting.Cfg, features featuremgmt.FeatureManager) *sqlStore {
	return &sqlStore{db: db, log: log.New("folder-store"), cfg: cfg, fm: features}
}

func (ss *sqlStore) Create(ctx context.Context, cmd folder.CreateFolderCommand) (*folder.Folder, error) {
	if cmd.UID == "" {
		return nil, folder.ErrBadRequest.Errorf("missing UID")
	}

	var foldr *folder.Folder
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		var sqlOrArgs []interface{}
		if cmd.ParentUID == "" {
			sql := "INSERT INTO folder(org_id, uid, title, description, created, updated) VALUES(?, ?, ?, ?, ?, ?)"
			sqlOrArgs = []interface{}{sql, cmd.OrgID, cmd.UID, cmd.Title, cmd.Description, time.Now(), time.Now()}
		} else {
			if cmd.ParentUID != folder.GeneralFolderUID {
				if _, err := ss.Get(ctx, folder.GetFolderQuery{
					UID:   &cmd.ParentUID,
					OrgID: cmd.OrgID,
				}); err != nil {
					return err
				}
			}
			sql := "INSERT INTO folder(org_id, uid, parent_uid, title, description, created, updated) VALUES(?, ?, ?, ?, ?, ?, ?)"
			sqlOrArgs = []interface{}{sql, cmd.OrgID, cmd.UID, cmd.ParentUID, cmd.Title, cmd.Description, time.Now(), time.Now()}
		}
		res, err := sess.Exec(sqlOrArgs...)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to insert folder: %w", err)
		}
		id, err := res.LastInsertId()
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get last inserted id: %w", err)
		}

		foldr, err = ss.Get(ctx, folder.GetFolderQuery{
			ID: &id,
		})
		if err != nil {
			return err
		}
		return nil
	})
	return foldr, err
}

func (ss *sqlStore) Delete(ctx context.Context, uid string, orgID int64) error {
	return ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		_, err := sess.Exec("DELETE FROM folder WHERE uid=? AND org_id=?", uid, orgID)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to delete folder: %w", err)
		}
		/*
			affected, err := res.RowsAffected()
			if err != nil {
				return folder.ErrDatabaseError.Errorf("failed to get affected rows: %w", err)
			}
				if affected == 0 {
					return folder.ErrFolderNotFound.Errorf("folder not found uid:%s org_id:%d", uid, orgID)
				}
		*/
		return nil
	})
}

func (ss *sqlStore) Update(ctx context.Context, cmd folder.UpdateFolderCommand) (*folder.Folder, error) {
	if cmd.Folder == nil {
		return nil, folder.ErrBadRequest.Errorf("invalid update command: missing folder")
	}

	cmd.Folder.Updated = time.Now()
	existingUID := cmd.Folder.UID
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		sql := strings.Builder{}
		sql.Write([]byte("UPDATE folder SET "))
		columnsToUpdate := []string{"updated = ?"}
		args := []interface{}{cmd.Folder.Updated}
		if cmd.NewDescription != nil {
			columnsToUpdate = append(columnsToUpdate, "description = ?")
			cmd.Folder.Description = *cmd.NewDescription
			args = append(args, cmd.Folder.Description)
		}

		if cmd.NewTitle != nil {
			columnsToUpdate = append(columnsToUpdate, "title = ?")
			cmd.Folder.Title = *cmd.NewTitle
			args = append(args, cmd.Folder.Title)
		}

		if cmd.NewUID != nil {
			columnsToUpdate = append(columnsToUpdate, "uid = ?")
			cmd.Folder.UID = *cmd.NewUID
			args = append(args, cmd.Folder.UID)
		}

		if len(columnsToUpdate) == 0 {
			return folder.ErrBadRequest.Errorf("no columns to update")
		}

		sql.Write([]byte(strings.Join(columnsToUpdate, ", ")))
		sql.Write([]byte(" WHERE uid = ? AND org_id = ?"))
		args = append(args, existingUID, cmd.Folder.OrgID)

		args = append([]interface{}{sql.String()}, args...)

		res, err := sess.Exec(args...)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to update folder: %w", err)
		}

		affected, err := res.RowsAffected()
		if err != nil {
			return folder.ErrInternal.Errorf("failed to get affected row: %w", err)
		}
		if affected == 0 {
			return folder.ErrInternal.Errorf("no folders are updated")
		}
		return nil
	})

	return cmd.Folder, err
}

func (ss *sqlStore) Get(ctx context.Context, q folder.GetFolderQuery) (*folder.Folder, error) {
	foldr := &folder.Folder{}
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		exists := false
		var err error

		switch {
		case q.ID != nil:
			exists, err = sess.SQL("SELECT * FROM folder WHERE id = ?", q.ID).Get(foldr)
		case q.Title != nil:
			exists, err = sess.SQL("SELECT * FROM folder WHERE title = ? AND org_id = ?", q.Title, q.OrgID).Get(foldr)
		case q.UID != nil:
			exists, err = sess.SQL("SELECT * FROM folder WHERE uid = ? AND org_id = ?", q.UID, q.OrgID).Get(foldr)
		default:
			return folder.ErrBadRequest.Errorf("one of ID, UID, or Title must be included in the command")
		}
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get folder: %w", err)
		}
		if !exists {
			return folder.ErrFolderNotFound.Errorf("folder not found")
		}
		return nil
	})
	return foldr, err
}

func (ss *sqlStore) GetParents(ctx context.Context, q folder.GetParentsQuery) ([]*folder.Folder, error) {
	var folders []*folder.Folder
	if ss.db.GetDBType() == migrator.MySQL {
		return ss.getParentsMySQL(ctx, q)
	}

	recQuery := `
		WITH RECURSIVE RecQry AS (
			SELECT * FROM folder WHERE uid = ? AND org_id = ?
			UNION ALL SELECT f.* FROM folder f INNER JOIN RecQry r ON f.uid = r.parent_uid and f.org_id = r.org_id
		)
		SELECT * FROM RecQry;
	`

	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		err := sess.SQL(recQuery, q.UID, q.OrgID).Find(&folders)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get folder parents: %w", err)
		}
		return nil
	})
	return util.Reverse(folders[1:]), err
}

func (ss *sqlStore) GetChildren(ctx context.Context, q folder.GetTreeQuery) ([]*folder.Folder, error) {
	var folders []*folder.Folder

	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		sql := strings.Builder{}
		sql.Write([]byte("SELECT * FROM folder WHERE parent_uid=? AND org_id=?"))

		if q.Limit != 0 {
			var offset int64 = 1
			if q.Page != 0 {
				offset = q.Page
			}
			sql.Write([]byte(ss.db.GetDialect().LimitOffset(q.Limit, offset)))
		}
		err := sess.SQL(sql.String(), q.UID, q.OrgID).Find(&folders)
		if err != nil {
			return folder.ErrDatabaseError.Errorf("failed to get folder children: %w", err)
		}
		return nil
	})
	return folders, err
}

func (ss *sqlStore) getParentsMySQL(ctx context.Context, cmd folder.GetParentsQuery) ([]*folder.Folder, error) {
	var foldrs []*folder.Folder
	var foldr *folder.Folder
	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
		uid := cmd.UID
		for uid != folder.GeneralFolderUID && len(foldrs) < 8 {
			err := sess.Where("uid=? AND org_id=>", uid, cmd.OrgID).Find(foldr)
			if err != nil {
				return folder.ErrDatabaseError.Errorf("failed to get folder parents: %w", err)
			}
			foldrs = append(foldrs, foldr)
			uid = foldr.ParentUID
		}
		return nil
	})
	return foldrs, err
}
<|endoftext|>