<commit_before>package privval

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

	"github.com/tendermint/tendermint/crypto/ed25519"
)

//-------------------------------------------
// helper funcs

func newPrivKey() ed25519.PrivKeyEd25519 {
	return ed25519.GenPrivKey()
}

//-------------------------------------------
// tests

type listenerTestCase struct {
	description string // For test reporting purposes.
	listener    net.Listener
	dialer      SocketDialer
}

// testUnixAddr will attempt to obtain a platform-independent temporary file
// name for a Unix socket
func testUnixAddr() (string, error) {
	f, err := ioutil.TempFile("", "tendermint-privval-test-*")
	if err != nil {
		return "", err
	}
	addr := f.Name()
	f.Close()
	os.Remove(addr)
	return addr, nil
}

func tcpListenerTestCase(t *testing.T, timeoutAccept, timeoutReadWrite time.Duration) listenerTestCase {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}

	tcpLn := NewTCPListener(ln, newPrivKey())
	TCPListenerTimeoutAccept(timeoutAccept)(tcpLn)
	TCPListenerTimeoutReadWrite(timeoutReadWrite)(tcpLn)
	return listenerTestCase{
		description: "TCP",
		listener:    tcpLn,
		dialer:      DialTCPFn(ln.Addr().String(), testTimeoutReadWrite, newPrivKey()),
	}
}

func unixListenerTestCase(t *testing.T, timeoutAccept, timeoutReadWrite time.Duration) listenerTestCase {
	addr, err := testUnixAddr()
	if err != nil {
		t.Fatal(err)
	}
	ln, err := net.Listen("unix", addr)
	if err != nil {
		t.Fatal(err)
	}

	unixLn := NewUnixListener(ln)
	UnixListenerTimeoutAccept(timeoutAccept)(unixLn)
	UnixListenerTimeoutReadWrite(timeoutReadWrite)(unixLn)
	return listenerTestCase{
		description: "Unix",
		listener:    unixLn,
		dialer:      DialUnixFn(addr),
	}
}

func listenerTestCases(t *testing.T, timeoutAccept, timeoutReadWrite time.Duration) []listenerTestCase {
	return []listenerTestCase{
		tcpListenerTestCase(t, timeoutAccept, timeoutReadWrite),
		unixListenerTestCase(t, timeoutAccept, timeoutReadWrite),
	}
}

func TestListenerTimeoutAccept(t *testing.T) {
	for _, tc := range listenerTestCases(t, time.Millisecond, time.Second) {
		_, err := tc.listener.Accept()
		opErr, ok := err.(*net.OpError)
		if !ok {
			t.Fatalf("for %s listener, have %v, want *net.OpError", tc.description, err)
		}

		if have, want := opErr.Op, "accept"; have != want {
			t.Errorf("for %s listener,  have %v, want %v", tc.description, have, want)
		}
	}
}

func TestListenerTimeoutReadWrite(t *testing.T) {
	for _, tc := range listenerTestCases(t, time.Second, time.Millisecond) {
		go func(dialer SocketDialer) {
			_, err := dialer()
			if err != nil {
				panic(err)
			}
		}(tc.dialer)

		c, err := tc.listener.Accept()
		if err != nil {
			t.Fatal(err)
		}

		time.Sleep(2 * time.Millisecond)

		msg := make([]byte, 200)
		_, err = c.Read(msg)
		opErr, ok := err.(*net.OpError)
		if !ok {
			t.Fatalf("for %s listener, have %v, want *net.OpError", tc.description, err)
		}

		if have, want := opErr.Op, "read"; have != want {
			t.Errorf("for %s listener, have %v, want %v", tc.description, have, want)
		}
	}
}
<commit_msg>privval: increase timeout to mitigate non-deterministic test failure (#3580)<commit_after>package privval

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

	"github.com/tendermint/tendermint/crypto/ed25519"
)

//-------------------------------------------
// helper funcs

func newPrivKey() ed25519.PrivKeyEd25519 {
	return ed25519.GenPrivKey()
}

//-------------------------------------------
// tests

type listenerTestCase struct {
	description string // For test reporting purposes.
	listener    net.Listener
	dialer      SocketDialer
}

// testUnixAddr will attempt to obtain a platform-independent temporary file
// name for a Unix socket
func testUnixAddr() (string, error) {
	f, err := ioutil.TempFile("", "tendermint-privval-test-*")
	if err != nil {
		return "", err
	}
	addr := f.Name()
	f.Close()
	os.Remove(addr)
	return addr, nil
}

func tcpListenerTestCase(t *testing.T, timeoutAccept, timeoutReadWrite time.Duration) listenerTestCase {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}

	tcpLn := NewTCPListener(ln, newPrivKey())
	TCPListenerTimeoutAccept(timeoutAccept)(tcpLn)
	TCPListenerTimeoutReadWrite(timeoutReadWrite)(tcpLn)
	return listenerTestCase{
		description: "TCP",
		listener:    tcpLn,
		dialer:      DialTCPFn(ln.Addr().String(), testTimeoutReadWrite, newPrivKey()),
	}
}

func unixListenerTestCase(t *testing.T, timeoutAccept, timeoutReadWrite time.Duration) listenerTestCase {
	addr, err := testUnixAddr()
	if err != nil {
		t.Fatal(err)
	}
	ln, err := net.Listen("unix", addr)
	if err != nil {
		t.Fatal(err)
	}

	unixLn := NewUnixListener(ln)
	UnixListenerTimeoutAccept(timeoutAccept)(unixLn)
	UnixListenerTimeoutReadWrite(timeoutReadWrite)(unixLn)
	return listenerTestCase{
		description: "Unix",
		listener:    unixLn,
		dialer:      DialUnixFn(addr),
	}
}

func listenerTestCases(t *testing.T, timeoutAccept, timeoutReadWrite time.Duration) []listenerTestCase {
	return []listenerTestCase{
		tcpListenerTestCase(t, timeoutAccept, timeoutReadWrite),
		unixListenerTestCase(t, timeoutAccept, timeoutReadWrite),
	}
}

func TestListenerTimeoutAccept(t *testing.T) {
	for _, tc := range listenerTestCases(t, time.Millisecond, time.Second) {
		_, err := tc.listener.Accept()
		opErr, ok := err.(*net.OpError)
		if !ok {
			t.Fatalf("for %s listener, have %v, want *net.OpError", tc.description, err)
		}

		if have, want := opErr.Op, "accept"; have != want {
			t.Errorf("for %s listener,  have %v, want %v", tc.description, have, want)
		}
	}
}

func TestListenerTimeoutReadWrite(t *testing.T) {
	const (
		// This needs to be long enough s.t. the Accept will definitely succeed:
		timeoutAccept = time.Second
		// This can be really short but in the TCP case, the accept can
		// also trigger a timeoutReadWrite. Hence, we need to give it some time.
		// Note: this controls how long this test actually runs.
		timeoutReadWrite = 10 * time.Millisecond
	)
	for _, tc := range listenerTestCases(t, timeoutAccept, timeoutReadWrite) {
		go func(dialer SocketDialer) {
			_, err := dialer()
			if err != nil {
				panic(err)
			}
		}(tc.dialer)

		c, err := tc.listener.Accept()
		if err != nil {
			t.Fatal(err)
		}

		// this will timeout because we don't write anything:
		msg := make([]byte, 200)
		_, err = c.Read(msg)
		opErr, ok := err.(*net.OpError)
		if !ok {
			t.Fatalf("for %s listener, have %v, want *net.OpError", tc.description, err)
		}

		if have, want := opErr.Op, "read"; have != want {
			t.Errorf("for %s listener, have %v, want %v", tc.description, have, want)
		}

		if have, want := opErr.Timeout(), true; have != want {
			t.Errorf("for %s listener, got unexpected error: have %v, want %v", tc.description, have, want)
		}
	}
}
<|endoftext|>


<commit_before>package ingress

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"

	"strconv"

	"github.com/rancher/norman/types/convert"
	util "github.com/rancher/rancher/pkg/controllers/user/workload"
	"github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/api/extensions/v1beta1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
)

//TODO fix the workload services cleanup

// This controller is responsible for monitoring ingress and
// creating services for them if the service is missing
// Creation would only happen if the service reference was put by Rancher API based on
// ingress.backend.workloadId. This information is stored in state field in the annotation

type Controller struct {
	serviceLister v1.ServiceLister
	services      v1.ServiceInterface
}

func Register(ctx context.Context, workload *config.UserOnlyContext) {
	c := &Controller{
		services:      workload.Core.Services(""),
		serviceLister: workload.Core.Services("").Controller().Lister(),
	}
	workload.Extensions.Ingresses("").AddHandler("ingressWorkloadController", c.sync)
}

func (c *Controller) sync(key string, obj *v1beta1.Ingress) error {
	if obj == nil || obj.DeletionTimestamp != nil {
		return nil
	}
	state := getIngressState(obj)
	if state == nil {
		return nil
	}
	serviceToPort := make(map[string]string)
	serviceToKey := make(map[string]string)
	for _, r := range obj.Spec.Rules {
		host := r.Host
		for _, b := range r.HTTP.Paths {
			path := b.Path
			port := b.Backend.ServicePort.IntVal
			key := getStateKey(host, path, convert.ToString(port))
			if _, ok := state[key]; ok {
				serviceToKey[b.Backend.ServiceName] = key
				serviceToPort[b.Backend.ServiceName] = convert.ToString(port)
			}
		}
	}
	if obj.Spec.Backend != nil {
		serviceName := obj.Spec.Backend.ServiceName
		portStr := convert.ToString(obj.Spec.Backend.ServicePort.IntVal)
		key := getStateKey("", "", portStr)
		if _, ok := state[key]; ok {
			serviceToKey[serviceName] = key
			serviceToPort[serviceName] = portStr
		}
	}

	for serviceName, portStr := range serviceToPort {
		workloadIDs := state[serviceToKey[serviceName]]
		existing, err := c.serviceLister.Get(obj.Namespace, serviceName)
		if err != nil && !apierrors.IsNotFound(err) {
			return err
		}
		if existing == nil {
			controller := true
			ownerRef := metav1.OwnerReference{
				Name:       serviceName,
				APIVersion: "v1beta1/extensions",
				UID:        obj.UID,
				Kind:       "Ingress",
				Controller: &controller,
			}
			port, err := strconv.ParseInt(portStr, 10, 64)
			if err != nil {
				return err
			}
			servicePorts := []corev1.ServicePort{
				{
					Port:       int32(port),
					TargetPort: intstr.Parse(portStr),
					Protocol:   "TCP",
				},
			}
			annotations := make(map[string]string)
			annotations[util.WorkloadAnnotation] = workloadIDs
			service := &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:            serviceName,
					OwnerReferences: []metav1.OwnerReference{ownerRef},
					Namespace:       obj.Namespace,
					Annotations:     annotations,
				},
				Spec: corev1.ServiceSpec{
					ClusterIP: "None",
					Type:      "ClusterIP",
					Ports:     servicePorts,
				},
			}
			logrus.Infof("Creating headless service %s for ingress %s, port %s", serviceName, key, portStr)
			if _, err := c.services.Create(service); err != nil {
				return err
			}
		} else {
			ingressManaged := false
			for _, o := range existing.OwnerReferences {
				if o.UID == obj.UID {
					ingressManaged = true
					break
				}
			}
			if ingressManaged {
				// TODO - fix so the update is done only when workload ids are changed
				toUpdate := existing.DeepCopy()
				if toUpdate.Annotations == nil {
					toUpdate.Annotations = map[string]string{}
				}
				toUpdate.Annotations[util.WorkloadAnnotation] = workloadIDs
				if _, err := c.services.Update(toUpdate); err != nil {
					return err
				}
			}
		}
	}

	return nil
}

func getStateKey(host string, path string, port string) string {
	key := fmt.Sprintf("%s/%s/%s", host, path, port)
	return base64.URLEncoding.EncodeToString([]byte(key))
}

func getIngressState(obj *v1beta1.Ingress) map[string]string {
	annotations := obj.Annotations
	if annotations == nil {
		return nil
	}
	if v, ok := annotations["ingress.cattle.io/state"]; ok {
		state := make(map[string]string)
		json.Unmarshal([]byte(convert.ToString(v)), &state)
		return state
	}
	return nil
}
<commit_msg>Fixed parsing ingress workloads<commit_after>package ingress

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"

	"strconv"

	"strings"

	"github.com/rancher/norman/types/convert"
	util "github.com/rancher/rancher/pkg/controllers/user/workload"
	"github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/api/extensions/v1beta1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
)

//TODO fix the workload services cleanup

// This controller is responsible for monitoring ingress and
// creating services for them if the service is missing
// Creation would only happen if the service reference was put by Rancher API based on
// ingress.backend.workloadId. This information is stored in state field in the annotation

type Controller struct {
	serviceLister v1.ServiceLister
	services      v1.ServiceInterface
}

func Register(ctx context.Context, workload *config.UserOnlyContext) {
	c := &Controller{
		services:      workload.Core.Services(""),
		serviceLister: workload.Core.Services("").Controller().Lister(),
	}
	workload.Extensions.Ingresses("").AddHandler("ingressWorkloadController", c.sync)
}

func (c *Controller) sync(key string, obj *v1beta1.Ingress) error {
	if obj == nil || obj.DeletionTimestamp != nil {
		return nil
	}
	state := getIngressState(obj)
	if state == nil {
		return nil
	}
	serviceToPort := make(map[string]string)
	serviceToKey := make(map[string]string)
	for _, r := range obj.Spec.Rules {
		host := r.Host
		for _, b := range r.HTTP.Paths {
			path := b.Path
			port := b.Backend.ServicePort.IntVal
			key := getStateKey(host, path, convert.ToString(port))
			if _, ok := state[key]; ok {
				serviceToKey[b.Backend.ServiceName] = key
				serviceToPort[b.Backend.ServiceName] = convert.ToString(port)
			}
		}
	}
	if obj.Spec.Backend != nil {
		serviceName := obj.Spec.Backend.ServiceName
		portStr := convert.ToString(obj.Spec.Backend.ServicePort.IntVal)
		key := getStateKey("", "", portStr)
		if _, ok := state[key]; ok {
			serviceToKey[serviceName] = key
			serviceToPort[serviceName] = portStr
		}
	}

	for serviceName, portStr := range serviceToPort {
		workloadIDsStr := state[serviceToKey[serviceName]]
		b, err := json.Marshal(strings.Split(workloadIDsStr, "/"))
		if err != nil {
			return err
		}
		workloadIDs := string(b)
		existing, err := c.serviceLister.Get(obj.Namespace, serviceName)
		if err != nil && !apierrors.IsNotFound(err) {
			return err
		}
		if existing == nil {
			controller := true
			ownerRef := metav1.OwnerReference{
				Name:       serviceName,
				APIVersion: "v1beta1/extensions",
				UID:        obj.UID,
				Kind:       "Ingress",
				Controller: &controller,
			}
			port, err := strconv.ParseInt(portStr, 10, 64)
			if err != nil {
				return err
			}
			servicePorts := []corev1.ServicePort{
				{
					Port:       int32(port),
					TargetPort: intstr.Parse(portStr),
					Protocol:   "TCP",
				},
			}
			annotations := make(map[string]string)

			annotations[util.WorkloadAnnotation] = workloadIDs
			service := &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:            serviceName,
					OwnerReferences: []metav1.OwnerReference{ownerRef},
					Namespace:       obj.Namespace,
					Annotations:     annotations,
				},
				Spec: corev1.ServiceSpec{
					ClusterIP: "None",
					Type:      "ClusterIP",
					Ports:     servicePorts,
				},
			}
			logrus.Infof("Creating headless service %s for ingress %s, port %s", serviceName, key, portStr)
			if _, err := c.services.Create(service); err != nil {
				return err
			}
		} else {
			ingressManaged := false
			for _, o := range existing.OwnerReferences {
				if o.UID == obj.UID {
					ingressManaged = true
					break
				}
			}
			if ingressManaged {
				// TODO - fix so the update is done only when workload ids are changed
				toUpdate := existing.DeepCopy()
				if toUpdate.Annotations == nil {
					toUpdate.Annotations = map[string]string{}
				}
				toUpdate.Annotations[util.WorkloadAnnotation] = workloadIDs
				if _, err := c.services.Update(toUpdate); err != nil {
					return err
				}
			}
		}
	}

	return nil
}

func getStateKey(host string, path string, port string) string {
	key := fmt.Sprintf("%s/%s/%s", host, path, port)
	return base64.URLEncoding.EncodeToString([]byte(key))
}

func getIngressState(obj *v1beta1.Ingress) map[string]string {
	annotations := obj.Annotations
	if annotations == nil {
		return nil
	}
	if v, ok := annotations["ingress.cattle.io/state"]; ok {
		state := make(map[string]string)
		json.Unmarshal([]byte(convert.ToString(v)), &state)
		return state
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 customresource

import (
	"context"
	"fmt"
	"strings"

	"k8s.io/apimachinery/pkg/api/errors"
	metainternalversion "k8s.io/apimachinery/pkg/apis/meta/internalversion"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/apiserver/pkg/registry/rest"
)

// Registry is an interface for things that know how to store CustomResources.
type Registry interface {
	ListCustomResources(ctx context.Context, options *metainternalversion.ListOptions) (*unstructured.UnstructuredList, error)
	WatchCustomResources(ctx context.Context, options *metainternalversion.ListOptions) (watch.Interface, error)
	GetCustomResource(ctx context.Context, customResourceID string, options *metav1.GetOptions) (*unstructured.Unstructured, error)
	CreateCustomResource(ctx context.Context, customResource *unstructured.Unstructured, createValidation rest.ValidateObjectFunc) (*unstructured.Unstructured, error)
	UpdateCustomResource(ctx context.Context, customResource *unstructured.Unstructured, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, options *metav1.UpdateOptions) (*unstructured.Unstructured, error)
	DeleteCustomResource(ctx context.Context, customResourceID string) error
}

// storage puts strong typing around storage calls
type storage struct {
	rest.StandardStorage
}

// NewRegistry returns a new Registry interface for the given Storage. Any mismatched
// types will panic.
func NewRegistry(s rest.StandardStorage) Registry {
	return &storage{s}
}

func (s *storage) ListCustomResources(ctx context.Context, options *metainternalversion.ListOptions) (*unstructured.UnstructuredList, error) {
	if options != nil && options.FieldSelector != nil && !options.FieldSelector.Empty() {
		return nil, fmt.Errorf("field selector not supported yet")
	}
	obj, err := s.List(ctx, options)
	if err != nil {
		return nil, err
	}
	return obj.(*unstructured.UnstructuredList), err
}

func (s *storage) WatchCustomResources(ctx context.Context, options *metainternalversion.ListOptions) (watch.Interface, error) {
	return s.Watch(ctx, options)
}

func (s *storage) GetCustomResource(ctx context.Context, customResourceID string, options *metav1.GetOptions) (*unstructured.Unstructured, error) {
	obj, err := s.Get(ctx, customResourceID, options)
	customResource, ok := obj.(*unstructured.Unstructured)
	if !ok {
		return nil, fmt.Errorf("custom resource must be of type Unstructured")
	}

	if err != nil {
		apiVersion := customResource.GetAPIVersion()
		groupVersion := strings.Split(apiVersion, "/")
		group := groupVersion[0]
		return nil, errors.NewNotFound(schema.GroupResource{Group: group, Resource: "scale"}, customResourceID)
	}
	return customResource, nil
}

func (s *storage) CreateCustomResource(ctx context.Context, customResource *unstructured.Unstructured, createValidation rest.ValidateObjectFunc) (*unstructured.Unstructured, error) {
	obj, err := s.Create(ctx, customResource, rest.ValidateAllObjectFunc, &metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return obj.(*unstructured.Unstructured), nil
}

func (s *storage) UpdateCustomResource(ctx context.Context, customResource *unstructured.Unstructured, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, options *metav1.UpdateOptions) (*unstructured.Unstructured, error) {
	obj, _, err := s.Update(ctx, customResource.GetName(), rest.DefaultUpdatedObjectInfo(customResource), createValidation, updateValidation, false, options)
	if err != nil {
		return nil, err
	}
	return obj.(*unstructured.Unstructured), nil
}

func (s *storage) DeleteCustomResource(ctx context.Context, customResourceID string) error {
	_, _, err := s.Delete(ctx, customResourceID, nil)
	return err
}
<commit_msg>Remove hand-written typed registries<commit_after><|endoftext|>


<commit_before>package accesscontrol

import (
	"bytes"
	"context"
	"fmt"
	"text/template"
	"time"

	"github.com/grafana/grafana/pkg/infra/localcache"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
)

const (
	ttl           = 30 * time.Second
	cleanInterval = 2 * time.Minute
)

func NewScopeResolvers() ScopeResolvers {
	return ScopeResolvers{
		keywordResolvers: map[string]ScopeKeywordResolver{
			"users:self": userSelfResolver,
		},
		attributeResolvers: map[string]ScopeAttributeResolver{},
		cache:              localcache.New(ttl, cleanInterval),
		log:                log.New("accesscontrol.resolver"),
	}
}

type ScopeResolvers struct {
	log                log.Logger
	cache              *localcache.CacheService
	keywordResolvers   map[string]ScopeKeywordResolver
	attributeResolvers map[string]ScopeAttributeResolver
}

func (s *ScopeResolvers) GetScopeAttributeMutator(orgID int64) ScopeAttributeMutator {
	return func(ctx context.Context, scope string) ([]string, error) {
		key := getScopeCacheKey(orgID, scope)
		// Check cache before computing the scope
		if cachedScope, ok := s.cache.Get(key); ok {
			scopes := cachedScope.([]string)
			s.log.Debug("used cache to resolve '%v' to '%v'", scope, scopes)
			return scopes, nil
		}

		prefix := ScopePrefix(scope)
		if resolver, ok := s.attributeResolvers[prefix]; ok {
			scopes, err := resolver.Resolve(ctx, orgID, scope)
			if err != nil {
				return nil, fmt.Errorf("could not resolve %v: %w", scope, err)
			}
			// Cache result
			s.cache.Set(key, scopes, ttl)
			s.log.Debug("resolved '%v' to '%v'", scope, scopes)
			return scopes, nil
		}
		return []string{scope}, nil
	}
}

func (s *ScopeResolvers) GetScopeKeywordMutator(user *models.SignedInUser) ScopeKeywordMutator {
	return func(ctx context.Context, scope string) (string, error) {
		if resolver, ok := s.keywordResolvers[scope]; ok {
			scopes, err := resolver.Resolve(ctx, user)
			if err != nil {
				return "", fmt.Errorf("could not resolve %v: %w", scope, err)
			}
			s.log.Debug("resolved '%v' to '%v'", scope, scopes)
			return scopes, nil
		}
		// By default, the scope remains unchanged
		return scope, nil
	}
}

func (s *ScopeResolvers) AddScopeKeywordResolver(keyword string, resolver ScopeKeywordResolver) {
	s.log.Debug("adding scope keyword resolver for '%v'", keyword)
	s.keywordResolvers[keyword] = resolver
}

func (s *ScopeResolvers) AddScopeAttributeResolver(prefix string, resolver ScopeAttributeResolver) {
	s.log.Debug("adding scope attribute resolver for '%v'", prefix)
	s.attributeResolvers[prefix] = resolver
}

// ScopeAttributeResolver is used to resolve attributes in scopes to one or more scopes that are
// evaluated by logical or. E.g. "dashboards:id:1" -> "dashboards:uid:test-dashboard" or "folder:uid:test-folder"
type ScopeAttributeResolver interface {
	Resolve(ctx context.Context, orgID int64, scope string) ([]string, error)
}

// ScopeAttributeResolverFunc is an adapter to allow functions to implement ScopeAttributeResolver interface
type ScopeAttributeResolverFunc func(ctx context.Context, orgID int64, scope string) ([]string, error)

func (f ScopeAttributeResolverFunc) Resolve(ctx context.Context, orgID int64, scope string) ([]string, error) {
	return f(ctx, orgID, scope)
}

type ScopeAttributeMutator func(context.Context, string) ([]string, error)

// ScopeKeywordResolver is used to resolve keywords in scopes e.g. "users:self" -> "user:id:1".
// These type of resolvers is used when fetching stored permissions
type ScopeKeywordResolver interface {
	Resolve(ctx context.Context, user *models.SignedInUser) (string, error)
}

// ScopeKeywordResolverFunc is an adapter to allow functions to implement ScopeKeywordResolver interface
type ScopeKeywordResolverFunc func(ctx context.Context, user *models.SignedInUser) (string, error)

func (f ScopeKeywordResolverFunc) Resolve(ctx context.Context, user *models.SignedInUser) (string, error) {
	return f(ctx, user)
}

type ScopeKeywordMutator func(context.Context, string) (string, error)

// getScopeCacheKey creates an identifier to fetch and store resolution of scopes in the cache
func getScopeCacheKey(orgID int64, scope string) string {
	return fmt.Sprintf("%s-%v", scope, orgID)
}

//ScopeInjector inject request params into the templated scopes. e.g. "settings:" + eval.Parameters(":id")
func ScopeInjector(params ScopeParams) ScopeAttributeMutator {
	return func(_ context.Context, scope string) ([]string, error) {
		tmpl, err := template.New("scope").Parse(scope)
		if err != nil {
			return nil, err
		}
		var buf bytes.Buffer
		if err = tmpl.Execute(&buf, params); err != nil {
			return nil, err
		}
		return []string{buf.String()}, nil
	}
}

var userSelfResolver = ScopeKeywordResolverFunc(func(ctx context.Context, user *models.SignedInUser) (string, error) {
	return Scope("users", "id", fmt.Sprintf("%v", user.UserId)), nil
})
<commit_msg>fix logging of resolving scopes (#50173)<commit_after>package accesscontrol

import (
	"bytes"
	"context"
	"fmt"
	"text/template"
	"time"

	"github.com/grafana/grafana/pkg/infra/localcache"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
)

const (
	ttl           = 30 * time.Second
	cleanInterval = 2 * time.Minute
)

func NewScopeResolvers() ScopeResolvers {
	return ScopeResolvers{
		keywordResolvers: map[string]ScopeKeywordResolver{
			"users:self": userSelfResolver,
		},
		attributeResolvers: map[string]ScopeAttributeResolver{},
		cache:              localcache.New(ttl, cleanInterval),
		log:                log.New("accesscontrol.resolver"),
	}
}

type ScopeResolvers struct {
	log                log.Logger
	cache              *localcache.CacheService
	keywordResolvers   map[string]ScopeKeywordResolver
	attributeResolvers map[string]ScopeAttributeResolver
}

func (s *ScopeResolvers) GetScopeAttributeMutator(orgID int64) ScopeAttributeMutator {
	return func(ctx context.Context, scope string) ([]string, error) {
		key := getScopeCacheKey(orgID, scope)
		// Check cache before computing the scope
		if cachedScope, ok := s.cache.Get(key); ok {
			scopes := cachedScope.([]string)
			s.log.Debug("used cache to resolve scope", "scope", scope, "resolved_scopes", scopes)
			return scopes, nil
		}

		prefix := ScopePrefix(scope)
		if resolver, ok := s.attributeResolvers[prefix]; ok {
			scopes, err := resolver.Resolve(ctx, orgID, scope)
			if err != nil {
				return nil, fmt.Errorf("could not resolve %v: %w", scope, err)
			}
			// Cache result
			s.cache.Set(key, scopes, ttl)
			s.log.Debug("resolved scope", "scope", scope, "resolved_scopes", scopes)
			return scopes, nil
		}
		return []string{scope}, nil
	}
}

func (s *ScopeResolvers) GetScopeKeywordMutator(user *models.SignedInUser) ScopeKeywordMutator {
	return func(ctx context.Context, scope string) (string, error) {
		if resolver, ok := s.keywordResolvers[scope]; ok {
			scopes, err := resolver.Resolve(ctx, user)
			if err != nil {
				return "", fmt.Errorf("could not resolve %v: %w", scope, err)
			}
			s.log.Debug("resolved scope", "scope", scope, "resolved_scopes", scopes)
			return scopes, nil
		}
		// By default, the scope remains unchanged
		return scope, nil
	}
}

func (s *ScopeResolvers) AddScopeKeywordResolver(keyword string, resolver ScopeKeywordResolver) {
	s.log.Debug("adding scope keyword resolver for '%v'", keyword)
	s.keywordResolvers[keyword] = resolver
}

func (s *ScopeResolvers) AddScopeAttributeResolver(prefix string, resolver ScopeAttributeResolver) {
	s.log.Debug("adding scope attribute resolver for '%v'", prefix)
	s.attributeResolvers[prefix] = resolver
}

// ScopeAttributeResolver is used to resolve attributes in scopes to one or more scopes that are
// evaluated by logical or. E.g. "dashboards:id:1" -> "dashboards:uid:test-dashboard" or "folder:uid:test-folder"
type ScopeAttributeResolver interface {
	Resolve(ctx context.Context, orgID int64, scope string) ([]string, error)
}

// ScopeAttributeResolverFunc is an adapter to allow functions to implement ScopeAttributeResolver interface
type ScopeAttributeResolverFunc func(ctx context.Context, orgID int64, scope string) ([]string, error)

func (f ScopeAttributeResolverFunc) Resolve(ctx context.Context, orgID int64, scope string) ([]string, error) {
	return f(ctx, orgID, scope)
}

type ScopeAttributeMutator func(context.Context, string) ([]string, error)

// ScopeKeywordResolver is used to resolve keywords in scopes e.g. "users:self" -> "user:id:1".
// These type of resolvers is used when fetching stored permissions
type ScopeKeywordResolver interface {
	Resolve(ctx context.Context, user *models.SignedInUser) (string, error)
}

// ScopeKeywordResolverFunc is an adapter to allow functions to implement ScopeKeywordResolver interface
type ScopeKeywordResolverFunc func(ctx context.Context, user *models.SignedInUser) (string, error)

func (f ScopeKeywordResolverFunc) Resolve(ctx context.Context, user *models.SignedInUser) (string, error) {
	return f(ctx, user)
}

type ScopeKeywordMutator func(context.Context, string) (string, error)

// getScopeCacheKey creates an identifier to fetch and store resolution of scopes in the cache
func getScopeCacheKey(orgID int64, scope string) string {
	return fmt.Sprintf("%s-%v", scope, orgID)
}

//ScopeInjector inject request params into the templated scopes. e.g. "settings:" + eval.Parameters(":id")
func ScopeInjector(params ScopeParams) ScopeAttributeMutator {
	return func(_ context.Context, scope string) ([]string, error) {
		tmpl, err := template.New("scope").Parse(scope)
		if err != nil {
			return nil, err
		}
		var buf bytes.Buffer
		if err = tmpl.Execute(&buf, params); err != nil {
			return nil, err
		}
		return []string{buf.String()}, nil
	}
}

var userSelfResolver = ScopeKeywordResolverFunc(func(ctx context.Context, user *models.SignedInUser) (string, error) {
	return Scope("users", "id", fmt.Sprintf("%v", user.UserId)), nil
})
<|endoftext|>


<commit_before>// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pai

type xgbPredictFiller struct {
	OSSModelDir      string
	DataSource       string
	PredSelect       string
	ResultTable      string
	HDFSNameNodeAddr string
	HiveLocation     string
	HDFSUser         string
	HDFSPass         string
	PAIPredictTable  string
}

const xgbPredTemplateText = `
import json
import sqlflow_submitter.xgboost as xgboost_extended
from sqlflow_submitter.xgboost.predict import pred
from sqlflow_submitter.pai import model
from sqlflow_submitter.tensorflow.pai_distributed import define_tf_flags, set_oss_environs

FLAGS = define_tf_flags()
set_oss_environs(FLAGS)

# NOTE(typhoonzero): the xgboost model file "my_model" is hard coded in xgboost/train.py
model.load_file("{{.OSSModelDir}}", "my_model")
(estimator,
model_params,
train_params,
feature_metas,
feature_column_names,
label_meta,
feature_column_code) = model.load_metas("{{.OSSModelDir}}", "xgboost_model_desc")

feature_column_transformers = eval('[{}]'.format(feature_column_code))
transform_fn = xgboost_extended.feature_column.ComposedColumnTransformer(feature_column_names, *feature_column_transformers)

pred(datasource='''{{.DataSource}}''',
    select='''{{.PredSelect}}''',
    feature_metas=feature_metas,
    feature_column_names=feature_column_names,
    label_meta=label_meta,
    result_table='''{{.ResultTable}}''',
    is_pai=True,
    hdfs_namenode_addr='''{{.HDFSNameNodeAddr}}''',
    hive_location='''{{.HiveLocation}}''',
    hdfs_user='''{{.HDFSUser}}''',
    hdfs_pass='''{{.HDFSPass}}''',
    pai_table='''{{.PAIPredictTable}}''',
    model_params=model_params,
    train_params=train_params,
    transform_fn=transform_fn,
    feature_column_code=feature_column_code)
`

type xgbExplainFiller struct {
	OSSModelDir       string
	DataSource        string
	DatasetSQL        string
	ResultTable       string
	IsPAI             bool
	PAIExplainTable   string
	HDFSNameNodeAddr  string
	HiveLocation      string
	HDFSUser          string
	HDFSPass          string
	ResultOSSDest     string
	ResultOSSAK       string
	ResultOSSSK       string
	ResultOSSEndpoint string
	ResultOSSBucket   string
}

const xgbExplainTemplateText = `
# Running on PAI
import os
import matplotlib
import sqlflow_submitter.xgboost as xgboost_extended

if os.environ.get('DISPLAY', '') == '':
    print('no display found. Using non-interactive Agg backend')
    matplotlib.use('Agg')

import json
from sqlflow_submitter.xgboost.explain import explain
from sqlflow_submitter.pai import model
from sqlflow_submitter.tensorflow.pai_distributed import define_tf_flags, set_oss_environs

FLAGS = define_tf_flags()
set_oss_environs(FLAGS)

# NOTE(typhoonzero): the xgboost model file "my_model" is hard coded in xgboost/train.py
model.load_file("{{.OSSModelDir}}", "my_model")

(estimator,
model_params,
train_params,
feature_field_meta,
feature_column_names,
label_field_meta,
feature_column_code) = model.load_metas("{{.OSSModelDir}}", "xgboost_model_desc")

feature_column_transformers = eval('[{}]'.format(feature_column_code))
transform_fn = xgboost_extended.feature_column.ComposedColumnTransformer(feature_column_names, *feature_column_transformers)

explain(
    datasource='''{{.DataSource}}''',
    select='''{{.DatasetSQL}}''',
	feature_field_meta=feature_field_meta,
	feature_column_names=feature_column_names,
	label_spec=label_field_meta,
	summary_params={},
	result_table='''{{.ResultTable}}''',
	is_pai="{{.IsPAI}}" == "true",
	pai_explain_table="{{.PAIExplainTable}}",
	hdfs_namenode_addr='''{{.HDFSNameNodeAddr}}''',
	hive_location='''{{.HiveLocation}}''',
	hdfs_user='''{{.HDFSUser}}''',
	hdfs_pass='''{{.HDFSPass}}''',
	oss_dest='''{{.ResultOSSDest}}''',
	oss_ak='''{{.ResultOSSAK}}''',
	oss_sk='''{{.ResultOSSSK}}''',
	oss_endpoint='''{{.ResultOSSEndpoint}}''',
	oss_bucket_name='''{{.ResultOSSBucket}}''',
	transform_fn=transform_fn,
	feature_column_code=feature_column_code)
`

type xgbEvaluateFiller struct {
	OSSModelDir      string
	DataSource       string
	PredSelect       string
	ResultTable      string
	MetricNames      string
	HDFSNameNodeAddr string
	HiveLocation     string
	HDFSUser         string
	HDFSPass         string
	PAIEvaluateTable string
}

const xgbEvalTemplateText = `
import json
import sqlflow_submitter.xgboost as xgboost_extended
from sqlflow_submitter.xgboost.evaluate import evaluate
from sqlflow_submitter.pai import model
from sqlflow_submitter.tensorflow.pai_distributed import define_tf_flags, set_oss_environs

FLAGS = define_tf_flags()
set_oss_environs(FLAGS)

# NOTE(typhoonzero): the xgboost model file "my_model" is hard coded in xgboost/train.py
model.load_file("{{.OSSModelDir}}", "my_model")
(estimator,
model_params,
train_params,
feature_metas,
feature_column_names,
label_meta,
feature_columns_code) = model.load_metas("{{.OSSModelDir}}", "xgboost_model_desc")

feature_column_transformers = eval('[{}]'.format(feature_column_code))
transform_fn = xgboost_extended.feature_column.ComposedColumnTransformer(feature_column_names, *feature_column_transformers)

evaluate(datasource='''{{.DataSource}}''',
         select='''{{.PredSelect}}''',
         feature_metas=feature_metas,
         feature_column_names=feature_column_names,
         label_meta=label_meta,
         result_table='''{{.ResultTable}}''',
         validation_metrics="{{.MetricNames}}".split(","),
         hdfs_namenode_addr='''{{.HDFSNameNodeAddr}}''',
         hive_location='''{{.HiveLocation}}''',
         hdfs_user='''{{.HDFSUser}}''',
         hdfs_pass='''{{.HDFSPass}}''',
         is_pai=True,
         pai_table="{{.PAIEvaluateTable}}",
         model_params=model_params,
         transform_fn=transform_fn,
         feature_column_code=feature_column_code)
`
<commit_msg>fix pai ci (#2450)<commit_after>// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pai

type xgbPredictFiller struct {
	OSSModelDir      string
	DataSource       string
	PredSelect       string
	ResultTable      string
	HDFSNameNodeAddr string
	HiveLocation     string
	HDFSUser         string
	HDFSPass         string
	PAIPredictTable  string
}

const xgbPredTemplateText = `
import json
import sqlflow_submitter.xgboost as xgboost_extended
from sqlflow_submitter.xgboost.predict import pred
from sqlflow_submitter.pai import model
from sqlflow_submitter.tensorflow.pai_distributed import define_tf_flags, set_oss_environs

FLAGS = define_tf_flags()
set_oss_environs(FLAGS)

# NOTE(typhoonzero): the xgboost model file "my_model" is hard coded in xgboost/train.py
model.load_file("{{.OSSModelDir}}", "my_model")
(estimator,
model_params,
train_params,
feature_metas,
feature_column_names,
label_meta,
feature_column_code) = model.load_metas("{{.OSSModelDir}}", "xgboost_model_desc")

feature_column_transformers = eval('[{}]'.format(feature_column_code))
transform_fn = xgboost_extended.feature_column.ComposedColumnTransformer(feature_column_names, *feature_column_transformers)

pred(datasource='''{{.DataSource}}''',
    select='''{{.PredSelect}}''',
    feature_metas=feature_metas,
    feature_column_names=feature_column_names,
    label_meta=label_meta,
    result_table='''{{.ResultTable}}''',
    is_pai=True,
    hdfs_namenode_addr='''{{.HDFSNameNodeAddr}}''',
    hive_location='''{{.HiveLocation}}''',
    hdfs_user='''{{.HDFSUser}}''',
    hdfs_pass='''{{.HDFSPass}}''',
    pai_table='''{{.PAIPredictTable}}''',
    model_params=model_params,
    train_params=train_params,
    transform_fn=transform_fn,
    feature_column_code=feature_column_code)
`

type xgbExplainFiller struct {
	OSSModelDir       string
	DataSource        string
	DatasetSQL        string
	ResultTable       string
	IsPAI             bool
	PAIExplainTable   string
	HDFSNameNodeAddr  string
	HiveLocation      string
	HDFSUser          string
	HDFSPass          string
	ResultOSSDest     string
	ResultOSSAK       string
	ResultOSSSK       string
	ResultOSSEndpoint string
	ResultOSSBucket   string
}

const xgbExplainTemplateText = `
# Running on PAI
import os
import matplotlib
import sqlflow_submitter.xgboost as xgboost_extended

if os.environ.get('DISPLAY', '') == '':
    print('no display found. Using non-interactive Agg backend')
    matplotlib.use('Agg')

import json
from sqlflow_submitter.xgboost.explain import explain
from sqlflow_submitter.pai import model
from sqlflow_submitter.tensorflow.pai_distributed import define_tf_flags, set_oss_environs

FLAGS = define_tf_flags()
set_oss_environs(FLAGS)

# NOTE(typhoonzero): the xgboost model file "my_model" is hard coded in xgboost/train.py
model.load_file("{{.OSSModelDir}}", "my_model")

(estimator,
model_params,
train_params,
feature_field_meta,
feature_column_names,
label_field_meta,
feature_column_code) = model.load_metas("{{.OSSModelDir}}", "xgboost_model_desc")

feature_column_transformers = eval('[{}]'.format(feature_column_code))
transform_fn = xgboost_extended.feature_column.ComposedColumnTransformer(feature_column_names, *feature_column_transformers)

explain(
    datasource='''{{.DataSource}}''',
    select='''{{.DatasetSQL}}''',
	feature_field_meta=feature_field_meta,
	feature_column_names=feature_column_names,
	label_spec=label_field_meta,
	summary_params={},
	result_table='''{{.ResultTable}}''',
	is_pai="{{.IsPAI}}" == "true",
	pai_explain_table="{{.PAIExplainTable}}",
	hdfs_namenode_addr='''{{.HDFSNameNodeAddr}}''',
	hive_location='''{{.HiveLocation}}''',
	hdfs_user='''{{.HDFSUser}}''',
	hdfs_pass='''{{.HDFSPass}}''',
	oss_dest='''{{.ResultOSSDest}}''',
	oss_ak='''{{.ResultOSSAK}}''',
	oss_sk='''{{.ResultOSSSK}}''',
	oss_endpoint='''{{.ResultOSSEndpoint}}''',
	oss_bucket_name='''{{.ResultOSSBucket}}''',
	transform_fn=transform_fn,
	feature_column_code=feature_column_code)
`

type xgbEvaluateFiller struct {
	OSSModelDir      string
	DataSource       string
	PredSelect       string
	ResultTable      string
	MetricNames      string
	HDFSNameNodeAddr string
	HiveLocation     string
	HDFSUser         string
	HDFSPass         string
	PAIEvaluateTable string
}

const xgbEvalTemplateText = `
import json
import sqlflow_submitter.xgboost as xgboost_extended
from sqlflow_submitter.xgboost.evaluate import evaluate
from sqlflow_submitter.pai import model
from sqlflow_submitter.tensorflow.pai_distributed import define_tf_flags, set_oss_environs

FLAGS = define_tf_flags()
set_oss_environs(FLAGS)

# NOTE(typhoonzero): the xgboost model file "my_model" is hard coded in xgboost/train.py
model.load_file("{{.OSSModelDir}}", "my_model")
(estimator,
model_params,
train_params,
feature_metas,
feature_column_names,
label_meta,
feature_column_code) = model.load_metas("{{.OSSModelDir}}", "xgboost_model_desc")

feature_column_transformers = eval('[{}]'.format(feature_column_code))
transform_fn = xgboost_extended.feature_column.ComposedColumnTransformer(feature_column_names, *feature_column_transformers)

evaluate(datasource='''{{.DataSource}}''',
         select='''{{.PredSelect}}''',
         feature_metas=feature_metas,
         feature_column_names=feature_column_names,
         label_meta=label_meta,
         result_table='''{{.ResultTable}}''',
         validation_metrics="{{.MetricNames}}".split(","),
         hdfs_namenode_addr='''{{.HDFSNameNodeAddr}}''',
         hive_location='''{{.HiveLocation}}''',
         hdfs_user='''{{.HDFSUser}}''',
         hdfs_pass='''{{.HDFSPass}}''',
         is_pai=True,
         pai_table="{{.PAIEvaluateTable}}",
         model_params=model_params,
         transform_fn=transform_fn,
         feature_column_code=feature_column_code)
`
<|endoftext|>


<commit_before>package es

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/Masterminds/semver"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/httpclient"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/tsdb/interval"

	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/plugins"
	"golang.org/x/net/context/ctxhttp"
)

const loggerName = "tsdb.elasticsearch.client"

var (
	clientLog = log.New(loggerName)
)

var newDatasourceHttpClient = func(httpClientProvider httpclient.Provider, ds *models.DataSource) (*http.Client, error) {
	return ds.GetHTTPClient(httpClientProvider)
}

// Client represents a client which can interact with elasticsearch api
type Client interface {
	GetVersion() *semver.Version
	GetTimeField() string
	GetMinInterval(queryInterval string) (time.Duration, error)
	ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error)
	MultiSearch() *MultiSearchRequestBuilder
	EnableDebug()
}

func coerceVersion(v *simplejson.Json) (*semver.Version, error) {
	versionString, err := v.String()

	if err != nil {
		versionNumber, err := v.Int()
		if err != nil {
			return nil, err
		}

		switch versionNumber {
		case 2:
			return semver.NewVersion("2.0.0")
		case 5:
			return semver.NewVersion("5.0.0")
		case 56:
			return semver.NewVersion("5.6.0")
		case 60:
			return semver.NewVersion("6.0.0")
		case 70:
			return semver.NewVersion("7.0.0")
		default:
			return nil, fmt.Errorf("elasticsearch version=%d is not supported", versionNumber)
		}
	}

	return semver.NewVersion(versionString)
}

// NewClient creates a new elasticsearch client
var NewClient = func(ctx context.Context, httpClientProvider httpclient.Provider, ds *models.DataSource, timeRange plugins.DataTimeRange) (Client, error) {
	version, err := coerceVersion(ds.JsonData.Get("esVersion"))

	if err != nil {
		return nil, fmt.Errorf("elasticsearch version is required, err=%v", err)
	}

	timeField, err := ds.JsonData.Get("timeField").String()
	if err != nil {
		return nil, fmt.Errorf("elasticsearch time field name is required, err=%v", err)
	}

	indexInterval := ds.JsonData.Get("interval").MustString()
	ip, err := newIndexPattern(indexInterval, ds.Database)
	if err != nil {
		return nil, err
	}

	indices, err := ip.GetIndices(timeRange)
	if err != nil {
		return nil, err
	}

	clientLog.Info("Creating new client", "version", version.String(), "timeField", timeField, "indices", strings.Join(indices, ", "))

	return &baseClientImpl{
		ctx:       ctx,
		ds:        ds,
		version:   version,
		timeField: timeField,
		indices:   indices,
		timeRange: timeRange,
	}, nil
}

type baseClientImpl struct {
	ctx                context.Context
	httpClientProvider httpclient.Provider
	ds                 *models.DataSource
	version            *semver.Version
	timeField          string
	indices            []string
	timeRange          plugins.DataTimeRange
	debugEnabled       bool
}

func (c *baseClientImpl) GetVersion() *semver.Version {
	return c.version
}

func (c *baseClientImpl) GetTimeField() string {
	return c.timeField
}

func (c *baseClientImpl) GetMinInterval(queryInterval string) (time.Duration, error) {
	return interval.GetIntervalFrom(c.ds, simplejson.NewFromAny(map[string]interface{}{
		"interval": queryInterval,
	}), 5*time.Second)
}

func (c *baseClientImpl) getSettings() *simplejson.Json {
	return c.ds.JsonData
}

type multiRequest struct {
	header   map[string]interface{}
	body     interface{}
	interval interval.Interval
}

func (c *baseClientImpl) executeBatchRequest(uriPath, uriQuery string, requests []*multiRequest) (*response, error) {
	bytes, err := c.encodeBatchRequests(requests)
	if err != nil {
		return nil, err
	}
	return c.executeRequest(http.MethodPost, uriPath, uriQuery, bytes)
}

func (c *baseClientImpl) encodeBatchRequests(requests []*multiRequest) ([]byte, error) {
	clientLog.Debug("Encoding batch requests to json", "batch requests", len(requests))
	start := time.Now()

	payload := bytes.Buffer{}
	for _, r := range requests {
		reqHeader, err := json.Marshal(r.header)
		if err != nil {
			return nil, err
		}
		payload.WriteString(string(reqHeader) + "\n")

		reqBody, err := json.Marshal(r.body)
		if err != nil {
			return nil, err
		}

		body := string(reqBody)
		body = strings.ReplaceAll(body, "$__interval_ms", strconv.FormatInt(r.interval.Milliseconds(), 10))
		body = strings.ReplaceAll(body, "$__interval", r.interval.Text)

		payload.WriteString(body + "\n")
	}

	elapsed := time.Since(start)
	clientLog.Debug("Encoded batch requests to json", "took", elapsed)

	return payload.Bytes(), nil
}

func (c *baseClientImpl) executeRequest(method, uriPath, uriQuery string, body []byte) (*response, error) {
	u, err := url.Parse(c.ds.Url)
	if err != nil {
		return nil, err
	}
	u.Path = path.Join(u.Path, uriPath)
	u.RawQuery = uriQuery

	var req *http.Request
	if method == http.MethodPost {
		req, err = http.NewRequest(http.MethodPost, u.String(), bytes.NewBuffer(body))
	} else {
		req, err = http.NewRequest(http.MethodGet, u.String(), nil)
	}
	if err != nil {
		return nil, err
	}

	clientLog.Debug("Executing request", "url", req.URL.String(), "method", method)

	var reqInfo *SearchRequestInfo
	if c.debugEnabled {
		reqInfo = &SearchRequestInfo{
			Method: req.Method,
			Url:    req.URL.String(),
			Data:   string(body),
		}
	}

	req.Header.Set("Content-Type", "application/x-ndjson")

	httpClient, err := newDatasourceHttpClient(c.httpClientProvider, c.ds)
	if err != nil {
		return nil, err
	}

	start := time.Now()
	defer func() {
		elapsed := time.Since(start)
		clientLog.Debug("Executed request", "took", elapsed)
	}()
	//nolint:bodyclose
	resp, err := ctxhttp.Do(c.ctx, httpClient, req)
	if err != nil {
		return nil, err
	}
	return &response{
		httpResponse: resp,
		reqInfo:      reqInfo,
	}, nil
}

func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error) {
	clientLog.Debug("Executing multisearch", "search requests", len(r.Requests))

	multiRequests := c.createMultiSearchRequests(r.Requests)
	queryParams := c.getMultiSearchQueryParameters()
	clientRes, err := c.executeBatchRequest("_msearch", queryParams, multiRequests)
	if err != nil {
		return nil, err
	}
	res := clientRes.httpResponse
	defer func() {
		if err := res.Body.Close(); err != nil {
			clientLog.Warn("Failed to close response body", "err", err)
		}
	}()

	clientLog.Debug("Received multisearch response", "code", res.StatusCode, "status", res.Status, "content-length", res.ContentLength)

	start := time.Now()
	clientLog.Debug("Decoding multisearch json response")

	var bodyBytes []byte
	if c.debugEnabled {
		tmpBytes, err := ioutil.ReadAll(res.Body)
		if err != nil {
			clientLog.Error("failed to read http response bytes", "error", err)
		} else {
			bodyBytes = make([]byte, len(tmpBytes))
			copy(bodyBytes, tmpBytes)
			res.Body = ioutil.NopCloser(bytes.NewBuffer(tmpBytes))
		}
	}

	var msr MultiSearchResponse
	dec := json.NewDecoder(res.Body)
	err = dec.Decode(&msr)
	if err != nil {
		return nil, err
	}

	elapsed := time.Since(start)
	clientLog.Debug("Decoded multisearch json response", "took", elapsed)

	msr.Status = res.StatusCode

	if c.debugEnabled {
		bodyJSON, err := simplejson.NewFromReader(bytes.NewBuffer(bodyBytes))
		var data *simplejson.Json
		if err != nil {
			clientLog.Error("failed to decode http response into json", "error", err)
		} else {
			data = bodyJSON
		}

		msr.DebugInfo = &SearchDebugInfo{
			Request: clientRes.reqInfo,
			Response: &SearchResponseInfo{
				Status: res.StatusCode,
				Data:   data,
			},
		}
	}

	return &msr, nil
}

func (c *baseClientImpl) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := []*multiRequest{}

	for _, searchReq := range searchRequests {
		mr := multiRequest{
			header: map[string]interface{}{
				"search_type":        "query_then_fetch",
				"ignore_unavailable": true,
				"index":              strings.Join(c.indices, ","),
			},
			body:     searchReq,
			interval: searchReq.Interval,
		}

		if c.version.Major() < 5 {
			mr.header["search_type"] = "count"
		} else {
			allowedVersionRange, _ := semver.NewConstraint(">=5.6.0, <7.0.0")

			if allowedVersionRange.Check(c.version) {
				maxConcurrentShardRequests := c.getSettings().Get("maxConcurrentShardRequests").MustInt(256)
				mr.header["max_concurrent_shard_requests"] = maxConcurrentShardRequests
			}
		}

		multiRequests = append(multiRequests, &mr)
	}

	return multiRequests
}

func (c *baseClientImpl) getMultiSearchQueryParameters() string {
	if c.version.Major() >= 7 {
		maxConcurrentShardRequests := c.getSettings().Get("maxConcurrentShardRequests").MustInt(5)
		return fmt.Sprintf("max_concurrent_shard_requests=%d", maxConcurrentShardRequests)
	}

	return ""
}

func (c *baseClientImpl) MultiSearch() *MultiSearchRequestBuilder {
	return NewMultiSearchRequestBuilder(c.GetVersion())
}

func (c *baseClientImpl) EnableDebug() {
	c.debugEnabled = true
}
<commit_msg>Elasticsearch: fix NewClient not passing httpClientProvider to client impl (#34539)<commit_after>package es

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/Masterminds/semver"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/httpclient"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/tsdb/interval"

	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/plugins"
	"golang.org/x/net/context/ctxhttp"
)

const loggerName = "tsdb.elasticsearch.client"

var (
	clientLog = log.New(loggerName)
)

var newDatasourceHttpClient = func(httpClientProvider httpclient.Provider, ds *models.DataSource) (*http.Client, error) {
	return ds.GetHTTPClient(httpClientProvider)
}

// Client represents a client which can interact with elasticsearch api
type Client interface {
	GetVersion() *semver.Version
	GetTimeField() string
	GetMinInterval(queryInterval string) (time.Duration, error)
	ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error)
	MultiSearch() *MultiSearchRequestBuilder
	EnableDebug()
}

func coerceVersion(v *simplejson.Json) (*semver.Version, error) {
	versionString, err := v.String()

	if err != nil {
		versionNumber, err := v.Int()
		if err != nil {
			return nil, err
		}

		switch versionNumber {
		case 2:
			return semver.NewVersion("2.0.0")
		case 5:
			return semver.NewVersion("5.0.0")
		case 56:
			return semver.NewVersion("5.6.0")
		case 60:
			return semver.NewVersion("6.0.0")
		case 70:
			return semver.NewVersion("7.0.0")
		default:
			return nil, fmt.Errorf("elasticsearch version=%d is not supported", versionNumber)
		}
	}

	return semver.NewVersion(versionString)
}

// NewClient creates a new elasticsearch client
var NewClient = func(ctx context.Context, httpClientProvider httpclient.Provider, ds *models.DataSource, timeRange plugins.DataTimeRange) (Client, error) {
	version, err := coerceVersion(ds.JsonData.Get("esVersion"))

	if err != nil {
		return nil, fmt.Errorf("elasticsearch version is required, err=%v", err)
	}

	timeField, err := ds.JsonData.Get("timeField").String()
	if err != nil {
		return nil, fmt.Errorf("elasticsearch time field name is required, err=%v", err)
	}

	indexInterval := ds.JsonData.Get("interval").MustString()
	ip, err := newIndexPattern(indexInterval, ds.Database)
	if err != nil {
		return nil, err
	}

	indices, err := ip.GetIndices(timeRange)
	if err != nil {
		return nil, err
	}

	clientLog.Info("Creating new client", "version", version.String(), "timeField", timeField, "indices", strings.Join(indices, ", "))

	return &baseClientImpl{
		ctx:                ctx,
		httpClientProvider: httpClientProvider,
		ds:                 ds,
		version:            version,
		timeField:          timeField,
		indices:            indices,
		timeRange:          timeRange,
	}, nil
}

type baseClientImpl struct {
	ctx                context.Context
	httpClientProvider httpclient.Provider
	ds                 *models.DataSource
	version            *semver.Version
	timeField          string
	indices            []string
	timeRange          plugins.DataTimeRange
	debugEnabled       bool
}

func (c *baseClientImpl) GetVersion() *semver.Version {
	return c.version
}

func (c *baseClientImpl) GetTimeField() string {
	return c.timeField
}

func (c *baseClientImpl) GetMinInterval(queryInterval string) (time.Duration, error) {
	return interval.GetIntervalFrom(c.ds, simplejson.NewFromAny(map[string]interface{}{
		"interval": queryInterval,
	}), 5*time.Second)
}

func (c *baseClientImpl) getSettings() *simplejson.Json {
	return c.ds.JsonData
}

type multiRequest struct {
	header   map[string]interface{}
	body     interface{}
	interval interval.Interval
}

func (c *baseClientImpl) executeBatchRequest(uriPath, uriQuery string, requests []*multiRequest) (*response, error) {
	bytes, err := c.encodeBatchRequests(requests)
	if err != nil {
		return nil, err
	}
	return c.executeRequest(http.MethodPost, uriPath, uriQuery, bytes)
}

func (c *baseClientImpl) encodeBatchRequests(requests []*multiRequest) ([]byte, error) {
	clientLog.Debug("Encoding batch requests to json", "batch requests", len(requests))
	start := time.Now()

	payload := bytes.Buffer{}
	for _, r := range requests {
		reqHeader, err := json.Marshal(r.header)
		if err != nil {
			return nil, err
		}
		payload.WriteString(string(reqHeader) + "\n")

		reqBody, err := json.Marshal(r.body)
		if err != nil {
			return nil, err
		}

		body := string(reqBody)
		body = strings.ReplaceAll(body, "$__interval_ms", strconv.FormatInt(r.interval.Milliseconds(), 10))
		body = strings.ReplaceAll(body, "$__interval", r.interval.Text)

		payload.WriteString(body + "\n")
	}

	elapsed := time.Since(start)
	clientLog.Debug("Encoded batch requests to json", "took", elapsed)

	return payload.Bytes(), nil
}

func (c *baseClientImpl) executeRequest(method, uriPath, uriQuery string, body []byte) (*response, error) {
	u, err := url.Parse(c.ds.Url)
	if err != nil {
		return nil, err
	}
	u.Path = path.Join(u.Path, uriPath)
	u.RawQuery = uriQuery

	var req *http.Request
	if method == http.MethodPost {
		req, err = http.NewRequest(http.MethodPost, u.String(), bytes.NewBuffer(body))
	} else {
		req, err = http.NewRequest(http.MethodGet, u.String(), nil)
	}
	if err != nil {
		return nil, err
	}

	clientLog.Debug("Executing request", "url", req.URL.String(), "method", method)

	var reqInfo *SearchRequestInfo
	if c.debugEnabled {
		reqInfo = &SearchRequestInfo{
			Method: req.Method,
			Url:    req.URL.String(),
			Data:   string(body),
		}
	}

	req.Header.Set("Content-Type", "application/x-ndjson")

	httpClient, err := newDatasourceHttpClient(c.httpClientProvider, c.ds)
	if err != nil {
		return nil, err
	}

	start := time.Now()
	defer func() {
		elapsed := time.Since(start)
		clientLog.Debug("Executed request", "took", elapsed)
	}()
	//nolint:bodyclose
	resp, err := ctxhttp.Do(c.ctx, httpClient, req)
	if err != nil {
		return nil, err
	}
	return &response{
		httpResponse: resp,
		reqInfo:      reqInfo,
	}, nil
}

func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error) {
	clientLog.Debug("Executing multisearch", "search requests", len(r.Requests))

	multiRequests := c.createMultiSearchRequests(r.Requests)
	queryParams := c.getMultiSearchQueryParameters()
	clientRes, err := c.executeBatchRequest("_msearch", queryParams, multiRequests)
	if err != nil {
		return nil, err
	}
	res := clientRes.httpResponse
	defer func() {
		if err := res.Body.Close(); err != nil {
			clientLog.Warn("Failed to close response body", "err", err)
		}
	}()

	clientLog.Debug("Received multisearch response", "code", res.StatusCode, "status", res.Status, "content-length", res.ContentLength)

	start := time.Now()
	clientLog.Debug("Decoding multisearch json response")

	var bodyBytes []byte
	if c.debugEnabled {
		tmpBytes, err := ioutil.ReadAll(res.Body)
		if err != nil {
			clientLog.Error("failed to read http response bytes", "error", err)
		} else {
			bodyBytes = make([]byte, len(tmpBytes))
			copy(bodyBytes, tmpBytes)
			res.Body = ioutil.NopCloser(bytes.NewBuffer(tmpBytes))
		}
	}

	var msr MultiSearchResponse
	dec := json.NewDecoder(res.Body)
	err = dec.Decode(&msr)
	if err != nil {
		return nil, err
	}

	elapsed := time.Since(start)
	clientLog.Debug("Decoded multisearch json response", "took", elapsed)

	msr.Status = res.StatusCode

	if c.debugEnabled {
		bodyJSON, err := simplejson.NewFromReader(bytes.NewBuffer(bodyBytes))
		var data *simplejson.Json
		if err != nil {
			clientLog.Error("failed to decode http response into json", "error", err)
		} else {
			data = bodyJSON
		}

		msr.DebugInfo = &SearchDebugInfo{
			Request: clientRes.reqInfo,
			Response: &SearchResponseInfo{
				Status: res.StatusCode,
				Data:   data,
			},
		}
	}

	return &msr, nil
}

func (c *baseClientImpl) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := []*multiRequest{}

	for _, searchReq := range searchRequests {
		mr := multiRequest{
			header: map[string]interface{}{
				"search_type":        "query_then_fetch",
				"ignore_unavailable": true,
				"index":              strings.Join(c.indices, ","),
			},
			body:     searchReq,
			interval: searchReq.Interval,
		}

		if c.version.Major() < 5 {
			mr.header["search_type"] = "count"
		} else {
			allowedVersionRange, _ := semver.NewConstraint(">=5.6.0, <7.0.0")

			if allowedVersionRange.Check(c.version) {
				maxConcurrentShardRequests := c.getSettings().Get("maxConcurrentShardRequests").MustInt(256)
				mr.header["max_concurrent_shard_requests"] = maxConcurrentShardRequests
			}
		}

		multiRequests = append(multiRequests, &mr)
	}

	return multiRequests
}

func (c *baseClientImpl) getMultiSearchQueryParameters() string {
	if c.version.Major() >= 7 {
		maxConcurrentShardRequests := c.getSettings().Get("maxConcurrentShardRequests").MustInt(5)
		return fmt.Sprintf("max_concurrent_shard_requests=%d", maxConcurrentShardRequests)
	}

	return ""
}

func (c *baseClientImpl) MultiSearch() *MultiSearchRequestBuilder {
	return NewMultiSearchRequestBuilder(c.GetVersion())
}

func (c *baseClientImpl) EnableDebug() {
	c.debugEnabled = true
}
<|endoftext|>


<commit_before>package tailor

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"

	"github.com/Masterminds/squirrel"
	"github.com/antihax/eve-axiom/attributes"
	"github.com/antihax/evedata/internal/datapackages"
	"github.com/antihax/evedata/internal/gobcoder"
	"github.com/antihax/goesi/esi"

	nsq "github.com/nsqio/go-nsq"
)

// KillmailAttributes provides a package of killmail, attributes, and id->name lookup
type KillmailAttributes struct {
	Attributes *attributes.Attributes                    `json:"attributes"`
	Killmail   *esi.GetKillmailsKillmailIdKillmailHashOk `json:"killmail"`
	NameMap    map[int32]string                          `json:"nameMap"`
}

var chanKillmailAttributes chan KillmailAttributes

func init() {
	chanKillmailAttributes = make(chan KillmailAttributes, 100)
}

func (s *Tailor) saveKillmail(pack *KillmailAttributes) error {
	j, err := json.MarshalIndent(pack, " ", "   ")
	if err != nil {
		log.Println(err)
		return err
	}
	fmt.Printf("%s\n", j)
	return nil
}

func (s *Tailor) killmailHandler(message *nsq.Message) error {
	killmail := datapackages.Killmail{}
	if err := gobcoder.GobDecoder(message.Body, &killmail); err != nil {
		log.Println(err)
		return err
	}

	attr, err := getAttributesForKillmail(&killmail.Kill)
	if err != nil {
		log.Println(err)
		return err
	}

	names, err := s.resolveNames(&killmail.Kill)
	if err != nil {
		log.Println(err)
		return err
	}
	pack := KillmailAttributes{
		Attributes: attr,
		Killmail:   &killmail.Kill,
		NameMap:    names,
	}
	s.saveKillmail(&pack)

	// Add the package to the list
	chanKillmailAttributes <- pack
	return nil
}

func joinInt32(a []int32) string {
	return strings.Trim(strings.Join(strings.Fields(fmt.Sprint(a)), ","), "[]")
}

func (s *Tailor) resolveNames(kill *esi.GetKillmailsKillmailIdKillmailHashOk) (map[int32]string, error) {
	// make a unique list of ids
	ids := make(map[int32]string)
	for _, a := range kill.Attackers {
		ids[a.AllianceId] = ""
		ids[a.CorporationId] = ""
		ids[a.CharacterId] = ""
		ids[a.ShipTypeId] = ""
		ids[a.WeaponTypeId] = ""
	}
	for _, a := range kill.Victim.Items {
		ids[a.ItemTypeId] = ""
		for _, a := range a.Items {
			ids[a.ItemTypeId] = ""
		}
	}
	ids[kill.Victim.AllianceId] = ""
	ids[kill.Victim.CorporationId] = ""
	ids[kill.Victim.CharacterId] = ""
	ids[kill.Victim.ShipTypeId] = ""

	// Delete 0 if we picked up no alliance
	delete(ids, 0)

	idList := []int32{}
	for id := range ids {
		if id > 0 {
			idList = append(idList, id)
		}
	}

	// Lookup the information
	rows, err := s.db.Query(`
		SELECT allianceID as id, name FROM evedata.alliances WHERE allianceID IN (` + joinInt32(idList) + `)
		UNION
		SELECT corporationID as id, name FROM evedata.corporations WHERE corporationID IN (` + joinInt32(idList) + `)
		UNION
		SELECT characterID as id, name FROM evedata.characters WHERE characterID IN (` + joinInt32(idList) + `)
		UNION
		SELECT typeID as id, typeName as name FROM eve.invTypes WHERE typeID IN  (` + joinInt32(idList) + `)
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var (
			id   int32
			name string
		)
		err := rows.Scan(&id, &name)
		if err != nil {
			return nil, err
		}
		ids[id] = name
	}

	// Double check we are not missing any
	for id, n := range ids {
		if id > 0 {
			if n == "" {
				return nil, fmt.Errorf("cannot find name for %d", id)
			}
		}
	}

	return ids, nil
}

// killmailConsumer receives killmails from NSQ and dumps attributes to SQL
func (s *Tailor) killmailConsumer() {
	for {
		a := <-chanKillmailAttributes
		b := a.Attributes.Ship

		sq := squirrel.Insert("evedata.killmailAttributes").Columns(
			"id", "eHP", "DPS", "Alpha", "scanResolution", "signatureRadiusNoMWD",
			"signatureRadius", "agility", "warpSpeed", "speedNoMWD", "speed", "remoteArmorRepair",
			"remoteShieldRepair", "remoteEnergyTransfer", "energyNeutralization", "sensorStrength",
			"RPS", "CPURemaining", "powerRemaining", "capacitorNoMWD", "capacitor",
			"capacitorTimeNoMWD", "capacitorTime",
		)

		sq = sq.Values(
			a.Killmail.KillmailId, b["avgEHP"], b["totalDPS"], b["totalAlphaDamage"], b["scanResolution"], b["signatureRadius"],
			b["signatureRadiusMWD"], b["agility"], b["warpSpeedMultiplier"], b["MaxVelocity"], b["MaxVelocityMWD"], b["remoteArmorRepairPerSecond"],
			b["remoteShieldBonusAmountPerSecond"], b["remotePowerTransferAmountPerSecond"], b["energyNeutralizerAmountPerSecond"],
			b["scanRadarStrength"]+b["scanLadarStrength"]+b["scanMagnetometricStrength"]+b["scanGravimetricStrength"],
			b["avgRPS"], b["cpuRemaining"], b["powerRemaining"], b["capacitorFraction"], b["capacitorFractionMWD"],
			b["capacitorDuration"]*1000000, b["capacitorDurationMWD"]*1000000,
		)

		sqlq, args, err := sq.ToSql()
		if err != nil {
			log.Println(err)
			continue
		}

		err = s.doSQL(sqlq+" ON DUPLICATE KEY UPDATE id = id", args...)
		if err != nil {
			log.Println(err)
			continue
		}

	}
}

func getAttributesForKillmail(km *esi.GetKillmailsKillmailIdKillmailHashOk) (*attributes.Attributes, error) {
	j, err := json.Marshal(km)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", "http://axiom.evedata:3005/killmail", bytes.NewBuffer(j))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	dec := json.NewDecoder(resp.Body)
	attributes := &attributes.Attributes{}
	if err = dec.Decode(attributes); err != nil {
		return nil, err
	}

	return attributes, nil
}
<commit_msg>add system info to killmails<commit_after>package tailor

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"

	"github.com/Masterminds/squirrel"
	"github.com/antihax/eve-axiom/attributes"
	"github.com/antihax/evedata/internal/datapackages"
	"github.com/antihax/evedata/internal/gobcoder"
	"github.com/antihax/goesi/esi"

	nsq "github.com/nsqio/go-nsq"
)

// KillmailAttributes provides a package of killmail, attributes, and id->name lookup
type KillmailAttributes struct {
	Attributes *attributes.Attributes                    `json:"attributes"`
	Killmail   *esi.GetKillmailsKillmailIdKillmailHashOk `json:"killmail"`
	NameMap    map[int32]string                          `json:"nameMap"`
	SystemInfo SystemInformation                         `json:"systemInfo"`
}

// SystemInformation for the killmail
type SystemInformation struct {
	CelestialID     int32  `db:"celestialID" json:"celestialID"`
	CelestialName   string `db:"celestialName" json:"celestialName"`
	RegionID        int32  `db:"regionID" json:"regionID"`
	RegionName      string `db:"regionName" json:"regionName"`
	SolarSystemName string `db:"solarSystemName" json:"solarSystemName"`
	SolarSystemID   int32  `db:"solarSystemID" json:"solarSystemID"`
}

var chanKillmailAttributes chan KillmailAttributes

func init() {
	chanKillmailAttributes = make(chan KillmailAttributes, 100)
}

func (s *Tailor) saveKillmail(pack *KillmailAttributes) error {
	j, err := json.MarshalIndent(pack, " ", "   ")
	if err != nil {
		log.Println(err)
		return err
	}
	return nil
}

func (s *Tailor) killmailHandler(message *nsq.Message) error {
	killmail := datapackages.Killmail{}
	if err := gobcoder.GobDecoder(message.Body, &killmail); err != nil {
		log.Println(err)
		return err
	}

	attr, err := getAttributesForKillmail(&killmail.Kill)
	if err != nil {
		log.Println(err)
		return err
	}

	names, err := s.resolveNames(&killmail.Kill)
	if err != nil {
		log.Println(err)
		return err
	}

	pos := killmail.Kill.Victim.Position
	sysinfo, err := s.getSystemInformation(killmail.Kill.SolarSystemId, pos.X, pos.Y, pos.Z)
	if err != nil {
		log.Println(err)
		return err
	}
	pack := KillmailAttributes{
		Attributes: attr,
		Killmail:   &killmail.Kill,
		NameMap:    names,
		SystemInfo: sysinfo,
	}
	s.saveKillmail(&pack)

	// Add the package to the list
	chanKillmailAttributes <- pack
	return nil
}

func joinInt32(a []int32) string {
	return strings.Trim(strings.Join(strings.Fields(fmt.Sprint(a)), ","), "[]")
}

func (s *Tailor) getSystemInformation(system int32, x, y, z float64) (SystemInformation, error) {
	sysInfo := []SystemInformation{}
	err := s.db.Select(&sysInfo,
		`SELECT itemName AS celestialName, itemID AS celestialID, S.solarSystemID, S.regionID, solarSystemName, regionName
		FROM  eve.mapDenormalize D
		INNER JOIN eve.mapSolarSystems S ON S.solarSystemID = D.solarSystemID
		INNER JOIN eve.mapRegions R ON R.regionID = D.regionID
		WHERE itemID = closestCelestial(?,?,?,?);`, system, x, y, z)
	if err != nil {
		return SystemInformation{}, err
	}
	return sysInfo[0], err
}

func (s *Tailor) resolveNames(kill *esi.GetKillmailsKillmailIdKillmailHashOk) (map[int32]string, error) {
	// make a unique list of ids
	ids := make(map[int32]string)
	for _, a := range kill.Attackers {
		ids[a.AllianceId] = ""
		ids[a.CorporationId] = ""
		ids[a.CharacterId] = ""
		ids[a.ShipTypeId] = ""
		ids[a.WeaponTypeId] = ""
	}
	for _, a := range kill.Victim.Items {
		ids[a.ItemTypeId] = ""
		for _, a := range a.Items {
			ids[a.ItemTypeId] = ""
		}
	}

	ids[kill.Victim.AllianceId] = ""
	ids[kill.Victim.CorporationId] = ""
	ids[kill.Victim.CharacterId] = ""
	ids[kill.Victim.ShipTypeId] = ""

	// Delete 0 if we picked up no alliance
	delete(ids, 0)

	idList := []int32{}
	for id := range ids {
		if id > 0 {
			idList = append(idList, id)
		}
	}

	// Lookup the information
	rows, err := s.db.Query(`
		SELECT allianceID as id, name FROM evedata.alliances WHERE allianceID IN (` + joinInt32(idList) + `)
		UNION
		SELECT corporationID as id, name FROM evedata.corporations WHERE corporationID IN (` + joinInt32(idList) + `)
		UNION
		SELECT characterID as id, name FROM evedata.characters WHERE characterID IN (` + joinInt32(idList) + `)
		UNION
		SELECT typeID as id, typeName as name FROM eve.invTypes WHERE typeID IN  (` + joinInt32(idList) + `)
		`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var (
			id   int32
			name string
		)
		err := rows.Scan(&id, &name)
		if err != nil {
			return nil, err
		}
		ids[id] = name
	}

	// Double check we are not missing any
	for id, n := range ids {
		if id > 0 {
			if n == "" {
				return nil, fmt.Errorf("cannot find name for %d", id)
			}
		}
	}

	return ids, nil
}

// killmailConsumer receives killmails from NSQ and dumps attributes to SQL
func (s *Tailor) killmailConsumer() {
	for {
		a := <-chanKillmailAttributes
		b := a.Attributes.Ship

		sq := squirrel.Insert("evedata.killmailAttributes").Columns(
			"id", "eHP", "DPS", "Alpha", "scanResolution", "signatureRadiusNoMWD",
			"signatureRadius", "agility", "warpSpeed", "speedNoMWD", "speed", "remoteArmorRepair",
			"remoteShieldRepair", "remoteEnergyTransfer", "energyNeutralization", "sensorStrength",
			"RPS", "CPURemaining", "powerRemaining", "capacitorNoMWD", "capacitor",
			"capacitorTimeNoMWD", "capacitorTime",
		)

		sq = sq.Values(
			a.Killmail.KillmailId, b["avgEHP"], b["totalDPS"], b["totalAlphaDamage"], b["scanResolution"], b["signatureRadius"],
			b["signatureRadiusMWD"], b["agility"], b["warpSpeedMultiplier"], b["MaxVelocity"], b["MaxVelocityMWD"], b["remoteArmorRepairPerSecond"],
			b["remoteShieldBonusAmountPerSecond"], b["remotePowerTransferAmountPerSecond"], b["energyNeutralizerAmountPerSecond"],
			b["scanRadarStrength"]+b["scanLadarStrength"]+b["scanMagnetometricStrength"]+b["scanGravimetricStrength"],
			b["avgRPS"], b["cpuRemaining"], b["powerRemaining"], b["capacitorFraction"], b["capacitorFractionMWD"],
			b["capacitorDuration"]*1000000, b["capacitorDurationMWD"]*1000000,
		)

		sqlq, args, err := sq.ToSql()
		if err != nil {
			log.Println(err)
			continue
		}

		err = s.doSQL(sqlq+" ON DUPLICATE KEY UPDATE id = id", args...)
		if err != nil {
			log.Println(err)
			continue
		}

	}
}

func getAttributesForKillmail(km *esi.GetKillmailsKillmailIdKillmailHashOk) (*attributes.Attributes, error) {
	j, err := json.Marshal(km)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", "http://axiom.evedata:3005/killmail", bytes.NewBuffer(j))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	dec := json.NewDecoder(resp.Body)
	attributes := &attributes.Attributes{}
	if err = dec.Decode(attributes); err != nil {
		return nil, err
	}

	return attributes, nil
}
<|endoftext|>


<commit_before>package secure_shell_test

import (
	"errors"
	"os"
	"syscall"

	"github.com/docker/docker/pkg/term"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"

	config_package "github.com/cloudfoundry-incubator/lattice/ltc/config"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell/fake_dialer"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell/fake_secure_session"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell/fake_term"
)

var _ = Describe("SecureShell", func() {
	var (
		fakeDialer  *fake_dialer.FakeDialer
		fakeSession *fake_secure_session.FakeSecureSession
		fakeTerm    *fake_term.FakeTerm
		fakeStdin   *gbytes.Buffer
		fakeStdout  *gbytes.Buffer
		fakeStderr  *gbytes.Buffer

		config      *config_package.Config
		secureShell *secure_shell.SecureShell

		oldTerm string
	)

	BeforeEach(func() {
		fakeDialer = &fake_dialer.FakeDialer{}
		fakeSession = &fake_secure_session.FakeSecureSession{}
		fakeTerm = &fake_term.FakeTerm{}
		fakeStdin = gbytes.NewBuffer()
		fakeStdout = gbytes.NewBuffer()
		fakeStderr = gbytes.NewBuffer()

		config = config_package.New(nil)
		config.SetTarget("10.0.12.34")
		config.SetLogin("user", "past")

		secureShell = &secure_shell.SecureShell{Dialer: fakeDialer, Term: fakeTerm}
		fakeDialer.DialReturns(fakeSession, nil)

		oldTerm = os.Getenv("TERM")
		os.Setenv("TERM", "defaultterm")
	})

	AfterEach(func() {
		os.Setenv("TERM", oldTerm)
	})

	Describe("#ConnectToShell", func() {
		It("connects to the correct server given app name, instance and config", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)
			fakeTerm.GetWinsizeReturns(1000, 2000)

			termState := &term.State{}
			fakeTerm.SetRawTerminalReturns(termState, nil)

			err := secureShell.ConnectToShell("app-name", 2, "", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeDialer.DialCallCount()).To(Equal(1))
			user, authUser, authPass, address := fakeDialer.DialArgsForCall(0)
			Expect(user).To(Equal("diego:app-name/2"))
			Expect(authUser).To(Equal("user"))
			Expect(authPass).To(Equal("past"))
			Expect(address).To(Equal("10.0.12.34:2222"))

			Expect(fakeTerm.SetRawTerminalCallCount()).To(Equal(1))
			Expect(fakeTerm.SetRawTerminalArgsForCall(0)).To(Equal(os.Stdin.Fd()))

			Expect(fakeTerm.GetWinsizeCallCount()).To(Equal(1))
			Expect(fakeTerm.GetWinsizeArgsForCall(0)).To(Equal(os.Stdout.Fd()))

			Expect(fakeSession.RequestPtyCallCount()).To(Equal(1))
			termType, height, width, _ := fakeSession.RequestPtyArgsForCall(0)
			Expect(termType).To(Equal("defaultterm"))
			Expect(width).To(Equal(1000))
			Expect(height).To(Equal(2000))

			Expect(fakeTerm.RestoreTerminalCallCount()).To(Equal(1))
			fd, state := fakeTerm.RestoreTerminalArgsForCall(0)
			Expect(fd).To(Equal(os.Stdin.Fd()))
			Expect(state).To(Equal(termState))

			Expect(fakeSession.ShellCallCount()).To(Equal(1))
			Expect(fakeSession.WaitCallCount()).To(Equal(1))
		})

		It("runs a remote command", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			err := secureShell.ConnectToShell("app-name", 2, "/bin/ls", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeSession.ShellCallCount()).To(Equal(0))
			Expect(fakeSession.WaitCallCount()).To(Equal(0))

			Expect(fakeSession.RunCallCount()).To(Equal(1))
			Expect(fakeSession.RunArgsForCall(0)).To(Equal("/bin/ls"))
		})

		It("respects the user's TERM environment variable", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			os.Setenv("TERM", "term2000")

			err := secureShell.ConnectToShell("app-name", 2, "", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeSession.RequestPtyCallCount()).To(Equal(1))
			termType, _, _, _ := fakeSession.RequestPtyArgsForCall(0)
			Expect(termType).To(Equal("term2000"))
		})

		It("defaults to xterm ifno TERM environment variable is set", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			os.Setenv("TERM", "")

			err := secureShell.ConnectToShell("app-name", 2, "", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeSession.RequestPtyCallCount()).To(Equal(1))
			termType, _, _, _ := fakeSession.RequestPtyArgsForCall(0)
			Expect(termType).To(Equal("xterm"))
		})

		It("resizes the remote terminal if the local terminal is resized", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			fakeTerm.GetWinsizeReturns(10, 20)

			waitChan := make(chan struct{})
			fakeSession.ShellStub = func() error {
				Expect(fakeSession.SendRequestCallCount()).To(Equal(0))
				Expect(fakeTerm.GetWinsizeCallCount()).To(Equal(1))
				fakeTerm.GetWinsizeReturns(30, 40)
				Eventually(waitChan, 3).Should(Receive())
				return nil
			}

			go func() {
				defer GinkgoRecover()

				Eventually(waitChan).Should(Receive())
				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).NotTo(HaveOccurred())
			}()

			waitChan <- struct{}{}

			err := syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)
			Expect(err).NotTo(HaveOccurred())

			Eventually(fakeTerm.GetWinsizeCallCount, 3).Should(Equal(2))
			Expect(fakeSession.SendRequestCallCount()).To(Equal(1))
			name, wantReply, payload := fakeSession.SendRequestArgsForCall(0)
			Expect(name).To(Equal("window-change"))
			Expect(wantReply).To(BeFalse())
			Expect(payload).To(Equal([]byte{0, 0, 0, 30, 0, 0, 0, 40, 0, 0, 0, 0, 0, 0, 0, 0}))

			waitChan <- struct{}{}
		})

		It("does not resize the remote terminal if SIGWINCH is received but the window is the same size", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			fakeTerm.GetWinsizeReturns(10, 20)

			waitChan := make(chan struct{})
			fakeSession.ShellStub = func() error {
				Expect(fakeSession.SendRequestCallCount()).To(Equal(0))
				Expect(fakeTerm.GetWinsizeCallCount()).To(Equal(1))
				Eventually(waitChan, 3).Should(Receive())
				return nil
			}

			go func() {
				defer GinkgoRecover()

				Eventually(waitChan).Should(Receive())
				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).NotTo(HaveOccurred())
			}()

			waitChan <- struct{}{}

			Expect(syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)).To(Succeed())

			Eventually(fakeTerm.GetWinsizeCallCount).Should(Equal(2))
			Expect(fakeSession.SendRequestCallCount()).To(Equal(0))

			waitChan <- struct{}{}
		})

		Context("when the SecureDialer#Dial fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(nil, errors.New("cannot dial error"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("cannot dial error"))
			})
		})

		Context("when the SecureSession#StdinPipe fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(nil, errors.New("put this in your pipe"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("put this in your pipe"))
			})
		})

		Context("when the SecureSession#StdoutPipe fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(nil, errors.New("put this in your pipe"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("put this in your pipe"))
			})
		})

		Context("when the SecureSession#StderrPipe fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(fakeStdout, nil)
				fakeSession.StderrPipeReturns(nil, errors.New("put this in your pipe"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("put this in your pipe"))
			})
		})

		Context("when the SecureSession#RequestPty fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(fakeStdout, nil)
				fakeSession.StderrPipeReturns(fakeStderr, nil)
				fakeSession.RequestPtyReturns(errors.New("no pty"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("no pty"))
			})
		})

		Context("when the SecureTerm#SetRawTerminal fails", func() {
			It("does not call RestoreTerminal", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(fakeStdout, nil)
				fakeSession.StderrPipeReturns(fakeStderr, nil)
				fakeTerm.SetRawTerminalReturns(nil, errors.New("can't set raw"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).NotTo(HaveOccurred())

				Expect(fakeTerm.RestoreTerminalCallCount()).To(Equal(0))
			})
		})
	})
})
<commit_msg>Fix remaining flakey secure_shell package test<commit_after>package secure_shell_test

import (
	"errors"
	"os"
	"syscall"

	"github.com/docker/docker/pkg/term"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"

	config_package "github.com/cloudfoundry-incubator/lattice/ltc/config"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell/fake_dialer"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell/fake_secure_session"
	"github.com/cloudfoundry-incubator/lattice/ltc/secure_shell/fake_term"
)

var _ = Describe("SecureShell", func() {
	var (
		fakeDialer  *fake_dialer.FakeDialer
		fakeSession *fake_secure_session.FakeSecureSession
		fakeTerm    *fake_term.FakeTerm
		fakeStdin   *gbytes.Buffer
		fakeStdout  *gbytes.Buffer
		fakeStderr  *gbytes.Buffer

		config      *config_package.Config
		secureShell *secure_shell.SecureShell

		oldTerm string
	)

	BeforeEach(func() {
		fakeDialer = &fake_dialer.FakeDialer{}
		fakeSession = &fake_secure_session.FakeSecureSession{}
		fakeTerm = &fake_term.FakeTerm{}
		fakeStdin = gbytes.NewBuffer()
		fakeStdout = gbytes.NewBuffer()
		fakeStderr = gbytes.NewBuffer()

		config = config_package.New(nil)
		config.SetTarget("10.0.12.34")
		config.SetLogin("user", "past")

		secureShell = &secure_shell.SecureShell{Dialer: fakeDialer, Term: fakeTerm}
		fakeDialer.DialReturns(fakeSession, nil)

		oldTerm = os.Getenv("TERM")
		os.Setenv("TERM", "defaultterm")
	})

	AfterEach(func() {
		os.Setenv("TERM", oldTerm)
	})

	Describe("#ConnectToShell", func() {
		It("connects to the correct server given app name, instance and config", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)
			fakeTerm.GetWinsizeReturns(1000, 2000)

			termState := &term.State{}
			fakeTerm.SetRawTerminalReturns(termState, nil)

			err := secureShell.ConnectToShell("app-name", 2, "", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeDialer.DialCallCount()).To(Equal(1))
			user, authUser, authPass, address := fakeDialer.DialArgsForCall(0)
			Expect(user).To(Equal("diego:app-name/2"))
			Expect(authUser).To(Equal("user"))
			Expect(authPass).To(Equal("past"))
			Expect(address).To(Equal("10.0.12.34:2222"))

			Expect(fakeTerm.SetRawTerminalCallCount()).To(Equal(1))
			Expect(fakeTerm.SetRawTerminalArgsForCall(0)).To(Equal(os.Stdin.Fd()))

			Expect(fakeTerm.GetWinsizeCallCount()).To(Equal(1))
			Expect(fakeTerm.GetWinsizeArgsForCall(0)).To(Equal(os.Stdout.Fd()))

			Expect(fakeSession.RequestPtyCallCount()).To(Equal(1))
			termType, height, width, _ := fakeSession.RequestPtyArgsForCall(0)
			Expect(termType).To(Equal("defaultterm"))
			Expect(width).To(Equal(1000))
			Expect(height).To(Equal(2000))

			Expect(fakeTerm.RestoreTerminalCallCount()).To(Equal(1))
			fd, state := fakeTerm.RestoreTerminalArgsForCall(0)
			Expect(fd).To(Equal(os.Stdin.Fd()))
			Expect(state).To(Equal(termState))

			Expect(fakeSession.ShellCallCount()).To(Equal(1))
			Expect(fakeSession.WaitCallCount()).To(Equal(1))
		})

		It("runs a remote command", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			err := secureShell.ConnectToShell("app-name", 2, "/bin/ls", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeSession.ShellCallCount()).To(Equal(0))
			Expect(fakeSession.WaitCallCount()).To(Equal(0))

			Expect(fakeSession.RunCallCount()).To(Equal(1))
			Expect(fakeSession.RunArgsForCall(0)).To(Equal("/bin/ls"))
		})

		It("respects the user's TERM environment variable", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			os.Setenv("TERM", "term2000")

			err := secureShell.ConnectToShell("app-name", 2, "", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeSession.RequestPtyCallCount()).To(Equal(1))
			termType, _, _, _ := fakeSession.RequestPtyArgsForCall(0)
			Expect(termType).To(Equal("term2000"))
		})

		It("defaults to xterm ifno TERM environment variable is set", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			os.Setenv("TERM", "")

			err := secureShell.ConnectToShell("app-name", 2, "", config)
			Expect(err).NotTo(HaveOccurred())

			Expect(fakeSession.RequestPtyCallCount()).To(Equal(1))
			termType, _, _, _ := fakeSession.RequestPtyArgsForCall(0)
			Expect(termType).To(Equal("xterm"))
		})

		It("resizes the remote terminal if the local terminal is resized", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			fakeTerm.GetWinsizeReturns(10, 20)

			waitChan := make(chan struct{})
			fakeSession.ShellStub = func() error {
				Expect(fakeSession.SendRequestCallCount()).To(Equal(0))
				Expect(fakeTerm.GetWinsizeCallCount()).To(Equal(1))
				fakeTerm.GetWinsizeReturns(30, 40)
				Eventually(waitChan, 3).Should(Receive())
				return nil
			}

			go func() {
				defer GinkgoRecover()

				Eventually(waitChan, 3).Should(Receive())
				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).NotTo(HaveOccurred())
			}()

			waitChan <- struct{}{}

			err := syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)
			Expect(err).NotTo(HaveOccurred())

			Eventually(fakeTerm.GetWinsizeCallCount, 3).Should(Equal(2))
			Expect(fakeSession.SendRequestCallCount()).To(Equal(1))
			name, wantReply, payload := fakeSession.SendRequestArgsForCall(0)
			Expect(name).To(Equal("window-change"))
			Expect(wantReply).To(BeFalse())
			Expect(payload).To(Equal([]byte{0, 0, 0, 30, 0, 0, 0, 40, 0, 0, 0, 0, 0, 0, 0, 0}))

			waitChan <- struct{}{}
		})

		It("does not resize the remote terminal if SIGWINCH is received but the window is the same size", func() {
			fakeDialer.DialReturns(fakeSession, nil)
			fakeSession.StdinPipeReturns(fakeStdin, nil)
			fakeSession.StdoutPipeReturns(fakeStdout, nil)
			fakeSession.StderrPipeReturns(fakeStderr, nil)

			fakeTerm.GetWinsizeReturns(10, 20)

			waitChan := make(chan struct{})
			fakeSession.ShellStub = func() error {
				Expect(fakeSession.SendRequestCallCount()).To(Equal(0))
				Expect(fakeTerm.GetWinsizeCallCount()).To(Equal(1))
				Eventually(waitChan, 3).Should(Receive())
				return nil
			}

			go func() {
				defer GinkgoRecover()

				Eventually(waitChan, 3).Should(Receive())
				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).NotTo(HaveOccurred())
			}()

			waitChan <- struct{}{}

			Expect(syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)).To(Succeed())

			Eventually(fakeTerm.GetWinsizeCallCount).Should(Equal(2))
			Expect(fakeSession.SendRequestCallCount()).To(Equal(0))

			waitChan <- struct{}{}
		})

		Context("when the SecureDialer#Dial fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(nil, errors.New("cannot dial error"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("cannot dial error"))
			})
		})

		Context("when the SecureSession#StdinPipe fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(nil, errors.New("put this in your pipe"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("put this in your pipe"))
			})
		})

		Context("when the SecureSession#StdoutPipe fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(nil, errors.New("put this in your pipe"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("put this in your pipe"))
			})
		})

		Context("when the SecureSession#StderrPipe fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(fakeStdout, nil)
				fakeSession.StderrPipeReturns(nil, errors.New("put this in your pipe"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("put this in your pipe"))
			})
		})

		Context("when the SecureSession#RequestPty fails", func() {
			It("returns an error", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(fakeStdout, nil)
				fakeSession.StderrPipeReturns(fakeStderr, nil)
				fakeSession.RequestPtyReturns(errors.New("no pty"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).To(MatchError("no pty"))
			})
		})

		Context("when the SecureTerm#SetRawTerminal fails", func() {
			It("does not call RestoreTerminal", func() {
				fakeDialer.DialReturns(fakeSession, nil)
				fakeSession.StdinPipeReturns(fakeStdin, nil)
				fakeSession.StdoutPipeReturns(fakeStdout, nil)
				fakeSession.StderrPipeReturns(fakeStderr, nil)
				fakeTerm.SetRawTerminalReturns(nil, errors.New("can't set raw"))

				err := secureShell.ConnectToShell("app-name", 2, "", config)
				Expect(err).NotTo(HaveOccurred())

				Expect(fakeTerm.RestoreTerminalCallCount()).To(Equal(0))
			})
		})
	})
})
<|endoftext|>