<commit_before>/*
   Copyright The containerd 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 integration

import (
	"fmt"
	goruntime "runtime"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
)

// Test to verify for a container ID
func TestContainerStats(t *testing.T) {
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "sandbox1", "stats")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run container in a pod")
	containerConfig := ContainerConfig(
		"container1",
		pauseImage,
		WithTestLabels(),
		WithTestAnnotations(),
	)
	cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
	require.NoError(t, err)
	defer func() {
		assert.NoError(t, runtimeService.RemoveContainer(cn))
	}()
	require.NoError(t, runtimeService.StartContainer(cn))
	defer func() {
		assert.NoError(t, runtimeService.StopContainer(cn, 10))
	}()

	t.Logf("Fetch stats for container")
	var s *runtime.ContainerStats
	require.NoError(t, Eventually(func() (bool, error) {
		s, err = runtimeService.ContainerStats(cn)
		if err != nil {
			return false, err
		}
		if s.GetWritableLayer().GetUsedBytes().GetValue() != 0 {
			return true, nil
		}
		return false, nil
	}, time.Second, 30*time.Second))

	t.Logf("Verify stats received for container %q", cn)
	testStats(t, s, containerConfig)
}

// Test to verify if the consumed stats are correct.
func TestContainerConsumedStats(t *testing.T) {
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "sandbox1", "stats")

	testImage := GetImage(ResourceConsumer)
	EnsureImageExists(t, testImage)

	t.Logf("Create a container config and run container in a pod")
	containerConfig := ContainerConfig(
		"container1",
		testImage,
		WithTestLabels(),
		WithTestAnnotations(),
	)
	cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
	require.NoError(t, err)
	defer func() {
		assert.NoError(t, runtimeService.RemoveContainer(cn))
	}()
	require.NoError(t, runtimeService.StartContainer(cn))
	defer func() {
		assert.NoError(t, runtimeService.StopContainer(cn, 10))
	}()

	t.Logf("Fetch initial stats for container")
	var s *runtime.ContainerStats
	require.NoError(t, Eventually(func() (bool, error) {
		s, err = runtimeService.ContainerStats(cn)
		if err != nil {
			return false, err
		}
		if s.GetMemory().GetWorkingSetBytes().GetValue() > 0 {
			return true, nil
		}
		return false, nil
	}, time.Second, 30*time.Second))

	initialMemory := s.GetMemory().GetWorkingSetBytes().GetValue()
	t.Logf("Initial container memory consumption is %f MB. Consume 100 MB and expect the reported stats to increase accordingly", float64(initialMemory)/(1024*1024))

	// consume 100 MB memory for 30 seconds.
	var command []string
	if goruntime.GOOS == "windows" {
		// -d: Leak and touch memory in specified MBs
		// -c: Count of number of objects to allocate
		command = []string{"testlimit.exe", "-accepteula", "-d", "25", "-c", "4"}
	} else {
		command = []string{"stress", "-m", "1", "--vm-bytes", "100M", "--vm-hang", "0", "-t", "30"}
	}

	go func() {
		_, _, err = runtimeService.ExecSync(cn, command, 30*time.Second)
	}()

	require.NoError(t, Eventually(func() (bool, error) {
		s, err = runtimeService.ContainerStats(cn)
		if err != nil {
			return false, err
		}
		if s.GetMemory().GetWorkingSetBytes().GetValue() > initialMemory+100*1024*1024 {
			return true, nil
		}
		return false, nil
	}, time.Second, 30*time.Second))
}

// Test to verify filtering without any filter
func TestContainerListStats(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		require.NoError(t, err)
		containerConfigMap[cn] = containerConfig
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}

	t.Logf("Fetch all container stats")
	require.NoError(t, Eventually(func() (bool, error) {
		stats, err = runtimeService.ListContainerStats(&runtime.ContainerStatsFilter{})
		if err != nil {
			return false, err
		}
		for _, s := range stats {
			if s.GetWritableLayer().GetUsedBytes().GetValue() == 0 {
				return false, nil
			}
		}
		return true, nil
	}, time.Second, 30*time.Second))

	t.Logf("Verify all container stats")
	for _, s := range stats {
		testStats(t, s, containerConfigMap[s.GetAttributes().GetId()])
	}
}

// Test to verify filtering given a specific container ID
// TODO Convert the filter tests into table driven tests and unit tests
func TestContainerListStatsWithIdFilter(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		containerConfigMap[cn] = containerConfig
		require.NoError(t, err)
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}

	t.Logf("Fetch container stats for each container with Filter")
	for id := range containerConfigMap {
		require.NoError(t, Eventually(func() (bool, error) {
			stats, err = runtimeService.ListContainerStats(
				&runtime.ContainerStatsFilter{Id: id})
			if err != nil {
				return false, err
			}
			if len(stats) != 1 {
				return false, errors.New("unexpected stats length")
			}
			if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
				return true, nil
			}
			return false, nil
		}, time.Second, 30*time.Second))

		t.Logf("Verify container stats for %s", id)
		for _, s := range stats {
			require.Equal(t, s.GetAttributes().GetId(), id)
			testStats(t, s, containerConfigMap[id])
		}
	}
}

// Test to verify filtering given a specific Sandbox ID. Stats for
// all the containers in a pod should be returned
func TestContainerListStatsWithSandboxIdFilter(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		containerConfigMap[cn] = containerConfig
		require.NoError(t, err)
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}

	t.Logf("Fetch container stats for each container with Filter")
	require.NoError(t, Eventually(func() (bool, error) {
		stats, err = runtimeService.ListContainerStats(
			&runtime.ContainerStatsFilter{PodSandboxId: sb})
		if err != nil {
			return false, err
		}
		if len(stats) != 3 {
			return false, errors.New("unexpected stats length")
		}
		if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
			return true, nil
		}
		return false, nil
	}, time.Second, 45*time.Second))
	// TODO(claudiub): Reduce the timer above to 30 seconds once Windows flakiness has been addressed.
	t.Logf("Verify container stats for sandbox %q", sb)
	for _, s := range stats {
		testStats(t, s, containerConfigMap[s.GetAttributes().GetId()])
	}
}

// Test to verify filtering given a specific container ID and
// sandbox ID
func TestContainerListStatsWithIdSandboxIdFilter(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		containerConfigMap[cn] = containerConfig
		require.NoError(t, err)
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}
	t.Logf("Fetch container stats for sandbox ID and container ID filter")
	for id, config := range containerConfigMap {
		require.NoError(t, Eventually(func() (bool, error) {
			stats, err = runtimeService.ListContainerStats(
				&runtime.ContainerStatsFilter{Id: id, PodSandboxId: sb})
			if err != nil {
				return false, err
			}
			if len(stats) != 1 {
				return false, errors.New("unexpected stats length")
			}
			if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
				return true, nil
			}
			return false, nil
		}, time.Second, 30*time.Second))
		t.Logf("Verify container stats for sandbox %q and container %q filter", sb, id)
		for _, s := range stats {
			testStats(t, s, config)
		}
	}

	t.Logf("Fetch container stats for sandbox truncID and container truncID filter ")
	for id, config := range containerConfigMap {
		require.NoError(t, Eventually(func() (bool, error) {
			stats, err = runtimeService.ListContainerStats(
				&runtime.ContainerStatsFilter{Id: id[:3], PodSandboxId: sb[:3]})
			if err != nil {
				return false, err
			}
			if len(stats) != 1 {
				return false, errors.New("unexpected stats length")
			}
			if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
				return true, nil
			}
			return false, nil
		}, time.Second, 30*time.Second))
		t.Logf("Verify container stats for sandbox %q and container %q filter", sb, id)
		for _, s := range stats {
			testStats(t, s, config)
		}
	}
}

// TODO make this as options to use for dead container tests
func testStats(t *testing.T,
	s *runtime.ContainerStats,
	config *runtime.ContainerConfig,
) {
	require.NotEmpty(t, s.GetAttributes().GetId())
	require.NotEmpty(t, s.GetAttributes().GetMetadata())
	require.NotEmpty(t, s.GetAttributes().GetAnnotations())
	require.Equal(t, s.GetAttributes().GetLabels(), config.Labels)
	require.Equal(t, s.GetAttributes().GetAnnotations(), config.Annotations)
	require.Equal(t, s.GetAttributes().GetMetadata().Name, config.Metadata.Name)
	require.NotEmpty(t, s.GetAttributes().GetLabels())
	require.NotEmpty(t, s.GetCpu().GetTimestamp())
	require.NotEmpty(t, s.GetCpu().GetUsageCoreNanoSeconds().GetValue())
	require.NotEmpty(t, s.GetMemory().GetTimestamp())
	require.NotEmpty(t, s.GetMemory().GetWorkingSetBytes().GetValue())
	require.NotEmpty(t, s.GetWritableLayer().GetTimestamp())
	require.NotEmpty(t, s.GetWritableLayer().GetFsId().GetMountpoint())
	require.NotEmpty(t, s.GetWritableLayer().GetUsedBytes().GetValue())

	// Windows does not collect inodes stats.
	if goruntime.GOOS != "windows" {
		require.NotEmpty(t, s.GetWritableLayer().GetInodesUsed().GetValue())
	}
}
<commit_msg>integration: add stats result in error message<commit_after>/*
   Copyright The containerd 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 integration

import (
	"fmt"
	goruntime "runtime"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
)

// Test to verify for a container ID
func TestContainerStats(t *testing.T) {
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "sandbox1", "stats")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run container in a pod")
	containerConfig := ContainerConfig(
		"container1",
		pauseImage,
		WithTestLabels(),
		WithTestAnnotations(),
	)
	cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
	require.NoError(t, err)
	defer func() {
		assert.NoError(t, runtimeService.RemoveContainer(cn))
	}()
	require.NoError(t, runtimeService.StartContainer(cn))
	defer func() {
		assert.NoError(t, runtimeService.StopContainer(cn, 10))
	}()

	t.Logf("Fetch stats for container")
	var s *runtime.ContainerStats
	require.NoError(t, Eventually(func() (bool, error) {
		s, err = runtimeService.ContainerStats(cn)
		if err != nil {
			return false, err
		}
		if s.GetWritableLayer().GetUsedBytes().GetValue() != 0 {
			return true, nil
		}
		return false, nil
	}, time.Second, 30*time.Second))

	t.Logf("Verify stats received for container %q", cn)
	testStats(t, s, containerConfig)
}

// Test to verify if the consumed stats are correct.
func TestContainerConsumedStats(t *testing.T) {
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "sandbox1", "stats")

	testImage := GetImage(ResourceConsumer)
	EnsureImageExists(t, testImage)

	t.Logf("Create a container config and run container in a pod")
	containerConfig := ContainerConfig(
		"container1",
		testImage,
		WithTestLabels(),
		WithTestAnnotations(),
	)
	cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
	require.NoError(t, err)
	defer func() {
		assert.NoError(t, runtimeService.RemoveContainer(cn))
	}()
	require.NoError(t, runtimeService.StartContainer(cn))
	defer func() {
		assert.NoError(t, runtimeService.StopContainer(cn, 10))
	}()

	t.Logf("Fetch initial stats for container")
	var s *runtime.ContainerStats
	require.NoError(t, Eventually(func() (bool, error) {
		s, err = runtimeService.ContainerStats(cn)
		if err != nil {
			return false, err
		}
		if s.GetMemory().GetWorkingSetBytes().GetValue() > 0 {
			return true, nil
		}
		return false, nil
	}, time.Second, 30*time.Second))

	initialMemory := s.GetMemory().GetWorkingSetBytes().GetValue()
	t.Logf("Initial container memory consumption is %f MB. Consume 100 MB and expect the reported stats to increase accordingly", float64(initialMemory)/(1024*1024))

	// consume 100 MB memory for 30 seconds.
	var command []string
	if goruntime.GOOS == "windows" {
		// -d: Leak and touch memory in specified MBs
		// -c: Count of number of objects to allocate
		command = []string{"testlimit.exe", "-accepteula", "-d", "25", "-c", "4"}
	} else {
		command = []string{"stress", "-m", "1", "--vm-bytes", "100M", "--vm-hang", "0", "-t", "30"}
	}

	go func() {
		_, _, err = runtimeService.ExecSync(cn, command, 30*time.Second)
	}()

	require.NoError(t, Eventually(func() (bool, error) {
		s, err = runtimeService.ContainerStats(cn)
		if err != nil {
			return false, err
		}
		if s.GetMemory().GetWorkingSetBytes().GetValue() > initialMemory+100*1024*1024 {
			return true, nil
		}
		return false, nil
	}, time.Second, 30*time.Second))
}

// Test to verify filtering without any filter
func TestContainerListStats(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		require.NoError(t, err)
		containerConfigMap[cn] = containerConfig
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}

	t.Logf("Fetch all container stats")
	require.NoError(t, Eventually(func() (bool, error) {
		stats, err = runtimeService.ListContainerStats(&runtime.ContainerStatsFilter{})
		if err != nil {
			return false, err
		}
		for _, s := range stats {
			if s.GetWritableLayer().GetUsedBytes().GetValue() == 0 {
				return false, nil
			}
		}
		return true, nil
	}, time.Second, 30*time.Second))

	t.Logf("Verify all container stats")
	for _, s := range stats {
		testStats(t, s, containerConfigMap[s.GetAttributes().GetId()])
	}
}

// Test to verify filtering given a specific container ID
// TODO Convert the filter tests into table driven tests and unit tests
func TestContainerListStatsWithIdFilter(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		containerConfigMap[cn] = containerConfig
		require.NoError(t, err)
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}

	t.Logf("Fetch container stats for each container with Filter")
	for id := range containerConfigMap {
		require.NoError(t, Eventually(func() (bool, error) {
			stats, err = runtimeService.ListContainerStats(
				&runtime.ContainerStatsFilter{Id: id})
			if err != nil {
				return false, err
			}
			if len(stats) != 1 {
				return false, errors.New("unexpected stats length")
			}
			if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
				return true, nil
			}
			return false, nil
		}, time.Second, 30*time.Second))

		t.Logf("Verify container stats for %s", id)
		for _, s := range stats {
			require.Equal(t, s.GetAttributes().GetId(), id)
			testStats(t, s, containerConfigMap[id])
		}
	}
}

// Test to verify filtering given a specific Sandbox ID. Stats for
// all the containers in a pod should be returned
func TestContainerListStatsWithSandboxIdFilter(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create a container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		containerConfigMap[cn] = containerConfig
		require.NoError(t, err)
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}

	t.Logf("Fetch container stats for each container with Filter")
	require.NoError(t, Eventually(func() (bool, error) {
		stats, err = runtimeService.ListContainerStats(
			&runtime.ContainerStatsFilter{PodSandboxId: sb})
		if err != nil {
			return false, err
		}
		if len(stats) != 3 {
			return false, errors.New("unexpected stats length")
		}
		if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
			return true, nil
		}
		return false, nil
	}, time.Second, 45*time.Second))
	// TODO(claudiub): Reduce the timer above to 30 seconds once Windows flakiness has been addressed.
	t.Logf("Verify container stats for sandbox %q", sb)
	for _, s := range stats {
		testStats(t, s, containerConfigMap[s.GetAttributes().GetId()])
	}
}

// Test to verify filtering given a specific container ID and
// sandbox ID
func TestContainerListStatsWithIdSandboxIdFilter(t *testing.T) {
	var (
		stats []*runtime.ContainerStats
		err   error
	)
	t.Logf("Create a pod config and run sandbox container")
	sb, sbConfig := PodSandboxConfigWithCleanup(t, "running-pod", "statsls")

	EnsureImageExists(t, pauseImage)

	t.Logf("Create container config and run containers in a pod")
	containerConfigMap := make(map[string]*runtime.ContainerConfig)
	for i := 0; i < 3; i++ {
		cName := fmt.Sprintf("container%d", i)
		containerConfig := ContainerConfig(
			cName,
			pauseImage,
			WithTestLabels(),
			WithTestAnnotations(),
		)
		cn, err := runtimeService.CreateContainer(sb, containerConfig, sbConfig)
		containerConfigMap[cn] = containerConfig
		require.NoError(t, err)
		defer func() {
			assert.NoError(t, runtimeService.RemoveContainer(cn))
		}()
		require.NoError(t, runtimeService.StartContainer(cn))
		defer func() {
			assert.NoError(t, runtimeService.StopContainer(cn, 10))
		}()
	}
	t.Logf("Fetch container stats for sandbox ID and container ID filter")
	for id, config := range containerConfigMap {
		require.NoError(t, Eventually(func() (bool, error) {
			stats, err = runtimeService.ListContainerStats(
				&runtime.ContainerStatsFilter{Id: id, PodSandboxId: sb})
			if err != nil {
				return false, err
			}
			if len(stats) != 1 {
				return false, errors.New("unexpected stats length")
			}
			if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
				return true, nil
			}
			return false, nil
		}, time.Second, 30*time.Second))
		t.Logf("Verify container stats for sandbox %q and container %q filter", sb, id)
		for _, s := range stats {
			testStats(t, s, config)
		}
	}

	t.Logf("Fetch container stats for sandbox truncID and container truncID filter ")
	for id, config := range containerConfigMap {
		require.NoError(t, Eventually(func() (bool, error) {
			stats, err = runtimeService.ListContainerStats(
				&runtime.ContainerStatsFilter{Id: id[:3], PodSandboxId: sb[:3]})
			if err != nil {
				return false, err
			}
			if len(stats) != 1 {
				return false, errors.Errorf("expected only one stat, but got %v", stats)
			}
			if stats[0].GetWritableLayer().GetUsedBytes().GetValue() != 0 {
				return true, nil
			}
			return false, nil
		}, time.Second, 30*time.Second))
		t.Logf("Verify container stats for sandbox %q and container %q filter", sb, id)
		for _, s := range stats {
			testStats(t, s, config)
		}
	}
}

// TODO make this as options to use for dead container tests
func testStats(t *testing.T,
	s *runtime.ContainerStats,
	config *runtime.ContainerConfig,
) {
	require.NotEmpty(t, s.GetAttributes().GetId())
	require.NotEmpty(t, s.GetAttributes().GetMetadata())
	require.NotEmpty(t, s.GetAttributes().GetAnnotations())
	require.Equal(t, s.GetAttributes().GetLabels(), config.Labels)
	require.Equal(t, s.GetAttributes().GetAnnotations(), config.Annotations)
	require.Equal(t, s.GetAttributes().GetMetadata().Name, config.Metadata.Name)
	require.NotEmpty(t, s.GetAttributes().GetLabels())
	require.NotEmpty(t, s.GetCpu().GetTimestamp())
	require.NotEmpty(t, s.GetCpu().GetUsageCoreNanoSeconds().GetValue())
	require.NotEmpty(t, s.GetMemory().GetTimestamp())
	require.NotEmpty(t, s.GetMemory().GetWorkingSetBytes().GetValue())
	require.NotEmpty(t, s.GetWritableLayer().GetTimestamp())
	require.NotEmpty(t, s.GetWritableLayer().GetFsId().GetMountpoint())
	require.NotEmpty(t, s.GetWritableLayer().GetUsedBytes().GetValue())

	// Windows does not collect inodes stats.
	if goruntime.GOOS != "windows" {
		require.NotEmpty(t, s.GetWritableLayer().GetInodesUsed().GetValue())
	}
}
<|endoftext|>


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

package ct

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"log"

	"github.com/google/certificate-transparency-go/tls"
	"github.com/google/certificate-transparency-go/x509"
)

// AllowVerificationWithNonCompliantKeys may be set to true in order to allow
// SignatureVerifier to use keys which are technically non-compliant with
// RFC6962.
var AllowVerificationWithNonCompliantKeys = false

// PublicKeyFromPEM parses a PEM formatted block and returns the public key contained within and any remaining unread bytes, or an error.
func PublicKeyFromPEM(b []byte) (crypto.PublicKey, SHA256Hash, []byte, error) {
	p, rest := pem.Decode(b)
	if p == nil {
		return nil, [sha256.Size]byte{}, rest, fmt.Errorf("no PEM block found in %s", string(b))
	}
	k, err := x509.ParsePKIXPublicKey(p.Bytes)
	return k, sha256.Sum256(p.Bytes), rest, err
}

// PublicKeyFromB64 parses a base64-encoded public key.
func PublicKeyFromB64(b64PubKey string) (crypto.PublicKey, error) {
	der, err := base64.StdEncoding.DecodeString(b64PubKey)
	if err != nil {
		return nil, fmt.Errorf("error decoding public key: %s", err)
	}
	return x509.ParsePKIXPublicKey(der)
}

// SignatureVerifier can verify signatures on SCTs and STHs
type SignatureVerifier struct {
	pubKey crypto.PublicKey
}

// NewSignatureVerifier creates a new SignatureVerifier using the passed in PublicKey.
func NewSignatureVerifier(pk crypto.PublicKey) (*SignatureVerifier, error) {
	switch pkType := pk.(type) {
	case *rsa.PublicKey:
		if pkType.N.BitLen() < 2048 {
			e := fmt.Errorf("public key is RSA with < 2048 bits (size:%d)", pkType.N.BitLen())
			if !AllowVerificationWithNonCompliantKeys {
				return nil, e
			}
			log.Printf("WARNING: %v", e)
		}
	case *ecdsa.PublicKey:
		params := *(pkType.Params())
		if params != *elliptic.P256().Params() {
			e := fmt.Errorf("public is ECDSA, but not on the P256 curve")
			if !AllowVerificationWithNonCompliantKeys {
				return nil, e
			}
			log.Printf("WARNING: %v", e)

		}
	default:
		return nil, fmt.Errorf("Unsupported public key type %v", pkType)
	}

	return &SignatureVerifier{
		pubKey: pk,
	}, nil
}

// VerifySignature verifies the given signature sig matches the data.
func (s SignatureVerifier) VerifySignature(data []byte, sig tls.DigitallySigned) error {
	return tls.VerifySignature(s.pubKey, data, sig)
}

// VerifySCTSignature verifies that the SCT's signature is valid for the given LogEntry.
func (s SignatureVerifier) VerifySCTSignature(sct SignedCertificateTimestamp, entry LogEntry) error {
	sctData, err := SerializeSCTSignatureInput(sct, entry)
	if err != nil {
		return err
	}
	return s.VerifySignature(sctData, tls.DigitallySigned(sct.Signature))
}

// VerifySTHSignature verifies that the STH's signature is valid.
func (s SignatureVerifier) VerifySTHSignature(sth SignedTreeHead) error {
	sthData, err := SerializeSTHSignatureInput(sth)
	if err != nil {
		return err
	}
	return s.VerifySignature(sthData, tls.DigitallySigned(sth.TreeHeadSignature))
}
<commit_msg>ct: export Verifier.PubKey for reuse<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ct

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"log"

	"github.com/google/certificate-transparency-go/tls"
	"github.com/google/certificate-transparency-go/x509"
)

// AllowVerificationWithNonCompliantKeys may be set to true in order to allow
// SignatureVerifier to use keys which are technically non-compliant with
// RFC6962.
var AllowVerificationWithNonCompliantKeys = false

// PublicKeyFromPEM parses a PEM formatted block and returns the public key contained within and any remaining unread bytes, or an error.
func PublicKeyFromPEM(b []byte) (crypto.PublicKey, SHA256Hash, []byte, error) {
	p, rest := pem.Decode(b)
	if p == nil {
		return nil, [sha256.Size]byte{}, rest, fmt.Errorf("no PEM block found in %s", string(b))
	}
	k, err := x509.ParsePKIXPublicKey(p.Bytes)
	return k, sha256.Sum256(p.Bytes), rest, err
}

// PublicKeyFromB64 parses a base64-encoded public key.
func PublicKeyFromB64(b64PubKey string) (crypto.PublicKey, error) {
	der, err := base64.StdEncoding.DecodeString(b64PubKey)
	if err != nil {
		return nil, fmt.Errorf("error decoding public key: %s", err)
	}
	return x509.ParsePKIXPublicKey(der)
}

// SignatureVerifier can verify signatures on SCTs and STHs
type SignatureVerifier struct {
	PubKey crypto.PublicKey
}

// NewSignatureVerifier creates a new SignatureVerifier using the passed in PublicKey.
func NewSignatureVerifier(pk crypto.PublicKey) (*SignatureVerifier, error) {
	switch pkType := pk.(type) {
	case *rsa.PublicKey:
		if pkType.N.BitLen() < 2048 {
			e := fmt.Errorf("public key is RSA with < 2048 bits (size:%d)", pkType.N.BitLen())
			if !AllowVerificationWithNonCompliantKeys {
				return nil, e
			}
			log.Printf("WARNING: %v", e)
		}
	case *ecdsa.PublicKey:
		params := *(pkType.Params())
		if params != *elliptic.P256().Params() {
			e := fmt.Errorf("public is ECDSA, but not on the P256 curve")
			if !AllowVerificationWithNonCompliantKeys {
				return nil, e
			}
			log.Printf("WARNING: %v", e)

		}
	default:
		return nil, fmt.Errorf("Unsupported public key type %v", pkType)
	}

	return &SignatureVerifier{PubKey: pk}, nil
}

// VerifySignature verifies the given signature sig matches the data.
func (s SignatureVerifier) VerifySignature(data []byte, sig tls.DigitallySigned) error {
	return tls.VerifySignature(s.PubKey, data, sig)
}

// VerifySCTSignature verifies that the SCT's signature is valid for the given LogEntry.
func (s SignatureVerifier) VerifySCTSignature(sct SignedCertificateTimestamp, entry LogEntry) error {
	sctData, err := SerializeSCTSignatureInput(sct, entry)
	if err != nil {
		return err
	}
	return s.VerifySignature(sctData, tls.DigitallySigned(sct.Signature))
}

// VerifySTHSignature verifies that the STH's signature is valid.
func (s SignatureVerifier) VerifySTHSignature(sth SignedTreeHead) error {
	sthData, err := SerializeSTHSignatureInput(sth)
	if err != nil {
		return err
	}
	return s.VerifySignature(sthData, tls.DigitallySigned(sth.TreeHeadSignature))
}
<|endoftext|>


<commit_before>package dry

// Max returns the larger of x or y.
func Max(x, y int64) int64 {
	if x > y {
		return x
	}
	return y
}

// Min returns the lower of x or y.
func Min(x, y int64) int64 {
	if x < y {
		return x
	}
	return y
}

// Ceil returns integer greater or equal to x and denominator d division.
// Works only with x >= 0 and m > 0. It returns 0 with other values.
func Ceil(x, d int64) int64 {
	if x <= 0 || d <= 0 {
		return int64(0)
	}
	return (x + d - 1) / d
}

// CeilToMultiplier returns the integer greater or equal to x and multiplier m product.
// Works only with x >= 0 and m > 0. It returns 0 with other values.
func CeilToMultiplier(x, m int64) int64 {
	return Ceil(x, m) * m
}

// FloorToMultiplier returns the integer less or equal to x and multiplier m product.
// Works only with x >= 0 and m > 0. It returns 0 with other values.
func FloorToMultiplier(x, m int64) int64 {
	if x <= 0 || m <= 0 {
		return int64(0)
	}
	return x / m * m
}

// GCD returns the absolute greatest common divisor calculated via Euclidean algorithm
func GCD(a, b int64) int64 {
	if b < 0 {
		b = -b
	}
	for b != 0 {
		t := b
		b = a % b
		a = t
	}
	return a
}

// LCM returns the absolute least common multiple of 2 integers via GDB
func LCM(a, b int64) int64 {
	if a*b < 0 {
		return -a / GCD(a, b) * b
	}
	return a / GCD(a, b) * b
}
<commit_msg>Update pkg/dry/math.go<commit_after>package dry

// Max returns the larger of x or y.
func Max(x, y int64) int64 {
	if x > y {
		return x
	}
	return y
}

// Min returns the lower of x or y.
func Min(x, y int64) int64 {
	if x < y {
		return x
	}
	return y
}

// Ceil returns integer greater or equal to x and denominator d division.
// Works only with x >= 0 and d > 0. It returns 0 with other values.
func Ceil(x, d int64) int64 {
	if x <= 0 || d <= 0 {
		return int64(0)
	}
	return (x + d - 1) / d
}

// CeilToMultiplier returns the integer greater or equal to x and multiplier m product.
// Works only with x >= 0 and m > 0. It returns 0 with other values.
func CeilToMultiplier(x, m int64) int64 {
	return Ceil(x, m) * m
}

// FloorToMultiplier returns the integer less or equal to x and multiplier m product.
// Works only with x >= 0 and m > 0. It returns 0 with other values.
func FloorToMultiplier(x, m int64) int64 {
	if x <= 0 || m <= 0 {
		return int64(0)
	}
	return x / m * m
}

// GCD returns the absolute greatest common divisor calculated via Euclidean algorithm
func GCD(a, b int64) int64 {
	if b < 0 {
		b = -b
	}
	for b != 0 {
		t := b
		b = a % b
		a = t
	}
	return a
}

// LCM returns the absolute least common multiple of 2 integers via GDB
func LCM(a, b int64) int64 {
	if a*b < 0 {
		return -a / GCD(a, b) * b
	}
	return a / GCD(a, b) * b
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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
//
//     https://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 mesh

import (
	"errors"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// StartApp uses the reminder of the command line to exec an app, using K8S_UID as UID, if present.
func (kr *KRun) StartApp() {
	var cmd *exec.Cmd
	if len(os.Args) == 1 {
		return
	} else if len(os.Args) == 2 {
		cmd = exec.Command(os.Args[1])
	} else {
		cmd = exec.Command(os.Args[1], os.Args[2:]...)
	}
	if os.Getuid() == 0 {
		uid := os.Getenv("K8S_UID")
		if uid != "" {
			uidi, err := strconv.Atoi(uid)
			if err == nil {
				cmd.SysProcAttr = &syscall.SysProcAttr{}
				cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uint32(uidi)}
			}
		}
	}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// Set port to 8080 - some apps use the PORT from knative to start.a
	cmd.Env = []string{"PORT=8080"}
	for _, e := range os.Environ() {
		if strings.HasPrefix(e, "PORT=") {
			continue
		}
		cmd.Env = append(cmd.Env, e)
	}
	if os.Getenv("GRPC_XDS_BOOTSTRAP") == "" {
		cmd.Env = append(cmd.Env, "GRPC_XDS_BOOTSTRAP=/var/run/grpc_bootstrap.json")
	}
	if kr.WhiteboxMode {
		cmd.Env = append(cmd.Env, "HTTP_PROXY=127.0.0.1:15007")
		cmd.Env = append(cmd.Env, "http_proxy=127.0.0.1:15007")
	}

	go func() {
		err := cmd.Start()
		if err != nil {
			log.Println("Failed to start ", cmd, err)
		}
		kr.appCmd = cmd
		err = cmd.Wait()
		if err != nil {
			log.Println("Application err exit ", err, cmd.ProcessState.ExitCode(), time.Since(kr.StartTime))
		} else {
			log.Println("Application clean exit ", err, cmd.ProcessState.ExitCode(), time.Since(kr.StartTime))
		}
		kr.Exit(cmd.ProcessState.ExitCode())
	}()

	kr.Signals()
}

// WaitTCPReady uses the same detection as CloudRun, i.e. TCP connect.
func (kr *KRun) WaitTCPReady(addr string, max time.Duration) error {
	t0 := time.Now()
	deadline := t0.Add(max)

	for {
		// if we cant connect, count as fail
		conn, err := net.DialTimeout("tcp", addr, deadline.Sub(time.Now()))
		if err != nil {
			if time.Now().After(deadline) {
				return err
			}
			time.Sleep(50 * time.Millisecond)
			if conn != nil {
				_ = conn.Close()
			}
			continue
		}
		err = conn.Close()
		if err != nil {
			log.Println("WaitTCP.Close()", err)
		}
		log.Println("Application ready", time.Since(t0), time.Since(kr.StartTime))
		return nil
	}
	return nil

}

func (kr *KRun) WaitHTTPReady(url string, max time.Duration) error {
	t0 := time.Now()
	for {
		res, _ := http.Get(url)
		if res != nil && res.StatusCode == 200 {
			log.Println("Ready")
			return nil
		}

		if time.Since(t0) > max {
			return errors.New("Timeout waiting for ready")
		}
		time.Sleep(100 * time.Millisecond)
	}
}

<commit_msg>Add envoy based health checks for ready-ness<commit_after>// Copyright 2021 Google LLC
//
// 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
//
//     https://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 mesh

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
)

const (
	serverStateKey        = "server.state"
	serverStateCheckRegex = "^server.state"
	listenerCheckKey      = "listener_manager.workers_started"
	listenerCheckRegex    = "^listener_manager.workers_started"
)

// StartApp uses the reminder of the command line to exec an app, using K8S_UID as UID, if present.
func (kr *KRun) StartApp() {
	var cmd *exec.Cmd
	if len(os.Args) == 1 {
		return
	} else if len(os.Args) == 2 {
		cmd = exec.Command(os.Args[1])
	} else {
		cmd = exec.Command(os.Args[1], os.Args[2:]...)
	}
	if os.Getuid() == 0 {
		uid := os.Getenv("K8S_UID")
		if uid != "" {
			uidi, err := strconv.Atoi(uid)
			if err == nil {
				cmd.SysProcAttr = &syscall.SysProcAttr{}
				cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uint32(uidi)}
			}
		}
	}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// Set port to 8080 - some apps use the PORT from knative to start.a
	cmd.Env = []string{"PORT=8080"}
	for _, e := range os.Environ() {
		if strings.HasPrefix(e, "PORT=") {
			continue
		}
		cmd.Env = append(cmd.Env, e)
	}
	if os.Getenv("GRPC_XDS_BOOTSTRAP") == "" {
		cmd.Env = append(cmd.Env, "GRPC_XDS_BOOTSTRAP=/var/run/grpc_bootstrap.json")
	}
	if kr.WhiteboxMode {
		cmd.Env = append(cmd.Env, "HTTP_PROXY=127.0.0.1:15007")
		cmd.Env = append(cmd.Env, "http_proxy=127.0.0.1:15007")
	}

	go func() {
		err := cmd.Start()
		if err != nil {
			log.Println("Failed to start ", cmd, err)
		}
		kr.appCmd = cmd
		err = cmd.Wait()
		if err != nil {
			log.Println("Application err exit ", err, cmd.ProcessState.ExitCode(), time.Since(kr.StartTime))
		} else {
			log.Println("Application clean exit ", err, cmd.ProcessState.ExitCode(), time.Since(kr.StartTime))
		}
		kr.Exit(cmd.ProcessState.ExitCode())
	}()

	kr.Signals()
}

// WaitTCPReady uses the same detection as CloudRun, i.e. TCP connect.
func (kr *KRun) WaitTCPReady(addr string, max time.Duration) error {
	t0 := time.Now()
	deadline := t0.Add(max)

	for {
		// if we cant connect, count as fail
		conn, err := net.DialTimeout("tcp", addr, deadline.Sub(time.Now()))
		if err != nil {
			if time.Now().After(deadline) {
				return err
			}
			time.Sleep(50 * time.Millisecond)
			if conn != nil {
				_ = conn.Close()
			}
			continue
		}
		err = conn.Close()
		if err != nil {
			log.Println("WaitTCP.Close()", err)
		}
		log.Println("Application ready", time.Since(t0), time.Since(kr.StartTime))
		return nil
	}
	return nil

}

func (kr *KRun) WaitHTTPReady(url string, max time.Duration) error {
	t0 := time.Now()
	for {
		res, _ := http.Get(url)
		if res != nil && res.StatusCode == 200 {
			log.Println("Ready")
			return nil
		}

		if time.Since(t0) > max {
			return errors.New("Timeout waiting for ready")
		}
		time.Sleep(100 * time.Millisecond)
	}
}

// WaitEnvoyReady waits for envoy to be ready until max is reached, otherwise returns a non-nil error.
func (kr *KRun) WaitEnvoyReady(addr string, max time.Duration) error {
	t0 := time.Now()
	for {
		serverStateReady, serverStateErr := kr.envoyServerStateCheck(addr)
		listenerReady, listenerErr := kr.envoyListenerWorkersStartedCheck(addr)
		if serverStateErr == nil && listenerErr == nil && serverStateReady && listenerReady {
			log.Println("Envoy is ready")
			return nil
		}

		if time.Since(t0) > max {
			return fmt.Errorf("Timeout waiting for ready from envoy")
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func (kr *KRun) envoyServerStateCheck(addr string) (bool, error) {
	checkURL := fmt.Sprintf("http://%s/stats?used_only&filter=%s", addr, serverStateCheckRegex)
	res, err := http.Get(checkURL)
	if err != nil {
		return false, fmt.Errorf("Unable to check envoy server state because of %s", err)
	}
	defer res.Body.Close()
	return kr.processHealthCheckResponse(res, serverStateKey, "0") // 0 indicates live.
}

func (kr *KRun) envoyListenerWorkersStartedCheck(addr string) (bool, error) {
	checkURL := fmt.Sprintf("http://%s/stats?used_only&filter=%s", addr, listenerCheckRegex)
	res, err := http.Get(checkURL)
	if err != nil {
		return false, fmt.Errorf("Unable to check envoy listener worker state because of : %s", err)
	}
	defer res.Body.Close()
	return kr.processHealthCheckResponse(res, listenerCheckKey, "1") // 1 indicates listener works have started.
}

// Checks the res has the following structure: "key: val" where key and val should be as specified.
func (kr *KRun) processHealthCheckResponse(res *http.Response, key string, val string) (bool, error) {
	if res == nil || res.StatusCode != 200 {
		return false, fmt.Errorf("Unable to check envoy for %s", key)
	}

	rawResponse, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return false, fmt.Errorf("Unable to check envoy for %s because of %s", key, err)
	}

	response := string(rawResponse)
	splits := strings.Split(response, ":")
	if len(splits) != 2 {
		return false, nil
	}

	// Check key
	if strings.TrimSpace(splits[0]) != key {
		return false, nil
	}
	// Check val
	if strings.TrimSpace(splits[1]) != val {
		return false, nil
	}
	return true, nil
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package uroot

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/imports"

	"github.com/u-root/u-root/pkg/cpio"
	"github.com/u-root/u-root/pkg/golang"
)

// Commands to skip building in bb mode.
var skip = map[string]struct{}{
	"bb": {},
}

type bbBuilder struct {
	opts  BuildOpts
	bbDir string
	af    ArchiveFiles

	init string

	importer types.Importer
}

// BBBuild is an implementation of Build for the busybox-like u-root initramfs.
//
// BBBuild rewrites the source files of the packages given to create one
// busybox-like binary containing all commands in `opts.Packages`.
func BBBuild(opts BuildOpts) (ArchiveFiles, error) {
	urootPkg, err := opts.Env.Package("github.com/u-root/u-root")
	if err != nil {
		return ArchiveFiles{}, err
	}

	bbDir := filepath.Join(urootPkg.Dir, "bb")
	// Blow bbsh away before trying to re-create it.
	if err := os.RemoveAll(bbDir); err != nil {
		return ArchiveFiles{}, err
	}
	if err := os.MkdirAll(bbDir, 0755); err != nil {
		return ArchiveFiles{}, err
	}

	builder := &bbBuilder{
		opts:  opts,
		bbDir: bbDir,
		af:    NewArchiveFiles(),
		init: `package main

		import (
`,
		importer: importer.For("source", nil),
	}

	// Move and rewrite package files.
	for _, pkg := range opts.Packages {
		if _, ok := skip[filepath.Base(pkg)]; ok {
			continue
		}

		if err := builder.moveCommand(pkg); err != nil {
			return ArchiveFiles{}, err
		}
	}

	// Move bb cmd over.
	p, err := opts.Env.Package("github.com/u-root/u-root/pkg/bb/cmd")
	if err != nil {
		return ArchiveFiles{}, err
	}

	if err := filepath.Walk(p.Dir, func(name string, fi os.FileInfo, err error) error {
		if err != nil || fi.IsDir() {
			return err
		}
		b, err := ioutil.ReadFile(name)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(filepath.Join(builder.bbDir, fi.Name()), b, 0644)
	}); err != nil {
		return ArchiveFiles{}, err
	}

	// Write init.go, which imports all command packages.
	//
	// I swear, I tried this with *ast.File and astutil.AddNamedImport, but
	// I just couldn't get it to work. Whatever.
	builder.init += ")\n"
	if err := writeGoFile(filepath.Join(builder.bbDir, "init.go"), []byte(builder.init)); err != nil {
		return ArchiveFiles{}, err
	}

	// Compile bb + commands to /bbin/bb.
	bbPath := filepath.Join(opts.TempDir, "bb")
	if err := opts.Env.Build("github.com/u-root/u-root/bb", bbPath, golang.BuildOpts{}); err != nil {
		return ArchiveFiles{}, err
	}
	if err := builder.af.AddFile(bbPath, "bbin/bb"); err != nil {
		return ArchiveFiles{}, err
	}

	// Symlink from /init to init.
	if err := builder.af.AddRecord(cpio.Symlink("init", "/bbin/init")); err != nil {
		return ArchiveFiles{}, err
	}
	return builder.af, nil
}

// Package is a Go package.
//
// It holds AST, type, file, and Go package information about a Go package.
type Package struct {
	name string

	pkg      *build.Package
	fset     *token.FileSet
	ast      *ast.Package
	typeInfo types.Info
	types    *types.Package

	// initCount keeps track of what the next init's index should be.
	initCount uint

	// init is the cmd.Init function that calls all other InitXs in the
	// right order.
	init *ast.FuncDecl

	// initAssigns is a map of assignment expression -> assignment
	// statement.
	//
	// types.Info.InitOrder keeps track of Initializations by Lhs name and
	// Rhs ast.Expr.  We reparent the Rhs in assignment statements in InitX
	// functions, so we use the Rhs as an easy key here.
	// types.Info.InitOrder + initAssigns can then easily be used to derive
	// the order of AssignStmts.
	//
	// The key Expr must also be the AssignStmt.Rhs[0].
	initAssigns map[ast.Expr]*ast.AssignStmt
}

func (p *Package) nextInit() *ast.Ident {
	i := ast.NewIdent(fmt.Sprintf("Init%d", p.initCount))
	p.init.Body.List = append(p.init.Body.List, &ast.ExprStmt{X: &ast.CallExpr{Fun: i}})
	p.initCount++
	return i
}

// TODO:
// - write an init name generator, in case InitN is already taken.
// - also rewrite all non-Go-stdlib dependencies.
func (p *Package) rewriteFile(f *ast.File) bool {
	hasMain := false

	// Change the package name declaration from main to the command's name.
	f.Name = ast.NewIdent(p.name)

	// Map of fully qualified package name -> imported alias in the file.
	importAliases := make(map[string]string)
	for _, impt := range f.Imports {
		if impt.Name != nil {
			importPath, err := strconv.Unquote(impt.Path.Value)
			if err != nil {
				panic(err)
			}
			importAliases[importPath] = impt.Name.Name
		}
	}

	// When the types.TypeString function translates package names, it uses
	// this function to map fully qualified package paths to a local alias,
	// if it exists.
	qualifier := func(pkg *types.Package) string {
		name, ok := importAliases[pkg.Path()]
		if ok {
			return name
		}
		// When referring to self, don't use any package name.
		if pkg == p.types {
			return ""
		}
		return pkg.Name()
	}

	for _, decl := range f.Decls {
		switch d := decl.(type) {
		case *ast.GenDecl:
			// We only care about vars.
			if d.Tok != token.VAR {
				break
			}
			for _, spec := range d.Specs {
				s := spec.(*ast.ValueSpec)
				if s.Values == nil {
					continue
				}

				// Add an assign statement for these values to init.
				for i, name := range s.Names {
					p.initAssigns[s.Values[i]] = &ast.AssignStmt{
						Lhs: []ast.Expr{name},
						Tok: token.ASSIGN,
						Rhs: []ast.Expr{s.Values[i]},
					}
				}

				// Add the type of the expression to the global
				// declaration instead.
				if s.Type == nil {
					typ := p.typeInfo.Types[s.Values[0]]
					s.Type = ast.NewIdent(types.TypeString(typ.Type, qualifier))
				}
				s.Values = nil
			}

		case *ast.FuncDecl:
			if d.Recv == nil && d.Name.Name == "main" {
				d.Name.Name = "Main"
				hasMain = true
			}
			if d.Recv == nil && d.Name.Name == "init" {
				d.Name = p.nextInit()
			}
		}
	}

	// Now we change any import names attached to package declarations. We
	// just upcase it for now; it makes it easy to look in bbsh for things
	// we changed, e.g. grep -r bbsh Import is useful.
	for _, cg := range f.Comments {
		for _, c := range cg.List {
			if strings.HasPrefix(c.Text, "// import") {
				c.Text = "// Import" + c.Text[9:]
			}
		}
	}
	return hasMain
}

func writeFile(path string, fset *token.FileSet, f *ast.File) error {
	var buf bytes.Buffer
	if err := format.Node(&buf, fset, f); err != nil {
		return fmt.Errorf("error formatting Go file %q: %v", path, err)
	}
	return writeGoFile(path, buf.Bytes())
}

func writeGoFile(path string, code []byte) error {
	// Fix up imports.
	opts := imports.Options{
		Fragment:  true,
		AllErrors: true,
		Comments:  true,
		TabIndent: true,
		TabWidth:  8,
	}
	code, err := imports.Process("commandline", code, &opts)
	if err != nil {
		return fmt.Errorf("bad parse while processing imports %q: %v", path, err)
	}

	if err := ioutil.WriteFile(path, code, 0644); err != nil {
		return fmt.Errorf("error writing Go file to %q: %v", path, err)
	}
	return nil
}

func getPackage(opts BuildOpts, importPath string, importer types.Importer) (*Package, error) {
	p, err := opts.Env.Package(importPath)
	if err != nil {
		return nil, err
	}

	name := filepath.Base(p.Dir)
	if !p.IsCommand() {
		return nil, fmt.Errorf("package %q is not a command and cannot be included in bb", name)
	}

	fset := token.NewFileSet()
	pars, err := parser.ParseDir(fset, p.Dir, func(fi os.FileInfo) bool {
		// Only parse Go files that match build tags of this package.
		for _, name := range p.GoFiles {
			if name == fi.Name() {
				return true
			}
		}
		return false
	}, parser.ParseComments)
	if err != nil {
		log.Printf("can't parsedir %q: %v", p.Dir, err)
		return nil, nil
	}

	pp := &Package{
		name: name,
		pkg:  p,
		fset: fset,
		ast:  pars[p.Name],
		typeInfo: types.Info{
			Types: make(map[ast.Expr]types.TypeAndValue),
		},
		initAssigns: make(map[ast.Expr]*ast.AssignStmt),
	}

	// This Init will hold calls to all other InitXs.
	pp.init = &ast.FuncDecl{
		Name: ast.NewIdent("Init"),
		Type: &ast.FuncType{
			Params:  &ast.FieldList{},
			Results: nil,
		},
		Body: &ast.BlockStmt{},
	}

	// The order of types.Info.InitOrder depends on this list of files
	// always being passed to conf.Check in the same order.
	filenames := make([]string, 0, len(pp.ast.Files))
	for name := range pp.ast.Files {
		filenames = append(filenames, name)
	}
	sort.Strings(filenames)

	sortedFiles := make([]*ast.File, 0, len(pp.ast.Files))
	for _, name := range filenames {
		sortedFiles = append(sortedFiles, pp.ast.Files[name])
	}
	// Type-check the package before we continue. We need types to rewrite
	// some statements.
	conf := types.Config{
		Importer: importer,

		// We only need global declarations' types.
		IgnoreFuncBodies: true,
	}
	tpkg, err := conf.Check(pp.pkg.ImportPath, pp.fset, sortedFiles, &pp.typeInfo)
	if err != nil {
		return nil, fmt.Errorf("type checking failed: %v", err)
	}
	pp.types = tpkg
	return pp, nil
}

func (b *bbBuilder) moveCommand(pkgPath string) error {
	p, err := getPackage(b.opts, pkgPath, b.importer)
	if err != nil {
		return err
	}
	if p == nil {
		return nil
	}

	pkgDir := filepath.Join(b.bbDir, "cmds", p.name)
	if err := os.MkdirAll(pkgDir, 0755); err != nil {
		return err
	}

	// This init holds all variable initializations.
	//
	// func [nextInitName]() {}
	varInit := &ast.FuncDecl{
		Name: p.nextInit(),
		Type: &ast.FuncType{
			Params:  &ast.FieldList{},
			Results: nil,
		},
		Body: &ast.BlockStmt{},
	}

	var mainPath string
	var hasMain bool
	for filePath, sourceFile := range p.ast.Files {
		if hasMainFile := p.rewriteFile(sourceFile); hasMainFile {
			hasMain = true
			mainPath = filePath
		}
	}
	if !hasMain {
		return os.RemoveAll(pkgDir)
	}

	// Add variable initializations to Init0 in the right order.
	for _, initStmt := range p.typeInfo.InitOrder {
		a, ok := p.initAssigns[initStmt.Rhs]
		if !ok {
			return fmt.Errorf("couldn't find init assignment %s", initStmt)
		}
		varInit.Body.List = append(varInit.Body.List, a)
	}

	astutil.AddImport(p.fset, p.ast.Files[mainPath], "github.com/u-root/u-root/pkg/bb")

	// func init() {
	//   bb.Register(p.name, Init, Main)
	// }
	bbRegisterInit := &ast.FuncDecl{
		Name: ast.NewIdent("init"),
		Type: &ast.FuncType{},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun: ast.NewIdent("bb.Register"),
					Args: []ast.Expr{
						// name=
						&ast.BasicLit{
							Kind:  token.STRING,
							Value: fmt.Sprintf("%#v", p.name),
						},
						// init=
						ast.NewIdent("Init"),
						// main=
						ast.NewIdent("Main"),
					},
				}},
			},
		},
	}

	// We could add these statements to any of the package files. We choose
	// the one that contains Main() to guarantee reproducibility of the
	// same bbsh binary.
	p.ast.Files[mainPath].Decls = append(p.ast.Files[mainPath].Decls, varInit, p.init, bbRegisterInit)

	for filePath, sourceFile := range p.ast.Files {
		path := filepath.Join(pkgDir, filepath.Base(filePath))
		if err := writeFile(path, p.fset, sourceFile); err != nil {
			return err
		}
	}

	// Add side-effect import to bb binary so init registers itself.
	b.init += fmt.Sprintf("_ \"github.com/u-root/u-root/bb/cmds/%s\"\n", p.name)

	// Add a symlink to our bbsh.
	return b.af.AddRecord(cpio.Symlink(filepath.Join("bbin", p.name), "/bbin/bb"))
}
<commit_msg>pkg/uroot/bb: write one bb main.go rather than main.go and init.go with string template.<commit_after>// Copyright 2015-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package uroot

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/imports"

	"github.com/u-root/u-root/pkg/cpio"
	"github.com/u-root/u-root/pkg/golang"
)

// Commands to skip building in bb mode.
var skip = map[string]struct{}{
	"bb": {},
}

type bbBuilder struct {
	opts  BuildOpts
	bbDir string
	af    ArchiveFiles

	bb *Package

	importer types.Importer
}

// BBBuild is an implementation of Build for the busybox-like u-root initramfs.
//
// BBBuild rewrites the source files of the packages given to create one
// busybox-like binary containing all commands in `opts.Packages`.
func BBBuild(opts BuildOpts) (ArchiveFiles, error) {
	urootPkg, err := opts.Env.Package("github.com/u-root/u-root")
	if err != nil {
		return ArchiveFiles{}, err
	}

	bbDir := filepath.Join(urootPkg.Dir, "bb")
	// Blow bbsh away before trying to re-create it.
	if err := os.RemoveAll(bbDir); err != nil {
		return ArchiveFiles{}, err
	}
	if err := os.MkdirAll(bbDir, 0755); err != nil {
		return ArchiveFiles{}, err
	}

	builder := &bbBuilder{
		opts:     opts,
		bbDir:    bbDir,
		af:       NewArchiveFiles(),
		importer: importer.For("source", nil),
	}

	p, err := getPackage(opts, "github.com/u-root/u-root/pkg/bb/cmd", builder.importer)
	if err != nil {
		return ArchiveFiles{}, err
	}
	if p == nil {
		return ArchiveFiles{}, fmt.Errorf("bb/cmd missing")
	}
	if len(p.ast.Files) != 1 {
		return ArchiveFiles{}, fmt.Errorf("bb/cmd is supposed to only have one file")
	}
	builder.bb = p

	// Move and rewrite package files.
	for _, pkg := range opts.Packages {
		if _, ok := skip[filepath.Base(pkg)]; ok {
			continue
		}

		if err := builder.moveCommand(pkg); err != nil {
			return ArchiveFiles{}, err
		}
	}

	// Write bb cmd out.
	for filePath, sourceFile := range p.ast.Files {
		path := filepath.Join(builder.bbDir, filepath.Base(filePath))
		if err := writeFile(path, p.fset, sourceFile); err != nil {
			return ArchiveFiles{}, err
		}
	}

	// Compile bb + commands to /bbin/bb.
	bbPath := filepath.Join(opts.TempDir, "bb")
	if err := opts.Env.Build("github.com/u-root/u-root/bb", bbPath, golang.BuildOpts{}); err != nil {
		return ArchiveFiles{}, err
	}
	if err := builder.af.AddFile(bbPath, "bbin/bb"); err != nil {
		return ArchiveFiles{}, err
	}

	// Symlink from /init to init.
	if err := builder.af.AddRecord(cpio.Symlink("init", "/bbin/init")); err != nil {
		return ArchiveFiles{}, err
	}
	return builder.af, nil
}

// Package is a Go package.
//
// It holds AST, type, file, and Go package information about a Go package.
type Package struct {
	name string

	pkg      *build.Package
	fset     *token.FileSet
	ast      *ast.Package
	typeInfo types.Info
	types    *types.Package

	// initCount keeps track of what the next init's index should be.
	initCount uint

	// init is the cmd.Init function that calls all other InitXs in the
	// right order.
	init *ast.FuncDecl

	// initAssigns is a map of assignment expression -> assignment
	// statement.
	//
	// types.Info.InitOrder keeps track of Initializations by Lhs name and
	// Rhs ast.Expr.  We reparent the Rhs in assignment statements in InitX
	// functions, so we use the Rhs as an easy key here.
	// types.Info.InitOrder + initAssigns can then easily be used to derive
	// the order of AssignStmts.
	//
	// The key Expr must also be the AssignStmt.Rhs[0].
	initAssigns map[ast.Expr]*ast.AssignStmt
}

func (p *Package) nextInit() *ast.Ident {
	i := ast.NewIdent(fmt.Sprintf("Init%d", p.initCount))
	p.init.Body.List = append(p.init.Body.List, &ast.ExprStmt{X: &ast.CallExpr{Fun: i}})
	p.initCount++
	return i
}

// TODO:
// - write an init name generator, in case InitN is already taken.
// - also rewrite all non-Go-stdlib dependencies.
func (p *Package) rewriteFile(f *ast.File) bool {
	hasMain := false

	// Change the package name declaration from main to the command's name.
	f.Name = ast.NewIdent(p.name)

	// Map of fully qualified package name -> imported alias in the file.
	importAliases := make(map[string]string)
	for _, impt := range f.Imports {
		if impt.Name != nil {
			importPath, err := strconv.Unquote(impt.Path.Value)
			if err != nil {
				panic(err)
			}
			importAliases[importPath] = impt.Name.Name
		}
	}

	// When the types.TypeString function translates package names, it uses
	// this function to map fully qualified package paths to a local alias,
	// if it exists.
	qualifier := func(pkg *types.Package) string {
		name, ok := importAliases[pkg.Path()]
		if ok {
			return name
		}
		// When referring to self, don't use any package name.
		if pkg == p.types {
			return ""
		}
		return pkg.Name()
	}

	for _, decl := range f.Decls {
		switch d := decl.(type) {
		case *ast.GenDecl:
			// We only care about vars.
			if d.Tok != token.VAR {
				break
			}
			for _, spec := range d.Specs {
				s := spec.(*ast.ValueSpec)
				if s.Values == nil {
					continue
				}

				// Add an assign statement for these values to init.
				for i, name := range s.Names {
					p.initAssigns[s.Values[i]] = &ast.AssignStmt{
						Lhs: []ast.Expr{name},
						Tok: token.ASSIGN,
						Rhs: []ast.Expr{s.Values[i]},
					}
				}

				// Add the type of the expression to the global
				// declaration instead.
				if s.Type == nil {
					typ := p.typeInfo.Types[s.Values[0]]
					s.Type = ast.NewIdent(types.TypeString(typ.Type, qualifier))
				}
				s.Values = nil
			}

		case *ast.FuncDecl:
			if d.Recv == nil && d.Name.Name == "main" {
				d.Name.Name = "Main"
				hasMain = true
			}
			if d.Recv == nil && d.Name.Name == "init" {
				d.Name = p.nextInit()
			}
		}
	}

	// Now we change any import names attached to package declarations. We
	// just upcase it for now; it makes it easy to look in bbsh for things
	// we changed, e.g. grep -r bbsh Import is useful.
	for _, cg := range f.Comments {
		for _, c := range cg.List {
			if strings.HasPrefix(c.Text, "// import") {
				c.Text = "// Import" + c.Text[9:]
			}
		}
	}
	return hasMain
}

func writeFile(path string, fset *token.FileSet, f *ast.File) error {
	var buf bytes.Buffer
	if err := format.Node(&buf, fset, f); err != nil {
		return fmt.Errorf("error formatting Go file %q: %v", path, err)
	}
	return writeGoFile(path, buf.Bytes())
}

func writeGoFile(path string, code []byte) error {
	// Fix up imports.
	opts := imports.Options{
		Fragment:  true,
		AllErrors: true,
		Comments:  true,
		TabIndent: true,
		TabWidth:  8,
	}
	code, err := imports.Process("commandline", code, &opts)
	if err != nil {
		return fmt.Errorf("bad parse while processing imports %q: %v", path, err)
	}

	if err := ioutil.WriteFile(path, code, 0644); err != nil {
		return fmt.Errorf("error writing Go file to %q: %v", path, err)
	}
	return nil
}

func getPackage(opts BuildOpts, importPath string, importer types.Importer) (*Package, error) {
	p, err := opts.Env.Package(importPath)
	if err != nil {
		return nil, err
	}

	name := filepath.Base(p.Dir)
	if !p.IsCommand() {
		return nil, fmt.Errorf("package %q is not a command and cannot be included in bb", name)
	}

	fset := token.NewFileSet()
	pars, err := parser.ParseDir(fset, p.Dir, func(fi os.FileInfo) bool {
		// Only parse Go files that match build tags of this package.
		for _, name := range p.GoFiles {
			if name == fi.Name() {
				return true
			}
		}
		return false
	}, parser.ParseComments)
	if err != nil {
		log.Printf("can't parsedir %q: %v", p.Dir, err)
		return nil, nil
	}

	pp := &Package{
		name: name,
		pkg:  p,
		fset: fset,
		ast:  pars[p.Name],
		typeInfo: types.Info{
			Types: make(map[ast.Expr]types.TypeAndValue),
		},
		initAssigns: make(map[ast.Expr]*ast.AssignStmt),
	}

	// This Init will hold calls to all other InitXs.
	pp.init = &ast.FuncDecl{
		Name: ast.NewIdent("Init"),
		Type: &ast.FuncType{
			Params:  &ast.FieldList{},
			Results: nil,
		},
		Body: &ast.BlockStmt{},
	}

	// The order of types.Info.InitOrder depends on this list of files
	// always being passed to conf.Check in the same order.
	filenames := make([]string, 0, len(pp.ast.Files))
	for name := range pp.ast.Files {
		filenames = append(filenames, name)
	}
	sort.Strings(filenames)

	sortedFiles := make([]*ast.File, 0, len(pp.ast.Files))
	for _, name := range filenames {
		sortedFiles = append(sortedFiles, pp.ast.Files[name])
	}
	// Type-check the package before we continue. We need types to rewrite
	// some statements.
	conf := types.Config{
		Importer: importer,

		// We only need global declarations' types.
		IgnoreFuncBodies: true,
	}
	tpkg, err := conf.Check(pp.pkg.ImportPath, pp.fset, sortedFiles, &pp.typeInfo)
	if err != nil {
		return nil, fmt.Errorf("type checking failed: %v", err)
	}
	pp.types = tpkg
	return pp, nil
}

func (b *bbBuilder) moveCommand(pkgPath string) error {
	p, err := getPackage(b.opts, pkgPath, b.importer)
	if err != nil {
		return err
	}
	if p == nil {
		return nil
	}

	pkgDir := filepath.Join(b.bbDir, "cmds", p.name)
	if err := os.MkdirAll(pkgDir, 0755); err != nil {
		return err
	}

	// This init holds all variable initializations.
	//
	// func [nextInitName]() {}
	varInit := &ast.FuncDecl{
		Name: p.nextInit(),
		Type: &ast.FuncType{
			Params:  &ast.FieldList{},
			Results: nil,
		},
		Body: &ast.BlockStmt{},
	}

	var mainPath string
	var hasMain bool
	for filePath, sourceFile := range p.ast.Files {
		if hasMainFile := p.rewriteFile(sourceFile); hasMainFile {
			hasMain = true
			mainPath = filePath
		}
	}
	if !hasMain {
		return os.RemoveAll(pkgDir)
	}

	// Add variable initializations to Init0 in the right order.
	for _, initStmt := range p.typeInfo.InitOrder {
		a, ok := p.initAssigns[initStmt.Rhs]
		if !ok {
			return fmt.Errorf("couldn't find init assignment %s", initStmt)
		}
		varInit.Body.List = append(varInit.Body.List, a)
	}

	astutil.AddImport(p.fset, p.ast.Files[mainPath], "github.com/u-root/u-root/pkg/bb")

	// func init() {
	//   bb.Register(p.name, Init, Main)
	// }
	bbRegisterInit := &ast.FuncDecl{
		Name: ast.NewIdent("init"),
		Type: &ast.FuncType{},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun: ast.NewIdent("bb.Register"),
					Args: []ast.Expr{
						// name=
						&ast.BasicLit{
							Kind:  token.STRING,
							Value: fmt.Sprintf("%#v", p.name),
						},
						// init=
						ast.NewIdent("Init"),
						// main=
						ast.NewIdent("Main"),
					},
				}},
			},
		},
	}

	// We could add these statements to any of the package files. We choose
	// the one that contains Main() to guarantee reproducibility of the
	// same bbsh binary.
	p.ast.Files[mainPath].Decls = append(p.ast.Files[mainPath].Decls, varInit, p.init, bbRegisterInit)

	for filePath, sourceFile := range p.ast.Files {
		path := filepath.Join(pkgDir, filepath.Base(filePath))
		if err := writeFile(path, p.fset, sourceFile); err != nil {
			return err
		}
	}

	// Add side-effect import to bb binary so init registers itself.
	for _, sourceFile := range b.bb.ast.Files {
		astutil.AddNamedImport(b.bb.fset, sourceFile, "_", fmt.Sprintf("github.com/u-root/u-root/bb/cmds/%s", p.name))
	}

	// Add a symlink to our bbsh.
	return b.af.AddRecord(cpio.Symlink(filepath.Join("bbin", p.name), "/bbin/bb"))
}
<|endoftext|>


<commit_before>package vfs

import (

	// #nosec
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
)

// FileDoc is a struct containing all the informations about a file.
// It implements the couchdb.Doc and jsonapi.Object interfaces.
type FileDoc struct {
	// Type of document. Useful to (de)serialize and filter the data
	// from couch.
	Type string `json:"type"`
	// Qualified file identifier
	DocID string `json:"_id,omitempty"`
	// File revision
	DocRev string `json:"_rev,omitempty"`
	// File name
	DocName string `json:"name"`
	// Parent directory identifier
	DirID       string `json:"dir_id,omitempty"`
	RestorePath string `json:"restore_path,omitempty"`

	CreatedAt time.Time `json:"created_at"`
	UpdatedAt time.Time `json:"updated_at"`

	ByteSize   int64    `json:"size,string"` // Serialized in JSON as a string, because JS has some issues with big numbers
	MD5Sum     []byte   `json:"md5sum"`
	Mime       string   `json:"mime"`
	Class      string   `json:"class"`
	Executable bool     `json:"executable"`
	Trashed    bool     `json:"trashed"`
	Tags       []string `json:"tags"`

	Metadata Metadata `json:"metadata,omitempty"`

	ReferencedBy []couchdb.DocReference `json:"referenced_by,omitempty"`

	// Cache of the fullpath of the file. Should not have to be invalidated
	// since we use FileDoc as immutable data-structures.
	fullpath string

	// NOTE: Do not forget to propagate changes made to this structure to the
	// structure DirOrFileDoc in pkg/vfs/vfs.go.
}

// ID returns the file qualified identifier
func (f *FileDoc) ID() string { return f.DocID }

// Rev returns the file revision
func (f *FileDoc) Rev() string { return f.DocRev }

// DocType returns the file document type
func (f *FileDoc) DocType() string { return consts.Files }

// Clone implements couchdb.Doc
func (f *FileDoc) Clone() couchdb.Doc { cloned := *f; return &cloned }

// SetID changes the file qualified identifier
func (f *FileDoc) SetID(id string) { f.DocID = id }

// SetRev changes the file revision
func (f *FileDoc) SetRev(rev string) { f.DocRev = rev }

// Path is used to generate the file path
func (f *FileDoc) Path(fs VFS) (string, error) {
	if f.fullpath != "" {
		return f.fullpath, nil
	}
	var err error
	f.fullpath, err = fs.FilePath(f)
	return f.fullpath, err
}

// Parent returns the parent directory document
func (f *FileDoc) Parent(fs VFS) (*DirDoc, error) {
	parent, err := fs.DirByID(f.DirID)
	if os.IsNotExist(err) {
		err = ErrParentDoesNotExist
	}
	return parent, err
}

// Name returns base name of the file
func (f *FileDoc) Name() string { return f.DocName }

// Size returns the length in bytes for regular files; system-dependent for others
func (f *FileDoc) Size() int64 { return f.ByteSize }

// Mode returns the file mode bits
func (f *FileDoc) Mode() os.FileMode { return getFileMode(f.Executable) }

// ModTime returns the modification time
func (f *FileDoc) ModTime() time.Time { return f.UpdatedAt }

// IsDir returns the abbreviation for Mode().IsDir()
func (f *FileDoc) IsDir() bool { return false }

// Sys returns the underlying data source (can return nil)
func (f *FileDoc) Sys() interface{} { return nil }

// AddReferencedBy adds referenced_by to the file
func (f *FileDoc) AddReferencedBy(ri ...couchdb.DocReference) {
	f.ReferencedBy = append(f.ReferencedBy, ri...)
}

func containsReferencedBy(haystack []couchdb.DocReference, needle couchdb.DocReference) bool {
	for _, ref := range haystack {
		if ref.ID == needle.ID && ref.Type == needle.Type {
			return true
		}
	}
	return false
}

// RemoveReferencedBy adds referenced_by to the file
func (f *FileDoc) RemoveReferencedBy(ri ...couchdb.DocReference) {
	// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
	referenced := f.ReferencedBy[:0]
	for _, ref := range f.ReferencedBy {
		if !containsReferencedBy(ri, ref) {
			referenced = append(referenced, ref)
		}
	}
	f.ReferencedBy = referenced
}

// NewFileDoc is the FileDoc constructor. The given name is validated.
func NewFileDoc(name, dirID string, size int64, md5Sum []byte, mime, class string, cdate time.Time, executable, trashed bool, tags []string) (*FileDoc, error) {
	if err := checkFileName(name); err != nil {
		return nil, err
	}

	if dirID == "" {
		dirID = consts.RootDirID
	}

	tags = uniqueTags(tags)

	doc := &FileDoc{
		Type:    consts.FileType,
		DocName: name,
		DirID:   dirID,

		CreatedAt:  cdate,
		UpdatedAt:  cdate,
		ByteSize:   size,
		MD5Sum:     md5Sum,
		Mime:       mime,
		Class:      class,
		Executable: executable,
		Trashed:    trashed,
		Tags:       tags,
	}

	return doc, nil
}

// ServeFileContent replies to a http request using the content of a
// file given its FileDoc.
//
// It uses internally http.ServeContent and benefits from it by
// offering support to Range, If-Modified-Since and If-None-Match
// requests. It uses the revision of the file as the Etag value for
// non-ranged requests
//
// The content disposition is inlined.
func ServeFileContent(fs VFS, doc *FileDoc, disposition string, req *http.Request, w http.ResponseWriter) error {
	header := w.Header()
	header.Set("Content-Type", doc.Mime)
	if disposition != "" {
		header.Set("Content-Disposition", ContentDisposition(disposition, doc.DocName))
	}

	if header.Get("Range") == "" {
		eTag := base64.StdEncoding.EncodeToString(doc.MD5Sum)
		header.Set("Etag", eTag)
	}

	content, err := fs.OpenFile(doc)
	if err != nil {
		return err
	}
	defer content.Close()

	http.ServeContent(w, req, doc.DocName, doc.UpdatedAt, content)
	return nil
}

// ModifyFileMetadata modify the metadata associated to a file. It can
// be used to rename or move the file in the VFS.
func ModifyFileMetadata(fs VFS, olddoc *FileDoc, patch *DocPatch) (*FileDoc, error) {
	var err error
	rename := patch.Name != nil
	cdate := olddoc.CreatedAt
	oname := olddoc.DocName
	trashed := olddoc.Trashed
	if patch.RestorePath != nil {
		trashed = *patch.RestorePath != ""
	}
	patch, err = normalizeDocPatch(&DocPatch{
		Name:        &oname,
		DirID:       &olddoc.DirID,
		RestorePath: &olddoc.RestorePath,
		Tags:        &olddoc.Tags,
		UpdatedAt:   &olddoc.UpdatedAt,
		Executable:  &olddoc.Executable,
	}, patch, cdate)
	if err != nil {
		return nil, err
	}

	// in case of a renaming of the file, if the extension of the file has
	// changed, we consider recalculating the mime and class attributes, using
	// the new extension.
	newname := *patch.Name
	oldname := olddoc.DocName
	var mime, class string
	if rename && path.Ext(newname) != path.Ext(oldname) {
		mime, class = ExtractMimeAndClassFromFilename(newname)
	} else {
		mime, class = olddoc.Mime, olddoc.Class
	}

	newdoc, err := NewFileDoc(
		newname,
		*patch.DirID,
		olddoc.Size(),
		olddoc.MD5Sum,
		mime,
		class,
		cdate,
		*patch.Executable,
		trashed,
		*patch.Tags,
	)
	if err != nil {
		return nil, err
	}

	newdoc.RestorePath = *patch.RestorePath
	newdoc.UpdatedAt = *patch.UpdatedAt
	newdoc.Metadata = olddoc.Metadata
	newdoc.ReferencedBy = olddoc.ReferencedBy

	if err = fs.UpdateFileDoc(olddoc, newdoc); err != nil {
		return nil, err
	}
	return newdoc, nil
}

// TrashFile is used to delete a file given its document
func TrashFile(fs VFS, olddoc *FileDoc) (*FileDoc, error) {
	oldpath, err := olddoc.Path(fs)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(oldpath, TrashDirName) {
		return nil, ErrFileInTrash
	}

	trashDirID := consts.TrashDirID
	restorePath := path.Dir(oldpath)

	var newdoc *FileDoc
	tryOrUseSuffix(olddoc.DocName, conflictFormat, func(name string) error {
		newdoc, err = ModifyFileMetadata(fs, olddoc, &DocPatch{
			DirID:       &trashDirID,
			RestorePath: &restorePath,
			Name:        &name,
		})
		return err
	})
	return newdoc, err
}

// RestoreFile is used to restore a trashed file given its document
func RestoreFile(fs VFS, olddoc *FileDoc) (*FileDoc, error) {
	oldpath, err := olddoc.Path(fs)
	if err != nil {
		return nil, err
	}
	restoreDir, err := getRestoreDir(fs, oldpath, olddoc.RestorePath)
	if err != nil {
		return nil, err
	}
	var newdoc *FileDoc
	var emptyStr string
	name := stripSuffix(olddoc.DocName, conflictSuffix)
	tryOrUseSuffix(name, "%s (%s)", func(name string) error {
		newdoc, err = ModifyFileMetadata(fs, olddoc, &DocPatch{
			DirID:       &restoreDir.DocID,
			RestorePath: &emptyStr,
			Name:        &name,
		})
		return err
	})
	return newdoc, err
}

func getFileMode(executable bool) os.FileMode {
	if executable {
		return 0755 // -rwxr-xr-x
	}
	return 0644 // -rw-r--r--
}

// ThumbPath gives the path to a thumbnail in the Thumbs hidden FS
func ThumbPath(img *FileDoc, format string) string {
	dir := ThumbDir(img)
	name := fmt.Sprintf("%s-%s.jpg", img.ID(), format)
	return path.Join(dir, name)
}

// ThumbDir gives the path to the parent directory of a thumbnail in the
// Thumbs hidden FS
func ThumbDir(img *FileDoc) string {
	return img.ID()[:4]
}

var (
	_ couchdb.Doc = &FileDoc{}
	_ os.FileInfo = &FileDoc{}
)
<commit_msg>Fix files tests<commit_after>package vfs

import (

	// #nosec
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
)

// FileDoc is a struct containing all the informations about a file.
// It implements the couchdb.Doc and jsonapi.Object interfaces.
type FileDoc struct {
	// Type of document. Useful to (de)serialize and filter the data
	// from couch.
	Type string `json:"type"`
	// Qualified file identifier
	DocID string `json:"_id,omitempty"`
	// File revision
	DocRev string `json:"_rev,omitempty"`
	// File name
	DocName string `json:"name"`
	// Parent directory identifier
	DirID       string `json:"dir_id,omitempty"`
	RestorePath string `json:"restore_path,omitempty"`

	CreatedAt time.Time `json:"created_at"`
	UpdatedAt time.Time `json:"updated_at"`

	ByteSize   int64    `json:"size,string"` // Serialized in JSON as a string, because JS has some issues with big numbers
	MD5Sum     []byte   `json:"md5sum"`
	Mime       string   `json:"mime"`
	Class      string   `json:"class"`
	Executable bool     `json:"executable"`
	Trashed    bool     `json:"trashed"`
	Tags       []string `json:"tags"`

	Metadata Metadata `json:"metadata,omitempty"`

	ReferencedBy []couchdb.DocReference `json:"referenced_by,omitempty"`

	// Cache of the fullpath of the file. Should not have to be invalidated
	// since we use FileDoc as immutable data-structures.
	fullpath string

	// NOTE: Do not forget to propagate changes made to this structure to the
	// structure DirOrFileDoc in pkg/vfs/vfs.go.
}

// ID returns the file qualified identifier
func (f *FileDoc) ID() string { return f.DocID }

// Rev returns the file revision
func (f *FileDoc) Rev() string { return f.DocRev }

// DocType returns the file document type
func (f *FileDoc) DocType() string { return consts.Files }

// Clone implements couchdb.Doc
func (f *FileDoc) Clone() couchdb.Doc {
	cloned := *f
	cloned.MD5Sum = make([]byte, len(f.MD5Sum))
	copy(cloned.MD5Sum, f.MD5Sum)
	cloned.Tags = make([]string, len(f.Tags))
	copy(cloned.Tags, f.Tags)
	cloned.ReferencedBy = make([]couchdb.DocReference, len(f.ReferencedBy))
	copy(cloned.ReferencedBy, f.ReferencedBy)
	return &cloned
}

// SetID changes the file qualified identifier
func (f *FileDoc) SetID(id string) { f.DocID = id }

// SetRev changes the file revision
func (f *FileDoc) SetRev(rev string) { f.DocRev = rev }

// Path is used to generate the file path
func (f *FileDoc) Path(fs VFS) (string, error) {
	if f.fullpath != "" {
		return f.fullpath, nil
	}
	var err error
	f.fullpath, err = fs.FilePath(f)
	return f.fullpath, err
}

// Parent returns the parent directory document
func (f *FileDoc) Parent(fs VFS) (*DirDoc, error) {
	parent, err := fs.DirByID(f.DirID)
	if os.IsNotExist(err) {
		err = ErrParentDoesNotExist
	}
	return parent, err
}

// Name returns base name of the file
func (f *FileDoc) Name() string { return f.DocName }

// Size returns the length in bytes for regular files; system-dependent for others
func (f *FileDoc) Size() int64 { return f.ByteSize }

// Mode returns the file mode bits
func (f *FileDoc) Mode() os.FileMode { return getFileMode(f.Executable) }

// ModTime returns the modification time
func (f *FileDoc) ModTime() time.Time { return f.UpdatedAt }

// IsDir returns the abbreviation for Mode().IsDir()
func (f *FileDoc) IsDir() bool { return false }

// Sys returns the underlying data source (can return nil)
func (f *FileDoc) Sys() interface{} { return nil }

// AddReferencedBy adds referenced_by to the file
func (f *FileDoc) AddReferencedBy(ri ...couchdb.DocReference) {
	f.ReferencedBy = append(f.ReferencedBy, ri...)
}

func containsReferencedBy(haystack []couchdb.DocReference, needle couchdb.DocReference) bool {
	for _, ref := range haystack {
		if ref.ID == needle.ID && ref.Type == needle.Type {
			return true
		}
	}
	return false
}

// RemoveReferencedBy adds referenced_by to the file
func (f *FileDoc) RemoveReferencedBy(ri ...couchdb.DocReference) {
	// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
	referenced := f.ReferencedBy[:0]
	for _, ref := range f.ReferencedBy {
		if !containsReferencedBy(ri, ref) {
			referenced = append(referenced, ref)
		}
	}
	f.ReferencedBy = referenced
}

// NewFileDoc is the FileDoc constructor. The given name is validated.
func NewFileDoc(name, dirID string, size int64, md5Sum []byte, mime, class string, cdate time.Time, executable, trashed bool, tags []string) (*FileDoc, error) {
	if err := checkFileName(name); err != nil {
		return nil, err
	}

	if dirID == "" {
		dirID = consts.RootDirID
	}

	tags = uniqueTags(tags)

	doc := &FileDoc{
		Type:    consts.FileType,
		DocName: name,
		DirID:   dirID,

		CreatedAt:  cdate,
		UpdatedAt:  cdate,
		ByteSize:   size,
		MD5Sum:     md5Sum,
		Mime:       mime,
		Class:      class,
		Executable: executable,
		Trashed:    trashed,
		Tags:       tags,
	}

	return doc, nil
}

// ServeFileContent replies to a http request using the content of a
// file given its FileDoc.
//
// It uses internally http.ServeContent and benefits from it by
// offering support to Range, If-Modified-Since and If-None-Match
// requests. It uses the revision of the file as the Etag value for
// non-ranged requests
//
// The content disposition is inlined.
func ServeFileContent(fs VFS, doc *FileDoc, disposition string, req *http.Request, w http.ResponseWriter) error {
	header := w.Header()
	header.Set("Content-Type", doc.Mime)
	if disposition != "" {
		header.Set("Content-Disposition", ContentDisposition(disposition, doc.DocName))
	}

	if header.Get("Range") == "" {
		eTag := base64.StdEncoding.EncodeToString(doc.MD5Sum)
		header.Set("Etag", eTag)
	}

	content, err := fs.OpenFile(doc)
	if err != nil {
		return err
	}
	defer content.Close()

	http.ServeContent(w, req, doc.DocName, doc.UpdatedAt, content)
	return nil
}

// ModifyFileMetadata modify the metadata associated to a file. It can
// be used to rename or move the file in the VFS.
func ModifyFileMetadata(fs VFS, olddoc *FileDoc, patch *DocPatch) (*FileDoc, error) {
	var err error
	rename := patch.Name != nil
	cdate := olddoc.CreatedAt
	oname := olddoc.DocName
	trashed := olddoc.Trashed
	if patch.RestorePath != nil {
		trashed = *patch.RestorePath != ""
	}
	patch, err = normalizeDocPatch(&DocPatch{
		Name:        &oname,
		DirID:       &olddoc.DirID,
		RestorePath: &olddoc.RestorePath,
		Tags:        &olddoc.Tags,
		UpdatedAt:   &olddoc.UpdatedAt,
		Executable:  &olddoc.Executable,
	}, patch, cdate)
	if err != nil {
		return nil, err
	}

	// in case of a renaming of the file, if the extension of the file has
	// changed, we consider recalculating the mime and class attributes, using
	// the new extension.
	newname := *patch.Name
	oldname := olddoc.DocName
	var mime, class string
	if rename && path.Ext(newname) != path.Ext(oldname) {
		mime, class = ExtractMimeAndClassFromFilename(newname)
	} else {
		mime, class = olddoc.Mime, olddoc.Class
	}

	newdoc, err := NewFileDoc(
		newname,
		*patch.DirID,
		olddoc.Size(),
		olddoc.MD5Sum,
		mime,
		class,
		cdate,
		*patch.Executable,
		trashed,
		*patch.Tags,
	)
	if err != nil {
		return nil, err
	}

	newdoc.RestorePath = *patch.RestorePath
	newdoc.UpdatedAt = *patch.UpdatedAt
	newdoc.Metadata = olddoc.Metadata
	newdoc.ReferencedBy = olddoc.ReferencedBy

	if err = fs.UpdateFileDoc(olddoc, newdoc); err != nil {
		return nil, err
	}
	return newdoc, nil
}

// TrashFile is used to delete a file given its document
func TrashFile(fs VFS, olddoc *FileDoc) (*FileDoc, error) {
	oldpath, err := olddoc.Path(fs)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(oldpath, TrashDirName) {
		return nil, ErrFileInTrash
	}

	trashDirID := consts.TrashDirID
	restorePath := path.Dir(oldpath)

	var newdoc *FileDoc
	tryOrUseSuffix(olddoc.DocName, conflictFormat, func(name string) error {
		newdoc, err = ModifyFileMetadata(fs, olddoc, &DocPatch{
			DirID:       &trashDirID,
			RestorePath: &restorePath,
			Name:        &name,
		})
		return err
	})
	return newdoc, err
}

// RestoreFile is used to restore a trashed file given its document
func RestoreFile(fs VFS, olddoc *FileDoc) (*FileDoc, error) {
	oldpath, err := olddoc.Path(fs)
	if err != nil {
		return nil, err
	}
	restoreDir, err := getRestoreDir(fs, oldpath, olddoc.RestorePath)
	if err != nil {
		return nil, err
	}
	var newdoc *FileDoc
	var emptyStr string
	name := stripSuffix(olddoc.DocName, conflictSuffix)
	tryOrUseSuffix(name, "%s (%s)", func(name string) error {
		newdoc, err = ModifyFileMetadata(fs, olddoc, &DocPatch{
			DirID:       &restoreDir.DocID,
			RestorePath: &emptyStr,
			Name:        &name,
		})
		return err
	})
	return newdoc, err
}

func getFileMode(executable bool) os.FileMode {
	if executable {
		return 0755 // -rwxr-xr-x
	}
	return 0644 // -rw-r--r--
}

// ThumbPath gives the path to a thumbnail in the Thumbs hidden FS
func ThumbPath(img *FileDoc, format string) string {
	dir := ThumbDir(img)
	name := fmt.Sprintf("%s-%s.jpg", img.ID(), format)
	return path.Join(dir, name)
}

// ThumbDir gives the path to the parent directory of a thumbnail in the
// Thumbs hidden FS
func ThumbDir(img *FileDoc) string {
	return img.ID()[:4]
}

var (
	_ couchdb.Doc = &FileDoc{}
	_ os.FileInfo = &FileDoc{}
)
<|endoftext|>