<commit_before>/*
Copyright 2017 The OpenEBS 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 storagepool provides a specific implementation of CAS template engine
package storagepool

import (
	"fmt"
	"strings"

	"github.com/ghodss/yaml"
	"github.com/openebs/maya/pkg/apis/openebs.io/v1alpha1"
	"github.com/openebs/maya/pkg/engine"
	"github.com/openebs/maya/pkg/util"
)

// casStoragePoolEngine is capable of creating a storagepool via CAS template
//
// It implements following interfaces:
// - engine.CASCreator
//
// NOTE:
//  It overrides the Create method exposed by generic CASEngine
type casStoragePoolEngine struct {
	// casEngine exposes generic CAS template operations
	casEngine *engine.CASEngine
	// defaultConfig is the default cas storagepool configurations found
	// in the CASTemplate
	defaultConfig []v1alpha1.Config
	// openebsConfig is the configurations that can be passes
	openebsConfig []v1alpha1.Config
}

func unMarshallToConfig(config string) (configs []v1alpha1.Config, err error) {
	err = yaml.Unmarshal([]byte(config), &configs)
	return
}

// NewCASStoragePoolEngine returns a new instance of casStoragePoolEngine based on
// the provided cas configs & runtime storagepool values
func NewCASStoragePoolEngine(
	casTemplate *v1alpha1.CASTemplate,
	openebsConfig string,
	runtimeKey string,
	runtimeStoragePoolValues map[string]interface{}) (storagePoolEngine *casStoragePoolEngine, err error) {

	if len(strings.TrimSpace(runtimeKey)) == 0 {
		err = fmt.Errorf("Failed to create cas template engine: nil runtime storagepool key was provided")
		return
	}

	if len(runtimeStoragePoolValues) == 0 {
		err = fmt.Errorf("Failed to create cas template engine: nil runtime storagepool values was provided")
		return
	}
	// CAS config from  storagepoolclaim
	openebsConf, err := unMarshallToConfig(openebsConfig)
	if err != nil {
		return
	}
	// make use of the generic CAS template engine
	cEngine, err := engine.NewCASEngine(casTemplate, runtimeKey, runtimeStoragePoolValues)
	if err != nil {
		return
	}

	storagePoolEngine = &casStoragePoolEngine{
		casEngine:     cEngine,
		defaultConfig: casTemplate.Spec.Defaults,
		openebsConfig: openebsConf,
	}

	return
}

// mergeConfig will merge the unique configuration elements of lowPriorityConfig
// into highPriorityConfig and return the result
func mergeConfig(highPriorityConfig, lowPriorityConfig []v1alpha1.Config) (final []v1alpha1.Config) {
	var prioritized []string

	for _, pc := range highPriorityConfig {
		final = append(final, pc)
		prioritized = append(prioritized, strings.TrimSpace(pc.Name))
	}

	for _, lc := range lowPriorityConfig {
		if !util.ContainsString(prioritized, strings.TrimSpace(lc.Name)) {
			final = append(final, lc)
		}
	}

	return
}
func (c *casStoragePoolEngine) prepareFinalConfig() (final []v1alpha1.Config) {
	// merge unique config elements from SC with config from PVC
	final = mergeConfig(c.openebsConfig, c.defaultConfig)
	return
}

// addConfigToConfigTLP will add final cas storagepool configurations to ConfigTLP.
//
// NOTE:
//  This will enable templating a run task template as follows:
//
// {{ .Config.<ConfigName>.enabled }}
// {{ .Config.<ConfigName>.value }}
//
// NOTE:
//  Above parsing scheme is translated by running `go template` against the run
// task template
func (c *casStoragePoolEngine) addConfigToConfigTLP() error {
	var configName string
	allConfigsHierarchy := map[string]interface{}{}
	allConfigs := c.prepareFinalConfig()

	for _, config := range allConfigs {
		configName = strings.TrimSpace(config.Name)
		if len(configName) == 0 {
			return fmt.Errorf("Failed to merge config '%#v': missing config name", config)
		}

		configHierarchy := map[string]interface{}{
			configName: map[string]string{
				string(v1alpha1.EnabledPTP): config.Enabled,
				string(v1alpha1.ValuePTP):   config.Value,
			},
		}

		isMerged := util.MergeMapOfObjects(allConfigsHierarchy, configHierarchy)
		if !isMerged {
			return fmt.Errorf("Failed to merge config: unable to add config '%s' to config hierarchy", configName)
		}
	}

	// update merged config as the top level property
	c.casEngine.SetConfig(allConfigsHierarchy)
	return nil
}

// Create creates a storagepool
func (c *casStoragePoolEngine) Create() ([]byte, error) {
	// set customized CAS config as a top level property
	err := c.addConfigToConfigTLP()
	if err != nil {
		return nil, err
	}

	// delegate to generic cas template engine
	return c.casEngine.Run()
}
<commit_msg>Remove the blank line in imports<commit_after>/*
Copyright 2017 The OpenEBS 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 storagepool provides a specific implementation of CAS template engine
package storagepool

import (
	"fmt"
	"strings"
	"github.com/ghodss/yaml"
	"github.com/openebs/maya/pkg/apis/openebs.io/v1alpha1"
	"github.com/openebs/maya/pkg/engine"
	"github.com/openebs/maya/pkg/util"
)

// casStoragePoolEngine is capable of creating a storagepool via CAS template
//
// It implements following interfaces:
// - engine.CASCreator
//
// NOTE:
//  It overrides the Create method exposed by generic CASEngine
type casStoragePoolEngine struct {
	// casEngine exposes generic CAS template operations
	casEngine *engine.CASEngine
	// defaultConfig is the default cas storagepool configurations found
	// in the CASTemplate
	defaultConfig []v1alpha1.Config
	// openebsConfig is the configurations that can be passes
	openebsConfig []v1alpha1.Config
}

func unMarshallToConfig(config string) (configs []v1alpha1.Config, err error) {
	err = yaml.Unmarshal([]byte(config), &configs)
	return
}

// NewCASStoragePoolEngine returns a new instance of casStoragePoolEngine based on
// the provided cas configs & runtime storagepool values
func NewCASStoragePoolEngine(
	casTemplate *v1alpha1.CASTemplate,
	openebsConfig string,
	runtimeKey string,
	runtimeStoragePoolValues map[string]interface{}) (storagePoolEngine *casStoragePoolEngine, err error) {

	if len(strings.TrimSpace(runtimeKey)) == 0 {
		err = fmt.Errorf("Failed to create cas template engine: nil runtime storagepool key was provided")
		return
	}

	if len(runtimeStoragePoolValues) == 0 {
		err = fmt.Errorf("Failed to create cas template engine: nil runtime storagepool values was provided")
		return
	}
	// CAS config from  storagepoolclaim
	openebsConf, err := unMarshallToConfig(openebsConfig)
	if err != nil {
		return
	}
	// make use of the generic CAS template engine
	cEngine, err := engine.NewCASEngine(casTemplate, runtimeKey, runtimeStoragePoolValues)
	if err != nil {
		return
	}

	storagePoolEngine = &casStoragePoolEngine{
		casEngine:     cEngine,
		defaultConfig: casTemplate.Spec.Defaults,
		openebsConfig: openebsConf,
	}

	return
}

// mergeConfig will merge the unique configuration elements of lowPriorityConfig
// into highPriorityConfig and return the result
func mergeConfig(highPriorityConfig, lowPriorityConfig []v1alpha1.Config) (final []v1alpha1.Config) {
	var prioritized []string

	for _, pc := range highPriorityConfig {
		final = append(final, pc)
		prioritized = append(prioritized, strings.TrimSpace(pc.Name))
	}

	for _, lc := range lowPriorityConfig {
		if !util.ContainsString(prioritized, strings.TrimSpace(lc.Name)) {
			final = append(final, lc)
		}
	}

	return
}
func (c *casStoragePoolEngine) prepareFinalConfig() (final []v1alpha1.Config) {
	// merge unique config elements from SC with config from PVC
	final = mergeConfig(c.openebsConfig, c.defaultConfig)
	return
}

// addConfigToConfigTLP will add final cas storagepool configurations to ConfigTLP.
//
// NOTE:
//  This will enable templating a run task template as follows:
//
// {{ .Config.<ConfigName>.enabled }}
// {{ .Config.<ConfigName>.value }}
//
// NOTE:
//  Above parsing scheme is translated by running `go template` against the run
// task template
func (c *casStoragePoolEngine) addConfigToConfigTLP() error {
	var configName string
	allConfigsHierarchy := map[string]interface{}{}
	allConfigs := c.prepareFinalConfig()

	for _, config := range allConfigs {
		configName = strings.TrimSpace(config.Name)
		if len(configName) == 0 {
			return fmt.Errorf("Failed to merge config '%#v': missing config name", config)
		}

		configHierarchy := map[string]interface{}{
			configName: map[string]string{
				string(v1alpha1.EnabledPTP): config.Enabled,
				string(v1alpha1.ValuePTP):   config.Value,
			},
		}

		isMerged := util.MergeMapOfObjects(allConfigsHierarchy, configHierarchy)
		if !isMerged {
			return fmt.Errorf("Failed to merge config: unable to add config '%s' to config hierarchy", configName)
		}
	}

	// update merged config as the top level property
	c.casEngine.SetConfig(allConfigsHierarchy)
	return nil
}

// Create creates a storagepool
func (c *casStoragePoolEngine) Create() ([]byte, error) {
	// set customized CAS config as a top level property
	err := c.addConfigToConfigTLP()
	if err != nil {
		return nil, err
	}

	// delegate to generic cas template engine
	return c.casEngine.Run()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Tigera, Inc. 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.

// Test cases:
// Test 1: AutoAssign 1 IPv4, 1 IPv6 - expect one of each to be returned.
// Test 2: AutoAssign 256 IPv4, 256 IPv6 - expect 256 IPv4 + IPv6 addresses
// Test 3: AutoAssign 257 IPv4, 0 IPv6 - expect 256 IPv4 addresses, no IPv6, and an error.
// Test 4: AutoAssign 0 IPv4, 257 IPv6 - expect 256 IPv6 addresses, no IPv6, and an error.
// Test 5: (use pool of size /25 so only two blocks are contained):
// - Assign 1 address on host A (Expect 1 address)
// - Assign 1 address on host B (Expect 1 address, different block)
// - Assign 64 more addresses on host A (Expect 63 addresses from host A's block, 1 address from host B's block)

package client_test

import (
	"fmt"
	"log"
	"os/exec"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/onsi/ginkgo/extensions/table"

	"github.com/tigera/libcalico-go/calicoctl/commands"
	"github.com/tigera/libcalico-go/lib/api"
	"github.com/tigera/libcalico-go/lib/backend/etcd"
	"github.com/tigera/libcalico-go/lib/client"
)

var etcdType api.BackendType

var _ = Describe("IPAM", func() {

	DescribeTable("Requested IPs vs returned IPs",
		func(host string, cleanEnv bool, pool string, inv4, inv6, expv4, expv6 int, expError error) {
			outv4, outv6, outError := testIPAM(inv4, inv6, host, cleanEnv, pool)
			Expect(outv4).To(Equal(expv4))
			Expect(outv6).To(Equal(expv6))
			if expError != nil {
				Ω(outError).Should(HaveOccurred())
			}
		},

		// Test 1: AutoAssign 1 IPv4, 1 IPv6 - expect one of each to be returned.
		Entry("1 v4 1 v6", "testHost", true, "pool1", 1, 1, 1, 1, nil),

		// Test 2: AutoAssign 256 IPv4, 256 IPv6 - expect 256 IPv4 + IPv6 addresses.
		Entry("256 v4 256 v6", "testHost", true, "pool1", 256, 256, 256, 256, nil),

		// Test 3: AutoAssign 257 IPv4, 0 IPv6 - expect 256 IPv4 addresses, no IPv6, and no error.
		Entry("257 v4 0 v6", "testHost", true, "pool1", 257, 0, 256, 0, nil),

		// Test 4: AutoAssign 0 IPv4, 257 IPv6 - expect 256 IPv6 addresses, no IPv6, and no error.
		Entry("0 v4 257 v6", "testHost", true, "pool1", 0, 257, 0, 256, nil),

		// Test 5: (use pool of size /25 (/test/pool2.yaml) so only two blocks are contained):
		// - Assign 1 address on host A (Expect 1 address)
		Entry("1 v4 0 v6 host-A", "host-A", true, "pool2", 1, 0, 1, 0, nil),

		// - Assign 1 address on host B (Expect 1 address, different block)
		Entry("1 v4 0 v6 host-B", "host-B", false, "pool2", 1, 0, 1, 0, nil),

		// - Assign 64 more addresses on host A (Expect 63 addresses from host A's block, 1 address from host B's block)
		Entry("64 v4 0 v6 host-A", "host-A", false, "pool2", 64, 0, 64, 0, nil),
	)
})

// testIPAM takes number of requested IPv4 and IPv6, and hostname, and setus up/cleans up client and etcd,
// then it calls AutoAssign (function under test) and returns the number of returned IPv4 and IPv6 addresses and returned error.
func testIPAM(inv4, inv6 int, host string, cleanEnv bool, pool string) (int, int, error) {

	etcdType = "etcdv2"

	etcdConfig := etcd.EtcdConfig{
		EtcdEndpoints: "http://127.0.0.1:2379",
	}
	ac := api.ClientConfig{BackendType: etcdType, BackendConfig: &etcdConfig}

	bc, err := client.New(ac)
	if err != nil {
		panic(err)
	}

	ic := bc.IPAM()
	if cleanEnv {
		ic.SetIPAMConfig(client.IPAMConfig{
			StrictAffinity:     false,
			AutoAllocateBlocks: true,
		})
	}

	entry := client.AutoAssignArgs{
		Num4:     inv4,
		Num6:     inv6,
		Hostname: host,
	}

	setupEnv(cleanEnv, pool)

	v4, v6, outErr := ic.AutoAssign(entry)

	if outErr != nil {
		log.Println(outErr)
	}

	return len(v4), len(v6), outErr

}

// setupEnv cleans up etcd if cleanEnv flag is passed and then creates IP pool based on the pool name passed to it.
func setupEnv(cleanEnv bool, pool string) {
	if cleanEnv {
		etcdArgs := strings.Fields("rm /calico --recursive")
		if err := exec.Command("etcdctl", etcdArgs...).Run(); err != nil {
			log.Println(err)
		}
	}

	argsPool := strings.Fields(fmt.Sprintf("create -f ../../test/%s.yaml", pool))
	if err := commands.Create(argsPool); err != nil {
		log.Println(err)
	}
}
<commit_msg>AssignIP test cases<commit_after>// Copyright (c) 2016 Tigera, Inc. 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.

// Test cases:
// Test 1: AutoAssign 1 IPv4, 1 IPv6 - expect one of each to be returned.
// Test 2: AutoAssign 256 IPv4, 256 IPv6 - expect 256 IPv4 + IPv6 addresses
// Test 3: AutoAssign 257 IPv4, 0 IPv6 - expect 256 IPv4 addresses, no IPv6, and an error.
// Test 4: AutoAssign 0 IPv4, 257 IPv6 - expect 256 IPv6 addresses, no IPv6, and an error.
// Test 5: (use pool of size /25 so only two blocks are contained):
// - Assign 1 address on host A (Expect 1 address)
// - Assign 1 address on host B (Expect 1 address, different block)
// - Assign 64 more addresses on host A (Expect 63 addresses from host A's block, 1 address from host B's block)

package client_test

import (
	"errors"
	"fmt"
	"log"
	"net"
	"os/exec"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/onsi/ginkgo/extensions/table"

	"github.com/tigera/libcalico-go/calicoctl/commands"
	"github.com/tigera/libcalico-go/lib/api"
	"github.com/tigera/libcalico-go/lib/backend/etcd"
	"github.com/tigera/libcalico-go/lib/client"
	cnet "github.com/tigera/libcalico-go/lib/net"
)

var etcdType api.BackendType = "etcdv2"
var etcdConfig = etcd.EtcdConfig{
	EtcdEndpoints: "http://127.0.0.1:2379",
}

var _ = Describe("IPAM", func() {

	DescribeTable("AutoAssign: requested IPs vs returned IPs",
		func(host string, cleanEnv bool, pool string, inv4, inv6, expv4, expv6 int, expError error) {
			outv4, outv6, outError := testIPAMAutoAssign(inv4, inv6, host, cleanEnv, pool)
			Expect(outv4).To(Equal(expv4))
			Expect(outv6).To(Equal(expv6))
			if expError != nil {
				Ω(outError).Should(HaveOccurred())
			}
		},

		// Test 1: AutoAssign 1 IPv4, 1 IPv6 - expect one of each to be returned.
		Entry("1 v4 1 v6", "testHost", true, "pool1", 1, 1, 1, 1, nil),

		// Test 2: AutoAssign 256 IPv4, 256 IPv6 - expect 256 IPv4 + IPv6 addresses.
		Entry("256 v4 256 v6", "testHost", true, "pool1", 256, 256, 256, 256, nil),

		// Test 3: AutoAssign 257 IPv4, 0 IPv6 - expect 256 IPv4 addresses, no IPv6, and no error.
		Entry("257 v4 0 v6", "testHost", true, "pool1", 257, 0, 256, 0, nil),

		// Test 4: AutoAssign 0 IPv4, 257 IPv6 - expect 256 IPv6 addresses, no IPv6, and no error.
		Entry("0 v4 257 v6", "testHost", true, "pool1", 0, 257, 0, 256, nil),

		// Test 5: (use pool of size /25 (/test/pool2.yaml) so only two blocks are contained):
		// - Assign 1 address on host A (Expect 1 address)
		Entry("1 v4 0 v6 host-A", "host-A", true, "pool2", 1, 0, 1, 0, nil),

		// - Assign 1 address on host B (Expect 1 address, different block)
		Entry("1 v4 0 v6 host-B", "host-B", false, "pool2", 1, 0, 1, 0, nil),

		// - Assign 64 more addresses on host A (Expect 63 addresses from host A's block, 1 address from host B's block)
		Entry("64 v4 0 v6 host-A", "host-A", false, "pool2", 64, 0, 64, 0, nil),
	)

	DescribeTable("AssignIP: requested IP vs returned error",
		func(inIP net.IP, host string, cleanEnv bool, pool string, expError error) {
			outError := testIPAMAssignIP(inIP, host, pool, cleanEnv)
			if expError != nil {
				Ω(outError).Should(HaveOccurred())
				Expect(outError).To(Equal(expError))
			}
		},

		// Test 1: Assign 1 IPv4 from a configured pool - expect no error returned.
		Entry("Assign 1 IPv4 from a configured pool", net.ParseIP("192.168.1.0"), "testHost", true, "pool1", nil),

		// Test 2: Assign 1 IPv6 from a configured pool - expect no error returned.
		Entry("Assign 1 IPv6 from a configured pool", net.ParseIP("fd80:24e2:f998:72d6::"), "testHost", true, "pool1", nil),

		// Test 3: Assign 1 IPv4 from a non-configured pool - expect an error returned.
		Entry("Assign 1 IPv4 from a non-configured pool", net.ParseIP("1.1.1.1"), "testHost", true, "pool1", errors.New("The provided IP address is not in a configured pool\n")),

		// Test 4: Assign 1 IPv4 from a configured pool twice:
		// - Expect no error returned while assigning the IP for the first time.
		Entry("Assign 1 IPv4 from a configured pool twice (first time)", net.ParseIP("192.168.1.0"), "testHost", true, "pool1", nil),

		// - Expect an error returned while assigning the SAME IP again.
		Entry("Assign 1 IPv4 from a configured pool twice (second time)", net.ParseIP("192.168.1.0"), "testHost", false, "pool1", errors.New("Address already assigned in block")),
	)
})

func testIPAMAssignIP(inIP net.IP, host, pool string, cleanEnv bool) error {
	args := client.AssignIPArgs{
		IP:       cnet.IP{inIP},
		Hostname: host,
	}
	setupEnv(cleanEnv, pool)
	ic := setupIPMAClient(cleanEnv)
	outErr := ic.(client.IPAMInterface).AssignIP(args)

	if outErr != nil {
		log.Println(outErr)
	}
	return outErr
}

// testIPAMAutoAssign takes number of requested IPv4 and IPv6, and hostname, and setus up/cleans up client and etcd,
// then it calls AutoAssign (function under test) and returns the number of returned IPv4 and IPv6 addresses and returned error.
func testIPAMAutoAssign(inv4, inv6 int, host string, cleanEnv bool, pool string) (int, int, error) {

	args := client.AutoAssignArgs{
		Num4:     inv4,
		Num6:     inv6,
		Hostname: host,
	}

	setupEnv(cleanEnv, pool)
	ic := setupIPMAClient(cleanEnv)
	v4, v6, outErr := ic.(client.IPAMInterface).AutoAssign(args)

	if outErr != nil {
		log.Println(outErr)
	}

	return len(v4), len(v6), outErr
}

// setupEnv cleans up etcd if cleanEnv flag is passed and then creates IP pool based on the pool name passed to it.
func setupEnv(cleanEnv bool, pool string) {
	if cleanEnv {
		etcdArgs := strings.Fields("rm /calico --recursive")
		if err := exec.Command("etcdctl", etcdArgs...).Run(); err != nil {
			log.Println(err)
		}
	}

	argsPool := strings.Fields(fmt.Sprintf("create -f ../../test/%s.yaml", pool))
	if err := commands.Create(argsPool); err != nil {
		log.Println(err)
	}
}

// add a comment here bro
func setupIPMAClient(cleanEnv bool) interface{} {
	ac := api.ClientConfig{BackendType: etcdType, BackendConfig: &etcdConfig}

	bc, err := client.New(ac)
	if err != nil {
		panic(err)
	}

	ic := bc.IPAM()
	if cleanEnv {
		ic.SetIPAMConfig(client.IPAMConfig{
			StrictAffinity:     false,
			AutoAllocateBlocks: true,
		})
	}
	return ic
}
<|endoftext|>


<commit_before>package endly

import (
	"fmt"
	"github.com/viant/endly/model/msg"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/data"
	"github.com/viant/toolbox/url"
	"reflect"
	"sync"
	"time"
)

//AbstractService represenst an abstract service.
type AbstractService struct {
	Service
	*sync.RWMutex
	routeByAction  map[string]*Route
	routeByRequest map[reflect.Type]*Route
	actions        []string
	id             string
	state          data.Map
}

//Mutex returns a mutex.
func (s *AbstractService) Mutex() *sync.RWMutex {
	return s.RWMutex
}

//Register register action routes
func (s *AbstractService) Register(routes ...*Route) {
	for _, route := range routes {
		s.routeByAction[route.Action] = route
		s.routeByRequest[reflect.TypeOf(route.RequestProvider())] = route
		s.actions = append(s.actions, route.Action)
	}
}

func (s *AbstractService) addRouteIfConvertible(request interface{}) *Route {
	var requestType = reflect.TypeOf(request)
	if requestType != nil {
		for k, v := range s.routeByRequest {
			if requestType.Kind() == reflect.Ptr && requestType.Elem().ConvertibleTo(k.Elem()) {

				s.routeByRequest[requestType] = &Route{
					Action:           v.Action,
					RequestInfo:      v.RequestInfo,
					ResponseInfo:     v.ResponseInfo,
					RequestProvider:  v.RequestProvider,
					ResponseProvider: v.ResponseProvider,
					Handler: func(context *Context, convertibleRequest interface{}) (interface{}, error) {
						var request = v.RequestProvider()
						var requestValue = reflect.ValueOf(request)
						var convertibleValue = reflect.ValueOf(convertibleRequest)
						requestValue.Elem().Set(convertibleValue.Elem().Convert(k.Elem()))
						return v.Handler(context, request)
					},
				}
				return s.routeByRequest[requestType]
			}
		}
	}
	return nil
}

//Run returns a service action for supplied action
func (s *AbstractService) Run(context *Context, request interface{}) (response *ServiceResponse) {
	response = &ServiceResponse{Status: "ok"}
	startEvent := s.Begin(context, request)
	var err error
	defer func() {
		s.End(context)(startEvent, response.Response)
		if err != nil {
			response.Err = err
			response.Status = "error"
			response.Error = fmt.Sprintf("%v", err)
		}
	}()
	service, ok := s.routeByRequest[reflect.TypeOf(request)]
	if !ok {

		service = s.addRouteIfConvertible(request)
		if service == nil {
			err = NewError(s.ID(), fmt.Sprintf("%T", request), fmt.Errorf("failed to lookup service route: %T", request))
			return response
		}
	}

	if initializer, ok := request.(Initializer); ok {
		if err = initializer.Init(); err != nil {
			err = NewError(s.ID(), service.Action, fmt.Errorf("init %T failed: %v", request, err))
			return response
		}
	}

	if validator, ok := request.(Validator); ok {
		if err = validator.Validate(); err != nil {
			err = NewError(s.ID(), service.Action, fmt.Errorf("validation %T failed: %v", request, err))
			return response
		}
	}

	response.Response, err = service.Handler(context, request)
	if err != nil {
		var previous = err
		err = NewError(s.ID(), service.Action, err)
		if previous != err {
			context.Publish(msg.NewErrorEvent(fmt.Sprintf("%v", err)))
		}
		response.Err = err
	}
	return response
}

//Route returns a service action route for supplied action
func (s *AbstractService) Route(action string) (*Route, error) {
	if result, ok := s.routeByAction[action]; ok {
		return result, nil
	}
	return nil, fmt.Errorf("unknown %v.%v service action", s.id, action)
}

//Sleep sleeps for provided time in ms
func (s *AbstractService) Sleep(context *Context, sleepTimeMs int) {
	if sleepTimeMs > 0 {
		context.Publish(msg.NewSleepEvent(sleepTimeMs))
		time.Sleep(time.Millisecond * time.Duration(sleepTimeMs))
	}
}

//GetHostAndSSHPort return host and ssh port
func (s *AbstractService) GetHostAndSSHPort(target *url.Resource) (string, int) {
	if target == nil {
		return "", 0
	}
	port := toolbox.AsInt(target.ParsedURL.Port())
	if port == 0 {
		port = 22
	}
	hostname := target.ParsedURL.Hostname()
	if hostname == "" {
		hostname = "127.0.0.1"
	}
	return hostname, port
}

//Actions returns service actions
func (s *AbstractService) Actions() []string {
	return s.actions
}

//Begin add starting event
func (s *AbstractService) Begin(context *Context, value interface{}) msg.Event {
	return context.Publish(value)
}

//End adds finishing event.
func (s *AbstractService) End(context *Context) func(startEvent msg.Event, value interface{}) msg.Event {
	return func(startEvent msg.Event, value interface{}) msg.Event {
		return context.PublishWithStartEvent(value, startEvent)
	}
}

//ID returns this service id.
func (s *AbstractService) ID() string {
	return s.id
}

//State returns this service state map.
func (s *AbstractService) State() data.Map {
	return s.state
}

//NewAbstractService creates a new abstract service.
func NewAbstractService(id string) *AbstractService {
	return &AbstractService{
		id:             id,
		actions:        make([]string, 0),
		RWMutex:        &sync.RWMutex{},
		state:          data.NewMap(),
		routeByAction:  make(map[string]*Route),
		routeByRequest: make(map[reflect.Type]*Route),
	}
}

//NopRequest represent no operation to be deprecated
type NopRequest struct {
	In interface{}
}

//nopService represents no operation nopService (deprecated, use workflow, nop instead)
type nopService struct {
	*AbstractService
}

func (s *nopService) registerRoutes() {
	s.Register(&Route{
		Action: "nop",
		RequestInfo: &ActionInfo{
			Description: "no operation action, helper for separating action.Init as self descriptive steps",
		},
		RequestProvider: func() interface{} {
			return &NopRequest{}
		},
		ResponseProvider: func() interface{} {
			return struct{}{}
		},
		Handler: func(context *Context, request interface{}) (interface{}, error) {
			if req, ok := request.(*NopRequest); ok {
				return req.In, nil
			}
			return nil, fmt.Errorf("unsupported request type: %T", request)
		},
	})
}

//newNopService creates a new NoOperation nopService.
func newNopService() Service {
	var result = &nopService{
		AbstractService: NewAbstractService("nop"),
	}
	result.AbstractService.Service = result
	result.registerRoutes()
	return result
}
<commit_msg>added logging control<commit_after>package endly

import (
	"fmt"
	"github.com/viant/endly/model/msg"
	_ "github.com/viant/endly/unsafe"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/data"
	"github.com/viant/toolbox/url"
	"reflect"
	"sync"
	"time"
)

//AbstractService represenst an abstract service.
type AbstractService struct {
	Service
	*sync.RWMutex
	routeByAction  map[string]*Route
	routeByRequest map[reflect.Type]*Route
	actions        []string
	id             string
	state          data.Map
}

//Mutex returns a mutex.
func (s *AbstractService) Mutex() *sync.RWMutex {
	return s.RWMutex
}

//Register register action routes
func (s *AbstractService) Register(routes ...*Route) {
	for _, route := range routes {
		s.routeByAction[route.Action] = route
		s.routeByRequest[reflect.TypeOf(route.RequestProvider())] = route
		s.actions = append(s.actions, route.Action)
	}
}

func (s *AbstractService) addRouteIfConvertible(request interface{}) *Route {
	var requestType = reflect.TypeOf(request)
	if requestType != nil {
		for k, v := range s.routeByRequest {
			if requestType.Kind() == reflect.Ptr && requestType.Elem().ConvertibleTo(k.Elem()) {

				s.routeByRequest[requestType] = &Route{
					Action:           v.Action,
					RequestInfo:      v.RequestInfo,
					ResponseInfo:     v.ResponseInfo,
					RequestProvider:  v.RequestProvider,
					ResponseProvider: v.ResponseProvider,
					Handler: func(context *Context, convertibleRequest interface{}) (interface{}, error) {
						var request = v.RequestProvider()
						var requestValue = reflect.ValueOf(request)
						var convertibleValue = reflect.ValueOf(convertibleRequest)
						requestValue.Elem().Set(convertibleValue.Elem().Convert(k.Elem()))
						return v.Handler(context, request)
					},
				}
				return s.routeByRequest[requestType]
			}
		}
	}
	return nil
}

//Run returns a service action for supplied action
func (s *AbstractService) Run(context *Context, request interface{}) (response *ServiceResponse) {
	response = &ServiceResponse{Status: "ok"}
	startEvent := s.Begin(context, request)
	var err error
	defer func() {
		s.End(context)(startEvent, response.Response)
		if err != nil {
			response.Err = err
			response.Status = "error"
			response.Error = fmt.Sprintf("%v", err)
		}
	}()
	service, ok := s.routeByRequest[reflect.TypeOf(request)]
	if !ok {

		service = s.addRouteIfConvertible(request)
		if service == nil {
			err = NewError(s.ID(), fmt.Sprintf("%T", request), fmt.Errorf("failed to lookup service route: %T", request))
			return response
		}
	}

	if initializer, ok := request.(Initializer); ok {
		if err = initializer.Init(); err != nil {
			err = NewError(s.ID(), service.Action, fmt.Errorf("init %T failed: %v", request, err))
			return response
		}
	}

	if validator, ok := request.(Validator); ok {
		if err = validator.Validate(); err != nil {
			err = NewError(s.ID(), service.Action, fmt.Errorf("validation %T failed: %v", request, err))
			return response
		}
	}

	response.Response, err = service.Handler(context, request)
	if err != nil {
		var previous = err
		err = NewError(s.ID(), service.Action, err)
		if previous != err {
			context.Publish(msg.NewErrorEvent(fmt.Sprintf("%v", err)))
		}
		response.Err = err
	}
	return response
}

//Route returns a service action route for supplied action
func (s *AbstractService) Route(action string) (*Route, error) {
	if result, ok := s.routeByAction[action]; ok {
		return result, nil
	}
	return nil, fmt.Errorf("unknown %v.%v service action", s.id, action)
}

//Sleep sleeps for provided time in ms
func (s *AbstractService) Sleep(context *Context, sleepTimeMs int) {
	if sleepTimeMs > 0 {
		if context.IsLoggingEnabled() {
			context.Publish(msg.NewSleepEvent(sleepTimeMs))
		}
		time.Sleep(time.Millisecond * time.Duration(sleepTimeMs))
	}
}

//GetHostAndSSHPort return host and ssh port
func (s *AbstractService) GetHostAndSSHPort(target *url.Resource) (string, int) {
	if target == nil {
		return "", 0
	}
	port := toolbox.AsInt(target.ParsedURL.Port())
	if port == 0 {
		port = 22
	}
	hostname := target.ParsedURL.Hostname()
	if hostname == "" {
		hostname = "127.0.0.1"
	}
	return hostname, port
}

//Actions returns service actions
func (s *AbstractService) Actions() []string {
	return s.actions
}

//Begin add starting event
func (s *AbstractService) Begin(context *Context, value interface{}) msg.Event {
	return context.Publish(value)
}

//End adds finishing event.
func (s *AbstractService) End(context *Context) func(startEvent msg.Event, value interface{}) msg.Event {
	return func(startEvent msg.Event, value interface{}) msg.Event {
		return context.PublishWithStartEvent(value, startEvent)
	}
}

//ID returns this service id.
func (s *AbstractService) ID() string {
	return s.id
}

//State returns this service state map.
func (s *AbstractService) State() data.Map {
	return s.state
}

//NewAbstractService creates a new abstract service.
func NewAbstractService(id string) *AbstractService {
	return &AbstractService{
		id:             id,
		actions:        make([]string, 0),
		RWMutex:        &sync.RWMutex{},
		state:          data.NewMap(),
		routeByAction:  make(map[string]*Route),
		routeByRequest: make(map[reflect.Type]*Route),
	}
}

//NopRequest represent no operation to be deprecated
type NopRequest struct {
	In interface{}
}

//nopService represents no operation nopService (deprecated, use workflow, nop instead)
type nopService struct {
	*AbstractService
}

func (s *nopService) registerRoutes() {
	s.Register(&Route{
		Action: "nop",
		RequestInfo: &ActionInfo{
			Description: "no operation action, helper for separating action.Init as self descriptive steps",
		},
		RequestProvider: func() interface{} {
			return &NopRequest{}
		},
		ResponseProvider: func() interface{} {
			return struct{}{}
		},
		Handler: func(context *Context, request interface{}) (interface{}, error) {
			if req, ok := request.(*NopRequest); ok {
				return req.In, nil
			}
			return nil, fmt.Errorf("unsupported request type: %T", request)
		},
	})
}

//newNopService creates a new NoOperation nopService.
func newNopService() Service {
	var result = &nopService{
		AbstractService: NewAbstractService("nop"),
	}
	result.AbstractService.Service = result
	result.registerRoutes()
	return result
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
)

type Service struct {
	RemotePath string // 從哪裡 clone
	LocalPath  string // 或是從哪裡 copy
	Name       string

	Commit string // 想要的 commit hash
	Tag    string // 或是想要的 version tag

	Path string // 最後 service 被存到哪了
}

var ErrLocalPathAlreadyExists = errors.New("local path already exists")

func (s *Service) Fetch() ([]byte, error) {
	os.MkdirAll(s.Path, os.ModePerm)
	fmt.Println(s.LocalPath, s.RemotePath, "->", s.Path)

	if _, err := os.Stat(s.Path); os.IsNotExist(err) {
		return nil, ErrLocalPathAlreadyExists
	}

	if s.RemotePath != "" {
		return Exec("git", "clone", "--depth=1", s.RemotePath, s.Path+s.Name)
	} else if s.LocalPath != "" {
		return Exec("cp", "-r", s.LocalPath, s.Path)
	}

	return nil, errors.New("Unknown dependency format " + s.Name)
}

func (s *Service) VersionIdentifier() string {
	if s.Commit != "" {
		return s.Commit
	}

	if s.Tag != "" {
		return s.Tag
	}

	return ""
}
<commit_msg>respect version indentifier<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
)

type Service struct {
	RemotePath string // 從哪裡 clone
	LocalPath  string // 或是從哪裡 copy
	Name       string

	Commit string // 想要的 commit hash
	Tag    string // 或是想要的 version tag

	Path string // 最後 service 被存到哪了
}

var ErrLocalPathAlreadyExists = errors.New("local path already exists")

func (s *Service) Fetch() ([]byte, error) {
	os.MkdirAll(s.Path, os.ModePerm)
	fmt.Println(s.LocalPath, s.RemotePath, "->", s.Path)

	if _, err := os.Stat(s.Path); os.IsNotExist(err) {
		return nil, ErrLocalPathAlreadyExists
	}

	if s.RemotePath != "" {
		out, err := Exec("git", "clone", "--depth=1", s.RemotePath, s.Path+s.Name)
		if err != nil {
			return out, err
		}
		if s.VersionIdentifier() != "" {
			out, err := Exec("git", "checkout", s.VersionIdentifier)
			if err != nil {
				return out, err
			}
		}

		return out, nil
	} else if s.LocalPath != "" {
		return Exec("cp", "-r", s.LocalPath, s.Path)
	}

	return nil, errors.New("Unknown dependency format " + s.Name)
}

func (s *Service) VersionIdentifier() string {
	if s.Commit != "" {
		return s.Commit
	}

	if s.Tag != "" {
		return s.Tag
	}

	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

DEFAULT_REGION = "us-east-1"

func SessionNew(config *Config) *s3.S3 {
    creds := credentials.NewStaticCredentials(config.AccessKey, config.SecretKey, "")

    // By default make sure a region is specified
	return s3.New(session.New(&aws.Config{ Credentials: creds, Region: aws.String(DEFAULT_REGION) }))
}

func SessionForBucket(svc *s3.S3, bucket string) (*s3.S3, error) {
    params := &s3.HeadBucketInput{ Bucket: aws.String(bucket) }
    _, err := svc.HeadBucket(params)
    if err != nil {
        return nil, err
    }

    if loc, err := svc.GetBucketLocation(&s3.GetBucketLocationInput{Bucket: &bucket}); err != nil {
        return nil, err
    } else if (loc.LocationConstraint != nil) {
        return s3.New(session.New(&svc.Client.Config, &aws.Config{Region: loc.LocationConstraint})), nil
    }
    return svc, nil
}
<commit_msg>fixed typo<commit_after>package main

import (
	"github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

const DEFAULT_REGION = "us-east-1"

func SessionNew(config *Config) *s3.S3 {
    creds := credentials.NewStaticCredentials(config.AccessKey, config.SecretKey, "")

    // By default make sure a region is specified
	return s3.New(session.New(&aws.Config{ Credentials: creds, Region: aws.String(DEFAULT_REGION) }))
}

func SessionForBucket(svc *s3.S3, bucket string) (*s3.S3, error) {
    params := &s3.HeadBucketInput{ Bucket: aws.String(bucket) }
    _, err := svc.HeadBucket(params)
    if err != nil {
        return nil, err
    }

    if loc, err := svc.GetBucketLocation(&s3.GetBucketLocationInput{Bucket: &bucket}); err != nil {
        return nil, err
    } else if (loc.LocationConstraint != nil) {
        return s3.New(session.New(&svc.Client.Config, &aws.Config{Region: loc.LocationConstraint})), nil
    }
    return svc, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"os"
	"time"
)

// A Session is the current monitoring session. It keeps track of the
// averages, how long the tracker has been running, etc.
type Session struct {
	// Total is the total average of all session.
	Total RollingAverage `json:"total"`

	// NoShort is an average that excludes 'short' sessions. See
	// flags.short.
	NoShort RollingAverage `json:"noshort"`

	// Runtime is a timestamp of the time that the tracker was started.
	// timeDiff is a wrapper around time.Time.
	//
	// TODO: Make this a slice that can track the lengths of multiple
	// uses of the same session.
	Runtime timeDiff `json:"runtime"`
}

// LoadSession loads a session from the file at path. It returns the
// session and an error, if any.
func LoadSession(path string) (s Session, err error) {
	file, err := os.Open(path)
	if err != nil {
		return s, err
	}
	defer file.Close()

	d := json.NewDecoder(file)
	err = d.Decode(&s)
	return s, err
}

// Save saves the session to a file at path. The file is created if it
// doesn't exist, and truncated if it does. It returns an error if any
// are encountered.
func (s Session) Save(path string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	e := json.NewEncoder(file)
	return e.Encode(&s)
}

// timeDiff is a light wrapper around time.Time that marshals to JSON
// as a duration since the time that the diff represents. For example,
// if time.Now() is 3 seconds after the time represented by the
// timeDiff, the JSON representation will be "3s".
type timeDiff time.Time

// Since returns the duration representing the difference between the
// current time and t.
func (t timeDiff) Since() time.Duration {
	return time.Now().Sub(time.Time(t))
}

func (t timeDiff) MarshalJSON() ([]byte, error) {
	str := (t.Since() / time.Minute * time.Minute).String()

	buf := bytes.NewBuffer(make([]byte, 0, len(str)+2))
	buf.WriteByte('"')
	buf.WriteString(str)
	buf.WriteByte('"')

	return buf.Bytes(), nil
}

func (t *timeDiff) UnmarshalJSON(data []byte) error {
	data = bytes.Trim(data, `"`)
	d, err := time.ParseDuration(string(data))
	if err != nil {
		return err
	}

	*t = timeDiff(time.Now().Add(-d))

	return nil
}
<commit_msg>Add another TODO.<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"os"
	"time"
)

// A Session is the current monitoring session. It keeps track of the
// averages, how long the tracker has been running, etc.
type Session struct {
	// Total is the total average of all session.
	Total RollingAverage `json:"total"`

	// NoShort is an average that excludes 'short' sessions. See
	// flags.short.
	NoShort RollingAverage `json:"noshort"`

	// TODO: Add another average that doesn't include repeat characters?

	// Runtime is a timestamp of the time that the tracker was started.
	// timeDiff is a wrapper around time.Time.
	//
	// TODO: Make this a slice that can track the lengths of multiple
	// uses of the same session.
	Runtime timeDiff `json:"runtime"`
}

// LoadSession loads a session from the file at path. It returns the
// session and an error, if any.
func LoadSession(path string) (s Session, err error) {
	file, err := os.Open(path)
	if err != nil {
		return s, err
	}
	defer file.Close()

	d := json.NewDecoder(file)
	err = d.Decode(&s)
	return s, err
}

// Save saves the session to a file at path. The file is created if it
// doesn't exist, and truncated if it does. It returns an error if any
// are encountered.
func (s Session) Save(path string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	e := json.NewEncoder(file)
	return e.Encode(&s)
}

// timeDiff is a light wrapper around time.Time that marshals to JSON
// as a duration since the time that the diff represents. For example,
// if time.Now() is 3 seconds after the time represented by the
// timeDiff, the JSON representation will be "3s".
type timeDiff time.Time

// Since returns the duration representing the difference between the
// current time and t.
func (t timeDiff) Since() time.Duration {
	return time.Now().Sub(time.Time(t))
}

func (t timeDiff) MarshalJSON() ([]byte, error) {
	str := (t.Since() / time.Minute * time.Minute).String()

	buf := bytes.NewBuffer(make([]byte, 0, len(str)+2))
	buf.WriteByte('"')
	buf.WriteString(str)
	buf.WriteByte('"')

	return buf.Bytes(), nil
}

func (t *timeDiff) UnmarshalJSON(data []byte) error {
	data = bytes.Trim(data, `"`)
	d, err := time.ParseDuration(string(data))
	if err != nil {
		return err
	}

	*t = timeDiff(time.Now().Add(-d))

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"io"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

type ptyRequestPayload struct {
	Term                                   string
	Width, Height, PixelWidth, PixelHeight uint32
	Modes                                  string
}

func (request ptyRequestPayload) reply() []byte {
	return nil
}
func (request ptyRequestPayload) logEntry(channelID int) logEntry {
	return ptyLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Terminal: request.Term,
		Width:    request.Width,
		Height:   request.Height,
	}
}

type shellRequestPayload struct{}

func (request shellRequestPayload) reply() []byte {
	return nil
}
func (request shellRequestPayload) logEntry(channelID int) logEntry {
	return shellLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
	}
}

type x11RequestPayload struct {
	SingleConnection         bool
	AuthProtocol, AuthCookie string
	ScreenNumber             uint32
}

func (request x11RequestPayload) reply() []byte {
	return nil
}
func (request x11RequestPayload) logEntry(channelID int) logEntry {
	return x11Log{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Screen: request.ScreenNumber,
	}
}

type envRequestPayload struct {
	Name, Value string
}

func (request envRequestPayload) reply() []byte {
	return nil
}
func (request envRequestPayload) logEntry(channelID int) logEntry {
	return envLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Name:  request.Name,
		Value: request.Value,
	}
}

type execRequestPayload struct {
	Command string
}

func (request execRequestPayload) reply() []byte {
	return nil
}
func (request execRequestPayload) logEntry(channelID int) logEntry {
	return execLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Command: request.Command,
	}
}

type subsystemRequestPayload struct {
	Subsystem string
}

func (request subsystemRequestPayload) reply() []byte {
	return nil
}
func (request subsystemRequestPayload) logEntry(channelID int) logEntry {
	return subsystemLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Subsystem: request.Subsystem,
	}
}

type windowChangeRequestPayload struct {
	Width, Height, PixelWidth, PixelHeight uint32
}

func (request windowChangeRequestPayload) reply() []byte {
	return nil
}
func (request windowChangeRequestPayload) logEntry(channelID int) logEntry {
	return windowChangeLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Width:  request.Width,
		Height: request.Height,
	}
}

type sessionContext struct {
	channelContext
	ssh.Channel
	inputChan chan string
	active    bool
	pty       bool
}

type scannerReadLiner struct {
	scanner   *bufio.Scanner
	inputChan chan<- string
}

func (r scannerReadLiner) ReadLine() (string, error) {
	if !r.scanner.Scan() {
		if err := r.scanner.Err(); err != nil {
			return "", err
		}
		return "", io.EOF
	}
	line := r.scanner.Text()
	r.inputChan <- line
	return line, nil
}

type terminalReadLiner struct {
	terminal  *term.Terminal
	inputChan chan<- string
}

type clientEOFError struct{}

var clientEOF = clientEOFError{}

func (clientEOFError) Error() string {
	return "Client EOF"
}

func (r terminalReadLiner) ReadLine() (string, error) {
	line, err := r.terminal.ReadLine()
	if err == nil || line != "" {
		r.inputChan <- line
	}
	if err == io.EOF {
		return line, clientEOF
	}
	return line, err
}

func (context *sessionContext) handleProgram(program []string) {
	context.active = true
	var stdin readLiner
	var stdout, stderr io.Writer
	if context.pty {
		terminal := term.NewTerminal(context, "")
		stdin = terminalReadLiner{terminal, context.inputChan}
		stdout = terminal
		stderr = terminal
	} else {
		stdin = scannerReadLiner{bufio.NewScanner(context), context.inputChan}
		stdout = context
		stderr = context.Stderr()
	}
	go func() {
		defer close(context.inputChan)

		result, err := executeProgram(commandContext{program, stdin, stdout, stderr, context.pty, context.User()})
		if err != nil && err != io.EOF && err != clientEOF {
			warningLogger.Printf("Error executing program: %s", err)
			return
		}

		if err == clientEOF && context.pty {
			if _, err := context.Write([]byte("\r\n")); err != nil {
				warningLogger.Printf("Error sending CRLF: %s", err)
				return
			}
		}

		if _, err := context.SendRequest("exit-status", false, ssh.Marshal(struct {
			ExitStatus uint32
		}{result})); err != nil {
			warningLogger.Printf("Error sending exit status: %s", err)
			return
		}

		if (context.pty && err == clientEOF) || err == nil {
			if _, err := context.SendRequest("eow@openssh.com", false, nil); err != nil {
				warningLogger.Printf("Error sending EOW: %s", err)
				return
			}
		}

		if err := context.CloseWrite(); err != nil {
			warningLogger.Printf("Error sending EOF: %s", err)
			return
		}

		if err := context.Close(); err != nil {
			warningLogger.Printf("Error closing channel: %s", err)
			return
		}
	}()
}

func (context *sessionContext) handleRequest(request *ssh.Request) error {
	switch request.Type {
	case "pty-req":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			if context.pty {
				return errors.New("a pty is already requested")
			}
			payload := &ptyRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.pty = true
			return nil
		}
	case "shell":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			if len(request.Payload) != 0 {
				return errors.New("invalid request payload")
			}
			payload := &shellRequestPayload{}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(shellProgram)
			return nil
		}
	case "x11-req":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &x11RequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "env":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &envRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "exec":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &execRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Command))
			return nil
		}
	case "subsystem":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &subsystemRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Subsystem))
		}
	case "window-change":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		payload := &windowChangeRequestPayload{}
		if err := ssh.Unmarshal(request.Payload, payload); err != nil {
			return err
		}
		context.logEvent(payload.logEntry(context.channelID))
		return request.Reply(true, payload.reply())
	}
	sessionChannelRequestsMetric.WithLabelValues("unknown").Inc()
	warningLogger.Printf("Rejected session request: %s", request.Type)
	return request.Reply(false, nil)
}

var (
	sessionChannelsMetric = promauto.NewCounter(prometheus.CounterOpts{
		Name: "sshesame_session_channels_total",
		Help: "Total number of session channels",
	})
	activeSessionChannelsMetric = promauto.NewGauge(prometheus.GaugeOpts{
		Name: "sshesame_active_session_channels",
		Help: "Number of active session channels",
	})
	sessionChannelRequestsMetric = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "sshesame_session_channel_requests_total",
		Help: "Total number of session channel requests",
	}, []string{"type"})
)

func handleSessionChannel(newChannel ssh.NewChannel, context channelContext) error {
	if context.noMoreSessions {
		return errors.New("nore more sessions were supposed to be requested")
	}
	if len(newChannel.ExtraData()) != 0 {
		return errors.New("invalid channel data")
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		return err
	}
	sessionChannelsMetric.Inc()
	activeSessionChannelsMetric.Inc()
	defer activeSessionChannelsMetric.Dec()
	context.logEvent(sessionLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})
	defer context.logEvent(sessionCloseLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})

	inputChan := make(chan string)
	session := sessionContext{context, channel, inputChan, false, false}

	for inputChan != nil || requests != nil {
		select {
		case input, ok := <-inputChan:
			if !ok {
				inputChan = nil
				continue
			}
			context.logEvent(sessionInputLog{
				channelLog: channelLog{
					ChannelID: context.channelID,
				},
				Input: input,
			})
		case request, ok := <-requests:
			if !ok {
				requests = nil
				if !session.active {
					close(inputChan)
				}
				continue
			}
			if err := session.handleRequest(request); err != nil {
				return err
			}
		}
	}

	return nil
}
<commit_msg>Don't log known but rejected session requests twice<commit_after>package main

import (
	"bufio"
	"errors"
	"io"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

type ptyRequestPayload struct {
	Term                                   string
	Width, Height, PixelWidth, PixelHeight uint32
	Modes                                  string
}

func (request ptyRequestPayload) reply() []byte {
	return nil
}
func (request ptyRequestPayload) logEntry(channelID int) logEntry {
	return ptyLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Terminal: request.Term,
		Width:    request.Width,
		Height:   request.Height,
	}
}

type shellRequestPayload struct{}

func (request shellRequestPayload) reply() []byte {
	return nil
}
func (request shellRequestPayload) logEntry(channelID int) logEntry {
	return shellLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
	}
}

type x11RequestPayload struct {
	SingleConnection         bool
	AuthProtocol, AuthCookie string
	ScreenNumber             uint32
}

func (request x11RequestPayload) reply() []byte {
	return nil
}
func (request x11RequestPayload) logEntry(channelID int) logEntry {
	return x11Log{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Screen: request.ScreenNumber,
	}
}

type envRequestPayload struct {
	Name, Value string
}

func (request envRequestPayload) reply() []byte {
	return nil
}
func (request envRequestPayload) logEntry(channelID int) logEntry {
	return envLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Name:  request.Name,
		Value: request.Value,
	}
}

type execRequestPayload struct {
	Command string
}

func (request execRequestPayload) reply() []byte {
	return nil
}
func (request execRequestPayload) logEntry(channelID int) logEntry {
	return execLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Command: request.Command,
	}
}

type subsystemRequestPayload struct {
	Subsystem string
}

func (request subsystemRequestPayload) reply() []byte {
	return nil
}
func (request subsystemRequestPayload) logEntry(channelID int) logEntry {
	return subsystemLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Subsystem: request.Subsystem,
	}
}

type windowChangeRequestPayload struct {
	Width, Height, PixelWidth, PixelHeight uint32
}

func (request windowChangeRequestPayload) reply() []byte {
	return nil
}
func (request windowChangeRequestPayload) logEntry(channelID int) logEntry {
	return windowChangeLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Width:  request.Width,
		Height: request.Height,
	}
}

type sessionContext struct {
	channelContext
	ssh.Channel
	inputChan chan string
	active    bool
	pty       bool
}

type scannerReadLiner struct {
	scanner   *bufio.Scanner
	inputChan chan<- string
}

func (r scannerReadLiner) ReadLine() (string, error) {
	if !r.scanner.Scan() {
		if err := r.scanner.Err(); err != nil {
			return "", err
		}
		return "", io.EOF
	}
	line := r.scanner.Text()
	r.inputChan <- line
	return line, nil
}

type terminalReadLiner struct {
	terminal  *term.Terminal
	inputChan chan<- string
}

type clientEOFError struct{}

var clientEOF = clientEOFError{}

func (clientEOFError) Error() string {
	return "Client EOF"
}

func (r terminalReadLiner) ReadLine() (string, error) {
	line, err := r.terminal.ReadLine()
	if err == nil || line != "" {
		r.inputChan <- line
	}
	if err == io.EOF {
		return line, clientEOF
	}
	return line, err
}

func (context *sessionContext) handleProgram(program []string) {
	context.active = true
	var stdin readLiner
	var stdout, stderr io.Writer
	if context.pty {
		terminal := term.NewTerminal(context, "")
		stdin = terminalReadLiner{terminal, context.inputChan}
		stdout = terminal
		stderr = terminal
	} else {
		stdin = scannerReadLiner{bufio.NewScanner(context), context.inputChan}
		stdout = context
		stderr = context.Stderr()
	}
	go func() {
		defer close(context.inputChan)

		result, err := executeProgram(commandContext{program, stdin, stdout, stderr, context.pty, context.User()})
		if err != nil && err != io.EOF && err != clientEOF {
			warningLogger.Printf("Error executing program: %s", err)
			return
		}

		if err == clientEOF && context.pty {
			if _, err := context.Write([]byte("\r\n")); err != nil {
				warningLogger.Printf("Error sending CRLF: %s", err)
				return
			}
		}

		if _, err := context.SendRequest("exit-status", false, ssh.Marshal(struct {
			ExitStatus uint32
		}{result})); err != nil {
			warningLogger.Printf("Error sending exit status: %s", err)
			return
		}

		if (context.pty && err == clientEOF) || err == nil {
			if _, err := context.SendRequest("eow@openssh.com", false, nil); err != nil {
				warningLogger.Printf("Error sending EOW: %s", err)
				return
			}
		}

		if err := context.CloseWrite(); err != nil {
			warningLogger.Printf("Error sending EOF: %s", err)
			return
		}

		if err := context.Close(); err != nil {
			warningLogger.Printf("Error closing channel: %s", err)
			return
		}
	}()
}

func (context *sessionContext) handleRequest(request *ssh.Request) error {
	switch request.Type {
	case "pty-req":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			if context.pty {
				return errors.New("a pty is already requested")
			}
			payload := &ptyRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.pty = true
			return nil
		}
	case "shell":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			if len(request.Payload) != 0 {
				return errors.New("invalid request payload")
			}
			payload := &shellRequestPayload{}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(shellProgram)
			return nil
		}
	case "x11-req":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &x11RequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "env":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &envRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "exec":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &execRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Command))
			return nil
		}
	case "subsystem":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		if !context.active {
			payload := &subsystemRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Subsystem))
		}
	case "window-change":
		sessionChannelRequestsMetric.WithLabelValues(request.Type).Inc()
		payload := &windowChangeRequestPayload{}
		if err := ssh.Unmarshal(request.Payload, payload); err != nil {
			return err
		}
		context.logEvent(payload.logEntry(context.channelID))
		return request.Reply(true, payload.reply())
	default:
		sessionChannelRequestsMetric.WithLabelValues("unknown").Inc()
	}
	warningLogger.Printf("Rejected session request: %s", request.Type)
	return request.Reply(false, nil)
}

var (
	sessionChannelsMetric = promauto.NewCounter(prometheus.CounterOpts{
		Name: "sshesame_session_channels_total",
		Help: "Total number of session channels",
	})
	activeSessionChannelsMetric = promauto.NewGauge(prometheus.GaugeOpts{
		Name: "sshesame_active_session_channels",
		Help: "Number of active session channels",
	})
	sessionChannelRequestsMetric = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "sshesame_session_channel_requests_total",
		Help: "Total number of session channel requests",
	}, []string{"type"})
)

func handleSessionChannel(newChannel ssh.NewChannel, context channelContext) error {
	if context.noMoreSessions {
		return errors.New("nore more sessions were supposed to be requested")
	}
	if len(newChannel.ExtraData()) != 0 {
		return errors.New("invalid channel data")
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		return err
	}
	sessionChannelsMetric.Inc()
	activeSessionChannelsMetric.Inc()
	defer activeSessionChannelsMetric.Dec()
	context.logEvent(sessionLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})
	defer context.logEvent(sessionCloseLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})

	inputChan := make(chan string)
	session := sessionContext{context, channel, inputChan, false, false}

	for inputChan != nil || requests != nil {
		select {
		case input, ok := <-inputChan:
			if !ok {
				inputChan = nil
				continue
			}
			context.logEvent(sessionInputLog{
				channelLog: channelLog{
					ChannelID: context.channelID,
				},
				Input: input,
			})
		case request, ok := <-requests:
			if !ok {
				requests = nil
				if !session.active {
					close(inputChan)
				}
				continue
			}
			if err := session.handleRequest(request); err != nil {
				return err
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package http2

import (
	"fmt"
	hpack "github.com/ami-GS/GoHPACK"
	"net"
)

type Session struct {
	Conn  net.Conn
	Table hpack.Table
}

func (self *Session) Parse(buf []byte) {
	info := Http2Header{}
	info.Parse(buf[:9])

	if info.Type == TYPE_DATA {
		data := Data{}
		data.Parse(buf[9:], info.Flag, info.Length)
		fmt.Printf("data: %s", data.Data)
	} else if info.Type == TYPE_HEADERS {
		headers := Headers{}
		headers.Parse(buf[9:], info.Flag, &self.Table)
		if info.Flag == FLAG_END_HEADERS {
			self.Send(NewData("Hello! DATA frame", 1, FLAG_PADDED, 5))
		}
		fmt.Println("headers")
	} else if info.Type == TYPE_PRIORITY {
		fmt.Println("priority")
	} else if info.Type == TYPE_RST_STREAM {
		fmt.Println("rst stream")
	} else if info.Type == TYPE_SETTINGS {
		settings := Settings{}
		settings.Parse(buf[9:], info.Flag)
		if info.Flag == FLAG_NO {
			self.Send(NewSettings(SETTINGS_NO, 0, FLAG_ACK))
		} else if info.Flag == FLAG_ACK {
			fmt.Println("recv ACK setting!")
		}
		fmt.Println("settings")
	} else if info.Type == TYPE_PING {
		fmt.Println("ping")
	} else if info.Type == TYPE_GOAWAY {
		goaway := GoAway{}
		goaway.Parse(buf[9:])
		fmt.Printf("goaway: %s", goaway.Debug)
	} else if info.Type == TYPE_WINDOW_UPDATE {
		fmt.Println("window update")
	} else if info.Type == TYPE_CONTINUATION {
		fmt.Println("continuation")
	} else {
		panic("undefined frame type")
	}
}

func (self *Session) Send(data []byte) {
	self.Conn.Write(data)
}

func (self *Session) RunReceiver() {
	var buf []byte
	for {
		buf = make([]byte, 1024)
		_, err := self.Conn.Read(buf)
		if err != nil {
			return //EOF?
		} else {
			self.Parse(buf)
		}
	}
}

func NewSession(conn net.Conn) (client Session) {
	client.Conn = conn
	client.Table = hpack.InitTable()
	return
}
<commit_msg>parse connection preface to detect new session<commit_after>package http2

import (
	"fmt"
	hpack "github.com/ami-GS/GoHPACK"
	"net"
	"reflect"
)

type Session struct {
	Conn  net.Conn
	Table hpack.Table
}

func (self *Session) Parse(buf []byte) {
	info := Http2Header{}
	info.Parse(buf[:9])

	if info.Type == TYPE_DATA {
		data := Data{}
		data.Parse(buf[9:], info.Flag, info.Length)
		fmt.Printf("data: %s", data.Data)
	} else if info.Type == TYPE_HEADERS {
		headers := Headers{}
		headers.Parse(buf[9:], info.Flag, &self.Table)
		if info.Flag == FLAG_END_HEADERS {
			self.Send(NewData("Hello! DATA frame", 1, FLAG_PADDED, 5))
		}
		fmt.Println("headers")
	} else if info.Type == TYPE_PRIORITY {
		fmt.Println("priority")
	} else if info.Type == TYPE_RST_STREAM {
		fmt.Println("rst stream")
	} else if info.Type == TYPE_SETTINGS {
		settings := Settings{}
		settings.Parse(buf[9:], info.Flag)
		if info.Flag == FLAG_NO {
			self.Send(NewSettings(SETTINGS_NO, 0, FLAG_ACK))
		} else if info.Flag == FLAG_ACK {
			fmt.Println("recv ACK setting!")
		}
		fmt.Println("settings")
	} else if info.Type == TYPE_PING {
		fmt.Println("ping")
	} else if info.Type == TYPE_GOAWAY {
		goaway := GoAway{}
		goaway.Parse(buf[9:])
		fmt.Printf("goaway: %s", goaway.Debug)
	} else if info.Type == TYPE_WINDOW_UPDATE {
		fmt.Println("window update")
	} else if info.Type == TYPE_CONTINUATION {
		fmt.Println("continuation")
	} else {
		panic("undefined frame type")
	}
}

func (self *Session) Send(data []byte) {
	self.Conn.Write(data)
}

func (self *Session) RunReceiver() {
	var buf []byte
	for {
		buf = make([]byte, 1024)
		_, err := self.Conn.Read(buf)
		if err != nil {
			return //EOF?
		} else {
			if reflect.DeepEqual(buf[:24], CONNECTION_PREFACE) {
				fmt.Printf("New connection from %v\n", self.Conn.RemoteAddr())
				continue
			}
			self.Parse(buf)
		}
	}
}

func NewSession(conn net.Conn) (client Session) {
	client.Conn = conn
	client.Table = hpack.InitTable()
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 testutil

import (
	"fmt"
	"io"
	"math"
	"reflect"
	"sort"
	"strings"

	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"

	"k8s.io/component-base/metrics"
)

var (
	// MetricNameLabel is label under which model.Sample stores metric name
	MetricNameLabel model.LabelName = model.MetricNameLabel
	// QuantileLabel is label under which model.Sample stores latency quantile value
	QuantileLabel model.LabelName = model.QuantileLabel
)

// Metrics is generic metrics for other specific metrics
type Metrics map[string]model.Samples

// Equal returns true if all metrics are the same as the arguments.
func (m *Metrics) Equal(o Metrics) bool {
	var leftKeySet []string
	var rightKeySet []string
	for k := range *m {
		leftKeySet = append(leftKeySet, k)
	}
	for k := range o {
		rightKeySet = append(rightKeySet, k)
	}
	if !reflect.DeepEqual(leftKeySet, rightKeySet) {
		return false
	}
	for _, k := range leftKeySet {
		if !(*m)[k].Equal(o[k]) {
			return false
		}
	}
	return true
}

// NewMetrics returns new metrics which are initialized.
func NewMetrics() Metrics {
	result := make(Metrics)
	return result
}

// ParseMetrics parses Metrics from data returned from prometheus endpoint
func ParseMetrics(data string, output *Metrics) error {
	dec := expfmt.NewDecoder(strings.NewReader(data), expfmt.FmtText)
	decoder := expfmt.SampleDecoder{
		Dec:  dec,
		Opts: &expfmt.DecodeOptions{},
	}

	for {
		var v model.Vector
		if err := decoder.Decode(&v); err != nil {
			if err == io.EOF {
				// Expected loop termination condition.
				return nil
			}
			continue
		}
		for _, metric := range v {
			name := string(metric.Metric[model.MetricNameLabel])
			(*output)[name] = append((*output)[name], metric)
		}
	}
}

// TextToMetricFamilies reads 'in' as the simple and flat text-based exchange
// format and creates MetricFamily proto messages. It returns the MetricFamily
// proto messages in a map where the metric names are the keys, along with any
// error encountered.
func TextToMetricFamilies(in io.Reader) (map[string]*dto.MetricFamily, error) {
	var textParser expfmt.TextParser
	return textParser.TextToMetricFamilies(in)
}

// ExtractMetricSamples parses the prometheus metric samples from the input string.
func ExtractMetricSamples(metricsBlob string) ([]*model.Sample, error) {
	dec := expfmt.NewDecoder(strings.NewReader(metricsBlob), expfmt.FmtText)
	decoder := expfmt.SampleDecoder{
		Dec:  dec,
		Opts: &expfmt.DecodeOptions{},
	}

	var samples []*model.Sample
	for {
		var v model.Vector
		if err := decoder.Decode(&v); err != nil {
			if err == io.EOF {
				// Expected loop termination condition.
				return samples, nil
			}
			return nil, err
		}
		samples = append(samples, v...)
	}
}

// PrintSample returns formatted representation of metric Sample
func PrintSample(sample *model.Sample) string {
	buf := make([]string, 0)
	// Id is a VERY special label. For 'normal' container it's useless, but it's necessary
	// for 'system' containers (e.g. /docker-daemon, /kubelet, etc.). We know if that's the
	// case by checking if there's a label "kubernetes_container_name" present. It's hacky
	// but it works...
	_, normalContainer := sample.Metric["kubernetes_container_name"]
	for k, v := range sample.Metric {
		if strings.HasPrefix(string(k), "__") {
			continue
		}

		if string(k) == "id" && normalContainer {
			continue
		}
		buf = append(buf, fmt.Sprintf("%v=%v", string(k), v))
	}
	return fmt.Sprintf("[%v] = %v", strings.Join(buf, ","), sample.Value)
}

// ComputeHistogramDelta computes the change in histogram metric for a selected label.
// Results are stored in after samples
func ComputeHistogramDelta(before, after model.Samples, label model.LabelName) {
	beforeSamplesMap := make(map[string]*model.Sample)
	for _, bSample := range before {
		beforeSamplesMap[makeKey(bSample.Metric[label], bSample.Metric["le"])] = bSample
	}
	for _, aSample := range after {
		if bSample, found := beforeSamplesMap[makeKey(aSample.Metric[label], aSample.Metric["le"])]; found {
			aSample.Value = aSample.Value - bSample.Value
		}
	}
}

func makeKey(a, b model.LabelValue) string {
	return string(a) + "___" + string(b)
}

// GetMetricValuesForLabel returns value of metric for a given dimension
func GetMetricValuesForLabel(ms Metrics, metricName, label string) map[string]int64 {
	samples, found := ms[metricName]
	result := make(map[string]int64, len(samples))
	if !found {
		return result
	}
	for _, sample := range samples {
		count := int64(sample.Value)
		dimensionName := string(sample.Metric[model.LabelName(label)])
		result[dimensionName] = count
	}
	return result
}

// ValidateMetrics verifies if every sample of metric has all expected labels
func ValidateMetrics(metrics Metrics, metricName string, expectedLabels ...string) error {
	samples, ok := metrics[metricName]
	if !ok {
		return fmt.Errorf("metric %q was not found in metrics", metricName)
	}
	for _, sample := range samples {
		for _, l := range expectedLabels {
			if _, ok := sample.Metric[model.LabelName(l)]; !ok {
				return fmt.Errorf("metric %q is missing label %q, sample: %q", metricName, l, sample.String())
			}
		}
	}
	return nil
}

// Histogram wraps prometheus histogram DTO (data transfer object)
type Histogram struct {
	*dto.Histogram
}

// GetHistogramFromGatherer collects a metric from a gatherer implementing k8s.io/component-base/metrics.Gatherer interface.
// Used only for testing purposes where we need to gather metrics directly from a running binary (without metrics endpoint).
func GetHistogramFromGatherer(gatherer metrics.Gatherer, metricName string) (Histogram, error) {
	var metricFamily *dto.MetricFamily
	m, err := gatherer.Gather()
	if err != nil {
		return Histogram{}, err
	}
	for _, mFamily := range m {
		if mFamily.GetName() == metricName {
			metricFamily = mFamily
			break
		}
	}

	if metricFamily == nil {
		return Histogram{}, fmt.Errorf("metric %q not found", metricName)
	}

	if metricFamily.GetMetric() == nil {
		return Histogram{}, fmt.Errorf("metric %q is empty", metricName)
	}

	if len(metricFamily.GetMetric()) == 0 {
		return Histogram{}, fmt.Errorf("metric %q is empty", metricName)
	}

	return Histogram{
		// Histograms are stored under the first index (based on observation).
		// Given there's only one histogram registered per each metric name, accessing
		// the first index is sufficient.
		metricFamily.GetMetric()[0].GetHistogram(),
	}, nil
}

func uint64Ptr(u uint64) *uint64 {
	return &u
}

// Bucket of a histogram
type bucket struct {
	upperBound float64
	count      float64
}

func bucketQuantile(q float64, buckets []bucket) float64 {
	if q < 0 {
		return math.Inf(-1)
	}
	if q > 1 {
		return math.Inf(+1)
	}

	if len(buckets) < 2 {
		return math.NaN()
	}

	rank := q * buckets[len(buckets)-1].count
	b := sort.Search(len(buckets)-1, func(i int) bool { return buckets[i].count >= rank })

	if b == 0 {
		return buckets[0].upperBound * (rank / buckets[0].count)
	}

	if b == len(buckets)-1 && math.IsInf(buckets[b].upperBound, 1) {
		return buckets[len(buckets)-2].upperBound
	}

	// linear approximation of b-th bucket
	brank := rank - buckets[b-1].count
	bSize := buckets[b].upperBound - buckets[b-1].upperBound
	bCount := buckets[b].count - buckets[b-1].count

	return buckets[b-1].upperBound + bSize*(brank/bCount)
}

// Quantile computes q-th quantile of a cumulative histogram.
// It's expected the histogram is valid (by calling Validate)
func (hist *Histogram) Quantile(q float64) float64 {
	var buckets []bucket

	for _, bckt := range hist.Bucket {
		buckets = append(buckets, bucket{
			count:      float64(bckt.GetCumulativeCount()),
			upperBound: bckt.GetUpperBound(),
		})
	}

	if len(buckets) == 0 || buckets[len(buckets)-1].upperBound != math.Inf(+1) {
		// The list of buckets in dto.Histogram doesn't include the final +Inf bucket, so we
		// add it here for the reset of the samples.
		buckets = append(buckets, bucket{
			count:      float64(hist.GetSampleCount()),
			upperBound: math.Inf(+1),
		})
	}

	return bucketQuantile(q, buckets)
}

// Average computes histogram's average value
func (hist *Histogram) Average() float64 {
	return hist.GetSampleSum() / float64(hist.GetSampleCount())
}

// Clear clears all fields of the wrapped histogram
func (hist *Histogram) Clear() {
	if hist.SampleCount != nil {
		*hist.SampleCount = 0
	}
	if hist.SampleSum != nil {
		*hist.SampleSum = 0
	}
	for _, b := range hist.Bucket {
		if b.CumulativeCount != nil {
			*b.CumulativeCount = 0
		}
	}
}

// Validate makes sure the wrapped histogram has all necessary fields set and with valid values.
func (hist *Histogram) Validate() error {
	if hist.SampleCount == nil || hist.GetSampleCount() == 0 {
		return fmt.Errorf("nil or empty histogram SampleCount")
	}

	if hist.SampleSum == nil || hist.GetSampleSum() == 0 {
		return fmt.Errorf("nil or empty histogram SampleSum")
	}

	for _, bckt := range hist.Bucket {
		if bckt == nil {
			return fmt.Errorf("empty histogram bucket")
		}
		if bckt.UpperBound == nil || bckt.GetUpperBound() < 0 {
			return fmt.Errorf("nil or negative histogram bucket UpperBound")
		}
	}

	return nil
}

// GetGaugeMetricValue extract metric value from GaugeMetric
func GetGaugeMetricValue(m metrics.GaugeMetric) (float64, error) {
	metricProto := &dto.Metric{}
	if err := m.Write(metricProto); err != nil {
		return 0, fmt.Errorf("error writing m: %v", err)
	}
	return metricProto.Gauge.GetValue(), nil
}

// GetCounterMetricValue extract metric value from CounterMetric
func GetCounterMetricValue(m metrics.CounterMetric) (float64, error) {
	metricProto := &dto.Metric{}
	if err := m.(metrics.Metric).Write(metricProto); err != nil {
		return 0, fmt.Errorf("error writing m: %v", err)
	}
	return metricProto.Counter.GetValue(), nil
}

// GetHistogramMetricValue extract sum of all samples from ObserverMetric
func GetHistogramMetricValue(m metrics.ObserverMetric) (float64, error) {
	metricProto := &dto.Metric{}
	if err := m.(metrics.Metric).Write(metricProto); err != nil {
		return 0, fmt.Errorf("error writing m: %v", err)
	}
	return metricProto.Histogram.GetSampleSum(), nil
}

// LabelsMatch returns true if metric has all expected labels otherwise false
func LabelsMatch(metric *dto.Metric, labelFilter map[string]string) bool {
	metricLabels := map[string]string{}

	for _, labelPair := range metric.Label {
		metricLabels[labelPair.GetName()] = labelPair.GetValue()
	}

	// length comparison then match key to values in the maps
	if len(labelFilter) > len(metricLabels) {
		return false
	}

	for labelName, labelValue := range labelFilter {
		if value, ok := metricLabels[labelName]; !ok || value != labelValue {
			return false
		}
	}

	return true
}
<commit_msg>clean up testutil/metrics content<commit_after>/*
Copyright 2019 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 testutil

import (
	"fmt"
	"io"
	"math"
	"reflect"
	"sort"
	"strings"

	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"

	"k8s.io/component-base/metrics"
)

var (
	// MetricNameLabel is label under which model.Sample stores metric name
	MetricNameLabel model.LabelName = model.MetricNameLabel
	// QuantileLabel is label under which model.Sample stores latency quantile value
	QuantileLabel model.LabelName = model.QuantileLabel
)

// Metrics is generic metrics for other specific metrics
type Metrics map[string]model.Samples

// Equal returns true if all metrics are the same as the arguments.
func (m *Metrics) Equal(o Metrics) bool {
	var leftKeySet []string
	var rightKeySet []string
	for k := range *m {
		leftKeySet = append(leftKeySet, k)
	}
	for k := range o {
		rightKeySet = append(rightKeySet, k)
	}
	if !reflect.DeepEqual(leftKeySet, rightKeySet) {
		return false
	}
	for _, k := range leftKeySet {
		if !(*m)[k].Equal(o[k]) {
			return false
		}
	}
	return true
}

// NewMetrics returns new metrics which are initialized.
func NewMetrics() Metrics {
	result := make(Metrics)
	return result
}

// ParseMetrics parses Metrics from data returned from prometheus endpoint
func ParseMetrics(data string, output *Metrics) error {
	dec := expfmt.NewDecoder(strings.NewReader(data), expfmt.FmtText)
	decoder := expfmt.SampleDecoder{
		Dec:  dec,
		Opts: &expfmt.DecodeOptions{},
	}

	for {
		var v model.Vector
		if err := decoder.Decode(&v); err != nil {
			if err == io.EOF {
				// Expected loop termination condition.
				return nil
			}
			continue
		}
		for _, metric := range v {
			name := string(metric.Metric[MetricNameLabel])
			(*output)[name] = append((*output)[name], metric)
		}
	}
}

// TextToMetricFamilies reads 'in' as the simple and flat text-based exchange
// format and creates MetricFamily proto messages. It returns the MetricFamily
// proto messages in a map where the metric names are the keys, along with any
// error encountered.
func TextToMetricFamilies(in io.Reader) (map[string]*dto.MetricFamily, error) {
	var textParser expfmt.TextParser
	return textParser.TextToMetricFamilies(in)
}

// PrintSample returns formatted representation of metric Sample
func PrintSample(sample *model.Sample) string {
	buf := make([]string, 0)
	// Id is a VERY special label. For 'normal' container it's useless, but it's necessary
	// for 'system' containers (e.g. /docker-daemon, /kubelet, etc.). We know if that's the
	// case by checking if there's a label "kubernetes_container_name" present. It's hacky
	// but it works...
	_, normalContainer := sample.Metric["kubernetes_container_name"]
	for k, v := range sample.Metric {
		if strings.HasPrefix(string(k), "__") {
			continue
		}

		if string(k) == "id" && normalContainer {
			continue
		}
		buf = append(buf, fmt.Sprintf("%v=%v", string(k), v))
	}
	return fmt.Sprintf("[%v] = %v", strings.Join(buf, ","), sample.Value)
}

// ComputeHistogramDelta computes the change in histogram metric for a selected label.
// Results are stored in after samples
func ComputeHistogramDelta(before, after model.Samples, label model.LabelName) {
	beforeSamplesMap := make(map[string]*model.Sample)
	for _, bSample := range before {
		beforeSamplesMap[makeKey(bSample.Metric[label], bSample.Metric["le"])] = bSample
	}
	for _, aSample := range after {
		if bSample, found := beforeSamplesMap[makeKey(aSample.Metric[label], aSample.Metric["le"])]; found {
			aSample.Value = aSample.Value - bSample.Value
		}
	}
}

func makeKey(a, b model.LabelValue) string {
	return string(a) + "___" + string(b)
}

// GetMetricValuesForLabel returns value of metric for a given dimension
func GetMetricValuesForLabel(ms Metrics, metricName, label string) map[string]int64 {
	samples, found := ms[metricName]
	result := make(map[string]int64, len(samples))
	if !found {
		return result
	}
	for _, sample := range samples {
		count := int64(sample.Value)
		dimensionName := string(sample.Metric[model.LabelName(label)])
		result[dimensionName] = count
	}
	return result
}

// ValidateMetrics verifies if every sample of metric has all expected labels
func ValidateMetrics(metrics Metrics, metricName string, expectedLabels ...string) error {
	samples, ok := metrics[metricName]
	if !ok {
		return fmt.Errorf("metric %q was not found in metrics", metricName)
	}
	for _, sample := range samples {
		for _, l := range expectedLabels {
			if _, ok := sample.Metric[model.LabelName(l)]; !ok {
				return fmt.Errorf("metric %q is missing label %q, sample: %q", metricName, l, sample.String())
			}
		}
	}
	return nil
}

// Histogram wraps prometheus histogram DTO (data transfer object)
type Histogram struct {
	*dto.Histogram
}

// GetHistogramFromGatherer collects a metric from a gatherer implementing k8s.io/component-base/metrics.Gatherer interface.
// Used only for testing purposes where we need to gather metrics directly from a running binary (without metrics endpoint).
func GetHistogramFromGatherer(gatherer metrics.Gatherer, metricName string) (Histogram, error) {
	var metricFamily *dto.MetricFamily
	m, err := gatherer.Gather()
	if err != nil {
		return Histogram{}, err
	}
	for _, mFamily := range m {
		if mFamily.GetName() == metricName {
			metricFamily = mFamily
			break
		}
	}

	if metricFamily == nil {
		return Histogram{}, fmt.Errorf("metric %q not found", metricName)
	}

	if metricFamily.GetMetric() == nil {
		return Histogram{}, fmt.Errorf("metric %q is empty", metricName)
	}

	if len(metricFamily.GetMetric()) == 0 {
		return Histogram{}, fmt.Errorf("metric %q is empty", metricName)
	}

	return Histogram{
		// Histograms are stored under the first index (based on observation).
		// Given there's only one histogram registered per each metric name, accessing
		// the first index is sufficient.
		metricFamily.GetMetric()[0].GetHistogram(),
	}, nil
}

func uint64Ptr(u uint64) *uint64 {
	return &u
}

// Bucket of a histogram
type bucket struct {
	upperBound float64
	count      float64
}

func bucketQuantile(q float64, buckets []bucket) float64 {
	if q < 0 {
		return math.Inf(-1)
	}
	if q > 1 {
		return math.Inf(+1)
	}

	if len(buckets) < 2 {
		return math.NaN()
	}

	rank := q * buckets[len(buckets)-1].count
	b := sort.Search(len(buckets)-1, func(i int) bool { return buckets[i].count >= rank })

	if b == 0 {
		return buckets[0].upperBound * (rank / buckets[0].count)
	}

	if b == len(buckets)-1 && math.IsInf(buckets[b].upperBound, 1) {
		return buckets[len(buckets)-2].upperBound
	}

	// linear approximation of b-th bucket
	brank := rank - buckets[b-1].count
	bSize := buckets[b].upperBound - buckets[b-1].upperBound
	bCount := buckets[b].count - buckets[b-1].count

	return buckets[b-1].upperBound + bSize*(brank/bCount)
}

// Quantile computes q-th quantile of a cumulative histogram.
// It's expected the histogram is valid (by calling Validate)
func (hist *Histogram) Quantile(q float64) float64 {
	var buckets []bucket

	for _, bckt := range hist.Bucket {
		buckets = append(buckets, bucket{
			count:      float64(bckt.GetCumulativeCount()),
			upperBound: bckt.GetUpperBound(),
		})
	}

	if len(buckets) == 0 || buckets[len(buckets)-1].upperBound != math.Inf(+1) {
		// The list of buckets in dto.Histogram doesn't include the final +Inf bucket, so we
		// add it here for the reset of the samples.
		buckets = append(buckets, bucket{
			count:      float64(hist.GetSampleCount()),
			upperBound: math.Inf(+1),
		})
	}

	return bucketQuantile(q, buckets)
}

// Average computes histogram's average value
func (hist *Histogram) Average() float64 {
	return hist.GetSampleSum() / float64(hist.GetSampleCount())
}

// Clear clears all fields of the wrapped histogram
func (hist *Histogram) Clear() {
	if hist.SampleCount != nil {
		*hist.SampleCount = 0
	}
	if hist.SampleSum != nil {
		*hist.SampleSum = 0
	}
	for _, b := range hist.Bucket {
		if b.CumulativeCount != nil {
			*b.CumulativeCount = 0
		}
	}
}

// Validate makes sure the wrapped histogram has all necessary fields set and with valid values.
func (hist *Histogram) Validate() error {
	if hist.SampleCount == nil || hist.GetSampleCount() == 0 {
		return fmt.Errorf("nil or empty histogram SampleCount")
	}

	if hist.SampleSum == nil || hist.GetSampleSum() == 0 {
		return fmt.Errorf("nil or empty histogram SampleSum")
	}

	for _, bckt := range hist.Bucket {
		if bckt == nil {
			return fmt.Errorf("empty histogram bucket")
		}
		if bckt.UpperBound == nil || bckt.GetUpperBound() < 0 {
			return fmt.Errorf("nil or negative histogram bucket UpperBound")
		}
	}

	return nil
}

// GetGaugeMetricValue extract metric value from GaugeMetric
func GetGaugeMetricValue(m metrics.GaugeMetric) (float64, error) {
	metricProto := &dto.Metric{}
	if err := m.Write(metricProto); err != nil {
		return 0, fmt.Errorf("error writing m: %v", err)
	}
	return metricProto.Gauge.GetValue(), nil
}

// GetCounterMetricValue extract metric value from CounterMetric
func GetCounterMetricValue(m metrics.CounterMetric) (float64, error) {
	metricProto := &dto.Metric{}
	if err := m.(metrics.Metric).Write(metricProto); err != nil {
		return 0, fmt.Errorf("error writing m: %v", err)
	}
	return metricProto.Counter.GetValue(), nil
}

// GetHistogramMetricValue extract sum of all samples from ObserverMetric
func GetHistogramMetricValue(m metrics.ObserverMetric) (float64, error) {
	metricProto := &dto.Metric{}
	if err := m.(metrics.Metric).Write(metricProto); err != nil {
		return 0, fmt.Errorf("error writing m: %v", err)
	}
	return metricProto.Histogram.GetSampleSum(), nil
}

// LabelsMatch returns true if metric has all expected labels otherwise false
func LabelsMatch(metric *dto.Metric, labelFilter map[string]string) bool {
	metricLabels := map[string]string{}

	for _, labelPair := range metric.Label {
		metricLabels[labelPair.GetName()] = labelPair.GetValue()
	}

	// length comparison then match key to values in the maps
	if len(labelFilter) > len(metricLabels) {
		return false
	}

	for labelName, labelValue := range labelFilter {
		if value, ok := metricLabels[labelName]; !ok || value != labelValue {
			return false
		}
	}

	return true
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Tigera, Inc. 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 conntrack

import (
	"net"
	"os/exec"
	"strings"

	log "github.com/sirupsen/logrus"
)

// For TCP/UDP, each conntrack entry holds two copies of the tuple
// (src addr, dst addr, src port, dst port).  One copy for the original direction and one copy for
// the reply direction.  This is how the kernel handles NAT: by looking up the tuple for a packet
// by its original tuple and mapping onto the corresponding reply direction tuple (or vice versa).
// The reply tuple is calculated when the original outgoing packet is processed (and possibly
// NATted).
//
// When we delete conntrack entries by IP address, we need to specify which element of the tuple
// to look in.  This slice holds the flags corresponding to the fields we care about.  Since we're
// deleting entries for local workload endpoints, either the endpoint originated the traffic, or it
// received the traffic and replied to it.  In the originating case, the "original source" will be
// set to the endpoint's IP; in the other case, the "reply source". Hence, it's sufficient to only
// look in those two fields.
var deleteDirections = []string{
	"--orig-src",
	"--reply-src",
}

const numRetries = 3

type Conntrack struct {
	newCmd newCmd
}

func New() *Conntrack {
	return NewWithCmdShim(func(name string, arg ...string) CmdIface {
		return exec.Command(name, arg...)
	})
}

// NewWithCmdShim is a test constructor that allows for shimming exec.Command.
func NewWithCmdShim(newCmd newCmd) *Conntrack {
	return &Conntrack{
		newCmd: newCmd,
	}
}

type newCmd func(name string, arg ...string) CmdIface

type CmdIface interface {
	CombinedOutput() ([]byte, error)
}

func (c Conntrack) RemoveConntrackFlows(ipVersion uint8, ipAddr net.IP) {
	var family string
	switch ipVersion {
	case 4:
		family = "ipv4"
	case 6:
		family = "ipv6"
	default:
		log.WithField("version", ipVersion).Panic("Unknown IP version")
	}
	log.WithField("ip", ipAddr).Info("Removing conntrack flows")
	for _, direction := range deleteDirections {
		logCxt := log.WithFields(log.Fields{"ip": ipAddr, "direction": direction})
		// Retry a few times because the conntrack command seems to fail at random.
		for retry := 0; retry <= numRetries; retry += 1 {
			cmd := c.newCmd("conntrack",
				"--family", family,
				"--delete", direction,
				ipAddr.String())
			output, err := cmd.CombinedOutput()
			if err == nil {
				logCxt.Debug("Successfully removed conntrack flows.")
				break
			}
			if strings.Contains(string(output), "0 flow entries") {
				// Success, there were no flows.
				logCxt.Debug("IP wasn't in conntrack")
				break
			}
			if retry == numRetries {
				logCxt.WithError(err).Error("Failed to remove conntrack flows after retries.")
			} else {
				logCxt.WithError(err).Warn("Failed to remove conntrack flows, will retry...")
			}
		}
	}
}
<commit_msg>Suppress conntrack cleanup errors when no flows are found.<commit_after>// Copyright (c) 2016-2017 Tigera, Inc. 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 conntrack

import (
	"bytes"
	"net"
	"os/exec"
	"strings"

	log "github.com/sirupsen/logrus"
)

// For TCP/UDP, each conntrack entry holds two copies of the tuple
// (src addr, dst addr, src port, dst port).  One copy for the original direction and one copy for
// the reply direction.  This is how the kernel handles NAT: by looking up the tuple for a packet
// by its original tuple and mapping onto the corresponding reply direction tuple (or vice versa).
// The reply tuple is calculated when the original outgoing packet is processed (and possibly
// NATted).
//
// When we delete conntrack entries by IP address, we need to specify which element of the tuple
// to look in.  This slice holds the flags corresponding to the fields we care about.  Since we're
// deleting entries for local workload endpoints, either the endpoint originated the traffic, or it
// received the traffic and replied to it.  In the originating case, the "original source" will be
// set to the endpoint's IP; in the other case, the "reply source". Hence, it's sufficient to only
// look in those two fields.
var deleteDirections = []string{
	"--orig-src",
	"--reply-src",
}

const numRetries = 3

type Conntrack struct {
	newCmd newCmd
}

func New() *Conntrack {
	return NewWithCmdShim(func(name string, arg ...string) CmdIface {
		return exec.Command(name, arg...)
	})
}

// NewWithCmdShim is a test constructor that allows for shimming exec.Command.
func NewWithCmdShim(newCmd newCmd) *Conntrack {
	return &Conntrack{
		newCmd: newCmd,
	}
}

type newCmd func(name string, arg ...string) CmdIface

type CmdIface interface {
	CombinedOutput() ([]byte, error)
}

func (c Conntrack) RemoveConntrackFlows(ipVersion uint8, ipAddr net.IP) {
	var family string
	switch ipVersion {
	case 4:
		family = "ipv4"
	case 6:
		family = "ipv6"
	default:
		log.WithField("version", ipVersion).Panic("Unknown IP version")
	}
	log.WithField("ip", ipAddr).Info("Removing conntrack flows")
	for _, direction := range deleteDirections {
		logCxt := log.WithFields(log.Fields{"ip": ipAddr, "direction": direction})
		// Retry a few times because the conntrack command seems to fail at random.
		for retry := 0; retry <= numRetries; retry += 1 {
			cmd := c.newCmd("conntrack",
				"--family", family,
				"--delete", direction,
				ipAddr.String())
			output, err := cmd.CombinedOutput()
			if err == nil {
				logCxt.Debug("Successfully removed conntrack flows.")
				break
			} else if bytes.Contains(output, []byte("0 flow entries have been deleted")) {
				logCxt.Debug("conntrack tool didn't find any flows.")
				break
			}

			if strings.Contains(string(output), "0 flow entries") {
				// Success, there were no flows.
				logCxt.Debug("IP wasn't in conntrack")
				break
			}
			if retry == numRetries {
				logCxt.WithError(err).Error("Failed to remove conntrack flows after retries.")
			} else {
				logCxt.WithError(err).Warn("Failed to remove conntrack flows, will retry...")
			}
		}
	}
}
<|endoftext|>