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

package linux

import (
	"time"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/sentry/arch"
	"gvisor.dev/gvisor/pkg/sentry/fs"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	ktime "gvisor.dev/gvisor/pkg/sentry/kernel/time"
	"gvisor.dev/gvisor/pkg/sentry/socket"
	"gvisor.dev/gvisor/pkg/sentry/usermem"
	"gvisor.dev/gvisor/pkg/syserror"
	"gvisor.dev/gvisor/pkg/waiter"
)

const (
	// EventMaskWrite contains events that can be triggered on writes.
	//
	// Note that EventHUp is not going to happen for pipes but may for
	// implementations of poll on some sockets, see net/core/datagram.c.
	EventMaskWrite = waiter.EventOut | waiter.EventHUp | waiter.EventErr
)

// Write implements linux syscall write(2).
func Write(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	size := args[2].SizeT()

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Check that the size is legitimate.
	si := int(size)
	if si < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Get the source of the write.
	src, err := t.SingleIOSequence(addr, si, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := writev(t, file, src)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "write", file)
}

// Pwrite64 implements linux syscall pwrite64(2).
func Pwrite64(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	size := args[2].SizeT()
	offset := args[3].Int64()

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the offset is legitimate.
	if offset < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Is writing at an offset supported?
	if !file.Flags().Pwrite {
		return 0, nil, syserror.ESPIPE
	}

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Check that the size is legitimate.
	si := int(size)
	if si < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Get the source of the write.
	src, err := t.SingleIOSequence(addr, si, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := pwritev(t, file, src, offset)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwrite64", file)
}

// Writev implements linux syscall writev(2).
func Writev(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	iovcnt := int(args[2].Int())

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Read the iovecs that specify the source of the write.
	src, err := t.IovecsIOSequence(addr, iovcnt, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := writev(t, file, src)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "writev", file)
}

// Pwritev implements linux syscall pwritev(2).
func Pwritev(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	iovcnt := int(args[2].Int())
	offset := args[3].Int64()

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the offset is legitimate.
	if offset < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Is writing at an offset supported?
	if !file.Flags().Pwrite {
		return 0, nil, syserror.ESPIPE
	}

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Read the iovecs that specify the source of the write.
	src, err := t.IovecsIOSequence(addr, iovcnt, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := pwritev(t, file, src, offset)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwritev", file)
}

// Pwritev2 implements linux syscall pwritev2(2).
// TODO(b/120161091): Implement O_SYNC and D_SYNC functionality.
func Pwritev2(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	// While the syscall is
	// pwritev2(int fd, struct iovec* iov, int iov_cnt, off_t offset, int flags)
	// the linux internal call
	// (https://elixir.bootlin.com/linux/v4.18/source/fs/read_write.c#L1354)
	// splits the offset argument into a high/low value for compatibility with
	// 32-bit architectures. The flags argument is the 5th argument.

	fd := args[0].Int()
	addr := args[1].Pointer()
	iovcnt := int(args[2].Int())
	offset := args[3].Int64()
	flags := int(args[5].Int())

	if int(args[4].Int())&0x4 == 1 {
		return 0, nil, syserror.EACCES
	}

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the offset is legitimate.
	if offset < -1 {
		return 0, nil, syserror.EINVAL
	}

	// Is writing at an offset supported?
	if offset > -1 && !file.Flags().Pwrite {
		return 0, nil, syserror.ESPIPE
	}

	// Note: gVisor does not implement the RWF_HIPRI feature, but the flag is
	// accepted as a valid flag argument for pwritev2.
	if flags&^linux.RWF_VALID != 0 {
		return uintptr(flags), nil, syserror.EOPNOTSUPP
	}

	// Check that the file is writeable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Read the iovecs that specify the source of the write.
	src, err := t.IovecsIOSequence(addr, iovcnt, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	// If pwritev2 is called with an offset of -1, writev is called.
	if offset == -1 {
		n, err := writev(t, file, src)
		t.IOUsage().AccountWriteSyscall(n)
		return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwritev2", file)
	}

	n, err := pwritev(t, file, src, offset)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwritev2", file)
}

func writev(t *kernel.Task, f *fs.File, src usermem.IOSequence) (int64, error) {
	n, err := f.Writev(t, src)
	if err != syserror.ErrWouldBlock || f.Flags().NonBlocking {
		if n > 0 {
			// Queue notification if we wrote anything.
			f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
		}
		return n, err
	}

	// Sockets support write timeouts.
	var haveDeadline bool
	var deadline ktime.Time
	if s, ok := f.FileOperations.(socket.Socket); ok {
		dl := s.SendTimeout()
		if dl < 0 && err == syserror.ErrWouldBlock {
			return n, err
		}
		if dl > 0 {
			deadline = t.Kernel().MonotonicClock().Now().Add(time.Duration(dl) * time.Nanosecond)
			haveDeadline = true
		}
	}

	// Register for notifications.
	w, ch := waiter.NewChannelEntry(nil)
	f.EventRegister(&w, EventMaskWrite)

	total := n
	for {
		// Shorten src to reflect bytes previously written.
		src = src.DropFirst64(n)

		// Issue the request and break out if it completes with
		// anything other than "would block".
		n, err = f.Writev(t, src)
		total += n
		if err != syserror.ErrWouldBlock {
			break
		}

		// Wait for a notification that we should retry.
		if err = t.BlockWithDeadline(ch, haveDeadline, deadline); err != nil {
			if err == syserror.ETIMEDOUT {
				err = syserror.ErrWouldBlock
			}
			break
		}
	}

	f.EventUnregister(&w)

	if total > 0 {
		// Queue notification if we wrote anything.
		f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
	}

	return total, err
}

func pwritev(t *kernel.Task, f *fs.File, src usermem.IOSequence, offset int64) (int64, error) {
	n, err := f.Pwritev(t, src, offset)
	if err != syserror.ErrWouldBlock || f.Flags().NonBlocking {
		if n > 0 {
			// Queue notification if we wrote anything.
			f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
		}
		return n, err
	}

	// Register for notifications.
	w, ch := waiter.NewChannelEntry(nil)
	f.EventRegister(&w, EventMaskWrite)

	total := n
	for {
		// Shorten src to reflect bytes previously written.
		src = src.DropFirst64(n)

		// Issue the request and break out if it completes with
		// anything other than "would block".
		n, err = f.Pwritev(t, src, offset+total)
		total += n
		if err != syserror.ErrWouldBlock {
			break
		}

		// Wait for a notification that we should retry.
		if err = t.Block(ch); err != nil {
			break
		}
	}

	f.EventUnregister(&w)

	if total > 0 {
		// Queue notification if we wrote anything.
		f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
	}

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

package linux

import (
	"time"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/sentry/arch"
	"gvisor.dev/gvisor/pkg/sentry/fs"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	ktime "gvisor.dev/gvisor/pkg/sentry/kernel/time"
	"gvisor.dev/gvisor/pkg/sentry/socket"
	"gvisor.dev/gvisor/pkg/sentry/usermem"
	"gvisor.dev/gvisor/pkg/syserror"
	"gvisor.dev/gvisor/pkg/waiter"
)

const (
	// EventMaskWrite contains events that can be triggered on writes.
	//
	// Note that EventHUp is not going to happen for pipes but may for
	// implementations of poll on some sockets, see net/core/datagram.c.
	EventMaskWrite = waiter.EventOut | waiter.EventHUp | waiter.EventErr
)

// Write implements linux syscall write(2).
func Write(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	size := args[2].SizeT()

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Check that the size is legitimate.
	si := int(size)
	if si < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Get the source of the write.
	src, err := t.SingleIOSequence(addr, si, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := writev(t, file, src)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "write", file)
}

// Pwrite64 implements linux syscall pwrite64(2).
func Pwrite64(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	size := args[2].SizeT()
	offset := args[3].Int64()

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the offset is legitimate.
	if offset < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Is writing at an offset supported?
	if !file.Flags().Pwrite {
		return 0, nil, syserror.ESPIPE
	}

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Check that the size is legitimate.
	si := int(size)
	if si < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Get the source of the write.
	src, err := t.SingleIOSequence(addr, si, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := pwritev(t, file, src, offset)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwrite64", file)
}

// Writev implements linux syscall writev(2).
func Writev(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	iovcnt := int(args[2].Int())

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Read the iovecs that specify the source of the write.
	src, err := t.IovecsIOSequence(addr, iovcnt, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := writev(t, file, src)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "writev", file)
}

// Pwritev implements linux syscall pwritev(2).
func Pwritev(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	fd := args[0].Int()
	addr := args[1].Pointer()
	iovcnt := int(args[2].Int())
	offset := args[3].Int64()

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the offset is legitimate.
	if offset < 0 {
		return 0, nil, syserror.EINVAL
	}

	// Is writing at an offset supported?
	if !file.Flags().Pwrite {
		return 0, nil, syserror.ESPIPE
	}

	// Check that the file is writable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Read the iovecs that specify the source of the write.
	src, err := t.IovecsIOSequence(addr, iovcnt, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	n, err := pwritev(t, file, src, offset)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwritev", file)
}

// Pwritev2 implements linux syscall pwritev2(2).
func Pwritev2(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	// While the syscall is
	// pwritev2(int fd, struct iovec* iov, int iov_cnt, off_t offset, int flags)
	// the linux internal call
	// (https://elixir.bootlin.com/linux/v4.18/source/fs/read_write.c#L1354)
	// splits the offset argument into a high/low value for compatibility with
	// 32-bit architectures. The flags argument is the 5th argument.

	fd := args[0].Int()
	addr := args[1].Pointer()
	iovcnt := int(args[2].Int())
	offset := args[3].Int64()
	flags := int(args[5].Int())

	if int(args[4].Int())&0x4 == 1 {
		return 0, nil, syserror.EACCES
	}

	file := t.GetFile(fd)
	if file == nil {
		return 0, nil, syserror.EBADF
	}
	defer file.DecRef()

	// Check that the offset is legitimate.
	if offset < -1 {
		return 0, nil, syserror.EINVAL
	}

	// Is writing at an offset supported?
	if offset > -1 && !file.Flags().Pwrite {
		return 0, nil, syserror.ESPIPE
	}

	// Note: gVisor does not implement the RWF_HIPRI feature, but the flag is
	// accepted as a valid flag argument for pwritev2.
	if flags&^linux.RWF_VALID != 0 {
		return uintptr(flags), nil, syserror.EOPNOTSUPP
	}

	// Check that the file is writeable.
	if !file.Flags().Write {
		return 0, nil, syserror.EBADF
	}

	// Read the iovecs that specify the source of the write.
	src, err := t.IovecsIOSequence(addr, iovcnt, usermem.IOOpts{
		AddressSpaceActive: true,
	})
	if err != nil {
		return 0, nil, err
	}

	// If pwritev2 is called with an offset of -1, writev is called.
	if offset == -1 {
		n, err := writev(t, file, src)
		t.IOUsage().AccountWriteSyscall(n)
		return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwritev2", file)
	}

	n, err := pwritev(t, file, src, offset)
	t.IOUsage().AccountWriteSyscall(n)
	return uintptr(n), nil, handleIOError(t, n != 0, err, kernel.ERESTARTSYS, "pwritev2", file)
}

func writev(t *kernel.Task, f *fs.File, src usermem.IOSequence) (int64, error) {
	n, err := f.Writev(t, src)
	if err != syserror.ErrWouldBlock || f.Flags().NonBlocking {
		if n > 0 {
			// Queue notification if we wrote anything.
			f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
		}
		return n, err
	}

	// Sockets support write timeouts.
	var haveDeadline bool
	var deadline ktime.Time
	if s, ok := f.FileOperations.(socket.Socket); ok {
		dl := s.SendTimeout()
		if dl < 0 && err == syserror.ErrWouldBlock {
			return n, err
		}
		if dl > 0 {
			deadline = t.Kernel().MonotonicClock().Now().Add(time.Duration(dl) * time.Nanosecond)
			haveDeadline = true
		}
	}

	// Register for notifications.
	w, ch := waiter.NewChannelEntry(nil)
	f.EventRegister(&w, EventMaskWrite)

	total := n
	for {
		// Shorten src to reflect bytes previously written.
		src = src.DropFirst64(n)

		// Issue the request and break out if it completes with
		// anything other than "would block".
		n, err = f.Writev(t, src)
		total += n
		if err != syserror.ErrWouldBlock {
			break
		}

		// Wait for a notification that we should retry.
		if err = t.BlockWithDeadline(ch, haveDeadline, deadline); err != nil {
			if err == syserror.ETIMEDOUT {
				err = syserror.ErrWouldBlock
			}
			break
		}
	}

	f.EventUnregister(&w)

	if total > 0 {
		// Queue notification if we wrote anything.
		f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
	}

	return total, err
}

func pwritev(t *kernel.Task, f *fs.File, src usermem.IOSequence, offset int64) (int64, error) {
	n, err := f.Pwritev(t, src, offset)
	if err != syserror.ErrWouldBlock || f.Flags().NonBlocking {
		if n > 0 {
			// Queue notification if we wrote anything.
			f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
		}
		return n, err
	}

	// Register for notifications.
	w, ch := waiter.NewChannelEntry(nil)
	f.EventRegister(&w, EventMaskWrite)

	total := n
	for {
		// Shorten src to reflect bytes previously written.
		src = src.DropFirst64(n)

		// Issue the request and break out if it completes with
		// anything other than "would block".
		n, err = f.Pwritev(t, src, offset+total)
		total += n
		if err != syserror.ErrWouldBlock {
			break
		}

		// Wait for a notification that we should retry.
		if err = t.Block(ch); err != nil {
			break
		}
	}

	f.EventUnregister(&w)

	if total > 0 {
		// Queue notification if we wrote anything.
		f.Dirent.InotifyEvent(linux.IN_MODIFY, 0)
	}

	return total, err
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package imageupload

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/spf13/cobra"
	"gopkg.in/cheggaaa/pb.v1"
	"k8s.io/api/core/v1"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"

	uploadcdiv1 "kubevirt.io/containerized-data-importer/pkg/apis/uploadcontroller/v1alpha1"
	cdiClientset "kubevirt.io/containerized-data-importer/pkg/client/clientset/versioned"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const (
	// PodPhaseAnnotation is the annotation on a PVC containing the upload pod phase
	PodPhaseAnnotation = "cdi.kubevirt.io/storage.pod.phase"

	uploadRequestAnnotation = "cdi.kubevirt.io/storage.upload.target"

	uploadPodWaitInterval = 2 * time.Second

	uploadProxyURI = "/v1alpha1/upload"
)

var (
	insecure       bool
	uploadProxyURL string
	pvcName        string
	pvcSize        string
	storageClass   string
	imagePath      string

	uploadPodWaitSecs uint = 60
	accessMode             = "ReadWriteOnce"
	noCreate               = false
)

// HTTPClientCreator is a function that creates http clients
type HTTPClientCreator func(bool) *http.Client

var httpClientCreatorFunc HTTPClientCreator

// SetHTTPClientCreator allows overriding the default http client
// useful for unit tests
func SetHTTPClientCreator(f HTTPClientCreator) {
	httpClientCreatorFunc = f
}

// SetDefaultHTTPClientCreator sets the http client creator back to default
func SetDefaultHTTPClientCreator() {
	httpClientCreatorFunc = getHTTPClient
}

func init() {
	SetDefaultHTTPClientCreator()
}

// NewImageUploadCommand returns a cobra.Command for handling the the uploading of VM images
func NewImageUploadCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "image-upload",
		Short:   "Upload a VM image to a PersistentVolumeClaim.",
		Example: usage(),
		Args:    cobra.ExactArgs(0),
		RunE: func(cmd *cobra.Command, args []string) error {
			v := command{clientConfig: clientConfig}
			return v.run(cmd, args)
		},
	}
	cmd.Flags().BoolVar(&insecure, "insecure", insecure, "Allow insecure server connections when using HTTPS.")
	cmd.Flags().StringVar(&uploadProxyURL, "uploadproxy-url", "", "The URL of the cdi-upload proxy service.")
	cmd.MarkFlagRequired("uploadproxy-url")
	cmd.Flags().StringVar(&pvcName, "pvc-name", "", "The destination PVC.")
	cmd.MarkFlagRequired("pvc-name")
	cmd.Flags().StringVar(&pvcSize, "pvc-size", "", "The size of the PVC to create (ex. 10Gi, 500Mi).")
	cmd.Flags().StringVar(&storageClass, "storage-class", "", "The storage class for the PVC.")
	cmd.Flags().StringVar(&accessMode, "access-mode", accessMode, "The access mode for the PVC.")
	cmd.Flags().StringVar(&imagePath, "image-path", "", "Path to the local VM image.")
	cmd.MarkFlagRequired("image-path")
	cmd.Flags().BoolVar(&noCreate, "no-create", noCreate, "Don't attempt to create a new PVC.")
	cmd.Flags().UintVar(&uploadPodWaitSecs, "wait-secs", uploadPodWaitSecs, "Seconds to wait for upload pod to start.")
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

func usage() string {
	usage := `  # Upload a local disk image to a newly created PersistentVolumeClaim:
	virtctl image-upload --upload-proxy-url=https://cdi-uploadproxy.mycluster.com --pvc-name=upload-pvc --pvc-size=10Gi --image-path=/images/fedora28.qcow2`
	return usage
}

type command struct {
	clientConfig clientcmd.ClientConfig
}

func (c *command) run(cmd *cobra.Command, args []string) error {
	file, err := os.Open(imagePath)
	if err != nil {
		return err
	}
	defer file.Close()

	namespace, _, err := c.clientConfig.Namespace()
	if err != nil {
		return err
	}

	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(c.clientConfig)
	if err != nil {
		return fmt.Errorf("Cannot obtain KubeVirt client: %v", err)
	}

	pvc, err := getAndValidateUploadPVC(virtClient, namespace, pvcName, noCreate)
	if err != nil {
		if !(k8serrors.IsNotFound(err) && !noCreate) {
			return err
		}

		pvc, err = createUploadPVC(virtClient, namespace, pvcName, pvcSize, storageClass, accessMode)
		if err != nil {
			return err
		}

		fmt.Printf("PVC %s/%s created\n", namespace, pvc.Name)
	} else {
		pvc, err = ensurePVCSupportsUpload(virtClient, pvc)
		if err != nil {
			return err
		}

		fmt.Printf("Using existing PVC %s/%s\n", namespace, pvc.Name)
	}

	err = waitUploadPodRunning(virtClient, namespace, pvcName, uploadPodWaitInterval, time.Duration(uploadPodWaitSecs)*time.Second)
	if err != nil {
		return err
	}

	token, err := getUploadToken(virtClient.CdiClient(), namespace, pvcName)
	if err != nil {
		return err
	}

	fmt.Printf("Uploading data to %s\n", uploadProxyURL)

	err = uploadData(uploadProxyURL, token, file, insecure)
	if err != nil {
		return err
	}

	fmt.Printf("Uploading %s completed successfully\n", imagePath)

	return nil
}

func getHTTPClient(insecure bool) *http.Client {
	client := &http.Client{}

	if insecure {
		client.Transport = &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
	}

	return client
}

func uploadData(uploadProxyURL, token string, file *os.File, insecure bool) error {
	url := uploadProxyURL + uploadProxyURI

	fi, err := file.Stat()
	if err != nil {
		return err
	}

	bar := pb.New64(fi.Size()).SetUnits(pb.U_BYTES)
	reader := bar.NewProxyReader(file)

	client := httpClientCreatorFunc(insecure)
	req, _ := http.NewRequest("POST", url, reader)

	req.Header.Add("Authorization", "Bearer "+token)
	req.Header.Add("Content-Type", "application/octet-stream")

	fmt.Println()
	bar.Start()

	resp, err := client.Do(req)

	bar.Finish()
	fmt.Println()

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("Unexpected return value %d", resp.StatusCode)
	}

	return nil
}

func getUploadToken(client cdiClientset.Interface, namespace, name string) (string, error) {
	request := &uploadcdiv1.UploadTokenRequest{
		ObjectMeta: metav1.ObjectMeta{
			Name: "token-for-virtctl",
		},
		Spec: uploadcdiv1.UploadTokenRequestSpec{
			PvcName: name,
		},
	}

	response, err := client.UploadV1alpha1().UploadTokenRequests(namespace).Create(request)
	if err != nil {
		return "", err
	}

	return response.Status.Token, nil
}

func waitUploadPodRunning(client kubernetes.Interface, namespace, name string, interval, timeout time.Duration) error {
	serviceName := "cdi-upload-" + name
	loggedStatus := false

	err := wait.PollImmediate(interval, timeout, func() (bool, error) {
		pvc, err := client.CoreV1().PersistentVolumeClaims(namespace).Get(name, metav1.GetOptions{})
		if err != nil {
			return false, err
		}

		endpoints, err := client.CoreV1().Endpoints(namespace).Get(serviceName, metav1.GetOptions{})
		if err != nil {
			if k8serrors.IsNotFound(err) {
				return false, nil
			}
			return false, err
		}

		podPhase, _ := pvc.Annotations[PodPhaseAnnotation]

		done := (podPhase == string(v1.PodRunning)) && (len(endpoints.Subsets) > 0)

		if !done && !loggedStatus {
			fmt.Printf("Waiting for PVC %s upload pod to be running...\n", name)
			loggedStatus = true
		}

		if done && loggedStatus {
			// be really sure
			time.Sleep(interval)
			fmt.Printf("Pod now running\n")
		}

		return done, nil
	})

	return err
}

func createUploadPVC(client kubernetes.Interface, namespace, name, size, storageClass, accessMode string) (*v1.PersistentVolumeClaim, error) {
	quantity, err := resource.ParseQuantity(size)
	if err != nil {
		return nil, err
	}

	pvc := &v1.PersistentVolumeClaim{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			Annotations: map[string]string{
				uploadRequestAnnotation: "",
			},
		},
		Spec: v1.PersistentVolumeClaimSpec{
			Resources: v1.ResourceRequirements{
				Requests: v1.ResourceList{
					v1.ResourceStorage: quantity,
				},
			},
		},
	}

	if storageClass != "" {
		pvc.Spec.StorageClassName = &storageClass
	}

	if accessMode != "" {
		pvc.Spec.AccessModes = []v1.PersistentVolumeAccessMode{v1.PersistentVolumeAccessMode(accessMode)}
	}

	pvc, err = client.CoreV1().PersistentVolumeClaims(namespace).Create(pvc)
	if err != nil {
		return nil, err
	}

	return pvc, nil
}

func ensurePVCSupportsUpload(client kubernetes.Interface, pvc *v1.PersistentVolumeClaim) (*v1.PersistentVolumeClaim, error) {
	var err error
	_, hasAnnotation := pvc.Annotations[uploadRequestAnnotation]

	if !hasAnnotation {
		pvc.Annotations[uploadRequestAnnotation] = ""
		pvc, err = client.CoreV1().PersistentVolumeClaims(pvc.Namespace).Update(pvc)
		if err != nil {
			return nil, err
		}
	}

	return pvc, nil
}

func getAndValidateUploadPVC(client kubernetes.Interface, namespace, name string, shouldExist bool) (*v1.PersistentVolumeClaim, error) {
	pvc, err := client.CoreV1().PersistentVolumeClaims(namespace).Get(pvcName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	// for PVCs that exist, we ony want to use them if
	// 1. They have not already been used AND EITHER
	//   a. shouldExist is true
	//   b. shouldExist is false AND the upload annotation exists

	_, isUploadPVC := pvc.Annotations[uploadRequestAnnotation]
	podPhase, _ := pvc.Annotations[PodPhaseAnnotation]

	if podPhase == string(v1.PodSucceeded) {
		return nil, fmt.Errorf("PVC %s already successfully imported/cloned/updated", name)
	}

	if !shouldExist && !isUploadPVC {
		return nil, fmt.Errorf("PVC %s not available for upload", name)
	}

	return pvc, nil
}
<commit_msg>Fix typo in virtctl imageupload usage example<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package imageupload

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/spf13/cobra"
	"gopkg.in/cheggaaa/pb.v1"
	"k8s.io/api/core/v1"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"

	uploadcdiv1 "kubevirt.io/containerized-data-importer/pkg/apis/uploadcontroller/v1alpha1"
	cdiClientset "kubevirt.io/containerized-data-importer/pkg/client/clientset/versioned"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const (
	// PodPhaseAnnotation is the annotation on a PVC containing the upload pod phase
	PodPhaseAnnotation = "cdi.kubevirt.io/storage.pod.phase"

	uploadRequestAnnotation = "cdi.kubevirt.io/storage.upload.target"

	uploadPodWaitInterval = 2 * time.Second

	uploadProxyURI = "/v1alpha1/upload"
)

var (
	insecure       bool
	uploadProxyURL string
	pvcName        string
	pvcSize        string
	storageClass   string
	imagePath      string

	uploadPodWaitSecs uint = 60
	accessMode             = "ReadWriteOnce"
	noCreate               = false
)

// HTTPClientCreator is a function that creates http clients
type HTTPClientCreator func(bool) *http.Client

var httpClientCreatorFunc HTTPClientCreator

// SetHTTPClientCreator allows overriding the default http client
// useful for unit tests
func SetHTTPClientCreator(f HTTPClientCreator) {
	httpClientCreatorFunc = f
}

// SetDefaultHTTPClientCreator sets the http client creator back to default
func SetDefaultHTTPClientCreator() {
	httpClientCreatorFunc = getHTTPClient
}

func init() {
	SetDefaultHTTPClientCreator()
}

// NewImageUploadCommand returns a cobra.Command for handling the the uploading of VM images
func NewImageUploadCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "image-upload",
		Short:   "Upload a VM image to a PersistentVolumeClaim.",
		Example: usage(),
		Args:    cobra.ExactArgs(0),
		RunE: func(cmd *cobra.Command, args []string) error {
			v := command{clientConfig: clientConfig}
			return v.run(cmd, args)
		},
	}
	cmd.Flags().BoolVar(&insecure, "insecure", insecure, "Allow insecure server connections when using HTTPS.")
	cmd.Flags().StringVar(&uploadProxyURL, "uploadproxy-url", "", "The URL of the cdi-upload proxy service.")
	cmd.MarkFlagRequired("uploadproxy-url")
	cmd.Flags().StringVar(&pvcName, "pvc-name", "", "The destination PVC.")
	cmd.MarkFlagRequired("pvc-name")
	cmd.Flags().StringVar(&pvcSize, "pvc-size", "", "The size of the PVC to create (ex. 10Gi, 500Mi).")
	cmd.Flags().StringVar(&storageClass, "storage-class", "", "The storage class for the PVC.")
	cmd.Flags().StringVar(&accessMode, "access-mode", accessMode, "The access mode for the PVC.")
	cmd.Flags().StringVar(&imagePath, "image-path", "", "Path to the local VM image.")
	cmd.MarkFlagRequired("image-path")
	cmd.Flags().BoolVar(&noCreate, "no-create", noCreate, "Don't attempt to create a new PVC.")
	cmd.Flags().UintVar(&uploadPodWaitSecs, "wait-secs", uploadPodWaitSecs, "Seconds to wait for upload pod to start.")
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

func usage() string {
	usage := `  # Upload a local disk image to a newly created PersistentVolumeClaim:
	virtctl image-upload --uploadproxy-url=https://cdi-uploadproxy.mycluster.com --pvc-name=upload-pvc --pvc-size=10Gi --image-path=/images/fedora28.qcow2`
	return usage
}

type command struct {
	clientConfig clientcmd.ClientConfig
}

func (c *command) run(cmd *cobra.Command, args []string) error {
	file, err := os.Open(imagePath)
	if err != nil {
		return err
	}
	defer file.Close()

	namespace, _, err := c.clientConfig.Namespace()
	if err != nil {
		return err
	}

	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(c.clientConfig)
	if err != nil {
		return fmt.Errorf("Cannot obtain KubeVirt client: %v", err)
	}

	pvc, err := getAndValidateUploadPVC(virtClient, namespace, pvcName, noCreate)
	if err != nil {
		if !(k8serrors.IsNotFound(err) && !noCreate) {
			return err
		}

		pvc, err = createUploadPVC(virtClient, namespace, pvcName, pvcSize, storageClass, accessMode)
		if err != nil {
			return err
		}

		fmt.Printf("PVC %s/%s created\n", namespace, pvc.Name)
	} else {
		pvc, err = ensurePVCSupportsUpload(virtClient, pvc)
		if err != nil {
			return err
		}

		fmt.Printf("Using existing PVC %s/%s\n", namespace, pvc.Name)
	}

	err = waitUploadPodRunning(virtClient, namespace, pvcName, uploadPodWaitInterval, time.Duration(uploadPodWaitSecs)*time.Second)
	if err != nil {
		return err
	}

	token, err := getUploadToken(virtClient.CdiClient(), namespace, pvcName)
	if err != nil {
		return err
	}

	fmt.Printf("Uploading data to %s\n", uploadProxyURL)

	err = uploadData(uploadProxyURL, token, file, insecure)
	if err != nil {
		return err
	}

	fmt.Printf("Uploading %s completed successfully\n", imagePath)

	return nil
}

func getHTTPClient(insecure bool) *http.Client {
	client := &http.Client{}

	if insecure {
		client.Transport = &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
	}

	return client
}

func uploadData(uploadProxyURL, token string, file *os.File, insecure bool) error {
	url := uploadProxyURL + uploadProxyURI

	fi, err := file.Stat()
	if err != nil {
		return err
	}

	bar := pb.New64(fi.Size()).SetUnits(pb.U_BYTES)
	reader := bar.NewProxyReader(file)

	client := httpClientCreatorFunc(insecure)
	req, _ := http.NewRequest("POST", url, reader)

	req.Header.Add("Authorization", "Bearer "+token)
	req.Header.Add("Content-Type", "application/octet-stream")

	fmt.Println()
	bar.Start()

	resp, err := client.Do(req)

	bar.Finish()
	fmt.Println()

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("Unexpected return value %d", resp.StatusCode)
	}

	return nil
}

func getUploadToken(client cdiClientset.Interface, namespace, name string) (string, error) {
	request := &uploadcdiv1.UploadTokenRequest{
		ObjectMeta: metav1.ObjectMeta{
			Name: "token-for-virtctl",
		},
		Spec: uploadcdiv1.UploadTokenRequestSpec{
			PvcName: name,
		},
	}

	response, err := client.UploadV1alpha1().UploadTokenRequests(namespace).Create(request)
	if err != nil {
		return "", err
	}

	return response.Status.Token, nil
}

func waitUploadPodRunning(client kubernetes.Interface, namespace, name string, interval, timeout time.Duration) error {
	serviceName := "cdi-upload-" + name
	loggedStatus := false

	err := wait.PollImmediate(interval, timeout, func() (bool, error) {
		pvc, err := client.CoreV1().PersistentVolumeClaims(namespace).Get(name, metav1.GetOptions{})
		if err != nil {
			return false, err
		}

		endpoints, err := client.CoreV1().Endpoints(namespace).Get(serviceName, metav1.GetOptions{})
		if err != nil {
			if k8serrors.IsNotFound(err) {
				return false, nil
			}
			return false, err
		}

		podPhase, _ := pvc.Annotations[PodPhaseAnnotation]

		done := (podPhase == string(v1.PodRunning)) && (len(endpoints.Subsets) > 0)

		if !done && !loggedStatus {
			fmt.Printf("Waiting for PVC %s upload pod to be running...\n", name)
			loggedStatus = true
		}

		if done && loggedStatus {
			// be really sure
			time.Sleep(interval)
			fmt.Printf("Pod now running\n")
		}

		return done, nil
	})

	return err
}

func createUploadPVC(client kubernetes.Interface, namespace, name, size, storageClass, accessMode string) (*v1.PersistentVolumeClaim, error) {
	quantity, err := resource.ParseQuantity(size)
	if err != nil {
		return nil, err
	}

	pvc := &v1.PersistentVolumeClaim{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			Annotations: map[string]string{
				uploadRequestAnnotation: "",
			},
		},
		Spec: v1.PersistentVolumeClaimSpec{
			Resources: v1.ResourceRequirements{
				Requests: v1.ResourceList{
					v1.ResourceStorage: quantity,
				},
			},
		},
	}

	if storageClass != "" {
		pvc.Spec.StorageClassName = &storageClass
	}

	if accessMode != "" {
		pvc.Spec.AccessModes = []v1.PersistentVolumeAccessMode{v1.PersistentVolumeAccessMode(accessMode)}
	}

	pvc, err = client.CoreV1().PersistentVolumeClaims(namespace).Create(pvc)
	if err != nil {
		return nil, err
	}

	return pvc, nil
}

func ensurePVCSupportsUpload(client kubernetes.Interface, pvc *v1.PersistentVolumeClaim) (*v1.PersistentVolumeClaim, error) {
	var err error
	_, hasAnnotation := pvc.Annotations[uploadRequestAnnotation]

	if !hasAnnotation {
		pvc.Annotations[uploadRequestAnnotation] = ""
		pvc, err = client.CoreV1().PersistentVolumeClaims(pvc.Namespace).Update(pvc)
		if err != nil {
			return nil, err
		}
	}

	return pvc, nil
}

func getAndValidateUploadPVC(client kubernetes.Interface, namespace, name string, shouldExist bool) (*v1.PersistentVolumeClaim, error) {
	pvc, err := client.CoreV1().PersistentVolumeClaims(namespace).Get(pvcName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	// for PVCs that exist, we ony want to use them if
	// 1. They have not already been used AND EITHER
	//   a. shouldExist is true
	//   b. shouldExist is false AND the upload annotation exists

	_, isUploadPVC := pvc.Annotations[uploadRequestAnnotation]
	podPhase, _ := pvc.Annotations[PodPhaseAnnotation]

	if podPhase == string(v1.PodSucceeded) {
		return nil, fmt.Errorf("PVC %s already successfully imported/cloned/updated", name)
	}

	if !shouldExist && !isUploadPVC {
		return nil, fmt.Errorf("PVC %s not available for upload", name)
	}

	return pvc, nil
}
<|endoftext|>


<commit_before>package ws

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"sync"
)

type Watcher interface {
	Watch(r *Res) error
	Close() error
}
type Op uint

const (
	Add Op = 1 << iota
	Change
	Remove
	_
	Create
	Modify
	Delete
	_
	WsMask Op = 0x0F
	FsMask Op = 0xF0
)

type Config struct {
	// CapHint hints the expected peak resource capacity.
	CapHint uint
	// Watcher returns a new watcher given workspace control.
	// Mounting a path results in a snapshot if no Watcher is configured.
	Watcher func(Controller) (Watcher, error)
	// Handler handles resource operation events if set.
	Handler func(Op, *Res)
	// Filter filters resources if set.
	// If filter returns true the resource is flagged with FlagIgnore but remains in the workspace.
	// Ignored directories are not read.
	Filter func(*Res) bool
}

func (c *Config) filldefaullts() {
	if c.Filter == nil {
		c.Filter = nilfilter
	}
	if c.Handler == nil {
		c.Handler = nilhandler
	}
}
func nilhandler(Op, *Res) {}
func nilfilter(*Res) bool { return false }

type Controller interface {
	Control(op Op, id Id, name string) error
}

type Ws struct {
	sync.RWMutex
	config  Config
	root    *Res
	all     map[Id]*Res
	watcher Watcher
}

// New creates a workspace configured with c.
func New(c Config) *Ws {
	r := &Res{Id: NewId("/")}
	c.filldefaullts()
	w := &Ws{config: c, root: r, all: make(map[Id]*Res, c.CapHint)}
	w.all[r.Id] = r
	return w
}

// Mount adds the directory tree at path to the workspace.
func (w *Ws) Mount(path string) (*Res, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("not a directory")
	}
	r, err := w.mount(path)
	if err != nil {
		return r, err
	}
	r.Lock()
	err = read(r, w.config.Filter)
	r.Unlock()
	if err != nil {
		return r, err
	}
	w.Lock()
	defer w.Unlock()
	w.addAllChildren(0, r)
	return r, nil
}
func (w *Ws) mount(path string) (*Res, error) {
	id := NewId(path)
	d, f := filepath.Split(path)
	w.Lock()
	defer w.Unlock()
	if w.watcher == nil && w.config.Watcher != nil {
		watcher, err := w.config.Watcher((*ctrl)(w))
		if err != nil {
			return nil, err
		}
		w.watcher = watcher
	}
	r, ok := w.all[id]
	if ok {
		return r, fmt.Errorf("duplicate")
	}
	r = &Res{Id: id, Name: f, Flag: FlagDir | FlagMount, Dir: &Dir{Path: path}}
	// add virtual parent
	r.Parent = w.addParent(d[:len(d)-1])
	w.all[id] = r
	w.config.Handler(Add, r)
	return r, nil
}

// Close closes the workspace.
// The workspace and all its resources are invalid afterwards.
func (w *Ws) Close() {
	w.Lock()
	defer w.Unlock()
	if w.watcher != nil {
		w.watcher.Close()
		w.watcher = nil
	}
	// scatter garbage
	for id, r := range w.all {
		r.Parent, r.Dir = nil, nil
		delete(w.all, id)
	}
	w.all = nil
	w.root = nil
}
func (w *Ws) addParent(path string) *Res {
	id := NewId(path)
	if r, ok := w.all[id]; ok {
		return r
	}
	parent, name := filepath.Split(path)
	var pa *Res
	if len(parent) > 0 && parent != "/" {
		pa = w.addParent(parent[:len(parent)-1])
	}
	r := &Res{Id: id, Name: name, Flag: FlagDir | FlagLogical, Parent: pa}
	w.all[id] = r
	return r
}
func read(r *Res, filter func(*Res) bool) error {
	f, err := os.Open(r.Dir.Path)
	if err != nil {
		return err
	}
	list, err := f.Readdir(-1)
	f.Close()
	if err != nil {
		return err
	}
	children := make([]*Res, 0, len(list))
	for _, fi := range list {
		children = append(children, newChild(r, fi))
	}
	sort.Sort(byTypeAndName(children))
	r.Children = children
	for _, c := range children {
		if filter(c) {
			c.Flag |= FlagIgnore
			continue
		}
		if c.Flag&FlagDir != 0 {
			if err := read(c, filter); err != nil {
				fmt.Println(err)
			}
		}
	}
	return nil
}
func (w *Ws) addAllChildren(fsop Op, r *Res) {
	for _, c := range r.Children {
		w.all[c.Id] = c
		w.config.Handler(fsop|Add, c)
		if c.Flag&(FlagDir|FlagIgnore) == FlagDir {
			w.addAllChildren(fsop, c)
		}
	}
	if w.watcher != nil {
		err := w.watcher.Watch(r)
		if err != nil {
			fmt.Println(err)
		}
	}
	w.config.Handler(fsop|Change, r)
}
<commit_msg>apply filter to mounted paths closing a workspace leaves resource parent links to reconstruct the path<commit_after>package ws

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"sync"
)

type Watcher interface {
	Watch(r *Res) error
	Close() error
}
type Op uint

const (
	Add Op = 1 << iota
	Change
	Remove
	_
	Create
	Modify
	Delete
	_
	WsMask Op = 0x0F
	FsMask Op = 0xF0
)

type Config struct {
	// CapHint hints the expected peak resource capacity.
	CapHint uint
	// Watcher returns a new watcher given workspace control.
	// Mounting a path results in a snapshot if no Watcher is configured.
	Watcher func(Controller) (Watcher, error)
	// Handler handles resource operation events if set.
	Handler func(Op, *Res)
	// Filter filters resources if set.
	// If filter returns true the resource is flagged with FlagIgnore but remains in the workspace.
	// Ignored directories are not read.
	Filter func(*Res) bool
}

func (c *Config) filldefaullts() {
	if c.Filter == nil {
		c.Filter = nilfilter
	}
	if c.Handler == nil {
		c.Handler = nilhandler
	}
}
func nilhandler(Op, *Res) {}
func nilfilter(*Res) bool { return false }

type Controller interface {
	Control(op Op, id Id, name string) error
}

type Ws struct {
	sync.RWMutex
	config  Config
	root    *Res
	all     map[Id]*Res
	watcher Watcher
}

// New creates a workspace configured with c.
func New(c Config) *Ws {
	r := &Res{Id: NewId("/")}
	c.filldefaullts()
	w := &Ws{config: c, root: r, all: make(map[Id]*Res, c.CapHint)}
	w.all[r.Id] = r
	return w
}

// Mount adds the directory tree at path to the workspace.
func (w *Ws) Mount(path string) (*Res, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("not a directory")
	}
	r, err := w.mount(path)
	if err != nil {
		return r, err
	}
	if w.config.Filter(r) {
		r.Flag |= FlagIgnore
		return r, nil
	}
	r.Lock()
	err = read(r, w.config.Filter)
	r.Unlock()
	if err != nil {
		return r, err
	}
	w.Lock()
	defer w.Unlock()
	w.addAllChildren(0, r)
	return r, nil
}
func (w *Ws) mount(path string) (*Res, error) {
	id := NewId(path)
	d, f := filepath.Split(path)
	w.Lock()
	defer w.Unlock()
	if w.watcher == nil && w.config.Watcher != nil {
		watcher, err := w.config.Watcher((*ctrl)(w))
		if err != nil {
			return nil, err
		}
		w.watcher = watcher
	}
	r, ok := w.all[id]
	if ok {
		return r, fmt.Errorf("duplicate")
	}
	r = &Res{Id: id, Name: f, Flag: FlagDir | FlagMount, Dir: &Dir{Path: path}}
	// add virtual parent
	r.Parent = w.addParent(d[:len(d)-1])
	w.all[id] = r
	w.config.Handler(Add, r)
	return r, nil
}

// Close closes the workspace.
// The workspace and all its resources are invalid afterwards.
func (w *Ws) Close() {
	w.Lock()
	defer w.Unlock()
	if w.watcher != nil {
		w.watcher.Close()
		w.watcher = nil
	}
	// scatter garbage
	for id, r := range w.all {
		r.Dir = nil
		delete(w.all, id)
	}
	w.all = nil
	w.root = nil
}
func (w *Ws) addParent(path string) *Res {
	id := NewId(path)
	if r, ok := w.all[id]; ok {
		return r
	}
	parent, name := filepath.Split(path)
	var pa *Res
	if len(parent) > 0 && parent != "/" {
		pa = w.addParent(parent[:len(parent)-1])
	}
	r := &Res{Id: id, Name: name, Flag: FlagDir | FlagLogical, Parent: pa}
	w.all[id] = r
	return r
}
func read(r *Res, filter func(*Res) bool) error {
	f, err := os.Open(r.Dir.Path)
	if err != nil {
		return err
	}
	list, err := f.Readdir(-1)
	f.Close()
	if err != nil {
		return err
	}
	children := make([]*Res, 0, len(list))
	for _, fi := range list {
		children = append(children, newChild(r, fi))
	}
	sort.Sort(byTypeAndName(children))
	r.Children = children
	for _, c := range children {
		if filter(c) {
			c.Flag |= FlagIgnore
			continue
		}
		if c.Flag&FlagDir != 0 {
			if err := read(c, filter); err != nil {
				fmt.Println(err)
			}
		}
	}
	return nil
}
func (w *Ws) addAllChildren(fsop Op, r *Res) {
	for _, c := range r.Children {
		w.all[c.Id] = c
		w.config.Handler(fsop|Add, c)
		if c.Flag&(FlagDir|FlagIgnore) == FlagDir {
			w.addAllChildren(fsop, c)
		}
	}
	if w.watcher != nil {
		err := w.watcher.Watch(r)
		if err != nil {
			fmt.Println(err)
		}
	}
	w.config.Handler(fsop|Change, r)
}
<|endoftext|>


<commit_before>package droppool

import (
	"math"
	"time"

	"github.com/Hucaru/Valhalla/mnet"
	"github.com/Hucaru/Valhalla/mpacket"
	"github.com/Hucaru/Valhalla/server/item"
	"github.com/Hucaru/Valhalla/server/pos"
)

const (
	SpawnDisappears      = 0
	SpawnNormal          = 1
	SpawnShow            = 2
	SpawnFadeAtTopOfDrop = 3
)

type field interface {
	Send(mpacket.Packet) error
	CalculateFinalDropPos(pos.Data) pos.Data
}

type controller interface {
	Send(mpacket.Packet)
	Conn() mnet.Client
}

// Data structure for the pool
type Data struct {
	instance field
	poolID   int32
	drops map[int32]drop
}

// CreateNewPool for drops
func CreateNewPool(inst field) Data {
	return Data{instance: inst, drops: make(map[int32]drop)}
}

func (pool *Data) nextID() int32 {
	pool.poolID++

	if pool.poolID == math.MaxInt32-1 {
		pool.poolID = math.MaxInt32 / 2
	} else if pool.poolID == 0 {
		pool.poolID = 1
	}

	return pool.poolID
}

// CanClose the pool down
func (pool Data) CanClose() bool {
	return false
}

// PlayerShowDrops when entering instance
func (pool Data) PlayerShowDrops(plr controller) {
	for _, drop := range pool.drops {
		plr.Send(packetShowDrop(SpawnShow, drop))
	}
}

func (pool *Data) RemoveDrop(instant bool, id ...int32) {
	for _, id := range id {
		pool.instance.Send(packetRemoveDrop(instant, id))

		if _, ok := pool.drops[id]; ok {
			delete(pool.drops, id)
		}
	}
}

// PlayerAttemptPickup of item
func (pool *Data) PlayerAttemptPickup(dropID int32, position pos.Data) (bool, item.Data) {
	return false, item.Data{}
}

const itemDistance = 20 // Between 15 and 20?
const itemDisppearTimeout = time.Minute * 2
const itemLootableByAllTimeout = time.Minute * 1

// CreateDrop into field
func (pool *Data) CreateDrop(spawnType byte, dropType byte, mesos int32, dropFrom pos.Data, expire bool, ownerID, partyID int32, items ...item.Data) {
	// TODO: Clean up separation logic, should pass in drop struct
	iCount := len(items)
	var offset int16 = 0

	if mesos > 0 {
		iCount++
	}

	if iCount > 0 {
		offset = int16(itemDistance * (iCount / 2))
	}

	currentTime := time.Now()
	expireTime := currentTime.Add(itemDisppearTimeout).Unix()
	var timeoutTime int64 = 0

	if dropType == DropTimeoutNonOwner || dropType == DropTimeoutNonOwnerParty {
		timeoutTime =  currentTime.Add(itemLootableByAllTimeout).Unix()
	}

	for i, item := range items {
		finalPos := pool.instance.CalculateFinalDropPos(dropFrom) // (dropFrom, xShift)

		finalPos.SetX(finalPos.X() - offset + int16(i*itemDistance)) // This calculation needs to be interpolated to be placed on correct position on ledge

		drop := drop{
			ID:      pool.nextID(),
			ownerID: ownerID,
			partyID: partyID,
			mesos:   0,
			item:    item,

			expireTime:  expireTime,
			timeoutTime: timeoutTime,
			neverExpire: false,

			originPos: dropFrom,
			finalPos:  finalPos,

			dropType: dropType,
		}

		pool.drops[drop.ID] = drop

		pool.instance.Send(packetShowDrop(spawnType, drop))
	}

	if mesos > 0 {
		finalPos := pool.instance.CalculateFinalDropPos(dropFrom)

		if iCount > 1 {
			finalPos.SetX(finalPos.X() - offset + int16((iCount-1)*itemDistance))
		}

		drop := drop{
			ID:      pool.nextID(),
			ownerID: ownerID,
			partyID: partyID,
			mesos:   mesos,

			expireTime:  expireTime,
			timeoutTime: timeoutTime,
			neverExpire: false,

			originPos: dropFrom,
			finalPos:  finalPos,

			dropType: dropType,
		}

		pool.drops[drop.ID] = drop

		pool.instance.Send(packetShowDrop(spawnType, drop))
	}

}

// Update logic for the pool e.g. drops disappear
func (pool *Data) Update(t time.Time) {
	id := make([]int32, 0, len(pool.drops))

	currentTime := time.Now().Unix()

	for _, v := range pool.drops {
		if v.expireTime <= currentTime {
			id = append(id, v.ID)
		}
	}

	if len(id) > 0 {
		pool.RemoveDrop(false, id...)
	}
}
<commit_msg>moved drops from slice to map based<commit_after>package droppool

import (
	"fmt"
	"math"
	"time"

	"github.com/Hucaru/Valhalla/mnet"
	"github.com/Hucaru/Valhalla/mpacket"
	"github.com/Hucaru/Valhalla/server/item"
	"github.com/Hucaru/Valhalla/server/pos"
)

const (
	SpawnDisappears      = 0
	SpawnNormal          = 1
	SpawnShow            = 2
	SpawnFadeAtTopOfDrop = 3
)

type field interface {
	Send(mpacket.Packet) error
	CalculateFinalDropPos(pos.Data) pos.Data
}

type controller interface {
	Send(mpacket.Packet)
	Conn() mnet.Client
}

// Data structure for the pool
type Data struct {
	instance field
	poolID   int32
	drops    map[int32]drop
}

// CreateNewPool for drops
func CreateNewPool(inst field) Data {
	return Data{instance: inst, drops: make(map[int32]drop)}
}

func (pool *Data) nextID() (int32, error) {
	for {
		pool.poolID++

		if pool.poolID == math.MaxInt32-1 {
			pool.poolID = math.MaxInt32 / 2
		} else if pool.poolID == 0 {
			pool.poolID = 1
		}

		if _, ok := pool.drops[pool.poolID]; !ok {
			return pool.poolID, nil
		}
	}

	return 0, fmt.Errorf("No space to generate id in drop pool")
}

// CanClose the pool down
func (pool Data) CanClose() bool {
	return false
}

// PlayerShowDrops when entering instance
func (pool Data) PlayerShowDrops(plr controller) {
	for _, drop := range pool.drops {
		plr.Send(packetShowDrop(SpawnShow, drop))
	}
}

func (pool *Data) RemoveDrop(instant bool, id ...int32) {
	for _, id := range id {
		pool.instance.Send(packetRemoveDrop(instant, id))

		if _, ok := pool.drops[id]; ok {
			delete(pool.drops, id)
		}
	}
}

// PlayerAttemptPickup of item
func (pool *Data) PlayerAttemptPickup(dropID int32, position pos.Data) (bool, item.Data) {
	return false, item.Data{}
}

const itemDistance = 20 // Between 15 and 20?
const itemDisppearTimeout = time.Minute * 2
const itemLootableByAllTimeout = time.Minute * 1

// CreateDrop into field
func (pool *Data) CreateDrop(spawnType byte, dropType byte, mesos int32, dropFrom pos.Data, expire bool, ownerID, partyID int32, items ...item.Data) {
	// TODO: Clean up separation logic, should pass in drop struct
	iCount := len(items)
	var offset int16 = 0

	if mesos > 0 {
		iCount++
	}

	if iCount > 0 {
		offset = int16(itemDistance * (iCount / 2))
	}

	currentTime := time.Now()
	expireTime := currentTime.Add(itemDisppearTimeout).Unix()
	var timeoutTime int64 = 0

	if dropType == DropTimeoutNonOwner || dropType == DropTimeoutNonOwnerParty {
		timeoutTime = currentTime.Add(itemLootableByAllTimeout).Unix()
	}

	for i, item := range items {
		finalPos := pool.instance.CalculateFinalDropPos(dropFrom) // (dropFrom, xShift)

		finalPos.SetX(finalPos.X() - offset + int16(i*itemDistance)) // This calculation needs to be interpolated to be placed on correct position on ledge

		if poolID, err := pool.nextID(); err == nil {
			drop := drop{
				ID:      poolID,
				ownerID: ownerID,
				partyID: partyID,
				mesos:   0,
				item:    item,

				expireTime:  expireTime,
				timeoutTime: timeoutTime,
				neverExpire: false,

				originPos: dropFrom,
				finalPos:  finalPos,

				dropType: dropType,
			}

			pool.drops[drop.ID] = drop

			pool.instance.Send(packetShowDrop(spawnType, drop))
		}
	}

	if mesos > 0 {
		finalPos := pool.instance.CalculateFinalDropPos(dropFrom)

		if iCount > 1 {
			finalPos.SetX(finalPos.X() - offset + int16((iCount-1)*itemDistance))
		}

		if poolID, err := pool.nextID(); err == nil {
			drop := drop{
				ID:      poolID,
				ownerID: ownerID,
				partyID: partyID,
				mesos:   mesos,

				expireTime:  expireTime,
				timeoutTime: timeoutTime,
				neverExpire: false,

				originPos: dropFrom,
				finalPos:  finalPos,

				dropType: dropType,
			}

			pool.drops[drop.ID] = drop

			pool.instance.Send(packetShowDrop(spawnType, drop))
		}
	}
}

// Update logic for the pool e.g. drops disappear
func (pool *Data) Update(t time.Time) {
	id := make([]int32, 0, len(pool.drops))

	currentTime := time.Now().Unix()

	for _, v := range pool.drops {
		if v.expireTime <= currentTime {
			id = append(id, v.ID)
		}
	}

	if len(id) > 0 {
		pool.RemoveDrop(false, id...)
	}
}
<|endoftext|>


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

package metrics

import (
	"time"

	"istio.io/pkg/monitoring"
)

var (
	typeTag = monitoring.MustCreateLabel("type")

	// StartupTime measures the time it takes for the agent to get ready Note: This
	// is dependant on readiness probes. This means our granularity is correlated to
	// the probing interval.
	startupTime = monitoring.NewGauge(
		"startup_duration_seconds",
		"The time from the process starting to being marked ready.",
	)

	// scrapeErrors records total number of failed scrapes.
	scrapeErrors = monitoring.NewSum(
		"scrape_failures_total",
		"The total number of failed scrapes.",
		monitoring.WithLabels(typeTag),
	)
	EnvoyScrapeErrors = scrapeErrors.With(typeTag.Value(ScrapeTypeEnvoy))
	AppScrapeErrors   = scrapeErrors.With(typeTag.Value(ScrapeTypeApp))
	AgentScrapeErrors = scrapeErrors.With(typeTag.Value(ScrapeTypeAgent))

	// ScrapeTotals records total number of scrapes.
	ScrapeTotals = monitoring.NewSum(
		"scrapes_total",
		"The total number of scrapes.",
	)
)

var (
	ScrapeTypeEnvoy = "envoy"
	ScrapeTypeApp   = "application"
	ScrapeTypeAgent = "agent"
)

var (
	processStartTime = time.Now()
)

func RecordStartupTime() {
	startupTime.Record(time.Since(processStartTime).Seconds())
}

func init() {
	monitoring.MustRegister(
		ScrapeTotals,
		scrapeErrors,
		startupTime,
	)
}
<commit_msg>Add logging for agent initialization time (#28416)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package metrics

import (
	"time"

	"istio.io/pkg/log"
	"istio.io/pkg/monitoring"
)

var (
	typeTag = monitoring.MustCreateLabel("type")

	// StartupTime measures the time it takes for the agent to get ready Note: This
	// is dependant on readiness probes. This means our granularity is correlated to
	// the probing interval.
	startupTime = monitoring.NewGauge(
		"startup_duration_seconds",
		"The time from the process starting to being marked ready.",
	)

	// scrapeErrors records total number of failed scrapes.
	scrapeErrors = monitoring.NewSum(
		"scrape_failures_total",
		"The total number of failed scrapes.",
		monitoring.WithLabels(typeTag),
	)
	EnvoyScrapeErrors = scrapeErrors.With(typeTag.Value(ScrapeTypeEnvoy))
	AppScrapeErrors   = scrapeErrors.With(typeTag.Value(ScrapeTypeApp))
	AgentScrapeErrors = scrapeErrors.With(typeTag.Value(ScrapeTypeAgent))

	// ScrapeTotals records total number of scrapes.
	ScrapeTotals = monitoring.NewSum(
		"scrapes_total",
		"The total number of scrapes.",
	)
)

var (
	ScrapeTypeEnvoy = "envoy"
	ScrapeTypeApp   = "application"
	ScrapeTypeAgent = "agent"
)

var (
	processStartTime = time.Now()
)

func RecordStartupTime() {
	delta := time.Since(processStartTime)
	startupTime.Record(delta.Seconds())
	log.Infof("Initialization took %v", delta)
}

func init() {
	monitoring.MustRegister(
		ScrapeTotals,
		scrapeErrors,
		startupTime,
	)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors

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

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

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

package statserver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	network "knative.dev/networking/pkg"
	"knative.dev/serving/pkg/autoscaler/metrics"

	"k8s.io/apimachinery/pkg/types"
)

var (
	msg1 = metrics.StatMessage{
		Key: types.NamespacedName{Namespace: "test-namespace", Name: "test-revision"},
		Stat: metrics.Stat{
			PodName:                   "activator1",
			AverageConcurrentRequests: 2.1,
			RequestCount:              51,
		},
	}
	msg2 = metrics.StatMessage{
		Key: types.NamespacedName{Namespace: "test-namespace", Name: "test-revision2"},
		Stat: metrics.Stat{
			PodName:                   "activator2",
			AverageConcurrentRequests: 2.2,
			RequestCount:              30,
		},
	}
	both = []metrics.StatMessage{msg1, msg2}
)

func TestServerLifecycle(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	eg := errgroup.Group{}
	eg.Go(server.listenAndServe)

	server.listenAddr()
	server.Shutdown(time.Second)

	if err := eg.Wait(); err != nil {
		t.Error("listenAndServe failed.", err)
	}
}

func TestProbe(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	defer server.Shutdown(0)
	go server.listenAndServe()
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/%s", server.listenAddr(), network.ProbePath), nil)
	if err != nil {
		t.Fatal("Error creating request:", err)
	}
	req.Header.Set("User-Agent", "kube-probe/1.15.i.wish")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal("Error roundtripping:", err)
	}
	defer resp.Body.Close()
	if got, want := resp.StatusCode, http.StatusOK; got != want {
		t.Errorf("StatusCode: %v, want: %v", got, want)
	}
}

func TestStatsReceived(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	defer server.Shutdown(0)
	go server.listenAndServe()

	statSink := dialOK(t, server.listenAddr())

	// protobuf
	assertReceivedProto(t, both, statSink, statsCh)

	// json encoding
	assertReceivedJSON(t, msg1, statSink, statsCh)
	assertReceivedJSON(t, msg2, statSink, statsCh)

	closeSink(t, statSink)
}

func TestServerShutdown(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	go server.listenAndServe()

	listenAddr := server.listenAddr()
	statSink := dialOK(t, listenAddr)

	assertReceivedProto(t, both, statSink, statsCh)

	server.Shutdown(time.Second)
	// We own the channel.
	close(statsCh)

	// Send a statistic to the server
	if err := sendProto(statSink, both); err != nil {
		t.Fatal("Expected send to succeed, got:", err)
	}

	// Check the statistic was not received
	_, ok := <-statsCh
	if ok {
		t.Fatal("Received statistic after shutdown")
	}

	// Check connection has been closed with a close control message with a "service restart" close code
	if _, _, err := statSink.NextReader(); err == nil {
		t.Fatal("Connection not closed")
	} else {
		err, ok := err.(*websocket.CloseError)
		if !ok {
			t.Fatal("CloseError not received")
		}
		if err.Code != 1012 {
			t.Fatalf("CloseError with unexpected close code %d received", err.Code)
		}
	}

	// Check that new connections are refused with some error
	if _, err := dial(listenAddr); err == nil {
		t.Fatal("Connection not refused")
	}

	closeSink(t, statSink)
}

func TestServerDoesNotLeakGoroutines(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)
	defer server.Shutdown(0)

	go server.listenAndServe()

	originalGoroutines := runtime.NumGoroutine()

	listenAddr := server.listenAddr()
	statSink := dialOK(t, listenAddr)

	assertReceivedProto(t, both, statSink, statsCh)

	closeSink(t, statSink)

	// Check the number of goroutines eventually reduces to the number there were before the connection was created
	for i := 1000; i >= 0; i-- {
		currentGoRoutines := runtime.NumGoroutine()
		if currentGoRoutines <= originalGoroutines {
			break
		}
		time.Sleep(5 * time.Millisecond)
		if i == 0 {
			t.Fatalf("Current number of goroutines %d is not equal to the original number %d", currentGoRoutines, originalGoroutines)
		}
	}
}

func TestServerNotBucketHost(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServerWithOwnerFunc(statsCh, alwaysFalse)
	defer server.Shutdown(0)

	go server.listenAndServe()

	// Verify the server behaves normally when the host is not a bucket host.
	statSink := dialOK(t, server.listenAddr())
	assertReceivedProto(t, both, statSink, statsCh)
	closeSink(t, statSink)
}

func TestServerOwnerForBucketHost(t *testing.T) {
	// Override the function to mock a bucket host.
	isBucketHost = alwaysTrue

	statsCh := make(chan metrics.StatMessage)
	server := newTestServerWithOwnerFunc(statsCh, alwaysTrue)
	defer server.Shutdown(0)

	go server.listenAndServe()

	// Verify the server behaves normally when the host is a bucket host
	// and the pod is the owner of the bucket.
	statSink := dialOK(t, server.listenAddr())
	assertReceivedProto(t, both, statSink, statsCh)
	closeSink(t, statSink)
}

func TestServerNotOwnerForBucketHost(t *testing.T) {
	// Override the function to mock a bucket host.
	isBucketHost = alwaysTrue

	statsCh := make(chan metrics.StatMessage)
	server := newTestServerWithOwnerFunc(statsCh, alwaysFalse)
	defer server.Shutdown(0)

	go server.listenAndServe()

	_, err := dial(server.listenAddr())
	if err == nil {
		t.Error("Want error from Dial but got none")
	}
}

func BenchmarkStatServer(b *testing.B) {
	statsCh := make(chan metrics.StatMessage, 100)
	server := newTestServer(statsCh)
	go server.listenAndServe()
	defer server.Shutdown(time.Second)

	statSink, err := dial(server.listenAddr())
	if err != nil {
		b.Fatal("Dial failed:", err)
	}

	// The activator sends a bunch of metrics at once usually. This simulates cases with
	// the respective number of active revisions, sending via the activator.
	for _, size := range []int{1, 2, 5, 10, 20, 50, 100} {
		msgs := make([]metrics.StatMessage, 0, size)
		for i := 0; i < size; i++ {
			msgs = append(msgs, msg1)
		}

		b.Run(fmt.Sprintf("json-encoding-%d-msgs", len(msgs)), func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				for _, msg := range msgs {
					if err := sendJSON(statSink, msg); err != nil {
						b.Fatal("Expected send to succeed, but got:", err)
					}
				}

				for range msgs {
					<-statsCh
				}
			}
		})

		b.Run(fmt.Sprintf("proto-encoding-%d-msgs", len(msgs)), func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				if err := sendProto(statSink, msgs); err != nil {
					b.Fatal("Expected send to succeed, but got:", err)
				}

				for range msgs {
					<-statsCh
				}
			}
		})
	}
}

func assertReceivedJSON(t *testing.T, sm metrics.StatMessage, statSink *websocket.Conn, statsCh <-chan metrics.StatMessage) {
	t.Helper()

	if err := sendJSON(statSink, sm); err != nil {
		t.Fatal("Expected send to succeed, got:", err)
	}

	recv := <-statsCh
	if !cmp.Equal(sm, recv) {
		t.Fatalf("StatMessage mismatch: diff (-got, +want) %s", cmp.Diff(recv, sm))
	}
}

func assertReceivedProto(t *testing.T, sms []metrics.StatMessage, statSink *websocket.Conn, statsCh <-chan metrics.StatMessage) {
	t.Helper()

	if err := sendProto(statSink, sms); err != nil {
		t.Fatal("Expected send to succeed, got:", err)
	}

	got := make([]metrics.StatMessage, 0, len(sms))
	for range sms {
		got = append(got, <-statsCh)
	}
	if !cmp.Equal(sms, got) {
		t.Fatalf("StatMessage mismatch: diff (-got, +want) %s", cmp.Diff(got, sms))
	}
}

func dialOK(t *testing.T, serverURL string) *websocket.Conn {
	t.Helper()

	statSink, err := dial(serverURL)
	if err != nil {
		t.Fatal("Dial failed:", err)
	}
	return statSink
}

func dial(serverURL string) (*websocket.Conn, error) {
	u, err := url.Parse(serverURL)
	if err != nil {
		return nil, err
	}
	u.Scheme = "ws"

	dialer := &websocket.Dialer{
		HandshakeTimeout: time.Second,
	}
	statSink, _, err := dialer.Dial(u.String(), nil)
	return statSink, err
}

func sendJSON(statSink *websocket.Conn, sm metrics.StatMessage) error {
	var b bytes.Buffer
	enc := json.NewEncoder(&b)
	if err := enc.Encode(sm); err != nil {
		return fmt.Errorf("failed to encode StatMessage: %w", err)
	}

	if err := statSink.WriteMessage(websocket.TextMessage, b.Bytes()); err != nil {
		return fmt.Errorf("failed to write to stat sink: %w", err)
	}
	return nil
}

func sendProto(statSink *websocket.Conn, sms []metrics.StatMessage) error {
	wsms := metrics.ToWireStatMessages(sms)
	msg, err := wsms.Marshal()
	if err != nil {
		return fmt.Errorf("failed to marshal StatMessage: %w", err)
	}

	if err := statSink.WriteMessage(websocket.BinaryMessage, msg); err != nil {
		return fmt.Errorf("failed to write to stat sink: %w", err)
	}

	return nil
}

func closeSink(t *testing.T, statSink *websocket.Conn) {
	t.Helper()

	if err := statSink.Close(); err != nil {
		t.Fatal("Failed to close", err)
	}
}

const testAddress = "127.0.0.1:0"

type testServer struct {
	*Server
	listenAddrCh chan string
}

func newTestServer(statsCh chan<- metrics.StatMessage) *testServer {
	return newTestServerWithOwnerFunc(statsCh, nil)
}

func newTestServerWithOwnerFunc(statsCh chan<- metrics.StatMessage, f func(bkt string) bool) *testServer {
	return &testServer{
		Server:       New(testAddress, statsCh, zap.NewNop().Sugar(), f),
		listenAddrCh: make(chan string, 1),
	}
}

// listenAddr returns the address on which the server is listening. Blocks until listenAndServe is called.
func (s *testServer) listenAddr() string {
	return <-s.listenAddrCh
}

func (s *testServer) listenAndServe() error {
	listener, err := s.listen()
	if err != nil {
		return err
	}
	return s.serve(&testListener{listener, s.listenAddrCh})
}

type testListener struct {
	net.Listener
	listenAddr chan string
}

func (t *testListener) Accept() (net.Conn, error) {
	t.listenAddr <- "http://" + t.Listener.Addr().String()
	return t.Listener.Accept()
}

func alwaysTrue(_ string) bool {
	return true
}

func alwaysFalse(_ string) bool {
	return false
}
<commit_msg>Micronit: make the URL correct. (#9563)<commit_after>/*
Copyright 2018 The Knative Authors

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

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

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

package statserver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	network "knative.dev/networking/pkg"
	"knative.dev/serving/pkg/autoscaler/metrics"

	"k8s.io/apimachinery/pkg/types"
)

var (
	msg1 = metrics.StatMessage{
		Key: types.NamespacedName{Namespace: "test-namespace", Name: "test-revision"},
		Stat: metrics.Stat{
			PodName:                   "activator1",
			AverageConcurrentRequests: 2.1,
			RequestCount:              51,
		},
	}
	msg2 = metrics.StatMessage{
		Key: types.NamespacedName{Namespace: "test-namespace", Name: "test-revision2"},
		Stat: metrics.Stat{
			PodName:                   "activator2",
			AverageConcurrentRequests: 2.2,
			RequestCount:              30,
		},
	}
	both = []metrics.StatMessage{msg1, msg2}
)

func TestServerLifecycle(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	eg := errgroup.Group{}
	eg.Go(server.listenAndServe)

	server.listenAddr()
	server.Shutdown(time.Second)

	if err := eg.Wait(); err != nil {
		t.Error("listenAndServe failed.", err)
	}
}

func TestProbe(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	defer server.Shutdown(0)
	go server.listenAndServe()
	req, err := http.NewRequest(http.MethodGet, server.listenAddr()+network.ProbePath, nil)
	if err != nil {
		t.Fatal("Error creating request:", err)
	}
	req.Header.Set("User-Agent", "kube-probe/1.15.i.wish")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal("Error roundtripping:", err)
	}
	defer resp.Body.Close()
	if got, want := resp.StatusCode, http.StatusOK; got != want {
		t.Errorf("StatusCode: %v, want: %v", got, want)
	}
}

func TestStatsReceived(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	defer server.Shutdown(0)
	go server.listenAndServe()

	statSink := dialOK(t, server.listenAddr())

	// protobuf
	assertReceivedProto(t, both, statSink, statsCh)

	// json encoding
	assertReceivedJSON(t, msg1, statSink, statsCh)
	assertReceivedJSON(t, msg2, statSink, statsCh)

	closeSink(t, statSink)
}

func TestServerShutdown(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)

	go server.listenAndServe()

	listenAddr := server.listenAddr()
	statSink := dialOK(t, listenAddr)

	assertReceivedProto(t, both, statSink, statsCh)

	server.Shutdown(time.Second)
	// We own the channel.
	close(statsCh)

	// Send a statistic to the server
	if err := sendProto(statSink, both); err != nil {
		t.Fatal("Expected send to succeed, got:", err)
	}

	// Check the statistic was not received
	_, ok := <-statsCh
	if ok {
		t.Fatal("Received statistic after shutdown")
	}

	// Check connection has been closed with a close control message with a "service restart" close code
	if _, _, err := statSink.NextReader(); err == nil {
		t.Fatal("Connection not closed")
	} else {
		err, ok := err.(*websocket.CloseError)
		if !ok {
			t.Fatal("CloseError not received")
		}
		if err.Code != 1012 {
			t.Fatalf("CloseError with unexpected close code %d received", err.Code)
		}
	}

	// Check that new connections are refused with some error
	if _, err := dial(listenAddr); err == nil {
		t.Fatal("Connection not refused")
	}

	closeSink(t, statSink)
}

func TestServerDoesNotLeakGoroutines(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServer(statsCh)
	defer server.Shutdown(0)

	go server.listenAndServe()

	originalGoroutines := runtime.NumGoroutine()

	listenAddr := server.listenAddr()
	statSink := dialOK(t, listenAddr)

	assertReceivedProto(t, both, statSink, statsCh)

	closeSink(t, statSink)

	// Check the number of goroutines eventually reduces to the number there were before the connection was created
	for i := 1000; i >= 0; i-- {
		currentGoRoutines := runtime.NumGoroutine()
		if currentGoRoutines <= originalGoroutines {
			break
		}
		time.Sleep(5 * time.Millisecond)
		if i == 0 {
			t.Fatalf("Current number of goroutines %d is not equal to the original number %d", currentGoRoutines, originalGoroutines)
		}
	}
}

func TestServerNotBucketHost(t *testing.T) {
	statsCh := make(chan metrics.StatMessage)
	server := newTestServerWithOwnerFunc(statsCh, alwaysFalse)
	defer server.Shutdown(0)

	go server.listenAndServe()

	// Verify the server behaves normally when the host is not a bucket host.
	statSink := dialOK(t, server.listenAddr())
	assertReceivedProto(t, both, statSink, statsCh)
	closeSink(t, statSink)
}

func TestServerOwnerForBucketHost(t *testing.T) {
	// Override the function to mock a bucket host.
	isBucketHost = alwaysTrue

	statsCh := make(chan metrics.StatMessage)
	server := newTestServerWithOwnerFunc(statsCh, alwaysTrue)
	defer server.Shutdown(0)

	go server.listenAndServe()

	// Verify the server behaves normally when the host is a bucket host
	// and the pod is the owner of the bucket.
	statSink := dialOK(t, server.listenAddr())
	assertReceivedProto(t, both, statSink, statsCh)
	closeSink(t, statSink)
}

func TestServerNotOwnerForBucketHost(t *testing.T) {
	// Override the function to mock a bucket host.
	isBucketHost = alwaysTrue

	statsCh := make(chan metrics.StatMessage)
	server := newTestServerWithOwnerFunc(statsCh, alwaysFalse)
	defer server.Shutdown(0)

	go server.listenAndServe()

	_, err := dial(server.listenAddr())
	if err == nil {
		t.Error("Want error from Dial but got none")
	}
}

func BenchmarkStatServer(b *testing.B) {
	statsCh := make(chan metrics.StatMessage, 100)
	server := newTestServer(statsCh)
	go server.listenAndServe()
	defer server.Shutdown(time.Second)

	statSink, err := dial(server.listenAddr())
	if err != nil {
		b.Fatal("Dial failed:", err)
	}

	// The activator sends a bunch of metrics at once usually. This simulates cases with
	// the respective number of active revisions, sending via the activator.
	for _, size := range []int{1, 2, 5, 10, 20, 50, 100} {
		msgs := make([]metrics.StatMessage, 0, size)
		for i := 0; i < size; i++ {
			msgs = append(msgs, msg1)
		}

		b.Run(fmt.Sprintf("json-encoding-%d-msgs", len(msgs)), func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				for _, msg := range msgs {
					if err := sendJSON(statSink, msg); err != nil {
						b.Fatal("Expected send to succeed, but got:", err)
					}
				}

				for range msgs {
					<-statsCh
				}
			}
		})

		b.Run(fmt.Sprintf("proto-encoding-%d-msgs", len(msgs)), func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				if err := sendProto(statSink, msgs); err != nil {
					b.Fatal("Expected send to succeed, but got:", err)
				}

				for range msgs {
					<-statsCh
				}
			}
		})
	}
}

func assertReceivedJSON(t *testing.T, sm metrics.StatMessage, statSink *websocket.Conn, statsCh <-chan metrics.StatMessage) {
	t.Helper()

	if err := sendJSON(statSink, sm); err != nil {
		t.Fatal("Expected send to succeed, got:", err)
	}

	recv := <-statsCh
	if !cmp.Equal(sm, recv) {
		t.Fatalf("StatMessage mismatch: diff (-got, +want) %s", cmp.Diff(recv, sm))
	}
}

func assertReceivedProto(t *testing.T, sms []metrics.StatMessage, statSink *websocket.Conn, statsCh <-chan metrics.StatMessage) {
	t.Helper()

	if err := sendProto(statSink, sms); err != nil {
		t.Fatal("Expected send to succeed, got:", err)
	}

	got := make([]metrics.StatMessage, 0, len(sms))
	for range sms {
		got = append(got, <-statsCh)
	}
	if !cmp.Equal(sms, got) {
		t.Fatalf("StatMessage mismatch: diff (-got, +want) %s", cmp.Diff(got, sms))
	}
}

func dialOK(t *testing.T, serverURL string) *websocket.Conn {
	t.Helper()

	statSink, err := dial(serverURL)
	if err != nil {
		t.Fatal("Dial failed:", err)
	}
	return statSink
}

func dial(serverURL string) (*websocket.Conn, error) {
	u, err := url.Parse(serverURL)
	if err != nil {
		return nil, err
	}
	u.Scheme = "ws"

	dialer := &websocket.Dialer{
		HandshakeTimeout: time.Second,
	}
	statSink, _, err := dialer.Dial(u.String(), nil)
	return statSink, err
}

func sendJSON(statSink *websocket.Conn, sm metrics.StatMessage) error {
	var b bytes.Buffer
	enc := json.NewEncoder(&b)
	if err := enc.Encode(sm); err != nil {
		return fmt.Errorf("failed to encode StatMessage: %w", err)
	}

	if err := statSink.WriteMessage(websocket.TextMessage, b.Bytes()); err != nil {
		return fmt.Errorf("failed to write to stat sink: %w", err)
	}
	return nil
}

func sendProto(statSink *websocket.Conn, sms []metrics.StatMessage) error {
	wsms := metrics.ToWireStatMessages(sms)
	msg, err := wsms.Marshal()
	if err != nil {
		return fmt.Errorf("failed to marshal StatMessage: %w", err)
	}

	if err := statSink.WriteMessage(websocket.BinaryMessage, msg); err != nil {
		return fmt.Errorf("failed to write to stat sink: %w", err)
	}

	return nil
}

func closeSink(t *testing.T, statSink *websocket.Conn) {
	t.Helper()

	if err := statSink.Close(); err != nil {
		t.Fatal("Failed to close", err)
	}
}

const testAddress = "127.0.0.1:0"

type testServer struct {
	*Server
	listenAddrCh chan string
}

func newTestServer(statsCh chan<- metrics.StatMessage) *testServer {
	return newTestServerWithOwnerFunc(statsCh, nil)
}

func newTestServerWithOwnerFunc(statsCh chan<- metrics.StatMessage, f func(bkt string) bool) *testServer {
	return &testServer{
		Server:       New(testAddress, statsCh, zap.NewNop().Sugar(), f),
		listenAddrCh: make(chan string, 1),
	}
}

// listenAddr returns the address on which the server is listening. Blocks until listenAndServe is called.
func (s *testServer) listenAddr() string {
	return <-s.listenAddrCh
}

func (s *testServer) listenAndServe() error {
	listener, err := s.listen()
	if err != nil {
		return err
	}
	return s.serve(&testListener{listener, s.listenAddrCh})
}

type testListener struct {
	net.Listener
	listenAddr chan string
}

func (t *testListener) Accept() (net.Conn, error) {
	t.listenAddr <- "http://" + t.Listener.Addr().String()
	return t.Listener.Accept()
}

func alwaysTrue(_ string) bool {
	return true
}

func alwaysFalse(_ string) bool {
	return false
}
<|endoftext|>


<commit_before>package dependency

import (
	"fmt"
	"log"
	"sync"
	"time"
)

// Secret is a vault secret.
type Secret struct {
	LeaseID       string
	LeaseDuration int
	Renewable     bool

	// Data is the actual contents of the secret. The format of the data
	// is arbitrary and up to the secret backend.
	Data map[string]interface{}
}

// VaultSecret is the dependency to Vault for a secret
type VaultSecret struct {
	sync.Mutex

	Path   string
	secret *Secret

	stopped bool
	stopCh  chan struct{}
}

// Fetch queries the Vault API
func (d *VaultSecret) Fetch(clients *ClientSet, opts *QueryOptions) (interface{}, *ResponseMetadata, error) {
	d.Lock()
	if d.stopped {
		defer d.Unlock()
		return nil, nil, ErrStopped
	}
	d.Unlock()

	if opts == nil {
		opts = &QueryOptions{}
	}

	log.Printf("[DEBUG] (%s) querying vault with %+v", d.Display(), opts)

	// If this is not the first query and we have a lease duration, sleep until we
	// try to renew.
	if opts.WaitIndex != 0 && d.secret != nil && d.secret.LeaseDuration != 0 {
		duration := time.Duration(d.secret.LeaseDuration/2.0) * time.Second
		log.Printf("[DEBUG] (%s) pretending to long-poll for %q",
			d.Display(), duration)
		select {
		case <-d.stopCh:
			log.Printf("[DEBUG] (%s) received interrupt", d.Display())
			return nil, nil, ErrStopped
		case <-time.After(duration):
		}
	}

	// Grab the vault client
	vault, err := clients.Vault()
	if err != nil {
		return nil, nil, fmt.Errorf("vault secret: %s", err)
	}

	// Attempt to renew the secret. If we do not have a secret or if that secret
	// is not renewable, we will attempt a (re-)read later.
	if d.secret != nil && d.secret.LeaseID != "" && d.secret.Renewable {
		renewal, err := vault.Sys().Renew(d.secret.LeaseID, 0)
		if err == nil {
			log.Printf("[DEBUG] (%s) successfully renewed", d.Display())

			secret := &Secret{
				LeaseID:       renewal.LeaseID,
				LeaseDuration: leaseDurationOrDefault(renewal.LeaseDuration),
				Renewable:     renewal.Renewable,
				Data:          d.secret.Data,
			}

			d.Lock()
			d.secret = secret
			d.Unlock()

			return respWithMetadata(secret)
		}

		// The renewal failed for some reason.
		log.Printf("[WARN] (%s) failed to renew, re-reading: %s", d.Display(), err)
	}

	// If we got this far, we either didn't have a secret to renew, the secret was
	// not renewable, or the renewal failed, so attempt a fresh read.
	vaultSecret, err := vault.Logical().Read(d.Path)
	if err != nil {
		return nil, nil, fmt.Errorf("error reading from vault: %s", err)
	}

	// The secret could be nil (maybe it does not exist yet). This is not an error
	// to Vault, but it is an error to Consul Template, so return an error
	// instead.
	if vaultSecret == nil {
		return nil, nil, fmt.Errorf("no secret exists at path %q", d.Display())
	}

	// Create our cloned secret
	secret := &Secret{
		LeaseID:       vaultSecret.LeaseID,
		LeaseDuration: leaseDurationOrDefault(vaultSecret.LeaseDuration),
		Renewable:     vaultSecret.Renewable,
		Data:          vaultSecret.Data,
	}

	d.Lock()
	d.secret = secret
	d.Unlock()

	log.Printf("[DEBUG] (%s) vault returned the secret", d.Display())

	return respWithMetadata(secret)
}

// CanShare returns if this dependency is shareable.
func (d *VaultSecret) CanShare() bool {
	return false
}

// HashCode returns the hash code for this dependency.
func (d *VaultSecret) HashCode() string {
	return fmt.Sprintf("VaultSecret|%s", d.Path)
}

// Display returns a string that should be displayed to the user in output (for
// example).
func (d *VaultSecret) Display() string {
	return fmt.Sprintf(`"secret(%s)"`, d.Path)
}

// Stop halts the given dependency's fetch.
func (d *VaultSecret) Stop() {
	d.Lock()
	defer d.Unlock()

	if !d.stopped {
		close(d.stopCh)
		d.stopped = true
	}
}

// ParseVaultSecret creates a new datacenter dependency.
func ParseVaultSecret(s string) (*VaultSecret, error) {
	vs := &VaultSecret{
		Path:   s,
		stopCh: make(chan struct{}),
	}
	return vs, nil
}

func leaseDurationOrDefault(d int) int {
	if d == 0 {
		return 5 * 60
	}
	return d
}
<commit_msg>Use original lease duration<commit_after>package dependency

import (
	"fmt"
	"log"
	"sync"
	"time"
)

// Secret is a vault secret.
type Secret struct {
	LeaseID       string
	LeaseDuration int
	Renewable     bool

	// Data is the actual contents of the secret. The format of the data
	// is arbitrary and up to the secret backend.
	Data map[string]interface{}
}

// VaultSecret is the dependency to Vault for a secret
type VaultSecret struct {
	sync.Mutex

	Path   string
	secret *Secret

	stopped bool
	stopCh  chan struct{}
}

// Fetch queries the Vault API
func (d *VaultSecret) Fetch(clients *ClientSet, opts *QueryOptions) (interface{}, *ResponseMetadata, error) {
	d.Lock()
	if d.stopped {
		defer d.Unlock()
		return nil, nil, ErrStopped
	}
	d.Unlock()

	if opts == nil {
		opts = &QueryOptions{}
	}

	log.Printf("[DEBUG] (%s) querying vault with %+v", d.Display(), opts)

	// If this is not the first query and we have a lease duration, sleep until we
	// try to renew.
	if opts.WaitIndex != 0 && d.secret != nil && d.secret.LeaseDuration != 0 {
		duration := time.Duration(d.secret.LeaseDuration/2.0) * time.Second
		log.Printf("[DEBUG] (%s) pretending to long-poll for %q",
			d.Display(), duration)
		select {
		case <-d.stopCh:
			log.Printf("[DEBUG] (%s) received interrupt", d.Display())
			return nil, nil, ErrStopped
		case <-time.After(duration):
		}
	}

	// Grab the vault client
	vault, err := clients.Vault()
	if err != nil {
		return nil, nil, fmt.Errorf("vault secret: %s", err)
	}

	// Attempt to renew the secret. If we do not have a secret or if that secret
	// is not renewable, we will attempt a (re-)read later.
	if d.secret != nil && d.secret.LeaseID != "" && d.secret.Renewable {
		renewal, err := vault.Sys().Renew(d.secret.LeaseID, 0)
		if err == nil {
			log.Printf("[DEBUG] (%s) successfully renewed", d.Display())

			log.Printf("[DEBUG] (%s) %#v", d.Display(), renewal)

			secret := &Secret{
				LeaseID:       renewal.LeaseID,
				LeaseDuration: d.secret.LeaseDuration,
				Renewable:     renewal.Renewable,
				Data:          d.secret.Data,
			}

			d.Lock()
			d.secret = secret
			d.Unlock()

			return respWithMetadata(secret)
		}

		// The renewal failed for some reason.
		log.Printf("[WARN] (%s) failed to renew, re-reading: %s", d.Display(), err)
	}

	// If we got this far, we either didn't have a secret to renew, the secret was
	// not renewable, or the renewal failed, so attempt a fresh read.
	vaultSecret, err := vault.Logical().Read(d.Path)
	if err != nil {
		return nil, nil, fmt.Errorf("error reading from vault: %s", err)
	}

	// The secret could be nil (maybe it does not exist yet). This is not an error
	// to Vault, but it is an error to Consul Template, so return an error
	// instead.
	if vaultSecret == nil {
		return nil, nil, fmt.Errorf("no secret exists at path %q", d.Display())
	}

	// Create our cloned secret
	secret := &Secret{
		LeaseID:       vaultSecret.LeaseID,
		LeaseDuration: leaseDurationOrDefault(vaultSecret.LeaseDuration),
		Renewable:     vaultSecret.Renewable,
		Data:          vaultSecret.Data,
	}

	d.Lock()
	d.secret = secret
	d.Unlock()

	log.Printf("[DEBUG] (%s) vault returned the secret", d.Display())

	return respWithMetadata(secret)
}

// CanShare returns if this dependency is shareable.
func (d *VaultSecret) CanShare() bool {
	return false
}

// HashCode returns the hash code for this dependency.
func (d *VaultSecret) HashCode() string {
	return fmt.Sprintf("VaultSecret|%s", d.Path)
}

// Display returns a string that should be displayed to the user in output (for
// example).
func (d *VaultSecret) Display() string {
	return fmt.Sprintf(`"secret(%s)"`, d.Path)
}

// Stop halts the given dependency's fetch.
func (d *VaultSecret) Stop() {
	d.Lock()
	defer d.Unlock()

	if !d.stopped {
		close(d.stopCh)
		d.stopped = true
	}
}

// ParseVaultSecret creates a new datacenter dependency.
func ParseVaultSecret(s string) (*VaultSecret, error) {
	vs := &VaultSecret{
		Path:   s,
		stopCh: make(chan struct{}),
	}
	return vs, nil
}

func leaseDurationOrDefault(d int) int {
	if d == 0 {
		return 5 * 60
	}
	return d
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package genericclioptions

import (
	"net/http"
	"strings"

	"github.com/google/uuid"
	"github.com/spf13/cobra"
)

const (
	kubectlCommandHeader = "Kubectl-Command"
	kubectlSessionHeader = "Kubectl-Session"
)

// CommandHeaderRoundTripper adds a layer around the standard
// round tripper to add Request headers before delegation. Implements
// the go standard library "http.RoundTripper" interface.
type CommandHeaderRoundTripper struct {
	Delegate http.RoundTripper
	Headers  map[string]string
}

// CommandHeaderRoundTripper adds Request headers before delegating to standard
// round tripper. These headers are kubectl command headers which
// detail the kubectl command. See SIG CLI KEP 859:
//   https://github.com/kubernetes/enhancements/tree/master/keps/sig-cli/859-kubectl-headers
func (c *CommandHeaderRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	for header, value := range c.Headers {
		req.Header.Set(header, value)
	}
	return c.Delegate.RoundTrip(req)
}

// ParseCommandHeaders fills in a map of custom headers into the CommandHeaderRoundTripper. These
// headers are then filled into each request. For details on the custom headers see:
//   https://github.com/kubernetes/enhancements/tree/master/keps/sig-cli/859-kubectl-headers
// Each call overwrites the previously parsed command headers (not additive).
// TODO(seans3): Parse/add flags removing PII from flag values.
func (c *CommandHeaderRoundTripper) ParseCommandHeaders(cmd *cobra.Command, args []string) {
	if cmd == nil {
		return
	}
	// Overwrites previously parsed command headers (headers not additive).
	c.Headers = map[string]string{}
	// Session identifier to aggregate multiple Requests from single kubectl command.
	uid := uuid.New().String()
	c.Headers[kubectlSessionHeader] = uid
	// Iterate up the hierarchy of commands from the leaf command to create
	// the full command string. Example: kubectl create secret generic
	cmdStrs := []string{}
	for cmd.HasParent() {
		parent := cmd.Parent()
		currName := strings.TrimSpace(cmd.Name())
		cmdStrs = append([]string{currName}, cmdStrs...)
		cmd = parent
	}
	currName := strings.TrimSpace(cmd.Name())
	cmdStrs = append([]string{currName}, cmdStrs...)
	if len(cmdStrs) > 0 {
		c.Headers[kubectlCommandHeader] = strings.Join(cmdStrs, " ")
	}
}

// CancelRequest is propagated to the Delegate RoundTripper within
// if the wrapped RoundTripper implements this function.
func (c *CommandHeaderRoundTripper) CancelRequest(req *http.Request) {
	type canceler interface {
		CancelRequest(*http.Request)
	}
	// If possible, call "CancelRequest" on the wrapped Delegate RoundTripper.
	if cr, ok := c.Delegate.(canceler); ok {
		cr.CancelRequest(req)
	}
}
<commit_msg>Generate and format files<commit_after>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package genericclioptions

import (
	"net/http"
	"strings"

	"github.com/google/uuid"
	"github.com/spf13/cobra"
)

const (
	kubectlCommandHeader = "Kubectl-Command"
	kubectlSessionHeader = "Kubectl-Session"
)

// CommandHeaderRoundTripper adds a layer around the standard
// round tripper to add Request headers before delegation. Implements
// the go standard library "http.RoundTripper" interface.
type CommandHeaderRoundTripper struct {
	Delegate http.RoundTripper
	Headers  map[string]string
}

// CommandHeaderRoundTripper adds Request headers before delegating to standard
// round tripper. These headers are kubectl command headers which
// detail the kubectl command. See SIG CLI KEP 859:
//
//	https://github.com/kubernetes/enhancements/tree/master/keps/sig-cli/859-kubectl-headers
func (c *CommandHeaderRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	for header, value := range c.Headers {
		req.Header.Set(header, value)
	}
	return c.Delegate.RoundTrip(req)
}

// ParseCommandHeaders fills in a map of custom headers into the CommandHeaderRoundTripper. These
// headers are then filled into each request. For details on the custom headers see:
//
//	https://github.com/kubernetes/enhancements/tree/master/keps/sig-cli/859-kubectl-headers
//
// Each call overwrites the previously parsed command headers (not additive).
// TODO(seans3): Parse/add flags removing PII from flag values.
func (c *CommandHeaderRoundTripper) ParseCommandHeaders(cmd *cobra.Command, args []string) {
	if cmd == nil {
		return
	}
	// Overwrites previously parsed command headers (headers not additive).
	c.Headers = map[string]string{}
	// Session identifier to aggregate multiple Requests from single kubectl command.
	uid := uuid.New().String()
	c.Headers[kubectlSessionHeader] = uid
	// Iterate up the hierarchy of commands from the leaf command to create
	// the full command string. Example: kubectl create secret generic
	cmdStrs := []string{}
	for cmd.HasParent() {
		parent := cmd.Parent()
		currName := strings.TrimSpace(cmd.Name())
		cmdStrs = append([]string{currName}, cmdStrs...)
		cmd = parent
	}
	currName := strings.TrimSpace(cmd.Name())
	cmdStrs = append([]string{currName}, cmdStrs...)
	if len(cmdStrs) > 0 {
		c.Headers[kubectlCommandHeader] = strings.Join(cmdStrs, " ")
	}
}

// CancelRequest is propagated to the Delegate RoundTripper within
// if the wrapped RoundTripper implements this function.
func (c *CommandHeaderRoundTripper) CancelRequest(req *http.Request) {
	type canceler interface {
		CancelRequest(*http.Request)
	}
	// If possible, call "CancelRequest" on the wrapped Delegate RoundTripper.
	if cr, ok := c.Delegate.(canceler); ok {
		cr.CancelRequest(req)
	}
}
<|endoftext|>