<commit_before>// Demonstrate how to use rest.RouteObjectMethod
//
// rest.RouteObjectMethod helps create a Route that points to
// an object method instead of just a function.
//
// The Curl Demo:
// curl -i -d '{"Name":"Antoine"}' http://127.0.0.1:8080/users
// curl -i http://127.0.0.1:8080/users/0
//
package main

import (
	"fmt"
	"github.com/ant0ine/go-json-rest"
	"net/http"
)

type User struct {
	Id   string
	Name string
}

type Users struct {
	Store map[string]*User
}

func (self *Users) GetUser(w *rest.ResponseWriter, r *rest.Request) {
	id := r.PathParam("id")
	user := self.Store[id]
	if user == nil {
		http.NotFound(w, r.Request)
		return
	}
	w.WriteJson(&user)
}

func (self *Users) PostUser(w *rest.ResponseWriter, r *rest.Request) {
	user := User{}
	err := r.DecodeJsonPayload(&user)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	id := fmt.Sprintf("%d", len(self.Store)) // stupid
	user.Id = id
	self.Store[id] = &user
	w.WriteJson(&user)
}

func main() {

	users := Users{
		Store: map[string]*User{},
	}

	handler := rest.ResourceHandler{}
	handler.SetRoutes(
		rest.RouteObjectMethod("GET", "/users/:id", &users, "GetUser"),
		rest.RouteObjectMethod("POST", "/users", &users, "PostUser"),
	)
	http.ListenAndServe(":8080", &handler)
}
<commit_msg>implement all the CRUD ops for this example.<commit_after>// Demonstrate how to use rest.RouteObjectMethod
//
// rest.RouteObjectMethod helps create a Route that points to
// an object method instead of just a function.
//
// The Curl Demo:
// curl -i -d '{"Name":"Antoine"}' http://127.0.0.1:8080/users
// curl -i http://127.0.0.1:8080/users/0
// curl -i -X PUT -d '{"Name":"Antoine Imbert"}' http://127.0.0.1:8080/users/0
// curl -i -X DELETE http://127.0.0.1:8080/users/0
// curl -i http://127.0.0.1:8080/users
//
package main

import (
	"fmt"
	"github.com/ant0ine/go-json-rest"
	"net/http"
)

type User struct {
	Id   string
	Name string
}

type Users struct {
	Store map[string]*User
}

func (self *Users) GetAllUsers(w *rest.ResponseWriter, r *rest.Request) {
	users := []*User{}
	for _, user := range self.Store {
		users = append(users, user)
	}
	w.WriteJson(&users)
}

func (self *Users) GetUser(w *rest.ResponseWriter, r *rest.Request) {
	id := r.PathParam("id")
	user := self.Store[id]
	if user == nil {
		http.NotFound(w, r.Request)
		return
	}
	w.WriteJson(&user)
}

func (self *Users) PostUser(w *rest.ResponseWriter, r *rest.Request) {
	user := User{}
	err := r.DecodeJsonPayload(&user)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	id := fmt.Sprintf("%d", len(self.Store)) // stupid
	user.Id = id
	self.Store[id] = &user
	w.WriteJson(&user)
}

func (self *Users) PutUser(w *rest.ResponseWriter, r *rest.Request) {
	id := r.PathParam("id")
	if self.Store[id] == nil {
		http.NotFound(w, r.Request)
		return
	}
	user := User{}
	err := r.DecodeJsonPayload(&user)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	user.Id = id
	self.Store[id] = &user
	w.WriteJson(&user)
}

func (self *Users) DeleteUser(w *rest.ResponseWriter, r *rest.Request) {
	id := r.PathParam("id")
	delete(self.Store, id)
}

func main() {

	users := Users{
		Store: map[string]*User{},
	}

	handler := rest.ResourceHandler{}
	handler.SetRoutes(
		rest.RouteObjectMethod("GET", "/users", &users, "GetAllUsers"),
		rest.RouteObjectMethod("POST", "/users", &users, "PostUser"),
		rest.RouteObjectMethod("GET", "/users/:id", &users, "GetUser"),
		rest.RouteObjectMethod("PUT", "/users/:id", &users, "PutUser"),
		rest.RouteObjectMethod("DELETE", "/users/:id", &users, "DeleteUser"),
	)
	http.ListenAndServe(":8080", &handler)
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 sdb

import (
)

// The name of an item within a SimpleDB domain. Item names must be UTF-8
// strings no longer than 1024 bytes. They must contain only characters that
// are valid in XML 1.0 documents, as defined by Section 2.2 of the XML 1.0
// spec. (Note that this is a more restrictive condition than imposed by
// SimpleDB itself, and is done for the sake of Go's XML 1.0 parser.)
//
// For more info:
//
//     http://goo.gl/Fkjnz
//     http://goo.gl/csem8
//
type ItemName string

// An attribute is a (name, value) pair possessed by an item. Items contain
// sets of attributes; they may contain multiple attributes with the same name,
// but not with the same (name, value) pair.
//
// Attribute names and values share the same restrictions as those on item
// names.
type Attribute struct {
	Name string
	Value string
}

// A precondition for a conditional Put or Delete operation. Preconditions may
// specify a value that an attribute must have or whether the attribute must
// exist or not.
type Precondition struct {
	// The name of the attribute to be inspected. Attributes with multiple values
	// are not supported.
	Name string

	// If present, the value that the attribute must possess at the time of the
	// update. Must be present iff Exists is not present.
	Value *string

	// If present, whether the attribute must exist at the time of the update.
	// Must be present iff Value is not present.
	Exists *bool
}

// An update to make to a particular attribute as part of a Put request.
type PutUpdate struct {
	// The name of the attribute.
	Name string

	// The value to set for the attribute.
	Value string

	// Whether to replace existing values for the attribute or to simply add a
	// new one.
	Replace bool
}

// An update to make to a particular attribute as part of a Delete request.
type DeleteUpdate struct {
	// The name of the attribute.
	Name string

	// Te requests, the particular value of the attribute to delete if present.
	// Otherwise, all values will be deleted.
	Value *string
}

// A domain represents a named domain within the SimpleDB service. It is a
// collection of named items, each of which possesses a set of attributes.
type Domain interface {
	// Atomically apply the supplied updates to the attributes of the named item,
	// but only if the supplied preconditions hold.
	//
	// The length of updates must be in [1, 256].
	PutAttributes(
		item ItemName,
		updates []PutUpdate,
		preconditions []Precondition) error

	// Atomically apply updates to multiple items simultaneously.
	//
	// The length of updates must be in [1, 25]. The length of each of its values
	// must be in [1, 256]. An error may be returned if the underlying request to
	// SimpleDB is too large.
	BatchPutAttributes(updates map[ItemName][]PutUpdate) error

	// Atomically delete attributes from the named item, but only if the supplied
	// preconditions hold.
	//
	// If deletes is empty, delete all attributes from the item. Otherwise
	// perform only the deletes is specifies. Deleting a non-existent attribute
	// does not result in an error.
	DeleteAttributes(
		item ItemName,
		deletes []DeleteUpdate,
		preconditions []Precondition) error

	// Atomically delete attributes from multiple items simultaneously.
	//
	// If no updates are supplied for a particular item, delete all of its
	// attributes.
	BatchDeleteAttributes(deletes map[ItemName][]DeleteUpdate) error

	GetAttributes(
		item ItemName,
		constistentRead bool,
		attributes []string) ([]Attribute, error)

	Select(
		query string,
		constistentRead bool,
		nextToken []byte) (res map[ItemName][]Attribute, tok []byte, err error)
}
<commit_msg>Added Get comments.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 sdb

import (
)

// The name of an item within a SimpleDB domain. Item names must be UTF-8
// strings no longer than 1024 bytes. They must contain only characters that
// are valid in XML 1.0 documents, as defined by Section 2.2 of the XML 1.0
// spec. (Note that this is a more restrictive condition than imposed by
// SimpleDB itself, and is done for the sake of Go's XML 1.0 parser.)
//
// For more info:
//
//     http://goo.gl/Fkjnz
//     http://goo.gl/csem8
//
type ItemName string

// An attribute is a (name, value) pair possessed by an item. Items contain
// sets of attributes; they may contain multiple attributes with the same name,
// but not with the same (name, value) pair.
//
// Attribute names and values share the same restrictions as those on item
// names.
type Attribute struct {
	Name string
	Value string
}

// A precondition for a conditional Put or Delete operation. Preconditions may
// specify a value that an attribute must have or whether the attribute must
// exist or not.
type Precondition struct {
	// The name of the attribute to be inspected. Attributes with multiple values
	// are not supported.
	Name string

	// If present, the value that the attribute must possess at the time of the
	// update. Must be present iff Exists is not present.
	Value *string

	// If present, whether the attribute must exist at the time of the update.
	// Must be present iff Value is not present.
	Exists *bool
}

// An update to make to a particular attribute as part of a Put request.
type PutUpdate struct {
	// The name of the attribute.
	Name string

	// The value to set for the attribute.
	Value string

	// Whether to replace existing values for the attribute or to simply add a
	// new one.
	Replace bool
}

// An update to make to a particular attribute as part of a Delete request.
type DeleteUpdate struct {
	// The name of the attribute.
	Name string

	// Te requests, the particular value of the attribute to delete if present.
	// Otherwise, all values will be deleted.
	Value *string
}

// A domain represents a named domain within the SimpleDB service. It is a
// collection of named items, each of which possesses a set of attributes.
type Domain interface {
	// Atomically apply the supplied updates to the attributes of the named item,
	// but only if the supplied preconditions hold.
	//
	// The length of updates must be in [1, 256].
	PutAttributes(
		item ItemName,
		updates []PutUpdate,
		preconditions []Precondition) error

	// Atomically apply updates to multiple items simultaneously.
	//
	// The length of updates must be in [1, 25]. The length of each of its values
	// must be in [1, 256]. An error may be returned if the underlying request to
	// SimpleDB is too large.
	BatchPutAttributes(updates map[ItemName][]PutUpdate) error

	// Atomically delete attributes from the named item, but only if the supplied
	// preconditions hold.
	//
	// If deletes is empty, delete all attributes from the item. Otherwise
	// perform only the deletes is specifies. Deleting a non-existent attribute
	// does not result in an error.
	DeleteAttributes(
		item ItemName,
		deletes []DeleteUpdate,
		preconditions []Precondition) error

	// Atomically delete attributes from multiple items simultaneously.
	//
	// If no updates are supplied for a particular item, delete all of its
	// attributes.
	BatchDeleteAttributes(deletes map[ItemName][]DeleteUpdate) error

	// Retrieve a set of attributes for the named item, or all attributes if the
	// attributes slice is empty.
	//
	// If the named item doesn't exist, the empty set is returned.
	//
	// constistentRead specifies whether completely fresh data is needed or not.
	GetAttributes(
		item ItemName,
		constistentRead bool,
		attributes []string) ([]Attribute, error)

	Select(
		query string,
		constistentRead bool,
		nextToken []byte) (res map[ItemName][]Attribute, tok []byte, err error)
}
<|endoftext|>


<commit_before>package cliutil

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	grpcCodes "google.golang.org/grpc/codes"
	grpcStatus "google.golang.org/grpc/status"
	empty "google.golang.org/protobuf/types/known/emptypb"
	"gopkg.in/yaml.v2"

	"github.com/datawire/dlib/dexec"
	"github.com/telepresenceio/telepresence/rpc/v2/connector"
	"github.com/telepresenceio/telepresence/v2/pkg/client"
	"github.com/telepresenceio/telepresence/v2/pkg/client/errcat"
	"github.com/telepresenceio/telepresence/v2/pkg/client/scout"
	"github.com/telepresenceio/telepresence/v2/pkg/client/userd/auth/authdata"
	"github.com/telepresenceio/telepresence/v2/pkg/filelocation"
)

// EnsureLoggedIn ensures that the user is logged in to Ambassador Cloud.  An error is returned if
// login fails.  The result code will indicate if this is a new login or if it resued an existing
// login.  If the `apikey` argument is empty an interactive login is performed; if it is non-empty
// the key is used instead of performing an interactive login.
func EnsureLoggedIn(ctx context.Context, apikey string) (connector.LoginResult_Code, error) {
	err := GetTelepresencePro(ctx)
	if err != nil {
		return connector.LoginResult_UNSPECIFIED, err
	}
	var code connector.LoginResult_Code
	err = WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		code, err = ClientEnsureLoggedIn(ctx, apikey, connectorClient)
		return err
	})
	return code, err
}

// ClientEnsureLoggedIn is like EnsureLoggedIn but uses an already acquired ConnectorClient.
func ClientEnsureLoggedIn(ctx context.Context, apikey string, connectorClient connector.ConnectorClient) (connector.LoginResult_Code, error) {
	resp, err := connectorClient.Login(ctx, &connector.LoginRequest{
		ApiKey: apikey,
	})
	if err != nil {
		if grpcStatus.Code(err) == grpcCodes.PermissionDenied {
			err = errcat.User.New(grpcStatus.Convert(err).Message())
		}
		return connector.LoginResult_UNSPECIFIED, err
	}
	return resp.GetCode(), nil
}

// Logout logs out of Ambassador Cloud.  Returns an error if not logged in.
func Logout(ctx context.Context) error {
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		_, err := connectorClient.Logout(ctx, &empty.Empty{})
		return err
	})
	if grpcStatus.Code(err) == grpcCodes.NotFound {
		err = errcat.User.New(grpcStatus.Convert(err).Message())
	}
	if err != nil {
		return err
	}
	return nil
}

// EnsureLoggedOut ensures that the user is logged out of Ambassador Cloud.  Returns nil if not
// logged in.
func EnsureLoggedOut(ctx context.Context) error {
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		_, err := connectorClient.Logout(ctx, &empty.Empty{})
		return err
	})
	if grpcStatus.Code(err) == grpcCodes.NotFound {
		err = nil
	}
	if err != nil {
		return err
	}
	return nil
}

// HasLoggedIn returns true if either the user has an active login session or an expired login
// session, and returns false if either the user has never logged in or has explicitly logged out.
func HasLoggedIn(ctx context.Context) bool {
	_, err := authdata.LoadUserInfoFromUserCache(ctx)
	return err == nil
}

func GetCloudUserInfo(ctx context.Context, autoLogin bool, refresh bool) (*connector.UserInfo, error) {
	var userInfo *connector.UserInfo
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		userInfo, err = connectorClient.GetCloudUserInfo(ctx, &connector.UserInfoRequest{
			AutoLogin: autoLogin,
			Refresh:   refresh,
		})
		return err
	})
	if err != nil {
		return nil, err
	}
	return userInfo, nil
}

func GetCloudAPIKey(ctx context.Context, description string, autoLogin bool) (string, error) {
	var keyData *connector.KeyData
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		keyData, err = connectorClient.GetCloudAPIKey(ctx, &connector.KeyRequest{
			AutoLogin:   autoLogin,
			Description: description,
		})
		return err
	})
	if err != nil {
		return "", err
	}
	return keyData.GetApiKey(), nil
}

// GetCloudLicense communicates with system a to get the jwt version of the
// license, puts it in a kubernetes secret, and then writes that secret to the
// output file for the user to apply to their cluster
func GetCloudLicense(ctx context.Context, outputFile, id string) (string, string, error) {
	var licenseData *connector.LicenseData
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		licenseData, err = connectorClient.GetCloudLicense(ctx, &connector.LicenseRequest{
			Id: id,
		})
		return err
	})
	if err != nil {
		return "", "", err
	}
	return licenseData.GetLicense(), licenseData.GetHostDomain(), nil
}

// GetTelepresencePro prompts the user to optionally install Telepresence Pro
// if it isn't installed. If the user installs it, it also asks the user to
// automatically update their configuration to use the new binary.
func GetTelepresencePro(ctx context.Context) error {
	dir, err := filelocation.AppUserConfigDir(ctx)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("unable to get path to config files: %w", err)
	}

	sc := scout.NewReporter(ctx, "cli")
	sc.Start(ctx)
	defer sc.Close()
	installRefused := false
	defer func() {
		if err != nil {
			sc.Report(ctx, "pro_connector_upgrade_fail", scout.Entry{Key: "error", Value: err.Error()})
		} else if installRefused {
			sc.Report(ctx, "pro_connector_upgrade_refusal")
		} else {
			sc.Report(ctx, "pro_connector_upgrade_success")
		}
	}()

	// If telepresence-pro doesn't exist, then we should ask the user
	// if they want to install it
	telProLocation := filepath.Join(dir, "telepresence-pro")
	if runtime.GOOS == "windows" {
		telProLocation += ".exe"
	}
	if _, err = os.Stat(telProLocation); os.IsNotExist(err) {
		sc.SetMetadatum(ctx, "first_install", true)
		reader := bufio.NewReader(os.Stdin)
		fmt.Printf("Telepresence Pro is recommended when using login features, can Telepresence install it? (y/n)")
		var reply string
		reply, err = reader.ReadString('\n')
		if err != nil {
			return errcat.User.Newf("error reading input: %w", err)
		}

		// If the user doesn't want to install it, then we we'll proceed
		// with launching the daemon normally
		reply = strings.TrimSpace(reply)
		if reply != "y" {
			installRefused = true
			return nil
		}

		err = installTelepresencePro(ctx, telProLocation)
		if err != nil {
			return err
		}

		// Ask the user if they want to automatically update their config
		// with the telepresence-pro binary.
		// TODO: This will remove any comments that exist in the config file
		// which it's yaml so that's _fine_ but it would be nice if we didn't
		// do that.
		fmt.Printf("Update your Telepresence config to use Telepresence Pro? (y/n)")
		reply, err = reader.ReadString('\n')
		if err != nil {
			return errcat.User.Newf("error reading input: %w", err)
		}
		reply = strings.TrimSpace(reply)
		if reply != "y" {
			return nil
		}
		err = updateConfig(ctx, telProLocation)
		if err != nil {
			return err
		}
	} else {
		// If the binary is present, we check its version to ensure it's compatible
		// with the CLI
		sc.SetMetadatum(ctx, "first_install", false)
		proCmd := dexec.CommandContext(ctx, telProLocation, "pro-version")
		proCmd.DisableLogging = true

		var output []byte
		output, err = proCmd.CombinedOutput()
		if err != nil {
			return errcat.NoDaemonLogs.Newf("Unable to get telepresence pro version")
		}

		if !strings.Contains(string(output), client.Version()) {
			reader := bufio.NewReader(os.Stdin)
			fmt.Printf("Telepresence Pro needs to be upgraded to work with CLI version %s, allow Telepresence to upgrade it? (y/n)",
				client.Version())
			var reply string
			reply, err = reader.ReadString('\n')
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error reading input: %w", err)
			}

			// If the user doesn't want to install it, then we we'll proceed
			// with launching the daemon normally
			reply = strings.TrimSpace(reply)
			if reply != "y" {
				installRefused = true
				return nil
			}
			err = os.Remove(telProLocation)
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error removing Telepresence Pro: %w", err)
			}
			// Since we've already asked the user for permission to upgrade,
			// we can run these functions without asking permission again.
			err = installTelepresencePro(ctx, telProLocation)
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error installing updated Telepresence Pro: %w",
					err)
			}

			// The users configuration is most likely correct if they are upgrading,
			// but we update it just to be extra sure.
			err = updateConfig(ctx, telProLocation)
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error updating config: %w",
					err)
			}
		}
	}
	return nil
}

// installTelepresencePro installs the binary. Users should be asked for
// permission before using this function
func installTelepresencePro(ctx context.Context, telProLocation string) error {
	// We install the correct version of telepresence-pro based on
	// the OSS version that is associated with this client since
	// daemon versions need to match
	clientVersion := strings.Trim(client.Version(), "v")
	systemAHost := client.GetConfig(ctx).Cloud.SystemaHost
	installString := fmt.Sprintf("https://%s/download/tel-pro/%s/%s/%s/latest/%s",
		systemAHost, runtime.GOOS, runtime.GOARCH, clientVersion, filepath.Base(telProLocation))

	resp, err := http.Get(installString)
	if err == nil {
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			err = errors.New(resp.Status)
		}
	}
	if err != nil {
		return errcat.NoDaemonLogs.Newf("unable to download Telepresence Pro: %w", err)
	}

	out, err := os.Create(telProLocation)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("unable to create file %q for Telepresence Pro: %w", telProLocation, err)
	}
	defer out.Close()

	if _, err = io.Copy(out, resp.Body); err != nil {
		return errcat.NoDaemonLogs.Newf("unable to copy Telepresence Pro to %q: %w", telProLocation, err)
	}

	if err = os.Chmod(telProLocation, 0755); err != nil {
		return errcat.NoDaemonLogs.Newf("unable to set permissions of %q to 755: %w", telProLocation, err)
	}
	return nil
}

// updateConfig updates the userDaemonBinary in the config to point to
// telProLocation. Users should be asked for permission before this is done.
func updateConfig(ctx context.Context, telProLocation string) error {
	cfg := client.GetConfig(ctx)
	cfg.Daemons.UserDaemonBinary = telProLocation

	b, err := yaml.Marshal(cfg)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("error marshaling updating config: %w", err)
	}
	cfgFile := client.GetConfigFile(ctx)
	_, err = os.OpenFile(cfgFile, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("error opening config file: %w", err)
	}
	err = os.WriteFile(cfgFile, b, 0644)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("error writing config file: %w", err)
	}
	return nil
}
<commit_msg>Fix linting<commit_after>package cliutil

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	grpcCodes "google.golang.org/grpc/codes"
	grpcStatus "google.golang.org/grpc/status"
	empty "google.golang.org/protobuf/types/known/emptypb"
	"gopkg.in/yaml.v2"

	"github.com/datawire/dlib/dexec"
	"github.com/telepresenceio/telepresence/rpc/v2/connector"
	"github.com/telepresenceio/telepresence/v2/pkg/client"
	"github.com/telepresenceio/telepresence/v2/pkg/client/errcat"
	"github.com/telepresenceio/telepresence/v2/pkg/client/scout"
	"github.com/telepresenceio/telepresence/v2/pkg/client/userd/auth/authdata"
	"github.com/telepresenceio/telepresence/v2/pkg/filelocation"
)

// EnsureLoggedIn ensures that the user is logged in to Ambassador Cloud.  An error is returned if
// login fails.  The result code will indicate if this is a new login or if it resued an existing
// login.  If the `apikey` argument is empty an interactive login is performed; if it is non-empty
// the key is used instead of performing an interactive login.
func EnsureLoggedIn(ctx context.Context, apikey string) (connector.LoginResult_Code, error) {
	err := GetTelepresencePro(ctx)
	if err != nil {
		return connector.LoginResult_UNSPECIFIED, err
	}
	var code connector.LoginResult_Code
	err = WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		code, err = ClientEnsureLoggedIn(ctx, apikey, connectorClient)
		return err
	})
	return code, err
}

// ClientEnsureLoggedIn is like EnsureLoggedIn but uses an already acquired ConnectorClient.
func ClientEnsureLoggedIn(ctx context.Context, apikey string, connectorClient connector.ConnectorClient) (connector.LoginResult_Code, error) {
	resp, err := connectorClient.Login(ctx, &connector.LoginRequest{
		ApiKey: apikey,
	})
	if err != nil {
		if grpcStatus.Code(err) == grpcCodes.PermissionDenied {
			err = errcat.User.New(grpcStatus.Convert(err).Message())
		}
		return connector.LoginResult_UNSPECIFIED, err
	}
	return resp.GetCode(), nil
}

// Logout logs out of Ambassador Cloud.  Returns an error if not logged in.
func Logout(ctx context.Context) error {
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		_, err := connectorClient.Logout(ctx, &empty.Empty{})
		return err
	})
	if grpcStatus.Code(err) == grpcCodes.NotFound {
		err = errcat.User.New(grpcStatus.Convert(err).Message())
	}
	if err != nil {
		return err
	}
	return nil
}

// EnsureLoggedOut ensures that the user is logged out of Ambassador Cloud.  Returns nil if not
// logged in.
func EnsureLoggedOut(ctx context.Context) error {
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		_, err := connectorClient.Logout(ctx, &empty.Empty{})
		return err
	})
	if grpcStatus.Code(err) == grpcCodes.NotFound {
		err = nil
	}
	if err != nil {
		return err
	}
	return nil
}

// HasLoggedIn returns true if either the user has an active login session or an expired login
// session, and returns false if either the user has never logged in or has explicitly logged out.
func HasLoggedIn(ctx context.Context) bool {
	_, err := authdata.LoadUserInfoFromUserCache(ctx)
	return err == nil
}

func GetCloudUserInfo(ctx context.Context, autoLogin bool, refresh bool) (*connector.UserInfo, error) {
	var userInfo *connector.UserInfo
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		userInfo, err = connectorClient.GetCloudUserInfo(ctx, &connector.UserInfoRequest{
			AutoLogin: autoLogin,
			Refresh:   refresh,
		})
		return err
	})
	if err != nil {
		return nil, err
	}
	return userInfo, nil
}

func GetCloudAPIKey(ctx context.Context, description string, autoLogin bool) (string, error) {
	var keyData *connector.KeyData
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		keyData, err = connectorClient.GetCloudAPIKey(ctx, &connector.KeyRequest{
			AutoLogin:   autoLogin,
			Description: description,
		})
		return err
	})
	if err != nil {
		return "", err
	}
	return keyData.GetApiKey(), nil
}

// GetCloudLicense communicates with system a to get the jwt version of the
// license, puts it in a kubernetes secret, and then writes that secret to the
// output file for the user to apply to their cluster
func GetCloudLicense(ctx context.Context, outputFile, id string) (string, string, error) {
	var licenseData *connector.LicenseData
	err := WithConnector(ctx, func(ctx context.Context, connectorClient connector.ConnectorClient) error {
		var err error
		licenseData, err = connectorClient.GetCloudLicense(ctx, &connector.LicenseRequest{
			Id: id,
		})
		return err
	})
	if err != nil {
		return "", "", err
	}
	return licenseData.GetLicense(), licenseData.GetHostDomain(), nil
}

// GetTelepresencePro prompts the user to optionally install Telepresence Pro
// if it isn't installed. If the user installs it, it also asks the user to
// automatically update their configuration to use the new binary.
func GetTelepresencePro(ctx context.Context) error {
	dir, err := filelocation.AppUserConfigDir(ctx)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("unable to get path to config files: %w", err)
	}

	sc := scout.NewReporter(ctx, "cli")
	sc.Start(ctx)
	defer sc.Close()
	installRefused := false
	defer func() {
		switch {
		case err != nil:
			sc.Report(ctx, "pro_connector_upgrade_fail", scout.Entry{Key: "error", Value: err.Error()})
		case installRefused:
			sc.Report(ctx, "pro_connector_upgrade_refusal")
		default:
			sc.Report(ctx, "pro_connector_upgrade_success")
		}
	}()

	// If telepresence-pro doesn't exist, then we should ask the user
	// if they want to install it
	telProLocation := filepath.Join(dir, "telepresence-pro")
	if runtime.GOOS == "windows" {
		telProLocation += ".exe"
	}
	if _, err = os.Stat(telProLocation); os.IsNotExist(err) {
		sc.SetMetadatum(ctx, "first_install", true)
		reader := bufio.NewReader(os.Stdin)
		fmt.Printf("Telepresence Pro is recommended when using login features, can Telepresence install it? (y/n)")
		var reply string
		reply, err = reader.ReadString('\n')
		if err != nil {
			return errcat.User.Newf("error reading input: %w", err)
		}

		// If the user doesn't want to install it, then we we'll proceed
		// with launching the daemon normally
		reply = strings.TrimSpace(reply)
		if reply != "y" {
			installRefused = true
			return nil
		}

		err = installTelepresencePro(ctx, telProLocation)
		if err != nil {
			return err
		}

		// Ask the user if they want to automatically update their config
		// with the telepresence-pro binary.
		// TODO: This will remove any comments that exist in the config file
		// which it's yaml so that's _fine_ but it would be nice if we didn't
		// do that.
		fmt.Printf("Update your Telepresence config to use Telepresence Pro? (y/n)")
		reply, err = reader.ReadString('\n')
		if err != nil {
			return errcat.User.Newf("error reading input: %w", err)
		}
		reply = strings.TrimSpace(reply)
		if reply != "y" {
			return nil
		}
		err = updateConfig(ctx, telProLocation)
		if err != nil {
			return err
		}
	} else {
		// If the binary is present, we check its version to ensure it's compatible
		// with the CLI
		sc.SetMetadatum(ctx, "first_install", false)
		proCmd := dexec.CommandContext(ctx, telProLocation, "pro-version")
		proCmd.DisableLogging = true

		var output []byte
		output, err = proCmd.CombinedOutput()
		if err != nil {
			return errcat.NoDaemonLogs.Newf("Unable to get telepresence pro version")
		}

		if !strings.Contains(string(output), client.Version()) {
			reader := bufio.NewReader(os.Stdin)
			fmt.Printf("Telepresence Pro needs to be upgraded to work with CLI version %s, allow Telepresence to upgrade it? (y/n)",
				client.Version())
			var reply string
			reply, err = reader.ReadString('\n')
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error reading input: %w", err)
			}

			// If the user doesn't want to install it, then we we'll proceed
			// with launching the daemon normally
			reply = strings.TrimSpace(reply)
			if reply != "y" {
				installRefused = true
				return nil
			}
			err = os.Remove(telProLocation)
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error removing Telepresence Pro: %w", err)
			}
			// Since we've already asked the user for permission to upgrade,
			// we can run these functions without asking permission again.
			err = installTelepresencePro(ctx, telProLocation)
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error installing updated Telepresence Pro: %w",
					err)
			}

			// The users configuration is most likely correct if they are upgrading,
			// but we update it just to be extra sure.
			err = updateConfig(ctx, telProLocation)
			if err != nil {
				return errcat.NoDaemonLogs.Newf("error updating config: %w",
					err)
			}
		}
	}
	return nil
}

// installTelepresencePro installs the binary. Users should be asked for
// permission before using this function
func installTelepresencePro(ctx context.Context, telProLocation string) error {
	// We install the correct version of telepresence-pro based on
	// the OSS version that is associated with this client since
	// daemon versions need to match
	clientVersion := strings.Trim(client.Version(), "v")
	systemAHost := client.GetConfig(ctx).Cloud.SystemaHost
	installString := fmt.Sprintf("https://%s/download/tel-pro/%s/%s/%s/latest/%s",
		systemAHost, runtime.GOOS, runtime.GOARCH, clientVersion, filepath.Base(telProLocation))

	resp, err := http.Get(installString)
	if err == nil {
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			err = errors.New(resp.Status)
		}
	}
	if err != nil {
		return errcat.NoDaemonLogs.Newf("unable to download Telepresence Pro: %w", err)
	}

	out, err := os.Create(telProLocation)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("unable to create file %q for Telepresence Pro: %w", telProLocation, err)
	}
	defer out.Close()

	if _, err = io.Copy(out, resp.Body); err != nil {
		return errcat.NoDaemonLogs.Newf("unable to copy Telepresence Pro to %q: %w", telProLocation, err)
	}

	if err = os.Chmod(telProLocation, 0755); err != nil {
		return errcat.NoDaemonLogs.Newf("unable to set permissions of %q to 755: %w", telProLocation, err)
	}
	return nil
}

// updateConfig updates the userDaemonBinary in the config to point to
// telProLocation. Users should be asked for permission before this is done.
func updateConfig(ctx context.Context, telProLocation string) error {
	cfg := client.GetConfig(ctx)
	cfg.Daemons.UserDaemonBinary = telProLocation

	b, err := yaml.Marshal(cfg)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("error marshaling updating config: %w", err)
	}
	cfgFile := client.GetConfigFile(ctx)
	_, err = os.OpenFile(cfgFile, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("error opening config file: %w", err)
	}
	err = os.WriteFile(cfgFile, b, 0644)
	if err != nil {
		return errcat.NoDaemonLogs.Newf("error writing config file: %w", err)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2019, EnMasse authors.
 * License: Apache License 2.0 (see the file LICENSE or http://apache.org/licenses/LICENSE-2.0.html).
 */

package agent

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/enmasseproject/enmasse/pkg/amqpcommand"

	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"pack.ag/amqp"
)

// The agent delegate is the facade to the address space's agent component.  It allows events broadcast by the
// agent to be received.  It also allows commands to be sent to the agent.

type Delegate interface {
	Collect(handler EventHandler) error
	CommandDelegate(bearerToken string, impersonateUser string) (CommandDelegate, error)
	Shutdown()
}

type CommandDelegate interface {
	PurgeAddress(address string) error
	CloseConnection(address v1.ObjectMeta) error
	Shutdown()
}

const amqpOverrideSaslFrameSize = 4096
const agentDataAddress = "agent_data"
const agentCommandAddress = "agent_command"
const agentCommandResponseAddress = "agent_command_response"

type EventHandler = func(event AgentEvent) error

type DelegateCreator = func(bearerToken, host string, port int32) Delegate

type commandDelegatePair struct {
	delegate CommandDelegate
	lastUsed time.Time
}

type amqpAgentDelegate struct {
	bearerToken                 string
	host                        string
	port                        int32
	tlsConfig                   *tls.Config
	addressSpaceNamespace       string
	addressSpace                string
	infraUuid                   string
	handler                     EventHandler
	stopchan                    chan struct{}
	stoppedchan                 chan struct{}
	commandDelegates            map[string]commandDelegatePair
	commandDelegatesMux         sync.Mutex
	commandDelegateExpiryPeriod time.Duration
	connectTimeout              time.Duration
	maxFrameSize                uint32
}

func NewAmqpAgentDelegate(bearerToken, host string, port int32, tlsConfig *tls.Config, addressSpaceNamespace, addressSpace, infraUuid string, expirePeriod, connectTimeout time.Duration, maxFrameSize uint32) Delegate {
	return &amqpAgentDelegate{
		bearerToken:                 bearerToken,
		host:                        host,
		port:                        port,
		tlsConfig:                   tlsConfig,
		addressSpaceNamespace:       addressSpaceNamespace,
		addressSpace:                addressSpace,
		infraUuid:                   infraUuid,
		stopchan:                    make(chan struct{}),
		stoppedchan:                 make(chan struct{}),
		commandDelegates:            make(map[string]commandDelegatePair),
		commandDelegateExpiryPeriod: expirePeriod,
		connectTimeout:              connectTimeout,
		maxFrameSize:                maxFrameSize,
	}
}

func (aad *amqpAgentDelegate) Collect(handler EventHandler) error {
	aad.handler = handler

	go func() {
		defer close(aad.stoppedchan)
		defer log.Printf("Agent Collector %s - stopped", aad.infraUuid)

		log.Printf("Agent Collector %s - starting", aad.infraUuid)
		for {
			select {
			case <-aad.stopchan:
				return
			default:
				err := aad.doCollect()

				if err == nil {
					return
				} else {
					backoff := computeBackoff(err)
					log.Printf("Agent Collector %s - restarting - backoff %s(s), %v", aad.infraUuid, backoff, err)
					if backoff > 0 {
						time.Sleep(backoff)
					}
				}
			}
		}
	}()

	return nil
}

func computeBackoff(err error) time.Duration {
	backoff := 5 * time.Second
	if _, ok := err.(net.Error); ok {
		backoff = 30 * time.Second
	}
	return backoff
}

func (aad *amqpAgentDelegate) Shutdown() {
	log.Printf("Agent Collector %s - Shutting down", aad.infraUuid)
	close(aad.stopchan)
	for key, commandDelegate := range aad.commandDelegates {
		commandDelegate.delegate.Shutdown()
		delete(aad.commandDelegates, key)
	}
	<-aad.stoppedchan
}

func (aad *amqpAgentDelegate) doCollect() error {

	addr := buildAmqpAddress(aad.host, aad.port)
	log.Printf("Agent Collector %s - connecting %s", aad.infraUuid, addr)

	client, err := amqp.Dial(addr,
		amqp.ConnTLSConfig(aad.tlsConfig),
		amqp.ConnSASLXOAUTH2("", aad.bearerToken, amqpOverrideSaslFrameSize),
		amqp.ConnServerHostname(aad.host),
		amqp.ConnProperty("product", "console-server"),
		amqp.ConnConnectTimeout(aad.connectTimeout),
		amqp.ConnMaxFrameSize(aad.maxFrameSize),
	)
	if err != nil {
		return err
	}
	defer func() {
		_ = client.Close()
	}()
	log.Printf("Agent Collector %s - connected %s", aad.infraUuid, addr)

	// Open a session
	session, err := client.NewSession()
	if err != nil {
		return err
	}

	ctx := context.Background()

	// Create a receiver
	receiver, err := session.NewReceiver(
		amqp.LinkSourceAddress(agentDataAddress),
	)
	if err != nil {
		return err
	}

	defer func() {
		ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
		_ = receiver.Close(ctx)
		cancel()
	}()

	log.Printf("Agent Collector %s - commencing collecting", aad.infraUuid)
	restart := AgentEvent{
		AddressSpaceNamespace: aad.addressSpaceNamespace,
		AddressSpace:          aad.addressSpace,
		InfraUuid:             aad.infraUuid,
		Type:                  AgentEventTypeRestart,
	}
	err = aad.handler(restart)
	if err != nil {
		return err
	}

	next := time.Now().Add(aad.commandDelegateExpiryPeriod)
	for {

		select {
		case <-aad.stopchan:
			return nil
		default:
			now := time.Now()
			if next.Before(now) {
				go aad.expireCommandDelegates(now)
				next = now.Add(aad.commandDelegateExpiryPeriod)
			}

			// Receive next message
			msg, err := receiveWithTimeout(ctx, receiver, 5*time.Second)
			if err != nil {
				return err
			} else if msg != nil {
				switch msg.Properties.Subject {

				case "connection":
					body := msg.Value.(map[string]interface{})
					agentcon, err := FromAgentConnectionBody(body)
					if err != nil {
						log.Printf("ignoring failure unmarsall connection message %s for infraUuid %s, %v", body, aad.infraUuid, err)
						break
					} else {
						evt := AgentEvent{
							AddressSpaceNamespace: aad.addressSpaceNamespace,
							AddressSpace:          aad.addressSpace,
							InfraUuid:             aad.infraUuid,
							Type:                  AgentEventInsertOrUpdateType,
							Object:                agentcon,
						}
						err = aad.handler(evt)
						if err != nil {
							return err
						}
					}
				case "connection_deleted":
					body := msg.Value.(map[string]interface{})
					agentcon, err := FromAgentConnectionBody(body)
					if err != nil {
						log.Printf("ignoring failure unmarsall connection message %s for infraUuid %s %v", body, aad.infraUuid, err)
						break
					} else {
						evt := AgentEvent{
							AddressSpaceNamespace: aad.addressSpaceNamespace,
							AddressSpace:          aad.addressSpace,
							InfraUuid:             aad.infraUuid,
							Type:                  AgentEventTypeDelete,
							Object:                agentcon,
						}
						err = aad.handler(evt)
						if err != nil {
							return err
						}
					}
				case "address":
					body := msg.Value.(map[string]interface{})
					agentaddr, err := FromAgentAddressBody(body)
					if err != nil {
						log.Printf("ignoring failure unmarsall address message %s for infraUuid %s, %v", body, aad.infraUuid, err)
						break
					} else {
						evt := AgentEvent{
							AddressSpaceNamespace: aad.addressSpaceNamespace,
							AddressSpace:          aad.addressSpace,
							InfraUuid:             aad.infraUuid,
							Type:                  AgentEventInsertOrUpdateType,
							Object:                agentaddr,
						}
						err = aad.handler(evt)
						if err != nil {
							return err
						}

					}

				case "address_deleted":
				default:
					// Ignore messages with other subjects
				}

				// Accept message
				err = msg.Accept()
				if err != nil {
					return err
				}
			}
		}
	}
}

func (aad *amqpAgentDelegate) CommandDelegate(bearerToken string, impersonateUser string) (CommandDelegate, error) {
	aad.commandDelegatesMux.Lock()
	defer aad.commandDelegatesMux.Unlock()

	key := getShaSum(bearerToken)

	now := time.Now()
	if pair, present := aad.commandDelegates[key]; !present {
		delegate := aad.newAgentDelegate(bearerToken, impersonateUser)
		aad.commandDelegates[key] = commandDelegatePair{
			delegate: delegate,
			lastUsed: now,
		}
		return delegate, nil
	} else {
		pair.lastUsed = now
		return pair.delegate, nil
	}
}

func (aad *amqpAgentDelegate) expireCommandDelegates(now time.Time) {
	findExpiredCommandDelegates := func() []CommandDelegate {
		aad.commandDelegatesMux.Lock()
		defer aad.commandDelegatesMux.Unlock()

		expired := make([]CommandDelegate, 0)

		horizon := now.Add(aad.commandDelegateExpiryPeriod)
		for key, commandDelegate := range aad.commandDelegates {
			if commandDelegate.lastUsed.Before(horizon) {
				expired = append(expired, commandDelegate.delegate)
				delete(aad.commandDelegates, key)
			}
		}

		return expired
	}

	expired := findExpiredCommandDelegates()
	for _, d := range expired {
		d.Shutdown()
	}
	if len(expired) > 0 {
		log.Printf("Shutdown %d expired command delegate(s)", len(expired))
	}
}

func receiveWithTimeout(ctx context.Context, receiver *amqp.Receiver, timeout time.Duration) (*amqp.Message, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	msg, err := receiver.Receive(ctx)
	if err != nil && err != context.DeadlineExceeded {
		return nil, err
	}
	return msg, nil
}

func getShaSum(token string) string {
	accessTokenSha := sha256.Sum256([]byte(token))
	return base64.StdEncoding.EncodeToString(accessTokenSha[:])
}

type agentCommandRequest struct {
	commandMessage *amqp.Message
	response       chan error
}

type amqpAgentCommandDelegate struct {
	aac           *amqpAgentDelegate
	commandClient amqpcommand.Client
	lastUsed      time.Time
}

func (aad *amqpAgentDelegate) newAgentDelegate(token string, impersonateUser string) CommandDelegate {
	commandClient := amqpcommand.NewCommandClient(buildAmqpAddress(aad.host, aad.port),
		agentCommandAddress,
		agentCommandResponseAddress,
		amqp.ConnTLSConfig(aad.tlsConfig),
		amqp.ConnSASLXOAUTH2(impersonateUser, token, amqpOverrideSaslFrameSize),
		amqp.ConnServerHostname(aad.host),
		amqp.ConnProperty("product", "command-delegate; console-server"),
		amqp.ConnConnectTimeout(aad.connectTimeout),
		amqp.ConnMaxFrameSize(aad.maxFrameSize))

	a := &amqpAgentCommandDelegate{
		aac:           aad,
		commandClient: commandClient,
		lastUsed:      time.Now(),
	}

	a.commandClient.Start()
	return a
}

func (ad *amqpAgentCommandDelegate) PurgeAddress(address string) error {
	request := &amqp.Message{
		Properties: &amqp.MessageProperties{
			Subject: "purge_address",
		},
		Value: map[interface{}]interface{}{
			"address": address,
		},
	}

	response, err := ad.commandClient.Request(request)
	if err != nil {
		return fmt.Errorf("failed to purge address %s : %s", address, err)
	}

	if outcome, present := response.ApplicationProperties["outcome"]; present {
		if oc, ok := outcome.(bool); ok && oc {
			return nil
		} else {
			if e, present := response.ApplicationProperties["error"]; present && e != nil {
				return fmt.Errorf("failed to purge address %s : %s", address, e)
			}
		}
	}
	return fmt.Errorf("failed to purge address %s : command %+v failed for unknown reason", address, request)
}

func (ad *amqpAgentCommandDelegate) CloseConnection(connection v1.ObjectMeta) error {
	request := &amqp.Message{
		Properties: &amqp.MessageProperties{
			Subject: "close_connection",
		},
		Value: map[interface{}]interface{}{
			"connectionUid": string(connection.UID),
		},
	}

	response, err := ad.commandClient.Request(request)
	if err != nil {
		return fmt.Errorf("failed to close connection %s : %s", connection.UID, err)
	}

	if outcome, present := response.ApplicationProperties["outcome"]; present {
		if oc, ok := outcome.(bool); ok && oc {
			return nil
		} else {
			if e, present := response.ApplicationProperties["error"]; present && e != nil {
				return fmt.Errorf("failed to close connection %s : %s", connection.UID, e)
			}
		}
	}
	return fmt.Errorf("failed to close connection %s : command %+v failed for unknown reason", connection.UID, request)
}

func (ad *amqpAgentCommandDelegate) LastUsed() time.Time {
	return ad.lastUsed
}

func (ad *amqpAgentCommandDelegate) Shutdown() {
	ad.commandClient.Stop()
}

func buildAmqpAddress(host string, port int32) string {
	addr := fmt.Sprintf("amqps://%s:%d", host, port)
	return addr
}
<commit_msg>Fix #5238: Console receiver grants too little credit to agent_data receiving link (#5239)<commit_after>/*
 * Copyright 2019, EnMasse authors.
 * License: Apache License 2.0 (see the file LICENSE or http://apache.org/licenses/LICENSE-2.0.html).
 */

package agent

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/enmasseproject/enmasse/pkg/amqpcommand"

	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"pack.ag/amqp"
)

// The agent delegate is the facade to the address space's agent component.  It allows events broadcast by the
// agent to be received.  It also allows commands to be sent to the agent.

type Delegate interface {
	Collect(handler EventHandler) error
	CommandDelegate(bearerToken string, impersonateUser string) (CommandDelegate, error)
	Shutdown()
}

type CommandDelegate interface {
	PurgeAddress(address string) error
	CloseConnection(address v1.ObjectMeta) error
	Shutdown()
}

const amqpOverrideSaslFrameSize = 4096
const agentDataAddress = "agent_data"
const agentCommandAddress = "agent_command"
const agentCommandResponseAddress = "agent_command_response"

type EventHandler = func(event AgentEvent) error

type DelegateCreator = func(bearerToken, host string, port int32) Delegate

type commandDelegatePair struct {
	delegate CommandDelegate
	lastUsed time.Time
}

type amqpAgentDelegate struct {
	bearerToken                 string
	host                        string
	port                        int32
	tlsConfig                   *tls.Config
	addressSpaceNamespace       string
	addressSpace                string
	infraUuid                   string
	handler                     EventHandler
	stopchan                    chan struct{}
	stoppedchan                 chan struct{}
	commandDelegates            map[string]commandDelegatePair
	commandDelegatesMux         sync.Mutex
	commandDelegateExpiryPeriod time.Duration
	connectTimeout              time.Duration
	maxFrameSize                uint32
}

func NewAmqpAgentDelegate(bearerToken, host string, port int32, tlsConfig *tls.Config, addressSpaceNamespace, addressSpace, infraUuid string, expirePeriod, connectTimeout time.Duration, maxFrameSize uint32) Delegate {
	return &amqpAgentDelegate{
		bearerToken:                 bearerToken,
		host:                        host,
		port:                        port,
		tlsConfig:                   tlsConfig,
		addressSpaceNamespace:       addressSpaceNamespace,
		addressSpace:                addressSpace,
		infraUuid:                   infraUuid,
		stopchan:                    make(chan struct{}),
		stoppedchan:                 make(chan struct{}),
		commandDelegates:            make(map[string]commandDelegatePair),
		commandDelegateExpiryPeriod: expirePeriod,
		connectTimeout:              connectTimeout,
		maxFrameSize:                maxFrameSize,
	}
}

func (aad *amqpAgentDelegate) Collect(handler EventHandler) error {
	aad.handler = handler

	go func() {
		defer close(aad.stoppedchan)
		defer log.Printf("Agent Collector %s - stopped", aad.infraUuid)

		log.Printf("Agent Collector %s - starting", aad.infraUuid)
		for {
			select {
			case <-aad.stopchan:
				return
			default:
				err := aad.doCollect()

				if err == nil {
					return
				} else {
					backoff := computeBackoff(err)
					log.Printf("Agent Collector %s - restarting - backoff %s(s), %v", aad.infraUuid, backoff, err)
					if backoff > 0 {
						time.Sleep(backoff)
					}
				}
			}
		}
	}()

	return nil
}

func computeBackoff(err error) time.Duration {
	backoff := 5 * time.Second
	if _, ok := err.(net.Error); ok {
		backoff = 30 * time.Second
	}
	return backoff
}

func (aad *amqpAgentDelegate) Shutdown() {
	log.Printf("Agent Collector %s - Shutting down", aad.infraUuid)
	close(aad.stopchan)
	for key, commandDelegate := range aad.commandDelegates {
		commandDelegate.delegate.Shutdown()
		delete(aad.commandDelegates, key)
	}
	<-aad.stoppedchan
}

func (aad *amqpAgentDelegate) doCollect() error {

	addr := buildAmqpAddress(aad.host, aad.port)
	log.Printf("Agent Collector %s - connecting %s", aad.infraUuid, addr)

	client, err := amqp.Dial(addr,
		amqp.ConnTLSConfig(aad.tlsConfig),
		amqp.ConnSASLXOAUTH2("", aad.bearerToken, amqpOverrideSaslFrameSize),
		amqp.ConnServerHostname(aad.host),
		amqp.ConnProperty("product", "console-server"),
		amqp.ConnConnectTimeout(aad.connectTimeout),
		amqp.ConnMaxFrameSize(aad.maxFrameSize),
	)
	if err != nil {
		return err
	}
	defer func() {
		_ = client.Close()
	}()
	log.Printf("Agent Collector %s - connected %s", aad.infraUuid, addr)

	// Open a session
	session, err := client.NewSession()
	if err != nil {
		return err
	}

	ctx := context.Background()

	// Create a receiver
	receiver, err := session.NewReceiver(
		amqp.LinkSourceAddress(agentDataAddress),
		amqp.LinkCredit(500),
	)
	if err != nil {
		return err
	}

	defer func() {
		ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
		_ = receiver.Close(ctx)
		cancel()
	}()

	log.Printf("Agent Collector %s - commencing collecting", aad.infraUuid)
	restart := AgentEvent{
		AddressSpaceNamespace: aad.addressSpaceNamespace,
		AddressSpace:          aad.addressSpace,
		InfraUuid:             aad.infraUuid,
		Type:                  AgentEventTypeRestart,
	}
	err = aad.handler(restart)
	if err != nil {
		return err
	}

	next := time.Now().Add(aad.commandDelegateExpiryPeriod)
	for {

		select {
		case <-aad.stopchan:
			return nil
		default:
			now := time.Now()
			if next.Before(now) {
				go aad.expireCommandDelegates(now)
				next = now.Add(aad.commandDelegateExpiryPeriod)
			}

			// Receive next message
			msg, err := receiveWithTimeout(ctx, receiver, 5*time.Second)
			if err != nil {
				return err
			} else if msg != nil {
				switch msg.Properties.Subject {

				case "connection":
					body := msg.Value.(map[string]interface{})
					agentcon, err := FromAgentConnectionBody(body)
					if err != nil {
						log.Printf("ignoring failure unmarsall connection message %s for infraUuid %s, %v", body, aad.infraUuid, err)
						break
					} else {
						evt := AgentEvent{
							AddressSpaceNamespace: aad.addressSpaceNamespace,
							AddressSpace:          aad.addressSpace,
							InfraUuid:             aad.infraUuid,
							Type:                  AgentEventInsertOrUpdateType,
							Object:                agentcon,
						}
						err = aad.handler(evt)
						if err != nil {
							return err
						}
					}
				case "connection_deleted":
					body := msg.Value.(map[string]interface{})
					agentcon, err := FromAgentConnectionBody(body)
					if err != nil {
						log.Printf("ignoring failure unmarsall connection message %s for infraUuid %s %v", body, aad.infraUuid, err)
						break
					} else {
						evt := AgentEvent{
							AddressSpaceNamespace: aad.addressSpaceNamespace,
							AddressSpace:          aad.addressSpace,
							InfraUuid:             aad.infraUuid,
							Type:                  AgentEventTypeDelete,
							Object:                agentcon,
						}
						err = aad.handler(evt)
						if err != nil {
							return err
						}
					}
				case "address":
					body := msg.Value.(map[string]interface{})
					agentaddr, err := FromAgentAddressBody(body)
					if err != nil {
						log.Printf("ignoring failure unmarsall address message %s for infraUuid %s, %v", body, aad.infraUuid, err)
						break
					} else {
						evt := AgentEvent{
							AddressSpaceNamespace: aad.addressSpaceNamespace,
							AddressSpace:          aad.addressSpace,
							InfraUuid:             aad.infraUuid,
							Type:                  AgentEventInsertOrUpdateType,
							Object:                agentaddr,
						}
						err = aad.handler(evt)
						if err != nil {
							return err
						}

					}

				case "address_deleted":
				default:
					// Ignore messages with other subjects
				}

				// Accept message
				err = msg.Accept()
				if err != nil {
					return err
				}
			}
		}
	}
}

func (aad *amqpAgentDelegate) CommandDelegate(bearerToken string, impersonateUser string) (CommandDelegate, error) {
	aad.commandDelegatesMux.Lock()
	defer aad.commandDelegatesMux.Unlock()

	key := getShaSum(bearerToken)

	now := time.Now()
	if pair, present := aad.commandDelegates[key]; !present {
		delegate := aad.newAgentDelegate(bearerToken, impersonateUser)
		aad.commandDelegates[key] = commandDelegatePair{
			delegate: delegate,
			lastUsed: now,
		}
		return delegate, nil
	} else {
		pair.lastUsed = now
		return pair.delegate, nil
	}
}

func (aad *amqpAgentDelegate) expireCommandDelegates(now time.Time) {
	findExpiredCommandDelegates := func() []CommandDelegate {
		aad.commandDelegatesMux.Lock()
		defer aad.commandDelegatesMux.Unlock()

		expired := make([]CommandDelegate, 0)

		horizon := now.Add(aad.commandDelegateExpiryPeriod)
		for key, commandDelegate := range aad.commandDelegates {
			if commandDelegate.lastUsed.Before(horizon) {
				expired = append(expired, commandDelegate.delegate)
				delete(aad.commandDelegates, key)
			}
		}

		return expired
	}

	expired := findExpiredCommandDelegates()
	for _, d := range expired {
		d.Shutdown()
	}
	if len(expired) > 0 {
		log.Printf("Shutdown %d expired command delegate(s)", len(expired))
	}
}

func receiveWithTimeout(ctx context.Context, receiver *amqp.Receiver, timeout time.Duration) (*amqp.Message, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	msg, err := receiver.Receive(ctx)
	if err != nil && err != context.DeadlineExceeded {
		return nil, err
	}
	return msg, nil
}

func getShaSum(token string) string {
	accessTokenSha := sha256.Sum256([]byte(token))
	return base64.StdEncoding.EncodeToString(accessTokenSha[:])
}

type agentCommandRequest struct {
	commandMessage *amqp.Message
	response       chan error
}

type amqpAgentCommandDelegate struct {
	aac           *amqpAgentDelegate
	commandClient amqpcommand.Client
	lastUsed      time.Time
}

func (aad *amqpAgentDelegate) newAgentDelegate(token string, impersonateUser string) CommandDelegate {
	commandClient := amqpcommand.NewCommandClient(buildAmqpAddress(aad.host, aad.port),
		agentCommandAddress,
		agentCommandResponseAddress,
		amqp.ConnTLSConfig(aad.tlsConfig),
		amqp.ConnSASLXOAUTH2(impersonateUser, token, amqpOverrideSaslFrameSize),
		amqp.ConnServerHostname(aad.host),
		amqp.ConnProperty("product", "command-delegate; console-server"),
		amqp.ConnConnectTimeout(aad.connectTimeout),
		amqp.ConnMaxFrameSize(aad.maxFrameSize))

	a := &amqpAgentCommandDelegate{
		aac:           aad,
		commandClient: commandClient,
		lastUsed:      time.Now(),
	}

	a.commandClient.Start()
	return a
}

func (ad *amqpAgentCommandDelegate) PurgeAddress(address string) error {
	request := &amqp.Message{
		Properties: &amqp.MessageProperties{
			Subject: "purge_address",
		},
		Value: map[interface{}]interface{}{
			"address": address,
		},
	}

	response, err := ad.commandClient.Request(request)
	if err != nil {
		return fmt.Errorf("failed to purge address %s : %s", address, err)
	}

	if outcome, present := response.ApplicationProperties["outcome"]; present {
		if oc, ok := outcome.(bool); ok && oc {
			return nil
		} else {
			if e, present := response.ApplicationProperties["error"]; present && e != nil {
				return fmt.Errorf("failed to purge address %s : %s", address, e)
			}
		}
	}
	return fmt.Errorf("failed to purge address %s : command %+v failed for unknown reason", address, request)
}

func (ad *amqpAgentCommandDelegate) CloseConnection(connection v1.ObjectMeta) error {
	request := &amqp.Message{
		Properties: &amqp.MessageProperties{
			Subject: "close_connection",
		},
		Value: map[interface{}]interface{}{
			"connectionUid": string(connection.UID),
		},
	}

	response, err := ad.commandClient.Request(request)
	if err != nil {
		return fmt.Errorf("failed to close connection %s : %s", connection.UID, err)
	}

	if outcome, present := response.ApplicationProperties["outcome"]; present {
		if oc, ok := outcome.(bool); ok && oc {
			return nil
		} else {
			if e, present := response.ApplicationProperties["error"]; present && e != nil {
				return fmt.Errorf("failed to close connection %s : %s", connection.UID, e)
			}
		}
	}
	return fmt.Errorf("failed to close connection %s : command %+v failed for unknown reason", connection.UID, request)
}

func (ad *amqpAgentCommandDelegate) LastUsed() time.Time {
	return ad.lastUsed
}

func (ad *amqpAgentCommandDelegate) Shutdown() {
	ad.commandClient.Stop()
}

func buildAmqpAddress(host string, port int32) string {
	addr := fmt.Sprintf("amqps://%s:%d", host, port)
	return addr
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 node

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

	"k8s.io/kubernetes/pkg/api"
	apierrors "k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/fake"
	unversionedcore "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/core/unversioned"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util/clock"
	utilnode "k8s.io/kubernetes/pkg/util/node"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/watch"
)

// FakeNodeHandler is a fake implementation of NodesInterface and NodeInterface. It
// allows test cases to have fine-grained control over mock behaviors. We also need
// PodsInterface and PodInterface to test list & delet pods, which is implemented in
// the embedded client.Fake field.
type FakeNodeHandler struct {
	*fake.Clientset

	// Input: Hooks determine if request is valid or not
	CreateHook func(*FakeNodeHandler, *api.Node) bool
	Existing   []*api.Node

	// Output
	CreatedNodes        []*api.Node
	DeletedNodes        []*api.Node
	UpdatedNodes        []*api.Node
	UpdatedNodeStatuses []*api.Node
	RequestCount        int

	// Synchronization
	lock           sync.Mutex
	deleteWaitChan chan struct{}
}

type FakeLegacyHandler struct {
	unversionedcore.CoreInterface
	n *FakeNodeHandler
}

func (c *FakeNodeHandler) getUpdatedNodesCopy() []*api.Node {
	c.lock.Lock()
	defer c.lock.Unlock()
	updatedNodesCopy := make([]*api.Node, len(c.UpdatedNodes), len(c.UpdatedNodes))
	for i, ptr := range c.UpdatedNodes {
		updatedNodesCopy[i] = ptr
	}
	return updatedNodesCopy
}

func (c *FakeNodeHandler) Core() unversionedcore.CoreInterface {
	return &FakeLegacyHandler{c.Clientset.Core(), c}
}

func (m *FakeLegacyHandler) Nodes() unversionedcore.NodeInterface {
	return m.n
}

func (m *FakeNodeHandler) Create(node *api.Node) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	for _, n := range m.Existing {
		if n.Name == node.Name {
			return nil, apierrors.NewAlreadyExists(api.Resource("nodes"), node.Name)
		}
	}
	if m.CreateHook == nil || m.CreateHook(m, node) {
		nodeCopy := *node
		m.CreatedNodes = append(m.CreatedNodes, &nodeCopy)
		return node, nil
	} else {
		return nil, errors.New("Create error.")
	}
}

func (m *FakeNodeHandler) Get(name string) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	for i := range m.Existing {
		if m.Existing[i].Name == name {
			nodeCopy := *m.Existing[i]
			return &nodeCopy, nil
		}
	}
	return nil, nil
}

func (m *FakeNodeHandler) List(opts api.ListOptions) (*api.NodeList, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	var nodes []*api.Node
	for i := 0; i < len(m.UpdatedNodes); i++ {
		if !contains(m.UpdatedNodes[i], m.DeletedNodes) {
			nodes = append(nodes, m.UpdatedNodes[i])
		}
	}
	for i := 0; i < len(m.Existing); i++ {
		if !contains(m.Existing[i], m.DeletedNodes) && !contains(m.Existing[i], nodes) {
			nodes = append(nodes, m.Existing[i])
		}
	}
	for i := 0; i < len(m.CreatedNodes); i++ {
		if !contains(m.Existing[i], m.DeletedNodes) && !contains(m.CreatedNodes[i], nodes) {
			nodes = append(nodes, m.CreatedNodes[i])
		}
	}
	nodeList := &api.NodeList{}
	for _, node := range nodes {
		nodeList.Items = append(nodeList.Items, *node)
	}
	return nodeList, nil
}

func (m *FakeNodeHandler) Delete(id string, opt *api.DeleteOptions) error {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		if m.deleteWaitChan != nil {
			m.deleteWaitChan <- struct{}{}
		}
		m.lock.Unlock()
	}()
	m.DeletedNodes = append(m.DeletedNodes, newNode(id))
	return nil
}

func (m *FakeNodeHandler) DeleteCollection(opt *api.DeleteOptions, listOpts api.ListOptions) error {
	return nil
}

func (m *FakeNodeHandler) Update(node *api.Node) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	nodeCopy := *node
	m.UpdatedNodes = append(m.UpdatedNodes, &nodeCopy)
	return node, nil
}

func (m *FakeNodeHandler) UpdateStatus(node *api.Node) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	nodeCopy := *node
	m.UpdatedNodeStatuses = append(m.UpdatedNodeStatuses, &nodeCopy)
	return node, nil
}

func (m *FakeNodeHandler) PatchStatus(nodeName string, data []byte) (*api.Node, error) {
	m.RequestCount++
	return &api.Node{}, nil
}

func (m *FakeNodeHandler) Watch(opts api.ListOptions) (watch.Interface, error) {
	return nil, nil
}

func (m *FakeNodeHandler) Patch(name string, pt api.PatchType, data []byte, subresources ...string) (*api.Node, error) {
	return nil, nil
}

// FakeRecorder is used as a fake during testing.
type FakeRecorder struct {
	source api.EventSource
	events []*api.Event
	clock  clock.Clock
}

func (f *FakeRecorder) Event(obj runtime.Object, eventtype, reason, message string) {
	f.generateEvent(obj, unversioned.Now(), eventtype, reason, message)
}

func (f *FakeRecorder) Eventf(obj runtime.Object, eventtype, reason, messageFmt string, args ...interface{}) {
	f.Event(obj, eventtype, reason, fmt.Sprintf(messageFmt, args...))
}

func (f *FakeRecorder) PastEventf(obj runtime.Object, timestamp unversioned.Time, eventtype, reason, messageFmt string, args ...interface{}) {
}

func (f *FakeRecorder) generateEvent(obj runtime.Object, timestamp unversioned.Time, eventtype, reason, message string) {
	ref, err := api.GetReference(obj)
	if err != nil {
		return
	}
	event := f.makeEvent(ref, eventtype, reason, message)
	event.Source = f.source
	if f.events != nil {
		fmt.Println("write event")
		f.events = append(f.events, event)
	}
}

func (f *FakeRecorder) makeEvent(ref *api.ObjectReference, eventtype, reason, message string) *api.Event {
	fmt.Println("make event")
	t := unversioned.Time{Time: f.clock.Now()}
	namespace := ref.Namespace
	if namespace == "" {
		namespace = api.NamespaceDefault
	}
	return &api.Event{
		ObjectMeta: api.ObjectMeta{
			Name:      fmt.Sprintf("%v.%x", ref.Name, t.UnixNano()),
			Namespace: namespace,
		},
		InvolvedObject: *ref,
		Reason:         reason,
		Message:        message,
		FirstTimestamp: t,
		LastTimestamp:  t,
		Count:          1,
		Type:           eventtype,
	}
}

func NewFakeRecorder() *FakeRecorder {
	return &FakeRecorder{
		source: api.EventSource{Component: "nodeControllerTest"},
		events: []*api.Event{},
		clock:  clock.NewFakeClock(time.Now()),
	}
}

func newNode(name string) *api.Node {
	return &api.Node{
		ObjectMeta: api.ObjectMeta{Name: name},
		Spec: api.NodeSpec{
			ExternalID: name,
		},
		Status: api.NodeStatus{
			Capacity: api.ResourceList{
				api.ResourceName(api.ResourceCPU):    resource.MustParse("10"),
				api.ResourceName(api.ResourceMemory): resource.MustParse("10G"),
			},
		},
	}
}

func newPod(name, host string) *api.Pod {
	pod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Namespace: "default",
			Name:      name,
		},
		Spec: api.PodSpec{
			NodeName: host,
		},
		Status: api.PodStatus{
			Conditions: []api.PodCondition{
				{
					Type:   api.PodReady,
					Status: api.ConditionTrue,
				},
			},
		},
	}

	return pod
}

func contains(node *api.Node, nodes []*api.Node) bool {
	for i := 0; i < len(nodes); i++ {
		if node.Name == nodes[i].Name {
			return true
		}
	}
	return false
}

// Returns list of zones for all Nodes stored in FakeNodeHandler
func getZones(nodeHandler *FakeNodeHandler) []string {
	nodes, _ := nodeHandler.List(api.ListOptions{})
	zones := sets.NewString()
	for _, node := range nodes.Items {
		zones.Insert(utilnode.GetZoneKey(&node))
	}
	return zones.List()
}

func createZoneID(region, zone string) string {
	return region + ":\x00:" + zone
}
<commit_msg>fix FakeNodeHandler List()<commit_after>/*
Copyright 2016 The Kubernetes 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 node

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

	"k8s.io/kubernetes/pkg/api"
	apierrors "k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/fake"
	unversionedcore "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/core/unversioned"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util/clock"
	utilnode "k8s.io/kubernetes/pkg/util/node"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/watch"
)

// FakeNodeHandler is a fake implementation of NodesInterface and NodeInterface. It
// allows test cases to have fine-grained control over mock behaviors. We also need
// PodsInterface and PodInterface to test list & delet pods, which is implemented in
// the embedded client.Fake field.
type FakeNodeHandler struct {
	*fake.Clientset

	// Input: Hooks determine if request is valid or not
	CreateHook func(*FakeNodeHandler, *api.Node) bool
	Existing   []*api.Node

	// Output
	CreatedNodes        []*api.Node
	DeletedNodes        []*api.Node
	UpdatedNodes        []*api.Node
	UpdatedNodeStatuses []*api.Node
	RequestCount        int

	// Synchronization
	lock           sync.Mutex
	deleteWaitChan chan struct{}
}

type FakeLegacyHandler struct {
	unversionedcore.CoreInterface
	n *FakeNodeHandler
}

func (c *FakeNodeHandler) getUpdatedNodesCopy() []*api.Node {
	c.lock.Lock()
	defer c.lock.Unlock()
	updatedNodesCopy := make([]*api.Node, len(c.UpdatedNodes), len(c.UpdatedNodes))
	for i, ptr := range c.UpdatedNodes {
		updatedNodesCopy[i] = ptr
	}
	return updatedNodesCopy
}

func (c *FakeNodeHandler) Core() unversionedcore.CoreInterface {
	return &FakeLegacyHandler{c.Clientset.Core(), c}
}

func (m *FakeLegacyHandler) Nodes() unversionedcore.NodeInterface {
	return m.n
}

func (m *FakeNodeHandler) Create(node *api.Node) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	for _, n := range m.Existing {
		if n.Name == node.Name {
			return nil, apierrors.NewAlreadyExists(api.Resource("nodes"), node.Name)
		}
	}
	if m.CreateHook == nil || m.CreateHook(m, node) {
		nodeCopy := *node
		m.CreatedNodes = append(m.CreatedNodes, &nodeCopy)
		return node, nil
	} else {
		return nil, errors.New("Create error.")
	}
}

func (m *FakeNodeHandler) Get(name string) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	for i := range m.Existing {
		if m.Existing[i].Name == name {
			nodeCopy := *m.Existing[i]
			return &nodeCopy, nil
		}
	}
	return nil, nil
}

func (m *FakeNodeHandler) List(opts api.ListOptions) (*api.NodeList, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	var nodes []*api.Node
	for i := 0; i < len(m.UpdatedNodes); i++ {
		if !contains(m.UpdatedNodes[i], m.DeletedNodes) {
			nodes = append(nodes, m.UpdatedNodes[i])
		}
	}
	for i := 0; i < len(m.Existing); i++ {
		if !contains(m.Existing[i], m.DeletedNodes) && !contains(m.Existing[i], nodes) {
			nodes = append(nodes, m.Existing[i])
		}
	}
	for i := 0; i < len(m.CreatedNodes); i++ {
		if !contains(m.CreatedNodes[i], m.DeletedNodes) && !contains(m.CreatedNodes[i], nodes) {
			nodes = append(nodes, m.CreatedNodes[i])
		}
	}
	nodeList := &api.NodeList{}
	for _, node := range nodes {
		nodeList.Items = append(nodeList.Items, *node)
	}
	return nodeList, nil
}

func (m *FakeNodeHandler) Delete(id string, opt *api.DeleteOptions) error {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		if m.deleteWaitChan != nil {
			m.deleteWaitChan <- struct{}{}
		}
		m.lock.Unlock()
	}()
	m.DeletedNodes = append(m.DeletedNodes, newNode(id))
	return nil
}

func (m *FakeNodeHandler) DeleteCollection(opt *api.DeleteOptions, listOpts api.ListOptions) error {
	return nil
}

func (m *FakeNodeHandler) Update(node *api.Node) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	nodeCopy := *node
	m.UpdatedNodes = append(m.UpdatedNodes, &nodeCopy)
	return node, nil
}

func (m *FakeNodeHandler) UpdateStatus(node *api.Node) (*api.Node, error) {
	m.lock.Lock()
	defer func() {
		m.RequestCount++
		m.lock.Unlock()
	}()
	nodeCopy := *node
	m.UpdatedNodeStatuses = append(m.UpdatedNodeStatuses, &nodeCopy)
	return node, nil
}

func (m *FakeNodeHandler) PatchStatus(nodeName string, data []byte) (*api.Node, error) {
	m.RequestCount++
	return &api.Node{}, nil
}

func (m *FakeNodeHandler) Watch(opts api.ListOptions) (watch.Interface, error) {
	return nil, nil
}

func (m *FakeNodeHandler) Patch(name string, pt api.PatchType, data []byte, subresources ...string) (*api.Node, error) {
	return nil, nil
}

// FakeRecorder is used as a fake during testing.
type FakeRecorder struct {
	source api.EventSource
	events []*api.Event
	clock  clock.Clock
}

func (f *FakeRecorder) Event(obj runtime.Object, eventtype, reason, message string) {
	f.generateEvent(obj, unversioned.Now(), eventtype, reason, message)
}

func (f *FakeRecorder) Eventf(obj runtime.Object, eventtype, reason, messageFmt string, args ...interface{}) {
	f.Event(obj, eventtype, reason, fmt.Sprintf(messageFmt, args...))
}

func (f *FakeRecorder) PastEventf(obj runtime.Object, timestamp unversioned.Time, eventtype, reason, messageFmt string, args ...interface{}) {
}

func (f *FakeRecorder) generateEvent(obj runtime.Object, timestamp unversioned.Time, eventtype, reason, message string) {
	ref, err := api.GetReference(obj)
	if err != nil {
		return
	}
	event := f.makeEvent(ref, eventtype, reason, message)
	event.Source = f.source
	if f.events != nil {
		fmt.Println("write event")
		f.events = append(f.events, event)
	}
}

func (f *FakeRecorder) makeEvent(ref *api.ObjectReference, eventtype, reason, message string) *api.Event {
	fmt.Println("make event")
	t := unversioned.Time{Time: f.clock.Now()}
	namespace := ref.Namespace
	if namespace == "" {
		namespace = api.NamespaceDefault
	}
	return &api.Event{
		ObjectMeta: api.ObjectMeta{
			Name:      fmt.Sprintf("%v.%x", ref.Name, t.UnixNano()),
			Namespace: namespace,
		},
		InvolvedObject: *ref,
		Reason:         reason,
		Message:        message,
		FirstTimestamp: t,
		LastTimestamp:  t,
		Count:          1,
		Type:           eventtype,
	}
}

func NewFakeRecorder() *FakeRecorder {
	return &FakeRecorder{
		source: api.EventSource{Component: "nodeControllerTest"},
		events: []*api.Event{},
		clock:  clock.NewFakeClock(time.Now()),
	}
}

func newNode(name string) *api.Node {
	return &api.Node{
		ObjectMeta: api.ObjectMeta{Name: name},
		Spec: api.NodeSpec{
			ExternalID: name,
		},
		Status: api.NodeStatus{
			Capacity: api.ResourceList{
				api.ResourceName(api.ResourceCPU):    resource.MustParse("10"),
				api.ResourceName(api.ResourceMemory): resource.MustParse("10G"),
			},
		},
	}
}

func newPod(name, host string) *api.Pod {
	pod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Namespace: "default",
			Name:      name,
		},
		Spec: api.PodSpec{
			NodeName: host,
		},
		Status: api.PodStatus{
			Conditions: []api.PodCondition{
				{
					Type:   api.PodReady,
					Status: api.ConditionTrue,
				},
			},
		},
	}

	return pod
}

func contains(node *api.Node, nodes []*api.Node) bool {
	for i := 0; i < len(nodes); i++ {
		if node.Name == nodes[i].Name {
			return true
		}
	}
	return false
}

// Returns list of zones for all Nodes stored in FakeNodeHandler
func getZones(nodeHandler *FakeNodeHandler) []string {
	nodes, _ := nodeHandler.List(api.ListOptions{})
	zones := sets.NewString()
	for _, node := range nodes.Items {
		zones.Insert(utilnode.GetZoneKey(&node))
	}
	return zones.List()
}

func createZoneID(region, zone string) string {
	return region + ":\x00:" + zone
}
<|endoftext|>


<commit_before>package gateway_test

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"testing"

	"github.com/datawire/ambassador/pkg/envoy"
	"github.com/datawire/ambassador/pkg/gateway"
	"github.com/datawire/ambassador/pkg/kates"
	"github.com/stretchr/testify/require"
	"gotest.tools/assert"
)

func TestGatewayMatches(t *testing.T) {
	envoy.SetupRequestLogger(t, ":9000", ":9002")
	e := envoy.SetupEnvoyController(t, ":8003")
	envoy.SetupEnvoy(t, envoy.GetLoopbackAddr(8003), "8080:8080")

	d := makeDispatcher(t)

	// One rule for each type of path match (exact, prefix, regex) and each type of header match
	// (exact and regex).
	err := d.UpsertYaml(`
---
kind: Gateway
apiVersion: networking.x-k8s.io/v1alpha1
metadata:
  name: my-gateway
spec:
  listeners:
  - protocol: HTTP
    port: 8080
---
kind: HTTPRoute
apiVersion: networking.x-k8s.io/v1alpha1
metadata:
  name: my-route
spec:
  rules:
  - matches:
    - path:
        type: Exact
        value: /exact
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - path:
        type: Prefix
        value: /prefix
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - path:
        type: RegularExpression
        value: "/regular_expression(_[aA]+)?"
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - headers:
        type: Exact
        values:
          exact: foo
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - headers:
        type: RegularExpression
        values:
          regular_expression: "foo(_[aA]+)?"
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
`)

	require.NoError(t, err)

	loopbackIp := envoy.GetLoopbackIp()

	err = d.Upsert(makeEndpoint("default", "foo-backend-1", loopbackIp, 9000))
	require.NoError(t, err)
	err = d.Upsert(makeEndpoint("default", "foo-backend-2", loopbackIp, 9001))
	require.NoError(t, err)

	version, snapshot := d.GetSnapshot()
	status := e.Configure("test-id", version, *snapshot)
	if status != nil {
		t.Fatalf("envoy error: %s", status.Message)
	}

	urlBase := fmt.Sprintf("http://%s:8080", loopbackIp)

	assertGet(t, urlBase+"/exact", 200, "Hello World")
	assertGet(t, urlBase+"/exact/foo", 404, "")
	assertGet(t, urlBase+"/prefix", 200, "Hello World")
	assertGet(t, urlBase+"/prefix/foo", 200, "Hello World")

	assertGet(t, urlBase+"/regular_expression", 200, "Hello World")
	assertGet(t, urlBase+"/regular_expression_a", 200, "Hello World")
	assertGet(t, urlBase+"/regular_expression_aaaaaaaa", 200, "Hello World")
	assertGet(t, urlBase+"/regular_expression_aaAaaaAa", 200, "Hello World")
	assertGet(t, urlBase+"/regular_expression_aaAaaaAab", 404, "")

	assertGetHeader(t, urlBase+"", "exact", "foo", 200, "Hello World")
	assertGetHeader(t, urlBase+"", "exact", "bar", 404, "")
	assertGetHeader(t, urlBase+"", "regular_expression", "foo", 200, "Hello World")
	assertGetHeader(t, urlBase+"", "regular_expression", "foo_aaaaAaaaa", 200, "Hello World")
	assertGetHeader(t, urlBase+"", "regular_expression", "foo_aaaaAaaaab", 404, "")
	assertGetHeader(t, urlBase+"", "regular_expression", "bar", 404, "")
}

func makeDispatcher(t *testing.T) *gateway.Dispatcher {
	d := gateway.NewDispatcher()
	err := d.Register("Gateway", gateway.Compile_Gateway)
	require.NoError(t, err)
	err = d.Register("HTTPRoute", gateway.Compile_HTTPRoute)
	require.NoError(t, err)
	err = d.Register("Endpoints", gateway.Compile_Endpoints)
	require.NoError(t, err)
	return d
}

func makeEndpoint(namespace, name, ip string, port int) *kates.Endpoints {
	ports := []kates.EndpointPort{{Port: int32(port)}}
	addrs := []kates.EndpointAddress{{IP: ip}}

	return &kates.Endpoints{
		TypeMeta:   kates.TypeMeta{Kind: "Endpoints"},
		ObjectMeta: kates.ObjectMeta{Namespace: namespace, Name: name},
		Subsets:    []kates.EndpointSubset{{Addresses: addrs, Ports: ports}},
	}
}

func assertGet(t *testing.T, url string, code int, expected string) {
	resp, err := http.Get(url)
	require.NoError(t, err)
	require.Equal(t, code, resp.StatusCode)
	actual, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Equal(t, expected, string(actual))
}

func assertGetHeader(t *testing.T, url, header, value string, code int, expected string) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	require.NoError(t, err)
	req.Header.Set(header, value)
	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, code, resp.StatusCode)
	actual, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Equal(t, expected, string(actual))
}
<commit_msg>(from AES) add checkReady to envoy test<commit_after>package gateway_test

import (
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/datawire/ambassador/pkg/envoy"
	"github.com/datawire/ambassador/pkg/gateway"
	"github.com/datawire/ambassador/pkg/kates"
	"github.com/stretchr/testify/require"
	"gotest.tools/assert"
)

func TestGatewayMatches(t *testing.T) {
	envoy.SetupRequestLogger(t, ":9000", ":9002")
	e := envoy.SetupEnvoyController(t, ":8003")
	envoy.SetupEnvoy(t, envoy.GetLoopbackAddr(8003), "8080:8080")

	d := makeDispatcher(t)

	// One rule for each type of path match (exact, prefix, regex) and each type of header match
	// (exact and regex).
	err := d.UpsertYaml(`
---
kind: Gateway
apiVersion: networking.x-k8s.io/v1alpha1
metadata:
  name: my-gateway
spec:
  listeners:
  - protocol: HTTP
    port: 8080
---
kind: HTTPRoute
apiVersion: networking.x-k8s.io/v1alpha1
metadata:
  name: my-route
spec:
  rules:
  - matches:
    - path:
        type: Exact
        value: /exact
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - path:
        type: Prefix
        value: /prefix
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - path:
        type: RegularExpression
        value: "/regular_expression(_[aA]+)?"
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - headers:
        type: Exact
        values:
          exact: foo
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
  - matches:
    - headers:
        type: RegularExpression
        values:
          regular_expression: "foo(_[aA]+)?"
    forwardTo:
    - serviceName: foo-backend-1
      weight: 100
`)

	require.NoError(t, err)

	loopbackIp := envoy.GetLoopbackIp()

	err = d.Upsert(makeEndpoint("default", "foo-backend-1", loopbackIp, 9000))
	require.NoError(t, err)
	err = d.Upsert(makeEndpoint("default", "foo-backend-2", loopbackIp, 9001))
	require.NoError(t, err)

	version, snapshot := d.GetSnapshot()
	status := e.Configure("test-id", version, *snapshot)
	if status != nil {
		t.Fatalf("envoy error: %s", status.Message)
	}

	// Sometimes envoy seems to acknowledge the configuration before listening on the port. (This is
	// weird because sometimes envoy sends back an error indicating that it cannot bind to the
	// port. Either way, we need to check that we can actually connect before running the rest of
	// the tests.
	checkReady(t, "http://127.0.0.1:8080/")

	assertGet(t, "http://127.0.0.1:8080/exact", 200, "Hello World")
	assertGet(t, "http://127.0.0.1:8080/exact/foo", 404, "")
	assertGet(t, "http://127.0.0.1:8080/prefix", 200, "Hello World")
	assertGet(t, "http://127.0.0.1:8080/prefix/foo", 200, "Hello World")

	assertGet(t, "http://127.0.0.1:8080/regular_expression", 200, "Hello World")
	assertGet(t, "http://127.0.0.1:8080/regular_expression_a", 200, "Hello World")
	assertGet(t, "http://127.0.0.1:8080/regular_expression_aaaaaaaa", 200, "Hello World")
	assertGet(t, "http://127.0.0.1:8080/regular_expression_aaAaaaAa", 200, "Hello World")
	assertGet(t, "http://127.0.0.1:8080/regular_expression_aaAaaaAab", 404, "")

	assertGetHeader(t, "http://127.0.0.1:8080", "exact", "foo", 200, "Hello World")
	assertGetHeader(t, "http://127.0.0.1:8080", "exact", "bar", 404, "")
	assertGetHeader(t, "http://127.0.0.1:8080", "regular_expression", "foo", 200, "Hello World")
	assertGetHeader(t, "http://127.0.0.1:8080", "regular_expression", "foo_aaaaAaaaa", 200, "Hello World")
	assertGetHeader(t, "http://127.0.0.1:8080", "regular_expression", "foo_aaaaAaaaab", 404, "")
	assertGetHeader(t, "http://127.0.0.1:8080", "regular_expression", "bar", 404, "")
}

func makeDispatcher(t *testing.T) *gateway.Dispatcher {
	d := gateway.NewDispatcher()
	err := d.Register("Gateway", gateway.Compile_Gateway)
	require.NoError(t, err)
	err = d.Register("HTTPRoute", gateway.Compile_HTTPRoute)
	require.NoError(t, err)
	err = d.Register("Endpoints", gateway.Compile_Endpoints)
	require.NoError(t, err)
	return d
}

func makeEndpoint(namespace, name, ip string, port int) *kates.Endpoints {
	ports := []kates.EndpointPort{{Port: int32(port)}}
	addrs := []kates.EndpointAddress{{IP: ip}}

	return &kates.Endpoints{
		TypeMeta:   kates.TypeMeta{Kind: "Endpoints"},
		ObjectMeta: kates.ObjectMeta{Namespace: namespace, Name: name},
		Subsets:    []kates.EndpointSubset{{Addresses: addrs, Ports: ports}},
	}
}

func checkReady(t *testing.T, url string) {
	delay := 10 * time.Millisecond
	for {
		if delay > 10*time.Second {
			require.Fail(t, "url never became ready", url)
		}
		_, err := http.Get(url)
		if err != nil {
			t.Logf("error %v, retrying...", err)
			time.Sleep(delay)
			delay = delay * 2
		}
		return
	}
}

func assertGet(t *testing.T, url string, code int, expected string) {
	resp, err := http.Get(url)
	require.NoError(t, err)
	require.Equal(t, code, resp.StatusCode)
	actual, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Equal(t, expected, string(actual))
}

func assertGetHeader(t *testing.T, url, header, value string, code int, expected string) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	require.NoError(t, err)
	req.Header.Set(header, value)
	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, code, resp.StatusCode)
	actual, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Equal(t, expected, string(actual))
}
<|endoftext|>


<commit_before>package middleware

import (
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/login"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAuthProxyWithLdapEnabled(t *testing.T) {
	Convey("When calling sync grafana user with ldap user", t, func() {

		setting.LdapEnabled = true
		setting.AuthProxyLdapSyncTtl = 60

		servers := []*login.LdapServerConf{{Host: "127.0.0.1"}}
		login.LdapCfg = login.LdapConfig{Servers: servers}
		mockLdapAuther := mockLdapAuthenticator{}

		login.NewLdapAuthenticator = func(server *login.LdapServerConf) login.ILdapAuther {
			return &mockLdapAuther
		}

		signedInUser := m.SignedInUser{}
		query := m.GetSignedInUserQuery{Result: &signedInUser}

		Convey("When session variable lastLdapSync not set, call syncSignedInUser and set lastLdapSync", func() {
			// arrange
			session := mockSession{}
			ctx := Context{Session: &session}
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldBeNil)

			// act
			syncGrafanaUserWithLdapUser(&ctx, &query)

			// assert
			So(mockLdapAuther.syncSignedInUserCalled, ShouldBeTrue)
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldBeGreaterThan, 0)
		})

		Convey("When session variable not expired, don't sync and don't change session var", func() {
			// arrange
			session := mockSession{}
			ctx := Context{Session: &session}
			now := time.Now().Unix()
			session.Set(SESS_KEY_LASTLDAPSYNC, now)

			// act
			syncGrafanaUserWithLdapUser(&ctx, &query)

			// assert
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldEqual, now)
			So(mockLdapAuther.syncSignedInUserCalled, ShouldBeFalse)
		})

		Convey("When lastldapsync is expired, session variable should be updated", func() {
			// arrange
			session := mockSession{}
			ctx := Context{Session: &session}
			expiredTime := time.Now().Add(time.Duration(-120) * time.Minute).Unix()
			session.Set(SESS_KEY_LASTLDAPSYNC, expiredTime)

			// act
			syncGrafanaUserWithLdapUser(&ctx, &query)

			// assert
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldBeGreaterThan, expiredTime)
			So(mockLdapAuther.syncSignedInUserCalled, ShouldBeTrue)
		})
	})
}

type mockSession struct {
	value interface{}
}

func (s *mockSession) Start(c *Context) error {
	return nil
}

func (s *mockSession) Set(k interface{}, v interface{}) error {
	s.value = v
	return nil
}

func (s *mockSession) Get(k interface{}) interface{} {
	return s.value
}

func (s *mockSession) ID() string {
	return ""
}

func (s *mockSession) Release() error {
	return nil
}

func (s *mockSession) Destory(c *Context) error {
	return nil
}

type mockLdapAuthenticator struct {
	syncSignedInUserCalled bool
}

func (a *mockLdapAuthenticator) Login(query *login.LoginUserQuery) error {
	return nil
}

func (a *mockLdapAuthenticator) SyncSignedInUser(signedInUser *m.SignedInUser) error {
	a.syncSignedInUserCalled = true
	return nil
}

func (a *mockLdapAuthenticator) GetGrafanaUserFor(ldapUser *login.LdapUserInfo) (*m.User, error) {
	return nil, nil
}
func (a *mockLdapAuthenticator) SyncOrgRoles(user *m.User, ldapUser *login.LdapUserInfo) error {
	return nil
}
<commit_msg>auth: fix warning in test<commit_after>package middleware

import (
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/login"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAuthProxyWithLdapEnabled(t *testing.T) {
	Convey("When calling sync grafana user with ldap user", t, func() {

		setting.LdapEnabled = true
		setting.AuthProxyLdapSyncTtl = 60

		servers := []*login.LdapServerConf{{Host: "127.0.0.1"}}
		login.LdapCfg = login.LdapConfig{Servers: servers}
		mockLdapAuther := mockLdapAuthenticator{}

		login.NewLdapAuthenticator = func(server *login.LdapServerConf) login.ILdapAuther {
			return &mockLdapAuther
		}

		signedInUser := m.SignedInUser{}
		query := m.GetSignedInUserQuery{Result: &signedInUser}

		Convey("When session variable lastLdapSync not set, call syncSignedInUser and set lastLdapSync", func() {
			// arrange
			session := mockSession{}
			ctx := Context{Session: &session}
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldBeNil)

			// act
			syncGrafanaUserWithLdapUser(&ctx, &query)

			// assert
			So(mockLdapAuther.syncSignedInUserCalled, ShouldBeTrue)
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldBeGreaterThan, 0)
		})

		Convey("When session variable not expired, don't sync and don't change session var", func() {
			// arrange
			session := mockSession{}
			ctx := Context{Session: &session}
			now := time.Now().Unix()
			session.Set(SESS_KEY_LASTLDAPSYNC, now)

			// act
			syncGrafanaUserWithLdapUser(&ctx, &query)

			// assert
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldEqual, now)
			So(mockLdapAuther.syncSignedInUserCalled, ShouldBeFalse)
		})

		Convey("When lastldapsync is expired, session variable should be updated", func() {
			// arrange
			session := mockSession{}
			ctx := Context{Session: &session}
			expiredTime := time.Now().Add(time.Duration(-120) * time.Minute).Unix()
			session.Set(SESS_KEY_LASTLDAPSYNC, expiredTime)

			// act
			syncGrafanaUserWithLdapUser(&ctx, &query)

			// assert
			So(session.Get(SESS_KEY_LASTLDAPSYNC), ShouldBeGreaterThan, expiredTime)
			So(mockLdapAuther.syncSignedInUserCalled, ShouldBeTrue)
		})
	})
}

type mockSession struct {
	value interface{}
}

func (s *mockSession) Start(c *Context) error {
	return nil
}

func (s *mockSession) Set(k interface{}, v interface{}) error {
	s.value = v
	return nil
}

func (s *mockSession) Get(k interface{}) interface{} {
	return s.value
}

func (s *mockSession) Delete(k interface{}) interface{} {
	return nil
}

func (s *mockSession) ID() string {
	return ""
}

func (s *mockSession) Release() error {
	return nil
}

func (s *mockSession) Destory(c *Context) error {
	return nil
}

type mockLdapAuthenticator struct {
	syncSignedInUserCalled bool
}

func (a *mockLdapAuthenticator) Login(query *login.LoginUserQuery) error {
	return nil
}

func (a *mockLdapAuthenticator) SyncSignedInUser(signedInUser *m.SignedInUser) error {
	a.syncSignedInUserCalled = true
	return nil
}

func (a *mockLdapAuthenticator) GetGrafanaUserFor(ldapUser *login.LdapUserInfo) (*m.User, error) {
	return nil, nil
}
func (a *mockLdapAuthenticator) SyncOrgRoles(user *m.User, ldapUser *login.LdapUserInfo) error {
	return nil
}
<|endoftext|>