<commit_before>package cluster

import (
	"fmt"
	"net/http"

	"time"

	"github.com/rancher/norman/api/access"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/parse"
	"github.com/rancher/norman/types"
	"github.com/rancher/rancher/pkg/controllers/management/rbac"
	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/ref"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	mgmtclient "github.com/rancher/types/client/management/v3"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	numberOfRetries   = 3
	retryIntervalInMs = 5
)

func (a ActionHandler) saveAsTemplate(actionName string, action *types.Action, apiContext *types.APIContext) error {

	var err error
	var clusterTemplate *v3.ClusterTemplate
	var clusterTemplateRevision *v3.ClusterTemplateRevision

	defer a.cleanup(err, clusterTemplate)

	clusterTempName, clusterTempRevName, err := a.validateTemplateInput(apiContext)
	if err != nil {
		return err
	}

	cluster, err := a.validateClusterState(apiContext)
	if err != nil {
		return err
	}

	//Copy the cluster rke spec to a template and revision object
	clusterTemplate, err = a.createNewClusterTemplate(apiContext, clusterTempName, cluster)
	if err != nil {
		return httperror.WrapAPIError(err, httperror.ServerError,
			fmt.Sprintf("failed to create RKE Template object"))
	}

	clusterTemplateRevision, err = a.createNewClusterTemplateRevision(apiContext, clusterTempRevName, clusterTemplate, cluster)
	if err != nil {
		return httperror.WrapAPIError(err, httperror.ServerError,
			fmt.Sprintf("failed to create RKE Template Revision object"))
	}

	err = a.updateCluster(apiContext, clusterTemplate, clusterTemplateRevision)
	if err != nil {
		return err
	}

	apiContext.WriteResponse(http.StatusNoContent, map[string]interface{}{})
	return nil
}

func (a ActionHandler) validateTemplateInput(apiContext *types.APIContext) (string, string, error) {
	actionInput, err := parse.ReadBody(apiContext.Request)
	if err != nil {
		return "", "", err
	}
	clusterTempName, ok := actionInput["clusterTemplateName"].(string)
	if !ok || len(clusterTempName) == 0 {
		return "", "", httperror.NewAPIError(httperror.InvalidBodyContent, "must specify a name for the RKE Template")
	}
	clusterTempRevName, ok := actionInput["clusterTemplateRevisionName"].(string)
	if !ok || len(clusterTempRevName) == 0 {
		return "", "", httperror.NewAPIError(httperror.InvalidBodyContent, "must specify a name for the RKE Template Revision")
	}
	return clusterTempName, clusterTempRevName, nil
}

func (a ActionHandler) validateClusterState(apiContext *types.APIContext) (*v3.Cluster, error) {
	var clusterForAccessCheck mgmtclient.Cluster
	var err error

	if err := access.ByID(apiContext, apiContext.Version, apiContext.Type, apiContext.ID, &clusterForAccessCheck); err != nil {
		return nil, httperror.NewAPIError(httperror.NotFound,
			fmt.Sprintf("failed to get cluster by id %v", apiContext.ID))
	}

	cluster, err := a.ClusterClient.Get(apiContext.ID, v1.GetOptions{})
	if err != nil {
		return nil, httperror.WrapAPIError(err, httperror.NotFound,
			fmt.Sprintf("cluster with id %v doesn't exist", apiContext.ID))
	}

	if cluster.DeletionTimestamp != nil {
		return nil, httperror.NewAPIError(httperror.InvalidType,
			fmt.Sprintf("cluster with id %v is being deleted", apiContext.ID))
	}
	/*if !v3.ClusterConditionReady.IsTrue(cluster) {
		return nil, httperror.WrapAPIError(err, httperror.ClusterUnavailable,
			fmt.Sprintf("cluster with id %v is not ready", apiContext.ID))
	}*/
	return cluster, nil
}

func (a ActionHandler) createNewClusterTemplate(apiContext *types.APIContext, clusterTempName string, cluster *v3.Cluster) (*v3.ClusterTemplate, error) {
	creatorID := apiContext.Request.Header.Get("Impersonate-User")
	newTemplateObj := a.newClusterTemplate(clusterTempName, cluster, creatorID)
	return a.ClusterTemplateClient.Create(newTemplateObj)
}

func (a ActionHandler) newClusterTemplate(clusterTempName string, cluster *v3.Cluster, creatorID string) *v3.ClusterTemplate {
	return &v3.ClusterTemplate{
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "ct-",
			Namespace:    namespace.GlobalNamespace,
			Labels: map[string]string{
				"cattle.io/creator": "norman",
			},
			Annotations: map[string]string{
				rbac.CreatorIDAnn: creatorID,
			},
		},
		Spec: v3.ClusterTemplateSpec{
			DisplayName: clusterTempName,
			Description: fmt.Sprintf("RKETemplate generated from cluster Id %v", cluster.Name),
		},
	}
}

func (a ActionHandler) createNewClusterTemplateRevision(apiContext *types.APIContext, clusterTempRevName string, clusterTemplate *v3.ClusterTemplate, cluster *v3.Cluster) (*v3.ClusterTemplateRevision, error) {
	creatorID := apiContext.Request.Header.Get("Impersonate-User")
	newTemplateRevObj := a.newClusterTemplateRevision(clusterTempRevName, cluster, clusterTemplate, creatorID)
	return a.ClusterTemplateRevisionClient.Create(newTemplateRevObj)
}

func (a ActionHandler) newClusterTemplateRevision(clusterTempRevisionName string, cluster *v3.Cluster, clusterTemplate *v3.ClusterTemplate, creatorID string) *v3.ClusterTemplateRevision {
	controller := true
	clusterConfig := cluster.Status.AppliedSpec.ClusterSpecBase.DeepCopy()

	return &v3.ClusterTemplateRevision{
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "ctr-",
			Namespace:    namespace.GlobalNamespace,
			OwnerReferences: []v1.OwnerReference{
				{
					Name:       clusterTemplate.Name,
					UID:        clusterTemplate.UID,
					APIVersion: "management.cattle.io/v3",
					Kind:       "ClusterTemplate",
					Controller: &controller,
				},
			},
			Labels: map[string]string{
				"cattle.io/creator":                 "norman",
				"io.cattle.field/clusterTemplateId": clusterTemplate.Name,
			},
			Annotations: map[string]string{
				rbac.CreatorIDAnn: creatorID,
			},
		},
		Spec: v3.ClusterTemplateRevisionSpec{
			DisplayName:         clusterTempRevisionName,
			ClusterTemplateName: ref.Ref(clusterTemplate),
			ClusterConfig:       clusterConfig,
		},
	}
}

func (a ActionHandler) updateCluster(apiContext *types.APIContext, clusterTemplate *v3.ClusterTemplate, clusterTemplateRevision *v3.ClusterTemplateRevision) error {
	// Can't add either too many retries or longer interval as this an API handler
	for i := 0; i < numberOfRetries; i++ {
		cluster, err := a.ClusterClient.Get(apiContext.ID, v1.GetOptions{})
		if err != nil {
			logrus.Errorf("error fetching cluster with id %v: %v", apiContext.ID, err)
			continue
		}
		updatedCluster := cluster.DeepCopy()
		updatedCluster.Spec.ClusterTemplateRevisionName = ref.Ref(clusterTemplateRevision)
		updatedCluster.Spec.ClusterTemplateName = ref.Ref(clusterTemplate)

		_, err = a.ClusterClient.Update(updatedCluster)
		if err == nil {
			return nil
		}
		time.Sleep(retryIntervalInMs * time.Millisecond)
	}
	return httperror.NewAPIError(httperror.Conflict,
		fmt.Sprintf("Error while updating cluster %v with RKE Template Id %v and RKE TemplateRevision Id %v", apiContext.ID, clusterTemplate.Name, clusterTemplateRevision.Name))
}

func (a ActionHandler) cleanup(err error, clusterTemplate *v3.ClusterTemplate) {
	if err == nil {
		return
	}
	if clusterTemplate != nil {
		for i := 0; i < numberOfRetries; i++ {
			//delete the clusterTemplate created, any revision will be deleted due to owner-ref
			deleteErr := a.ClusterTemplateClient.Delete(clusterTemplate.Name, &metav1.DeleteOptions{})
			if deleteErr != nil {
				logrus.Errorf("Failed to delete the RKE Template %v", clusterTemplate.Name)
				continue
			}
			time.Sleep(retryIntervalInMs * time.Millisecond)
		}
	}
}
<commit_msg>Adding response to the SaveAsTemplate action on cluster<commit_after>package cluster

import (
	"fmt"
	"net/http"

	"time"

	"github.com/rancher/norman/api/access"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/parse"
	"github.com/rancher/norman/types"
	"github.com/rancher/rancher/pkg/controllers/management/rbac"
	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/ref"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	mgmtclient "github.com/rancher/types/client/management/v3"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	numberOfRetries   = 3
	retryIntervalInMs = 5
)

func (a ActionHandler) saveAsTemplate(actionName string, action *types.Action, apiContext *types.APIContext) error {

	var err error
	var clusterTemplate *v3.ClusterTemplate
	var clusterTemplateRevision *v3.ClusterTemplateRevision

	defer a.cleanup(err, clusterTemplate)

	clusterTempName, clusterTempRevName, err := a.validateTemplateInput(apiContext)
	if err != nil {
		return err
	}

	cluster, err := a.validateClusterState(apiContext)
	if err != nil {
		return err
	}

	//Copy the cluster rke spec to a template and revision object
	clusterTemplate, err = a.createNewClusterTemplate(apiContext, clusterTempName, cluster)
	if err != nil {
		return httperror.WrapAPIError(err, httperror.ServerError,
			fmt.Sprintf("failed to create RKE Template object"))
	}

	clusterTemplateRevision, err = a.createNewClusterTemplateRevision(apiContext, clusterTempRevName, clusterTemplate, cluster)
	if err != nil {
		return httperror.WrapAPIError(err, httperror.ServerError,
			fmt.Sprintf("failed to create RKE Template Revision object"))
	}

	err = a.updateCluster(apiContext, clusterTemplate, clusterTemplateRevision)
	if err != nil {
		return err
	}

	response := map[string]interface{}{
		"clusterTemplateName":         ref.Ref(clusterTemplate),
		"clusterTemplateRevisionName": ref.Ref(clusterTemplateRevision),
		"type":                        "saveAsTemplateOutput",
	}

	apiContext.WriteResponse(http.StatusOK, response)
	return nil
}

func (a ActionHandler) validateTemplateInput(apiContext *types.APIContext) (string, string, error) {
	actionInput, err := parse.ReadBody(apiContext.Request)
	if err != nil {
		return "", "", err
	}
	clusterTempName, ok := actionInput["clusterTemplateName"].(string)
	if !ok || len(clusterTempName) == 0 {
		return "", "", httperror.NewAPIError(httperror.InvalidBodyContent, "must specify a name for the RKE Template")
	}
	clusterTempRevName, ok := actionInput["clusterTemplateRevisionName"].(string)
	if !ok || len(clusterTempRevName) == 0 {
		return "", "", httperror.NewAPIError(httperror.InvalidBodyContent, "must specify a name for the RKE Template Revision")
	}
	return clusterTempName, clusterTempRevName, nil
}

func (a ActionHandler) validateClusterState(apiContext *types.APIContext) (*v3.Cluster, error) {
	var clusterForAccessCheck mgmtclient.Cluster
	var err error

	if err := access.ByID(apiContext, apiContext.Version, apiContext.Type, apiContext.ID, &clusterForAccessCheck); err != nil {
		return nil, httperror.NewAPIError(httperror.NotFound,
			fmt.Sprintf("failed to get cluster by id %v", apiContext.ID))
	}

	cluster, err := a.ClusterClient.Get(apiContext.ID, v1.GetOptions{})
	if err != nil {
		return nil, httperror.WrapAPIError(err, httperror.NotFound,
			fmt.Sprintf("cluster with id %v doesn't exist", apiContext.ID))
	}

	if cluster.DeletionTimestamp != nil {
		return nil, httperror.NewAPIError(httperror.InvalidType,
			fmt.Sprintf("cluster with id %v is being deleted", apiContext.ID))
	}
	/*if !v3.ClusterConditionReady.IsTrue(cluster) {
		return nil, httperror.WrapAPIError(err, httperror.ClusterUnavailable,
			fmt.Sprintf("cluster with id %v is not ready", apiContext.ID))
	}*/
	return cluster, nil
}

func (a ActionHandler) createNewClusterTemplate(apiContext *types.APIContext, clusterTempName string, cluster *v3.Cluster) (*v3.ClusterTemplate, error) {
	creatorID := apiContext.Request.Header.Get("Impersonate-User")
	newTemplateObj := a.newClusterTemplate(clusterTempName, cluster, creatorID)
	return a.ClusterTemplateClient.Create(newTemplateObj)
}

func (a ActionHandler) newClusterTemplate(clusterTempName string, cluster *v3.Cluster, creatorID string) *v3.ClusterTemplate {
	return &v3.ClusterTemplate{
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "ct-",
			Namespace:    namespace.GlobalNamespace,
			Labels: map[string]string{
				"cattle.io/creator": "norman",
			},
			Annotations: map[string]string{
				rbac.CreatorIDAnn: creatorID,
			},
		},
		Spec: v3.ClusterTemplateSpec{
			DisplayName: clusterTempName,
			Description: fmt.Sprintf("RKETemplate generated from cluster Id %v", cluster.Name),
		},
	}
}

func (a ActionHandler) createNewClusterTemplateRevision(apiContext *types.APIContext, clusterTempRevName string, clusterTemplate *v3.ClusterTemplate, cluster *v3.Cluster) (*v3.ClusterTemplateRevision, error) {
	creatorID := apiContext.Request.Header.Get("Impersonate-User")
	newTemplateRevObj := a.newClusterTemplateRevision(clusterTempRevName, cluster, clusterTemplate, creatorID)
	return a.ClusterTemplateRevisionClient.Create(newTemplateRevObj)
}

func (a ActionHandler) newClusterTemplateRevision(clusterTempRevisionName string, cluster *v3.Cluster, clusterTemplate *v3.ClusterTemplate, creatorID string) *v3.ClusterTemplateRevision {
	controller := true
	clusterConfig := cluster.Status.AppliedSpec.ClusterSpecBase.DeepCopy()

	return &v3.ClusterTemplateRevision{
		ObjectMeta: v1.ObjectMeta{
			GenerateName: "ctr-",
			Namespace:    namespace.GlobalNamespace,
			OwnerReferences: []v1.OwnerReference{
				{
					Name:       clusterTemplate.Name,
					UID:        clusterTemplate.UID,
					APIVersion: "management.cattle.io/v3",
					Kind:       "ClusterTemplate",
					Controller: &controller,
				},
			},
			Labels: map[string]string{
				"cattle.io/creator":                 "norman",
				"io.cattle.field/clusterTemplateId": clusterTemplate.Name,
			},
			Annotations: map[string]string{
				rbac.CreatorIDAnn: creatorID,
			},
		},
		Spec: v3.ClusterTemplateRevisionSpec{
			DisplayName:         clusterTempRevisionName,
			ClusterTemplateName: ref.Ref(clusterTemplate),
			ClusterConfig:       clusterConfig,
		},
	}
}

func (a ActionHandler) updateCluster(apiContext *types.APIContext, clusterTemplate *v3.ClusterTemplate, clusterTemplateRevision *v3.ClusterTemplateRevision) error {
	// Can't add either too many retries or longer interval as this an API handler
	for i := 0; i < numberOfRetries; i++ {
		cluster, err := a.ClusterClient.Get(apiContext.ID, v1.GetOptions{})
		if err != nil {
			logrus.Errorf("error fetching cluster with id %v: %v", apiContext.ID, err)
			continue
		}
		updatedCluster := cluster.DeepCopy()
		updatedCluster.Spec.ClusterTemplateRevisionName = ref.Ref(clusterTemplateRevision)
		updatedCluster.Spec.ClusterTemplateName = ref.Ref(clusterTemplate)

		_, err = a.ClusterClient.Update(updatedCluster)
		if err == nil {
			return nil
		}
		time.Sleep(retryIntervalInMs * time.Millisecond)
	}
	return httperror.NewAPIError(httperror.Conflict,
		fmt.Sprintf("Error while updating cluster %v with RKE Template Id %v and RKE TemplateRevision Id %v", apiContext.ID, clusterTemplate.Name, clusterTemplateRevision.Name))
}

func (a ActionHandler) cleanup(err error, clusterTemplate *v3.ClusterTemplate) {
	if err == nil {
		return
	}
	if clusterTemplate != nil {
		for i := 0; i < numberOfRetries; i++ {
			//delete the clusterTemplate created, any revision will be deleted due to owner-ref
			deleteErr := a.ClusterTemplateClient.Delete(clusterTemplate.Name, &metav1.DeleteOptions{})
			if deleteErr != nil {
				logrus.Errorf("Failed to delete the RKE Template %v", clusterTemplate.Name)
				continue
			}
			time.Sleep(retryIntervalInMs * time.Millisecond)
		}
	}
}
<|endoftext|>


<commit_before>package rkeworkerupgrader

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/pkg/errors"
	kd "github.com/rancher/rancher/pkg/controllers/management/kontainerdrivermetadata"
	"github.com/rancher/rancher/pkg/librke"
	nodeserver "github.com/rancher/rancher/pkg/rkenodeconfigserver"
	rkehosts "github.com/rancher/rke/hosts"
	rkeservices "github.com/rancher/rke/services"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
)

func (uh *upgradeHandler) nonWorkerPlan(node *v3.Node, cluster *v3.Cluster) (*v3.RKEConfigNodePlan, error) {
	rkeConfig := cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.DeepCopy()
	rkeConfig.Nodes = []v3.RKEConfigNode{
		*node.Status.NodeConfig,
	}
	rkeConfig.Nodes[0].Role = []string{rkeservices.WorkerRole, rkeservices.ETCDRole, rkeservices.ControlRole}

	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	hostAddress := node.Status.NodeConfig.Address
	hostDockerInfo := infos[hostAddress]

	logrus.Debugf("getDockerInfo for node [%s] dockerInfo [%s]", node.Name, hostDockerInfo.DockerRootDir)

	svcOptions, err := uh.getServiceOptions(rkeConfig.Version, hostDockerInfo.OSType)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(uh.ctx, rkeConfig, infos, svcOptions)
	if err != nil {
		return nil, err
	}

	token, err := uh.systemAccountManager.GetOrCreateSystemClusterToken(cluster.Name)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create or get cluster token for share-mnt")
	}

	np := &v3.RKEConfigNodePlan{}

	for _, tempNode := range plan.Nodes {
		if tempNode.Address == hostAddress {

			b2d := strings.Contains(infos[tempNode.Address].OperatingSystem, rkehosts.B2DOS)
			np.Processes = nodeserver.AugmentProcesses(token, tempNode.Processes, false, b2d,
				node.Status.NodeConfig.HostnameOverride, cluster)

			np.Processes = nodeserver.AppendTaintsToKubeletArgs(np.Processes, node.Status.NodeConfig.Taints)

			return np, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for %s", hostAddress)
}

func (uh *upgradeHandler) workerPlan(node *v3.Node, cluster *v3.Cluster) (*v3.RKEConfigNodePlan, error) {
	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	hostAddress := node.Status.NodeConfig.Address
	hostDockerInfo := infos[hostAddress]

	rkeConfig := cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.DeepCopy()
	nodeserver.FilterHostForSpec(rkeConfig, node)

	logrus.Debugf("The number of nodes sent to the plan: %v", len(rkeConfig.Nodes))
	svcOptions, err := uh.getServiceOptions(rkeConfig.Version, hostDockerInfo.OSType)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(uh.ctx, rkeConfig, infos, svcOptions)
	if err != nil {
		return nil, err
	}

	logrus.Debugf("getDockerInfo for node [%s] dockerInfo [%s]", node.Name, hostDockerInfo.DockerRootDir)

	np := &v3.RKEConfigNodePlan{}

	token, err := uh.systemAccountManager.GetOrCreateSystemClusterToken(cluster.Name)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create or get cluster token for share-mnt")
	}
	for _, tempNode := range plan.Nodes {
		if tempNode.Address == hostAddress {
			if hostDockerInfo.OSType == "windows" { // compatible with Windows
				np.Processes = nodeserver.EnhanceWindowsProcesses(tempNode.Processes)
			} else {
				b2d := strings.Contains(infos[tempNode.Address].OperatingSystem, rkehosts.B2DOS)
				np.Processes = nodeserver.AugmentProcesses(token, tempNode.Processes, true, b2d,
					node.Status.NodeConfig.HostnameOverride, cluster)
			}
			np.Processes = nodeserver.AppendTaintsToKubeletArgs(np.Processes, node.Status.NodeConfig.Taints)
			np.Files = tempNode.Files

			return np, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for %s", hostAddress)

}

func (uh *upgradeHandler) getServiceOptions(k8sVersion string, osType string) (map[string]interface{}, error) {
	data := map[string]interface{}{}
	svcOptions, err := kd.GetRKEK8sServiceOptions(k8sVersion, uh.serviceOptionsLister,
		uh.serviceOptions, uh.sysImagesLister, uh.sysImages, kd.Linux)
	if err != nil {
		logrus.Errorf("getK8sServiceOptions: k8sVersion %s [%v]", k8sVersion, err)
		return data, err
	}
	if svcOptions != nil {
		data["k8s-service-options"] = svcOptions
	}
	if osType == "windows" {
		svcOptionsWindows, err := kd.GetRKEK8sServiceOptions(k8sVersion, uh.serviceOptionsLister,
			uh.serviceOptions, uh.sysImagesLister, uh.sysImages, kd.Windows)
		if err != nil {
			logrus.Errorf("getK8sServiceOptionsWindows: k8sVersion %s [%v]", k8sVersion, err)
			return data, err
		}
		if svcOptionsWindows != nil {
			data["k8s-windows-service-options"] = svcOptionsWindows
		}
	}
	return data, nil
}

func planChangedForUpgrade(newPlan *v3.RKEConfigNodePlan, oldPlan *v3.RKEConfigNodePlan) bool {
	if newPlan == nil || oldPlan == nil {
		return true
	}
	newProcesses := newPlan.Processes
	oldProcesses := oldPlan.Processes

	if len(newProcesses) != len(oldProcesses) {
		logrus.Infof("number of processes changed: old: %v new: %v", len(oldProcesses), len(newProcesses))
		return true
	}

	for k, newProcess := range newProcesses {
		if strings.Contains(k, "share-mnt") {
			// don't need to upgrade if share-mnt changed
			continue
		}
		oldProcess, ok := oldProcesses[k]
		if !ok {
			return true
		}

		if processChanged(newProcess, oldProcess) {
			return true
		}
	}

	return false
}

func processChanged(newp v3.Process, oldp v3.Process) bool {
	name := newp.Name

	if oldp.Image != newp.Image {
		logrus.Infof("image changed for [%s] old: %s new: %s", name, oldp.Image, newp.Image)
		return true
	}

	if sliceChangedUnordered(oldp.Command, newp.Command) {
		logrus.Infof("command changed for [%s] old: %v new: %v", name, oldp.Command, newp.Command)
		return true
	}

	if sliceChangedUnordered(oldp.Env, newp.Env) {
		logrus.Infof("env changed for [%s] old: %v new %v", name, oldp.Env, newp.Env)
		return true
	}

	if sliceChangedUnordered(oldp.Args, newp.Args) {
		logrus.Infof("args changed for [%s] old: %v new %v", name, oldp.Args, newp.Args)
		return true
	}

	if sliceChanged(oldp.Binds, newp.Binds) {
		logrus.Infof("binds changed for [%s] old: %v new %v", name, oldp.Binds, newp.Binds)
		return true
	}

	if sliceChanged(oldp.VolumesFrom, newp.VolumesFrom) {
		logrus.Infof("volumesFrom changed for [%s] old: %v new %v", name, oldp.VolumesFrom, newp.VolumesFrom)
		return true
	}

	if sliceChanged(oldp.Publish, newp.Publish) {
		logrus.Infof("publish changed for [%s] old: %v new %v", name, oldp.Publish, newp.Publish)
		return true
	}

	oldProcess := forCompareProcess(oldp)
	newProcess := forCompareProcess(newp)

	if !reflect.DeepEqual(oldProcess, newProcess) {
		logrus.Infof("process changed for [%s] old: %#v new %#v", name, oldProcess, newProcess)
		return true
	}

	return false

}

type compareProcess struct {
	Labels      map[string]string
	NetworkMode string
	PidMode     string
	Privileged  bool
}

func forCompareProcess(p v3.Process) compareProcess {
	return compareProcess{
		Labels:      p.Labels,
		NetworkMode: p.NetworkMode,
		PidMode:     p.PidMode,
		Privileged:  p.Privileged,
	}
}

func sliceChangedUnordered(olds, news []string) bool {
	oldMap := sliceToMap(olds)
	newMap := sliceToMap(news)

	if !reflect.DeepEqual(oldMap, newMap) {
		return true
	}
	return false
}

func sliceChanged(olds, news []string) bool {
	if len(olds) == 0 && len(news) == 0 {
		// DeepEqual considers []string{} and []string(nil) as different
		return false
	}

	if !reflect.DeepEqual(olds, news) {
		return true
	}

	return false
}

func planChangedForUpdate(newPlan, oldPlan *v3.RKEConfigNodePlan) bool {
	// files passed in node config
	if !reflect.DeepEqual(newPlan.Files, oldPlan.Files) {
		return true
	}
	// things that aren't included to restart container
	for k, newProcess := range newPlan.Processes {
		if oldProcess, ok := oldPlan.Processes[k]; ok {
			if oldProcess.Name != newProcess.Name {
				return true
			}

			if oldProcess.HealthCheck.URL != newProcess.HealthCheck.URL {
				return true
			}

			if oldProcess.RestartPolicy != newProcess.RestartPolicy {
				return true
			}

			if oldProcess.ImageRegistryAuthConfig != newProcess.ImageRegistryAuthConfig {
				return true
			}

			if strings.Contains(k, "share-mnt") {
				// don't need to upgrade if share-mnt changed
				if processChanged(newProcess, oldProcess) {
					return true
				}
			}
		}
	}

	return false

}

func sliceToMap(s []string) map[string]bool {
	m := map[string]bool{}
	for _, v := range s {
		m[v] = true
	}
	return m
}
<commit_msg>Compare binds unordered for worker upgrade plan<commit_after>package rkeworkerupgrader

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/pkg/errors"
	kd "github.com/rancher/rancher/pkg/controllers/management/kontainerdrivermetadata"
	"github.com/rancher/rancher/pkg/librke"
	nodeserver "github.com/rancher/rancher/pkg/rkenodeconfigserver"
	rkehosts "github.com/rancher/rke/hosts"
	rkeservices "github.com/rancher/rke/services"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
)

func (uh *upgradeHandler) nonWorkerPlan(node *v3.Node, cluster *v3.Cluster) (*v3.RKEConfigNodePlan, error) {
	rkeConfig := cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.DeepCopy()
	rkeConfig.Nodes = []v3.RKEConfigNode{
		*node.Status.NodeConfig,
	}
	rkeConfig.Nodes[0].Role = []string{rkeservices.WorkerRole, rkeservices.ETCDRole, rkeservices.ControlRole}

	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	hostAddress := node.Status.NodeConfig.Address
	hostDockerInfo := infos[hostAddress]

	logrus.Debugf("getDockerInfo for node [%s] dockerInfo [%s]", node.Name, hostDockerInfo.DockerRootDir)

	svcOptions, err := uh.getServiceOptions(rkeConfig.Version, hostDockerInfo.OSType)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(uh.ctx, rkeConfig, infos, svcOptions)
	if err != nil {
		return nil, err
	}

	token, err := uh.systemAccountManager.GetOrCreateSystemClusterToken(cluster.Name)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create or get cluster token for share-mnt")
	}

	np := &v3.RKEConfigNodePlan{}

	for _, tempNode := range plan.Nodes {
		if tempNode.Address == hostAddress {

			b2d := strings.Contains(infos[tempNode.Address].OperatingSystem, rkehosts.B2DOS)
			np.Processes = nodeserver.AugmentProcesses(token, tempNode.Processes, false, b2d,
				node.Status.NodeConfig.HostnameOverride, cluster)

			np.Processes = nodeserver.AppendTaintsToKubeletArgs(np.Processes, node.Status.NodeConfig.Taints)

			return np, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for %s", hostAddress)
}

func (uh *upgradeHandler) workerPlan(node *v3.Node, cluster *v3.Cluster) (*v3.RKEConfigNodePlan, error) {
	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	hostAddress := node.Status.NodeConfig.Address
	hostDockerInfo := infos[hostAddress]

	rkeConfig := cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.DeepCopy()
	nodeserver.FilterHostForSpec(rkeConfig, node)

	logrus.Debugf("The number of nodes sent to the plan: %v", len(rkeConfig.Nodes))
	svcOptions, err := uh.getServiceOptions(rkeConfig.Version, hostDockerInfo.OSType)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(uh.ctx, rkeConfig, infos, svcOptions)
	if err != nil {
		return nil, err
	}

	logrus.Debugf("getDockerInfo for node [%s] dockerInfo [%s]", node.Name, hostDockerInfo.DockerRootDir)

	np := &v3.RKEConfigNodePlan{}

	token, err := uh.systemAccountManager.GetOrCreateSystemClusterToken(cluster.Name)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create or get cluster token for share-mnt")
	}
	for _, tempNode := range plan.Nodes {
		if tempNode.Address == hostAddress {
			if hostDockerInfo.OSType == "windows" { // compatible with Windows
				np.Processes = nodeserver.EnhanceWindowsProcesses(tempNode.Processes)
			} else {
				b2d := strings.Contains(infos[tempNode.Address].OperatingSystem, rkehosts.B2DOS)
				np.Processes = nodeserver.AugmentProcesses(token, tempNode.Processes, true, b2d,
					node.Status.NodeConfig.HostnameOverride, cluster)
			}
			np.Processes = nodeserver.AppendTaintsToKubeletArgs(np.Processes, node.Status.NodeConfig.Taints)
			np.Files = tempNode.Files

			return np, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for %s", hostAddress)

}

func (uh *upgradeHandler) getServiceOptions(k8sVersion string, osType string) (map[string]interface{}, error) {
	data := map[string]interface{}{}
	svcOptions, err := kd.GetRKEK8sServiceOptions(k8sVersion, uh.serviceOptionsLister,
		uh.serviceOptions, uh.sysImagesLister, uh.sysImages, kd.Linux)
	if err != nil {
		logrus.Errorf("getK8sServiceOptions: k8sVersion %s [%v]", k8sVersion, err)
		return data, err
	}
	if svcOptions != nil {
		data["k8s-service-options"] = svcOptions
	}
	if osType == "windows" {
		svcOptionsWindows, err := kd.GetRKEK8sServiceOptions(k8sVersion, uh.serviceOptionsLister,
			uh.serviceOptions, uh.sysImagesLister, uh.sysImages, kd.Windows)
		if err != nil {
			logrus.Errorf("getK8sServiceOptionsWindows: k8sVersion %s [%v]", k8sVersion, err)
			return data, err
		}
		if svcOptionsWindows != nil {
			data["k8s-windows-service-options"] = svcOptionsWindows
		}
	}
	return data, nil
}

func planChangedForUpgrade(newPlan *v3.RKEConfigNodePlan, oldPlan *v3.RKEConfigNodePlan) bool {
	if newPlan == nil || oldPlan == nil {
		return true
	}
	newProcesses := newPlan.Processes
	oldProcesses := oldPlan.Processes

	if len(newProcesses) != len(oldProcesses) {
		logrus.Infof("number of processes changed: old: %v new: %v", len(oldProcesses), len(newProcesses))
		return true
	}

	for k, newProcess := range newProcesses {
		if strings.Contains(k, "share-mnt") {
			// don't need to upgrade if share-mnt changed
			continue
		}
		oldProcess, ok := oldProcesses[k]
		if !ok {
			return true
		}

		if processChanged(newProcess, oldProcess) {
			return true
		}
	}

	return false
}

func processChanged(newp v3.Process, oldp v3.Process) bool {
	name := newp.Name

	if oldp.Image != newp.Image {
		logrus.Infof("image changed for [%s] old: %s new: %s", name, oldp.Image, newp.Image)
		return true
	}

	if sliceChangedUnordered(oldp.Command, newp.Command) {
		logrus.Infof("command changed for [%s] old: %v new: %v", name, oldp.Command, newp.Command)
		return true
	}

	if sliceChangedUnordered(oldp.Env, newp.Env) {
		logrus.Infof("env changed for [%s] old: %v new %v", name, oldp.Env, newp.Env)
		return true
	}

	if sliceChangedUnordered(oldp.Args, newp.Args) {
		logrus.Infof("args changed for [%s] old: %v new %v", name, oldp.Args, newp.Args)
		return true
	}

	if sliceChangedUnordered(oldp.Binds, newp.Binds) {
		logrus.Infof("binds changed for [%s] old: %v new %v", name, oldp.Binds, newp.Binds)
		return true
	}

	if sliceChanged(oldp.VolumesFrom, newp.VolumesFrom) {
		logrus.Infof("volumesFrom changed for [%s] old: %v new %v", name, oldp.VolumesFrom, newp.VolumesFrom)
		return true
	}

	if sliceChanged(oldp.Publish, newp.Publish) {
		logrus.Infof("publish changed for [%s] old: %v new %v", name, oldp.Publish, newp.Publish)
		return true
	}

	oldProcess := forCompareProcess(oldp)
	newProcess := forCompareProcess(newp)

	if !reflect.DeepEqual(oldProcess, newProcess) {
		logrus.Infof("process changed for [%s] old: %#v new %#v", name, oldProcess, newProcess)
		return true
	}

	return false

}

type compareProcess struct {
	Labels      map[string]string
	NetworkMode string
	PidMode     string
	Privileged  bool
}

func forCompareProcess(p v3.Process) compareProcess {
	return compareProcess{
		Labels:      p.Labels,
		NetworkMode: p.NetworkMode,
		PidMode:     p.PidMode,
		Privileged:  p.Privileged,
	}
}

func sliceChangedUnordered(olds, news []string) bool {
	oldMap := sliceToMap(olds)
	newMap := sliceToMap(news)

	if !reflect.DeepEqual(oldMap, newMap) {
		return true
	}
	return false
}

func sliceChanged(olds, news []string) bool {
	if len(olds) == 0 && len(news) == 0 {
		// DeepEqual considers []string{} and []string(nil) as different
		return false
	}

	if !reflect.DeepEqual(olds, news) {
		return true
	}

	return false
}

func planChangedForUpdate(newPlan, oldPlan *v3.RKEConfigNodePlan) bool {
	// files passed in node config
	if !reflect.DeepEqual(newPlan.Files, oldPlan.Files) {
		return true
	}
	// things that aren't included to restart container
	for k, newProcess := range newPlan.Processes {
		if oldProcess, ok := oldPlan.Processes[k]; ok {
			if oldProcess.Name != newProcess.Name {
				return true
			}

			if oldProcess.HealthCheck.URL != newProcess.HealthCheck.URL {
				return true
			}

			if oldProcess.RestartPolicy != newProcess.RestartPolicy {
				return true
			}

			if oldProcess.ImageRegistryAuthConfig != newProcess.ImageRegistryAuthConfig {
				return true
			}

			if strings.Contains(k, "share-mnt") {
				// don't need to upgrade if share-mnt changed
				if processChanged(newProcess, oldProcess) {
					return true
				}
			}
		}
	}

	return false

}

func sliceToMap(s []string) map[string]bool {
	m := map[string]bool{}
	for _, v := range s {
		m[v] = true
	}
	return m
}
<|endoftext|>


<commit_before>// Package myers implements the Myers diff algorithm.
package myers

import (
	"context"
	"fmt"

	"github.com/pkg/diff/edit"
)

// A Pair is two things that can be diffed using the Myers diff algorithm.
// A is the initial state; B is the final state.
type Pair interface {
	// LenA returns the number of initial elements.
	LenA() int
	// LenB returns the number of final elements.
	LenB() int
	// Equal reports whether the aᵢ'th element of A is equal to the bᵢ'th element of B.
	Equal(ai, bi int) bool
}

// Diff calculates an edit.Script for ab using the Myers diff algorithm.
// Because diff calculation can be expensive, Myers supports cancellation via ctx.
func Diff(ctx context.Context, ab Pair) edit.Script {
	aLen := ab.LenA()
	bLen := ab.LenB()
	if aLen == 0 {
		return edit.NewScript(edit.Range{HighB: bLen})
	}
	if bLen == 0 {
		return edit.NewScript(edit.Range{HighA: aLen})
	}

	max := aLen + bLen
	if max < 0 {
		panic("overflow in diff.Myers")
	}
	v := make([]int, 2*max+1) // indices: -max .. 0 .. max

	var trace [][]int
search:
	for d := 0; d < max; d++ {
		// Only check context every 16th iteration to reduce overhead.
		if ctx != nil && uint(d)%16 == 0 && ctx.Err() != nil {
			return edit.Script{}
		}

		// TODO: this seems like it will frequently be bigger than necessary.
		// Use sparse lookup? prefixes?
		vc := make([]int, 2*max+1)
		copy(vc, v)
		trace = append(trace, vc)

		for k := -d; k <= d; k += 2 {
			var x int
			if k == -d || (k != d && v[max+k-1] < v[max+k+1]) {
				x = v[max+k+1]
			} else {
				x = v[max+k-1] + 1
			}

			y := x - k
			for x < aLen && y < bLen && ab.Equal(x, y) {
				x++
				y++
			}
			v[max+k] = x

			if x == aLen && y == bLen {
				break search
			}
		}
	}

	if len(trace) == max {
		// No commonality at all, delete everything and then insert everything.
		// This is handled as a special case to avoid complicating the logic below.
		return edit.NewScript(edit.Range{HighA: aLen}, edit.Range{HighB: bLen})
	}

	// Create reversed edit script.
	x := aLen
	y := bLen
	var e edit.Script
	for d := len(trace) - 1; d >= 0; d-- {
		v := trace[d]
		k := x - y
		var prevk int
		if k == -d || (k != d && v[max+k-1] < v[max+k+1]) {
			prevk = k + 1
		} else {
			prevk = k - 1
		}
		prevx := v[max+prevk]
		prevy := prevx - prevk
		for x > prevx && y > prevy {
			appendToReversed(&e, edit.Range{LowA: x - 1, LowB: y - 1, HighA: x, HighB: y})
			x--
			y--
		}
		if d > 0 {
			appendToReversed(&e, edit.Range{LowA: prevx, LowB: prevy, HighA: x, HighB: y})
		}
		x, y = prevx, prevy
	}

	// Reverse reversed edit script, to return to natural order.
	reverse(e)

	// Sanity check
	for i := 1; i < len(e.Ranges); i++ {
		prevop := e.Ranges[i-1].Op()
		currop := e.Ranges[i].Op()
		if (prevop == currop) || (prevop == edit.Ins && currop != edit.Eq) || (currop == edit.Del && prevop != edit.Eq) {
			panic(fmt.Errorf("bad script: %v -> %v", prevop, currop))
		}
	}

	return e
}

func reverse(e edit.Script) {
	for i := 0; i < len(e.Ranges)/2; i++ {
		j := len(e.Ranges) - i - 1
		e.Ranges[i], e.Ranges[j] = e.Ranges[j], e.Ranges[i]
	}
}

func appendToReversed(e *edit.Script, seg edit.Range) {
	if len(e.Ranges) == 0 {
		e.Ranges = append(e.Ranges, seg)
		return
	}
	u, ok := combineRanges(seg, e.Ranges[len(e.Ranges)-1])
	if !ok {
		e.Ranges = append(e.Ranges, seg)
		return
	}
	e.Ranges[len(e.Ranges)-1] = u
	return
}

// combineRanges combines s and t into a single edit.Range if possible
// and reports whether it succeeded.
func combineRanges(s, t edit.Range) (u edit.Range, ok bool) {
	if t.Len() == 0 {
		return s, true
	}
	if s.Len() == 0 {
		return t, true
	}
	if s.Op() != t.Op() {
		return edit.Range{LowA: -1, HighA: -1, LowB: -1, HighB: -1}, false
	}
	switch s.Op() {
	case edit.Ins:
		s.HighB = t.HighB
	case edit.Del:
		s.HighA = t.HighA
	case edit.Eq:
		s.HighA = t.HighA
		s.HighB = t.HighB
	default:
		panic("bad op")
	}
	return s, true
}

func rangeString(r edit.Range) string {
	// This output is helpful when hacking on a Myers diff.
	// In other contexts it is usually more natural to group LowA, HighA and LowB, HighB.
	return fmt.Sprintf("(%d, %d) -- %s %d --> (%d, %d)", r.LowA, r.LowB, r.Op(), r.Len(), r.HighA, r.HighB)
}
<commit_msg>myers: fix stupid amounts of allocation and copying<commit_after>// Package myers implements the Myers diff algorithm.
package myers

import (
	"context"
	"fmt"

	"github.com/pkg/diff/edit"
)

// A Pair is two things that can be diffed using the Myers diff algorithm.
// A is the initial state; B is the final state.
type Pair interface {
	// LenA returns the number of initial elements.
	LenA() int
	// LenB returns the number of final elements.
	LenB() int
	// Equal reports whether the aᵢ'th element of A is equal to the bᵢ'th element of B.
	Equal(ai, bi int) bool
}

// Diff calculates an edit.Script for ab using the Myers diff algorithm.
// Because diff calculation can be expensive, Myers supports cancellation via ctx.
func Diff(ctx context.Context, ab Pair) edit.Script {
	aLen := ab.LenA()
	bLen := ab.LenB()
	if aLen == 0 {
		return edit.NewScript(edit.Range{HighB: bLen})
	}
	if bLen == 0 {
		return edit.NewScript(edit.Range{HighA: aLen})
	}

	max := aLen + bLen
	if max < 0 {
		panic("overflow in myers.Diff")
	}
	// v has indices -max .. 0 .. max
	// access to elements of v have the form max + actual offset
	v := make([]int, 2*max+1)

	var trace [][]int
search:
	for d := 0; d < max; d++ {
		// Only check context every 16th iteration to reduce overhead.
		if ctx != nil && uint(d)%16 == 0 && ctx.Err() != nil {
			return edit.Script{}
		}

		// append the middle (populated) elements of v to trace
		middle := v[max-d : max+d+1]
		vcopy := make([]int, len(middle))
		copy(vcopy, middle)
		trace = append(trace, vcopy)

		for k := -d; k <= d; k += 2 {
			var x int
			if k == -d || (k != d && v[max+k-1] < v[max+k+1]) {
				x = v[max+k+1]
			} else {
				x = v[max+k-1] + 1
			}

			y := x - k
			for x < aLen && y < bLen && ab.Equal(x, y) {
				x++
				y++
			}
			v[max+k] = x

			if x == aLen && y == bLen {
				break search
			}
		}
	}

	if len(trace) == max {
		// No commonality at all, delete everything and then insert everything.
		// This is handled as a special case to avoid complicating the logic below.
		return edit.NewScript(edit.Range{HighA: aLen}, edit.Range{HighB: bLen})
	}

	// Create reversed edit script.
	x := aLen
	y := bLen
	var e edit.Script
	for d := len(trace) - 1; d >= 0; d-- {
		// v has indices -d .. 0 .. d
		// access to elements of v have the form d + actual offset
		v := trace[d]
		k := x - y
		var prevk int
		if k == -d || (k != d && v[d+k-1] < v[d+k+1]) {
			prevk = k + 1
		} else {
			prevk = k - 1
		}
		var prevx int
		if idx := d + prevk; 0 <= idx && idx < len(v) {
			prevx = v[idx]
		}
		prevy := prevx - prevk
		for x > prevx && y > prevy {
			appendToReversed(&e, edit.Range{LowA: x - 1, LowB: y - 1, HighA: x, HighB: y})
			x--
			y--
		}
		if d > 0 {
			appendToReversed(&e, edit.Range{LowA: prevx, LowB: prevy, HighA: x, HighB: y})
		}
		x, y = prevx, prevy
	}

	// Reverse reversed edit script, to return to natural order.
	reverse(e)

	// Sanity check
	for i := 1; i < len(e.Ranges); i++ {
		prevop := e.Ranges[i-1].Op()
		currop := e.Ranges[i].Op()
		if (prevop == currop) || (prevop == edit.Ins && currop != edit.Eq) || (currop == edit.Del && prevop != edit.Eq) {
			panic(fmt.Errorf("bad script: %v -> %v", prevop, currop))
		}
	}

	return e
}

func reverse(e edit.Script) {
	for i := 0; i < len(e.Ranges)/2; i++ {
		j := len(e.Ranges) - i - 1
		e.Ranges[i], e.Ranges[j] = e.Ranges[j], e.Ranges[i]
	}
}

func appendToReversed(e *edit.Script, seg edit.Range) {
	if len(e.Ranges) == 0 {
		e.Ranges = append(e.Ranges, seg)
		return
	}
	u, ok := combineRanges(seg, e.Ranges[len(e.Ranges)-1])
	if !ok {
		e.Ranges = append(e.Ranges, seg)
		return
	}
	e.Ranges[len(e.Ranges)-1] = u
	return
}

// combineRanges combines s and t into a single edit.Range if possible
// and reports whether it succeeded.
func combineRanges(s, t edit.Range) (u edit.Range, ok bool) {
	if t.Len() == 0 {
		return s, true
	}
	if s.Len() == 0 {
		return t, true
	}
	if s.Op() != t.Op() {
		return edit.Range{LowA: -1, HighA: -1, LowB: -1, HighB: -1}, false
	}
	switch s.Op() {
	case edit.Ins:
		s.HighB = t.HighB
	case edit.Del:
		s.HighA = t.HighA
	case edit.Eq:
		s.HighA = t.HighA
		s.HighB = t.HighB
	default:
		panic("bad op")
	}
	return s, true
}

func rangeString(r edit.Range) string {
	// This output is helpful when hacking on a Myers diff.
	// In other contexts it is usually more natural to group LowA, HighA and LowB, HighB.
	return fmt.Sprintf("(%d, %d) -- %s %d --> (%d, %d)", r.LowA, r.LowB, r.Op(), r.Len(), r.HighA, r.HighB)
}
<|endoftext|>


<commit_before>package naivebayes

/*
NaiveBayes text classifier.
Based on: http://nlp.stanford.edu/IR-book/html/htmledition/naive-bayes-text-classification-1.html

TODO:
	* smarter text parsing ( stemming, synonyms )
*/

import (
	"encoding/json"
	"io/ioutil"
	"math"
	"strings"
)

type Observation struct {
	Classes    []string
	WordCounts map[string]int
}

func NewObservationFromText(classes []string, text string) *Observation {
	counts := make(map[string]int)
	for _, word := range strings.Split(text, " ") {
		_, ok := counts[word]
		if !ok {
			counts[word] = 0
		}
		counts[word]++
	}
	return &Observation{Classes: classes, WordCounts: counts}
}

type Prediction map[string]float64

func (p *Prediction) BestFit() (bestClassName string, bestProbability float64) {
	for className, probability := range *p {
		if bestClassName == "" || probability > bestProbability {
			bestClassName = className
			bestProbability = probability
		}
	}
	return bestClassName, bestProbability
}

/*
Class struct
*/
type Class struct {
	Name             string
	ObservationCount int
	WordCounts       map[string]int
	TotalCount       int
}

// Build an empty class struct
func NewClass(name string) *Class {
	return &Class{Name: name, WordCounts: make(map[string]int), TotalCount: 0}
}

// Increment the count for the given word on the class
func (c *Class) addWord(word string, count int) {
	c.WordCounts[word] += count
	c.TotalCount += count
}

/*
Model struct
*/
type Model struct {
	Name             string
	Classes          map[string]*Class
	ObservationCount int
	Vocabulary       map[string]int
}

func NewModel(name string) *Model {
	return &Model{Name: name, Classes: make(map[string]*Class), ObservationCount: 0, Vocabulary: make(map[string]int)}
}

func NewModelFromFile(path string) (m *Model, err error) {
	// read the whole thing at once
	m = &Model{}
	jsonModel, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(jsonModel, m)
	if err != nil {
		return nil, err
	}
	return m, nil
}

func (m *Model) SaveToFile(path string) (err error) {
	jsonModel, _ := json.Marshal(m)
	// write the whole thing at once
	err = ioutil.WriteFile(path, jsonModel, 0644)
	return err
}

/*
 P( class ) - prior probability of this class within the model
 Number of occurences of the class / Number of trained observations
*/
func (m *Model) classPriorProbability(class *Class) (p float64) {
	p = math.Log(float64(class.ObservationCount) / float64(m.ObservationCount))
	return p

}

/*
 P( observation | class ) - conditional probability of this observation
 given the class.
 Multiply the probabilities of each word being in this class (sum logs)
 P( word | class ) - probability of a word, given a class
 Occurences of the word within the class, divided by total number of words
 Laplace smoothing (always add one so new words don't break everything)
 and also add unique word count to denominator
*/
func (m *Model) classConditionalProbability(class *Class, o *Observation) (p float64) {
	p = 0
	for word, count := range o.WordCounts {
		classWordCount, _ := class.WordCounts[word]
		raw := float64(classWordCount+1) / float64(class.TotalCount+len(m.Vocabulary))
		p = p + (math.Log(raw) * float64(count))
	}
	return p
}

// train the model with the given observation
func (m *Model) Train(o *Observation) {
	for _, className := range o.Classes {
		class, ok := m.Classes[className]
		if !ok {
			class = NewClass(className)
			m.Classes[className] = class
		}
		class.ObservationCount++
		for word, count := range o.WordCounts {
			class.addWord(word, count)
			m.Vocabulary[word] = 1
		}
	}

	m.ObservationCount++
}

// calculate observation membership probability for each class
func (m *Model) Predict(o *Observation) (p Prediction) {
	p = make(map[string]float64)

	for _, class := range m.Classes {
		classConditionalProbability := m.classPriorProbability(class) + m.classConditionalProbability(class, o)
		p[class.Name] = math.Exp(classConditionalProbability)
	}
	return p
}
<commit_msg>add todo note<commit_after>package naivebayes

/*
NaiveBayes text classifier.
Based on: http://nlp.stanford.edu/IR-book/html/htmledition/naive-bayes-text-classification-1.html

TODO:
	* add documentation
	* smarter text parsing ( stemming, synonyms )
*/

import (
	"encoding/json"
	"io/ioutil"
	"math"
	"strings"
)

type Observation struct {
	Classes    []string
	WordCounts map[string]int
}

func NewObservationFromText(classes []string, text string) *Observation {
	counts := make(map[string]int)
	for _, word := range strings.Split(text, " ") {
		_, ok := counts[word]
		if !ok {
			counts[word] = 0
		}
		counts[word]++
	}
	return &Observation{Classes: classes, WordCounts: counts}
}

type Prediction map[string]float64

func (p *Prediction) BestFit() (bestClassName string, bestProbability float64) {
	for className, probability := range *p {
		if bestClassName == "" || probability > bestProbability {
			bestClassName = className
			bestProbability = probability
		}
	}
	return bestClassName, bestProbability
}

/*
Class struct
*/
type Class struct {
	Name             string
	ObservationCount int
	WordCounts       map[string]int
	TotalCount       int
}

// Build an empty class struct
func NewClass(name string) *Class {
	return &Class{Name: name, WordCounts: make(map[string]int), TotalCount: 0}
}

// Increment the count for the given word on the class
func (c *Class) addWord(word string, count int) {
	c.WordCounts[word] += count
	c.TotalCount += count
}

/*
Model struct
*/
type Model struct {
	Name             string
	Classes          map[string]*Class
	ObservationCount int
	Vocabulary       map[string]int
}

func NewModel(name string) *Model {
	return &Model{Name: name, Classes: make(map[string]*Class), ObservationCount: 0, Vocabulary: make(map[string]int)}
}

func NewModelFromFile(path string) (m *Model, err error) {
	// read the whole thing at once
	m = &Model{}
	jsonModel, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(jsonModel, m)
	if err != nil {
		return nil, err
	}
	return m, nil
}

func (m *Model) SaveToFile(path string) (err error) {
	jsonModel, _ := json.Marshal(m)
	// write the whole thing at once
	err = ioutil.WriteFile(path, jsonModel, 0644)
	return err
}

/*
 P( class ) - prior probability of this class within the model
 Number of occurences of the class / Number of trained observations
*/
func (m *Model) classPriorProbability(class *Class) (p float64) {
	p = math.Log(float64(class.ObservationCount) / float64(m.ObservationCount))
	return p

}

/*
 P( observation | class ) - conditional probability of this observation
 given the class.
 Multiply the probabilities of each word being in this class (sum logs)
 P( word | class ) - probability of a word, given a class
 Occurences of the word within the class, divided by total number of words
 Laplace smoothing (always add one so new words don't break everything)
 and also add unique word count to denominator
*/
func (m *Model) classConditionalProbability(class *Class, o *Observation) (p float64) {
	p = 0
	for word, count := range o.WordCounts {
		classWordCount, _ := class.WordCounts[word]
		raw := float64(classWordCount+1) / float64(class.TotalCount+len(m.Vocabulary))
		p = p + (math.Log(raw) * float64(count))
	}
	return p
}

// train the model with the given observation
func (m *Model) Train(o *Observation) {
	for _, className := range o.Classes {
		class, ok := m.Classes[className]
		if !ok {
			class = NewClass(className)
			m.Classes[className] = class
		}
		class.ObservationCount++
		for word, count := range o.WordCounts {
			class.addWord(word, count)
			m.Vocabulary[word] = 1
		}
	}

	m.ObservationCount++
}

// calculate observation membership probability for each class
func (m *Model) Predict(o *Observation) (p Prediction) {
	p = make(map[string]float64)

	for _, class := range m.Classes {
		classConditionalProbability := m.classPriorProbability(class) + m.classConditionalProbability(class, o)
		p[class.Name] = math.Exp(classConditionalProbability)
	}
	return p
}
<|endoftext|>


<commit_before>package nameserver

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"sync"
	"time"

	"github.com/miekg/dns"

	. "github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/common/docker"
	"github.com/weaveworks/weave/net/address"
	"github.com/weaveworks/weave/router"
)

const (
	// Tombstones do not need to survive long periods of peer disconnection, as
	// we delete entries for disconnected peers.  Therefore they just need to hang
	// around to account for propagation delay through gossip.  10 minutes sounds
	// long enough.
	tombstoneTimeout = time.Minute * 30

	// Used by prog/weaver/main.go and proxy/create_container_interceptor.go
	DefaultDomain = "weave.local."

	// Maximum age of acceptable gossip messages (to account for clock skew)
	gossipWindow = int64(tombstoneTimeout/time.Second) / 2
)

// Nameserver: gossip-based, in memory nameserver.
// - Holds a sorted list of (hostname, peer, container id, ip) tuples for the whole cluster.
// - This list is gossiped & merged around the cluser.
// - Lookup-by-hostname are O(nlogn), and return a (copy of a) slice of the entries
// - Update is O(n) for now
type Nameserver struct {
	sync.RWMutex
	ourName router.PeerName
	domain  string
	gossip  router.Gossip
	docker  *docker.Client
	entries Entries
	peers   *router.Peers
	quit    chan struct{}
}

func New(ourName router.PeerName, peers *router.Peers, docker *docker.Client, domain string) *Nameserver {
	ns := &Nameserver{
		ourName: ourName,
		domain:  dns.Fqdn(domain),
		peers:   peers,
		docker:  docker,
		quit:    make(chan struct{}),
	}
	if peers != nil {
		peers.OnGC(ns.PeerGone)
	}
	return ns
}

func (n *Nameserver) SetGossip(gossip router.Gossip) {
	n.gossip = gossip
}

func (n *Nameserver) Start() {
	go func() {
		ticker := time.Tick(tombstoneTimeout)
		for {
			select {
			case <-n.quit:
				return
			case <-ticker:
				n.deleteTombstones()
			}
		}
	}()
}

func (n *Nameserver) Stop() {
	n.quit <- struct{}{}
}

func (n *Nameserver) broadcastEntries(es ...Entry) error {
	if n.gossip != nil {
		return n.gossip.GossipBroadcast(&GossipData{
			Entries:   Entries(es),
			Timestamp: now(),
		})
	}
	return nil
}

func (n *Nameserver) AddEntry(hostname, containerid string, origin router.PeerName, addr address.Address) error {
	n.infof("adding entry %s -> %s", hostname, addr.String())
	n.Lock()
	entry := n.entries.add(hostname, containerid, origin, addr)
	n.Unlock()
	return n.broadcastEntries(entry)
}

func (n *Nameserver) Lookup(hostname string) []address.Address {
	n.RLock()
	defer n.RUnlock()

	entries := n.entries.lookup(hostname)
	result := []address.Address{}
	for _, e := range entries {
		if e.Tombstone > 0 {
			continue
		}
		result = append(result, e.Addr)
	}
	n.debugf("lookup %s -> %s", hostname, &result)
	return result
}

func (n *Nameserver) ReverseLookup(ip address.Address) (string, error) {
	n.RLock()
	defer n.RUnlock()

	match, err := n.entries.first(func(e *Entry) bool {
		return e.Addr == ip
	})
	if err != nil {
		return "", err
	}
	n.debugf("reverse lookup %s -> %s", ip, match.Hostname)
	return match.Hostname, nil
}

func (n *Nameserver) ContainerDied(ident string) {
	n.Lock()
	entries := n.entries.tombstone(n.ourName, func(e *Entry) bool {
		if e.ContainerID == ident {
			n.infof("container %s died, tombstoning entry %s", ident, e.String())
			return true
		}
		return false
	})
	n.Unlock()
	if len(entries) > 0 {
		if err := n.broadcastEntries(entries...); err != nil {
			n.errorf("Failed to broadcast container '%s' death: %v", ident, err)
		}
	}
}

func (n *Nameserver) PeerGone(peer *router.Peer) {
	n.infof("peer gone %s", peer.String())
	n.Lock()
	defer n.Unlock()
	n.entries.filter(func(e *Entry) bool {
		return e.Origin != peer.Name
	})
}

func (n *Nameserver) Delete(hostname, containerid, ipStr string, ip address.Address) error {
	n.Lock()
	n.infof("tombstoning hostname=%s, containerid=%s, ip=%s", hostname, containerid, ipStr)
	entries := n.entries.tombstone(n.ourName, func(e *Entry) bool {
		if hostname != "*" && e.Hostname != hostname {
			return false
		}

		if containerid != "*" && e.ContainerID != containerid {
			return false
		}

		if ipStr != "*" && e.Addr != ip {
			return false
		}

		n.infof("tombstoning entry %v", e)
		return true
	})
	n.Unlock()
	return n.broadcastEntries(entries...)
}

func (n *Nameserver) deleteTombstones() {
	n.Lock()
	defer n.Unlock()
	now := time.Now().Unix()
	n.entries.filter(func(e *Entry) bool {
		return e.Tombstone == 0 || now-e.Tombstone <= int64(tombstoneTimeout/time.Second)
	})
}

func (n *Nameserver) Gossip() router.GossipData {
	n.RLock()
	defer n.RUnlock()
	gossip := &GossipData{
		Entries:   make(Entries, len(n.entries)),
		Timestamp: now(),
	}
	copy(gossip.Entries, n.entries)
	return gossip
}

func (n *Nameserver) OnGossipUnicast(sender router.PeerName, msg []byte) error {
	return nil
}

func (n *Nameserver) receiveGossip(msg []byte) (router.GossipData, router.GossipData, error) {
	var gossip GossipData
	if err := gob.NewDecoder(bytes.NewReader(msg)).Decode(&gossip); err != nil {
		return nil, nil, err
	}

	if delta := gossip.Timestamp - now(); delta > gossipWindow || delta < -gossipWindow {
		return nil, nil, fmt.Errorf("clock skew of %d detected", delta)
	}

	if err := gossip.Entries.check(); err != nil {
		return nil, nil, err
	}

	n.Lock()
	defer n.Unlock()

	if n.peers != nil {
		gossip.Entries.filter(func(e *Entry) bool {
			return n.peers.Fetch(e.Origin) != nil
		})
	}

	newEntries := n.entries.merge(gossip.Entries)
	if len(newEntries) > 0 {
		return &GossipData{Entries: newEntries, Timestamp: now()}, &gossip, nil
	}
	return nil, &gossip, nil
}

// merge received data into state and return "everything new I've
// just learnt", or nil if nothing in the received data was new
func (n *Nameserver) OnGossip(msg []byte) (router.GossipData, error) {
	newEntries, _, err := n.receiveGossip(msg)
	return newEntries, err
}

// merge received data into state and return a representation of
// the received data, for further propagation
func (n *Nameserver) OnGossipBroadcast(_ router.PeerName, msg []byte) (router.GossipData, error) {
	_, entries, err := n.receiveGossip(msg)
	return entries, err
}

func (n *Nameserver) infof(fmt string, args ...interface{}) {
	Log.Infof("[nameserver %s] "+fmt, append([]interface{}{n.ourName}, args...)...)
}
func (n *Nameserver) debugf(fmt string, args ...interface{}) {
	Log.Debugf("[nameserver %s] "+fmt, append([]interface{}{n.ourName}, args...)...)
}
func (n *Nameserver) errorf(fmt string, args ...interface{}) {
	Log.Errorf("[nameserver %s] "+fmt, append([]interface{}{n.ourName}, args...)...)
}
<commit_msg>Update comment.<commit_after>package nameserver

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"sync"
	"time"

	"github.com/miekg/dns"

	. "github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/common/docker"
	"github.com/weaveworks/weave/net/address"
	"github.com/weaveworks/weave/router"
)

const (
	// Tombstones do not need to survive long periods of peer disconnection, as
	// we delete entries for disconnected peers.  Therefore they just need to hang
	// around to account for propagation delay through gossip.  We do need them
	// to hang around longer than the allowed clock skew (defined next).  30mins
	// tombstone timeout allows for 15mins max clock skew.
	tombstoneTimeout = time.Minute * 30

	// Maximum age of acceptable gossip messages (to account for clock skew)
	gossipWindow = int64(tombstoneTimeout/time.Second) / 2

	// Used by prog/weaver/main.go and proxy/create_container_interceptor.go
	DefaultDomain = "weave.local."
)

// Nameserver: gossip-based, in memory nameserver.
// - Holds a sorted list of (hostname, peer, container id, ip) tuples for the whole cluster.
// - This list is gossiped & merged around the cluser.
// - Lookup-by-hostname are O(nlogn), and return a (copy of a) slice of the entries
// - Update is O(n) for now
type Nameserver struct {
	sync.RWMutex
	ourName router.PeerName
	domain  string
	gossip  router.Gossip
	docker  *docker.Client
	entries Entries
	peers   *router.Peers
	quit    chan struct{}
}

func New(ourName router.PeerName, peers *router.Peers, docker *docker.Client, domain string) *Nameserver {
	ns := &Nameserver{
		ourName: ourName,
		domain:  dns.Fqdn(domain),
		peers:   peers,
		docker:  docker,
		quit:    make(chan struct{}),
	}
	if peers != nil {
		peers.OnGC(ns.PeerGone)
	}
	return ns
}

func (n *Nameserver) SetGossip(gossip router.Gossip) {
	n.gossip = gossip
}

func (n *Nameserver) Start() {
	go func() {
		ticker := time.Tick(tombstoneTimeout)
		for {
			select {
			case <-n.quit:
				return
			case <-ticker:
				n.deleteTombstones()
			}
		}
	}()
}

func (n *Nameserver) Stop() {
	n.quit <- struct{}{}
}

func (n *Nameserver) broadcastEntries(es ...Entry) error {
	if n.gossip != nil {
		return n.gossip.GossipBroadcast(&GossipData{
			Entries:   Entries(es),
			Timestamp: now(),
		})
	}
	return nil
}

func (n *Nameserver) AddEntry(hostname, containerid string, origin router.PeerName, addr address.Address) error {
	n.infof("adding entry %s -> %s", hostname, addr.String())
	n.Lock()
	entry := n.entries.add(hostname, containerid, origin, addr)
	n.Unlock()
	return n.broadcastEntries(entry)
}

func (n *Nameserver) Lookup(hostname string) []address.Address {
	n.RLock()
	defer n.RUnlock()

	entries := n.entries.lookup(hostname)
	result := []address.Address{}
	for _, e := range entries {
		if e.Tombstone > 0 {
			continue
		}
		result = append(result, e.Addr)
	}
	n.debugf("lookup %s -> %s", hostname, &result)
	return result
}

func (n *Nameserver) ReverseLookup(ip address.Address) (string, error) {
	n.RLock()
	defer n.RUnlock()

	match, err := n.entries.first(func(e *Entry) bool {
		return e.Addr == ip
	})
	if err != nil {
		return "", err
	}
	n.debugf("reverse lookup %s -> %s", ip, match.Hostname)
	return match.Hostname, nil
}

func (n *Nameserver) ContainerDied(ident string) {
	n.Lock()
	entries := n.entries.tombstone(n.ourName, func(e *Entry) bool {
		if e.ContainerID == ident {
			n.infof("container %s died, tombstoning entry %s", ident, e.String())
			return true
		}
		return false
	})
	n.Unlock()
	if len(entries) > 0 {
		if err := n.broadcastEntries(entries...); err != nil {
			n.errorf("Failed to broadcast container '%s' death: %v", ident, err)
		}
	}
}

func (n *Nameserver) PeerGone(peer *router.Peer) {
	n.infof("peer gone %s", peer.String())
	n.Lock()
	defer n.Unlock()
	n.entries.filter(func(e *Entry) bool {
		return e.Origin != peer.Name
	})
}

func (n *Nameserver) Delete(hostname, containerid, ipStr string, ip address.Address) error {
	n.Lock()
	n.infof("tombstoning hostname=%s, containerid=%s, ip=%s", hostname, containerid, ipStr)
	entries := n.entries.tombstone(n.ourName, func(e *Entry) bool {
		if hostname != "*" && e.Hostname != hostname {
			return false
		}

		if containerid != "*" && e.ContainerID != containerid {
			return false
		}

		if ipStr != "*" && e.Addr != ip {
			return false
		}

		n.infof("tombstoning entry %v", e)
		return true
	})
	n.Unlock()
	return n.broadcastEntries(entries...)
}

func (n *Nameserver) deleteTombstones() {
	n.Lock()
	defer n.Unlock()
	now := time.Now().Unix()
	n.entries.filter(func(e *Entry) bool {
		return e.Tombstone == 0 || now-e.Tombstone <= int64(tombstoneTimeout/time.Second)
	})
}

func (n *Nameserver) Gossip() router.GossipData {
	n.RLock()
	defer n.RUnlock()
	gossip := &GossipData{
		Entries:   make(Entries, len(n.entries)),
		Timestamp: now(),
	}
	copy(gossip.Entries, n.entries)
	return gossip
}

func (n *Nameserver) OnGossipUnicast(sender router.PeerName, msg []byte) error {
	return nil
}

func (n *Nameserver) receiveGossip(msg []byte) (router.GossipData, router.GossipData, error) {
	var gossip GossipData
	if err := gob.NewDecoder(bytes.NewReader(msg)).Decode(&gossip); err != nil {
		return nil, nil, err
	}

	if delta := gossip.Timestamp - now(); delta > gossipWindow || delta < -gossipWindow {
		return nil, nil, fmt.Errorf("clock skew of %d detected", delta)
	}

	if err := gossip.Entries.check(); err != nil {
		return nil, nil, err
	}

	n.Lock()
	defer n.Unlock()

	if n.peers != nil {
		gossip.Entries.filter(func(e *Entry) bool {
			return n.peers.Fetch(e.Origin) != nil
		})
	}

	newEntries := n.entries.merge(gossip.Entries)
	if len(newEntries) > 0 {
		return &GossipData{Entries: newEntries, Timestamp: now()}, &gossip, nil
	}
	return nil, &gossip, nil
}

// merge received data into state and return "everything new I've
// just learnt", or nil if nothing in the received data was new
func (n *Nameserver) OnGossip(msg []byte) (router.GossipData, error) {
	newEntries, _, err := n.receiveGossip(msg)
	return newEntries, err
}

// merge received data into state and return a representation of
// the received data, for further propagation
func (n *Nameserver) OnGossipBroadcast(_ router.PeerName, msg []byte) (router.GossipData, error) {
	_, entries, err := n.receiveGossip(msg)
	return entries, err
}

func (n *Nameserver) infof(fmt string, args ...interface{}) {
	Log.Infof("[nameserver %s] "+fmt, append([]interface{}{n.ourName}, args...)...)
}
func (n *Nameserver) debugf(fmt string, args ...interface{}) {
	Log.Debugf("[nameserver %s] "+fmt, append([]interface{}{n.ourName}, args...)...)
}
func (n *Nameserver) errorf(fmt string, args ...interface{}) {
	Log.Errorf("[nameserver %s] "+fmt, append([]interface{}{n.ourName}, args...)...)
}
<|endoftext|>


<commit_before>package netlink

import (
	"net"
	"syscall"
	"unsafe"

	"github.com/vishvananda/netlink/nl"
	"github.com/vishvananda/netns"
	"golang.org/x/sys/unix"
)

const (
	NDA_UNSPEC = iota
	NDA_DST
	NDA_LLADDR
	NDA_CACHEINFO
	NDA_PROBES
	NDA_VLAN
	NDA_PORT
	NDA_VNI
	NDA_IFINDEX
	NDA_MAX = NDA_IFINDEX
)

// Neighbor Cache Entry States.
const (
	NUD_NONE       = 0x00
	NUD_INCOMPLETE = 0x01
	NUD_REACHABLE  = 0x02
	NUD_STALE      = 0x04
	NUD_DELAY      = 0x08
	NUD_PROBE      = 0x10
	NUD_FAILED     = 0x20
	NUD_NOARP      = 0x40
	NUD_PERMANENT  = 0x80
)

// Neighbor Flags
const (
	NTF_USE    = 0x01
	NTF_SELF   = 0x02
	NTF_MASTER = 0x04
	NTF_PROXY  = 0x08
	NTF_ROUTER = 0x80
)

type Ndmsg struct {
	Family uint8
	Index  uint32
	State  uint16
	Flags  uint8
	Type   uint8
}

func deserializeNdmsg(b []byte) *Ndmsg {
	var dummy Ndmsg
	return (*Ndmsg)(unsafe.Pointer(&b[0:unsafe.Sizeof(dummy)][0]))
}

func (msg *Ndmsg) Serialize() []byte {
	return (*(*[unsafe.Sizeof(*msg)]byte)(unsafe.Pointer(msg)))[:]
}

func (msg *Ndmsg) Len() int {
	return int(unsafe.Sizeof(*msg))
}

// NeighAdd will add an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh add ....`
func NeighAdd(neigh *Neigh) error {
	return pkgHandle.NeighAdd(neigh)
}

// NeighAdd will add an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh add ....`
func (h *Handle) NeighAdd(neigh *Neigh) error {
	return h.neighAdd(neigh, unix.NLM_F_CREATE|unix.NLM_F_EXCL)
}

// NeighSet will add or replace an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh replace....`
func NeighSet(neigh *Neigh) error {
	return pkgHandle.NeighSet(neigh)
}

// NeighSet will add or replace an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh replace....`
func (h *Handle) NeighSet(neigh *Neigh) error {
	return h.neighAdd(neigh, unix.NLM_F_CREATE|unix.NLM_F_REPLACE)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func NeighAppend(neigh *Neigh) error {
	return pkgHandle.NeighAppend(neigh)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func (h *Handle) NeighAppend(neigh *Neigh) error {
	return h.neighAdd(neigh, unix.NLM_F_CREATE|unix.NLM_F_APPEND)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func neighAdd(neigh *Neigh, mode int) error {
	return pkgHandle.neighAdd(neigh, mode)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func (h *Handle) neighAdd(neigh *Neigh, mode int) error {
	req := h.newNetlinkRequest(unix.RTM_NEWNEIGH, mode|unix.NLM_F_ACK)
	return neighHandle(neigh, req)
}

// NeighDel will delete an IP address from a link device.
// Equivalent to: `ip addr del $addr dev $link`
func NeighDel(neigh *Neigh) error {
	return pkgHandle.NeighDel(neigh)
}

// NeighDel will delete an IP address from a link device.
// Equivalent to: `ip addr del $addr dev $link`
func (h *Handle) NeighDel(neigh *Neigh) error {
	req := h.newNetlinkRequest(unix.RTM_DELNEIGH, unix.NLM_F_ACK)
	return neighHandle(neigh, req)
}

func neighHandle(neigh *Neigh, req *nl.NetlinkRequest) error {
	var family int

	if neigh.Family > 0 {
		family = neigh.Family
	} else {
		family = nl.GetIPFamily(neigh.IP)
	}

	msg := Ndmsg{
		Family: uint8(family),
		Index:  uint32(neigh.LinkIndex),
		State:  uint16(neigh.State),
		Type:   uint8(neigh.Type),
		Flags:  uint8(neigh.Flags),
	}
	req.AddData(&msg)

	ipData := neigh.IP.To4()
	if ipData == nil {
		ipData = neigh.IP.To16()
	}

	dstData := nl.NewRtAttr(NDA_DST, ipData)
	req.AddData(dstData)

	if neigh.LLIPAddr != nil {
		llIPData := nl.NewRtAttr(NDA_LLADDR, neigh.LLIPAddr.To4())
		req.AddData(llIPData)
	} else if neigh.Flags != NTF_PROXY || neigh.HardwareAddr != nil {
		hwData := nl.NewRtAttr(NDA_LLADDR, []byte(neigh.HardwareAddr))
		req.AddData(hwData)
	}

	if neigh.Vlan != 0 {
		vlanData := nl.NewRtAttr(NDA_VLAN, nl.Uint16Attr(uint16(neigh.Vlan)))
		req.AddData(vlanData)
	}

	if neigh.VNI != 0 {
		vniData := nl.NewRtAttr(NDA_VNI, nl.Uint32Attr(uint32(neigh.VNI)))
		req.AddData(vniData)
	}

	_, err := req.Execute(unix.NETLINK_ROUTE, 0)
	return err
}

// NeighList gets a list of IP-MAC mappings in the system (ARP table).
// Equivalent to: `ip neighbor show`.
// The list can be filtered by link and ip family.
func NeighList(linkIndex, family int) ([]Neigh, error) {
	return pkgHandle.NeighList(linkIndex, family)
}

// NeighProxyList gets a list of neighbor proxies in the system.
// Equivalent to: `ip neighbor show proxy`.
// The list can be filtered by link and ip family.
func NeighProxyList(linkIndex, family int) ([]Neigh, error) {
	return pkgHandle.NeighProxyList(linkIndex, family)
}

// NeighList gets a list of IP-MAC mappings in the system (ARP table).
// Equivalent to: `ip neighbor show`.
// The list can be filtered by link and ip family.
func (h *Handle) NeighList(linkIndex, family int) ([]Neigh, error) {
	return h.neighList(linkIndex, family, 0)
}

// NeighProxyList gets a list of neighbor proxies in the system.
// Equivalent to: `ip neighbor show proxy`.
// The list can be filtered by link, ip family.
func (h *Handle) NeighProxyList(linkIndex, family int) ([]Neigh, error) {
	return h.neighList(linkIndex, family, NTF_PROXY)
}

func (h *Handle) neighList(linkIndex, family, flags int) ([]Neigh, error) {
	req := h.newNetlinkRequest(unix.RTM_GETNEIGH, unix.NLM_F_DUMP)
	msg := Ndmsg{
		Family: uint8(family),
		Index:  uint32(linkIndex),
		Flags:  uint8(flags),
	}
	req.AddData(&msg)

	msgs, err := req.Execute(unix.NETLINK_ROUTE, unix.RTM_NEWNEIGH)
	if err != nil {
		return nil, err
	}

	var res []Neigh
	for _, m := range msgs {
		ndm := deserializeNdmsg(m)
		if linkIndex != 0 && int(ndm.Index) != linkIndex {
			// Ignore messages from other interfaces
			continue
		}

		neigh, err := NeighDeserialize(m)
		if err != nil {
			continue
		}

		res = append(res, *neigh)
	}

	return res, nil
}

func NeighDeserialize(m []byte) (*Neigh, error) {
	msg := deserializeNdmsg(m)

	neigh := Neigh{
		LinkIndex: int(msg.Index),
		Family:    int(msg.Family),
		State:     int(msg.State),
		Type:      int(msg.Type),
		Flags:     int(msg.Flags),
	}

	attrs, err := nl.ParseRouteAttr(m[msg.Len():])
	if err != nil {
		return nil, err
	}

	// This should be cached for perfomance
	// once per table dump
	link, err := LinkByIndex(neigh.LinkIndex)
	if err != nil {
		return nil, err
	}
	encapType := link.Attrs().EncapType

	for _, attr := range attrs {
		switch attr.Attr.Type {
		case NDA_DST:
			neigh.IP = net.IP(attr.Value)
		case NDA_LLADDR:
			// BUG: Is this a bug in the netlink library?
			// #define RTA_LENGTH(len) (RTA_ALIGN(sizeof(struct rtattr)) + (len))
			// #define RTA_PAYLOAD(rta) ((int)((rta)->rta_len) - RTA_LENGTH(0))
			attrLen := attr.Attr.Len - unix.SizeofRtAttr
			if attrLen == 4 && (encapType == "ipip" ||
				encapType == "sit" ||
				encapType == "gre") {
				neigh.LLIPAddr = net.IP(attr.Value)
			} else if attrLen == 16 &&
				encapType == "tunnel6" {
				neigh.IP = net.IP(attr.Value)
			} else {
				neigh.HardwareAddr = net.HardwareAddr(attr.Value)
			}
		case NDA_VLAN:
			neigh.Vlan = int(native.Uint16(attr.Value[0:2]))
		case NDA_VNI:
			neigh.VNI = int(native.Uint32(attr.Value[0:4]))
		}
	}

	return &neigh, nil
}

// NeighSubscribe takes a chan down which notifications will be sent
// when neighbors are added or deleted. Close the 'done' chan to stop subscription.
func NeighSubscribe(ch chan<- NeighUpdate, done <-chan struct{}) error {
	return neighSubscribeAt(netns.None(), netns.None(), ch, done, nil, false)
}

// NeighSubscribeAt works like NeighSubscribe plus it allows the caller
// to choose the network namespace in which to subscribe (ns).
func NeighSubscribeAt(ns netns.NsHandle, ch chan<- NeighUpdate, done <-chan struct{}) error {
	return neighSubscribeAt(ns, netns.None(), ch, done, nil, false)
}

// NeighSubscribeOptions contains a set of options to use with
// NeighSubscribeWithOptions.
type NeighSubscribeOptions struct {
	Namespace     *netns.NsHandle
	ErrorCallback func(error)
	ListExisting  bool
}

// NeighSubscribeWithOptions work like NeighSubscribe but enable to
// provide additional options to modify the behavior. Currently, the
// namespace can be provided as well as an error callback.
func NeighSubscribeWithOptions(ch chan<- NeighUpdate, done <-chan struct{}, options NeighSubscribeOptions) error {
	if options.Namespace == nil {
		none := netns.None()
		options.Namespace = &none
	}
	return neighSubscribeAt(*options.Namespace, netns.None(), ch, done, options.ErrorCallback, options.ListExisting)
}

func neighSubscribeAt(newNs, curNs netns.NsHandle, ch chan<- NeighUpdate, done <-chan struct{}, cberr func(error), listExisting bool) error {
	s, err := nl.SubscribeAt(newNs, curNs, unix.NETLINK_ROUTE, unix.RTNLGRP_NEIGH)
	if err != nil {
		return err
	}
	if done != nil {
		go func() {
			<-done
			s.Close()
		}()
	}
	if listExisting {
		req := pkgHandle.newNetlinkRequest(unix.RTM_GETNEIGH,
			unix.NLM_F_DUMP)
		infmsg := nl.NewIfInfomsg(unix.AF_UNSPEC)
		req.AddData(infmsg)
		if err := s.Send(req); err != nil {
			return err
		}
	}
	go func() {
		defer close(ch)
		for {
			msgs, err := s.Receive()
			if err != nil {
				if cberr != nil {
					cberr(err)
				}
				return
			}
			for _, m := range msgs {
				if m.Header.Type == unix.NLMSG_DONE {
					continue
				}
				if m.Header.Type == unix.NLMSG_ERROR {
					native := nl.NativeEndian()
					error := int32(native.Uint32(m.Data[0:4]))
					if error == 0 {
						continue
					}
					if cberr != nil {
						cberr(syscall.Errno(-error))
					}
					return
				}
				neigh, err := NeighDeserialize(m.Data)
				if err != nil {
					if cberr != nil {
						cberr(err)
					}
					return
				}
				ch <- NeighUpdate{Type: m.Header.Type, Neigh: *neigh}
			}
		}
	}()

	return nil
}
<commit_msg>Pass Ndmsg to NeighListExecute<commit_after>package netlink

import (
	"net"
	"syscall"
	"unsafe"

	"github.com/vishvananda/netlink/nl"
	"github.com/vishvananda/netns"
	"golang.org/x/sys/unix"
)

const (
	NDA_UNSPEC = iota
	NDA_DST
	NDA_LLADDR
	NDA_CACHEINFO
	NDA_PROBES
	NDA_VLAN
	NDA_PORT
	NDA_VNI
	NDA_IFINDEX
	NDA_MAX = NDA_IFINDEX
)

// Neighbor Cache Entry States.
const (
	NUD_NONE       = 0x00
	NUD_INCOMPLETE = 0x01
	NUD_REACHABLE  = 0x02
	NUD_STALE      = 0x04
	NUD_DELAY      = 0x08
	NUD_PROBE      = 0x10
	NUD_FAILED     = 0x20
	NUD_NOARP      = 0x40
	NUD_PERMANENT  = 0x80
)

// Neighbor Flags
const (
	NTF_USE    = 0x01
	NTF_SELF   = 0x02
	NTF_MASTER = 0x04
	NTF_PROXY  = 0x08
	NTF_ROUTER = 0x80
)

// Ndmsg is for adding, removing or receiving information about a neighbor table entry
type Ndmsg struct {
	Family uint8
	Index  uint32
	State  uint16
	Flags  uint8
	Type   uint8
}

func deserializeNdmsg(b []byte) *Ndmsg {
	var dummy Ndmsg
	return (*Ndmsg)(unsafe.Pointer(&b[0:unsafe.Sizeof(dummy)][0]))
}

func (msg *Ndmsg) Serialize() []byte {
	return (*(*[unsafe.Sizeof(*msg)]byte)(unsafe.Pointer(msg)))[:]
}

func (msg *Ndmsg) Len() int {
	return int(unsafe.Sizeof(*msg))
}

// NeighAdd will add an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh add ....`
func NeighAdd(neigh *Neigh) error {
	return pkgHandle.NeighAdd(neigh)
}

// NeighAdd will add an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh add ....`
func (h *Handle) NeighAdd(neigh *Neigh) error {
	return h.neighAdd(neigh, unix.NLM_F_CREATE|unix.NLM_F_EXCL)
}

// NeighSet will add or replace an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh replace....`
func NeighSet(neigh *Neigh) error {
	return pkgHandle.NeighSet(neigh)
}

// NeighSet will add or replace an IP to MAC mapping to the ARP table
// Equivalent to: `ip neigh replace....`
func (h *Handle) NeighSet(neigh *Neigh) error {
	return h.neighAdd(neigh, unix.NLM_F_CREATE|unix.NLM_F_REPLACE)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func NeighAppend(neigh *Neigh) error {
	return pkgHandle.NeighAppend(neigh)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func (h *Handle) NeighAppend(neigh *Neigh) error {
	return h.neighAdd(neigh, unix.NLM_F_CREATE|unix.NLM_F_APPEND)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func neighAdd(neigh *Neigh, mode int) error {
	return pkgHandle.neighAdd(neigh, mode)
}

// NeighAppend will append an entry to FDB
// Equivalent to: `bridge fdb append...`
func (h *Handle) neighAdd(neigh *Neigh, mode int) error {
	req := h.newNetlinkRequest(unix.RTM_NEWNEIGH, mode|unix.NLM_F_ACK)
	return neighHandle(neigh, req)
}

// NeighDel will delete an IP address from a link device.
// Equivalent to: `ip addr del $addr dev $link`
func NeighDel(neigh *Neigh) error {
	return pkgHandle.NeighDel(neigh)
}

// NeighDel will delete an IP address from a link device.
// Equivalent to: `ip addr del $addr dev $link`
func (h *Handle) NeighDel(neigh *Neigh) error {
	req := h.newNetlinkRequest(unix.RTM_DELNEIGH, unix.NLM_F_ACK)
	return neighHandle(neigh, req)
}

func neighHandle(neigh *Neigh, req *nl.NetlinkRequest) error {
	var family int

	if neigh.Family > 0 {
		family = neigh.Family
	} else {
		family = nl.GetIPFamily(neigh.IP)
	}

	msg := Ndmsg{
		Family: uint8(family),
		Index:  uint32(neigh.LinkIndex),
		State:  uint16(neigh.State),
		Type:   uint8(neigh.Type),
		Flags:  uint8(neigh.Flags),
	}
	req.AddData(&msg)

	ipData := neigh.IP.To4()
	if ipData == nil {
		ipData = neigh.IP.To16()
	}

	dstData := nl.NewRtAttr(NDA_DST, ipData)
	req.AddData(dstData)

	if neigh.LLIPAddr != nil {
		llIPData := nl.NewRtAttr(NDA_LLADDR, neigh.LLIPAddr.To4())
		req.AddData(llIPData)
	} else if neigh.Flags != NTF_PROXY || neigh.HardwareAddr != nil {
		hwData := nl.NewRtAttr(NDA_LLADDR, []byte(neigh.HardwareAddr))
		req.AddData(hwData)
	}

	if neigh.Vlan != 0 {
		vlanData := nl.NewRtAttr(NDA_VLAN, nl.Uint16Attr(uint16(neigh.Vlan)))
		req.AddData(vlanData)
	}

	if neigh.VNI != 0 {
		vniData := nl.NewRtAttr(NDA_VNI, nl.Uint32Attr(uint32(neigh.VNI)))
		req.AddData(vniData)
	}

	_, err := req.Execute(unix.NETLINK_ROUTE, 0)
	return err
}

// NeighList returns a list of IP-MAC mappings in the system (ARP table).
// Equivalent to: `ip neighbor show`.
// The list can be filtered by link and ip family.
func NeighList(linkIndex, family int) ([]Neigh, error) {
	return pkgHandle.NeighList(linkIndex, family)
}

// NeighProxyList returns a list of neighbor proxies in the system.
// Equivalent to: `ip neighbor show proxy`.
// The list can be filtered by link and ip family.
func NeighProxyList(linkIndex, family int) ([]Neigh, error) {
	return pkgHandle.NeighProxyList(linkIndex, family)
}

// NeighList returns a list of IP-MAC mappings in the system (ARP table).
// Equivalent to: `ip neighbor show`.
// The list can be filtered by link and ip family.
func (h *Handle) NeighList(linkIndex, family int) ([]Neigh, error) {
	return h.NeighListExecute(Ndmsg{
		Family: uint8(family),
		Index:  uint32(linkIndex),
	})
}

// NeighProxyList returns a list of neighbor proxies in the system.
// Equivalent to: `ip neighbor show proxy`.
// The list can be filtered by link, ip family.
func (h *Handle) NeighProxyList(linkIndex, family int) ([]Neigh, error) {
	return h.NeighListExecute(Ndmsg{
		Family: uint8(family),
		Index:  uint32(linkIndex),
		Flags:  NTF_PROXY,
	})
}

// NeighListExecute returns a list of neighbour entries filtered by link, ip family, flag and state.
func NeighListExecute(msg Ndmsg) ([]Neigh, error) {
	return pkgHandle.NeighListExecute(msg)
}

// NeighListExecute returns a list of neighbour entries filtered by link, ip family, flag and state.
func (h *Handle) NeighListExecute(msg Ndmsg) ([]Neigh, error) {
	req := h.newNetlinkRequest(unix.RTM_GETNEIGH, unix.NLM_F_DUMP)
	req.AddData(&msg)

	msgs, err := req.Execute(unix.NETLINK_ROUTE, unix.RTM_NEWNEIGH)
	if err != nil {
		return nil, err
	}

	var res []Neigh
	for _, m := range msgs {
		ndm := deserializeNdmsg(m)
		if msg.Index != 0 && ndm.Index != msg.Index {
			// Ignore messages from other interfaces
			continue
		}

		neigh, err := NeighDeserialize(m)
		if err != nil {
			continue
		}

		res = append(res, *neigh)
	}

	return res, nil
}

func NeighDeserialize(m []byte) (*Neigh, error) {
	msg := deserializeNdmsg(m)

	neigh := Neigh{
		LinkIndex: int(msg.Index),
		Family:    int(msg.Family),
		State:     int(msg.State),
		Type:      int(msg.Type),
		Flags:     int(msg.Flags),
	}

	attrs, err := nl.ParseRouteAttr(m[msg.Len():])
	if err != nil {
		return nil, err
	}

	// This should be cached for perfomance
	// once per table dump
	link, err := LinkByIndex(neigh.LinkIndex)
	if err != nil {
		return nil, err
	}
	encapType := link.Attrs().EncapType

	for _, attr := range attrs {
		switch attr.Attr.Type {
		case NDA_DST:
			neigh.IP = net.IP(attr.Value)
		case NDA_LLADDR:
			// BUG: Is this a bug in the netlink library?
			// #define RTA_LENGTH(len) (RTA_ALIGN(sizeof(struct rtattr)) + (len))
			// #define RTA_PAYLOAD(rta) ((int)((rta)->rta_len) - RTA_LENGTH(0))
			attrLen := attr.Attr.Len - unix.SizeofRtAttr
			if attrLen == 4 && (encapType == "ipip" ||
				encapType == "sit" ||
				encapType == "gre") {
				neigh.LLIPAddr = net.IP(attr.Value)
			} else if attrLen == 16 &&
				encapType == "tunnel6" {
				neigh.IP = net.IP(attr.Value)
			} else {
				neigh.HardwareAddr = net.HardwareAddr(attr.Value)
			}
		case NDA_VLAN:
			neigh.Vlan = int(native.Uint16(attr.Value[0:2]))
		case NDA_VNI:
			neigh.VNI = int(native.Uint32(attr.Value[0:4]))
		}
	}

	return &neigh, nil
}

// NeighSubscribe takes a chan down which notifications will be sent
// when neighbors are added or deleted. Close the 'done' chan to stop subscription.
func NeighSubscribe(ch chan<- NeighUpdate, done <-chan struct{}) error {
	return neighSubscribeAt(netns.None(), netns.None(), ch, done, nil, false)
}

// NeighSubscribeAt works like NeighSubscribe plus it allows the caller
// to choose the network namespace in which to subscribe (ns).
func NeighSubscribeAt(ns netns.NsHandle, ch chan<- NeighUpdate, done <-chan struct{}) error {
	return neighSubscribeAt(ns, netns.None(), ch, done, nil, false)
}

// NeighSubscribeOptions contains a set of options to use with
// NeighSubscribeWithOptions.
type NeighSubscribeOptions struct {
	Namespace     *netns.NsHandle
	ErrorCallback func(error)
	ListExisting  bool
}

// NeighSubscribeWithOptions work like NeighSubscribe but enable to
// provide additional options to modify the behavior. Currently, the
// namespace can be provided as well as an error callback.
func NeighSubscribeWithOptions(ch chan<- NeighUpdate, done <-chan struct{}, options NeighSubscribeOptions) error {
	if options.Namespace == nil {
		none := netns.None()
		options.Namespace = &none
	}
	return neighSubscribeAt(*options.Namespace, netns.None(), ch, done, options.ErrorCallback, options.ListExisting)
}

func neighSubscribeAt(newNs, curNs netns.NsHandle, ch chan<- NeighUpdate, done <-chan struct{}, cberr func(error), listExisting bool) error {
	s, err := nl.SubscribeAt(newNs, curNs, unix.NETLINK_ROUTE, unix.RTNLGRP_NEIGH)
	if err != nil {
		return err
	}
	if done != nil {
		go func() {
			<-done
			s.Close()
		}()
	}
	if listExisting {
		req := pkgHandle.newNetlinkRequest(unix.RTM_GETNEIGH,
			unix.NLM_F_DUMP)
		infmsg := nl.NewIfInfomsg(unix.AF_UNSPEC)
		req.AddData(infmsg)
		if err := s.Send(req); err != nil {
			return err
		}
	}
	go func() {
		defer close(ch)
		for {
			msgs, err := s.Receive()
			if err != nil {
				if cberr != nil {
					cberr(err)
				}
				return
			}
			for _, m := range msgs {
				if m.Header.Type == unix.NLMSG_DONE {
					continue
				}
				if m.Header.Type == unix.NLMSG_ERROR {
					native := nl.NativeEndian()
					error := int32(native.Uint32(m.Data[0:4]))
					if error == 0 {
						continue
					}
					if cberr != nil {
						cberr(syscall.Errno(-error))
					}
					return
				}
				neigh, err := NeighDeserialize(m.Data)
				if err != nil {
					if cberr != nil {
						cberr(err)
					}
					return
				}
				ch <- NeighUpdate{Type: m.Header.Type, Neigh: *neigh}
			}
		}
	}()

	return nil
}
<|endoftext|>


<commit_before>// +build !linux

package btrfs

import "github.com/pachyderm/pachyderm/src/pfs/drive"

// NewDriver constructs a new Driver for btrfs.
func NewDriver(rootDir string) drive.Driver {
	return nil
}
<commit_msg>fix btrfs.NewDriver for builds other than linux<commit_after>// +build !linux

package btrfs

import "github.com/pachyderm/pachyderm/src/pfs/drive"

// NewDriver constructs a new Driver for btrfs.
func NewDriver(rootDir string, namespace string) (drive.Driver, error) {
	return nil
}
<|endoftext|>


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

package zip

import (
	"bytes"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
)

// TODO(adg): a more sophisticated test suite

type WriteTest struct {
	Name   string
	Data   []byte
	Method uint16
	Mode   os.FileMode
}

var writeTests = []WriteTest{
	{
		Name:   "foo",
		Data:   []byte("Rabbits, guinea pigs, gophers, marsupial rats, and quolls."),
		Method: Store,
		Mode:   0666,
	},
	{
		Name:   "bar",
		Data:   nil, // large data set in the test
		Method: Deflate,
		Mode:   0644,
	},
	{
		Name:   "setuid",
		Data:   []byte("setuid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetuid,
	},
	{
		Name:   "setgid",
		Data:   []byte("setgid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetgid,
	},
	{
		Name:   "symlink",
		Data:   []byte("../link/target"),
		Method: Deflate,
		Mode:   0755 | os.ModeSymlink,
	},
}

func TestWriter(t *testing.T) {
	largeData := make([]byte, 1<<17)
	for i := range largeData {
		largeData[i] = byte(rand.Int())
	}
	writeTests[1].Data = largeData
	defer func() {
		writeTests[1].Data = nil
	}()

	// write a zip file
	buf := new(bytes.Buffer)
	w := NewWriter(buf)

	for _, wt := range writeTests {
		testCreate(t, w, &wt)
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// read it back
	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		t.Fatal(err)
	}
	for i, wt := range writeTests {
		testReadFile(t, r.File[i], &wt)
	}
}

func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
	header := &FileHeader{
		Name:   wt.Name,
		Method: wt.Method,
	}
	if wt.Mode != 0 {
		header.SetMode(wt.Mode)
	}
	f, err := w.CreateHeader(header)
	if err != nil {
		t.Fatal(err)
	}
	_, err = f.Write(wt.Data)
	if err != nil {
		t.Fatal(err)
	}
}

func testReadFile(t *testing.T, f *File, wt *WriteTest) {
	if f.Name != wt.Name {
		t.Fatalf("File name: got %q, want %q", f.Name, wt.Name)
	}
	testFileMode(t, wt.Name, f, wt.Mode)
	rc, err := f.Open()
	if err != nil {
		t.Fatal("opening:", err)
	}
	b, err := ioutil.ReadAll(rc)
	if err != nil {
		t.Fatal("reading:", err)
	}
	err = rc.Close()
	if err != nil {
		t.Fatal("closing:", err)
	}
	if !bytes.Equal(b, wt.Data) {
		t.Errorf("File contents %q, want %q", b, wt.Data)
	}
}
<commit_msg>archive/zip: add flate writing benchmark<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zip

import (
	"bytes"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
)

// TODO(adg): a more sophisticated test suite

type WriteTest struct {
	Name   string
	Data   []byte
	Method uint16
	Mode   os.FileMode
}

var writeTests = []WriteTest{
	{
		Name:   "foo",
		Data:   []byte("Rabbits, guinea pigs, gophers, marsupial rats, and quolls."),
		Method: Store,
		Mode:   0666,
	},
	{
		Name:   "bar",
		Data:   nil, // large data set in the test
		Method: Deflate,
		Mode:   0644,
	},
	{
		Name:   "setuid",
		Data:   []byte("setuid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetuid,
	},
	{
		Name:   "setgid",
		Data:   []byte("setgid file"),
		Method: Deflate,
		Mode:   0755 | os.ModeSetgid,
	},
	{
		Name:   "symlink",
		Data:   []byte("../link/target"),
		Method: Deflate,
		Mode:   0755 | os.ModeSymlink,
	},
}

func TestWriter(t *testing.T) {
	largeData := make([]byte, 1<<17)
	for i := range largeData {
		largeData[i] = byte(rand.Int())
	}
	writeTests[1].Data = largeData
	defer func() {
		writeTests[1].Data = nil
	}()

	// write a zip file
	buf := new(bytes.Buffer)
	w := NewWriter(buf)

	for _, wt := range writeTests {
		testCreate(t, w, &wt)
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// read it back
	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		t.Fatal(err)
	}
	for i, wt := range writeTests {
		testReadFile(t, r.File[i], &wt)
	}
}

func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
	header := &FileHeader{
		Name:   wt.Name,
		Method: wt.Method,
	}
	if wt.Mode != 0 {
		header.SetMode(wt.Mode)
	}
	f, err := w.CreateHeader(header)
	if err != nil {
		t.Fatal(err)
	}
	_, err = f.Write(wt.Data)
	if err != nil {
		t.Fatal(err)
	}
}

func testReadFile(t *testing.T, f *File, wt *WriteTest) {
	if f.Name != wt.Name {
		t.Fatalf("File name: got %q, want %q", f.Name, wt.Name)
	}
	testFileMode(t, wt.Name, f, wt.Mode)
	rc, err := f.Open()
	if err != nil {
		t.Fatal("opening:", err)
	}
	b, err := ioutil.ReadAll(rc)
	if err != nil {
		t.Fatal("reading:", err)
	}
	err = rc.Close()
	if err != nil {
		t.Fatal("closing:", err)
	}
	if !bytes.Equal(b, wt.Data) {
		t.Errorf("File contents %q, want %q", b, wt.Data)
	}
}

func BenchmarkCompressedZipGarbage(b *testing.B) {
	b.ReportAllocs()
	var buf bytes.Buffer
	bigBuf := bytes.Repeat([]byte("a"), 1<<20)
	for i := 0; i < b.N; i++ {
		buf.Reset()
		zw := NewWriter(&buf)
		for j := 0; j < 3; j++ {
			w, _ := zw.CreateHeader(&FileHeader{
				Name:   "foo",
				Method: Deflate,
			})
			w.Write(bigBuf)
		}
		zw.Close()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/dominikh/simple-router/conntrack"
	"github.com/dominikh/simple-router/lookup"

	"flag"
	"fmt"
	"os"
	"text/tabwriter"
)

// TODO implement the following flags
//       -p <protocol>        : display connections by protocol
//       -s <source-host>     : display connections by source
//       -d <destination-host>: display connections by destination
//       -x: extended hostnames view
//       -r src | dst | src-port | dst-port | state : sort connections
//       -N: display NAT box connection information (only valid with SNAT & DNAT)
//       -v: print version

var onlySNAT = flag.Bool("S", false, "Display only SNAT connections")
var onlyDNAT = flag.Bool("D", false, "Display only DNAT connections")
var onlyLocal = flag.Bool("L", false, "Display only local connections (originating from or going to the router)")
var onlyRouted = flag.Bool("R", false, "Display only connections routed through the router")
var noResolve = flag.Bool("n", false, "Do not resolve hostnames") // TODO resolve port names as well
var noHeader = flag.Bool("o", false, "Strip output header")

func main() {
	flag.Parse()

	which := conntrack.SNATFilter | conntrack.DNATFilter

	if *onlySNAT {
		which = conntrack.SNATFilter
	}

	if *onlyDNAT {
		which = conntrack.DNATFilter
	}

	if *onlyLocal {
		which = conntrack.LocalFilter
	}

	if *onlyRouted {
		which = conntrack.RoutedFilter
	}

	flows, err := conntrack.Flows()
	if err != nil {
		panic(err)
	}

	tabWriter := &tabwriter.Writer{}
	tabWriter.Init(os.Stdout, 0, 0, 4, ' ', 0)

	if !*noHeader {
		fmt.Fprintln(tabWriter, "Proto\tSource Address\tDestination Address\tState")
	}

	filteredFlows := flows.FilterByType(which)

	for _, flow := range filteredFlows {
		sHostname := lookup.Resolve(flow.Original.Source, *noResolve)
		dHostname := lookup.Resolve(flow.Original.Destination, *noResolve)

		fmt.Fprintf(tabWriter, "%s\t%s:%d\t%s:%d\t%s\n",
			flow.Protocol,
			sHostname,
			flow.Original.SPort,
			dHostname,
			flow.Original.DPort,
			flow.State,
		)
	}
	tabWriter.Flush()
}
<commit_msg>conntrack/netstat-nat: filter by protocol<commit_after>package main

import (
	"github.com/dominikh/simple-router/conntrack"
	"github.com/dominikh/simple-router/lookup"

	"flag"
	"fmt"
	"os"
	"text/tabwriter"
)

// TODO implement the following flags
//       -s <source-host>     : display connections by source
//       -d <destination-host>: display connections by destination
//       -x: extended hostnames view
//       -r src | dst | src-port | dst-port | state : sort connections
//       -N: display NAT box connection information (only valid with SNAT & DNAT)
//       -v: print version

var onlySNAT = flag.Bool("S", false, "Display only SNAT connections")
var onlyDNAT = flag.Bool("D", false, "Display only DNAT connections")
var onlyLocal = flag.Bool("L", false, "Display only local connections (originating from or going to the router)")
var onlyRouted = flag.Bool("R", false, "Display only connections routed through the router")
var noResolve = flag.Bool("n", false, "Do not resolve hostnames") // TODO resolve port names as well
var noHeader = flag.Bool("o", false, "Strip output header")
var protocol = flag.String("p", "", "Filter connections by protocol")

func main() {
	flag.Parse()

	which := conntrack.SNATFilter | conntrack.DNATFilter

	if *onlySNAT {
		which = conntrack.SNATFilter
	}

	if *onlyDNAT {
		which = conntrack.DNATFilter
	}

	if *onlyLocal {
		which = conntrack.LocalFilter
	}

	if *onlyRouted {
		which = conntrack.RoutedFilter
	}

	flows, err := conntrack.Flows()
	if err != nil {
		panic(err)
	}

	tabWriter := &tabwriter.Writer{}
	tabWriter.Init(os.Stdout, 0, 0, 4, ' ', 0)

	if !*noHeader {
		fmt.Fprintln(tabWriter, "Proto\tSource Address\tDestination Address\tState")
	}

	filteredFlows := flows.FilterByType(which)
	if *protocol != "" {
		filteredFlows = filteredFlows.FilterByProtocol(*protocol)
	}

	for _, flow := range filteredFlows {
		sHostname := lookup.Resolve(flow.Original.Source, *noResolve)
		dHostname := lookup.Resolve(flow.Original.Destination, *noResolve)

		fmt.Fprintf(tabWriter, "%s\t%s:%d\t%s:%d\t%s\n",
			flow.Protocol,
			sHostname,
			flow.Original.SPort,
			dHostname,
			flow.Original.DPort,
			flow.State,
		)
	}
	tabWriter.Flush()
}
<|endoftext|>


<commit_before>package resolver

import (
	"fmt"
	"github.com/miekg/dns"
	"golang.org/x/net/context"
	"sync"
)

func Query(q dns.Question, ctx context.Context) (*dns.Msg, error) {
	cl := dns.Client{
		Net: "tcp",
	}

	m := &dns.Msg{
		MsgHdr: dns.MsgHdr{
			Id:               dns.Id(),
			RecursionDesired: true,
		},
		Compress: true,
		Question: []dns.Question{q},
	}

	m = m.SetEdns0(4096, false)

	servers, err := getServers()
	if err != nil {
		return nil, err
	}

	var r *dns.Msg
	for _, s := range servers {
		// TODO: use context
		r, _, err = cl.Exchange(m, s+":53")
		if err == nil {
			return r, nil
		}
	}

	return nil, err
}

func getServers() ([]string, error) {
	err := loadConfig()
	if err != nil {
		return nil, err
	}

	if len(clientConfig.Servers) == 0 {
		return nil, fmt.Errorf("no DNS resolvers configured")
	}

	return clientConfig.Servers, nil
}

var loadConfigOnce sync.Once
var clientConfig *dns.ClientConfig
var configLoadError error

func loadConfig() error {
	loadConfigOnce.Do(func() {
		clientConfig, configLoadError = dns.ClientConfigFromFile("/etc/resolv.conf")
	})

	return configLoadError
}
<commit_msg>resolver<commit_after>package resolver

import (
	"fmt"
	"github.com/miekg/dns"
	"golang.org/x/net/context"
	"sync"
)

func Query(q dns.Question, ctx context.Context) (*dns.Msg, error) {
	cl := dns.Client{
		Net: "tcp",
	}

	m := &dns.Msg{
		MsgHdr: dns.MsgHdr{
			Id:               dns.Id(),
			RecursionDesired: true,
		},
		Compress: true,
		Question: []dns.Question{q},
	}

	m = m.SetEdns0(4096, false)

	servers, err := getServers()
	if err != nil {
		return nil, err
	}

	type txResult struct {
		Response *dns.Msg
		Err      error
	}

	maxTries := len(servers)
	if maxTries < 3 {
		maxTries = 3
	}

	for i := 0; i < maxTries; i++ {
		s := servers[i%len(servers)]

		txResultChan := make(chan txResult, 1)

		go func() {
			r, _, err := cl.Exchange(m, s+":53")
			txResultChan <- txResult{r, err}
		}()

		select {
		case <-ctx.Done():
			return nil, ctx.Err()

		case txResult := <-txResultChan:
			if txResult.Err == nil {
				return txResult.Response, nil
			}

			err = txResult.Err
		}
	}

	return nil, err
}

func getServers() ([]string, error) {
	err := loadConfig()
	if err != nil {
		return nil, err
	}

	if len(clientConfig.Servers) == 0 {
		return nil, fmt.Errorf("no DNS resolvers configured")
	}

	return clientConfig.Servers, nil
}

var loadConfigOnce sync.Once
var clientConfig *dns.ClientConfig
var configLoadError error

func loadConfig() error {
	loadConfigOnce.Do(func() {
		clientConfig, configLoadError = dns.ClientConfigFromFile("/etc/resolv.conf")
	})

	return configLoadError
}
<|endoftext|>