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

package w32

import (
	"errors"
	"fmt"
	"syscall"
	"unsafe"
)

var (
	modadvapi32 = syscall.NewLazyDLL("advapi32.dll")

	procRegCreateKeyEx = modadvapi32.NewProc("RegCreateKeyExW")
	procRegOpenKeyEx   = modadvapi32.NewProc("RegOpenKeyExW")
	procRegCloseKey    = modadvapi32.NewProc("RegCloseKey")
	procRegGetValue    = modadvapi32.NewProc("RegGetValueW")
	procRegEnumKeyEx   = modadvapi32.NewProc("RegEnumKeyExW")
	//	procRegSetKeyValue     = modadvapi32.NewProc("RegSetKeyValueW")
	procRegSetValueEx      = modadvapi32.NewProc("RegSetValueExW")
	procOpenEventLog       = modadvapi32.NewProc("OpenEventLogW")
	procReadEventLog       = modadvapi32.NewProc("ReadEventLogW")
	procCloseEventLog      = modadvapi32.NewProc("CloseEventLog")
	procOpenSCManager      = modadvapi32.NewProc("OpenSCManagerW")
	procCloseServiceHandle = modadvapi32.NewProc("CloseServiceHandle")
	procOpenService        = modadvapi32.NewProc("OpenServiceW")
	procStartService       = modadvapi32.NewProc("StartServiceW")
	procControlService     = modadvapi32.NewProc("ControlService")
)

func RegCreateKey(hKey HKEY, subKey string) HKEY {
	var result HKEY
	ret, _, _ := procRegCreateKeyEx.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(0),
		uintptr(0),
		uintptr(0),
		uintptr(KEY_ALL_ACCESS),
		uintptr(0),
		uintptr(unsafe.Pointer(&result)),
		uintptr(0))
	_ = ret
	return result
}

func RegOpenKeyEx(hKey HKEY, subKey string, samDesired uint32) HKEY {
	var result HKEY
	ret, _, _ := procRegOpenKeyEx.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(0),
		uintptr(samDesired),
		uintptr(unsafe.Pointer(&result)))

	if ret != ERROR_SUCCESS {
		panic(fmt.Sprintf("RegOpenKeyEx(%d, %s, %d) failed", hKey, subKey, samDesired))
	}
	return result
}

func RegCloseKey(hKey HKEY) error {
	var err error
	ret, _, _ := procRegCloseKey.Call(
		uintptr(hKey))

	if ret != ERROR_SUCCESS {
		err = errors.New("RegCloseKey failed")
	}
	return err
}

func RegGetRaw(hKey HKEY, subKey string, value string) []byte {
	var bufLen uint32
	var valptr unsafe.Pointer
	if len(value) > 0 {
		valptr = unsafe.Pointer(syscall.StringToUTF16Ptr(value))
	}
	procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(valptr),
		uintptr(RRF_RT_ANY),
		0,
		0,
		uintptr(unsafe.Pointer(&bufLen)))

	if bufLen == 0 {
		return nil
	}

	buf := make([]byte, bufLen)
	ret, _, _ := procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(valptr),
		uintptr(RRF_RT_ANY),
		0,
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&bufLen)))

	if ret != ERROR_SUCCESS {
		return nil
	}

	return buf
}

func RegSetBinary(hKey HKEY, subKey string, value []byte) (errno int) {
	var lptr, vptr unsafe.Pointer
	if len(subKey) > 0 {
		lptr = unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))
	}
	if len(value) > 0 {
		vptr = unsafe.Pointer(&value[0])
	}
	ret, _, _ := procRegSetValueEx.Call(
		uintptr(hKey),
		uintptr(lptr),
		uintptr(0),
		uintptr(REG_BINARY),
		uintptr(vptr),
		uintptr(len(value)))

	return int(ret)
}

func RegSetString(hKey HKEY, subKey string, value string) (errno int) {
	var lptr, vptr unsafe.Pointer
	if len(subKey) > 0 {
		lptr = unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))
	}
	var buf []uint16
	if len(value) > 0 {
		buf, err := syscall.UTF16FromString(value)
		if err != nil {
			return ERROR_BAD_FORMAT
		}
		vptr = unsafe.Pointer(&buf[0])
	}
	ret, _, _ := procRegSetValueEx.Call(
		uintptr(hKey),
		uintptr(lptr),
		uintptr(0),
		uintptr(REG_SZ),
		uintptr(vptr),
		uintptr(unsafe.Sizeof(buf) + 2)) // 2 is the size of the terminating null character

	return int(ret)
}

func RegSetUint32(hKey HKEY, subKey string, value uint32) (errno int) {
	var lptr unsafe.Pointer
	if len(subKey) > 0 {
		lptr = unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))
	}
	vptr := unsafe.Pointer(&value)
	ret, _, _ := procRegSetValueEx.Call(
		uintptr(hKey),
		uintptr(lptr),
		uintptr(0),
		uintptr(REG_DWORD),
		uintptr(vptr),
		uintptr(unsafe.Sizeof(value)))

	return int(ret)
}

func RegGetString(hKey HKEY, subKey string, value string) string {
	var bufLen uint32
	procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(value))),
		uintptr(RRF_RT_REG_SZ),
		0,
		0,
		uintptr(unsafe.Pointer(&bufLen)))

	if bufLen == 0 {
		return ""
	}

	buf := make([]uint16, bufLen)
	ret, _, _ := procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(value))),
		uintptr(RRF_RT_REG_SZ),
		0,
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&bufLen)))

	if ret != ERROR_SUCCESS {
		return ""
	}

	return syscall.UTF16ToString(buf)
}

func RegGetUint32(hKey HKEY, subKey string, value string) (data uint32, errno int) {
	var dataLen uint32 = uint32(unsafe.Sizeof(data))
	ret, _, _ := procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(value))),
		uintptr(RRF_RT_REG_DWORD),
		0,
		uintptr(unsafe.Pointer(&data)),
		uintptr(unsafe.Pointer(&dataLen)))
	errno = int(ret)
	return
}

/*
func RegSetKeyValue(hKey HKEY, subKey string, valueName string, dwType uint32, data uintptr, cbData uint16) (errno int) {
	ret, _, _ := procRegSetKeyValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(valueName))),
		uintptr(dwType),
		data,
		uintptr(cbData))

	return int(ret)
}
*/

func RegEnumKeyEx(hKey HKEY, index uint32) string {
	var bufLen uint32 = 255
	buf := make([]uint16, bufLen)
	procRegEnumKeyEx.Call(
		uintptr(hKey),
		uintptr(index),
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&bufLen)),
		0,
		0,
		0,
		0)
	return syscall.UTF16ToString(buf)
}

func OpenEventLog(servername string, sourcename string) HANDLE {
	ret, _, _ := procOpenEventLog.Call(
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(servername))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(sourcename))))

	return HANDLE(ret)
}

func ReadEventLog(eventlog HANDLE, readflags, recordoffset uint32, buffer []byte, numberofbytestoread uint32, bytesread, minnumberofbytesneeded *uint32) bool {
	ret, _, _ := procReadEventLog.Call(
		uintptr(eventlog),
		uintptr(readflags),
		uintptr(recordoffset),
		uintptr(unsafe.Pointer(&buffer[0])),
		uintptr(numberofbytestoread),
		uintptr(unsafe.Pointer(bytesread)),
		uintptr(unsafe.Pointer(minnumberofbytesneeded)))

	return ret != 0
}

func CloseEventLog(eventlog HANDLE) bool {
	ret, _, _ := procCloseEventLog.Call(
		uintptr(eventlog))

	return ret != 0
}

func OpenSCManager(lpMachineName, lpDatabaseName string, dwDesiredAccess uint32) (HANDLE, error) {
	var p1, p2 uintptr
	if len(lpMachineName) > 0 {
		p1 = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpMachineName)))
	}
	if len(lpDatabaseName) > 0 {
		p2 = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpDatabaseName)))
	}
	ret, _, _ := procOpenSCManager.Call(
		p1,
		p2,
		uintptr(dwDesiredAccess))

	if ret == 0 {
		return 0, syscall.GetLastError()
	}

	return HANDLE(ret), nil
}

func CloseServiceHandle(hSCObject HANDLE) error {
	ret, _, _ := procCloseServiceHandle.Call(uintptr(hSCObject))
	if ret == 0 {
		return syscall.GetLastError()
	}
	return nil
}

func OpenService(hSCManager HANDLE, lpServiceName string, dwDesiredAccess uint32) (HANDLE, error) {
	ret, _, _ := procOpenService.Call(
		uintptr(hSCManager),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpServiceName))),
		uintptr(dwDesiredAccess))

	if ret == 0 {
		return 0, syscall.GetLastError()
	}

	return HANDLE(ret), nil
}

func StartService(hService HANDLE, lpServiceArgVectors []string) error {
	l := len(lpServiceArgVectors)
	var ret uintptr
	if l == 0 {
		ret, _, _ = procStartService.Call(
			uintptr(hService),
			0,
			0)
	} else {
		lpArgs := make([]uintptr, l)
		for i := 0; i < l; i++ {
			lpArgs[i] = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpServiceArgVectors[i])))
		}

		ret, _, _ = procStartService.Call(
			uintptr(hService),
			uintptr(l),
			uintptr(unsafe.Pointer(&lpArgs[0])))
	}

	if ret == 0 {
		return syscall.GetLastError()
	}

	return nil
}

func ControlService(hService HANDLE, dwControl uint32, lpServiceStatus *SERVICE_STATUS) bool {
	if lpServiceStatus == nil {
		panic("ControlService:lpServiceStatus cannot be nil")
	}

	ret, _, _ := procControlService.Call(
		uintptr(hService),
		uintptr(dwControl),
		uintptr(unsafe.Pointer(lpServiceStatus)))

	return ret != 0
}
<commit_msg>Add RegDeleteKeyValue, RegDeleteValue, and RegDeleteTree<commit_after>// Copyright 2010-2012 The W32 Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package w32

import (
	"errors"
	"fmt"
	"syscall"
	"unsafe"
)

var (
	modadvapi32 = syscall.NewLazyDLL("advapi32.dll")

	procRegCreateKeyEx = modadvapi32.NewProc("RegCreateKeyExW")
	procRegOpenKeyEx   = modadvapi32.NewProc("RegOpenKeyExW")
	procRegCloseKey    = modadvapi32.NewProc("RegCloseKey")
	procRegGetValue    = modadvapi32.NewProc("RegGetValueW")
	procRegEnumKeyEx   = modadvapi32.NewProc("RegEnumKeyExW")
	//	procRegSetKeyValue     = modadvapi32.NewProc("RegSetKeyValueW")
	procRegSetValueEx      = modadvapi32.NewProc("RegSetValueExW")
	procRegDeleteKeyValue  = modadvapi32.NewProc("RegDeleteKeyValueW")
	procRegDeleteValue     = modadvapi32.NewProc("RegDeleteValueW")
	procRegDeleteTree      = modadvapi32.NewProc("RegDeleteTreeW")
	procOpenEventLog       = modadvapi32.NewProc("OpenEventLogW")
	procReadEventLog       = modadvapi32.NewProc("ReadEventLogW")
	procCloseEventLog      = modadvapi32.NewProc("CloseEventLog")
	procOpenSCManager      = modadvapi32.NewProc("OpenSCManagerW")
	procCloseServiceHandle = modadvapi32.NewProc("CloseServiceHandle")
	procOpenService        = modadvapi32.NewProc("OpenServiceW")
	procStartService       = modadvapi32.NewProc("StartServiceW")
	procControlService     = modadvapi32.NewProc("ControlService")
)

func RegCreateKey(hKey HKEY, subKey string) HKEY {
	var result HKEY
	ret, _, _ := procRegCreateKeyEx.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(0),
		uintptr(0),
		uintptr(0),
		uintptr(KEY_ALL_ACCESS),
		uintptr(0),
		uintptr(unsafe.Pointer(&result)),
		uintptr(0))
	_ = ret
	return result
}

func RegOpenKeyEx(hKey HKEY, subKey string, samDesired uint32) HKEY {
	var result HKEY
	ret, _, _ := procRegOpenKeyEx.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(0),
		uintptr(samDesired),
		uintptr(unsafe.Pointer(&result)))

	if ret != ERROR_SUCCESS {
		panic(fmt.Sprintf("RegOpenKeyEx(%d, %s, %d) failed", hKey, subKey, samDesired))
	}
	return result
}

func RegCloseKey(hKey HKEY) error {
	var err error
	ret, _, _ := procRegCloseKey.Call(
		uintptr(hKey))

	if ret != ERROR_SUCCESS {
		err = errors.New("RegCloseKey failed")
	}
	return err
}

func RegGetRaw(hKey HKEY, subKey string, value string) []byte {
	var bufLen uint32
	var valptr unsafe.Pointer
	if len(value) > 0 {
		valptr = unsafe.Pointer(syscall.StringToUTF16Ptr(value))
	}
	procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(valptr),
		uintptr(RRF_RT_ANY),
		0,
		0,
		uintptr(unsafe.Pointer(&bufLen)))

	if bufLen == 0 {
		return nil
	}

	buf := make([]byte, bufLen)
	ret, _, _ := procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(valptr),
		uintptr(RRF_RT_ANY),
		0,
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&bufLen)))

	if ret != ERROR_SUCCESS {
		return nil
	}

	return buf
}

func RegSetBinary(hKey HKEY, subKey string, value []byte) (errno int) {
	var lptr, vptr unsafe.Pointer
	if len(subKey) > 0 {
		lptr = unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))
	}
	if len(value) > 0 {
		vptr = unsafe.Pointer(&value[0])
	}
	ret, _, _ := procRegSetValueEx.Call(
		uintptr(hKey),
		uintptr(lptr),
		uintptr(0),
		uintptr(REG_BINARY),
		uintptr(vptr),
		uintptr(len(value)))

	return int(ret)
}

func RegSetString(hKey HKEY, subKey string, value string) (errno int) {
	var lptr, vptr unsafe.Pointer
	if len(subKey) > 0 {
		lptr = unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))
	}
	var buf []uint16
	if len(value) > 0 {
		buf, err := syscall.UTF16FromString(value)
		if err != nil {
			return ERROR_BAD_FORMAT
		}
		vptr = unsafe.Pointer(&buf[0])
	}
	ret, _, _ := procRegSetValueEx.Call(
		uintptr(hKey),
		uintptr(lptr),
		uintptr(0),
		uintptr(REG_SZ),
		uintptr(vptr),
		uintptr(unsafe.Sizeof(buf) + 2)) // 2 is the size of the terminating null character

	return int(ret)
}

func RegSetUint32(hKey HKEY, subKey string, value uint32) (errno int) {
	var lptr unsafe.Pointer
	if len(subKey) > 0 {
		lptr = unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))
	}
	vptr := unsafe.Pointer(&value)
	ret, _, _ := procRegSetValueEx.Call(
		uintptr(hKey),
		uintptr(lptr),
		uintptr(0),
		uintptr(REG_DWORD),
		uintptr(vptr),
		uintptr(unsafe.Sizeof(value)))

	return int(ret)
}

func RegGetString(hKey HKEY, subKey string, value string) string {
	var bufLen uint32
	procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(value))),
		uintptr(RRF_RT_REG_SZ),
		0,
		0,
		uintptr(unsafe.Pointer(&bufLen)))

	if bufLen == 0 {
		return ""
	}

	buf := make([]uint16, bufLen)
	ret, _, _ := procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(value))),
		uintptr(RRF_RT_REG_SZ),
		0,
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&bufLen)))

	if ret != ERROR_SUCCESS {
		return ""
	}

	return syscall.UTF16ToString(buf)
}

func RegGetUint32(hKey HKEY, subKey string, value string) (data uint32, errno int) {
	var dataLen uint32 = uint32(unsafe.Sizeof(data))
	ret, _, _ := procRegGetValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(value))),
		uintptr(RRF_RT_REG_DWORD),
		0,
		uintptr(unsafe.Pointer(&data)),
		uintptr(unsafe.Pointer(&dataLen)))
	errno = int(ret)
	return
}

/*
func RegSetKeyValue(hKey HKEY, subKey string, valueName string, dwType uint32, data uintptr, cbData uint16) (errno int) {
	ret, _, _ := procRegSetKeyValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(valueName))),
		uintptr(dwType),
		data,
		uintptr(cbData))

	return int(ret)
}
*/

func RegDeleteKeyValue(hKey HKEY, subKey string, valueName string) (errno int) {
	ret, _, _ := procRegDeleteKeyValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(valueName))))

	return int(ret)
}

func RegDeleteValue(hKey HKEY, valueName string) (errno int) {
	ret, _, _ := procRegDeleteValue.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(valueName))))

	return int(ret)
}

func RegDeleteTree(hKey HKEY, subKey string) (errno int) {
	ret, _, _ := procRegDeleteTree.Call(
		uintptr(hKey),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(subKey))))

	return int(ret)
}

func RegEnumKeyEx(hKey HKEY, index uint32) string {
	var bufLen uint32 = 255
	buf := make([]uint16, bufLen)
	procRegEnumKeyEx.Call(
		uintptr(hKey),
		uintptr(index),
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&bufLen)),
		0,
		0,
		0,
		0)
	return syscall.UTF16ToString(buf)
}

func OpenEventLog(servername string, sourcename string) HANDLE {
	ret, _, _ := procOpenEventLog.Call(
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(servername))),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(sourcename))))

	return HANDLE(ret)
}

func ReadEventLog(eventlog HANDLE, readflags, recordoffset uint32, buffer []byte, numberofbytestoread uint32, bytesread, minnumberofbytesneeded *uint32) bool {
	ret, _, _ := procReadEventLog.Call(
		uintptr(eventlog),
		uintptr(readflags),
		uintptr(recordoffset),
		uintptr(unsafe.Pointer(&buffer[0])),
		uintptr(numberofbytestoread),
		uintptr(unsafe.Pointer(bytesread)),
		uintptr(unsafe.Pointer(minnumberofbytesneeded)))

	return ret != 0
}

func CloseEventLog(eventlog HANDLE) bool {
	ret, _, _ := procCloseEventLog.Call(
		uintptr(eventlog))

	return ret != 0
}

func OpenSCManager(lpMachineName, lpDatabaseName string, dwDesiredAccess uint32) (HANDLE, error) {
	var p1, p2 uintptr
	if len(lpMachineName) > 0 {
		p1 = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpMachineName)))
	}
	if len(lpDatabaseName) > 0 {
		p2 = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpDatabaseName)))
	}
	ret, _, _ := procOpenSCManager.Call(
		p1,
		p2,
		uintptr(dwDesiredAccess))

	if ret == 0 {
		return 0, syscall.GetLastError()
	}

	return HANDLE(ret), nil
}

func CloseServiceHandle(hSCObject HANDLE) error {
	ret, _, _ := procCloseServiceHandle.Call(uintptr(hSCObject))
	if ret == 0 {
		return syscall.GetLastError()
	}
	return nil
}

func OpenService(hSCManager HANDLE, lpServiceName string, dwDesiredAccess uint32) (HANDLE, error) {
	ret, _, _ := procOpenService.Call(
		uintptr(hSCManager),
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpServiceName))),
		uintptr(dwDesiredAccess))

	if ret == 0 {
		return 0, syscall.GetLastError()
	}

	return HANDLE(ret), nil
}

func StartService(hService HANDLE, lpServiceArgVectors []string) error {
	l := len(lpServiceArgVectors)
	var ret uintptr
	if l == 0 {
		ret, _, _ = procStartService.Call(
			uintptr(hService),
			0,
			0)
	} else {
		lpArgs := make([]uintptr, l)
		for i := 0; i < l; i++ {
			lpArgs[i] = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(lpServiceArgVectors[i])))
		}

		ret, _, _ = procStartService.Call(
			uintptr(hService),
			uintptr(l),
			uintptr(unsafe.Pointer(&lpArgs[0])))
	}

	if ret == 0 {
		return syscall.GetLastError()
	}

	return nil
}

func ControlService(hService HANDLE, dwControl uint32, lpServiceStatus *SERVICE_STATUS) bool {
	if lpServiceStatus == nil {
		panic("ControlService:lpServiceStatus cannot be nil")
	}

	ret, _, _ := procControlService.Call(
		uintptr(hService),
		uintptr(dwControl),
		uintptr(unsafe.Pointer(lpServiceStatus)))

	return ret != 0
}
<|endoftext|>


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

package backend

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"gopkg.in/redis.v3"
)

var (
	errBackendIdxNotFound = errors.New("backend not in backends list")
)

type redisMonitor struct {
	hostID      string
	quit        chan struct{}
	done        chan struct{}
	limiter     chan struct{}
	redisClient *redis.Client
	httpClient  *http.Client
}

func newRedisMonitor(redisClient *redis.Client) (*redisMonitor, error) {
	hostID, err := os.Hostname()
	if err != nil {
		return nil, err
	}
	mon := &redisMonitor{
		hostID:      hostID,
		quit:        make(chan struct{}),
		done:        make(chan struct{}),
		limiter:     make(chan struct{}, 5),
		redisClient: redisClient,
		httpClient: &http.Client{
			Transport: &http.Transport{
				Dial: (&net.Dialer{
					Timeout:   10 * time.Second,
					KeepAlive: 10 * time.Second,
				}).Dial,
				// Disable connections reuse for efective dial timeouts.
				DisableKeepAlives:   true,
				MaxIdleConnsPerHost: -1,
				TLSHandshakeTimeout: 10 * time.Second,
			},
			Timeout: time.Minute,
		},
	}
	err = mon.start()
	if err != nil {
		return nil, err
	}
	return mon, nil
}

func (b *redisMonitor) start() error {
	pubsub, err := b.redisClient.Subscribe("dead")
	if err != nil {
		return err
	}
	go b.loop(pubsub)
	return nil
}

func (b *redisMonitor) loop(pubsub *redis.PubSub) {
	wg := sync.WaitGroup{}
	defer close(b.done)
	defer wg.Wait()
	defer pubsub.Close()
	msgCh := make(chan string)
	for {
		go func() {
			msg, _ := pubsub.ReceiveMessage()
			if msg == nil {
				msgCh <- ""
			} else {
				msgCh <- msg.Payload
			}
		}()
		select {
		case <-b.quit:
			return
		case msg := <-msgCh:
			if msg == "" {
				continue
			}
			wg.Add(1)
			go func(msg string) {
				defer wg.Done()
				b.watch(msg)
			}(msg)
		}
	}
}

func (b *redisMonitor) reserve(host, backend string) bool {
	key := "dead:" + host + ":" + backend
	tx, err := b.redisClient.Watch(key)
	if err != nil {
		return false
	}
	defer tx.Close()
	watchKey := tx.Get(key).Val()
	if watchKey != "" && watchKey != b.hostID {
		return false
	}
	_, err = tx.Exec(func() error {
		tx.Set(key, b.hostID, 30*time.Second)
		return nil
	})
	return err != redis.TxFailedErr
}

func (b *redisMonitor) free(host, backend string) {
	key := "dead:" + host + ":" + backend
	b.redisClient.Del(key)
}

func (b *redisMonitor) watch(msg string) {
	parts := strings.Split(msg, ";")
	if len(parts) != 4 {
		return
	}
	host := parts[0]
	backend := parts[1]
	if !b.reserve(host, backend) {
		return
	}
out:
	for {
		select {
		case <-b.quit:
			break out
		case <-time.After(time.Second):
		}
		b.limiter <- struct{}{}
		if !b.reserve(host, backend) {
			<-b.limiter
			return
		}
		isOk := b.check(host, backend)
		err := b.updateDead(host, backend, isOk)
		<-b.limiter
		if (err == nil && isOk) || err == errBackendIdxNotFound {
			break out
		}
	}
	b.free(host, backend)
}

func (b *redisMonitor) updateDead(host, backend string, isOk bool) error {
	frontend := "frontend:" + host
	tx, err := b.redisClient.Watch(frontend)
	if err != nil {
		return err
	}
	defer tx.Close()
	entries, err := tx.LRange(frontend, 1, -1).Result()
	if err != nil {
		if err == redis.Nil {
			return errBackendIdxNotFound
		}
		return err
	}
	var idx string
	for i := range entries {
		if entries[i] == backend {
			idx = strconv.Itoa(i)
			break
		}
	}
	if idx == "" {
		return errBackendIdxNotFound
	}
	deadKey := "dead:" + host
	_, err = tx.Exec(func() error {
		if isOk {
			tx.SRem(deadKey, idx)
		} else {
			tx.SAdd(deadKey, idx)
			tx.Expire(deadKey, 30*time.Second)
		}
		return nil
	})
	return err
}

type hcData struct {
	path   string
	body   string
	status int
}

func (b *redisMonitor) hcData(host string) (hcData, error) {
	mapData, err := b.redisClient.HGetAllMap("healthcheck:" + host).Result()
	if err != nil && err != redis.Nil {
		return hcData{}, err
	}
	status, _ := strconv.Atoi(mapData["status"])
	return hcData{
		path:   mapData["path"],
		body:   mapData["body"],
		status: status,
	}, nil
}

func (b *redisMonitor) check(host, backend string) bool {
	hcData, err := b.hcData(host)
	if err != nil {
		return false
	}
	url := fmt.Sprintf("%s/%s", backend, strings.TrimLeft(hcData.path, "/"))
	if !strings.HasPrefix(url, "http") {
		url = "http://" + url
	}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return false
	}
	rsp, err := b.httpClient.Do(req)
	if err != nil {
		return false
	}
	defer rsp.Body.Close()
	if hcData.status != 0 && rsp.StatusCode != hcData.status {
		return false
	}
	if hcData.body != "" {
		data, _ := ioutil.ReadAll(rsp.Body)
		return strings.Contains(string(data), hcData.body)
	}
	return true
}

func (b *redisMonitor) stop() {
	if b.quit != nil {
		close(b.quit)
	}
	if b.done != nil {
		<-b.done
	}
}
<commit_msg>backend: active monitor only monitors same backend once<commit_after>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package backend

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"gopkg.in/redis.v3"
)

var (
	errBackendIdxNotFound = errors.New("backend not in backends list")
)

type redisMonitor struct {
	mu          sync.Mutex
	reserved    map[string]struct{}
	hostID      string
	quit        chan struct{}
	done        chan struct{}
	limiter     chan struct{}
	redisClient *redis.Client
	httpClient  *http.Client
}

func newRedisMonitor(redisClient *redis.Client) (*redisMonitor, error) {
	hostID, err := os.Hostname()
	if err != nil {
		return nil, err
	}
	mon := &redisMonitor{
		hostID:      hostID,
		quit:        make(chan struct{}),
		done:        make(chan struct{}),
		limiter:     make(chan struct{}, 5),
		reserved:    make(map[string]struct{}),
		redisClient: redisClient,
		httpClient: &http.Client{
			Transport: &http.Transport{
				Dial: (&net.Dialer{
					Timeout: 10 * time.Second,
				}).Dial,
				// Disable connections reuse for effective dial timeouts.
				DisableKeepAlives:   true,
				MaxIdleConnsPerHost: -1,
				TLSHandshakeTimeout: 10 * time.Second,
			},
			Timeout: 15 * time.Second,
		},
	}
	err = mon.start()
	if err != nil {
		return nil, err
	}
	return mon, nil
}

func (b *redisMonitor) start() error {
	pubsub, err := b.redisClient.Subscribe("dead")
	if err != nil {
		return err
	}
	go b.loop(pubsub)
	return nil
}

func (b *redisMonitor) loop(pubsub *redis.PubSub) {
	wg := sync.WaitGroup{}
	defer close(b.done)
	defer wg.Wait()
	defer pubsub.Close()
	msgCh := make(chan string)
	for {
		go func() {
			msg, _ := pubsub.ReceiveMessage()
			if msg == nil {
				msgCh <- ""
			} else {
				msgCh <- msg.Payload
			}
		}()
		select {
		case <-b.quit:
			return
		case msg := <-msgCh:
			if msg == "" {
				continue
			}
			wg.Add(1)
			go func(msg string) {
				defer wg.Done()
				b.watch(msg)
			}(msg)
		}
	}
}

func (b *redisMonitor) reserve(host, backend string) bool {
	key := "dead:" + host + ":" + backend
	tx, err := b.redisClient.Watch(key)
	if err != nil {
		return false
	}
	defer tx.Close()
	watchKey := tx.Get(key).Val()
	if watchKey != "" && watchKey != b.hostID {
		return false
	}
	_, err = tx.Exec(func() error {
		tx.Set(key, b.hostID, 30*time.Second)
		return nil
	})
	return err != redis.TxFailedErr
}

func (b *redisMonitor) free(host, backend string) {
	key := "dead:" + host + ":" + backend
	b.redisClient.Del(key)
}

func (b *redisMonitor) watch(msg string) {
	parts := strings.Split(msg, ";")
	if len(parts) != 4 {
		return
	}
	host := parts[0]
	backend := parts[1]
	localKey := host + "-" + backend
	b.mu.Lock()
	if _, ok := b.reserved[localKey]; ok {
		b.mu.Unlock()
		return
	}
	b.reserved[localKey] = struct{}{}
	b.mu.Unlock()
	defer func() {
		b.mu.Lock()
		delete(b.reserved, localKey)
		b.mu.Unlock()
	}()
	if !b.reserve(host, backend) {
		return
	}
out:
	for {
		select {
		case <-b.quit:
			break out
		case <-time.After(time.Second):
		}
		b.limiter <- struct{}{}
		if !b.reserve(host, backend) {
			<-b.limiter
			return
		}
		isOk := b.check(host, backend)
		err := b.updateDead(host, backend, isOk)
		<-b.limiter
		if (err == nil && isOk) || err == errBackendIdxNotFound {
			break out
		}
	}
	b.free(host, backend)
}

func (b *redisMonitor) updateDead(host, backend string, isOk bool) error {
	frontend := "frontend:" + host
	tx, err := b.redisClient.Watch(frontend)
	if err != nil {
		return err
	}
	defer tx.Close()
	entries, err := tx.LRange(frontend, 1, -1).Result()
	if err != nil {
		if err == redis.Nil {
			return errBackendIdxNotFound
		}
		return err
	}
	var idx string
	for i := range entries {
		if entries[i] == backend {
			idx = strconv.Itoa(i)
			break
		}
	}
	if idx == "" {
		return errBackendIdxNotFound
	}
	deadKey := "dead:" + host
	_, err = tx.Exec(func() error {
		if isOk {
			tx.SRem(deadKey, idx)
		} else {
			tx.SAdd(deadKey, idx)
			tx.Expire(deadKey, 30*time.Second)
		}
		return nil
	})
	return err
}

type hcData struct {
	path   string
	body   string
	status int
}

func (b *redisMonitor) hcData(host string) (hcData, error) {
	mapData, err := b.redisClient.HGetAllMap("healthcheck:" + host).Result()
	if err != nil && err != redis.Nil {
		return hcData{}, err
	}
	status, _ := strconv.Atoi(mapData["status"])
	return hcData{
		path:   mapData["path"],
		body:   mapData["body"],
		status: status,
	}, nil
}

func (b *redisMonitor) check(host, backend string) bool {
	hcData, err := b.hcData(host)
	if err != nil {
		return false
	}
	url := fmt.Sprintf("%s/%s", backend, strings.TrimLeft(hcData.path, "/"))
	if !strings.HasPrefix(url, "http") {
		url = "http://" + url
	}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return false
	}
	rsp, err := b.httpClient.Do(req)
	if err != nil {
		return false
	}
	defer rsp.Body.Close()
	if hcData.status != 0 && rsp.StatusCode != hcData.status {
		return false
	}
	if hcData.body != "" {
		data, _ := ioutil.ReadAll(rsp.Body)
		return strings.Contains(string(data), hcData.body)
	}
	return true
}

func (b *redisMonitor) stop() {
	if b.quit != nil {
		close(b.quit)
	}
	if b.done != nil {
		<-b.done
	}
}
<|endoftext|>


<commit_before>//  Copyright 2016 The goscope 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 triggers

import (
	"time"

	"github.com/zagrodzki/goscope/scope"
)

// RisingEdge represents the trigger edge type, rising or falling
type RisingEdge bool

// RisingEdge values for readability.
const (
	Rising  RisingEdge = true
	Falling RisingEdge = false
)

// Mode represents the triggering mode, see comments in the constants below.
type Mode int

// Mode values.
const (
	// ModeSingle means trigger once and never again.
	ModeSingle Mode = iota
	// ModeNormal means trigger on every condition, but don't ever trigger
	// without the condition present. Might result in long intervals where
	// data is discarded.
	ModeNormal
	// ModeAuto is like ModeNormal, but will also trigger after some time
	// (currently hardcoded to 0.5s) has passed without the trigger.
	ModeAuto
)

// Trigger represents a filter running on the data channel, waiting for
// a triggering event and then allowing a set of samples equal to the
// configured timebase.
type Trigger struct {
	source  scope.ChanID
	slope   RisingEdge
	lvl     scope.Voltage
	rec     scope.DataRecorder
	tbCount int
	mode    Mode
}

// New returns an initialized Trigger.
func New(rec scope.DataRecorder) *Trigger {
	return &Trigger{
		rec:  rec,
		mode: ModeAuto,
	}
}

// TimeBase returns the trigger timebase, which is the same as the underlying recorder timebase.
func (t *Trigger) TimeBase() scope.Duration {
	return t.rec.TimeBase()
}

// Reset initializes the recording.
func (t *Trigger) Reset(i scope.Duration, ch <-chan []scope.ChannelData) {
	out := make(chan []scope.ChannelData, 20)
	t.tbCount = int(t.rec.TimeBase() / i)
	t.rec.Reset(i, out)
	go t.run(ch, out)
}

// Error passes the error down to the underlying recorder.
func (t *Trigger) Error(err error) {
	t.rec.Error(err)
}

// Source sets the source for the trigger. If received data doesn't contain
// samples for specified source, the trigger allows all samples without filtering.
func (t *Trigger) Source(id scope.ChanID) {
	t.source = id
}

// Edge configures the type of edge (rising/falling) that is the triggering condition.
func (t *Trigger) Edge(e RisingEdge) {
	t.slope = e
}

// Level configures the level that the edge has to cross for the triggering condition.
func (t *Trigger) Level(l scope.Voltage) {
	t.lvl = l
}

// Mode sets the trigger mode.
func (t *Trigger) Mode(m Mode) {
	t.mode = m
}

type thresholdState int

const (
	belowThreshold        thresholdState = -1
	unknownThresholdState thresholdState = 0
	aboveThreshold        thresholdState = 1
)

type slice struct {
	begin int
	end   int
}

func (t *Trigger) run(in <-chan []scope.ChannelData, out chan<- []scope.ChannelData) {
	var left, source int
	var trg, scanned, found bool
	var newState, prevState thresholdState
	var lastTrg time.Time
	for d := range in {
		if !scanned {
			scanned = true
			for i := range d {
				if d[i].ID == t.source {
					source = i
					found = true
					break
				}
			}
		}
		if !found {
			out <- d
			continue
		}
		num := len(d[source].Samples)
		// slices keeps indices of the samples that should be pushed out
		var outSlices []slice
		var curSlice slice
		for i, v := range d[source].Samples {
			switch {
			case v > t.lvl:
				newState = aboveThreshold
			case v < t.lvl:
				newState = belowThreshold
			}
			// if the previous state was uninitialized, do not trigger.
			// Once state is initialized, it's always either above or below, never unknown.
			if newState != prevState && prevState == unknownThresholdState {
				prevState = newState
			}
			// newState > prevState means we moved from below threshold to above threshold, i.e. rising slope.
			if !trg {
				switch {
				// mode single and triggered once already. Don't trigger.
				case t.mode == ModeSingle && !lastTrg.IsZero():
				// crossed the threshold
				case newState != prevState && RisingEdge(newState > prevState) == t.slope:
					trg = true
				// mode auto and time elapsed since last trigger.
				case t.mode == ModeAuto && time.Since(lastTrg) > 500*time.Millisecond:
					trg = true
				}
				if trg {
					lastTrg = time.Now()
					left = t.tbCount
					curSlice.begin = i
				}
			}
			if trg {
				curSlice.end = i + 1
				left--
				if left == 0 {
					outSlices = append(outSlices, curSlice)
					curSlice = slice{}
					trg = false
				}
			}
			prevState = newState
			num--
		}
		if trg {
			outSlices = append(outSlices, curSlice)
		}
		// flush samples
		if len(outSlices) > 0 {
			for _, b := range outSlices {
				chunk := make([]scope.ChannelData, len(d))
				for ch := range d {
					chunk[ch].ID = d[ch].ID
					chunk[ch].Samples = d[ch].Samples[b.begin:b.end]
				}
				out <- chunk
			}
			outSlices = outSlices[:0]
		}
	}
	close(out)
}
<commit_msg>Define an almost-unused default value for edge and mode, intended as an "invalid value" case.<commit_after>//  Copyright 2016 The goscope 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 triggers

import (
	"time"

	"github.com/zagrodzki/goscope/scope"
)

// RisingEdge represents the trigger edge type, rising or falling
type RisingEdge int

const (
	// EdgeNone represents unknown edge type.
	EdgeNone RisingEdge = iota
	// EdgeRising represents a signal crossing from below to above the threshold.
	EdgeRising
	// EdgeFalling represents a signal crossing from above to below the threshold.
	EdgeFalling
)

// Mode represents the triggering mode, see comments in the constants below.
type Mode int

const (
	// ModeNone means unknown mode.
	ModeNone = iota
	// ModeSingle means trigger once and never again.
	ModeSingle
	// ModeNormal means trigger on every condition, but don't ever trigger
	// without the condition present. Might result in long intervals where
	// data is discarded.
	ModeNormal
	// ModeAuto is like ModeNormal, but will also trigger after some time
	// (currently hardcoded to 0.5s) has passed without the trigger.
	ModeAuto
)

// Trigger represents a filter running on the data channel, waiting for
// a triggering event and then allowing a set of samples equal to the
// configured timebase.
type Trigger struct {
	source  scope.ChanID
	slope   RisingEdge
	lvl     scope.Voltage
	rec     scope.DataRecorder
	tbCount int
	mode    Mode
}

// New returns an initialized Trigger.
func New(rec scope.DataRecorder) *Trigger {
	return &Trigger{
		rec:  rec,
		mode: ModeAuto,
	}
}

// TimeBase returns the trigger timebase, which is the same as the underlying recorder timebase.
func (t *Trigger) TimeBase() scope.Duration {
	return t.rec.TimeBase()
}

// Reset initializes the recording.
func (t *Trigger) Reset(i scope.Duration, ch <-chan []scope.ChannelData) {
	out := make(chan []scope.ChannelData, 20)
	t.tbCount = int(t.rec.TimeBase() / i)
	t.rec.Reset(i, out)
	go t.run(ch, out)
}

// Error passes the error down to the underlying recorder.
func (t *Trigger) Error(err error) {
	t.rec.Error(err)
}

// Source sets the source for the trigger. If received data doesn't contain
// samples for specified source, the trigger allows all samples without filtering.
func (t *Trigger) Source(id scope.ChanID) {
	t.source = id
}

// Edge configures the type of edge (rising/falling) that is the triggering condition.
func (t *Trigger) Edge(e RisingEdge) {
	t.slope = e
}

// Level configures the level that the edge has to cross for the triggering condition.
func (t *Trigger) Level(l scope.Voltage) {
	t.lvl = l
}

// Mode sets the trigger mode.
func (t *Trigger) Mode(m Mode) {
	t.mode = m
}

type thresholdState int

const (
	belowThreshold        thresholdState = -1
	unknownThresholdState thresholdState = 0
	aboveThreshold        thresholdState = 1
)

func edgeType(prevState, newState thresholdState) RisingEdge {
	switch {
	case prevState == newState:
		return EdgeNone
	case prevState < newState:
		return EdgeRising
	case prevState > newState:
		return EdgeFalling
	}
	return EdgeNone
}

type slice struct {
	begin int
	end   int
}

func (t *Trigger) run(in <-chan []scope.ChannelData, out chan<- []scope.ChannelData) {
	var left, source int
	var trg, scanned, found bool
	var newState, prevState thresholdState
	var lastTrg time.Time
	for d := range in {
		if !scanned {
			scanned = true
			for i := range d {
				if d[i].ID == t.source {
					source = i
					found = true
					break
				}
			}
		}
		if !found {
			out <- d
			continue
		}
		num := len(d[source].Samples)
		// slices keeps indices of the samples that should be pushed out
		var outSlices []slice
		var curSlice slice
		for i, v := range d[source].Samples {
			switch {
			case v > t.lvl:
				newState = aboveThreshold
			case v < t.lvl:
				newState = belowThreshold
			}
			// if the previous state was uninitialized, do not trigger.
			// Once state is initialized, it's always either above or below, never unknown.
			if newState != prevState && prevState == unknownThresholdState {
				prevState = newState
			}
			// newState > prevState means we moved from below threshold to above threshold, i.e. rising slope.
			if !trg {
				switch {
				// mode single and triggered once already. Don't trigger.
				case t.mode == ModeSingle && !lastTrg.IsZero():
				// crossed the threshold
				case edgeType(prevState, newState) == t.slope:
					trg = true
				// mode auto and time elapsed since last trigger.
				case t.mode == ModeAuto && time.Since(lastTrg) > 500*time.Millisecond:
					trg = true
				}
				if trg {
					lastTrg = time.Now()
					left = t.tbCount
					curSlice.begin = i
				}
			}
			if trg {
				curSlice.end = i + 1
				left--
				if left == 0 {
					outSlices = append(outSlices, curSlice)
					curSlice = slice{}
					trg = false
				}
			}
			prevState = newState
			num--
		}
		if trg {
			outSlices = append(outSlices, curSlice)
		}
		// flush samples
		if len(outSlices) > 0 {
			for _, b := range outSlices {
				chunk := make([]scope.ChannelData, len(d))
				for ch := range d {
					chunk[ch].ID = d[ch].ID
					chunk[ch].Samples = d[ch].Samples[b.begin:b.end]
				}
				out <- chunk
			}
			outSlices = outSlices[:0]
		}
	}
	close(out)
}
<|endoftext|>


<commit_before>package revel

import (
	"bytes"
	"errors"
	"fmt"
	"html"
	"html/template"
	"reflect"
	"strings"
	"time"

	"github.com/xeonx/timeago"
)

var (
	// The functions available for use in the templates.
	TemplateFuncs = map[string]interface{}{
		"url": ReverseURL,
		"set": func(viewArgs map[string]interface{}, key string, value interface{}) template.JS {
			viewArgs[key] = value
			return template.JS("")
		},
		"append": func(viewArgs map[string]interface{}, key string, value interface{}) template.JS {
			if viewArgs[key] == nil {
				viewArgs[key] = []interface{}{value}
			} else {
				viewArgs[key] = append(viewArgs[key].([]interface{}), value)
			}
			return template.JS("")
		},
		"field": NewField,
		"firstof": func(args ...interface{}) interface{} {
			for _, val := range args {
				switch val.(type) {
				case nil:
					continue
				case string:
					if val == "" {
						continue
					}
					return val
				default:
					return val
				}
			}
			return nil
		},
		"option": func(f *Field, val interface{}, label string) template.HTML {
			selected := ""
			if f.Flash() == val || (f.Flash() == "" && f.Value() == val) {
				selected = " selected"
			}

			return template.HTML(fmt.Sprintf(`<option value="%s"%s>%s</option>`,
				html.EscapeString(fmt.Sprintf("%v", val)), selected, html.EscapeString(label)))
		},
		"radio": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="radio" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		"checkbox": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="checkbox" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		// Pads the given string with  's up to the given width.
		"pad": func(str string, width int) template.HTML {
			if len(str) >= width {
				return template.HTML(html.EscapeString(str))
			}
			return template.HTML(html.EscapeString(str) + strings.Repeat(" ", width-len(str)))
		},

		"errorClass": func(name string, viewArgs map[string]interface{}) template.HTML {
			errorMap, ok := viewArgs["errors"].(map[string]*ValidationError)
			if !ok || errorMap == nil {
				templateLog.Warn("errorClass: Called 'errorClass' without 'errors' in the view args.")
				return template.HTML("")
			}
			valError, ok := errorMap[name]
			if !ok || valError == nil {
				return template.HTML("")
			}
			return template.HTML(ErrorCSSClass)
		},

		"msg": func(viewArgs map[string]interface{}, message string, args ...interface{}) template.HTML {
			str, ok := viewArgs[CurrentLocaleViewArg].(string)
			if !ok {
				return ""
			}
			return template.HTML(MessageFunc(str, message, args...))
		},

		// Replaces newlines with

		"nl2br": func(text string) template.HTML {
			return template.HTML(strings.Replace(template.HTMLEscapeString(text), "\n", "
", -1))
		},

		// Skips sanitation on the parameter.  Do not use with dynamic data.
		"raw": func(text string) template.HTML {
			return template.HTML(text)
		},

		// Pluralize, a helper for pluralizing words to correspond to data of dynamic length.
		// items - a slice of items, or an integer indicating how many items there are.
		// pluralOverrides - optional arguments specifying the output in the
		//     singular and plural cases.  by default "" and "s"
		"pluralize": func(items interface{}, pluralOverrides ...string) string {
			singular, plural := "", "s"
			if len(pluralOverrides) >= 1 {
				singular = pluralOverrides[0]
				if len(pluralOverrides) == 2 {
					plural = pluralOverrides[1]
				}
			}

			switch v := reflect.ValueOf(items); v.Kind() {
			case reflect.Int:
				if items.(int) != 1 {
					return plural
				}
			case reflect.Slice:
				if v.Len() != 1 {
					return plural
				}
			default:
				templateLog.Error("pluralize: unexpected type: ", "value", v)
			}
			return singular
		},

		// Format a date according to the application's default date(time) format.
		"date": func(date time.Time) string {
			return date.Local().Format(DateFormat)
		},
		"datetime": func(date time.Time) string {
			return date.Local().Format(DateTimeFormat)
		},
		"slug": Slug,
		"even": func(a int) bool { return (a % 2) == 0 },

		// Using https://github.com/xeonx/timeago
		"timeago": TimeAgo,
		"i18ntemplate": func(args ...interface{}) (template.HTML, error) {
			templateName, lang := "", ""
			var viewArgs interface{}
			switch len(args) {
			case 0:
				templateLog.Error("i18ntemplate: No arguments passed to template call")
			case 1:
				// Assume only the template name is passed in
				templateName = args[0].(string)
			case 2:
				// Assume template name and viewArgs is passed in
				templateName = args[0].(string)
				viewArgs = args[1]
				// Try to extract language from the view args
				if viewargsmap, ok := viewArgs.(map[string]interface{}); ok {
					lang, _ = viewargsmap[CurrentLocaleViewArg].(string)
				}
			default:
				// Assume third argument is the region
				templateName = args[0].(string)
				viewArgs = args[1]
				lang, _ = args[2].(string)
				if len(args) > 3 {
					templateLog.Error("i18ntemplate: Received more parameters then needed for", "template", templateName)
				}
			}

			var buf bytes.Buffer
			// Get template
			tmpl, err := MainTemplateLoader.TemplateLang(templateName, lang)
			if err == nil {
				err = tmpl.Render(&buf, viewArgs)
			} else {
				templateLog.Error("i18ntemplate: Failed to render i18ntemplate ", "name", templateName, "error", err)
			}
			return template.HTML(buf.String()), err
		},
	}
)

/////////////////////
// Template functions
/////////////////////

// ReverseURL returns a url capable of invoking a given controller method:
// "Application.ShowApp 123" => "/app/123"
func ReverseURL(args ...interface{}) (template.URL, error) {
	if len(args) == 0 {
		return "", errors.New("no arguments provided to reverse route")
	}

	action := args[0].(string)
	if action == "Root" {
		return template.URL(AppRoot), nil
	}

	pathData, found := splitActionPath(nil, action, true)

	if !found {
		return "", fmt.Errorf("reversing '%s', expected 'Controller.Action'", action)
	}

	// Look up the types.

	if pathData.TypeOfController == nil {
		return "", fmt.Errorf("Failed reversing %s: controller not found %#v", action, pathData)
	}

	// Note method name is case insensitive search
	methodType := pathData.TypeOfController.Method(pathData.MethodName)
	if methodType == nil {
		return "", errors.New("revel/controller: In " + action + " failed to find function " + pathData.MethodName)
	}

	if len(methodType.Args) < len(args)-1 {
		return "", fmt.Errorf("reversing %s: route defines %d args, but received %d",
			action, len(methodType.Args), len(args)-1)
	}
	// Unbind the arguments.
	argsByName := make(map[string]string)
	// Bind any static args first
	fixedParams := len(pathData.FixedParamsByName)

	for i, argValue := range args[1:] {
		if methodType.Args[i+fixedParams] == nil {
			return "", fmt.Errorf("reversing '%s', args[%d] is unknown type", action, i+fixedParams)
		}

		if argValue == nil {
			return "", fmt.Errorf("reversing '%s', args[%d] is nil", action, i+fixedParams)
		}
		Unbind(argsByName, methodType.Args[i+fixedParams].Name, argValue)
	}

	return template.URL(MainRouter.Reverse(args[0].(string), argsByName).URL), nil
}

func Slug(text string) string {
	separator := "-"
	text = strings.ToLower(text)
	text = invalidSlugPattern.ReplaceAllString(text, "")
	text = whiteSpacePattern.ReplaceAllString(text, separator)
	text = strings.Trim(text, separator)
	return text
}

var timeAgoLangs = map[string]timeago.Config{}

func TimeAgo(args ...interface{}) string {

	datetime := time.Now()
	lang := ""
	var viewArgs interface{}
	switch len(args) {
	case 0:
		templateLog.Error("TimeAgo: No arguements passed to timeago")
	case 1:
		// only the time is passed in
		datetime = args[0].(time.Time)
	case 2:
		// time and region is passed in
		datetime = args[0].(time.Time)
		switch v := reflect.ValueOf(args[1]); v.Kind() {
		case reflect.String:
			// second params type string equals region
			lang, _ = args[1].(string)
		case reflect.Map:
			// second params type map equals viewArgs
			viewArgs = args[1]
			if viewargsmap, ok := viewArgs.(map[string]interface{}); ok {
				lang, _ = viewargsmap[CurrentLocaleViewArg].(string)
			}
		default:
			templateLog.Error("TimeAgo: unexpected type: ", "value", v)
		}
	default:
		// Assume third argument is the region
		datetime = args[0].(time.Time)
		if reflect.ValueOf(args[1]).Kind() != reflect.Map {
			templateLog.Error("TimeAgo: unexpected type", "value", args[1])
		}
		if reflect.ValueOf(args[2]).Kind() != reflect.String {
			templateLog.Error("TimeAgo: unexpected type: ", "value", args[2])
		}
		viewArgs = args[1]
		lang, _ = args[2].(string)
		if len(args) > 3 {
			templateLog.Error("TimeAgo: Received more parameters then needed for timeago")
		}
	}
	if lang == "" {
		lang, _ = Config.String(defaultLanguageOption)
		if lang == "en" {
			timeAgoLangs[lang] = timeago.English
		}
	}
	_, ok := timeAgoLangs[lang]
	if !ok {
		timeAgoLangs[lang] = timeago.Config{
			PastPrefix:   "",
			PastSuffix:   " " + MessageFunc(lang, "ago"),
			FuturePrefix: MessageFunc(lang, "in") + " ",
			FutureSuffix: "",
			Periods: []timeago.FormatPeriod{
				{time.Second, MessageFunc(lang, "about a second"), MessageFunc(lang, "%d seconds")},
				{time.Minute, MessageFunc(lang, "about a minute"), MessageFunc(lang, "%d minutes")},
				{time.Hour, MessageFunc(lang, "about an hour"), MessageFunc(lang, "%d hours")},
				{timeago.Day, MessageFunc(lang, "one day"), MessageFunc(lang, "%d days")},
				{timeago.Month, MessageFunc(lang, "one month"), MessageFunc(lang, "%d months")},
				{timeago.Year, MessageFunc(lang, "one year"), MessageFunc(lang, "%d years")},
			},
			Zero:          MessageFunc(lang, "about a second"),
			Max:           73 * time.Hour,
			DefaultLayout: "2006-01-02",
		}

	}
	return timeAgoLangs[lang].Format(datetime)
}
<commit_msg>fix date result is 0001-01-01<commit_after>package revel

import (
	"bytes"
	"errors"
	"fmt"
	"html"
	"html/template"
	"reflect"
	"strings"
	"time"

	"github.com/xeonx/timeago"
)

var (
	// The functions available for use in the templates.
	TemplateFuncs = map[string]interface{}{
		"url": ReverseURL,
		"set": func(viewArgs map[string]interface{}, key string, value interface{}) template.JS {
			viewArgs[key] = value
			return template.JS("")
		},
		"append": func(viewArgs map[string]interface{}, key string, value interface{}) template.JS {
			if viewArgs[key] == nil {
				viewArgs[key] = []interface{}{value}
			} else {
				viewArgs[key] = append(viewArgs[key].([]interface{}), value)
			}
			return template.JS("")
		},
		"field": NewField,
		"firstof": func(args ...interface{}) interface{} {
			for _, val := range args {
				switch val.(type) {
				case nil:
					continue
				case string:
					if val == "" {
						continue
					}
					return val
				default:
					return val
				}
			}
			return nil
		},
		"option": func(f *Field, val interface{}, label string) template.HTML {
			selected := ""
			if f.Flash() == val || (f.Flash() == "" && f.Value() == val) {
				selected = " selected"
			}

			return template.HTML(fmt.Sprintf(`<option value="%s"%s>%s</option>`,
				html.EscapeString(fmt.Sprintf("%v", val)), selected, html.EscapeString(label)))
		},
		"radio": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="radio" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		"checkbox": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="checkbox" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		// Pads the given string with  's up to the given width.
		"pad": func(str string, width int) template.HTML {
			if len(str) >= width {
				return template.HTML(html.EscapeString(str))
			}
			return template.HTML(html.EscapeString(str) + strings.Repeat(" ", width-len(str)))
		},

		"errorClass": func(name string, viewArgs map[string]interface{}) template.HTML {
			errorMap, ok := viewArgs["errors"].(map[string]*ValidationError)
			if !ok || errorMap == nil {
				templateLog.Warn("errorClass: Called 'errorClass' without 'errors' in the view args.")
				return template.HTML("")
			}
			valError, ok := errorMap[name]
			if !ok || valError == nil {
				return template.HTML("")
			}
			return template.HTML(ErrorCSSClass)
		},

		"msg": func(viewArgs map[string]interface{}, message string, args ...interface{}) template.HTML {
			str, ok := viewArgs[CurrentLocaleViewArg].(string)
			if !ok {
				return ""
			}
			return template.HTML(MessageFunc(str, message, args...))
		},

		// Replaces newlines with

		"nl2br": func(text string) template.HTML {
			return template.HTML(strings.Replace(template.HTMLEscapeString(text), "\n", "
", -1))
		},

		// Skips sanitation on the parameter.  Do not use with dynamic data.
		"raw": func(text string) template.HTML {
			return template.HTML(text)
		},

		// Pluralize, a helper for pluralizing words to correspond to data of dynamic length.
		// items - a slice of items, or an integer indicating how many items there are.
		// pluralOverrides - optional arguments specifying the output in the
		//     singular and plural cases.  by default "" and "s"
		"pluralize": func(items interface{}, pluralOverrides ...string) string {
			singular, plural := "", "s"
			if len(pluralOverrides) >= 1 {
				singular = pluralOverrides[0]
				if len(pluralOverrides) == 2 {
					plural = pluralOverrides[1]
				}
			}

			switch v := reflect.ValueOf(items); v.Kind() {
			case reflect.Int:
				if items.(int) != 1 {
					return plural
				}
			case reflect.Slice:
				if v.Len() != 1 {
					return plural
				}
			default:
				templateLog.Error("pluralize: unexpected type: ", "value", v)
			}
			return singular
		},

		// Format a date according to the application's default date(time) format.
		"date": func(date time.Time) string {
			if date.IsZero() {
				return ""
			}
			return date.Local().Format(DateFormat)
		},
		"datetime": func(date time.Time) string {
			if date.IsZero() {
				return ""
			}
			return date.Local().Format(DateTimeFormat)
		},
		"slug": Slug,
		"even": func(a int) bool { return (a % 2) == 0 },

		// Using https://github.com/xeonx/timeago
		"timeago": TimeAgo,
		"i18ntemplate": func(args ...interface{}) (template.HTML, error) {
			templateName, lang := "", ""
			var viewArgs interface{}
			switch len(args) {
			case 0:
				templateLog.Error("i18ntemplate: No arguments passed to template call")
			case 1:
				// Assume only the template name is passed in
				templateName = args[0].(string)
			case 2:
				// Assume template name and viewArgs is passed in
				templateName = args[0].(string)
				viewArgs = args[1]
				// Try to extract language from the view args
				if viewargsmap, ok := viewArgs.(map[string]interface{}); ok {
					lang, _ = viewargsmap[CurrentLocaleViewArg].(string)
				}
			default:
				// Assume third argument is the region
				templateName = args[0].(string)
				viewArgs = args[1]
				lang, _ = args[2].(string)
				if len(args) > 3 {
					templateLog.Error("i18ntemplate: Received more parameters then needed for", "template", templateName)
				}
			}

			var buf bytes.Buffer
			// Get template
			tmpl, err := MainTemplateLoader.TemplateLang(templateName, lang)
			if err == nil {
				err = tmpl.Render(&buf, viewArgs)
			} else {
				templateLog.Error("i18ntemplate: Failed to render i18ntemplate ", "name", templateName, "error", err)
			}
			return template.HTML(buf.String()), err
		},
	}
)

/////////////////////
// Template functions
/////////////////////

// ReverseURL returns a url capable of invoking a given controller method:
// "Application.ShowApp 123" => "/app/123"
func ReverseURL(args ...interface{}) (template.URL, error) {
	if len(args) == 0 {
		return "", errors.New("no arguments provided to reverse route")
	}

	action := args[0].(string)
	if action == "Root" {
		return template.URL(AppRoot), nil
	}

	pathData, found := splitActionPath(nil, action, true)

	if !found {
		return "", fmt.Errorf("reversing '%s', expected 'Controller.Action'", action)
	}

	// Look up the types.

	if pathData.TypeOfController == nil {
		return "", fmt.Errorf("Failed reversing %s: controller not found %#v", action, pathData)
	}

	// Note method name is case insensitive search
	methodType := pathData.TypeOfController.Method(pathData.MethodName)
	if methodType == nil {
		return "", errors.New("revel/controller: In " + action + " failed to find function " + pathData.MethodName)
	}

	if len(methodType.Args) < len(args)-1 {
		return "", fmt.Errorf("reversing %s: route defines %d args, but received %d",
			action, len(methodType.Args), len(args)-1)
	}
	// Unbind the arguments.
	argsByName := make(map[string]string)
	// Bind any static args first
	fixedParams := len(pathData.FixedParamsByName)

	for i, argValue := range args[1:] {
		if methodType.Args[i+fixedParams] == nil {
			return "", fmt.Errorf("reversing '%s', args[%d] is unknown type", action, i+fixedParams)
		}

		if argValue == nil {
			return "", fmt.Errorf("reversing '%s', args[%d] is nil", action, i+fixedParams)
		}
		Unbind(argsByName, methodType.Args[i+fixedParams].Name, argValue)
	}

	return template.URL(MainRouter.Reverse(args[0].(string), argsByName).URL), nil
}

func Slug(text string) string {
	separator := "-"
	text = strings.ToLower(text)
	text = invalidSlugPattern.ReplaceAllString(text, "")
	text = whiteSpacePattern.ReplaceAllString(text, separator)
	text = strings.Trim(text, separator)
	return text
}

var timeAgoLangs = map[string]timeago.Config{}

func TimeAgo(args ...interface{}) string {

	datetime := time.Now()
	lang := ""
	var viewArgs interface{}
	switch len(args) {
	case 0:
		templateLog.Error("TimeAgo: No arguements passed to timeago")
	case 1:
		// only the time is passed in
		datetime = args[0].(time.Time)
	case 2:
		// time and region is passed in
		datetime = args[0].(time.Time)
		switch v := reflect.ValueOf(args[1]); v.Kind() {
		case reflect.String:
			// second params type string equals region
			lang, _ = args[1].(string)
		case reflect.Map:
			// second params type map equals viewArgs
			viewArgs = args[1]
			if viewargsmap, ok := viewArgs.(map[string]interface{}); ok {
				lang, _ = viewargsmap[CurrentLocaleViewArg].(string)
			}
		default:
			templateLog.Error("TimeAgo: unexpected type: ", "value", v)
		}
	default:
		// Assume third argument is the region
		datetime = args[0].(time.Time)
		if reflect.ValueOf(args[1]).Kind() != reflect.Map {
			templateLog.Error("TimeAgo: unexpected type", "value", args[1])
		}
		if reflect.ValueOf(args[2]).Kind() != reflect.String {
			templateLog.Error("TimeAgo: unexpected type: ", "value", args[2])
		}
		viewArgs = args[1]
		lang, _ = args[2].(string)
		if len(args) > 3 {
			templateLog.Error("TimeAgo: Received more parameters then needed for timeago")
		}
	}
	if lang == "" {
		lang, _ = Config.String(defaultLanguageOption)
		if lang == "en" {
			timeAgoLangs[lang] = timeago.English
		}
	}
	_, ok := timeAgoLangs[lang]
	if !ok {
		timeAgoLangs[lang] = timeago.Config{
			PastPrefix:   "",
			PastSuffix:   " " + MessageFunc(lang, "ago"),
			FuturePrefix: MessageFunc(lang, "in") + " ",
			FutureSuffix: "",
			Periods: []timeago.FormatPeriod{
				{time.Second, MessageFunc(lang, "about a second"), MessageFunc(lang, "%d seconds")},
				{time.Minute, MessageFunc(lang, "about a minute"), MessageFunc(lang, "%d minutes")},
				{time.Hour, MessageFunc(lang, "about an hour"), MessageFunc(lang, "%d hours")},
				{timeago.Day, MessageFunc(lang, "one day"), MessageFunc(lang, "%d days")},
				{timeago.Month, MessageFunc(lang, "one month"), MessageFunc(lang, "%d months")},
				{timeago.Year, MessageFunc(lang, "one year"), MessageFunc(lang, "%d years")},
			},
			Zero:          MessageFunc(lang, "about a second"),
			Max:           73 * time.Hour,
			DefaultLayout: "2006-01-02",
		}

	}
	return timeAgoLangs[lang].Format(datetime)
}
<|endoftext|>


<commit_before>package drivers

import (
	"database/sql"
	"fmt"

	// Side-effect import sql driver
	_ "github.com/lib/pq"
	"github.com/nullbio/sqlboiler/bdb"
)

// PostgresDriver holds the database connection string and a handle
// to the database connection.
type PostgresDriver struct {
	connStr string
	dbConn  *sql.DB
}

// NewPostgresDriver takes the database connection details as parameters and
// returns a pointer to a PostgresDriver object. Note that it is required to
// call PostgresDriver.Open() and PostgresDriver.Close() to open and close
// the database connection once an object has been obtained.
func NewPostgresDriver(user, pass, dbname, host string, port int) *PostgresDriver {
	driver := PostgresDriver{
		connStr: fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%d",
			user, pass, dbname, host, port),
	}

	return &driver
}

// Open opens the database connection using the connection string
func (p *PostgresDriver) Open() error {
	var err error
	p.dbConn, err = sql.Open("postgres", p.connStr)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the database connection
func (p *PostgresDriver) Close() {
	p.dbConn.Close()
}

// TableNames connects to the postgres database and
// retrieves all table names from the information_schema where the
// table schema is public. It excludes common migration tool tables
// such as gorp_migrations
func (p *PostgresDriver) TableNames() ([]string, error) {
	var names []string

	rows, err := p.dbConn.Query(`
		select table_name from information_schema.tables
		where table_schema = 'public' and table_name not like '%migrations%'
	`)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		names = append(names, name)
	}

	return names, nil
}

// Columns takes a table name and attempts to retrieve the table information
// from the database information_schema.columns. It retrieves the column names
// and column types and returns those as a []Column after TranslateColumnType()
// converts the SQL types to Go types, for example: "varchar" to "string"
func (p *PostgresDriver) Columns(tableName string) ([]bdb.Column, error) {
	var columns []bdb.Column

	rows, err := p.dbConn.Query(`
	select column_name, data_type, column_default, is_nullable
	from information_schema.columns
	where table_name=$1 and table_schema = 'public'
	`, tableName)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var colName, colType, colDefault, isNullable string
		var defaultPtr *string
		if err := rows.Scan(&colName, &colType, &defaultPtr, &isNullable); err != nil {
			return nil, fmt.Errorf("unable to scan for table %s: %s", tableName, err)
		}

		if defaultPtr == nil {
			colDefault = ""
		} else {
			colDefault = *defaultPtr
		}

		column := bdb.Column{
			Name:       colName,
			Type:       colType,
			Default:    colDefault,
			IsNullable: isNullable == "YES",
		}
		columns = append(columns, column)
	}

	return columns, nil
}

// PrimaryKeyInfo looks up the primary key for a table.
func (p *PostgresDriver) PrimaryKeyInfo(tableName string) (*bdb.PrimaryKey, error) {
	pkey := &bdb.PrimaryKey{}
	var err error

	query := `
	select tc.constraint_name
	from information_schema.table_constraints as tc
	where tc.table_name = $1 and tc.constraint_type = 'PRIMARY KEY' and tc.table_schema = 'public';`

	row := p.dbConn.QueryRow(query, tableName)
	if err = row.Scan(&pkey.Name); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	queryColumns := `
	select kcu.column_name
	from   information_schema.key_column_usage as kcu
	where  constraint_name = $1 and table_schema = 'public';`

	var rows *sql.Rows
	if rows, err = p.dbConn.Query(queryColumns, pkey.Name); err != nil {
		return nil, err
	}

	var columns []string
	for rows.Next() {
		var column string

		err = rows.Scan(&column)
		if err != nil {
			return nil, err
		}

		columns = append(columns, column)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	pkey.Columns = columns

	return pkey, nil
}

// ForeignKeyInfo retrieves the foreign keys for a given table name.
func (p *PostgresDriver) ForeignKeyInfo(tableName string) ([]bdb.ForeignKey, error) {
	var fkeys []bdb.ForeignKey

	query := `
	select
		tc.constraint_name,
		kcu.table_name as source_table,
		kcu.column_name as source_column,
		ccu.table_name as dest_table,
		ccu.column_name as dest_column
	from information_schema.table_constraints as tc
		inner join information_schema.key_column_usage as kcu ON tc.constraint_name = kcu.constraint_name
		inner join information_schema.constraint_column_usage as ccu ON tc.constraint_name = ccu.constraint_name
	where tc.table_name = $1 and tc.constraint_type = 'FOREIGN KEY' and tc.table_schema = 'information_schema';`

	var rows *sql.Rows
	var err error
	if rows, err = p.dbConn.Query(query, tableName); err != nil {
		return nil, err
	}

	for rows.Next() {
		var fkey bdb.ForeignKey
		var sourceTable string

		err = rows.Scan(&fkey.Name, &sourceTable, &fkey.Column, &fkey.ForeignTable, &fkey.ForeignColumn)
		if err != nil {
			return nil, err
		}

		fkeys = append(fkeys, fkey)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return fkeys, nil
}

// TranslateColumnType converts postgres database types to Go types, for example
// "varchar" to "string" and "bigint" to "int64". It returns this parsed data
// as a Column object.
func (p *PostgresDriver) TranslateColumnType(c bdb.Column) bdb.Column {
	if c.IsNullable {
		switch c.Type {
		case "bigint", "bigserial":
			c.Type = "null.Int64"
		case "integer", "serial":
			c.Type = "null.Int32"
		case "smallint", "smallserial":
			c.Type = "null.Int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "null.Float64"
		case "real":
			c.Type = "null.Float32"
		case "bit", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "null.String"
		case "boolean":
			c.Type = "null.Bool"
		case "date", "interval", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "null.Time"
		default:
			c.Type = "null.String"
		}
	} else {
		switch c.Type {
		case "bigint", "bigserial":
			c.Type = "int64"
		case "integer", "serial":
			c.Type = "int32"
		case "smallint", "smallserial":
			c.Type = "int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "float64"
		case "real":
			c.Type = "float32"
		case "bit", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "string"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "bool"
		case "date", "interval", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "time.Time"
		default:
			c.Type = "string"
		}
	}

	return c
}
<commit_msg>Fix a foreign key lookup bug.<commit_after>package drivers

import (
	"database/sql"
	"fmt"

	// Side-effect import sql driver
	_ "github.com/lib/pq"
	"github.com/nullbio/sqlboiler/bdb"
)

// PostgresDriver holds the database connection string and a handle
// to the database connection.
type PostgresDriver struct {
	connStr string
	dbConn  *sql.DB
}

// NewPostgresDriver takes the database connection details as parameters and
// returns a pointer to a PostgresDriver object. Note that it is required to
// call PostgresDriver.Open() and PostgresDriver.Close() to open and close
// the database connection once an object has been obtained.
func NewPostgresDriver(user, pass, dbname, host string, port int) *PostgresDriver {
	driver := PostgresDriver{
		connStr: fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%d",
			user, pass, dbname, host, port),
	}

	return &driver
}

// Open opens the database connection using the connection string
func (p *PostgresDriver) Open() error {
	var err error
	p.dbConn, err = sql.Open("postgres", p.connStr)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the database connection
func (p *PostgresDriver) Close() {
	p.dbConn.Close()
}

// TableNames connects to the postgres database and
// retrieves all table names from the information_schema where the
// table schema is public. It excludes common migration tool tables
// such as gorp_migrations
func (p *PostgresDriver) TableNames() ([]string, error) {
	var names []string

	rows, err := p.dbConn.Query(`
		select table_name from information_schema.tables
		where table_schema = 'public' and table_name not like '%migrations%'
	`)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		names = append(names, name)
	}

	return names, nil
}

// Columns takes a table name and attempts to retrieve the table information
// from the database information_schema.columns. It retrieves the column names
// and column types and returns those as a []Column after TranslateColumnType()
// converts the SQL types to Go types, for example: "varchar" to "string"
func (p *PostgresDriver) Columns(tableName string) ([]bdb.Column, error) {
	var columns []bdb.Column

	rows, err := p.dbConn.Query(`
	select column_name, data_type, column_default, is_nullable
	from information_schema.columns
	where table_name=$1 and table_schema = 'public'
	`, tableName)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var colName, colType, colDefault, isNullable string
		var defaultPtr *string
		if err := rows.Scan(&colName, &colType, &defaultPtr, &isNullable); err != nil {
			return nil, fmt.Errorf("unable to scan for table %s: %s", tableName, err)
		}

		if defaultPtr == nil {
			colDefault = ""
		} else {
			colDefault = *defaultPtr
		}

		column := bdb.Column{
			Name:       colName,
			Type:       colType,
			Default:    colDefault,
			IsNullable: isNullable == "YES",
		}
		columns = append(columns, column)
	}

	return columns, nil
}

// PrimaryKeyInfo looks up the primary key for a table.
func (p *PostgresDriver) PrimaryKeyInfo(tableName string) (*bdb.PrimaryKey, error) {
	pkey := &bdb.PrimaryKey{}
	var err error

	query := `
	select tc.constraint_name
	from information_schema.table_constraints as tc
	where tc.table_name = $1 and tc.constraint_type = 'PRIMARY KEY' and tc.table_schema = 'public';`

	row := p.dbConn.QueryRow(query, tableName)
	if err = row.Scan(&pkey.Name); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	queryColumns := `
	select kcu.column_name
	from   information_schema.key_column_usage as kcu
	where  constraint_name = $1 and table_schema = 'public';`

	var rows *sql.Rows
	if rows, err = p.dbConn.Query(queryColumns, pkey.Name); err != nil {
		return nil, err
	}

	var columns []string
	for rows.Next() {
		var column string

		err = rows.Scan(&column)
		if err != nil {
			return nil, err
		}

		columns = append(columns, column)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	pkey.Columns = columns

	return pkey, nil
}

// ForeignKeyInfo retrieves the foreign keys for a given table name.
func (p *PostgresDriver) ForeignKeyInfo(tableName string) ([]bdb.ForeignKey, error) {
	var fkeys []bdb.ForeignKey

	query := `
	select
		tc.constraint_name,
		kcu.table_name as source_table,
		kcu.column_name as source_column,
		ccu.table_name as dest_table,
		ccu.column_name as dest_column
	from information_schema.table_constraints as tc
		inner join information_schema.key_column_usage as kcu ON tc.constraint_name = kcu.constraint_name
		inner join information_schema.constraint_column_usage as ccu ON tc.constraint_name = ccu.constraint_name
	where tc.table_name = $1 and tc.constraint_type = 'FOREIGN KEY' and tc.table_schema = 'public';`

	var rows *sql.Rows
	var err error
	if rows, err = p.dbConn.Query(query, tableName); err != nil {
		return nil, err
	}

	for rows.Next() {
		var fkey bdb.ForeignKey
		var sourceTable string

		err = rows.Scan(&fkey.Name, &sourceTable, &fkey.Column, &fkey.ForeignTable, &fkey.ForeignColumn)
		if err != nil {
			return nil, err
		}

		fkeys = append(fkeys, fkey)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return fkeys, nil
}

// TranslateColumnType converts postgres database types to Go types, for example
// "varchar" to "string" and "bigint" to "int64". It returns this parsed data
// as a Column object.
func (p *PostgresDriver) TranslateColumnType(c bdb.Column) bdb.Column {
	if c.IsNullable {
		switch c.Type {
		case "bigint", "bigserial":
			c.Type = "null.Int64"
		case "integer", "serial":
			c.Type = "null.Int32"
		case "smallint", "smallserial":
			c.Type = "null.Int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "null.Float64"
		case "real":
			c.Type = "null.Float32"
		case "bit", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "null.String"
		case "boolean":
			c.Type = "null.Bool"
		case "date", "interval", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "null.Time"
		default:
			c.Type = "null.String"
		}
	} else {
		switch c.Type {
		case "bigint", "bigserial":
			c.Type = "int64"
		case "integer", "serial":
			c.Type = "int32"
		case "smallint", "smallserial":
			c.Type = "int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "float64"
		case "real":
			c.Type = "float32"
		case "bit", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "string"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "bool"
		case "date", "interval", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "time.Time"
		default:
			c.Type = "string"
		}
	}

	return c
}
<|endoftext|>


<commit_before>package ttlcache

import (
	"testing"
	"time"
)

func TestGetExisting(t *testing.T) {
	c := New(1 * time.Minute)
	key, value := "key", "value"
	c.Set(key, value)

	gotValue, ok := c.Get(key)

	if !ok {
		t.Error("ok = false, want true")
	}

	if got, want := gotValue, value; got != want {
		t.Errorf("c.Get(%q) = %v, want %q", key, got, want)
	}
}

func TestGetNonExistent(t *testing.T) {
	c := New(1 * time.Minute)
	key := "key"
	c.Set(key, "value")

	gotValue, ok := c.Get("no-key")

	if ok {
		t.Error("ok = true, want false")
	}

	if gotValue != nil {
		t.Errorf("c.Get(%q) = %v, want nil", key, gotValue)
	}
}

func TestGetTTL(t *testing.T) {
	c := New(10 * time.Millisecond)
	key, value := "key", "value"

	c.Set(key, value)
	time.Sleep(20 * time.Millisecond)

	gotValue, ok := c.Get(key)

	if gotValue != nil || ok {
		t.Errorf("c.Get(%q) = %v, %v; want <nil>, false", key, gotValue, ok)
	}
}

func TestExpire(t *testing.T) {
	c := New(1 * time.Minute)
	key, value := "key", "value"

	c.Set(key, value)
	c.Expire(key)

	gotValue, ok := c.Get(key)

	if gotValue != nil || ok {
		t.Errorf("c.Get(%q) = %v, %v; want <nil>, false", key, gotValue, ok)
	}
}

func TestExpireAll(t *testing.T) {
	c := New(1 * time.Minute)

	for i := 0; i < 10; i++ {
		c.Set(i, "value")
	}

	c.ExpireAll()

	key := 0
	gotValue, ok := c.Get(key)

	if gotValue != nil || ok {
		t.Errorf("c.Get(%d) = %q, %t; want <nil>, false", key, gotValue, ok)
	}

	if got, want := len(c.(*cache).items), 0; got != want {
		t.Errorf("cache has %d items, want %d", got, want)
	}
}

func TestSetTTLReset(t *testing.T) {
	c := New(20 * time.Millisecond)
	key, value := "key", "value"

	for i := 0; i < 10; i++ {
		c.Set(key, value)
		time.Sleep(10 * time.Millisecond)
	}

	gotValue, ok := c.Get(key)

	if gotValue != value || !ok {
		t.Errorf("c.Get(%q) = %v, %v; want %q, true", key, gotValue, ok, value)
	}
}

func BenchmarkGetExisting(b *testing.B) {
	c := New(5 * time.Minute)

	const numKeys = 100000

	for key := 0; key < numKeys; key++ {
		c.Set(key, "value")
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		c.Get(i % numKeys)
	}

	b.StopTimer()

	c.ExpireAll()
}

func BenchmarkGetNonExistent(b *testing.B) {
	c := New(5 * time.Minute)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		c.Get(i)
	}

	b.StopTimer()

	c.ExpireAll()
}

func BenchmarkSetExisting(b *testing.B) {
	c := New(5 * time.Minute)

	const numKeys = 100000

	for key := 0; key < numKeys; key++ {
		c.Set(key, "value")
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		c.Set(i%numKeys, "value")
	}

	b.StopTimer()

	c.ExpireAll()
}

func BenchmarkSetNonExistent(b *testing.B) {
	c := New(5 * time.Minute)

	const numKeys = 100000

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if i%numKeys == 0 {
			b.StopTimer()
			c.ExpireAll()
			b.StartTimer()
		}
		c.Set(i%numKeys, "value")
	}

	b.StopTimer()

	c.ExpireAll()
}
<commit_msg>Use correct print verbs in test output<commit_after>package ttlcache

import (
	"testing"
	"time"
)

func TestGetExisting(t *testing.T) {
	c := New(1 * time.Minute)
	key, value := "key", "value"
	c.Set(key, value)

	gotValue, ok := c.Get(key)

	if !ok {
		t.Error("ok = false, want true")
	}

	if got, want := gotValue, value; got != want {
		t.Errorf("c.Get(%q) = %q; want %q", key, got, want)
	}
}

func TestGetNonExistent(t *testing.T) {
	c := New(1 * time.Minute)
	key := "key"
	c.Set(key, "value")

	gotValue, ok := c.Get("no-key")

	if ok {
		t.Error("ok = true, want false")
	}

	if gotValue != nil {
		t.Errorf("c.Get(%q) = %q, want nil", key, gotValue)
	}
}

func TestGetTTL(t *testing.T) {
	c := New(10 * time.Millisecond)
	key, value := "key", "value"

	c.Set(key, value)
	time.Sleep(20 * time.Millisecond)

	gotValue, ok := c.Get(key)

	if gotValue != nil || ok {
		t.Errorf("c.Get(%q) = %q, %t; want <nil>, false", key, gotValue, ok)
	}
}

func TestExpire(t *testing.T) {
	c := New(1 * time.Minute)
	key, value := "key", "value"

	c.Set(key, value)
	c.Expire(key)

	gotValue, ok := c.Get(key)

	if gotValue != nil || ok {
		t.Errorf("c.Get(%q) = %q, %t; want <nil>, false", key, gotValue, ok)
	}
}

func TestExpireAll(t *testing.T) {
	c := New(1 * time.Minute)

	for i := 0; i < 10; i++ {
		c.Set(i, "value")
	}

	c.ExpireAll()

	key := 0
	gotValue, ok := c.Get(key)

	if gotValue != nil || ok {
		t.Errorf("c.Get(%d) = %q, %t; want <nil>, false", key, gotValue, ok)
	}

	if got, want := len(c.(*cache).items), 0; got != want {
		t.Errorf("cache has %d items, want %d", got, want)
	}
}

func TestSetTTLReset(t *testing.T) {
	c := New(20 * time.Millisecond)
	key, value := "key", "value"

	for i := 0; i < 10; i++ {
		c.Set(key, value)
		time.Sleep(10 * time.Millisecond)
	}

	gotValue, ok := c.Get(key)

	if gotValue != value || !ok {
		t.Errorf("c.Get(%q) = %q, %t; want %q, true", key, gotValue, ok, value)
	}
}

func BenchmarkGetExisting(b *testing.B) {
	c := New(5 * time.Minute)

	const numKeys = 100000

	for key := 0; key < numKeys; key++ {
		c.Set(key, "value")
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		c.Get(i % numKeys)
	}

	b.StopTimer()

	c.ExpireAll()
}

func BenchmarkGetNonExistent(b *testing.B) {
	c := New(5 * time.Minute)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		c.Get(i)
	}

	b.StopTimer()

	c.ExpireAll()
}

func BenchmarkSetExisting(b *testing.B) {
	c := New(5 * time.Minute)

	const numKeys = 100000

	for key := 0; key < numKeys; key++ {
		c.Set(key, "value")
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		c.Set(i%numKeys, "value")
	}

	b.StopTimer()

	c.ExpireAll()
}

func BenchmarkSetNonExistent(b *testing.B) {
	c := New(5 * time.Minute)

	const numKeys = 100000

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if i%numKeys == 0 {
			b.StopTimer()
			c.ExpireAll()
			b.StartTimer()
		}
		c.Set(i%numKeys, "value")
	}

	b.StopTimer()

	c.ExpireAll()
}
<|endoftext|>


<commit_before>package tutum

import "encoding/json"

type TriggerListResponse struct {
	Objects []Trigger `json:"objects"`
}

type Trigger struct {
	Url          string `json:"url"`
	Name         string `json:"name"`
	Operation    string `json:"operation"`
	Resource_uri string `json:"resource_uri"`
}

/*
func ListTriggers
Returns : Array of Trigger objects
*/
func (self *Service) ListTriggers() (TriggerListResponse, error) {
	url := "service/" + self.Uuid + "/trigger/"
	request := "GET"
	//Empty Body Request
	body := []byte(`{}`)
	var response TriggerListResponse

	data, err := TutumCall(url, request, body)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}

	return response, nil
}

/*
func GetTrigger
Argument : service uuid and Trigger uuid
Returns : Trigger JSON object
*/
func (self *Service) GetTrigger(trigger_uuid string) (Trigger, error) {

	url := ""
	if string(trigger_uuid[0]) == "/" {
		url = trigger_uuid[8:]
	} else {
		url = "service/" + self.Uuid + "/trigger/" + trigger_uuid + "/"
	}

	request := "GET"
	body := []byte(`{}`)
	var response Trigger

	data, err := TutumCall(url, request, body)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}

	return response, nil
}

/*
func CreateTrigger
Argument : service uuid and Trigger JSON object
Returns : Array of Trigger objects
*/
func (self *Service) CreateTrigger(requestBody string) ([]Trigger, error) {

	url := "service/" + self.Uuid + "/trigger/handler/"
	request := "POST"

	newTrigger := []byte(requestBody)

	var response []Trigger

	data, err := TutumCall(url, request, newTrigger)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}

	return response, nil
}

/*
func DeleteTrigger
Argument : service uuid and Trigger uuid
*/
func (self *Service) DeleteTrigger(trigger_uuid string) error {
	url := ""
	if string(trigger_uuid[0]) == "/" {
		url = trigger_uuid[8:]
	} else {
		url = "service/" + self.Uuid + "/trigger/handler/" + trigger_uuid + "/"
	}

	request := "DELETE"
	body := []byte(`{}`)
	var response Trigger

	data, err := TutumCall(url, request, body)
	if err != nil {
		return err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Add CallTrigger Method<commit_after>package tutum

import "encoding/json"

type TriggerListResponse struct {
	Objects []Trigger `json:"objects"`
}

type Trigger struct {
	Url          string `json:"url"`
	Name         string `json:"name"`
	Operation    string `json:"operation"`
	Resource_uri string `json:"resource_uri"`
}

/*
func ListTriggers
Returns : Array of Trigger objects
*/
func (self *Service) ListTriggers() (TriggerListResponse, error) {
	url := "service/" + self.Uuid + "/trigger/"
	request := "GET"
	//Empty Body Request
	body := []byte(`{}`)
	var response TriggerListResponse

	data, err := TutumCall(url, request, body)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}

	return response, nil
}

/*
func GetTrigger
Argument : service uuid and Trigger uuid
Returns : Trigger JSON object
*/
func (self *Service) GetTrigger(trigger_uuid string) (Trigger, error) {

	url := ""
	if string(trigger_uuid[0]) == "/" {
		url = trigger_uuid[8:]
	} else {
		url = "service/" + self.Uuid + "/trigger/" + trigger_uuid + "/"
	}

	request := "GET"
	body := []byte(`{}`)
	var response Trigger

	data, err := TutumCall(url, request, body)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}

	return response, nil
}

/*
func CreateTrigger
Argument : service uuid and Trigger JSON object
Returns : Array of Trigger objects
*/
func (self *Service) CreateTrigger(requestBody string) ([]Trigger, error) {

	url := "service/" + self.Uuid + "/trigger/"
	request := "POST"

	newTrigger := []byte(requestBody)

	var response []Trigger

	data, err := TutumCall(url, request, newTrigger)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}

	return response, nil
}

/*
func DeleteTrigger
Argument : service uuid and Trigger uuid
*/
func (self *Service) DeleteTrigger(trigger_uuid string) error {
	url := ""
	if string(trigger_uuid[0]) == "/" {
		url = trigger_uuid[8:]
	} else {
		url = "service/" + self.Uuid + "/trigger/" + trigger_uuid + "/"
	}

	request := "DELETE"
	body := []byte(`{}`)
	var response Trigger

	data, err := TutumCall(url, request, body)
	if err != nil {
		return err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return err
	}

	return nil
}

/*
func CallTrigger
Argument : service uuid and Trigger uuid
Returns : Trigger JSON object
*/
func (self *Service) CallTrigger(trigger_uuid string) (Trigger, error) {
	url := ""
	if string(trigger_uuid[0]) == "/" {
		url = trigger_uuid[8:]
	} else {
		url = "service/" + self.Uuid + "/trigger/" + trigger_uuid + "/call/"
	}

	request := "POST"
	body := []byte(`{}`)
	var response Trigger

	data, err := TutumCall(url, request, body)
	if err != nil {
		return response, err
	}

	err = json.Unmarshal(data, &response)
	if err != nil {
		return response, err
	}
	return response, nil
}
<|endoftext|>


<commit_before>package test

import (
	"errors"
	"os"
	"sync"
	"testing"

	"github.com/anacrolix/log"

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

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/internal/testutil"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

func justOneNetwork(cc *torrent.ClientConfig) {
	cc.DisableTCP = true
	cc.DisableIPv4 = true
}

func TestReceiveChunkStorageFailure(t *testing.T) {
	seederDataDir, metainfo := testutil.GreetingTestTorrent()
	defer os.RemoveAll(seederDataDir)
	seederClientConfig := torrent.TestingConfig()
	seederClientConfig.Debug = true
	justOneNetwork(seederClientConfig)
	seederClientStorage := storage.NewMMap(seederDataDir)
	defer seederClientStorage.Close()
	seederClientConfig.DefaultStorage = seederClientStorage
	seederClientConfig.Seed = true
	seederClientConfig.Debug = true
	seederClient, err := torrent.NewClient(seederClientConfig)
	require.NoError(t, err)
	defer testutil.ExportStatusWriter(seederClient, "s")()
	leecherClientConfig := torrent.TestingConfig()
	leecherClientConfig.Debug = true
	justOneNetwork(leecherClientConfig)
	leecherClient, err := torrent.NewClient(leecherClientConfig)
	require.NoError(t, err)
	defer testutil.ExportStatusWriter(leecherClient, "l")()
	info, err := metainfo.UnmarshalInfo()
	require.NoError(t, err)
	leecherStorage := diskFullStorage{
		pieces: make([]pieceState, info.NumPieces()),
		data:   make([]byte, info.TotalLength()),
	}
	defer leecherStorage.Close()
	leecherTorrent, new, err := leecherClient.AddTorrentSpec(&torrent.TorrentSpec{
		InfoHash: metainfo.HashInfoBytes(),
		Storage:  &leecherStorage,
	})
	leecherStorage.t = leecherTorrent
	require.NoError(t, err)
	assert.True(t, new)
	seederTorrent, err := seederClient.AddTorrent(metainfo)
	require.NoError(t, err)
	// Tell the seeder to find the leecher. Is it guaranteed seeders will always try to do this?
	seederTorrent.AddClientPeer(leecherClient)
	<-leecherTorrent.GotInfo()
	assertReadAllGreeting(t, leecherTorrent.NewReader())
}

type pieceState struct {
	complete bool
}

type diskFullStorage struct {
	pieces                        []pieceState
	t                             *torrent.Torrent
	defaultHandledWriteChunkError bool
	data                          []byte

	mu          sync.Mutex
	diskNotFull bool
}

func (me *diskFullStorage) Piece(p metainfo.Piece) storage.PieceImpl {
	return pieceImpl{
		mip:             p,
		diskFullStorage: me,
	}
}

func (me diskFullStorage) Close() error {
	return nil
}

func (d diskFullStorage) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash) (storage.TorrentImpl, error) {
	return &d, nil
}

type pieceImpl struct {
	mip metainfo.Piece
	*diskFullStorage
}

func (me pieceImpl) state() *pieceState {
	return &me.diskFullStorage.pieces[me.mip.Index()]
}

func (me pieceImpl) ReadAt(p []byte, off int64) (n int, err error) {
	off += me.mip.Offset()
	return copy(p, me.data[off:]), nil
}

func (me pieceImpl) WriteAt(p []byte, off int64) (int, error) {
	off += me.mip.Offset()
	if !me.defaultHandledWriteChunkError {
		go func() {
			me.t.SetOnWriteChunkError(func(err error) {
				log.Printf("got write chunk error to custom handler: %v", err)
				me.mu.Lock()
				me.diskNotFull = true
				me.mu.Unlock()
				me.t.AllowDataDownload()
			})
			me.t.AllowDataDownload()
		}()
		me.defaultHandledWriteChunkError = true
	}
	me.mu.Lock()
	defer me.mu.Unlock()
	if me.diskNotFull {
		return copy(me.data[off:], p), nil
	}
	return copy(me.data[off:], p[:1]), errors.New("disk full")
}

func (me pieceImpl) MarkComplete() error {
	me.state().complete = true
	return nil
}

func (me pieceImpl) MarkNotComplete() error {
	panic("implement me")
}

func (me pieceImpl) Completion() storage.Completion {
	return storage.Completion{
		Complete: me.state().complete,
		Ok:       true,
	}
}
<commit_msg>Close leaked Clients in test<commit_after>package test

import (
	"errors"
	"os"
	"sync"
	"testing"

	"github.com/anacrolix/log"

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

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/internal/testutil"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

func justOneNetwork(cc *torrent.ClientConfig) {
	cc.DisableTCP = true
	cc.DisableIPv4 = true
}

func TestReceiveChunkStorageFailure(t *testing.T) {
	seederDataDir, metainfo := testutil.GreetingTestTorrent()
	defer os.RemoveAll(seederDataDir)
	seederClientConfig := torrent.TestingConfig()
	seederClientConfig.Debug = true
	justOneNetwork(seederClientConfig)
	seederClientStorage := storage.NewMMap(seederDataDir)
	defer seederClientStorage.Close()
	seederClientConfig.DefaultStorage = seederClientStorage
	seederClientConfig.Seed = true
	seederClientConfig.Debug = true
	seederClient, err := torrent.NewClient(seederClientConfig)
	require.NoError(t, err)
	defer seederClient.Close()
	defer testutil.ExportStatusWriter(seederClient, "s")()
	leecherClientConfig := torrent.TestingConfig()
	leecherClientConfig.Debug = true
	justOneNetwork(leecherClientConfig)
	leecherClient, err := torrent.NewClient(leecherClientConfig)
	require.NoError(t, err)
	defer leecherClient.Close()
	defer testutil.ExportStatusWriter(leecherClient, "l")()
	info, err := metainfo.UnmarshalInfo()
	require.NoError(t, err)
	leecherStorage := diskFullStorage{
		pieces: make([]pieceState, info.NumPieces()),
		data:   make([]byte, info.TotalLength()),
	}
	defer leecherStorage.Close()
	leecherTorrent, new, err := leecherClient.AddTorrentSpec(&torrent.TorrentSpec{
		InfoHash: metainfo.HashInfoBytes(),
		Storage:  &leecherStorage,
	})
	leecherStorage.t = leecherTorrent
	require.NoError(t, err)
	assert.True(t, new)
	seederTorrent, err := seederClient.AddTorrent(metainfo)
	require.NoError(t, err)
	// Tell the seeder to find the leecher. Is it guaranteed seeders will always try to do this?
	seederTorrent.AddClientPeer(leecherClient)
	<-leecherTorrent.GotInfo()
	assertReadAllGreeting(t, leecherTorrent.NewReader())
}

type pieceState struct {
	complete bool
}

type diskFullStorage struct {
	pieces                        []pieceState
	t                             *torrent.Torrent
	defaultHandledWriteChunkError bool
	data                          []byte

	mu          sync.Mutex
	diskNotFull bool
}

func (me *diskFullStorage) Piece(p metainfo.Piece) storage.PieceImpl {
	return pieceImpl{
		mip:             p,
		diskFullStorage: me,
	}
}

func (me diskFullStorage) Close() error {
	return nil
}

func (d diskFullStorage) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash) (storage.TorrentImpl, error) {
	return &d, nil
}

type pieceImpl struct {
	mip metainfo.Piece
	*diskFullStorage
}

func (me pieceImpl) state() *pieceState {
	return &me.diskFullStorage.pieces[me.mip.Index()]
}

func (me pieceImpl) ReadAt(p []byte, off int64) (n int, err error) {
	off += me.mip.Offset()
	return copy(p, me.data[off:]), nil
}

func (me pieceImpl) WriteAt(p []byte, off int64) (int, error) {
	off += me.mip.Offset()
	if !me.defaultHandledWriteChunkError {
		go func() {
			me.t.SetOnWriteChunkError(func(err error) {
				log.Printf("got write chunk error to custom handler: %v", err)
				me.mu.Lock()
				me.diskNotFull = true
				me.mu.Unlock()
				me.t.AllowDataDownload()
			})
			me.t.AllowDataDownload()
		}()
		me.defaultHandledWriteChunkError = true
	}
	me.mu.Lock()
	defer me.mu.Unlock()
	if me.diskNotFull {
		return copy(me.data[off:], p), nil
	}
	return copy(me.data[off:], p[:1]), errors.New("disk full")
}

func (me pieceImpl) MarkComplete() error {
	me.state().complete = true
	return nil
}

func (me pieceImpl) MarkNotComplete() error {
	panic("implement me")
}

func (me pieceImpl) Completion() storage.Completion {
	return storage.Completion{
		Complete: me.state().complete,
		Ok:       true,
	}
}
<|endoftext|>