<commit_before>package migrate

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"strings"
)

func bindata_read(data []byte, name string) ([]byte, error) {
	gz, err := gzip.NewReader(bytes.NewBuffer(data))
	if err != nil {
		return nil, fmt.Errorf("Read %q: %v", name, err)
	}

	var buf bytes.Buffer
	_, err = io.Copy(&buf, gz)
	gz.Close()

	if err != nil {
		return nil, fmt.Errorf("Read %q: %v", name, err)
	}

	return buf.Bytes(), nil
}

func test_migrations_1_initial_sql() ([]byte, error) {
	return bindata_read([]byte{
		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x00, 0xff, 0x8c, 0xcd,
		0x3d, 0x0e, 0x82, 0x40, 0x10, 0x05, 0xe0, 0x7e, 0x4e, 0xf1, 0x3a, 0x34,
		0x86, 0x13, 0x50, 0xa1, 0xd0, 0x91, 0xa8, 0x08, 0x07, 0x40, 0x76, 0x22,
		0x13, 0xd7, 0xdd, 0x09, 0xac, 0xc1, 0xe3, 0xbb, 0xc4, 0x68, 0xb4, 0xb3,
		0x7c, 0x6f, 0x7e, 0xbe, 0x34, 0xc5, 0xe6, 0x26, 0x97, 0xb1, 0x0b, 0x8c,
		0x56, 0x29, 0xc6, 0xd3, 0xb1, 0x82, 0x38, 0x4c, 0xdc, 0x07, 0xf1, 0x0e,
		0x49, 0xab, 0x09, 0x64, 0x02, 0x3f, 0xb8, 0xbf, 0x07, 0x36, 0x98, 0x07,
		0x76, 0x08, 0x43, 0xac, 0x5e, 0x77, 0xcb, 0x52, 0x0c, 0x9d, 0xaa, 0x15,
		0x36, 0xb4, 0xab, 0xcb, 0xbc, 0x29, 0xd1, 0xe4, 0xdb, 0xaa, 0x84, 0xb2,
		0x57, 0xcb, 0x58, 0x89, 0x89, 0x2f, 0xc3, 0x3a, 0x23, 0xa2, 0x6f, 0xb0,
		0xf0, 0xb3, 0x7b, 0x93, 0x1f, 0x6f, 0x29, 0xff, 0x12, 0x47, 0x6f, 0x6d,
		0x9c, 0x9e, 0xbb, 0xfe, 0x4a, 0x45, 0xbd, 0x3f, 0xfc, 0x98, 0x19, 0x3d,
		0x03, 0x00, 0x00, 0xff, 0xff, 0x0d, 0x70, 0x5e, 0xf9, 0xda, 0x00, 0x00,
		0x00,
	},
		"test-migrations/1_initial.sql",
	)
}

func test_migrations_2_record_sql() ([]byte, error) {
	return bindata_read([]byte{
		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x00, 0xff, 0xd2, 0xd5,
		0x55, 0xd0, 0xce, 0xcd, 0x4c, 0x2f, 0x4a, 0x2c, 0x49, 0x55, 0x08, 0x2d,
		0xe0, 0xf2, 0xf4, 0x0b, 0x76, 0x0d, 0x0a, 0x51, 0xf0, 0xf4, 0x0b, 0xf1,
		0x57, 0x28, 0x48, 0xcd, 0x2f, 0xc8, 0x49, 0x55, 0xd0, 0xc8, 0x4c, 0xd1,
		0x54, 0x08, 0x73, 0xf4, 0x09, 0x75, 0x0d, 0x56, 0xd0, 0x30, 0xd4, 0xb4,
		0xe6, 0xe2, 0x42, 0xd6, 0xe3, 0x92, 0x5f, 0x9e, 0xc7, 0xe5, 0xe2, 0xea,
		0xe3, 0x1a, 0xe2, 0xaa, 0xe0, 0x16, 0xe4, 0xef, 0x0b, 0xd3, 0x15, 0xee,
		0xe1, 0x1a, 0xe4, 0xaa, 0x90, 0x99, 0x62, 0x6b, 0x68, 0xcd, 0x05, 0x08,
		0x00, 0x00, 0xff, 0xff, 0xf4, 0x3a, 0x7b, 0xae, 0x64, 0x00, 0x00, 0x00,
	},
		"test-migrations/2_record.sql",
	)
}

// Asset loads and returns the asset for the given name.
// It returns an error if the asset could not be found or
// could not be loaded.
func Asset(name string) ([]byte, error) {
	cannonicalName := strings.Replace(name, "\\", "/", -1)
	if f, ok := _bindata[cannonicalName]; ok {
		return f()
	}
	return nil, fmt.Errorf("Asset %s not found", name)
}

// AssetNames returns the names of the assets.
func AssetNames() []string {
	names := make([]string, 0, len(_bindata))
	for name := range _bindata {
		names = append(names, name)
	}
	return names
}

// _bindata is a table, holding each asset generator, mapped to its name.
var _bindata = map[string]func() ([]byte, error){
	"test-migrations/1_initial.sql": test_migrations_1_initial_sql,
	"test-migrations/2_record.sql": test_migrations_2_record_sql,
}
// AssetDir returns the file names below a certain
// directory embedded in the file by go-bindata.
// For example if you run go-bindata on data/... and data contains the
// following hierarchy:
//     data/
//       foo.txt
//       img/
//         a.png
//         b.png
// then AssetDir("data") would return []string{"foo.txt", "img"}
// AssetDir("data/img") would return []string{"a.png", "b.png"}
// AssetDir("foo.txt") and AssetDir("notexist") would return an error
// AssetDir("") will return []string{"data"}.
func AssetDir(name string) ([]string, error) {
	node := _bintree
	if len(name) != 0 {
		cannonicalName := strings.Replace(name, "\\", "/", -1)
		pathList := strings.Split(cannonicalName, "/")
		for _, p := range pathList {
			node = node.Children[p]
			if node == nil {
				return nil, fmt.Errorf("Asset %s not found", name)
			}
		}
	}
	if node.Func != nil {
		return nil, fmt.Errorf("Asset %s not found", name)
	}
	rv := make([]string, 0, len(node.Children))
	for name := range node.Children {
		rv = append(rv, name)
	}
	return rv, nil
}

type _bintree_t struct {
	Func func() ([]byte, error)
	Children map[string]*_bintree_t
}
var _bintree = &_bintree_t{nil, map[string]*_bintree_t{
	"test-migrations": &_bintree_t{nil, map[string]*_bintree_t{
		"1_initial.sql": &_bintree_t{test_migrations_1_initial_sql, map[string]*_bintree_t{
		}},
		"2_record.sql": &_bintree_t{test_migrations_2_record_sql, map[string]*_bintree_t{
		}},
	}},
}}
<commit_msg>Handle error on gz.Close()<commit_after>package migrate

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"strings"
)

func bindata_read(data []byte, name string) ([]byte, error) {
	gz, err := gzip.NewReader(bytes.NewBuffer(data))
	if err != nil {
		return nil, fmt.Errorf("Read %q: %v", name, err)
	}

	var buf bytes.Buffer
	_, err = io.Copy(&buf, gz)
	if err := gz.Close(); err != nil {
		return nil, err
	}

	if err != nil {
		return nil, fmt.Errorf("Read %q: %v", name, err)
	}

	return buf.Bytes(), nil
}

func test_migrations_1_initial_sql() ([]byte, error) {
	return bindata_read([]byte{
		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x00, 0xff, 0x8c, 0xcd,
		0x3d, 0x0e, 0x82, 0x40, 0x10, 0x05, 0xe0, 0x7e, 0x4e, 0xf1, 0x3a, 0x34,
		0x86, 0x13, 0x50, 0xa1, 0xd0, 0x91, 0xa8, 0x08, 0x07, 0x40, 0x76, 0x22,
		0x13, 0xd7, 0xdd, 0x09, 0xac, 0xc1, 0xe3, 0xbb, 0xc4, 0x68, 0xb4, 0xb3,
		0x7c, 0x6f, 0x7e, 0xbe, 0x34, 0xc5, 0xe6, 0x26, 0x97, 0xb1, 0x0b, 0x8c,
		0x56, 0x29, 0xc6, 0xd3, 0xb1, 0x82, 0x38, 0x4c, 0xdc, 0x07, 0xf1, 0x0e,
		0x49, 0xab, 0x09, 0x64, 0x02, 0x3f, 0xb8, 0xbf, 0x07, 0x36, 0x98, 0x07,
		0x76, 0x08, 0x43, 0xac, 0x5e, 0x77, 0xcb, 0x52, 0x0c, 0x9d, 0xaa, 0x15,
		0x36, 0xb4, 0xab, 0xcb, 0xbc, 0x29, 0xd1, 0xe4, 0xdb, 0xaa, 0x84, 0xb2,
		0x57, 0xcb, 0x58, 0x89, 0x89, 0x2f, 0xc3, 0x3a, 0x23, 0xa2, 0x6f, 0xb0,
		0xf0, 0xb3, 0x7b, 0x93, 0x1f, 0x6f, 0x29, 0xff, 0x12, 0x47, 0x6f, 0x6d,
		0x9c, 0x9e, 0xbb, 0xfe, 0x4a, 0x45, 0xbd, 0x3f, 0xfc, 0x98, 0x19, 0x3d,
		0x03, 0x00, 0x00, 0xff, 0xff, 0x0d, 0x70, 0x5e, 0xf9, 0xda, 0x00, 0x00,
		0x00,
	},
		"test-migrations/1_initial.sql",
	)
}

func test_migrations_2_record_sql() ([]byte, error) {
	return bindata_read([]byte{
		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x00, 0xff, 0xd2, 0xd5,
		0x55, 0xd0, 0xce, 0xcd, 0x4c, 0x2f, 0x4a, 0x2c, 0x49, 0x55, 0x08, 0x2d,
		0xe0, 0xf2, 0xf4, 0x0b, 0x76, 0x0d, 0x0a, 0x51, 0xf0, 0xf4, 0x0b, 0xf1,
		0x57, 0x28, 0x48, 0xcd, 0x2f, 0xc8, 0x49, 0x55, 0xd0, 0xc8, 0x4c, 0xd1,
		0x54, 0x08, 0x73, 0xf4, 0x09, 0x75, 0x0d, 0x56, 0xd0, 0x30, 0xd4, 0xb4,
		0xe6, 0xe2, 0x42, 0xd6, 0xe3, 0x92, 0x5f, 0x9e, 0xc7, 0xe5, 0xe2, 0xea,
		0xe3, 0x1a, 0xe2, 0xaa, 0xe0, 0x16, 0xe4, 0xef, 0x0b, 0xd3, 0x15, 0xee,
		0xe1, 0x1a, 0xe4, 0xaa, 0x90, 0x99, 0x62, 0x6b, 0x68, 0xcd, 0x05, 0x08,
		0x00, 0x00, 0xff, 0xff, 0xf4, 0x3a, 0x7b, 0xae, 0x64, 0x00, 0x00, 0x00,
	},
		"test-migrations/2_record.sql",
	)
}

// Asset loads and returns the asset for the given name.
// It returns an error if the asset could not be found or
// could not be loaded.
func Asset(name string) ([]byte, error) {
	cannonicalName := strings.Replace(name, "\\", "/", -1)
	if f, ok := _bindata[cannonicalName]; ok {
		return f()
	}
	return nil, fmt.Errorf("Asset %s not found", name)
}

// AssetNames returns the names of the assets.
func AssetNames() []string {
	names := make([]string, 0, len(_bindata))
	for name := range _bindata {
		names = append(names, name)
	}
	return names
}

// _bindata is a table, holding each asset generator, mapped to its name.
var _bindata = map[string]func() ([]byte, error){
	"test-migrations/1_initial.sql": test_migrations_1_initial_sql,
	"test-migrations/2_record.sql": test_migrations_2_record_sql,
}
// AssetDir returns the file names below a certain
// directory embedded in the file by go-bindata.
// For example if you run go-bindata on data/... and data contains the
// following hierarchy:
//     data/
//       foo.txt
//       img/
//         a.png
//         b.png
// then AssetDir("data") would return []string{"foo.txt", "img"}
// AssetDir("data/img") would return []string{"a.png", "b.png"}
// AssetDir("foo.txt") and AssetDir("notexist") would return an error
// AssetDir("") will return []string{"data"}.
func AssetDir(name string) ([]string, error) {
	node := _bintree
	if len(name) != 0 {
		cannonicalName := strings.Replace(name, "\\", "/", -1)
		pathList := strings.Split(cannonicalName, "/")
		for _, p := range pathList {
			node = node.Children[p]
			if node == nil {
				return nil, fmt.Errorf("Asset %s not found", name)
			}
		}
	}
	if node.Func != nil {
		return nil, fmt.Errorf("Asset %s not found", name)
	}
	rv := make([]string, 0, len(node.Children))
	for name := range node.Children {
		rv = append(rv, name)
	}
	return rv, nil
}

type _bintree_t struct {
	Func func() ([]byte, error)
	Children map[string]*_bintree_t
}
var _bintree = &_bintree_t{nil, map[string]*_bintree_t{
	"test-migrations": &_bintree_t{nil, map[string]*_bintree_t{
		"1_initial.sql": &_bintree_t{test_migrations_1_initial_sql, map[string]*_bintree_t{
		}},
		"2_record.sql": &_bintree_t{test_migrations_2_record_sql, map[string]*_bintree_t{
		}},
	}},
}}
<|endoftext|>


<commit_before>package mutator

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	admission "k8s.io/api/admission/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"

	"github.com/datawire/ambassador/v2/pkg/kates"
	"github.com/datawire/dlib/dlog"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/managerutil"
	"github.com/telepresenceio/telepresence/v2/pkg/install"
)

var podResource = metav1.GroupVersionResource{Version: "v1", Group: "", Resource: "pods"}
var findMatchingService = install.FindMatchingService

func agentInjector(ctx context.Context, req *admission.AdmissionRequest) ([]patchOperation, error) {
	// This handler should only get called on Pod objects as per the MutatingWebhookConfiguration in the YAML file.
	// Pod objects are immutable, hence we only care about the CREATE event.
	// Applying patches to Pods instead of Deployments means we don't have side effects on
	// user-managed Deployments and Services. It also means we don't have to manage update flows
	// such as removing or updating the sidecar in Deployments objects... a new Pod just gets created instead!

	// If (for whatever reason) this handler is invoked on an object different than a Pod,
	// issue a log message and let the object request pass through.
	if req.Resource != podResource {
		dlog.Debugf(ctx, "expect resource to be %s, got %s; skipping", podResource, req.Resource)
		return nil, nil
	}

	// Parse the Pod object.
	raw := req.Object.Raw
	pod := corev1.Pod{}
	if _, _, err := universalDeserializer.Decode(raw, nil, &pod); err != nil {
		return nil, fmt.Errorf("could not deserialize pod object: %v", err)
	}

	podNamespace := pod.Namespace
	if podNamespace == "" {
		// It is very probable the pod was not yet assigned a namespace,
		// in which case we should use the AdmissionRequest namespace.
		podNamespace = req.Namespace
	}
	podName := pod.Name
	if podName == "" {
		// It is very probable the pod was not yet assigned a name,
		// in which case we should use the metadata generated name.
		podName = pod.ObjectMeta.GenerateName
	}

	// Validate traffic-agent injection preconditions.
	refPodName := fmt.Sprintf("%s.%s", podName, podNamespace)
	if podName == "" || podNamespace == "" {
		dlog.Debugf(ctx, "Unable to extract pod name and/or namespace (got %q); skipping", refPodName)
		return nil, nil
	}

	if pod.Annotations[install.InjectAnnotation] != "enabled" {
		dlog.Debugf(ctx, `The %s pod has not enabled %s container injection through %q annotation; skipping`,
			refPodName, install.AgentContainerName, install.InjectAnnotation)
		return nil, nil
	}

	// Make the kates client available in the context
	// TODO: Use the kubernetes SharedInformerFactory instead
	client, err := kates.NewClient(kates.ClientConfig{})
	if err != nil {
		return nil, fmt.Errorf("failed to create new kates client: %w", err)
	}

	svc, err := findMatchingService(ctx, client, "", "", podNamespace, pod.Labels)
	if err != nil {
		dlog.Error(ctx, err)
		return nil, nil
	}

	// The ServicePortAnnotation is expected to contain a string that identifies the service port.
	portNameOrNumber := pod.Annotations[install.ServicePortAnnotation]
	servicePort, appContainer, containerPortIndex, err := install.FindMatchingPort(pod.Spec.Containers, portNameOrNumber, svc)
	if err != nil {
		dlog.Error(ctx, err)
		return nil, nil
	}
	if appContainer.Name == install.AgentContainerName {
		dlog.Infof(ctx, "service %s/%s is already pointing at agent container %s; skipping", svc.Namespace, svc.Name, appContainer.Name)
		return nil, nil
	}

	env := managerutil.GetEnv(ctx)
	ports := appContainer.Ports
	for i := range ports {
		if ports[i].ContainerPort == env.AgentPort {
			dlog.Infof(ctx, "the %s pod container %s is exposing the same port (%d) as the %s sidecar; skipping",
				refPodName, appContainer.Name, env.AgentPort, install.AgentContainerName)
			return nil, nil
		}
	}

	if svc.Spec.ClusterIP == "None" {
		return nil, fmt.Errorf("intercepts of headless service: %s.%s won't work "+
			"see https://github.com/telepresenceio/telepresence/issues/1632",
			svc.Name, svc.Namespace)
	}

	var appPort corev1.ContainerPort
	switch {
	case containerPortIndex >= 0:
		appPort = appContainer.Ports[containerPortIndex]
	case servicePort.TargetPort.Type == intstr.Int:
		appPort = corev1.ContainerPort{
			Protocol:      servicePort.Protocol,
			ContainerPort: servicePort.TargetPort.IntVal,
		}
	default:
		// This really shouldn't have happened: the target port is a string, but we weren't able to
		// find a corresponding container port. This should've been caught in FindMatchingPort, but in
		// case it isn't, just return an error.
		return nil, fmt.Errorf("container port unexpectedly not found in %s", refPodName)
	}

	// Create patch operations to add the traffic-agent sidecar
	dlog.Infof(ctx, "Injecting %s into pod %s", install.AgentContainerName, refPodName)

	var patches []patchOperation
	patches, err = addAgentContainer(ctx, svc, &pod, servicePort, appContainer, &appPort, podName, podNamespace, patches)
	if err != nil {
		return nil, err
	}
	if servicePort.TargetPort.Type == intstr.Int {
		patches = addInitContainer(ctx, &pod, servicePort, &appPort, patches)
	} else {
		patches = hidePorts(&pod, appContainer, servicePort.TargetPort.StrVal, patches)
	}
	patches = addAgentVolume(&pod, patches)
	return patches, nil
}

func addInitContainer(ctx context.Context, pod *corev1.Pod, svcPort *corev1.ServicePort, appPort *corev1.ContainerPort, patches []patchOperation) []patchOperation {
	env := managerutil.GetEnv(ctx)
	proto := svcPort.Protocol
	if proto == "" {
		proto = appPort.Protocol
	}
	containerPort := corev1.ContainerPort{
		Protocol:      proto,
		ContainerPort: env.AgentPort,
	}
	container := install.InitContainer(
		env.AgentRegistry+"/"+env.AgentImage,
		containerPort,
		int(appPort.ContainerPort),
	)

	if pod.Spec.InitContainers == nil {
		patches = append(patches, patchOperation{
			Op:    "add",
			Path:  "/spec/initContainers",
			Value: []corev1.Container{},
		})
	} else {
		for i, container := range pod.Spec.InitContainers {
			if container.Name == install.InitContainerName {
				if i == len(pod.Spec.InitContainers)-1 {
					return patches
				}
				// If the container isn't the last one, remove it so it can be appended at the end.
				patches = append(patches, patchOperation{
					Op:   "remove",
					Path: fmt.Sprintf("/spec/initContainers/%d", i),
				})
			}
		}
	}

	return append(patches, patchOperation{
		Op:    "add",
		Path:  "/spec/initContainers/-",
		Value: container,
	})
}

func addAgentVolume(pod *corev1.Pod, patches []patchOperation) []patchOperation {
	for _, vol := range pod.Spec.Volumes {
		if vol.Name == install.AgentAnnotationVolumeName {
			return patches
		}
	}
	return append(patches, patchOperation{
		Op:    "add",
		Path:  "/spec/volumes/-",
		Value: install.AgentVolume(),
	})
}

// addAgentContainer creates a patch operation to add the traffic-agent container
func addAgentContainer(
	ctx context.Context,
	svc *corev1.Service,
	pod *corev1.Pod,
	svcPort *corev1.ServicePort,
	appContainer *corev1.Container,
	appPort *corev1.ContainerPort,
	podName, namespace string,
	patches []patchOperation,
) ([]patchOperation, error) {
	env := managerutil.GetEnv(ctx)

	refPodName := podName + "." + namespace
	for _, container := range pod.Spec.Containers {
		if container.Name == install.AgentContainerName {
			dlog.Infof(ctx, "Pod %s already has container %s", refPodName, install.AgentContainerName)
			return patches, nil
		}
	}

	dlog.Debugf(ctx, "using service %q port %q when intercepting %s",
		svc.Name,
		func() string {
			if svcPort.Name != "" {
				return svcPort.Name
			}
			return strconv.Itoa(int(svcPort.Port))
		}(), refPodName)

	agentName := podName
	if strings.HasSuffix(agentName, "-") {
		// Transform a generated name "my-echo-697464c6c5-" into an agent service name "my-echo"
		tokens := strings.Split(podName, "-")
		agentName = strings.Join(tokens[:len(tokens)-2], "-")
	}

	proto := svcPort.Protocol
	if proto == "" {
		proto = appPort.Protocol
	}
	containerPort := corev1.ContainerPort{
		Protocol:      proto,
		ContainerPort: env.AgentPort,
	}
	if svcPort.TargetPort.Type == intstr.String {
		containerPort.Name = svcPort.TargetPort.StrVal
	}
	patches = append(patches, patchOperation{
		Op:   "add",
		Path: "/spec/containers/-",
		Value: install.AgentContainer(
			agentName,
			env.AgentRegistry+"/"+env.AgentImage,
			appContainer,
			containerPort,
			int(appPort.ContainerPort),
			env.ManagerNamespace)})

	return patches, nil
}

// hidePorts  will replace the symbolic name of a container port with a generated name. It will perform
// the same replacement on all references to that port from the probes of the container
func hidePorts(pod *corev1.Pod, cn *corev1.Container, portName string, patches []patchOperation) []patchOperation {
	cns := pod.Spec.Containers
	var containerPath string
	for i := range cns {
		if &cns[i] == cn {
			containerPath = fmt.Sprintf("/spec/containers/%d", i)
			break
		}
	}

	hiddenPortName := install.HiddenPortName(portName, 0)
	hidePort := func(path string) {
		patches = append(patches, patchOperation{
			Op:    "replace",
			Path:  fmt.Sprintf("%s/%s", containerPath, path),
			Value: hiddenPortName,
		})
	}

	for i, p := range cn.Ports {
		if p.Name == portName {
			hidePort(fmt.Sprintf("ports/%d/name", i))
			break
		}
	}

	probes := []*corev1.Probe{cn.LivenessProbe, cn.ReadinessProbe, cn.StartupProbe}
	probeNames := []string{"livenessProbe/", "readinessProbe/", "startupProbe/"}

	for i, probe := range probes {
		if probe == nil {
			continue
		}
		if h := probe.HTTPGet; h != nil && h.Port.StrVal == portName {
			hidePort(probeNames[i] + "httpGet/port")
		}
		if t := probe.TCPSocket; t != nil && t.Port.StrVal == portName {
			hidePort(probeNames[i] + "tcpSocket/port")
		}
	}
	return patches
}
<commit_msg>Allow webhook injector to inject into headless services<commit_after>package mutator

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	admission "k8s.io/api/admission/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"

	"github.com/datawire/ambassador/v2/pkg/kates"
	"github.com/datawire/dlib/dlog"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/managerutil"
	"github.com/telepresenceio/telepresence/v2/pkg/install"
)

var podResource = metav1.GroupVersionResource{Version: "v1", Group: "", Resource: "pods"}
var findMatchingService = install.FindMatchingService

func agentInjector(ctx context.Context, req *admission.AdmissionRequest) ([]patchOperation, error) {
	// This handler should only get called on Pod objects as per the MutatingWebhookConfiguration in the YAML file.
	// Pod objects are immutable, hence we only care about the CREATE event.
	// Applying patches to Pods instead of Deployments means we don't have side effects on
	// user-managed Deployments and Services. It also means we don't have to manage update flows
	// such as removing or updating the sidecar in Deployments objects... a new Pod just gets created instead!

	// If (for whatever reason) this handler is invoked on an object different than a Pod,
	// issue a log message and let the object request pass through.
	if req.Resource != podResource {
		dlog.Debugf(ctx, "expect resource to be %s, got %s; skipping", podResource, req.Resource)
		return nil, nil
	}

	// Parse the Pod object.
	raw := req.Object.Raw
	pod := corev1.Pod{}
	if _, _, err := universalDeserializer.Decode(raw, nil, &pod); err != nil {
		return nil, fmt.Errorf("could not deserialize pod object: %v", err)
	}

	podNamespace := pod.Namespace
	if podNamespace == "" {
		// It is very probable the pod was not yet assigned a namespace,
		// in which case we should use the AdmissionRequest namespace.
		podNamespace = req.Namespace
	}
	podName := pod.Name
	if podName == "" {
		// It is very probable the pod was not yet assigned a name,
		// in which case we should use the metadata generated name.
		podName = pod.ObjectMeta.GenerateName
	}

	// Validate traffic-agent injection preconditions.
	refPodName := fmt.Sprintf("%s.%s", podName, podNamespace)
	if podName == "" || podNamespace == "" {
		dlog.Debugf(ctx, "Unable to extract pod name and/or namespace (got %q); skipping", refPodName)
		return nil, nil
	}

	if pod.Annotations[install.InjectAnnotation] != "enabled" {
		dlog.Debugf(ctx, `The %s pod has not enabled %s container injection through %q annotation; skipping`,
			refPodName, install.AgentContainerName, install.InjectAnnotation)
		return nil, nil
	}

	// Make the kates client available in the context
	// TODO: Use the kubernetes SharedInformerFactory instead
	client, err := kates.NewClient(kates.ClientConfig{})
	if err != nil {
		return nil, fmt.Errorf("failed to create new kates client: %w", err)
	}

	svc, err := findMatchingService(ctx, client, "", "", podNamespace, pod.Labels)
	if err != nil {
		dlog.Error(ctx, err)
		return nil, nil
	}

	// The ServicePortAnnotation is expected to contain a string that identifies the service port.
	portNameOrNumber := pod.Annotations[install.ServicePortAnnotation]
	servicePort, appContainer, containerPortIndex, err := install.FindMatchingPort(pod.Spec.Containers, portNameOrNumber, svc)
	if err != nil {
		dlog.Error(ctx, err)
		return nil, nil
	}
	if appContainer.Name == install.AgentContainerName {
		dlog.Infof(ctx, "service %s/%s is already pointing at agent container %s; skipping", svc.Namespace, svc.Name, appContainer.Name)
		return nil, nil
	}

	env := managerutil.GetEnv(ctx)
	ports := appContainer.Ports
	for i := range ports {
		if ports[i].ContainerPort == env.AgentPort {
			dlog.Infof(ctx, "the %s pod container %s is exposing the same port (%d) as the %s sidecar; skipping",
				refPodName, appContainer.Name, env.AgentPort, install.AgentContainerName)
			return nil, nil
		}
	}

	var appPort corev1.ContainerPort
	switch {
	case containerPortIndex >= 0:
		appPort = appContainer.Ports[containerPortIndex]
	case servicePort.TargetPort.Type == intstr.Int:
		appPort = corev1.ContainerPort{
			Protocol:      servicePort.Protocol,
			ContainerPort: servicePort.TargetPort.IntVal,
		}
	default:
		// This really shouldn't have happened: the target port is a string, but we weren't able to
		// find a corresponding container port. This should've been caught in FindMatchingPort, but in
		// case it isn't, just return an error.
		return nil, fmt.Errorf("container port unexpectedly not found in %s", refPodName)
	}

	// Create patch operations to add the traffic-agent sidecar
	dlog.Infof(ctx, "Injecting %s into pod %s", install.AgentContainerName, refPodName)

	var patches []patchOperation
	patches, err = addAgentContainer(ctx, svc, &pod, servicePort, appContainer, &appPort, podName, podNamespace, patches)
	if err != nil {
		return nil, err
	}
	if servicePort.TargetPort.Type == intstr.Int || svc.Spec.ClusterIP == "None" {
		patches = addInitContainer(ctx, &pod, servicePort, &appPort, patches)
	} else {
		patches = hidePorts(&pod, appContainer, servicePort.TargetPort.StrVal, patches)
	}
	patches = addAgentVolume(&pod, patches)
	return patches, nil
}

func addInitContainer(ctx context.Context, pod *corev1.Pod, svcPort *corev1.ServicePort, appPort *corev1.ContainerPort, patches []patchOperation) []patchOperation {
	env := managerutil.GetEnv(ctx)
	proto := svcPort.Protocol
	if proto == "" {
		proto = appPort.Protocol
	}
	containerPort := corev1.ContainerPort{
		Protocol:      proto,
		ContainerPort: env.AgentPort,
	}
	container := install.InitContainer(
		env.AgentRegistry+"/"+env.AgentImage,
		containerPort,
		int(appPort.ContainerPort),
	)

	if pod.Spec.InitContainers == nil {
		patches = append(patches, patchOperation{
			Op:    "add",
			Path:  "/spec/initContainers",
			Value: []corev1.Container{},
		})
	} else {
		for i, container := range pod.Spec.InitContainers {
			if container.Name == install.InitContainerName {
				if i == len(pod.Spec.InitContainers)-1 {
					return patches
				}
				// If the container isn't the last one, remove it so it can be appended at the end.
				patches = append(patches, patchOperation{
					Op:   "remove",
					Path: fmt.Sprintf("/spec/initContainers/%d", i),
				})
			}
		}
	}

	return append(patches, patchOperation{
		Op:    "add",
		Path:  "/spec/initContainers/-",
		Value: container,
	})
}

func addAgentVolume(pod *corev1.Pod, patches []patchOperation) []patchOperation {
	for _, vol := range pod.Spec.Volumes {
		if vol.Name == install.AgentAnnotationVolumeName {
			return patches
		}
	}
	return append(patches, patchOperation{
		Op:    "add",
		Path:  "/spec/volumes/-",
		Value: install.AgentVolume(),
	})
}

// addAgentContainer creates a patch operation to add the traffic-agent container
func addAgentContainer(
	ctx context.Context,
	svc *corev1.Service,
	pod *corev1.Pod,
	svcPort *corev1.ServicePort,
	appContainer *corev1.Container,
	appPort *corev1.ContainerPort,
	podName, namespace string,
	patches []patchOperation,
) ([]patchOperation, error) {
	env := managerutil.GetEnv(ctx)

	refPodName := podName + "." + namespace
	for _, container := range pod.Spec.Containers {
		if container.Name == install.AgentContainerName {
			dlog.Infof(ctx, "Pod %s already has container %s", refPodName, install.AgentContainerName)
			return patches, nil
		}
	}

	dlog.Debugf(ctx, "using service %q port %q when intercepting %s",
		svc.Name,
		func() string {
			if svcPort.Name != "" {
				return svcPort.Name
			}
			return strconv.Itoa(int(svcPort.Port))
		}(), refPodName)

	agentName := podName
	if strings.HasSuffix(agentName, "-") {
		// Transform a generated name "my-echo-697464c6c5-" into an agent service name "my-echo"
		tokens := strings.Split(podName, "-")
		agentName = strings.Join(tokens[:len(tokens)-2], "-")
	}

	proto := svcPort.Protocol
	if proto == "" {
		proto = appPort.Protocol
	}
	containerPort := corev1.ContainerPort{
		Protocol:      proto,
		ContainerPort: env.AgentPort,
	}
	if svcPort.TargetPort.Type == intstr.String {
		containerPort.Name = svcPort.TargetPort.StrVal
	}
	patches = append(patches, patchOperation{
		Op:   "add",
		Path: "/spec/containers/-",
		Value: install.AgentContainer(
			agentName,
			env.AgentRegistry+"/"+env.AgentImage,
			appContainer,
			containerPort,
			int(appPort.ContainerPort),
			env.ManagerNamespace)})

	return patches, nil
}

// hidePorts  will replace the symbolic name of a container port with a generated name. It will perform
// the same replacement on all references to that port from the probes of the container
func hidePorts(pod *corev1.Pod, cn *corev1.Container, portName string, patches []patchOperation) []patchOperation {
	cns := pod.Spec.Containers
	var containerPath string
	for i := range cns {
		if &cns[i] == cn {
			containerPath = fmt.Sprintf("/spec/containers/%d", i)
			break
		}
	}

	hiddenPortName := install.HiddenPortName(portName, 0)
	hidePort := func(path string) {
		patches = append(patches, patchOperation{
			Op:    "replace",
			Path:  fmt.Sprintf("%s/%s", containerPath, path),
			Value: hiddenPortName,
		})
	}

	for i, p := range cn.Ports {
		if p.Name == portName {
			hidePort(fmt.Sprintf("ports/%d/name", i))
			break
		}
	}

	probes := []*corev1.Probe{cn.LivenessProbe, cn.ReadinessProbe, cn.StartupProbe}
	probeNames := []string{"livenessProbe/", "readinessProbe/", "startupProbe/"}

	for i, probe := range probes {
		if probe == nil {
			continue
		}
		if h := probe.HTTPGet; h != nil && h.Port.StrVal == portName {
			hidePort(probeNames[i] + "httpGet/port")
		}
		if t := probe.TCPSocket; t != nil && t.Port.StrVal == portName {
			hidePort(probeNames[i] + "tcpSocket/port")
		}
	}
	return patches
}
<|endoftext|>


<commit_before>package vault

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

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

const (
	// expirationSubPath is the sub-path used for the expiration manager
	// view. This is nested under the system view.
	expirationSubPath = "expire/"

	// maxRevokeAttempts limits how many revoke attempts are made
	maxRevokeAttempts = 6

	// revokeRetryBase is a baseline retry time
	revokeRetryBase = 10 * time.Second

	// minRevokeDelay is used to prevent an instant revoke on restore
	minRevokeDelay = 5 * time.Second
)

// ExpirationManager is used by the Core to manage leases. Secrets
// can provide a lease, meaning that they can be renewed or revoked.
// If a secret is not renewed in timely manner, it may be expired, and
// the ExpirationManager will handle doing automatic revocation.
type ExpirationManager struct {
	router *Router
	view   *BarrierView
	logger *log.Logger

	pending     map[string]*time.Timer
	pendingLock sync.Mutex
}

// NewExpirationManager creates a new ExpirationManager that is backed
// using a given view, and uses the provided router for revocation.
func NewExpirationManager(router *Router, view *BarrierView, logger *log.Logger) *ExpirationManager {
	if logger == nil {
		logger = log.New(os.Stderr, "", log.LstdFlags)
	}
	exp := &ExpirationManager{
		router:  router,
		view:    view,
		logger:  logger,
		pending: make(map[string]*time.Timer),
	}
	return exp
}

// setupExpiration is invoked after we've loaded the mount table to
// initialize the expiration manager
func (c *Core) setupExpiration() error {
	// Create a sub-view
	view := c.systemView.SubView(expirationSubPath)

	// Create the manager
	mgr := NewExpirationManager(c.router, view, c.logger)
	c.expiration = mgr

	// Restore the existing state
	if err := c.expiration.Restore(); err != nil {
		return fmt.Errorf("expiration state restore failed: %v", err)
	}
	return nil
}

// stopExpiration is used to stop the expiration manager before
// sealing the Vault.
func (c *Core) stopExpiration() error {
	if err := c.expiration.Stop(); err != nil {
		return err
	}
	c.expiration = nil
	return nil
}

// Restore is used to recover the lease states when starting.
// This is used after starting the vault.
func (m *ExpirationManager) Restore() error {
	m.pendingLock.Lock()
	defer m.pendingLock.Unlock()

	// Accumulate existing leases
	var existing []string
	cb := func(path string) {
		existing = append(existing, path)
	}

	// Scan for all the leases
	if err := ScanView(m.view, cb); err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Restore each key
	for _, vaultID := range existing {
		// Load the entry
		le, err := m.loadEntry(vaultID)
		if err != nil {
			return err
		}

		// If there is no entry, nothing to restore
		if le == nil {
			continue
		}

		// Determine the remaining time to expiration
		expires := le.ExpireTime.Sub(time.Now().UTC())
		if expires <= 0 {
			expires = minRevokeDelay
		}

		// Setup revocation timer
		m.pending[le.VaultID] = time.AfterFunc(expires, func() {
			m.expireID(le.VaultID)
		})
	}
	m.logger.Printf("[INFO] expire: restored %d leases", len(m.pending))
	return nil
}

// Stop is used to prevent further automatic revocations.
// This must be called before sealing the view.
func (m *ExpirationManager) Stop() error {
	// Stop all the pending expiration timers
	m.pendingLock.Lock()
	for _, timer := range m.pending {
		timer.Stop()
	}
	m.pending = make(map[string]*time.Timer)
	m.pendingLock.Unlock()
	return nil
}

// Revoke is used to revoke a secret named by the given vaultID
func (m *ExpirationManager) Revoke(vaultID string) error {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return err
	}

	// If there is no entry, nothing to revoke
	if le == nil {
		return nil
	}

	// Revoke the entry
	if err := m.revokeEntry(le); err != nil {
		return err
	}

	// Delete the entry
	if err := m.deleteEntry(vaultID); err != nil {
		return err
	}

	// Clear the expiration handler
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Stop()
		delete(m.pending, vaultID)
	}
	m.pendingLock.Unlock()
	return nil
}

// RevokePrefix is used to revoke all secrets with a given prefix.
// The prefix maps to that of the mount table to make this simpler
// to reason about.
func (m *ExpirationManager) RevokePrefix(prefix string) error {
	// Ensure there is a trailing slash
	if !strings.HasSuffix(prefix, "/") {
		prefix = prefix + "/"
	}

	// Accumulate existing leases
	var existing []string
	cb := func(path string) {
		existing = append(existing, path)
	}

	// Scan for all the leases in the prefix
	sub := m.view.SubView(prefix)
	if err := ScanView(sub, cb); err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Revoke all the keys
	for idx, suffix := range existing {
		vaultID := prefix + suffix
		if err := m.Revoke(vaultID); err != nil {
			return fmt.Errorf("failed to revoke '%s' (%d / %d): %v",
				vaultID, idx+1, len(existing), err)
		}
	}
	return nil
}

// Renew is used to renew a secret using the given vaultID
// and a renew interval. The increment may be ignored.
func (m *ExpirationManager) Renew(vaultID string, increment time.Duration) (*logical.Response, error) {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return nil, err
	}

	// If there is no entry, cannot review
	if le == nil {
		return nil, fmt.Errorf("lease not found")
	}

	// Determine if the lease is expired
	if le.ExpireTime.Before(time.Now().UTC()) {
		return nil, fmt.Errorf("lease expired")
	}

	// Attempt to renew the entry
	resp, err := m.renewEntry(le, increment)
	if err != nil {
		return nil, err
	}

	// Fast-path if there is no lease
	if resp == nil || resp.Lease == nil || !resp.IsSecret {
		return resp, nil
	}

	// Validate the lease
	if err := resp.Lease.Validate(); err != nil {
		return nil, err
	}

	// Attach the VaultID
	resp.Lease.VaultID = vaultID

	// Update the lease entry
	le.Data = resp.Data
	le.Lease = resp.Lease
	le.ExpireTime = time.Now().UTC().Add(resp.Lease.Duration)
	if err := m.persistEntry(le); err != nil {
		return nil, err
	}

	// Update the expiration time
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Reset(resp.Lease.Duration)
	}
	m.pendingLock.Unlock()

	// Return the response
	return resp, nil
}

// Register is used to take a request and response with an associated
// lease. The secret gets assigned a vaultId and the management of
// of lease is assumed by the expiration manager.
func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Response) (string, error) {
	// Ignore if there is no lease
	if resp == nil || resp.Lease == nil {
		return "", nil
	}

	// Validate the lease
	if err := resp.Lease.Validate(); err != nil {
		return "", err
	}

	// Cannot register a non-secret (e.g. a policy or configuration key)
	if !resp.IsSecret {
		return "", fmt.Errorf("cannot attach lease to non-secret")
	}

	// Create a lease entry
	now := time.Now().UTC()
	le := leaseEntry{
		VaultID:    path.Join(req.Path, generateUUID()),
		Path:       req.Path,
		Data:       resp.Data,
		Lease:      resp.Lease,
		IssueTime:  now,
		ExpireTime: now.Add(resp.Lease.Duration),
	}

	// Encode the entry
	if err := m.persistEntry(&le); err != nil {
		return "", err
	}

	// Setup revocation timer
	m.pendingLock.Lock()
	m.pending[le.VaultID] = time.AfterFunc(resp.Lease.Duration, func() {
		m.expireID(le.VaultID)
	})
	m.pendingLock.Unlock()

	// Done
	return le.VaultID, nil
}

// expireID is invoked when a given ID is expired
func (m *ExpirationManager) expireID(vaultID string) {
	// Clear from the pending expiration
	m.pendingLock.Lock()
	delete(m.pending, vaultID)
	m.pendingLock.Unlock()

	for attempt := uint(0); attempt < maxRevokeAttempts; attempt++ {
		err := m.Revoke(vaultID)
		if err == nil {
			m.logger.Printf("[INFO] expire: revoked '%s'", vaultID)
			return
		}
		m.logger.Printf("[ERR] expire: failed to revoke '%s': %v", vaultID, err)
		time.Sleep((1 << attempt) * revokeRetryBase)
	}
	m.logger.Printf("[ERR] expire: maximum revoke attempts for '%s' reached", vaultID)
}

// revokeEntry is used to attempt revocation of an internal entry
func (m *ExpirationManager) revokeEntry(le *leaseEntry) error {
	data := map[string]interface{}{
		"previous": le.Data,
	}
	req := &logical.Request{
		Operation: logical.RevokeOperation,
		Path:      le.Path,
		Data:      data,
	}
	_, err := m.router.Route(req)
	if err != nil {
		return fmt.Errorf("failed to revoke entry: %v", err)
	}
	return nil
}

// renewEntry is used to attempt renew of an internal entry
func (m *ExpirationManager) renewEntry(le *leaseEntry, increment time.Duration) (*logical.Response, error) {
	data := map[string]interface{}{
		"previous":  le.Data,
		"increment": increment,
	}
	req := &logical.Request{
		Operation: logical.RenewOperation,
		Path:      le.Path,
		Data:      data,
	}
	resp, err := m.router.Route(req)
	if err != nil {
		return nil, fmt.Errorf("failed to renew entry: %v", err)
	}
	return resp, nil
}

// loadEntry is used to read a lease entry
func (m *ExpirationManager) loadEntry(vaultID string) (*leaseEntry, error) {
	out, err := m.view.Get(vaultID)
	if err != nil {
		return nil, fmt.Errorf("failed to read lease entry: %v", err)
	}
	if out == nil {
		return nil, nil
	}
	le, err := decodeLeaseEntry(out.Value)
	if err != nil {
		return nil, fmt.Errorf("failed to decode lease entry: %v", err)
	}
	return le, nil
}

// persistEntry is used to persist a lease entry
func (m *ExpirationManager) persistEntry(le *leaseEntry) error {
	// Encode the entry
	buf, err := le.encode()
	if err != nil {
		return fmt.Errorf("failed to encode lease entry: %v", err)
	}

	// Write out to the view
	ent := logical.StorageEntry{
		Key:   le.VaultID,
		Value: buf,
	}
	if err := m.view.Put(&ent); err != nil {
		return fmt.Errorf("failed to persist lease entry: %v", err)
	}
	return nil
}

// deleteEntry is used to delete a lease entry
func (m *ExpirationManager) deleteEntry(vaultID string) error {
	if err := m.view.Delete(vaultID); err != nil {
		return fmt.Errorf("failed to delete lease entry: %v", err)
	}
	return nil
}

// leaseEntry is used to structure the values the expiration
// manager stores. This is used to handle renew and revocation.
type leaseEntry struct {
	VaultID    string                 `json:"vault_id"`
	Path       string                 `json:"path"`
	Data       map[string]interface{} `json:"data"`
	Lease      *logical.Lease         `json:"lease"`
	IssueTime  time.Time              `json:"issue_time"`
	ExpireTime time.Time              `json:"expire_time"`
}

// encode is used to JSON encode the lease entry
func (l *leaseEntry) encode() ([]byte, error) {
	return json.Marshal(l)
}

// decodeLeaseEntry is used to reverse encode and return a new entry
func decodeLeaseEntry(buf []byte) (*leaseEntry, error) {
	out := new(leaseEntry)
	return out, json.Unmarshal(buf, out)
}
<commit_msg>vault: Respect grace period for revocation<commit_after>package vault

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

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

const (
	// expirationSubPath is the sub-path used for the expiration manager
	// view. This is nested under the system view.
	expirationSubPath = "expire/"

	// maxRevokeAttempts limits how many revoke attempts are made
	maxRevokeAttempts = 6

	// revokeRetryBase is a baseline retry time
	revokeRetryBase = 10 * time.Second

	// minRevokeDelay is used to prevent an instant revoke on restore
	minRevokeDelay = 5 * time.Second
)

// ExpirationManager is used by the Core to manage leases. Secrets
// can provide a lease, meaning that they can be renewed or revoked.
// If a secret is not renewed in timely manner, it may be expired, and
// the ExpirationManager will handle doing automatic revocation.
type ExpirationManager struct {
	router *Router
	view   *BarrierView
	logger *log.Logger

	pending     map[string]*time.Timer
	pendingLock sync.Mutex
}

// NewExpirationManager creates a new ExpirationManager that is backed
// using a given view, and uses the provided router for revocation.
func NewExpirationManager(router *Router, view *BarrierView, logger *log.Logger) *ExpirationManager {
	if logger == nil {
		logger = log.New(os.Stderr, "", log.LstdFlags)
	}
	exp := &ExpirationManager{
		router:  router,
		view:    view,
		logger:  logger,
		pending: make(map[string]*time.Timer),
	}
	return exp
}

// setupExpiration is invoked after we've loaded the mount table to
// initialize the expiration manager
func (c *Core) setupExpiration() error {
	// Create a sub-view
	view := c.systemView.SubView(expirationSubPath)

	// Create the manager
	mgr := NewExpirationManager(c.router, view, c.logger)
	c.expiration = mgr

	// Restore the existing state
	if err := c.expiration.Restore(); err != nil {
		return fmt.Errorf("expiration state restore failed: %v", err)
	}
	return nil
}

// stopExpiration is used to stop the expiration manager before
// sealing the Vault.
func (c *Core) stopExpiration() error {
	if err := c.expiration.Stop(); err != nil {
		return err
	}
	c.expiration = nil
	return nil
}

// Restore is used to recover the lease states when starting.
// This is used after starting the vault.
func (m *ExpirationManager) Restore() error {
	m.pendingLock.Lock()
	defer m.pendingLock.Unlock()

	// Accumulate existing leases
	var existing []string
	cb := func(path string) {
		existing = append(existing, path)
	}

	// Scan for all the leases
	if err := ScanView(m.view, cb); err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Restore each key
	for _, vaultID := range existing {
		// Load the entry
		le, err := m.loadEntry(vaultID)
		if err != nil {
			return err
		}

		// If there is no entry, nothing to restore
		if le == nil {
			continue
		}

		// Determine the remaining time to expiration
		expires := le.ExpireTime.Sub(time.Now().UTC())
		if expires <= 0 {
			expires = minRevokeDelay
		}

		// Setup revocation timer
		m.pending[le.VaultID] = time.AfterFunc(expires, func() {
			m.expireID(le.VaultID)
		})
	}
	m.logger.Printf("[INFO] expire: restored %d leases", len(m.pending))
	return nil
}

// Stop is used to prevent further automatic revocations.
// This must be called before sealing the view.
func (m *ExpirationManager) Stop() error {
	// Stop all the pending expiration timers
	m.pendingLock.Lock()
	for _, timer := range m.pending {
		timer.Stop()
	}
	m.pending = make(map[string]*time.Timer)
	m.pendingLock.Unlock()
	return nil
}

// Revoke is used to revoke a secret named by the given vaultID
func (m *ExpirationManager) Revoke(vaultID string) error {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return err
	}

	// If there is no entry, nothing to revoke
	if le == nil {
		return nil
	}

	// Revoke the entry
	if err := m.revokeEntry(le); err != nil {
		return err
	}

	// Delete the entry
	if err := m.deleteEntry(vaultID); err != nil {
		return err
	}

	// Clear the expiration handler
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Stop()
		delete(m.pending, vaultID)
	}
	m.pendingLock.Unlock()
	return nil
}

// RevokePrefix is used to revoke all secrets with a given prefix.
// The prefix maps to that of the mount table to make this simpler
// to reason about.
func (m *ExpirationManager) RevokePrefix(prefix string) error {
	// Ensure there is a trailing slash
	if !strings.HasSuffix(prefix, "/") {
		prefix = prefix + "/"
	}

	// Accumulate existing leases
	var existing []string
	cb := func(path string) {
		existing = append(existing, path)
	}

	// Scan for all the leases in the prefix
	sub := m.view.SubView(prefix)
	if err := ScanView(sub, cb); err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Revoke all the keys
	for idx, suffix := range existing {
		vaultID := prefix + suffix
		if err := m.Revoke(vaultID); err != nil {
			return fmt.Errorf("failed to revoke '%s' (%d / %d): %v",
				vaultID, idx+1, len(existing), err)
		}
	}
	return nil
}

// Renew is used to renew a secret using the given vaultID
// and a renew interval. The increment may be ignored.
func (m *ExpirationManager) Renew(vaultID string, increment time.Duration) (*logical.Response, error) {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return nil, err
	}

	// If there is no entry, cannot review
	if le == nil {
		return nil, fmt.Errorf("lease not found")
	}

	// Determine if the lease is expired
	if le.ExpireTime.Before(time.Now().UTC()) {
		return nil, fmt.Errorf("lease expired")
	}

	// Attempt to renew the entry
	resp, err := m.renewEntry(le, increment)
	if err != nil {
		return nil, err
	}

	// Fast-path if there is no lease
	if resp == nil || resp.Lease == nil || !resp.IsSecret {
		return resp, nil
	}

	// Validate the lease
	if err := resp.Lease.Validate(); err != nil {
		return nil, err
	}

	// Attach the VaultID
	resp.Lease.VaultID = vaultID

	// Update the lease entry
	leaseTotal := resp.Lease.Duration + resp.Lease.GracePeriod
	le.Data = resp.Data
	le.Lease = resp.Lease
	le.ExpireTime = time.Now().UTC().Add(leaseTotal)
	if err := m.persistEntry(le); err != nil {
		return nil, err
	}

	// Update the expiration time
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Reset(leaseTotal)
	}
	m.pendingLock.Unlock()

	// Return the response
	return resp, nil
}

// Register is used to take a request and response with an associated
// lease. The secret gets assigned a vaultId and the management of
// of lease is assumed by the expiration manager.
func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Response) (string, error) {
	// Ignore if there is no lease
	if resp == nil || resp.Lease == nil {
		return "", nil
	}

	// Validate the lease
	if err := resp.Lease.Validate(); err != nil {
		return "", err
	}

	// Cannot register a non-secret (e.g. a policy or configuration key)
	if !resp.IsSecret {
		return "", fmt.Errorf("cannot attach lease to non-secret")
	}

	// Create a lease entry
	now := time.Now().UTC()
	leaseTotal := resp.Lease.Duration + resp.Lease.GracePeriod
	le := leaseEntry{
		VaultID:    path.Join(req.Path, generateUUID()),
		Path:       req.Path,
		Data:       resp.Data,
		Lease:      resp.Lease,
		IssueTime:  now,
		ExpireTime: now.Add(leaseTotal),
	}

	// Encode the entry
	if err := m.persistEntry(&le); err != nil {
		return "", err
	}

	// Setup revocation timer
	m.pendingLock.Lock()
	m.pending[le.VaultID] = time.AfterFunc(leaseTotal, func() {
		m.expireID(le.VaultID)
	})
	m.pendingLock.Unlock()

	// Done
	return le.VaultID, nil
}

// expireID is invoked when a given ID is expired
func (m *ExpirationManager) expireID(vaultID string) {
	// Clear from the pending expiration
	m.pendingLock.Lock()
	delete(m.pending, vaultID)
	m.pendingLock.Unlock()

	for attempt := uint(0); attempt < maxRevokeAttempts; attempt++ {
		err := m.Revoke(vaultID)
		if err == nil {
			m.logger.Printf("[INFO] expire: revoked '%s'", vaultID)
			return
		}
		m.logger.Printf("[ERR] expire: failed to revoke '%s': %v", vaultID, err)
		time.Sleep((1 << attempt) * revokeRetryBase)
	}
	m.logger.Printf("[ERR] expire: maximum revoke attempts for '%s' reached", vaultID)
}

// revokeEntry is used to attempt revocation of an internal entry
func (m *ExpirationManager) revokeEntry(le *leaseEntry) error {
	data := map[string]interface{}{
		"previous": le.Data,
	}
	req := &logical.Request{
		Operation: logical.RevokeOperation,
		Path:      le.Path,
		Data:      data,
	}
	_, err := m.router.Route(req)
	if err != nil {
		return fmt.Errorf("failed to revoke entry: %v", err)
	}
	return nil
}

// renewEntry is used to attempt renew of an internal entry
func (m *ExpirationManager) renewEntry(le *leaseEntry, increment time.Duration) (*logical.Response, error) {
	data := map[string]interface{}{
		"previous":  le.Data,
		"increment": increment,
	}
	req := &logical.Request{
		Operation: logical.RenewOperation,
		Path:      le.Path,
		Data:      data,
	}
	resp, err := m.router.Route(req)
	if err != nil {
		return nil, fmt.Errorf("failed to renew entry: %v", err)
	}
	return resp, nil
}

// loadEntry is used to read a lease entry
func (m *ExpirationManager) loadEntry(vaultID string) (*leaseEntry, error) {
	out, err := m.view.Get(vaultID)
	if err != nil {
		return nil, fmt.Errorf("failed to read lease entry: %v", err)
	}
	if out == nil {
		return nil, nil
	}
	le, err := decodeLeaseEntry(out.Value)
	if err != nil {
		return nil, fmt.Errorf("failed to decode lease entry: %v", err)
	}
	return le, nil
}

// persistEntry is used to persist a lease entry
func (m *ExpirationManager) persistEntry(le *leaseEntry) error {
	// Encode the entry
	buf, err := le.encode()
	if err != nil {
		return fmt.Errorf("failed to encode lease entry: %v", err)
	}

	// Write out to the view
	ent := logical.StorageEntry{
		Key:   le.VaultID,
		Value: buf,
	}
	if err := m.view.Put(&ent); err != nil {
		return fmt.Errorf("failed to persist lease entry: %v", err)
	}
	return nil
}

// deleteEntry is used to delete a lease entry
func (m *ExpirationManager) deleteEntry(vaultID string) error {
	if err := m.view.Delete(vaultID); err != nil {
		return fmt.Errorf("failed to delete lease entry: %v", err)
	}
	return nil
}

// leaseEntry is used to structure the values the expiration
// manager stores. This is used to handle renew and revocation.
type leaseEntry struct {
	VaultID    string                 `json:"vault_id"`
	Path       string                 `json:"path"`
	Data       map[string]interface{} `json:"data"`
	Lease      *logical.Lease         `json:"lease"`
	IssueTime  time.Time              `json:"issue_time"`
	ExpireTime time.Time              `json:"expire_time"`
}

// encode is used to JSON encode the lease entry
func (l *leaseEntry) encode() ([]byte, error) {
	return json.Marshal(l)
}

// decodeLeaseEntry is used to reverse encode and return a new entry
func decodeLeaseEntry(buf []byte) (*leaseEntry, error) {
	out := new(leaseEntry)
	return out, json.Unmarshal(buf, out)
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2019 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cloud

import (
	"fmt"
	"net/url"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"gopkg.in/guregu/null.v3"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/netext/httpext"
	"github.com/loadimpact/k6/lib/types"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats"
)

func BenchmarkAggregateHTTP(b *testing.B) {
	script := &loader.SourceData{
		Data: []byte(""),
		URL:  &url.URL{Path: "/script.js"},
	}

	options := lib.Options{
		Duration: types.NullDurationFrom(1 * time.Second),
	}

	config := NewConfig().Apply(Config{
		NoCompress:              null.BoolFrom(true),
		AggregationCalcInterval: types.NullDurationFrom(time.Millisecond * 200),
		AggregationPeriod:       types.NullDurationFrom(time.Millisecond * 200),
	})
	collector, err := New(config, script, options, []lib.ExecutionStep{}, "1.0")
	require.NoError(b, err)
	now := time.Now()
	collector.referenceID = "something"

	for _, tagCount := range []int{1, 5, 10, 100, 1000} {
		tagCount := tagCount
		b.Run(fmt.Sprintf("tags:%d", tagCount), func(b *testing.B) {
			tags := make([]*stats.SampleTags, tagCount)
			for i := range tags {
				tags[i] = stats.IntoSampleTags(&map[string]string{
					"test": "mest", "a": "b",
					"url":  fmt.Sprintf("something%d", i),
					"name": fmt.Sprintf("else%d", i),
				})
			}
			b.ResetTimer()
			for s := 0; s < b.N; s++ {
				for j := time.Duration(1); j <= 200; j++ {
					var container = make([]stats.SampleContainer, 0, 500)
					for i := time.Duration(1); i <= 500; i++ {
						container = append(container, &httpext.Trail{
							Blocked:        i % 200 * 100 * time.Millisecond,
							Connecting:     i % 200 * 200 * time.Millisecond,
							TLSHandshaking: i % 200 * 300 * time.Millisecond,
							Sending:        i * i * 400 * time.Millisecond,
							Waiting:        500 * time.Millisecond,
							Receiving:      600 * time.Millisecond,

							EndTime:      now.Add(i * 100),
							ConnDuration: 500 * time.Millisecond,
							Duration:     j * i * 1500 * time.Millisecond,
							Tags:         stats.NewSampleTags(tags[int(i+j)%len(tags)].CloneTags()),
						})
					}
					collector.Collect(container)
				}
				collector.aggregateHTTPTrails(time.Millisecond * 200)
			}
		})
	}
}
<commit_msg>Even better cloud aggregation benchmark<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2019 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cloud

import (
	"fmt"
	"net/url"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"gopkg.in/guregu/null.v3"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/netext/httpext"
	"github.com/loadimpact/k6/lib/types"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats"
)

func BenchmarkAggregateHTTP(b *testing.B) {
	script := &loader.SourceData{
		Data: []byte(""),
		URL:  &url.URL{Path: "/script.js"},
	}

	options := lib.Options{
		Duration: types.NullDurationFrom(1 * time.Second),
	}

	config := NewConfig().Apply(Config{
		NoCompress:              null.BoolFrom(true),
		AggregationCalcInterval: types.NullDurationFrom(time.Millisecond * 200),
		AggregationPeriod:       types.NullDurationFrom(time.Millisecond * 200),
	})
	collector, err := New(config, script, options, []lib.ExecutionStep{}, "1.0")
	require.NoError(b, err)
	now := time.Now()
	collector.referenceID = "something"
	var containersCount = time.Duration(500000)

	for _, tagCount := range []int{1, 5, 35, 315, 3645} {
		tagCount := tagCount
		b.Run(fmt.Sprintf("tags:%d", tagCount), func(b *testing.B) {
			b.ResetTimer()
			for s := 0; s < b.N; s++ {
				b.StopTimer()
				var container = make([]stats.SampleContainer, containersCount)
				for i := time.Duration(1); i <= containersCount; i++ {
					var status = "200"
					if int(i)%tagCount%7 == 6 {
						status = "404"
					} else if int(i)%tagCount%7 == 5 {
						status = "500"
					}

					container[i-1] = &httpext.Trail{
						Blocked:        i % 200 * 100 * time.Millisecond,
						Connecting:     i % 200 * 200 * time.Millisecond,
						TLSHandshaking: i % 200 * 300 * time.Millisecond,
						Sending:        i % 200 * 400 * time.Millisecond,
						Waiting:        500 * time.Millisecond,
						Receiving:      600 * time.Millisecond,
						EndTime:        now.Add(i % 100 * 100),
						ConnDuration:   500 * time.Millisecond,
						Duration:       i % 150 * 1500 * time.Millisecond,
						Tags: stats.IntoSampleTags(&map[string]string{
							"test": "mest", "a": "b",
							"custom": fmt.Sprintf("group%d", int(i)%tagCount%9),
							"group":  fmt.Sprintf("group%d", int(i)%tagCount%5),
							"status": status,
							"url":    fmt.Sprintf("something%d", int(i)%tagCount%11),
							"name":   fmt.Sprintf("else%d", int(i)%tagCount%11),
						}),
					}
				}
				collector.Collect(container)
				b.StartTimer()
				collector.aggregateHTTPTrails(time.Millisecond * 200)
				collector.bufferSamples = nil
			}
		})
	}
}
<|endoftext|>


<commit_before>// Package httputil provides additional helper functions for http services
// NOTE: Subject to change, do not rely on this package from outside git-lfs source
package httputil

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/github/git-lfs/config"
	"github.com/rubyist/tracerx"
)

type httpTransferStats struct {
	HeaderSize int
	BodySize   int
	Start      time.Time
	Stop       time.Time
}

type httpTransfer struct {
	requestStats  *httpTransferStats
	responseStats *httpTransferStats
}

var (
	// TODO should use some locks
	httpTransfers           = make(map[*http.Response]*httpTransfer)
	httpTransferBuckets     = make(map[string][]*http.Response)
	httpTransfersLock       sync.Mutex
	httpTransferBucketsLock sync.Mutex
	httpClients             map[string]*HttpClient
	httpClientsMutex        sync.Mutex
	UserAgent               string
)

func LogTransfer(cfg *config.Configuration, key string, res *http.Response) {
	if cfg.IsLoggingStats {
		httpTransferBucketsLock.Lock()
		httpTransferBuckets[key] = append(httpTransferBuckets[key], res)
		httpTransferBucketsLock.Unlock()
	}
}

type HttpClient struct {
	Config *config.Configuration
	*http.Client
}

func (c *HttpClient) Do(req *http.Request) (*http.Response, error) {
	traceHttpRequest(c.Config, req)

	crc := countingRequest(c.Config, req)
	if req.Body != nil {
		// Only set the body if we have a body, but create the countingRequest
		// anyway to make using zeroed stats easier.
		req.Body = crc
	}

	start := time.Now()
	res, err := c.Client.Do(req)
	if err != nil {
		return res, err
	}

	traceHttpResponse(c.Config, res)

	cresp := countingResponse(c.Config, res)
	res.Body = cresp

	if c.Config.IsLoggingStats {
		reqHeaderSize := 0
		resHeaderSize := 0

		if dump, err := httputil.DumpRequest(req, false); err == nil {
			reqHeaderSize = len(dump)
		}

		if dump, err := httputil.DumpResponse(res, false); err == nil {
			resHeaderSize = len(dump)
		}

		reqstats := &httpTransferStats{HeaderSize: reqHeaderSize, BodySize: crc.Count}

		// Response body size cannot be figured until it is read. Do not rely on a Content-Length
		// header because it may not exist or be -1 in the case of chunked responses.
		resstats := &httpTransferStats{HeaderSize: resHeaderSize, Start: start}
		t := &httpTransfer{requestStats: reqstats, responseStats: resstats}
		httpTransfersLock.Lock()
		httpTransfers[res] = t
		httpTransfersLock.Unlock()
	}

	return res, err
}

// NewHttpClient returns a new HttpClient for the given host (which may be "host:port")
func NewHttpClient(c *config.Configuration, host string) *HttpClient {
	httpClientsMutex.Lock()
	defer httpClientsMutex.Unlock()

	if httpClients == nil {
		httpClients = make(map[string]*HttpClient)
	}
	if client, ok := httpClients[host]; ok {
		return client
	}

	dialtime := c.Git.Int("lfs.dialtimeout", 30)
	keepalivetime := c.Git.Int("lfs.keepalive", 1800) // 30 minutes
	tlstime := c.Git.Int("lfs.tlstimeout", 30)

	tr := &http.Transport{
		Proxy: ProxyFromGitConfigOrEnvironment(c),
		Dial: (&net.Dialer{
			Timeout:   time.Duration(dialtime) * time.Second,
			KeepAlive: time.Duration(keepalivetime) * time.Second,
		}).Dial,
		TLSHandshakeTimeout: time.Duration(tlstime) * time.Second,
		MaxIdleConnsPerHost: c.ConcurrentTransfers(),
	}

	tr.TLSClientConfig = &tls.Config{}
	if isCertVerificationDisabledForHost(c, host) {
		tr.TLSClientConfig.InsecureSkipVerify = true
	} else {
		tr.TLSClientConfig.RootCAs = getRootCAsForHost(c, host)
	}

	client := &HttpClient{
		Config: c,
		Client: &http.Client{Transport: tr, CheckRedirect: CheckRedirect},
	}
	httpClients[host] = client

	return client
}

func CheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= 3 {
		return errors.New("stopped after 3 redirects")
	}

	oldest := via[0]
	for key, _ := range oldest.Header {
		if key == "Authorization" {
			if req.URL.Scheme != oldest.URL.Scheme || req.URL.Host != oldest.URL.Host {
				continue
			}
		}
		req.Header.Set(key, oldest.Header.Get(key))
	}

	oldestUrl := strings.SplitN(oldest.URL.String(), "?", 2)[0]
	newUrl := strings.SplitN(req.URL.String(), "?", 2)[0]
	tracerx.Printf("api: redirect %s %s to %s", oldest.Method, oldestUrl, newUrl)

	return nil
}

var tracedTypes = []string{"json", "text", "xml", "html"}

func traceHttpRequest(cfg *config.Configuration, req *http.Request) {
	tracerx.Printf("HTTP: %s", TraceHttpReq(req))

	if cfg.IsTracingHttp == false {
		return
	}

	dump, err := httputil.DumpRequest(req, false)
	if err != nil {
		return
	}

	traceHttpDump(cfg, ">", dump)
}

func traceHttpResponse(cfg *config.Configuration, res *http.Response) {
	if res == nil {
		return
	}

	tracerx.Printf("HTTP: %d", res.StatusCode)

	if cfg.IsTracingHttp == false {
		return
	}

	dump, err := httputil.DumpResponse(res, false)
	if err != nil {
		return
	}

	if isTraceableContent(res.Header) {
		fmt.Fprintf(os.Stderr, "\n\n")
	} else {
		fmt.Fprintf(os.Stderr, "\n")
	}

	traceHttpDump(cfg, "<", dump)
}

func traceHttpDump(cfg *config.Configuration, direction string, dump []byte) {
	scanner := bufio.NewScanner(bytes.NewBuffer(dump))

	for scanner.Scan() {
		line := scanner.Text()
		if !cfg.IsDebuggingHttp && strings.HasPrefix(strings.ToLower(line), "authorization: basic") {
			fmt.Fprintf(os.Stderr, "%s Authorization: Basic * * * * *\n", direction)
		} else {
			fmt.Fprintf(os.Stderr, "%s %s\n", direction, line)
		}
	}
}

func isTraceableContent(h http.Header) bool {
	ctype := strings.ToLower(strings.SplitN(h.Get("Content-Type"), ";", 2)[0])
	for _, tracedType := range tracedTypes {
		if strings.Contains(ctype, tracedType) {
			return true
		}
	}
	return false
}

func countingRequest(cfg *config.Configuration, req *http.Request) *CountingReadCloser {
	return &CountingReadCloser{
		request:         req,
		cfg:             cfg,
		ReadCloser:      req.Body,
		isTraceableType: isTraceableContent(req.Header),
		useGitTrace:     false,
	}
}

func countingResponse(cfg *config.Configuration, res *http.Response) *CountingReadCloser {
	return &CountingReadCloser{
		response:        res,
		cfg:             cfg,
		ReadCloser:      res.Body,
		isTraceableType: isTraceableContent(res.Header),
		useGitTrace:     true,
	}
}

type CountingReadCloser struct {
	Count           int
	request         *http.Request
	response        *http.Response
	cfg             *config.Configuration
	isTraceableType bool
	useGitTrace     bool
	io.ReadCloser
}

func (c *CountingReadCloser) Read(b []byte) (int, error) {
	n, err := c.ReadCloser.Read(b)
	if err != nil && err != io.EOF {
		return n, err
	}

	c.Count += n

	if c.isTraceableType {
		chunk := string(b[0:n])
		if c.useGitTrace {
			tracerx.Printf("HTTP: %s", chunk)
		}

		if c.cfg.IsTracingHttp {
			fmt.Fprint(os.Stderr, chunk)
		}
	}

	if err == io.EOF && c.cfg.IsLoggingStats {
		// This httpTransfer is done, we're checking it this way so we can also
		// catch httpTransfers where the caller forgets to Close() the Body.
		if c.response != nil {
			httpTransfersLock.Lock()
			if httpTransfer, ok := httpTransfers[c.response]; ok {
				httpTransfer.responseStats.BodySize = c.Count
				httpTransfer.responseStats.Stop = time.Now()
			}
			httpTransfersLock.Unlock()
		}
	}
	return n, err
}

// LogHttpStats is intended to be called after all HTTP operations for the
// commmand have finished. It dumps k/v logs, one line per httpTransfer into
// a log file with the current timestamp.
func LogHttpStats(cfg *config.Configuration) {
	if !cfg.IsLoggingStats {
		return
	}

	file, err := statsLogFile()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error logging http stats: %s\n", err)
		return
	}

	fmt.Fprintf(file, "concurrent=%d batch=%v time=%d version=%s\n", cfg.ConcurrentTransfers(), cfg.BatchTransfer(), time.Now().Unix(), config.Version)

	for key, responses := range httpTransferBuckets {
		for _, response := range responses {
			stats := httpTransfers[response]
			fmt.Fprintf(file, "key=%s reqheader=%d reqbody=%d resheader=%d resbody=%d restime=%d status=%d url=%s\n",
				key,
				stats.requestStats.HeaderSize,
				stats.requestStats.BodySize,
				stats.responseStats.HeaderSize,
				stats.responseStats.BodySize,
				stats.responseStats.Stop.Sub(stats.responseStats.Start).Nanoseconds(),
				response.StatusCode,
				response.Request.URL)
		}
	}

	fmt.Fprintf(os.Stderr, "HTTP Stats logged to file %s\n", file.Name())
}

func statsLogFile() (*os.File, error) {
	logBase := filepath.Join(config.LocalLogDir, "http")
	if err := os.MkdirAll(logBase, 0755); err != nil {
		return nil, err
	}

	logFile := fmt.Sprintf("http-%d.log", time.Now().Unix())
	return os.Create(filepath.Join(logBase, logFile))
}

func TraceHttpReq(req *http.Request) string {
	return fmt.Sprintf("%s %s", req.Method, strings.SplitN(req.URL.String(), "?", 2)[0])
}

func init() {
	UserAgent = config.VersionDesc
}
<commit_msg>tweak http tracing so it doesn't trace an empty byte<commit_after>// Package httputil provides additional helper functions for http services
// NOTE: Subject to change, do not rely on this package from outside git-lfs source
package httputil

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/github/git-lfs/config"
	"github.com/rubyist/tracerx"
)

type httpTransferStats struct {
	HeaderSize int
	BodySize   int
	Start      time.Time
	Stop       time.Time
}

type httpTransfer struct {
	requestStats  *httpTransferStats
	responseStats *httpTransferStats
}

var (
	// TODO should use some locks
	httpTransfers           = make(map[*http.Response]*httpTransfer)
	httpTransferBuckets     = make(map[string][]*http.Response)
	httpTransfersLock       sync.Mutex
	httpTransferBucketsLock sync.Mutex
	httpClients             map[string]*HttpClient
	httpClientsMutex        sync.Mutex
	UserAgent               string
)

func LogTransfer(cfg *config.Configuration, key string, res *http.Response) {
	if cfg.IsLoggingStats {
		httpTransferBucketsLock.Lock()
		httpTransferBuckets[key] = append(httpTransferBuckets[key], res)
		httpTransferBucketsLock.Unlock()
	}
}

type HttpClient struct {
	Config *config.Configuration
	*http.Client
}

func (c *HttpClient) Do(req *http.Request) (*http.Response, error) {
	traceHttpRequest(c.Config, req)

	crc := countingRequest(c.Config, req)
	if req.Body != nil {
		// Only set the body if we have a body, but create the countingRequest
		// anyway to make using zeroed stats easier.
		req.Body = crc
	}

	start := time.Now()
	res, err := c.Client.Do(req)
	if err != nil {
		return res, err
	}

	traceHttpResponse(c.Config, res)

	cresp := countingResponse(c.Config, res)
	res.Body = cresp

	if c.Config.IsLoggingStats {
		reqHeaderSize := 0
		resHeaderSize := 0

		if dump, err := httputil.DumpRequest(req, false); err == nil {
			reqHeaderSize = len(dump)
		}

		if dump, err := httputil.DumpResponse(res, false); err == nil {
			resHeaderSize = len(dump)
		}

		reqstats := &httpTransferStats{HeaderSize: reqHeaderSize, BodySize: crc.Count}

		// Response body size cannot be figured until it is read. Do not rely on a Content-Length
		// header because it may not exist or be -1 in the case of chunked responses.
		resstats := &httpTransferStats{HeaderSize: resHeaderSize, Start: start}
		t := &httpTransfer{requestStats: reqstats, responseStats: resstats}
		httpTransfersLock.Lock()
		httpTransfers[res] = t
		httpTransfersLock.Unlock()
	}

	return res, err
}

// NewHttpClient returns a new HttpClient for the given host (which may be "host:port")
func NewHttpClient(c *config.Configuration, host string) *HttpClient {
	httpClientsMutex.Lock()
	defer httpClientsMutex.Unlock()

	if httpClients == nil {
		httpClients = make(map[string]*HttpClient)
	}
	if client, ok := httpClients[host]; ok {
		return client
	}

	dialtime := c.Git.Int("lfs.dialtimeout", 30)
	keepalivetime := c.Git.Int("lfs.keepalive", 1800) // 30 minutes
	tlstime := c.Git.Int("lfs.tlstimeout", 30)

	tr := &http.Transport{
		Proxy: ProxyFromGitConfigOrEnvironment(c),
		Dial: (&net.Dialer{
			Timeout:   time.Duration(dialtime) * time.Second,
			KeepAlive: time.Duration(keepalivetime) * time.Second,
		}).Dial,
		TLSHandshakeTimeout: time.Duration(tlstime) * time.Second,
		MaxIdleConnsPerHost: c.ConcurrentTransfers(),
	}

	tr.TLSClientConfig = &tls.Config{}
	if isCertVerificationDisabledForHost(c, host) {
		tr.TLSClientConfig.InsecureSkipVerify = true
	} else {
		tr.TLSClientConfig.RootCAs = getRootCAsForHost(c, host)
	}

	client := &HttpClient{
		Config: c,
		Client: &http.Client{Transport: tr, CheckRedirect: CheckRedirect},
	}
	httpClients[host] = client

	return client
}

func CheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= 3 {
		return errors.New("stopped after 3 redirects")
	}

	oldest := via[0]
	for key, _ := range oldest.Header {
		if key == "Authorization" {
			if req.URL.Scheme != oldest.URL.Scheme || req.URL.Host != oldest.URL.Host {
				continue
			}
		}
		req.Header.Set(key, oldest.Header.Get(key))
	}

	oldestUrl := strings.SplitN(oldest.URL.String(), "?", 2)[0]
	newUrl := strings.SplitN(req.URL.String(), "?", 2)[0]
	tracerx.Printf("api: redirect %s %s to %s", oldest.Method, oldestUrl, newUrl)

	return nil
}

var tracedTypes = []string{"json", "text", "xml", "html"}

func traceHttpRequest(cfg *config.Configuration, req *http.Request) {
	tracerx.Printf("HTTP: %s", TraceHttpReq(req))

	if cfg.IsTracingHttp == false {
		return
	}

	dump, err := httputil.DumpRequest(req, false)
	if err != nil {
		return
	}

	traceHttpDump(cfg, ">", dump)
}

func traceHttpResponse(cfg *config.Configuration, res *http.Response) {
	if res == nil {
		return
	}

	tracerx.Printf("HTTP: %d", res.StatusCode)

	if cfg.IsTracingHttp == false {
		return
	}

	dump, err := httputil.DumpResponse(res, false)
	if err != nil {
		return
	}

	if isTraceableContent(res.Header) {
		fmt.Fprintf(os.Stderr, "\n\n")
	} else {
		fmt.Fprintf(os.Stderr, "\n")
	}

	traceHttpDump(cfg, "<", dump)
}

func traceHttpDump(cfg *config.Configuration, direction string, dump []byte) {
	scanner := bufio.NewScanner(bytes.NewBuffer(dump))

	for scanner.Scan() {
		line := scanner.Text()
		if !cfg.IsDebuggingHttp && strings.HasPrefix(strings.ToLower(line), "authorization: basic") {
			fmt.Fprintf(os.Stderr, "%s Authorization: Basic * * * * *\n", direction)
		} else {
			fmt.Fprintf(os.Stderr, "%s %s\n", direction, line)
		}
	}
}

func isTraceableContent(h http.Header) bool {
	ctype := strings.ToLower(strings.SplitN(h.Get("Content-Type"), ";", 2)[0])
	for _, tracedType := range tracedTypes {
		if strings.Contains(ctype, tracedType) {
			return true
		}
	}
	return false
}

func countingRequest(cfg *config.Configuration, req *http.Request) *CountingReadCloser {
	return &CountingReadCloser{
		request:         req,
		cfg:             cfg,
		ReadCloser:      req.Body,
		isTraceableType: isTraceableContent(req.Header),
		useGitTrace:     false,
	}
}

func countingResponse(cfg *config.Configuration, res *http.Response) *CountingReadCloser {
	return &CountingReadCloser{
		response:        res,
		cfg:             cfg,
		ReadCloser:      res.Body,
		isTraceableType: isTraceableContent(res.Header),
		useGitTrace:     true,
	}
}

type CountingReadCloser struct {
	Count           int
	request         *http.Request
	response        *http.Response
	cfg             *config.Configuration
	isTraceableType bool
	useGitTrace     bool
	io.ReadCloser
}

func (c *CountingReadCloser) Read(b []byte) (int, error) {
	n, err := c.ReadCloser.Read(b)
	if err != nil && err != io.EOF {
		return n, err
	}

	c.Count += n

	if n > 0 && c.isTraceableType {
		chunk := string(b[0:n])
		if c.useGitTrace {
			tracerx.Printf("HTTP: %s", chunk)
		}

		if c.cfg.IsTracingHttp {
			fmt.Fprint(os.Stderr, chunk)
		}
	}

	if err == io.EOF && c.cfg.IsLoggingStats {
		// This httpTransfer is done, we're checking it this way so we can also
		// catch httpTransfers where the caller forgets to Close() the Body.
		if c.response != nil {
			httpTransfersLock.Lock()
			if httpTransfer, ok := httpTransfers[c.response]; ok {
				httpTransfer.responseStats.BodySize = c.Count
				httpTransfer.responseStats.Stop = time.Now()
			}
			httpTransfersLock.Unlock()
		}
	}
	return n, err
}

// LogHttpStats is intended to be called after all HTTP operations for the
// commmand have finished. It dumps k/v logs, one line per httpTransfer into
// a log file with the current timestamp.
func LogHttpStats(cfg *config.Configuration) {
	if !cfg.IsLoggingStats {
		return
	}

	file, err := statsLogFile()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error logging http stats: %s\n", err)
		return
	}

	fmt.Fprintf(file, "concurrent=%d batch=%v time=%d version=%s\n", cfg.ConcurrentTransfers(), cfg.BatchTransfer(), time.Now().Unix(), config.Version)

	for key, responses := range httpTransferBuckets {
		for _, response := range responses {
			stats := httpTransfers[response]
			fmt.Fprintf(file, "key=%s reqheader=%d reqbody=%d resheader=%d resbody=%d restime=%d status=%d url=%s\n",
				key,
				stats.requestStats.HeaderSize,
				stats.requestStats.BodySize,
				stats.responseStats.HeaderSize,
				stats.responseStats.BodySize,
				stats.responseStats.Stop.Sub(stats.responseStats.Start).Nanoseconds(),
				response.StatusCode,
				response.Request.URL)
		}
	}

	fmt.Fprintf(os.Stderr, "HTTP Stats logged to file %s\n", file.Name())
}

func statsLogFile() (*os.File, error) {
	logBase := filepath.Join(config.LocalLogDir, "http")
	if err := os.MkdirAll(logBase, 0755); err != nil {
		return nil, err
	}

	logFile := fmt.Sprintf("http-%d.log", time.Now().Unix())
	return os.Create(filepath.Join(logBase, logFile))
}

func TraceHttpReq(req *http.Request) string {
	return fmt.Sprintf("%s %s", req.Method, strings.SplitN(req.URL.String(), "?", 2)[0])
}

func init() {
	UserAgent = config.VersionDesc
}
<|endoftext|>


<commit_before>//Copyright 2017 Huawei Technologies Co., Ltd
//
//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 notification

import (
	"container/list"
	"encoding/json"
	"errors"
	apt "github.com/servicecomb/service-center/server/core"
	pb "github.com/servicecomb/service-center/server/core/proto"
	"github.com/servicecomb/service-center/server/core/registry"
	"github.com/servicecomb/service-center/server/service/dependency"
	"github.com/servicecomb/service-center/server/service/microservice"
	"github.com/servicecomb/service-center/util"
	"strings"
	"sync"
	"time"
)

const (
	DEFAULT_MAX_QUEUE           = 100
	DEFAULT_TIMEOUT             = 30 * time.Second
	DEFAULT_LISTWATCH_TIMEOUT   = 30 * time.Second
	NOTIFY_SERVER_CHECKER_NAME  = "__HealthChecker__"
	NOTIFY_SERVER_CHECK_SUBJECT = "__NotifyServerHealthCheck__"
)

type NotifyServerConfig struct {
	AddTimeout    time.Duration
	NotifyTimeout time.Duration
	MaxQueue      int64
}

type NotifyService struct {
	Config *NotifyServerConfig

	cacherMap map[string]registry.Cacher
	notifiers map[string]map[string]*list.List
	queue     chan NotifyJob
	wait      chan struct{}
	err       chan error
	mux       sync.Mutex
	closeMux  sync.RWMutex
	isClose   bool
}

func (s *NotifyService) Err() <-chan error {
	return s.err
}

func (s *NotifyService) ListWatcher(key string) registry.Cacher {
	return s.cacherMap[key]
}

func (s *NotifyService) AddNotifier(n Worker) error {
	if s.closed() {
		return errors.New("server is shutting down")
	}

	s.mux.Lock()
	_, ok := s.notifiers[n.Subject()]
	if !ok {
		s.notifiers[n.Subject()] = map[string]*list.List{}
	}

	m := s.notifiers[n.Subject()]
	ns, ok := m[n.Id()]
	if !ok {
		ns = list.New()
	}
	ns.PushBack(n)
	m[n.Id()] = ns

	n.SetService(s)
	s.mux.Unlock()

	n.OnAccept()
	return nil
}

func (s *NotifyService) RemoveNotifier(r Worker) {
	s.mux.Lock()
	m, ok := s.notifiers[r.Subject()]
	if ok {
		ns, ok := m[r.Id()]
		if ok {
			for n := ns.Front(); n != nil; n = n.Next() {
				if n.Value == r {
					ns.Remove(n)
					r.Close()
					break
				}
			}
		}
	}
	s.mux.Unlock()
}

func (s *NotifyService) RemoveAllNotifiers() {
	s.mux.Lock()
	for subject := range s.notifiers {
		for key := range s.notifiers[subject] {
			ns := s.notifiers[subject][key]
			for n := ns.Front(); n != nil; n = n.Next() {
				n.Value.(Worker).Close()
			}
		}
	}
	s.mux.Unlock()
}

//通知内容塞到队列里
func (s *NotifyService) AddJob(job NotifyJob) error {
	if s.closed() {
		return errors.New("add notify job failed for server shutdown")
	}
	select {
	case s.queue <- job:
		return nil
	case <-time.After(s.Config.AddTimeout):
		util.LOGGER.Errorf(nil, "Add job failed.%s")
		return errors.New("add notify job timeout")
	}
}

func (s *NotifyService) publish2Subscriber() {
	defer close(s.wait)
	for job := range s.queue {
		util.LOGGER.Infof("notification server got a job %s %s", job.Subject(), job.Id())

		s.mux.Lock()

		if s.closed() && len(s.notifiers) == 0 {
			s.mux.Unlock()
			return
		}

		m, ok := s.notifiers[job.Subject()]
		if ok {
			// publish的subject如果带上id，则单播，否则广播
			if len(job.Id()) != 0 {
				ns, ok := m[job.Id()]
				if ok {
					for n := ns.Front(); n != nil; n = n.Next() {
						go n.Value.(Worker).OnMessage(job)
					}
				}
				s.mux.Unlock()
				continue
			}
			for key := range m {
				ns := m[key]
				for n := ns.Front(); n != nil; n = n.Next() {
					go n.Value.(Worker).OnMessage(job)
				}
			}
		}

		s.mux.Unlock()
	}
}

func (s *NotifyService) StartNotifyService() {
	util.LOGGER.Info("starting notify service")

	if s.Config.AddTimeout <= 0 {
		s.Config.AddTimeout = DEFAULT_TIMEOUT
	}
	if s.Config.NotifyTimeout <= 0 {
		s.Config.NotifyTimeout = DEFAULT_TIMEOUT
	}
	if s.Config.MaxQueue <= 0 || s.Config.MaxQueue > DEFAULT_MAX_QUEUE {
		s.Config.MaxQueue = DEFAULT_MAX_QUEUE
	}

	s.cacherMap = make(map[string]registry.Cacher)
	s.notifiers = map[string]map[string]*list.List{}
	s.err = make(chan error, 1)
	s.queue = make(chan NotifyJob, s.Config.MaxQueue)
	s.wait = make(chan struct{})
	s.isClose = false

	go s.publish2Subscriber()

	// 错误notifier清理
	s.AddNotifier(NewNotifyServiceHealthChecker())
	s.WatchTenants()
	util.LOGGER.Info("notify service is ready")
}

func (s *NotifyService) WatchTenants() {
	key := apt.GenerateTenantKey("")

	c := registry.NewKvCacher(&registry.KvCacherConfig{
		Key:     key[:len(key)-1],
		Timeout: DEFAULT_LISTWATCH_TIMEOUT,
		Period:  time.Second,
		OnEvent: func(evt *registry.KvEvent) error {
			kv := evt.KV
			action := evt.Action
			tenant := pb.GetInfoFromTenantChangeEvent(kv)
			if len(tenant) == 0 {
				util.LOGGER.Errorf(nil,
					"unmarshal tenant info failed, key %s [%s] event", string(kv.Key), action)
				return nil
			}
			if action != pb.EVT_CREATE {
				return nil
			}

			if s.closed() {
				util.LOGGER.Warnf(nil,
					"caught tenant info %s [%] event, but service closed", string(kv.Key), action)
				return nil
			}

			util.LOGGER.Warnf(nil, "new tenant %s instances watcher is created", tenant)
			s.WatchInstance(apt.GetInstanceRootKey(tenant))
			return nil
		},
	})
	s.cacherMap[key] = c
	c.Run()
}

//SC 负责监控所有实例变化
func (s *NotifyService) WatchInstance(instanceWatchByTenantKey string) {
	c := registry.NewKvCacher(&registry.KvCacherConfig{
		Key:     instanceWatchByTenantKey,
		Timeout: DEFAULT_LISTWATCH_TIMEOUT,
		Period:  time.Second,
		OnEvent: func(evt *registry.KvEvent) error {
			kv := evt.KV
			action := evt.Action
			providerId, providerInstanceId, tenantProject, data := pb.GetInfoFromInstChangedEvent(kv)
			if data == nil {
				util.LOGGER.Errorf(nil,
					"unmarshal provider service instance file failed, instance %s/%s [%s] event, data is nil",
					providerId, providerInstanceId, action)
				return nil
			}

			if s.closed() {
				util.LOGGER.Warnf(nil, "caught instance %s/%s [%] event, but service closed",
					providerId, providerInstanceId, action)
				return nil
			} else {
				util.LOGGER.Warnf(nil, "caught instance %s/%s [%s] event",
					providerId, providerInstanceId, action)
			}

			var instance pb.MicroServiceInstance
			err := json.Unmarshal(data, &instance)
			if err != nil {
				util.LOGGER.Errorf(err, "unmarshal provider service instance %s/%s file failed",
					providerId, providerInstanceId)
				return nil
			}
			// 查询服务版本信息
			ms, err := microservice.GetByIdInCache(tenantProject, providerId)
			if ms == nil {
				util.LOGGER.Errorf(err, "get provider service %s/%s id in cache failed",
					providerId, providerInstanceId)
				return nil
			}

			// 查询所有consumer
			Kvs, err := dependency.GetConsumersInCache(tenantProject, providerId)
			if err != nil {
				util.LOGGER.Errorf(err, "query service %s consumers failed", providerId)
				return nil
			}

			response := &pb.WatchInstanceResponse{
				Response: pb.CreateResponse(pb.Response_SUCCESS, "watch instance successfully"),
				Action:   string(action),
				Key: &pb.MicroServiceKey{
					AppId:       ms.AppId,
					ServiceName: ms.ServiceName,
					Version:     ms.Version,
				},
				Instance: &instance,
			}
			for _, dependence := range Kvs {
				consumer := string(dependence.Key)
				consumer = consumer[strings.LastIndex(consumer, "/")+1:]
				job := NewWatchJob(consumer, apt.GetInstanceRootKey(tenantProject)+"/",
					evt.Revision, response)
				util.LOGGER.Debugf("publish event to notify server, %v", job)

				// TODO add超时怎么处理？
				s.AddJob(job)
			}
			return nil
		},
	})
	s.cacherMap[instanceWatchByTenantKey+"/"] = c
	c.Run()
}

func (s *NotifyService) closed() (b bool) {
	s.closeMux.RLock()
	b = s.isClose
	s.closeMux.RUnlock()
	return
}

func (s *NotifyService) Close() {
	if s.closed() {
		return
	}

	s.closeMux.Lock()
	s.isClose = true
	s.closeMux.Unlock()

	close(s.queue)

	<-s.wait

	s.RemoveAllNotifiers()

	close(s.err)

	util.LOGGER.Info("notify service stopped.")
}
<commit_msg>Fixed notify service log print bug.<commit_after>//Copyright 2017 Huawei Technologies Co., Ltd
//
//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 notification

import (
	"container/list"
	"encoding/json"
	"errors"
	apt "github.com/servicecomb/service-center/server/core"
	pb "github.com/servicecomb/service-center/server/core/proto"
	"github.com/servicecomb/service-center/server/core/registry"
	"github.com/servicecomb/service-center/server/service/dependency"
	"github.com/servicecomb/service-center/server/service/microservice"
	"github.com/servicecomb/service-center/util"
	"strings"
	"sync"
	"time"
)

const (
	DEFAULT_MAX_QUEUE           = 100
	DEFAULT_TIMEOUT             = 30 * time.Second
	DEFAULT_LISTWATCH_TIMEOUT   = 30 * time.Second
	NOTIFY_SERVER_CHECKER_NAME  = "__HealthChecker__"
	NOTIFY_SERVER_CHECK_SUBJECT = "__NotifyServerHealthCheck__"
)

type NotifyServerConfig struct {
	AddTimeout    time.Duration
	NotifyTimeout time.Duration
	MaxQueue      int64
}

type NotifyService struct {
	Config *NotifyServerConfig

	cacherMap map[string]registry.Cacher
	notifiers map[string]map[string]*list.List
	queue     chan NotifyJob
	wait      chan struct{}
	err       chan error
	mux       sync.Mutex
	closeMux  sync.RWMutex
	isClose   bool
}

func (s *NotifyService) Err() <-chan error {
	return s.err
}

func (s *NotifyService) ListWatcher(key string) registry.Cacher {
	return s.cacherMap[key]
}

func (s *NotifyService) AddNotifier(n Worker) error {
	if s.closed() {
		return errors.New("server is shutting down")
	}

	s.mux.Lock()
	_, ok := s.notifiers[n.Subject()]
	if !ok {
		s.notifiers[n.Subject()] = map[string]*list.List{}
	}

	m := s.notifiers[n.Subject()]
	ns, ok := m[n.Id()]
	if !ok {
		ns = list.New()
	}
	ns.PushBack(n)
	m[n.Id()] = ns

	n.SetService(s)
	s.mux.Unlock()

	n.OnAccept()
	return nil
}

func (s *NotifyService) RemoveNotifier(r Worker) {
	s.mux.Lock()
	m, ok := s.notifiers[r.Subject()]
	if ok {
		ns, ok := m[r.Id()]
		if ok {
			for n := ns.Front(); n != nil; n = n.Next() {
				if n.Value == r {
					ns.Remove(n)
					r.Close()
					break
				}
			}
		}
	}
	s.mux.Unlock()
}

func (s *NotifyService) RemoveAllNotifiers() {
	s.mux.Lock()
	for subject := range s.notifiers {
		for key := range s.notifiers[subject] {
			ns := s.notifiers[subject][key]
			for n := ns.Front(); n != nil; n = n.Next() {
				n.Value.(Worker).Close()
			}
		}
	}
	s.mux.Unlock()
}

//通知内容塞到队列里
func (s *NotifyService) AddJob(job NotifyJob) error {
	if s.closed() {
		return errors.New("add notify job failed for server shutdown")
	}
	select {
	case s.queue <- job:
		return nil
	case <-time.After(s.Config.AddTimeout):
		util.LOGGER.Errorf(nil, "Add job failed.%s")
		return errors.New("add notify job timeout")
	}
}

func (s *NotifyService) publish2Subscriber() {
	defer close(s.wait)
	for job := range s.queue {
		util.LOGGER.Infof("notification server got a job %s %s", job.Subject(), job.Id())

		s.mux.Lock()

		if s.closed() && len(s.notifiers) == 0 {
			s.mux.Unlock()
			return
		}

		m, ok := s.notifiers[job.Subject()]
		if ok {
			// publish的subject如果带上id，则单播，否则广播
			if len(job.Id()) != 0 {
				ns, ok := m[job.Id()]
				if ok {
					for n := ns.Front(); n != nil; n = n.Next() {
						go n.Value.(Worker).OnMessage(job)
					}
				}
				s.mux.Unlock()
				continue
			}
			for key := range m {
				ns := m[key]
				for n := ns.Front(); n != nil; n = n.Next() {
					go n.Value.(Worker).OnMessage(job)
				}
			}
		}

		s.mux.Unlock()
	}
}

func (s *NotifyService) StartNotifyService() {
	util.LOGGER.Info("starting notify service")

	if s.Config.AddTimeout <= 0 {
		s.Config.AddTimeout = DEFAULT_TIMEOUT
	}
	if s.Config.NotifyTimeout <= 0 {
		s.Config.NotifyTimeout = DEFAULT_TIMEOUT
	}
	if s.Config.MaxQueue <= 0 || s.Config.MaxQueue > DEFAULT_MAX_QUEUE {
		s.Config.MaxQueue = DEFAULT_MAX_QUEUE
	}

	s.cacherMap = make(map[string]registry.Cacher)
	s.notifiers = map[string]map[string]*list.List{}
	s.err = make(chan error, 1)
	s.queue = make(chan NotifyJob, s.Config.MaxQueue)
	s.wait = make(chan struct{})
	s.isClose = false

	go s.publish2Subscriber()

	// 错误notifier清理
	s.AddNotifier(NewNotifyServiceHealthChecker())
	s.WatchTenants()
	util.LOGGER.Info("notify service is ready")
}

func (s *NotifyService) WatchTenants() {
	key := apt.GenerateTenantKey("")

	c := registry.NewKvCacher(&registry.KvCacherConfig{
		Key:     key[:len(key)-1],
		Timeout: DEFAULT_LISTWATCH_TIMEOUT,
		Period:  time.Second,
		OnEvent: func(evt *registry.KvEvent) error {
			kv := evt.KV
			action := evt.Action
			tenant := pb.GetInfoFromTenantChangeEvent(kv)
			if len(tenant) == 0 {
				util.LOGGER.Errorf(nil,
					"unmarshal tenant info failed, key %s [%s] event", string(kv.Key), action)
				return nil
			}
			if action != pb.EVT_CREATE {
				return nil
			}

			if s.closed() {
				util.LOGGER.Warnf(nil,
					"caught tenant info %s [%s] event, but service closed", string(kv.Key), action)
				return nil
			}

			util.LOGGER.Warnf(nil, "new tenant %s instances watcher is created", tenant)
			s.WatchInstance(apt.GetInstanceRootKey(tenant))
			return nil
		},
	})
	s.cacherMap[key] = c
	c.Run()
}

//SC 负责监控所有实例变化
func (s *NotifyService) WatchInstance(instanceWatchByTenantKey string) {
	c := registry.NewKvCacher(&registry.KvCacherConfig{
		Key:     instanceWatchByTenantKey,
		Timeout: DEFAULT_LISTWATCH_TIMEOUT,
		Period:  time.Second,
		OnEvent: func(evt *registry.KvEvent) error {
			kv := evt.KV
			action := evt.Action
			providerId, providerInstanceId, tenantProject, data := pb.GetInfoFromInstChangedEvent(kv)
			if data == nil {
				util.LOGGER.Errorf(nil,
					"unmarshal provider service instance file failed, instance %s/%s [%s] event, data is nil",
					providerId, providerInstanceId, action)
				return nil
			}

			if s.closed() {
				util.LOGGER.Warnf(nil, "caught instance %s/%s [%s] event, but service closed",
					providerId, providerInstanceId, action)
				return nil
			} else {
				util.LOGGER.Warnf(nil, "caught instance %s/%s [%s] event",
					providerId, providerInstanceId, action)
			}

			var instance pb.MicroServiceInstance
			err := json.Unmarshal(data, &instance)
			if err != nil {
				util.LOGGER.Errorf(err, "unmarshal provider service instance %s/%s file failed",
					providerId, providerInstanceId)
				return nil
			}
			// 查询服务版本信息
			ms, err := microservice.GetByIdInCache(tenantProject, providerId)
			if ms == nil {
				util.LOGGER.Errorf(err, "get provider service %s/%s id in cache failed",
					providerId, providerInstanceId)
				return nil
			}

			// 查询所有consumer
			Kvs, err := dependency.GetConsumersInCache(tenantProject, providerId)
			if err != nil {
				util.LOGGER.Errorf(err, "query service %s consumers failed", providerId)
				return nil
			}

			response := &pb.WatchInstanceResponse{
				Response: pb.CreateResponse(pb.Response_SUCCESS, "watch instance successfully"),
				Action:   string(action),
				Key: &pb.MicroServiceKey{
					AppId:       ms.AppId,
					ServiceName: ms.ServiceName,
					Version:     ms.Version,
				},
				Instance: &instance,
			}
			for _, dependence := range Kvs {
				consumer := string(dependence.Key)
				consumer = consumer[strings.LastIndex(consumer, "/")+1:]
				job := NewWatchJob(consumer, apt.GetInstanceRootKey(tenantProject)+"/",
					evt.Revision, response)
				util.LOGGER.Debugf("publish event to notify server, %v", job)

				// TODO add超时怎么处理？
				s.AddJob(job)
			}
			return nil
		},
	})
	s.cacherMap[instanceWatchByTenantKey+"/"] = c
	c.Run()
}

func (s *NotifyService) closed() (b bool) {
	s.closeMux.RLock()
	b = s.isClose
	s.closeMux.RUnlock()
	return
}

func (s *NotifyService) Close() {
	if s.closed() {
		return
	}

	s.closeMux.Lock()
	s.isClose = true
	s.closeMux.Unlock()

	close(s.queue)

	<-s.wait

	s.RemoveAllNotifiers()

	close(s.err)

	util.LOGGER.Info("notify service stopped.")
}
<|endoftext|>


<commit_before>package cmds

import (
	"errors"
	"fmt"

	"github.com/pachyderm/pachyderm/src/client/pkg/config"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmdutil"

	"github.com/spf13/cobra"
)

// Cmds returns a slice containing admin commands.
func Cmds() []*cobra.Command {
	var commands []*cobra.Command

	getMetrics := &cobra.Command{
		Short: "Gets whether metrics are enabled.",
		Long:  "Gets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(0, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			fmt.Printf("%v\n", !cfg.V2.NoMetrics)
			return nil
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(getMetrics, "config get metrics"))

	setMetrics := &cobra.Command{
		Short: "Sets whether metrics are enabled.",
		Long:  "Sets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			metrics := true
			if args[0] == "false" {
				metrics = false
			} else if args[0] != "true" {
				return errors.New("invalid argument; use either `true` or `false`")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			cfg.V2.NoMetrics = !metrics
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(setMetrics, "config set metrics"))

	useContext := &cobra.Command{
		Short: "Sets the currently active context.",
		Long:  "Sets the currently active context.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			if _, ok := cfg.V2.Contexts[args[0]]; !ok {
				return fmt.Errorf("context does not exist: %s", args[0])
			}
			cfg.V2.ActiveContext = args[0]
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(useContext, "config use context"))

	return commands
}
<commit_msg>Command for getting a context config<commit_after>package cmds

import (
	"errors"
	"fmt"
	"os"

	"github.com/pachyderm/pachyderm/src/client/pkg/config"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmdutil"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/spf13/cobra"
)

// Cmds returns a slice containing admin commands.
func Cmds() []*cobra.Command {
	marshaller := &jsonpb.Marshaler{
		Indent:   "  ",
		OrigName: true,
	}

	var commands []*cobra.Command

	getMetrics := &cobra.Command{
		Short: "Gets whether metrics are enabled.",
		Long:  "Gets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(0, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			fmt.Printf("%v\n", !cfg.V2.NoMetrics)
			return nil
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(getMetrics, "config get metrics"))

	setMetrics := &cobra.Command{
		Short: "Sets whether metrics are enabled.",
		Long:  "Sets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			metrics := true
			if args[0] == "false" {
				metrics = false
			} else if args[0] != "true" {
				return errors.New("invalid argument; use either `true` or `false`")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			cfg.V2.NoMetrics = !metrics
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(setMetrics, "config set metrics"))

	useContext := &cobra.Command{
		Short: "Sets the currently active context.",
		Long:  "Sets the currently active context.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			if _, ok := cfg.V2.Contexts[args[0]]; !ok {
				return fmt.Errorf("context does not exist: %s", args[0])
			}
			cfg.V2.ActiveContext = args[0]
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(useContext, "config use context"))

	var active bool
	getContext := &cobra.Command{
		Short: "Gets a context.",
		Long:  "Gets the config of a context by its name, or the current active one.",
		Run: cmdutil.RunBoundedArgs(0, 1, func(args []string) (retErr error) {
			if active && len(args) == 1 {
				return errors.New("cannot get both the active context, and a specific one by its name")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			var name string
			if active {
				name = cfg.V2.ActiveContext
			} else {
				name = args[0]
			}

			context, ok := cfg.V2.Contexts[name]
			if !ok {
				return fmt.Errorf("context does not exist: %s", name)
			}

			return marshaller.Marshal(os.Stdout, context)
		}),
	}
	getContext.Flags().BoolVar(&active, "active", false, "Get the active context.")
	commands = append(commands, cmdutil.CreateAlias(getContext, "config get context"))

	return commands
}
<|endoftext|>