<commit_before>package sandbox

import (
	"crypto/rand"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/containernetworking/plugins/pkg/ns"
	"github.com/docker/docker/pkg/mount"
	"github.com/docker/docker/pkg/symlink"
	"github.com/kubernetes-incubator/cri-o/oci"
	"github.com/sirupsen/logrus"
	"golang.org/x/sys/unix"
	"k8s.io/apimachinery/pkg/fields"
	pb "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
	"k8s.io/kubernetes/pkg/kubelet/network/hostport"
)

// NetNs handles data pertaining a network namespace
type NetNs struct {
	sync.Mutex
	ns       ns.NetNS
	symlink  *os.File
	closed   bool
	restored bool
}

func (ns *NetNs) symlinkCreate(name string) error {
	b := make([]byte, 4)
	_, randErr := rand.Reader.Read(b)
	if randErr != nil {
		return randErr
	}

	nsName := fmt.Sprintf("%s-%x", name, b)
	symlinkPath := filepath.Join(NsRunDir, nsName)

	if err := os.Symlink(ns.ns.Path(), symlinkPath); err != nil {
		return err
	}

	fd, err := os.Open(symlinkPath)
	if err != nil {
		if removeErr := os.RemoveAll(symlinkPath); removeErr != nil {
			return removeErr
		}

		return err
	}

	ns.symlink = fd

	return nil
}

func (ns *NetNs) symlinkRemove() error {
	if err := ns.symlink.Close(); err != nil {
		return err
	}

	return os.RemoveAll(ns.symlink.Name())
}

func isSymbolicLink(path string) (bool, error) {
	fi, err := os.Lstat(path)
	if err != nil {
		return false, err
	}

	return fi.Mode()&os.ModeSymlink == os.ModeSymlink, nil
}

// NetNsGet returns the NetNs associated with the given nspath and name
func NetNsGet(nspath, name string) (*NetNs, error) {
	if err := ns.IsNSorErr(nspath); err != nil {
		return nil, ErrClosedNetNS
	}

	symlink, symlinkErr := isSymbolicLink(nspath)
	if symlinkErr != nil {
		return nil, symlinkErr
	}

	var resolvedNsPath string
	if symlink {
		path, err := os.Readlink(nspath)
		if err != nil {
			return nil, err
		}
		resolvedNsPath = path
	} else {
		resolvedNsPath = nspath
	}

	netNS, err := ns.GetNS(resolvedNsPath)
	if err != nil {
		return nil, err
	}

	netNs := &NetNs{ns: netNS, closed: false, restored: true}

	if symlink {
		fd, err := os.Open(nspath)
		if err != nil {
			return nil, err
		}

		netNs.symlink = fd
	} else {
		if err := netNs.symlinkCreate(name); err != nil {
			return nil, err
		}
	}

	return netNs, nil
}

// HostNetNsPath returns the current network namespace for the host
func HostNetNsPath() (string, error) {
	netNS, err := ns.GetCurrentNS()
	if err != nil {
		return "", err
	}

	defer netNS.Close()
	return netNS.Path(), nil
}

// Sandbox contains data surrounding kubernetes sandboxes on the server
type Sandbox struct {
	id        string
	namespace string
	// OCI pod name (eg "<namespace>-<name>-<attempt>")
	name string
	// Kubernetes pod name (eg, "<name>")
	kubeName       string
	logDir         string
	labels         fields.Set
	annotations    map[string]string
	infraContainer *oci.Container
	containers     oci.ContainerStorer
	processLabel   string
	mountLabel     string
	netns          *NetNs
	metadata       *pb.PodSandboxMetadata
	nsOpts         *pb.NamespaceOption
	shmPath        string
	cgroupParent   string
	privileged     bool
	trusted        bool
	resolvPath     string
	hostnamePath   string
	hostname       string
	portMappings   []*hostport.PortMapping
	stopped        bool
	// ipv4 or ipv6 cache
	ip                 string
	seccompProfilePath string
	created            time.Time
	hostNetwork        bool
}

const (
	// DefaultShmSize is the default shm size
	DefaultShmSize = 64 * 1024 * 1024
	// NsRunDir is the default directory in which running network namespaces
	// are stored
	NsRunDir = "/var/run/netns"
	// PodInfraCommand is the default command when starting a pod infrastructure
	// container
	PodInfraCommand = "/pause"
)

var (
	// ErrIDEmpty is the erro returned when the id of the sandbox is empty
	ErrIDEmpty = errors.New("PodSandboxId should not be empty")
	// ErrClosedNetNS is the error returned when the network namespace of the
	// sandbox is closed
	ErrClosedNetNS = errors.New("PodSandbox networking namespace is closed")
)

// New creates and populates a new pod sandbox
// New sandboxes have no containers, no infra container, and no network namespaces associated with them
// An infra container must be attached before the sandbox is added to the state
func New(id, namespace, name, kubeName, logDir string, labels, annotations map[string]string, processLabel, mountLabel string, metadata *pb.PodSandboxMetadata, shmPath, cgroupParent string, privileged, trusted bool, resolvPath, hostname string, portMappings []*hostport.PortMapping) (*Sandbox, error) {
	sb := new(Sandbox)
	sb.id = id
	sb.namespace = namespace
	sb.name = name
	sb.kubeName = kubeName
	sb.logDir = logDir
	sb.labels = labels
	sb.annotations = annotations
	sb.containers = oci.NewMemoryStore()
	sb.processLabel = processLabel
	sb.mountLabel = mountLabel
	sb.metadata = metadata
	sb.shmPath = shmPath
	sb.cgroupParent = cgroupParent
	sb.privileged = privileged
	sb.trusted = trusted
	sb.resolvPath = resolvPath
	sb.hostname = hostname
	sb.portMappings = portMappings
	sb.created = time.Now()

	return sb, nil
}

// SetSeccompProfilePath sets the seccomp profile path
func (s *Sandbox) SetSeccompProfilePath(pp string) {
	s.seccompProfilePath = pp
}

// SeccompProfilePath returns the seccomp profile path
func (s *Sandbox) SeccompProfilePath() string {
	return s.seccompProfilePath
}

// AddIP stores the ip in the sandbox
func (s *Sandbox) AddIP(ip string) {
	s.ip = ip
}

// SetNamespaceOptions sets whether the pod is running using host network
func (s *Sandbox) SetNamespaceOptions(nsOpts *pb.NamespaceOption) {
	s.nsOpts = nsOpts
}

// NamespaceOptions returns the namespace options for the sandbox
func (s *Sandbox) NamespaceOptions() *pb.NamespaceOption {
	return s.nsOpts
}

// IP returns the ip of the sandbox
func (s *Sandbox) IP() string {
	return s.ip
}

// ID returns the id of the sandbox
func (s *Sandbox) ID() string {
	return s.id
}

// Namespace returns the namespace for the sandbox
func (s *Sandbox) Namespace() string {
	return s.namespace
}

// Name returns the name of the sandbox
func (s *Sandbox) Name() string {
	return s.name
}

// KubeName returns the kubernetes name for the sandbox
func (s *Sandbox) KubeName() string {
	return s.kubeName
}

// LogDir returns the location of the logging directory for the sandbox
func (s *Sandbox) LogDir() string {
	return s.logDir
}

// Labels returns the labels associated with the sandbox
func (s *Sandbox) Labels() fields.Set {
	return s.labels
}

// Annotations returns a list of annotations for the sandbox
func (s *Sandbox) Annotations() map[string]string {
	return s.annotations
}

// InfraContainer returns the infrastructure container for the sandbox
func (s *Sandbox) InfraContainer() *oci.Container {
	return s.infraContainer
}

// Containers returns the ContainerStorer that contains information on all
// of the containers in the sandbox
func (s *Sandbox) Containers() oci.ContainerStorer {
	return s.containers
}

// ProcessLabel returns the process label for the sandbox
func (s *Sandbox) ProcessLabel() string {
	return s.processLabel
}

// MountLabel returns the mount label for the sandbox
func (s *Sandbox) MountLabel() string {
	return s.mountLabel
}

// Metadata returns a set of metadata about the sandbox
func (s *Sandbox) Metadata() *pb.PodSandboxMetadata {
	return s.metadata
}

// ShmPath returns the shm path of the sandbox
func (s *Sandbox) ShmPath() string {
	return s.shmPath
}

// CgroupParent returns the cgroup parent of the sandbox
func (s *Sandbox) CgroupParent() string {
	return s.cgroupParent
}

// Privileged returns whether or not the containers in the sandbox are
// privileged containers
func (s *Sandbox) Privileged() bool {
	return s.privileged
}

// Trusted returns whether or not the containers in the sandbox are trusted
func (s *Sandbox) Trusted() bool {
	return s.trusted
}

// ResolvPath returns the resolv path for the sandbox
func (s *Sandbox) ResolvPath() string {
	return s.resolvPath
}

// AddHostnamePath adds the hostname path to the sandbox
func (s *Sandbox) AddHostnamePath(hostname string) {
	s.hostnamePath = hostname
}

// HostnamePath retrieves the hostname path from a sandbox
func (s *Sandbox) HostnamePath() string {
	return s.hostnamePath
}

// Hostname returns the hsotname of the sandbox
func (s *Sandbox) Hostname() string {
	return s.hostname
}

// PortMappings returns a list of port mappings between the host and the sandbox
func (s *Sandbox) PortMappings() []*hostport.PortMapping {
	return s.portMappings
}

// AddContainer adds a container to the sandbox
func (s *Sandbox) AddContainer(c *oci.Container) {
	s.containers.Add(c.Name(), c)
}

// GetContainer retrieves a container from the sandbox
func (s *Sandbox) GetContainer(name string) *oci.Container {
	return s.containers.Get(name)
}

// RemoveContainer deletes a container from the sandbox
func (s *Sandbox) RemoveContainer(c *oci.Container) {
	s.containers.Delete(c.Name())
}

// SetInfraContainer sets the infrastructure container of a sandbox
// Attempts to set the infrastructure container after one is already present will throw an error
func (s *Sandbox) SetInfraContainer(infraCtr *oci.Container) error {
	if s.infraContainer != nil {
		return fmt.Errorf("sandbox already has an infra container")
	} else if infraCtr == nil {
		return fmt.Errorf("must provide non-nil infra container")
	}

	s.infraContainer = infraCtr

	return nil
}

// RemoveInfraContainer removes the infrastructure container of a sandbox
func (s *Sandbox) RemoveInfraContainer() {
	s.infraContainer = nil
}

// NetNs retrieves the network namespace of the sandbox
// If the sandbox uses the host namespace, nil is returned
func (s *Sandbox) NetNs() ns.NetNS {
	if s.netns == nil {
		return nil
	}

	return s.netns.ns
}

// NetNsPath returns the path to the network namespace of the sandbox.
// If the sandbox uses the host namespace, nil is returned
func (s *Sandbox) NetNsPath() string {
	if s.netns == nil {
		return ""
	}

	return s.netns.symlink.Name()
}

// NetNsCreate creates a new network namespace for the sandbox
func (s *Sandbox) NetNsCreate() error {
	if s.netns != nil {
		return fmt.Errorf("net NS already created")
	}

	netNS, err := ns.NewNS()
	if err != nil {
		return err
	}

	s.netns = &NetNs{
		ns:     netNS,
		closed: false,
	}

	if err := s.netns.symlinkCreate(s.name); err != nil {
		logrus.Warnf("Could not create nentns symlink %v", err)

		if err1 := s.netns.ns.Close(); err1 != nil {
			return err1
		}

		return err
	}

	return nil
}

// SetStopped sets the sandbox state to stopped.
// This should be set after a stop operation succeeds
// so that subsequent stops can return fast.
func (s *Sandbox) SetStopped() {
	s.stopped = true
}

// Stopped returns whether the sandbox state has been
// set to stopped.
func (s *Sandbox) Stopped() bool {
	return s.stopped
}

// NetNsJoin attempts to join the sandbox to an existing network namespace
// This will fail if the sandbox is already part of a network namespace
func (s *Sandbox) NetNsJoin(nspath, name string) error {
	if s.netns != nil {
		return fmt.Errorf("sandbox already has a network namespace, cannot join another")
	}

	netNS, err := NetNsGet(nspath, name)
	if err != nil {
		return err
	}

	s.netns = netNS

	return nil
}

// NetNsRemove removes the network namespace associated with the sandbox
func (s *Sandbox) NetNsRemove() error {
	if s.netns == nil {
		logrus.Warn("no networking namespace")
		return nil
	}

	s.netns.Lock()
	defer s.netns.Unlock()

	if s.netns.closed {
		// netNsRemove() can be called multiple
		// times without returning an error.
		return nil
	}

	if err := s.netns.symlinkRemove(); err != nil {
		return err
	}

	if err := s.netns.ns.Close(); err != nil {
		return err
	}

	s.netns.closed = true

	if s.netns.restored {
		// we got namespaces in the form of
		// /var/run/netns/cni-0d08effa-06eb-a963-f51a-e2b0eceffc5d
		// but /var/run on most system is symlinked to /run so we first resolve
		// the symlink and then try and see if it's mounted
		fp, err := symlink.FollowSymlinkInScope(s.netns.ns.Path(), "/")
		if err != nil {
			return err
		}
		if mounted, err := mount.Mounted(fp); err == nil && mounted {
			if err := unix.Unmount(fp, unix.MNT_DETACH); err != nil {
				return err
			}
		}

		if err := os.RemoveAll(s.netns.ns.Path()); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Add context to net ns symlink removal errors<commit_after>package sandbox

import (
	"crypto/rand"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/containernetworking/plugins/pkg/ns"
	"github.com/docker/docker/pkg/mount"
	"github.com/docker/docker/pkg/symlink"
	"github.com/kubernetes-incubator/cri-o/oci"
	"github.com/sirupsen/logrus"
	"golang.org/x/sys/unix"
	"k8s.io/apimachinery/pkg/fields"
	pb "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
	"k8s.io/kubernetes/pkg/kubelet/network/hostport"
)

// NetNs handles data pertaining a network namespace
type NetNs struct {
	sync.Mutex
	ns       ns.NetNS
	symlink  *os.File
	closed   bool
	restored bool
}

func (ns *NetNs) symlinkCreate(name string) error {
	b := make([]byte, 4)
	_, randErr := rand.Reader.Read(b)
	if randErr != nil {
		return randErr
	}

	nsName := fmt.Sprintf("%s-%x", name, b)
	symlinkPath := filepath.Join(NsRunDir, nsName)

	if err := os.Symlink(ns.ns.Path(), symlinkPath); err != nil {
		return err
	}

	fd, err := os.Open(symlinkPath)
	if err != nil {
		if removeErr := os.RemoveAll(symlinkPath); removeErr != nil {
			return removeErr
		}

		return err
	}

	ns.symlink = fd

	return nil
}

func (ns *NetNs) symlinkRemove() error {
	if err := ns.symlink.Close(); err != nil {
		return fmt.Errorf("failed to close net ns symlink: %v", err)
	}

	if err := os.RemoveAll(ns.symlink.Name()); err != nil {
		return fmt.Errorf("failed to remove net ns symlink: %v", err)
	}

	return nil
}

func isSymbolicLink(path string) (bool, error) {
	fi, err := os.Lstat(path)
	if err != nil {
		return false, err
	}

	return fi.Mode()&os.ModeSymlink == os.ModeSymlink, nil
}

// NetNsGet returns the NetNs associated with the given nspath and name
func NetNsGet(nspath, name string) (*NetNs, error) {
	if err := ns.IsNSorErr(nspath); err != nil {
		return nil, ErrClosedNetNS
	}

	symlink, symlinkErr := isSymbolicLink(nspath)
	if symlinkErr != nil {
		return nil, symlinkErr
	}

	var resolvedNsPath string
	if symlink {
		path, err := os.Readlink(nspath)
		if err != nil {
			return nil, err
		}
		resolvedNsPath = path
	} else {
		resolvedNsPath = nspath
	}

	netNS, err := ns.GetNS(resolvedNsPath)
	if err != nil {
		return nil, err
	}

	netNs := &NetNs{ns: netNS, closed: false, restored: true}

	if symlink {
		fd, err := os.Open(nspath)
		if err != nil {
			return nil, err
		}

		netNs.symlink = fd
	} else {
		if err := netNs.symlinkCreate(name); err != nil {
			return nil, err
		}
	}

	return netNs, nil
}

// HostNetNsPath returns the current network namespace for the host
func HostNetNsPath() (string, error) {
	netNS, err := ns.GetCurrentNS()
	if err != nil {
		return "", err
	}

	defer netNS.Close()
	return netNS.Path(), nil
}

// Sandbox contains data surrounding kubernetes sandboxes on the server
type Sandbox struct {
	id        string
	namespace string
	// OCI pod name (eg "<namespace>-<name>-<attempt>")
	name string
	// Kubernetes pod name (eg, "<name>")
	kubeName       string
	logDir         string
	labels         fields.Set
	annotations    map[string]string
	infraContainer *oci.Container
	containers     oci.ContainerStorer
	processLabel   string
	mountLabel     string
	netns          *NetNs
	metadata       *pb.PodSandboxMetadata
	nsOpts         *pb.NamespaceOption
	shmPath        string
	cgroupParent   string
	privileged     bool
	trusted        bool
	resolvPath     string
	hostnamePath   string
	hostname       string
	portMappings   []*hostport.PortMapping
	stopped        bool
	// ipv4 or ipv6 cache
	ip                 string
	seccompProfilePath string
	created            time.Time
	hostNetwork        bool
}

const (
	// DefaultShmSize is the default shm size
	DefaultShmSize = 64 * 1024 * 1024
	// NsRunDir is the default directory in which running network namespaces
	// are stored
	NsRunDir = "/var/run/netns"
	// PodInfraCommand is the default command when starting a pod infrastructure
	// container
	PodInfraCommand = "/pause"
)

var (
	// ErrIDEmpty is the erro returned when the id of the sandbox is empty
	ErrIDEmpty = errors.New("PodSandboxId should not be empty")
	// ErrClosedNetNS is the error returned when the network namespace of the
	// sandbox is closed
	ErrClosedNetNS = errors.New("PodSandbox networking namespace is closed")
)

// New creates and populates a new pod sandbox
// New sandboxes have no containers, no infra container, and no network namespaces associated with them
// An infra container must be attached before the sandbox is added to the state
func New(id, namespace, name, kubeName, logDir string, labels, annotations map[string]string, processLabel, mountLabel string, metadata *pb.PodSandboxMetadata, shmPath, cgroupParent string, privileged, trusted bool, resolvPath, hostname string, portMappings []*hostport.PortMapping) (*Sandbox, error) {
	sb := new(Sandbox)
	sb.id = id
	sb.namespace = namespace
	sb.name = name
	sb.kubeName = kubeName
	sb.logDir = logDir
	sb.labels = labels
	sb.annotations = annotations
	sb.containers = oci.NewMemoryStore()
	sb.processLabel = processLabel
	sb.mountLabel = mountLabel
	sb.metadata = metadata
	sb.shmPath = shmPath
	sb.cgroupParent = cgroupParent
	sb.privileged = privileged
	sb.trusted = trusted
	sb.resolvPath = resolvPath
	sb.hostname = hostname
	sb.portMappings = portMappings
	sb.created = time.Now()

	return sb, nil
}

// SetSeccompProfilePath sets the seccomp profile path
func (s *Sandbox) SetSeccompProfilePath(pp string) {
	s.seccompProfilePath = pp
}

// SeccompProfilePath returns the seccomp profile path
func (s *Sandbox) SeccompProfilePath() string {
	return s.seccompProfilePath
}

// AddIP stores the ip in the sandbox
func (s *Sandbox) AddIP(ip string) {
	s.ip = ip
}

// SetNamespaceOptions sets whether the pod is running using host network
func (s *Sandbox) SetNamespaceOptions(nsOpts *pb.NamespaceOption) {
	s.nsOpts = nsOpts
}

// NamespaceOptions returns the namespace options for the sandbox
func (s *Sandbox) NamespaceOptions() *pb.NamespaceOption {
	return s.nsOpts
}

// IP returns the ip of the sandbox
func (s *Sandbox) IP() string {
	return s.ip
}

// ID returns the id of the sandbox
func (s *Sandbox) ID() string {
	return s.id
}

// Namespace returns the namespace for the sandbox
func (s *Sandbox) Namespace() string {
	return s.namespace
}

// Name returns the name of the sandbox
func (s *Sandbox) Name() string {
	return s.name
}

// KubeName returns the kubernetes name for the sandbox
func (s *Sandbox) KubeName() string {
	return s.kubeName
}

// LogDir returns the location of the logging directory for the sandbox
func (s *Sandbox) LogDir() string {
	return s.logDir
}

// Labels returns the labels associated with the sandbox
func (s *Sandbox) Labels() fields.Set {
	return s.labels
}

// Annotations returns a list of annotations for the sandbox
func (s *Sandbox) Annotations() map[string]string {
	return s.annotations
}

// InfraContainer returns the infrastructure container for the sandbox
func (s *Sandbox) InfraContainer() *oci.Container {
	return s.infraContainer
}

// Containers returns the ContainerStorer that contains information on all
// of the containers in the sandbox
func (s *Sandbox) Containers() oci.ContainerStorer {
	return s.containers
}

// ProcessLabel returns the process label for the sandbox
func (s *Sandbox) ProcessLabel() string {
	return s.processLabel
}

// MountLabel returns the mount label for the sandbox
func (s *Sandbox) MountLabel() string {
	return s.mountLabel
}

// Metadata returns a set of metadata about the sandbox
func (s *Sandbox) Metadata() *pb.PodSandboxMetadata {
	return s.metadata
}

// ShmPath returns the shm path of the sandbox
func (s *Sandbox) ShmPath() string {
	return s.shmPath
}

// CgroupParent returns the cgroup parent of the sandbox
func (s *Sandbox) CgroupParent() string {
	return s.cgroupParent
}

// Privileged returns whether or not the containers in the sandbox are
// privileged containers
func (s *Sandbox) Privileged() bool {
	return s.privileged
}

// Trusted returns whether or not the containers in the sandbox are trusted
func (s *Sandbox) Trusted() bool {
	return s.trusted
}

// ResolvPath returns the resolv path for the sandbox
func (s *Sandbox) ResolvPath() string {
	return s.resolvPath
}

// AddHostnamePath adds the hostname path to the sandbox
func (s *Sandbox) AddHostnamePath(hostname string) {
	s.hostnamePath = hostname
}

// HostnamePath retrieves the hostname path from a sandbox
func (s *Sandbox) HostnamePath() string {
	return s.hostnamePath
}

// Hostname returns the hsotname of the sandbox
func (s *Sandbox) Hostname() string {
	return s.hostname
}

// PortMappings returns a list of port mappings between the host and the sandbox
func (s *Sandbox) PortMappings() []*hostport.PortMapping {
	return s.portMappings
}

// AddContainer adds a container to the sandbox
func (s *Sandbox) AddContainer(c *oci.Container) {
	s.containers.Add(c.Name(), c)
}

// GetContainer retrieves a container from the sandbox
func (s *Sandbox) GetContainer(name string) *oci.Container {
	return s.containers.Get(name)
}

// RemoveContainer deletes a container from the sandbox
func (s *Sandbox) RemoveContainer(c *oci.Container) {
	s.containers.Delete(c.Name())
}

// SetInfraContainer sets the infrastructure container of a sandbox
// Attempts to set the infrastructure container after one is already present will throw an error
func (s *Sandbox) SetInfraContainer(infraCtr *oci.Container) error {
	if s.infraContainer != nil {
		return fmt.Errorf("sandbox already has an infra container")
	} else if infraCtr == nil {
		return fmt.Errorf("must provide non-nil infra container")
	}

	s.infraContainer = infraCtr

	return nil
}

// RemoveInfraContainer removes the infrastructure container of a sandbox
func (s *Sandbox) RemoveInfraContainer() {
	s.infraContainer = nil
}

// NetNs retrieves the network namespace of the sandbox
// If the sandbox uses the host namespace, nil is returned
func (s *Sandbox) NetNs() ns.NetNS {
	if s.netns == nil {
		return nil
	}

	return s.netns.ns
}

// NetNsPath returns the path to the network namespace of the sandbox.
// If the sandbox uses the host namespace, nil is returned
func (s *Sandbox) NetNsPath() string {
	if s.netns == nil {
		return ""
	}

	return s.netns.symlink.Name()
}

// NetNsCreate creates a new network namespace for the sandbox
func (s *Sandbox) NetNsCreate() error {
	if s.netns != nil {
		return fmt.Errorf("net NS already created")
	}

	netNS, err := ns.NewNS()
	if err != nil {
		return err
	}

	s.netns = &NetNs{
		ns:     netNS,
		closed: false,
	}

	if err := s.netns.symlinkCreate(s.name); err != nil {
		logrus.Warnf("Could not create nentns symlink %v", err)

		if err1 := s.netns.ns.Close(); err1 != nil {
			return err1
		}

		return err
	}

	return nil
}

// SetStopped sets the sandbox state to stopped.
// This should be set after a stop operation succeeds
// so that subsequent stops can return fast.
func (s *Sandbox) SetStopped() {
	s.stopped = true
}

// Stopped returns whether the sandbox state has been
// set to stopped.
func (s *Sandbox) Stopped() bool {
	return s.stopped
}

// NetNsJoin attempts to join the sandbox to an existing network namespace
// This will fail if the sandbox is already part of a network namespace
func (s *Sandbox) NetNsJoin(nspath, name string) error {
	if s.netns != nil {
		return fmt.Errorf("sandbox already has a network namespace, cannot join another")
	}

	netNS, err := NetNsGet(nspath, name)
	if err != nil {
		return err
	}

	s.netns = netNS

	return nil
}

// NetNsRemove removes the network namespace associated with the sandbox
func (s *Sandbox) NetNsRemove() error {
	if s.netns == nil {
		logrus.Warn("no networking namespace")
		return nil
	}

	s.netns.Lock()
	defer s.netns.Unlock()

	if s.netns.closed {
		// netNsRemove() can be called multiple
		// times without returning an error.
		return nil
	}

	if err := s.netns.symlinkRemove(); err != nil {
		return err
	}

	if err := s.netns.ns.Close(); err != nil {
		return err
	}

	s.netns.closed = true

	if s.netns.restored {
		// we got namespaces in the form of
		// /var/run/netns/cni-0d08effa-06eb-a963-f51a-e2b0eceffc5d
		// but /var/run on most system is symlinked to /run so we first resolve
		// the symlink and then try and see if it's mounted
		fp, err := symlink.FollowSymlinkInScope(s.netns.ns.Path(), "/")
		if err != nil {
			return err
		}
		if mounted, err := mount.Mounted(fp); err == nil && mounted {
			if err := unix.Unmount(fp, unix.MNT_DETACH); err != nil {
				return err
			}
		}

		if err := os.RemoveAll(s.netns.ns.Path()); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package repr

import (
	"testing"

	"github.com/alecthomas/assert"
)

type anotherStruct struct {
	A []int
}

type testStruct struct {
	S string
	I *int
	A anotherStruct
}

func TestRepr(t *testing.T) {
	assert.Equal(t, "[]string{\n  \"a\",\n  \"b\",\n}", Repr([]string{"a", "b"}))
	assert.Equal(t, "[]int{\n  1,\n  2,\n}", Repr([]int{1, 2}))
	pi := new(int)
	*pi = 13
	assert.Equal(t, `&13`, Repr(pi))

	ch := make(<-chan map[string]*testStruct, 1)
	assert.Equal(t, `make(<-chan map[string]*repr.testStruct, 1)`, Repr(ch))

	m := map[string]int{"a": 1}
	assert.Equal(t, "map[string]int{\n  \"a\": 1,\n}", Repr(m))

	s := &testStruct{
		S: "String",
		I: pi,
		A: anotherStruct{
			A: []int{1, 2, 3},
		},
	}
	assert.Equal(t, `&repr.testStruct{
  S: "String",
  I: &13,
  A: repr.anotherStruct{
    A: []int{
      1,
      2,
      3,
    },
  },
}`, Repr(s))

	b := []uint8{1, 2, 3}
	assert.Equal(t, `[]uint8{
  1,
  2,
  3,
}`, Repr(b))
}
<commit_msg>Use testify/assert to avoid cyclical import.<commit_after>package repr

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

type anotherStruct struct {
	A []int
}

type testStruct struct {
	S string
	I *int
	A anotherStruct
}

func TestRepr(t *testing.T) {
	assert.Equal(t, "[]string{\n  \"a\",\n  \"b\",\n}", Repr([]string{"a", "b"}))
	assert.Equal(t, "[]int{\n  1,\n  2,\n}", Repr([]int{1, 2}))
	pi := new(int)
	*pi = 13
	assert.Equal(t, `&13`, Repr(pi))

	ch := make(<-chan map[string]*testStruct, 1)
	assert.Equal(t, `make(<-chan map[string]*repr.testStruct, 1)`, Repr(ch))

	m := map[string]int{"a": 1}
	assert.Equal(t, "map[string]int{\n  \"a\": 1,\n}", Repr(m))

	s := &testStruct{
		S: "String",
		I: pi,
		A: anotherStruct{
			A: []int{1, 2, 3},
		},
	}
	assert.Equal(t, `&repr.testStruct{
  S: "String",
  I: &13,
  A: repr.anotherStruct{
    A: []int{
      1,
      2,
      3,
    },
  },
}`, Repr(s))

	b := []uint8{1, 2, 3}
	assert.Equal(t, `[]uint8{
  1,
  2,
  3,
}`, Repr(b))
}
<|endoftext|>


<commit_before>package rules

import (
	"fmt"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/skia-dev/glog"
	"github.com/skia-dev/influxdb/client"
	"go.skia.org/infra/alertserver/go/alerting"
	"golang.org/x/tools/go/exact"
	"golang.org/x/tools/go/types"
)

/*
	Rules for triggering alerts.
*/

// Rule is an object used for triggering Alerts.
type Rule struct {
	Name        string        `json:"name"`
	Query       string        `json:"query"`
	Category    string        `json:"category"`
	Condition   string        `json:"condition"`
	Message     string        `json:"message"`
	Nag         time.Duration `json:"nag"`
	client      queryable
	AutoDismiss bool `json:"autoDismiss"`
	Actions     []string
}

// Fire causes the Alert to become Active() and not Snoozed(), and causes each
// action to be performed. Active Alerts do not perform new queries.
func (r *Rule) fire(am *alerting.AlertManager, message string) error {
	actions, err := alerting.ParseActions(r.Actions)
	if err != nil {
		return fmt.Errorf("Could not fire alert: %v", err)
	}
	a := alerting.Alert{
		Name:     r.Name,
		Category: r.Category,
		Message:  message,
		Nag:      int64(r.Nag),
		Actions:  actions,
	}
	return am.AddAlert(&a)
}

func (r *Rule) queryExecutionAlert(queryErr error, am *alerting.AlertManager) error {
	actions, err := alerting.ParseActions([]string{"Email(infra-alerts@skia.org)"})
	if err != nil {
		return err
	}
	name := fmt.Sprintf("Failed to execute query (%s)", r.Name)
	if am.ActiveAlert(name) == 0 {
		return am.AddAlert(&alerting.Alert{
			Name:     name,
			Category: r.Category, // Should the category be "internal error" or something?
			Message:  fmt.Sprintf("Failed to execute query for rule \"%s\": [ %s ]\nFull error:\n%v", r.Name, r.Query, queryErr),
			Nag:      int64(1 * time.Hour),
			Actions:  actions,
		})
	}
	return nil
}

func (r *Rule) queryEvaluationAlert(queryErr error, am *alerting.AlertManager) error {
	actions, err := alerting.ParseActions([]string{"Email(infra-alerts@skia.org)"})
	if err != nil {
		return err
	}
	name := fmt.Sprintf("Failed to evaluate query (%s)", r.Name)
	if am.ActiveAlert(name) == 0 {
		return am.AddAlert(&alerting.Alert{
			Name:     name,
			Category: r.Category, // Should the category be "internal error" or something?
			Message:  fmt.Sprintf("Failed to evaluate query for rule \"%s\": [ %s ]\nFull error:\n%v", r.Name, r.Condition, queryErr),
			Nag:      int64(1 * time.Hour),
			Actions:  actions,
		})
	}
	return nil
}

func (r *Rule) tick(am *alerting.AlertManager) error {
	a := am.ActiveAlert(r.Name)
	if a == 0 || r.AutoDismiss {
		d, err := executeQuery(r.client, r.Query)
		if err != nil {
			// We shouldn't fail to execute a query. Trigger an alert.
			return r.queryExecutionAlert(err, am)
		}
		doAlert, err := r.evaluate(d)
		if err != nil {
			return r.queryEvaluationAlert(err, am)
		}
		if a != 0 && r.AutoDismiss && !doAlert {
			return am.Dismiss(a, "AlertServer", "Condition is no longer true.")
		} else if a == 0 && doAlert {
			return r.fire(am, r.Message)
		}
	}
	return nil
}

func (r *Rule) evaluate(d float64) (bool, error) {
	pkg := types.NewPackage("evaluateme", "evaluateme")
	v := exact.MakeFloat64(d)
	pkg.Scope().Insert(types.NewConst(0, pkg, "x", types.Typ[types.Float64], v))
	tv, err := types.Eval(r.Condition, pkg, pkg.Scope())
	if err != nil {
		return false, fmt.Errorf("Failed to evaluate condition %q: %s", r.Condition, err)
	}
	if tv.Type.String() != "untyped bool" {
		return false, fmt.Errorf("Rule \"%v\" does not return boolean type.", r.Condition)
	}
	return exact.BoolVal(tv.Value), nil
}

type parsedRule map[string]interface{}

func newRule(r parsedRule, client *client.Client, testing bool) (*Rule, error) {
	errString := "Alert rule missing field %q"
	name, ok := r["name"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "name")
	}
	query, ok := r["query"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "query")
	}
	condition, ok := r["condition"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "condition")
	}
	message, ok := r["message"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "message")
	}
	AutoDismiss, ok := r["auto-dismiss"].(bool)
	if !ok {
		return nil, fmt.Errorf(errString, "auto-dismiss")
	}
	actionsInterface, ok := r["actions"]
	if !ok {
		return nil, fmt.Errorf(errString, "actions")
	}
	actionsInterfaceList := actionsInterface.([]interface{})
	actionStrings := make([]string, 0, len(actionsInterfaceList))
	for _, iface := range actionsInterfaceList {
		actionStrings = append(actionStrings, iface.(string))
	}
	nagDuration := time.Duration(0)
	nag, ok := r["nag"].(string)
	if ok {
		var err error
		nagDuration, err = time.ParseDuration(nag)
		if err != nil {
			return nil, fmt.Errorf("Invalid nag duration %q: %v", nag, err)
		}
	}
	rule := Rule{
		Name:        name,
		Query:       query,
		Condition:   condition,
		Message:     message,
		Nag:         nagDuration,
		client:      client,
		AutoDismiss: AutoDismiss,
		Actions:     actionStrings,
	}
	// Verify that the condition can be evaluated.
	_, err := rule.evaluate(0.0)
	if err != nil {
		return nil, err
	}
	return &rule, nil
}

func parseAlertRules(cfgFile string) ([]parsedRule, error) {
	var cfg struct {
		Rule []parsedRule
	}
	_, err := toml.DecodeFile(cfgFile, &cfg)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse %s: %s", cfgFile, err)
	}
	return cfg.Rule, nil
}

func MakeRules(cfgFile string, dbClient *client.Client, tickInterval time.Duration, am *alerting.AlertManager, testing bool) ([]*Rule, error) {
	parsedRules, err := parseAlertRules(cfgFile)
	if err != nil {
		return nil, err
	}
	rules := []*Rule{}
	for _, r := range parsedRules {
		r, err := newRule(r, dbClient, testing)
		if err != nil {
			return nil, err
		}
		rules = append(rules, r)
	}

	// Start the goroutines.
	for _, r := range rules {
		go func(rule *Rule) {
			if err := rule.tick(am); err != nil {
				glog.Error(err)
			}
			for _ = range time.Tick(tickInterval) {
				if err := rule.tick(am); err != nil {
					glog.Error(err)
				}
			}
		}(r)
	}

	return rules, nil
}

type queryable interface {
	Query(string, ...client.TimePrecision) ([]*client.Series, error)
}

func executeQuery(c queryable, q string) (float64, error) {
	results, err := c.Query(q)
	if err != nil {
		return 0.0, fmt.Errorf("Failed to query InfluxDB with query %q: %s", q, err)
	}
	if len(results) < 1 {
		return 0.0, fmt.Errorf("Query returned no data: %q", q)
	}
	points := results[0].Points
	if len(points) < 1 {
		return 0.0, fmt.Errorf("Query returned no points: %q", q)
	}
	if len(points) > 1 {
		return 0.0, fmt.Errorf("Query returned more than one point: %q", q)
	}
	valueColumn := 0
	for _, label := range results[0].Columns {
		if label == "time" || label == "sequence_number" {
			valueColumn++
		} else {
			break
		}
	}
	if len(results[0].Columns) != valueColumn+1 {
		return 0.0, fmt.Errorf("Query returned an incorrect set of columns: %q %v", q, results[0].Columns)
	}
	if len(results[0].Columns) != len(points[0]) {
		return 0.0, fmt.Errorf("Invalid data from InfluxDB: Point data does not match column spec.")
	}
	return points[0][valueColumn].(float64), nil
}
<commit_msg>Alertserver: only one "failed to execute query" alert<commit_after>package rules

import (
	"fmt"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/skia-dev/glog"
	"github.com/skia-dev/influxdb/client"
	"go.skia.org/infra/alertserver/go/alerting"
	"golang.org/x/tools/go/exact"
	"golang.org/x/tools/go/types"
)

/*
	Rules for triggering alerts.
*/

// Rule is an object used for triggering Alerts.
type Rule struct {
	Name        string        `json:"name"`
	Query       string        `json:"query"`
	Category    string        `json:"category"`
	Condition   string        `json:"condition"`
	Message     string        `json:"message"`
	Nag         time.Duration `json:"nag"`
	client      queryable
	AutoDismiss bool `json:"autoDismiss"`
	Actions     []string
}

// Fire causes the Alert to become Active() and not Snoozed(), and causes each
// action to be performed. Active Alerts do not perform new queries.
func (r *Rule) fire(am *alerting.AlertManager, message string) error {
	actions, err := alerting.ParseActions(r.Actions)
	if err != nil {
		return fmt.Errorf("Could not fire alert: %v", err)
	}
	a := alerting.Alert{
		Name:     r.Name,
		Category: r.Category,
		Message:  message,
		Nag:      int64(r.Nag),
		Actions:  actions,
	}
	return am.AddAlert(&a)
}

func (r *Rule) queryExecutionAlert(queryErr error, am *alerting.AlertManager) error {
	actions, err := alerting.ParseActions([]string{"Email(infra-alerts@skia.org)"})
	if err != nil {
		return err
	}
	name := "Failed to execute query"
	if am.ActiveAlert(name) == 0 {
		return am.AddAlert(&alerting.Alert{
			Name:     name,
			Category: r.Category, // Should the category be "internal error" or something?
			Message:  fmt.Sprintf("Failed to execute query for rule \"%s\": [ %s ]\nFull error:\n%v", r.Name, r.Query, queryErr),
			Nag:      int64(1 * time.Hour),
			Actions:  actions,
		})
	}
	return nil
}

func (r *Rule) queryEvaluationAlert(queryErr error, am *alerting.AlertManager) error {
	actions, err := alerting.ParseActions([]string{"Email(infra-alerts@skia.org)"})
	if err != nil {
		return err
	}
	name := "Failed to evaluate query"
	if am.ActiveAlert(name) == 0 {
		return am.AddAlert(&alerting.Alert{
			Name:     name,
			Category: r.Category, // Should the category be "internal error" or something?
			Message:  fmt.Sprintf("Failed to evaluate query for rule \"%s\": [ %s ]\nFull error:\n%v", r.Name, r.Condition, queryErr),
			Nag:      int64(1 * time.Hour),
			Actions:  actions,
		})
	}
	return nil
}

func (r *Rule) tick(am *alerting.AlertManager) error {
	a := am.ActiveAlert(r.Name)
	if a == 0 || r.AutoDismiss {
		d, err := executeQuery(r.client, r.Query)
		if err != nil {
			// We shouldn't fail to execute a query. Trigger an alert.
			return r.queryExecutionAlert(err, am)
		}
		doAlert, err := r.evaluate(d)
		if err != nil {
			return r.queryEvaluationAlert(err, am)
		}
		if a != 0 && r.AutoDismiss && !doAlert {
			return am.Dismiss(a, "AlertServer", "Condition is no longer true.")
		} else if a == 0 && doAlert {
			return r.fire(am, r.Message)
		}
	}
	return nil
}

func (r *Rule) evaluate(d float64) (bool, error) {
	pkg := types.NewPackage("evaluateme", "evaluateme")
	v := exact.MakeFloat64(d)
	pkg.Scope().Insert(types.NewConst(0, pkg, "x", types.Typ[types.Float64], v))
	tv, err := types.Eval(r.Condition, pkg, pkg.Scope())
	if err != nil {
		return false, fmt.Errorf("Failed to evaluate condition %q: %s", r.Condition, err)
	}
	if tv.Type.String() != "untyped bool" {
		return false, fmt.Errorf("Rule \"%v\" does not return boolean type.", r.Condition)
	}
	return exact.BoolVal(tv.Value), nil
}

type parsedRule map[string]interface{}

func newRule(r parsedRule, client *client.Client, testing bool) (*Rule, error) {
	errString := "Alert rule missing field %q"
	name, ok := r["name"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "name")
	}
	query, ok := r["query"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "query")
	}
	condition, ok := r["condition"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "condition")
	}
	message, ok := r["message"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "message")
	}
	AutoDismiss, ok := r["auto-dismiss"].(bool)
	if !ok {
		return nil, fmt.Errorf(errString, "auto-dismiss")
	}
	actionsInterface, ok := r["actions"]
	if !ok {
		return nil, fmt.Errorf(errString, "actions")
	}
	actionsInterfaceList := actionsInterface.([]interface{})
	actionStrings := make([]string, 0, len(actionsInterfaceList))
	for _, iface := range actionsInterfaceList {
		actionStrings = append(actionStrings, iface.(string))
	}
	nagDuration := time.Duration(0)
	nag, ok := r["nag"].(string)
	if ok {
		var err error
		nagDuration, err = time.ParseDuration(nag)
		if err != nil {
			return nil, fmt.Errorf("Invalid nag duration %q: %v", nag, err)
		}
	}
	rule := Rule{
		Name:        name,
		Query:       query,
		Condition:   condition,
		Message:     message,
		Nag:         nagDuration,
		client:      client,
		AutoDismiss: AutoDismiss,
		Actions:     actionStrings,
	}
	// Verify that the condition can be evaluated.
	_, err := rule.evaluate(0.0)
	if err != nil {
		return nil, err
	}
	return &rule, nil
}

func parseAlertRules(cfgFile string) ([]parsedRule, error) {
	var cfg struct {
		Rule []parsedRule
	}
	_, err := toml.DecodeFile(cfgFile, &cfg)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse %s: %s", cfgFile, err)
	}
	return cfg.Rule, nil
}

func MakeRules(cfgFile string, dbClient *client.Client, tickInterval time.Duration, am *alerting.AlertManager, testing bool) ([]*Rule, error) {
	parsedRules, err := parseAlertRules(cfgFile)
	if err != nil {
		return nil, err
	}
	rules := []*Rule{}
	for _, r := range parsedRules {
		r, err := newRule(r, dbClient, testing)
		if err != nil {
			return nil, err
		}
		rules = append(rules, r)
	}

	// Start the goroutines.
	for _, r := range rules {
		go func(rule *Rule) {
			if err := rule.tick(am); err != nil {
				glog.Error(err)
			}
			for _ = range time.Tick(tickInterval) {
				if err := rule.tick(am); err != nil {
					glog.Error(err)
				}
			}
		}(r)
	}

	return rules, nil
}

type queryable interface {
	Query(string, ...client.TimePrecision) ([]*client.Series, error)
}

func executeQuery(c queryable, q string) (float64, error) {
	results, err := c.Query(q)
	if err != nil {
		return 0.0, fmt.Errorf("Failed to query InfluxDB with query %q: %s", q, err)
	}
	if len(results) < 1 {
		return 0.0, fmt.Errorf("Query returned no data: %q", q)
	}
	points := results[0].Points
	if len(points) < 1 {
		return 0.0, fmt.Errorf("Query returned no points: %q", q)
	}
	if len(points) > 1 {
		return 0.0, fmt.Errorf("Query returned more than one point: %q", q)
	}
	valueColumn := 0
	for _, label := range results[0].Columns {
		if label == "time" || label == "sequence_number" {
			valueColumn++
		} else {
			break
		}
	}
	if len(results[0].Columns) != valueColumn+1 {
		return 0.0, fmt.Errorf("Query returned an incorrect set of columns: %q %v", q, results[0].Columns)
	}
	if len(results[0].Columns) != len(points[0]) {
		return 0.0, fmt.Errorf("Invalid data from InfluxDB: Point data does not match column spec.")
	}
	return points[0][valueColumn].(float64), nil
}
<|endoftext|>


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

package server

import (
	"crypto/rand"
	"fmt"
	"github.com/uniqush/uniqush-conn/msgcache"
	"github.com/uniqush/uniqush-conn/proto"
	"github.com/uniqush/uniqush-conn/proto/client"
	"io"
	"sync"
	"testing"
	"time"
)

func sendTestMessages(s2c, c2s proto.Conn, serverToClient bool, msgs ...*proto.Message) error {
	wg := new(sync.WaitGroup)
	wg.Add(2)

	var src proto.Conn
	var dst proto.Conn

	if serverToClient {
		src = s2c
		dst = c2s
	} else {
		src = c2s
		dst = s2c
	}

	var es error
	var ed error

	go func() {
		defer wg.Done()
		for _, msg := range msgs {
			es = src.WriteMessage(msg, true, true)
			if es != nil {
				return
			}
		}
	}()

	go func() {
		defer wg.Done()
		var m *proto.Message
		for _, msg := range msgs {
			m, ed = dst.ReadMessage()
			if ed != nil {
				return
			}
			if !m.Eq(msg) {
				ed = fmt.Errorf("corrupted data")
				return
			}
		}
	}()
	wg.Wait()
	if es != nil {
		return es
	}
	if ed != nil {
		return ed
	}
	return nil
}

func randomMessage() *proto.Message {
	msg := new(proto.Message)
	msg.Body = make([]byte, 10)
	io.ReadFull(rand.Reader, msg.Body)
	msg.Header = make(map[string]string, 2)
	msg.Header["aaa"] = "hello"
	msg.Header["aa"] = "hell"
	return msg
}

func TestMessageSendServerToClient(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	N := 100
	msgs := make([]*proto.Message, N)

	for i := 0; i < N; i++ {
		msgs[i] = randomMessage()
	}

	err = sendTestMessages(servConn, cliConn, true, msgs...)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	if servConn != nil {
		servConn.Close()
	}
	if cliConn != nil {
		cliConn.Close()
	}
}

func TestMessageSendClientToServer(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	N := 100
	msgs := make([]*proto.Message, N)

	for i := 0; i < N; i++ {
		msgs[i] = randomMessage()
	}

	err = sendTestMessages(servConn, cliConn, false, msgs...)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	if servConn != nil {
		servConn.Close()
	}
	if cliConn != nil {
		cliConn.Close()
	}
}

func TestDigestSetting(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	defer servConn.Close()
	defer cliConn.Close()

	// We always want to receive digest
	err = cliConn.Config(0, 512, true, nil)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	// Wait it to be effect
	time.Sleep(1 * time.Second)
	mcache := msgcache.NewRedisMessageCache("", "", 1)
	servConn.SetMessageCache(mcache)
	diChan := make(chan *client.Digest)
	cliConn.SetDigestChannel(diChan)
	msg := randomMessage()

	wg := new(sync.WaitGroup)
	wg.Add(2)

	reqOK := make(chan bool)
	// Server:
	go func() {
		err := servConn.SendOrBox(msg, nil, 0*time.Second)
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		reqOK <- true
		wg.Done()
	}()

	// Client:
	go func() {
		digest := <-diChan
		if nil == digest {
			t.Errorf("Error: Empty digest")
		}
		<-reqOK
		cliConn.RequestMessage(digest.MsgId)
		m, err := cliConn.ReadMessage()
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		m.Id = ""
		if !msg.Eq(m) {
			t.Errorf("Error: should same: %v != %v", msg, m)
		}
		wg.Done()
	}()
	wg.Wait()
}

func TestDigestSettingWithFields(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	defer servConn.Close()
	defer cliConn.Close()

	fields := []string{"digest"}
	// We always want to receive digest
	err = cliConn.Config(0, 512, true, fields)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	// Wait it to be effect
	time.Sleep(1 * time.Second)
	mcache := msgcache.NewRedisMessageCache("", "", 1)
	servConn.SetMessageCache(mcache)
	diChan := make(chan *client.Digest)
	cliConn.SetDigestChannel(diChan)
	msg := randomMessage()
	msg.Header[fields[0]] = "new"

	wg := new(sync.WaitGroup)
	wg.Add(2)

	reqOK := make(chan bool)
	// Server:
	go func() {
		err := servConn.SendOrBox(msg, nil, 0*time.Second)
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		reqOK <- true
		wg.Done()
	}()

	// Client:
	go func() {
		digest := <-diChan
		if nil == digest {
			t.Errorf("Error: Empty digest")
		}
		if digest.Info[fields[0]] != msg.Header[fields[0]] {
			t.Errorf("Error: field not match")
		}
		<-reqOK
		cliConn.RequestMessage(digest.MsgId)
		m, err := cliConn.ReadMessage()
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		m.Id = ""
		if !msg.Eq(m) {
			t.Errorf("Error: should same: %v != %v", msg, m)
		}
		wg.Done()
	}()
	wg.Wait()
}

func TestDigestSettingWithMessageQueue(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	defer servConn.Close()
	defer cliConn.Close()

	// We always want to receive digest
	err = cliConn.Config(0, 512, true, nil)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	// Wait it to be effect
	time.Sleep(1 * time.Second)
	mcache := msgcache.NewRedisMessageCache("", "", 1)
	mcache.Clrqueue("service", "username")
	servConn.SetMessageCache(mcache)
	diChan := make(chan *client.Digest)
	cliConn.SetDigestChannel(diChan)
	msg := randomMessage()

	wg := new(sync.WaitGroup)
	wg.Add(2)

	reqOK := make(chan bool)
	// Server:
	go func() {
		_, err := servConn.SendOrQueue(msg, nil)
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		reqOK <- true
		wg.Done()
	}()

	// Client:
	go func() {
		digest := <-diChan
		if nil == digest {
			t.Errorf("Error: Empty digest")
		}
		<-reqOK
		cliConn.RequestMessage(digest.MsgId)
		m, err := cliConn.ReadMessage()
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		m.Id = ""
		if !msg.Eq(m) {
			t.Errorf("Error: should same: %v != %v", msg, m)
		}
		wg.Done()
	}()
	wg.Wait()
}
<commit_msg>Changed the test cases.<commit_after>/*
 * Copyright 2012 Nan Deng
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package server

import (
	"crypto/rand"
	"fmt"
	"github.com/uniqush/uniqush-conn/msgcache"
	"github.com/uniqush/uniqush-conn/proto"
	"github.com/uniqush/uniqush-conn/proto/client"
	"io"
	"sync"
	"testing"
	"time"
)

func sendTestMessages(s2c, c2s proto.Conn, serverToClient bool, msgs ...*proto.Message) error {
	wg := new(sync.WaitGroup)
	wg.Add(2)

	var src proto.Conn
	var dst proto.Conn

	if serverToClient {
		src = s2c
		dst = c2s
	} else {
		src = c2s
		dst = s2c
	}

	var es error
	var ed error

	go func() {
		defer wg.Done()
		for _, msg := range msgs {
			es = src.WriteMessage(msg, true, true)
			if es != nil {
				return
			}
		}
	}()

	go func() {
		defer wg.Done()
		var m *proto.Message
		for _, msg := range msgs {
			m, ed = dst.ReadMessage()
			if ed != nil {
				return
			}
			if !m.Eq(msg) {
				ed = fmt.Errorf("corrupted data")
				return
			}
		}
	}()
	wg.Wait()
	if es != nil {
		return es
	}
	if ed != nil {
		return ed
	}
	return nil
}

func randomMessage() *proto.Message {
	msg := new(proto.Message)
	msg.Body = make([]byte, 10)
	io.ReadFull(rand.Reader, msg.Body)
	msg.Header = make(map[string]string, 2)
	msg.Header["aaa"] = "hello"
	msg.Header["aa"] = "hell"
	return msg
}

func TestMessageSendServerToClient(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	N := 100
	msgs := make([]*proto.Message, N)

	for i := 0; i < N; i++ {
		msgs[i] = randomMessage()
	}

	err = sendTestMessages(servConn, cliConn, true, msgs...)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	if servConn != nil {
		servConn.Close()
	}
	if cliConn != nil {
		cliConn.Close()
	}
}

func TestMessageSendClientToServer(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	N := 100
	msgs := make([]*proto.Message, N)

	for i := 0; i < N; i++ {
		msgs[i] = randomMessage()
	}

	err = sendTestMessages(servConn, cliConn, false, msgs...)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	if servConn != nil {
		servConn.Close()
	}
	if cliConn != nil {
		cliConn.Close()
	}
}

func TestDigestSetting(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	defer servConn.Close()
	defer cliConn.Close()

	// We always want to receive digest
	err = cliConn.Config(0, 512, true, nil)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	// Wait it to be effect
	time.Sleep(1 * time.Second)
	mcache := msgcache.NewRedisMessageCache("", "", 1)
	servConn.SetMessageCache(mcache)
	diChan := make(chan *client.Digest)
	cliConn.SetDigestChannel(diChan)
	msg := randomMessage()

	wg := new(sync.WaitGroup)
	wg.Add(2)

	// Server:
	go func() {
		err := servConn.SendOrBox(msg, nil, 0*time.Second)
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		wg.Done()
	}()

	// Client:
	go func() {
		digest := <-diChan
		if nil == digest {
			t.Errorf("Error: Empty digest")
		}
		cliConn.RequestMessage(digest.MsgId)
		m, err := cliConn.ReadMessage()
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		m.Id = ""
		if !msg.Eq(m) {
			t.Errorf("Error: should same: %v != %v", msg, m)
		}
		wg.Done()
	}()
	wg.Wait()
}

func TestDigestSettingWithFields(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	defer servConn.Close()
	defer cliConn.Close()

	fields := []string{"digest"}
	// We always want to receive digest
	err = cliConn.Config(0, 512, true, fields)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	// Wait it to be effect
	time.Sleep(1 * time.Second)
	mcache := msgcache.NewRedisMessageCache("", "", 1)
	servConn.SetMessageCache(mcache)
	diChan := make(chan *client.Digest)
	cliConn.SetDigestChannel(diChan)
	msg := randomMessage()
	msg.Header[fields[0]] = "new"

	wg := new(sync.WaitGroup)
	wg.Add(2)

	// Server:
	go func() {
		err := servConn.SendOrBox(msg, nil, 0*time.Second)
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		wg.Done()
	}()

	// Client:
	go func() {
		digest := <-diChan
		if nil == digest {
			t.Errorf("Error: Empty digest")
		}
		if digest.Info[fields[0]] != msg.Header[fields[0]] {
			t.Errorf("Error: field not match")
		}
		cliConn.RequestMessage(digest.MsgId)
		m, err := cliConn.ReadMessage()
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		m.Id = ""
		if !msg.Eq(m) {
			t.Errorf("Error: should same: %v != %v", msg, m)
		}
		wg.Done()
	}()
	wg.Wait()
}

func TestDigestSettingWithMessageQueue(t *testing.T) {
	addr := "127.0.0.1:8088"
	token := "token"
	servConn, cliConn, err := buildServerClientConns(addr, token, 3*time.Second)
	defer servConn.Close()
	defer cliConn.Close()

	// We always want to receive digest
	err = cliConn.Config(0, 512, true, nil)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
	// Wait it to be effect
	time.Sleep(1 * time.Second)
	mcache := msgcache.NewRedisMessageCache("", "", 1)
	mcache.Clrqueue("service", "username")
	servConn.SetMessageCache(mcache)
	diChan := make(chan *client.Digest)
	cliConn.SetDigestChannel(diChan)
	msg := randomMessage()

	wg := new(sync.WaitGroup)
	wg.Add(2)

	// Server:
	go func() {
		_, err := servConn.SendOrQueue(msg, nil)
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		wg.Done()
	}()

	// Client:
	go func() {
		digest := <-diChan
		if nil == digest {
			t.Errorf("Error: Empty digest")
		}
		cliConn.RequestMessage(digest.MsgId)
		m, err := cliConn.ReadMessage()
		if err != nil {
			t.Errorf("Error: %v", err)
		}
		m.Id = ""
		if !msg.Eq(m) {
			t.Errorf("Error: should same: %v != %v", msg, m)
		}
		wg.Done()
	}()
	wg.Wait()
}
<|endoftext|>


<commit_before>package validation

import (
	"encoding/hex"
	"reflect"
	"strings"

	"chain/errors"
	"chain/protocol/bc"
	"chain/protocol/state"
	"chain/protocol/vm"
	"chain/protocol/vmutil"
)

var (
	// ErrBadTx is returned for transactions failing validation
	ErrBadTx = errors.New("invalid transaction")

	// ErrFalseVMResult is one of the ways for a transaction to fail validation
	ErrFalseVMResult = errors.New("false VM result")
)

// ConfirmTx validates the given transaction against the given state tree
// before it's added to a block. If tx is invalid, it returns a non-nil
// error describing why.
//
// Tx should have already been validated (with `ValidateTx`) when the tx
// was added to the pool.
func ConfirmTx(snapshot *state.Snapshot, tx *bc.Tx, timestampMS uint64) error {
	if timestampMS < tx.MinTime {
		return errors.WithDetail(ErrBadTx, "block time is before transaction min time")
	}
	if tx.MaxTime > 0 && timestampMS > tx.MaxTime {
		return errors.WithDetail(ErrBadTx, "block time is after transaction max time")
	}
	for inIndex, txin := range tx.Inputs {
		if ic, ok := txin.InputCommitment.(*bc.IssuanceInputCommitment); ok {
			// txin is an issuance

			if inIndex == 0 {
				if ic.MinTimeMS > timestampMS || ic.MaxTimeMS < timestampMS {
					return errors.WithDetailf(ErrBadTx, "block time is outside issuance time window (input %d)", inIndex)
				}
				for txhash, expireMS := range snapshot.Issuances {
					if timestampMS > expireMS {
						delete(snapshot.Issuances, txhash)
					}
				}
				if _, ok2 := snapshot.Issuances[tx.Hash]; ok2 {
					return errors.WithDetail(ErrBadTx, "duplicate issuance transaction")
				}
			}
			continue
		}

		// txin is a spend

		// Lookup the prevout in the blockchain state tree.
		k, val := state.OutputTreeItem(state.Prevout(txin))
		n := snapshot.Tree.Lookup(k)
		if n == nil || n.Hash() != val.Hash() {
			return errors.WithDetailf(ErrBadTx, "output %s for input %d is invalid", txin.Outpoint().String(), inIndex)
		}
	}
	return nil
}

// ValidateTx checks whether tx passes context-free validation:
// - inputs and outputs balance
// - no duplicate prevouts
// - input scripts pass
//
// If tx is well formed and valid, it returns a nil error; otherwise, it
// returns an error describing why tx is invalid.
func ValidateTx(tx *bc.Tx) error {
	if len(tx.Inputs) == 0 {
		return errors.WithDetail(ErrBadTx, "inputs are missing")
	}

	if tx.MaxTime > 0 && tx.MaxTime < tx.MinTime {
		return errors.WithDetail(ErrBadTx, "positive maxtime must be >= mintime")
	}

	issued := make(map[bc.AssetID]bool)
	parity := make(map[bc.AssetID]int64)

	for i, txin := range tx.Inputs {
		assetID := txin.AssetID()
		parity[assetID] += int64(txin.Amount())
		if txin.IsIssuance() {
			issued[assetID] = true
			if i == 0 {
				ic := txin.InputCommitment.(*bc.IssuanceInputCommitment)
				if ic.MaxTimeMS < ic.MinTimeMS {
					return errors.WithDetail(ErrBadTx, "input 0 is an issuance with maxtime < mintime")
				}
				// TODO(bobg): test that issuance maxtime - issuance mintime <= issuance window limit
			}
		}

		for j := 0; j < i; j++ {
			other := tx.Inputs[j]
			if reflect.DeepEqual(txin.InputCommitment, other.InputCommitment) {
				return errors.WithDetailf(ErrBadTx, "input %d is a duplicate of %d", j, i)
			}
		}
	}

	// Check that every output has a valid value.
	for _, txout := range tx.Outputs {
		// Transactions cannot have zero-value outputs.
		// If all inputs have zero value, tx therefore must have no outputs.
		if txout.Amount == 0 {
			return errors.WithDetailf(ErrBadTx, "output value must be greater than 0")
		}
		parity[txout.AssetID] -= int64(txout.Amount)
	}

	for asset, val := range parity {
		if val != 0 {
			return errors.WithDetailf(ErrBadTx, "amounts for asset %s are not balanced on inputs and outputs", asset)
		}
	}

	for i := range tx.Inputs {
		ok, err := vm.VerifyTxInput(tx, uint32(i))
		if err == nil && !ok {
			err = ErrFalseVMResult
		}
		if err != nil {
			input := tx.Inputs[i]
			var program []byte
			if input.IsIssuance() {
				program = input.IssuanceProgram()
			} else {
				program = input.ControlProgram()
			}
			scriptStr, _ := vm.Decompile(program)
			hexArgs := make([]string, 0, len(input.InputWitness))
			for _, arg := range input.InputWitness {
				hexArgs = append(hexArgs, hex.EncodeToString(arg))
			}
			return errors.WithDetailf(ErrBadTx, "validation failed in script execution, input %d (program [%s] args [%s]): %s", i, scriptStr, strings.Join(hexArgs, " "), err)
		}
	}
	return nil
}

// ApplyTx updates the state tree with all the changes to the ledger.
func ApplyTx(snapshot *state.Snapshot, tx *bc.Tx) error {
	for i, in := range tx.Inputs {
		if ic, ok := in.InputCommitment.(*bc.IssuanceInputCommitment); ok {
			// issuance input
			if i == 0 {
				snapshot.Issuances[tx.Hash] = ic.MaxTimeMS
			}
			continue
		}

		// Remove the consumed output from the state tree.
		prevoutKey, _ := state.OutputTreeItem(state.Prevout(in))
		err := snapshot.Tree.Delete(prevoutKey)
		if err != nil {
			return err
		}
	}

	for i, out := range tx.Outputs {
		if vmutil.IsUnspendable(out.ControlProgram) {
			continue
		}
		// Insert new outputs into the state tree.
		o := state.NewOutput(*out, bc.Outpoint{Hash: tx.Hash, Index: uint32(i)})
		err := snapshot.Tree.Insert(state.OutputTreeItem(o))
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>protocol/validation: check for int64 overflow when verifying tx amounts<commit_after>package validation

import (
	"encoding/hex"
	"math"
	"reflect"
	"strings"

	"chain/errors"
	"chain/protocol/bc"
	"chain/protocol/state"
	"chain/protocol/vm"
	"chain/protocol/vmutil"
)

var (
	// ErrBadTx is returned for transactions failing validation
	ErrBadTx = errors.New("invalid transaction")

	// ErrFalseVMResult is one of the ways for a transaction to fail validation
	ErrFalseVMResult = errors.New("false VM result")
)

// ConfirmTx validates the given transaction against the given state tree
// before it's added to a block. If tx is invalid, it returns a non-nil
// error describing why.
//
// Tx should have already been validated (with `ValidateTx`) when the tx
// was added to the pool.
func ConfirmTx(snapshot *state.Snapshot, tx *bc.Tx, timestampMS uint64) error {
	if timestampMS < tx.MinTime {
		return errors.WithDetail(ErrBadTx, "block time is before transaction min time")
	}
	if tx.MaxTime > 0 && timestampMS > tx.MaxTime {
		return errors.WithDetail(ErrBadTx, "block time is after transaction max time")
	}
	for inIndex, txin := range tx.Inputs {
		if ic, ok := txin.InputCommitment.(*bc.IssuanceInputCommitment); ok {
			// txin is an issuance

			if inIndex == 0 {
				if ic.MinTimeMS > timestampMS || ic.MaxTimeMS < timestampMS {
					return errors.WithDetailf(ErrBadTx, "block time is outside issuance time window (input %d)", inIndex)
				}
				for txhash, expireMS := range snapshot.Issuances {
					if timestampMS > expireMS {
						delete(snapshot.Issuances, txhash)
					}
				}
				if _, ok2 := snapshot.Issuances[tx.Hash]; ok2 {
					return errors.WithDetail(ErrBadTx, "duplicate issuance transaction")
				}
			}
			continue
		}

		// txin is a spend

		// Lookup the prevout in the blockchain state tree.
		k, val := state.OutputTreeItem(state.Prevout(txin))
		n := snapshot.Tree.Lookup(k)
		if n == nil || n.Hash() != val.Hash() {
			return errors.WithDetailf(ErrBadTx, "output %s for input %d is invalid", txin.Outpoint().String(), inIndex)
		}
	}
	return nil
}

// ValidateTx checks whether tx passes context-free validation:
// - inputs and outputs balance
// - no duplicate input commitments
// - input scripts pass
//
// If tx is well formed and valid, it returns a nil error; otherwise, it
// returns an error describing why tx is invalid.
func ValidateTx(tx *bc.Tx) error {
	if len(tx.Inputs) == 0 {
		return errors.WithDetail(ErrBadTx, "inputs are missing")
	}

	// Check that the transaction maximum time is greater than or equal to the
	// minimum time, if it is greater than 0.
	if tx.MaxTime > 0 && tx.MaxTime < tx.MinTime {
		return errors.WithDetail(ErrBadTx, "positive maxtime must be >= mintime")
	}

	// Check that each input commitment appears only once. Also check that sums
	// of inputs and outputs balance, and check that both input and output sums
	// are less than 2^63 so that they don't overflow their int64 representation.
	issued := make(map[bc.AssetID]bool)
	parityIns := make(map[bc.AssetID]int64)
	parityOuts := make(map[bc.AssetID]int64)

	for i, txin := range tx.Inputs {
		assetID := txin.AssetID()

		if txin.Amount() > math.MaxInt64 {
			return errors.WithDetail(ErrBadTx, "input value exceeds maximum value of int64")
		}

		sum, err := safeSum(parityIns[assetID], int64(txin.Amount()))
		if err != nil {
			return errors.WithDetailf(ErrBadTx, "adding input %d overflows the allowed asset amount", i)
		}
		parityIns[assetID] = sum

		if txin.IsIssuance() {
			issued[assetID] = true
			if i == 0 {
				ic := txin.InputCommitment.(*bc.IssuanceInputCommitment)
				if ic.MaxTimeMS < ic.MinTimeMS {
					return errors.WithDetail(ErrBadTx, "input 0 is an issuance with maxtime < mintime")
				}
				// TODO(bobg): test that issuance maxtime - issuance mintime <= issuance window limit
			}
		}

		for j := 0; j < i; j++ {
			other := tx.Inputs[j]
			if reflect.DeepEqual(txin.InputCommitment, other.InputCommitment) {
				return errors.WithDetailf(ErrBadTx, "input %d is a duplicate of %d", j, i)
			}
		}
	}

	// Check that every output has a valid value.
	for i, txout := range tx.Outputs {
		// Transactions cannot have zero-value outputs.
		// If all inputs have zero value, tx therefore must have no outputs.
		if txout.Amount == 0 {
			return errors.WithDetail(ErrBadTx, "output value must be greater than 0")
		}

		if txout.Amount > math.MaxInt64 {
			return errors.WithDetail(ErrBadTx, "output value exceeds maximum value of int64")
		}

		sum, err := safeSum(parityOuts[txout.AssetID], int64(txout.Amount))
		if err != nil {
			return errors.WithDetailf(ErrBadTx, "adding output %d overflows the allowed asset amount", i)
		}
		parityOuts[txout.AssetID] = sum
	}

	for asset, val := range parityIns {
		if val != parityOuts[asset] {
			return errors.WithDetailf(ErrBadTx, "amounts for asset %s are not balanced on inputs and outputs", asset)
		}
	}

	for i := range tx.Inputs {
		ok, err := vm.VerifyTxInput(tx, uint32(i))
		if err == nil && !ok {
			err = ErrFalseVMResult
		}
		if err != nil {
			input := tx.Inputs[i]
			var program []byte
			if input.IsIssuance() {
				program = input.IssuanceProgram()
			} else {
				program = input.ControlProgram()
			}
			scriptStr, _ := vm.Decompile(program)
			hexArgs := make([]string, 0, len(input.InputWitness))
			for _, arg := range input.InputWitness {
				hexArgs = append(hexArgs, hex.EncodeToString(arg))
			}
			return errors.WithDetailf(ErrBadTx, "validation failed in script execution, input %d (program [%s] args [%s]): %s", i, scriptStr, strings.Join(hexArgs, " "), err)
		}
	}
	return nil
}

// safeSum is a helper for ValidateTx which adds two int64s. An error is returned
// if the result overflows.
func safeSum(x, y int64) (int64, error) {
	sum := x + y
	if (x > 0 && y > 0 && sum < 0) || (x < 0 && y < 0 && sum > 0) {
		return 0, errors.New("int64 overflow")
	}
	return sum, nil
}

// ApplyTx updates the state tree with all the changes to the ledger.
func ApplyTx(snapshot *state.Snapshot, tx *bc.Tx) error {
	for i, in := range tx.Inputs {
		if ic, ok := in.InputCommitment.(*bc.IssuanceInputCommitment); ok {
			// issuance input
			if i == 0 {
				snapshot.Issuances[tx.Hash] = ic.MaxTimeMS
			}
			continue
		}

		// Remove the consumed output from the state tree.
		prevoutKey, _ := state.OutputTreeItem(state.Prevout(in))
		err := snapshot.Tree.Delete(prevoutKey)
		if err != nil {
			return err
		}
	}

	for i, out := range tx.Outputs {
		if vmutil.IsUnspendable(out.ControlProgram) {
			continue
		}
		// Insert new outputs into the state tree.
		o := state.NewOutput(*out, bc.Outpoint{Hash: tx.Hash, Index: uint32(i)})
		err := snapshot.Tree.Insert(state.OutputTreeItem(o))
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package turnpike

import (
	"sync"
)

// Broker is the interface implemented by an object that handles routing EVENTS
// from Publishers to Subscribers.
type Broker interface {
	// Publishes a message to all Subscribers.
	Publish(*Session, *Publish)
	// Subscribes to messages on a URI.
	Subscribe(*Session, *Subscribe)
	// Unsubscribes from messages on a URI.
	Unsubscribe(*Session, *Unsubscribe)
}

// A super simple broker that matches URIs to Subscribers.
type defaultBroker struct {
	routes        map[URI]map[ID]Sender
	subscriptions map[ID]URI

	lastRequestId ID

	sync.RWMutex
}

// NewDefaultBroker initializes and returns a simple broker that matches URIs to
// Subscribers.
func NewDefaultBroker() Broker {
	return &defaultBroker{
		routes:        make(map[URI]map[ID]Sender),
		subscriptions: make(map[ID]URI),
	}
}

func (br *defaultBroker) nextRequestId() ID {
	br.lastRequestId++
	if br.lastRequestId > MAX_REQUEST_ID {
		br.lastRequestId = 1
	}

	return br.lastRequestId
}

// Publish sends a message to all subscribed clients except for the sender.
//
// If msg.Options["acknowledge"] == true, the publisher receives a Published event
// after the message has been sent to all subscribers.
func (br *defaultBroker) Publish(sess *Session, msg *Publish) {
	br.RLock()
	defer br.RUnlock()

	pub := sess.Peer
	pubID := NewID()
	evtTemplate := Event{
		Publication: pubID,
		Arguments:   msg.Arguments,
		ArgumentsKw: msg.ArgumentsKw,
		Details:     make(map[string]interface{}),
	}

	excludePublisher := true
	if exclude, ok := msg.Options["exclude_me"].(bool); ok {
		excludePublisher = exclude
	}

	for id, sub := range br.routes[msg.Topic] {
		// shallow-copy the template
		event := evtTemplate
		event.Subscription = id
		// don't send event to publisher
		if sub != pub || !excludePublisher {
			sub.Send(&event)
		}
	}

	// only send published message if acknowledge is present and set to true
	if doPub, _ := msg.Options["acknowledge"].(bool); doPub {
		pub.Send(&Published{Request: msg.Request, Publication: pubID})
	}
}

// Subscribe subscribes the client to the given topic.
func (br *defaultBroker) Subscribe(sess *Session, msg *Subscribe) {
	br.Lock()
	defer br.Unlock()

	if _, ok := br.routes[msg.Topic]; !ok {
		br.routes[msg.Topic] = make(map[ID]Sender)
	}
	id := br.nextRequestId()
	br.routes[msg.Topic][id] = sess.Peer

	br.subscriptions[id] = msg.Topic

	sess.Peer.Send(&Subscribed{Request: msg.Request, Subscription: id})
}

func (br *defaultBroker) Unsubscribe(sess *Session, msg *Unsubscribe) {
	br.Lock()
	defer br.Unlock()

	topic, ok := br.subscriptions[msg.Subscription]
	if !ok {
		err := &Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Error:   ErrNoSuchSubscription,
		}
		sess.Peer.Send(err)
		log.Printf("Error unsubscribing: no such subscription %v", msg.Subscription)
		return
	}
	delete(br.subscriptions, msg.Subscription)

	if r, ok := br.routes[topic]; !ok {
		log.Printf("Error unsubscribing: unable to find routes for %s topic", topic)
	} else if _, ok := r[msg.Subscription]; !ok {
		log.Printf("Error unsubscribing: %s route does not exist for %v subscription", topic, msg.Subscription)
	} else {
		delete(r, msg.Subscription)
		if len(r) == 0 {
			delete(br.routes, topic)
		}
	}
	sess.Peer.Send(&Unsubscribed{Request: msg.Request})
}
<commit_msg>Send messages in a goroutine<commit_after>package turnpike

import (
	"sync"
)

// Broker is the interface implemented by an object that handles routing EVENTS
// from Publishers to Subscribers.
type Broker interface {
	// Publishes a message to all Subscribers.
	Publish(*Session, *Publish)
	// Subscribes to messages on a URI.
	Subscribe(*Session, *Subscribe)
	// Unsubscribes from messages on a URI.
	Unsubscribe(*Session, *Unsubscribe)
}

// A super simple broker that matches URIs to Subscribers.
type defaultBroker struct {
	routes        map[URI]map[ID]Sender
	subscriptions map[ID]URI

	lastRequestId ID

	sync.RWMutex
}

// NewDefaultBroker initializes and returns a simple broker that matches URIs to
// Subscribers.
func NewDefaultBroker() Broker {
	return &defaultBroker{
		routes:        make(map[URI]map[ID]Sender),
		subscriptions: make(map[ID]URI),
	}
}

func (br *defaultBroker) nextRequestId() ID {
	br.lastRequestId++
	if br.lastRequestId > MAX_REQUEST_ID {
		br.lastRequestId = 1
	}

	return br.lastRequestId
}

// Publish sends a message to all subscribed clients except for the sender.
//
// If msg.Options["acknowledge"] == true, the publisher receives a Published event
// after the message has been sent to all subscribers.
func (br *defaultBroker) Publish(sess *Session, msg *Publish) {
	br.RLock()
	defer br.RUnlock()

	pub := sess.Peer
	pubID := NewID()
	evtTemplate := Event{
		Publication: pubID,
		Arguments:   msg.Arguments,
		ArgumentsKw: msg.ArgumentsKw,
		Details:     make(map[string]interface{}),
	}

	excludePublisher := true
	if exclude, ok := msg.Options["exclude_me"].(bool); ok {
		excludePublisher = exclude
	}

	for id, sub := range br.routes[msg.Topic] {
		// shallow-copy the template
		event := evtTemplate
		event.Subscription = id
		// don't send event to publisher
		if sub != pub || !excludePublisher {
			go sub.Send(&event)
		}
	}

	// only send published message if acknowledge is present and set to true
	if doPub, _ := msg.Options["acknowledge"].(bool); doPub {
		go pub.Send(&Published{Request: msg.Request, Publication: pubID})
	}
}

// Subscribe subscribes the client to the given topic.
func (br *defaultBroker) Subscribe(sess *Session, msg *Subscribe) {
	br.Lock()
	defer br.Unlock()

	if _, ok := br.routes[msg.Topic]; !ok {
		br.routes[msg.Topic] = make(map[ID]Sender)
	}
	id := br.nextRequestId()
	br.routes[msg.Topic][id] = sess.Peer

	br.subscriptions[id] = msg.Topic

	go sess.Peer.Send(&Subscribed{Request: msg.Request, Subscription: id})
}

func (br *defaultBroker) Unsubscribe(sess *Session, msg *Unsubscribe) {
	br.Lock()
	defer br.Unlock()

	topic, ok := br.subscriptions[msg.Subscription]
	if !ok {
		err := &Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Error:   ErrNoSuchSubscription,
		}
		sess.Peer.Send(err)
		log.Printf("Error unsubscribing: no such subscription %v", msg.Subscription)
		return
	}
	delete(br.subscriptions, msg.Subscription)

	if r, ok := br.routes[topic]; !ok {
		log.Printf("Error unsubscribing: unable to find routes for %s topic", topic)
	} else if _, ok := r[msg.Subscription]; !ok {
		log.Printf("Error unsubscribing: %s route does not exist for %v subscription", topic, msg.Subscription)
	} else {
		delete(r, msg.Subscription)
		if len(r) == 0 {
			delete(br.routes, topic)
		}
	}
	go sess.Peer.Send(&Unsubscribed{Request: msg.Request})
}
<|endoftext|>


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

package rlog

import (
	"bufio"
	"fmt"
	"os"
	"path"
	"runtime"
	"testing"
	"time"
)

var logfile string

// These two flags are used, so that during development of the tests we can
// quickly change behaviour to manually try or check things. The settings here
// are the correct behaviour for normal test runs.
var removeLogfile bool = true
var fixedLogfileName bool = false

// setup is called at the start of each test and prepares a new log file. It
// also returns a new configuration, which can be used by this test.
func setup() rlogConfig {
	if fixedLogfileName {
		logfile = "/tmp/rlog-test.log"
	} else {
		logfile = fmt.Sprintf("/tmp/rlog-test-%d.log", time.Now().UnixNano())
	}
	// If there's a logfile with that name already, remove it so that our tests
	// always start from scratch.
	os.Remove(logfile)

	return rlogConfig{
		logLevel:       "",
		traceLevel:     "",
		logTimeFormat:  "",
		logFile:        logfile,
		logStream:      "NONE",
		logTimeDate:    false,
		showCallerInfo: false,
	}
}

// cleanup is called at the end of each test.
func cleanup() {
	if removeLogfile {
		os.Remove(logfile)
	}
}

// fileMatch compares entries in the logfile with expected entries provided as
// a list of strings (one for each line). If a timeLayout string is provided
// then we will assume the first part of the line is a timestamp, which we will
// check to see if it's correctly formatted according to the specified time
// layout.
func fileMatch(t *testing.T, checkLines []string, timeLayout string) {
	// We need to know how many characters at the start of the line we should
	// assume to belong to the time stamp. The formatted time stamp can
	// actually be of different length than the time layout string, because
	// actual timezone names can have more characters than the TZ specified in
	// the layout. So we create the current time in the specified layout, which
	// should be very similar to the timestamps in the log lines.
	timeStampLen := len(time.Now().Format(timeLayout))

	// Scan over the logfile, line by line and compare to the lines provided in
	// checkLines.
	file, err := os.Open(logfile)
	if err != nil {
		t.Fatal(err)
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	i := 0
	for scanner.Scan() {
		line := scanner.Text()
		// Process and strip off the time stamp at the start, if a time layout
		// string was provided.
		if timeLayout != "" {
			dateTime := line[:timeStampLen]
			line = line[timeStampLen+1:]
			_, err := time.Parse(timeLayout, dateTime)
			if err != nil {
				t.Fatalf("Incorrect date/time format.\nSHOULD: %s\nIS:     %s\n", timeLayout, dateTime)
			}
		}
		t.Logf("\n-- fileMatch: SHOULD %s\n              IS     %s\n",
			checkLines[i], line)
		if i >= len(checkLines) {
			t.Fatal("Not enough lines provided in checkLines.")
		}
		if line != checkLines[i] {
			t.Fatalf("Log line %d does not match check line.\nSHOULD: %s\nIS:     %s\n", i, checkLines[i], line)
		}
		i++
	}
	if len(checkLines) > i {
		t.Fatalf("Only %d of %d checklines found in output file.", i, len(checkLines))
	}
	if i == 0 {
		t.Fatal("No input scanned")
	}
}

func TestLogLevels(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "DEBUG"
	Initialize(conf)

	Debug("Test Debug")
	Info("Test Info")
	Warn("Test Warning")
	Error("Test Error")
	Critical("Test Critical")

	checkLines := []string{
		"DEBUG    : Test Debug",
		"INFO     : Test Info",
		"WARN     : Test Warning",
		"ERROR    : Test Error",
		"CRITICAL : Test Critical",
	}
	fileMatch(t, checkLines, "")
}

func TestLogLevelsLimited(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "WARN"
	conf.traceLevel = "3"
	Initialize(conf)

	Debug("Test Debug")
	Info("Test Info")
	Warn("Test Warning")
	Error("Test Error")
	Critical("Test Critical")
	Trace(1, "Trace 1")
	Trace(2, "Trace 2")
	Trace(3, "Trace 3")
	Trace(4, "Trace 4")
	checkLines := []string{
		"WARN     : Test Warning",
		"ERROR    : Test Error",
		"CRITICAL : Test Critical",
		"TRACE(1) : Trace 1",
		"TRACE(2) : Trace 2",
		"TRACE(3) : Trace 3",
	}
	fileMatch(t, checkLines, "")
}

func TestLogFormatted(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "DEBUG"
	conf.traceLevel = "1"
	Initialize(conf)

	Debugf("Test Debug %d", 123)
	Infof("Test Info %d", 123)
	Warnf("Test Warning %d", 123)
	Errorf("Test Error %d", 123)
	Criticalf("Test Critical %d", 123)
	Tracef(1, "Trace 1 %d", 123)
	Tracef(2, "Trace 2 %d", 123)
	checkLines := []string{
		"DEBUG    : Test Debug 123",
		"INFO     : Test Info 123",
		"WARN     : Test Warning 123",
		"ERROR    : Test Error 123",
		"CRITICAL : Test Critical 123",
		"TRACE(1) : Trace 1 123",
	}
	fileMatch(t, checkLines, "")
}

// TestLogTimestamp checks that the time stamp format can be changed and that
// we indeed get a properly formatted timestamp output.
func TestLogTimestamp(t *testing.T) {
	conf := setup()
	conf.logTimeDate = true
	defer cleanup()

	checkLines := []string{
		"INFO     : Test Info",
	}

	// Map of various 'user specified' time layouts and the actual time layout
	// to which they should be mapped. Some of the capitalization in the well
	// known time stamps is off, to show that those names can be specified in a
	// case insensitive manner.
	checkTimeStamps := map[string]string{
		"ansIC":       time.ANSIC,
		"UNIXDATE":    time.UnixDate,
		"rubydate":    time.RubyDate,
		"rfc822":      time.RFC822,
		"rfc822z":     time.RFC822Z,
		"rfc1123":     time.RFC1123,
		"rfc1123z":    time.RFC1123Z,
		"RFC3339":     time.RFC3339,
		"RFC3339Nano": time.RFC3339Nano,
		"Kitchen":     time.Kitchen,
		"":            time.RFC3339, // If nothing specified, default is RFC3339
		"2006/01/02 15:04:05": "2006/01/02 15:04:05", // custom format
	}

	for tsUserSpecified, tsActualFormat := range checkTimeStamps {
		os.Remove(logfile)

		// Specify a time layout...
		conf.logTimeFormat = tsUserSpecified
		Initialize(conf)

		Info("Test Info")
		// We can specify a time layout to fileMatch, which then performs the extra
		// check for the correctly formatted time stamp.
		fileMatch(t, checkLines, tsActualFormat)
	}
}

// TestLogCallerInfo manually figures out the caller info, which should be
// displayed by rlog. The code that's creating the expected caller info
// within the test is pretty much exactly the code that should be at work
// within rlog.
func TestLogCallerInfo(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.showCallerInfo = true
	Initialize(conf)

	Info("Test Info")
	pc, fullFilePath, line, _ := runtime.Caller(0)
	line-- // The log was called in the line before, so... -1

	// The following lines simply format the caller info in the way that it
	// should be formatted by rlog
	callingFuncName := runtime.FuncForPC(pc).Name()
	dirPath, fileName := path.Split(fullFilePath)
	var moduleName string = ""
	if dirPath != "" {
		dirPath = dirPath[:len(dirPath)-1]
		dirPath, moduleName = path.Split(dirPath)
	}
	moduleAndFileName := moduleName + "/" + fileName
	shouldLine := fmt.Sprintf("INFO     : [%s:%d (%s)] Test Info",
		moduleAndFileName, line, callingFuncName)

	checkLines := []string{shouldLine}
	fileMatch(t, checkLines, "")
}

// TestLogLevelsFiltered checks whether the per-module filtering works
// correctly. For that, we provide a log-level filter that names this
// executable here, so that log messages should be displayed, and a trace level
// filter for a non-existent module, so that trace messages should not be
// displayed.
func TestLogLevelsFiltered(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "rlog_test.go=WARN"
	conf.traceLevel = "foobar.go=2" // should not see any of those
	Initialize(conf)

	Debug("Test Debug")
	Info("Test Info")
	Warn("Test Warning")
	Error("Test Error")
	Critical("Test Critical")
	Trace(1, "Trace 1")
	Trace(2, "Trace 2")
	Trace(3, "Trace 3")
	Trace(4, "Trace 4")
	checkLines := []string{
		"WARN     : Test Warning",
		"ERROR    : Test Error",
		"CRITICAL : Test Critical",
	}
	fileMatch(t, checkLines, "")
}
<commit_msg>Updated comments.<commit_after>// Copyright (c) 2016 Pani Networks
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package rlog

import (
	"bufio"
	"fmt"
	"os"
	"path"
	"runtime"
	"testing"
	"time"
)

var logfile string

// These two flags are used to quickly change behaviour of our tests, so that
// we can manually check and test things during development of those tests.
// The settings here reflect the correct behaviour for normal test runs.
var removeLogfile bool = true
var fixedLogfileName bool = false

// setup is called at the start of each test and prepares a new log file. It
// also returns a new configuration, which can be used by this test.
func setup() rlogConfig {
	if fixedLogfileName {
		logfile = "/tmp/rlog-test.log"
	} else {
		logfile = fmt.Sprintf("/tmp/rlog-test-%d.log", time.Now().UnixNano())
	}

	// If there's a logfile with that name already, remove it so that our tests
	// always start from scratch.
	os.Remove(logfile)

	// Provide a default config, which can be used or modified by the tests
	return rlogConfig{
		logLevel:       "",
		traceLevel:     "",
		logTimeFormat:  "",
		logFile:        logfile,
		logStream:      "NONE",
		logTimeDate:    false,
		showCallerInfo: false,
	}
}

// cleanup is called at the end of each test.
func cleanup() {
	if removeLogfile {
		os.Remove(logfile)
	}
}

// fileMatch compares entries in the logfile with expected entries provided as
// a list of strings (one for each line). If a timeLayout string is provided
// then we will assume the first part of the line is a timestamp, which we will
// check to see if it's correctly formatted according to the specified time
// layout.
func fileMatch(t *testing.T, checkLines []string, timeLayout string) {
	// We need to know how many characters at the start of the line we should
	// assume to belong to the time stamp. The formatted time stamp can
	// actually be of different length than the time layout string, because
	// actual timezone names can have more characters than the TZ specified in
	// the layout. So we create the current time in the specified layout, which
	// should be very similar to the timestamps in the log lines.
	timeStampLen := len(time.Now().Format(timeLayout))

	// Scan over the logfile, line by line and compare to the lines provided in
	// checkLines.
	file, err := os.Open(logfile)
	if err != nil {
		t.Fatal(err)
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	i := 0
	for scanner.Scan() {
		line := scanner.Text()
		// Process and strip off the time stamp at the start, if a time layout
		// string was provided.
		if timeLayout != "" {
			dateTime := line[:timeStampLen]
			line = line[timeStampLen+1:]
			_, err := time.Parse(timeLayout, dateTime)
			if err != nil {
				t.Fatalf("Incorrect date/time format.\nSHOULD: %s\nIS:     %s\n", timeLayout, dateTime)
			}
		}
		t.Logf("\n-- fileMatch: SHOULD %s\n              IS     %s\n",
			checkLines[i], line)
		if i >= len(checkLines) {
			t.Fatal("Not enough lines provided in checkLines.")
		}
		if line != checkLines[i] {
			t.Fatalf("Log line %d does not match check line.\nSHOULD: %s\nIS:     %s\n", i, checkLines[i], line)
		}
		i++
	}
	if len(checkLines) > i {
		t.Fatalf("Only %d of %d checklines found in output file.", i, len(checkLines))
	}
	if i == 0 {
		t.Fatal("No input scanned")
	}
}

// ---------- Tests -----------

// TestLogLevels performs some basic tests for each known log level.
func TestLogLevels(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "DEBUG"
	Initialize(conf)

	Debug("Test Debug")
	Info("Test Info")
	Warn("Test Warning")
	Error("Test Error")
	Critical("Test Critical")

	checkLines := []string{
		"DEBUG    : Test Debug",
		"INFO     : Test Info",
		"WARN     : Test Warning",
		"ERROR    : Test Error",
		"CRITICAL : Test Critical",
	}
	fileMatch(t, checkLines, "")
}

// TestLogLevelsLimited checks that we can limit the output of log and trace
// messages that don't meed the minimum configured logging levels.
func TestLogLevelsLimited(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "WARN"
	conf.traceLevel = "3"
	Initialize(conf)

	Debug("Test Debug")
	Info("Test Info")
	Warn("Test Warning")
	Error("Test Error")
	Critical("Test Critical")
	Trace(1, "Trace 1")
	Trace(2, "Trace 2")
	Trace(3, "Trace 3")
	Trace(4, "Trace 4")
	checkLines := []string{
		"WARN     : Test Warning",
		"ERROR    : Test Error",
		"CRITICAL : Test Critical",
		"TRACE(1) : Trace 1",
		"TRACE(2) : Trace 2",
		"TRACE(3) : Trace 3",
	}
	fileMatch(t, checkLines, "")
}

// TestLogFormatted checks whether the *f functions for formatted output work
// as expected.
func TestLogFormatted(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "DEBUG"
	conf.traceLevel = "1"
	Initialize(conf)

	Debugf("Test Debug %d", 123)
	Infof("Test Info %d", 123)
	Warnf("Test Warning %d", 123)
	Errorf("Test Error %d", 123)
	Criticalf("Test Critical %d", 123)
	Tracef(1, "Trace 1 %d", 123)
	Tracef(2, "Trace 2 %d", 123)
	checkLines := []string{
		"DEBUG    : Test Debug 123",
		"INFO     : Test Info 123",
		"WARN     : Test Warning 123",
		"ERROR    : Test Error 123",
		"CRITICAL : Test Critical 123",
		"TRACE(1) : Trace 1 123",
	}
	fileMatch(t, checkLines, "")
}

// TestLogTimestamp checks that the time stamp format can be changed and that
// we indeed get a properly formatted timestamp output.
func TestLogTimestamp(t *testing.T) {
	conf := setup()
	conf.logTimeDate = true
	defer cleanup()

	checkLines := []string{
		"INFO     : Test Info",
	}

	// Map of various 'user specified' time layouts and the actual time layout
	// to which they should be mapped. Some of the capitalization in the well
	// known time stamps is off, to show that those names can be specified in a
	// case insensitive manner.
	checkTimeStamps := map[string]string{
		"ansIC":       time.ANSIC,
		"UNIXDATE":    time.UnixDate,
		"rubydate":    time.RubyDate,
		"rfc822":      time.RFC822,
		"rfc822z":     time.RFC822Z,
		"rfc1123":     time.RFC1123,
		"rfc1123z":    time.RFC1123Z,
		"RFC3339":     time.RFC3339,
		"RFC3339Nano": time.RFC3339Nano,
		"Kitchen":     time.Kitchen,
		"":            time.RFC3339, // If nothing specified, default is RFC3339
		"2006/01/02 15:04:05": "2006/01/02 15:04:05", // custom format
	}

	for tsUserSpecified, tsActualFormat := range checkTimeStamps {
		os.Remove(logfile)

		// Specify a time layout...
		conf.logTimeFormat = tsUserSpecified
		Initialize(conf)

		Info("Test Info")
		// We can specify a time layout to fileMatch, which then performs the extra
		// check for the correctly formatted time stamp.
		fileMatch(t, checkLines, tsActualFormat)
	}
}

// TestLogCallerInfo manually figures out the caller info, which should be
// displayed by rlog. The code that's creating the expected caller info
// within the test is pretty much exactly the code that should be at work
// within rlog.
func TestLogCallerInfo(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.showCallerInfo = true
	Initialize(conf)

	Info("Test Info")
	pc, fullFilePath, line, _ := runtime.Caller(0)
	line-- // The log was called in the line before, so... -1

	// The following lines simply format the caller info in the way that it
	// should be formatted by rlog
	callingFuncName := runtime.FuncForPC(pc).Name()
	dirPath, fileName := path.Split(fullFilePath)
	var moduleName string = ""
	if dirPath != "" {
		dirPath = dirPath[:len(dirPath)-1]
		dirPath, moduleName = path.Split(dirPath)
	}
	moduleAndFileName := moduleName + "/" + fileName
	shouldLine := fmt.Sprintf("INFO     : [%s:%d (%s)] Test Info",
		moduleAndFileName, line, callingFuncName)

	checkLines := []string{shouldLine}
	fileMatch(t, checkLines, "")
}

// TestLogLevelsFiltered checks whether the per-module filtering works
// correctly. For that, we provide a log-level filter that names this
// executable here, so that log messages should be displayed, and a trace level
// filter for a non-existent module, so that trace messages should not be
// displayed.
func TestLogLevelsFiltered(t *testing.T) {
	conf := setup()
	defer cleanup()

	conf.logLevel = "rlog_test.go=WARN"
	conf.traceLevel = "foobar.go=2" // should not see any of those
	Initialize(conf)

	Debug("Test Debug")
	Info("Test Info")
	Warn("Test Warning")
	Error("Test Error")
	Critical("Test Critical")
	Trace(1, "Trace 1")
	Trace(2, "Trace 2")
	Trace(3, "Trace 3")
	Trace(4, "Trace 4")
	checkLines := []string{
		"WARN     : Test Warning",
		"ERROR    : Test Error",
		"CRITICAL : Test Critical",
	}
	fileMatch(t, checkLines, "")
}
<|endoftext|>


<commit_before>package turnpike

import "sync"

// Broker is the interface implemented by an object that handles routing EVENTS
// from Publishers to Subscribers.
type Broker interface {
	// Publishes a message to all Subscribers.
	Publish(*Session, *Publish)
	// Subscribes to messages on a URI.
	Subscribe(*Session, *Subscribe)
	// Unsubscribes from messages on a URI.
	Unsubscribe(*Session, *Unsubscribe)
	// Removes all subscriptions of the subscriber.
	RemoveSession(*Session)
}

// A super simple broker that matches URIs to Subscribers.
type defaultBroker struct {
	routes        map[URI]map[ID]*Session
	subscriptions map[ID]URI
	sessions      map[*Session]map[ID]struct{}
	lock          sync.RWMutex
}

// NewDefaultBroker initializes and returns a simple broker that matches URIs to
// Subscribers.
func NewDefaultBroker() Broker {
	return &defaultBroker{
		routes:        make(map[URI]map[ID]*Session),
		subscriptions: make(map[ID]URI),
		sessions:      make(map[*Session]map[ID]struct{}),
	}
}

// Publish sends a message to all subscribed clients except for the sender.
//
// If msg.Options["acknowledge"] == true, the publisher receives a Published event
// after the message has been sent to all subscribers.
func (br *defaultBroker) Publish(pub *Session, msg *Publish) {
	pubID := NewID()
	evtTemplate := Event{
		Publication: pubID,
		Arguments:   msg.Arguments,
		ArgumentsKw: msg.ArgumentsKw,
		Details:     make(map[string]interface{}),
	}

	excludePublisher := true
	if exclude, ok := msg.Options["exclude_me"].(bool); ok {
		excludePublisher = exclude
	}

	br.lock.RLock()
	for id, sub := range br.routes[msg.Topic] {
		// don't send event to publisher
		if sub == pub && excludePublisher {
			continue
		}

		// shallow-copy the template
		event := evtTemplate
		event.Subscription = id
		sub.Send(&event)
	}
	br.lock.RUnlock()

	// only send published message if acknowledge is present and set to true
	if doPub, _ := msg.Options["acknowledge"].(bool); doPub {
		pub.Send(&Published{Request: msg.Request, Publication: pubID})
	}
}

// Subscribe subscribes the client to the given topic.
func (br *defaultBroker) Subscribe(sub *Session, msg *Subscribe) {
	id := NewID()

	br.lock.Lock()
	r, ok := br.routes[msg.Topic]
	if !ok {
		r = make(map[ID]*Session)
		br.routes[msg.Topic] = r
	}
	r[id] = sub

	s, ok := br.sessions[sub]
	if !ok {
		s = make(map[ID]struct{})
		br.sessions[sub] = s
	}
	s[id] = struct{}{}

	br.subscriptions[id] = msg.Topic
	br.lock.Unlock()

	sub.Send(&Subscribed{Request: msg.Request, Subscription: id})
}

func (br *defaultBroker) Unsubscribe(sub *Session, msg *Unsubscribe) {
	br.lock.Lock()
	topic, ok := br.subscriptions[msg.Subscription]
	if !ok {
		br.lock.Unlock()
		err := &Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Error:   ErrNoSuchSubscription,
		}
		sub.Send(err)
		log.Printf("Error unsubscribing: no such subscription %v", msg.Subscription)
		return
	}
	delete(br.subscriptions, msg.Subscription)

	// clean up routes
	if r, ok := br.routes[topic]; !ok {
		log.Printf("Error unsubscribing: unable to find routes for %s topic", topic)
	} else if _, ok := r[msg.Subscription]; !ok {
		log.Printf("Error unsubscribing: %s route does not exist for %v subscription", topic, msg.Subscription)
	} else {
		delete(r, msg.Subscription)
		if len(r) == 0 {
			delete(br.routes, topic)
		}
	}

	// clean up sender's subscription
	if s, ok := br.sessions[sub]; !ok {
		log.Println("Error unsubscribing: unable to find sender's subscriptions")
	} else if _, ok := s[msg.Subscription]; !ok {
		log.Printf("Error unsubscribing: sender does not contain %s subscription", msg.Subscription)
	} else {
		delete(s, msg.Subscription)
		if len(s) == 0 {
			delete(br.sessions, sub)
		}
	}
	br.lock.Unlock()

	sub.Send(&Unsubscribed{Request: msg.Request})
}

func (br *defaultBroker) RemoveSession(sub *Session) {
	br.lock.Lock()
	defer br.lock.Unlock()

	for id, _ := range br.sessions[sub] {
		topic, ok := br.subscriptions[id]
		if !ok {
			continue
		}
		delete(br.subscriptions, id)

		// clean up routes
		if r, ok := br.routes[topic]; ok {
			if _, ok := r[id]; ok {
				delete(r, id)
				if len(r) == 0 {
					delete(br.routes, topic)
				}
			}
		}
	}
	delete(br.sessions, sub)
}
<commit_msg>Fix 'go vet' warning<commit_after>package turnpike

import "sync"

// Broker is the interface implemented by an object that handles routing EVENTS
// from Publishers to Subscribers.
type Broker interface {
	// Publishes a message to all Subscribers.
	Publish(*Session, *Publish)
	// Subscribes to messages on a URI.
	Subscribe(*Session, *Subscribe)
	// Unsubscribes from messages on a URI.
	Unsubscribe(*Session, *Unsubscribe)
	// Removes all subscriptions of the subscriber.
	RemoveSession(*Session)
}

// A super simple broker that matches URIs to Subscribers.
type defaultBroker struct {
	routes        map[URI]map[ID]*Session
	subscriptions map[ID]URI
	sessions      map[*Session]map[ID]struct{}
	lock          sync.RWMutex
}

// NewDefaultBroker initializes and returns a simple broker that matches URIs to
// Subscribers.
func NewDefaultBroker() Broker {
	return &defaultBroker{
		routes:        make(map[URI]map[ID]*Session),
		subscriptions: make(map[ID]URI),
		sessions:      make(map[*Session]map[ID]struct{}),
	}
}

// Publish sends a message to all subscribed clients except for the sender.
//
// If msg.Options["acknowledge"] == true, the publisher receives a Published event
// after the message has been sent to all subscribers.
func (br *defaultBroker) Publish(pub *Session, msg *Publish) {
	pubID := NewID()
	evtTemplate := Event{
		Publication: pubID,
		Arguments:   msg.Arguments,
		ArgumentsKw: msg.ArgumentsKw,
		Details:     make(map[string]interface{}),
	}

	excludePublisher := true
	if exclude, ok := msg.Options["exclude_me"].(bool); ok {
		excludePublisher = exclude
	}

	br.lock.RLock()
	for id, sub := range br.routes[msg.Topic] {
		// don't send event to publisher
		if sub == pub && excludePublisher {
			continue
		}

		// shallow-copy the template
		event := evtTemplate
		event.Subscription = id
		sub.Send(&event)
	}
	br.lock.RUnlock()

	// only send published message if acknowledge is present and set to true
	if doPub, _ := msg.Options["acknowledge"].(bool); doPub {
		pub.Send(&Published{Request: msg.Request, Publication: pubID})
	}
}

// Subscribe subscribes the client to the given topic.
func (br *defaultBroker) Subscribe(sub *Session, msg *Subscribe) {
	id := NewID()

	br.lock.Lock()
	r, ok := br.routes[msg.Topic]
	if !ok {
		r = make(map[ID]*Session)
		br.routes[msg.Topic] = r
	}
	r[id] = sub

	s, ok := br.sessions[sub]
	if !ok {
		s = make(map[ID]struct{})
		br.sessions[sub] = s
	}
	s[id] = struct{}{}

	br.subscriptions[id] = msg.Topic
	br.lock.Unlock()

	sub.Send(&Subscribed{Request: msg.Request, Subscription: id})
}

func (br *defaultBroker) Unsubscribe(sub *Session, msg *Unsubscribe) {
	br.lock.Lock()
	topic, ok := br.subscriptions[msg.Subscription]
	if !ok {
		br.lock.Unlock()
		err := &Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Error:   ErrNoSuchSubscription,
		}
		sub.Send(err)
		log.Printf("Error unsubscribing: no such subscription %v", msg.Subscription)
		return
	}
	delete(br.subscriptions, msg.Subscription)

	// clean up routes
	if r, ok := br.routes[topic]; !ok {
		log.Printf("Error unsubscribing: unable to find routes for %s topic", topic)
	} else if _, ok := r[msg.Subscription]; !ok {
		log.Printf("Error unsubscribing: %s route does not exist for %v subscription", topic, msg.Subscription)
	} else {
		delete(r, msg.Subscription)
		if len(r) == 0 {
			delete(br.routes, topic)
		}
	}

	// clean up sender's subscription
	if s, ok := br.sessions[sub]; !ok {
		log.Println("Error unsubscribing: unable to find sender's subscriptions")
	} else if _, ok := s[msg.Subscription]; !ok {
		log.Printf("Error unsubscribing: sender does not contain %v subscription", msg.Subscription)
	} else {
		delete(s, msg.Subscription)
		if len(s) == 0 {
			delete(br.sessions, sub)
		}
	}
	br.lock.Unlock()

	sub.Send(&Unsubscribed{Request: msg.Request})
}

func (br *defaultBroker) RemoveSession(sub *Session) {
	br.lock.Lock()
	defer br.lock.Unlock()

	for id, _ := range br.sessions[sub] {
		topic, ok := br.subscriptions[id]
		if !ok {
			continue
		}
		delete(br.subscriptions, id)

		// clean up routes
		if r, ok := br.routes[topic]; ok {
			if _, ok := r[id]; ok {
				delete(r, id)
				if len(r) == 0 {
					delete(br.routes, topic)
				}
			}
		}
	}
	delete(br.sessions, sub)
}
<|endoftext|>