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

package isolation

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

/*
 ATTENTION: Rerun code generators when interface signatures are modified.
*/

import (
	"bufio"
	"encoding/csv"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"unsafe"

	"golang.org/x/sys/unix"

	ps "github.com/mitchellh/go-ps"
	"k8s.io/apimachinery/pkg/api/resource"

	v1 "kubevirt.io/client-go/api/v1"
	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/util"
	cmdclient "kubevirt.io/kubevirt/pkg/virt-handler/cmd-client"
)

// PodIsolationDetector helps detecting cgroups, namespaces and PIDs of Pods from outside of them.
// Different strategies may be applied to do that.
type PodIsolationDetector interface {
	// Detect takes a vm, looks up a socket based the VM and detects pid, cgroups and namespaces of the owner of that socket.
	// It returns an IsolationResult containing all isolation information
	Detect(vm *v1.VirtualMachineInstance) (*IsolationResult, error)

	DetectForSocket(vm *v1.VirtualMachineInstance, socket string) (*IsolationResult, error)

	// Whitelist allows specifying cgroup controller which should be considered to detect the cgroup slice
	// It returns a PodIsolationDetector to allow configuring the PodIsolationDetector via the builder pattern.
	Whitelist(controller []string) PodIsolationDetector

	// Adjust system resources to run the passed VM
	AdjustResources(vm *v1.VirtualMachineInstance) error
}

type MountInfo struct {
	DeviceContainingFile string
	Root                 string
	MountPoint           string
}

type socketBasedIsolationDetector struct {
	socketDir  string
	controller []string
}

func (s *socketBasedIsolationDetector) DetectForSocket(vm *v1.VirtualMachineInstance, socket string) (*IsolationResult, error) {
	var pid int
	var slice string
	var err error
	var controller []string

	if pid, err = s.getPid(socket); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get owner Pid of socket %s", socket)
		return nil, err

	}

	// Look up the cgroup slice based on the whitelisted controller
	if controller, slice, err = s.getSlice(pid); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get cgroup slice for Pid %d", pid)
		return nil, err
	}

	return NewIsolationResult(pid, slice, controller), nil
}

// NewSocketBasedIsolationDetector takes socketDir and creates a socket based IsolationDetector
// It returns a PodIsolationDetector which detects pid, cgroups and namespaces of the socket owner.
func NewSocketBasedIsolationDetector(socketDir string) PodIsolationDetector {
	return &socketBasedIsolationDetector{socketDir: socketDir, controller: []string{"devices"}}
}

func (s *socketBasedIsolationDetector) Whitelist(controller []string) PodIsolationDetector {
	s.controller = controller
	return s
}

func (s *socketBasedIsolationDetector) Detect(vm *v1.VirtualMachineInstance) (*IsolationResult, error) {
	var pid int
	var slice string
	var err error
	var controller []string

	// Look up the socket of the virt-launcher Pod which was created for that VM, and extract the PID from it
	socket := cmdclient.SocketFromUID(s.socketDir, string(vm.UID))
	if pid, err = s.getPid(socket); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get owner Pid of socket %s", socket)
		return nil, err
	}

	// Look up the cgroup slice based on the whitelisted controller
	if controller, slice, err = s.getSlice(pid); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get cgroup slice for Pid %d", pid)
		return nil, err
	}

	return NewIsolationResult(pid, slice, controller), nil
}

// standard golang libraries don't provide API to set runtime limits
// for other processes, so we have to directly call to kernel
func prLimit(pid int, limit uintptr, rlimit *unix.Rlimit) error {
	_, _, errno := unix.RawSyscall6(unix.SYS_PRLIMIT64,
		uintptr(pid),
		limit,
		uintptr(unsafe.Pointer(rlimit)),
		0, 0, 0)
	if errno != 0 {
		return fmt.Errorf("Error setting prlimit: %v", errno)
	}
	return nil
}

func (s *socketBasedIsolationDetector) AdjustResources(vm *v1.VirtualMachineInstance) error {
	// only VFIO attached domains require MEMLOCK adjustment
	if !util.IsSRIOVVmi(vm) {
		return nil
	}

	// bump memlock ulimit for libvirtd
	res, err := s.Detect(vm)
	if err != nil {
		return err
	}
	launcherPid := res.Pid()

	processes, err := ps.Processes()
	if err != nil {
		return fmt.Errorf("failed to get all processes: %v", err)
	}

	for _, process := range processes {
		// consider all processes that are virt-launcher children
		if process.PPid() != launcherPid {
			continue
		}

		// libvirtd process sets the memory lock limit before fork/exec-ing into qemu
		if process.Executable() != "libvirtd" {
			continue
		}

		// make the best estimate for memory required by libvirt
		memlockSize, err := getMemlockSize(vm)
		if err != nil {
			return err
		}
		rLimit := unix.Rlimit{
			Max: uint64(memlockSize),
			Cur: uint64(memlockSize),
		}
		err = prLimit(process.Pid(), unix.RLIMIT_MEMLOCK, &rLimit)
		if err != nil {
			return fmt.Errorf("failed to set rlimit for memory lock: %v", err)
		}
		// we assume a single process should match
		break
	}
	return nil
}

// consider reusing getMemoryOverhead()
func getMemlockSize(vm *v1.VirtualMachineInstance) (int64, error) {
	memlockSize := resource.NewQuantity(0, resource.DecimalSI)

	// start with base memory requested for the VM
	vmiMemoryReq := vm.Spec.Domain.Resources.Requests.Memory()
	memlockSize.Add(*resource.NewScaledQuantity(vmiMemoryReq.ScaledValue(resource.Kilo), resource.Kilo))

	// allocate 1Gb for VFIO needs
	memlockSize.Add(resource.MustParse("1G"))

	// add some more memory for NUMA / CPU topology, platform memory alignment and other needs
	memlockSize.Add(resource.MustParse("256M"))

	bytes_, ok := memlockSize.AsInt64()
	if !ok {
		return 0, fmt.Errorf("could not calculate memory lock size")
	}
	return bytes_, nil
}

func NewIsolationResult(pid int, slice string, controller []string) *IsolationResult {
	return &IsolationResult{pid: pid, slice: slice, controller: controller}
}

type IsolationResult struct {
	pid        int
	slice      string
	controller []string
}

func (r *IsolationResult) Slice() string {
	return r.slice
}

func (r *IsolationResult) PIDNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/pid", r.pid)
}

func (r *IsolationResult) MountNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/mnt", r.pid)
}

func (r *IsolationResult) mountInfo() string {
	return fmt.Sprintf("/proc/%d/mountinfo", r.pid)
}

// MountInfoRoot returns information about the root entry in /proc/mountinfo
func (r *IsolationResult) MountInfoRoot() (*MountInfo, error) {
	in, err := os.Open(r.mountInfo())
	if err != nil {
		return nil, fmt.Errorf("could not open mountinfo: %v", err)
	}
	defer in.Close()
	c := csv.NewReader(in)
	c.Comma = ' '
	c.LazyQuotes = true
	for {
		record, err := c.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			if e, ok := err.(*csv.ParseError); ok {
				if e.Err != csv.ErrFieldCount {
					return nil, err
				}
			} else {
				return nil, err
			}
		}

		if record[3] == "/" && record[4] == "/" {
			return &MountInfo{
				DeviceContainingFile: record[2],
				Root:                 record[3],
				MountPoint:           record[4],
			}, nil
		}
	}

	//impossible
	return nil, fmt.Errorf("process has no root entry")
}

// IsMounted checks if a path in the mount namespace of a
// given process isolation result is a mount point. Works with symlinks.
func (r *IsolationResult) IsMounted(mountPoint string) (bool, error) {
	mountPoint, err := filepath.EvalSymlinks(mountPoint)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("could not resolve mount point path: %v", err)
	}
	in, err := os.Open(r.mountInfo())
	if err != nil {
		return false, fmt.Errorf("could not open mountinfo: %v", err)
	}
	defer in.Close()
	c := csv.NewReader(in)
	c.Comma = ' '
	c.LazyQuotes = true
	for {
		record, err := c.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			if e, ok := err.(*csv.ParseError); ok {
				if e.Err != csv.ErrFieldCount {
					return false, err
				}
			} else {
				return false, err
			}
		}

		if record[4] == mountPoint {
			return true, nil
		}
	}
	return false, nil
}

// ParentMountInfoFor takes the mount info from a container, and looks the corresponding
// entry in /proc/mountinfo of the isolation result of the given process.
func (r *IsolationResult) ParentMountInfoFor(mountInfo *MountInfo) (*MountInfo, error) {
	in, err := os.Open(r.mountInfo())
	if err != nil {
		return nil, fmt.Errorf("could not open mountinfo: %v", err)
	}
	defer in.Close()
	c := csv.NewReader(in)
	c.Comma = ' '
	c.LazyQuotes = true
	for {
		record, err := c.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			if e, ok := err.(*csv.ParseError); ok {
				if e.Err != csv.ErrFieldCount {
					return nil, err
				}
			} else {
				return nil, err
			}
		}

		if record[2] == mountInfo.DeviceContainingFile {
			return &MountInfo{
				DeviceContainingFile: record[2],
				Root:                 record[3],
				MountPoint:           record[4],
			}, nil
		}
	}
	return nil, fmt.Errorf("no parent entry for %v found in the mount namespace of %d", mountInfo.DeviceContainingFile, r.pid)
}

func (r *IsolationResult) NetNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/net", r.pid)
}

func (r *IsolationResult) MountRoot() string {
	return fmt.Sprintf("/proc/%d/root", r.pid)
}

func (r *IsolationResult) Pid() int {
	return r.pid
}

func (r *IsolationResult) Controller() []string {
	return r.controller
}

func (s *socketBasedIsolationDetector) getPid(socket string) (int, error) {
	sock, err := net.Dial("unix", socket)
	if err != nil {
		return -1, err
	}
	defer sock.Close()

	ufile, err := sock.(*net.UnixConn).File()
	if err != nil {
		return -1, err
	}
	// This is the tricky part, which will give us the PID of the owning socket
	ucreds, err := syscall.GetsockoptUcred(int(ufile.Fd()), syscall.SOL_SOCKET, syscall.SO_PEERCRED)
	if err != nil {
		return -1, err
	}

	if int(ucreds.Pid) == 0 {
		return -1, fmt.Errorf("The detected PID is 0. Is the isolation detector running in the host PID namespace?")
	}

	return int(ucreds.Pid), nil
}

func (s *socketBasedIsolationDetector) getSlice(pid int) (controller []string, slice string, err error) {
	cgroups, err := os.Open(fmt.Sprintf("/proc/%d/cgroup", pid))
	if err != nil {
		return
	}
	defer cgroups.Close()

	scanner := bufio.NewScanner(cgroups)
	for scanner.Scan() {
		cgEntry := strings.Split(scanner.Text(), ":")
		// Check if we have a sane cgroup line
		if len(cgEntry) != 3 {
			err = fmt.Errorf("Could not extract slice from cgroup line: %s", scanner.Text())
			return
		}
		// Skip not supported cgroup controller
		if !sliceContains(s.controller, cgEntry[1]) {
			continue
		}

		// Set and check cgroup slice
		if slice == "" {
			slice = cgEntry[2]
		} else if slice != cgEntry[2] {
			err = fmt.Errorf("Process is part of more than one slice. Expected %s, found %s", slice, cgEntry[2])
			return
		}
		// Add controller
		controller = append(controller, cgEntry[1])
	}
	// Check if we encountered a read error
	if scanner.Err() != nil {
		err = scanner.Err()
		return
	}

	if slice == "" {
		err = fmt.Errorf("Could not detect slice of whitelisted controller: %v", s.controller)
		return
	}
	return
}

func sliceContains(controllers []string, value string) bool {
	for _, c := range controllers {
		if c == value {
			return true
		}
	}
	return false
}

func NodeIsolationResult() *IsolationResult {
	return &IsolationResult{
		pid: 1,
	}
}
<commit_msg>Document how we estimate memory lock size for libvirtd<commit_after>/*
 * This file is part of the kubevirt project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package isolation

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

/*
 ATTENTION: Rerun code generators when interface signatures are modified.
*/

import (
	"bufio"
	"encoding/csv"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"unsafe"

	"golang.org/x/sys/unix"

	ps "github.com/mitchellh/go-ps"
	"k8s.io/apimachinery/pkg/api/resource"

	v1 "kubevirt.io/client-go/api/v1"
	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/util"
	cmdclient "kubevirt.io/kubevirt/pkg/virt-handler/cmd-client"
)

// PodIsolationDetector helps detecting cgroups, namespaces and PIDs of Pods from outside of them.
// Different strategies may be applied to do that.
type PodIsolationDetector interface {
	// Detect takes a vm, looks up a socket based the VM and detects pid, cgroups and namespaces of the owner of that socket.
	// It returns an IsolationResult containing all isolation information
	Detect(vm *v1.VirtualMachineInstance) (*IsolationResult, error)

	DetectForSocket(vm *v1.VirtualMachineInstance, socket string) (*IsolationResult, error)

	// Whitelist allows specifying cgroup controller which should be considered to detect the cgroup slice
	// It returns a PodIsolationDetector to allow configuring the PodIsolationDetector via the builder pattern.
	Whitelist(controller []string) PodIsolationDetector

	// Adjust system resources to run the passed VM
	AdjustResources(vm *v1.VirtualMachineInstance) error
}

type MountInfo struct {
	DeviceContainingFile string
	Root                 string
	MountPoint           string
}

type socketBasedIsolationDetector struct {
	socketDir  string
	controller []string
}

func (s *socketBasedIsolationDetector) DetectForSocket(vm *v1.VirtualMachineInstance, socket string) (*IsolationResult, error) {
	var pid int
	var slice string
	var err error
	var controller []string

	if pid, err = s.getPid(socket); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get owner Pid of socket %s", socket)
		return nil, err

	}

	// Look up the cgroup slice based on the whitelisted controller
	if controller, slice, err = s.getSlice(pid); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get cgroup slice for Pid %d", pid)
		return nil, err
	}

	return NewIsolationResult(pid, slice, controller), nil
}

// NewSocketBasedIsolationDetector takes socketDir and creates a socket based IsolationDetector
// It returns a PodIsolationDetector which detects pid, cgroups and namespaces of the socket owner.
func NewSocketBasedIsolationDetector(socketDir string) PodIsolationDetector {
	return &socketBasedIsolationDetector{socketDir: socketDir, controller: []string{"devices"}}
}

func (s *socketBasedIsolationDetector) Whitelist(controller []string) PodIsolationDetector {
	s.controller = controller
	return s
}

func (s *socketBasedIsolationDetector) Detect(vm *v1.VirtualMachineInstance) (*IsolationResult, error) {
	var pid int
	var slice string
	var err error
	var controller []string

	// Look up the socket of the virt-launcher Pod which was created for that VM, and extract the PID from it
	socket := cmdclient.SocketFromUID(s.socketDir, string(vm.UID))
	if pid, err = s.getPid(socket); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get owner Pid of socket %s", socket)
		return nil, err
	}

	// Look up the cgroup slice based on the whitelisted controller
	if controller, slice, err = s.getSlice(pid); err != nil {
		log.Log.Object(vm).Reason(err).Errorf("Could not get cgroup slice for Pid %d", pid)
		return nil, err
	}

	return NewIsolationResult(pid, slice, controller), nil
}

// standard golang libraries don't provide API to set runtime limits
// for other processes, so we have to directly call to kernel
func prLimit(pid int, limit uintptr, rlimit *unix.Rlimit) error {
	_, _, errno := unix.RawSyscall6(unix.SYS_PRLIMIT64,
		uintptr(pid),
		limit,
		uintptr(unsafe.Pointer(rlimit)),
		0, 0, 0)
	if errno != 0 {
		return fmt.Errorf("Error setting prlimit: %v", errno)
	}
	return nil
}

func (s *socketBasedIsolationDetector) AdjustResources(vm *v1.VirtualMachineInstance) error {
	// only VFIO attached domains require MEMLOCK adjustment
	if !util.IsSRIOVVmi(vm) {
		return nil
	}

	// bump memlock ulimit for libvirtd
	res, err := s.Detect(vm)
	if err != nil {
		return err
	}
	launcherPid := res.Pid()

	processes, err := ps.Processes()
	if err != nil {
		return fmt.Errorf("failed to get all processes: %v", err)
	}

	for _, process := range processes {
		// consider all processes that are virt-launcher children
		if process.PPid() != launcherPid {
			continue
		}

		// libvirtd process sets the memory lock limit before fork/exec-ing into qemu
		if process.Executable() != "libvirtd" {
			continue
		}

		// make the best estimate for memory required by libvirt
		memlockSize, err := getMemlockSize(vm)
		if err != nil {
			return err
		}
		rLimit := unix.Rlimit{
			Max: uint64(memlockSize),
			Cur: uint64(memlockSize),
		}
		err = prLimit(process.Pid(), unix.RLIMIT_MEMLOCK, &rLimit)
		if err != nil {
			return fmt.Errorf("failed to set rlimit for memory lock: %v", err)
		}
		// we assume a single process should match
		break
	}
	return nil
}

// consider reusing getMemoryOverhead()
// This is not scientific, but neither what libvirtd does is. See details in:
// https://www.redhat.com/archives/libvirt-users/2019-August/msg00051.html
func getMemlockSize(vm *v1.VirtualMachineInstance) (int64, error) {
	memlockSize := resource.NewQuantity(0, resource.DecimalSI)

	// start with base memory requested for the VM
	vmiMemoryReq := vm.Spec.Domain.Resources.Requests.Memory()
	memlockSize.Add(*resource.NewScaledQuantity(vmiMemoryReq.ScaledValue(resource.Kilo), resource.Kilo))

	// allocate 1Gb for VFIO needs
	memlockSize.Add(resource.MustParse("1G"))

	// add some more memory for NUMA / CPU topology, platform memory alignment and other needs
	memlockSize.Add(resource.MustParse("256M"))

	bytes_, ok := memlockSize.AsInt64()
	if !ok {
		return 0, fmt.Errorf("could not calculate memory lock size")
	}
	return bytes_, nil
}

func NewIsolationResult(pid int, slice string, controller []string) *IsolationResult {
	return &IsolationResult{pid: pid, slice: slice, controller: controller}
}

type IsolationResult struct {
	pid        int
	slice      string
	controller []string
}

func (r *IsolationResult) Slice() string {
	return r.slice
}

func (r *IsolationResult) PIDNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/pid", r.pid)
}

func (r *IsolationResult) MountNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/mnt", r.pid)
}

func (r *IsolationResult) mountInfo() string {
	return fmt.Sprintf("/proc/%d/mountinfo", r.pid)
}

// MountInfoRoot returns information about the root entry in /proc/mountinfo
func (r *IsolationResult) MountInfoRoot() (*MountInfo, error) {
	in, err := os.Open(r.mountInfo())
	if err != nil {
		return nil, fmt.Errorf("could not open mountinfo: %v", err)
	}
	defer in.Close()
	c := csv.NewReader(in)
	c.Comma = ' '
	c.LazyQuotes = true
	for {
		record, err := c.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			if e, ok := err.(*csv.ParseError); ok {
				if e.Err != csv.ErrFieldCount {
					return nil, err
				}
			} else {
				return nil, err
			}
		}

		if record[3] == "/" && record[4] == "/" {
			return &MountInfo{
				DeviceContainingFile: record[2],
				Root:                 record[3],
				MountPoint:           record[4],
			}, nil
		}
	}

	//impossible
	return nil, fmt.Errorf("process has no root entry")
}

// IsMounted checks if a path in the mount namespace of a
// given process isolation result is a mount point. Works with symlinks.
func (r *IsolationResult) IsMounted(mountPoint string) (bool, error) {
	mountPoint, err := filepath.EvalSymlinks(mountPoint)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("could not resolve mount point path: %v", err)
	}
	in, err := os.Open(r.mountInfo())
	if err != nil {
		return false, fmt.Errorf("could not open mountinfo: %v", err)
	}
	defer in.Close()
	c := csv.NewReader(in)
	c.Comma = ' '
	c.LazyQuotes = true
	for {
		record, err := c.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			if e, ok := err.(*csv.ParseError); ok {
				if e.Err != csv.ErrFieldCount {
					return false, err
				}
			} else {
				return false, err
			}
		}

		if record[4] == mountPoint {
			return true, nil
		}
	}
	return false, nil
}

// ParentMountInfoFor takes the mount info from a container, and looks the corresponding
// entry in /proc/mountinfo of the isolation result of the given process.
func (r *IsolationResult) ParentMountInfoFor(mountInfo *MountInfo) (*MountInfo, error) {
	in, err := os.Open(r.mountInfo())
	if err != nil {
		return nil, fmt.Errorf("could not open mountinfo: %v", err)
	}
	defer in.Close()
	c := csv.NewReader(in)
	c.Comma = ' '
	c.LazyQuotes = true
	for {
		record, err := c.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			if e, ok := err.(*csv.ParseError); ok {
				if e.Err != csv.ErrFieldCount {
					return nil, err
				}
			} else {
				return nil, err
			}
		}

		if record[2] == mountInfo.DeviceContainingFile {
			return &MountInfo{
				DeviceContainingFile: record[2],
				Root:                 record[3],
				MountPoint:           record[4],
			}, nil
		}
	}
	return nil, fmt.Errorf("no parent entry for %v found in the mount namespace of %d", mountInfo.DeviceContainingFile, r.pid)
}

func (r *IsolationResult) NetNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/net", r.pid)
}

func (r *IsolationResult) MountRoot() string {
	return fmt.Sprintf("/proc/%d/root", r.pid)
}

func (r *IsolationResult) Pid() int {
	return r.pid
}

func (r *IsolationResult) Controller() []string {
	return r.controller
}

func (s *socketBasedIsolationDetector) getPid(socket string) (int, error) {
	sock, err := net.Dial("unix", socket)
	if err != nil {
		return -1, err
	}
	defer sock.Close()

	ufile, err := sock.(*net.UnixConn).File()
	if err != nil {
		return -1, err
	}
	// This is the tricky part, which will give us the PID of the owning socket
	ucreds, err := syscall.GetsockoptUcred(int(ufile.Fd()), syscall.SOL_SOCKET, syscall.SO_PEERCRED)
	if err != nil {
		return -1, err
	}

	if int(ucreds.Pid) == 0 {
		return -1, fmt.Errorf("The detected PID is 0. Is the isolation detector running in the host PID namespace?")
	}

	return int(ucreds.Pid), nil
}

func (s *socketBasedIsolationDetector) getSlice(pid int) (controller []string, slice string, err error) {
	cgroups, err := os.Open(fmt.Sprintf("/proc/%d/cgroup", pid))
	if err != nil {
		return
	}
	defer cgroups.Close()

	scanner := bufio.NewScanner(cgroups)
	for scanner.Scan() {
		cgEntry := strings.Split(scanner.Text(), ":")
		// Check if we have a sane cgroup line
		if len(cgEntry) != 3 {
			err = fmt.Errorf("Could not extract slice from cgroup line: %s", scanner.Text())
			return
		}
		// Skip not supported cgroup controller
		if !sliceContains(s.controller, cgEntry[1]) {
			continue
		}

		// Set and check cgroup slice
		if slice == "" {
			slice = cgEntry[2]
		} else if slice != cgEntry[2] {
			err = fmt.Errorf("Process is part of more than one slice. Expected %s, found %s", slice, cgEntry[2])
			return
		}
		// Add controller
		controller = append(controller, cgEntry[1])
	}
	// Check if we encountered a read error
	if scanner.Err() != nil {
		err = scanner.Err()
		return
	}

	if slice == "" {
		err = fmt.Errorf("Could not detect slice of whitelisted controller: %v", s.controller)
		return
	}
	return
}

func sliceContains(controllers []string, value string) bool {
	for _, c := range controllers {
		if c == value {
			return true
		}
	}
	return false
}

func NodeIsolationResult() *IsolationResult {
	return &IsolationResult{
		pid: 1,
	}
}
<|endoftext|>


<commit_before>// +build !providerless

/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package awsebs

import (
	"os"
	"path/filepath"
	"testing"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	utiltesting "k8s.io/client-go/util/testing"
	"k8s.io/kubernetes/pkg/volume"
	volumetest "k8s.io/kubernetes/pkg/volume/testing"
)

const (
	testVolName    = "vol-1234"
	testPVName     = "pv1"
	testGlobalPath = "plugins/kubernetes.io/aws-ebs/volumeDevices/vol-1234"
	testPodPath    = "pods/poduid/volumeDevices/kubernetes.io~aws-ebs"
)

func TestGetVolumeSpecFromGlobalMapPath(t *testing.T) {
	// make our test path for fake GlobalMapPath
	// /tmp symbolized our pluginDir
	// /tmp/testGlobalPathXXXXX/plugins/kubernetes.io/gce-pd/volumeDevices/pdVol1
	tmpVDir, err := utiltesting.MkTmpdir("awsBlockTest")
	if err != nil {
		t.Fatalf("can't make a temp dir: %v", err)
	}
	//deferred clean up
	defer os.RemoveAll(tmpVDir)

	expectedGlobalPath := filepath.Join(tmpVDir, testGlobalPath)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpVDir, nil, nil))
	plug, err := plugMgr.FindMapperPluginByName(awsElasticBlockStorePluginName)
	if err != nil {
		os.RemoveAll(tmpVDir)
		t.Fatalf("Can't find the plugin by name: %q", awsElasticBlockStorePluginName)
	}

	//Bad Path
	badspec, err := plug.(*awsElasticBlockStorePlugin).getVolumeSpecFromGlobalMapPath("", "")
	if badspec != nil || err == nil {
		t.Fatalf("Expected not to get spec from GlobalMapPath but did")
	}

	// Good Path
	spec, err := plug.(*awsElasticBlockStorePlugin).getVolumeSpecFromGlobalMapPath("myVolume", expectedGlobalPath)
	if spec == nil || err != nil {
		t.Fatalf("Failed to get spec from GlobalMapPath: %v", err)
	}
	if spec.PersistentVolume.Name != "myVolume" {
		t.Errorf("Invalid PV name from GlobalMapPath spec: %s", spec.PersistentVolume.Name)
	}
	if spec.PersistentVolume.Spec.AWSElasticBlockStore.VolumeID != testVolName {
		t.Errorf("Invalid volumeID from GlobalMapPath spec: %s", spec.PersistentVolume.Spec.AWSElasticBlockStore.VolumeID)
	}
	block := v1.PersistentVolumeBlock
	specMode := spec.PersistentVolume.Spec.VolumeMode
	if specMode == nil {
		t.Fatalf("Failed to get spec from GlobalMapPath: %v", err)
	}
	if *specMode != block {
		t.Errorf("Invalid volumeMode from GlobalMapPath spec: %v - %v", *specMode, block)
	}
}

func getTestVolume(readOnly bool, isBlock bool) *volume.Spec {
	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: testPVName,
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeSource: v1.PersistentVolumeSource{
				AWSElasticBlockStore: &v1.AWSElasticBlockStoreVolumeSource{
					VolumeID: testVolName,
				},
			},
		},
	}

	if isBlock {
		blockMode := v1.PersistentVolumeBlock
		pv.Spec.VolumeMode = &blockMode
	}
	return volume.NewSpecFromPersistentVolume(pv, readOnly)
}

func TestGetPodAndPluginMapPaths(t *testing.T) {
	tmpVDir, err := utiltesting.MkTmpdir("awsBlockTest")
	if err != nil {
		t.Fatalf("can't make a temp dir: %v", err)
	}
	//deferred clean up
	defer os.RemoveAll(tmpVDir)

	expectedGlobalPath := filepath.Join(tmpVDir, testGlobalPath)
	expectedPodPath := filepath.Join(tmpVDir, testPodPath)

	spec := getTestVolume(false, true /*isBlock*/)
	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpVDir, nil, nil))
	plug, err := plugMgr.FindMapperPluginByName(awsElasticBlockStorePluginName)
	if err != nil {
		os.RemoveAll(tmpVDir)
		t.Fatalf("Can't find the plugin by name: %q", awsElasticBlockStorePluginName)
	}
	if plug.GetPluginName() != awsElasticBlockStorePluginName {
		t.Fatalf("Wrong name: %s", plug.GetPluginName())
	}
	pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{UID: types.UID("poduid")}}
	mapper, err := plug.NewBlockVolumeMapper(spec, pod, volume.VolumeOptions{})
	if err != nil {
		t.Fatalf("Failed to make a new Mounter: %v", err)
	}
	if mapper == nil {
		t.Fatalf("Got a nil Mounter")
	}

	//GetGlobalMapPath
	gMapPath, err := mapper.GetGlobalMapPath(spec)
	if err != nil || len(gMapPath) == 0 {
		t.Fatalf("Invalid path from GlobalMapPath spec: %s", spec.PersistentVolume.Spec.GCEPersistentDisk.PDName)
	}
	if gMapPath != expectedGlobalPath {
		t.Fatalf("Failed to get GlobalMapPath: %s %s", gMapPath, expectedGlobalPath)
	}

	//GetPodDeviceMapPath
	gDevicePath, gVolName := mapper.GetPodDeviceMapPath()
	if gDevicePath != expectedPodPath {
		t.Errorf("Got unexpected pod path: %s, expected %s", gDevicePath, expectedPodPath)
	}
	if gVolName != testPVName {
		t.Errorf("Got unexpected volNamne: %s, expected %s", gVolName, testPVName)
	}
}
<commit_msg>Code review changes<commit_after>// +build !providerless

/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package awsebs

import (
	"os"
	"path/filepath"
	"testing"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	utiltesting "k8s.io/client-go/util/testing"
	"k8s.io/kubernetes/pkg/volume"
	volumetest "k8s.io/kubernetes/pkg/volume/testing"
)

const (
	testVolName    = "vol-1234"
	testPVName     = "pv1"
	testGlobalPath = "plugins/kubernetes.io/aws-ebs/volumeDevices/vol-1234"
	testPodPath    = "pods/poduid/volumeDevices/kubernetes.io~aws-ebs"
)

func TestGetVolumeSpecFromGlobalMapPath(t *testing.T) {
	// make our test path for fake GlobalMapPath
	// /tmp symbolized our pluginDir
	// /tmp/testGlobalPathXXXXX/plugins/kubernetes.io/gce-pd/volumeDevices/pdVol1
	tmpVDir, err := utiltesting.MkTmpdir("awsBlockTest")
	if err != nil {
		t.Fatalf("can't make a temp dir: %v", err)
	}
	//deferred clean up
	defer os.RemoveAll(tmpVDir)

	expectedGlobalPath := filepath.Join(tmpVDir, testGlobalPath)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpVDir, nil, nil))
	plug, err := plugMgr.FindMapperPluginByName(awsElasticBlockStorePluginName)
	if err != nil {
		os.RemoveAll(tmpVDir)
		t.Fatalf("Can't find the plugin by name: %q", awsElasticBlockStorePluginName)
	}

	//Bad Path
	badspec, err := plug.(*awsElasticBlockStorePlugin).getVolumeSpecFromGlobalMapPath("", "")
	if badspec != nil || err == nil {
		t.Fatalf("Expected not to get spec from GlobalMapPath but did")
	}

	// Good Path
	spec, err := plug.(*awsElasticBlockStorePlugin).getVolumeSpecFromGlobalMapPath("myVolume", expectedGlobalPath)
	if spec == nil || err != nil {
		t.Errorf("Invalid volumeMode from GlobalMapPath spec: %v - %v", &specMode, block)
	}
	if spec.PersistentVolume.Name != "myVolume" {
		t.Errorf("Invalid PV name from GlobalMapPath spec: %s", spec.PersistentVolume.Name)
	}
	if spec.PersistentVolume.Spec.AWSElasticBlockStore.VolumeID != testVolName {
		t.Errorf("Invalid volumeID from GlobalMapPath spec: %s", spec.PersistentVolume.Spec.AWSElasticBlockStore.VolumeID)
	}
	block := v1.PersistentVolumeBlock
	specMode := spec.PersistentVolume.Spec.VolumeMode
	if specMode == nil {
		t.Fatalf("Failed to get spec from GlobalMapPath: %v", err)
	}
	if *specMode != block {
		t.Errorf("Invalid volumeMode from GlobalMapPath spec: %v - %v", *specMode, block)
	}
}

func getTestVolume(readOnly bool, isBlock bool) *volume.Spec {
	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: testPVName,
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeSource: v1.PersistentVolumeSource{
				AWSElasticBlockStore: &v1.AWSElasticBlockStoreVolumeSource{
					VolumeID: testVolName,
				},
			},
		},
	}

	if isBlock {
		blockMode := v1.PersistentVolumeBlock
		pv.Spec.VolumeMode = &blockMode
	}
	return volume.NewSpecFromPersistentVolume(pv, readOnly)
}

func TestGetPodAndPluginMapPaths(t *testing.T) {
	tmpVDir, err := utiltesting.MkTmpdir("awsBlockTest")
	if err != nil {
		t.Fatalf("can't make a temp dir: %v", err)
	}
	//deferred clean up
	defer os.RemoveAll(tmpVDir)

	expectedGlobalPath := filepath.Join(tmpVDir, testGlobalPath)
	expectedPodPath := filepath.Join(tmpVDir, testPodPath)

	spec := getTestVolume(false, true /*isBlock*/)
	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpVDir, nil, nil))
	plug, err := plugMgr.FindMapperPluginByName(awsElasticBlockStorePluginName)
	if err != nil {
		os.RemoveAll(tmpVDir)
		t.Fatalf("Can't find the plugin by name: %q", awsElasticBlockStorePluginName)
	}
	if plug.GetPluginName() != awsElasticBlockStorePluginName {
		t.Fatalf("Wrong name: %s", plug.GetPluginName())
	}
	pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{UID: types.UID("poduid")}}
	mapper, err := plug.NewBlockVolumeMapper(spec, pod, volume.VolumeOptions{})
	if err != nil {
		t.Fatalf("Failed to make a new Mounter: %v", err)
	}
	if mapper == nil {
		t.Fatalf("Got a nil Mounter")
	}

	//GetGlobalMapPath
	gMapPath, err := mapper.GetGlobalMapPath(spec)
	if err != nil || len(gMapPath) == 0 {
		t.Fatalf("Invalid path from GlobalMapPath spec: %s", spec.PersistentVolume.Spec.GCEPersistentDisk.PDName)
	}
	if gMapPath != expectedGlobalPath {
		t.Fatalf("Failed to get GlobalMapPath: %s %s", gMapPath, expectedGlobalPath)
	}

	//GetPodDeviceMapPath
	gDevicePath, gVolName := mapper.GetPodDeviceMapPath()
	if gDevicePath != expectedPodPath {
		t.Errorf("Got unexpected pod path: %s, expected %s", gDevicePath, expectedPodPath)
	}
	if gVolName != testPVName {
		t.Errorf("Got unexpected volNamne: %s, expected %s", gVolName, testPVName)
	}
}
<|endoftext|>


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

package godata

import (
	"context"
	"sync"
	"testing"

	"golang.org/x/build/maintner"
)

func BenchmarkGet(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_, err := Get(context.Background())
		if err != nil {
			b.Fatal(err)
		}
	}
}

var (
	corpusMu    sync.Mutex
	corpusCache *maintner.Corpus
)

func getGoData(tb testing.TB) *maintner.Corpus {
	corpusMu.Lock()
	defer corpusMu.Unlock()
	if corpusCache != nil {
		return corpusCache
	}
	var err error
	corpusCache, err = Get(context.Background())
	if err != nil {
		tb.Fatalf("getting corpus: %v", err)
	}
	return corpusCache
}

func TestCorpusCheck(t *testing.T) {
	c := getGoData(t)
	if err := c.Check(); err != nil {
		t.Fatal(err)
	}
}

func TestGerritForeachNonChangeRef(t *testing.T) {
	c := getGoData(t)
	c.Gerrit().ForeachProjectUnsorted(func(gp *maintner.GerritProject) error {
		t.Logf("%s:", gp.ServerSlashProject())
		gp.ForeachNonChangeRef(func(ref string, hash maintner.GitHash) error {
			t.Logf("  %s %s", hash, ref)
			return nil
		})
		return nil
	})
}

// In the past, some Gerrit ref changes came before the git in the log.
// This tests that we handle Gerrit meta changes that happen before
// the referenced git commit is known.
func TestGerritOutOfOrderMetaChanges(t *testing.T) {
	c := getGoData(t)

	// Merged:
	goProj := c.Gerrit().Project("go.googlesource.com", "go")
	cl := goProj.CL(38634)
	if cl == nil {
		t.Fatal("CL 38634 not found")
	}
	if g, w := cl.Status, "merged"; g != w {
		t.Errorf("CL status = %q; want %q", g, w)
	}

	// Deleted:
	gddo := c.Gerrit().Project("go.googlesource.com", "gddo")
	cl = gddo.CL(37452)
	if cl == nil {
		t.Fatal("CL 37452 not found")
	}
	t.Logf("Got: %+v", *cl)
}

func TestGerritSkipPrivateCLs(t *testing.T) {
	c := getGoData(t)
	proj := c.Gerrit().Project("go.googlesource.com", "gddo")
	proj.ForeachOpenCL(func(cl *maintner.GerritCL) error {
		if cl.Number == 37452 {
			t.Error("unexpected private CL 37452")
		}
		return nil
	})
}

func TestGitAncestor(t *testing.T) {
	c := getGoData(t)
	tests := []struct {
		subject, ancestor string
		want              bool
	}{
		{"3b5637ff2bd5c03479780995e7a35c48222157c1", "0bb0b61d6a85b2a1a33dcbc418089656f2754d32", true},
		{"0bb0b61d6a85b2a1a33dcbc418089656f2754d32", "3b5637ff2bd5c03479780995e7a35c48222157c1", false},

		{"8f06e217eac10bae4993ca371ade35fecd26270e", "22f1b56dab29d397d2bdbdd603d85e60fb678089", true},
		{"22f1b56dab29d397d2bdbdd603d85e60fb678089", "8f06e217eac10bae4993ca371ade35fecd26270e", false},

		// Was crashing. Issue 22753.
		{"3a181dc7bc8fd0c61d6090a85f87c934f1874802", "f65abf6ddc8d1f3d403a9195fd74eaffa022b07f", true},
		// The reverse of the above, to try to reproduce the
		// panic if I got the order backwards:
		{"f65abf6ddc8d1f3d403a9195fd74eaffa022b07f", "3a181dc7bc8fd0c61d6090a85f87c934f1874802", false},

		// Same on both sides:
		{"0bb0b61d6a85b2a1a33dcbc418089656f2754d32", "0bb0b61d6a85b2a1a33dcbc418089656f2754d32", false},
		{"3b5637ff2bd5c03479780995e7a35c48222157c1", "3b5637ff2bd5c03479780995e7a35c48222157c1", false},
	}
	for i, tt := range tests {
		subject := c.GitCommit(tt.subject)
		if subject == nil {
			t.Errorf("%d. missing subject commit %q", i, tt.subject)
			continue
		}
		anc := c.GitCommit(tt.ancestor)
		if anc == nil {
			t.Errorf("%d. missing ancestor commit %q", i, tt.ancestor)
			continue
		}
		got := subject.HasAncestor(anc)
		if got != tt.want {
			t.Errorf("HasAncestor(%q, %q) = %v; want %v", tt.subject, tt.ancestor, got, tt.want)
		}
	}
}

func BenchmarkGitAncestor(b *testing.B) {
	c := getGoData(b)
	subject := c.GitCommit("3b5637ff2bd5c03479780995e7a35c48222157c1")
	anc := c.GitCommit("0bb0b61d6a85b2a1a33dcbc418089656f2754d32")
	if subject == nil || anc == nil {
		b.Fatal("missing commit(s)")
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !subject.HasAncestor(anc) {
			b.Fatal("wrong answer")
		}
	}
}
<commit_msg>maintner/godata: skip tests in short mode<commit_after>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package godata

import (
	"context"
	"sync"
	"testing"

	"golang.org/x/build/maintner"
)

func BenchmarkGet(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_, err := Get(context.Background())
		if err != nil {
			b.Fatal(err)
		}
	}
}

var (
	corpusMu    sync.Mutex
	corpusCache *maintner.Corpus
)

func getGoData(tb testing.TB) *maintner.Corpus {
	if testing.Short() {
		tb.Skip("not running tests requiring large download in short mode")
	}
	corpusMu.Lock()
	defer corpusMu.Unlock()
	if corpusCache != nil {
		return corpusCache
	}
	var err error
	corpusCache, err = Get(context.Background())
	if err != nil {
		tb.Fatalf("getting corpus: %v", err)
	}
	return corpusCache
}

func TestCorpusCheck(t *testing.T) {
	c := getGoData(t)
	if err := c.Check(); err != nil {
		t.Fatal(err)
	}
}

func TestGerritForeachNonChangeRef(t *testing.T) {
	c := getGoData(t)
	c.Gerrit().ForeachProjectUnsorted(func(gp *maintner.GerritProject) error {
		t.Logf("%s:", gp.ServerSlashProject())
		gp.ForeachNonChangeRef(func(ref string, hash maintner.GitHash) error {
			t.Logf("  %s %s", hash, ref)
			return nil
		})
		return nil
	})
}

// In the past, some Gerrit ref changes came before the git in the log.
// This tests that we handle Gerrit meta changes that happen before
// the referenced git commit is known.
func TestGerritOutOfOrderMetaChanges(t *testing.T) {
	c := getGoData(t)

	// Merged:
	goProj := c.Gerrit().Project("go.googlesource.com", "go")
	cl := goProj.CL(38634)
	if cl == nil {
		t.Fatal("CL 38634 not found")
	}
	if g, w := cl.Status, "merged"; g != w {
		t.Errorf("CL status = %q; want %q", g, w)
	}

	// Deleted:
	gddo := c.Gerrit().Project("go.googlesource.com", "gddo")
	cl = gddo.CL(37452)
	if cl == nil {
		t.Fatal("CL 37452 not found")
	}
	t.Logf("Got: %+v", *cl)
}

func TestGerritSkipPrivateCLs(t *testing.T) {
	c := getGoData(t)
	proj := c.Gerrit().Project("go.googlesource.com", "gddo")
	proj.ForeachOpenCL(func(cl *maintner.GerritCL) error {
		if cl.Number == 37452 {
			t.Error("unexpected private CL 37452")
		}
		return nil
	})
}

func TestGitAncestor(t *testing.T) {
	c := getGoData(t)
	tests := []struct {
		subject, ancestor string
		want              bool
	}{
		{"3b5637ff2bd5c03479780995e7a35c48222157c1", "0bb0b61d6a85b2a1a33dcbc418089656f2754d32", true},
		{"0bb0b61d6a85b2a1a33dcbc418089656f2754d32", "3b5637ff2bd5c03479780995e7a35c48222157c1", false},

		{"8f06e217eac10bae4993ca371ade35fecd26270e", "22f1b56dab29d397d2bdbdd603d85e60fb678089", true},
		{"22f1b56dab29d397d2bdbdd603d85e60fb678089", "8f06e217eac10bae4993ca371ade35fecd26270e", false},

		// Was crashing. Issue 22753.
		{"3a181dc7bc8fd0c61d6090a85f87c934f1874802", "f65abf6ddc8d1f3d403a9195fd74eaffa022b07f", true},
		// The reverse of the above, to try to reproduce the
		// panic if I got the order backwards:
		{"f65abf6ddc8d1f3d403a9195fd74eaffa022b07f", "3a181dc7bc8fd0c61d6090a85f87c934f1874802", false},

		// Same on both sides:
		{"0bb0b61d6a85b2a1a33dcbc418089656f2754d32", "0bb0b61d6a85b2a1a33dcbc418089656f2754d32", false},
		{"3b5637ff2bd5c03479780995e7a35c48222157c1", "3b5637ff2bd5c03479780995e7a35c48222157c1", false},
	}
	for i, tt := range tests {
		subject := c.GitCommit(tt.subject)
		if subject == nil {
			t.Errorf("%d. missing subject commit %q", i, tt.subject)
			continue
		}
		anc := c.GitCommit(tt.ancestor)
		if anc == nil {
			t.Errorf("%d. missing ancestor commit %q", i, tt.ancestor)
			continue
		}
		got := subject.HasAncestor(anc)
		if got != tt.want {
			t.Errorf("HasAncestor(%q, %q) = %v; want %v", tt.subject, tt.ancestor, got, tt.want)
		}
	}
}

func BenchmarkGitAncestor(b *testing.B) {
	c := getGoData(b)
	subject := c.GitCommit("3b5637ff2bd5c03479780995e7a35c48222157c1")
	anc := c.GitCommit("0bb0b61d6a85b2a1a33dcbc418089656f2754d32")
	if subject == nil || anc == nil {
		b.Fatal("missing commit(s)")
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !subject.HasAncestor(anc) {
			b.Fatal("wrong answer")
		}
	}
}
<|endoftext|>


<commit_before>package acceptance_test

import (
	"cf-pusher/cf_cli_adapter"
	"encoding/json"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

const Timeout_Task_Curl = 1 * time.Minute

type ProxyResponse struct {
	ListenAddresses []string ""
	Port            int
}

var _ = Describe("task connectivity on the overlay network", func() {
	Describe("networking policy", func() {
		var (
			prefix  string
			domain  string
			orgName string
			cfCli   *cf_cli_adapter.Adapter
			proxy1  string
			proxy2  string
		)

		BeforeEach(func() {
			cfCli = &cf_cli_adapter.Adapter{CfCliPath: "cf"}
			prefix = testConfig.Prefix
			domain = config.AppsDomain

			orgName = prefix + "task-org"

			Expect(cf.Cf("delete-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

			spaceName := prefix + "space"
			Expect(cf.Cf("create-space", spaceName, "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName, "-s", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))

			proxy1 = "proxy-task-connectivity-1"
			proxy2 = "proxy-task-connectivity-2"

			pushProxy(proxy1)
			pushProxy(proxy2)

			cfCli.AddNetworkPolicy(proxy1, proxy2, 8080, "tcp")

		})

		AfterEach(func() {
			Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))
		})

		It("allows tasks to talk to app instances", func(done Done) {
			By("getting the overlay ip of proxy2")
			cmd := exec.Command("curl", "--fail", proxy2+"."+domain)
			sess, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess, 5*time.Second).Should(gexec.Exit(0))
			var proxy2Response ProxyResponse
			Expect(json.Unmarshal(sess.Out.Contents(), &proxy2Response)).To(Succeed())
			containerIP := getContainerIP(proxy2Response.ListenAddresses)

			By("Checking that the task associated with proxy1 can connect to proxy2")
			Expect(cf.Cf("run-task", proxy1, `
			while true; do
				if curl --fail "`+containerIP+`:`+strconv.Itoa(proxy2Response.Port)+`" ; then
					exit 0
				fi
			done;
			exit 1
			`).Wait(5 * time.Second)).To(gexec.Exit(0))

			Eventually(func() *gbytes.Buffer {
				return cf.Cf("tasks", proxy1).Wait(5 * time.Second).Out
			}, Timeout_Task_Curl).Should(gbytes.Say("SUCCEEDED"))

			close(done)
		}, 30*60 /* <-- overall spec timeout in seconds */)
	})
})

func getContainerIP(listenAddresses []string) string {
	for _, listenAddr := range listenAddresses {
		if !strings.HasPrefix(listenAddr, "127.0.0.1") {
			return listenAddr
		}
	}

	return ""
}
<commit_msg>Revert "Increase the odds of running acceptance successfully in subsequent runs"<commit_after>package acceptance_test

import (
	"cf-pusher/cf_cli_adapter"
	"encoding/json"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

const Timeout_Task_Curl = 1 * time.Minute

type ProxyResponse struct {
	ListenAddresses []string ""
	Port            int
}

var _ = Describe("task connectivity on the overlay network", func() {
	Describe("networking policy", func() {
		var (
			prefix  string
			domain  string
			orgName string
			cfCli   *cf_cli_adapter.Adapter
			proxy1  string
			proxy2  string
		)

		BeforeEach(func() {
			cfCli = &cf_cli_adapter.Adapter{CfCliPath: "cf"}
			prefix = testConfig.Prefix
			domain = config.AppsDomain

			orgName = prefix + "task-org"
			Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

			spaceName := prefix + "space"
			Expect(cf.Cf("create-space", spaceName, "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName, "-s", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))

			proxy1 = "proxy-task-connectivity-1"
			proxy2 = "proxy-task-connectivity-2"

			pushProxy(proxy1)
			pushProxy(proxy2)

			cfCli.AddNetworkPolicy(proxy1, proxy2, 8080, "tcp")

		})

		AfterEach(func() {
			Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))
		})

		It("allows tasks to talk to app instances", func(done Done) {
			By("getting the overlay ip of proxy2")
			cmd := exec.Command("curl", "--fail", proxy2+"."+domain)
			sess, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess, 5*time.Second).Should(gexec.Exit(0))
			var proxy2Response ProxyResponse
			Expect(json.Unmarshal(sess.Out.Contents(), &proxy2Response)).To(Succeed())
			containerIP := getContainerIP(proxy2Response.ListenAddresses)

			By("Checking that the task associated with proxy1 can connect to proxy2")
			Expect(cf.Cf("run-task", proxy1, `
			while true; do
				if curl --fail "`+containerIP+`:`+strconv.Itoa(proxy2Response.Port)+`" ; then
					exit 0
				fi
			done;
			exit 1
			`).Wait(5 * time.Second)).To(gexec.Exit(0))

			Eventually(func() *gbytes.Buffer {
				return cf.Cf("tasks", proxy1).Wait(5 * time.Second).Out
			}, Timeout_Task_Curl).Should(gbytes.Say("SUCCEEDED"))

			close(done)
		}, 30*60 /* <-- overall spec timeout in seconds */)
	})
})

func getContainerIP(listenAddresses []string) string {
	for _, listenAddr := range listenAddresses {
		if !strings.HasPrefix(listenAddr, "127.0.0.1") {
			return listenAddr
		}
	}

	return ""
}
<|endoftext|>


<commit_before>package flow

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"sync"
	"time"

	"encoding/json"

	"github.com/codelingo/lingo/app/util"
	"github.com/common-nighthawk/go-figure"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/urfave/cli"
)

type DecoratedResult struct {

	// Ctx are the flags and args parsed from the decorated fact
	Ctx *cli.Context

	// Payload is the struct returned form the flow server for the matched
	// fact.
	Payload proto.Message
}

type flowRunner struct {
	cliCtx           *cli.Context
	cliApp           *CLIApp
	decoratorApp     *DecoratorApp
	decoratedResultc chan *DecoratedResult
	errc             chan error
}

type CLIApp struct {
	cli.App
	Request func(*cli.Context) (chan proto.Message, <-chan *UserVar, chan error, func(), error)

	// Help data
	Tagline string
}

type DecoratorApp struct {
	cli.App
	ConfirmDecorated func(*cli.Context, proto.Message) (bool, error)
	SetUserVar       func(*UserVar)

	// Help info
	DecoratorUsage   string
	DecoratorExample string
}

// NewFlow creates a flowRunner from a CLIApp.
// For flows that query the flow server it overrides the action with command function, which
// runs the cliApp's Request function and listens on the channels it returns.
// TODO: move away from flowRunner model by explicitly defining the command function as the
// action in the flow.
func NewFlow(cliApp *CLIApp, decoratorApp *DecoratorApp) *flowRunner {

	// setBaseApp overrides Action with the help action. We don't want this
	// and either want nil or the action that was already set on
	// CLIApp.Action.
	action := cliApp.Action
	setBaseApp(cliApp)
	cliApp.Action = action

	fRunner := &flowRunner{
		cliApp:       cliApp,
		decoratorApp: decoratorApp,
		errc:         make(chan error),
	}
	if fRunner.decoratorApp != nil {
		fRunner.decoratorApp.Setup()
	}
	fRunner.setHelp()

	if fRunner.cliApp.Action == nil {

		// If Action is set, decoratedResultc is not used, nor is it closed.
		// Only set it if we're calling the internal flowRunner.action.
		fRunner.decoratedResultc = make(chan *DecoratedResult)
		fRunner.cliApp.Action = fRunner.action
	}
	return fRunner
}

func (f *flowRunner) setHelp() {
	cli.HelpPrinter = func(w io.Writer, templ string, data interface{}) {
		defer func() {

			// decoratedResultc is nil if the Flow is using it's own action.
			if f.decoratedResultc != nil {
				close(f.decoratedResultc)
			}
		}()
		figure.NewFigure("codelingo", "larry3d", false).Print()
		printHelp(w, CLIAPPHELPTMP, f.cliApp)
		if f.decoratorApp != nil {
			printHelp(w, DECAPPHELPTMP, f.decoratorApp)
		}
		printHelp(w, INFOTMP, f.cliApp)
	}

	f.cliApp.Commands = append(f.cliApp.Commands, cli.Command{
		Name: "_genJSONHelp",
		Action: func(ctx *cli.Context) {

			defer func() {
				// decoratedResultc is nil if the Flow is using it's own action.
				if f.decoratedResultc != nil {
					close(f.decoratedResultc)
				}
			}()

			var cliFlags, decFlags []string
			for _, flag := range f.cliApp.Flags {
				cliFlags = append(cliFlags, flag.String())
			}
			if f.decoratorApp != nil {
				for _, flag := range f.decoratorApp.Flags {
					decFlags = append(decFlags, flag.String())
				}
			}

			helpData := helpData{
				Name:         f.cliApp.Name,
				Tagline:      f.cliApp.Tagline,
				Description:  f.cliApp.Usage,
				LastCompiled: f.cliApp.Compiled.String(),
				Version:      f.cliApp.Version,

				Options: decFlags,
			}

			if f.decoratorApp != nil {
				helpData.Decorator = decHelp{
					Description: f.decoratorApp.Usage,
					// Example: f.decoratorApp.Example,
					Options: decFlags,
				}
			}

			b, err := json.Marshal(helpData)
			if err != nil {
				fmt.Println(err)
				return
			}

			outPath := ctx.Args()[0]
			if err := ioutil.WriteFile(outPath, b, 0644); err != nil {
				fmt.Println(err)
				return
			}
			fmt.Println("Success! JSON Help written to:", outPath)

		},
	})

}

func setBaseApp(cliApp *CLIApp) {

	app := CLIApp{
		App:     *cli.NewApp(),
		Request: cliApp.Request,
	}

	// base settings
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "lingo-file",
			Usage: "the Tenet `FILE` to run rewrite over. If the flag is not set, codelingo.yaml files are read from the branch being rewritten.",
			//	Destination: &language,
		},
	}
	app.Compiled = time.Now()
	app.EnableBashCompletion = true

	// user settings
	app.Name = cliApp.Name
	app.Tagline = cliApp.Tagline
	app.Usage = cliApp.Usage
	app.Flags = append(app.Flags, cliApp.Flags...)
	app.Action = cliApp.Action
	app.Version = cliApp.Version

	*cliApp = app
}

// TODO(waigani) incorrect usage func

// Run runs the CLI app. We assume that cliApp's uses f.command as its action to query
// the flow server and stream back decorated/confirmed results on f.decoratedResultc.
// Special cliApps that don't use the flow server have their own custom actions, in which
// case the chans returned from here will not be closed.
func (f *flowRunner) Run() (chan *DecoratedResult, chan error) {

	go func() {
		defer close(f.errc)
		if err := f.cliApp.Run(os.Args); err != nil {
			f.errc <- err
		}
	}()

	return f.decoratedResultc, f.errc
}

func (f *flowRunner) action(ctx *cli.Context) {
	if err := f.command(ctx); err != nil {
		util.Logger.Debugw(errors.ErrorStack(err))
		util.FatalOSErr(err)
	}

}

func (f *flowRunner) command(ctx *cli.Context) (err error) {
	wg := sync.WaitGroup{}
	wg.Add(1)
	defer func() {
		wg.Wait()
		close(f.decoratedResultc)
	}()
	defer wg.Done()

	resultc, userVarC, errc, cancel, err := f.RunCLI()
	if err != nil {
		return errors.Trace(err)
	}

	// If user is manually confirming results, set a long timeout.
	timeout := time.After(time.Hour * 1)
l:
	for {
		select {
		case err, ok := <-errc:
			if !ok {
				errc = nil
				break
			}

			util.Logger.Debugf("Result error: %s", errors.ErrorStack(err))
			return errors.Trace(err)
		case v, ok := <-userVarC:
			if !ok {
				userVarC = nil
				break
			}

			f.SetUserVar(v)
		case result, ok := <-resultc:
			if !ok {
				resultc = nil
				break
			}

			// TODO(waigani) this is brittle and expects the result struct to have a
			// DecoratorOptions field. We need to refactor the Flow server to return a
			// tuple of [<decorator>, <result>].
			decorator := reflect.Indirect(reflect.ValueOf(result)).FieldByName("DecoratorOptions").String()

			var keep bool
			if ctx.Bool("keep-all") {
				keep = true
			} else {
				keep, err = f.ConfirmDecorated(decorator, result)
				if err != nil {
					cancel()
					return errors.Trace(err)
				}
			}

			if keep {
				wg.Add(1)
				go func(string, proto.Message) {
					defer wg.Done()
					ctx, err := NewCtx(&f.decoratorApp.App, strings.Split(decorator, " ")[1:]...)
					if err != nil {
						cancel()
						f.errc <- err
						util.Logger.Fatalf("error getting decorated context: %q", err)
						return
					}
					util.Logger.Debug("sending result to Flow...")
					f.decoratedResultc <- &DecoratedResult{
						Ctx:     ctx,
						Payload: result,
					}
					util.Logger.Debug("...result sent to Flow")

				}(decorator, result)

			}

		case <-timeout:
			cancel()
			return errors.New("timed out waiting for issue")
		}
		if resultc == nil && errc == nil && userVarC == nil {
			break l
		}
	}

	return
}

func (f *flowRunner) CliCtx() (*cli.Context, error) {
	if f.cliCtx == nil {
		ctx, err := NewCtx(&f.cliApp.App, os.Args[1:]...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		f.cliCtx = ctx
	}
	return f.cliCtx, nil
}

func (f *flowRunner) RunCLI() (chan proto.Message, <-chan *UserVar, chan error, func(), error) {
	ctx, err := f.CliCtx()
	if err != nil {
		return nil, nil, nil, nil, errors.Trace(err)
	}
	if ctx.Bool("debug") {
		err := util.SetDebugLogger()
		if err != nil {
			return nil, nil, nil, nil, errors.Trace(err)
		}
	}

	return f.cliApp.Request(ctx)
}

func (f *flowRunner) ConfirmDecorated(decorator string, payload proto.Message) (bool, error) {
	ctx, err := NewCtx(&f.decoratorApp.App, strings.Split(decorator, " ")...)
	if err != nil {
		return false, errors.Trace(err)
	}

	return f.decoratorApp.ConfirmDecorated(ctx, payload)
}

func (f *flowRunner) SetUserVar(userVar *UserVar) {
	f.decoratorApp.SetUserVar(userVar)
}

// DecoratorArgs parses a decorator string and returns the CLI arguments
func DecoratorArgs(decStr string) []string {
	return strings.Split(decStr, " ")[1:]
}

func NewCtx(app *cli.App, input ...string) (*cli.Context, error) {

	fSet := flag.NewFlagSet(app.Name, flag.ContinueOnError)
	for _, flag := range app.Flags {
		flag.Apply(fSet)
	}

	if err := fSet.Parse(input); err != nil {
		return nil, errors.Trace(err)
	}
	if err := normalizeFlags(app.Flags, fSet); err != nil {
		return nil, errors.Trace(err)
	}

	ctx := cli.NewContext(nil, fSet, nil)

	return ctx, nil
}

func normalizeFlags(flags []cli.Flag, set *flag.FlagSet) error {
	visited := make(map[string]bool)
	set.Visit(func(f *flag.Flag) {
		visited[f.Name] = true
	})
	for _, f := range flags {
		parts := strings.Split(f.GetName(), ",")
		if len(parts) == 1 {
			continue
		}
		var ff *flag.Flag
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if visited[name] {
				if ff != nil {
					return errors.New("Cannot use two forms of the same flag: " + name + " " + ff.Name)
				}
				ff = set.Lookup(name)
			}
		}
		if ff == nil {
			continue
		}
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if !visited[name] {
				copyFlag(name, ff, set)
			}
		}
	}
	return nil
}

func copyFlag(name string, ff *flag.Flag, set *flag.FlagSet) {
	switch ff.Value.(type) {
	case *cli.StringSlice:
	default:
		set.Set(name, ff.Value.String())
	}
}
<commit_msg>Move arg parser to util.<commit_after>package flow

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"sync"
	"time"

	"encoding/json"

	"github.com/codelingo/lingo/app/util"
	"github.com/common-nighthawk/go-figure"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/urfave/cli"
)

type DecoratedResult struct {

	// Ctx are the flags and args parsed from the decorated fact
	Ctx *cli.Context

	// Payload is the struct returned form the flow server for the matched
	// fact.
	Payload proto.Message
}

type flowRunner struct {
	cliCtx           *cli.Context
	cliApp           *CLIApp
	decoratorApp     *DecoratorApp
	decoratedResultc chan *DecoratedResult
	errc             chan error
}

type CLIApp struct {
	cli.App
	Request func(*cli.Context) (chan proto.Message, <-chan *UserVar, chan error, func(), error)

	// Help data
	Tagline string
}

type DecoratorApp struct {
	cli.App
	ConfirmDecorated func(*cli.Context, proto.Message) (bool, error)
	SetUserVar       func(*UserVar)

	// Help info
	DecoratorUsage   string
	DecoratorExample string
}

// NewFlow creates a flowRunner from a CLIApp.
// For flows that query the flow server it overrides the action with command function, which
// runs the cliApp's Request function and listens on the channels it returns.
// TODO: move away from flowRunner model by explicitly defining the command function as the
// action in the flow.
func NewFlow(cliApp *CLIApp, decoratorApp *DecoratorApp) *flowRunner {

	// setBaseApp overrides Action with the help action. We don't want this
	// and either want nil or the action that was already set on
	// CLIApp.Action.
	action := cliApp.Action
	setBaseApp(cliApp)
	cliApp.Action = action

	fRunner := &flowRunner{
		cliApp:       cliApp,
		decoratorApp: decoratorApp,
		errc:         make(chan error),
	}
	if fRunner.decoratorApp != nil {
		fRunner.decoratorApp.Setup()
	}
	fRunner.setHelp()

	if fRunner.cliApp.Action == nil {

		// If Action is set, decoratedResultc is not used, nor is it closed.
		// Only set it if we're calling the internal flowRunner.action.
		fRunner.decoratedResultc = make(chan *DecoratedResult)
		fRunner.cliApp.Action = fRunner.action
	}
	return fRunner
}

func (f *flowRunner) setHelp() {
	cli.HelpPrinter = func(w io.Writer, templ string, data interface{}) {
		defer func() {

			// decoratedResultc is nil if the Flow is using it's own action.
			if f.decoratedResultc != nil {
				close(f.decoratedResultc)
			}
		}()
		figure.NewFigure("codelingo", "larry3d", false).Print()
		printHelp(w, CLIAPPHELPTMP, f.cliApp)
		if f.decoratorApp != nil {
			printHelp(w, DECAPPHELPTMP, f.decoratorApp)
		}
		printHelp(w, INFOTMP, f.cliApp)
	}

	f.cliApp.Commands = append(f.cliApp.Commands, cli.Command{
		Name: "_genJSONHelp",
		Action: func(ctx *cli.Context) {

			defer func() {
				// decoratedResultc is nil if the Flow is using it's own action.
				if f.decoratedResultc != nil {
					close(f.decoratedResultc)
				}
			}()

			var cliFlags, decFlags []string
			for _, flag := range f.cliApp.Flags {
				cliFlags = append(cliFlags, flag.String())
			}
			if f.decoratorApp != nil {
				for _, flag := range f.decoratorApp.Flags {
					decFlags = append(decFlags, flag.String())
				}
			}

			helpData := helpData{
				Name:         f.cliApp.Name,
				Tagline:      f.cliApp.Tagline,
				Description:  f.cliApp.Usage,
				LastCompiled: f.cliApp.Compiled.String(),
				Version:      f.cliApp.Version,

				Options: decFlags,
			}

			if f.decoratorApp != nil {
				helpData.Decorator = decHelp{
					Description: f.decoratorApp.Usage,
					// Example: f.decoratorApp.Example,
					Options: decFlags,
				}
			}

			b, err := json.Marshal(helpData)
			if err != nil {
				fmt.Println(err)
				return
			}

			outPath := ctx.Args()[0]
			if err := ioutil.WriteFile(outPath, b, 0644); err != nil {
				fmt.Println(err)
				return
			}
			fmt.Println("Success! JSON Help written to:", outPath)

		},
	})

}

func setBaseApp(cliApp *CLIApp) {

	app := CLIApp{
		App:     *cli.NewApp(),
		Request: cliApp.Request,
	}

	// base settings
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "lingo-file",
			Usage: "the Tenet `FILE` to run rewrite over. If the flag is not set, codelingo.yaml files are read from the branch being rewritten.",
			//	Destination: &language,
		},
	}
	app.Compiled = time.Now()
	app.EnableBashCompletion = true

	// user settings
	app.Name = cliApp.Name
	app.Tagline = cliApp.Tagline
	app.Usage = cliApp.Usage
	app.Flags = append(app.Flags, cliApp.Flags...)
	app.Action = cliApp.Action
	app.Version = cliApp.Version

	*cliApp = app
}

// TODO(waigani) incorrect usage func

// Run runs the CLI app. We assume that cliApp's uses f.command as its action to query
// the flow server and stream back decorated/confirmed results on f.decoratedResultc.
// Special cliApps that don't use the flow server have their own custom actions, in which
// case the chans returned from here will not be closed.
func (f *flowRunner) Run() (chan *DecoratedResult, chan error) {

	go func() {
		defer close(f.errc)
		if err := f.cliApp.Run(os.Args); err != nil {
			f.errc <- err
		}
	}()

	return f.decoratedResultc, f.errc
}

func (f *flowRunner) action(ctx *cli.Context) {
	if err := f.command(ctx); err != nil {
		util.Logger.Debugw(errors.ErrorStack(err))
		util.FatalOSErr(err)
	}

}

func (f *flowRunner) command(ctx *cli.Context) (err error) {
	wg := sync.WaitGroup{}
	wg.Add(1)
	defer func() {
		wg.Wait()
		close(f.decoratedResultc)
	}()
	defer wg.Done()

	resultc, userVarC, errc, cancel, err := f.RunCLI()
	if err != nil {
		return errors.Trace(err)
	}

	// If user is manually confirming results, set a long timeout.
	timeout := time.After(time.Hour * 1)
l:
	for {
		select {
		case err, ok := <-errc:
			if !ok {
				errc = nil
				break
			}

			util.Logger.Debugf("Result error: %s", errors.ErrorStack(err))
			return errors.Trace(err)
		case v, ok := <-userVarC:
			if !ok {
				userVarC = nil
				break
			}

			f.SetUserVar(v)
		case result, ok := <-resultc:
			if !ok {
				resultc = nil
				break
			}

			// TODO(waigani) this is brittle and expects the result struct to have a
			// DecoratorOptions field. We need to refactor the Flow server to return a
			// tuple of [<decorator>, <result>].
			decorator := reflect.Indirect(reflect.ValueOf(result)).FieldByName("DecoratorOptions").String()

			var keep bool
			if ctx.Bool("keep-all") {
				keep = true
			} else {
				keep, err = f.ConfirmDecorated(decorator, result)
				if err != nil {
					cancel()
					return errors.Trace(err)
				}
			}

			if keep {
				wg.Add(1)
				go func(string, proto.Message) {
					defer wg.Done()
					ctx, err := NewCtx(&f.decoratorApp.App, strings.Split(decorator, " ")[1:]...)
					if err != nil {
						cancel()
						f.errc <- err
						util.Logger.Fatalf("error getting decorated context: %q", err)
						return
					}
					util.Logger.Debug("sending result to Flow...")
					f.decoratedResultc <- &DecoratedResult{
						Ctx:     ctx,
						Payload: result,
					}
					util.Logger.Debug("...result sent to Flow")

				}(decorator, result)

			}

		case <-timeout:
			cancel()
			return errors.New("timed out waiting for issue")
		}
		if resultc == nil && errc == nil && userVarC == nil {
			break l
		}
	}

	return
}

func (f *flowRunner) CliCtx() (*cli.Context, error) {
	if f.cliCtx == nil {
		ctx, err := NewCtx(&f.cliApp.App, os.Args[1:]...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		f.cliCtx = ctx
	}
	return f.cliCtx, nil
}

func (f *flowRunner) RunCLI() (chan proto.Message, <-chan *UserVar, chan error, func(), error) {
	ctx, err := f.CliCtx()
	if err != nil {
		return nil, nil, nil, nil, errors.Trace(err)
	}
	if ctx.Bool("debug") {
		err := util.SetDebugLogger()
		if err != nil {
			return nil, nil, nil, nil, errors.Trace(err)
		}
	}

	return f.cliApp.Request(ctx)
}

func (f *flowRunner) ConfirmDecorated(decorator string, payload proto.Message) (bool, error) {
	ctx, err := NewCtx(&f.decoratorApp.App, strings.Split(decorator, " ")...)
	if err != nil {
		return false, errors.Trace(err)
	}

	return f.decoratorApp.ConfirmDecorated(ctx, payload)
}

func (f *flowRunner) SetUserVar(userVar *UserVar) {
	f.decoratorApp.SetUserVar(userVar)
}

// DecoratorArgs parses a decorator string and returns the CLI arguments
func DecoratorArgs(decStr string) []string {
	return strings.Split(decStr, " ")[1:]
}

// ParseArgs formats args for urfave/cli
// urfave/cli doesn't distinguish between literal and variable args e.g. "this
// arg" vs thisArg. "this arg" gets parsed as two args. The function below
// takes all args as the input. The cli lib pads {} chars which we undo below.
// A better cli lib should be used.
func ParseArgs(args []string) string {
	var finalArg string
	noSpace := map[string]bool{
		"{": true,
		"}": true,
	}

	for _, arg := range args {

		finalArg += arg

		if !noSpace[arg] {
			finalArg += " "
		}

	}
	return strings.TrimRight(finalArg, " ")
}

func NewCtx(app *cli.App, input ...string) (*cli.Context, error) {

	fSet := flag.NewFlagSet(app.Name, flag.ContinueOnError)
	for _, flag := range app.Flags {
		flag.Apply(fSet)
	}

	if err := fSet.Parse(input); err != nil {
		return nil, errors.Trace(err)
	}
	if err := normalizeFlags(app.Flags, fSet); err != nil {
		return nil, errors.Trace(err)
	}

	ctx := cli.NewContext(nil, fSet, nil)

	return ctx, nil
}

func normalizeFlags(flags []cli.Flag, set *flag.FlagSet) error {
	visited := make(map[string]bool)
	set.Visit(func(f *flag.Flag) {
		visited[f.Name] = true
	})
	for _, f := range flags {
		parts := strings.Split(f.GetName(), ",")
		if len(parts) == 1 {
			continue
		}
		var ff *flag.Flag
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if visited[name] {
				if ff != nil {
					return errors.New("Cannot use two forms of the same flag: " + name + " " + ff.Name)
				}
				ff = set.Lookup(name)
			}
		}
		if ff == nil {
			continue
		}
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if !visited[name] {
				copyFlag(name, ff, set)
			}
		}
	}
	return nil
}

func copyFlag(name string, ff *flag.Flag, set *flag.FlagSet) {
	switch ff.Value.(type) {
	case *cli.StringSlice:
	default:
		set.Set(name, ff.Value.String())
	}
}
<|endoftext|>


<commit_before>package mysql

import (
	"time"

	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/dbus/pkg/model"
	"github.com/funkygao/dbus/pkg/myslave"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
)

// MysqlbinlogInput is an input plugin that pretends to be a mysql instance's
// slave and consumes mysql binlog events.
type MysqlbinlogInput struct {
	maxEventLength int
	stopChan       chan struct{}

	slave *myslave.MySlave
	dsn   string
}

func (this *MysqlbinlogInput) Init(config *conf.Conf) {
	this.dsn = config.String("dsn", "")
	this.maxEventLength = config.Int("max_event_length", (1<<20)-100)
	this.stopChan = make(chan struct{})
	this.slave = myslave.New().LoadConfig(config)
	if err := this.slave.AssertValidRowFormat(); err != nil {
		panic(err)
	}
}

func (this *MysqlbinlogInput) Stop(r engine.InputRunner) {
	log.Trace("[%s] stopping...", r.Name())

	close(this.stopChan)
	this.slave.StopReplication()
}

func (this *MysqlbinlogInput) MySlave() *myslave.MySlave {
	return this.slave
}

func (this *MysqlbinlogInput) OnAck(pack *engine.Packet) error {
	return this.slave.MarkAsProcessed(pack.Payload.(*model.RowsEvent))
}

func (this *MysqlbinlogInput) Run(r engine.InputRunner, h engine.PluginHelper) error {
	name := r.Name()
	backoff := time.Second * 5

	if err := r.DeclareResource(this.dsn); err != nil {
		return err
	}

	for {
	RESTART_REPLICATION:

		clusterRebalance := r.RebalanceChannel()
		<-clusterRebalance

		log.Debug("%+v", r.LeadingResources())

		log.Trace("[%s] starting replication...", name)
		if img, err := this.slave.BinlogRowImage(); err != nil {
			log.Error("[%s] %v", name, err)
		} else {
			log.Trace("[%s] binlog row image=%s", name, img)
		}

		ready := make(chan struct{})
		go this.slave.StartReplication(ready)
		select {
		case <-ready:
		case <-this.stopChan:
			log.Trace("[%s] yes sir!", name)
			return nil
		}

		rows := this.slave.Events()
		errors := this.slave.Errors()
		for {
			select {
			case <-this.stopChan:
				log.Trace("[%s] yes sir!", name)
				return nil

			case err := <-errors:
				// e,g.
				// ERROR 1236 (HY000): Could not find first log file name in binary log index file
				// ERROR 1236 (HY000): Could not open log file
				log.Error("[%s] backoff %s: %v", name, backoff, err)
				this.slave.StopReplication()

				select {
				case <-time.After(backoff):
				case <-this.stopChan:
					log.Trace("[%s] yes sir!", name)
					return nil
				}
				goto RESTART_REPLICATION

			case pack, ok := <-r.InChan():
				if !ok {
					log.Trace("[%s] yes sir!", name)
					return nil
				}

				select {
				case err := <-errors:
					// TODO is this necessary?
					log.Error("[%s] backoff %s: %v", name, backoff, err)
					this.slave.StopReplication()

					select {
					case <-time.After(backoff):
					case <-this.stopChan:
						log.Trace("[%s] yes sir!", name)
						return nil
					}
					goto RESTART_REPLICATION

				case row, ok := <-rows:
					if !ok {
						log.Info("[%s] event stream closed", name)
						return nil
					}

					if row.Length() < this.maxEventLength {
						pack.Payload = row
						r.Inject(pack)
					} else {
						// TODO this.slave.MarkAsProcessed(r), also consider batcher partial failure
						log.Warn("[%s] ignored len=%d %s", name, row.Length(), row.MetaInfo())
						pack.Recycle()
					}

				case <-this.stopChan:
					log.Trace("[%s] yes sir!", name)
					return nil
				}
			}
		}
	}

	return nil
}
<commit_msg>MysqlbinlogInput integration with the cluster rebalance<commit_after>package mysql

import (
	"time"

	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/dbus/pkg/model"
	"github.com/funkygao/dbus/pkg/myslave"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
)

// MysqlbinlogInput is an input plugin that pretends to be a mysql instance's
// slave and consumes mysql binlog events.
type MysqlbinlogInput struct {
	maxEventLength int
	stopChan       chan struct{}

	slave *myslave.MySlave
	dsn   string
}

func (this *MysqlbinlogInput) Init(config *conf.Conf) {
	this.dsn = config.String("dsn", "")
	this.maxEventLength = config.Int("max_event_length", (1<<20)-100)
	this.stopChan = make(chan struct{})
	this.slave = myslave.New().LoadConfig(config)
	if err := this.slave.AssertValidRowFormat(); err != nil {
		panic(err)
	}
}

func (this *MysqlbinlogInput) Stop(r engine.InputRunner) {
	log.Trace("[%s] stopping...", r.Name())

	close(this.stopChan)
	this.slave.StopReplication()
}

func (this *MysqlbinlogInput) MySlave() *myslave.MySlave {
	return this.slave
}

func (this *MysqlbinlogInput) OnAck(pack *engine.Packet) error {
	return this.slave.MarkAsProcessed(pack.Payload.(*model.RowsEvent))
}

func (this *MysqlbinlogInput) Run(r engine.InputRunner, h engine.PluginHelper) error {
	name := r.Name()
	backoff := time.Second * 5

	if err := r.DeclareResource(this.dsn); err != nil {
		return err
	}

	clusterRebalance := r.RebalanceChannel()
	<-clusterRebalance

	for {
	RESTART_REPLICATION:

		log.Trace("[%s] starting replication %+v...", name, r.LeadingResources()[0])
		if img, err := this.slave.BinlogRowImage(); err != nil {
			log.Error("[%s] %v", name, err)
		} else {
			log.Trace("[%s] binlog row image=%s", name, img)
		}

		ready := make(chan struct{})
		go this.slave.StartReplication(ready)
		select {
		case <-ready:
		case <-this.stopChan:
			log.Trace("[%s] yes sir!", name)
			return nil
		}

		rows := this.slave.Events()
		errors := this.slave.Errors()
		for {
			select {
			case <-this.stopChan:
				log.Trace("[%s] yes sir!", name)
				return nil

			case err := <-errors:
				// e,g.
				// ERROR 1236 (HY000): Could not find first log file name in binary log index file
				// ERROR 1236 (HY000): Could not open log file
				log.Error("[%s] backoff %s: %v", name, backoff, err)
				this.slave.StopReplication()

				select {
				case <-time.After(backoff):
				case <-this.stopChan:
					log.Trace("[%s] yes sir!", name)
					return nil
				}
				goto RESTART_REPLICATION

			case pack, ok := <-r.InChan():
				if !ok {
					log.Trace("[%s] yes sir!", name)
					return nil
				}

				select {
				case err := <-errors:
					// TODO is this necessary?
					log.Error("[%s] backoff %s: %v", name, backoff, err)
					this.slave.StopReplication()

					select {
					case <-time.After(backoff):
					case <-this.stopChan:
						log.Trace("[%s] yes sir!", name)
						return nil
					}
					goto RESTART_REPLICATION

				case <-clusterRebalance:
					this.slave.StopReplication()
					goto RESTART_REPLICATION

				case row, ok := <-rows:
					if !ok {
						log.Info("[%s] event stream closed", name)
						return nil
					}

					if row.Length() < this.maxEventLength {
						pack.Payload = row
						r.Inject(pack)
					} else {
						// TODO this.slave.MarkAsProcessed(r), also consider batcher partial failure
						log.Warn("[%s] ignored len=%d %s", name, row.Length(), row.MetaInfo())
						pack.Recycle()
					}

				case <-this.stopChan:
					log.Trace("[%s] yes sir!", name)
					return nil
				}
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gdb

import (
	"context"
	"fmt"
	"github.com/gogf/gf/text/gregex"
	"time"

	"github.com/gogf/gf/text/gstr"
)

// Model is the DAO for ORM.
type Model struct {
	db            DB             // Underlying DB interface.
	tx            *TX            // Underlying TX interface.
	schema        string         // Custom database schema.
	linkType      int            // Mark for operation on master or slave.
	tablesInit    string         // Table names when model initialization.
	tables        string         // Operation table names, which can be more than one table names and aliases, like: "user", "user u", "user u, user_detail ud".
	fields        string         // Operation fields, multiple fields joined using char ','.
	fieldsEx      string         // Excluded operation fields, multiple fields joined using char ','.
	extraArgs     []interface{}  // Extra custom arguments for sql.
	whereHolder   []*whereHolder // Condition strings for where operation.
	groupBy       string         // Used for "group by" statement.
	orderBy       string         // Used for "order by" statement.
	having        []interface{}  // Used for "having..." statement.
	start         int            // Used for "select ... start, limit ..." statement.
	limit         int            // Used for "select ... start, limit ..." statement.
	option        int            // Option for extra operation features.
	offset        int            // Offset statement for some databases grammar.
	data          interface{}    // Data for operation, which can be type of map/[]map/struct/*struct/string, etc.
	batch         int            // Batch number for batch Insert/Replace/Save operations.
	filter        bool           // Filter data and where key-value pairs according to the fields of the table.
	lockInfo      string         // Lock for update or in shared lock.
	cacheEnabled  bool           // Enable sql result cache feature.
	cacheDuration time.Duration  // Cache TTL duration.
	cacheName     string         // Cache name for custom operation.
	unscoped      bool           // Disables soft deleting features when select/delete operations.
	safe          bool           // If true, it clones and returns a new model object whenever operation done; or else it changes the attribute of current model.
}

// whereHolder is the holder for where condition preparing.
type whereHolder struct {
	operator int           // Operator for this holder.
	where    interface{}   // Where parameter.
	args     []interface{} // Arguments for where parameter.
}

const (
	OPTION_OMITEMPTY  = 1
	OPTION_ALLOWEMPTY = 2

	linkTypeMaster   = 1
	linkTypeMSlave   = 2
	whereHolderWhere = 1
	whereHolderAnd   = 2
	whereHolderOr    = 3
)

// Table creates and returns a new ORM model from given schema.
// The parameter <table> can be more than one table names, and also alias name, like:
// 1. Table names:
//    Table("user")
//    Table("user u")
//    Table("user, user_detail")
//    Table("user u, user_detail ud")
// 2. Table name with alias: Table("user", "u")
func (c *Core) Table(table ...string) *Model {
	tables := ""
	if len(table) > 1 {
		tables = fmt.Sprintf(
			`%s AS %s`, c.DB.QuotePrefixTableName(table[0]), c.DB.QuoteWord(table[1]),
		)
	} else if len(table) == 1 {
		tables = c.DB.QuotePrefixTableName(table[0])
	} else {
		panic("table cannot be empty")
	}
	return &Model{
		db:         c.DB,
		tablesInit: tables,
		tables:     tables,
		fields:     "*",
		start:      -1,
		offset:     -1,
		option:     OPTION_ALLOWEMPTY,
	}
}

// Model is alias of Core.Table.
// See Core.Table.
func (c *Core) Model(table ...string) *Model {
	return c.DB.Table(table...)
}

// Table acts like Core.Table except it operates on transaction.
// See Core.Table.
func (tx *TX) Table(table ...string) *Model {
	model := tx.db.Table(table...)
	model.db = tx.db
	model.tx = tx
	return model
}

// Model is alias of tx.Table.
// See tx.Table.
func (tx *TX) Model(table ...string) *Model {
	return tx.Table(table...)
}

//
func (m *Model) Ctx(ctx context.Context) *Model {
	if ctx == nil {
		return m
	}
	model := m.getModel()
	model.db = model.db.Ctx(ctx)
	return model
}

// As sets an alias name for current table.
func (m *Model) As(as string) *Model {
	if m.tables != "" {
		model := m.getModel()
		split := " JOIN "
		if gstr.Contains(model.tables, split) {
			// For join table.
			array := gstr.Split(model.tables, split)
			array[len(array)-1], _ = gregex.ReplaceString(`(.+) ON`, fmt.Sprintf(`$1 AS %s ON`, as), array[len(array)-1])
			model.tables = gstr.Join(array, split)
		} else {
			// For base table.
			model.tables = gstr.TrimRight(model.tables) + " AS " + as
		}
		return model
	}
	return m
}

// DB sets/changes the db object for current operation.
func (m *Model) DB(db DB) *Model {
	model := m.getModel()
	model.db = db
	return model
}

// TX sets/changes the transaction for current operation.
func (m *Model) TX(tx *TX) *Model {
	model := m.getModel()
	model.db = tx.db
	model.tx = tx
	return model
}

// Schema sets the schema for current operation.
func (m *Model) Schema(schema string) *Model {
	model := m.getModel()
	model.schema = schema
	return model
}

// Clone creates and returns a new model which is a clone of current model.
// Note that it uses deep-copy for the clone.
func (m *Model) Clone() *Model {
	newModel := (*Model)(nil)
	if m.tx != nil {
		newModel = m.tx.Table(m.tablesInit)
	} else {
		newModel = m.db.Table(m.tablesInit)
	}
	*newModel = *m
	// Deep copy slice attributes.
	if n := len(m.extraArgs); n > 0 {
		newModel.extraArgs = make([]interface{}, n)
		copy(newModel.extraArgs, m.extraArgs)
	}
	if n := len(m.whereHolder); n > 0 {
		newModel.whereHolder = make([]*whereHolder, n)
		copy(newModel.whereHolder, m.whereHolder)
	}
	return newModel
}

// Master marks the following operation on master node.
func (m *Model) Master() *Model {
	model := m.getModel()
	model.linkType = linkTypeMaster
	return model
}

// Slave marks the following operation on slave node.
// Note that it makes sense only if there's any slave node configured.
func (m *Model) Slave() *Model {
	model := m.getModel()
	model.linkType = linkTypeMSlave
	return model
}

// Safe marks this model safe or unsafe. If safe is true, it clones and returns a new model object
// whenever the operation done, or else it changes the attribute of current model.
func (m *Model) Safe(safe ...bool) *Model {
	if len(safe) > 0 {
		m.safe = safe[0]
	} else {
		m.safe = true
	}
	return m
}

// Args sets custom arguments for model operation.
func (m *Model) Args(args ...interface{}) *Model {
	model := m.getModel()
	model.extraArgs = append(model.extraArgs, args)
	return model
}
<commit_msg>add context feature for package gdb<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gdb

import (
	"context"
	"fmt"
	"github.com/gogf/gf/text/gregex"
	"time"

	"github.com/gogf/gf/text/gstr"
)

// Model is the DAO for ORM.
type Model struct {
	db            DB             // Underlying DB interface.
	tx            *TX            // Underlying TX interface.
	schema        string         // Custom database schema.
	linkType      int            // Mark for operation on master or slave.
	tablesInit    string         // Table names when model initialization.
	tables        string         // Operation table names, which can be more than one table names and aliases, like: "user", "user u", "user u, user_detail ud".
	fields        string         // Operation fields, multiple fields joined using char ','.
	fieldsEx      string         // Excluded operation fields, multiple fields joined using char ','.
	extraArgs     []interface{}  // Extra custom arguments for sql.
	whereHolder   []*whereHolder // Condition strings for where operation.
	groupBy       string         // Used for "group by" statement.
	orderBy       string         // Used for "order by" statement.
	having        []interface{}  // Used for "having..." statement.
	start         int            // Used for "select ... start, limit ..." statement.
	limit         int            // Used for "select ... start, limit ..." statement.
	option        int            // Option for extra operation features.
	offset        int            // Offset statement for some databases grammar.
	data          interface{}    // Data for operation, which can be type of map/[]map/struct/*struct/string, etc.
	batch         int            // Batch number for batch Insert/Replace/Save operations.
	filter        bool           // Filter data and where key-value pairs according to the fields of the table.
	lockInfo      string         // Lock for update or in shared lock.
	cacheEnabled  bool           // Enable sql result cache feature.
	cacheDuration time.Duration  // Cache TTL duration.
	cacheName     string         // Cache name for custom operation.
	unscoped      bool           // Disables soft deleting features when select/delete operations.
	safe          bool           // If true, it clones and returns a new model object whenever operation done; or else it changes the attribute of current model.
}

// whereHolder is the holder for where condition preparing.
type whereHolder struct {
	operator int           // Operator for this holder.
	where    interface{}   // Where parameter.
	args     []interface{} // Arguments for where parameter.
}

const (
	OPTION_OMITEMPTY  = 1
	OPTION_ALLOWEMPTY = 2

	linkTypeMaster   = 1
	linkTypeMSlave   = 2
	whereHolderWhere = 1
	whereHolderAnd   = 2
	whereHolderOr    = 3
)

// Table creates and returns a new ORM model from given schema.
// The parameter <table> can be more than one table names, and also alias name, like:
// 1. Table names:
//    Table("user")
//    Table("user u")
//    Table("user, user_detail")
//    Table("user u, user_detail ud")
// 2. Table name with alias: Table("user", "u")
func (c *Core) Table(table ...string) *Model {
	tables := ""
	if len(table) > 1 {
		tables = fmt.Sprintf(
			`%s AS %s`, c.DB.QuotePrefixTableName(table[0]), c.DB.QuoteWord(table[1]),
		)
	} else if len(table) == 1 {
		tables = c.DB.QuotePrefixTableName(table[0])
	} else {
		panic("table cannot be empty")
	}
	return &Model{
		db:         c.DB,
		tablesInit: tables,
		tables:     tables,
		fields:     "*",
		start:      -1,
		offset:     -1,
		option:     OPTION_ALLOWEMPTY,
	}
}

// Model is alias of Core.Table.
// See Core.Table.
func (c *Core) Model(table ...string) *Model {
	return c.DB.Table(table...)
}

// Table acts like Core.Table except it operates on transaction.
// See Core.Table.
func (tx *TX) Table(table ...string) *Model {
	model := tx.db.Table(table...)
	model.db = tx.db
	model.tx = tx
	return model
}

// Model is alias of tx.Table.
// See tx.Table.
func (tx *TX) Model(table ...string) *Model {
	return tx.Table(table...)
}

// Ctx sets the context for current operation.
func (m *Model) Ctx(ctx context.Context) *Model {
	if ctx == nil {
		return m
	}
	model := m.getModel()
	model.db = model.db.Ctx(ctx)
	return model
}

// As sets an alias name for current table.
func (m *Model) As(as string) *Model {
	if m.tables != "" {
		model := m.getModel()
		split := " JOIN "
		if gstr.Contains(model.tables, split) {
			// For join table.
			array := gstr.Split(model.tables, split)
			array[len(array)-1], _ = gregex.ReplaceString(`(.+) ON`, fmt.Sprintf(`$1 AS %s ON`, as), array[len(array)-1])
			model.tables = gstr.Join(array, split)
		} else {
			// For base table.
			model.tables = gstr.TrimRight(model.tables) + " AS " + as
		}
		return model
	}
	return m
}

// DB sets/changes the db object for current operation.
func (m *Model) DB(db DB) *Model {
	model := m.getModel()
	model.db = db
	return model
}

// TX sets/changes the transaction for current operation.
func (m *Model) TX(tx *TX) *Model {
	model := m.getModel()
	model.db = tx.db
	model.tx = tx
	return model
}

// Schema sets the schema for current operation.
func (m *Model) Schema(schema string) *Model {
	model := m.getModel()
	model.schema = schema
	return model
}

// Clone creates and returns a new model which is a clone of current model.
// Note that it uses deep-copy for the clone.
func (m *Model) Clone() *Model {
	newModel := (*Model)(nil)
	if m.tx != nil {
		newModel = m.tx.Table(m.tablesInit)
	} else {
		newModel = m.db.Table(m.tablesInit)
	}
	*newModel = *m
	// Deep copy slice attributes.
	if n := len(m.extraArgs); n > 0 {
		newModel.extraArgs = make([]interface{}, n)
		copy(newModel.extraArgs, m.extraArgs)
	}
	if n := len(m.whereHolder); n > 0 {
		newModel.whereHolder = make([]*whereHolder, n)
		copy(newModel.whereHolder, m.whereHolder)
	}
	return newModel
}

// Master marks the following operation on master node.
func (m *Model) Master() *Model {
	model := m.getModel()
	model.linkType = linkTypeMaster
	return model
}

// Slave marks the following operation on slave node.
// Note that it makes sense only if there's any slave node configured.
func (m *Model) Slave() *Model {
	model := m.getModel()
	model.linkType = linkTypeMSlave
	return model
}

// Safe marks this model safe or unsafe. If safe is true, it clones and returns a new model object
// whenever the operation done, or else it changes the attribute of current model.
func (m *Model) Safe(safe ...bool) *Model {
	if len(safe) > 0 {
		m.safe = safe[0]
	} else {
		m.safe = true
	}
	return m
}

// Args sets custom arguments for model operation.
func (m *Model) Args(args ...interface{}) *Model {
	model := m.getModel()
	model.extraArgs = append(model.extraArgs, args)
	return model
}
<|endoftext|>


<commit_before>package pgx

import (
	"bytes"
	"encoding/binary"
	"errors"
	"net"

	"github.com/jackc/pgx/chunkreader"
)

// msgReader is a helper that reads values from a PostgreSQL message.
type msgReader struct {
	cr        *chunkreader.ChunkReader
	msgType   byte
	msgBody   []byte
	rp        int // read position
	err       error
	log       func(lvl int, msg string, ctx ...interface{})
	shouldLog func(lvl int) bool
}

// fatal tells rc that a Fatal error has occurred
func (r *msgReader) fatal(err error) {
	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.fatal", "error", err, "msgType", r.msgType, "msgBody", r.msgBody, "rp", r.rp)
	}
	r.err = err
}

// rxMsg reads the type and size of the next message.
func (r *msgReader) rxMsg() (byte, error) {
	if r.err != nil {
		return 0, r.err
	}

	header, err := r.cr.Next(5)
	if err != nil {
		if netErr, ok := err.(net.Error); !(ok && netErr.Timeout()) {
			r.fatal(err)
		}
		return 0, err
	}

	r.msgType = header[0]
	bodyLen := int(binary.BigEndian.Uint32(header[1:])) - 4

	r.msgBody, err = r.cr.Next(bodyLen)
	if err != nil {
		if netErr, ok := err.(net.Error); !(ok && netErr.Timeout()) {
			r.fatal(err)
		}
		return 0, err
	}

	r.rp = 0

	return r.msgType, nil
}

func (r *msgReader) readByte() byte {
	if r.err != nil {
		return 0
	}

	if len(r.msgBody)-r.rp < 1 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.msgBody[r.rp]
	r.rp++

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readByte", "value", b, "byteAsString", string(b), "msgType", r.msgType, "rp", r.rp)
	}

	return b
}

func (r *msgReader) readInt16() int16 {
	if r.err != nil {
		return 0
	}

	if len(r.msgBody)-r.rp < 2 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	n := int16(binary.BigEndian.Uint16(r.msgBody[r.rp:]))
	r.rp += 2

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readInt16", "value", n, "msgType", r.msgType, "rp", r.rp)
	}

	return n
}

func (r *msgReader) readInt32() int32 {
	if r.err != nil {
		return 0
	}

	if len(r.msgBody)-r.rp < 4 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	n := int32(binary.BigEndian.Uint32(r.msgBody[r.rp:]))
	r.rp += 4

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readInt32", "value", n, "msgType", r.msgType, "rp", r.rp)
	}

	return n
}

func (r *msgReader) readUint16() uint16 {
	if r.err != nil {
		return 0
	}

	if len(r.msgBody)-r.rp < 2 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	n := binary.BigEndian.Uint16(r.msgBody[r.rp:])
	r.rp += 2

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readUint16", "value", n, "msgType", r.msgType, "rp", r.rp)
	}

	return n
}

func (r *msgReader) readUint32() uint32 {
	if r.err != nil {
		return 0
	}

	if len(r.msgBody)-r.rp < 4 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	n := binary.BigEndian.Uint32(r.msgBody[r.rp:])
	r.rp += 4

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readUint32", "value", n, "msgType", r.msgType, "rp", r.rp)
	}

	return n
}

func (r *msgReader) readInt64() int64 {
	if r.err != nil {
		return 0
	}

	if len(r.msgBody)-r.rp < 8 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	n := int64(binary.BigEndian.Uint64(r.msgBody[r.rp:]))
	r.rp += 8

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readInt64", "value", n, "msgType", r.msgType, "rp", r.rp)
	}

	return n
}

// readCString reads a null terminated string
func (r *msgReader) readCString() string {
	if r.err != nil {
		return ""
	}

	nullIdx := bytes.IndexByte(r.msgBody[r.rp:], 0)
	if nullIdx == -1 {
		r.fatal(errors.New("null terminated string not found"))
		return ""
	}

	s := string(r.msgBody[r.rp : r.rp+nullIdx])
	r.rp += nullIdx + 1

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readCString", "value", s, "msgType", r.msgType, "rp", r.rp)
	}

	return s
}

// readString reads count bytes and returns as string
func (r *msgReader) readString(countI32 int32) string {
	if r.err != nil {
		return ""
	}

	count := int(countI32)

	if len(r.msgBody)-r.rp < count {
		r.fatal(errors.New("read past end of message"))
		return ""
	}

	s := string(r.msgBody[r.rp : r.rp+count])
	r.rp += count

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readString", "value", s, "msgType", r.msgType, "rp", r.rp)
	}

	return s
}

// readBytes reads count bytes and returns as []byte
func (r *msgReader) readBytes(countI32 int32) []byte {
	if r.err != nil {
		return nil
	}

	count := int(countI32)

	if len(r.msgBody)-r.rp < count {
		r.fatal(errors.New("read past end of message"))
		return nil
	}

	b := r.msgBody[r.rp : r.rp+count]
	r.rp += count

	r.cr.KeepLast()

	if r.shouldLog(LogLevelTrace) {
		r.log(LogLevelTrace, "msgReader.readBytes", "value", b, r.msgType, "rp", r.rp)
	}

	return b
}
<commit_msg>Remove unused msg_reader.go<commit_after><|endoftext|>