<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 2022 Red Hat, Inc.
 *
 */

package testsuite

import (
	"context"
	"fmt"
	"time"

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

	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/pkg/controller"

	"kubevirt.io/kubevirt/tests/flags"
	"kubevirt.io/kubevirt/tests/framework/cleanup"
	"kubevirt.io/kubevirt/tests/libnet"
	"kubevirt.io/kubevirt/tests/libstorage"
	"kubevirt.io/kubevirt/tests/util"
)

// NamespaceTestAlternative is used to test controller-namespace independently.
var NamespaceTestAlternative = "kubevirt-test-alternative"

// NamespaceTestOperator is used to test if namespaces can still be deleted when kubevirt is uninstalled
var NamespaceTestOperator = "kubevirt-test-operator"

var TestNamespaces = []string{util.NamespaceTestDefault, NamespaceTestAlternative, NamespaceTestOperator}

func CleanNamespaces() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)

	for _, namespace := range TestNamespaces {

		_, err := virtCli.CoreV1().Namespaces().Get(context.Background(), namespace, metav1.GetOptions{})
		if err != nil {
			continue
		}

		// Clean namespace labels
		err = libnet.RemoveAllLabelsFromNamespace(virtCli, namespace)
		util.PanicOnError(err)

		//Remove all Jobs
		jobDeleteStrategy := metav1.DeletePropagationOrphan
		jobDeleteOptions := metav1.DeleteOptions{PropagationPolicy: &jobDeleteStrategy}
		util.PanicOnError(virtCli.BatchV1().RESTClient().Delete().Namespace(namespace).Resource("jobs").Body(&jobDeleteOptions).Do(context.Background()).Error())
		//Remove all HPA
		util.PanicOnError(virtCli.AutoscalingV1().RESTClient().Delete().Namespace(namespace).Resource("horizontalpodautoscalers").Do(context.Background()).Error())

		// Remove all VirtualMachinePools
		util.PanicOnError(virtCli.VirtualMachinePool(namespace).DeleteCollection(context.Background(), metav1.DeleteOptions{}, metav1.ListOptions{}))

		// Remove all VirtualMachines
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachines").Do(context.Background()).Error())

		// Remove all VirtualMachineReplicaSets
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstancereplicasets").Do(context.Background()).Error())

		// Remove all VMIs
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstances").Do(context.Background()).Error())
		vmis, err := virtCli.VirtualMachineInstance(namespace).List(&metav1.ListOptions{})
		util.PanicOnError(err)
		for _, vmi := range vmis.Items {
			if controller.HasFinalizer(&vmi, v1.VirtualMachineInstanceFinalizer) {
				_, err := virtCli.VirtualMachineInstance(vmi.Namespace).Patch(vmi.Name, types.JSONPatchType, []byte("[{ \"op\": \"remove\", \"path\": \"/metadata/finalizers\" }]"), &metav1.PatchOptions{})
				if !errors.IsNotFound(err) {
					util.PanicOnError(err)
				}
			}
		}

		// Remove all Pods
		podList, err := virtCli.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{})
		util.PanicOnError(err)
		var gracePeriod int64 = 0
		for _, pod := range podList.Items {
			err := virtCli.CoreV1().Pods(namespace).Delete(context.Background(), pod.Name, metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod})
			if errors.IsNotFound(err) {
				continue
			}
			Expect(err).ToNot(HaveOccurred())
		}

		// Remove all Services
		svcList, err := virtCli.CoreV1().Services(namespace).List(context.Background(), metav1.ListOptions{})
		util.PanicOnError(err)
		for _, svc := range svcList.Items {
			util.PanicOnError(virtCli.CoreV1().Services(namespace).Delete(context.Background(), svc.Name, metav1.DeleteOptions{}))
		}

		// Remove PVCs
		util.PanicOnError(virtCli.CoreV1().RESTClient().Delete().Namespace(namespace).Resource("persistentvolumeclaims").Do(context.Background()).Error())
		if libstorage.HasCDI() {
			// Remove DataVolumes
			util.PanicOnError(virtCli.CdiClient().CdiV1beta1().RESTClient().Delete().Namespace(namespace).Resource("datavolumes").Do(context.Background()).Error())
		}
		// Remove PVs
		pvs, err := virtCli.CoreV1().PersistentVolumes().List(context.Background(), metav1.ListOptions{
			LabelSelector: fmt.Sprintf("%s", cleanup.TestLabelForNamespace(namespace)),
		})
		util.PanicOnError(err)
		for _, pv := range pvs.Items {
			err := virtCli.CoreV1().PersistentVolumes().Delete(context.Background(), pv.Name, metav1.DeleteOptions{})
			if err != nil && !errors.IsNotFound(err) {
				util.PanicOnError(err)
			}
		}

		// Remove all VirtualMachineInstance Secrets
		labelSelector := fmt.Sprintf("%s", SecretLabel)
		util.PanicOnError(
			virtCli.CoreV1().Secrets(namespace).DeleteCollection(context.Background(),
				metav1.DeleteOptions{}, metav1.ListOptions{LabelSelector: labelSelector},
			),
		)

		// Remove all VirtualMachineInstance Presets
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstancepresets").Do(context.Background()).Error())
		// Remove all limit ranges
		util.PanicOnError(virtCli.CoreV1().RESTClient().Delete().Namespace(namespace).Resource("limitranges").Do(context.Background()).Error())

		// Remove all Migration Objects
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstancemigrations").Do(context.Background()).Error())
		migrations, err := virtCli.VirtualMachineInstanceMigration(namespace).List(&metav1.ListOptions{})
		util.PanicOnError(err)
		for _, migration := range migrations.Items {
			if controller.HasFinalizer(&migration, v1.VirtualMachineInstanceMigrationFinalizer) {
				_, err := virtCli.VirtualMachineInstanceMigration(namespace).Patch(migration.Name, types.JSONPatchType, []byte("[{ \"op\": \"remove\", \"path\": \"/metadata/finalizers\" }]"))
				if !errors.IsNotFound(err) {
					util.PanicOnError(err)
				}
			}
		}
		// Remove all NetworkAttachmentDefinitions
		nets, err := virtCli.NetworkClient().K8sCniCncfIoV1().NetworkAttachmentDefinitions(namespace).List(context.Background(), metav1.ListOptions{})
		if err != nil && !errors.IsNotFound(err) {
			util.PanicOnError(err)
		}
		for _, netDef := range nets.Items {
			util.PanicOnError(virtCli.NetworkClient().K8sCniCncfIoV1().NetworkAttachmentDefinitions(namespace).Delete(context.Background(), netDef.GetName(), metav1.DeleteOptions{}))
		}

		// Remove all Istio Sidecars, VirtualServices, DestinationRules and Gateways
		for _, res := range []string{"sidecars", "virtualservices", "destinationrules", "gateways"} {
			util.PanicOnError(removeAllGroupVersionResourceFromNamespace(schema.GroupVersionResource{Group: "networking.istio.io", Version: "v1beta1", Resource: res}, namespace))
		}

		// Remove all Istio PeerAuthentications
		util.PanicOnError(removeAllGroupVersionResourceFromNamespace(schema.GroupVersionResource{Group: "security.istio.io", Version: "v1beta1", Resource: "peerauthentications"}, namespace))

		// Remove migration policies
		migrationPolicyList, err := virtCli.MigrationPolicy().List(context.Background(), metav1.ListOptions{
			LabelSelector: fmt.Sprintf("%s", cleanup.TestLabelForNamespace(namespace)),
		})
		util.PanicOnError(err)
		for _, policy := range migrationPolicyList.Items {
			util.PanicOnError(virtCli.MigrationPolicy().Delete(context.Background(), policy.Name, metav1.DeleteOptions{}))
		}
	}
}

func removeNamespaces() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)

	// First send an initial delete to every namespace
	for _, namespace := range TestNamespaces {
		err := virtCli.CoreV1().Namespaces().Delete(context.Background(), namespace, metav1.DeleteOptions{})
		if !errors.IsNotFound(err) {
			util.PanicOnError(err)
		}
	}

	// Wait until the namespaces are terminated
	fmt.Println("")
	for _, namespace := range TestNamespaces {
		fmt.Printf("Waiting for namespace %s to be removed, this can take a while ...\n", namespace)
		EventuallyWithOffset(1, func() error {
			return virtCli.CoreV1().Namespaces().Delete(context.Background(), namespace, metav1.DeleteOptions{})
		}, 240*time.Second, 1*time.Second).Should(SatisfyAll(HaveOccurred(), WithTransform(errors.IsNotFound, BeTrue())), fmt.Sprintf("should successfully delete namespace '%s'", namespace))
	}
}

func removeAllGroupVersionResourceFromNamespace(groupVersionResource schema.GroupVersionResource, namespace string) error {
	virtCli, err := kubecli.GetKubevirtClient()
	if err != nil {
		return err
	}

	gvr, err := virtCli.DynamicClient().Resource(groupVersionResource).Namespace(namespace).List(context.Background(), metav1.ListOptions{})
	if errors.IsNotFound(err) {
		return nil
	}
	if err != nil {
		return err
	}

	for _, r := range gvr.Items {
		err = virtCli.DynamicClient().Resource(groupVersionResource).Namespace(namespace).Delete(context.Background(), r.GetName(), metav1.DeleteOptions{})
		if err != nil {
			return err
		}
	}
	return nil
}

func detectInstallNamespace() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)
	kvs, err := virtCli.KubeVirt("").List(&metav1.ListOptions{})
	util.PanicOnError(err)
	if len(kvs.Items) == 0 {
		util.PanicOnError(fmt.Errorf("Could not detect a kubevirt installation"))
	}
	if len(kvs.Items) > 1 {
		util.PanicOnError(fmt.Errorf("Invalid kubevirt installation, more than one KubeVirt resource found"))
	}
	flags.KubeVirtInstallNamespace = kvs.Items[0].Namespace
}

func createNamespaces() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)

	// Create a Test Namespaces
	for _, namespace := range TestNamespaces {
		ns := &k8sv1.Namespace{
			ObjectMeta: metav1.ObjectMeta{
				Name: namespace,
				Labels: map[string]string{
					cleanup.TestLabelForNamespace(namespace): "",
				},
			},
		}
		_, err = virtCli.CoreV1().Namespaces().Create(context.Background(), ns, metav1.CreateOptions{})
		if !errors.IsAlreadyExists(err) {
			util.PanicOnError(err)
		}
	}
}

// CalculateNamespaces checks on which ginkgo gest node the tests are run and sets the namespaces accordingly
func CalculateNamespaces() {
	worker := GinkgoParallelProcess()
	util.NamespaceTestDefault = fmt.Sprintf("%s%d", util.NamespaceTestDefault, worker)
	NamespaceTestAlternative = fmt.Sprintf("%s%d", NamespaceTestAlternative, worker)
	// TODO, that is not needed, just a shortcut to not have to treat this namespace
	// differently when running in parallel
	NamespaceTestOperator = fmt.Sprintf("%s%d", NamespaceTestOperator, worker)
	TestNamespaces = []string{util.NamespaceTestDefault, NamespaceTestAlternative, NamespaceTestOperator}
}
<commit_msg>test, Fail on namespace creation failure<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 2022 Red Hat, Inc.
 *
 */

package testsuite

import (
	"context"
	"fmt"
	"time"

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

	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/pkg/controller"

	"kubevirt.io/kubevirt/tests/flags"
	"kubevirt.io/kubevirt/tests/framework/cleanup"
	"kubevirt.io/kubevirt/tests/libnet"
	"kubevirt.io/kubevirt/tests/libstorage"
	"kubevirt.io/kubevirt/tests/util"
)

// NamespaceTestAlternative is used to test controller-namespace independently.
var NamespaceTestAlternative = "kubevirt-test-alternative"

// NamespaceTestOperator is used to test if namespaces can still be deleted when kubevirt is uninstalled
var NamespaceTestOperator = "kubevirt-test-operator"

var TestNamespaces = []string{util.NamespaceTestDefault, NamespaceTestAlternative, NamespaceTestOperator}

func CleanNamespaces() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)

	for _, namespace := range TestNamespaces {

		_, err := virtCli.CoreV1().Namespaces().Get(context.Background(), namespace, metav1.GetOptions{})
		if err != nil {
			continue
		}

		// Clean namespace labels
		err = libnet.RemoveAllLabelsFromNamespace(virtCli, namespace)
		util.PanicOnError(err)

		//Remove all Jobs
		jobDeleteStrategy := metav1.DeletePropagationOrphan
		jobDeleteOptions := metav1.DeleteOptions{PropagationPolicy: &jobDeleteStrategy}
		util.PanicOnError(virtCli.BatchV1().RESTClient().Delete().Namespace(namespace).Resource("jobs").Body(&jobDeleteOptions).Do(context.Background()).Error())
		//Remove all HPA
		util.PanicOnError(virtCli.AutoscalingV1().RESTClient().Delete().Namespace(namespace).Resource("horizontalpodautoscalers").Do(context.Background()).Error())

		// Remove all VirtualMachinePools
		util.PanicOnError(virtCli.VirtualMachinePool(namespace).DeleteCollection(context.Background(), metav1.DeleteOptions{}, metav1.ListOptions{}))

		// Remove all VirtualMachines
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachines").Do(context.Background()).Error())

		// Remove all VirtualMachineReplicaSets
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstancereplicasets").Do(context.Background()).Error())

		// Remove all VMIs
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstances").Do(context.Background()).Error())
		vmis, err := virtCli.VirtualMachineInstance(namespace).List(&metav1.ListOptions{})
		util.PanicOnError(err)
		for _, vmi := range vmis.Items {
			if controller.HasFinalizer(&vmi, v1.VirtualMachineInstanceFinalizer) {
				_, err := virtCli.VirtualMachineInstance(vmi.Namespace).Patch(vmi.Name, types.JSONPatchType, []byte("[{ \"op\": \"remove\", \"path\": \"/metadata/finalizers\" }]"), &metav1.PatchOptions{})
				if !errors.IsNotFound(err) {
					util.PanicOnError(err)
				}
			}
		}

		// Remove all Pods
		podList, err := virtCli.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{})
		util.PanicOnError(err)
		var gracePeriod int64 = 0
		for _, pod := range podList.Items {
			err := virtCli.CoreV1().Pods(namespace).Delete(context.Background(), pod.Name, metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod})
			if errors.IsNotFound(err) {
				continue
			}
			Expect(err).ToNot(HaveOccurred())
		}

		// Remove all Services
		svcList, err := virtCli.CoreV1().Services(namespace).List(context.Background(), metav1.ListOptions{})
		util.PanicOnError(err)
		for _, svc := range svcList.Items {
			util.PanicOnError(virtCli.CoreV1().Services(namespace).Delete(context.Background(), svc.Name, metav1.DeleteOptions{}))
		}

		// Remove PVCs
		util.PanicOnError(virtCli.CoreV1().RESTClient().Delete().Namespace(namespace).Resource("persistentvolumeclaims").Do(context.Background()).Error())
		if libstorage.HasCDI() {
			// Remove DataVolumes
			util.PanicOnError(virtCli.CdiClient().CdiV1beta1().RESTClient().Delete().Namespace(namespace).Resource("datavolumes").Do(context.Background()).Error())
		}
		// Remove PVs
		pvs, err := virtCli.CoreV1().PersistentVolumes().List(context.Background(), metav1.ListOptions{
			LabelSelector: fmt.Sprintf("%s", cleanup.TestLabelForNamespace(namespace)),
		})
		util.PanicOnError(err)
		for _, pv := range pvs.Items {
			err := virtCli.CoreV1().PersistentVolumes().Delete(context.Background(), pv.Name, metav1.DeleteOptions{})
			if err != nil && !errors.IsNotFound(err) {
				util.PanicOnError(err)
			}
		}

		// Remove all VirtualMachineInstance Secrets
		labelSelector := fmt.Sprintf("%s", SecretLabel)
		util.PanicOnError(
			virtCli.CoreV1().Secrets(namespace).DeleteCollection(context.Background(),
				metav1.DeleteOptions{}, metav1.ListOptions{LabelSelector: labelSelector},
			),
		)

		// Remove all VirtualMachineInstance Presets
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstancepresets").Do(context.Background()).Error())
		// Remove all limit ranges
		util.PanicOnError(virtCli.CoreV1().RESTClient().Delete().Namespace(namespace).Resource("limitranges").Do(context.Background()).Error())

		// Remove all Migration Objects
		util.PanicOnError(virtCli.RestClient().Delete().Namespace(namespace).Resource("virtualmachineinstancemigrations").Do(context.Background()).Error())
		migrations, err := virtCli.VirtualMachineInstanceMigration(namespace).List(&metav1.ListOptions{})
		util.PanicOnError(err)
		for _, migration := range migrations.Items {
			if controller.HasFinalizer(&migration, v1.VirtualMachineInstanceMigrationFinalizer) {
				_, err := virtCli.VirtualMachineInstanceMigration(namespace).Patch(migration.Name, types.JSONPatchType, []byte("[{ \"op\": \"remove\", \"path\": \"/metadata/finalizers\" }]"))
				if !errors.IsNotFound(err) {
					util.PanicOnError(err)
				}
			}
		}
		// Remove all NetworkAttachmentDefinitions
		nets, err := virtCli.NetworkClient().K8sCniCncfIoV1().NetworkAttachmentDefinitions(namespace).List(context.Background(), metav1.ListOptions{})
		if err != nil && !errors.IsNotFound(err) {
			util.PanicOnError(err)
		}
		for _, netDef := range nets.Items {
			util.PanicOnError(virtCli.NetworkClient().K8sCniCncfIoV1().NetworkAttachmentDefinitions(namespace).Delete(context.Background(), netDef.GetName(), metav1.DeleteOptions{}))
		}

		// Remove all Istio Sidecars, VirtualServices, DestinationRules and Gateways
		for _, res := range []string{"sidecars", "virtualservices", "destinationrules", "gateways"} {
			util.PanicOnError(removeAllGroupVersionResourceFromNamespace(schema.GroupVersionResource{Group: "networking.istio.io", Version: "v1beta1", Resource: res}, namespace))
		}

		// Remove all Istio PeerAuthentications
		util.PanicOnError(removeAllGroupVersionResourceFromNamespace(schema.GroupVersionResource{Group: "security.istio.io", Version: "v1beta1", Resource: "peerauthentications"}, namespace))

		// Remove migration policies
		migrationPolicyList, err := virtCli.MigrationPolicy().List(context.Background(), metav1.ListOptions{
			LabelSelector: fmt.Sprintf("%s", cleanup.TestLabelForNamespace(namespace)),
		})
		util.PanicOnError(err)
		for _, policy := range migrationPolicyList.Items {
			util.PanicOnError(virtCli.MigrationPolicy().Delete(context.Background(), policy.Name, metav1.DeleteOptions{}))
		}
	}
}

func removeNamespaces() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)

	// First send an initial delete to every namespace
	for _, namespace := range TestNamespaces {
		err := virtCli.CoreV1().Namespaces().Delete(context.Background(), namespace, metav1.DeleteOptions{})
		if !errors.IsNotFound(err) {
			util.PanicOnError(err)
		}
	}

	// Wait until the namespaces are terminated
	fmt.Println("")
	for _, namespace := range TestNamespaces {
		fmt.Printf("Waiting for namespace %s to be removed, this can take a while ...\n", namespace)
		EventuallyWithOffset(1, func() error {
			return virtCli.CoreV1().Namespaces().Delete(context.Background(), namespace, metav1.DeleteOptions{})
		}, 240*time.Second, 1*time.Second).Should(SatisfyAll(HaveOccurred(), WithTransform(errors.IsNotFound, BeTrue())), fmt.Sprintf("should successfully delete namespace '%s'", namespace))
	}
}

func removeAllGroupVersionResourceFromNamespace(groupVersionResource schema.GroupVersionResource, namespace string) error {
	virtCli, err := kubecli.GetKubevirtClient()
	if err != nil {
		return err
	}

	gvr, err := virtCli.DynamicClient().Resource(groupVersionResource).Namespace(namespace).List(context.Background(), metav1.ListOptions{})
	if errors.IsNotFound(err) {
		return nil
	}
	if err != nil {
		return err
	}

	for _, r := range gvr.Items {
		err = virtCli.DynamicClient().Resource(groupVersionResource).Namespace(namespace).Delete(context.Background(), r.GetName(), metav1.DeleteOptions{})
		if err != nil {
			return err
		}
	}
	return nil
}

func detectInstallNamespace() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)
	kvs, err := virtCli.KubeVirt("").List(&metav1.ListOptions{})
	util.PanicOnError(err)
	if len(kvs.Items) == 0 {
		util.PanicOnError(fmt.Errorf("Could not detect a kubevirt installation"))
	}
	if len(kvs.Items) > 1 {
		util.PanicOnError(fmt.Errorf("Invalid kubevirt installation, more than one KubeVirt resource found"))
	}
	flags.KubeVirtInstallNamespace = kvs.Items[0].Namespace
}

func createNamespaces() {
	virtCli, err := kubecli.GetKubevirtClient()
	util.PanicOnError(err)

	// Create a Test Namespaces
	for _, namespace := range TestNamespaces {
		ns := &k8sv1.Namespace{
			ObjectMeta: metav1.ObjectMeta{
				Name: namespace,
				Labels: map[string]string{
					cleanup.TestLabelForNamespace(namespace): "",
				},
			},
		}
		_, err = virtCli.CoreV1().Namespaces().Create(context.Background(), ns, metav1.CreateOptions{})
		if err != nil {
			util.PanicOnError(err)
		}
	}
}

// CalculateNamespaces checks on which ginkgo gest node the tests are run and sets the namespaces accordingly
func CalculateNamespaces() {
	worker := GinkgoParallelProcess()
	util.NamespaceTestDefault = fmt.Sprintf("%s%d", util.NamespaceTestDefault, worker)
	NamespaceTestAlternative = fmt.Sprintf("%s%d", NamespaceTestAlternative, worker)
	// TODO, that is not needed, just a shortcut to not have to treat this namespace
	// differently when running in parallel
	NamespaceTestOperator = fmt.Sprintf("%s%d", NamespaceTestOperator, worker)
	TestNamespaces = []string{util.NamespaceTestDefault, NamespaceTestAlternative, NamespaceTestOperator}
}
<|endoftext|>


<commit_before>package of10

import . "github.com/oshothebig/goflow/openflow"

type ActionType uint16

const (
	OFPAT_OUTPUT ActionType = iota
	OFPAT_SET_VLAN_VID
	OFPAT_SET_VLAN_PCP
	OFPAT_SET_STRIP_VLAN
	OFPAT_SET_DL_SRC
	OFPAT_SET_DL_DST
	OFPAT_SET_NW_SRC
	OFPAT_SET_NW_DST
	OFPAT_SET_NW_TOS
	OFPAT_SET_TP_SRC
	OFPAT_SET_TP_DST
	OFPAT_ENQUEUE
	OFPAT_VENDOR ActionType = 0xffff
)

var ActionTypes = struct {
	Output        ActionType
	SetVlanId     ActionType
	SetVlanPcp    ActionType
	StripVlan     ActionType
	SetEtherSrc   ActionType
	SetEtherDst   ActionType
	SetIpSrc      ActionType
	SetIpDst      ActionType
	SetIpTos      ActionType
	SetNetworkSrc ActionType
	SetNetworkDst ActionType
	Enqueue       ActionType
	Vendor        ActionType
}{
	OFPAT_OUTPUT,
	OFPAT_SET_VLAN_VID,
	OFPAT_SET_VLAN_PCP,
	OFPAT_SET_STRIP_VLAN,
	OFPAT_SET_DL_SRC,
	OFPAT_SET_DL_DST,
	OFPAT_SET_NW_SRC,
	OFPAT_SET_NW_DST,
	OFPAT_SET_NW_TOS,
	OFPAT_SET_TP_SRC,
	OFPAT_SET_TP_DST,
	OFPAT_ENQUEUE,
	OFPAT_VENDOR,
}

type Action interface {
	Packetizable
	GetType() ActionType
}

type ActionHeader struct {
	Type   ActionType
	Length uint16
	pad    [4]uint8
}
<commit_msg>Define GetType() function in ActionHeader<commit_after>package of10

import . "github.com/oshothebig/goflow/openflow"

type ActionType uint16

const (
	OFPAT_OUTPUT ActionType = iota
	OFPAT_SET_VLAN_VID
	OFPAT_SET_VLAN_PCP
	OFPAT_SET_STRIP_VLAN
	OFPAT_SET_DL_SRC
	OFPAT_SET_DL_DST
	OFPAT_SET_NW_SRC
	OFPAT_SET_NW_DST
	OFPAT_SET_NW_TOS
	OFPAT_SET_TP_SRC
	OFPAT_SET_TP_DST
	OFPAT_ENQUEUE
	OFPAT_VENDOR ActionType = 0xffff
)

var ActionTypes = struct {
	Output        ActionType
	SetVlanId     ActionType
	SetVlanPcp    ActionType
	StripVlan     ActionType
	SetEtherSrc   ActionType
	SetEtherDst   ActionType
	SetIpSrc      ActionType
	SetIpDst      ActionType
	SetIpTos      ActionType
	SetNetworkSrc ActionType
	SetNetworkDst ActionType
	Enqueue       ActionType
	Vendor        ActionType
}{
	OFPAT_OUTPUT,
	OFPAT_SET_VLAN_VID,
	OFPAT_SET_VLAN_PCP,
	OFPAT_SET_STRIP_VLAN,
	OFPAT_SET_DL_SRC,
	OFPAT_SET_DL_DST,
	OFPAT_SET_NW_SRC,
	OFPAT_SET_NW_DST,
	OFPAT_SET_NW_TOS,
	OFPAT_SET_TP_SRC,
	OFPAT_SET_TP_DST,
	OFPAT_ENQUEUE,
	OFPAT_VENDOR,
}

type Action interface {
	Packetizable
	GetType() ActionType
}

type ActionHeader struct {
	Type   ActionType
	Length uint16
	pad    [4]uint8
}

func (header *ActionHeader) GetType() ActionType {
	return header.Type
}
<|endoftext|>


<commit_before>package openssh_test

import (
	"testing"

	"github.com/mlafeldt/chef-runner/openssh"
	"github.com/stretchr/testify/assert"
)

var newClientTests = []struct {
	host      string
	client    *openssh.Client
	errString string
}{
	{
		host: "some-host",
		client: &openssh.Client{
			Host: "some-host",
		},
	},
	{
		host: "some-user@some-host",
		client: &openssh.Client{
			Host: "some-host",
			User: "some-user",
		},
	},
	{
		host: "some-host:1234",
		client: &openssh.Client{
			Host: "some-host",
			Port: 1234,
		},
	},
	{
		host: "some-user@some-host:1234",
		client: &openssh.Client{
			Host: "some-host",
			User: "some-user",
			Port: 1234,
		},
	},
	{
		host:      "some-host:abc",
		client:    nil,
		errString: "invalid SSH port",
	},
}

func TestNewClient(t *testing.T) {
	for _, test := range newClientTests {
		c, err := openssh.NewClient(test.host)
		if test.errString == "" {
			assert.NoError(t, err)
		} else {
			assert.EqualError(t, err, test.errString)
		}
		assert.Equal(t, test.client, c)
	}
}

func TestString(t *testing.T) {
	expect := "OpenSSH (host: some-host)"
	c, err := openssh.NewClient("some-host")
	assert.NoError(t, err)
	assert.Equal(t, expect, c.String())
}

var commandTests = []struct {
	client openssh.Client
	cmd    []string
}{
	{
		openssh.Client{
			Host: "some-host",
		},
		[]string{"ssh", "some-host", "uname -a"},
	},
	{
		openssh.Client{
			Host: "some-host",
			User: "some-user",
		},
		[]string{"ssh", "-l", "some-user", "some-host", "uname -a"},
	},
	{
		openssh.Client{
			Host: "some-host",
			Port: 1234,
		},
		[]string{"ssh", "-p", "1234", "some-host", "uname -a"},
	},
	{
		openssh.Client{
			Host:        "some-host",
			PrivateKeys: []string{"some-key", "another-key"},
		},
		[]string{"ssh", "-i", "some-key", "-i", "another-key",
			"some-host", "uname -a"},
	},
	{
		openssh.Client{
			Host: "some-host",
			Options: map[string]string{
				"SomeOption":    "yes",
				"AnotherOption": "no",
			},
		},
		[]string{"ssh", "-o", "AnotherOption=no", "-o", "SomeOption=yes",
			"some-host", "uname -a"},
	},
	{
		openssh.Client{
			Host:        "some-host",
			User:        "some-user",
			Port:        1234,
			PrivateKeys: []string{"some-key"},
			Options:     map[string]string{"SomeOption": "yes"},
		},
		[]string{"ssh", "-l", "some-user", "-p", "1234", "-i", "some-key",
			"-o", "SomeOption=yes", "some-host", "uname -a"},
	},
	{
		openssh.Client{
			Host:       "some-host",
			ConfigFile: "some/config/file",
		},
		[]string{"ssh", "-F", "some/config/file", "some-host", "uname -a"},
	},
}

func TestCommand(t *testing.T) {
	for _, test := range commandTests {
		cmd, err := test.client.Command("uname -a")
		if assert.NoError(t, err) {
			assert.Equal(t, test.cmd, cmd)
		}
	}
}

func TestRunCommand_MissingCommand(t *testing.T) {
	err := openssh.Client{Host: "some-host"}.RunCommand("")
	assert.EqualError(t, err, "no command given")
}

func TestRunCommand_MissingHost(t *testing.T) {
	err := openssh.Client{}.RunCommand("uname -a")
	assert.EqualError(t, err, "no host given")
}
<commit_msg>Consolidate openssh tests<commit_after>package openssh_test

import (
	"testing"

	"github.com/mlafeldt/chef-runner/openssh"
	"github.com/stretchr/testify/assert"
)

var newClientTests = []struct {
	host      string
	client    *openssh.Client
	errString string
}{
	{
		host: "some-host",
		client: &openssh.Client{
			Host: "some-host",
		},
	},
	{
		host: "some-user@some-host",
		client: &openssh.Client{
			Host: "some-host",
			User: "some-user",
		},
	},
	{
		host: "some-host:1234",
		client: &openssh.Client{
			Host: "some-host",
			Port: 1234,
		},
	},
	{
		host: "some-user@some-host:1234",
		client: &openssh.Client{
			Host: "some-host",
			User: "some-user",
			Port: 1234,
		},
	},
	// Check for errors
	{
		host:      "some-host:abc",
		client:    nil,
		errString: "invalid SSH port",
	},
}

func TestNewClient(t *testing.T) {
	for _, test := range newClientTests {
		c, err := openssh.NewClient(test.host)
		if test.errString == "" {
			assert.NoError(t, err)
		} else {
			assert.EqualError(t, err, test.errString)
		}
		assert.Equal(t, test.client, c)
	}
}

func TestString(t *testing.T) {
	expect := "OpenSSH (host: some-host)"
	c, _ := openssh.NewClient("some-host")
	assert.Equal(t, expect, c.String())
}

var commandTests = []struct {
	client    openssh.Client
	cmd       string
	result    []string
	errString string
}{
	{
		client: openssh.Client{
			Host: "some-host",
		},
		cmd:    "uname -a",
		result: []string{"ssh", "some-host", "uname -a"},
	},
	{
		client: openssh.Client{
			Host: "some-host",
			User: "some-user",
		},
		cmd:    "uname -a",
		result: []string{"ssh", "-l", "some-user", "some-host", "uname -a"},
	},
	{
		client: openssh.Client{
			Host: "some-host",
			Port: 1234,
		},
		cmd:    "uname -a",
		result: []string{"ssh", "-p", "1234", "some-host", "uname -a"},
	},
	{
		client: openssh.Client{
			Host:        "some-host",
			PrivateKeys: []string{"some-key", "another-key"},
		},
		cmd: "uname -a",
		result: []string{"ssh", "-i", "some-key", "-i", "another-key",
			"some-host", "uname -a"},
	},
	{
		client: openssh.Client{
			Host: "some-host",
			Options: map[string]string{
				"SomeOption":    "yes",
				"AnotherOption": "no",
			},
		},
		cmd: "uname -a",
		result: []string{"ssh", "-o", "AnotherOption=no",
			"-o", "SomeOption=yes", "some-host", "uname -a"},
	},
	{
		client: openssh.Client{
			Host:        "some-host",
			User:        "some-user",
			Port:        1234,
			PrivateKeys: []string{"some-key"},
			Options:     map[string]string{"SomeOption": "yes"},
		},
		cmd: "uname -a",
		result: []string{"ssh", "-l", "some-user", "-p", "1234",
			"-i", "some-key", "-o", "SomeOption=yes", "some-host",
			"uname -a"},
	},
	{
		client: openssh.Client{
			Host:       "some-host",
			ConfigFile: "some/config/file",
		},
		cmd: "uname -a",
		result: []string{"ssh", "-F", "some/config/file", "some-host",
			"uname -a"},
	},
	// Check for errors
	{
		client:    openssh.Client{},
		cmd:       "uname -a",
		result:    nil,
		errString: "no host given",
	},
	{
		client:    openssh.Client{Host: "some-host"},
		cmd:       "",
		result:    nil,
		errString: "no command given",
	},
}

func TestCommand(t *testing.T) {
	for _, test := range commandTests {
		result, err := test.client.Command(test.cmd)
		if test.errString == "" {
			assert.NoError(t, err)
		} else {
			assert.EqualError(t, err, test.errString)
		}
		assert.Equal(t, test.result, result)
	}
}
<|endoftext|>


<commit_before>package blankhost

import (
	"context"
	"io"

	"github.com/libp2p/go-libp2p-core/connmgr"
	"github.com/libp2p/go-libp2p-core/event"
	"github.com/libp2p/go-libp2p-core/host"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/peerstore"
	"github.com/libp2p/go-libp2p-core/protocol"

	"github.com/libp2p/go-eventbus"

	logging "github.com/ipfs/go-log"

	ma "github.com/multiformats/go-multiaddr"
	mstream "github.com/multiformats/go-multistream"
)

var log = logging.Logger("blankhost")

// BlankHost is the thinnest implementation of the host.Host interface
type BlankHost struct {
	n        network.Network
	mux      *mstream.MultistreamMuxer
	cmgr     connmgr.ConnManager
	eventbus event.Bus
	emitters struct {
		evtLocalProtocolsUpdated event.Emitter
	}
}

func NewBlankHost(n network.Network) *BlankHost {
	bh := &BlankHost{
		n:        n,
		cmgr:     &connmgr.NullConnMgr{},
		mux:      mstream.NewMultistreamMuxer(),
		eventbus: eventbus.NewBus(),
	}

	var err error
	if bh.emitters.evtLocalProtocolsUpdated, err = bh.eventbus.Emitter(&event.EvtLocalProtocolsUpdated{}); err != nil {
		return nil
	}

	n.SetStreamHandler(bh.newStreamHandler)
	return bh
}

var _ host.Host = (*BlankHost)(nil)

func (bh *BlankHost) Addrs() []ma.Multiaddr {
	addrs, err := bh.n.InterfaceListenAddresses()
	if err != nil {
		log.Debug("error retrieving network interface addrs: ", err)
		return nil
	}

	return addrs
}

func (bh *BlankHost) Close() error {
	return bh.n.Close()
}

func (bh *BlankHost) Connect(ctx context.Context, ai peer.AddrInfo) error {
	// absorb addresses into peerstore
	bh.Peerstore().AddAddrs(ai.ID, ai.Addrs, peerstore.TempAddrTTL)

	cs := bh.n.ConnsToPeer(ai.ID)
	if len(cs) > 0 {
		return nil
	}

	_, err := bh.Network().DialPeer(ctx, ai.ID)
	return err
}

func (bh *BlankHost) Peerstore() peerstore.Peerstore {
	return bh.n.Peerstore()
}

func (bh *BlankHost) ID() peer.ID {
	return bh.n.LocalPeer()
}

func (bh *BlankHost) NewStream(ctx context.Context, p peer.ID, protos ...protocol.ID) (network.Stream, error) {
	s, err := bh.n.NewStream(ctx, p)
	if err != nil {
		return nil, err
	}

	var protoStrs []string
	for _, pid := range protos {
		protoStrs = append(protoStrs, string(pid))
	}

	selected, err := mstream.SelectOneOf(protoStrs, s)
	if err != nil {
		s.Close()
		return nil, err
	}

	selpid := protocol.ID(selected)
	s.SetProtocol(selpid)
	bh.Peerstore().AddProtocols(p, selected)

	return s, nil
}

func (bh *BlankHost) RemoveStreamHandler(pid protocol.ID) {
	bh.Mux().RemoveHandler(string(pid))
	bh.emitters.evtLocalProtocolsUpdated.Emit(event.EvtLocalProtocolsUpdated{
		Removed: []protocol.ID{pid},
	})
}

func (bh *BlankHost) SetStreamHandler(pid protocol.ID, handler network.StreamHandler) {
	bh.Mux().AddHandler(string(pid), func(p string, rwc io.ReadWriteCloser) error {
		is := rwc.(network.Stream)
		is.SetProtocol(protocol.ID(p))
		handler(is)
		return nil
	})
	bh.emitters.evtLocalProtocolsUpdated.Emit(event.EvtLocalProtocolsUpdated{
		Added: []protocol.ID{pid},
	})
}

func (bh *BlankHost) SetStreamHandlerMatch(pid protocol.ID, m func(string) bool, handler network.StreamHandler) {
	bh.Mux().AddHandlerWithFunc(string(pid), m, func(p string, rwc io.ReadWriteCloser) error {
		is := rwc.(network.Stream)
		is.SetProtocol(protocol.ID(p))
		handler(is)
		return nil
	})
	bh.emitters.evtLocalProtocolsUpdated.Emit(event.EvtLocalProtocolsUpdated{
		Added: []protocol.ID{pid},
	})
}

// newStreamHandler is the remote-opened stream handler for network.Network
func (bh *BlankHost) newStreamHandler(s network.Stream) {
	protoID, handle, err := bh.Mux().Negotiate(s)
	if err != nil {
		log.Warning("protocol mux failed: %s", err)
		s.Close()
		return
	}

	s.SetProtocol(protocol.ID(protoID))

	go handle(protoID, s)
}

// TODO: i'm not sure this really needs to be here
func (bh *BlankHost) Mux() protocol.Switch {
	return bh.mux
}

// TODO: also not sure this fits... Might be better ways around this (leaky abstractions)
func (bh *BlankHost) Network() network.Network {
	return bh.n
}

func (bh *BlankHost) ConnManager() connmgr.ConnManager {
	return bh.cmgr
}

func (bh *BlankHost) EventBus() event.Bus {
	return bh.eventbus
}
<commit_msg>add WithConnectionManager option to blankhost<commit_after>package blankhost

import (
	"context"
	"io"

	"github.com/libp2p/go-libp2p-core/connmgr"
	"github.com/libp2p/go-libp2p-core/event"
	"github.com/libp2p/go-libp2p-core/host"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/peerstore"
	"github.com/libp2p/go-libp2p-core/protocol"

	"github.com/libp2p/go-eventbus"

	logging "github.com/ipfs/go-log"

	ma "github.com/multiformats/go-multiaddr"
	mstream "github.com/multiformats/go-multistream"
)

var log = logging.Logger("blankhost")

// BlankHost is the thinnest implementation of the host.Host interface
type BlankHost struct {
	n        network.Network
	mux      *mstream.MultistreamMuxer
	cmgr     connmgr.ConnManager
	eventbus event.Bus
	emitters struct {
		evtLocalProtocolsUpdated event.Emitter
	}
}

type config struct {
	cmgr connmgr.ConnManager
}

type Option = func(cfg *config)

func WithConnectionManager(cmgr connmgr.ConnManager) Option {
	return func(cfg *config) {
		cfg.cmgr = cmgr
	}
}

func NewBlankHost(n network.Network, options ...Option) *BlankHost {
	cfg := config{
		cmgr: &connmgr.NullConnMgr{},
	}
	for _, opt := range options {
		opt(&cfg)
	}

	bh := &BlankHost{
		n:        n,
		cmgr:     cfg.cmgr,
		mux:      mstream.NewMultistreamMuxer(),
		eventbus: eventbus.NewBus(),
	}

	var err error
	if bh.emitters.evtLocalProtocolsUpdated, err = bh.eventbus.Emitter(&event.EvtLocalProtocolsUpdated{}); err != nil {
		return nil
	}

	n.SetStreamHandler(bh.newStreamHandler)
	return bh
}

var _ host.Host = (*BlankHost)(nil)

func (bh *BlankHost) Addrs() []ma.Multiaddr {
	addrs, err := bh.n.InterfaceListenAddresses()
	if err != nil {
		log.Debug("error retrieving network interface addrs: ", err)
		return nil
	}

	return addrs
}

func (bh *BlankHost) Close() error {
	return bh.n.Close()
}

func (bh *BlankHost) Connect(ctx context.Context, ai peer.AddrInfo) error {
	// absorb addresses into peerstore
	bh.Peerstore().AddAddrs(ai.ID, ai.Addrs, peerstore.TempAddrTTL)

	cs := bh.n.ConnsToPeer(ai.ID)
	if len(cs) > 0 {
		return nil
	}

	_, err := bh.Network().DialPeer(ctx, ai.ID)
	return err
}

func (bh *BlankHost) Peerstore() peerstore.Peerstore {
	return bh.n.Peerstore()
}

func (bh *BlankHost) ID() peer.ID {
	return bh.n.LocalPeer()
}

func (bh *BlankHost) NewStream(ctx context.Context, p peer.ID, protos ...protocol.ID) (network.Stream, error) {
	s, err := bh.n.NewStream(ctx, p)
	if err != nil {
		return nil, err
	}

	var protoStrs []string
	for _, pid := range protos {
		protoStrs = append(protoStrs, string(pid))
	}

	selected, err := mstream.SelectOneOf(protoStrs, s)
	if err != nil {
		s.Close()
		return nil, err
	}

	selpid := protocol.ID(selected)
	s.SetProtocol(selpid)
	bh.Peerstore().AddProtocols(p, selected)

	return s, nil
}

func (bh *BlankHost) RemoveStreamHandler(pid protocol.ID) {
	bh.Mux().RemoveHandler(string(pid))
	bh.emitters.evtLocalProtocolsUpdated.Emit(event.EvtLocalProtocolsUpdated{
		Removed: []protocol.ID{pid},
	})
}

func (bh *BlankHost) SetStreamHandler(pid protocol.ID, handler network.StreamHandler) {
	bh.Mux().AddHandler(string(pid), func(p string, rwc io.ReadWriteCloser) error {
		is := rwc.(network.Stream)
		is.SetProtocol(protocol.ID(p))
		handler(is)
		return nil
	})
	bh.emitters.evtLocalProtocolsUpdated.Emit(event.EvtLocalProtocolsUpdated{
		Added: []protocol.ID{pid},
	})
}

func (bh *BlankHost) SetStreamHandlerMatch(pid protocol.ID, m func(string) bool, handler network.StreamHandler) {
	bh.Mux().AddHandlerWithFunc(string(pid), m, func(p string, rwc io.ReadWriteCloser) error {
		is := rwc.(network.Stream)
		is.SetProtocol(protocol.ID(p))
		handler(is)
		return nil
	})
	bh.emitters.evtLocalProtocolsUpdated.Emit(event.EvtLocalProtocolsUpdated{
		Added: []protocol.ID{pid},
	})
}

// newStreamHandler is the remote-opened stream handler for network.Network
func (bh *BlankHost) newStreamHandler(s network.Stream) {
	protoID, handle, err := bh.Mux().Negotiate(s)
	if err != nil {
		log.Warning("protocol mux failed: %s", err)
		s.Close()
		return
	}

	s.SetProtocol(protocol.ID(protoID))

	go handle(protoID, s)
}

// TODO: i'm not sure this really needs to be here
func (bh *BlankHost) Mux() protocol.Switch {
	return bh.mux
}

// TODO: also not sure this fits... Might be better ways around this (leaky abstractions)
func (bh *BlankHost) Network() network.Network {
	return bh.n
}

func (bh *BlankHost) ConnManager() connmgr.ConnManager {
	return bh.cmgr
}

func (bh *BlankHost) EventBus() event.Bus {
	return bh.eventbus
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"github.com/skeswa/gophr/common"
	"io/ioutil"
	"log"
	"strings"
	"sync"
	"time"
)

type GoPackage struct {
	Description    string
	GitHubURL      string
	GoDocURL       string
	IndexTime      string
	HttpStatusCode int
	AwesomeGo      bool
	Versions       []string
}

func main() {
	start := time.Now()
	log.Println("Started Download of Godoc/index")

	doc, err := goquery.NewDocument("https://godoc.org/-/index")
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Finished Download of Godoc/index\n")
	log.Println("Started scraping all github packages from Godoc/index")

	var goPackageMap = make(map[string]*GoPackage)

	// For each tr element on the page
	doc.Find("tr").Each(func(i int, s *goquery.Selection) {
		children := s.Children()
		goPackage := new(GoPackage)

		// For each child in the tr element
		children.Each(func(i int, s2 *goquery.Selection) {
			childURL, childURLexists := s.Find("a").Attr("href")
			childDescription := s.Text()

			if childURLexists == true {
				goPackage.GitHubURL = strings.Trim(childURL, "/")
			}

			if len(childDescription) > 0 {
				// TODO check if description isn't just the url, if so dont set it
				goPackage.Description = childDescription
			}
		})

		// Only continue if this goPackage contains a GitHub URL
		// TODO check to make sure github.com is the prefix
		gitHubUrlTokens := strings.Split(goPackage.GitHubURL, "/")

		if strings.Contains(goPackage.GitHubURL, "github.com") && len(gitHubUrlTokens) == 3 {
			// Build go doc url
			goPackage.GoDocURL = ("https://godoc.org/" + goPackage.GitHubURL)

			// Create Index Time
			t := time.Now()
			time := t.String()
			goPackage.IndexTime = time

			goPackageMap[goPackage.GitHubURL] = goPackage
		}
	})

	log.Println("Finished scraping ", len(goPackageMap), " github packages from Godoc/index\n")
	log.Println("Started Download of awesome-go/README.md")

	doc, err = goquery.NewDocument("https://godoc.org/-/index")
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Finished Download of awesome-go/README.md\n")
	log.Println("Started scraping awesome-go/README.md")

	doc.Find("a").Each(func(i int, s *goquery.Selection) {
		childURL, _ := s.Attr("href")
		childURL = strings.Trim(childURL, "/")
		goPackage, goPackageExists := goPackageMap[childURL]

		if goPackageExists == true {
			goPackage.AwesomeGo = true
			goPackageMap[childURL] = goPackage
		}
	})

	// TODO ADD A FUCKING COUNT THAT FUCKING WORKS, FUCK
	log.Println("Finished scraping awesome-go/README.md\n")
	log.Println("Started retrieving versions for go packages\n")

	var goErrorPackageMap = make(map[string]*GoPackage)

	// TODO Write a constant for this
	nbConcurrentGet := 20
	urls := make(chan *GoPackage, nbConcurrentGet)
	var wg sync.WaitGroup
	for i := 0; i < nbConcurrentGet; i++ {
		wg.Add(1)
		go func() {
			for goPackage := range urls {
				refs, err := common.FetchRefs(goPackage.GitHubURL)
				if err != nil {
					log.Println("ERROR", goPackage.GitHubURL, " failed to return.\n", err)
					delete(goPackageMap, goPackage.GitHubURL)
					goErrorPackageMap[goPackage.GitHubURL] = goPackage
				} else {
					var versions []string
					for _, version := range refs.Candidates {
						versions = append(versions, version.String())
					}
					goPackage.Versions = versions
					goPackageMap[goPackage.GitHubURL] = goPackage
				}
			}
			wg.Done()
		}()
	}

	for _, goPackage := range goPackageMap {
		urls <- goPackage
	}

	close(urls)
	wg.Wait()

	log.Println("SUCCESS: ", len(goPackageMap), " GoPackages were successfully built")
	log.Println("Creating JSON dump of ", len(goPackageMap), " valid go packages")
	createJSONDump(goPackageMap, "valid-go-packages")
	log.Println("Finished creating JSON dump\n")

	log.Println("WARNING: ", len(goErrorPackageMap), " GoPackages were not found on github")
	log.Println("Creating JSON dump of ", len(goErrorPackageMap), " err packages")
	createJSONDump(goErrorPackageMap, "invalid-go-packages")
	log.Println("Finished creating JSON dump\n")

	elapsed := time.Since(start)
	log.Printf("Program took %s to fully execute", elapsed)
}

// Create a tmp JSON dump of all serialized goPackageData
func createJSONDump(goPackageMap map[string]*GoPackage, fileName string) {
	var buffer bytes.Buffer

	for _, goPackage := range goPackageMap {
		buffer.WriteString("{\"url\": \"" + goPackage.GitHubURL + "\", \"description\": \"" + goPackage.Description + "\", \"index_time\": \"" + goPackage.IndexTime + "\" }, \"versions\": \"" + fmt.Sprintf("%v", goPackage.Versions) + "\" \n")
	}

	t := time.Now()
	time := strings.Replace(t.String(), " ", "", -1)
	jsonData := []byte(buffer.String())
	_ = ioutil.WriteFile("./"+time+"-"+fileName+".json", jsonData, 0644)
}

func printGoPackage(goPackage *GoPackage) {
	fmt.Printf("%v", goPackage)
}
<commit_msg>now we log all the go packages regardless if they're 404 and denote them by their httpstatus<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"github.com/skeswa/gophr/common"
	"io/ioutil"
	"log"
	"strings"
	"sync"
	"time"
)

type GoPackage struct {
	Description    string
	GitHubURL      string
	GoDocURL       string
	IndexTime      string
	HttpStatusCode int
	AwesomeGo      bool
	Versions       []string
}

func main() {
	start := time.Now()
	log.Println("Started Download of Godoc/index")

	doc, err := goquery.NewDocument("https://godoc.org/-/index")
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Finished Download of Godoc/index\n")
	log.Println("Started scraping all github packages from Godoc/index")

	var goPackageMap = make(map[string]*GoPackage)

	// For each tr element on the page
	doc.Find("tr").Each(func(i int, s *goquery.Selection) {
		children := s.Children()
		goPackage := new(GoPackage)

		// For each child in the tr element
		children.Each(func(i int, s2 *goquery.Selection) {
			childURL, childURLexists := s.Find("a").Attr("href")
			childDescription := s.Text()

			if childURLexists == true {
				goPackage.GitHubURL = strings.Trim(childURL, "/")
			}

			if len(childDescription) > 0 {
				// TODO check if description isn't just the url, if so dont set it
				goPackage.Description = childDescription
			}
		})

		// Only continue if this goPackage contains a GitHub URL
		// TODO check to make sure github.com is the prefix
		gitHubURLTokens := strings.Split(goPackage.GitHubURL, "/")

		if strings.Contains(goPackage.GitHubURL, "github.com") && len(gitHubURLTokens) == 3 {
			// Build go doc url
			goPackage.GoDocURL = ("https://godoc.org/" + goPackage.GitHubURL)

			// Create Index Time
			t := time.Now()
			time := t.String()
			goPackage.IndexTime = time

			goPackageMap[goPackage.GitHubURL] = goPackage
		}
	})

	log.Println("Finished scraping ", len(goPackageMap), " github packages from Godoc/index\n")
	log.Println("Started Download of awesome-go/README.md")

	doc, err = goquery.NewDocument("https://godoc.org/-/index")
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Finished Download of awesome-go/README.md\n")
	log.Println("Started scraping awesome-go/README.md")

	doc.Find("a").Each(func(i int, s *goquery.Selection) {
		childURL, _ := s.Attr("href")
		childURL = strings.Trim(childURL, "/")
		goPackage, goPackageExists := goPackageMap[childURL]

		if goPackageExists == true {
			goPackage.AwesomeGo = true
			goPackageMap[childURL] = goPackage
		}
	})

	// TODO ADD A FUCKING COUNT THAT FUCKING WORKS, FUCK
	log.Println("Finished scraping awesome-go/README.md\n")
	log.Println("Started retrieving versions for go packages\n")

	var goErrorPackageMap = make(map[string]*GoPackage)

	// TODO Write a constant for this
	nbConcurrentGet := 20
	urls := make(chan *GoPackage, nbConcurrentGet)
	var wg sync.WaitGroup
	for i := 0; i < nbConcurrentGet; i++ {
		wg.Add(1)
		go func() {
			for goPackage := range urls {
				refs, err := common.FetchRefs(goPackage.GitHubURL)
				if err != nil {
					log.Println("ERROR", goPackage.GitHubURL, " failed to return.\n", err)
					goPackage.HttpStatusCode = 404
					goErrorPackageMap[goPackage.GitHubURL] = goPackage
				} else {
					goPackage.HttpStatusCode = 200
					var versions []string
					for _, version := range refs.Candidates {
						versions = append(versions, version.String())
					}
					goPackage.Versions = versions
				}
				goPackageMap[goPackage.GitHubURL] = goPackage
			}
			wg.Done()
		}()
	}

	for _, goPackage := range goPackageMap {
		urls <- goPackage
	}

	close(urls)
	wg.Wait()

	successfulNumPackages := len(goPackageMap) - len(goErrorPackageMap)
	log.Println("SUCCESS: ", successfulNumPackages, " GoPackages were successfully built")
	log.Println("Creating JSON dump of ", len(goPackageMap), " go packages")
	createJSONDump(goPackageMap, "valid-go-packages")
	log.Println("Finished creating JSON dump\n")

	log.Println("WARNING: ", len(goErrorPackageMap), " GoPackages were not found on github")
	log.Println("Creating JSON dump of ", len(goErrorPackageMap), " err packages")
	createJSONDump(goErrorPackageMap, "invalid-go-packages")
	log.Println("Finished creating JSON dump\n")

	elapsed := time.Since(start)
	log.Printf("Program took %s to fully execute", elapsed)
}

// Create a tmp JSON dump of all serialized goPackageData
func createJSONDump(goPackageMap map[string]*GoPackage, fileName string) {
	var buffer bytes.Buffer

	for _, goPackage := range goPackageMap {
		buffer.WriteString("{\"url\": \"" + goPackage.GitHubURL + "\", \"description\": \"" + goPackage.Description + "\", \"index_time\": \"" + goPackage.IndexTime + "\" }, \"versions\": \"" + fmt.Sprintf("%v", goPackage.Versions) + "\" \n")
	}

	t := time.Now()
	time := strings.Replace(t.String(), " ", "", -1)
	jsonData := []byte(buffer.String())
	_ = ioutil.WriteFile("./"+time+"-"+fileName+".json", jsonData, 0644)
}

func printGoPackage(goPackage *GoPackage) {
	fmt.Printf("%v", goPackage)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	log "github.com/Sirupsen/logrus"
	"github.com/stefanprodan/syros/models"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func main() {
	var config = &Config{}
	flag.StringVar(&config.LogLevel, "LogLevel", "debug", "logging threshold level: debug|info|warn|error|fatal|panic")
	flag.IntVar(&config.Port, "Port", 8887, "HTTP port to listen on")
	flag.StringVar(&config.Nats, "Nats", "nats://localhost:4222", "Nats server addresses comma delimited")
	flag.StringVar(&config.CollectorTopic, "CollectorTopic", "docker", "Nats collector topic name")
	flag.StringVar(&config.CollectorQueue, "CollectorQueue", "syros", "Nats collector queue name")
	flag.StringVar(&config.RegistryTopic, "RegistryTopic", "registry", "Nats registry topic name")
	flag.StringVar(&config.RegistryQueue, "RegistryQueue", "syros", "Nats registry queue name")
	flag.StringVar(&config.RethinkDB, "RethinkDB", "localhost:28015", "RethinkDB server addresses comma delimited")
	flag.StringVar(&config.Database, "Database", "syros", "RethinkDB database name")
	flag.IntVar(&config.DatabaseStale, "DatabaseStale", 5, "Deletes database records older than specified value in minutes, set 0 to disable")
	flag.Parse()

	setLogLevel(config.LogLevel)
	log.Infof("Starting with config: %+v", config)

	repo, err := NewRepository(config)
	if err != nil {
		log.Fatalf("RethinkDB connection error %v", err)
	}

	repo.Initialize()
	log.Infof("Connected to RethinkDB cluster %v database initialization done", config.RethinkDB)

	repo.RunGarbageCollector([]string{"containers", "hosts", "syros_services"})

	nc, err := NewNatsConnection(config.Nats)
	if err != nil {
		log.Fatalf("Nats connection error %v", err)
	}
	defer nc.Close()

	log.Infof("Connected to NATS server %v status %v", nc.ConnectedUrl(), nc.Status())

	registry := NewRegistry(config, nc, repo)
	registry.WatchForAgents()

	indexer := models.SyrosService{
		Environment: "all",
		Type:        "indexer",
	}
	indexer.Config, _ = models.ConfigToMap(config, "m")
	indexer.Hostname, _ = os.Hostname()
	uuid, _ := models.NewUUID()
	indexer.Id = models.Hash(indexer.Hostname + uuid)
	go func(a models.SyrosService) {
		for true {
			indexer.Collected = time.Now().UTC()
			registry.SelfRegister(indexer)
			time.Sleep(10 * time.Second)
		}
	}(indexer)

	consumer, err := NewConsumer(config, nc, repo)
	if err != nil {
		log.Fatalf("Consumer init error %v", err)
	}
	consumer.Consume()

	server := &HttpServer{
		Config:     config,
		Registry:   registry,
		Repository: repo,
	}
	log.Infof("Starting HTTP server on port %v", config.Port)
	go server.Start()

	//wait for SIGINT (Ctrl+C) or SIGTERM (docker stop)
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	sig := <-sigChan
	log.Infof("Shuting down %v signal received", sig)
	nc.Close()
}

func setLogLevel(levelName string) {
	level, err := log.ParseLevel(levelName)
	if err != nil {
		log.Fatal(err)
	}
	log.SetLevel(level)
}
<commit_msg>GC checks<commit_after>package main

import (
	"flag"
	log "github.com/Sirupsen/logrus"
	"github.com/stefanprodan/syros/models"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func main() {
	var config = &Config{}
	flag.StringVar(&config.LogLevel, "LogLevel", "debug", "logging threshold level: debug|info|warn|error|fatal|panic")
	flag.IntVar(&config.Port, "Port", 8887, "HTTP port to listen on")
	flag.StringVar(&config.Nats, "Nats", "nats://localhost:4222", "Nats server addresses comma delimited")
	flag.StringVar(&config.CollectorTopic, "CollectorTopic", "docker", "Nats collector topic name")
	flag.StringVar(&config.CollectorQueue, "CollectorQueue", "syros", "Nats collector queue name")
	flag.StringVar(&config.RegistryTopic, "RegistryTopic", "registry", "Nats registry topic name")
	flag.StringVar(&config.RegistryQueue, "RegistryQueue", "syros", "Nats registry queue name")
	flag.StringVar(&config.RethinkDB, "RethinkDB", "localhost:28015", "RethinkDB server addresses comma delimited")
	flag.StringVar(&config.Database, "Database", "syros", "RethinkDB database name")
	flag.IntVar(&config.DatabaseStale, "DatabaseStale", 5, "Deletes database records older than specified value in minutes, set 0 to disable")
	flag.Parse()

	setLogLevel(config.LogLevel)
	log.Infof("Starting with config: %+v", config)

	repo, err := NewRepository(config)
	if err != nil {
		log.Fatalf("RethinkDB connection error %v", err)
	}

	repo.Initialize()
	log.Infof("Connected to RethinkDB cluster %v database initialization done", config.RethinkDB)

	repo.RunGarbageCollector([]string{"containers", "hosts", "checks", "syros_services"})

	nc, err := NewNatsConnection(config.Nats)
	if err != nil {
		log.Fatalf("Nats connection error %v", err)
	}
	defer nc.Close()

	log.Infof("Connected to NATS server %v status %v", nc.ConnectedUrl(), nc.Status())

	registry := NewRegistry(config, nc, repo)
	registry.WatchForAgents()

	indexer := models.SyrosService{
		Environment: "all",
		Type:        "indexer",
	}
	indexer.Config, _ = models.ConfigToMap(config, "m")
	indexer.Hostname, _ = os.Hostname()
	uuid, _ := models.NewUUID()
	indexer.Id = models.Hash(indexer.Hostname + uuid)
	go func(a models.SyrosService) {
		for true {
			indexer.Collected = time.Now().UTC()
			registry.SelfRegister(indexer)
			time.Sleep(10 * time.Second)
		}
	}(indexer)

	consumer, err := NewConsumer(config, nc, repo)
	if err != nil {
		log.Fatalf("Consumer init error %v", err)
	}
	consumer.Consume()

	server := &HttpServer{
		Config:     config,
		Registry:   registry,
		Repository: repo,
	}
	log.Infof("Starting HTTP server on port %v", config.Port)
	go server.Start()

	//wait for SIGINT (Ctrl+C) or SIGTERM (docker stop)
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	sig := <-sigChan
	log.Infof("Shuting down %v signal received", sig)
	nc.Close()
}

func setLogLevel(levelName string) {
	level, err := log.ParseLevel(levelName)
	if err != nil {
		log.Fatal(err)
	}
	log.SetLevel(level)
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. 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. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 utils

import (
	"encoding/json"
	"testing"
)

func TestJsonKeys(t *testing.T) {
	keys, err := JsonKeys([]byte(`{"key1": 1, "key2": "val", "key3": {}}`))
	if err != nil {
		t.Error(err)
	}
	if !SlicesDeepEqual(keys, []string{"key1", "key2", "key3"}) {
		t.Errorf("Incorrect json keys: Got %v", keys)
	}

	_, err = JsonKeys([]byte(`Invalid json }{`))
	if err == nil {
		t.Error("Could find keys for invalid json")
	}

	keys, err = JsonKeys([]byte(`{}`))
	if err != nil {
		t.Error(err)
	}
	if len(keys) != 0 {
		t.Error("Keys for empty objet should be empty")
	}

	_, err = JsonKeys([]byte(`[]`))
	if err == nil {
		t.Error("Trying to get keys of a non-object wasn't an error")
	}
}

type testStruct struct {
	Field1 string
	Field2 int
	Field3 int `json:"f3"`
	Field4 int `json:"f4,omitempty,string"`
}

type futureTestStruct struct {
	Field1 string
	Field2 int
	Field3 int `json:"f3"`
	Field4 int `json:"f4,omitempty,string"`
	Field5 int `json:"f5,omitempty"`
}

func TestCompleteJsonUnmarshal(t *testing.T) {
	test := testStruct{Field1: "str", Field2: 1, Field3: 1, Field4: 1}
	jsonString, _ := json.Marshal(test)

	var unmarshaledTest testStruct
	err := json.Unmarshal(jsonString, &unmarshaledTest)
	if err != nil {
		t.Error(err)
	}

	if err = CompleteJsonUnmarshal(jsonString, unmarshaledTest); err != nil {
		t.Errorf("Unmarshal should have been complete: %v", err)
	}

	future := futureTestStruct{Field1: "str", Field2: 1, Field5: 1}
	futureStr, _ := json.Marshal(future)

	var outOfDate testStruct
	err = json.Unmarshal(futureStr, &outOfDate)
	if err != nil {
		t.Error(err)
	}

	if err = CompleteJsonUnmarshal(futureStr, outOfDate); err == nil {
		t.Error("Shouldn't be complete with an unrecognized field5")
	}

	var empty struct{}
	if err = CompleteJsonUnmarshal([]byte(`{}`), empty); err != nil {
		t.Error("The empty object can losslessly unmarshal into anything")
	}
	if err := CompleteJsonUnmarshal([]byte(`{}`), test); err != nil {
		t.Error("The empty object can losslessly unmarshal into anything")
	}

	if err = CompleteJsonUnmarshal([]byte(`{"key":"val"}`), empty); err == nil {
		t.Error("The non-empty object can't unmarshal into the empty struct losslessly")
	}
}
<commit_msg>Fixed typo: objet -> object<commit_after>// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. 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. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 utils

import (
	"encoding/json"
	"testing"
)

func TestJsonKeys(t *testing.T) {
	keys, err := JsonKeys([]byte(`{"key1": 1, "key2": "val", "key3": {}}`))
	if err != nil {
		t.Error(err)
	}
	if !SlicesDeepEqual(keys, []string{"key1", "key2", "key3"}) {
		t.Errorf("Incorrect json keys: Got %v", keys)
	}

	_, err = JsonKeys([]byte(`Invalid json }{`))
	if err == nil {
		t.Error("Could find keys for invalid json")
	}

	keys, err = JsonKeys([]byte(`{}`))
	if err != nil {
		t.Error(err)
	}
	if len(keys) != 0 {
		t.Error("Keys for empty object should be empty")
	}

	_, err = JsonKeys([]byte(`[]`))
	if err == nil {
		t.Error("Trying to get keys of a non-object wasn't an error")
	}
}

type testStruct struct {
	Field1 string
	Field2 int
	Field3 int `json:"f3"`
	Field4 int `json:"f4,omitempty,string"`
}

type futureTestStruct struct {
	Field1 string
	Field2 int
	Field3 int `json:"f3"`
	Field4 int `json:"f4,omitempty,string"`
	Field5 int `json:"f5,omitempty"`
}

func TestCompleteJsonUnmarshal(t *testing.T) {
	test := testStruct{Field1: "str", Field2: 1, Field3: 1, Field4: 1}
	jsonString, _ := json.Marshal(test)

	var unmarshaledTest testStruct
	err := json.Unmarshal(jsonString, &unmarshaledTest)
	if err != nil {
		t.Error(err)
	}

	if err = CompleteJsonUnmarshal(jsonString, unmarshaledTest); err != nil {
		t.Errorf("Unmarshal should have been complete: %v", err)
	}

	future := futureTestStruct{Field1: "str", Field2: 1, Field5: 1}
	futureStr, _ := json.Marshal(future)

	var outOfDate testStruct
	err = json.Unmarshal(futureStr, &outOfDate)
	if err != nil {
		t.Error(err)
	}

	if err = CompleteJsonUnmarshal(futureStr, outOfDate); err == nil {
		t.Error("Shouldn't be complete with an unrecognized field5")
	}

	var empty struct{}
	if err = CompleteJsonUnmarshal([]byte(`{}`), empty); err != nil {
		t.Error("The empty object can losslessly unmarshal into anything")
	}
	if err := CompleteJsonUnmarshal([]byte(`{}`), test); err != nil {
		t.Error("The empty object can losslessly unmarshal into anything")
	}

	if err = CompleteJsonUnmarshal([]byte(`{"key":"val"}`), empty); err == nil {
		t.Error("The non-empty object can't unmarshal into the empty struct losslessly")
	}
}
<|endoftext|>


<commit_before>// Package rest contains the code to use the REST-based Google API
package rest

import (
	"context"
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"

	"github.com/google/go-tpm-tools/launcher/verifier"

	v1alpha1 "google.golang.org/api/confidentialcomputing/v1alpha1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/option"
)

// NewClient creates a new REST client which is configured to perform
// attestations in a particular project and region.
func NewClient(ctx context.Context, projectID string, region string, opts ...option.ClientOption) (verifier.Client, error) {
	service, err := v1alpha1.NewService(ctx, opts...)
	if err != nil {
		return nil, fmt.Errorf("can't create ConfidentialComputing v1alpha1 API client: %w", err)
	}

	projectName := fmt.Sprintf("projects/%s", projectID)
	locationName := fmt.Sprintf("%s/locations/%v", projectName, region)

	location, err := service.Projects.Locations.Get(locationName).Do()
	if err == nil {
		return &restClient{service, location}, nil
	}

	// Check if the error was due to a bad region name
	if apiErr, ok := err.(*googleapi.Error); ok && apiErr.Code == 403 {
		// In this case, inform the user about the allowed regions
		if list, listErr := service.Projects.Locations.List(projectName).Do(); listErr == nil {
			locations := make([]string, len(list.Locations))
			for i, loc := range list.Locations {
				locations[i] = loc.LocationId
			}
			return nil, fmt.Errorf(
				"unable to find region %q, available regions are [%s]: %w",
				region, strings.Join(locations, ", "), err,
			)
		}
	}

	return nil, fmt.Errorf("unable to use project %q and region %q: %w", projectID, region, err)
}

type restClient struct {
	service  *v1alpha1.Service
	location *v1alpha1.Location
}

// CreateChallenge implements verifier.Client
func (c *restClient) CreateChallenge(ctx context.Context) (*verifier.Challenge, error) {
	// Pass an empty Challenge for the input (all params are output-only)
	chal, err := c.service.Projects.Locations.Challenges.Create(
		c.location.Name,
		&v1alpha1.Challenge{},
	).Do()
	if err != nil {
		return nil, fmt.Errorf("calling v1alpha1.CreateChallenge: %w", err)
	}
	return convertChallengeFromREST(chal)
}

// VerifyAttestation implements verifier.Client
func (c *restClient) VerifyAttestation(ctx context.Context, request verifier.VerifyAttestationRequest) (*verifier.VerifyAttestationResponse, error) {
	if request.Challenge == nil || request.Attestation == nil {
		return nil, fmt.Errorf("nil value provided in challenge")
	}
	response, err := c.service.Projects.Locations.Challenges.VerifyAttestation(
		request.Challenge.Name,
		convertRequestToREST(request),
	).Do()
	if err != nil {
		return nil, fmt.Errorf("calling v1alpha1.VerifyAttestation: %w", err)
	}
	return convertResponseFromREST(response)
}

var encoding = base64.StdEncoding

func convertChallengeFromREST(chal *v1alpha1.Challenge) (*verifier.Challenge, error) {
	nonce, err := encoding.DecodeString(chal.Nonce)
	if err != nil {
		return nil, fmt.Errorf("failed to decode Challenge.Nonce: %w", err)
	}
	return &verifier.Challenge{
		Name:  chal.Name,
		Nonce: nonce,
	}, nil
}

func convertRequestToREST(request verifier.VerifyAttestationRequest) *v1alpha1.VerifyAttestationRequest {
	idTokens := make([]string, len(request.GcpCredentials))
	for i, token := range request.GcpCredentials {
		idTokens[i] = encoding.EncodeToString(token)
	}

	quotes := make([]*v1alpha1.Quote, len(request.Attestation.GetQuotes()))
	for i, quote := range request.Attestation.GetQuotes() {
		pcrVals := map[string]string{}
		for idx, val := range quote.GetPcrs().GetPcrs() {
			strIdx := strconv.FormatUint(uint64(idx), 10)
			pcrVals[strIdx] = encoding.EncodeToString(val)
		}

		quotes[i] = &v1alpha1.Quote{
			RawQuote:     encoding.EncodeToString(quote.GetQuote()),
			RawSignature: encoding.EncodeToString(quote.GetRawSig()),
			HashAlgo:     int64(quote.GetPcrs().GetHash()),
			PcrValues:    pcrVals,
		}
	}

	certs := make([]string, len(request.Attestation.GetIntermediateCerts()))
	for i, cert := range request.Attestation.GetIntermediateCerts() {
		certs[i] = encoding.EncodeToString(cert)
	}

	return &v1alpha1.VerifyAttestationRequest{
		GcpCredentials: &v1alpha1.GcpCredentials{
			IdTokens: idTokens,
		},
		TpmAttestation: &v1alpha1.TpmAttestation{
			Quotes:            quotes,
			TcgEventLog:       encoding.EncodeToString(request.Attestation.GetEventLog()),
			CanonicalEventLog: encoding.EncodeToString(request.Attestation.GetCanonicalEventLog()),
			AkCert:            encoding.EncodeToString(request.Attestation.GetAkCert()),
			CertChain:         certs,
		},
	}
}

func convertResponseFromREST(resp *v1alpha1.VerifyAttestationResponse) (*verifier.VerifyAttestationResponse, error) {
	token, err := encoding.DecodeString(resp.ClaimsToken)
	if err != nil {
		return nil, fmt.Errorf("failed to decode VerifyAttestationResponse.ClaimsToken: %w", err)
	}
	return &verifier.VerifyAttestationResponse{
		ClaimsToken: token,
	}, nil
}
<commit_msg>Improve rest.NewClient error handling. (#218)<commit_after>// Package rest contains the code to use the REST-based Google API
package rest

import (
	"context"
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"

	"github.com/google/go-tpm-tools/launcher/verifier"

	v1alpha1 "google.golang.org/api/confidentialcomputing/v1alpha1"
	"google.golang.org/api/option"
)

// BadRegionError indicates that:
//   - the requested Region cannot be used with this API
//   - other Regions _can_ be used with this API
type BadRegionError struct {
	RequestedRegion  string
	AvailableRegions []string
	err              error
}

func (e *BadRegionError) Error() string {
	return fmt.Sprintf(
		"invalid region %q, available regions are [%s]: %v",
		e.RequestedRegion, strings.Join(e.AvailableRegions, ", "), e.err,
	)
}

func (e *BadRegionError) Unwrap() error {
	return e.err
}

// NewClient creates a new REST client which is configured to perform
// attestations in a particular project and region. Returns a *BadRegionError
// if the requested project is valid, but the region is invalid.
func NewClient(ctx context.Context, projectID string, region string, opts ...option.ClientOption) (verifier.Client, error) {
	service, err := v1alpha1.NewService(ctx, opts...)
	if err != nil {
		return nil, fmt.Errorf("can't create ConfidentialComputing v1alpha1 API client: %w", err)
	}

	projectName := fmt.Sprintf("projects/%s", projectID)
	locationName := fmt.Sprintf("%s/locations/%v", projectName, region)

	location, getErr := service.Projects.Locations.Get(locationName).Do()
	if getErr == nil {
		return &restClient{service, location}, nil
	}

	// If we can't get the location, try to list the locations. This handles
	// situations where the projectID is invalid.
	list, listErr := service.Projects.Locations.List(projectName).Do()
	if listErr != nil {
		return nil, fmt.Errorf("listing regions in project %q: %w", projectID, listErr)
	}

	// The project is valid, but can't get the desired region.
	regions := make([]string, len(list.Locations))
	for i, loc := range list.Locations {
		regions[i] = loc.LocationId
	}
	return nil, &BadRegionError{
		RequestedRegion:  region,
		AvailableRegions: regions,
		err:              getErr,
	}
}

type restClient struct {
	service  *v1alpha1.Service
	location *v1alpha1.Location
}

// CreateChallenge implements verifier.Client
func (c *restClient) CreateChallenge(ctx context.Context) (*verifier.Challenge, error) {
	// Pass an empty Challenge for the input (all params are output-only)
	chal, err := c.service.Projects.Locations.Challenges.Create(
		c.location.Name,
		&v1alpha1.Challenge{},
	).Do()
	if err != nil {
		return nil, fmt.Errorf("calling v1alpha1.CreateChallenge: %w", err)
	}
	return convertChallengeFromREST(chal)
}

// VerifyAttestation implements verifier.Client
func (c *restClient) VerifyAttestation(ctx context.Context, request verifier.VerifyAttestationRequest) (*verifier.VerifyAttestationResponse, error) {
	if request.Challenge == nil || request.Attestation == nil {
		return nil, fmt.Errorf("nil value provided in challenge")
	}
	response, err := c.service.Projects.Locations.Challenges.VerifyAttestation(
		request.Challenge.Name,
		convertRequestToREST(request),
	).Do()
	if err != nil {
		return nil, fmt.Errorf("calling v1alpha1.VerifyAttestation: %w", err)
	}
	return convertResponseFromREST(response)
}

var encoding = base64.StdEncoding

func convertChallengeFromREST(chal *v1alpha1.Challenge) (*verifier.Challenge, error) {
	nonce, err := encoding.DecodeString(chal.Nonce)
	if err != nil {
		return nil, fmt.Errorf("failed to decode Challenge.Nonce: %w", err)
	}
	return &verifier.Challenge{
		Name:  chal.Name,
		Nonce: nonce,
	}, nil
}

func convertRequestToREST(request verifier.VerifyAttestationRequest) *v1alpha1.VerifyAttestationRequest {
	idTokens := make([]string, len(request.GcpCredentials))
	for i, token := range request.GcpCredentials {
		idTokens[i] = encoding.EncodeToString(token)
	}

	quotes := make([]*v1alpha1.Quote, len(request.Attestation.GetQuotes()))
	for i, quote := range request.Attestation.GetQuotes() {
		pcrVals := map[string]string{}
		for idx, val := range quote.GetPcrs().GetPcrs() {
			strIdx := strconv.FormatUint(uint64(idx), 10)
			pcrVals[strIdx] = encoding.EncodeToString(val)
		}

		quotes[i] = &v1alpha1.Quote{
			RawQuote:     encoding.EncodeToString(quote.GetQuote()),
			RawSignature: encoding.EncodeToString(quote.GetRawSig()),
			HashAlgo:     int64(quote.GetPcrs().GetHash()),
			PcrValues:    pcrVals,
		}
	}

	certs := make([]string, len(request.Attestation.GetIntermediateCerts()))
	for i, cert := range request.Attestation.GetIntermediateCerts() {
		certs[i] = encoding.EncodeToString(cert)
	}

	return &v1alpha1.VerifyAttestationRequest{
		GcpCredentials: &v1alpha1.GcpCredentials{
			IdTokens: idTokens,
		},
		TpmAttestation: &v1alpha1.TpmAttestation{
			Quotes:            quotes,
			TcgEventLog:       encoding.EncodeToString(request.Attestation.GetEventLog()),
			CanonicalEventLog: encoding.EncodeToString(request.Attestation.GetCanonicalEventLog()),
			AkCert:            encoding.EncodeToString(request.Attestation.GetAkCert()),
			CertChain:         certs,
		},
	}
}

func convertResponseFromREST(resp *v1alpha1.VerifyAttestationResponse) (*verifier.VerifyAttestationResponse, error) {
	token, err := encoding.DecodeString(resp.ClaimsToken)
	if err != nil {
		return nil, fmt.Errorf("failed to decode VerifyAttestationResponse.ClaimsToken: %w", err)
	}
	return &verifier.VerifyAttestationResponse{
		ClaimsToken: token,
	}, nil
}
<|endoftext|>


<commit_before>package internal

import (
	"encoding/hex"
	"fmt"
	"strconv"
	"time"
)

func AppendArg(b []byte, v interface{}) []byte {
	switch v := v.(type) {
	case nil:
		return append(b, "<nil>"...)
	case string:
		return appendUTF8String(b, Bytes(v))
	case []byte:
		return appendUTF8String(b, v)
	case int:
		return strconv.AppendInt(b, int64(v), 10)
	case int8:
		return strconv.AppendInt(b, int64(v), 10)
	case int16:
		return strconv.AppendInt(b, int64(v), 10)
	case int32:
		return strconv.AppendInt(b, int64(v), 10)
	case int64:
		return strconv.AppendInt(b, v, 10)
	case uint:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint8:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint16:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint32:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint64:
		return strconv.AppendUint(b, v, 10)
	case float32:
		return strconv.AppendFloat(b, float64(v), 'f', -1, 64)
	case float64:
		return strconv.AppendFloat(b, v, 'f', -1, 64)
	case bool:
		if v {
			return append(b, "true"...)
		}
		return append(b, "false"...)
	case time.Time:
		return v.AppendFormat(b, time.RFC3339Nano)
	default:
		return append(b, fmt.Sprint(v)...)
	}
}

func appendUTF8String(dst []byte, src []byte) []byte {
	if isSimple(src) {
		dst = append(dst, src...)
		return dst
	}

	s := len(dst)
	dst = append(dst, make([]byte, hex.EncodedLen(len(src)))...)
	hex.Encode(dst[s:], src)
	return dst
}

func isSimple(b []byte) bool {
	for _, c := range b {
		if !isSimpleByte(c) {
			return false
		}
	}
	return true
}

func isSimpleByte(c byte) bool {
	return simple[c]
}

var simple = [256]bool{
	'-': true,
	'_': true,

	'0': true,
	'1': true,
	'2': true,
	'3': true,
	'4': true,
	'5': true,
	'6': true,
	'7': true,
	'8': true,
	'9': true,

	'a': true,
	'b': true,
	'c': true,
	'd': true,
	'e': true,
	'f': true,
	'g': true,
	'h': true,
	'i': true,
	'j': true,
	'k': true,
	'l': true,
	'm': true,
	'n': true,
	'o': true,
	'p': true,
	'q': true,
	'r': true,
	's': true,
	't': true,
	'u': true,
	'v': true,
	'w': true,
	'x': true,
	'y': true,
	'z': true,

	'A': true,
	'B': true,
	'C': true,
	'D': true,
	'E': true,
	'F': true,
	'G': true,
	'H': true,
	'I': true,
	'J': true,
	'K': true,
	'L': true,
	'M': true,
	'N': true,
	'O': true,
	'P': true,
	'Q': true,
	'R': true,
	'S': true,
	'T': true,
	'U': true,
	'V': true,
	'W': true,
	'X': true,
	'Y': true,
	'Z': true,
}
<commit_msg>Remove arg hex encoding<commit_after>package internal

import (
	"fmt"
	"strconv"
	"time"
)

func AppendArg(b []byte, v interface{}) []byte {
	switch v := v.(type) {
	case nil:
		return append(b, "<nil>"...)
	case string:
		return appendUTF8String(b, Bytes(v))
	case []byte:
		return appendUTF8String(b, v)
	case int:
		return strconv.AppendInt(b, int64(v), 10)
	case int8:
		return strconv.AppendInt(b, int64(v), 10)
	case int16:
		return strconv.AppendInt(b, int64(v), 10)
	case int32:
		return strconv.AppendInt(b, int64(v), 10)
	case int64:
		return strconv.AppendInt(b, v, 10)
	case uint:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint8:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint16:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint32:
		return strconv.AppendUint(b, uint64(v), 10)
	case uint64:
		return strconv.AppendUint(b, v, 10)
	case float32:
		return strconv.AppendFloat(b, float64(v), 'f', -1, 64)
	case float64:
		return strconv.AppendFloat(b, v, 'f', -1, 64)
	case bool:
		if v {
			return append(b, "true"...)
		}
		return append(b, "false"...)
	case time.Time:
		return v.AppendFormat(b, time.RFC3339Nano)
	default:
		return append(b, fmt.Sprint(v)...)
	}
}

func appendUTF8String(dst []byte, src []byte) []byte {
	dst = append(dst, src...)
	return dst
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"regexp"
	"sort"
	"strings"
	"time"
)

type Server struct {
	eventChan  chan Event
	running    bool
	name       string
	clientMap  map[string]*Client  //Map of nicks → clients
	channelMap map[string]*Channel //Map of channel names → channels
}

type Client struct {
	server     *Server
	connection net.Conn
	signalChan chan int
	outputChan chan string
	nick       string
	channelMap map[string]*Channel
}

type Channel struct {
	name      string
	topic     string
	clientMap map[string]*Client
}

type Event struct {
	client *Client
	input  string
}

const (
	signalStop int = iota
)

const (
	rplWelcome int = iota
	rplJoin
	rplPart
	rplTopic
	rplNoTopic
	rplNames
	rplNickChange
	errMoreArgs
	errNoNick
	errInvalidNick
	errNickInUse
	errAlreadyReg
	errNoSuchNick
)

var (
	nickRegexp    = regexp.MustCompile(`^[a-zA-Z\[\]_^{|}][a-zA-Z0-9\[\]_^{|}]*$`)
	channelRegexp = regexp.MustCompile(`^#[a-z0-9_\-]+$`)
)

func NewServer() *Server {
	return &Server{eventChan: make(chan Event),
		name:       "rosella",
		clientMap:  make(map[string]*Client),
		channelMap: make(map[string]*Channel)}
}

func (s *Server) Run() {
	go func() {
		for {
			s.handleEvent(<-s.eventChan)
		}
	}()
}

func (s *Server) HandleConnection(conn net.Conn) {

	client := &Client{server: s,
		connection: conn,
		outputChan: make(chan string),
		signalChan: make(chan int),
		channelMap: make(map[string]*Channel)}

	go client.clientThread()
}

func (s *Server) handleEvent(e Event) {
	fields := strings.Fields(e.input)

	if len(fields) < 1 {
		return
	}

	if strings.HasPrefix(fields[0], ":") {
		fields = fields[1:]
	}

	command := strings.ToUpper(fields[0])
	args := fields[1:]

	switch {
	case command == "NICK":
		if len(args) < 1 {
			e.client.reply(errNoNick)
			return
		}

		newNick := args[0]

		//Check newNick is of valid formatting (regex)
		if nickRegexp.MatchString(newNick) == false {
			e.client.reply(errInvalidNick)
			return
		}

		if _, exists := s.clientMap[newNick]; exists {
			e.client.reply(errNickInUse)
			return
		}

		e.client.setNick(newNick)

	case command == "USER":
		if e.client.nick == "" {
			//Give them a unique Guest nick
			newNick := fmt.Sprintf("Guest%d", rand.Int())
			for s.clientMap[newNick] != nil {
				newNick = fmt.Sprintf("Guest%d", rand.Int())
			}
		}
		e.client.reply(rplWelcome)

	case command == "JOIN":
		if len(args) < 1 {
			e.client.reply(errMoreArgs)
			return
		}

		if args[0] == "0" {
			//Quit all channels
			for channel := range e.client.channelMap {
				s.partChannel(e.client, channel)
			}
			return
		}

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Join the channel if it's valid
			if channelRegexp.Match([]byte(channel)) {
				s.joinChannel(e.client, channel)
			}
		}

	case command == "PART":
		if len(args) < 1 {
			e.client.reply(errMoreArgs)
			return
		}

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Part the channel if it's valid
			if channelRegexp.Match([]byte(channel)) {
				s.partChannel(e.client, channel)
			}
		}

	case command == "PRIVMSG":
		if len(args) < 2 {
			e.client.reply(errMoreArgs)
			return
		}

		message := strings.Join(args[1:], " ")

		channel, chanExists := s.channelMap[args[0]]
		client, clientExists := s.clientMap[args[0]]

		if chanExists {
			for _, c := range channel.clientMap {
				if c != e.client {
					c.outputChan <- fmt.Sprintf(":%s PRIVMSG %s %s", e.client.nick, args[0], message)
				}
			}
		} else if clientExists {
			client.outputChan <- fmt.Sprintf(":%s PRIVMSG %s %s", e.client.nick, client.nick, message)
		} else {
			e.client.reply(errNoSuchNick)
		}

	case command == "QUIT":
		//Stop the client, which will auto part channels and quit
		e.client.signalChan <- signalStop
	case command == "TOPIC":
		if len(args) < 1 {
			e.client.reply(errMoreArgs)
			return
		}

		channel, exists := s.channelMap[args[0]]
		if exists == false {
			e.client.reply(errNoSuchNick)
			return
		}

		channelName := args[0]

		if len(args) == 1 {
			e.client.reply(rplTopic, channelName, channel.topic)
			return
		}

		if args[1] == ":" {
			channel.topic = ""
			for _, client := range channel.clientMap {
				client.reply(rplNoTopic, channelName)
			}
		} else {
			topic := strings.Join(args[1:], " ")
			topic = strings.TrimPrefix(topic, ":")
			channel.topic = topic

			for _, client := range channel.clientMap {
				client.reply(rplTopic, channelName, channel.topic)
			}
		}

	default:
		log.Printf("Unknown command: %q", command)
	}
}

func (s *Server) joinChannel(client *Client, channelName string) {
	channel, exists := s.channelMap[channelName]
	if exists == false {
		channel = &Channel{name: channelName,
			topic:     "",
			clientMap: make(map[string]*Client)}
		s.channelMap[channelName] = channel
	}

	channel.clientMap[client.nick] = client
	client.channelMap[channelName] = channel

	for _, c := range channel.clientMap {
		c.reply(rplJoin, client.nick, channelName)
	}

	if channel.topic != "" {
		client.reply(rplTopic, channelName, channel.topic)
	} else {
		client.reply(rplNoTopic, channelName)
	}

	nicks := make([]string, 0, 100)
	for nick := range channel.clientMap {
		nicks = append(nicks, nick)
	}

	client.reply(rplNames, channelName, strings.Join(nicks, " "))
}

func (s *Server) partChannel(client *Client, channelName string) {
	channel, exists := s.channelMap[channelName]
	if exists == false {
		return
	}

	//Notify clients of the part
	for _, c := range channel.clientMap {
		c.reply(rplPart, client.nick, channelName)
	}

	delete(channel.clientMap, client.nick)
	delete(client.channelMap, channelName)
}

func (c *Client) clientThread() {
	defer c.connection.Close()

	readSignalChan := make(chan int, 1)
	writeSignalChan := make(chan int, 1)
	writeChan := make(chan string, 100)

	go c.readThread(readSignalChan)
	go c.writeThread(writeSignalChan, writeChan)

	for {
		select {
		case signal := <-c.signalChan:
			//Do stuff
			if signal == signalStop {
				readSignalChan <- signalStop
				writeSignalChan <- signalStop
				break
			}
		case line := <-c.outputChan:
			select {
			case writeChan <- line:
				//It worked
			default:
				log.Printf("Dropped a line for client: %q", c.nick)
				//Do nothing, dropping the line
			}
		}
	}

	//Part from all channels
	for channelName := range c.channelMap {
		c.server.partChannel(c, channelName)
	}

	//Remove from client list
	delete(c.server.clientMap, c.nick)
}

func (c *Client) readThread(signalChan chan int) {
	for {
		select {
		case signal := <-signalChan:
			if signal == signalStop {
				return
			}
		default:
			c.connection.SetReadDeadline(time.Now().Add(time.Second * 3))
			buf := make([]byte, 512)
			ln, err := c.connection.Read(buf)
			if err != nil {
				if err == io.EOF {
					//They must have dc'd
					c.signalChan <- signalStop
					return
				}
				continue
			}

			rawLines := buf[:ln]
			lines := bytes.Split(rawLines, []byte("\r\n"))
			for _, line := range lines {
				if len(line) > 0 {
					c.server.eventChan <- Event{client: c, input: string(line)}
				}
			}
		}
	}
}

func (c *Client) writeThread(signalChan chan int, outputChan chan string) {
	for {
		select {
		case signal := <-signalChan:
			if signal == signalStop {
				return
			}
		case output := <-outputChan:
			line := []byte(fmt.Sprintf("%s\r\n", output))

			c.connection.SetWriteDeadline(time.Now().Add(time.Second * 30))
			_, err := c.connection.Write(line)
			if err != nil {
				log.Printf("Write err: %q", err.Error())
				c.signalChan <- signalStop
				return
			}
		}
	}
}

//Send a reply to a user with the code specified
func (c *Client) reply(code int, args ...string) {
	switch code {
	case rplWelcome:
		c.outputChan <- fmt.Sprintf(":%s 001 %s :Welcome to %s", c.server.name, c.nick, c.server.name)
	case rplJoin:
		c.outputChan <- fmt.Sprintf(":%s JOIN %s", args[0], args[1])
	case rplPart:
		c.outputChan <- fmt.Sprintf(":%s PART %s", args[0], args[1])
	case rplTopic:
		c.outputChan <- fmt.Sprintf(":%s 332 %s %s :%s", c.server.name, c.nick, args[0], args[1])
	case rplNoTopic:
		c.outputChan <- fmt.Sprintf(":%s 331 %s %s :No topic is set", c.server.name, c.nick, args[0])
	case rplNames:
		//TODO: break long lists up into multiple messages
		c.outputChan <- fmt.Sprintf(":%s 353 %s = %s :%s", c.server.name, c.nick, args[0], args[1])
		c.outputChan <- fmt.Sprintf(":%s 366", c.server.name)
	case rplNickChange:
		c.outputChan <- fmt.Sprintf(":%s NICK %s", args[0], args[1])
	case errMoreArgs:
		c.outputChan <- fmt.Sprintf(":%s 461 :Not enough params", c.server.name)
	case errNoNick:
		c.outputChan <- fmt.Sprintf(":%s 431 :No nickanme given", c.server.name)
	case errInvalidNick:
		c.outputChan <- fmt.Sprintf(":%s 432 :Erronenous nickname", c.server.name)
	case errNickInUse:
		c.outputChan <- fmt.Sprintf(":%s 433 :Nick already in use", c.server.name)
	case errAlreadyReg:
		c.outputChan <- fmt.Sprintf(":%s 462 :You need a valid nick first", c.server.name)
	case errNoSuchNick:
		c.outputChan <- fmt.Sprintf("%s 401 :No such nick/channel", c.server.name)
	}
}

func (c *Client) setNick(nick string) {
	if c.nick != "" {
		delete(c.server.clientMap, c.nick)
		for _, channel := range c.channelMap {
			delete(channel.clientMap, c.nick)
		}
	}

	//Set up new nick
	oldNick := c.nick
	c.nick = nick
	c.server.clientMap[c.nick] = c

	clients := make([]string, 0, 100)

	for _, channel := range c.channelMap {
		channel.clientMap[c.nick] = c

		//Collect list of client nicks who can see us
		for client := range channel.clientMap {
			clients = append(clients, client)
		}
	}

	//By sorting the nicks and skipping duplicates we send each client one message
	sort.Strings(clients)
	prevNick := ""
	for _, nick := range clients {
		if nick == prevNick {
			continue
		}
		prevNick = nick

		client, exists := c.server.clientMap[nick]
		if exists {
			client.reply(rplNickChange, oldNick, c.nick)
		}
	}
}
<commit_msg>Added unknown command reply<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"regexp"
	"sort"
	"strings"
	"time"
)

type Server struct {
	eventChan  chan Event
	running    bool
	name       string
	clientMap  map[string]*Client  //Map of nicks → clients
	channelMap map[string]*Channel //Map of channel names → channels
}

type Client struct {
	server     *Server
	connection net.Conn
	signalChan chan int
	outputChan chan string
	nick       string
	channelMap map[string]*Channel
}

type Channel struct {
	name      string
	topic     string
	clientMap map[string]*Client
}

type Event struct {
	client *Client
	input  string
}

const (
	signalStop int = iota
)

const (
	rplWelcome int = iota
	rplJoin
	rplPart
	rplTopic
	rplNoTopic
	rplNames
	rplNickChange
	errMoreArgs
	errNoNick
	errInvalidNick
	errNickInUse
	errAlreadyReg
	errNoSuchNick
	errUnknownCommand
)

var (
	nickRegexp    = regexp.MustCompile(`^[a-zA-Z\[\]_^{|}][a-zA-Z0-9\[\]_^{|}]*$`)
	channelRegexp = regexp.MustCompile(`^#[a-z0-9_\-]+$`)
)

func NewServer() *Server {
	return &Server{eventChan: make(chan Event),
		name:       "rosella",
		clientMap:  make(map[string]*Client),
		channelMap: make(map[string]*Channel)}
}

func (s *Server) Run() {
	go func() {
		for {
			s.handleEvent(<-s.eventChan)
		}
	}()
}

func (s *Server) HandleConnection(conn net.Conn) {

	client := &Client{server: s,
		connection: conn,
		outputChan: make(chan string),
		signalChan: make(chan int),
		channelMap: make(map[string]*Channel)}

	go client.clientThread()
}

func (s *Server) handleEvent(e Event) {
	fields := strings.Fields(e.input)

	if len(fields) < 1 {
		return
	}

	if strings.HasPrefix(fields[0], ":") {
		fields = fields[1:]
	}

	command := strings.ToUpper(fields[0])
	args := fields[1:]

	switch {
	case command == "NICK":
		if len(args) < 1 {
			e.client.reply(errNoNick)
			return
		}

		newNick := args[0]

		//Check newNick is of valid formatting (regex)
		if nickRegexp.MatchString(newNick) == false {
			e.client.reply(errInvalidNick)
			return
		}

		if _, exists := s.clientMap[newNick]; exists {
			e.client.reply(errNickInUse)
			return
		}

		e.client.setNick(newNick)

	case command == "USER":
		if e.client.nick == "" {
			//Give them a unique Guest nick
			newNick := fmt.Sprintf("Guest%d", rand.Int())
			for s.clientMap[newNick] != nil {
				newNick = fmt.Sprintf("Guest%d", rand.Int())
			}
		}
		e.client.reply(rplWelcome)

	case command == "JOIN":
		if len(args) < 1 {
			e.client.reply(errMoreArgs)
			return
		}

		if args[0] == "0" {
			//Quit all channels
			for channel := range e.client.channelMap {
				s.partChannel(e.client, channel)
			}
			return
		}

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Join the channel if it's valid
			if channelRegexp.Match([]byte(channel)) {
				s.joinChannel(e.client, channel)
			}
		}

	case command == "PART":
		if len(args) < 1 {
			e.client.reply(errMoreArgs)
			return
		}

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Part the channel if it's valid
			if channelRegexp.Match([]byte(channel)) {
				s.partChannel(e.client, channel)
			}
		}

	case command == "PRIVMSG":
		if len(args) < 2 {
			e.client.reply(errMoreArgs)
			return
		}

		message := strings.Join(args[1:], " ")

		channel, chanExists := s.channelMap[args[0]]
		client, clientExists := s.clientMap[args[0]]

		if chanExists {
			for _, c := range channel.clientMap {
				if c != e.client {
					c.outputChan <- fmt.Sprintf(":%s PRIVMSG %s %s", e.client.nick, args[0], message)
				}
			}
		} else if clientExists {
			client.outputChan <- fmt.Sprintf(":%s PRIVMSG %s %s", e.client.nick, client.nick, message)
		} else {
			e.client.reply(errNoSuchNick)
		}

	case command == "QUIT":
		//Stop the client, which will auto part channels and quit
		e.client.signalChan <- signalStop
	case command == "TOPIC":
		if len(args) < 1 {
			e.client.reply(errMoreArgs)
			return
		}

		channel, exists := s.channelMap[args[0]]
		if exists == false {
			e.client.reply(errNoSuchNick)
			return
		}

		channelName := args[0]

		if len(args) == 1 {
			e.client.reply(rplTopic, channelName, channel.topic)
			return
		}

		if args[1] == ":" {
			channel.topic = ""
			for _, client := range channel.clientMap {
				client.reply(rplNoTopic, channelName)
			}
		} else {
			topic := strings.Join(args[1:], " ")
			topic = strings.TrimPrefix(topic, ":")
			channel.topic = topic

			for _, client := range channel.clientMap {
				client.reply(rplTopic, channelName, channel.topic)
			}
		}

	default:
		e.client.reply(errUnknownCommand, command)
	}
}

func (s *Server) joinChannel(client *Client, channelName string) {
	channel, exists := s.channelMap[channelName]
	if exists == false {
		channel = &Channel{name: channelName,
			topic:     "",
			clientMap: make(map[string]*Client)}
		s.channelMap[channelName] = channel
	}

	channel.clientMap[client.nick] = client
	client.channelMap[channelName] = channel

	for _, c := range channel.clientMap {
		c.reply(rplJoin, client.nick, channelName)
	}

	if channel.topic != "" {
		client.reply(rplTopic, channelName, channel.topic)
	} else {
		client.reply(rplNoTopic, channelName)
	}

	nicks := make([]string, 0, 100)
	for nick := range channel.clientMap {
		nicks = append(nicks, nick)
	}

	client.reply(rplNames, channelName, strings.Join(nicks, " "))
}

func (s *Server) partChannel(client *Client, channelName string) {
	channel, exists := s.channelMap[channelName]
	if exists == false {
		return
	}

	//Notify clients of the part
	for _, c := range channel.clientMap {
		c.reply(rplPart, client.nick, channelName)
	}

	delete(channel.clientMap, client.nick)
	delete(client.channelMap, channelName)
}

func (c *Client) clientThread() {
	defer c.connection.Close()

	readSignalChan := make(chan int, 1)
	writeSignalChan := make(chan int, 1)
	writeChan := make(chan string, 100)

	go c.readThread(readSignalChan)
	go c.writeThread(writeSignalChan, writeChan)

	for {
		select {
		case signal := <-c.signalChan:
			//Do stuff
			if signal == signalStop {
				readSignalChan <- signalStop
				writeSignalChan <- signalStop
				break
			}
		case line := <-c.outputChan:
			select {
			case writeChan <- line:
				//It worked
			default:
				log.Printf("Dropped a line for client: %q", c.nick)
				//Do nothing, dropping the line
			}
		}
	}

	//Part from all channels
	for channelName := range c.channelMap {
		c.server.partChannel(c, channelName)
	}

	//Remove from client list
	delete(c.server.clientMap, c.nick)
}

func (c *Client) readThread(signalChan chan int) {
	for {
		select {
		case signal := <-signalChan:
			if signal == signalStop {
				return
			}
		default:
			c.connection.SetReadDeadline(time.Now().Add(time.Second * 3))
			buf := make([]byte, 512)
			ln, err := c.connection.Read(buf)
			if err != nil {
				if err == io.EOF {
					//They must have dc'd
					c.signalChan <- signalStop
					return
				}
				continue
			}

			rawLines := buf[:ln]
			lines := bytes.Split(rawLines, []byte("\r\n"))
			for _, line := range lines {
				if len(line) > 0 {
					c.server.eventChan <- Event{client: c, input: string(line)}
				}
			}
		}
	}
}

func (c *Client) writeThread(signalChan chan int, outputChan chan string) {
	for {
		select {
		case signal := <-signalChan:
			if signal == signalStop {
				return
			}
		case output := <-outputChan:
			line := []byte(fmt.Sprintf("%s\r\n", output))

			c.connection.SetWriteDeadline(time.Now().Add(time.Second * 30))
			_, err := c.connection.Write(line)
			if err != nil {
				log.Printf("Write err: %q", err.Error())
				c.signalChan <- signalStop
				return
			}
		}
	}
}

//Send a reply to a user with the code specified
func (c *Client) reply(code int, args ...string) {
	switch code {
	case rplWelcome:
		c.outputChan <- fmt.Sprintf(":%s 001 %s :Welcome to %s", c.server.name, c.nick, c.server.name)
	case rplJoin:
		c.outputChan <- fmt.Sprintf(":%s JOIN %s", args[0], args[1])
	case rplPart:
		c.outputChan <- fmt.Sprintf(":%s PART %s", args[0], args[1])
	case rplTopic:
		c.outputChan <- fmt.Sprintf(":%s 332 %s %s :%s", c.server.name, c.nick, args[0], args[1])
	case rplNoTopic:
		c.outputChan <- fmt.Sprintf(":%s 331 %s %s :No topic is set", c.server.name, c.nick, args[0])
	case rplNames:
		//TODO: break long lists up into multiple messages
		c.outputChan <- fmt.Sprintf(":%s 353 %s = %s :%s", c.server.name, c.nick, args[0], args[1])
		c.outputChan <- fmt.Sprintf(":%s 366", c.server.name)
	case rplNickChange:
		c.outputChan <- fmt.Sprintf(":%s NICK %s", args[0], args[1])
	case errMoreArgs:
		c.outputChan <- fmt.Sprintf(":%s 461 :Not enough params", c.server.name)
	case errNoNick:
		c.outputChan <- fmt.Sprintf(":%s 431 :No nickanme given", c.server.name)
	case errInvalidNick:
		c.outputChan <- fmt.Sprintf(":%s 432 :Erronenous nickname", c.server.name)
	case errNickInUse:
		c.outputChan <- fmt.Sprintf(":%s 433 :Nick already in use", c.server.name)
	case errAlreadyReg:
		c.outputChan <- fmt.Sprintf(":%s 462 :You need a valid nick first", c.server.name)
	case errNoSuchNick:
		c.outputChan <- fmt.Sprintf(":%s 401 :No such nick/channel", c.server.name)
	case errUnknownCommand:
		c.outputChan <- fmt.Sprintf(":%s 421 %s :Unknown command", c.server.name, args[0])
	}
}

func (c *Client) setNick(nick string) {
	if c.nick != "" {
		delete(c.server.clientMap, c.nick)
		for _, channel := range c.channelMap {
			delete(channel.clientMap, c.nick)
		}
	}

	//Set up new nick
	oldNick := c.nick
	c.nick = nick
	c.server.clientMap[c.nick] = c

	clients := make([]string, 0, 100)

	for _, channel := range c.channelMap {
		channel.clientMap[c.nick] = c

		//Collect list of client nicks who can see us
		for client := range channel.clientMap {
			clients = append(clients, client)
		}
	}

	//By sorting the nicks and skipping duplicates we send each client one message
	sort.Strings(clients)
	prevNick := ""
	for _, nick := range clients {
		if nick == prevNick {
			continue
		}
		prevNick = nick

		client, exists := c.server.clientMap[nick]
		if exists {
			client.reply(rplNickChange, oldNick, c.nick)
		}
	}
}
<|endoftext|>