<commit_before>package x86_16

import (
	"fmt"
	"strings"
	"syscall"

	"github.com/lunixbochs/argjoy"
	"github.com/pkg/errors"

	co "github.com/lunixbochs/usercorn/go/kernel/common"
	"github.com/lunixbochs/usercorn/go/models"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
)

const (
	STACK_BASE = 0x8000
	STACK_SIZE = 0x1000
	NUM_FDS    = 256

	// Registers
	AH    = uc.X86_REG_AH
	AL    = uc.X86_REG_AL
	AX    = uc.X86_REG_AX
	BH    = uc.X86_REG_BH
	BL    = uc.X86_REG_BL
	BP    = uc.X86_REG_BP
	BX    = uc.X86_REG_BX
	CH    = uc.X86_REG_CH
	CL    = uc.X86_REG_CL
	CS    = uc.X86_REG_CS
	CX    = uc.X86_REG_CX
	DH    = uc.X86_REG_DH
	DI    = uc.X86_REG_DI
	DL    = uc.X86_REG_DL
	DS    = uc.X86_REG_DS
	DX    = uc.X86_REG_DX
	ES    = uc.X86_REG_ES
	FS    = uc.X86_REG_FS
	GS    = uc.X86_REG_GS
	IP    = uc.X86_REG_IP
	SI    = uc.X86_REG_SI
	SP    = uc.X86_REG_SP
	SS    = uc.X86_REG_SS
	FLAGS = uc.X86_REG_EFLAGS
)

func (k *DosKernel) reg16(enum int) uint16 {
	r, _ := k.U.RegRead(enum)
	return uint16(r)
}
func (k *DosKernel) reg8(enum int) uint8 {
	r, _ := k.U.RegRead(enum)
	return uint8(r)
}
func (k *DosKernel) wreg16(enum int, val uint16) {
	k.U.RegWrite(enum, uint64(val))
}
func (k *DosKernel) wreg8(enum int, val uint8) {
	k.U.RegWrite(enum, uint64(val))
}
func (k *DosKernel) setFlagC(set bool) {
	// TODO: Write setFlagX with enum for each flag
	// Unicorn doesn't have the non-extended FLAGS register, so we're
	// dealing with 32 bits here
	flags, _ := k.U.RegRead(FLAGS)
	if set {
		flags |= 1 // CF = 1
	} else {
		flags &= 0xfffffffe // CF = 0
	}
	k.U.RegWrite(FLAGS, flags)
}

var dosSysNum = map[int]string{
	0x00: "terminate",
	0x01: "char_in",
	0x02: "char_out",
	0x09: "display",
	0x30: "get_dos_version",
	0x3C: "create_or_truncate",
	0x3D: "open",
	0x3E: "close",
	0x3F: "read",
	0x40: "write",
	0x41: "unlink",
	0x4C: "terminate_with_code",
}

type abi []int

// ABI to syscall number mapping
var abiMap = map[*abi][]int{
	&abi{BX, DX, CX}: {0x00, 0x30, 0x3E, 0x3F, 0x40},
	&abi{DX, CX}:     {0x01, 0x02, 0x09, 0x3C, 0x41},
	&abi{DX, AL}:     {0x3D},
	&abi{AL}:         {0x4C},
}

var syscallAbis = map[int][]int{}

type PSP struct {
	CPMExit                     [2]byte
	FirstFreeSegment            uint16
	Reserved1                   uint8
	CPMCall5Compat              [5]byte
	OldTSRAddress               uint32
	OldBreakAddress             uint32
	CriticalErrorHandlerAddress uint32
	CallerPSPSegment            uint16
	JobFileTable                [20]byte
	EnvironmentSegment          uint16
	INT21SSSP                   uint32
	JobFileTableSize            uint16
	JobFileTablePointer         uint32
	PreviousPSP                 uint32
	Reserved2                   uint32
	DOSVersion                  uint16
	Reserved3                   [14]byte
	DOSFarCall                  [3]byte
	Reserved4                   uint16
	ExtendedFCB1                [7]byte
	FCB1                        [16]byte
	FCB2                        [20]byte
	CommandLineLength           uint8
	CommandLine                 [127]byte
}

type DosKernel struct {
	*co.KernelBase
	fds [NUM_FDS]int
}

func initPsp(argc int, argv []string) *PSP {
	psp := &PSP{
		CPMExit:    [2]byte{0xcd, 0x20},       // int 0x20
		DOSFarCall: [3]byte{0xcd, 0x21, 0xcd}, // int 0x21 + retf
	}

	psp.FCB1[0] = 0x01
	psp.FCB1[1] = 0x20

	// Combine all args into one string
	commandline := strings.Join(argv, " ")
	copy(psp.CommandLine[:126], commandline)
	if len(commandline) > 126 {
		psp.CommandLineLength = 126
	} else {
		psp.CommandLineLength = uint8(len(commandline))
	}

	return psp
}

func (k *DosKernel) readUntilChar(addr uint64, c byte) []byte {
	var mem []byte
	var i uint64
	var char byte = 0

	// TODO: Read ahead? This'll be slow
	for i = 1; char != c || i == 1; i++ {
		mem, _ = k.U.MemRead(addr, i)
		char = mem[i-1]
	}
	return mem[:i-2]
}

func (k *DosKernel) getFd(fd int) (uint16, error) {
	for i := uint16(0); i < NUM_FDS; i++ {
		if k.fds[i] == -1 {
			k.fds[i] = fd
			return i, nil
		}
	}
	return 0xFFFF, errors.New("DOS FD table exhausted")
}

func (k *DosKernel) freeFd(fd int) (int, error) {
	realfd := k.fds[fd]
	if realfd == -1 {
		return 0xFFFF, errors.New("FD not found in FD table")
	}
	k.fds[fd] = -1
	return realfd, nil
}

func (k *DosKernel) Terminate() {
	k.U.Exit(models.ExitStatus(0))
}

func (k *DosKernel) CharIn(buf co.Buf) byte {
	var char byte
	fmt.Scanf("%c", &char)
	k.U.MemWrite(buf.Addr, []byte{char})
	return char
}

func (k *DosKernel) CharOut(char uint16) byte {
	fmt.Printf("%c", byte(char&0xFF))
	return byte(char & 0xFF)
}

func (k *DosKernel) Display(buf co.Buf) int {
	mem := k.readUntilChar(buf.Addr, '$')

	syscall.Write(1, mem)
	k.wreg8(AL, 0x24)
	return 0x24
}

func (k *DosKernel) GetDosVersion() int {
	k.wreg16(AX, 0x7)
	return 0x7
}

func (k *DosKernel) openFile(filename string, mode int) co.Fd {
	realfd, err := syscall.Open(filename, mode, 0666)
	if err != nil {
		k.wreg16(AX, 0xFFFF)
		k.setFlagC(true)
		return 0xFFFF
	}

	// Find an internal fd number
	dosfd, err := k.getFd(realfd)
	if err != nil {
		k.wreg16(AX, dosfd)
		k.setFlagC(true)
		return 0xFFFF
	}
	k.setFlagC(false)
	k.wreg16(AX, dosfd)
	return co.Fd(dosfd)
}

func (k *DosKernel) CreateOrTruncate(buf co.Buf, attr int) co.Fd {
	filename := string(k.readUntilChar(buf.Addr, '$'))
	return k.openFile(filename, syscall.O_CREAT|syscall.O_TRUNC|syscall.O_RDWR)
}

func (k *DosKernel) Open(filename string, mode int) co.Fd {
	return k.openFile(filename, mode)
}

func (k *DosKernel) Close(fd co.Fd) {
	// Find and free the internal fd
	realfd, _ := k.freeFd(int(fd))
	err := syscall.Close(realfd)
	if err != nil {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	k.setFlagC(false)
	k.wreg16(AX, 0)
}

func (k *DosKernel) Read(fd co.Fd, buf co.Obuf, len co.Len) int {
	mem := make([]byte, len)
	n, err := syscall.Read(int(fd), mem)
	if err != nil {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	k.U.MemWrite(buf.Addr, mem)
	k.setFlagC(false)
	k.wreg16(AX, uint16(n))
	return n
}

func (k *DosKernel) Write(fd co.Fd, buf co.Buf, n co.Len) int {
	mem, _ := k.U.MemRead(buf.Addr, uint64(n))
	written, err := syscall.Write(k.fds[fd], mem)
	if err != nil {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	k.setFlagC(false)
	k.wreg16(AX, uint16(written))
	return written
}

func (k *DosKernel) Unlink(filename string, attr int) int {
	err := syscall.Unlink(filename)
	if err != nil {
		k.setFlagC(true)
		k.wreg16(AX, 0xFFFF)
		return 0xFFFF
	}
	k.setFlagC(false)
	k.wreg16(AX, 0)
	return 0
}

func (k *DosKernel) TerminateWithCode(code int) {
	k.U.Exit(models.ExitStatus(code))
}

func NewKernel() *DosKernel {
	k := &DosKernel{
		KernelBase: &co.KernelBase{},
	}

	// Init FDs
	for i := 0; i < NUM_FDS; i++ {
		k.fds[i] = -1
	}
	k.fds[0] = 0
	k.fds[1] = 1
	k.fds[2] = 2

	// Invert the syscall map
	for abi, syscalls := range abiMap {
		for _, syscall := range syscalls {
			syscallAbis[syscall] = *abi
		}
	}

	k.Argjoy.Register(k.getDosArgCodec())
	return k
}

func DosInit(u models.Usercorn, args, env []string) error {
	// Setup PSP
	// TODO: Setup args
	psp := initPsp(0, nil)
	u.StrucAt(0).Pack(psp)

	u.SetEntry(0x100)
	return u.MapStack(STACK_BASE, STACK_SIZE, false)
}

func DosSyscall(u models.Usercorn) {
	num, _ := u.RegRead(AH)
	name, _ := dosSysNum[int(num)]
	// TODO: How are registers numbered from here?
	u.Syscall(int(num), name, dosArgs(u, int(num)))
	// TODO: Set error
}

func (k *DosKernel) getDosArgCodec() func(interface{}, []interface{}) error {
	return func(arg interface{}, vals []interface{}) error {
		// DOS takes address as DS+DX
		if reg, ok := vals[0].(uint64); ok && len(vals) > 1 {
			ds, _ := k.U.RegRead(DS)
			reg += ds
			switch v := arg.(type) {
			case *co.Buf:
				*v = co.NewBuf(k, reg)
			case *co.Obuf:
				*v = co.Obuf{co.NewBuf(k, reg)}
			case *co.Ptr:
				*v = co.Ptr(reg)
			case *string:
				s, err := k.U.Mem().ReadStrAt(reg)
				if err != nil {
					return errors.Wrapf(err, "ReadStrAt(%#x) failed", reg)
				}
				*v = s
			default:
				return argjoy.NoMatch
			}
			return nil
		}
		return argjoy.NoMatch
	}
}

func dosArgs(u models.Usercorn, num int) func(int) ([]uint64, error) {
	// Return a closure over the correct arglist based on the syscall number
	return co.RegArgs(u, syscallAbis[num])
}

func DosInterrupt(u models.Usercorn, cause uint32) {
	intno := cause & 0xFF
	if intno == 0x21 {
		DosSyscall(u)
	} else if intno == 0x20 {
		u.Syscall(0, "terminate", func(int) ([]uint64, error) { return []uint64{}, nil })
	} else {
		panic(fmt.Sprintf("unhandled X86 interrupt %#X", intno))
	}
}
func DosKernels(u models.Usercorn) []interface{} {
	return []interface{}{NewKernel()}
}

func init() {
	Arch.RegisterOS(&models.OS{
		Name:      "DOS",
		Init:      DosInit,
		Interrupt: DosInterrupt,
		Kernels:   DosKernels,
	})
}
<commit_msg>fix dos out-of-range fds<commit_after>package x86_16

import (
	"fmt"
	"strings"
	"syscall"

	"github.com/lunixbochs/argjoy"
	"github.com/pkg/errors"

	co "github.com/lunixbochs/usercorn/go/kernel/common"
	"github.com/lunixbochs/usercorn/go/models"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
)

const (
	STACK_BASE = 0x8000
	STACK_SIZE = 0x1000
	NUM_FDS    = 256

	// Registers
	AH    = uc.X86_REG_AH
	AL    = uc.X86_REG_AL
	AX    = uc.X86_REG_AX
	BH    = uc.X86_REG_BH
	BL    = uc.X86_REG_BL
	BP    = uc.X86_REG_BP
	BX    = uc.X86_REG_BX
	CH    = uc.X86_REG_CH
	CL    = uc.X86_REG_CL
	CS    = uc.X86_REG_CS
	CX    = uc.X86_REG_CX
	DH    = uc.X86_REG_DH
	DI    = uc.X86_REG_DI
	DL    = uc.X86_REG_DL
	DS    = uc.X86_REG_DS
	DX    = uc.X86_REG_DX
	ES    = uc.X86_REG_ES
	FS    = uc.X86_REG_FS
	GS    = uc.X86_REG_GS
	IP    = uc.X86_REG_IP
	SI    = uc.X86_REG_SI
	SP    = uc.X86_REG_SP
	SS    = uc.X86_REG_SS
	FLAGS = uc.X86_REG_EFLAGS
)

func (k *DosKernel) reg16(enum int) uint16 {
	r, _ := k.U.RegRead(enum)
	return uint16(r)
}
func (k *DosKernel) reg8(enum int) uint8 {
	r, _ := k.U.RegRead(enum)
	return uint8(r)
}
func (k *DosKernel) wreg16(enum int, val uint16) {
	k.U.RegWrite(enum, uint64(val))
}
func (k *DosKernel) wreg8(enum int, val uint8) {
	k.U.RegWrite(enum, uint64(val))
}
func (k *DosKernel) setFlagC(set bool) {
	// TODO: Write setFlagX with enum for each flag
	// Unicorn doesn't have the non-extended FLAGS register, so we're
	// dealing with 32 bits here
	flags, _ := k.U.RegRead(FLAGS)
	if set {
		flags |= 1 // CF = 1
	} else {
		flags &= 0xfffffffe // CF = 0
	}
	k.U.RegWrite(FLAGS, flags)
}

var dosSysNum = map[int]string{
	0x00: "terminate",
	0x01: "char_in",
	0x02: "char_out",
	0x09: "display",
	0x30: "get_dos_version",
	0x3C: "create_or_truncate",
	0x3D: "open",
	0x3E: "close",
	0x3F: "read",
	0x40: "write",
	0x41: "unlink",
	0x4C: "terminate_with_code",
}

type abi []int

// ABI to syscall number mapping
var abiMap = map[*abi][]int{
	&abi{BX, DX, CX}: {0x00, 0x30, 0x3E, 0x3F, 0x40},
	&abi{DX, CX}:     {0x01, 0x02, 0x09, 0x3C, 0x41},
	&abi{DX, AL}:     {0x3D},
	&abi{AL}:         {0x4C},
}

var syscallAbis = map[int][]int{}

type PSP struct {
	CPMExit                     [2]byte
	FirstFreeSegment            uint16
	Reserved1                   uint8
	CPMCall5Compat              [5]byte
	OldTSRAddress               uint32
	OldBreakAddress             uint32
	CriticalErrorHandlerAddress uint32
	CallerPSPSegment            uint16
	JobFileTable                [20]byte
	EnvironmentSegment          uint16
	INT21SSSP                   uint32
	JobFileTableSize            uint16
	JobFileTablePointer         uint32
	PreviousPSP                 uint32
	Reserved2                   uint32
	DOSVersion                  uint16
	Reserved3                   [14]byte
	DOSFarCall                  [3]byte
	Reserved4                   uint16
	ExtendedFCB1                [7]byte
	FCB1                        [16]byte
	FCB2                        [20]byte
	CommandLineLength           uint8
	CommandLine                 [127]byte
}

type DosKernel struct {
	*co.KernelBase
	fds map[int]int
}

func initPsp(argc int, argv []string) *PSP {
	psp := &PSP{
		CPMExit:    [2]byte{0xcd, 0x20},       // int 0x20
		DOSFarCall: [3]byte{0xcd, 0x21, 0xcd}, // int 0x21 + retf
	}

	psp.FCB1[0] = 0x01
	psp.FCB1[1] = 0x20

	// Combine all args into one string
	commandline := strings.Join(argv, " ")
	copy(psp.CommandLine[:126], commandline)
	if len(commandline) > 126 {
		psp.CommandLineLength = 126
	} else {
		psp.CommandLineLength = uint8(len(commandline))
	}

	return psp
}

func (k *DosKernel) readUntilChar(addr uint64, c byte) []byte {
	var mem []byte
	var i uint64
	var char byte = 0

	// TODO: Read ahead? This'll be slow
	for i = 1; char != c || i == 1; i++ {
		mem, _ = k.U.MemRead(addr, i)
		char = mem[i-1]
	}
	return mem[:i-2]
}

func (k *DosKernel) getFd(fd int) (uint16, error) {
	for i := 0; i < NUM_FDS; i++ {
		if _, ok := k.fds[i]; !ok {
			k.fds[i] = fd
			return uint16(i), nil
		}
	}
	return 0xFFFF, errors.New("DOS FD table exhausted")
}

func (k *DosKernel) freeFd(fd int) (int, error) {
	realfd, ok := k.fds[fd]
	if !ok {
		return 0xFFFF, errors.New("FD not found in FD table")
	}
	delete(k.fds, fd)
	return realfd, nil
}

func (k *DosKernel) Terminate() {
	k.U.Exit(models.ExitStatus(0))
}

func (k *DosKernel) CharIn(buf co.Buf) byte {
	var char byte
	fmt.Scanf("%c", &char)
	k.U.MemWrite(buf.Addr, []byte{char})
	return char
}

func (k *DosKernel) CharOut(char uint16) byte {
	fmt.Printf("%c", byte(char&0xFF))
	return byte(char & 0xFF)
}

func (k *DosKernel) Display(buf co.Buf) int {
	mem := k.readUntilChar(buf.Addr, '$')

	syscall.Write(1, mem)
	k.wreg8(AL, 0x24)
	return 0x24
}

func (k *DosKernel) GetDosVersion() int {
	k.wreg16(AX, 0x7)
	return 0x7
}

func (k *DosKernel) openFile(filename string, mode int) co.Fd {
	realfd, err := syscall.Open(filename, mode, 0666)
	if err != nil {
		k.wreg16(AX, 0xFFFF)
		k.setFlagC(true)
		return 0xFFFF
	}

	// Find an internal fd number
	dosfd, err := k.getFd(realfd)
	if err != nil {
		k.wreg16(AX, dosfd)
		k.setFlagC(true)
		return 0xFFFF
	}
	k.setFlagC(false)
	k.wreg16(AX, dosfd)
	return co.Fd(dosfd)
}

func (k *DosKernel) CreateOrTruncate(buf co.Buf, attr int) co.Fd {
	filename := string(k.readUntilChar(buf.Addr, '$'))
	return k.openFile(filename, syscall.O_CREAT|syscall.O_TRUNC|syscall.O_RDWR)
}

func (k *DosKernel) Open(filename string, mode int) co.Fd {
	return k.openFile(filename, mode)
}

func (k *DosKernel) Close(fd co.Fd) {
	// Find and free the internal fd
	realfd, _ := k.freeFd(int(fd))
	err := syscall.Close(realfd)
	if err != nil {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	k.setFlagC(false)
	k.wreg16(AX, 0)
}

func (k *DosKernel) Read(fd co.Fd, buf co.Obuf, len co.Len) int {
	mem := make([]byte, len)
	n, err := syscall.Read(int(fd), mem)
	if err != nil {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	k.U.MemWrite(buf.Addr, mem)
	k.setFlagC(false)
	k.wreg16(AX, uint16(n))
	return n
}

func (k *DosKernel) Write(fd co.Fd, buf co.Buf, n co.Len) int {
	mem, _ := k.U.MemRead(buf.Addr, uint64(n))
	realfd, ok := k.fds[int(fd)]
	if !ok {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	written, err := syscall.Write(realfd, mem)
	if err != nil {
		k.setFlagC(true)
		// TODO: Set AX to error code
		k.wreg16(AX, 0xFFFF)
	}
	k.setFlagC(false)
	k.wreg16(AX, uint16(written))
	return written
}

func (k *DosKernel) Unlink(filename string, attr int) int {
	err := syscall.Unlink(filename)
	if err != nil {
		k.setFlagC(true)
		k.wreg16(AX, 0xFFFF)
		return 0xFFFF
	}
	k.setFlagC(false)
	k.wreg16(AX, 0)
	return 0
}

func (k *DosKernel) TerminateWithCode(code int) {
	k.U.Exit(models.ExitStatus(code))
}

func NewKernel() *DosKernel {
	k := &DosKernel{
		KernelBase: &co.KernelBase{},
		fds: map[int]int{
			0: 0,
			1: 1,
			2: 2,
		},
	}
	// Invert the syscall map
	for abi, syscalls := range abiMap {
		for _, syscall := range syscalls {
			syscallAbis[syscall] = *abi
		}
	}
	k.Argjoy.Register(k.getDosArgCodec())
	return k
}

func DosInit(u models.Usercorn, args, env []string) error {
	// Setup PSP
	// TODO: Setup args
	psp := initPsp(0, nil)
	u.StrucAt(0).Pack(psp)

	u.SetEntry(0x100)
	return u.MapStack(STACK_BASE, STACK_SIZE, false)
}

func DosSyscall(u models.Usercorn) {
	num, _ := u.RegRead(AH)
	name, _ := dosSysNum[int(num)]
	// TODO: How are registers numbered from here?
	u.Syscall(int(num), name, dosArgs(u, int(num)))
	// TODO: Set error
}

func (k *DosKernel) getDosArgCodec() func(interface{}, []interface{}) error {
	return func(arg interface{}, vals []interface{}) error {
		// DOS takes address as DS+DX
		if reg, ok := vals[0].(uint64); ok && len(vals) > 1 {
			ds, _ := k.U.RegRead(DS)
			reg += ds
			switch v := arg.(type) {
			case *co.Buf:
				*v = co.NewBuf(k, reg)
			case *co.Obuf:
				*v = co.Obuf{co.NewBuf(k, reg)}
			case *co.Ptr:
				*v = co.Ptr(reg)
			case *string:
				s, err := k.U.Mem().ReadStrAt(reg)
				if err != nil {
					return errors.Wrapf(err, "ReadStrAt(%#x) failed", reg)
				}
				*v = s
			default:
				return argjoy.NoMatch
			}
			return nil
		}
		return argjoy.NoMatch
	}
}

func dosArgs(u models.Usercorn, num int) func(int) ([]uint64, error) {
	// Return a closure over the correct arglist based on the syscall number
	return co.RegArgs(u, syscallAbis[num])
}

func DosInterrupt(u models.Usercorn, cause uint32) {
	intno := cause & 0xFF
	if intno == 0x21 {
		DosSyscall(u)
	} else if intno == 0x20 {
		u.Syscall(0, "terminate", func(int) ([]uint64, error) { return []uint64{}, nil })
	} else {
		panic(fmt.Sprintf("unhandled X86 interrupt %#X", intno))
	}
}
func DosKernels(u models.Usercorn) []interface{} {
	return []interface{}{NewKernel()}
}

func init() {
	Arch.RegisterOS(&models.OS{
		Name:      "DOS",
		Init:      DosInit,
		Interrupt: DosInterrupt,
		Kernels:   DosKernels,
	})
}
<|endoftext|>


<commit_before>package baseapp

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"path/filepath"
	"runtime"
	"time"

	"github.com/gorilla/mux"
	"github.com/unrolled/secure"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/sklog"
)

var (
	Local        = flag.Bool("local", false, "Running locally if true. As opposed to in production.")
	Port         = flag.String("port", ":8000", "HTTP service address (e.g., ':8000')")
	PromPort     = flag.String("prom_port", ":20000", "Metrics service address (e.g., ':10110')")
	ResourcesDir = flag.String("resources_dir", "", "The directory to find templates, JS, and CSS files. If blank the current directory will be used.")
)

const (
	SERVER_READ_TIMEOUT  = 5 * time.Minute
	SERVER_WRITE_TIMEOUT = 5 * time.Minute
)

// Applications that want to use Serve() must conform the App interface.
type App interface {
	// AddHandlers is called by Serve and the receiver must add all handlers
	// to the passed in mux.Router.
	AddHandlers(*mux.Router)

	// AddMiddleware returns a list of mux.Middleware's to add to the router.
	// This is a good place to add auth middleware.
	AddMiddleware() []mux.MiddlewareFunc
}

// Constructor is a function that builds an App instance.
//
// Used as a parameter to Serve.
type Constructor func() (App, error)

// cspReporter takes csp failure reports and turns them into structured log
// entries.
func cspReporter(w http.ResponseWriter, r *http.Request) {
	var body interface{}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		sklog.Errorf("Failed to decode csp report: %s", err)
		return
	}
	c := struct {
		Type string      `json:"type"`
		Body interface{} `json:"body"`
	}{
		Type: "csp",
		Body: body,
	}
	b, err := json.Marshal(c)
	if err != nil {
		sklog.Errorf("Failed to marshal csp log entry: %s", err)
		return
	}
	fmt.Println(string(b))
	return
}

func securityMiddleware(allowedHosts []string) mux.MiddlewareFunc {
	// Configure Content Security Policy (CSP).
	addScriptSrc := ""
	if *Local {
		// webpack uses eval() in development mode, so allow unsafe-eval when local.
		addScriptSrc = "'unsafe-eval'"
	}
	// This non-local CSP string passes the tests at https://csp-evaluator.withgoogle.com/.
	//
	// See also: https://csp.withgoogle.com/docs/strict-csp.html
	cspString := fmt.Sprintf("base-uri 'none';  img-src 'self' ; object-src 'none' ; style-src 'self'  https://fonts.googleapis.com/ 'unsafe-inline' ; script-src 'strict-dynamic' $NONCE 'unsafe-inline' %s https: http: ; report-uri /cspreport ;", addScriptSrc)

	// Apply CSP and other security minded headers.
	secureMiddleware := secure.New(secure.Options{
		AllowedHosts:          allowedHosts,
		HostsProxyHeaders:     []string{"X-Forwarded-Host"},
		SSLRedirect:           true,
		SSLProxyHeaders:       map[string]string{"X-Forwarded-Proto": "https"},
		STSSeconds:            60 * 60 * 24 * 365,
		STSIncludeSubdomains:  true,
		ContentSecurityPolicy: cspString,
		IsDevelopment:         *Local,
	})

	return secureMiddleware.Handler
}

// Serve builds and runs the App in a secure manner in out kubernetes cluster.
//
// The constructor builds an App instance. Note that we don't pass in an App
// instance directly, because we want the constructor called after the
// common.Init*() functions are called, i.e. after flags are parsed.
//
// The allowedHosts are the list of domains that are allowed to make requests
// to this app. Make sure to include the domain name of the app itself. For
// example; []string{"am.skia.org"}.
//
// See https://csp.withgoogle.com/docs/strict-csp.html for more information on
// Strict CSP in general.
//
// For this to work every script and style tag must have a nonce attribute
// whose value matches the one sent in the Content-Security-Policy: header. You
// can have webpack inject an attribute with a template for the nonce by adding
// the HtmlWebPackInjectAttributesPlugin to your plugins, i.e.
//
//    config.plugins.push(
//      new HtmlWebpackInjectAttributesPlugin({
//        nonce: "{%.nonce%}",
//      }),
//    );
//
// And then include that nonce when expanding any pages:
//
//    if err := srv.templates.ExecuteTemplate(w, "index.html", map[string]string{
//      "nonce": secure.CSPNonce(r.Context()),
//    }); err != nil {
//     sklog.Errorf("Failed to expand template: %s", err)
//   }
//
// Since our audience is small and only uses modern browsers we shouldn't need
// any further XSS protection. For example, even if a user is logged into
// another Google site that is compromised, while they can request the main
// index page and get both the csrf token and value, they couldn't POST it back
// to the site we are serving since that site wouldn't be listed in
// allowedHosts.
//
// CSP failures will be logged as structured log events.
//
// Static resources, e.g. webpack output, will be served at '/static/' and will
// serve the contents of the '/dist' directory.
func Serve(constructor Constructor, allowedHosts []string) {
	// Do common init.
	common.InitWithMust(
		"generic-k8s-app", // The app name is only used by ../go/sklog/cloud_logging, and we don't use that on k8s.
		common.PrometheusOpt(PromPort),
		common.MetricsLoggingOpt(),
	)

	// Fix up flag values.
	if *ResourcesDir == "" {
		_, filename, _, _ := runtime.Caller(1)
		*ResourcesDir = filepath.Join(filepath.Dir(filename), "../../dist")
	}

	// Build App instance.
	app, err := constructor()
	if err != nil {
		sklog.Fatal(err)
	}

	// Add all routing.
	r := mux.NewRouter()
	r.HandleFunc("/cspreport", cspReporter).Methods("POST")
	r.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.HandlerFunc(httputils.MakeResourceHandler(*ResourcesDir))))
	app.AddHandlers(r)

	// Layer on all the middleware.
	middleware := []mux.MiddlewareFunc{}
	if !*Local {
		middleware = append(middleware, httputils.HealthzAndHTTPS)
	}
	middleware = append(middleware, app.AddMiddleware()...)
	middleware = append(middleware,
		httputils.LoggingGzipRequestResponse,
		securityMiddleware(allowedHosts),
	)
	r.Use(middleware...)

	// Start serving.
	sklog.Info("Ready to serve.")
	server := &http.Server{
		Addr:           *Port,
		Handler:        r,
		ReadTimeout:    SERVER_READ_TIMEOUT,
		WriteTimeout:   SERVER_WRITE_TIMEOUT,
		MaxHeaderBytes: 1 << 20,
	}
	sklog.Fatal(server.ListenAndServe())
}
<commit_msg>Allow styles from gstatic in baseapp.go<commit_after>package baseapp

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"path/filepath"
	"runtime"
	"time"

	"github.com/gorilla/mux"
	"github.com/unrolled/secure"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/sklog"
)

var (
	Local        = flag.Bool("local", false, "Running locally if true. As opposed to in production.")
	Port         = flag.String("port", ":8000", "HTTP service address (e.g., ':8000')")
	PromPort     = flag.String("prom_port", ":20000", "Metrics service address (e.g., ':10110')")
	ResourcesDir = flag.String("resources_dir", "", "The directory to find templates, JS, and CSS files. If blank the current directory will be used.")
)

const (
	SERVER_READ_TIMEOUT  = 5 * time.Minute
	SERVER_WRITE_TIMEOUT = 5 * time.Minute
)

// Applications that want to use Serve() must conform the App interface.
type App interface {
	// AddHandlers is called by Serve and the receiver must add all handlers
	// to the passed in mux.Router.
	AddHandlers(*mux.Router)

	// AddMiddleware returns a list of mux.Middleware's to add to the router.
	// This is a good place to add auth middleware.
	AddMiddleware() []mux.MiddlewareFunc
}

// Constructor is a function that builds an App instance.
//
// Used as a parameter to Serve.
type Constructor func() (App, error)

// cspReporter takes csp failure reports and turns them into structured log
// entries.
func cspReporter(w http.ResponseWriter, r *http.Request) {
	var body interface{}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		sklog.Errorf("Failed to decode csp report: %s", err)
		return
	}
	c := struct {
		Type string      `json:"type"`
		Body interface{} `json:"body"`
	}{
		Type: "csp",
		Body: body,
	}
	b, err := json.Marshal(c)
	if err != nil {
		sklog.Errorf("Failed to marshal csp log entry: %s", err)
		return
	}
	fmt.Println(string(b))
	return
}

func securityMiddleware(allowedHosts []string) mux.MiddlewareFunc {
	// Configure Content Security Policy (CSP).
	addScriptSrc := ""
	if *Local {
		// webpack uses eval() in development mode, so allow unsafe-eval when local.
		addScriptSrc = "'unsafe-eval'"
	}
	// This non-local CSP string passes the tests at https://csp-evaluator.withgoogle.com/.
	//
	// See also: https://csp.withgoogle.com/docs/strict-csp.html
	cspString := fmt.Sprintf("base-uri 'none';  img-src 'self' ; object-src 'none' ; style-src 'self'  https://fonts.googleapis.com/ https://www.gstatic.com/ 'unsafe-inline' ; script-src 'strict-dynamic' $NONCE 'unsafe-inline' %s https: http: ; report-uri /cspreport ;", addScriptSrc)

	// Apply CSP and other security minded headers.
	secureMiddleware := secure.New(secure.Options{
		AllowedHosts:          allowedHosts,
		HostsProxyHeaders:     []string{"X-Forwarded-Host"},
		SSLRedirect:           true,
		SSLProxyHeaders:       map[string]string{"X-Forwarded-Proto": "https"},
		STSSeconds:            60 * 60 * 24 * 365,
		STSIncludeSubdomains:  true,
		ContentSecurityPolicy: cspString,
		IsDevelopment:         *Local,
	})

	return secureMiddleware.Handler
}

// Serve builds and runs the App in a secure manner in out kubernetes cluster.
//
// The constructor builds an App instance. Note that we don't pass in an App
// instance directly, because we want the constructor called after the
// common.Init*() functions are called, i.e. after flags are parsed.
//
// The allowedHosts are the list of domains that are allowed to make requests
// to this app. Make sure to include the domain name of the app itself. For
// example; []string{"am.skia.org"}.
//
// See https://csp.withgoogle.com/docs/strict-csp.html for more information on
// Strict CSP in general.
//
// For this to work every script and style tag must have a nonce attribute
// whose value matches the one sent in the Content-Security-Policy: header. You
// can have webpack inject an attribute with a template for the nonce by adding
// the HtmlWebPackInjectAttributesPlugin to your plugins, i.e.
//
//    config.plugins.push(
//      new HtmlWebpackInjectAttributesPlugin({
//        nonce: "{%.nonce%}",
//      }),
//    );
//
// And then include that nonce when expanding any pages:
//
//    if err := srv.templates.ExecuteTemplate(w, "index.html", map[string]string{
//      "nonce": secure.CSPNonce(r.Context()),
//    }); err != nil {
//     sklog.Errorf("Failed to expand template: %s", err)
//   }
//
// Since our audience is small and only uses modern browsers we shouldn't need
// any further XSS protection. For example, even if a user is logged into
// another Google site that is compromised, while they can request the main
// index page and get both the csrf token and value, they couldn't POST it back
// to the site we are serving since that site wouldn't be listed in
// allowedHosts.
//
// CSP failures will be logged as structured log events.
//
// Static resources, e.g. webpack output, will be served at '/static/' and will
// serve the contents of the '/dist' directory.
func Serve(constructor Constructor, allowedHosts []string) {
	// Do common init.
	common.InitWithMust(
		"generic-k8s-app", // The app name is only used by ../go/sklog/cloud_logging, and we don't use that on k8s.
		common.PrometheusOpt(PromPort),
		common.MetricsLoggingOpt(),
	)

	// Fix up flag values.
	if *ResourcesDir == "" {
		_, filename, _, _ := runtime.Caller(1)
		*ResourcesDir = filepath.Join(filepath.Dir(filename), "../../dist")
	}

	// Build App instance.
	app, err := constructor()
	if err != nil {
		sklog.Fatal(err)
	}

	// Add all routing.
	r := mux.NewRouter()
	r.HandleFunc("/cspreport", cspReporter).Methods("POST")
	r.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.HandlerFunc(httputils.MakeResourceHandler(*ResourcesDir))))
	app.AddHandlers(r)

	// Layer on all the middleware.
	middleware := []mux.MiddlewareFunc{}
	if !*Local {
		middleware = append(middleware, httputils.HealthzAndHTTPS)
	}
	middleware = append(middleware, app.AddMiddleware()...)
	middleware = append(middleware,
		httputils.LoggingGzipRequestResponse,
		securityMiddleware(allowedHosts),
	)
	r.Use(middleware...)

	// Start serving.
	sklog.Info("Ready to serve.")
	server := &http.Server{
		Addr:           *Port,
		Handler:        r,
		ReadTimeout:    SERVER_READ_TIMEOUT,
		WriteTimeout:   SERVER_WRITE_TIMEOUT,
		MaxHeaderBytes: 1 << 20,
	}
	sklog.Fatal(server.ListenAndServe())
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
)

type PGPPullEngineArg struct {
	UserAsserts []string
}

type PGPPullEngine struct {
	listTrackingEngine *ListTrackingEngine
	userAsserts        []string
	libkb.Contextified
}

func NewPGPPullEngine(arg *PGPPullEngineArg, g *libkb.GlobalContext) *PGPPullEngine {
	return &PGPPullEngine{
		listTrackingEngine: NewListTrackingEngine(&ListTrackingEngineArg{}, g),
		userAsserts:        arg.UserAsserts,
		Contextified:       libkb.NewContextified(g),
	}
}

func (e *PGPPullEngine) Name() string {
	return "PGPPull"
}

func (e *PGPPullEngine) Prereqs() Prereqs {
	return Prereqs{
		Session: true,
	}
}

func (e *PGPPullEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.LogUIKind,
	}
}

func (e *PGPPullEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{e.listTrackingEngine}
}

func proofSetFromUserSummary(summary keybase1.UserSummary) *libkb.ProofSet {
	// TODO: This logic is kind of duplicated from User.ToOkProofSet(). Clean that up?
	proofs := []libkb.Proof{
		{Key: "keybase", Value: summary.Username},
		{Key: "uid", Value: summary.Uid.String()},
	}
	for _, socialProof := range summary.Proofs.Social {
		proofs = append(proofs, libkb.Proof{
			Key:   socialProof.ProofType,
			Value: socialProof.ProofName,
		})
	}
	for _, webProof := range summary.Proofs.Web {
		for _, protocol := range webProof.Protocols {
			proofs = append(proofs, libkb.Proof{
				Key:   protocol,
				Value: webProof.Hostname,
			})
		}
	}
	return libkb.NewProofSet(proofs)
}

func (e *PGPPullEngine) getTrackedUserSummaries(ctx *Context) ([]keybase1.UserSummary, error) {
	err := RunEngine(e.listTrackingEngine, ctx)
	if err != nil {
		return nil, err
	}
	allTrackedSummaries := e.listTrackingEngine.TableResult()

	// Without any userAsserts specified, just return everything.
	if e.userAsserts == nil || len(e.userAsserts) == 0 {
		return allTrackedSummaries, nil
	}

	// With userAsserts specified, return only those summaries. If an assert
	// doesn't match any tracked users, that's an error. If an assert matches
	// more than one tracked user, that is also an error. If multiple
	// assertions match the same user, that's fine.

	// First parse all the assertion expressions.
	parsedAsserts := make(map[string]libkb.AssertionExpression)
	for _, assertString := range e.userAsserts {
		assertExpr, err := libkb.AssertionParseAndOnly(assertString)
		if err != nil {
			return nil, err
		}
		parsedAsserts[assertString] = assertExpr
	}

	// Then loop over all the tracked users, keeping track of which expressions
	// have matched before.
	matchedSummaries := make(map[string]keybase1.UserSummary)
	assertionsUsed := make(map[string]bool)
	for _, summary := range allTrackedSummaries {
		proofSet := proofSetFromUserSummary(summary)
		for assertStr, parsedAssert := range parsedAsserts {
			if parsedAssert.MatchSet(*proofSet) {
				if assertionsUsed[assertStr] {
					return nil, fmt.Errorf("Assertion \"%s\" matched more than one tracked user.", assertStr)
				}
				assertionsUsed[assertStr] = true
				matchedSummaries[summary.Username] = summary
			}
		}
	}

	// Make sure every assertion found a match.
	for _, assertString := range e.userAsserts {
		if !assertionsUsed[assertString] {
			return nil, fmt.Errorf("Assertion \"%s\" did not match any tracked users.", assertString)
		}
	}

	matchedList := []keybase1.UserSummary{}
	for _, summary := range matchedSummaries {
		matchedList = append(matchedList, summary)
	}
	return matchedList, nil
}

func (e *PGPPullEngine) Run(ctx *Context) error {
	summaries, err := e.getTrackedUserSummaries(ctx)
	if err != nil {
		return err
	}

	gpgClient := libkb.NewGpgCLI(libkb.GpgCLIArg{
		LogUI: ctx.LogUI,
	})
	_, err = gpgClient.Configure()
	if err != nil {
		return err
	}

	// Loop over the list of all users we track.
	for _, userSummary := range summaries {
		// Compute the set of tracked pgp fingerprints. LoadUser will fetch key
		// data from the server, and we will compare it against this.
		trackedFingerprints := make(map[string]bool)
		for _, pubKey := range userSummary.Proofs.PublicKeys {
			if pubKey.PGPFingerprint != "" {
				trackedFingerprints[pubKey.PGPFingerprint] = true
			}
		}

		// Get user data from the server.
		loadUserArg := libkb.LoadUserArg{
			Name: userSummary.Username,
		}
		user, err := libkb.LoadUser(loadUserArg)
		if err != nil {
			ctx.LogUI.Errorf("Failed to load user %s: %s", userSummary.Username, err)
			continue
		}

		for _, bundle := range user.GetActivePGPKeys(false) {
			// Check each key against the tracked set.
			if !trackedFingerprints[bundle.GetFingerprint().String()] {
				ctx.LogUI.Warning("Keybase says that %s owns key %s, but you have not tracked this fingerprint before.", user.GetName(), bundle.GetFingerprint())
				continue
			}

			err = gpgClient.ExportKey(*bundle)
			if err != nil {
				return err
			}
			ctx.LogUI.Info("Imported key for %s.", user.GetName())
		}
	}
	return nil
}
<commit_msg>remove an out-of-date TODO<commit_after>package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
)

type PGPPullEngineArg struct {
	UserAsserts []string
}

type PGPPullEngine struct {
	listTrackingEngine *ListTrackingEngine
	userAsserts        []string
	libkb.Contextified
}

func NewPGPPullEngine(arg *PGPPullEngineArg, g *libkb.GlobalContext) *PGPPullEngine {
	return &PGPPullEngine{
		listTrackingEngine: NewListTrackingEngine(&ListTrackingEngineArg{}, g),
		userAsserts:        arg.UserAsserts,
		Contextified:       libkb.NewContextified(g),
	}
}

func (e *PGPPullEngine) Name() string {
	return "PGPPull"
}

func (e *PGPPullEngine) Prereqs() Prereqs {
	return Prereqs{
		Session: true,
	}
}

func (e *PGPPullEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.LogUIKind,
	}
}

func (e *PGPPullEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{e.listTrackingEngine}
}

func proofSetFromUserSummary(summary keybase1.UserSummary) *libkb.ProofSet {
	proofs := []libkb.Proof{
		{Key: "keybase", Value: summary.Username},
		{Key: "uid", Value: summary.Uid.String()},
	}
	for _, socialProof := range summary.Proofs.Social {
		proofs = append(proofs, libkb.Proof{
			Key:   socialProof.ProofType,
			Value: socialProof.ProofName,
		})
	}
	for _, webProof := range summary.Proofs.Web {
		for _, protocol := range webProof.Protocols {
			proofs = append(proofs, libkb.Proof{
				Key:   protocol,
				Value: webProof.Hostname,
			})
		}
	}
	return libkb.NewProofSet(proofs)
}

func (e *PGPPullEngine) getTrackedUserSummaries(ctx *Context) ([]keybase1.UserSummary, error) {
	err := RunEngine(e.listTrackingEngine, ctx)
	if err != nil {
		return nil, err
	}
	allTrackedSummaries := e.listTrackingEngine.TableResult()

	// Without any userAsserts specified, just return everything.
	if e.userAsserts == nil || len(e.userAsserts) == 0 {
		return allTrackedSummaries, nil
	}

	// With userAsserts specified, return only those summaries. If an assert
	// doesn't match any tracked users, that's an error. If an assert matches
	// more than one tracked user, that is also an error. If multiple
	// assertions match the same user, that's fine.

	// First parse all the assertion expressions.
	parsedAsserts := make(map[string]libkb.AssertionExpression)
	for _, assertString := range e.userAsserts {
		assertExpr, err := libkb.AssertionParseAndOnly(assertString)
		if err != nil {
			return nil, err
		}
		parsedAsserts[assertString] = assertExpr
	}

	// Then loop over all the tracked users, keeping track of which expressions
	// have matched before.
	matchedSummaries := make(map[string]keybase1.UserSummary)
	assertionsUsed := make(map[string]bool)
	for _, summary := range allTrackedSummaries {
		proofSet := proofSetFromUserSummary(summary)
		for assertStr, parsedAssert := range parsedAsserts {
			if parsedAssert.MatchSet(*proofSet) {
				if assertionsUsed[assertStr] {
					return nil, fmt.Errorf("Assertion \"%s\" matched more than one tracked user.", assertStr)
				}
				assertionsUsed[assertStr] = true
				matchedSummaries[summary.Username] = summary
			}
		}
	}

	// Make sure every assertion found a match.
	for _, assertString := range e.userAsserts {
		if !assertionsUsed[assertString] {
			return nil, fmt.Errorf("Assertion \"%s\" did not match any tracked users.", assertString)
		}
	}

	matchedList := []keybase1.UserSummary{}
	for _, summary := range matchedSummaries {
		matchedList = append(matchedList, summary)
	}
	return matchedList, nil
}

func (e *PGPPullEngine) Run(ctx *Context) error {
	summaries, err := e.getTrackedUserSummaries(ctx)
	if err != nil {
		return err
	}

	gpgClient := libkb.NewGpgCLI(libkb.GpgCLIArg{
		LogUI: ctx.LogUI,
	})
	_, err = gpgClient.Configure()
	if err != nil {
		return err
	}

	// Loop over the list of all users we track.
	for _, userSummary := range summaries {
		// Compute the set of tracked pgp fingerprints. LoadUser will fetch key
		// data from the server, and we will compare it against this.
		trackedFingerprints := make(map[string]bool)
		for _, pubKey := range userSummary.Proofs.PublicKeys {
			if pubKey.PGPFingerprint != "" {
				trackedFingerprints[pubKey.PGPFingerprint] = true
			}
		}

		// Get user data from the server.
		loadUserArg := libkb.LoadUserArg{
			Name: userSummary.Username,
		}
		user, err := libkb.LoadUser(loadUserArg)
		if err != nil {
			ctx.LogUI.Errorf("Failed to load user %s: %s", userSummary.Username, err)
			continue
		}

		for _, bundle := range user.GetActivePGPKeys(false) {
			// Check each key against the tracked set.
			if !trackedFingerprints[bundle.GetFingerprint().String()] {
				ctx.LogUI.Warning("Keybase says that %s owns key %s, but you have not tracked this fingerprint before.", user.GetName(), bundle.GetFingerprint())
				continue
			}

			err = gpgClient.ExportKey(*bundle)
			if err != nil {
				return err
			}
			ctx.LogUI.Info("Imported key for %s.", user.GetName())
		}
	}
	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 (
	"fmt"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

type PGPPullEngineArg struct {
	UserAsserts []string
}

type PGPPullEngine struct {
	listTrackingEngine *ListTrackingEngine
	userAsserts        []string
	gpgClient          *libkb.GpgCLI
	libkb.Contextified
}

func NewPGPPullEngine(arg *PGPPullEngineArg, g *libkb.GlobalContext) *PGPPullEngine {
	return &PGPPullEngine{
		listTrackingEngine: NewListTrackingEngine(&ListTrackingEngineArg{}, g),
		userAsserts:        arg.UserAsserts,
		Contextified:       libkb.NewContextified(g),
	}
}

func (e *PGPPullEngine) Name() string {
	return "PGPPull"
}

func (e *PGPPullEngine) Prereqs() Prereqs {
	return Prereqs{}
}

func (e *PGPPullEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.LogUIKind,
	}
}

func (e *PGPPullEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{e.listTrackingEngine}
}

func proofSetFromUserSummary(summary keybase1.UserSummary) *libkb.ProofSet {
	proofs := []libkb.Proof{
		{Key: "keybase", Value: summary.Username},
		{Key: "uid", Value: summary.Uid.String()},
	}
	for _, socialProof := range summary.Proofs.Social {
		proofs = append(proofs, libkb.Proof{
			Key:   socialProof.ProofType,
			Value: socialProof.ProofName,
		})
	}
	for _, webProof := range summary.Proofs.Web {
		for _, protocol := range webProof.Protocols {
			proofs = append(proofs, libkb.Proof{
				Key:   protocol,
				Value: webProof.Hostname,
			})
		}
	}
	return libkb.NewProofSet(proofs)
}

func (e *PGPPullEngine) getTrackedUserSummaries(ctx *Context) ([]keybase1.UserSummary, error) {
	err := RunEngine(e.listTrackingEngine, ctx)
	if err != nil {
		return nil, err
	}
	allTrackedSummaries := e.listTrackingEngine.TableResult()

	// Without any userAsserts specified, just return everything.
	if e.userAsserts == nil || len(e.userAsserts) == 0 {
		return allTrackedSummaries, nil
	}

	// With userAsserts specified, return only those summaries. If an assert
	// doesn't match any tracked users, that's an error. If an assert matches
	// more than one tracked user, that is also an error. If multiple
	// assertions match the same user, that's fine.

	// First parse all the assertion expressions.
	parsedAsserts := make(map[string]libkb.AssertionExpression)
	for _, assertString := range e.userAsserts {
		assertExpr, err := libkb.AssertionParseAndOnly(assertString)
		if err != nil {
			return nil, err
		}
		parsedAsserts[assertString] = assertExpr
	}

	// Then loop over all the tracked users, keeping track of which expressions
	// have matched before.
	matchedSummaries := make(map[string]keybase1.UserSummary)
	assertionsUsed := make(map[string]bool)
	for _, summary := range allTrackedSummaries {
		proofSet := proofSetFromUserSummary(summary)
		for assertStr, parsedAssert := range parsedAsserts {
			if parsedAssert.MatchSet(*proofSet) {
				if assertionsUsed[assertStr] {
					return nil, fmt.Errorf("Assertion \"%s\" matched more than one tracked user.", assertStr)
				}
				assertionsUsed[assertStr] = true
				matchedSummaries[summary.Username] = summary
			}
		}
	}

	// Make sure every assertion found a match.
	for _, assertString := range e.userAsserts {
		if !assertionsUsed[assertString] {
			return nil, fmt.Errorf("Assertion \"%s\" did not match any tracked users.", assertString)
		}
	}

	matchedList := []keybase1.UserSummary{}
	for _, summary := range matchedSummaries {
		matchedList = append(matchedList, summary)
	}
	return matchedList, nil
}

func (e *PGPPullEngine) runLoggedOut(ctx *Context) error {
	if len(e.userAsserts) == 0 {
		return libkb.PGPPullLoggedOutError{}
	}
	t := time.Now()
	for i, assertString := range e.userAsserts {
		t = e.rateLimit(t, i)
		if err := e.processUserWhenLoggedOut(ctx, assertString); err != nil {
			return err
		}
	}
	return nil
}

func (e *PGPPullEngine) processUserWhenLoggedOut(ctx *Context, u string) error {
	iarg := NewIdentifyTrackArg(u, false, true, keybase1.TrackOptions{
		LocalOnly: true,
	})
	ieng := NewIdentify(iarg, e.G())
	if err := RunEngine(ieng, ctx); err != nil {
		e.G().Log.Info("identify run err: %s", err)
		return err
	}

	// prompt if the identify is correct
	outcome := ieng.Outcome().Export()
	outcome.ForPGPPull = true
	result, err := ctx.IdentifyUI.Confirm(outcome)
	if err != nil {
		return err
	}

	if !result.IdentityConfirmed {
		e.G().Log.Warning("Not confirmed; skipping key import")
		return nil
	}

	if err = e.exportKeysToGPG(ctx, ieng.User(), nil); err != nil {
		return err
	}

	return nil
}

func (e *PGPPullEngine) Run(ctx *Context) error {

	e.gpgClient = libkb.NewGpgCLI(e.G(), ctx.LogUI)
	err := e.gpgClient.Configure()
	if err != nil {
		return err
	}

	if ok, _, _ := IsLoggedIn(e, ctx); !ok {
		return e.runLoggedOut(ctx)
	}

	summaries, err := e.getTrackedUserSummaries(ctx)
	if err != nil {
		return err
	}

	return e.runLoggedIn(ctx, summaries)
}

func (e *PGPPullEngine) runLoggedIn(ctx *Context, summaries []keybase1.UserSummary) error {

	// Loop over the list of all users we track.
	t := time.Now()
	for i, userSummary := range summaries {
		t = e.rateLimit(t, i)
		// Compute the set of tracked pgp fingerprints. LoadUser will fetch key
		// data from the server, and we will compare it against this.
		trackedFingerprints := make(map[string]bool)
		for _, pubKey := range userSummary.Proofs.PublicKeys {
			if pubKey.PGPFingerprint != "" {
				trackedFingerprints[pubKey.PGPFingerprint] = true
			}
		}

		// Get user data from the server.
		user, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(e.G(), userSummary.Username))
		if err != nil {
			ctx.LogUI.Errorf("Failed to load user %s: %s", userSummary.Username, err)
			continue
		}

		if err = e.exportKeysToGPG(ctx, user, trackedFingerprints); err != nil {
			return err
		}
	}
	return nil
}

func (e *PGPPullEngine) exportKeysToGPG(ctx *Context, user *libkb.User, tfp map[string]bool) error {
	for _, bundle := range user.GetActivePGPKeys(false) {
		// Check each key against the tracked set.
		if tfp != nil && !tfp[bundle.GetFingerprint().String()] {
			ctx.LogUI.Warning("Keybase says that %s owns key %s, but you have not tracked this fingerprint before.", user.GetName(), bundle.GetFingerprint())
			continue
		}

		if err := e.gpgClient.ExportKey(*bundle, false /* export public key only */); err != nil {
			return err
		}

		ctx.LogUI.Info("Imported key for %s.", user.GetName())
	}
	return nil
}

func (e *PGPPullEngine) rateLimit(start time.Time, index int) time.Time {
	// server currently limiting to 32 req/s, but there can be 4 requests for each loaduser call.
	const loadUserPerSec = 8
	if index == 0 {
		return start
	}
	if index%loadUserPerSec != 0 {
		return start
	}
	d := time.Second - time.Since(start)
	if d > 0 {
		e.G().Log.Debug("sleeping for %s to slow down api requests", d)
		time.Sleep(d)
	}
	return time.Now()
}
<commit_msg>Change load user per sec from 8 to 4<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"fmt"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

type PGPPullEngineArg struct {
	UserAsserts []string
}

type PGPPullEngine struct {
	listTrackingEngine *ListTrackingEngine
	userAsserts        []string
	gpgClient          *libkb.GpgCLI
	libkb.Contextified
}

func NewPGPPullEngine(arg *PGPPullEngineArg, g *libkb.GlobalContext) *PGPPullEngine {
	return &PGPPullEngine{
		listTrackingEngine: NewListTrackingEngine(&ListTrackingEngineArg{}, g),
		userAsserts:        arg.UserAsserts,
		Contextified:       libkb.NewContextified(g),
	}
}

func (e *PGPPullEngine) Name() string {
	return "PGPPull"
}

func (e *PGPPullEngine) Prereqs() Prereqs {
	return Prereqs{}
}

func (e *PGPPullEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.LogUIKind,
	}
}

func (e *PGPPullEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{e.listTrackingEngine}
}

func proofSetFromUserSummary(summary keybase1.UserSummary) *libkb.ProofSet {
	proofs := []libkb.Proof{
		{Key: "keybase", Value: summary.Username},
		{Key: "uid", Value: summary.Uid.String()},
	}
	for _, socialProof := range summary.Proofs.Social {
		proofs = append(proofs, libkb.Proof{
			Key:   socialProof.ProofType,
			Value: socialProof.ProofName,
		})
	}
	for _, webProof := range summary.Proofs.Web {
		for _, protocol := range webProof.Protocols {
			proofs = append(proofs, libkb.Proof{
				Key:   protocol,
				Value: webProof.Hostname,
			})
		}
	}
	return libkb.NewProofSet(proofs)
}

func (e *PGPPullEngine) getTrackedUserSummaries(ctx *Context) ([]keybase1.UserSummary, error) {
	err := RunEngine(e.listTrackingEngine, ctx)
	if err != nil {
		return nil, err
	}
	allTrackedSummaries := e.listTrackingEngine.TableResult()

	// Without any userAsserts specified, just return everything.
	if e.userAsserts == nil || len(e.userAsserts) == 0 {
		return allTrackedSummaries, nil
	}

	// With userAsserts specified, return only those summaries. If an assert
	// doesn't match any tracked users, that's an error. If an assert matches
	// more than one tracked user, that is also an error. If multiple
	// assertions match the same user, that's fine.

	// First parse all the assertion expressions.
	parsedAsserts := make(map[string]libkb.AssertionExpression)
	for _, assertString := range e.userAsserts {
		assertExpr, err := libkb.AssertionParseAndOnly(assertString)
		if err != nil {
			return nil, err
		}
		parsedAsserts[assertString] = assertExpr
	}

	// Then loop over all the tracked users, keeping track of which expressions
	// have matched before.
	matchedSummaries := make(map[string]keybase1.UserSummary)
	assertionsUsed := make(map[string]bool)
	for _, summary := range allTrackedSummaries {
		proofSet := proofSetFromUserSummary(summary)
		for assertStr, parsedAssert := range parsedAsserts {
			if parsedAssert.MatchSet(*proofSet) {
				if assertionsUsed[assertStr] {
					return nil, fmt.Errorf("Assertion \"%s\" matched more than one tracked user.", assertStr)
				}
				assertionsUsed[assertStr] = true
				matchedSummaries[summary.Username] = summary
			}
		}
	}

	// Make sure every assertion found a match.
	for _, assertString := range e.userAsserts {
		if !assertionsUsed[assertString] {
			return nil, fmt.Errorf("Assertion \"%s\" did not match any tracked users.", assertString)
		}
	}

	matchedList := []keybase1.UserSummary{}
	for _, summary := range matchedSummaries {
		matchedList = append(matchedList, summary)
	}
	return matchedList, nil
}

func (e *PGPPullEngine) runLoggedOut(ctx *Context) error {
	if len(e.userAsserts) == 0 {
		return libkb.PGPPullLoggedOutError{}
	}
	t := time.Now()
	for i, assertString := range e.userAsserts {
		t = e.rateLimit(t, i)
		if err := e.processUserWhenLoggedOut(ctx, assertString); err != nil {
			return err
		}
	}
	return nil
}

func (e *PGPPullEngine) processUserWhenLoggedOut(ctx *Context, u string) error {
	iarg := NewIdentifyTrackArg(u, false, true, keybase1.TrackOptions{
		LocalOnly: true,
	})
	ieng := NewIdentify(iarg, e.G())
	if err := RunEngine(ieng, ctx); err != nil {
		e.G().Log.Info("identify run err: %s", err)
		return err
	}

	// prompt if the identify is correct
	outcome := ieng.Outcome().Export()
	outcome.ForPGPPull = true
	result, err := ctx.IdentifyUI.Confirm(outcome)
	if err != nil {
		return err
	}

	if !result.IdentityConfirmed {
		e.G().Log.Warning("Not confirmed; skipping key import")
		return nil
	}

	if err = e.exportKeysToGPG(ctx, ieng.User(), nil); err != nil {
		return err
	}

	return nil
}

func (e *PGPPullEngine) Run(ctx *Context) error {

	e.gpgClient = libkb.NewGpgCLI(e.G(), ctx.LogUI)
	err := e.gpgClient.Configure()
	if err != nil {
		return err
	}

	if ok, _, _ := IsLoggedIn(e, ctx); !ok {
		return e.runLoggedOut(ctx)
	}

	summaries, err := e.getTrackedUserSummaries(ctx)
	if err != nil {
		return err
	}

	return e.runLoggedIn(ctx, summaries)
}

func (e *PGPPullEngine) runLoggedIn(ctx *Context, summaries []keybase1.UserSummary) error {

	// Loop over the list of all users we track.
	t := time.Now()
	for i, userSummary := range summaries {
		t = e.rateLimit(t, i)
		// Compute the set of tracked pgp fingerprints. LoadUser will fetch key
		// data from the server, and we will compare it against this.
		trackedFingerprints := make(map[string]bool)
		for _, pubKey := range userSummary.Proofs.PublicKeys {
			if pubKey.PGPFingerprint != "" {
				trackedFingerprints[pubKey.PGPFingerprint] = true
			}
		}

		// Get user data from the server.
		user, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(e.G(), userSummary.Username))
		if err != nil {
			ctx.LogUI.Errorf("Failed to load user %s: %s", userSummary.Username, err)
			continue
		}

		if err = e.exportKeysToGPG(ctx, user, trackedFingerprints); err != nil {
			return err
		}
	}
	return nil
}

func (e *PGPPullEngine) exportKeysToGPG(ctx *Context, user *libkb.User, tfp map[string]bool) error {
	for _, bundle := range user.GetActivePGPKeys(false) {
		// Check each key against the tracked set.
		if tfp != nil && !tfp[bundle.GetFingerprint().String()] {
			ctx.LogUI.Warning("Keybase says that %s owns key %s, but you have not tracked this fingerprint before.", user.GetName(), bundle.GetFingerprint())
			continue
		}

		if err := e.gpgClient.ExportKey(*bundle, false /* export public key only */); err != nil {
			return err
		}

		ctx.LogUI.Info("Imported key for %s.", user.GetName())
	}
	return nil
}

func (e *PGPPullEngine) rateLimit(start time.Time, index int) time.Time {
	// server currently limiting to 32 req/s, but there can be 4 requests for each loaduser call.
	const loadUserPerSec = 4
	if index == 0 {
		return start
	}
	if index%loadUserPerSec != 0 {
		return start
	}
	d := time.Second - time.Since(start)
	if d > 0 {
		e.G().Log.Debug("sleeping for %s to slow down api requests", d)
		time.Sleep(d)
	}
	return time.Now()
}
<|endoftext|>


<commit_before>// Copyright 2012, 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 key

import (
	"encoding/json"
	"testing"
)

func TestKey(t *testing.T) {
	k0 := Uint64Key(0).KeyspaceId()
	k1 := Uint64Key(1).KeyspaceId()
	k2 := Uint64Key(0x7FFFFFFFFFFFFFFF).KeyspaceId()
	k3 := Uint64Key(0x8000000000000000).KeyspaceId()
	k4 := Uint64Key(0xFFFFFFFFFFFFFFFF).KeyspaceId()

	f := func(k KeyspaceId, x string) {
		if x != string(k) {
			t.Errorf("byte mismatch %#v != %#v", k, x)
		}
		data, err := json.MarshalIndent(k, "  ", "  ")
		if err != nil {
			t.Errorf("serialize error: %v", err)
		} else {
			t.Logf("json: %v", string(data))
		}

		k_r := new(KeyspaceId)
		err = json.Unmarshal(data, k_r)
		if err != nil {
			t.Errorf("reserialize error: %v", err)
		}

		if k != *k_r {
			t.Errorf("keyspace compare failed: %#v != %#v", k, k_r)
		}
	}

	f(MinKey, "")
	f(k0, "\x00\x00\x00\x00\x00\x00\x00\x00")
	f(k1, "\x00\x00\x00\x00\x00\x00\x00\x01")
	f(k2, "\x7f\xff\xff\xff\xff\xff\xff\xff")
	f(k3, "\x80\x00\x00\x00\x00\x00\x00\x00")
	f(k4, "\xff\xff\xff\xff\xff\xff\xff\xff")

	hv := k4.Hex()
	if hv != "FFFFFFFFFFFFFFFF" {
		t.Errorf("Was expecting FFFFFFFFFFFFFFFF but got %v", hv)
	}
}

func TestKeyUint64Sort(t *testing.T) {
	k0 := Uint64Key(0).KeyspaceId()
	k1 := Uint64Key(1).KeyspaceId()
	k2 := Uint64Key(0x7FFFFFFFFFFFFFFF).KeyspaceId()
	k3 := Uint64Key(0x8000000000000000).KeyspaceId()
	k4 := Uint64Key(0xFFFFFFFFFFFFFFFF).KeyspaceId()
	kl := make([]KeyspaceId, 0, 16)
	klSorted := make([]KeyspaceId, 0, 16)
	kl = append(kl, MinKey, MaxKey, k4, k3, k2, k1, k0)
	klSorted = append(kl, MinKey, k0, k1, k2, k3, k4, MaxKey)
	KeyspaceIdArray(kl).Sort()

	for i, k := range kl {
		if k != klSorted[i] {
			t.Errorf("key order error: %d %v %v", i, k, klSorted[i])
		}
	}
}

func TestKeyStringSort(t *testing.T) {
	k0 := KeyspaceId("0")
	k1 := KeyspaceId("9")
	k2 := KeyspaceId("Zzzz")
	k3 := KeyspaceId("a")
	k4 := KeyspaceId("z")
	kl := make([]KeyspaceId, 0, 16)
	klSorted := make([]KeyspaceId, 0, 16)
	kl = append(kl, MinKey, MaxKey, k4, k3, k2, k1, k0)
	klSorted = append(kl, MinKey, k0, k1, k2, k3, k4, MaxKey)
	KeyspaceIdArray(kl).Sort()

	for i, k := range kl {
		if k != klSorted[i] {
			t.Errorf("key order error: %d %v %v", i, k, klSorted[i])
		}
	}
}

func TestParseShardingSpec(t *testing.T) {
	x40, err := HexKeyspaceId("4000000000000000").Unhex()
	if err != nil {
		t.Errorf("Unexpected error: %v.", err)
	}
	x80, err := HexKeyspaceId("8000000000000000").Unhex()
	if err != nil {
		t.Errorf("Unexpected error: %v.", err)
	}

	goodTable := map[string][]KeyRange{
		"-": {{Start: MinKey, End: MaxKey}},
		"-4000000000000000-8000000000000000-": {
			{Start: MinKey, End: x40},
			{Start: x40, End: x80},
			{Start: x80, End: MaxKey},
		},
	}
	badTable := []string{
		"4000000000000000",
		"---",
		"4000000000000000--8000000000000000",
		"4000000000000000-3000000000000000", // not in order
	}
	for key, wanted := range goodTable {
		r, err := ParseShardingSpec(key)
		if err != nil {
			t.Errorf("Unexpected error: %v.", err)
		}
		if len(r) != len(wanted) {
			t.Errorf("Wrong result: wanted %v, got %v", wanted, r)
			continue
		}
		for i, w := range wanted {
			if r[i] != w {
				t.Errorf("Wrong result: wanted %v, got %v", wanted, r)
				break
			}
		}
	}
	for _, bad := range badTable {
		_, err := ParseShardingSpec(bad)
		if err == nil {
			t.Errorf("Didn't get expected error for %v.", bad)
		}
	}
}

func TestContains(t *testing.T) {
	var table = []struct {
		kid       string
		start     string
		end       string
		contained bool
	}{
		{kid: "3000000000000000", start: "3000000000000000", end: "", contained: true},
		{kid: "3000000000000000", start: "", end: "3000000000000000", contained: false},
		{kid: "4000000000000000", start: "3000000000000000", end: "", contained: true},
		{kid: "2000000000000000", start: "3000000000000000", end: "", contained: false},
	}

	for _, el := range table {
		s, err := HexKeyspaceId(el.start).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		e, err := HexKeyspaceId(el.end).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		kr := KeyRange{Start: s, End: e}
		k, err := HexKeyspaceId(el.kid).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		if c := kr.Contains(k); c != el.contained {
			t.Errorf("Unexpected result: contains for %v and (%v-%v) yields %v.", el.kid, el.start, el.end, c)
		}
	}
}

func TestIntersectOverlap(t *testing.T) {
	var table = []struct {
		a          string
		b          string
		c          string
		d          string
		intersects bool
		overlap    string
	}{
		{a: "40", b: "80", c: "c0", d: "d0", intersects: false},
		{a: "", b: "80", c: "80", d: "", intersects: false},
		{a: "", b: "80", c: "", d: "40", intersects: true, overlap: "-40"},
		{a: "80", b: "", c: "C0", d: "", intersects: true, overlap: "C0-"},
		{a: "", b: "80", c: "40", d: "80", intersects: true, overlap: "40-80"},
		{a: "40", b: "80", c: "60", d: "A0", intersects: true, overlap: "60-80"},
		{a: "40", b: "80", c: "50", d: "60", intersects: true, overlap: "50-60"},
		{a: "40", b: "80", c: "10", d: "50", intersects: true, overlap: "40-50"},
		{a: "40", b: "80", c: "40", d: "80", intersects: true, overlap: "40-80"},
		{a: "", b: "80", c: "", d: "80", intersects: true, overlap: "-80"},
		{a: "40", b: "", c: "40", d: "", intersects: true, overlap: "40-"},
		{a: "40", b: "80", c: "20", d: "40", intersects: false},
		{a: "80", b: "", c: "80", d: "C0", intersects: true, overlap: "80-C0"},
	}

	for _, el := range table {
		a, err := HexKeyspaceId(el.a).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		b, err := HexKeyspaceId(el.b).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		left := KeyRange{Start: a, End: b}
		c, err := HexKeyspaceId(el.c).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		d, err := HexKeyspaceId(el.d).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		right := KeyRange{Start: c, End: d}
		if c := KeyRangesIntersect(left, right); c != el.intersects {
			t.Errorf("Unexpected result: KeyRangesIntersect for %v and %v yields %v.", left, right, c)
		}
		overlap, err := KeyRangesOverlap(left, right)
		if el.intersects {
			if err != nil {
				t.Errorf("Unexpected result: KeyRangesOverlap for overlapping %v and %v returned an error: %v", left, right, err)
			} else {
				got := string(overlap.Start.Hex()) + "-" + string(overlap.End.Hex())
				if got != el.overlap {
					t.Errorf("Unexpected result: KeyRangesOverlap for overlapping %v and %v should have returned: %v but got: %v", left, right, el.overlap, got)
				}
			}
		} else {
			if err == nil {
				t.Errorf("Unexpected result: KeyRangesOverlap for non-overlapping %v and %v should have returned an error", left, right)
			}
		}
	}
}
<commit_msg>One extra test case.<commit_after>// Copyright 2012, 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 key

import (
	"encoding/json"
	"testing"
)

func TestKey(t *testing.T) {
	k0 := Uint64Key(0).KeyspaceId()
	k1 := Uint64Key(1).KeyspaceId()
	k2 := Uint64Key(0x7FFFFFFFFFFFFFFF).KeyspaceId()
	k3 := Uint64Key(0x8000000000000000).KeyspaceId()
	k4 := Uint64Key(0xFFFFFFFFFFFFFFFF).KeyspaceId()

	f := func(k KeyspaceId, x string) {
		if x != string(k) {
			t.Errorf("byte mismatch %#v != %#v", k, x)
		}
		data, err := json.MarshalIndent(k, "  ", "  ")
		if err != nil {
			t.Errorf("serialize error: %v", err)
		} else {
			t.Logf("json: %v", string(data))
		}

		k_r := new(KeyspaceId)
		err = json.Unmarshal(data, k_r)
		if err != nil {
			t.Errorf("reserialize error: %v", err)
		}

		if k != *k_r {
			t.Errorf("keyspace compare failed: %#v != %#v", k, k_r)
		}
	}

	f(MinKey, "")
	f(k0, "\x00\x00\x00\x00\x00\x00\x00\x00")
	f(k1, "\x00\x00\x00\x00\x00\x00\x00\x01")
	f(k2, "\x7f\xff\xff\xff\xff\xff\xff\xff")
	f(k3, "\x80\x00\x00\x00\x00\x00\x00\x00")
	f(k4, "\xff\xff\xff\xff\xff\xff\xff\xff")

	hv := k4.Hex()
	if hv != "FFFFFFFFFFFFFFFF" {
		t.Errorf("Was expecting FFFFFFFFFFFFFFFF but got %v", hv)
	}
}

func TestKeyUint64Sort(t *testing.T) {
	k0 := Uint64Key(0).KeyspaceId()
	k1 := Uint64Key(1).KeyspaceId()
	k2 := Uint64Key(0x7FFFFFFFFFFFFFFF).KeyspaceId()
	k3 := Uint64Key(0x8000000000000000).KeyspaceId()
	k4 := Uint64Key(0xFFFFFFFFFFFFFFFF).KeyspaceId()
	kl := make([]KeyspaceId, 0, 16)
	klSorted := make([]KeyspaceId, 0, 16)
	kl = append(kl, MinKey, MaxKey, k4, k3, k2, k1, k0)
	klSorted = append(kl, MinKey, k0, k1, k2, k3, k4, MaxKey)
	KeyspaceIdArray(kl).Sort()

	for i, k := range kl {
		if k != klSorted[i] {
			t.Errorf("key order error: %d %v %v", i, k, klSorted[i])
		}
	}
}

func TestKeyStringSort(t *testing.T) {
	k0 := KeyspaceId("0")
	k1 := KeyspaceId("9")
	k2 := KeyspaceId("Zzzz")
	k3 := KeyspaceId("a")
	k4 := KeyspaceId("z")
	kl := make([]KeyspaceId, 0, 16)
	klSorted := make([]KeyspaceId, 0, 16)
	kl = append(kl, MinKey, MaxKey, k4, k3, k2, k1, k0)
	klSorted = append(kl, MinKey, k0, k1, k2, k3, k4, MaxKey)
	KeyspaceIdArray(kl).Sort()

	for i, k := range kl {
		if k != klSorted[i] {
			t.Errorf("key order error: %d %v %v", i, k, klSorted[i])
		}
	}
}

func TestParseShardingSpec(t *testing.T) {
	x40, err := HexKeyspaceId("4000000000000000").Unhex()
	if err != nil {
		t.Errorf("Unexpected error: %v.", err)
	}
	x80, err := HexKeyspaceId("8000000000000000").Unhex()
	if err != nil {
		t.Errorf("Unexpected error: %v.", err)
	}

	goodTable := map[string][]KeyRange{
		"-": {{Start: MinKey, End: MaxKey}},
		"-4000000000000000-8000000000000000-": {
			{Start: MinKey, End: x40},
			{Start: x40, End: x80},
			{Start: x80, End: MaxKey},
		},
	}
	badTable := []string{
		"4000000000000000",
		"---",
		"4000000000000000--8000000000000000",
		"4000000000000000-3000000000000000", // not in order
	}
	for key, wanted := range goodTable {
		r, err := ParseShardingSpec(key)
		if err != nil {
			t.Errorf("Unexpected error: %v.", err)
		}
		if len(r) != len(wanted) {
			t.Errorf("Wrong result: wanted %v, got %v", wanted, r)
			continue
		}
		for i, w := range wanted {
			if r[i] != w {
				t.Errorf("Wrong result: wanted %v, got %v", wanted, r)
				break
			}
		}
	}
	for _, bad := range badTable {
		_, err := ParseShardingSpec(bad)
		if err == nil {
			t.Errorf("Didn't get expected error for %v.", bad)
		}
	}
}

func TestContains(t *testing.T) {
	var table = []struct {
		kid       string
		start     string
		end       string
		contained bool
	}{
		{kid: "3000000000000000", start: "3000000000000000", end: "", contained: true},
		{kid: "3000000000000000", start: "", end: "3000000000000000", contained: false},
		{kid: "4000000000000000", start: "3000000000000000", end: "", contained: true},
		{kid: "2000000000000000", start: "3000000000000000", end: "", contained: false},
	}

	for _, el := range table {
		s, err := HexKeyspaceId(el.start).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		e, err := HexKeyspaceId(el.end).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		kr := KeyRange{Start: s, End: e}
		k, err := HexKeyspaceId(el.kid).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		if c := kr.Contains(k); c != el.contained {
			t.Errorf("Unexpected result: contains for %v and (%v-%v) yields %v.", el.kid, el.start, el.end, c)
		}
	}
}

func TestIntersectOverlap(t *testing.T) {
	var table = []struct {
		a          string
		b          string
		c          string
		d          string
		intersects bool
		overlap    string
	}{
		{a: "40", b: "80", c: "c0", d: "d0", intersects: false},
		{a: "", b: "80", c: "80", d: "", intersects: false},
		{a: "", b: "80", c: "", d: "40", intersects: true, overlap: "-40"},
		{a: "80", b: "", c: "C0", d: "", intersects: true, overlap: "C0-"},
		{a: "", b: "80", c: "40", d: "80", intersects: true, overlap: "40-80"},
		{a: "40", b: "80", c: "60", d: "A0", intersects: true, overlap: "60-80"},
		{a: "40", b: "80", c: "50", d: "60", intersects: true, overlap: "50-60"},
		{a: "40", b: "80", c: "10", d: "50", intersects: true, overlap: "40-50"},
		{a: "40", b: "80", c: "40", d: "80", intersects: true, overlap: "40-80"},
		{a: "", b: "80", c: "", d: "80", intersects: true, overlap: "-80"},
		{a: "40", b: "", c: "40", d: "", intersects: true, overlap: "40-"},
		{a: "40", b: "80", c: "20", d: "40", intersects: false},
		{a: "80", b: "", c: "80", d: "C0", intersects: true, overlap: "80-C0"},
		{a: "", b: "", c: "C0", d: "D0", intersects: true, overlap: "C0-D0"},
	}

	for _, el := range table {
		a, err := HexKeyspaceId(el.a).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		b, err := HexKeyspaceId(el.b).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		left := KeyRange{Start: a, End: b}
		c, err := HexKeyspaceId(el.c).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		d, err := HexKeyspaceId(el.d).Unhex()
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
		}
		right := KeyRange{Start: c, End: d}
		if c := KeyRangesIntersect(left, right); c != el.intersects {
			t.Errorf("Unexpected result: KeyRangesIntersect for %v and %v yields %v.", left, right, c)
		}
		overlap, err := KeyRangesOverlap(left, right)
		if el.intersects {
			if err != nil {
				t.Errorf("Unexpected result: KeyRangesOverlap for overlapping %v and %v returned an error: %v", left, right, err)
			} else {
				got := string(overlap.Start.Hex()) + "-" + string(overlap.End.Hex())
				if got != el.overlap {
					t.Errorf("Unexpected result: KeyRangesOverlap for overlapping %v and %v should have returned: %v but got: %v", left, right, el.overlap, got)
				}
			}
		} else {
			if err == nil {
				t.Errorf("Unexpected result: KeyRangesOverlap for non-overlapping %v and %v should have returned an error", left, right)
			}
		}
	}
}
<|endoftext|>


<commit_before>package gopush

import (
	"github.com/appleboy/gofight"
	"github.com/buger/jsonparser"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"net/http"
	"os"
	"runtime"
	"testing"
	"time"
)

var go_version = runtime.Version()

func initTest() {
	PushConf = BuildDefaultPushConf()
	PushConf.Core.Mode = "test"
}

func TestPrintGoPushVersion(t *testing.T) {
	PrintGoPushVersion()
}

func TestRunNormalServer(t *testing.T) {
	initTest()

	gin.SetMode(gin.TestMode)

	router := gin.New()

	go func() {
		assert.NoError(t, RunHTTPServer())
	}()
	// have to wait for the goroutine to start and run the server
	// otherwise the main thread will complete
	time.Sleep(5 * time.Millisecond)

	assert.Error(t, router.Run(":8088"))
	gofight.TestRequest(t, "http://localhost:8088/api/status")
}

// func TestRunTLSServer(t *testing.T) {
// 	initTest()

// 	PushConf.Core.SSL = true
// 	PushConf.Core.Port = "8087"
// 	PushConf.Core.CertPath = "../certificate/localhost.cert"
// 	PushConf.Core.KeyPath = "../certificate/localhost.key"
// 	router := gin.New()

// 	go func() {
// 		assert.NoError(t, RunHTTPServer())
// 	}()
// 	// have to wait for the goroutine to start and run the server
// 	// otherwise the main thread will complete
// 	time.Sleep(5 * time.Millisecond)

// 	assert.Error(t, router.Run(":8087"))
// 	testRequest(t, "https://localhost:8087/api/status")
// }

func TestRootHandler(t *testing.T) {
	initTest()

	r := gofight.New()

	r.GET("/").
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {
			data := []byte(r.Body.String())

			value, _ := jsonparser.GetString(data, "text")

			assert.Equal(t, "Welcome to notification server.", value)
			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestAPIStatusHandler(t *testing.T) {
	initTest()

	r := gofight.New()

	r.GET("/api/status").
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {
			data := []byte(r.Body.String())

			value, _ := jsonparser.GetString(data, "go_version")

			assert.Equal(t, go_version, value)
			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestMissingParameterPushHandler(t *testing.T) {
	initTest()

	r := gofight.New()

	// missing some parameter.
	r.POST("/api/push").
		SetJSON(gofight.D{
			"platform": 1,
		}).
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {

			assert.Equal(t, http.StatusBadRequest, r.Code)
		})
}

func TestDisabledIosPushHandler(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = false
	InitAPNSClient()

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"11aa01229f15f0f0c52029d8cf8cd0aeaf2365fe4cebc4af26cd6d76b7919ef7"},
			"platform": 1,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestMissingIosCertificate(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = true
	PushConf.Ios.PemKeyPath = "test"
	err := InitAPNSClient()

	assert.Error(t, err)
}

func TestIosPushDevelopment(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = true
	PushConf.Ios.PemKeyPath = "../certificate/certificate-valid.pem"
	InitAPNSClient()

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"11aa01229f15f0f0c52029d8cf8cd0aeaf2365fe4cebc4af26cd6d76b7919ef7"},
			"platform": 1,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestIosPushProduction(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = true
	PushConf.Ios.Production = true
	PushConf.Ios.PemKeyPath = "../certificate/certificate-valid.pem"
	InitAPNSClient()

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"11aa01229f15f0f0c52029d8cf8cd0aeaf2365fe4cebc4af26cd6d76b7919ef7"},
			"platform": 1,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestDisabledAndroidPushHandler(t *testing.T) {
	initTest()

	PushConf.Android.Enabled = false

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"aaaaaa", "bbbbb"},
			"platform": 2,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestAndroidPushHandler(t *testing.T) {
	initTest()

	PushConf.Android.Enabled = true
	PushConf.Android.ApiKey = os.Getenv("ANDROID_API_KEY")

	android_token := os.Getenv("ANDROID_TEST_TOKEN")

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{android_token, "bbbbb"},
			"platform": 2,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HttpResponse, rq gofight.HttpRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}
<commit_msg>update server testing.<commit_after>package gopush

import (
	"github.com/appleboy/gofight"
	"github.com/buger/jsonparser"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"net/http"
	"os"
	"runtime"
	"testing"
	"time"
)

var go_version = runtime.Version()

func initTest() {
	PushConf = BuildDefaultPushConf()
	PushConf.Core.Mode = "test"
}

func TestPrintGoPushVersion(t *testing.T) {
	PrintGoPushVersion()
}

func TestRunNormalServer(t *testing.T) {
	initTest()

	gin.SetMode(gin.TestMode)

	router := gin.New()

	go func() {
		assert.NoError(t, RunHTTPServer())
	}()
	// have to wait for the goroutine to start and run the server
	// otherwise the main thread will complete
	time.Sleep(5 * time.Millisecond)

	assert.Error(t, router.Run(":8088"))
	gofight.TestRequest(t, "http://localhost:8088/api/status")
}

// func TestRunTLSServer(t *testing.T) {
// 	initTest()

// 	PushConf.Core.SSL = true
// 	PushConf.Core.Port = "8087"
// 	PushConf.Core.CertPath = "../certificate/localhost.cert"
// 	PushConf.Core.KeyPath = "../certificate/localhost.key"
// 	router := gin.New()

// 	go func() {
// 		assert.NoError(t, RunHTTPServer())
// 	}()
// 	// have to wait for the goroutine to start and run the server
// 	// otherwise the main thread will complete
// 	time.Sleep(5 * time.Millisecond)

// 	assert.Error(t, router.Run(":8087"))
// 	testRequest(t, "https://localhost:8087/api/status")
// }

func TestRootHandler(t *testing.T) {
	initTest()

	r := gofight.New()

	r.GET("/").
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
			data := []byte(r.Body.String())

			value, _ := jsonparser.GetString(data, "text")

			assert.Equal(t, "Welcome to notification server.", value)
			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestAPIStatusHandler(t *testing.T) {
	initTest()

	r := gofight.New()

	r.GET("/api/status").
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
			data := []byte(r.Body.String())

			value, _ := jsonparser.GetString(data, "go_version")

			assert.Equal(t, go_version, value)
			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestMissingParameterPushHandler(t *testing.T) {
	initTest()

	r := gofight.New()

	// missing some parameter.
	r.POST("/api/push").
		SetJSON(gofight.D{
			"platform": 1,
		}).
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {

			assert.Equal(t, http.StatusBadRequest, r.Code)
		})
}

func TestDisabledIosPushHandler(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = false
	InitAPNSClient()

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"11aa01229f15f0f0c52029d8cf8cd0aeaf2365fe4cebc4af26cd6d76b7919ef7"},
			"platform": 1,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestMissingIosCertificate(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = true
	PushConf.Ios.PemKeyPath = "test"
	err := InitAPNSClient()

	assert.Error(t, err)
}

func TestIosPushDevelopment(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = true
	PushConf.Ios.PemKeyPath = "../certificate/certificate-valid.pem"
	InitAPNSClient()

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"11aa01229f15f0f0c52029d8cf8cd0aeaf2365fe4cebc4af26cd6d76b7919ef7"},
			"platform": 1,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestIosPushProduction(t *testing.T) {
	initTest()

	PushConf.Ios.Enabled = true
	PushConf.Ios.Production = true
	PushConf.Ios.PemKeyPath = "../certificate/certificate-valid.pem"
	InitAPNSClient()

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"11aa01229f15f0f0c52029d8cf8cd0aeaf2365fe4cebc4af26cd6d76b7919ef7"},
			"platform": 1,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestDisabledAndroidPushHandler(t *testing.T) {
	initTest()

	PushConf.Android.Enabled = false

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{"aaaaaa", "bbbbb"},
			"platform": 2,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}

func TestAndroidPushHandler(t *testing.T) {
	initTest()

	PushConf.Android.Enabled = true
	PushConf.Android.ApiKey = os.Getenv("ANDROID_API_KEY")

	android_token := os.Getenv("ANDROID_TEST_TOKEN")

	r := gofight.New()

	r.POST("/api/push").
		SetJSON(gofight.D{
			"tokens":   []string{android_token, "bbbbb"},
			"platform": 2,
			"message":  "Welcome",
		}).
		Run(GetMainEngine(), func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {

			assert.Equal(t, http.StatusOK, r.Code)
		})
}
<|endoftext|>