<commit_before>// Copyright (c) 2012 VMware, Inc.

package sigar

/*
#include <stdlib.h>
#include <sys/sysctl.h>
#include <sys/mount.h>
#include <mach/mach_init.h>
#include <mach/mach_host.h>
#include <mach/host_info.h>
#include <libproc.h>
#include <mach/processor_info.h>
#include <mach/vm_map.h>
*/
import "C"

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"syscall"
	"time"
	"unsafe"
)

func (self *LoadAverage) Get() error {
	avg := []C.double{0, 0, 0}

	C.getloadavg(&avg[0], C.int(len(avg)))

	self.One = float64(avg[0])
	self.Five = float64(avg[1])
	self.Fifteen = float64(avg[2])

	return nil
}

func (self *Uptime) Get() error {
	tv := syscall.Timeval32{}

	if err := sysctlbyname("kern.boottime", &tv); err != nil {
		return err
	}

	self.Length = time.Since(time.Unix(int64(tv.Sec), int64(tv.Usec)*1000)).Seconds()

	return nil
}

func (self *Mem) Get() error {
	var vmstat C.vm_statistics_data_t

	if err := sysctlbyname("hw.memsize", &self.Total); err != nil {
		return err
	}

	if err := vm_info(&vmstat); err != nil {
		return err
	}

	kern := uint64(vmstat.inactive_count) << 12
	self.Free = uint64(vmstat.free_count) << 12

	self.Used = self.Total - self.Free
	self.ActualFree = self.Free + kern
	self.ActualUsed = self.Used - kern

	return nil
}

type xsw_usage struct {
	Total, Avail, Used uint64
}

func (self *Swap) Get() error {
	sw_usage := xsw_usage{}

	if err := sysctlbyname("vm.swapusage", &sw_usage); err != nil {
		return err
	}

	self.Total = sw_usage.Total
	self.Used = sw_usage.Used
	self.Free = sw_usage.Avail

	return nil
}

func (self *Cpu) Get() error {
	var count C.mach_msg_type_number_t = C.HOST_CPU_LOAD_INFO_COUNT
	var cpuload C.host_cpu_load_info_data_t

	status := C.host_statistics(C.host_t(C.mach_host_self()),
		C.HOST_CPU_LOAD_INFO,
		C.host_info_t(unsafe.Pointer(&cpuload)),
		&count)

	if status != C.KERN_SUCCESS {
		return fmt.Errorf("host_statistics error=%d", status)
	}

	self.User = uint64(cpuload.cpu_ticks[C.CPU_STATE_USER])
	self.Sys = uint64(cpuload.cpu_ticks[C.CPU_STATE_SYSTEM])
	self.Idle = uint64(cpuload.cpu_ticks[C.CPU_STATE_IDLE])
	self.Nice = uint64(cpuload.cpu_ticks[C.CPU_STATE_NICE])

	return nil
}

func (self *CpuList) Get() error {
	var count C.mach_msg_type_number_t
	var cpuload *C.processor_cpu_load_info_data_t
	var ncpu C.natural_t

	status := C.host_processor_info(C.host_t(C.mach_host_self()),
		C.PROCESSOR_CPU_LOAD_INFO,
		&ncpu,
		(*C.processor_info_array_t)(unsafe.Pointer(&cpuload)),
		&count)

	if status != C.KERN_SUCCESS {
		return fmt.Errorf("host_processor_info error=%d", status)
	}

	// jump through some cgo casting hoops and ensure we properly free
	// the memory that cpuload points to
	target := C.vm_map_t(C.mach_task_self_)
	address := C.vm_address_t(uintptr(unsafe.Pointer(cpuload)))
	defer C.vm_deallocate(target, address, C.vm_size_t(ncpu))

	// the body of struct processor_cpu_load_info
	// aka processor_cpu_load_info_data_t
	var cpu_ticks [C.CPU_STATE_MAX]uint32

	// copy the cpuload array to a []byte buffer
	// where we can binary.Read the data
	size := int(ncpu) * binary.Size(cpu_ticks)
	buf := C.GoBytes(unsafe.Pointer(cpuload), C.int(size))

	bbuf := bytes.NewBuffer(buf)

	self.List = make([]Cpu, 0, ncpu)

	for i := 0; i < int(ncpu); i++ {
		cpu := Cpu{}

		err := binary.Read(bbuf, binary.LittleEndian, &cpu_ticks)
		if err != nil {
			return err
		}

		cpu.User = uint64(cpu_ticks[C.CPU_STATE_USER])
		cpu.Sys = uint64(cpu_ticks[C.CPU_STATE_SYSTEM])
		cpu.Idle = uint64(cpu_ticks[C.CPU_STATE_IDLE])
		cpu.Nice = uint64(cpu_ticks[C.CPU_STATE_NICE])

		self.List = append(self.List, cpu)
	}

	return nil
}

func (self *FileSystemList) Get() error {
	num, err := getfsstat(nil, C.MNT_NOWAIT)
	if num < 0 {
		return err
	}

	buf := make([]syscall.Statfs_t, num)

	num, err = getfsstat(buf, C.MNT_NOWAIT)
	if err != nil {
		return err
	}

	fslist := make([]FileSystem, 0, num)

	for i := 0; i < num; i++ {
		fs := FileSystem{}

		fs.DirName = bytePtrToString(&buf[i].Mntonname[0])
		fs.DevName = bytePtrToString(&buf[i].Mntfromname[0])
		fs.SysTypeName = bytePtrToString(&buf[i].Fstypename[0])

		fslist = append(fslist, fs)
	}

	self.List = fslist

	return err
}

func (self *DiskList) Get() error {
	return notImplemented()
}

func (self *ProcList) Get() error {
	n := C.proc_listpids(C.PROC_ALL_PIDS, 0, nil, 0)
	if n <= 0 {
		return syscall.EINVAL
	}
	buf := make([]byte, n)
	n = C.proc_listpids(C.PROC_ALL_PIDS, 0, unsafe.Pointer(&buf[0]), n)
	if n <= 0 {
		return syscall.ENOMEM
	}

	var pid int32
	num := int(n) / binary.Size(pid)
	list := make([]int, 0, num)
	bbuf := bytes.NewBuffer(buf)

	for i := 0; i < num; i++ {
		if err := binary.Read(bbuf, binary.LittleEndian, &pid); err != nil {
			return err
		}
		if pid == 0 {
			continue
		}

		list = append(list, int(pid))
	}

	self.List = list

	return nil
}

func (self *ProcState) Get(pid int) error {
	info := C.struct_proc_taskallinfo{}

	if err := task_info(pid, &info); err != nil {
		return err
	}

	self.Name = C.GoString(&info.pbsd.pbi_comm[0])

	switch info.pbsd.pbi_status {
	case C.SIDL:
		self.State = RunStateIdle
	case C.SRUN:
		self.State = RunStateRun
	case C.SSLEEP:
		self.State = RunStateSleep
	case C.SSTOP:
		self.State = RunStateStop
	case C.SZOMB:
		self.State = RunStateZombie
	default:
		self.State = RunStateUnknown
	}

	self.Ppid = int(info.pbsd.pbi_ppid)

	self.Tty = int(info.pbsd.e_tdev)

	self.Priority = int(info.ptinfo.pti_priority)

	self.Nice = int(info.pbsd.pbi_nice)

	return nil
}

func (self *ProcMem) Get(pid int) error {
	info := C.struct_proc_taskallinfo{}

	if err := task_info(pid, &info); err != nil {
		return err
	}

	self.Size = uint64(info.ptinfo.pti_virtual_size)
	self.Resident = uint64(info.ptinfo.pti_resident_size)
	self.PageFaults = uint64(info.ptinfo.pti_faults)

	return nil
}

func (self *ProcTime) Get(pid int) error {
	info := C.struct_proc_taskallinfo{}

	if err := task_info(pid, &info); err != nil {
		return err
	}

	self.User =
		uint64(info.ptinfo.pti_total_user) / uint64(time.Millisecond)

	self.Sys =
		uint64(info.ptinfo.pti_total_system) / uint64(time.Millisecond)

	self.Total = self.User + self.Sys

	self.StartTime = (uint64(info.pbsd.pbi_start_tvsec) * 1000) +
		(uint64(info.pbsd.pbi_start_tvusec) / 1000)

	return nil
}

func (self *ProcArgs) Get(pid int) error {
	var args []string

	argv := func(arg string) {
		args = append(args, arg)
	}

	err := kern_procargs(pid, nil, argv, nil)

	self.List = args

	return err
}

func (self *ProcExe) Get(pid int) error {
	exe := func(arg string) {
		self.Name = arg
	}

	return kern_procargs(pid, exe, nil, nil)
}

// wrapper around sysctl KERN_PROCARGS2
// callbacks params are optional,
// up to the caller as to which pieces of data they want
func kern_procargs(pid int,
	exe func(string),
	argv func(string),
	env func(string, string)) error {

	mib := []C.int{C.CTL_KERN, C.KERN_PROCARGS2, C.int(pid)}
	argmax := uintptr(C.ARG_MAX)
	buf := make([]byte, argmax)
	err := sysctl(mib, &buf[0], &argmax, nil, 0)
	if err != nil {
		return nil
	}

	bbuf := bytes.NewBuffer(buf)
	bbuf.Truncate(int(argmax))

	var argc int32
	binary.Read(bbuf, binary.LittleEndian, &argc)

	path, err := bbuf.ReadBytes(0)
	if exe != nil {
		exe(string(chop(path)))
	}

	// skip trailing \0's
	for {
		c, _ := bbuf.ReadByte()
		if c != 0 {
			bbuf.UnreadByte()
			break // start of argv[0]
		}
	}

	for i := 0; i < int(argc); i++ {
		arg, err := bbuf.ReadBytes(0)
		if err == io.EOF {
			break
		}
		if argv != nil {
			argv(string(chop(arg)))
		}
	}

	if env == nil {
		return nil
	}

	delim := []byte{61} // "="

	for {
		line, err := bbuf.ReadBytes(0)
		if err == io.EOF || line[0] == 0 {
			break
		}
		pair := bytes.SplitN(chop(line), delim, 2)
		env(string(pair[0]), string(pair[1]))
	}

	return nil
}

// XXX copied from zsyscall_darwin_amd64.go
func sysctl(mib []C.int, old *byte, oldlen *uintptr,
	new *byte, newlen uintptr) (err error) {
	var p0 unsafe.Pointer
	p0 = unsafe.Pointer(&mib[0])
	_, _, e1 := syscall.Syscall6(syscall.SYS___SYSCTL, uintptr(p0),
		uintptr(len(mib)),
		uintptr(unsafe.Pointer(old)), uintptr(unsafe.Pointer(oldlen)),
		uintptr(unsafe.Pointer(new)), uintptr(newlen))
	if e1 != 0 {
		err = e1
	}
	return
}

func vm_info(vmstat *C.vm_statistics_data_t) error {
	var count C.mach_msg_type_number_t = C.HOST_VM_INFO_COUNT

	status := C.host_statistics(
		C.host_t(C.mach_host_self()),
		C.HOST_VM_INFO,
		C.host_info_t(unsafe.Pointer(vmstat)),
		&count)

	if status != C.KERN_SUCCESS {
		return fmt.Errorf("host_statistics=%d", status)
	}

	return nil
}

// generic Sysctl buffer unmarshalling
func sysctlbyname(name string, data interface{}) (err error) {
	val, err := syscall.Sysctl(name)
	if err != nil {
		return err
	}

	buf := []byte(val)

	switch v := data.(type) {
	case *uint64:
		*v = *(*uint64)(unsafe.Pointer(&buf[0]))
		return
	}

	bbuf := bytes.NewBuffer([]byte(val))
	return binary.Read(bbuf, binary.LittleEndian, data)
}

// syscall.Getfsstat() wrapper is broken, roll our own to workaround.
func getfsstat(buf []syscall.Statfs_t, flags int) (n int, err error) {
	var ptr uintptr
	var size uintptr

	if len(buf) > 0 {
		ptr = uintptr(unsafe.Pointer(&buf[0]))
		size = unsafe.Sizeof(buf[0]) * uintptr(len(buf))
	} else {
		ptr = uintptr(0)
		size = uintptr(0)
	}

	trap := uintptr(syscall.SYS_GETFSSTAT64)
	ret, _, errno := syscall.Syscall(trap, ptr, size, uintptr(flags))

	n = int(ret)
	if errno != 0 {
		err = errno
	}

	return
}

func task_info(pid int, info *C.struct_proc_taskallinfo) error {
	size := C.int(unsafe.Sizeof(*info))
	ptr := unsafe.Pointer(info)

	n := C.proc_pidinfo(C.int(pid), C.PROC_PIDTASKALLINFO, 0, ptr, size)
	if n != size {
		return syscall.ENOMEM
	}

	return nil
}

func notImplemented() error {
	panic("Not Implemented")
	return nil
}
<commit_msg>ROCANA-3338 Remove unsupported SIGAR methods on Darwin<commit_after>// Copyright (c) 2012 VMware, Inc.

package sigar

/*
#include <stdlib.h>
#include <sys/sysctl.h>
#include <sys/mount.h>
#include <mach/mach_init.h>
#include <mach/mach_host.h>
#include <mach/host_info.h>
#include <libproc.h>
#include <mach/processor_info.h>
#include <mach/vm_map.h>
*/
import "C"

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"syscall"
	"time"
	"unsafe"
)

func (self *LoadAverage) Get() error {
	avg := []C.double{0, 0, 0}

	C.getloadavg(&avg[0], C.int(len(avg)))

	self.One = float64(avg[0])
	self.Five = float64(avg[1])
	self.Fifteen = float64(avg[2])

	return nil
}

func (self *Uptime) Get() error {
	tv := syscall.Timeval32{}

	if err := sysctlbyname("kern.boottime", &tv); err != nil {
		return err
	}

	self.Length = time.Since(time.Unix(int64(tv.Sec), int64(tv.Usec)*1000)).Seconds()

	return nil
}

func (self *Mem) Get() error {
	var vmstat C.vm_statistics_data_t

	if err := sysctlbyname("hw.memsize", &self.Total); err != nil {
		return err
	}

	if err := vm_info(&vmstat); err != nil {
		return err
	}

	kern := uint64(vmstat.inactive_count) << 12
	self.Free = uint64(vmstat.free_count) << 12

	self.Used = self.Total - self.Free
	self.ActualFree = self.Free + kern
	self.ActualUsed = self.Used - kern

	return nil
}

type xsw_usage struct {
	Total, Avail, Used uint64
}

func (self *Swap) Get() error {
	sw_usage := xsw_usage{}

	if err := sysctlbyname("vm.swapusage", &sw_usage); err != nil {
		return err
	}

	self.Total = sw_usage.Total
	self.Used = sw_usage.Used
	self.Free = sw_usage.Avail

	return nil
}

func (self *Cpu) Get() error {
	return notImplemented()
}

func (self *CpuList) Get() error {
	var count C.mach_msg_type_number_t
	var cpuload *C.processor_cpu_load_info_data_t
	var ncpu C.natural_t

	status := C.host_processor_info(C.host_t(C.mach_host_self()),
		C.PROCESSOR_CPU_LOAD_INFO,
		&ncpu,
		(*C.processor_info_array_t)(unsafe.Pointer(&cpuload)),
		&count)

	if status != C.KERN_SUCCESS {
		return fmt.Errorf("host_processor_info error=%d", status)
	}

	// jump through some cgo casting hoops and ensure we properly free
	// the memory that cpuload points to
	target := C.vm_map_t(C.mach_task_self_)
	address := C.vm_address_t(uintptr(unsafe.Pointer(cpuload)))
	defer C.vm_deallocate(target, address, C.vm_size_t(ncpu))

	// the body of struct processor_cpu_load_info
	// aka processor_cpu_load_info_data_t
	var cpu_ticks [C.CPU_STATE_MAX]uint32

	// copy the cpuload array to a []byte buffer
	// where we can binary.Read the data
	size := int(ncpu) * binary.Size(cpu_ticks)
	buf := C.GoBytes(unsafe.Pointer(cpuload), C.int(size))

	bbuf := bytes.NewBuffer(buf)

	self.List = make([]Cpu, 0, ncpu)

	for i := 0; i < int(ncpu); i++ {
		cpu := Cpu{}

		err := binary.Read(bbuf, binary.LittleEndian, &cpu_ticks)
		if err != nil {
			return err
		}

		cpu.User = uint64(cpu_ticks[C.CPU_STATE_USER])
		cpu.Sys = uint64(cpu_ticks[C.CPU_STATE_SYSTEM])
		cpu.Idle = uint64(cpu_ticks[C.CPU_STATE_IDLE])
		cpu.Nice = uint64(cpu_ticks[C.CPU_STATE_NICE])

		self.List = append(self.List, cpu)
	}

	return nil
}

func (self *FileSystemList) Get() error {
	num, err := getfsstat(nil, C.MNT_NOWAIT)
	if num < 0 {
		return err
	}

	buf := make([]syscall.Statfs_t, num)

	num, err = getfsstat(buf, C.MNT_NOWAIT)
	if err != nil {
		return err
	}

	fslist := make([]FileSystem, 0, num)

	for i := 0; i < num; i++ {
		fs := FileSystem{}

		fs.DirName = bytePtrToString(&buf[i].Mntonname[0])
		fs.DevName = bytePtrToString(&buf[i].Mntfromname[0])
		fs.SysTypeName = bytePtrToString(&buf[i].Fstypename[0])

		fslist = append(fslist, fs)
	}

	self.List = fslist

	return err
}

func (self *DiskList) Get() error {
	return notImplemented()
}

func (self *ProcList) Get() error {
	n := C.proc_listpids(C.PROC_ALL_PIDS, 0, nil, 0)
	if n <= 0 {
		return syscall.EINVAL
	}
	buf := make([]byte, n)
	n = C.proc_listpids(C.PROC_ALL_PIDS, 0, unsafe.Pointer(&buf[0]), n)
	if n <= 0 {
		return syscall.ENOMEM
	}

	var pid int32
	num := int(n) / binary.Size(pid)
	list := make([]int, 0, num)
	bbuf := bytes.NewBuffer(buf)

	for i := 0; i < num; i++ {
		if err := binary.Read(bbuf, binary.LittleEndian, &pid); err != nil {
			return err
		}
		if pid == 0 {
			continue
		}

		list = append(list, int(pid))
	}

	self.List = list

	return nil
}

func (self *ProcState) Get(pid int) error {
  return notImplemented()
}

func (self *ProcMem) Get(pid int) error {
	info := C.struct_proc_taskallinfo{}

	if err := task_info(pid, &info); err != nil {
		return err
	}

	self.Size = uint64(info.ptinfo.pti_virtual_size)
	self.Resident = uint64(info.ptinfo.pti_resident_size)
	self.PageFaults = uint64(info.ptinfo.pti_faults)

	return nil
}

func (self *ProcTime) Get(pid int) error {
	info := C.struct_proc_taskallinfo{}

	if err := task_info(pid, &info); err != nil {
		return err
	}

	self.User =
		uint64(info.ptinfo.pti_total_user) / uint64(time.Millisecond)

	self.Sys =
		uint64(info.ptinfo.pti_total_system) / uint64(time.Millisecond)

	self.Total = self.User + self.Sys

	self.StartTime = (uint64(info.pbsd.pbi_start_tvsec) * 1000) +
		(uint64(info.pbsd.pbi_start_tvusec) / 1000)

	return nil
}

func (self *ProcArgs) Get(pid int) error {
	var args []string

	argv := func(arg string) {
		args = append(args, arg)
	}

	err := kern_procargs(pid, nil, argv, nil)

	self.List = args

	return err
}

func (self *ProcExe) Get(pid int) error {
	exe := func(arg string) {
		self.Name = arg
	}

	return kern_procargs(pid, exe, nil, nil)
}

// wrapper around sysctl KERN_PROCARGS2
// callbacks params are optional,
// up to the caller as to which pieces of data they want
func kern_procargs(pid int,
	exe func(string),
	argv func(string),
	env func(string, string)) error {

	mib := []C.int{C.CTL_KERN, C.KERN_PROCARGS2, C.int(pid)}
	argmax := uintptr(C.ARG_MAX)
	buf := make([]byte, argmax)
	err := sysctl(mib, &buf[0], &argmax, nil, 0)
	if err != nil {
		return nil
	}

	bbuf := bytes.NewBuffer(buf)
	bbuf.Truncate(int(argmax))

	var argc int32
	binary.Read(bbuf, binary.LittleEndian, &argc)

	path, err := bbuf.ReadBytes(0)
	if exe != nil {
		exe(string(chop(path)))
	}

	// skip trailing \0's
	for {
		c, _ := bbuf.ReadByte()
		if c != 0 {
			bbuf.UnreadByte()
			break // start of argv[0]
		}
	}

	for i := 0; i < int(argc); i++ {
		arg, err := bbuf.ReadBytes(0)
		if err == io.EOF {
			break
		}
		if argv != nil {
			argv(string(chop(arg)))
		}
	}

	if env == nil {
		return nil
	}

	delim := []byte{61} // "="

	for {
		line, err := bbuf.ReadBytes(0)
		if err == io.EOF || line[0] == 0 {
			break
		}
		pair := bytes.SplitN(chop(line), delim, 2)
		env(string(pair[0]), string(pair[1]))
	}

	return nil
}

// XXX copied from zsyscall_darwin_amd64.go
func sysctl(mib []C.int, old *byte, oldlen *uintptr,
	new *byte, newlen uintptr) (err error) {
	var p0 unsafe.Pointer
	p0 = unsafe.Pointer(&mib[0])
	_, _, e1 := syscall.Syscall6(syscall.SYS___SYSCTL, uintptr(p0),
		uintptr(len(mib)),
		uintptr(unsafe.Pointer(old)), uintptr(unsafe.Pointer(oldlen)),
		uintptr(unsafe.Pointer(new)), uintptr(newlen))
	if e1 != 0 {
		err = e1
	}
	return
}

func vm_info(vmstat *C.vm_statistics_data_t) error {
	var count C.mach_msg_type_number_t = C.HOST_VM_INFO_COUNT

	status := C.host_statistics(
		C.host_t(C.mach_host_self()),
		C.HOST_VM_INFO,
		C.host_info_t(unsafe.Pointer(vmstat)),
		&count)

	if status != C.KERN_SUCCESS {
		return fmt.Errorf("host_statistics=%d", status)
	}

	return nil
}

// generic Sysctl buffer unmarshalling
func sysctlbyname(name string, data interface{}) (err error) {
	val, err := syscall.Sysctl(name)
	if err != nil {
		return err
	}

	buf := []byte(val)

	switch v := data.(type) {
	case *uint64:
		*v = *(*uint64)(unsafe.Pointer(&buf[0]))
		return
	}

	bbuf := bytes.NewBuffer([]byte(val))
	return binary.Read(bbuf, binary.LittleEndian, data)
}

// syscall.Getfsstat() wrapper is broken, roll our own to workaround.
func getfsstat(buf []syscall.Statfs_t, flags int) (n int, err error) {
	var ptr uintptr
	var size uintptr

	if len(buf) > 0 {
		ptr = uintptr(unsafe.Pointer(&buf[0]))
		size = unsafe.Sizeof(buf[0]) * uintptr(len(buf))
	} else {
		ptr = uintptr(0)
		size = uintptr(0)
	}

	trap := uintptr(syscall.SYS_GETFSSTAT64)
	ret, _, errno := syscall.Syscall(trap, ptr, size, uintptr(flags))

	n = int(ret)
	if errno != 0 {
		err = errno
	}

	return
}

func task_info(pid int, info *C.struct_proc_taskallinfo) error {
	size := C.int(unsafe.Sizeof(*info))
	ptr := unsafe.Pointer(info)

	n := C.proc_pidinfo(C.int(pid), C.PROC_PIDTASKALLINFO, 0, ptr, size)
	if n != size {
		return syscall.ENOMEM
	}

	return nil
}

func notImplemented() error {
	panic("Not Implemented")
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"testing"

	"github.com/keybase/client/go/externals"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	insecureTriplesec "github.com/keybase/go-triplesec-insecure"
)

func SetupEngineTest(tb testing.TB, name string) libkb.TestContext {
	tc := externals.SetupTest(tb, name, 2)
	tc.G.NewTriplesec = func(passphrase []byte, salt []byte) (libkb.Triplesec, error) {
		warner := func() { tc.G.Log.Warning("Installing insecure Triplesec with weak stretch parameters") }
		isProduction := func() bool {
			return tc.G.Env.GetRunMode() == libkb.ProductionRunMode
		}
		return insecureTriplesec.NewCipher(passphrase, salt, warner, isProduction)
	}
	return tc
}

func SetupEngineTestRealTriplesec(tb testing.TB, name string) libkb.TestContext {
	tc := externals.SetupTest(tb, name, 2)
	tc.G.NewTriplesec = libkb.NewSecureTriplesec
	return tc
}

type FakeUser struct {
	Username      string
	Email         string
	Passphrase    string
	User          *libkb.User
	EncryptionKey libkb.GenericKey
}

func NewFakeUser(prefix string) (fu *FakeUser, err error) {
	buf := make([]byte, 5)
	if _, err = rand.Read(buf); err != nil {
		return
	}
	username := fmt.Sprintf("%s_%s", prefix, hex.EncodeToString(buf))
	email := fmt.Sprintf("%s@noemail.keybase.io", username)
	buf = make([]byte, 12)
	if _, err = rand.Read(buf); err != nil {
		return
	}
	passphrase := hex.EncodeToString(buf)
	fu = &FakeUser{Username: username, Email: email, Passphrase: passphrase}
	return
}

func (fu FakeUser) NormalizedUsername() libkb.NormalizedUsername {
	return libkb.NewNormalizedUsername(fu.Username)
}

func (fu FakeUser) UID() keybase1.UID {
	// All new-style names will have a 1-to-1 mapping
	return libkb.UsernameToUID(fu.Username)
}

func NewFakeUserOrBust(tb testing.TB, prefix string) (fu *FakeUser) {
	var err error
	if fu, err = NewFakeUser(prefix); err != nil {
		tb.Fatal(err)
	}
	return fu
}

const defaultDeviceName = "my device"

// MakeTestSignupEngineRunArg fills a SignupEngineRunArg with the most
// common parameters for testing and returns it.
func MakeTestSignupEngineRunArg(fu *FakeUser) SignupEngineRunArg {
	return SignupEngineRunArg{
		Username:    fu.Username,
		Email:       fu.Email,
		InviteCode:  libkb.TestInvitationCode,
		Passphrase:  fu.Passphrase,
		StoreSecret: false,
		DeviceName:  defaultDeviceName,
		SkipGPG:     true,
		SkipMail:    true,
		SkipPaper:   true,
	}
}

func SignupFakeUserWithArg(tc libkb.TestContext, fu *FakeUser, arg SignupEngineRunArg) *SignupEngine {
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, tc.G)
	err := RunEngine(s, ctx)
	if err != nil {
		tc.T.Fatal(err)
	}
	fu.EncryptionKey = s.encryptionKey
	return s
}

func CreateAndSignupFakeUser(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	tc.G.Log.Debug("New test user: %s / %s", fu.Username, fu.Email)
	arg := MakeTestSignupEngineRunArg(fu)
	_ = SignupFakeUserWithArg(tc, fu, arg)
	return fu
}

func CreateAndSignupFakeUserPaper(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	tc.G.Log.Debug("New test user: %s / %s", fu.Username, fu.Email)
	arg := MakeTestSignupEngineRunArg(fu)
	arg.SkipPaper = false
	_ = SignupFakeUserWithArg(tc, fu, arg)
	return fu
}

func CreateAndSignupFakeUserSafe(g *libkb.GlobalContext, prefix string) (*FakeUser, error) {
	fu, err := NewFakeUser(prefix)
	if err != nil {
		return nil, err
	}

	arg := MakeTestSignupEngineRunArg(fu)
	ctx := &Context{
		LogUI:    g.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, g)
	err = RunEngine(s, ctx)
	if err != nil {
		return nil, err
	}
	return fu, nil
}

func CreateAndSignupFakeUserGPG(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	if err := tc.GenerateGPGKeyring(fu.Email); err != nil {
		tc.T.Fatal(err)
	}
	arg := MakeTestSignupEngineRunArg(fu)
	arg.SkipGPG = false
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, tc.G)
	err := RunEngine(s, ctx)
	if err != nil {
		tc.T.Fatal(err)
	}
	return fu
}

func SignupFakeUserStoreSecret(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	tc.G.Log.Debug("New test user: %s / %s", fu.Username, fu.Email)
	arg := MakeTestSignupEngineRunArg(fu)
	arg.SkipPaper = false
	arg.StoreSecret = true
	_ = SignupFakeUserWithArg(tc, fu, arg)
	return fu
}

func CreateAndSignupFakeUserCustomArg(tc libkb.TestContext, prefix string, fmod func(*SignupEngineRunArg)) (*FakeUser, libkb.GenericKey) {
	fu := NewFakeUserOrBust(tc.T, prefix)
	arg := MakeTestSignupEngineRunArg(fu)
	fmod(&arg)
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, tc.G)
	err := RunEngine(s, ctx)
	if err != nil {
		tc.T.Fatal(err)
	}
	return fu, s.signingKey
}

func (fu *FakeUser) LoginWithSecretUI(secui libkb.SecretUI, g *libkb.GlobalContext) error {
	ctx := &Context{
		ProvisionUI: newTestProvisionUI(),
		LogUI:       g.UI.GetLogUI(),
		GPGUI:       &gpgtestui{},
		SecretUI:    secui,
		LoginUI:     &libkb.TestLoginUI{Username: fu.Username},
	}
	li := NewLogin(g, libkb.DeviceTypeDesktop, fu.Username, keybase1.ClientType_CLI)
	return RunEngine(li, ctx)
}

func (fu *FakeUser) Login(g *libkb.GlobalContext) error {
	s := fu.NewSecretUI()
	return fu.LoginWithSecretUI(s, g)
}

func (fu *FakeUser) LoginOrBust(tc libkb.TestContext) {
	if err := fu.Login(tc.G); err != nil {
		tc.T.Fatal(err)
	}
}

func (fu *FakeUser) NewSecretUI() *libkb.TestSecretUI {
	return &libkb.TestSecretUI{Passphrase: fu.Passphrase}
}

func AssertProvisioned(tc libkb.TestContext) error {
	prov, err := tc.G.LoginState().LoggedInProvisionedLoad()
	if err != nil {
		return err
	}
	if !prov {
		return libkb.LoginRequiredError{}
	}
	return nil
}

func AssertNotProvisioned(tc libkb.TestContext) error {
	prov, err := tc.G.LoginState().LoggedInProvisionedLoad()
	if err != nil {
		return err
	}
	if prov {
		return errors.New("AssertNotProvisioned failed:  user is provisioned")
	}
	return nil
}

func AssertLoggedIn(tc libkb.TestContext) error {
	if !LoggedIn(tc) {
		return libkb.LoginRequiredError{}
	}
	return nil
}

func AssertLoggedOut(tc libkb.TestContext) error {
	if LoggedIn(tc) {
		return libkb.LogoutError{}
	}
	return nil
}

func LoggedIn(tc libkb.TestContext) bool {
	lin, _ := tc.G.LoginState().LoggedInLoad()
	return lin
}

func Logout(tc libkb.TestContext) {
	if err := tc.G.Logout(); err != nil {
		tc.T.Fatalf("logout error: %s", err)
	}
}

// TODO: Add tests that use testEngineWithSecretStore for every engine
// that should work with the secret store.

// testEngineWithSecretStore takes a given engine-running function and
// makes sure that it works with the secret store, i.e. that it stores
// data into it when told to and reads data out from it.
func testEngineWithSecretStore(
	t *testing.T,
	runEngine func(libkb.TestContext, *FakeUser, libkb.SecretUI)) {

	tc := SetupEngineTest(t, "wss")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "wss")
	tc.ResetLoginState()

	testSecretUI := libkb.TestSecretUI{
		Passphrase:  fu.Passphrase,
		StoreSecret: true,
	}
	runEngine(tc, fu, &testSecretUI)

	if !testSecretUI.CalledGetPassphrase {
		t.Fatal("GetPassphrase() unexpectedly not called")
	}

	tc.ResetLoginState()

	testSecretUI = libkb.TestSecretUI{}
	runEngine(tc, fu, &testSecretUI)

	if testSecretUI.CalledGetPassphrase {
		t.Fatal("GetPassphrase() unexpectedly called")
	}
}

func SetupTwoDevices(t *testing.T, nm string) (user *FakeUser, dev1 libkb.TestContext, dev2 libkb.TestContext, cleanup func()) {
	return SetupTwoDevicesWithHook(t, nm, nil)
}

func SetupTwoDevicesWithHook(t *testing.T, nm string, hook func(tc *libkb.TestContext)) (user *FakeUser, dev1 libkb.TestContext, dev2 libkb.TestContext, cleanup func()) {
	if len(nm) > 5 {
		t.Fatalf("Sorry, test name must be fewer than 6 chars (got %q)", nm)
	}

	// device X (provisioner) context:
	dev1 = SetupEngineTest(t, nm)

	// device Y (provisionee) context:
	dev2 = SetupEngineTest(t, nm)
	if hook != nil {
		hook(&dev2)
	}

	user = NewFakeUserOrBust(t, nm)
	arg := MakeTestSignupEngineRunArg(user)
	arg.SkipPaper = false
	loginUI := &paperLoginUI{Username: user.Username}
	ctx := &Context{
		LogUI:    dev1.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: user.NewSecretUI(),
		LoginUI:  loginUI,
	}
	s := NewSignupEngine(&arg, dev1.G)
	err := RunEngine(s, ctx)
	if err != nil {
		t.Fatal(err)
	}

	assertNumDevicesAndKeys(dev1, user, 2, 4)

	if len(loginUI.PaperPhrase) == 0 {
		t.Fatal("login ui has no paper key phrase")
	}

	secUI := user.NewSecretUI()
	secUI.Passphrase = loginUI.PaperPhrase
	provUI := newTestProvisionUIPaper()
	provLoginUI := &libkb.TestLoginUI{Username: user.Username}
	ctx = &Context{
		ProvisionUI: provUI,
		LogUI:       dev2.G.UI.GetLogUI(),
		SecretUI:    secUI,
		LoginUI:     provLoginUI,
		GPGUI:       &gpgtestui{},
	}
	eng := NewLogin(dev2.G, libkb.DeviceTypeDesktop, "", keybase1.ClientType_CLI)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	testUserHasDeviceKey(dev2)

	assertNumDevicesAndKeys(dev2, user, 3, 6)

	if err := AssertProvisioned(dev2); err != nil {
		t.Fatal(err)
	}

	cleanup = func() {
		dev1.Cleanup()
		dev2.Cleanup()
	}

	return user, dev1, dev2, cleanup
}

func ResetAccount(tc libkb.TestContext, u *FakeUser) {
	tc.T.Skip("reset account tests broken")
	err := tc.G.LoginState().ResetAccount(u.Username)
	if err != nil {
		tc.T.Fatalf("In account reset: %s", err)
	}
	tc.T.Logf("Account reset for user %s", u.Username)
	Logout(tc)
}
<commit_msg>Undo ResetAccount skip<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"testing"

	"github.com/keybase/client/go/externals"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	insecureTriplesec "github.com/keybase/go-triplesec-insecure"
)

func SetupEngineTest(tb testing.TB, name string) libkb.TestContext {
	tc := externals.SetupTest(tb, name, 2)
	tc.G.NewTriplesec = func(passphrase []byte, salt []byte) (libkb.Triplesec, error) {
		warner := func() { tc.G.Log.Warning("Installing insecure Triplesec with weak stretch parameters") }
		isProduction := func() bool {
			return tc.G.Env.GetRunMode() == libkb.ProductionRunMode
		}
		return insecureTriplesec.NewCipher(passphrase, salt, warner, isProduction)
	}
	return tc
}

func SetupEngineTestRealTriplesec(tb testing.TB, name string) libkb.TestContext {
	tc := externals.SetupTest(tb, name, 2)
	tc.G.NewTriplesec = libkb.NewSecureTriplesec
	return tc
}

type FakeUser struct {
	Username      string
	Email         string
	Passphrase    string
	User          *libkb.User
	EncryptionKey libkb.GenericKey
}

func NewFakeUser(prefix string) (fu *FakeUser, err error) {
	buf := make([]byte, 5)
	if _, err = rand.Read(buf); err != nil {
		return
	}
	username := fmt.Sprintf("%s_%s", prefix, hex.EncodeToString(buf))
	email := fmt.Sprintf("%s@noemail.keybase.io", username)
	buf = make([]byte, 12)
	if _, err = rand.Read(buf); err != nil {
		return
	}
	passphrase := hex.EncodeToString(buf)
	fu = &FakeUser{Username: username, Email: email, Passphrase: passphrase}
	return
}

func (fu FakeUser) NormalizedUsername() libkb.NormalizedUsername {
	return libkb.NewNormalizedUsername(fu.Username)
}

func (fu FakeUser) UID() keybase1.UID {
	// All new-style names will have a 1-to-1 mapping
	return libkb.UsernameToUID(fu.Username)
}

func NewFakeUserOrBust(tb testing.TB, prefix string) (fu *FakeUser) {
	var err error
	if fu, err = NewFakeUser(prefix); err != nil {
		tb.Fatal(err)
	}
	return fu
}

const defaultDeviceName = "my device"

// MakeTestSignupEngineRunArg fills a SignupEngineRunArg with the most
// common parameters for testing and returns it.
func MakeTestSignupEngineRunArg(fu *FakeUser) SignupEngineRunArg {
	return SignupEngineRunArg{
		Username:    fu.Username,
		Email:       fu.Email,
		InviteCode:  libkb.TestInvitationCode,
		Passphrase:  fu.Passphrase,
		StoreSecret: false,
		DeviceName:  defaultDeviceName,
		SkipGPG:     true,
		SkipMail:    true,
		SkipPaper:   true,
	}
}

func SignupFakeUserWithArg(tc libkb.TestContext, fu *FakeUser, arg SignupEngineRunArg) *SignupEngine {
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, tc.G)
	err := RunEngine(s, ctx)
	if err != nil {
		tc.T.Fatal(err)
	}
	fu.EncryptionKey = s.encryptionKey
	return s
}

func CreateAndSignupFakeUser(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	tc.G.Log.Debug("New test user: %s / %s", fu.Username, fu.Email)
	arg := MakeTestSignupEngineRunArg(fu)
	_ = SignupFakeUserWithArg(tc, fu, arg)
	return fu
}

func CreateAndSignupFakeUserPaper(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	tc.G.Log.Debug("New test user: %s / %s", fu.Username, fu.Email)
	arg := MakeTestSignupEngineRunArg(fu)
	arg.SkipPaper = false
	_ = SignupFakeUserWithArg(tc, fu, arg)
	return fu
}

func CreateAndSignupFakeUserSafe(g *libkb.GlobalContext, prefix string) (*FakeUser, error) {
	fu, err := NewFakeUser(prefix)
	if err != nil {
		return nil, err
	}

	arg := MakeTestSignupEngineRunArg(fu)
	ctx := &Context{
		LogUI:    g.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, g)
	err = RunEngine(s, ctx)
	if err != nil {
		return nil, err
	}
	return fu, nil
}

func CreateAndSignupFakeUserGPG(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	if err := tc.GenerateGPGKeyring(fu.Email); err != nil {
		tc.T.Fatal(err)
	}
	arg := MakeTestSignupEngineRunArg(fu)
	arg.SkipGPG = false
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, tc.G)
	err := RunEngine(s, ctx)
	if err != nil {
		tc.T.Fatal(err)
	}
	return fu
}

func SignupFakeUserStoreSecret(tc libkb.TestContext, prefix string) *FakeUser {
	fu := NewFakeUserOrBust(tc.T, prefix)
	tc.G.Log.Debug("New test user: %s / %s", fu.Username, fu.Email)
	arg := MakeTestSignupEngineRunArg(fu)
	arg.SkipPaper = false
	arg.StoreSecret = true
	_ = SignupFakeUserWithArg(tc, fu, arg)
	return fu
}

func CreateAndSignupFakeUserCustomArg(tc libkb.TestContext, prefix string, fmod func(*SignupEngineRunArg)) (*FakeUser, libkb.GenericKey) {
	fu := NewFakeUserOrBust(tc.T, prefix)
	arg := MakeTestSignupEngineRunArg(fu)
	fmod(&arg)
	ctx := &Context{
		LogUI:    tc.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: fu.NewSecretUI(),
		LoginUI:  &libkb.TestLoginUI{Username: fu.Username},
	}
	s := NewSignupEngine(&arg, tc.G)
	err := RunEngine(s, ctx)
	if err != nil {
		tc.T.Fatal(err)
	}
	return fu, s.signingKey
}

func (fu *FakeUser) LoginWithSecretUI(secui libkb.SecretUI, g *libkb.GlobalContext) error {
	ctx := &Context{
		ProvisionUI: newTestProvisionUI(),
		LogUI:       g.UI.GetLogUI(),
		GPGUI:       &gpgtestui{},
		SecretUI:    secui,
		LoginUI:     &libkb.TestLoginUI{Username: fu.Username},
	}
	li := NewLogin(g, libkb.DeviceTypeDesktop, fu.Username, keybase1.ClientType_CLI)
	return RunEngine(li, ctx)
}

func (fu *FakeUser) Login(g *libkb.GlobalContext) error {
	s := fu.NewSecretUI()
	return fu.LoginWithSecretUI(s, g)
}

func (fu *FakeUser) LoginOrBust(tc libkb.TestContext) {
	if err := fu.Login(tc.G); err != nil {
		tc.T.Fatal(err)
	}
}

func (fu *FakeUser) NewSecretUI() *libkb.TestSecretUI {
	return &libkb.TestSecretUI{Passphrase: fu.Passphrase}
}

func AssertProvisioned(tc libkb.TestContext) error {
	prov, err := tc.G.LoginState().LoggedInProvisionedLoad()
	if err != nil {
		return err
	}
	if !prov {
		return libkb.LoginRequiredError{}
	}
	return nil
}

func AssertNotProvisioned(tc libkb.TestContext) error {
	prov, err := tc.G.LoginState().LoggedInProvisionedLoad()
	if err != nil {
		return err
	}
	if prov {
		return errors.New("AssertNotProvisioned failed:  user is provisioned")
	}
	return nil
}

func AssertLoggedIn(tc libkb.TestContext) error {
	if !LoggedIn(tc) {
		return libkb.LoginRequiredError{}
	}
	return nil
}

func AssertLoggedOut(tc libkb.TestContext) error {
	if LoggedIn(tc) {
		return libkb.LogoutError{}
	}
	return nil
}

func LoggedIn(tc libkb.TestContext) bool {
	lin, _ := tc.G.LoginState().LoggedInLoad()
	return lin
}

func Logout(tc libkb.TestContext) {
	if err := tc.G.Logout(); err != nil {
		tc.T.Fatalf("logout error: %s", err)
	}
}

// TODO: Add tests that use testEngineWithSecretStore for every engine
// that should work with the secret store.

// testEngineWithSecretStore takes a given engine-running function and
// makes sure that it works with the secret store, i.e. that it stores
// data into it when told to and reads data out from it.
func testEngineWithSecretStore(
	t *testing.T,
	runEngine func(libkb.TestContext, *FakeUser, libkb.SecretUI)) {

	tc := SetupEngineTest(t, "wss")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "wss")
	tc.ResetLoginState()

	testSecretUI := libkb.TestSecretUI{
		Passphrase:  fu.Passphrase,
		StoreSecret: true,
	}
	runEngine(tc, fu, &testSecretUI)

	if !testSecretUI.CalledGetPassphrase {
		t.Fatal("GetPassphrase() unexpectedly not called")
	}

	tc.ResetLoginState()

	testSecretUI = libkb.TestSecretUI{}
	runEngine(tc, fu, &testSecretUI)

	if testSecretUI.CalledGetPassphrase {
		t.Fatal("GetPassphrase() unexpectedly called")
	}
}

func SetupTwoDevices(t *testing.T, nm string) (user *FakeUser, dev1 libkb.TestContext, dev2 libkb.TestContext, cleanup func()) {
	return SetupTwoDevicesWithHook(t, nm, nil)
}

func SetupTwoDevicesWithHook(t *testing.T, nm string, hook func(tc *libkb.TestContext)) (user *FakeUser, dev1 libkb.TestContext, dev2 libkb.TestContext, cleanup func()) {
	if len(nm) > 5 {
		t.Fatalf("Sorry, test name must be fewer than 6 chars (got %q)", nm)
	}

	// device X (provisioner) context:
	dev1 = SetupEngineTest(t, nm)

	// device Y (provisionee) context:
	dev2 = SetupEngineTest(t, nm)
	if hook != nil {
		hook(&dev2)
	}

	user = NewFakeUserOrBust(t, nm)
	arg := MakeTestSignupEngineRunArg(user)
	arg.SkipPaper = false
	loginUI := &paperLoginUI{Username: user.Username}
	ctx := &Context{
		LogUI:    dev1.G.UI.GetLogUI(),
		GPGUI:    &gpgtestui{},
		SecretUI: user.NewSecretUI(),
		LoginUI:  loginUI,
	}
	s := NewSignupEngine(&arg, dev1.G)
	err := RunEngine(s, ctx)
	if err != nil {
		t.Fatal(err)
	}

	assertNumDevicesAndKeys(dev1, user, 2, 4)

	if len(loginUI.PaperPhrase) == 0 {
		t.Fatal("login ui has no paper key phrase")
	}

	secUI := user.NewSecretUI()
	secUI.Passphrase = loginUI.PaperPhrase
	provUI := newTestProvisionUIPaper()
	provLoginUI := &libkb.TestLoginUI{Username: user.Username}
	ctx = &Context{
		ProvisionUI: provUI,
		LogUI:       dev2.G.UI.GetLogUI(),
		SecretUI:    secUI,
		LoginUI:     provLoginUI,
		GPGUI:       &gpgtestui{},
	}
	eng := NewLogin(dev2.G, libkb.DeviceTypeDesktop, "", keybase1.ClientType_CLI)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	testUserHasDeviceKey(dev2)

	assertNumDevicesAndKeys(dev2, user, 3, 6)

	if err := AssertProvisioned(dev2); err != nil {
		t.Fatal(err)
	}

	cleanup = func() {
		dev1.Cleanup()
		dev2.Cleanup()
	}

	return user, dev1, dev2, cleanup
}

func ResetAccount(tc libkb.TestContext, u *FakeUser) {
	err := tc.G.LoginState().ResetAccount(u.Username)
	if err != nil {
		tc.T.Fatalf("In account reset: %s", err)
	}
	tc.T.Logf("Account reset for user %s", u.Username)
	Logout(tc)
}
<|endoftext|>


<commit_before>package scrape

import (
	"io"
	"io/ioutil"
	"reflect"
	"strings"
	"testing"

	"golang.org/x/text/encoding/japanese"
	"golang.org/x/text/transform"

	"github.com/PuerkitoBio/goquery"
)

var kyukoDoc, noKyukoDoc *goquery.Document

const (
	KYUKOFILE   = "../testdata/kyuko.html"
	NOKYUKOFILE = "../testdata/not_kyuko.html"
)

func SjisToUtf8(str string) (string, error) {
	ret, err := ioutil.ReadAll(transform.NewReader(strings.NewReader(str), japanese.ShiftJIS.NewDecoder()))
	if err != nil {
		return "", err
	}
	return string(ret), err
}

func EncodeTestFile(fileName string) (io.Reader, error) {
	//testfileのenocde
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, err
	}
	utfFile, err := SjisToUtf8(string(file))
	if err != nil {
		return nil, err
	}
	stringReader := strings.NewReader(utfFile)

	return stringReader, nil
}

func init() {
	//休講ある
	kyukoReader, _ := EncodeTestFile(KYUKOFILE)
	kyukoDoc, _ = goquery.NewDocumentFromReader(kyukoReader)

	//休講ない
	noKyukoReader, _ := EncodeTestFile(NOKYUKOFILE)
	noKyukoDoc, _ = goquery.NewDocumentFromReader(noKyukoReader)
}

func TestSetUrl(t *testing.T) {
	if url, err := SetUrl(1, 1); url != "http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=1&kouchi=1" {
		t.Fatalf("urlの生成がうまくできていないようです\n err: %s", err)
	}

	if url, err := SetUrl(2, 5); url != "http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=5&kouchi=2" || err != nil {
		t.Fatalf("urlの生成がうまくできていないようです\n err: %s", err)
	}

	if url, err := SetUrl(3, 1); err == nil {
		t.Fatalf("存在しない校地のurlが生成されています\n created url: %s", url)
	}

	if url, err := SetUrl(1, 7); err == nil {
		t.Fatalf("日曜日のurlは必要ありません\n created url: %s", url)
	}

}

func TestScrapePeriod(t *testing.T) {
	periods, err := ScrapePeriod(kyukoDoc)
	if err != nil {
		t.Fatal("periodをスクレイピングできませんでした\n%s", err)
	}

	testSlice := []int{2, 2, 2, 5}
	if reflect.DeepEqual(periods, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %d\n got:  %d", testSlice, periods)
	}

	periods, err = ScrapePeriod(noKyukoDoc)
	if err != nil {
		t.Fatal("periodをスクレイピングできませんでした\n%s", err)
	}
	testSlice = []int{}
	if !reflect.DeepEqual(periods, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, periods)
	}
}

func TestScrapeReason(t *testing.T) {

	/*
		//httpでやるとき
			stringReader, err := Get("http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=2&kouchi=2")
			if err != nil {
				t.Fatal("hoge\n%v", err)
			}
	*/
	reasons, err := ScrapeReason(kyukoDoc)
	if err != nil {
		t.Fatalf("reasonをスクレイピングできませんでした\n%s", err)
	}

	testSlice := []string{"公務", "出張", "公務", ""}
	if !reflect.DeepEqual(reasons, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, reasons)
	}
}

func TestScrapeNameAndInstructor(t *testing.T) {
	names, instructors, err := ScrapeNameAndInstructor(kyukoDoc)
	if err != nil {
		t.Fatalf("Nameのスクレイピングに失敗したようです\n%s", err)
	}

	testSlice := []string{"環境生理学", "電気・電子計測Ｉ－１", "応用数学ＩＩ－１", "イングリッシュ・セミナー２－７０２"}
	if !reflect.DeepEqual(names, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, names)
	}

	testSlice = []string{"福岡義之", "松川真美", "大川領", "稲垣俊史"}
	if !reflect.DeepEqual(instructors, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, instructors)
	}

}

//まだできてない
func testScrape(t *testing.T) {

	//r, err := Scrape("http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=4&kouchi=2")

	/*
		file, err := os.Open("../testdata/kyuko.html")
		if err != nil {
			t.Fatalf("テストデータを開けませんでした\n%s", err)
		}
		defer file.Close()

		r, err := Scrape("", file)
		if err != nil {
			t.Fatalf("hoge\n%s", err)
		}

		fmt.Printf("%d\nhoge\n%d", r, err)
	*/
}
<commit_msg>!!!!!!!!!!!!!!!<commit_after>package scrape

import (
	"io"
	"io/ioutil"
	"reflect"
	"strings"
	"testing"

	"golang.org/x/text/encoding/japanese"
	"golang.org/x/text/transform"

	"github.com/PuerkitoBio/goquery"
)

var kyukoDoc, noKyukoDoc *goquery.Document

const (
	KYUKOFILE   = "../testdata/kyuko.html"
	NOKYUKOFILE = "../testdata/not_kyuko.html"
)

func SjisToUtf8(str string) (string, error) {
	ret, err := ioutil.ReadAll(transform.NewReader(strings.NewReader(str), japanese.ShiftJIS.NewDecoder()))
	if err != nil {
		return "", err
	}
	return string(ret), err
}

func EncodeTestFile(fileName string) (io.Reader, error) {
	//testfileのenocde
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, err
	}
	utfFile, err := SjisToUtf8(string(file))
	if err != nil {
		return nil, err
	}
	stringReader := strings.NewReader(utfFile)

	return stringReader, nil
}

func init() {
	//休講ある
	kyukoReader, _ := EncodeTestFile(KYUKOFILE)
	kyukoDoc, _ = goquery.NewDocumentFromReader(kyukoReader)

	//休講ない
	noKyukoReader, _ := EncodeTestFile(NOKYUKOFILE)
	noKyukoDoc, _ = goquery.NewDocumentFromReader(noKyukoReader)
}

func TestSetUrl(t *testing.T) {
	if url, err := SetUrl(1, 1); url != "http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=1&kouchi=1" {
		t.Fatalf("urlの生成がうまくできていないようです\n err: %s", err)
	}

	if url, err := SetUrl(2, 5); url != "http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=5&kouchi=2" || err != nil {
		t.Fatalf("urlの生成がうまくできていないようです\n err: %s", err)
	}

	if url, err := SetUrl(3, 1); err == nil {
		t.Fatalf("存在しない校地のurlが生成されています\n created url: %s", url)
	}

	if url, err := SetUrl(1, 7); err == nil {
		t.Fatalf("日曜日のurlは必要ありません\n created url: %s", url)
	}

}

func TestScrapePeriod(t *testing.T) {
	periods, err := ScrapePeriod(kyukoDoc)
	if err != nil {
		t.Fatal("periodをスクレイピングできませんでした\n%s", err)
	}

	testSlice := []int{2, 2, 2, 5}
	if !reflect.DeepEqual(periods, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %d\n got:  %d", testSlice, periods)
	}

	periods, err = ScrapePeriod(noKyukoDoc)
	if err != nil {
		t.Fatal("periodをスクレイピングできませんでした\n%s", err)
	}
	if len(periods) != 0 {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, periods)
	}
}

func TestScrapeReason(t *testing.T) {

	/*
		//httpでやるとき
			stringReader, err := Get("http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=2&kouchi=2")
			if err != nil {
				t.Fatal("hoge\n%v", err)
			}
	*/
	reasons, err := ScrapeReason(kyukoDoc)
	if err != nil {
		t.Fatalf("reasonをスクレイピングできませんでした\n%s", err)
	}

	testSlice := []string{"公務", "出張", "公務", ""}
	if !reflect.DeepEqual(reasons, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, reasons)
	}
}

func TestScrapeNameAndInstructor(t *testing.T) {
	names, instructors, err := ScrapeNameAndInstructor(kyukoDoc)
	if err != nil {
		t.Fatalf("Nameのスクレイピングに失敗したようです\n%s", err)
	}

	testSlice := []string{"環境生理学", "電気・電子計測Ｉ－１", "応用数学ＩＩ－１", "イングリッシュ・セミナー２－７０２"}
	if !reflect.DeepEqual(names, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, names)
	}

	testSlice = []string{"福岡義之", "松川真美", "大川領", "稲垣俊史"}
	if !reflect.DeepEqual(instructors, testSlice) {
		t.Fatalf("取得した結果が求めるものと違ったようです\n want: %v\n got:  %v", testSlice, instructors)
	}

}

//まだできてない
func testScrape(t *testing.T) {

	//r, err := Scrape("http://duet.doshisha.ac.jp/info/KK1000.jsp?katei=1&youbi=4&kouchi=2")

	/*
		file, err := os.Open("../testdata/kyuko.html")
		if err != nil {
			t.Fatalf("テストデータを開けませんでした\n%s", err)
		}
		defer file.Close()

		r, err := Scrape("", file)
		if err != nil {
			t.Fatalf("hoge\n%s", err)
		}

		fmt.Printf("%d\nhoge\n%d", r, err)
	*/
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess 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 vtctl

import (
	"context"
	"flag"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/olekukonko/tablewriter"
	"google.golang.org/protobuf/encoding/prototext"

	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/throttler"
	"vitess.io/vitess/go/vt/throttler/throttlerclient"
	"vitess.io/vitess/go/vt/wrangler"

	throttlerdatapb "vitess.io/vitess/go/vt/proto/throttlerdata"
)

const (
	throttlerGroupName = "Resharding Throttler"
	shortTimeout       = 15 * time.Second
)

// This file contains the commands to control the throttler which is used during
// resharding (vtworker) and by filtered replication (vttablet).

func init() {
	addCommandGroup(throttlerGroupName)

	addCommand(throttlerGroupName, command{
		name:   "ThrottlerMaxRates",
		method: commandThrottlerMaxRates,
		params: "-server <vtworker or vttablet>",
		help:   "Returns the current max rate of all active resharding throttlers on the server.",
	})
	addCommand(throttlerGroupName, command{
		name:   "ThrottlerSetMaxRate",
		method: commandThrottlerSetMaxRate,
		params: "-server <vtworker or vttablet> <rate>",
		help:   "Sets the max rate for all active resharding throttlers on the server.",
	})

	addCommand(throttlerGroupName, command{
		name:   "GetThrottlerConfiguration",
		method: commandGetThrottlerConfiguration,
		params: "-server <vtworker or vttablet> [<throttler name>]",
		help:   "Returns the current configuration of the MaxReplicationLag module. If no throttler name is specified, the configuration of all throttlers will be returned.",
	})
	addCommand(throttlerGroupName, command{
		name:   "UpdateThrottlerConfiguration",
		method: commandUpdateThrottlerConfiguration,
		// Note: <configuration protobuf text> is put in quotes to tell the user
		// that the value must be quoted such that it's one argument only.
		params: `-server <vtworker or vttablet> [-copy_zero_values] "<configuration protobuf text>" [<throttler name>]`,
		help:   "Updates the configuration of the MaxReplicationLag module. The configuration must be specified as protobuf text. If a field is omitted or has a zero value, it will be ignored unless -copy_zero_values is specified. If no throttler name is specified, all throttlers will be updated.",
	})
	addCommand(throttlerGroupName, command{
		name:   "ResetThrottlerConfiguration",
		method: commandResetThrottlerConfiguration,
		params: "-server <vtworker or vttablet> [<throttler name>]",
		help:   "Resets the current configuration of the MaxReplicationLag module. If no throttler name is specified, the configuration of all throttlers will be reset.",
	})
}

func commandThrottlerMaxRates(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 0 {
		return fmt.Errorf("the ThrottlerSetMaxRate command does not accept any positional parameters")
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	rates, err := client.MaxRates(ctx)
	if err != nil {
		return fmt.Errorf("failed to get the throttler rate from server '%v': %v", *server, err)
	}

	if len(rates) == 0 {
		wr.Logger().Printf("There are no active throttlers on server '%v'.\n", *server)
		return nil
	}

	table := tablewriter.NewWriter(loggerWriter{wr.Logger()})
	table.SetAutoFormatHeaders(false)
	table.SetHeader([]string{"Name", "Rate"})
	for name, rate := range rates {
		rateText := strconv.FormatInt(rate, 10)
		if rate == throttler.MaxRateModuleDisabled {
			rateText = "unlimited"
		}
		table.Append([]string{name, rateText})
	}
	table.Render()
	wr.Logger().Printf("%d active throttler(s) on server '%v'.\n", len(rates), *server)
	return nil
}

func commandThrottlerSetMaxRate(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 1 {
		return fmt.Errorf("the <rate> argument is required for the ThrottlerSetMaxRate command")
	}
	var rate int64
	if strings.ToLower(subFlags.Arg(0)) == "unlimited" {
		rate = throttler.MaxRateModuleDisabled
	} else {
		var err error
		rate, err = strconv.ParseInt(subFlags.Arg(0), 0, 64)
		if err != nil {
			return fmt.Errorf("failed to parse rate '%v' as integer value: %v", subFlags.Arg(0), err)
		}
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	names, err := client.SetMaxRate(ctx, rate)
	if err != nil {
		return fmt.Errorf("failed to set the throttler rate on server '%v': %v", *server, err)
	}

	if len(names) == 0 {
		wr.Logger().Printf("ThrottlerSetMaxRate did nothing because server '%v' has no active throttlers.\n", *server)
		return nil
	}

	printUpdatedThrottlers(wr.Logger(), *server, names)
	return nil
}

func commandGetThrottlerConfiguration(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() > 1 {
		return fmt.Errorf("the GetThrottlerConfiguration command accepts only <throttler name> as optional positional parameter")
	}

	var throttlerName string
	if subFlags.NArg() == 1 {
		throttlerName = subFlags.Arg(0)
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	configurations, err := client.GetConfiguration(ctx, throttlerName)
	if err != nil {
		return fmt.Errorf("failed to get the throttler configuration from server '%v': %v", *server, err)
	}

	if len(configurations) == 0 {
		wr.Logger().Printf("There are no active throttlers on server '%v'.\n", *server)
		return nil
	}

	table := tablewriter.NewWriter(loggerWriter{wr.Logger()})
	table.SetAutoFormatHeaders(false)
	// The full protobuf text will span more than one terminal line. Do not wrap
	// it to make it easy to copy and paste it.
	table.SetAutoWrapText(false)
	table.SetHeader([]string{"Name", "Configuration (protobuf text, fields with a zero value are omitted)"})
	for name, c := range configurations {
		pcfg, _ := prototext.Marshal(c)
		table.Append([]string{name, string(pcfg)})
	}
	table.Render()
	wr.Logger().Printf("%d active throttler(s) on server '%v'.\n", len(configurations), *server)
	return nil
}

func commandUpdateThrottlerConfiguration(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	copyZeroValues := subFlags.Bool("copy_zero_values", false, "If true, fields with zero values will be copied as well")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() > 2 {
		return fmt.Errorf(`the "<configuration protobuf text>" argument is required for the UpdateThrottlerConfiguration command. The <throttler name> is an optional positional parameter`)
	}

	var throttlerName string
	if subFlags.NArg() == 2 {
		throttlerName = subFlags.Arg(1)
	}

	protoText := subFlags.Arg(0)
	configuration := &throttlerdatapb.Configuration{}
	if err := prototext.Unmarshal([]byte(protoText), configuration); err != nil {
		return fmt.Errorf("failed to unmarshal the configuration protobuf text (%v) into a protobuf instance: %v", protoText, err)
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	names, err := client.UpdateConfiguration(ctx, throttlerName, configuration, *copyZeroValues)
	if err != nil {
		return fmt.Errorf("failed to update the throttler configuration on server '%v': %v", *server, err)
	}

	if len(names) == 0 {
		wr.Logger().Printf("UpdateThrottlerConfiguration did nothing because server '%v' has no active throttlers.\n", *server)
		return nil
	}

	printUpdatedThrottlers(wr.Logger(), *server, names)
	wr.Logger().Printf("The new configuration will become effective with the next recalculation event.\n")
	return nil
}

func commandResetThrottlerConfiguration(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() > 1 {
		return fmt.Errorf("the ResetThrottlerConfiguration command accepts only <throttler name> as optional positional parameter")
	}

	var throttlerName string
	if subFlags.NArg() == 1 {
		throttlerName = subFlags.Arg(0)
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	names, err := client.ResetConfiguration(ctx, throttlerName)
	if err != nil {
		return fmt.Errorf("failed to get the throttler configuration from server '%v': %v", *server, err)
	}

	if len(names) == 0 {
		wr.Logger().Printf("ResetThrottlerConfiguration did nothing because server '%v' has no active throttlers.\n", *server)
		return nil
	}

	printUpdatedThrottlers(wr.Logger(), *server, names)
	wr.Logger().Printf("The reset initial configuration will become effective with the next recalculation event.\n")
	return nil
}

func printUpdatedThrottlers(logger logutil.Logger, server string, names []string) {
	table := tablewriter.NewWriter(loggerWriter{logger})
	table.SetAutoFormatHeaders(false)
	table.SetHeader([]string{"Name"})
	for _, name := range names {
		table.Append([]string{name})
	}
	table.Render()
	logger.Printf("%d active throttler(s) on server '%v' were updated.\n", len(names), server)
}
<commit_msg>Deprecate Throttler RPCs (#9962)<commit_after>/*
Copyright 2019 The Vitess 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 vtctl

import (
	"context"
	"flag"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/olekukonko/tablewriter"
	"google.golang.org/protobuf/encoding/prototext"

	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/throttler"
	"vitess.io/vitess/go/vt/throttler/throttlerclient"
	"vitess.io/vitess/go/vt/wrangler"

	throttlerdatapb "vitess.io/vitess/go/vt/proto/throttlerdata"
)

const (
	throttlerGroupName = "Resharding Throttler"
	shortTimeout       = 15 * time.Second
)

// This file contains the commands to control the throttler which is used during
// resharding (vtworker) and by filtered replication (vttablet).

func init() {
	addCommandGroup(throttlerGroupName)

	addCommand(throttlerGroupName, command{
		name:         "ThrottlerMaxRates",
		method:       commandThrottlerMaxRates,
		params:       "-server <vtworker or vttablet>",
		help:         "Returns the current max rate of all active resharding throttlers on the server.",
		deprecated:   true,
		deprecatedBy: "the new Reshard/MoveTables workflows",
	})
	addCommand(throttlerGroupName, command{
		name:         "ThrottlerSetMaxRate",
		method:       commandThrottlerSetMaxRate,
		params:       "-server <vtworker or vttablet> <rate>",
		help:         "Sets the max rate for all active resharding throttlers on the server.",
		deprecated:   true,
		deprecatedBy: "the new Reshard/MoveTables workflows",
	})

	addCommand(throttlerGroupName, command{
		name:         "GetThrottlerConfiguration",
		method:       commandGetThrottlerConfiguration,
		params:       "-server <vtworker or vttablet> [<throttler name>]",
		help:         "Returns the current configuration of the MaxReplicationLag module. If no throttler name is specified, the configuration of all throttlers will be returned.",
		deprecated:   true,
		deprecatedBy: "the new Reshard/MoveTables workflows",
	})
	addCommand(throttlerGroupName, command{
		name:   "UpdateThrottlerConfiguration",
		method: commandUpdateThrottlerConfiguration,
		// Note: <configuration protobuf text> is put in quotes to tell the user
		// that the value must be quoted such that it's one argument only.
		params:       `-server <vtworker or vttablet> [-copy_zero_values] "<configuration protobuf text>" [<throttler name>]`,
		help:         "Updates the configuration of the MaxReplicationLag module. The configuration must be specified as protobuf text. If a field is omitted or has a zero value, it will be ignored unless -copy_zero_values is specified. If no throttler name is specified, all throttlers will be updated.",
		deprecated:   true,
		deprecatedBy: "the new Reshard/MoveTables workflows",
	})
	addCommand(throttlerGroupName, command{
		name:         "ResetThrottlerConfiguration",
		method:       commandResetThrottlerConfiguration,
		params:       "-server <vtworker or vttablet> [<throttler name>]",
		help:         "Resets the current configuration of the MaxReplicationLag module. If no throttler name is specified, the configuration of all throttlers will be reset.",
		deprecated:   true,
		deprecatedBy: "the new Reshard/MoveTables workflows",
	})
}

func commandThrottlerMaxRates(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 0 {
		return fmt.Errorf("the ThrottlerSetMaxRate command does not accept any positional parameters")
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	rates, err := client.MaxRates(ctx)
	if err != nil {
		return fmt.Errorf("failed to get the throttler rate from server '%v': %v", *server, err)
	}

	if len(rates) == 0 {
		wr.Logger().Printf("There are no active throttlers on server '%v'.\n", *server)
		return nil
	}

	table := tablewriter.NewWriter(loggerWriter{wr.Logger()})
	table.SetAutoFormatHeaders(false)
	table.SetHeader([]string{"Name", "Rate"})
	for name, rate := range rates {
		rateText := strconv.FormatInt(rate, 10)
		if rate == throttler.MaxRateModuleDisabled {
			rateText = "unlimited"
		}
		table.Append([]string{name, rateText})
	}
	table.Render()
	wr.Logger().Printf("%d active throttler(s) on server '%v'.\n", len(rates), *server)
	return nil
}

func commandThrottlerSetMaxRate(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 1 {
		return fmt.Errorf("the <rate> argument is required for the ThrottlerSetMaxRate command")
	}
	var rate int64
	if strings.ToLower(subFlags.Arg(0)) == "unlimited" {
		rate = throttler.MaxRateModuleDisabled
	} else {
		var err error
		rate, err = strconv.ParseInt(subFlags.Arg(0), 0, 64)
		if err != nil {
			return fmt.Errorf("failed to parse rate '%v' as integer value: %v", subFlags.Arg(0), err)
		}
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	names, err := client.SetMaxRate(ctx, rate)
	if err != nil {
		return fmt.Errorf("failed to set the throttler rate on server '%v': %v", *server, err)
	}

	if len(names) == 0 {
		wr.Logger().Printf("ThrottlerSetMaxRate did nothing because server '%v' has no active throttlers.\n", *server)
		return nil
	}

	printUpdatedThrottlers(wr.Logger(), *server, names)
	return nil
}

func commandGetThrottlerConfiguration(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() > 1 {
		return fmt.Errorf("the GetThrottlerConfiguration command accepts only <throttler name> as optional positional parameter")
	}

	var throttlerName string
	if subFlags.NArg() == 1 {
		throttlerName = subFlags.Arg(0)
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	configurations, err := client.GetConfiguration(ctx, throttlerName)
	if err != nil {
		return fmt.Errorf("failed to get the throttler configuration from server '%v': %v", *server, err)
	}

	if len(configurations) == 0 {
		wr.Logger().Printf("There are no active throttlers on server '%v'.\n", *server)
		return nil
	}

	table := tablewriter.NewWriter(loggerWriter{wr.Logger()})
	table.SetAutoFormatHeaders(false)
	// The full protobuf text will span more than one terminal line. Do not wrap
	// it to make it easy to copy and paste it.
	table.SetAutoWrapText(false)
	table.SetHeader([]string{"Name", "Configuration (protobuf text, fields with a zero value are omitted)"})
	for name, c := range configurations {
		pcfg, _ := prototext.Marshal(c)
		table.Append([]string{name, string(pcfg)})
	}
	table.Render()
	wr.Logger().Printf("%d active throttler(s) on server '%v'.\n", len(configurations), *server)
	return nil
}

func commandUpdateThrottlerConfiguration(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	copyZeroValues := subFlags.Bool("copy_zero_values", false, "If true, fields with zero values will be copied as well")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() > 2 {
		return fmt.Errorf(`the "<configuration protobuf text>" argument is required for the UpdateThrottlerConfiguration command. The <throttler name> is an optional positional parameter`)
	}

	var throttlerName string
	if subFlags.NArg() == 2 {
		throttlerName = subFlags.Arg(1)
	}

	protoText := subFlags.Arg(0)
	configuration := &throttlerdatapb.Configuration{}
	if err := prototext.Unmarshal([]byte(protoText), configuration); err != nil {
		return fmt.Errorf("failed to unmarshal the configuration protobuf text (%v) into a protobuf instance: %v", protoText, err)
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	names, err := client.UpdateConfiguration(ctx, throttlerName, configuration, *copyZeroValues)
	if err != nil {
		return fmt.Errorf("failed to update the throttler configuration on server '%v': %v", *server, err)
	}

	if len(names) == 0 {
		wr.Logger().Printf("UpdateThrottlerConfiguration did nothing because server '%v' has no active throttlers.\n", *server)
		return nil
	}

	printUpdatedThrottlers(wr.Logger(), *server, names)
	wr.Logger().Printf("The new configuration will become effective with the next recalculation event.\n")
	return nil
}

func commandResetThrottlerConfiguration(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	server := subFlags.String("server", "", "vtworker or vttablet to connect to")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() > 1 {
		return fmt.Errorf("the ResetThrottlerConfiguration command accepts only <throttler name> as optional positional parameter")
	}

	var throttlerName string
	if subFlags.NArg() == 1 {
		throttlerName = subFlags.Arg(0)
	}

	// Connect to the server.
	ctx, cancel := context.WithTimeout(ctx, shortTimeout)
	defer cancel()
	client, err := throttlerclient.New(*server)
	if err != nil {
		return fmt.Errorf("error creating a throttler client for server '%v': %v", *server, err)
	}
	defer client.Close()

	names, err := client.ResetConfiguration(ctx, throttlerName)
	if err != nil {
		return fmt.Errorf("failed to get the throttler configuration from server '%v': %v", *server, err)
	}

	if len(names) == 0 {
		wr.Logger().Printf("ResetThrottlerConfiguration did nothing because server '%v' has no active throttlers.\n", *server)
		return nil
	}

	printUpdatedThrottlers(wr.Logger(), *server, names)
	wr.Logger().Printf("The reset initial configuration will become effective with the next recalculation event.\n")
	return nil
}

func printUpdatedThrottlers(logger logutil.Logger, server string, names []string) {
	table := tablewriter.NewWriter(loggerWriter{logger})
	table.SetAutoFormatHeaders(false)
	table.SetHeader([]string{"Name"})
	for _, name := range names {
		table.Append([]string{name})
	}
	table.Render()
	logger.Printf("%d active throttler(s) on server '%v' were updated.\n", len(names), server)
}
<|endoftext|>


<commit_before>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vttablet

// This file handles the binlog players launched on masters for filtered
// replication

import (
	"time"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/relog"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/topo"
)

// BinlogPlayerController controls one player
type BinlogPlayerController struct {
	ts       topo.Server
	vtClient mysqlctl.VtClient
	keyspace string
	source   topo.SourceShard

	// Player is the BinlogPlayer when we have one
	player *mysqlctl.BinlogPlayer

	// interrupted is the channel to close to stop the playback
	interrupted chan struct{}

	// TODO(alainjobart): add state
	// TODO(alainjobart): add statsJson, include stats and player stats if any
	// TODO(alainjobart): figure out if we need a lock on structure (for stats)
}

func NewBinlogController(ts topo.Server, vtClient mysqlctl.VtClient, keyspace string, source topo.SourceShard) *BinlogPlayerController {
	return &BinlogPlayerController{
		ts:          ts,
		vtClient:    vtClient,
		keyspace:    keyspace,
		source:      source,
		interrupted: make(chan struct{}, 1),
	}
}

func (bpc *BinlogPlayerController) Start() {
	relog.Info("Starting binlog player for %v", bpc.source)
	go bpc.Loop()
}

func (bpc *BinlogPlayerController) Stop() {
	relog.Info("Stopping binlog player for %v", bpc.source)
	close(bpc.interrupted)
}

func (bpc *BinlogPlayerController) Loop() {
	for {
		// Read the start position
		startPosition, err := mysqlctl.ReadStartPosition(bpc.vtClient, string(bpc.source.KeyRange.Start.Hex()), string(bpc.source.KeyRange.End.Hex()))
		if err != nil {
			relog.Warning("BinlogPlayerController: can't read startPosition: %v", err)
			time.Sleep(5)
			continue
		}

		// TODO(alainjobart): Find the server list

		// TODO(alainjobart): Pick a server (same if it's available,
		// if not clear master file / pos and keep only group id)

		// Create the player.
		bpc.player, err = mysqlctl.NewBinlogPlayer(bpc.vtClient, startPosition, nil /*tables*/, 1 /*txnBatch*/, 30*time.Second /*maxTxnInterval*/, false /*execDdl*/)
		if err != nil {
			relog.Warning("BinlogPlayerController: can't create player: %v", err)
			time.Sleep(5)
			continue
		}

		// Run player loop until it's done.
		err = bpc.player.ApplyBinlogEvents(bpc.interrupted)
		if err != nil {
			relog.Error("Error in applying binlog events, err %v", err)
			time.Sleep(5)
		} else {
			// We were interrupted.
			break
		}
	}
	relog.Info("Exited main binlog player loop for %v", bpc.source)
}

// BinlogPlayerMap controls all the players
type BinlogPlayerMap struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	players  map[topo.SourceShard]*BinlogPlayerController
}

func NewBinlogPlayerMap(ts topo.Server, dbConfig *mysql.ConnectionParams) *BinlogPlayerMap {
	return &BinlogPlayerMap{
		ts:       ts,
		dbConfig: dbConfig,
		players:  make(map[topo.SourceShard]*BinlogPlayerController),
	}
}

// TODO(alainjobart) add stats, register them

func (blm *BinlogPlayerMap) AddPlayer(keyspace string, source topo.SourceShard) {
	bpc, ok := blm.players[source]
	if ok {
		relog.Info("Already playing logs for %v", source)
		return
	}

	// create the db connection, connect it
	vtClient := mysqlctl.NewDbClient(blm.dbConfig)
	if err := vtClient.Connect(); err != nil {
		relog.Error("BinlogPlayerMap: can't connect to database: %v", err)
		return
	}

	bpc = NewBinlogController(blm.ts, vtClient, keyspace, source)
	blm.players[source] = bpc
	bpc.Start()
}

func (blm *BinlogPlayerMap) StopAllPlayers() {
	for _, bpc := range blm.players {
		bpc.Stop()
	}
	blm.players = make(map[topo.SourceShard]*BinlogPlayerController)
}

// RefreshMap reads the right data from topo.Server and makes sure
// we're playing the right logs
func (blm *BinlogPlayerMap) RefreshMap(tablet topo.Tablet) {
	relog.Info("Refreshing map of binlog players")

	// read the shard to get SourceShards
	shardInfo, err := blm.ts.GetShard(tablet.Keyspace, tablet.Shard)
	if err != nil {
		relog.Error("Cannot read shard for this tablet: %v", tablet.Alias())
		return
	}

	// get the existing sources and build a map of sources to remove
	toRemove := make(map[topo.SourceShard]bool)
	for source, _ := range blm.players {
		toRemove[source] = true
	}

	// for each source, add it if not there, and delete from toRemove
	for _, source := range shardInfo.SourceShards {
		blm.AddPlayer(tablet.Keyspace, source)
		delete(toRemove, source)
	}

	// remove all entries from toRemove
	for source, _ := range toRemove {
		blm.players[source].Stop()
		delete(blm.players, source)
	}
}
<commit_msg>Moving connection logic so it's in the loop.<commit_after>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vttablet

// This file handles the binlog players launched on masters for filtered
// replication

import (
	"time"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/relog"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/topo"
)

// BinlogPlayerController controls one player
type BinlogPlayerController struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	keyspace string
	source   topo.SourceShard

	// Player is the BinlogPlayer when we have one
	player *mysqlctl.BinlogPlayer

	// interrupted is the channel to close to stop the playback
	interrupted chan struct{}

	// TODO(alainjobart): add state
	// TODO(alainjobart): add statsJson, include stats and player stats if any
	// TODO(alainjobart): figure out if we need a lock on structure (for stats)
}

func NewBinlogController(ts topo.Server, dbConfig *mysql.ConnectionParams, keyspace string, source topo.SourceShard) *BinlogPlayerController {
	return &BinlogPlayerController{
		ts:          ts,
		dbConfig:    dbConfig,
		keyspace:    keyspace,
		source:      source,
		interrupted: make(chan struct{}, 1),
	}
}

func (bpc *BinlogPlayerController) Start() {
	relog.Info("Starting binlog player for %v", bpc.source)
	go bpc.Loop()
}

func (bpc *BinlogPlayerController) Stop() {
	relog.Info("Stopping binlog player for %v", bpc.source)
	close(bpc.interrupted)
}

func (bpc *BinlogPlayerController) Loop() {
	for {
		// create the db connection, connect it
		vtClient := mysqlctl.NewDbClient(bpc.dbConfig)
		if err := vtClient.Connect(); err != nil {
			relog.Error("BinlogPlayerMap: can't connect to database: %v", err)
			continue
		}

		// Read the start position
		startPosition, err := mysqlctl.ReadStartPosition(vtClient, string(bpc.source.KeyRange.Start.Hex()), string(bpc.source.KeyRange.End.Hex()))
		if err != nil {
			relog.Warning("BinlogPlayerController: can't read startPosition: %v", err)
			time.Sleep(5)
			continue
		}

		// TODO(alainjobart): Find the server list

		// TODO(alainjobart): Pick a server (same if it's available,
		// if not clear master file / pos and keep only group id)

		// Create the player.
		bpc.player, err = mysqlctl.NewBinlogPlayer(vtClient, startPosition, nil /*tables*/, 1 /*txnBatch*/, 30*time.Second /*maxTxnInterval*/, false /*execDdl*/)
		if err != nil {
			relog.Warning("BinlogPlayerController: can't create player: %v", err)
			time.Sleep(5)
			continue
		}

		// Run player loop until it's done.
		err = bpc.player.ApplyBinlogEvents(bpc.interrupted)
		if err != nil {
			relog.Error("Error in applying binlog events, err %v", err)
			time.Sleep(5)
		} else {
			// We were interrupted.
			break
		}
	}
	relog.Info("Exited main binlog player loop for %v", bpc.source)
}

// BinlogPlayerMap controls all the players
type BinlogPlayerMap struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	players  map[topo.SourceShard]*BinlogPlayerController
}

func NewBinlogPlayerMap(ts topo.Server, dbConfig *mysql.ConnectionParams) *BinlogPlayerMap {
	return &BinlogPlayerMap{
		ts:       ts,
		dbConfig: dbConfig,
		players:  make(map[topo.SourceShard]*BinlogPlayerController),
	}
}

// TODO(alainjobart) add stats, register them

func (blm *BinlogPlayerMap) AddPlayer(keyspace string, source topo.SourceShard) {
	bpc, ok := blm.players[source]
	if ok {
		relog.Info("Already playing logs for %v", source)
		return
	}

	bpc = NewBinlogController(blm.ts, blm.dbConfig, keyspace, source)
	blm.players[source] = bpc
	bpc.Start()
}

func (blm *BinlogPlayerMap) StopAllPlayers() {
	for _, bpc := range blm.players {
		bpc.Stop()
	}
	blm.players = make(map[topo.SourceShard]*BinlogPlayerController)
}

// RefreshMap reads the right data from topo.Server and makes sure
// we're playing the right logs
func (blm *BinlogPlayerMap) RefreshMap(tablet topo.Tablet) {
	relog.Info("Refreshing map of binlog players")

	// read the shard to get SourceShards
	shardInfo, err := blm.ts.GetShard(tablet.Keyspace, tablet.Shard)
	if err != nil {
		relog.Error("Cannot read shard for this tablet: %v", tablet.Alias())
		return
	}

	// get the existing sources and build a map of sources to remove
	toRemove := make(map[topo.SourceShard]bool)
	for source, _ := range blm.players {
		toRemove[source] = true
	}

	// for each source, add it if not there, and delete from toRemove
	for _, source := range shardInfo.SourceShards {
		blm.AddPlayer(tablet.Keyspace, source)
		delete(toRemove, source)
	}

	// remove all entries from toRemove
	for source, _ := range toRemove {
		blm.players[source].Stop()
		delete(blm.players, source)
	}
}
<|endoftext|>


<commit_before>package git

/*
extern int _govcs_gcrypt_init();
#cgo LDFLAGS: -lgcrypt
*/
import "C"
import (
	"fmt"
	"log"
	"os"
	"os/user"
	"strings"

	"crypto/md5"

	"code.google.com/p/go.crypto/ssh"

	git2go "github.com/libgit2/git2go"
	"github.com/sourcegraph/go-vcs/vcs"
	"github.com/sourcegraph/go-vcs/vcs/gitcmd"
	sshutil "github.com/sourcegraph/go-vcs/vcs/ssh"
	"github.com/sourcegraph/go-vcs/vcs/util"
)

func init() {
	// Overwrite the git cloner to use the faster libgit2
	// implementation.
	vcs.RegisterCloner("git", func(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
		return Clone(url, dir, opt)
	})
}

func init() {
	// Initialize gcrypt for multithreaded operation. See
	// gcrypt_init.c for more information.
	rv := C._govcs_gcrypt_init()
	if rv != 0 {
		log.Fatal("gcrypt multithreaded init failed (see gcrypt_init.c)")
	}
}

func Clone(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
	clopt := git2go.CloneOptions{Bare: opt.Bare}

	rc, cfs, err := makeRemoteCallbacks(url, opt.RemoteOpts)
	if err != nil {
		return nil, err
	}
	if cfs != nil {
		defer cfs.run()
	}
	clopt.RemoteCallbacks = rc

	u, err := git2go.Clone(url, dir, &clopt)
	if err != nil {
		return nil, err
	}
	cr, err := gitcmd.Open(dir)
	if err != nil {
		return nil, err
	}
	return &Repository{cr, u}, nil
}

func (r *Repository) UpdateEverything(opt vcs.RemoteOpts) error {
	// TODO(sqs): allow use of a remote other than "origin"
	rm, err := r.u.LoadRemote("origin")
	if err != nil {
		return err
	}

	rc, cfs, err := makeRemoteCallbacks(rm.Url(), opt)
	if err != nil {
		return err
	}
	if cfs != nil {
		defer cfs.run()
	}
	rm.SetCallbacks(rc)

	if err := rm.Fetch(nil, nil, ""); err != nil {
		return err
	}

	return nil
}

type cleanupFuncs []func() error

func (f cleanupFuncs) run() error {
	for _, cf := range f {
		if err := cf(); err != nil {
			return err
		}
	}
	return nil
}

// makeRemoteCallbacks constructs the remote callbacks for libgit2
// remote operations. Currently the remote callbacks are trivial
// (empty) except when using an SSH remote.
//
// cleanupFuncs's run method should be called when the RemoteCallbacks
// struct is done being used. It is OK to ignore the error return.
func makeRemoteCallbacks(url string, opt vcs.RemoteOpts) (rc *git2go.RemoteCallbacks, cfs cleanupFuncs, err error) {
	defer func() {
		// Clean up if error; don't expect the caller to clean up if
		// we have a non-nil error.
		if err != nil {
			cfs.run()
		}
	}()

	if opt.SSH != nil {
		privkeyFilename, privkeyFile, err := util.WriteKeyTempFile(url, opt.SSH.PrivateKey)
		if err != nil {
			return nil, nil, err
		}
		cfs = append(cfs, privkeyFile.Close)
		cfs = append(cfs, func() error { return os.Remove(privkeyFile.Name()) })

		// Derive public key from private key if empty.
		if opt.SSH.PublicKey == nil {
			privKey, err := ssh.ParsePrivateKey(opt.SSH.PrivateKey)
			if err != nil {
				return nil, cfs, err
			}
			opt.SSH.PublicKey = ssh.MarshalAuthorizedKey(privKey.PublicKey())
		}

		pubkeyFilename, pubkeyFile, err := util.WriteKeyTempFile(url, opt.SSH.PublicKey)
		if err != nil {
			return nil, cfs, err
		}
		cfs = append(cfs, pubkeyFile.Close)
		cfs = append(cfs, func() error { return os.Remove(pubkeyFile.Name()) })

		rc = &git2go.RemoteCallbacks{
			CredentialsCallback: func(url string, usernameFromURL string, allowedTypes git2go.CredType) (int, *git2go.Cred) {
				var username string
				if usernameFromURL != "" {
					username = usernameFromURL
				} else if opt.SSH.User != "" {
					username = opt.SSH.User
				} else {
					if username == "" {
						u, err := user.Current()
						if err == nil {
							username = u.Username
						}
					}
				}
				if allowedTypes&git2go.CredTypeSshKey != 0 {
					rv, cred := git2go.NewCredSshKey(username, pubkeyFilename, privkeyFilename, "")
					return rv, &cred
				}
				log.Printf("No authentication available for git URL %q.", url)
				return 1, nil
			},
			CertificateCheckCallback: func(cert *git2go.Certificate, valid bool, hostname string) int {
				// libgit2 currently always returns valid=false. It
				// may return valid=true in the future if it checks
				// host keys using known_hosts, but let's ignore valid
				// so we don't get that behavior unexpectedly.

				if InsecureSkipCheckVerifySSH {
					return 0
				}

				if cert == nil {
					return -1
				}

				if cert.Hostkey.Kind&git2go.HostkeyMD5 > 0 {
					keys, found := standardKnownHosts.Lookup(hostname)
					if found {
						hostFingerprint := md5String(cert.Hostkey.HashMD5)
						for _, key := range keys {
							knownFingerprint := md5String(md5.Sum(key.Marshal()))
							if hostFingerprint == knownFingerprint {
								return 0
							}
						}
					}
				}

				log.Printf("Invalid certificate for SSH host %s: %v.", hostname, cert)
				return -1
			},
		}
	}

	return rc, cfs, nil
}

// InsecureSkipCheckVerifySSH controls whether the client verifies the
// SSH server's certificate or host key. If InsecureSkipCheckVerifySSH
// is true, the program is susceptible to a man-in-the-middle
// attack. This should only be used for testing.
var InsecureSkipCheckVerifySSH bool

// standardKnownHosts contains known_hosts from the system known_hosts
// file and the user's known_hosts file.
var standardKnownHosts sshutil.KnownHosts

func init() {
	var err error
	standardKnownHosts, err = sshutil.ReadStandardKnownHostsFiles()
	if err != nil {
		log.Printf("Warning: failed to read standard SSH known_hosts files (%s). SSH host key checking will fail.")
	}
}

// md5String returns a formatted string representing the given md5Sum in hex
func md5String(md5Sum [16]byte) string {
	md5Str := fmt.Sprintf("% x", md5Sum)
	md5Str = strings.Replace(md5Str, " ", ":", -1)
	return md5Str
}
<commit_msg>only set callbacks if needed<commit_after>package git

/*
extern int _govcs_gcrypt_init();
#cgo LDFLAGS: -lgcrypt
*/
import "C"
import (
	"fmt"
	"log"
	"os"
	"os/user"
	"strings"

	"crypto/md5"

	"code.google.com/p/go.crypto/ssh"

	git2go "github.com/libgit2/git2go"
	"github.com/sourcegraph/go-vcs/vcs"
	"github.com/sourcegraph/go-vcs/vcs/gitcmd"
	sshutil "github.com/sourcegraph/go-vcs/vcs/ssh"
	"github.com/sourcegraph/go-vcs/vcs/util"
)

func init() {
	// Overwrite the git cloner to use the faster libgit2
	// implementation.
	vcs.RegisterCloner("git", func(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
		return Clone(url, dir, opt)
	})
}

func init() {
	// Initialize gcrypt for multithreaded operation. See
	// gcrypt_init.c for more information.
	rv := C._govcs_gcrypt_init()
	if rv != 0 {
		log.Fatal("gcrypt multithreaded init failed (see gcrypt_init.c)")
	}
}

func Clone(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
	clopt := git2go.CloneOptions{Bare: opt.Bare}

	rc, cfs, err := makeRemoteCallbacks(url, opt.RemoteOpts)
	if err != nil {
		return nil, err
	}
	if cfs != nil {
		defer cfs.run()
	}
	clopt.RemoteCallbacks = rc

	u, err := git2go.Clone(url, dir, &clopt)
	if err != nil {
		return nil, err
	}
	cr, err := gitcmd.Open(dir)
	if err != nil {
		return nil, err
	}
	return &Repository{cr, u}, nil
}

func (r *Repository) UpdateEverything(opt vcs.RemoteOpts) error {
	// TODO(sqs): allow use of a remote other than "origin"
	rm, err := r.u.LoadRemote("origin")
	if err != nil {
		return err
	}

	rc, cfs, err := makeRemoteCallbacks(rm.Url(), opt)
	if err != nil {
		return err
	}
	if cfs != nil {
		defer cfs.run()
	}
	if rc != nil {
		rm.SetCallbacks(rc)
	}

	if err := rm.Fetch(nil, nil, ""); err != nil {
		return err
	}

	return nil
}

type cleanupFuncs []func() error

func (f cleanupFuncs) run() error {
	for _, cf := range f {
		if err := cf(); err != nil {
			return err
		}
	}
	return nil
}

// makeRemoteCallbacks constructs the remote callbacks for libgit2
// remote operations. Currently the remote callbacks are trivial
// (empty) except when using an SSH remote.
//
// cleanupFuncs's run method should be called when the RemoteCallbacks
// struct is done being used. It is OK to ignore the error return.
func makeRemoteCallbacks(url string, opt vcs.RemoteOpts) (rc *git2go.RemoteCallbacks, cfs cleanupFuncs, err error) {
	defer func() {
		// Clean up if error; don't expect the caller to clean up if
		// we have a non-nil error.
		if err != nil {
			cfs.run()
		}
	}()

	if opt.SSH != nil {
		privkeyFilename, privkeyFile, err := util.WriteKeyTempFile(url, opt.SSH.PrivateKey)
		if err != nil {
			return nil, nil, err
		}
		cfs = append(cfs, privkeyFile.Close)
		cfs = append(cfs, func() error { return os.Remove(privkeyFile.Name()) })

		// Derive public key from private key if empty.
		if opt.SSH.PublicKey == nil {
			privKey, err := ssh.ParsePrivateKey(opt.SSH.PrivateKey)
			if err != nil {
				return nil, cfs, err
			}
			opt.SSH.PublicKey = ssh.MarshalAuthorizedKey(privKey.PublicKey())
		}

		pubkeyFilename, pubkeyFile, err := util.WriteKeyTempFile(url, opt.SSH.PublicKey)
		if err != nil {
			return nil, cfs, err
		}
		cfs = append(cfs, pubkeyFile.Close)
		cfs = append(cfs, func() error { return os.Remove(pubkeyFile.Name()) })

		rc = &git2go.RemoteCallbacks{
			CredentialsCallback: func(url string, usernameFromURL string, allowedTypes git2go.CredType) (int, *git2go.Cred) {
				var username string
				if usernameFromURL != "" {
					username = usernameFromURL
				} else if opt.SSH.User != "" {
					username = opt.SSH.User
				} else {
					if username == "" {
						u, err := user.Current()
						if err == nil {
							username = u.Username
						}
					}
				}
				if allowedTypes&git2go.CredTypeSshKey != 0 {
					rv, cred := git2go.NewCredSshKey(username, pubkeyFilename, privkeyFilename, "")
					return rv, &cred
				}
				log.Printf("No authentication available for git URL %q.", url)
				return 1, nil
			},
			CertificateCheckCallback: func(cert *git2go.Certificate, valid bool, hostname string) int {
				// libgit2 currently always returns valid=false. It
				// may return valid=true in the future if it checks
				// host keys using known_hosts, but let's ignore valid
				// so we don't get that behavior unexpectedly.

				if InsecureSkipCheckVerifySSH {
					return 0
				}

				if cert == nil {
					return -1
				}

				if cert.Hostkey.Kind&git2go.HostkeyMD5 > 0 {
					keys, found := standardKnownHosts.Lookup(hostname)
					if found {
						hostFingerprint := md5String(cert.Hostkey.HashMD5)
						for _, key := range keys {
							knownFingerprint := md5String(md5.Sum(key.Marshal()))
							if hostFingerprint == knownFingerprint {
								return 0
							}
						}
					}
				}

				log.Printf("Invalid certificate for SSH host %s: %v.", hostname, cert)
				return -1
			},
		}
	}

	return rc, cfs, nil
}

// InsecureSkipCheckVerifySSH controls whether the client verifies the
// SSH server's certificate or host key. If InsecureSkipCheckVerifySSH
// is true, the program is susceptible to a man-in-the-middle
// attack. This should only be used for testing.
var InsecureSkipCheckVerifySSH bool

// standardKnownHosts contains known_hosts from the system known_hosts
// file and the user's known_hosts file.
var standardKnownHosts sshutil.KnownHosts

func init() {
	var err error
	standardKnownHosts, err = sshutil.ReadStandardKnownHostsFiles()
	if err != nil {
		log.Printf("Warning: failed to read standard SSH known_hosts files (%s). SSH host key checking will fail.")
	}
}

// md5String returns a formatted string representing the given md5Sum in hex
func md5String(md5Sum [16]byte) string {
	md5Str := fmt.Sprintf("% x", md5Sum)
	md5Str = strings.Replace(md5Str, " ", ":", -1)
	return md5Str
}
<|endoftext|>


<commit_before>package article

import (
	"regexp"
	"strconv"
	"strings"

	"time"

	"fmt"

	"github.com/garyburd/redigo/redis"
	log "github.com/meifamily/logrus"
	"github.com/meifamily/ptt-alertor/connections"
	"github.com/meifamily/ptt-alertor/models/pushsum"
	"github.com/meifamily/ptt-alertor/myutil"
)

const prefix = "article:"
const subsSuffix = ":subs"

type Article struct {
	ID               int    `json:"ID,omitempty"`
	Code             string `json:"code,omitempty"`
	Title            string
	Link             string
	Date             string    `json:"Date,omitempty"`
	Author           string    `json:"Author,omitempty"`
	Comments         Comments  `json:"pushList,omitempty"` // TODO: rename json key to comments
	LastPushDateTime time.Time `json:"lastPushDateTime,omitempty"`
	Board            string    `json:"board,omitempty"`
	PushSum          int       `json:"pushSum,omitempty"`
	drive            Driver
}

type Driver interface {
	Find(code string) Article
	Save(a Article) error
	Delete(code string) error
}

func NewArticle(drive Driver) *Article {
	return &Article{
		drive: drive,
	}
}

func (a Article) ParseID(Link string) (id int) {
	reg, err := regexp.Compile("https?://www.ptt.cc/bbs/.*/[GM]\\.(\\d+)\\..*")
	if err != nil {
		log.Fatal(err)
	}
	strs := reg.FindStringSubmatch(Link)
	if len(strs) < 2 {
		return 0
	}
	id, err = strconv.Atoi(strs[1])
	if err != nil {
		return 0
	}
	return id
}

func (a Article) MatchKeyword(keyword string) bool {
	if strings.Contains(keyword, "&") {
		keywords := strings.Split(keyword, "&")
		for _, keyword := range keywords {
			if !matchKeyword(a.Title, keyword) {
				return false
			}
		}
		return true
	}
	if strings.HasPrefix(keyword, "regexp:") {
		return matchRegex(a.Title, keyword)
	}
	return matchKeyword(a.Title, keyword)
}

// Exist check article exist or not
func (a Article) Exist() (bool, error) {
	conn := connections.Redis()
	defer conn.Close()

	bl, err := redis.Bool(conn.Do("EXISTS", prefix+a.Code+subsSuffix, "board"))
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return bl, err
}

func (a Article) Find(code string) Article {
	return a.drive.Find(code)
}

func (a Article) Save() error {
	return a.drive.Save(a)
}

func (a Article) Destroy() error {
	if err := a.drive.Delete(a.Code); err != nil {
		return err
	}

	conn := connections.Redis()
	defer conn.Close()

	_, err := conn.Do("DEL", prefix+a.Code+subsSuffix)
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return err
}

func (a Article) AddSubscriber(account string) error {
	conn := connections.Redis()
	defer conn.Close()

	_, err := conn.Do("SADD", prefix+a.Code+subsSuffix, account)
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return err
}

func (a Article) Subscribers() ([]string, error) {
	conn := connections.Redis()
	defer conn.Close()

	accounts, err := redis.Strings(conn.Do("SMEMBERS", prefix+a.Code+subsSuffix))
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return accounts, err
}

func (a Article) RemoveSubscriber(sub string) error {
	conn := connections.Redis()
	defer conn.Close()

	_, err := conn.Do("SREM", prefix+a.Code+subsSuffix, sub)
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return err
}

func (a Article) String() string {
	return a.Title + "\r\n" + a.Link
}

func (a Article) StringWithPushSum() string {
	sumStr := strconv.Itoa(a.PushSum)
	if text, ok := pushsum.NumTextMap[a.PushSum]; ok {
		sumStr = text
	}
	return fmt.Sprintf("%s %s\r\n%s", sumStr, a.Title, a.Link)
}

func matchRegex(title string, regex string) bool {
	pattern := strings.TrimPrefix(regex, "regexp:")
	b, err := regexp.MatchString(pattern, title)
	if err != nil {
		return false
	}
	return b
}

func matchKeyword(title string, keyword string) bool {
	if strings.HasPrefix(keyword, "!") {
		excludeKeyword := strings.Trim(keyword, "!")
		return !containKeyword(title, excludeKeyword)
	}
	return containKeyword(title, keyword)
}

func containKeyword(title string, keyword string) bool {
	return strings.Contains(strings.ToLower(title), strings.ToLower(keyword))
}
<commit_msg>:art: rename pushlist to comments<commit_after>package article

import (
	"regexp"
	"strconv"
	"strings"

	"time"

	"fmt"

	"github.com/garyburd/redigo/redis"
	log "github.com/meifamily/logrus"
	"github.com/meifamily/ptt-alertor/connections"
	"github.com/meifamily/ptt-alertor/models/pushsum"
	"github.com/meifamily/ptt-alertor/myutil"
)

const prefix = "article:"
const subsSuffix = ":subs"

type Article struct {
	ID               int    `json:"ID,omitempty"`
	Code             string `json:"code,omitempty"`
	Title            string
	Link             string
	Date             string    `json:"Date,omitempty"`
	Author           string    `json:"Author,omitempty"`
	Comments         Comments  `json:"comments,omitempty"`
	LastPushDateTime time.Time `json:"lastPushDateTime,omitempty"`
	Board            string    `json:"board,omitempty"`
	PushSum          int       `json:"pushSum,omitempty"`
	drive            Driver
}

type Driver interface {
	Find(code string) Article
	Save(a Article) error
	Delete(code string) error
}

func NewArticle(drive Driver) *Article {
	return &Article{
		drive: drive,
	}
}

func (a Article) ParseID(Link string) (id int) {
	reg, err := regexp.Compile("https?://www.ptt.cc/bbs/.*/[GM]\\.(\\d+)\\..*")
	if err != nil {
		log.Fatal(err)
	}
	strs := reg.FindStringSubmatch(Link)
	if len(strs) < 2 {
		return 0
	}
	id, err = strconv.Atoi(strs[1])
	if err != nil {
		return 0
	}
	return id
}

func (a Article) MatchKeyword(keyword string) bool {
	if strings.Contains(keyword, "&") {
		keywords := strings.Split(keyword, "&")
		for _, keyword := range keywords {
			if !matchKeyword(a.Title, keyword) {
				return false
			}
		}
		return true
	}
	if strings.HasPrefix(keyword, "regexp:") {
		return matchRegex(a.Title, keyword)
	}
	return matchKeyword(a.Title, keyword)
}

// Exist check article exist or not
func (a Article) Exist() (bool, error) {
	conn := connections.Redis()
	defer conn.Close()

	bl, err := redis.Bool(conn.Do("EXISTS", prefix+a.Code+subsSuffix, "board"))
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return bl, err
}

func (a Article) Find(code string) Article {
	return a.drive.Find(code)
}

func (a Article) Save() error {
	return a.drive.Save(a)
}

func (a Article) Destroy() error {
	if err := a.drive.Delete(a.Code); err != nil {
		return err
	}

	conn := connections.Redis()
	defer conn.Close()

	_, err := conn.Do("DEL", prefix+a.Code+subsSuffix)
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return err
}

func (a Article) AddSubscriber(account string) error {
	conn := connections.Redis()
	defer conn.Close()

	_, err := conn.Do("SADD", prefix+a.Code+subsSuffix, account)
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return err
}

func (a Article) Subscribers() ([]string, error) {
	conn := connections.Redis()
	defer conn.Close()

	accounts, err := redis.Strings(conn.Do("SMEMBERS", prefix+a.Code+subsSuffix))
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return accounts, err
}

func (a Article) RemoveSubscriber(sub string) error {
	conn := connections.Redis()
	defer conn.Close()

	_, err := conn.Do("SREM", prefix+a.Code+subsSuffix, sub)
	if err != nil {
		log.WithField("runtime", myutil.BasicRuntimeInfo()).WithError(err).Error()
	}
	return err
}

func (a Article) String() string {
	return a.Title + "\r\n" + a.Link
}

func (a Article) StringWithPushSum() string {
	sumStr := strconv.Itoa(a.PushSum)
	if text, ok := pushsum.NumTextMap[a.PushSum]; ok {
		sumStr = text
	}
	return fmt.Sprintf("%s %s\r\n%s", sumStr, a.Title, a.Link)
}

func matchRegex(title string, regex string) bool {
	pattern := strings.TrimPrefix(regex, "regexp:")
	b, err := regexp.MatchString(pattern, title)
	if err != nil {
		return false
	}
	return b
}

func matchKeyword(title string, keyword string) bool {
	if strings.HasPrefix(keyword, "!") {
		excludeKeyword := strings.Trim(keyword, "!")
		return !containKeyword(title, excludeKeyword)
	}
	return containKeyword(title, keyword)
}

func containKeyword(title string, keyword string) bool {
	return strings.Contains(strings.ToLower(title), strings.ToLower(keyword))
}
<|endoftext|>