<commit_before>package bootstrappolicy

import (
	kapi "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/serviceaccount"
)

const (
	// SecurityContextConstraintPrivileged is used as the name for the system default privileged scc.
	SecurityContextConstraintPrivileged     = "privileged"
	SecurityContextConstraintPrivilegedDesc = "privileged allows access to all privileged and host features and the ability to run as any user, any group, any fsGroup, and with any SELinux context.  WARNING: this is the most relaxed SCC and should be used only for cluster administration. Grant with caution."

	// SecurityContextConstraintRestricted is used as the name for the system default restricted scc.
	SecurityContextConstraintRestricted     = "restricted"
	SecurityContextConstraintRestrictedDesc = "restricted denies access to all host features and requires pods to be run with a UID, and SELinux context that are allocated to the namespace.  This is the most restrictive SCC."

	// SecurityContextConstraintNonRoot is used as the name for the system default non-root scc.
	SecurityContextConstraintNonRoot     = "nonroot"
	SecurityContextConstraintNonRootDesc = "nonroot provides all features of the restricted SCC but allows users to run with any non-root UID.  The user must specify the UID or it must be specified on the by the manifest of the container runtime."

	// SecurityContextConstraintHostMountAndAnyUID is used as the name for the system default host mount + any UID scc.
	SecurityContextConstraintHostMountAndAnyUID     = "hostmount-anyuid"
	SecurityContextConstraintHostMountAndAnyUIDDesc = "hostmount-anyuid provides all the features of the restricted SCC but allows host mounts and any UID by a pod.  This is primarily used by the persistent volume recycler. WARNING: this SCC allows host file system access as any UID, including UID 0.  Grant with caution."

	// SecurityContextConstraintHostNS is used as the name for the system default scc
	// that grants access to all host ns features.
	SecurityContextConstraintHostNS     = "hostaccess"
	SecurityContextConstraintHostNSDesc = "hostaccess allows access to all host namespaces but still requires pods to be run with a UID and SELinux context that are allocated to the namespace. WARNING: this SCC allows host access to namespaces, file systems, and PIDS.  It should only be used by trusted pods.  Grant with caution."

	// SecurityContextConstraintsAnyUID is used as the name for the system default scc that
	// grants access to run as any uid but is still restricted to specific SELinux contexts.
	SecurityContextConstraintsAnyUID     = "anyuid"
	SecurityContextConstraintsAnyUIDDesc = "anyuid provides all features of the restricted SCC but allows users to run with any UID and any GID.  This is the default SCC for authenticated users."

	// DescriptionAnnotation is the annotation used for attaching descriptions.
	DescriptionAnnotation = "kubernetes.io/description"
)

// GetBootstrapSecurityContextConstraints returns the slice of default SecurityContextConstraints
// for system bootstrapping.  This method takes additional users and groups that should be added
// to the strategies.  Use GetBoostrapSCCAccess to produce the default set of mappings.
func GetBootstrapSecurityContextConstraints(sccNameToAdditionalGroups map[string][]string, sccNameToAdditionalUsers map[string][]string) []kapi.SecurityContextConstraints {
	// define priorities here and reference them below so it is easy to see, at a glance
	// what we're setting
	var (
		// this is set to 10 to allow wiggle room for admins to set other priorities without
		// having to adjust anyUID.
		securityContextConstraintsAnyUIDPriority = 10
	)

	constraints := []kapi.SecurityContextConstraints{
		// SecurityContextConstraintPrivileged allows all access for every field
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintPrivileged,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintPrivilegedDesc,
				},
			},
			AllowPrivilegedContainer:  true,
			AllowHostDirVolumePlugin:  true,
			AllowEmptyDirVolumePlugin: true,
			AllowHostNetwork:          true,
			AllowHostPorts:            true,
			AllowHostPID:              true,
			AllowHostIPC:              true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				Type: kapi.SELinuxStrategyRunAsAny,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				Type: kapi.RunAsUserStrategyRunAsAny,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintNonRoot does not allow host access, allocates SELinux labels
		// and allows the user to request a specific UID or provide the default in the dockerfile.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintNonRoot,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintNonRootDesc,
				},
			},
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that the user request to run as a specific UID or that
				// the docker file contain a USER directive.
				Type: kapi.RunAsUserStrategyMustRunAsNonRoot,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintHostMountAndAnyUID is the same as the restricted scc but allows host mounts and running as any UID.
		// Used by the PV recycler.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintHostMountAndAnyUID,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintHostMountAndAnyUIDDesc,
				},
			},
			AllowHostDirVolumePlugin:  true,
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.RunAsUserStrategyRunAsAny,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintHostNS allows access to everything except privileged on the host
		// but still allocates UIDs and SELinux.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintHostNS,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintHostNSDesc,
				},
			},
			AllowHostDirVolumePlugin:  true,
			AllowEmptyDirVolumePlugin: true,
			AllowHostNetwork:          true,
			AllowHostPorts:            true,
			AllowHostPID:              true,
			AllowHostIPC:              true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.RunAsUserStrategyMustRunAsRange,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintRestricted allows no host access and allocates UIDs and SELinux.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintRestricted,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintRestrictedDesc,
				},
			},
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.RunAsUserStrategyMustRunAsRange,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
			// drops unsafe caps
			RequiredDropCapabilities: []kapi.Capability{"KILL", "MKNOD", "SYS_CHROOT", "SETUID", "SETGID"},
		},
		// SecurityContextConstraintsAnyUID allows no host access and allocates SELinux.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintsAnyUID,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintsAnyUIDDesc,
				},
			},
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				Type: kapi.RunAsUserStrategyRunAsAny,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
			// prefer the anyuid SCC over ones that force a uid
			Priority: &securityContextConstraintsAnyUIDPriority,
			// drops unsafe caps
			RequiredDropCapabilities: []kapi.Capability{"KILL", "MKNOD", "SYS_CHROOT", "SETUID", "SETGID"},
		},
	}

	// add default access
	for i, constraint := range constraints {
		if usersToAdd, ok := sccNameToAdditionalUsers[constraint.Name]; ok {
			constraints[i].Users = append(constraints[i].Users, usersToAdd...)
		}
		if groupsToAdd, ok := sccNameToAdditionalGroups[constraint.Name]; ok {
			constraints[i].Groups = append(constraints[i].Groups, groupsToAdd...)
		}
	}
	return constraints
}

// GetBoostrapSCCAccess provides the default set of access that should be passed to GetBootstrapSecurityContextConstraints.
func GetBoostrapSCCAccess(infraNamespace string) (map[string][]string, map[string][]string) {
	groups := map[string][]string{
		SecurityContextConstraintPrivileged: {ClusterAdminGroup, NodesGroup},
		SecurityContextConstraintsAnyUID:    {ClusterAdminGroup},
		SecurityContextConstraintRestricted: {AuthenticatedGroup},
	}

	buildControllerUsername := serviceaccount.MakeUsername(infraNamespace, InfraBuildControllerServiceAccountName)
	pvControllerUsername := serviceaccount.MakeUsername(infraNamespace, InfraPersistentVolumeBinderControllerServiceAccountName)
	users := map[string][]string{
		SecurityContextConstraintPrivileged:         {buildControllerUsername},
		SecurityContextConstraintHostMountAndAnyUID: {pvControllerUsername},
	}
	return groups, users
}
<commit_msg>Many anyuid programs fail due to SETGID/SETUID caps<commit_after>package bootstrappolicy

import (
	kapi "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/serviceaccount"
)

const (
	// SecurityContextConstraintPrivileged is used as the name for the system default privileged scc.
	SecurityContextConstraintPrivileged     = "privileged"
	SecurityContextConstraintPrivilegedDesc = "privileged allows access to all privileged and host features and the ability to run as any user, any group, any fsGroup, and with any SELinux context.  WARNING: this is the most relaxed SCC and should be used only for cluster administration. Grant with caution."

	// SecurityContextConstraintRestricted is used as the name for the system default restricted scc.
	SecurityContextConstraintRestricted     = "restricted"
	SecurityContextConstraintRestrictedDesc = "restricted denies access to all host features and requires pods to be run with a UID, and SELinux context that are allocated to the namespace.  This is the most restrictive SCC."

	// SecurityContextConstraintNonRoot is used as the name for the system default non-root scc.
	SecurityContextConstraintNonRoot     = "nonroot"
	SecurityContextConstraintNonRootDesc = "nonroot provides all features of the restricted SCC but allows users to run with any non-root UID.  The user must specify the UID or it must be specified on the by the manifest of the container runtime."

	// SecurityContextConstraintHostMountAndAnyUID is used as the name for the system default host mount + any UID scc.
	SecurityContextConstraintHostMountAndAnyUID     = "hostmount-anyuid"
	SecurityContextConstraintHostMountAndAnyUIDDesc = "hostmount-anyuid provides all the features of the restricted SCC but allows host mounts and any UID by a pod.  This is primarily used by the persistent volume recycler. WARNING: this SCC allows host file system access as any UID, including UID 0.  Grant with caution."

	// SecurityContextConstraintHostNS is used as the name for the system default scc
	// that grants access to all host ns features.
	SecurityContextConstraintHostNS     = "hostaccess"
	SecurityContextConstraintHostNSDesc = "hostaccess allows access to all host namespaces but still requires pods to be run with a UID and SELinux context that are allocated to the namespace. WARNING: this SCC allows host access to namespaces, file systems, and PIDS.  It should only be used by trusted pods.  Grant with caution."

	// SecurityContextConstraintsAnyUID is used as the name for the system default scc that
	// grants access to run as any uid but is still restricted to specific SELinux contexts.
	SecurityContextConstraintsAnyUID     = "anyuid"
	SecurityContextConstraintsAnyUIDDesc = "anyuid provides all features of the restricted SCC but allows users to run with any UID and any GID.  This is the default SCC for authenticated users."

	// DescriptionAnnotation is the annotation used for attaching descriptions.
	DescriptionAnnotation = "kubernetes.io/description"
)

// GetBootstrapSecurityContextConstraints returns the slice of default SecurityContextConstraints
// for system bootstrapping.  This method takes additional users and groups that should be added
// to the strategies.  Use GetBoostrapSCCAccess to produce the default set of mappings.
func GetBootstrapSecurityContextConstraints(sccNameToAdditionalGroups map[string][]string, sccNameToAdditionalUsers map[string][]string) []kapi.SecurityContextConstraints {
	// define priorities here and reference them below so it is easy to see, at a glance
	// what we're setting
	var (
		// this is set to 10 to allow wiggle room for admins to set other priorities without
		// having to adjust anyUID.
		securityContextConstraintsAnyUIDPriority = 10
	)

	constraints := []kapi.SecurityContextConstraints{
		// SecurityContextConstraintPrivileged allows all access for every field
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintPrivileged,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintPrivilegedDesc,
				},
			},
			AllowPrivilegedContainer:  true,
			AllowHostDirVolumePlugin:  true,
			AllowEmptyDirVolumePlugin: true,
			AllowHostNetwork:          true,
			AllowHostPorts:            true,
			AllowHostPID:              true,
			AllowHostIPC:              true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				Type: kapi.SELinuxStrategyRunAsAny,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				Type: kapi.RunAsUserStrategyRunAsAny,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintNonRoot does not allow host access, allocates SELinux labels
		// and allows the user to request a specific UID or provide the default in the dockerfile.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintNonRoot,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintNonRootDesc,
				},
			},
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that the user request to run as a specific UID or that
				// the docker file contain a USER directive.
				Type: kapi.RunAsUserStrategyMustRunAsNonRoot,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintHostMountAndAnyUID is the same as the restricted scc but allows host mounts and running as any UID.
		// Used by the PV recycler.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintHostMountAndAnyUID,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintHostMountAndAnyUIDDesc,
				},
			},
			AllowHostDirVolumePlugin:  true,
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.RunAsUserStrategyRunAsAny,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintHostNS allows access to everything except privileged on the host
		// but still allocates UIDs and SELinux.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintHostNS,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintHostNSDesc,
				},
			},
			AllowHostDirVolumePlugin:  true,
			AllowEmptyDirVolumePlugin: true,
			AllowHostNetwork:          true,
			AllowHostPorts:            true,
			AllowHostPID:              true,
			AllowHostIPC:              true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.RunAsUserStrategyMustRunAsRange,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
		},
		// SecurityContextConstraintRestricted allows no host access and allocates UIDs and SELinux.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintRestricted,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintRestrictedDesc,
				},
			},
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.RunAsUserStrategyMustRunAsRange,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
			// drops unsafe caps
			RequiredDropCapabilities: []kapi.Capability{"KILL", "MKNOD", "SYS_CHROOT", "SETUID", "SETGID"},
		},
		// SecurityContextConstraintsAnyUID allows no host access and allocates SELinux.
		{
			ObjectMeta: kapi.ObjectMeta{
				Name: SecurityContextConstraintsAnyUID,
				Annotations: map[string]string{
					DescriptionAnnotation: SecurityContextConstraintsAnyUIDDesc,
				},
			},
			AllowEmptyDirVolumePlugin: true,
			SELinuxContext: kapi.SELinuxContextStrategyOptions{
				// This strategy requires that annotations on the namespace which will be populated
				// by the admission controller.  If namespaces are not annotated creating the strategy
				// will fail.
				Type: kapi.SELinuxStrategyMustRunAs,
			},
			RunAsUser: kapi.RunAsUserStrategyOptions{
				Type: kapi.RunAsUserStrategyRunAsAny,
			},
			FSGroup: kapi.FSGroupStrategyOptions{
				Type: kapi.FSGroupStrategyRunAsAny,
			},
			SupplementalGroups: kapi.SupplementalGroupsStrategyOptions{
				Type: kapi.SupplementalGroupsStrategyRunAsAny,
			},
			// prefer the anyuid SCC over ones that force a uid
			Priority: &securityContextConstraintsAnyUIDPriority,
			// drops unsafe caps
			RequiredDropCapabilities: []kapi.Capability{"KILL", "MKNOD", "SYS_CHROOT"},
		},
	}

	// add default access
	for i, constraint := range constraints {
		if usersToAdd, ok := sccNameToAdditionalUsers[constraint.Name]; ok {
			constraints[i].Users = append(constraints[i].Users, usersToAdd...)
		}
		if groupsToAdd, ok := sccNameToAdditionalGroups[constraint.Name]; ok {
			constraints[i].Groups = append(constraints[i].Groups, groupsToAdd...)
		}
	}
	return constraints
}

// GetBoostrapSCCAccess provides the default set of access that should be passed to GetBootstrapSecurityContextConstraints.
func GetBoostrapSCCAccess(infraNamespace string) (map[string][]string, map[string][]string) {
	groups := map[string][]string{
		SecurityContextConstraintPrivileged: {ClusterAdminGroup, NodesGroup},
		SecurityContextConstraintsAnyUID:    {ClusterAdminGroup},
		SecurityContextConstraintRestricted: {AuthenticatedGroup},
	}

	buildControllerUsername := serviceaccount.MakeUsername(infraNamespace, InfraBuildControllerServiceAccountName)
	pvControllerUsername := serviceaccount.MakeUsername(infraNamespace, InfraPersistentVolumeBinderControllerServiceAccountName)
	users := map[string][]string{
		SecurityContextConstraintPrivileged:         {buildControllerUsername},
		SecurityContextConstraintHostMountAndAnyUID: {pvControllerUsername},
	}
	return groups, users
}
<|endoftext|>


<commit_before>package cgoroonga

/*
#cgo LDFLAGS: -lgroonga
#include "cgoroonga.h"
*/
import "C"
import "unsafe"

func (c *Ctx) ExprParse(expr *Obj, str string, defaultColumn *Obj, defaultMode, defaultOp Operator, flags ExprFlags) error {
	cStr := C.CString(str)
	defer C.free(unsafe.Pointer(cStr))

	cStrLen := C.strlen(cStr)
	rc := C.grn_expr_parse(
		(*C.struct__grn_ctx)(unsafe.Pointer(c)),
		(*C.struct__grn_obj)(unsafe.Pointer(expr)),
		cStr, C.uint(cStrLen),
		(*C.struct__grn_obj)(unsafe.Pointer(defaultColumn)),
		C.grn_operator(defaultMode),
		C.grn_operator(defaultOp),
		C.grn_expr_flags(flags))
	if rc != SUCCESS {
		return errorFromRc(rc)
	}
	return nil
}

func (c *Ctx) ExprCreateForQuery(table *Obj) (expr, var_ *Obj) {
	cCtx := (*C.struct__grn_ctx)(unsafe.Pointer(c))
	cExpr := C.grn_expr_create(cCtx, nil, 0)
	var cVar *C.grn_obj
	if cExpr != nil {
		cVar = C.grn_expr_add_var(cCtx, cExpr, nil, 0)
		if cVar != nil {
			cTable := (*C.struct__grn_obj)(unsafe.Pointer(table))
			C.go_grn_record_init(cVar, 0, C.grn_obj_id(cCtx, cTable))
		}
	}
	return (*Obj)(cExpr), (*Obj)(cVar)
}
<commit_msg>Add error to return values of Ctx.ExprCreateForQuery<commit_after>package cgoroonga

/*
#cgo LDFLAGS: -lgroonga
#include "cgoroonga.h"
*/
import "C"
import "unsafe"

func (c *Ctx) ExprParse(expr *Obj, str string, defaultColumn *Obj, defaultMode, defaultOp Operator, flags ExprFlags) error {
	cStr := C.CString(str)
	defer C.free(unsafe.Pointer(cStr))

	cStrLen := C.strlen(cStr)
	rc := C.grn_expr_parse(
		(*C.struct__grn_ctx)(unsafe.Pointer(c)),
		(*C.struct__grn_obj)(unsafe.Pointer(expr)),
		cStr, C.uint(cStrLen),
		(*C.struct__grn_obj)(unsafe.Pointer(defaultColumn)),
		C.grn_operator(defaultMode),
		C.grn_operator(defaultOp),
		C.grn_expr_flags(flags))
	if rc != SUCCESS {
		return errorFromRc(rc)
	}
	return nil
}

func (c *Ctx) ExprCreateForQuery(table *Obj) (expr, var_ *Obj, err error) {
	cCtx := (*C.struct__grn_ctx)(unsafe.Pointer(c))
	cExpr := C.grn_expr_create(cCtx, nil, 0)
	var cVar *C.grn_obj
	if cExpr != nil {
		cVar = C.grn_expr_add_var(cCtx, cExpr, nil, 0)
		if cVar != nil {
			cTable := (*C.struct__grn_obj)(unsafe.Pointer(table))
			C.go_grn_record_init(cVar, 0, C.grn_obj_id(cCtx, cTable))
		}
	}
	if cExpr == nil || cVar == nil {
		return nil, nil, errorFromRc(c.rc)
	}
	return (*Obj)(cExpr), (*Obj)(cVar), nil
}
<|endoftext|>


<commit_before>package diff

import (
	"fmt"
	"sort"

	"github.com/gobwas/glob"

	"github.com/StackExchange/dnscontrol/models"
	"github.com/StackExchange/dnscontrol/pkg/printer"
)

// Correlation stores a difference between two domains.
type Correlation struct {
	d        *differ
	Existing *models.RecordConfig
	Desired  *models.RecordConfig
}

// Changeset stores many Correlation.
type Changeset []Correlation

// Differ is an interface for computing the difference between two zones.
type Differ interface {
	// IncrementalDiff performs a diff on a record-by-record basis, and returns a sets for which records need to be created, deleted, or modified.
	IncrementalDiff(existing []*models.RecordConfig) (unchanged, create, toDelete, modify Changeset)
	// ChangedGroups performs a diff more appropriate for providers with a "RecordSet" model, where all records with the same name and type are grouped.
	// Individual record changes are often not useful in such scenarios. Instead we return a map of record keys to a list of change descriptions within that group.
	ChangedGroups(existing []*models.RecordConfig) map[models.RecordKey][]string
}

// New is a constructor for a Differ.
func New(dc *models.DomainConfig, extraValues ...func(*models.RecordConfig) map[string]string) Differ {
	return &differ{
		dc:          dc,
		extraValues: extraValues,

		// compile IGNORE glob patterns
		compiledIgnoredLabels: compileIgnoredLabels(dc.IgnoredLabels),
	}
}

type differ struct {
	dc          *models.DomainConfig
	extraValues []func(*models.RecordConfig) map[string]string

	compiledIgnoredLabels []glob.Glob
}

// get normalized content for record. target, ttl, mxprio, and specified metadata
func (d *differ) content(r *models.RecordConfig) string {
	// NB(tlim): This function will eventually be replaced by calling
	// r.GetTargetDiffable().  In the meanwhile, this function compares
	// its output with r.GetTargetDiffable() to make sure the same
	// results are generated.  Once we have confidence, this function will go away.
	content := fmt.Sprintf("%v ttl=%d", r.GetTargetCombined(), r.TTL)
	var allMaps []map[string]string
	for _, f := range d.extraValues {
		// sort the extra values map keys to perform a deterministic
		// comparison since Golang maps iteration order is not guaranteed
		valueMap := f(r)
		allMaps = append(allMaps, valueMap)
		keys := make([]string, 0)
		for k := range valueMap {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			v := valueMap[k]
			content += fmt.Sprintf(" %s=%s", k, v)
		}
	}
	control := r.ToDiffable(allMaps...)
	if control != content {
		panic("OOPS! control != content")
	}
	return content
}

func (d *differ) IncrementalDiff(existing []*models.RecordConfig) (unchanged, create, toDelete, modify Changeset) {
	unchanged = Changeset{}
	create = Changeset{}
	toDelete = Changeset{}
	modify = Changeset{}
	desired := d.dc.Records

	// sort existing and desired by name

	existingByNameAndType := map[models.RecordKey][]*models.RecordConfig{}
	desiredByNameAndType := map[models.RecordKey][]*models.RecordConfig{}
	for _, e := range existing {
		if d.matchIgnored(e.GetLabel()) {
			printer.Debugf("Ignoring record %s %s due to IGNORE\n", e.GetLabel(), e.Type)
		} else {
			k := e.Key()
			existingByNameAndType[k] = append(existingByNameAndType[k], e)
		}
	}
	for _, dr := range desired {
		if d.matchIgnored(dr.GetLabel()) {
			panic(fmt.Sprintf("Trying to update/add IGNOREd record: %s %s", dr.GetLabel(), dr.Type))
		} else {
			k := dr.Key()
			desiredByNameAndType[k] = append(desiredByNameAndType[k], dr)
		}
	}
	// if NO_PURGE is set, just remove anything that is only in existing.
	if d.dc.KeepUnknown {
		for k := range existingByNameAndType {
			if _, ok := desiredByNameAndType[k]; !ok {
				printer.Debugf("Ignoring record set %s %s due to NO_PURGE\n", k.Type, k.NameFQDN)
				delete(existingByNameAndType, k)
			}
		}
	}
	// Look through existing records. This will give us changes and deletions and some additions.
	// Each iteration is only for a single type/name record set
	for key, existingRecords := range existingByNameAndType {
		desiredRecords := desiredByNameAndType[key]
		// first look through records that are the same target on both sides. Those are either modifications or unchanged
		for i := len(existingRecords) - 1; i >= 0; i-- {
			ex := existingRecords[i]
			for j, de := range desiredRecords {
				if de.GetTargetField() == ex.GetTargetField() {
					// they're either identical or should be a modification of each other (ttl or metadata changes)
					if d.content(de) == d.content(ex) {
						unchanged = append(unchanged, Correlation{d, ex, de})
					} else {
						modify = append(modify, Correlation{d, ex, de})
					}
					// remove from both slices by index
					existingRecords = existingRecords[:i+copy(existingRecords[i:], existingRecords[i+1:])]
					desiredRecords = desiredRecords[:j+copy(desiredRecords[j:], desiredRecords[j+1:])]
					break
				}
			}
		}

		desiredLookup := map[string]*models.RecordConfig{}
		existingLookup := map[string]*models.RecordConfig{}
		// build index based on normalized content data
		for _, ex := range existingRecords {
			normalized := d.content(ex)
			if existingLookup[normalized] != nil {
				panic(fmt.Sprintf("DUPLICATE E_RECORD FOUND: %s %s", key, normalized))
			}
			existingLookup[normalized] = ex
		}
		for _, de := range desiredRecords {
			normalized := d.content(de)
			if desiredLookup[normalized] != nil {
				panic(fmt.Sprintf("DUPLICATE D_RECORD FOUND: %s %s", key, normalized))
			}
			desiredLookup[normalized] = de
		}
		// if a record is in both, it is unchanged
		for norm, ex := range existingLookup {
			if de, ok := desiredLookup[norm]; ok {
				unchanged = append(unchanged, Correlation{d, ex, de})
				delete(existingLookup, norm)
				delete(desiredLookup, norm)
			}
		}
		// sort records by normalized text. Keeps behaviour deterministic
		existingStrings, desiredStrings := sortedKeys(existingLookup), sortedKeys(desiredLookup)
		// Modifications. Take 1 from each side.
		for len(desiredStrings) > 0 && len(existingStrings) > 0 {
			modify = append(modify, Correlation{d, existingLookup[existingStrings[0]], desiredLookup[desiredStrings[0]]})
			existingStrings = existingStrings[1:]
			desiredStrings = desiredStrings[1:]
		}
		// If desired still has things they are additions
		for _, norm := range desiredStrings {
			rec := desiredLookup[norm]
			create = append(create, Correlation{d, nil, rec})
		}
		// if found , but not desired, delete it
		for _, norm := range existingStrings {
			rec := existingLookup[norm]
			toDelete = append(toDelete, Correlation{d, rec, nil})
		}
		// remove this set from the desired list to indicate we have processed it.
		delete(desiredByNameAndType, key)
	}

	// any name/type sets not already processed are pure additions
	for name := range existingByNameAndType {
		delete(desiredByNameAndType, name)
	}
	for _, desiredList := range desiredByNameAndType {
		for _, rec := range desiredList {
			create = append(create, Correlation{d, nil, rec})
		}
	}
	return
}

func (d *differ) ChangedGroups(existing []*models.RecordConfig) map[models.RecordKey][]string {
	changedKeys := map[models.RecordKey][]string{}
	_, create, delete, modify := d.IncrementalDiff(existing)
	for _, c := range create {
		changedKeys[c.Desired.Key()] = append(changedKeys[c.Desired.Key()], c.String())
	}
	for _, d := range delete {
		changedKeys[d.Existing.Key()] = append(changedKeys[d.Existing.Key()], d.String())
	}
	for _, m := range modify {
		changedKeys[m.Desired.Key()] = append(changedKeys[m.Desired.Key()], m.String())
	}
	return changedKeys
}

func (c Correlation) String() string {
	if c.Existing == nil {
		return fmt.Sprintf("CREATE %s %s %s", c.Desired.Type, c.Desired.GetLabelFQDN(), c.d.content(c.Desired))
	}
	if c.Desired == nil {
		return fmt.Sprintf("DELETE %s %s %s", c.Existing.Type, c.Existing.GetLabelFQDN(), c.d.content(c.Existing))
	}
	return fmt.Sprintf("MODIFY %s %s: (%s) -> (%s)", c.Existing.Type, c.Existing.GetLabelFQDN(), c.d.content(c.Existing), c.d.content(c.Desired))
}

func sortedKeys(m map[string]*models.RecordConfig) []string {
	s := []string{}
	for v := range m {
		s = append(s, v)
	}
	sort.Strings(s)
	return s
}

func compileIgnoredLabels(ignoredLabels []string) []glob.Glob {
	result := make([]glob.Glob, 0, len(ignoredLabels))

	for _, tst := range ignoredLabels {
		g, err := glob.Compile(tst, '.')
		if err != nil {
			panic(fmt.Sprintf("Failed to compile IGNORE pattern %q: %v", tst, err))
		}

		result = append(result, g)
	}

	return result
}

func (d *differ) matchIgnored(name string) bool {
	for _, tst := range d.compiledIgnoredLabels {
		if tst.Match(name) {
			return true
		}
	}
	return false
}
<commit_msg>pull out records with identical content first. Saves from ordering issues.<commit_after>package diff

import (
	"fmt"
	"sort"

	"github.com/gobwas/glob"

	"github.com/StackExchange/dnscontrol/models"
	"github.com/StackExchange/dnscontrol/pkg/printer"
)

// Correlation stores a difference between two domains.
type Correlation struct {
	d        *differ
	Existing *models.RecordConfig
	Desired  *models.RecordConfig
}

// Changeset stores many Correlation.
type Changeset []Correlation

// Differ is an interface for computing the difference between two zones.
type Differ interface {
	// IncrementalDiff performs a diff on a record-by-record basis, and returns a sets for which records need to be created, deleted, or modified.
	IncrementalDiff(existing []*models.RecordConfig) (unchanged, create, toDelete, modify Changeset)
	// ChangedGroups performs a diff more appropriate for providers with a "RecordSet" model, where all records with the same name and type are grouped.
	// Individual record changes are often not useful in such scenarios. Instead we return a map of record keys to a list of change descriptions within that group.
	ChangedGroups(existing []*models.RecordConfig) map[models.RecordKey][]string
}

// New is a constructor for a Differ.
func New(dc *models.DomainConfig, extraValues ...func(*models.RecordConfig) map[string]string) Differ {
	return &differ{
		dc:          dc,
		extraValues: extraValues,

		// compile IGNORE glob patterns
		compiledIgnoredLabels: compileIgnoredLabels(dc.IgnoredLabels),
	}
}

type differ struct {
	dc          *models.DomainConfig
	extraValues []func(*models.RecordConfig) map[string]string

	compiledIgnoredLabels []glob.Glob
}

// get normalized content for record. target, ttl, mxprio, and specified metadata
func (d *differ) content(r *models.RecordConfig) string {
	// NB(tlim): This function will eventually be replaced by calling
	// r.GetTargetDiffable().  In the meanwhile, this function compares
	// its output with r.GetTargetDiffable() to make sure the same
	// results are generated.  Once we have confidence, this function will go away.
	content := fmt.Sprintf("%v ttl=%d", r.GetTargetCombined(), r.TTL)
	var allMaps []map[string]string
	for _, f := range d.extraValues {
		// sort the extra values map keys to perform a deterministic
		// comparison since Golang maps iteration order is not guaranteed
		valueMap := f(r)
		allMaps = append(allMaps, valueMap)
		keys := make([]string, 0)
		for k := range valueMap {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			v := valueMap[k]
			content += fmt.Sprintf(" %s=%s", k, v)
		}
	}
	control := r.ToDiffable(allMaps...)
	if control != content {
		panic("OOPS! control != content")
	}
	return content
}

func (d *differ) IncrementalDiff(existing []*models.RecordConfig) (unchanged, create, toDelete, modify Changeset) {
	unchanged = Changeset{}
	create = Changeset{}
	toDelete = Changeset{}
	modify = Changeset{}
	desired := d.dc.Records

	// sort existing and desired by name

	existingByNameAndType := map[models.RecordKey][]*models.RecordConfig{}
	desiredByNameAndType := map[models.RecordKey][]*models.RecordConfig{}
	for _, e := range existing {
		if d.matchIgnored(e.GetLabel()) {
			printer.Debugf("Ignoring record %s %s due to IGNORE\n", e.GetLabel(), e.Type)
		} else {
			k := e.Key()
			existingByNameAndType[k] = append(existingByNameAndType[k], e)
		}
	}
	for _, dr := range desired {
		if d.matchIgnored(dr.GetLabel()) {
			panic(fmt.Sprintf("Trying to update/add IGNOREd record: %s %s", dr.GetLabel(), dr.Type))
		} else {
			k := dr.Key()
			desiredByNameAndType[k] = append(desiredByNameAndType[k], dr)
		}
	}
	// if NO_PURGE is set, just remove anything that is only in existing.
	if d.dc.KeepUnknown {
		for k := range existingByNameAndType {
			if _, ok := desiredByNameAndType[k]; !ok {
				printer.Debugf("Ignoring record set %s %s due to NO_PURGE\n", k.Type, k.NameFQDN)
				delete(existingByNameAndType, k)
			}
		}
	}
	// Look through existing records. This will give us changes and deletions and some additions.
	// Each iteration is only for a single type/name record set
	for key, existingRecords := range existingByNameAndType {
		desiredRecords := desiredByNameAndType[key]

		// Very first, get rid of any identical records. Easy.
		for i := len(existingRecords) - 1; i >= 0; i-- {
			ex := existingRecords[i]
			for j, de := range desiredRecords {
				if d.content(de) != d.content(ex) {
					continue
				}
				unchanged = append(unchanged, Correlation{d, ex, de})
				existingRecords = existingRecords[:i+copy(existingRecords[i:], existingRecords[i+1:])]
				desiredRecords = desiredRecords[:j+copy(desiredRecords[j:], desiredRecords[j+1:])]
				break
			}
		}

		// Next, match by target. This will give the most natural modifications.
		for i := len(existingRecords) - 1; i >= 0; i-- {
			ex := existingRecords[i]
			for j, de := range desiredRecords {
				if de.GetTargetField() == ex.GetTargetField() {
					// two records share a target, but different content (ttl or metadata changes)
					modify = append(modify, Correlation{d, ex, de})
					// remove from both slices by index
					existingRecords = existingRecords[:i+copy(existingRecords[i:], existingRecords[i+1:])]
					desiredRecords = desiredRecords[:j+copy(desiredRecords[j:], desiredRecords[j+1:])]
					break
				}
			}
		}

		desiredLookup := map[string]*models.RecordConfig{}
		existingLookup := map[string]*models.RecordConfig{}
		// build index based on normalized content data
		for _, ex := range existingRecords {
			normalized := d.content(ex)
			if existingLookup[normalized] != nil {
				panic(fmt.Sprintf("DUPLICATE E_RECORD FOUND: %s %s", key, normalized))
			}
			existingLookup[normalized] = ex
		}
		for _, de := range desiredRecords {
			normalized := d.content(de)
			if desiredLookup[normalized] != nil {
				panic(fmt.Sprintf("DUPLICATE D_RECORD FOUND: %s %s", key, normalized))
			}
			desiredLookup[normalized] = de
		}
		// if a record is in both, it is unchanged
		for norm, ex := range existingLookup {
			if de, ok := desiredLookup[norm]; ok {
				unchanged = append(unchanged, Correlation{d, ex, de})
				delete(existingLookup, norm)
				delete(desiredLookup, norm)
			}
		}
		// sort records by normalized text. Keeps behaviour deterministic
		existingStrings, desiredStrings := sortedKeys(existingLookup), sortedKeys(desiredLookup)
		// Modifications. Take 1 from each side.
		for len(desiredStrings) > 0 && len(existingStrings) > 0 {
			modify = append(modify, Correlation{d, existingLookup[existingStrings[0]], desiredLookup[desiredStrings[0]]})
			existingStrings = existingStrings[1:]
			desiredStrings = desiredStrings[1:]
		}
		// If desired still has things they are additions
		for _, norm := range desiredStrings {
			rec := desiredLookup[norm]
			create = append(create, Correlation{d, nil, rec})
		}
		// if found , but not desired, delete it
		for _, norm := range existingStrings {
			rec := existingLookup[norm]
			toDelete = append(toDelete, Correlation{d, rec, nil})
		}
		// remove this set from the desired list to indicate we have processed it.
		delete(desiredByNameAndType, key)
	}

	// any name/type sets not already processed are pure additions
	for name := range existingByNameAndType {
		delete(desiredByNameAndType, name)
	}
	for _, desiredList := range desiredByNameAndType {
		for _, rec := range desiredList {
			create = append(create, Correlation{d, nil, rec})
		}
	}
	return
}

func (d *differ) ChangedGroups(existing []*models.RecordConfig) map[models.RecordKey][]string {
	changedKeys := map[models.RecordKey][]string{}
	_, create, delete, modify := d.IncrementalDiff(existing)
	for _, c := range create {
		changedKeys[c.Desired.Key()] = append(changedKeys[c.Desired.Key()], c.String())
	}
	for _, d := range delete {
		changedKeys[d.Existing.Key()] = append(changedKeys[d.Existing.Key()], d.String())
	}
	for _, m := range modify {
		changedKeys[m.Desired.Key()] = append(changedKeys[m.Desired.Key()], m.String())
	}
	return changedKeys
}

func (c Correlation) String() string {
	if c.Existing == nil {
		return fmt.Sprintf("CREATE %s %s %s", c.Desired.Type, c.Desired.GetLabelFQDN(), c.d.content(c.Desired))
	}
	if c.Desired == nil {
		return fmt.Sprintf("DELETE %s %s %s", c.Existing.Type, c.Existing.GetLabelFQDN(), c.d.content(c.Existing))
	}
	return fmt.Sprintf("MODIFY %s %s: (%s) -> (%s)", c.Existing.Type, c.Existing.GetLabelFQDN(), c.d.content(c.Existing), c.d.content(c.Desired))
}

func sortedKeys(m map[string]*models.RecordConfig) []string {
	s := []string{}
	for v := range m {
		s = append(s, v)
	}
	sort.Strings(s)
	return s
}

func compileIgnoredLabels(ignoredLabels []string) []glob.Glob {
	result := make([]glob.Glob, 0, len(ignoredLabels))

	for _, tst := range ignoredLabels {
		g, err := glob.Compile(tst, '.')
		if err != nil {
			panic(fmt.Sprintf("Failed to compile IGNORE pattern %q: %v", tst, err))
		}

		result = append(result, g)
	}

	return result
}

func (d *differ) matchIgnored(name string) bool {
	for _, tst := range d.compiledIgnoredLabels {
		if tst.Match(name) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package dnsdisco_test

import (
	"fmt"
	"net"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
	"github.com/rafaeljusto/dnsdisco"
)

func TestDiscover(t *testing.T) {
	scenarios := []struct {
		description    string
		service        string
		proto          string
		name           string
		expectedTarget string
		expectedPort   uint16
		expectedError  error
	}{
		{
			description:    "it should retrieve the target correctly",
			service:        "jabber",
			proto:          "tcp",
			name:           "registro.br",
			expectedTarget: "jabber.registro.br.",
			expectedPort:   5269,
		},
		{
			description: "it should fail when the protocol is invalid",
			service:     "jabber",
			proto:       "xxx",
			name:        "registro.br",
			expectedError: &net.DNSError{
				Err:  "no such host",
				Name: "_jabber._xxx.registro.br",
			},
		},
	}

	for i, item := range scenarios {
		target, port, err := dnsdisco.Discover(item.service, item.proto, item.name)

		if target != item.expectedTarget {
			t.Errorf("scenario %d, “%s”: mismatch targets. Expecting: “%s”; found “%s”",
				i, item.description, item.expectedTarget, target)
		}

		if port != item.expectedPort {
			t.Errorf("scenario %d, “%s”: mismatch ports. Expecting: “%d”; found “%d”",
				i, item.description, item.expectedPort, port)
		}

		// As the resolver change between machines, we can't guess the DNSError name's attribute. So we
		// need to inject the value on the expected error
		dnsError, ok1 := err.(*net.DNSError)
		expectedDNSError, ok2 := item.expectedError.(*net.DNSError)

		if ok1 && ok2 {
			expectedDNSError.Server = dnsError.Server
		}

		if !reflect.DeepEqual(err, item.expectedError) {
			t.Errorf("scenario %d, “%s”: mismatch errors. Expecting: “%v”; found “%v”",
				i, item.description, item.expectedError, err)
		}
	}
}

func TestDiscoverDefaultBalancer(t *testing.T) {
	scenarios := []struct {
		description    string
		service        string
		proto          string
		name           string
		retriever      dnsdisco.RetrieverFunc
		healthChecker  dnsdisco.HealthCheckerFunc
		rerun          int
		expectedTarget string
		expectedPort   uint16
	}{
		{
			description: "it should retrieve the target correctly (fallback inside priority group)",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   20,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   10,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return false, nil
				case "server2.example.com.":
					return true, nil
				}

				return false, nil
			}),
			expectedTarget: "server2.example.com.",
			expectedPort:   2222,
		},
		{
			description: "it should retrieve the target correctly (fallback to other priority group)",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   20,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   10,
					},
					&net.SRV{
						Target:   "server3.example.com.",
						Port:     3333,
						Priority: 20,
						Weight:   20,
					},
					&net.SRV{
						Target:   "server4.example.com.",
						Port:     4444,
						Priority: 20,
						Weight:   10,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return false, nil
				case "server2.example.com.":
					return false, nil
				case "server3.example.com.":
					return false, nil
				case "server4.example.com.":
					return true, nil
				}

				return false, nil
			}),
			expectedTarget: "server4.example.com.",
			expectedPort:   4444,
		},
		{
			description: "it should retrieve the target correctly (use the less used server)",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   200,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   0,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return true, nil
				case "server2.example.com.":
					return true, nil
				}

				return false, nil
			}),
			rerun:          1,
			expectedTarget: "server2.example.com.",
			expectedPort:   2222,
		},
		{
			description: "it should not select any target",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   200,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   0,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return false, nil
				case "server2.example.com.":
					return false, nil
				}

				return false, nil
			}),
			expectedTarget: "",
			expectedPort:   0,
		},
	}

	for i, item := range scenarios {
		discovery := dnsdisco.NewDiscovery(item.service, item.proto, item.name)
		discovery.Retriever = item.retriever
		discovery.HealthChecker = item.healthChecker

		if err := discovery.Refresh(); err != nil {
			t.Errorf("scenario %d, “%s”: unexpected error while retrieving DNS records. Details: %s",
				i, item.description, item.expectedTarget, err)
		}

		var target string
		var port uint16

		for i := 0; i <= item.rerun; i++ {
			target, port = discovery.Choose()
		}

		if target != item.expectedTarget {
			t.Errorf("scenario %d, “%s”: mismatch targets. Expecting: “%s”; found “%s”",
				i, item.description, item.expectedTarget, target)
		}

		if port != item.expectedPort {
			t.Errorf("scenario %d, “%s”: mismatch ports. Expecting: “%d”; found “%d”",
				i, item.description, item.expectedPort, port)
		}
	}
}

func ExampleDiscover() {
	target, port, err := dnsdisco.Discover("jabber", "tcp", "registro.br")
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("Target: %s\nPort: %d\n", target, port)

	// Output:
	// Target: jabber.registro.br.
	// Port: 5269
}

func ExampleRetrieverFunc() {
	discovery := dnsdisco.NewDiscovery("jabber", "tcp", "registro.br")
	discovery.Retriever = dnsdisco.RetrieverFunc(func(service, proto, name string) (servers []*net.SRV, err error) {
		client := dns.Client{
			ReadTimeout:  2 * time.Second,
			WriteTimeout: 2 * time.Second,
		}

		name = strings.TrimRight(name, ".")
		z := fmt.Sprintf("_%s._%s.%s.", service, proto, name)

		var request dns.Msg
		request.SetQuestion(z, dns.TypeSRV)
		request.RecursionDesired = true

		response, _, err := client.Exchange(&request, "8.8.8.8:53")
		if err != nil {
			return nil, err
		}

		for _, rr := range response.Answer {
			if srv, ok := rr.(*dns.SRV); ok {
				servers = append(servers, &net.SRV{
					Target:   srv.Target,
					Port:     srv.Port,
					Priority: srv.Priority,
					Weight:   srv.Weight,
				})
			}
		}

		return
	})

	// Retrieve the servers
	if err := discovery.Refresh(); err != nil {
		fmt.Println(err)
		return
	}

	target, port := discovery.Choose()
	fmt.Printf("Target: %s\nPort: %d\n", target, port)

	// Output:
	// Target: jabber.registro.br.
	// Port: 5269
}
<commit_msg>Add example for a custom load balancer algorithm<commit_after>package dnsdisco_test

import (
	"fmt"
	"net"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
	"github.com/rafaeljusto/dnsdisco"
)

func TestDiscover(t *testing.T) {
	scenarios := []struct {
		description    string
		service        string
		proto          string
		name           string
		expectedTarget string
		expectedPort   uint16
		expectedError  error
	}{
		{
			description:    "it should retrieve the target correctly",
			service:        "jabber",
			proto:          "tcp",
			name:           "registro.br",
			expectedTarget: "jabber.registro.br.",
			expectedPort:   5269,
		},
		{
			description: "it should fail when the protocol is invalid",
			service:     "jabber",
			proto:       "xxx",
			name:        "registro.br",
			expectedError: &net.DNSError{
				Err:  "no such host",
				Name: "_jabber._xxx.registro.br",
			},
		},
	}

	for i, item := range scenarios {
		target, port, err := dnsdisco.Discover(item.service, item.proto, item.name)

		if target != item.expectedTarget {
			t.Errorf("scenario %d, “%s”: mismatch targets. Expecting: “%s”; found “%s”",
				i, item.description, item.expectedTarget, target)
		}

		if port != item.expectedPort {
			t.Errorf("scenario %d, “%s”: mismatch ports. Expecting: “%d”; found “%d”",
				i, item.description, item.expectedPort, port)
		}

		// As the resolver change between machines, we can't guess the DNSError name's attribute. So we
		// need to inject the value on the expected error
		dnsError, ok1 := err.(*net.DNSError)
		expectedDNSError, ok2 := item.expectedError.(*net.DNSError)

		if ok1 && ok2 {
			expectedDNSError.Server = dnsError.Server
		}

		if !reflect.DeepEqual(err, item.expectedError) {
			t.Errorf("scenario %d, “%s”: mismatch errors. Expecting: “%v”; found “%v”",
				i, item.description, item.expectedError, err)
		}
	}
}

func TestDiscoverDefaultBalancer(t *testing.T) {
	scenarios := []struct {
		description    string
		service        string
		proto          string
		name           string
		retriever      dnsdisco.RetrieverFunc
		healthChecker  dnsdisco.HealthCheckerFunc
		rerun          int
		expectedTarget string
		expectedPort   uint16
	}{
		{
			description: "it should retrieve the target correctly (fallback inside priority group)",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   20,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   10,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return false, nil
				case "server2.example.com.":
					return true, nil
				}

				return false, nil
			}),
			expectedTarget: "server2.example.com.",
			expectedPort:   2222,
		},
		{
			description: "it should retrieve the target correctly (fallback to other priority group)",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   20,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   10,
					},
					&net.SRV{
						Target:   "server3.example.com.",
						Port:     3333,
						Priority: 20,
						Weight:   20,
					},
					&net.SRV{
						Target:   "server4.example.com.",
						Port:     4444,
						Priority: 20,
						Weight:   10,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return false, nil
				case "server2.example.com.":
					return false, nil
				case "server3.example.com.":
					return false, nil
				case "server4.example.com.":
					return true, nil
				}

				return false, nil
			}),
			expectedTarget: "server4.example.com.",
			expectedPort:   4444,
		},
		{
			description: "it should retrieve the target correctly (use the less used server)",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   200,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   0,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return true, nil
				case "server2.example.com.":
					return true, nil
				}

				return false, nil
			}),
			rerun:          1,
			expectedTarget: "server2.example.com.",
			expectedPort:   2222,
		},
		{
			description: "it should not select any target",
			service:     "jabber",
			proto:       "tcp",
			name:        "registro.br",
			retriever: dnsdisco.RetrieverFunc(func(service, proto, name string) ([]*net.SRV, error) {
				return []*net.SRV{
					&net.SRV{
						Target:   "server1.example.com.",
						Port:     1111,
						Priority: 10,
						Weight:   200,
					},
					&net.SRV{
						Target:   "server2.example.com.",
						Port:     2222,
						Priority: 10,
						Weight:   0,
					},
				}, nil
			}),
			healthChecker: dnsdisco.HealthCheckerFunc(func(target string, port uint16, proto string) (ok bool, err error) {
				switch target {
				case "server1.example.com.":
					return false, nil
				case "server2.example.com.":
					return false, nil
				}

				return false, nil
			}),
			expectedTarget: "",
			expectedPort:   0,
		},
	}

	for i, item := range scenarios {
		discovery := dnsdisco.NewDiscovery(item.service, item.proto, item.name)
		discovery.Retriever = item.retriever
		discovery.HealthChecker = item.healthChecker

		if err := discovery.Refresh(); err != nil {
			t.Errorf("scenario %d, “%s”: unexpected error while retrieving DNS records. Details: %s",
				i, item.description, item.expectedTarget, err)
		}

		var target string
		var port uint16

		for i := 0; i <= item.rerun; i++ {
			target, port = discovery.Choose()
		}

		if target != item.expectedTarget {
			t.Errorf("scenario %d, “%s”: mismatch targets. Expecting: “%s”; found “%s”",
				i, item.description, item.expectedTarget, target)
		}

		if port != item.expectedPort {
			t.Errorf("scenario %d, “%s”: mismatch ports. Expecting: “%d”; found “%d”",
				i, item.description, item.expectedPort, port)
		}
	}
}

func ExampleDiscover() {
	target, port, err := dnsdisco.Discover("jabber", "tcp", "registro.br")
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("Target: %s\nPort: %d\n", target, port)

	// Output:
	// Target: jabber.registro.br.
	// Port: 5269
}

func ExampleRetrieverFunc() {
	discovery := dnsdisco.NewDiscovery("jabber", "tcp", "registro.br")
	discovery.Retriever = dnsdisco.RetrieverFunc(func(service, proto, name string) (servers []*net.SRV, err error) {
		// Using a specific resolver with custom timeouts

		client := dns.Client{
			ReadTimeout:  2 * time.Second,
			WriteTimeout: 2 * time.Second,
		}

		name = strings.TrimRight(name, ".")
		z := fmt.Sprintf("_%s._%s.%s.", service, proto, name)

		var request dns.Msg
		request.SetQuestion(z, dns.TypeSRV)
		request.RecursionDesired = true

		response, _, err := client.Exchange(&request, "8.8.8.8:53")
		if err != nil {
			return nil, err
		}

		for _, rr := range response.Answer {
			if srv, ok := rr.(*dns.SRV); ok {
				servers = append(servers, &net.SRV{
					Target:   srv.Target,
					Port:     srv.Port,
					Priority: srv.Priority,
					Weight:   srv.Weight,
				})
			}
		}

		return
	})

	// Retrieve the servers
	if err := discovery.Refresh(); err != nil {
		fmt.Println(err)
		return
	}

	target, port := discovery.Choose()
	fmt.Printf("Target: %s\nPort: %d\n", target, port)

	// Output:
	// Target: jabber.registro.br.
	// Port: 5269
}

func ExampleBalancerFunc() {
	discovery := dnsdisco.NewDiscovery("jabber", "tcp", "registro.br")
	discovery.Balancer = dnsdisco.BalancerFunc(func(servers []dnsdisco.Server) (index int) {
		// Using a round-robin algorithm. As we don't known which server position
		// was used in the last time, we try to select using the Used attribute.

		minimum := -1
		for _, server := range servers {
			if server.Used < minimum || minimum == -1 {
				minimum = server.Used
			}
		}

		for i, server := range servers {
			if server.Used == minimum {
				return i
			}
		}

		return -1
	})

	// Retrieve the servers
	if err := discovery.Refresh(); err != nil {
		fmt.Println(err)
		return
	}

	target, port := discovery.Choose()
	fmt.Printf("Target: %s\nPort: %d\n", target, port)

	// Output:
	// Target: jabber.registro.br.
	// Port: 5269
}
<|endoftext|>


<commit_before>package ndn

import (
	"bytes"
	"errors"
	"io"
	"time"

	"github.com/paulbellamy/go-ndn/encoding"
	"github.com/paulbellamy/go-ndn/encoding/tlv"
	"github.com/paulbellamy/go-ndn/name"
	"github.com/paulbellamy/go-ndn/packets"
)

var ErrUnexpectedEventTypeReceived = errors.New("unexpected event type received")
var ErrInvalidPacket = errors.New("invalid packet received")

func NewFace(transport Transport) *Face {
	return &Face{
		transport:            transport,
		encoder:              tlv.NewEncoder(transport),
		pendingInterestTable: newPendingInterestTable(),
	}
}

type Face struct {
	transport                Transport
	encoder                  encoding.Encoder
	pendingInterestTable     *pendingInterestTable
	commandKeyChain          *KeyChain
	commandCertificateName   name.Name
	commandInterestGenerator commandInterestGenerator
}

func (f *Face) SetCommandSigningInfo(keyChain *KeyChain, certificateName name.Name) {
	f.commandKeyChain = keyChain
	f.commandCertificateName = certificateName.Copy()
}

func (f *Face) ExpressInterest(i *packets.Interest) (*pendingInterest, error) {
	pendingInterest := f.pendingInterestTable.AddInterest(i)
	err := f.encoder.Encode(i)
	if err != nil {
		return nil, err
	}
	return pendingInterest, nil
}

func (f *Face) RemovePendingInterest(id uint64) {
	f.pendingInterestTable.RemovePendingInterest(id)
}

func (f *Face) ProcessEvents() error {
	r := tlv.NewDecoder(f.transport)

	for {
		var packet interface{}
		err := r.Decode(&packet)
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		}

		err = f.dispatchEvent(packet)
		if err != nil {
			return err
		}
	}
}

func (f *Face) dispatchEvent(packet interface{}) error {
	switch packet := packet.(type) {
	case *packets.Interest:
		// TODO: implement this
	case *packets.Data:
		// parent, ok := event.(tlv.ParentTLV)
		// if !ok /*|| len(parent.V) != 4 */ {
		// 	return ErrInvalidPacket
		// }
		// name, err := name.FromTLV(parent.V[0])
		// if err != nil {
		// 	return err
		// }
		f.pendingInterestTable.DispatchData(packet)
	default:
		return ErrUnexpectedEventTypeReceived
	}
	return nil
}

func (f *Face) Put(d *packets.Data) error {
	return f.encoder.Encode(d)
}

func (f *Face) RegisterPrefix(prefix name.Name) (<-chan *packets.Interest, error) {
	if f.commandKeyChain == nil {
		return nil, ErrCommandKeyChainNotSet
	}
	if f.commandCertificateName.Size() == 0 {
		return nil, ErrCommandCertificateNameNotSet
	}

	n := name.New(name.Component{"localhost"}, name.Component{"nfd"}, name.Component{"rib"}, name.Component{"register"})
	n = n.AppendBytes(f.registerPrefixControlParameters(prefix))
	commandInterest := &packets.Interest{}
	commandInterest.SetName(n)
	f.makeCommandInterest(commandInterest)
	// The interest is answered by the local host, so set a short timeout.
	commandInterest.SetInterestLifetime(2 * time.Second)
	// make command interest
	return nil, nil
}

func (f *Face) makeCommandInterest(i *packets.Interest) {
	f.commandInterestGenerator.generate(i, f.commandKeyChain, f.commandCertificateName, tlv.NewEncoder)
}

func (f *Face) registerPrefixControlParameters(prefix name.Name) []byte {
	buf := &bytes.Buffer{}
	(&controlParameters{name: prefix}).WriteTo(buf)
	return buf.Bytes()
}

func (f *Face) Close() error {
	return f.transport.Close()
}
<commit_msg>updating call to Decode<commit_after>package ndn

import (
	"bytes"
	"errors"
	"io"
	"time"

	"github.com/paulbellamy/go-ndn/encoding"
	"github.com/paulbellamy/go-ndn/encoding/tlv"
	"github.com/paulbellamy/go-ndn/name"
	"github.com/paulbellamy/go-ndn/packets"
)

var ErrUnexpectedEventTypeReceived = errors.New("unexpected event type received")
var ErrInvalidPacket = errors.New("invalid packet received")

func NewFace(transport Transport) *Face {
	return &Face{
		transport:            transport,
		encoder:              tlv.NewEncoder(transport),
		pendingInterestTable: newPendingInterestTable(),
	}
}

type Face struct {
	transport                Transport
	encoder                  encoding.Encoder
	pendingInterestTable     *pendingInterestTable
	commandKeyChain          *KeyChain
	commandCertificateName   name.Name
	commandInterestGenerator commandInterestGenerator
}

func (f *Face) SetCommandSigningInfo(keyChain *KeyChain, certificateName name.Name) {
	f.commandKeyChain = keyChain
	f.commandCertificateName = certificateName.Copy()
}

func (f *Face) ExpressInterest(i *packets.Interest) (*pendingInterest, error) {
	pendingInterest := f.pendingInterestTable.AddInterest(i)
	err := f.encoder.Encode(i)
	if err != nil {
		return nil, err
	}
	return pendingInterest, nil
}

func (f *Face) RemovePendingInterest(id uint64) {
	f.pendingInterestTable.RemovePendingInterest(id)
}

func (f *Face) ProcessEvents() error {
	r := tlv.NewDecoder(f.transport)

	for {
		packet, err := r.Decode()
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		}

		err = f.dispatchEvent(packet)
		if err != nil {
			return err
		}
	}
}

func (f *Face) dispatchEvent(packet interface{}) error {
	switch packet := packet.(type) {
	case *packets.Interest:
		// TODO: implement this
	case *packets.Data:
		// parent, ok := event.(tlv.ParentTLV)
		// if !ok /*|| len(parent.V) != 4 */ {
		// 	return ErrInvalidPacket
		// }
		// name, err := name.FromTLV(parent.V[0])
		// if err != nil {
		// 	return err
		// }
		f.pendingInterestTable.DispatchData(packet)
	default:
		return ErrUnexpectedEventTypeReceived
	}
	return nil
}

func (f *Face) Put(d *packets.Data) error {
	return f.encoder.Encode(d)
}

func (f *Face) RegisterPrefix(prefix name.Name) (<-chan *packets.Interest, error) {
	if f.commandKeyChain == nil {
		return nil, ErrCommandKeyChainNotSet
	}
	if f.commandCertificateName.Size() == 0 {
		return nil, ErrCommandCertificateNameNotSet
	}

	n := name.New(name.Component{"localhost"}, name.Component{"nfd"}, name.Component{"rib"}, name.Component{"register"})
	n = n.AppendBytes(f.registerPrefixControlParameters(prefix))
	commandInterest := &packets.Interest{}
	commandInterest.SetName(n)
	f.makeCommandInterest(commandInterest)
	// The interest is answered by the local host, so set a short timeout.
	commandInterest.SetInterestLifetime(2 * time.Second)
	// make command interest
	return nil, nil
}

func (f *Face) makeCommandInterest(i *packets.Interest) {
	f.commandInterestGenerator.generate(i, f.commandKeyChain, f.commandCertificateName, tlv.NewEncoder)
}

func (f *Face) registerPrefixControlParameters(prefix name.Name) []byte {
	buf := &bytes.Buffer{}
	(&controlParameters{name: prefix}).WriteTo(buf)
	return buf.Bytes()
}

func (f *Face) Close() error {
	return f.transport.Close()
}
<|endoftext|>


<commit_before>/*
Copyright 2011 Google Inc.

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

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

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

// runsit runs stuff.
//
// Author: Brad Fitzpatrick <brad@danga.com>

package main

import (
	"bufio"
	"container/list"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"danga.com/runsit/jsonconfig"
)

// Flags.
var (
	httpPort  = flag.Int("http_port", 4762, "HTTP localhost admin port.")
	configDir = flag.String("config_dir", "config", "Directory containing per-task *.json config files.")
)

var (
	logBuf = new(logBuffer)
	logger = log.New(io.MultiWriter(os.Stderr, logBuf), "", log.Lmicroseconds|log.Lshortfile)

	tasksMu sync.Mutex
	tasks   = make(map[string]*Task)
)

const systemLogSize = 64 << 10

// logBuffer is a ring buffer.
type logBuffer struct {
	mu   sync.Mutex
	i    int
	full bool
	buf  [systemLogSize]byte
}

func (b *logBuffer) Write(p []byte) (ntot int, err error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for len(p) > 0 {
		n := copy(b.buf[b.i:], p)
		p = p[n:]
		ntot += n
		b.i += n
		if b.i == len(b.buf) {
			b.i = 0
			b.full = true
		}
	}
	return
}

func (b *logBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	if !b.full {
		return string(b.buf[:b.i])
	}
	s := string(b.buf[b.i:]) + string(b.buf[:b.i])
	if nl := strings.Index(s, "\n"); nl != -1 {
		// Remove first line, since it's probably truncated
		s = s[nl+1:]
	}
	return "...\n" + s
}

// A Task is a named daemon. A single instance of Task exists for the
// life of the runsit daemon, despite how many times the task has
// failed and restarted.
type Task struct {
	// Immutable:
	Name     string
	controlc chan interface{}

	// State owned by loop's goroutine:
	config   jsonconfig.Obj // last valid config
	running  *TaskInstance
	failures []*TaskInstance // last few failures, oldest first.
}

// TaskInstance is a particular instance of a running Task.
type TaskInstance struct {
	task      *Task          // set once; not goroutine safe (may only call public methods)
	startTime time.Time      // set once; immutable
	config    jsonconfig.Obj // set once; immutable
	cmd       *exec.Cmd      // set once; immutable
	output    TaskOutput     // internal locking, safe for concurrent access
}

// ID returns a unique ID string for this task instance.
func (in *TaskInstance) ID() string {
	return fmt.Sprintf("%s/%d-pid%d", in.task.Name, in.startTime.Unix(), in.Pid())
}

func (in *TaskInstance) Printf(format string, args ...interface{}) {
	logger.Printf(fmt.Sprintf("Task %s: %s", in.ID(), format), args...)
}

func (in *TaskInstance) Pid() int {
	if in.cmd == nil || in.cmd.Process == nil {
		return 0
	}
	return in.cmd.Process.Pid
}

// TaskOutput is the output
type TaskOutput struct {
	mu    sync.Mutex
	lines list.List // of *outputLine
}

func (to *TaskOutput) Add(o *outputLine) {
	to.mu.Lock()
	defer to.mu.Unlock()
	to.lines.PushBack(o)
	const maxKeepLines = 5000
	if to.lines.Len() > maxKeepLines {
		to.lines.Remove(to.lines.Front())
	}
}

func NewTask(name string) *Task {
	t := &Task{
		Name:     name,
		controlc: make(chan interface{}),
	}
	go t.loop()
	return t
}

func (t *Task) Printf(format string, args ...interface{}) {
	logger.Printf(fmt.Sprintf("Task %q: %s", t.Name, format), args...)
}

func (t *Task) loop() {
	t.Printf("Starting")
	for cm := range t.controlc {
		switch m := cm.(type) {
		case statusRequestMessage:
			m.resCh <- t.status()
		case runningRequestMessage:
			m.resCh <- t.running
		case updateMessage:
			t.update(m.tf)
		case stopMessage:
			t.stop()
		case waitMessage:
			t.onTaskFinished(m)
		}
	}
}

type waitMessage struct {
	instance *TaskInstance
	err      error // return from cmd.Wait(), nil, *exec.ExitError, or other type
}

type outputLine struct {
	t        time.Time
	instance *TaskInstance
	name     string // "stdout" or "stderr"
	isPrefix bool   // truncated line? (too long)
	data     string // line or prefix of line
}

type updateMessage struct {
	tf TaskFile
}

type stopMessage struct{}

type statusRequestMessage struct {
	resCh chan<- string
}

type runningRequestMessage struct {
	resCh chan<- *TaskInstance
}

func (t *Task) Update(tf TaskFile) {
	t.controlc <- updateMessage{tf}
}

// run in Task.loop
func (t *Task) onTaskFinished(m waitMessage) {
	m.instance.Printf("Task exited; err=%v", m.err)
	if m.instance == t.running {
		t.running = nil
	}
	const keepFailures = 5
	if len(t.failures) == keepFailures {
		copy(t.failures, t.failures[1:])
		t.failures = t.failures[:keepFailures-1]
	}
	t.failures = append(t.failures, m.instance)

	if m.err == nil {
		// TODO: vary sleep time (but not in this goroutine)
		// based on how process ended and when it was last
		// started (prevent crash/restart loops)
	}
	if t.config != nil {
		t.Printf("Restarting")
		t.updateFromConfig(t.config)
	}
}

// run in Task.loop
func (t *Task) update(tf TaskFile) {
	t.config = nil
	jc, err := jsonconfig.ReadFile(tf.ConfigFileName())
	t.stop()
	if err != nil {
		t.Printf("Bad config file: %v", err)
		return
	}
	t.updateFromConfig(jc)
}

// run in Task.loop
func (t *Task) updateFromConfig(jc jsonconfig.Obj) {
	t.config = nil
	t.stop()

	ports := jc.OptionalObject("ports")
	_ = ports
	user := jc.OptionalString("user", "")
	curUser := os.Getenv("USER")
	if user == "" {
		user = curUser
	}
	if user != curUser {
		panic("TODO: switch user")
	}

	env := []string{}
	stdEnv := jc.OptionalBool("standardEnv", true)
	if stdEnv {
		env = append(env, fmt.Sprintf("USER=%s", user))
	}
	envMap := jc.OptionalObject("env")
	for k, v := range envMap {
		env = append(env, fmt.Sprintf("%s=%s", k, v))
	}

	bin := jc.RequiredString("binary")
	cwd := jc.OptionalString("cwd", "")
	args := jc.OptionalList("args")
	if err := jc.Validate(); err != nil {
		t.Printf("configuration error: %v", err)
		return
	}
	t.config = jc

	_, err := os.Stat(bin)
	if err != nil {
		t.Printf("stat of binary %q failed: %v", bin, err)
		return
	}

	instance := &TaskInstance{
		task:      t,
		config:    jc,
		startTime: time.Now(),
		cmd:       exec.Command(bin, args...),
	}

	cmd := instance.cmd
	cmd.Dir = cwd
	cmd.Env = env

	outPipe, err := cmd.StdoutPipe()
	if err != nil {
		t.Printf("StdoutPipe: %v", err)
		return
	}
	errPipe, err := cmd.StderrPipe()
	if err != nil {
		t.Printf("StderrPipe: %v", err)
		outPipe.Close()
		return
	}

	err = cmd.Start()
	if err != nil {
		outPipe.Close()
		errPipe.Close()
		t.Printf("Error starting: %v", err)
		return
	}
	t.Printf("started with PID %d", instance.Pid())
	t.running = instance
	go instance.watchPipe(outPipe, "stdout")
	go instance.watchPipe(errPipe, "stderr")
	go instance.awaitDeath()
}

// run in its own goroutine
func (in *TaskInstance) awaitDeath() {
	err := in.cmd.Wait()
	in.task.controlc <- waitMessage{instance: in, err: err}
}

// run in its own goroutine
func (in *TaskInstance) watchPipe(r io.Reader, name string) {
	br := bufio.NewReader(r)
	for {
		sl, isPrefix, err := br.ReadLine()
		if err == io.EOF {
			// Not worth logging about.
			return
		}
		if err != nil {
			in.Printf("pipe %q closed: %v", name, err)
			return
		}
		in.output.Add(&outputLine{
			t:        time.Now(),
			instance: in,
			name:     name,
			isPrefix: isPrefix,
			data:     string(sl),
		})
	}
	panic("unreachable")
}

func (t *Task) Stop() {
	t.controlc <- stopMessage{}
}

// runs in Task.loop
func (t *Task) stop() {
	in := t.running
	if in == nil {
		return
	}
	in.Printf("sending SIGKILL")
	// TODO: more graceful kill types
	in.cmd.Process.Kill()
	t.running = nil
}

func (t *Task) Status() string {
	ch := make(chan string, 1)
	t.controlc <- statusRequestMessage{resCh: ch}
	return <-ch
}

func (t *Task) RunningInstance() (*TaskInstance, bool) {
	ch := make(chan *TaskInstance, 1)
	t.controlc <- runningRequestMessage{resCh: ch}
	in := <-ch
	return in, in != nil
}

// runs in Task.loop
func (t *Task) status() string {
	in := t.running
	if in != nil {
		d := time.Now().Sub(in.startTime)
		return fmt.Sprintf("running; for %v", d)
	}
	if t.config == nil {
		return "not running, no valid config"
	}
	// TODO: flesh these not running states out.
	// e.g. intentionaly stopped, how long we're pausing before
	// next re-start attempt, etc.
	return "not running; valid config"
}

func watchConfigDir() {
	for tf := range dirWatcher().Updates() {
		t := GetOrMakeTask(tf.Name())
		go t.Update(tf)
	}
}

func GetTask(name string) (*Task, bool) {
	tasksMu.Lock()
	defer tasksMu.Unlock()
	t, ok := tasks[name]
	return t, ok
}

// GetOrMakeTask returns or create the named task.
func GetOrMakeTask(name string) *Task {
	tasksMu.Lock()
	defer tasksMu.Unlock()
	t, ok := tasks[name]
	if !ok {
		t = NewTask(name)
		tasks[name] = t
	}
	return t
}

// GetTasks returns all known tasks.
func GetTasks() []*Task {
	ts := []*Task{}
	tasksMu.Lock()
	defer tasksMu.Unlock()
	for _, t := range tasks {
		ts = append(ts, t)
	}
	return ts
}

func main() {
	flag.Parse()

	ln, err := net.Listen("tcp", "localhost:"+strconv.Itoa(*httpPort))
	if err != nil {
		logger.Printf("Error listening on port %d: %v", *httpPort, err)
		os.Exit(1)
		return
	}
	logger.Printf("Listening on port %d", *httpPort)

	go watchConfigDir()
	go runWebServer(ln)
	select {}
}
<commit_msg>make outputLine more accessible<commit_after>/*
Copyright 2011 Google Inc.

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

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

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

// runsit runs stuff.
//
// Author: Brad Fitzpatrick <brad@danga.com>

package main

import (
	"bufio"
	"container/list"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"danga.com/runsit/jsonconfig"
)

// Flags.
var (
	httpPort  = flag.Int("http_port", 4762, "HTTP localhost admin port.")
	configDir = flag.String("config_dir", "config", "Directory containing per-task *.json config files.")
)

var (
	logBuf = new(logBuffer)
	logger = log.New(io.MultiWriter(os.Stderr, logBuf), "", log.Lmicroseconds|log.Lshortfile)

	tasksMu sync.Mutex
	tasks   = make(map[string]*Task)
)

const systemLogSize = 64 << 10

// logBuffer is a ring buffer.
type logBuffer struct {
	mu   sync.Mutex
	i    int
	full bool
	buf  [systemLogSize]byte
}

func (b *logBuffer) Write(p []byte) (ntot int, err error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for len(p) > 0 {
		n := copy(b.buf[b.i:], p)
		p = p[n:]
		ntot += n
		b.i += n
		if b.i == len(b.buf) {
			b.i = 0
			b.full = true
		}
	}
	return
}

func (b *logBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	if !b.full {
		return string(b.buf[:b.i])
	}
	s := string(b.buf[b.i:]) + string(b.buf[:b.i])
	if nl := strings.Index(s, "\n"); nl != -1 {
		// Remove first line, since it's probably truncated
		s = s[nl+1:]
	}
	return "...\n" + s
}

// A Task is a named daemon. A single instance of Task exists for the
// life of the runsit daemon, despite how many times the task has
// failed and restarted.
type Task struct {
	// Immutable:
	Name     string
	controlc chan interface{}

	// State owned by loop's goroutine:
	config   jsonconfig.Obj // last valid config
	running  *TaskInstance
	failures []*TaskInstance // last few failures, oldest first.
}

// TaskInstance is a particular instance of a running Task.
type TaskInstance struct {
	task      *Task          // set once; not goroutine safe (may only call public methods)
	startTime time.Time      // set once; immutable
	config    jsonconfig.Obj // set once; immutable
	cmd       *exec.Cmd      // set once; immutable
	output    TaskOutput     // internal locking, safe for concurrent access
}

// ID returns a unique ID string for this task instance.
func (in *TaskInstance) ID() string {
	return fmt.Sprintf("%s/%d-pid%d", in.task.Name, in.startTime.Unix(), in.Pid())
}

func (in *TaskInstance) Printf(format string, args ...interface{}) {
	logger.Printf(fmt.Sprintf("Task %s: %s", in.ID(), format), args...)
}

func (in *TaskInstance) Pid() int {
	if in.cmd == nil || in.cmd.Process == nil {
		return 0
	}
	return in.cmd.Process.Pid
}

// TaskOutput is the output
type TaskOutput struct {
	mu    sync.Mutex
	lines list.List // of *outputLine
}

func (to *TaskOutput) Add(o *outputLine) {
	to.mu.Lock()
	defer to.mu.Unlock()
	to.lines.PushBack(o)
	const maxKeepLines = 5000
	if to.lines.Len() > maxKeepLines {
		to.lines.Remove(to.lines.Front())
	}
}

func (to *TaskOutput) lineSlice() []*outputLine {
	to.mu.Lock()
	defer to.mu.Unlock()
	var lines []*outputLine
	for e := to.lines.Front(); e != nil; e = e.Next() {
		lines = append(lines, e.Value.(*outputLine))
	}
	return lines
}

func NewTask(name string) *Task {
	t := &Task{
		Name:     name,
		controlc: make(chan interface{}),
	}
	go t.loop()
	return t
}

func (t *Task) Printf(format string, args ...interface{}) {
	logger.Printf(fmt.Sprintf("Task %q: %s", t.Name, format), args...)
}

func (t *Task) loop() {
	t.Printf("Starting")
	for cm := range t.controlc {
		switch m := cm.(type) {
		case statusRequestMessage:
			m.resCh <- t.status()
		case runningRequestMessage:
			m.resCh <- t.running
		case updateMessage:
			t.update(m.tf)
		case stopMessage:
			t.stop()
		case waitMessage:
			t.onTaskFinished(m)
		}
	}
}

type waitMessage struct {
	instance *TaskInstance
	err      error // return from cmd.Wait(), nil, *exec.ExitError, or other type
}

type outputLine struct {
	T        time.Time
	instance *TaskInstance
	Name     string // "stdout" or "stderr"
	isPrefix bool   // truncated line? (too long)
	Data     string // line or prefix of line
}

type updateMessage struct {
	tf TaskFile
}

type stopMessage struct{}

type statusRequestMessage struct {
	resCh chan<- string
}

type runningRequestMessage struct {
	resCh chan<- *TaskInstance
}

func (t *Task) Update(tf TaskFile) {
	t.controlc <- updateMessage{tf}
}

// run in Task.loop
func (t *Task) onTaskFinished(m waitMessage) {
	m.instance.Printf("Task exited; err=%v", m.err)
	if m.instance == t.running {
		t.running = nil
	}
	const keepFailures = 5
	if len(t.failures) == keepFailures {
		copy(t.failures, t.failures[1:])
		t.failures = t.failures[:keepFailures-1]
	}
	t.failures = append(t.failures, m.instance)

	if m.err == nil {
		// TODO: vary sleep time (but not in this goroutine)
		// based on how process ended and when it was last
		// started (prevent crash/restart loops)
	}
	if t.config != nil {
		t.Printf("Restarting")
		t.updateFromConfig(t.config)
	}
}

// run in Task.loop
func (t *Task) update(tf TaskFile) {
	t.config = nil
	jc, err := jsonconfig.ReadFile(tf.ConfigFileName())
	t.stop()
	if err != nil {
		t.Printf("Bad config file: %v", err)
		return
	}
	t.updateFromConfig(jc)
}

// run in Task.loop
func (t *Task) updateFromConfig(jc jsonconfig.Obj) {
	t.config = nil
	t.stop()

	ports := jc.OptionalObject("ports")
	_ = ports
	user := jc.OptionalString("user", "")
	curUser := os.Getenv("USER")
	if user == "" {
		user = curUser
	}
	if user != curUser {
		panic("TODO: switch user")
	}

	env := []string{}
	stdEnv := jc.OptionalBool("standardEnv", true)
	if stdEnv {
		env = append(env, fmt.Sprintf("USER=%s", user))
	}
	envMap := jc.OptionalObject("env")
	for k, v := range envMap {
		env = append(env, fmt.Sprintf("%s=%s", k, v))
	}

	bin := jc.RequiredString("binary")
	cwd := jc.OptionalString("cwd", "")
	args := jc.OptionalList("args")
	if err := jc.Validate(); err != nil {
		t.Printf("configuration error: %v", err)
		return
	}
	t.config = jc

	_, err := os.Stat(bin)
	if err != nil {
		t.Printf("stat of binary %q failed: %v", bin, err)
		return
	}

	instance := &TaskInstance{
		task:      t,
		config:    jc,
		startTime: time.Now(),
		cmd:       exec.Command(bin, args...),
	}

	cmd := instance.cmd
	cmd.Dir = cwd
	cmd.Env = env

	outPipe, err := cmd.StdoutPipe()
	if err != nil {
		t.Printf("StdoutPipe: %v", err)
		return
	}
	errPipe, err := cmd.StderrPipe()
	if err != nil {
		t.Printf("StderrPipe: %v", err)
		outPipe.Close()
		return
	}

	err = cmd.Start()
	if err != nil {
		outPipe.Close()
		errPipe.Close()
		t.Printf("Error starting: %v", err)
		return
	}
	t.Printf("started with PID %d", instance.Pid())
	t.running = instance
	go instance.watchPipe(outPipe, "stdout")
	go instance.watchPipe(errPipe, "stderr")
	go instance.awaitDeath()
}

// run in its own goroutine
func (in *TaskInstance) awaitDeath() {
	err := in.cmd.Wait()
	in.task.controlc <- waitMessage{instance: in, err: err}
}

// run in its own goroutine
func (in *TaskInstance) watchPipe(r io.Reader, name string) {
	br := bufio.NewReader(r)
	for {
		sl, isPrefix, err := br.ReadLine()
		if err == io.EOF {
			// Not worth logging about.
			return
		}
		if err != nil {
			in.Printf("pipe %q closed: %v", name, err)
			return
		}
		in.output.Add(&outputLine{
			T:        time.Now(),
			instance: in,
			Name:     name,
			isPrefix: isPrefix,
			Data:     string(sl),
		})
	}
	panic("unreachable")
}

func (t *Task) Stop() {
	t.controlc <- stopMessage{}
}

// runs in Task.loop
func (t *Task) stop() {
	in := t.running
	if in == nil {
		return
	}
	in.Printf("sending SIGKILL")
	// TODO: more graceful kill types
	in.cmd.Process.Kill()
	t.running = nil
}

func (t *Task) Status() string {
	ch := make(chan string, 1)
	t.controlc <- statusRequestMessage{resCh: ch}
	return <-ch
}

func (t *Task) RunningInstance() (*TaskInstance, bool) {
	ch := make(chan *TaskInstance, 1)
	t.controlc <- runningRequestMessage{resCh: ch}
	in := <-ch
	return in, in != nil
}

// runs in Task.loop
func (t *Task) status() string {
	in := t.running
	if in != nil {
		d := time.Now().Sub(in.startTime)
		return fmt.Sprintf("running; for %v", d)
	}
	if t.config == nil {
		return "not running, no valid config"
	}
	// TODO: flesh these not running states out.
	// e.g. intentionaly stopped, how long we're pausing before
	// next re-start attempt, etc.
	return "not running; valid config"
}

func watchConfigDir() {
	for tf := range dirWatcher().Updates() {
		t := GetOrMakeTask(tf.Name())
		go t.Update(tf)
	}
}

func GetTask(name string) (*Task, bool) {
	tasksMu.Lock()
	defer tasksMu.Unlock()
	t, ok := tasks[name]
	return t, ok
}

// GetOrMakeTask returns or create the named task.
func GetOrMakeTask(name string) *Task {
	tasksMu.Lock()
	defer tasksMu.Unlock()
	t, ok := tasks[name]
	if !ok {
		t = NewTask(name)
		tasks[name] = t
	}
	return t
}

// GetTasks returns all known tasks.
func GetTasks() []*Task {
	ts := []*Task{}
	tasksMu.Lock()
	defer tasksMu.Unlock()
	for _, t := range tasks {
		ts = append(ts, t)
	}
	return ts
}

func main() {
	flag.Parse()

	ln, err := net.Listen("tcp", "localhost:"+strconv.Itoa(*httpPort))
	if err != nil {
		logger.Printf("Error listening on port %d: %v", *httpPort, err)
		os.Exit(1)
		return
	}
	logger.Printf("Listening on port %d", *httpPort)

	go watchConfigDir()
	go runWebServer(ln)
	select {}
}
<|endoftext|>


<commit_before>package bag

import "sync"

type Bag struct {
	m     sync.RWMutex
	items []interface{}
}

func NewBag() *Bag {
	return &Bag{
		items: make([]interface{}, 0),
	}
}

func (b *Bag) Add(v interface{}) {
	b.m.Lock()
	defer b.m.Unlock()
	b.items = append(b.items, v)
}

func (b *Bag) Size() int {
	b.m.RLock()
	defer b.m.RUnlock()
	return len(b.items)
}

func (b *Bag) Empty() bool {
	return b.Size() == 0
}

func (b *Bag) Items() []interface{} {
	b.m.RLock()
	defer b.m.RUnlock()
	return b.items
}
<commit_msg>ssh config test<commit_after>package bag

import "sync"

type Bag struct {
	m     sync.RWMutex
	items []interface{}
}

func NewBag() *Bag {
	return &Bag{
		items: make([]interface{}, 0),
	}
}

func (b *Bag) Add(v interface{}) {
	b.m.Lock()
	defer b.m.Unlock()
	b.items = append(b.items, v)
}

func (b *Bag) Size() int {
	b.m.RLock()
	defer b.m.RUnlock()
	return len(b.items)
}

func (b *Bag) Empty() bool {
	return b.Size() ==0
}

func (b *Bag) Items() []interface{} {
	b.m.RLock()
	defer b.m.RUnlock()
	return b.items
}
<|endoftext|>


<commit_before>package main

import(
    "encoding/xml"
    "log"
    "time"
    "net/http"
)

func FeedEntryToRss2Item(entry *FeedEntry) (item Rss2Item) {
    item.Title = entry.Title
    item.Link = entry.Link
    item.Description = string(entry.Content)
    if "" == entry.Cache.LastModified {
        item.PubDate = time.Now().Format(http.TimeFormat)
    }
    item.PubDate = entry.Cache.LastModified
    item.Guid = entry.Link

    return
}

func GenerateRss2Feed(indexCache HtmlCache, entries []FeedEntry) (feedStr []byte, err error) {
    if "" == indexCache.LastModified {
        indexCache.LastModified = time.Now().Format(http.TimeFormat)
    }

    feed := Rss2Feed { Version: FEED_VERSION }
    feed.Channel = Rss2Channel {
        Title: indexCache.Title,
        Link: indexCache.URL,
        Description: "",
        PubDate: indexCache.LastModified,
        Generator: GOFEED_NAME + " " + GOFEED_VERSION,
    }

    feed.Channel.Items = make([]Rss2Item, len(entries))
    for itemInd, entry := range entries {
        feed.Channel.Items = append(feed.Channel.Items[:itemInd], FeedEntryToRss2Item(&entry))
    }

    feedStr, err = xml.MarshalIndent(&feed, "  ", "    ")
    if nil != err {
        log.Printf("failed to marshal rss2 feed: %s", err)
    }

    return
}

<commit_msg>fix emtpy item.pubDate issue<commit_after>package main

import(
    "encoding/xml"
    "log"
    "time"
    "net/http"
)

func FeedEntryToRss2Item(entry *FeedEntry) (item Rss2Item) {
    item.Title = entry.Title
    item.Link = entry.Link
    item.Description = string(entry.Content)
    if "" == entry.Cache.LastModified {
        item.PubDate = time.Now().Format(http.TimeFormat)
    } else {
        item.PubDate = entry.Cache.LastModified
    }
    item.Guid = entry.Link

    return
}

func GenerateRss2Feed(indexCache HtmlCache, entries []FeedEntry) (feedStr []byte, err error) {
    if "" == indexCache.LastModified {
        indexCache.LastModified = time.Now().Format(http.TimeFormat)
    }

    feed := Rss2Feed { Version: FEED_VERSION }
    feed.Channel = Rss2Channel {
        Title: indexCache.Title,
        Link: indexCache.URL,
        Description: "",
        PubDate: indexCache.LastModified,
        Generator: GOFEED_NAME + " " + GOFEED_VERSION,
    }

    feed.Channel.Items = make([]Rss2Item, len(entries))
    for itemInd, entry := range entries {
        feed.Channel.Items = append(feed.Channel.Items[:itemInd], FeedEntryToRss2Item(&entry))
    }

    feedStr, err = xml.MarshalIndent(&feed, "  ", "    ")
    if nil != err {
        log.Printf("failed to marshal rss2 feed: %s", err)
    }

    return
}

<|endoftext|>


<commit_before>package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"regexp"
)

var (
	rComment = regexp.MustCompile(`^//\s*@inject_tag:\s*(\S+)$`)
	rInject  = regexp.MustCompile("`$")
)

type textArea struct {
	Start int
	End   int
	Tag   string
}

func parseFile(inputPath string) (areas []textArea, err error) {
	log.Printf("parsing file %q for inject tag comments", inputPath)
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		return
	}

	for _, decl := range f.Decls {
		// check if is generic declaration
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}

		var typeSpec *ast.TypeSpec
		for _, spec := range genDecl.Specs {
			if ts, tsOK := spec.(*ast.TypeSpec); tsOK {
				typeSpec = ts
				break
			}
		}

		// skip if can't get type spec
		if typeSpec == nil {
			continue
		}

		// not a struct, skip
		structDecl, ok := typeSpec.Type.(*ast.StructType)
		if !ok {
			continue
		}

		for _, field := range structDecl.Fields.List {
			// skip if field has no doc
			if field.Doc == nil {
				continue
			}
			for _, comment := range field.Doc.List {
				tag := tagFromComment(comment.Text)
				if tag == "" {
					continue
				}
				area := textArea{
					Start: int(field.Pos()),
					End:   int(field.End()),
					Tag:   tag,
				}
				areas = append(areas, area)
			}
		}
	}
	log.Printf("parsed file %q, number of fields to inject custom tags: %d", inputPath, len(areas))
	return
}

func writeFile(inputPath string, areas []textArea) (err error) {
	f, err := os.Open(inputPath)
	if err != nil {
		return
	}

	contents, err := ioutil.ReadAll(f)
	if err != nil {
		return
	}

	if err = f.Close(); err != nil {
		return
	}

	// inject custom tags from tail of file first to preserve order
	for i := range areas {
		area := areas[len(areas)-i-1]
		log.Printf("inject custom tag %q to expression %q", area.Tag, string(contents[area.Start-1:area.End-1]))
		contents = injectTag(contents, area)
	}
	if err = ioutil.WriteFile(inputPath, contents, 0644); err != nil {
		return
	}

	if len(areas) > 0 {
		log.Printf("file %q is injected with custom tags", inputPath)
	}
	return
}
<commit_msg>fixes comment regex<commit_after>package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"regexp"
)

var (
	rComment = regexp.MustCompile(`^//\s*@inject_tag:\s*(.*)$`)
	rInject  = regexp.MustCompile("`$")
)

type textArea struct {
	Start int
	End   int
	Tag   string
}

func parseFile(inputPath string) (areas []textArea, err error) {
	log.Printf("parsing file %q for inject tag comments", inputPath)
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		return
	}

	for _, decl := range f.Decls {
		// check if is generic declaration
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}

		var typeSpec *ast.TypeSpec
		for _, spec := range genDecl.Specs {
			if ts, tsOK := spec.(*ast.TypeSpec); tsOK {
				typeSpec = ts
				break
			}
		}

		// skip if can't get type spec
		if typeSpec == nil {
			continue
		}

		// not a struct, skip
		structDecl, ok := typeSpec.Type.(*ast.StructType)
		if !ok {
			continue
		}

		for _, field := range structDecl.Fields.List {
			// skip if field has no doc
			if field.Doc == nil {
				continue
			}
			for _, comment := range field.Doc.List {
				tag := tagFromComment(comment.Text)
				if tag == "" {
					continue
				}
				area := textArea{
					Start: int(field.Pos()),
					End:   int(field.End()),
					Tag:   tag,
				}
				areas = append(areas, area)
			}
		}
	}
	log.Printf("parsed file %q, number of fields to inject custom tags: %d", inputPath, len(areas))
	return
}

func writeFile(inputPath string, areas []textArea) (err error) {
	f, err := os.Open(inputPath)
	if err != nil {
		return
	}

	contents, err := ioutil.ReadAll(f)
	if err != nil {
		return
	}

	if err = f.Close(); err != nil {
		return
	}

	// inject custom tags from tail of file first to preserve order
	for i := range areas {
		area := areas[len(areas)-i-1]
		log.Printf("inject custom tag %q to expression %q", area.Tag, string(contents[area.Start-1:area.End-1]))
		contents = injectTag(contents, area)
	}
	if err = ioutil.WriteFile(inputPath, contents, 0644); err != nil {
		return
	}

	if len(areas) > 0 {
		log.Printf("file %q is injected with custom tags", inputPath)
	}
	return
}
<|endoftext|>