<commit_before>// Copyright 2018 The OpenSDS 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 connector

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

// ExecCmd Log and convert the result of exec.Command
func ExecCmd(name string, arg ...string) (string, error) {
	log.Printf("Command: %s %s:\n", name, strings.Join(arg, " "))
	info, err := exec.Command(name, arg...).CombinedOutput()
	return string(info), err
}

// GetFSType returns the File System Type of device
func GetFSType(device string) (string, error) {
	log.Printf("GetFSType: %s\n", device)

	var fsType string
	blkidCmd := "blkid"
	out, err := ExecCmd("blkid", device)
	if err != nil {
		log.Printf("failed to GetFSType: %v cmd: %s output: %s\n",
			err, blkidCmd, string(out))
		return fsType, nil
	}

	for _, v := range strings.Split(string(out), " ") {
		if strings.Contains(v, "TYPE=") {
			fsType = strings.Split(v, "=")[1]
			fsType = strings.Replace(fsType, "\"", "", -1)
			fsType = strings.Replace(fsType, "\n", "", -1)
			fsType = strings.Replace(fsType, "\r", "", -1)
			return fsType, nil
		}
	}

	return fsType, nil
}

// Format device by File System Type
func Format(device string, fsType string) error {
	log.Printf("Format device: %s fstype: %s\n", device, fsType)

	mkfsCmd := fmt.Sprintf("mkfs.%s", fsType)

	_, err := exec.LookPath(mkfsCmd)
	if err != nil {
		if err == exec.ErrNotFound {
			return fmt.Errorf("%q executable not found in $PATH", mkfsCmd)
		}
		return err
	}

	mkfsArgs := []string{}
	mkfsArgs = append(mkfsArgs, device)
	if fsType == "ext4" || fsType == "ext3" {
		mkfsArgs = []string{"-F", device}
	}

	out, err := ExecCmd(mkfsCmd, mkfsArgs...)
	if err != nil {
		return fmt.Errorf("formatting disk failed: %v cmd: '%s %s' output: %q",
			err, mkfsCmd, strings.Join(mkfsArgs, " "), string(out))
	}

	return nil
}

// Mount device into mount point
func Mount(device, mountpoint, fsType string, mountFlags []string) error {
	log.Printf("Mount device: %s mountpoint: %s, fsType: %s, mountFlags:　%v\n", device, mountpoint, fsType, mountFlags)

	_, err := ExecCmd("mkdir", "-p", mountpoint)
	if err != nil {
		log.Printf("failed to mkdir: %v\n", err)
		return err
	}

	mountArgs := []string{}

	mountArgs = append(mountArgs, "-t", fsType)

	if len(mountFlags) > 0 {
		mountArgs = append(mountArgs, "-o", strings.Join(mountFlags, ","))
	}

	mountArgs = append(mountArgs, device)
	mountArgs = append(mountArgs, mountpoint)

	_, err = exec.Command("mount", mountArgs...).CombinedOutput()
	if err != nil {
		log.Printf("failed to mount: %v\n", err)
		return err
	}

	// Make sure the mount is not lost after the host reboots
	cmd := fmt.Sprintf("echo \"%s %s %s defaults 0 0\" >> /etc/fstab", device, mountpoint, fsType)
	_, err = ExecCmd("/bin/bash", "-c", cmd)
	if err != nil {
		return err
	}

	return nil
}

// Umount from mountpoint
func Umount(mountpoint string) error {
	log.Printf("Umount mountpoint: %s\n", mountpoint)

	_, err := ExecCmd("umount", mountpoint)
	if err != nil {
		log.Printf("failed to Umount: %v\n", err)
		return err
	}

	cmd := fmt.Sprintf("cat -n /etc/fstab | grep -w '%s' | awk -F ' ' '{ print $1 }'| xargs -i sed -i '{}d' /etc/fstab", mountpoint)
	_, err = ExecCmd("/bin/bash", "-c", cmd)
	if err != nil {
		return err
	}
	return nil
}

// GetHostIP return Host IP
func GetHostIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return "127.0.0.1"
	}

	for _, address := range addrs {
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			return ipnet.IP.String()
		}
	}

	return "127.0.0.1"
}

// GetHostName ...
func GetHostName() (string, error) {
	hostName, err := ExecCmd("hostname")
	if err != nil {
		log.Printf("failed to get host name: %v\n", err)
		return "", err
	}
	hostName = strings.Replace(hostName, "\n", "", -1)
	return hostName, nil
}

// IsMounted ...
func IsMounted(target string) (bool, error) {
	findmntCmd := "findmnt"
	_, err := exec.LookPath(findmntCmd)
	if err != nil {
		if err == exec.ErrNotFound {
			msg := fmt.Sprintf("%s executable not found in $PATH, err: %v\n", findmntCmd, err)
			log.Printf(msg)
			return false, fmt.Errorf(msg)
		}
		log.Printf("failed to check IsMounted %v\n", err)
		return false, err
	}

	findmntArgs := []string{"--target", target}

	log.Printf("findmnt args is %s\n", findmntArgs)

	out, err := ExecCmd(findmntCmd, findmntArgs...)
	if err != nil {
		// findmnt exits with non zero exit status if it couldn't find anything
		if strings.TrimSpace(string(out)) == "" {
			return false, nil
		}

		errIsMounted := fmt.Errorf("checking mounted failed: %v cmd: %s output: %s",
			err, findmntCmd, string(out))

		log.Printf("checking mounted failed: %v\n", errIsMounted)
		return false, errIsMounted
	}

	log.Printf("checking mounted result is %s\n", strings.TrimSpace(string(out)))
	if strings.TrimSpace(string(out)) == "" {
		return false, nil
	}

	line := strings.Split(string(out), "\n")

	if strings.Split(line[1], " ")[0] != target {
		return false, nil
	}

	return true, nil
}
<commit_msg>Remove fstab file writing code (#20)<commit_after>// Copyright 2018 The OpenSDS 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 connector

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

// ExecCmd Log and convert the result of exec.Command
func ExecCmd(name string, arg ...string) (string, error) {
	log.Printf("Command: %s %s:\n", name, strings.Join(arg, " "))
	info, err := exec.Command(name, arg...).CombinedOutput()
	return string(info), err
}

// GetFSType returns the File System Type of device
func GetFSType(device string) (string, error) {
	log.Printf("GetFSType: %s\n", device)

	var fsType string
	blkidCmd := "blkid"
	out, err := ExecCmd("blkid", device)
	if err != nil {
		log.Printf("failed to GetFSType: %v cmd: %s output: %s\n",
			err, blkidCmd, string(out))
		return fsType, nil
	}

	for _, v := range strings.Split(string(out), " ") {
		if strings.Contains(v, "TYPE=") {
			fsType = strings.Split(v, "=")[1]
			fsType = strings.Replace(fsType, "\"", "", -1)
			fsType = strings.Replace(fsType, "\n", "", -1)
			fsType = strings.Replace(fsType, "\r", "", -1)
			return fsType, nil
		}
	}

	return fsType, nil
}

// Format device by File System Type
func Format(device string, fsType string) error {
	log.Printf("Format device: %s fstype: %s\n", device, fsType)

	mkfsCmd := fmt.Sprintf("mkfs.%s", fsType)

	_, err := exec.LookPath(mkfsCmd)
	if err != nil {
		if err == exec.ErrNotFound {
			return fmt.Errorf("%q executable not found in $PATH", mkfsCmd)
		}
		return err
	}

	mkfsArgs := []string{}
	mkfsArgs = append(mkfsArgs, device)
	if fsType == "ext4" || fsType == "ext3" {
		mkfsArgs = []string{"-F", device}
	}

	out, err := ExecCmd(mkfsCmd, mkfsArgs...)
	if err != nil {
		return fmt.Errorf("formatting disk failed: %v cmd: '%s %s' output: %q",
			err, mkfsCmd, strings.Join(mkfsArgs, " "), string(out))
	}

	return nil
}

// Mount device into mount point
func Mount(device, mountpoint, fsType string, mountFlags []string) error {
	log.Printf("Mount device: %s mountpoint: %s, fsType: %s, mountFlags:　%v\n", device, mountpoint, fsType, mountFlags)

	_, err := ExecCmd("mkdir", "-p", mountpoint)
	if err != nil {
		log.Printf("failed to mkdir: %v\n", err)
		return err
	}

	mountArgs := []string{}

	mountArgs = append(mountArgs, "-t", fsType)

	if len(mountFlags) > 0 {
		mountArgs = append(mountArgs, "-o", strings.Join(mountFlags, ","))
	}

	mountArgs = append(mountArgs, device)
	mountArgs = append(mountArgs, mountpoint)

	_, err = exec.Command("mount", mountArgs...).CombinedOutput()
	if err != nil {
		log.Printf("failed to mount: %v\n", err)
		return err
	}

	return nil
}

// Umount from mountpoint
func Umount(mountpoint string) error {
	log.Printf("Umount mountpoint: %s\n", mountpoint)

	_, err := ExecCmd("umount", mountpoint)
	if err != nil {
		log.Printf("failed to Umount: %v\n", err)
		return err
	}

	return nil
}

// GetHostIP return Host IP
func GetHostIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return "127.0.0.1"
	}

	for _, address := range addrs {
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			return ipnet.IP.String()
		}
	}

	return "127.0.0.1"
}

// GetHostName ...
func GetHostName() (string, error) {
	hostName, err := ExecCmd("hostname")
	if err != nil {
		log.Printf("failed to get host name: %v\n", err)
		return "", err
	}
	hostName = strings.Replace(hostName, "\n", "", -1)
	return hostName, nil
}

// IsMounted ...
func IsMounted(target string) (bool, error) {
	findmntCmd := "findmnt"
	_, err := exec.LookPath(findmntCmd)
	if err != nil {
		if err == exec.ErrNotFound {
			msg := fmt.Sprintf("%s executable not found in $PATH, err: %v\n", findmntCmd, err)
			log.Printf(msg)
			return false, fmt.Errorf(msg)
		}
		log.Printf("failed to check IsMounted %v\n", err)
		return false, err
	}

	findmntArgs := []string{"--target", target}

	log.Printf("findmnt args is %s\n", findmntArgs)

	out, err := ExecCmd(findmntCmd, findmntArgs...)
	if err != nil {
		// findmnt exits with non zero exit status if it couldn't find anything
		if strings.TrimSpace(string(out)) == "" {
			return false, nil
		}

		errIsMounted := fmt.Errorf("checking mounted failed: %v cmd: %s output: %s",
			err, findmntCmd, string(out))

		log.Printf("checking mounted failed: %v\n", errIsMounted)
		return false, errIsMounted
	}

	log.Printf("checking mounted result is %s\n", strings.TrimSpace(string(out)))
	if strings.TrimSpace(string(out)) == "" {
		return false, nil
	}

	line := strings.Split(string(out), "\n")

	if strings.Split(line[1], " ")[0] != target {
		return false, nil
	}

	return true, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 cm

import (
	"flag"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/uuid"
	pluginapi "k8s.io/kubernetes/pkg/kubelet/apis/deviceplugin/v1alpha"
)

func TestUpdateCapacity(t *testing.T) {
	var expected = v1.ResourceList{}
	as := assert.New(t)
	verifyCapacityFunc := func(updates v1.ResourceList) {
		as.Equal(expected, updates)
	}
	testDevicePluginHandler, err := NewDevicePluginHandlerImpl(verifyCapacityFunc)
	as.NotNil(testDevicePluginHandler)
	as.Nil(err)

	devs := []*pluginapi.Device{
		{ID: "Device1", Health: pluginapi.Healthy},
		{ID: "Device2", Health: pluginapi.Healthy},
		{ID: "Device3", Health: pluginapi.Unhealthy},
	}

	resourceName := "resource1"
	// Adds three devices for resource1, two healthy and one unhealthy.
	// Expects capacity for resource1 to be 2.
	expected[v1.ResourceName(resourceName)] = *resource.NewQuantity(int64(2), resource.DecimalSI)
	testDevicePluginHandler.devicePluginManagerMonitorCallback(resourceName, devs, []*pluginapi.Device{}, []*pluginapi.Device{})
	// Deletes an unhealthy device should NOT change capacity.
	testDevicePluginHandler.devicePluginManagerMonitorCallback(resourceName, []*pluginapi.Device{}, []*pluginapi.Device{}, []*pluginapi.Device{devs[2]})
	// Updates a healthy device to unhealthy should reduce capacity by 1.
	expected[v1.ResourceName(resourceName)] = *resource.NewQuantity(int64(1), resource.DecimalSI)
	// Deletes a healthy device should reduce capacity by 1.
	expected[v1.ResourceName(resourceName)] = *resource.NewQuantity(int64(0), resource.DecimalSI)
	// Tests adding another resource.
	delete(expected, v1.ResourceName(resourceName))
	resourceName2 := "resource2"
	expected[v1.ResourceName(resourceName2)] = *resource.NewQuantity(int64(2), resource.DecimalSI)
	testDevicePluginHandler.devicePluginManagerMonitorCallback(resourceName2, devs, []*pluginapi.Device{}, []*pluginapi.Device{})
}

type stringPairType struct {
	value1 string
	value2 string
}

// DevicePluginManager stub to test device Allocation behavior.
type DevicePluginManagerTestStub struct {
	// All data structs are keyed by resourceName+DevId
	devRuntimeDevices map[string][]stringPairType
	devRuntimeMounts  map[string][]stringPairType
	devRuntimeEnvs    map[string][]stringPairType
}

func NewDevicePluginManagerTestStub() (*DevicePluginManagerTestStub, error) {
	return &DevicePluginManagerTestStub{
		devRuntimeDevices: make(map[string][]stringPairType),
		devRuntimeMounts:  make(map[string][]stringPairType),
		devRuntimeEnvs:    make(map[string][]stringPairType),
	}, nil
}

func (m *DevicePluginManagerTestStub) Start() error {
	return nil
}

func (m *DevicePluginManagerTestStub) Devices() map[string][]*pluginapi.Device {
	return make(map[string][]*pluginapi.Device)
}

func (m *DevicePluginManagerTestStub) Allocate(resourceName string, devIds []string) (*pluginapi.AllocateResponse, error) {
	resp := new(pluginapi.AllocateResponse)
	for _, id := range devIds {
		key := resourceName + id
		fmt.Printf("Alloc device %q for resource %q\n", id, resourceName)
		for _, dev := range m.devRuntimeDevices[key] {
			resp.Devices = append(resp.Devices, &pluginapi.DeviceSpec{
				ContainerPath: dev.value1,
				HostPath:      dev.value2,
				Permissions:   "mrw",
			})
		}
		for _, mount := range m.devRuntimeMounts[key] {
			fmt.Printf("Add mount %q %q\n", mount.value1, mount.value2)
			resp.Mounts = append(resp.Mounts, &pluginapi.Mount{
				ContainerPath: mount.value1,
				HostPath:      mount.value2,
				ReadOnly:      true,
			})
		}
		resp.Envs = make(map[string]string)
		for _, env := range m.devRuntimeEnvs[key] {
			resp.Envs[env.value1] = env.value2
		}
	}
	return resp, nil
}

func (m *DevicePluginManagerTestStub) Stop() error {
	return nil
}

func (m *DevicePluginManagerTestStub) CheckpointFile() string {
	return "/tmp/device-plugin-checkpoint"
}

func TestCheckpoint(t *testing.T) {
	resourceName1 := "domain1.com/resource1"
	resourceName2 := "domain2.com/resource2"

	m, err := NewDevicePluginManagerTestStub()
	as := assert.New(t)
	as.Nil(err)

	testDevicePluginHandler := &DevicePluginHandlerImpl{
		devicePluginManager: m,
		allDevices:          make(map[string]sets.String),
		allocatedDevices:    make(map[string]podDevices),
	}
	testDevicePluginHandler.allocatedDevices[resourceName1] = make(podDevices)
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod1", "con1", "dev1")
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod1", "con1", "dev2")
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod1", "con2", "dev1")
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod2", "con1", "dev1")
	testDevicePluginHandler.allocatedDevices[resourceName2] = make(podDevices)
	testDevicePluginHandler.allocatedDevices[resourceName2].insert("pod1", "con1", "dev3")
	testDevicePluginHandler.allocatedDevices[resourceName2].insert("pod1", "con1", "dev4")

	err = testDevicePluginHandler.writeCheckpoint()
	as.Nil(err)
	expected := testDevicePluginHandler.allocatedDevices
	testDevicePluginHandler.allocatedDevices = make(map[string]podDevices)
	err = testDevicePluginHandler.readCheckpoint()
	as.Nil(err)
	as.Equal(expected, testDevicePluginHandler.allocatedDevices)
}

func TestPodContainerDeviceAllocation(t *testing.T) {
	flag.Set("alsologtostderr", fmt.Sprintf("%t", true))
	var logLevel string
	flag.StringVar(&logLevel, "logLevel", "4", "test")
	flag.Lookup("v").Value.Set(logLevel)

	var activePods []*v1.Pod
	resourceName1 := "domain1.com/resource1"
	resourceQuantity1 := *resource.NewQuantity(int64(2), resource.DecimalSI)
	devId1 := "dev1"
	devId2 := "dev2"
	resourceName2 := "domain2.com/resource2"
	resourceQuantity2 := *resource.NewQuantity(int64(1), resource.DecimalSI)
	devId3 := "dev3"
	devId4 := "dev4"

	m, err := NewDevicePluginManagerTestStub()
	as := assert.New(t)
	as.Nil(err)
	monitorCallback := func(resourceName string, added, updated, deleted []*pluginapi.Device) {}

	testDevicePluginHandler := &DevicePluginHandlerImpl{
		devicePluginManager:                m,
		devicePluginManagerMonitorCallback: monitorCallback,
		allDevices:                         make(map[string]sets.String),
		allocatedDevices:                   make(map[string]podDevices),
	}
	testDevicePluginHandler.allDevices[resourceName1] = sets.NewString()
	testDevicePluginHandler.allDevices[resourceName1].Insert(devId1)
	testDevicePluginHandler.allDevices[resourceName1].Insert(devId2)
	testDevicePluginHandler.allDevices[resourceName2] = sets.NewString()
	testDevicePluginHandler.allDevices[resourceName2].Insert(devId3)
	testDevicePluginHandler.allDevices[resourceName2].Insert(devId4)

	m.devRuntimeDevices[resourceName1+devId1] = append(m.devRuntimeDevices[resourceName1+devId1], stringPairType{"/dev/aaa", "/dev/aaa"})
	m.devRuntimeDevices[resourceName1+devId1] = append(m.devRuntimeDevices[resourceName1+devId1], stringPairType{"/dev/bbb", "/dev/bbb"})
	m.devRuntimeDevices[resourceName1+devId2] = append(m.devRuntimeDevices[resourceName1+devId2], stringPairType{"/dev/ccc", "/dev/ccc"})
	m.devRuntimeMounts[resourceName1+devId1] = append(m.devRuntimeMounts[resourceName1+devId1], stringPairType{"/container_dir1/file1", "host_dir1/file1"})
	m.devRuntimeMounts[resourceName1+devId2] = append(m.devRuntimeMounts[resourceName1+devId2], stringPairType{"/container_dir1/file1", "host_dir1/file1"})
	m.devRuntimeEnvs[resourceName1+devId2] = append(m.devRuntimeEnvs[resourceName1+devId2], stringPairType{"key1", "val1"})
	m.devRuntimeEnvs[resourceName2+devId3] = append(m.devRuntimeEnvs[resourceName2+devId3], stringPairType{"key2", "val2"})
	m.devRuntimeEnvs[resourceName2+devId4] = append(m.devRuntimeEnvs[resourceName2+devId4], stringPairType{"key2", "val2"})

	pod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			UID: uuid.NewUUID(),
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name: string(uuid.NewUUID()),
					Resources: v1.ResourceRequirements{
						Limits: v1.ResourceList{
							v1.ResourceName(resourceName1): resourceQuantity1,
							v1.ResourceName("cpu"):         resourceQuantity1,
							v1.ResourceName(resourceName2): resourceQuantity2,
						},
					},
				},
			},
		},
	}

	cm := &containerManagerImpl{
		devicePluginHandler: testDevicePluginHandler,
	}
	activePods = append(activePods, pod)
	runContainerOpts, err := cm.GetResources(pod, &pod.Spec.Containers[0], activePods)
	as.Equal(len(runContainerOpts.Devices), 3)
	// Two devices require to mount the same path. Expects a single mount entry to be created.
	as.Equal(len(runContainerOpts.Mounts), 1)
	as.Equal(runContainerOpts.Mounts[0].ContainerPath, "/container_dir1/file1")
	as.Equal(len(runContainerOpts.Envs), 2)

	// Requesting to create a pod without enough resources should fail.
	failPod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			UID: uuid.NewUUID(),
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name: string(uuid.NewUUID()),
					Resources: v1.ResourceRequirements{
						Limits: v1.ResourceList{
							v1.ResourceName(resourceName1): resourceQuantity1,
						},
					},
				},
			},
		},
	}
	runContainerOpts2, err := cm.GetResources(failPod, &failPod.Spec.Containers[0], activePods)
	as.NotNil(err)
	as.Equal(len(runContainerOpts2.Devices), 0)
	as.Equal(len(runContainerOpts2.Mounts), 0)
	as.Equal(len(runContainerOpts2.Envs), 0)

	// Requesting to create a new pod with a single resourceName2 should succeed.
	newPod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			UID: uuid.NewUUID(),
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name: string(uuid.NewUUID()),
					Resources: v1.ResourceRequirements{
						Limits: v1.ResourceList{
							v1.ResourceName(resourceName2): resourceQuantity2,
						},
					},
				},
			},
		},
	}
	runContainerOpts3, err := cm.GetResources(newPod, &newPod.Spec.Containers[0], activePods)
	as.Nil(err)
	as.Equal(len(runContainerOpts3.Envs), 1)
}
<commit_msg>Disabling k8s.io/kubernetes/pkg/kubelet/cm TestPodContainerDeviceAllocation due to #54100<commit_after>/*
Copyright 2017 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 cm

import (
	"flag"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/uuid"
	pluginapi "k8s.io/kubernetes/pkg/kubelet/apis/deviceplugin/v1alpha"
)

func TestUpdateCapacity(t *testing.T) {
	var expected = v1.ResourceList{}
	as := assert.New(t)
	verifyCapacityFunc := func(updates v1.ResourceList) {
		as.Equal(expected, updates)
	}
	testDevicePluginHandler, err := NewDevicePluginHandlerImpl(verifyCapacityFunc)
	as.NotNil(testDevicePluginHandler)
	as.Nil(err)

	devs := []*pluginapi.Device{
		{ID: "Device1", Health: pluginapi.Healthy},
		{ID: "Device2", Health: pluginapi.Healthy},
		{ID: "Device3", Health: pluginapi.Unhealthy},
	}

	resourceName := "resource1"
	// Adds three devices for resource1, two healthy and one unhealthy.
	// Expects capacity for resource1 to be 2.
	expected[v1.ResourceName(resourceName)] = *resource.NewQuantity(int64(2), resource.DecimalSI)
	testDevicePluginHandler.devicePluginManagerMonitorCallback(resourceName, devs, []*pluginapi.Device{}, []*pluginapi.Device{})
	// Deletes an unhealthy device should NOT change capacity.
	testDevicePluginHandler.devicePluginManagerMonitorCallback(resourceName, []*pluginapi.Device{}, []*pluginapi.Device{}, []*pluginapi.Device{devs[2]})
	// Updates a healthy device to unhealthy should reduce capacity by 1.
	expected[v1.ResourceName(resourceName)] = *resource.NewQuantity(int64(1), resource.DecimalSI)
	// Deletes a healthy device should reduce capacity by 1.
	expected[v1.ResourceName(resourceName)] = *resource.NewQuantity(int64(0), resource.DecimalSI)
	// Tests adding another resource.
	delete(expected, v1.ResourceName(resourceName))
	resourceName2 := "resource2"
	expected[v1.ResourceName(resourceName2)] = *resource.NewQuantity(int64(2), resource.DecimalSI)
	testDevicePluginHandler.devicePluginManagerMonitorCallback(resourceName2, devs, []*pluginapi.Device{}, []*pluginapi.Device{})
}

type stringPairType struct {
	value1 string
	value2 string
}

// DevicePluginManager stub to test device Allocation behavior.
type DevicePluginManagerTestStub struct {
	// All data structs are keyed by resourceName+DevId
	devRuntimeDevices map[string][]stringPairType
	devRuntimeMounts  map[string][]stringPairType
	devRuntimeEnvs    map[string][]stringPairType
}

func NewDevicePluginManagerTestStub() (*DevicePluginManagerTestStub, error) {
	return &DevicePluginManagerTestStub{
		devRuntimeDevices: make(map[string][]stringPairType),
		devRuntimeMounts:  make(map[string][]stringPairType),
		devRuntimeEnvs:    make(map[string][]stringPairType),
	}, nil
}

func (m *DevicePluginManagerTestStub) Start() error {
	return nil
}

func (m *DevicePluginManagerTestStub) Devices() map[string][]*pluginapi.Device {
	return make(map[string][]*pluginapi.Device)
}

func (m *DevicePluginManagerTestStub) Allocate(resourceName string, devIds []string) (*pluginapi.AllocateResponse, error) {
	resp := new(pluginapi.AllocateResponse)
	for _, id := range devIds {
		key := resourceName + id
		fmt.Printf("Alloc device %q for resource %q\n", id, resourceName)
		for _, dev := range m.devRuntimeDevices[key] {
			resp.Devices = append(resp.Devices, &pluginapi.DeviceSpec{
				ContainerPath: dev.value1,
				HostPath:      dev.value2,
				Permissions:   "mrw",
			})
		}
		for _, mount := range m.devRuntimeMounts[key] {
			fmt.Printf("Add mount %q %q\n", mount.value1, mount.value2)
			resp.Mounts = append(resp.Mounts, &pluginapi.Mount{
				ContainerPath: mount.value1,
				HostPath:      mount.value2,
				ReadOnly:      true,
			})
		}
		resp.Envs = make(map[string]string)
		for _, env := range m.devRuntimeEnvs[key] {
			resp.Envs[env.value1] = env.value2
		}
	}
	return resp, nil
}

func (m *DevicePluginManagerTestStub) Stop() error {
	return nil
}

func (m *DevicePluginManagerTestStub) CheckpointFile() string {
	return "/tmp/device-plugin-checkpoint"
}

func TestCheckpoint(t *testing.T) {
	resourceName1 := "domain1.com/resource1"
	resourceName2 := "domain2.com/resource2"

	m, err := NewDevicePluginManagerTestStub()
	as := assert.New(t)
	as.Nil(err)

	testDevicePluginHandler := &DevicePluginHandlerImpl{
		devicePluginManager: m,
		allDevices:          make(map[string]sets.String),
		allocatedDevices:    make(map[string]podDevices),
	}
	testDevicePluginHandler.allocatedDevices[resourceName1] = make(podDevices)
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod1", "con1", "dev1")
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod1", "con1", "dev2")
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod1", "con2", "dev1")
	testDevicePluginHandler.allocatedDevices[resourceName1].insert("pod2", "con1", "dev1")
	testDevicePluginHandler.allocatedDevices[resourceName2] = make(podDevices)
	testDevicePluginHandler.allocatedDevices[resourceName2].insert("pod1", "con1", "dev3")
	testDevicePluginHandler.allocatedDevices[resourceName2].insert("pod1", "con1", "dev4")

	err = testDevicePluginHandler.writeCheckpoint()
	as.Nil(err)
	expected := testDevicePluginHandler.allocatedDevices
	testDevicePluginHandler.allocatedDevices = make(map[string]podDevices)
	err = testDevicePluginHandler.readCheckpoint()
	as.Nil(err)
	as.Equal(expected, testDevicePluginHandler.allocatedDevices)
}

func TestPodContainerDeviceAllocation(t *testing.T) {
	t.Skip("Skipping due to issue #54100")
	flag.Set("alsologtostderr", fmt.Sprintf("%t", true))
	var logLevel string
	flag.StringVar(&logLevel, "logLevel", "4", "test")
	flag.Lookup("v").Value.Set(logLevel)

	var activePods []*v1.Pod
	resourceName1 := "domain1.com/resource1"
	resourceQuantity1 := *resource.NewQuantity(int64(2), resource.DecimalSI)
	devId1 := "dev1"
	devId2 := "dev2"
	resourceName2 := "domain2.com/resource2"
	resourceQuantity2 := *resource.NewQuantity(int64(1), resource.DecimalSI)
	devId3 := "dev3"
	devId4 := "dev4"

	m, err := NewDevicePluginManagerTestStub()
	as := assert.New(t)
	as.Nil(err)
	monitorCallback := func(resourceName string, added, updated, deleted []*pluginapi.Device) {}

	testDevicePluginHandler := &DevicePluginHandlerImpl{
		devicePluginManager:                m,
		devicePluginManagerMonitorCallback: monitorCallback,
		allDevices:                         make(map[string]sets.String),
		allocatedDevices:                   make(map[string]podDevices),
	}
	testDevicePluginHandler.allDevices[resourceName1] = sets.NewString()
	testDevicePluginHandler.allDevices[resourceName1].Insert(devId1)
	testDevicePluginHandler.allDevices[resourceName1].Insert(devId2)
	testDevicePluginHandler.allDevices[resourceName2] = sets.NewString()
	testDevicePluginHandler.allDevices[resourceName2].Insert(devId3)
	testDevicePluginHandler.allDevices[resourceName2].Insert(devId4)

	m.devRuntimeDevices[resourceName1+devId1] = append(m.devRuntimeDevices[resourceName1+devId1], stringPairType{"/dev/aaa", "/dev/aaa"})
	m.devRuntimeDevices[resourceName1+devId1] = append(m.devRuntimeDevices[resourceName1+devId1], stringPairType{"/dev/bbb", "/dev/bbb"})
	m.devRuntimeDevices[resourceName1+devId2] = append(m.devRuntimeDevices[resourceName1+devId2], stringPairType{"/dev/ccc", "/dev/ccc"})
	m.devRuntimeMounts[resourceName1+devId1] = append(m.devRuntimeMounts[resourceName1+devId1], stringPairType{"/container_dir1/file1", "host_dir1/file1"})
	m.devRuntimeMounts[resourceName1+devId2] = append(m.devRuntimeMounts[resourceName1+devId2], stringPairType{"/container_dir1/file1", "host_dir1/file1"})
	m.devRuntimeEnvs[resourceName1+devId2] = append(m.devRuntimeEnvs[resourceName1+devId2], stringPairType{"key1", "val1"})
	m.devRuntimeEnvs[resourceName2+devId3] = append(m.devRuntimeEnvs[resourceName2+devId3], stringPairType{"key2", "val2"})
	m.devRuntimeEnvs[resourceName2+devId4] = append(m.devRuntimeEnvs[resourceName2+devId4], stringPairType{"key2", "val2"})

	pod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			UID: uuid.NewUUID(),
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name: string(uuid.NewUUID()),
					Resources: v1.ResourceRequirements{
						Limits: v1.ResourceList{
							v1.ResourceName(resourceName1): resourceQuantity1,
							v1.ResourceName("cpu"):         resourceQuantity1,
							v1.ResourceName(resourceName2): resourceQuantity2,
						},
					},
				},
			},
		},
	}

	cm := &containerManagerImpl{
		devicePluginHandler: testDevicePluginHandler,
	}
	activePods = append(activePods, pod)
	runContainerOpts, err := cm.GetResources(pod, &pod.Spec.Containers[0], activePods)
	as.Equal(len(runContainerOpts.Devices), 3)
	// Two devices require to mount the same path. Expects a single mount entry to be created.
	as.Equal(len(runContainerOpts.Mounts), 1)
	as.Equal(runContainerOpts.Mounts[0].ContainerPath, "/container_dir1/file1")
	as.Equal(len(runContainerOpts.Envs), 2)

	// Requesting to create a pod without enough resources should fail.
	failPod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			UID: uuid.NewUUID(),
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name: string(uuid.NewUUID()),
					Resources: v1.ResourceRequirements{
						Limits: v1.ResourceList{
							v1.ResourceName(resourceName1): resourceQuantity1,
						},
					},
				},
			},
		},
	}
	runContainerOpts2, err := cm.GetResources(failPod, &failPod.Spec.Containers[0], activePods)
	as.NotNil(err)
	as.Equal(len(runContainerOpts2.Devices), 0)
	as.Equal(len(runContainerOpts2.Mounts), 0)
	as.Equal(len(runContainerOpts2.Envs), 0)

	// Requesting to create a new pod with a single resourceName2 should succeed.
	newPod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			UID: uuid.NewUUID(),
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name: string(uuid.NewUUID()),
					Resources: v1.ResourceRequirements{
						Limits: v1.ResourceList{
							v1.ResourceName(resourceName2): resourceQuantity2,
						},
					},
				},
			},
		},
	}
	runContainerOpts3, err := cm.GetResources(newPod, &newPod.Spec.Containers[0], activePods)
	as.Nil(err)
	as.Equal(len(runContainerOpts3.Envs), 1)
}
<|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.
 *
 */

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

package network

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"regexp"
	"strings"

	"github.com/vishvananda/netlink"

	"kubevirt.io/kubevirt/pkg/log"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/network/dhcp"

	lmf "github.com/subgraph/libmacouflage"
)

type VIF struct {
	Name    string
	IP      netlink.Addr
	MAC     net.HardwareAddr
	Gateway net.IP
	Routes  *[]netlink.Route
	Mtu     uint16
}

type NetworkHandler interface {
	LinkByName(name string) (netlink.Link, error)
	AddrList(link netlink.Link, family int) ([]netlink.Addr, error)
	RouteList(link netlink.Link, family int) ([]netlink.Route, error)
	AddrDel(link netlink.Link, addr *netlink.Addr) error
	AddrAdd(link netlink.Link, addr *netlink.Addr) error
	LinkSetDown(link netlink.Link) error
	LinkSetUp(link netlink.Link) error
	LinkAdd(link netlink.Link) error
	ParseAddr(s string) (*netlink.Addr, error)
	SetRandomMac(iface string) (net.HardwareAddr, error)
	GetMacDetails(iface string) (net.HardwareAddr, error)
	StartDHCP(nic *VIF, serverAddr *netlink.Addr)
}

type NetworkUtilsHandler struct{}

var Handler NetworkHandler

func (h *NetworkUtilsHandler) LinkByName(name string) (netlink.Link, error) {
	return netlink.LinkByName(name)
}
func (h *NetworkUtilsHandler) AddrList(link netlink.Link, family int) ([]netlink.Addr, error) {
	return netlink.AddrList(link, family)
}
func (h *NetworkUtilsHandler) RouteList(link netlink.Link, family int) ([]netlink.Route, error) {
	return netlink.RouteList(link, family)
}
func (h *NetworkUtilsHandler) AddrDel(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrDel(link, addr)
}
func (h *NetworkUtilsHandler) LinkSetDown(link netlink.Link) error {
	return netlink.LinkSetDown(link)
}
func (h *NetworkUtilsHandler) LinkSetUp(link netlink.Link) error {
	return netlink.LinkSetUp(link)
}
func (h *NetworkUtilsHandler) LinkAdd(link netlink.Link) error {
	return netlink.LinkAdd(link)
}
func (h *NetworkUtilsHandler) ParseAddr(s string) (*netlink.Addr, error) {
	return netlink.ParseAddr(s)
}
func (h *NetworkUtilsHandler) AddrAdd(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrAdd(link, addr)
}

// GetMacDetails from an interface
func (h *NetworkUtilsHandler) GetMacDetails(iface string) (net.HardwareAddr, error) {
	currentMac, err := lmf.GetCurrentMac(iface)
	if err != nil {
		log.Log.Reason(err).Errorf("failed to get mac information for interface: %s", iface)
		return nil, err
	}
	return currentMac, nil
}

// SetRandomMac changes the MAC address for a agiven interface to a randomly generated, preserving the vendor prefix
func (h *NetworkUtilsHandler) SetRandomMac(iface string) (net.HardwareAddr, error) {
	var mac net.HardwareAddr

	currentMac, err := Handler.GetMacDetails(iface)
	if err != nil {
		return nil, err
	}

	changed, err := lmf.SpoofMacSameVendor(iface, false)
	if err != nil {
		log.Log.Reason(err).Errorf("failed to spoof MAC for an interface: %s", iface)
		return nil, err
	}

	if changed {
		mac, err = Handler.GetMacDetails(iface)
		if err != nil {
			return nil, err
		}
		log.Log.Reason(err).Errorf("updated MAC for interface: %s - %s", iface, mac)
	}
	return currentMac, nil
}

func (h *NetworkUtilsHandler) StartDHCP(nic *VIF, serverAddr *netlink.Addr) {
	nameservers, searchDomains, err := getResolvConfDetailsFromPod()
	if err != nil {
		log.Log.Errorf("Failed to get DNS servers from resolv.conf: %v", err)
		panic(err)
	}

	// panic in case the DHCP server failed during the vm creation
	// but ignore dhcp errors when the vm is destroyed or shutting down
	go func() {
		if err = DHCPServer(
			nic.MAC,
			nic.IP.IP,
			nic.IP.Mask,
			api.DefaultBridgeName,
			serverAddr.IP,
			nic.Gateway,
			nameservers,
			nic.Routes,
			searchDomains,
			nic.Mtu,
		); err != nil {
			log.Log.Errorf("failed to run DHCP: %v", err)
			panic(err)
		}
	}()
}

// Allow mocking for tests
var SetupPodNetwork = SetupNetworkInterfaces
var DHCPServer = dhcp.SingleClientDHCPServer

func initHandler() {
	if Handler == nil {
		Handler = &NetworkUtilsHandler{}
	}
}

func writeToCachedFile(inter interface{}, fileName, name string) error {
	buf, err := json.MarshalIndent(&inter, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshaling interface cache: %v", err)
	}
	err = ioutil.WriteFile(getInterfaceCacheFile(fileName, name), buf, 0644)
	if err != nil {
		return fmt.Errorf("error writing interface cache %v", err)
	}
	return nil
}

func readFromCachedFile(name, fileName string, inter interface{}) (bool, error) {
	buf, err := ioutil.ReadFile(getInterfaceCacheFile(fileName, name))
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}

	err = json.Unmarshal(buf, &inter)
	if err != nil {
		return false, fmt.Errorf("error unmarshaling command line: %v", err)
	}
	return true, nil
}

func getInterfaceCacheFile(filePath, name string) string {
	return fmt.Sprintf(filePath, name)
}

// func setCachedInterface(name string, ifconf *api.Interface) error {
// 	buf, err := json.MarshalIndent(&ifconf, "", "  ")
// 	if err != nil {
// 		return fmt.Errorf("error marshaling interface cache: %v", err)
// 	}
// 	err = ioutil.WriteFile(getInterfaceCacheFile(interfaceCacheFile, name), buf, 0644)
// 	if err != nil {
// 		return fmt.Errorf("error writing interface cache %v", err)
// 	}
// 	return nil
// }

// func getCachedInterface(name string) (*api.Interface, error) {
// 	buf, err := ioutil.ReadFile(getInterfaceCacheFile(interfaceCacheFile, name))
// 	if err != nil {
// 		if os.IsNotExist(err) {
// 			return nil, nil
// 		}
// 		return nil, err
// 	}
// 	ifconf := api.Interface{}
// 	err = json.Unmarshal(buf, &ifconf)
// 	if err != nil {
// 		return nil, fmt.Errorf("error unmarshaling interface: %v", err)
// 	}
// 	return &ifconf, nil
// }

// filter out irrelevant routes
func filterPodNetworkRoutes(routes []netlink.Route, nic *VIF) (filteredRoutes []netlink.Route) {
	for _, route := range routes {
		// don't create empty static routes
		if route.Dst == nil && route.Src.Equal(nil) && route.Gw.Equal(nil) {
			continue
		}

		// don't create static route for src == nic
		if route.Src != nil && route.Src.Equal(nic.IP.IP) {
			continue
		}

		filteredRoutes = append(filteredRoutes, route)
	}
	return
}

// only used by unit test suite
func setInterfaceCacheFile(path string) {
	interfaceCacheFile = path
}

// returns nameservers [][]byte, searchdomains []string, error
func getResolvConfDetailsFromPod() ([][]byte, []string, error) {
	b, err := ioutil.ReadFile(resolvConf)
	if err != nil {
		return nil, nil, err
	}

	nameservers, err := ParseNameservers(string(b))
	if err != nil {
		return nil, nil, err
	}

	searchDomains, err := ParseSearchDomains(string(b))
	if err != nil {
		return nil, nil, err
	}

	log.Log.Reason(err).Infof("Found nameservers in %s: %s", resolvConf, bytes.Join(nameservers, []byte{' '}))
	log.Log.Reason(err).Infof("Found search domains in %s: %s", resolvConf, strings.Join(searchDomains, " "))

	return nameservers, searchDomains, err
}

func ParseNameservers(content string) ([][]byte, error) {
	var nameservers [][]byte

	re, err := regexp.Compile("([0-9]{1,3}.?){4}")
	if err != nil {
		return nameservers, err
	}

	scanner := bufio.NewScanner(strings.NewReader(content))

	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, nameserverPrefix) {
			nameserver := re.FindString(line)
			if nameserver != "" {
				nameservers = append(nameservers, net.ParseIP(nameserver).To4())
			}
		}
	}

	if err = scanner.Err(); err != nil {
		return nameservers, err
	}

	// apply a default DNS if none found from pod
	if len(nameservers) == 0 {
		nameservers = append(nameservers, net.ParseIP(defaultDNS).To4())
	}

	return nameservers, nil
}

func ParseSearchDomains(content string) ([]string, error) {
	var searchDomains []string

	scanner := bufio.NewScanner(strings.NewReader(content))

	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, domainSearchPrefix) {
			doms := strings.Fields(strings.TrimPrefix(line, domainSearchPrefix))
			for _, dom := range doms {
				searchDomains = append(searchDomains, dom)
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	if len(searchDomains) == 0 {
		searchDomains = append(searchDomains, defaultSearchDomain)
	}

	return searchDomains, nil
}
<commit_msg>remove comment section<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.
 *
 */

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

package network

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"regexp"
	"strings"

	"github.com/vishvananda/netlink"

	"kubevirt.io/kubevirt/pkg/log"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/network/dhcp"

	lmf "github.com/subgraph/libmacouflage"
)

type VIF struct {
	Name    string
	IP      netlink.Addr
	MAC     net.HardwareAddr
	Gateway net.IP
	Routes  *[]netlink.Route
	Mtu     uint16
}

type NetworkHandler interface {
	LinkByName(name string) (netlink.Link, error)
	AddrList(link netlink.Link, family int) ([]netlink.Addr, error)
	RouteList(link netlink.Link, family int) ([]netlink.Route, error)
	AddrDel(link netlink.Link, addr *netlink.Addr) error
	AddrAdd(link netlink.Link, addr *netlink.Addr) error
	LinkSetDown(link netlink.Link) error
	LinkSetUp(link netlink.Link) error
	LinkAdd(link netlink.Link) error
	ParseAddr(s string) (*netlink.Addr, error)
	SetRandomMac(iface string) (net.HardwareAddr, error)
	GetMacDetails(iface string) (net.HardwareAddr, error)
	StartDHCP(nic *VIF, serverAddr *netlink.Addr)
}

type NetworkUtilsHandler struct{}

var Handler NetworkHandler

func (h *NetworkUtilsHandler) LinkByName(name string) (netlink.Link, error) {
	return netlink.LinkByName(name)
}
func (h *NetworkUtilsHandler) AddrList(link netlink.Link, family int) ([]netlink.Addr, error) {
	return netlink.AddrList(link, family)
}
func (h *NetworkUtilsHandler) RouteList(link netlink.Link, family int) ([]netlink.Route, error) {
	return netlink.RouteList(link, family)
}
func (h *NetworkUtilsHandler) AddrDel(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrDel(link, addr)
}
func (h *NetworkUtilsHandler) LinkSetDown(link netlink.Link) error {
	return netlink.LinkSetDown(link)
}
func (h *NetworkUtilsHandler) LinkSetUp(link netlink.Link) error {
	return netlink.LinkSetUp(link)
}
func (h *NetworkUtilsHandler) LinkAdd(link netlink.Link) error {
	return netlink.LinkAdd(link)
}
func (h *NetworkUtilsHandler) ParseAddr(s string) (*netlink.Addr, error) {
	return netlink.ParseAddr(s)
}
func (h *NetworkUtilsHandler) AddrAdd(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrAdd(link, addr)
}

// GetMacDetails from an interface
func (h *NetworkUtilsHandler) GetMacDetails(iface string) (net.HardwareAddr, error) {
	currentMac, err := lmf.GetCurrentMac(iface)
	if err != nil {
		log.Log.Reason(err).Errorf("failed to get mac information for interface: %s", iface)
		return nil, err
	}
	return currentMac, nil
}

// SetRandomMac changes the MAC address for a agiven interface to a randomly generated, preserving the vendor prefix
func (h *NetworkUtilsHandler) SetRandomMac(iface string) (net.HardwareAddr, error) {
	var mac net.HardwareAddr

	currentMac, err := Handler.GetMacDetails(iface)
	if err != nil {
		return nil, err
	}

	changed, err := lmf.SpoofMacSameVendor(iface, false)
	if err != nil {
		log.Log.Reason(err).Errorf("failed to spoof MAC for an interface: %s", iface)
		return nil, err
	}

	if changed {
		mac, err = Handler.GetMacDetails(iface)
		if err != nil {
			return nil, err
		}
		log.Log.Reason(err).Errorf("updated MAC for interface: %s - %s", iface, mac)
	}
	return currentMac, nil
}

func (h *NetworkUtilsHandler) StartDHCP(nic *VIF, serverAddr *netlink.Addr) {
	nameservers, searchDomains, err := getResolvConfDetailsFromPod()
	if err != nil {
		log.Log.Errorf("Failed to get DNS servers from resolv.conf: %v", err)
		panic(err)
	}

	// panic in case the DHCP server failed during the vm creation
	// but ignore dhcp errors when the vm is destroyed or shutting down
	go func() {
		if err = DHCPServer(
			nic.MAC,
			nic.IP.IP,
			nic.IP.Mask,
			api.DefaultBridgeName,
			serverAddr.IP,
			nic.Gateway,
			nameservers,
			nic.Routes,
			searchDomains,
			nic.Mtu,
		); err != nil {
			log.Log.Errorf("failed to run DHCP: %v", err)
			panic(err)
		}
	}()
}

// Allow mocking for tests
var SetupPodNetwork = SetupNetworkInterfaces
var DHCPServer = dhcp.SingleClientDHCPServer

func initHandler() {
	if Handler == nil {
		Handler = &NetworkUtilsHandler{}
	}
}

func writeToCachedFile(inter interface{}, fileName, name string) error {
	buf, err := json.MarshalIndent(&inter, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshaling interface cache: %v", err)
	}
	err = ioutil.WriteFile(getInterfaceCacheFile(fileName, name), buf, 0644)
	if err != nil {
		return fmt.Errorf("error writing interface cache %v", err)
	}
	return nil
}

func readFromCachedFile(name, fileName string, inter interface{}) (bool, error) {
	buf, err := ioutil.ReadFile(getInterfaceCacheFile(fileName, name))
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}

	err = json.Unmarshal(buf, &inter)
	if err != nil {
		return false, fmt.Errorf("error unmarshaling command line: %v", err)
	}
	return true, nil
}

func getInterfaceCacheFile(filePath, name string) string {
	return fmt.Sprintf(filePath, name)
}

// filter out irrelevant routes
func filterPodNetworkRoutes(routes []netlink.Route, nic *VIF) (filteredRoutes []netlink.Route) {
	for _, route := range routes {
		// don't create empty static routes
		if route.Dst == nil && route.Src.Equal(nil) && route.Gw.Equal(nil) {
			continue
		}

		// don't create static route for src == nic
		if route.Src != nil && route.Src.Equal(nic.IP.IP) {
			continue
		}

		filteredRoutes = append(filteredRoutes, route)
	}
	return
}

// only used by unit test suite
func setInterfaceCacheFile(path string) {
	interfaceCacheFile = path
}

// returns nameservers [][]byte, searchdomains []string, error
func getResolvConfDetailsFromPod() ([][]byte, []string, error) {
	b, err := ioutil.ReadFile(resolvConf)
	if err != nil {
		return nil, nil, err
	}

	nameservers, err := ParseNameservers(string(b))
	if err != nil {
		return nil, nil, err
	}

	searchDomains, err := ParseSearchDomains(string(b))
	if err != nil {
		return nil, nil, err
	}

	log.Log.Reason(err).Infof("Found nameservers in %s: %s", resolvConf, bytes.Join(nameservers, []byte{' '}))
	log.Log.Reason(err).Infof("Found search domains in %s: %s", resolvConf, strings.Join(searchDomains, " "))

	return nameservers, searchDomains, err
}

func ParseNameservers(content string) ([][]byte, error) {
	var nameservers [][]byte

	re, err := regexp.Compile("([0-9]{1,3}.?){4}")
	if err != nil {
		return nameservers, err
	}

	scanner := bufio.NewScanner(strings.NewReader(content))

	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, nameserverPrefix) {
			nameserver := re.FindString(line)
			if nameserver != "" {
				nameservers = append(nameservers, net.ParseIP(nameserver).To4())
			}
		}
	}

	if err = scanner.Err(); err != nil {
		return nameservers, err
	}

	// apply a default DNS if none found from pod
	if len(nameservers) == 0 {
		nameservers = append(nameservers, net.ParseIP(defaultDNS).To4())
	}

	return nameservers, nil
}

func ParseSearchDomains(content string) ([]string, error) {
	var searchDomains []string

	scanner := bufio.NewScanner(strings.NewReader(content))

	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, domainSearchPrefix) {
			doms := strings.Fields(strings.TrimPrefix(line, domainSearchPrefix))
			for _, dom := range doms {
				searchDomains = append(searchDomains, dom)
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	if len(searchDomains) == 0 {
		searchDomains = append(searchDomains, defaultSearchDomain)
	}

	return searchDomains, nil
}
<|endoftext|>


<commit_before>package ts3

// ServerMethods groups server methods.
type ServerMethods struct {
	*Client
}

// Server represents a TeamSpeak 3 virtual server.
type Server struct {
	AnitFloodPointsNeededCommandBlock      int     `ms:"virtualserver_antiflood_points_needed_command_block"`
	AnitFloodPointsNeededIPBlock           int     `ms:"virtualserver_antiflood_points_needed_ip_block"`
	AnitFloodPointsTickReduce              int     `ms:"virtualserver_antiflood_points_tick_reduce"`
	ChannelsOnline                         int     `ms:"virtualserver_channelsonline"`
	ClientsOnline                          int     `ms:"virtualserver_clientsonline"`
	CodecEncryptionMode                    int     `ms:"virtualserver_codec_encryption_mode"`
	ComplainAutoBanCount                   int     `ms:"virtualserver_complain_autoban_count"`
	ComplainAutoBanTime                    int     `ms:"virtualserver_complain_autoban_time"`
	ComplainRemoveTime                     int     `ms:"virtualserver_complain_remove_time"`
	Created                                int     `ms:"virtualserver_created"`
	DefaultChannelAdminGroup               int     `ms:"virtualserver_default_channel_admin_group"`
	DefaultChannelGroup                    int     `ms:"virtualserver_default_channel_group"`
	DefaultServerGroup                     int     `ms:"virtualserver_default_server_group"`
	HostBannerGFXInterval                  int     `ms:"virtualserver_hostbanner_gfx_interval"`
	HostMessageMode                        int     `ms:"virtualserver_hostmessage_mode"`
	ID                                     int     `ms:"virtualserver_id"`
	MachineID                              string  `ms:"virtualserver_machine_id"`
	MaxClients                             int     `ms:"virtualserver_maxclients"`
	MinClientsInChannelBeforeForcedSilence int     `ms:"virtualserver_min_clients_in_channel_before_forced_silence"`
	NeededIdentitySecurityLevel            int     `ms:"virtualserver_needed_identity_security_level"`
	Port                                   int     `ms:"virtualserver_port"`
	QueryClientsOnline                     int     `ms:"virtualserver_queryclientsonline"`
	Uptime                                 int     `ms:"virtualserver_uptime"` // TODO(steve): use time.Duration
	AskForPrivilegeKey                     bool    `ms:"virtualserver_ask_for_privilegekey"`
	AutoStart                              bool    `ms:"virtualserver_autostart"`
	FlagPassword                           bool    `ms:"virtualserver_flag_password"`
	LogChannel                             bool    `ms:"virtualserver_log_channel"`
	LogClient                              bool    `ms:"virtualserver_log_client"`
	LogFileTransfer                        bool    `ms:"virtualserver_log_filetransfer"`
	LogPermissions                         bool    `ms:"virtualserver_log_permissions"`
	LogQuery                               bool    `ms:"virtualserver_log_client"`
	LogServer                              bool    `ms:"virtualserver_log_server"`
	WebListEnabled                         bool    `ms:"virtualserver_web_list_enabled"`
	PrioritySpeakerDimmModificator         float32 `ms:"virtualserver_priority_speaker_dimm_modificator"`
	BandwidthReceivedLastMinuteTotal       int     `ms:"virtualserver_bandwidth_received_last_minute_total"`
	BandwidthReceivedLastSecondTotal       int     `ms:"virtualserver_bandwidth_received_last_second_total"`
	BandwidthSentLastMinuteTotal           int     `ms:"virtualserver_bandwidth_sent_last_minute_total"`
	BandwidthSentLastSecondTotal           int     `ms:"virtualserver_bandwidth_sent_last_second_total"`
	ChannelTempDeleteDelayDefault          int     `ms:"virtualserver_channel_temp_delete_delay_default"`
	HostBannerMode                         int     `ms:"virtualserver_hostbanner_mode"`
	IconID                                 int     `ms:"virtualserver_icon_id"`
	MinAndroidVersion                      int     `ms:"virtualserver_min_android_version"`
	MinClientVersion                       int     `ms:"virtualserver_min_client_version"`
	MiniOSVersion                          int     `ms:"virtualserver_min_ios_version"`
	ReserverSlots                          int     `ms:"virtualserver_reserved_slots"`
	TotalPing                              float32 `ms:"virtualserver_total_ping"`
	MaxDownloadTotalBandwidth              uint64  `ms:"virtualserver_max_download_total_bandwidth"`
	MaxUploadTotalBandwidth                uint64  `ms:"virtualserver_max_upload_total_bandwidth"`
	MonthBytesDownloaded                   int64   `ms:"virtualserver_MonthBytesDownloaded"`
	MonthBytesUploaded                     int64   `ms:"virtualserver_MonthBytesUploaded"`
	TotalBytesDownloaded                   int64   `ms:"virtualserver_TotalBytesDownloaded"`
	TotalBytesUploaded                     int64   `ms:"virtualserver_TotalBytesUploaded"`
	TotalPacketLossControl                 int64   `ms:"virtualserver_total_packet_loss_control"`
	TotalPacketLossKeepalive               int64   `ms:"virtualserver_total_packet_loss_keepalive"`
	TotalPacketLossSpeech                  int64   `ms:"virtualserver_total_packet_loss_speech"`
	TotalPacketLossTotal                   int64   `ms:"virtualserver_total_packet_loss_total"`
	VirtualServerDownloadQuota             int64   `ms:"virtualserver_virtualserver_download_quota"`
	VirtualServerUploadQuota               int64   `ms:"virtualserver_virtualserver_Upload_quota"`
	FileBase                               string  `ms:"virtualserver_filebase"`
	HostBannerGFXURL                       string  `ms:"virtualserver_hostbanner_gfx_url"`
	HostBannerURL                          string  `ms:"virtualserver_hostbanner_url"`
	HostButtonGFXURL                       string  `ms:"virtualserver_hostbutton_gfx_url"`
	HostButtonToolTip                      string  `ms:"virtualserver_hostbutton_tooltip"`
	HostButtonURL                          string  `ms:"virtualserver_hostbutton_url"`
	HostMessage                            string  `ms:"virtualserver_hostmessage"`
	Name                                   string  `ms:"virtualserver_name"`
	NamePhonetic                           string  `ms:"virtualserver_name_phonetic"`
	Password                               string  `ms:"virtualserver_password"`
	Platform                               string  `ms:"virtualserver_platform"`
	Status                                 string  `ms:"virtualserver_status"`
	UniqueIdentifier                       string  `ms:"virtualserver_unique_identifier"`
	Version                                string  `ms:"virtualserver_version"`
	WelcomeMessage                         string  `ms:"virtualserver_welcomemessage"`
}

// List lists virtual servers.
func (s *ServerMethods) List(options ...string) ([]*Server, error) {
	var servers []*Server
	if _, err := s.ExecCmd(NewCmd("serverlist").WithOptions(options...).WithResponse(&servers)); err != nil {
		return nil, err
	}

	return servers, nil
}

// IDGetByPort returns the database id of the virtual server running on UDP port.
func (s *ServerMethods) IDGetByPort(port uint16) (int, error) {
	r := struct {
		ID int `ms:"server_id"`
	}{}
	_, err := s.ExecCmd(NewCmd("serveridgetbyport").WithArgs(NewArg("virtualserver_port", port)).WithResponse(&r))
	return r.ID, err
}

// Info returns detailed configuration information about the selected server.
func (s *ServerMethods) Info() (*Server, error) {
	r := &Server{}
	if _, err := s.ExecCmd(NewCmd("serverinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// Edit changes the selected virtual servers configuration using the given args.
func (s *ServerMethods) Edit(args ...CmdArg) error {
	_, err := s.ExecCmd(NewCmd("serveredit").WithArgs(args...))
	return err
}

// Delete deletes the virtual server specified by id.
// Only virtual server in a stopped state can be deleted.
func (s *ServerMethods) Delete(id int) error {
	_, err := s.ExecCmd(NewCmd("serverdelete").WithArgs(NewArg("sid", id)))
	return err
}

// CreatedServer is the details returned by a server create.
type CreatedServer struct {
	ID    int    `ms:"sid"`
	Port  uint16 `ms:"virtualserver_port"`
	Token string
}

// Create creates a new virtual server using the given properties and returns
// its ID, port and initial administrator privilege key.
// If virtualserver_port arg is not specified, the server will use the first unused
// UDP port.
func (s *ServerMethods) Create(name string, args ...CmdArg) (*CreatedServer, error) {
	r := &CreatedServer{}
	args = append(args, NewArg("virtualserver_name", name))
	if _, err := s.ExecCmd(NewCmd("servercreate").WithArgs(args...).WithResponse(r)); err != nil {
		return nil, err
	}

	return r, nil
}

// Start starts the virtual server specified by id.
func (s *ServerMethods) Start(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstart").WithArgs(NewArg("sid", id)))
	return err
}

// Stop stops the virtual server specified by id.
func (s *ServerMethods) Stop(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstop").WithArgs(NewArg("sid", id)))
	return err
}

// Group represents a virtual server group.
type Group struct {
	ID                int `ms:"sgid"`
	Name              string
	Type              int
	IconID            int
	Saved             bool `ms:"savedb"`
	SortID            int
	NameMode          int
	ModifyPower       int `ms:"n_modifyp"`
	MemberAddPower    int `ms:"n_member_addp"`
	MemberRemovePower int `ms:"n_member_addp"`
}

// GroupList returns a list of available groups for the selected server.
func (s *ServerMethods) GroupList() ([]*Group, error) {
	var groups []*Group
	if _, err := s.ExecCmd(NewCmd("servergrouplist").WithResponse(&groups)); err != nil {
		return nil, err
	}

	return groups, nil
}

// PrivilegeKey represents a server privilege key.
type PrivilegeKey struct {
	Token       string
	Type        int    `ms:"token_type"`
	ID1         int    `ms:"token_id1"`
	ID2         int    `ms:"token_id2"`
	Created     int    `ms:"token_created"`
	Description string `ms:"token_description"`
}

// PrivilegeKeyList returns a list of available privilege keys for the selected server,
// including their type and group IDs.
func (s *ServerMethods) PrivilegeKeyList() ([]*PrivilegeKey, error) {
	var keys []*PrivilegeKey
	if _, err := s.ExecCmd(NewCmd("privilegekeylist").WithResponse(&keys)); err != nil {
		return nil, err
	}

	return keys, nil
}

// PrivilegeKeyAdd creates a new privilege token to the selected server and returns it.
// If tokentype is set to 0, the ID specified with id1 will be a server group ID.
// Otherwise, id1 is used as a channel group ID and you need to provide a valid channel ID using id2.
func (s *ServerMethods) PrivilegeKeyAdd(ttype, id1, id2 int, options ...CmdArg) (string, error) {
	t := struct {
		Token string
	}{}
	options = append(options, NewArg("tokentype", ttype), NewArg("tokenid1", id1), NewArg("tokenid2", id2))
	_, err := s.ExecCmd(NewCmd("privilegekeylist").WithArgs(options...).WithResponse(&t))
	return t.Token, err
}
<commit_msg>Adding a few methods to gather server statistics<commit_after>package ts3

type ServerMethods struct {
	*Client
}

type Channel struct {
	ID                   int    `ms:"cid"`
	ParentID             int    `ms:"pid"`
	ChannelOrder         int    `ms:"channel_order"`
	ChannelName          string `ms:"channel_name"`
	TotalClients         int    `ms:"total_clients"`
	NeededSubscribePower int    `ms:"channel_needed_subscribe_power"`
}

type Instance struct {
	DatabaseVersion             int    `ms:"serverinstance_database_version"`
	FileTransferPort            int    `ms:"serverinstance_filetransfer_port"`
	MaxTotalDownloadBandwidth   uint64 `ms:"serverinstance_max_download_total_bandwidth"`
	MaxTotalUploadBandwidth     uint64 `ms:"serverinstance_max_upload_total_bandwidth"`
	GuestServerQueryGroup       int    `ms:"serverinstance_guest_serverquery_group"`
	ServerQueryFloodCommands    int    `ms:"serverinstance_serverquery_flood_commands"`
	ServerQueryFloodTime        int    `ms:"serverinstance_serverquery_flood_time"`
	ServerQueryBanTime          int    `ms:"serverinstance_serverquery_ban_time"`
	TemplateServerAdminGroup    int    `ms:"serverinstance_template_serveradmin_group"`
	TemplateServerDefaultGroup  int    `ms:"serverinstance_template_serverdefault_group"`
	TemplateChannelAdminGroup   int    `ms:"serverinstance_template_channeladmin_group"`
	TemplateChannelDefaultGroup int    `ms:"serverinstance_template_channeldefault_group"`
}

type ServerConnectionInfo struct {
	FileTransferBandwidthSent     uint64  `ms:"connection_filetransfer_bandwidth_sent"`
	FileTransferBandwidthReceived uint64  `ms:"connection_filetransfer_bandwidth_received"`
	FileTransferTotalSent         uint64  `ms:"connection_filetransfer_bytes_sent_total"`
	FileTransferTotalReceived     uint64  `ms:"connection_filetransfer_bytes_received_total"`
	PacketsSentTotal              uint64  `ms:"connection_packets_sent_total"`
	PacketsReceivedTotal          uint64  `ms:"connection_packets_received_total"`
	BytesSentTotal                uint64  `ms:"connection_bytes_sent_total"`
	BytesReceivedTotal            uint64  `ms:"connection_bytes_received_total"`
	BandwidthSentLastSecond       uint64  `ms:"connection_bandwidth_sent_last_second_total"`
	BandwidthReceivedLastSecond   uint64  `ms:"connection_bandwidth_received_last_second_total"`
	BandwidthSentLastMinute       uint64  `ms:"connection_bandwidth_sent_last_minute_total"`
	BandwidthReceivedLastMinute   uint64  `ms:"connection_bandwidth_received_last_minute_total"`
	ConnectedTime                 uint32  `ms:"connection_connected_time"`
	PacketLossTotalAvg            float32 `ms:"connection_packetloss_total"`
	PingTotalAvg                  float32 `ms:"connection_ping"`
}

// Server represents a TeamSpeak 3 virtual server.
type Server struct {
	AntiFloodPointsNeededCommandBlock      int     `ms:"virtualserver_antiflood_points_needed_command_block"`
	AntiFloodPointsNeededIPBlock           int     `ms:"virtualserver_antiflood_points_needed_ip_block"`
	AntiFloodPointsTickReduce              int     `ms:"virtualserver_antiflood_points_tick_reduce"`
	ChannelsOnline                         int     `ms:"virtualserver_channelsonline"`
	ClientsOnline                          int     `ms:"virtualserver_clientsonline"`
	CodecEncryptionMode                    int     `ms:"virtualserver_codec_encryption_mode"`
	ComplainAutoBanCount                   int     `ms:"virtualserver_complain_autoban_count"`
	ComplainAutoBanTime                    int     `ms:"virtualserver_complain_autoban_time"`
	ComplainRemoveTime                     int     `ms:"virtualserver_complain_remove_time"`
	Created                                int     `ms:"virtualserver_created"`
	DefaultChannelAdminGroup               int     `ms:"virtualserver_default_channel_admin_group"`
	DefaultChannelGroup                    int     `ms:"virtualserver_default_channel_group"`
	DefaultServerGroup                     int     `ms:"virtualserver_default_server_group"`
	HostBannerGFXInterval                  int     `ms:"virtualserver_hostbanner_gfx_interval"`
	HostMessageMode                        int     `ms:"virtualserver_hostmessage_mode"`
	ID                                     int     `ms:"virtualserver_id"`
	MachineID                              string  `ms:"virtualserver_machine_id"`
	MaxClients                             int     `ms:"virtualserver_maxclients"`
	MinClientsInChannelBeforeForcedSilence int     `ms:"virtualserver_min_clients_in_channel_before_forced_silence"`
	NeededIdentitySecurityLevel            int     `ms:"virtualserver_needed_identity_security_level"`
	Port                                   int     `ms:"virtualserver_port"`
	QueryClientsOnline                     int     `ms:"virtualserver_queryclientsonline"`
	Uptime                                 int     `ms:"virtualserver_uptime"` // TODO(steve): use time.Duration
	AskForPrivilegeKey                     bool    `ms:"virtualserver_ask_for_privilegekey"`
	AutoStart                              bool    `ms:"virtualserver_autostart"`
	FlagPassword                           bool    `ms:"virtualserver_flag_password"`
	LogChannel                             bool    `ms:"virtualserver_log_channel"`
	LogClient                              bool    `ms:"virtualserver_log_client"`
	LogFileTransfer                        bool    `ms:"virtualserver_log_filetransfer"`
	LogPermissions                         bool    `ms:"virtualserver_log_permissions"`
	LogQuery                               bool    `ms:"virtualserver_log_client"`
	LogServer                              bool    `ms:"virtualserver_log_server"`
	WebListEnabled                         bool    `ms:"virtualserver_web_list_enabled"`
	PrioritySpeakerDimmModificator         float32 `ms:"virtualserver_priority_speaker_dimm_modificator"`
	BandwidthReceivedLastMinuteTotal       int     `ms:"virtualserver_bandwidth_received_last_minute_total"`
	BandwidthReceivedLastSecondTotal       int     `ms:"virtualserver_bandwidth_received_last_second_total"`
	BandwidthSentLastMinuteTotal           int     `ms:"virtualserver_bandwidth_sent_last_minute_total"`
	BandwidthSentLastSecondTotal           int     `ms:"virtualserver_bandwidth_sent_last_second_total"`
	ChannelTempDeleteDelayDefault          int     `ms:"virtualserver_channel_temp_delete_delay_default"`
	HostBannerMode                         int     `ms:"virtualserver_hostbanner_mode"`
	IconID                                 int     `ms:"virtualserver_icon_id"`
	MinAndroidVersion                      int     `ms:"virtualserver_min_android_version"`
	MinClientVersion                       int     `ms:"virtualserver_min_client_version"`
	MiniOSVersion                          int     `ms:"virtualserver_min_ios_version"`
	ReservedSlots                          int     `ms:"virtualserver_reserved_slots"`
	TotalPing                              float32 `ms:"virtualserver_total_ping"`
	MaxDownloadTotalBandwidth              uint64  `ms:"virtualserver_max_download_total_bandwidth"`
	MaxUploadTotalBandwidth                uint64  `ms:"virtualserver_max_upload_total_bandwidth"`
	MonthBytesDownloaded                   int64   `ms:"virtualserver_MonthBytesDownloaded"`
	MonthBytesUploaded                     int64   `ms:"virtualserver_MonthBytesUploaded"`
	TotalBytesDownloaded                   int64   `ms:"virtualserver_TotalBytesDownloaded"`
	TotalBytesUploaded                     int64   `ms:"virtualserver_TotalBytesUploaded"`
	TotalPacketLossControl                 int64   `ms:"virtualserver_total_packet_loss_control"`
	TotalPacketLossKeepalive               int64   `ms:"virtualserver_total_packet_loss_keepalive"`
	TotalPacketLossSpeech                  int64   `ms:"virtualserver_total_packet_loss_speech"`
	TotalPacketLossTotal                   int64   `ms:"virtualserver_total_packet_loss_total"`
	VirtualServerDownloadQuota             int64   `ms:"virtualserver_virtualserver_download_quota"`
	VirtualServerUploadQuota               int64   `ms:"virtualserver_virtualserver_Upload_quota"`
	FileBase                               string  `ms:"virtualserver_filebase"`
	HostBannerGFXURL                       string  `ms:"virtualserver_hostbanner_gfx_url"`
	HostBannerURL                          string  `ms:"virtualserver_hostbanner_url"`
	HostButtonGFXURL                       string  `ms:"virtualserver_hostbutton_gfx_url"`
	HostButtonToolTip                      string  `ms:"virtualserver_hostbutton_tooltip"`
	HostButtonURL                          string  `ms:"virtualserver_hostbutton_url"`
	HostMessage                            string  `ms:"virtualserver_hostmessage"`
	Name                                   string  `ms:"virtualserver_name"`
	NamePhonetic                           string  `ms:"virtualserver_name_phonetic"`
	Password                               string  `ms:"virtualserver_password"`
	Platform                               string  `ms:"virtualserver_platform"`
	Status                                 string  `ms:"virtualserver_status"`
	UniqueIdentifier                       string  `ms:"virtualserver_unique_identifier"`
	Version                                string  `ms:"virtualserver_version"`
	WelcomeMessage                         string  `ms:"virtualserver_welcomemessage"`
}

// List lists virtual servers.
func (s *ServerMethods) List(options ...string) ([]*Server, error) {
	var servers []*Server
	if _, err := s.ExecCmd(NewCmd("serverlist").WithOptions(options...).WithResponse(&servers)); err != nil {
		return nil, err
	}

	return servers, nil
}

// IDGetByPort returns the database id of the virtual server running on UDP port.
func (s *ServerMethods) IDGetByPort(port uint16) (int, error) {
	r := struct {
		ID int `ms:"server_id"`
	}{}
	_, err := s.ExecCmd(NewCmd("serveridgetbyport").WithArgs(NewArg("virtualserver_port", port)).WithResponse(&r))
	return r.ID, err
}

// Info returns detailed configuration information about the selected server.
func (s *ServerMethods) Info() (*Server, error) {
	r := &Server{}
	if _, err := s.ExecCmd(NewCmd("serverinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// InstanceInfo returns detailed information about the selected instance.
func (s *ServerMethods) InstanceInfo() (*Instance, error) {
	r := &Instance{}
	if _, err := s.ExecCmd(NewCmd("instanceinfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil
}

// ServerConnectionInfo returns detailed bandwidth and transfer information about the selected instance.
func (s *ServerMethods) ServerConnectionInfo() (*ServerConnectionInfo, error) {
	r := &ServerConnectionInfo{}
	if _, err := s.ExecCmd(NewCmd("serverrequestconnectioninfo").WithResponse(&r)); err != nil {
		return nil, err
	}

	return r, nil

}

// Edit changes the selected virtual servers configuration using the given args.
func (s *ServerMethods) Edit(args ...CmdArg) error {
	_, err := s.ExecCmd(NewCmd("serveredit").WithArgs(args...))
	return err
}

// Delete deletes the virtual server specified by id.
// Only virtual server in a stopped state can be deleted.
func (s *ServerMethods) Delete(id int) error {
	_, err := s.ExecCmd(NewCmd("serverdelete").WithArgs(NewArg("sid", id)))
	return err
}

// CreatedServer is the details returned by a server create.
type CreatedServer struct {
	ID    int    `ms:"sid"`
	Port  uint16 `ms:"virtualserver_port"`
	Token string
}

// Create creates a new virtual server using the given properties and returns
// its ID, port and initial administrator privilege key.
// If virtualserver_port arg is not specified, the server will use the first unused
// UDP port.
func (s *ServerMethods) Create(name string, args ...CmdArg) (*CreatedServer, error) {
	r := &CreatedServer{}
	args = append(args, NewArg("virtualserver_name", name))
	if _, err := s.ExecCmd(NewCmd("servercreate").WithArgs(args...).WithResponse(r)); err != nil {
		return nil, err
	}

	return r, nil
}

// Start starts the virtual server specified by id.
func (s *ServerMethods) Start(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstart").WithArgs(NewArg("sid", id)))
	return err
}

// Stop stops the virtual server specified by id.
func (s *ServerMethods) Stop(id int) error {
	_, err := s.ExecCmd(NewCmd("serverstop").WithArgs(NewArg("sid", id)))
	return err
}

// Group represents a virtual server group.
type Group struct {
	ID                int `ms:"sgid"`
	Name              string
	Type              int
	IconID            int
	Saved             bool `ms:"savedb"`
	SortID            int
	NameMode          int
	ModifyPower       int `ms:"n_modifyp"`
	MemberAddPower    int `ms:"n_member_addp"`
	MemberRemovePower int `ms:"n_member_addp"`
}

// GroupList returns a list of available groups for the selected server.
func (s *ServerMethods) GroupList() ([]*Group, error) {
	var groups []*Group
	if _, err := s.ExecCmd(NewCmd("servergrouplist").WithResponse(&groups)); err != nil {
		return nil, err
	}

	return groups, nil
}

// ChannelList returns a list of channels for the selected server.
func (s *ServerMethods) ChannelList() ([]*Channel, error) {
	var channels []*Channel
	if _, err := s.ExecCmd(NewCmd("channellist").WithResponse(&channels)); err != nil {
		return nil, err
	}

	return channels, nil
}

// PrivilegeKey represents a server privilege key.
type PrivilegeKey struct {
	Token       string
	Type        int    `ms:"token_type"`
	ID1         int    `ms:"token_id1"`
	ID2         int    `ms:"token_id2"`
	Created     int    `ms:"token_created"`
	Description string `ms:"token_description"`
}

// PrivilegeKeyList returns a list of available privilege keys for the selected server,
// including their type and group IDs.
func (s *ServerMethods) PrivilegeKeyList() ([]*PrivilegeKey, error) {
	var keys []*PrivilegeKey
	if _, err := s.ExecCmd(NewCmd("privilegekeylist").WithResponse(&keys)); err != nil {
		return nil, err
	}

	return keys, nil
}

// PrivilegeKeyAdd creates a new privilege token to the selected server and returns it.
// If tokentype is set to 0, the ID specified with id1 will be a server group ID.
// Otherwise, id1 is used as a channel group ID and you need to provide a valid channel ID using id2.
func (s *ServerMethods) PrivilegeKeyAdd(ttype, id1, id2 int, options ...CmdArg) (string, error) {
	t := struct {
		Token string
	}{}
	options = append(options, NewArg("tokentype", ttype), NewArg("tokenid1", id1), NewArg("tokenid2", id2))
	_, err := s.ExecCmd(NewCmd("privilegekeylist").WithArgs(options...).WithResponse(&t))
	return t.Token, err
}
<|endoftext|>


<commit_before>package gunfish

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

	"github.com/kayac/Gunfish/apns"
	"github.com/sirupsen/logrus"
)

func init() {
	InitErrorResponseHandler(DefaultResponseHandler{hook: `cat `})
	InitSuccessResponseHandler(DefaultResponseHandler{})
	logrus.SetLevel(logrus.WarnLevel)
	config.Apns.Host = MockServer
}

func TestInvalidCertification(t *testing.T) {
	c, _ := LoadConfig("./test/gunfish_test.toml")
	c.Apns.CertFile = "./test/invalid.crt"
	c.Apns.KeyFile = "./test/invalid.key"
	ss, err := StartSupervisor(&c)
	if err != nil {
		t.Errorf("Expected supervisor cannot start because of using invalid certification files.: %v", ss)
	}
}

func TestSuccessToPostJson(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	// application/json
	jsons := createJSONPostedData(3)
	w := httptest.NewRecorder()
	r, err := newRequest(jsons, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusOK {
		t.Errorf("Expected status code is 200 but got %d", w.Code)
	}

	// application/x-www-form-urlencoded
	data := createFormPostedData(3)
	w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.
	r, err = newRequest(data, "POST", ApplicationXW3FormURLEncoded)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusOK {
		t.Errorf("Expected status code is 200 but got %d", w.Code)
	}

	sup.Shutdown()
}

func TestFailedToPostMalformedJson(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	jsons := []string{
		`{"test":"test"}`,
		"[{\"payload\": {\"aps\": {\"alert\":\"msg\", \"sound\":\"default\" }}}]",
		"[{\"token\":\"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\"}]",
	}
	for _, s := range jsons {
		v := url.Values{}
		v.Add("json", s)

		// application/x-www-form-urlencoded
		r, err := newRequest([]byte(v.Encode()), "POST", ApplicationXW3FormURLEncoded)
		if err != nil {
			t.Errorf("%s", err)
		}

		w := httptest.NewRecorder()
		handler.ServeHTTP(w, r)
		if w.Code == http.StatusOK {
			t.Errorf("Expected status code is NOT 200 but got %d", w.Code)
		}

		// application/json
		r, err = newRequest([]byte(s), "POST", ApplicationJSON)
		if err != nil {
			t.Errorf("%s", err)
		}

		w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.
		handler.ServeHTTP(w, r)
		if w.Code == http.StatusOK {
			t.Errorf("Expected status code is NOT 200 but got %d", w.Code)
		}
	}

	sup.Shutdown()
}

func TestEnqueueTooManyRequest(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	srvStats = NewStats(config)
	handler := prov.pushAPNsHandler()

	// When queue stack is full, return 503
	var manyNum int
	tp := ((config.Provider.RequestQueueSize * int(AverageResponseTime/time.Millisecond)) / 1000) / SenderNum
	dif := (RequestPerSec - config.Provider.RequestQueueSize/tp)
	if dif > 0 {
		manyNum = dif * int(FlowRateInterval/time.Second) * 2
	} else {
		manyNum = -1 * dif * int(FlowRateInterval/time.Second) * 2
	}

	var jsons [][]byte
	for i := 0; i < manyNum; i++ {
		jsons = append(jsons, createJSONPostedData(1)) // Too many requests
	}

	// Test 503 returns
	check503 := false
	w := httptest.NewRecorder() // creates new Recoder because cannot overwrite header.
	var ra string
	for _, json := range jsons {
		r, err := newRequest(json, "POST", ApplicationJSON)
		if err != nil {
			t.Errorf("%s", err)
		}

		handler.ServeHTTP(w, r)
		if w.Code == http.StatusServiceUnavailable {
			check503 = true
			ra = w.Header().Get("Retry-After")
			break
		} else {
			w = httptest.NewRecorder()
		}
	}
	if check503 == false {
		t.Errorf("Expected status code is 503 but got %d", w.Code)
	}
	if w.Header().Get("Retry-After") == "" {
		t.Error("Not set Retry-After correctlly")
	}

	// Test Retry-After value increases
	w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.
	r, err := newRequest(jsons[0], "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusServiceUnavailable {
		t.Errorf("Expected status code is 503 but got %d", w.Code)
	}
	if w.Header().Get("Retry-After") == ra {
		t.Errorf("Retry-After should become different: old[ %s ], new[ %s ]", ra, w.Header().Get("Retry-After"))
	}

	sup.Shutdown()
}

func TestTooLargeRequest(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	srvStats = NewStats(config)
	handler := prov.pushAPNsHandler()

	jsons := createJSONPostedData(MaxRequestSize + 1) // Too many requests
	r, err := newRequest(jsons, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	w := httptest.NewRecorder()
	handler.ServeHTTP(w, r)
	if w.Code != http.StatusBadRequest {
		t.Errorf("Expected status code is %d but got %d", http.StatusBadRequest, w.Code)
	}

	sup.Shutdown()
}

func TestMethodNotAllowed(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	jsons := createJSONPostedData(1)
	w := httptest.NewRecorder()
	r, err := newRequest(jsons, "GET", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status code is 405 but got %d", w.Code)
	}

	sup.Shutdown()
}

func TestUnsupportedMediaType(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	jsons := createPostedData(1)
	r, err := http.NewRequest(
		"POST",
		"",
		bytes.NewBuffer([]byte(jsons)),
	)
	r.Header.Set("Content-Type", "plain/text")
	w := httptest.NewRecorder()
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusUnsupportedMediaType {
		t.Errorf("Expected status code is 415 but got %d", w.Code)
	}

	sup.Shutdown()
}

func TestStats(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	srvStats = NewStats(config)
	pushh := prov.pushAPNsHandler()
	statsh := prov.statsHandler()

	// Updates stat
	jsons := createJSONPostedData(1)
	w := httptest.NewRecorder()
	r, err := newRequest(jsons, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}
	pushh.ServeHTTP(w, r)
	w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.

	// GET status
	r, err = newRequest([]byte(""), "GET", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}
	time.Sleep(time.Second * 1) // for update uptime of stats
	statsh.ServeHTTP(w, r)

	// check stat
	var resStat Stats
	de := json.NewDecoder(w.Body)
	de.Decode(&resStat)

	if !reflect.DeepEqual(srvStats, resStat) {
		t.Errorf("Expected total conenctions \"%v\" but got \"%v\"", srvStats, resStat)
	}

	if resStat.CertificateExpireUntil < 0 {
		t.Errorf("Certificate expired %s %d", resStat.CertificateNotAfter, resStat.CertificateExpireUntil)
	}

	sup.Shutdown()
}

func newRequest(data []byte, method string, c string) (*http.Request, error) {
	req, err := http.NewRequest(
		method,
		"",
		bytes.NewBuffer([]byte(data)),
	)
	req.Header.Set("Content-Type", c)
	return req, err
}

func createJSONPostedData(num int) []byte {
	return createPostedData(num)
}

func createFormPostedData(num int) []byte {
	jsonStr := createPostedData(num)
	v := url.Values{}
	v.Add("json", string(jsonStr))
	return []byte(v.Encode())
}

func createPostedData(num int) []byte {
	pds := make([]PostedData, num)
	tokens := make([]string, num)
	for i := 0; i < num; i++ {
		tokens[i] = fmt.Sprintf("%032d", i)
	}
	for i, v := range tokens {
		payload := apns.Payload{}

		payload.APS = &apns.APS{
			Alert: apns.Alert{
				Title: "test",
				Body:  "message",
			},
			Sound: "default",
		}

		pds[i] = PostedData{
			Payload: payload,
			Token:   v,
		}
	}

	jsonStr, _ := json.Marshal(pds)

	return jsonStr
}
<commit_msg>add repeatability test<commit_after>package gunfish

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

	"github.com/kayac/Gunfish/apns"
	"github.com/sirupsen/logrus"
)

func init() {
	InitErrorResponseHandler(DefaultResponseHandler{hook: `cat `})
	InitSuccessResponseHandler(DefaultResponseHandler{})
	logrus.SetLevel(logrus.WarnLevel)
	config.Apns.Host = MockServer
}

func TestInvalidCertification(t *testing.T) {
	c, _ := LoadConfig("./test/gunfish_test.toml")
	c.Apns.CertFile = "./test/invalid.crt"
	c.Apns.KeyFile = "./test/invalid.key"
	ss, err := StartSupervisor(&c)
	if err != nil {
		t.Errorf("Expected supervisor cannot start because of using invalid certification files.: %v", ss)
	}
}

func TestSuccessToPostJson(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	// application/json
	jsons := createJSONPostedData(3)
	w := httptest.NewRecorder()
	r, err := newRequest(jsons, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusOK {
		t.Errorf("Expected status code is 200 but got %d", w.Code)
	}

	// application/x-www-form-urlencoded
	data := createFormPostedData(3)
	w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.
	r, err = newRequest(data, "POST", ApplicationXW3FormURLEncoded)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusOK {
		t.Errorf("Expected status code is 200 but got %d", w.Code)
	}

	sup.Shutdown()
}

func TestFailedToPostInvalidJson(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushFCMHandler()

	// missing `}`
	invalidJson := []byte(`{"registration_ids": ["xxxxxxxxx"], "data": {"message":"test"`)

	w := httptest.NewRecorder()
	r, err := newRequest(invalidJson, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)

	invalidResponse := bytes.NewBufferString("{\"reason\":\"unexpected EOF\"}{\"result\": \"ok\"}").String()
	if w.Body.String() == invalidResponse {
		t.Errorf("Invalid Json responce: '%s'", w.Body)
	}

	sup.Shutdown()
}

func TestFailedToPostMalformedJson(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	jsons := []string{
		`{"test":"test"}`,
		"[{\"payload\": {\"aps\": {\"alert\":\"msg\", \"sound\":\"default\" }}}]",
		"[{\"token\":\"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\"}]",
	}
	for _, s := range jsons {
		v := url.Values{}
		v.Add("json", s)

		// application/x-www-form-urlencoded
		r, err := newRequest([]byte(v.Encode()), "POST", ApplicationXW3FormURLEncoded)
		if err != nil {
			t.Errorf("%s", err)
		}

		w := httptest.NewRecorder()
		handler.ServeHTTP(w, r)
		if w.Code == http.StatusOK {
			t.Errorf("Expected status code is NOT 200 but got %d", w.Code)
		}

		// application/json
		r, err = newRequest([]byte(s), "POST", ApplicationJSON)
		if err != nil {
			t.Errorf("%s", err)
		}

		w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.
		handler.ServeHTTP(w, r)
		if w.Code == http.StatusOK {
			t.Errorf("Expected status code is NOT 200 but got %d", w.Code)
		}
	}

	sup.Shutdown()
}

func TestEnqueueTooManyRequest(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	srvStats = NewStats(config)
	handler := prov.pushAPNsHandler()

	// When queue stack is full, return 503
	var manyNum int
	tp := ((config.Provider.RequestQueueSize * int(AverageResponseTime/time.Millisecond)) / 1000) / SenderNum
	dif := (RequestPerSec - config.Provider.RequestQueueSize/tp)
	if dif > 0 {
		manyNum = dif * int(FlowRateInterval/time.Second) * 2
	} else {
		manyNum = -1 * dif * int(FlowRateInterval/time.Second) * 2
	}

	var jsons [][]byte
	for i := 0; i < manyNum; i++ {
		jsons = append(jsons, createJSONPostedData(1)) // Too many requests
	}

	// Test 503 returns
	check503 := false
	w := httptest.NewRecorder() // creates new Recoder because cannot overwrite header.
	var ra string
	for _, json := range jsons {
		r, err := newRequest(json, "POST", ApplicationJSON)
		if err != nil {
			t.Errorf("%s", err)
		}

		handler.ServeHTTP(w, r)
		if w.Code == http.StatusServiceUnavailable {
			check503 = true
			ra = w.Header().Get("Retry-After")
			break
		} else {
			w = httptest.NewRecorder()
		}
	}
	if check503 == false {
		t.Errorf("Expected status code is 503 but got %d", w.Code)
	}
	if w.Header().Get("Retry-After") == "" {
		t.Error("Not set Retry-After correctlly")
	}

	// Test Retry-After value increases
	w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.
	r, err := newRequest(jsons[0], "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusServiceUnavailable {
		t.Errorf("Expected status code is 503 but got %d", w.Code)
	}
	if w.Header().Get("Retry-After") == ra {
		t.Errorf("Retry-After should become different: old[ %s ], new[ %s ]", ra, w.Header().Get("Retry-After"))
	}

	sup.Shutdown()
}

func TestTooLargeRequest(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	srvStats = NewStats(config)
	handler := prov.pushAPNsHandler()

	jsons := createJSONPostedData(MaxRequestSize + 1) // Too many requests
	r, err := newRequest(jsons, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	w := httptest.NewRecorder()
	handler.ServeHTTP(w, r)
	if w.Code != http.StatusBadRequest {
		t.Errorf("Expected status code is %d but got %d", http.StatusBadRequest, w.Code)
	}

	sup.Shutdown()
}

func TestMethodNotAllowed(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	jsons := createJSONPostedData(1)
	w := httptest.NewRecorder()
	r, err := newRequest(jsons, "GET", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status code is 405 but got %d", w.Code)
	}

	sup.Shutdown()
}

func TestUnsupportedMediaType(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	handler := prov.pushAPNsHandler()

	jsons := createPostedData(1)
	r, err := http.NewRequest(
		"POST",
		"",
		bytes.NewBuffer([]byte(jsons)),
	)
	r.Header.Set("Content-Type", "plain/text")
	w := httptest.NewRecorder()
	if err != nil {
		t.Errorf("%s", err)
	}

	handler.ServeHTTP(w, r)
	if w.Code != http.StatusUnsupportedMediaType {
		t.Errorf("Expected status code is 415 but got %d", w.Code)
	}

	sup.Shutdown()
}

func TestStats(t *testing.T) {
	sup, _ := StartSupervisor(&config)
	prov := &Provider{sup: sup}
	srvStats = NewStats(config)
	pushh := prov.pushAPNsHandler()
	statsh := prov.statsHandler()

	// Updates stat
	jsons := createJSONPostedData(1)
	w := httptest.NewRecorder()
	r, err := newRequest(jsons, "POST", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}
	pushh.ServeHTTP(w, r)
	w = httptest.NewRecorder() // re-creates Recoder because cannot overwrite header after to write body.

	// GET status
	r, err = newRequest([]byte(""), "GET", ApplicationJSON)
	if err != nil {
		t.Errorf("%s", err)
	}
	time.Sleep(time.Second * 1) // for update uptime of stats
	statsh.ServeHTTP(w, r)

	// check stat
	var resStat Stats
	de := json.NewDecoder(w.Body)
	de.Decode(&resStat)

	if !reflect.DeepEqual(srvStats, resStat) {
		t.Errorf("Expected total conenctions \"%v\" but got \"%v\"", srvStats, resStat)
	}

	if resStat.CertificateExpireUntil < 0 {
		t.Errorf("Certificate expired %s %d", resStat.CertificateNotAfter, resStat.CertificateExpireUntil)
	}

	sup.Shutdown()
}

func newRequest(data []byte, method string, c string) (*http.Request, error) {
	req, err := http.NewRequest(
		method,
		"",
		bytes.NewBuffer([]byte(data)),
	)
	req.Header.Set("Content-Type", c)
	return req, err
}

func createJSONPostedData(num int) []byte {
	return createPostedData(num)
}

func createFormPostedData(num int) []byte {
	jsonStr := createPostedData(num)
	v := url.Values{}
	v.Add("json", string(jsonStr))
	return []byte(v.Encode())
}

func createPostedData(num int) []byte {
	pds := make([]PostedData, num)
	tokens := make([]string, num)
	for i := 0; i < num; i++ {
		tokens[i] = fmt.Sprintf("%032d", i)
	}
	for i, v := range tokens {
		payload := apns.Payload{}

		payload.APS = &apns.APS{
			Alert: apns.Alert{
				Title: "test",
				Body:  "message",
			},
			Sound: "default",
		}

		pds[i] = PostedData{
			Payload: payload,
			Token:   v,
		}
	}

	jsonStr, _ := json.Marshal(pds)

	return jsonStr
}
<|endoftext|>