<commit_before>// Package flags provides a convenient way to configure the client connection from command line
// flags.
package flags

import (
	"context"
	"flag"

	"github.com/bazelbuild/remote-apis-sdks/go/client"
)

var (
	// The flags credential_file, use_application_default_credentials, and use_gce_credentials
	// determine the client identity that is used to authenticate with remote execution.
	// One of the following must be true for client authentication to work, and they are used in
	// this order of preference:
	// - the use_application_default_credentials flag must be set to true, or
	// - the use_gce_credentials must be set to true, or
	// - the credential_file flag must be set to point to a valid credential file

	// CredFile is the name of a file that contains service account credentials to use when calling
	// remote execution. Used only if --use_application_default_credentials and --use_gce_credentials
	// are false.
	CredFile = flag.String("credential_file", "", "The name of a file that contains service account credentials to use when calling remote execution. Used only if --use_application_default_credentials and --use_gce_credentials are false.")
	// UseApplicationDefaultCreds is whether to use application default credentials to connect to
	// remote execution. See
	// https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login
	UseApplicationDefaultCreds = flag.Bool("use_application_default_credentials", false, "If true, use application default credentials to connect to remote execution. See https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login")
	// UseGCECredentials is whether to use the default GCE credentials to authenticate with remote
	// execution. --use_application_default_credentials must be false.
	UseGCECredentials = flag.Bool("use_gce_credentials", false, "If true (and --use_application_default_credentials is false), use the default GCE credentials to authenticate with remote execution.")
	// Service represents the host (and, if applicable, port) of the remote execution service.
	Service = flag.String("service", "", "The remote execution service to dial when calling via gRPC, including port, such as 'localhost:8790' or 'remotebuildexecution.googleapis.com:443'")
	// Instance gives the instance of remote execution to test (in
	// projects/[PROJECT_ID]/instances/[INSTANCE_NAME] format for Google RBE).
	Instance = flag.String("instance", "", "The instance ID to target when calling remote execution via gRPC (e.g., projects/$PROJECT/instances/default_instance for Google RBE).")
)

// DialFromFlags dials a remote execution service and returns a client suitable for higher-level
// functionality. It uses the flags from above to configure the connection to remote execution.
func DialFromFlags(ctx context.Context, opts ...client.Opt) (*client.Client, error) {
	return client.Dial(ctx, *Instance, client.DialParams{
		Service:               *Service,
		CredFile:              *CredFile,
		UseApplicationDefault: *UseApplicationDefaultCreds,
		UseComputeEngine:      *UseGCECredentials,
	}, opts...)
}
<commit_msg>Demo commit (#9)<commit_after>// Package flags provides a convenient way to initialize the remote client from flags.
package flags

import (
	"context"
	"flag"

	"github.com/bazelbuild/remote-apis-sdks/go/client"
)

var (
	// The flags credential_file, use_application_default_credentials, and use_gce_credentials
	// determine the client identity that is used to authenticate with remote execution.
	// One of the following must be true for client authentication to work, and they are used in
	// this order of preference:
	// - the use_application_default_credentials flag must be set to true, or
	// - the use_gce_credentials must be set to true, or
	// - the credential_file flag must be set to point to a valid credential file

	// CredFile is the name of a file that contains service account credentials to use when calling
	// remote execution. Used only if --use_application_default_credentials and --use_gce_credentials
	// are false.
	CredFile = flag.String("credential_file", "", "The name of a file that contains service account credentials to use when calling remote execution. Used only if --use_application_default_credentials and --use_gce_credentials are false.")
	// UseApplicationDefaultCreds is whether to use application default credentials to connect to
	// remote execution. See
	// https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login
	UseApplicationDefaultCreds = flag.Bool("use_application_default_credentials", false, "If true, use application default credentials to connect to remote execution. See https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login")
	// UseGCECredentials is whether to use the default GCE credentials to authenticate with remote
	// execution. --use_application_default_credentials must be false.
	UseGCECredentials = flag.Bool("use_gce_credentials", false, "If true (and --use_application_default_credentials is false), use the default GCE credentials to authenticate with remote execution.")
	// Service represents the host (and, if applicable, port) of the remote execution service.
	Service = flag.String("service", "", "The remote execution service to dial when calling via gRPC, including port, such as 'localhost:8790' or 'remotebuildexecution.googleapis.com:443'")
	// Instance gives the instance of remote execution to test (in
	// projects/[PROJECT_ID]/instances/[INSTANCE_NAME] format for Google RBE).
	Instance = flag.String("instance", "", "The instance ID to target when calling remote execution via gRPC (e.g., projects/$PROJECT/instances/default_instance for Google RBE).")
)

// DialFromFlags dials a remote execution service and returns a client suitable for higher-level
// functionality. It uses the flags from above to configure the connection to remote execution.
func DialFromFlags(ctx context.Context, opts ...client.Opt) (*client.Client, error) {
	return client.Dial(ctx, *Instance, client.DialParams{
		Service:               *Service,
		CredFile:              *CredFile,
		UseApplicationDefault: *UseApplicationDefaultCreds,
		UseComputeEngine:      *UseGCECredentials,
	}, opts...)
}
<|endoftext|>


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

package mysql

/*
#cgo pkg-config: gomysql
#include <stdlib.h>
#include <mysqld_error.h>
#include "vtmysql.h"
*/
import "C"

import (
	"fmt"
	"unsafe"

	"github.com/youtube/vitess/go/hack"
	"github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/sqltypes"
)

const (
	// NOTE(szopa): maxSize used to be 1 << 30, but that causes
	// compiler errors in some situations.
	maxSize = 1 << 20
)

func init() {
	// This needs to be called before threads begin to spawn.
	C.vt_library_init()
}

const (
	DUP_ENTRY         = C.ER_DUP_ENTRY
	LOCK_WAIT_TIMEOUT = C.ER_LOCK_WAIT_TIMEOUT
	LOCK_DEADLOCK     = C.ER_LOCK_DEADLOCK
)

type SqlError struct {
	Num     int
	Message string
	Query   string
}

func NewSqlError(number int, format string, args ...interface{}) *SqlError {
	return &SqlError{Num: number, Message: fmt.Sprintf(format, args...)}
}

func (se *SqlError) Error() string {
	if se.Query == "" {
		return fmt.Sprintf("%v (errno %v)", se.Message, se.Num)
	}
	return fmt.Sprintf("%v (errno %v) during query: %s", se.Message, se.Num, se.Query)
}

func (se *SqlError) Number() int {
	return se.Num
}

func handleError(err *error) {
	if x := recover(); x != nil {
		terr := x.(*SqlError)
		*err = terr
	}
}

type ConnectionParams struct {
	Host       string `json:"host"`
	Port       int    `json:"port"`
	Uname      string `json:"uname"`
	Pass       string `json:"pass"`
	DbName     string `json:"dbname"`
	UnixSocket string `json:"unix_socket"`
	Charset    string `json:"charset"`
	Flags      uint64 `json:"flags"`

	// the following flags are only used for 'Change Master' command
	// for now (along with flags |= 2048 for CLIENT_SSL)
	SslCa     string `json:"ssl_ca"`
	SslCaPath string `json:"ssl_ca_path"`
	SslCert   string `json:"ssl_cert"`
	SslKey    string `json:"ssl_key"`
}

func (c *ConnectionParams) EnableMultiStatements() {
	c.Flags |= C.CLIENT_MULTI_STATEMENTS
}

func (c *ConnectionParams) SslEnabled() bool {
	return (c.Flags & C.CLIENT_SSL) != 0
}

func (c ConnectionParams) Redacted() ConnectionParams {
	c.Pass = "****"
	return c
}

type Connection struct {
	c C.VT_CONN
}

func Connect(params ConnectionParams) (conn *Connection, err error) {
	defer handleError(&err)

	host := C.CString(params.Host)
	defer cfree(host)
	port := C.uint(params.Port)
	uname := C.CString(params.Uname)
	defer cfree(uname)
	pass := C.CString(params.Pass)
	defer cfree(pass)
	dbname := C.CString(params.DbName)
	defer cfree(dbname)
	unix_socket := C.CString(params.UnixSocket)
	defer cfree(unix_socket)
	charset := C.CString(params.Charset)
	defer cfree(charset)
	flags := C.ulong(params.Flags)

	conn = &Connection{}
	if C.vt_connect(&conn.c, host, uname, pass, dbname, port, unix_socket, charset, flags) != 0 {
		defer conn.Close()
		return nil, conn.lastError("")
	}
	return conn, nil
}

func (conn *Connection) Close() {
	C.vt_close(&conn.c)
}

func (conn *Connection) IsClosed() bool {
	return conn.c.mysql == nil
}

func (conn *Connection) ExecuteFetch(query string, maxrows int, wantfields bool) (qr *proto.QueryResult, err error) {
	if conn.IsClosed() {
		return nil, NewSqlError(2006, "Connection is closed")
	}

	if C.vt_execute(&conn.c, (*C.char)(hack.StringPointer(query)), C.ulong(len(query)), 0) != 0 {
		return nil, conn.lastError(query)
	}
	defer conn.CloseResult()

	qr = &proto.QueryResult{}
	qr.RowsAffected = uint64(conn.c.affected_rows)
	qr.InsertId = uint64(conn.c.insert_id)
	if conn.c.num_fields == 0 {
		return qr, nil
	}

	if qr.RowsAffected > uint64(maxrows) {
		return nil, &SqlError{0, fmt.Sprintf("Row count exceeded %d", maxrows), string(query)}
	}
	if wantfields {
		qr.Fields = conn.Fields()
	}
	qr.Rows, err = conn.fetchAll()
	return qr, err
}

// when using ExecuteStreamFetch, use FetchNext on the Connection until it returns nil or error
func (conn *Connection) ExecuteStreamFetch(query string) (err error) {
	if conn.IsClosed() {
		return NewSqlError(2006, "Connection is closed")
	}
	if C.vt_execute(&conn.c, (*C.char)(hack.StringPointer(query)), C.ulong(len(query)), 1) != 0 {
		return conn.lastError(query)
	}
	return nil
}

func (conn *Connection) Fields() (fields []proto.Field) {
	nfields := int(conn.c.num_fields)
	if nfields == 0 {
		return nil
	}
	cfields := (*[maxSize]C.MYSQL_FIELD)(unsafe.Pointer(conn.c.fields))
	totalLength := uint64(0)
	for i := 0; i < nfields; i++ {
		totalLength += uint64(cfields[i].name_length)
	}
	fields = make([]proto.Field, nfields)
	for i := 0; i < nfields; i++ {
		length := cfields[i].name_length
		fname := (*[maxSize]byte)(unsafe.Pointer(cfields[i].name))[:length]
		fields[i].Name = string(fname)
		fields[i].Type = int64(cfields[i]._type)
	}
	return fields
}

func (conn *Connection) fetchAll() (rows [][]sqltypes.Value, err error) {
	rowCount := int(conn.c.affected_rows)
	if rowCount == 0 {
		return nil, nil
	}
	rows = make([][]sqltypes.Value, rowCount)
	for i := 0; i < rowCount; i++ {
		rows[i], err = conn.FetchNext()
		if err != nil {
			return nil, err
		}
	}
	return rows, nil
}

func (conn *Connection) FetchNext() (row []sqltypes.Value, err error) {
	vtrow := C.vt_fetch_next(&conn.c)
	if vtrow.has_error != 0 {
		return nil, conn.lastError("")
	}
	rowPtr := (*[maxSize]*[maxSize]byte)(unsafe.Pointer(vtrow.mysql_row))
	if rowPtr == nil {
		return nil, nil
	}
	colCount := int(conn.c.num_fields)
	cfields := (*[maxSize]C.MYSQL_FIELD)(unsafe.Pointer(conn.c.fields))
	row = make([]sqltypes.Value, colCount)
	lengths := (*[maxSize]uint64)(unsafe.Pointer(vtrow.lengths))
	totalLength := uint64(0)
	for i := 0; i < colCount; i++ {
		totalLength += lengths[i]
	}
	arena := make([]byte, 0, int(totalLength))
	for i := 0; i < colCount; i++ {
		colLength := lengths[i]
		colPtr := rowPtr[i]
		if colPtr == nil {
			continue
		}
		start := len(arena)
		arena = append(arena, colPtr[:colLength]...)
		row[i] = BuildValue(arena[start:start+int(colLength)], cfields[i]._type)
	}
	return row, nil
}

func (conn *Connection) CloseResult() {
	C.vt_close_result(&conn.c)
}

func (conn *Connection) Id() int64 {
	if conn.c.mysql == nil {
		return 0
	}
	return int64(C.vt_thread_id(&conn.c))
}

func (conn *Connection) lastError(query string) error {
	if err := C.vt_error(&conn.c); *err != 0 {
		return &SqlError{Num: int(C.vt_errno(&conn.c)), Message: C.GoString(err), Query: query}
	}
	return &SqlError{0, "Dummy", string(query)}
}

func BuildValue(bytes []byte, fieldType uint32) sqltypes.Value {
	switch fieldType {
	case C.MYSQL_TYPE_DECIMAL, C.MYSQL_TYPE_FLOAT, C.MYSQL_TYPE_DOUBLE, C.MYSQL_TYPE_NEWDECIMAL:
		return sqltypes.MakeFractional(bytes)
	case C.MYSQL_TYPE_TIMESTAMP:
		return sqltypes.MakeString(bytes)
	}
	// The below condition represents the following list of values:
	// C.MYSQL_TYPE_TINY, C.MYSQL_TYPE_SHORT, C.MYSQL_TYPE_LONG, C.MYSQL_TYPE_LONGLONG, C.MYSQL_TYPE_INT24, C.MYSQL_TYPE_YEAR:
	if fieldType <= C.MYSQL_TYPE_INT24 || fieldType == C.MYSQL_TYPE_YEAR {
		return sqltypes.MakeNumeric(bytes)
	}
	return sqltypes.MakeString(bytes)
}

func cfree(str *C.char) {
	if str != nil {
		C.free(unsafe.Pointer(str))
	}
}
<commit_msg>Fix handling of NULL in mysql row values<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mysql

/*
#cgo pkg-config: gomysql
#include <stdlib.h>
#include <mysqld_error.h>
#include "vtmysql.h"
*/
import "C"

import (
	"fmt"
	"unsafe"

	"github.com/youtube/vitess/go/hack"
	"github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/sqltypes"
)

const (
	// NOTE(szopa): maxSize used to be 1 << 30, but that causes
	// compiler errors in some situations.
	maxSize = 1 << 20
)

func init() {
	// This needs to be called before threads begin to spawn.
	C.vt_library_init()
}

const (
	DUP_ENTRY         = C.ER_DUP_ENTRY
	LOCK_WAIT_TIMEOUT = C.ER_LOCK_WAIT_TIMEOUT
	LOCK_DEADLOCK     = C.ER_LOCK_DEADLOCK
)

type SqlError struct {
	Num     int
	Message string
	Query   string
}

func NewSqlError(number int, format string, args ...interface{}) *SqlError {
	return &SqlError{Num: number, Message: fmt.Sprintf(format, args...)}
}

func (se *SqlError) Error() string {
	if se.Query == "" {
		return fmt.Sprintf("%v (errno %v)", se.Message, se.Num)
	}
	return fmt.Sprintf("%v (errno %v) during query: %s", se.Message, se.Num, se.Query)
}

func (se *SqlError) Number() int {
	return se.Num
}

func handleError(err *error) {
	if x := recover(); x != nil {
		terr := x.(*SqlError)
		*err = terr
	}
}

type ConnectionParams struct {
	Host       string `json:"host"`
	Port       int    `json:"port"`
	Uname      string `json:"uname"`
	Pass       string `json:"pass"`
	DbName     string `json:"dbname"`
	UnixSocket string `json:"unix_socket"`
	Charset    string `json:"charset"`
	Flags      uint64 `json:"flags"`

	// the following flags are only used for 'Change Master' command
	// for now (along with flags |= 2048 for CLIENT_SSL)
	SslCa     string `json:"ssl_ca"`
	SslCaPath string `json:"ssl_ca_path"`
	SslCert   string `json:"ssl_cert"`
	SslKey    string `json:"ssl_key"`
}

func (c *ConnectionParams) EnableMultiStatements() {
	c.Flags |= C.CLIENT_MULTI_STATEMENTS
}

func (c *ConnectionParams) SslEnabled() bool {
	return (c.Flags & C.CLIENT_SSL) != 0
}

func (c ConnectionParams) Redacted() ConnectionParams {
	c.Pass = "****"
	return c
}

type Connection struct {
	c C.VT_CONN
}

func Connect(params ConnectionParams) (conn *Connection, err error) {
	defer handleError(&err)

	host := C.CString(params.Host)
	defer cfree(host)
	port := C.uint(params.Port)
	uname := C.CString(params.Uname)
	defer cfree(uname)
	pass := C.CString(params.Pass)
	defer cfree(pass)
	dbname := C.CString(params.DbName)
	defer cfree(dbname)
	unix_socket := C.CString(params.UnixSocket)
	defer cfree(unix_socket)
	charset := C.CString(params.Charset)
	defer cfree(charset)
	flags := C.ulong(params.Flags)

	conn = &Connection{}
	if C.vt_connect(&conn.c, host, uname, pass, dbname, port, unix_socket, charset, flags) != 0 {
		defer conn.Close()
		return nil, conn.lastError("")
	}
	return conn, nil
}

func (conn *Connection) Close() {
	C.vt_close(&conn.c)
}

func (conn *Connection) IsClosed() bool {
	return conn.c.mysql == nil
}

func (conn *Connection) ExecuteFetch(query string, maxrows int, wantfields bool) (qr *proto.QueryResult, err error) {
	if conn.IsClosed() {
		return nil, NewSqlError(2006, "Connection is closed")
	}

	if C.vt_execute(&conn.c, (*C.char)(hack.StringPointer(query)), C.ulong(len(query)), 0) != 0 {
		return nil, conn.lastError(query)
	}
	defer conn.CloseResult()

	qr = &proto.QueryResult{}
	qr.RowsAffected = uint64(conn.c.affected_rows)
	qr.InsertId = uint64(conn.c.insert_id)
	if conn.c.num_fields == 0 {
		return qr, nil
	}

	if qr.RowsAffected > uint64(maxrows) {
		return nil, &SqlError{0, fmt.Sprintf("Row count exceeded %d", maxrows), string(query)}
	}
	if wantfields {
		qr.Fields = conn.Fields()
	}
	qr.Rows, err = conn.fetchAll()
	return qr, err
}

// when using ExecuteStreamFetch, use FetchNext on the Connection until it returns nil or error
func (conn *Connection) ExecuteStreamFetch(query string) (err error) {
	if conn.IsClosed() {
		return NewSqlError(2006, "Connection is closed")
	}
	if C.vt_execute(&conn.c, (*C.char)(hack.StringPointer(query)), C.ulong(len(query)), 1) != 0 {
		return conn.lastError(query)
	}
	return nil
}

func (conn *Connection) Fields() (fields []proto.Field) {
	nfields := int(conn.c.num_fields)
	if nfields == 0 {
		return nil
	}
	cfields := (*[maxSize]C.MYSQL_FIELD)(unsafe.Pointer(conn.c.fields))
	totalLength := uint64(0)
	for i := 0; i < nfields; i++ {
		totalLength += uint64(cfields[i].name_length)
	}
	fields = make([]proto.Field, nfields)
	for i := 0; i < nfields; i++ {
		length := cfields[i].name_length
		fname := (*[maxSize]byte)(unsafe.Pointer(cfields[i].name))[:length]
		fields[i].Name = string(fname)
		fields[i].Type = int64(cfields[i]._type)
	}
	return fields
}

func (conn *Connection) fetchAll() (rows [][]sqltypes.Value, err error) {
	rowCount := int(conn.c.affected_rows)
	if rowCount == 0 {
		return nil, nil
	}
	rows = make([][]sqltypes.Value, rowCount)
	for i := 0; i < rowCount; i++ {
		rows[i], err = conn.FetchNext()
		if err != nil {
			return nil, err
		}
	}
	return rows, nil
}

func (conn *Connection) FetchNext() (row []sqltypes.Value, err error) {
	vtrow := C.vt_fetch_next(&conn.c)
	if vtrow.has_error != 0 {
		return nil, conn.lastError("")
	}
	rowPtr := (*[maxSize]*[maxSize]byte)(unsafe.Pointer(vtrow.mysql_row))
	if rowPtr == nil {
		return nil, nil
	}
	colCount := int(conn.c.num_fields)
	cfields := (*[maxSize]C.MYSQL_FIELD)(unsafe.Pointer(conn.c.fields))
	row = make([]sqltypes.Value, colCount)
	lengths := (*[maxSize]uint64)(unsafe.Pointer(vtrow.lengths))
	totalLength := uint64(0)
	for i := 0; i < colCount; i++ {
		totalLength += lengths[i]
	}
	arena := make([]byte, 0, int(totalLength))
	for i := 0; i < colCount; i++ {
		colLength := lengths[i]
		colPtr := rowPtr[i]
		if colPtr == nil {
			continue
		}
		start := len(arena)
		arena = append(arena, colPtr[:colLength]...)
		row[i] = BuildValue(arena[start:start+int(colLength)], cfields[i]._type)
	}
	return row, nil
}

func (conn *Connection) CloseResult() {
	C.vt_close_result(&conn.c)
}

func (conn *Connection) Id() int64 {
	if conn.c.mysql == nil {
		return 0
	}
	return int64(C.vt_thread_id(&conn.c))
}

func (conn *Connection) lastError(query string) error {
	if err := C.vt_error(&conn.c); *err != 0 {
		return &SqlError{Num: int(C.vt_errno(&conn.c)), Message: C.GoString(err), Query: query}
	}
	return &SqlError{0, "Dummy", string(query)}
}

func BuildValue(bytes []byte, fieldType uint32) sqltypes.Value {
	if bytes == nil {
		return sqltypes.NULL
	}
	switch fieldType {
	case C.MYSQL_TYPE_DECIMAL, C.MYSQL_TYPE_FLOAT, C.MYSQL_TYPE_DOUBLE, C.MYSQL_TYPE_NEWDECIMAL:
		return sqltypes.MakeFractional(bytes)
	case C.MYSQL_TYPE_TIMESTAMP:
		return sqltypes.MakeString(bytes)
	}
	// The below condition represents the following list of values:
	// C.MYSQL_TYPE_TINY, C.MYSQL_TYPE_SHORT, C.MYSQL_TYPE_LONG, C.MYSQL_TYPE_LONGLONG, C.MYSQL_TYPE_INT24, C.MYSQL_TYPE_YEAR:
	if fieldType <= C.MYSQL_TYPE_INT24 || fieldType == C.MYSQL_TYPE_YEAR {
		return sqltypes.MakeNumeric(bytes)
	}
	return sqltypes.MakeString(bytes)
}

func cfree(str *C.char) {
	if str != nil {
		C.free(unsafe.Pointer(str))
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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 power provides utilities for managing system power state.
package power

import (
	"golang.org/x/sys/windows"
)

// Flag mirrors Windows exit flags
// https://docs.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-exitwindowsex
type Flag uint32

const (
	EWX_LOGOFF   = Flag(0x00000000)
	EWX_POWEROFF = Flag(0x00000008)
	EWX_REBOOT   = Flag(0x00000002)
	EWX_SHUTDOWN = Flag(0x00000001)
)

// Reason Codes for system shutdown
// https://docs.microsoft.com/en-us/windows/win32/shutdown/system-shutdown-reason-codes
type Reason uint32

const (
	SHTDN_REASON_MINOR_OTHER        = Reason(0x00000000)
	SHTDN_REASON_MINOR_MAINTENANCE  = Reason(0x00000001)
	SHTDN_REASON_MINOR_INSTALLATION = Reason(0x00000002)
	SHTDN_REASON_MINOR_UPGRADE      = Reason(0x00000003)
	SHTDN_REASON_MINOR_RECONFIG     = Reason(0x00000004)
)

// Exit exits the active session using custom settings.
//
// Example: Exit(EWX_LOGOFF, SHTDN_REASON_MINOR_MAINTENANCE)
func Exit(flag Flag, reason Reason) error {
	return windows.ExitWindowsEx(uint32(flag), uint32(reason))
}

// Reboot reboots the system.
//
// Example: Reboot(SHTDN_REASON_MINOR_MAINTENANCE)
func Reboot(reason Reason) error {
	return windows.ExitWindowsEx(uint32(EWX_REBOOT), uint32(reason))
}

// Shutdown shuts down the system.
//
// Example: Shutdown(SHTDN_REASON_MINOR_MAINTENANCE)
func Shutdown(reason Reason) error {
	return windows.ExitWindowsEx(uint32(EWX_SHUTDOWN), uint32(reason))
}
<commit_msg>Obtain the proper privileges to enable reboot/shutdown API calls.<commit_after>// Copyright 2021 Google LLC
//
// 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 power provides utilities for managing system power state.
package power

import (
	"fmt"
	"syscall"

	"golang.org/x/sys/windows"
)

// Flag mirrors Windows exit flags
// https://docs.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-exitwindowsex
type Flag uint32

const (
	EWX_LOGOFF   = Flag(0x00000000)
	EWX_SHUTDOWN = Flag(0x00000001)
	EWX_REBOOT   = Flag(0x00000002)
	EWX_FORCE    = Flag(0x00000004)
	EWX_POWEROFF = Flag(0x00000008)
)

// Reason Codes for system shutdown
// https://docs.microsoft.com/en-us/windows/win32/shutdown/system-shutdown-reason-codes
type Reason uint32

const (
	SHTDN_REASON_MINOR_OTHER        = Reason(0x00000000)
	SHTDN_REASON_MINOR_MAINTENANCE  = Reason(0x00000001)
	SHTDN_REASON_MINOR_INSTALLATION = Reason(0x00000002)
	SHTDN_REASON_MINOR_UPGRADE      = Reason(0x00000003)
	SHTDN_REASON_MINOR_RECONFIG     = Reason(0x00000004)
)

// Privileges
const (
	SE_SHUTDOWN_NAME = "SeShutdownPrivilege"
)

// To shut down or restart the system, the calling process must use the AdjustTokenPrivileges function to enable the SE_SHUTDOWN_NAME privilege.
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-exitwindowsex
func setPrivToken() error {
	var hToken windows.Token
	err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &hToken)
	if err != nil {
		return fmt.Errorf("windows.OpenProcessToken: %w", err)
	}
	defer hToken.Close()

	tkp := windows.Tokenprivileges{}
	err = windows.LookupPrivilegeValue(nil, syscall.StringToUTF16Ptr(SE_SHUTDOWN_NAME), &tkp.Privileges[0].Luid)
	if err != nil {
		return fmt.Errorf("windows.LookupPrivilegeValue: %w", err)
	}

	tkp.PrivilegeCount = 1
	tkp.Privileges[0].Attributes = windows.SE_PRIVILEGE_ENABLED

	err = windows.AdjustTokenPrivileges(hToken, false, &tkp, 0, nil, nil)
	if err != nil {
		return fmt.Errorf("windows.AdjustTokenPrivileges: %w", err)
	}
	return nil
}

// Exit exits the active session using custom settings.
//
// Example: Exit(EWX_LOGOFF, SHTDN_REASON_MINOR_MAINTENANCE)
func Exit(flag Flag, reason Reason) error {
	return windows.ExitWindowsEx(uint32(flag), uint32(reason))
}

// Reboot reboots the system.
//
// Example: Reboot(SHTDN_REASON_MINOR_MAINTENANCE, true)
func Reboot(reason Reason, force bool) error {
	if err := setPrivToken(); err != nil {
		return err
	}
	fl := uint32(EWX_REBOOT)
	if force {
		fl = fl | uint32(EWX_FORCE)
	}
	return windows.ExitWindowsEx(fl, uint32(reason))
}

// Shutdown shuts down the system.
//
// Example: Shutdown(SHTDN_REASON_MINOR_MAINTENANCE, true)
func Shutdown(reason Reason, force bool) error {

	if err := setPrivToken(); err != nil {
		return err
	}
	fl := uint32(EWX_SHUTDOWN)
	if force {
		fl = fl | uint32(EWX_FORCE)
	}

	return windows.ExitWindowsEx(fl, uint32(reason))
}
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"code.google.com/p/weed-fs/go/glog"
	"code.google.com/p/weed-fs/go/storage"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"
	"text/template"
	"time"
)

func init() {
	cmdExport.Run = runExport // break init cycle
	cmdExport.IsDebug = cmdExport.Flag.Bool("debug", false, "enable debug mode")
}

const (
	defaultFnFormat = `{{.Mime}}/{{.Id}}:{{.Name}}`
)

var cmdExport = &Command{
	UsageLine: "export -dir=/tmp -volumeId=234 -o=/dir/name.tar -fileNameFormat={{.Name}}",
	Short:     "list or export files from one volume data file",
	Long: `List all files in a volume, or Export all files in a volume to a tar file if the output is specified.

	The format of file name in the tar file can be customized. Default is {{.Mime}}/{{.Id}}:{{.Name}}. Also available is {{.Key}}.

  `,
}

var (
	exportVolumePath = cmdExport.Flag.String("dir", "/tmp", "input data directory to store volume data files")
	exportVolumeId   = cmdExport.Flag.Int("volumeId", -1, "a volume id. The volume should already exist in the dir. The volume index file should not exist.")
	dest             = cmdExport.Flag.String("o", "", "output tar file name, must ends with .tar, or just a \"-\" for stdout")
	format           = cmdExport.Flag.String("fileNameFormat", defaultFnFormat, "filename format, default to {{.Mime}}/{{.Id}}:{{.Name}}")
	tarFh            *tar.Writer
	tarHeader        tar.Header
	fnTmpl           *template.Template
	fnTmplBuf        = bytes.NewBuffer(nil)
)

func runExport(cmd *Command, args []string) bool {

	if *exportVolumeId == -1 {
		return false
	}

	var err error
	if *dest != "" {
		if *dest != "-" && !strings.HasSuffix(*dest, ".tar") {
			fmt.Println("the output file", *dest, "should be '-' or end with .tar")
			return false
		}

		if fnTmpl, err = template.New("name").Parse(*format); err != nil {
			fmt.Println("cannot parse format " + *format + ": " + err.Error())
			return false
		}

		var fh *os.File
		if *dest == "-" {
			fh = os.Stdout
		} else {
			if fh, err = os.Create(*dest); err != nil {
				glog.Fatalf("cannot open output tar %s: %s", *dest, err)
			}
		}
		defer fh.Close()
		tarFh = tar.NewWriter(fh)
		defer tarFh.Close()
		t := time.Now()
		tarHeader = tar.Header{Mode: 0644,
			ModTime: t, Uid: os.Getuid(), Gid: os.Getgid(),
			Typeflag:   tar.TypeReg,
			AccessTime: t, ChangeTime: t}
	}

	fileName := strconv.Itoa(*exportVolumeId)
	vid := storage.VolumeId(*exportVolumeId)
	indexFile, err := os.OpenFile(path.Join(*exportVolumePath, fileName+".idx"), os.O_RDONLY, 0644)
	if err != nil {
		glog.Fatalf("Create Volume Index [ERROR] %s\n", err)
	}
	defer indexFile.Close()

	nm, err := storage.LoadNeedleMap(indexFile)
	if err != nil {
		glog.Fatalf("cannot load needle map from %s: %s", indexFile, err)
	}

	var version storage.Version

	err = storage.ScanVolumeFile(*exportVolumePath, vid, func(superBlock storage.SuperBlock) error {
		version = superBlock.Version
		return nil
	}, func(n *storage.Needle, offset int64) error {
		debug("key", n.Id, "offset", offset, "size", n.Size, "disk_size", n.DiskSize(), "gzip", n.IsGzipped())
		nv, ok := nm.Get(n.Id)
		if ok && nv.Size > 0 {
			return walker(vid, n, version)
		} else {
			if !ok {
				debug("This seems deleted", n.Id)
			} else {
				debug("Id", n.Id, "size", n.Size)
			}
		}
		return nil
	})
	if err != nil {
		glog.Fatalf("Export Volume File [ERROR] %s\n", err)
	}
	return true
}

type nameParams struct {
	Name string
	Id   uint64
	Mime string
	Key  string
}

func walker(vid storage.VolumeId, n *storage.Needle, version storage.Version) (err error) {
	key := storage.NewFileId(vid, n.Id, n.Cookie).String()
	if tarFh != nil {
		fnTmplBuf.Reset()
		if err = fnTmpl.Execute(fnTmplBuf,
			nameParams{Name: string(n.Name),
				Id:   n.Id,
				Mime: string(n.Mime),
				Key:  key,
			},
		); err != nil {
			return err
		}
		nm := fnTmplBuf.String()

		if n.IsGzipped() && path.Ext(nm) != ".gz" {
			nm = nm + ".gz"
		}

		tarHeader.Name, tarHeader.Size = nm, int64(len(n.Data))
		if err = tarFh.WriteHeader(&tarHeader); err != nil {
			return err
		}
		_, err = tarFh.Write(n.Data)
	} else {
		size := n.DataSize
		if version == storage.Version1 {
			size = n.Size
		}
		fmt.Printf("key=%s Name=%s Size=%d gzip=%t mime=%s\n",
			key,
			n.Name,
			size,
			n.IsGzipped(),
			n.Mime,
		)
	}
	return
}
<commit_msg>debug message<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"code.google.com/p/weed-fs/go/glog"
	"code.google.com/p/weed-fs/go/storage"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"
	"text/template"
	"time"
)

func init() {
	cmdExport.Run = runExport // break init cycle
	cmdExport.IsDebug = cmdExport.Flag.Bool("debug", false, "enable debug mode")
}

const (
	defaultFnFormat = `{{.Mime}}/{{.Id}}:{{.Name}}`
)

var cmdExport = &Command{
	UsageLine: "export -dir=/tmp -volumeId=234 -o=/dir/name.tar -fileNameFormat={{.Name}}",
	Short:     "list or export files from one volume data file",
	Long: `List all files in a volume, or Export all files in a volume to a tar file if the output is specified.

	The format of file name in the tar file can be customized. Default is {{.Mime}}/{{.Id}}:{{.Name}}. Also available is {{.Key}}.

  `,
}

var (
	exportVolumePath = cmdExport.Flag.String("dir", "/tmp", "input data directory to store volume data files")
	exportVolumeId   = cmdExport.Flag.Int("volumeId", -1, "a volume id. The volume should already exist in the dir. The volume index file should not exist.")
	dest             = cmdExport.Flag.String("o", "", "output tar file name, must ends with .tar, or just a \"-\" for stdout")
	format           = cmdExport.Flag.String("fileNameFormat", defaultFnFormat, "filename format, default to {{.Mime}}/{{.Id}}:{{.Name}}")
	tarFh            *tar.Writer
	tarHeader        tar.Header
	fnTmpl           *template.Template
	fnTmplBuf        = bytes.NewBuffer(nil)
)

func runExport(cmd *Command, args []string) bool {

	if *exportVolumeId == -1 {
		return false
	}

	var err error
	if *dest != "" {
		if *dest != "-" && !strings.HasSuffix(*dest, ".tar") {
			fmt.Println("the output file", *dest, "should be '-' or end with .tar")
			return false
		}

		if fnTmpl, err = template.New("name").Parse(*format); err != nil {
			fmt.Println("cannot parse format " + *format + ": " + err.Error())
			return false
		}

		var fh *os.File
		if *dest == "-" {
			fh = os.Stdout
		} else {
			if fh, err = os.Create(*dest); err != nil {
				glog.Fatalf("cannot open output tar %s: %s", *dest, err)
			}
		}
		defer fh.Close()
		tarFh = tar.NewWriter(fh)
		defer tarFh.Close()
		t := time.Now()
		tarHeader = tar.Header{Mode: 0644,
			ModTime: t, Uid: os.Getuid(), Gid: os.Getgid(),
			Typeflag:   tar.TypeReg,
			AccessTime: t, ChangeTime: t}
	}

	fileName := strconv.Itoa(*exportVolumeId)
	vid := storage.VolumeId(*exportVolumeId)
	indexFile, err := os.OpenFile(path.Join(*exportVolumePath, fileName+".idx"), os.O_RDONLY, 0644)
	if err != nil {
		glog.Fatalf("Create Volume Index [ERROR] %s\n", err)
	}
	defer indexFile.Close()

	nm, err := storage.LoadNeedleMap(indexFile)
	if err != nil {
		glog.Fatalf("cannot load needle map from %s: %s", indexFile, err)
	}

	var version storage.Version

	err = storage.ScanVolumeFile(*exportVolumePath, vid, func(superBlock storage.SuperBlock) error {
		version = superBlock.Version
		return nil
	}, func(n *storage.Needle, offset int64) error {
		debug("key", n.Id, "offset", offset, "size", n.Size, "disk_size", n.DiskSize(), "gzip", n.IsGzipped())
		nv, ok := nm.Get(n.Id)
		if ok && nv.Size > 0 {
			return walker(vid, n, version)
		} else {
			if !ok {
				debug("This seems deleted", n.Id, "size", n.Size)
			} else {
				debug("Id", n.Id, "size", n.Size)
			}
		}
		return nil
	})
	if err != nil {
		glog.Fatalf("Export Volume File [ERROR] %s\n", err)
	}
	return true
}

type nameParams struct {
	Name string
	Id   uint64
	Mime string
	Key  string
}

func walker(vid storage.VolumeId, n *storage.Needle, version storage.Version) (err error) {
	key := storage.NewFileId(vid, n.Id, n.Cookie).String()
	if tarFh != nil {
		fnTmplBuf.Reset()
		if err = fnTmpl.Execute(fnTmplBuf,
			nameParams{Name: string(n.Name),
				Id:   n.Id,
				Mime: string(n.Mime),
				Key:  key,
			},
		); err != nil {
			return err
		}
		nm := fnTmplBuf.String()

		if n.IsGzipped() && path.Ext(nm) != ".gz" {
			nm = nm + ".gz"
		}

		tarHeader.Name, tarHeader.Size = nm, int64(len(n.Data))
		if err = tarFh.WriteHeader(&tarHeader); err != nil {
			return err
		}
		_, err = tarFh.Write(n.Data)
	} else {
		size := n.DataSize
		if version == storage.Version1 {
			size = n.Size
		}
		fmt.Printf("key=%s Name=%s Size=%d gzip=%t mime=%s\n",
			key,
			n.Name,
			size,
			n.IsGzipped(),
			n.Mime,
		)
	}
	return
}
<|endoftext|>


<commit_before>package goat

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"log"
	"net"
	"net/url"
	"strconv"
	"strings"
	"sync/atomic"
)

// Handshake for UDP tracker protocol
const udpInitID = 4497486125440

// UDP errors
var (
	// udpActionError is returned when a client requests an invalid tracker action
	udpActionError = errors.New("udp: client did not send a valid UDP tracker action")
	// udpHandshakeError is returned when a client does not send the proper handshake ID
	udpHandshakeError = errors.New("udp: client did not send proper UDP tracker handshake")
	// udpIntegerError is returned when a client sends an invalid integer parameter
	udpIntegerError = errors.New("udp: client sent an invalid integer parameter")
	// udpWriteError is returned when the tracker cannot generate a proper response
	udpWriteError = errors.New("udp: tracker cannot generate UDP tracker response")
)

// UDP address to connection ID map
var udpAddrToID map[string]uint64 = map[string]uint64{}

// Handle incoming UDP connections and return response
func handleUDP(l *net.UDPConn, sendChan chan bool, recvChan chan bool) {
	// Create shutdown function
	go func(l *net.UDPConn, sendChan chan bool, recvChan chan bool) {
		// Wait for done signal
		<-sendChan

		// Close listener
		if err := l.Close(); err != nil {
			log.Println(err.Error())
		}

		log.Println("UDP listener stopped")
		recvChan <- true
	}(l, sendChan, recvChan)

	// Loop and read connections
	for {
		buf := make([]byte, 2048)
		rlen, addr, err := l.ReadFromUDP(buf)

		// Count incoming connections
		atomic.AddInt64(&static.UDP.Current, 1)
		atomic.AddInt64(&static.UDP.Total, 1)

		// Triggered on graceful shutdown
		if err != nil {
			log.Println(err.Error())
			return
		}

		// Verify length is at least 16 bytes
		if rlen < 16 {
			log.Println("Invalid length")
			continue
		}

		// Spawn a goroutine to handle the connection and send back the response
		go func(l *net.UDPConn, buf []byte, addr *net.UDPAddr) {
			// Capture initial response from buffer
			res, err := parseUDP(buf, addr)
			if err != nil {
				// Client sent a malformed UDP handshake
				log.Println(err.Error())

				// If error, client did not handshake correctly, so boot them with error message
				_, err2 := l.WriteToUDP(res, addr)
				if err2 != nil {
					log.Println(err2.Error())
				}

				return
			}

			// Write response
			_, err = l.WriteToUDP(res, addr)
			if err != nil {
				log.Println(err.Error())
			}

			return
		}(l, buf, addr)
	}
}

// Parse a UDP byte buffer, return response from tracker
func parseUDP(buf []byte, addr *net.UDPAddr) ([]byte, error) {
	// Current connection ID (initially handshake, then generated by tracker)
	connID := binary.BigEndian.Uint64(buf[0:8])
	// Action integer (connect: 0, announce: 1)
	action := binary.BigEndian.Uint32(buf[8:12])
	// Transaction ID, to match between requests
	transID := buf[12:16]

	// Action switch
	switch action {
	// Connect
	case 0:
		// Validate UDP tracker handshake
		if connID != udpInitID {
			return udpTrackerError("Invalid UDP tracker handshake", transID), udpHandshakeError
		}

		res := bytes.NewBuffer(make([]byte, 0))

		// Action
		err := binary.Write(res, binary.BigEndian, uint32(0))
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Could not generate UDP tracker response", transID), udpWriteError
		}

		// Transaction ID
		err = binary.Write(res, binary.BigEndian, transID)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Could not generate UDP tracker response", transID), udpWriteError
		}

		// Generate a connection ID, which will be expected for this client next call
		expID := uint64(randRange(1, 1000000000))

		// Store this client's address and ID in map
		udpAddrToID[addr.String()] = expID

		// Connection ID, generated for this session
		err = binary.Write(res, binary.BigEndian, expID)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Could not generate UDP tracker response", transID), udpWriteError
		}

		return res.Bytes(), nil
	// Announce
	case 1:
		// Validate expected connection ID using map
		if connID != udpAddrToID[addr.String()] {
			return udpTrackerError("Invalid UDP connection ID", transID), udpHandshakeError
		}

		// Clear this IP from the connection map
		delete(udpAddrToID, addr.String())

		// Generate connection query
		query := url.Values{}

		// Mark client as UDP
		query.Set("udp", "1")

		// Transaction ID
		transID := buf[12:16]

		// Info hash
		query.Set("info_hash", string(buf[16:36]))

		// Skipped: peer_id: buf[36:56]

		// Downloaded
		t, err := strconv.ParseInt(hex.EncodeToString(buf[56:64]), 16, 64)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: downloaded", transID), udpIntegerError
		}
		query.Set("downloaded", strconv.FormatInt(t, 10))

		// Left
		t, err = strconv.ParseInt(hex.EncodeToString(buf[64:72]), 16, 64)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: left", transID), udpIntegerError
		}
		query.Set("left", strconv.FormatInt(t, 10))

		// Uploaded
		t, err = strconv.ParseInt(hex.EncodeToString(buf[72:80]), 16, 64)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: uploaded", transID), udpIntegerError
		}
		query.Set("uploaded", strconv.FormatInt(t, 10))

		// Event
		t, err = strconv.ParseInt(hex.EncodeToString(buf[80:84]), 16, 32)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: event", transID), udpIntegerError
		}
		event := strconv.FormatInt(t, 10)

		// Convert event to actual string
		switch event {
		case "0":
			query.Set("event", "")
		case "1":
			query.Set("event", "completed")
		case "2":
			query.Set("event", "started")
		case "3":
			query.Set("event", "stopped")
		}

		// IP address
		t, err = strconv.ParseInt(hex.EncodeToString(buf[84:88]), 16, 32)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: ip", transID), udpIntegerError
		}
		query.Set("ip", strconv.FormatInt(t, 10))

		// If no IP address set, use the UDP source
		if query.Get("ip") == "0" {
			query.Set("ip", strings.Split(addr.String(), ":")[0])
		}

		// Key
		query.Set("key", hex.EncodeToString(buf[88:92]))

		// Numwant
		query.Set("numwant", hex.EncodeToString(buf[92:96]))

		// If numwant is hex max value, default to 50
		if query.Get("numwant") == "ffffffff" {
			query.Set("numwant", "50")
		}

		// Port
		t, err = strconv.ParseInt(hex.EncodeToString(buf[96:98]), 16, 32)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: port", transID), udpIntegerError
		}
		query.Set("port", strconv.FormatInt(t, 10))

		// Trigger an anonymous announce
		return trackerAnnounce(userRecord{}, query, transID), nil
	default:
		return udpTrackerError("Invalid action", transID), udpActionError
	}
}
<commit_msg>Ignore connection closing error in UDP router<commit_after>package goat

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"log"
	"net"
	"net/url"
	"strconv"
	"strings"
	"sync/atomic"
)

// Handshake for UDP tracker protocol
const udpInitID = 4497486125440

// UDP errors
var (
	// udpActionError is returned when a client requests an invalid tracker action
	udpActionError = errors.New("udp: client did not send a valid UDP tracker action")
	// udpHandshakeError is returned when a client does not send the proper handshake ID
	udpHandshakeError = errors.New("udp: client did not send proper UDP tracker handshake")
	// udpIntegerError is returned when a client sends an invalid integer parameter
	udpIntegerError = errors.New("udp: client sent an invalid integer parameter")
	// udpWriteError is returned when the tracker cannot generate a proper response
	udpWriteError = errors.New("udp: tracker cannot generate UDP tracker response")
)

// UDP address to connection ID map
var udpAddrToID map[string]uint64 = map[string]uint64{}

// Handle incoming UDP connections and return response
func handleUDP(l *net.UDPConn, sendChan chan bool, recvChan chan bool) {
	// Create shutdown function
	go func(l *net.UDPConn, sendChan chan bool, recvChan chan bool) {
		// Wait for done signal
		<-sendChan

		// Close listener
		if err := l.Close(); err != nil {
			log.Println(err.Error())
		}

		log.Println("UDP listener stopped")
		recvChan <- true
	}(l, sendChan, recvChan)

	// Loop and read connections
	for {
		buf := make([]byte, 2048)
		rlen, addr, err := l.ReadFromUDP(buf)

		// Count incoming connections
		atomic.AddInt64(&static.UDP.Current, 1)
		atomic.AddInt64(&static.UDP.Total, 1)

		// Triggered on graceful shutdown
		if err != nil {
			// Ignore connection closing error, caused by stopping network listener
			if !strings.Contains(err.Error(), "use of closed network connection") {
				log.Println(err.Error())
				panic(err)
			}

			return
		}

		// Verify length is at least 16 bytes
		if rlen < 16 {
			log.Println("Invalid length")
			continue
		}

		// Spawn a goroutine to handle the connection and send back the response
		go func(l *net.UDPConn, buf []byte, addr *net.UDPAddr) {
			// Capture initial response from buffer
			res, err := parseUDP(buf, addr)
			if err != nil {
				// Client sent a malformed UDP handshake
				log.Println(err.Error())

				// If error, client did not handshake correctly, so boot them with error message
				_, err2 := l.WriteToUDP(res, addr)
				if err2 != nil {
					log.Println(err2.Error())
				}

				return
			}

			// Write response
			_, err = l.WriteToUDP(res, addr)
			if err != nil {
				log.Println(err.Error())
			}

			return
		}(l, buf, addr)
	}
}

// Parse a UDP byte buffer, return response from tracker
func parseUDP(buf []byte, addr *net.UDPAddr) ([]byte, error) {
	// Current connection ID (initially handshake, then generated by tracker)
	connID := binary.BigEndian.Uint64(buf[0:8])
	// Action integer (connect: 0, announce: 1)
	action := binary.BigEndian.Uint32(buf[8:12])
	// Transaction ID, to match between requests
	transID := buf[12:16]

	// Action switch
	switch action {
	// Connect
	case 0:
		// Validate UDP tracker handshake
		if connID != udpInitID {
			return udpTrackerError("Invalid UDP tracker handshake", transID), udpHandshakeError
		}

		res := bytes.NewBuffer(make([]byte, 0))

		// Action
		err := binary.Write(res, binary.BigEndian, uint32(0))
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Could not generate UDP tracker response", transID), udpWriteError
		}

		// Transaction ID
		err = binary.Write(res, binary.BigEndian, transID)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Could not generate UDP tracker response", transID), udpWriteError
		}

		// Generate a connection ID, which will be expected for this client next call
		expID := uint64(randRange(1, 1000000000))

		// Store this client's address and ID in map
		udpAddrToID[addr.String()] = expID

		// Connection ID, generated for this session
		err = binary.Write(res, binary.BigEndian, expID)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Could not generate UDP tracker response", transID), udpWriteError
		}

		return res.Bytes(), nil
	// Announce
	case 1:
		// Validate expected connection ID using map
		if connID != udpAddrToID[addr.String()] {
			return udpTrackerError("Invalid UDP connection ID", transID), udpHandshakeError
		}

		// Clear this IP from the connection map
		delete(udpAddrToID, addr.String())

		// Generate connection query
		query := url.Values{}

		// Mark client as UDP
		query.Set("udp", "1")

		// Transaction ID
		transID := buf[12:16]

		// Info hash
		query.Set("info_hash", string(buf[16:36]))

		// Skipped: peer_id: buf[36:56]

		// Downloaded
		t, err := strconv.ParseInt(hex.EncodeToString(buf[56:64]), 16, 64)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: downloaded", transID), udpIntegerError
		}
		query.Set("downloaded", strconv.FormatInt(t, 10))

		// Left
		t, err = strconv.ParseInt(hex.EncodeToString(buf[64:72]), 16, 64)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: left", transID), udpIntegerError
		}
		query.Set("left", strconv.FormatInt(t, 10))

		// Uploaded
		t, err = strconv.ParseInt(hex.EncodeToString(buf[72:80]), 16, 64)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: uploaded", transID), udpIntegerError
		}
		query.Set("uploaded", strconv.FormatInt(t, 10))

		// Event
		t, err = strconv.ParseInt(hex.EncodeToString(buf[80:84]), 16, 32)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: event", transID), udpIntegerError
		}
		event := strconv.FormatInt(t, 10)

		// Convert event to actual string
		switch event {
		case "0":
			query.Set("event", "")
		case "1":
			query.Set("event", "completed")
		case "2":
			query.Set("event", "started")
		case "3":
			query.Set("event", "stopped")
		}

		// IP address
		t, err = strconv.ParseInt(hex.EncodeToString(buf[84:88]), 16, 32)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: ip", transID), udpIntegerError
		}
		query.Set("ip", strconv.FormatInt(t, 10))

		// If no IP address set, use the UDP source
		if query.Get("ip") == "0" {
			query.Set("ip", strings.Split(addr.String(), ":")[0])
		}

		// Key
		query.Set("key", hex.EncodeToString(buf[88:92]))

		// Numwant
		query.Set("numwant", hex.EncodeToString(buf[92:96]))

		// If numwant is hex max value, default to 50
		if query.Get("numwant") == "ffffffff" {
			query.Set("numwant", "50")
		}

		// Port
		t, err = strconv.ParseInt(hex.EncodeToString(buf[96:98]), 16, 32)
		if err != nil {
			log.Println(err.Error())
			return udpTrackerError("Invalid integer parameter: port", transID), udpIntegerError
		}
		query.Set("port", strconv.FormatInt(t, 10))

		// Trigger an anonymous announce
		return trackerAnnounce(userRecord{}, query, transID), nil
	default:
		return udpTrackerError("Invalid action", transID), udpActionError
	}
}
<|endoftext|>


<commit_before>// Command gracedemo implements a demo server showing how to gracefully
// terminate an HTTP server using finmonitor and goagain.
package main

import (
	"flag"
	"fmt"
	"github.com/daaku/go.grace/gracehttp"
	"net/http"
	"time"
)

var (
	address0 = flag.String("a0", ":48567", "Zero address to bind to.")
	address1 = flag.String("a1", ":48568", "First address to bind to.")
	address2 = flag.String("a2", ":48569", "Second address to bind to.")
	now      = time.Now()
)

func main() {
	flag.Parse()
	gracehttp.Serve(
		gracehttp.Handler{*address0, newHandler("Zero  ")},
		gracehttp.Handler{*address1, newHandler("First ")},
		gracehttp.Handler{*address2, newHandler("Second")},
	)
}

func newHandler(name string) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/sleep/", func(w http.ResponseWriter, r *http.Request) {
		duration, err := time.ParseDuration(r.FormValue("duration"))
		if err != nil {
			http.Error(w, err.Error(), 400)
		}
		time.Sleep(duration)
		w.Write([]byte(fmt.Sprintf(
			"%s started at %s slept for %d nanoseconds.\n",
			name,
			now,
			duration.Nanoseconds())))
	})
	return mux
}
<commit_msg>demo code smell<commit_after>// Command gracedemo implements a demo server showing how to gracefully
// terminate an HTTP server using finmonitor and goagain.
package main

import (
	"flag"
	"fmt"
	"github.com/daaku/go.grace/gracehttp"
	"net/http"
	"time"
)

var (
	address0 = flag.String("a0", ":48567", "Zero address to bind to.")
	address1 = flag.String("a1", ":48568", "First address to bind to.")
	address2 = flag.String("a2", ":48569", "Second address to bind to.")
	now      = time.Now()
)

func main() {
	flag.Parse()
	gracehttp.Serve(
		gracehttp.Handler{*address0, newHandler("Zero  ")},
		gracehttp.Handler{*address1, newHandler("First ")},
		gracehttp.Handler{*address2, newHandler("Second")},
	)
}

func newHandler(name string) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/sleep/", func(w http.ResponseWriter, r *http.Request) {
		duration, err := time.ParseDuration(r.FormValue("duration"))
		if err != nil {
			http.Error(w, err.Error(), 400)
		}
		time.Sleep(duration)
		fmt.Fprintf(
			w,
			"%s started at %s slept for %d nanoseconds.\n",
			name,
			now,
			duration.Nanoseconds(),
		)
	})
	return mux
}
<|endoftext|>


<commit_before>package qlang

import (
	"qlang.io/exec.v2"
	"qiniupkg.com/text/tpl.v1/interpreter.util"
)

// -----------------------------------------------------------------------------

func (p *Compiler) Clear() {

	p.code.Block(exec.Clear)
}

func (p *Compiler) Unset(name string) {

	p.code.Block(exec.Unset(name))
}

func (p *Compiler) Inc(name string) {

	p.code.Block(exec.Inc(name))
}

func (p *Compiler) Dec(name string) {

	p.code.Block(exec.Dec(name))
}

func (p *Compiler) AddAssign(name string) {

	p.code.Block(exec.AddAssign(name))
}

func (p *Compiler) SubAssign(name string) {

	p.code.Block(exec.SubAssign(name))
}

func (p *Compiler) MulAssign(name string) {

	p.code.Block(exec.MulAssign(name))
}

func (p *Compiler) QuoAssign(name string) {

	p.code.Block(exec.QuoAssign(name))
}

func (p *Compiler) ModAssign(name string) {

	p.code.Block(exec.ModAssign(name))
}

func (p *Compiler) MultiAssign() {

	arity := p.popArity() + 1
	names := p.gstk.PopFnArgs(arity)
	p.code.Block(exec.MultiAssign(names))
}

func (p *Compiler) Assign(name string) {

	p.code.Block(exec.Assign(name))
}

func (p *Compiler) Ref(name string) {

	if val, ok := p.gvars[name]; ok {
		p.code.Block(exec.Push(val))
	} else {
		p.code.Block(exec.Ref(name))
	}
}

func (p *Compiler) Function(e interpreter.Engine) {

	fnb, _ := p.gstk.Pop()
	variadic := p.popArity()
	arity := p.popArity()
	args := p.gstk.PopFnArgs(arity)
	instr := p.code.Reserve()
	p.exits = append(p.exits, func() {
		start, end := p.cl(e, "doc", fnb)
		instr.Set(exec.Func(nil, start, end, args, variadic != 0))
	})
}

func (p *Compiler) Main(e interpreter.Engine) {

	fnb, _ := p.gstk.Pop()
	instr := p.code.Reserve()
	p.exits = append(p.exits, func() {
		start, end := p.cl(e, "doc", fnb)
		instr.Set(exec.Func(nil, start, end, nil, false))
	})
	p.code.Block(
		exec.CallFn(0),
		exec.Exit,
	)
}

func (p *Compiler) Include(file string) {

	code := p.code
	instr := code.Reserve()
	p.exits = append(p.exits, func() {
		start := code.Len()
		end := p.Incl(file)
		instr.Set(exec.Block(start, end))
	})
}

func (p *Compiler) Return(e interpreter.Engine) {

	arity := p.popArity()
	p.code.Block(exec.Return(arity))
}

func (p *Compiler) Done() {

	for {
		n := len(p.exits)
		if n == 0 {
			break
		}
		onExit := p.exits[n-1]
		p.exits = p.exits[:n-1]
		onExit()
	}
}

func (p *Compiler) Defer(e interpreter.Engine) {

	src, _ := p.gstk.Pop()
	instr := p.code.Reserve()
	p.exits = append(p.exits, func() {
		start, end := p.cl(e, "expr", src)
		instr.Set(exec.Defer(start, end))
	})
}

func (p *Compiler) Recover() {

	p.code.Block(exec.Recover)
}

func (p *Compiler) cl(e interpreter.Engine, g string, src interface{}) (start, end int) {

	start = p.code.Len()
	if src != nil {
		if err := e.EvalCode(p, g, src); err != nil {
			panic(err)
		}
	}
	end = p.code.Len()
	return
}

// -----------------------------------------------------------------------------
<commit_msg>include bugfix: unquote literal<commit_after>package qlang

import (
	"strconv"

	"qlang.io/exec.v2"
	"qiniupkg.com/text/tpl.v1/interpreter.util"
)

// -----------------------------------------------------------------------------

func (p *Compiler) Clear() {

	p.code.Block(exec.Clear)
}

func (p *Compiler) Unset(name string) {

	p.code.Block(exec.Unset(name))
}

func (p *Compiler) Inc(name string) {

	p.code.Block(exec.Inc(name))
}

func (p *Compiler) Dec(name string) {

	p.code.Block(exec.Dec(name))
}

func (p *Compiler) AddAssign(name string) {

	p.code.Block(exec.AddAssign(name))
}

func (p *Compiler) SubAssign(name string) {

	p.code.Block(exec.SubAssign(name))
}

func (p *Compiler) MulAssign(name string) {

	p.code.Block(exec.MulAssign(name))
}

func (p *Compiler) QuoAssign(name string) {

	p.code.Block(exec.QuoAssign(name))
}

func (p *Compiler) ModAssign(name string) {

	p.code.Block(exec.ModAssign(name))
}

func (p *Compiler) MultiAssign() {

	arity := p.popArity() + 1
	names := p.gstk.PopFnArgs(arity)
	p.code.Block(exec.MultiAssign(names))
}

func (p *Compiler) Assign(name string) {

	p.code.Block(exec.Assign(name))
}

func (p *Compiler) Ref(name string) {

	if val, ok := p.gvars[name]; ok {
		p.code.Block(exec.Push(val))
	} else {
		p.code.Block(exec.Ref(name))
	}
}

func (p *Compiler) Function(e interpreter.Engine) {

	fnb, _ := p.gstk.Pop()
	variadic := p.popArity()
	arity := p.popArity()
	args := p.gstk.PopFnArgs(arity)
	instr := p.code.Reserve()
	p.exits = append(p.exits, func() {
		start, end := p.cl(e, "doc", fnb)
		instr.Set(exec.Func(nil, start, end, args, variadic != 0))
	})
}

func (p *Compiler) Main(e interpreter.Engine) {

	fnb, _ := p.gstk.Pop()
	instr := p.code.Reserve()
	p.exits = append(p.exits, func() {
		start, end := p.cl(e, "doc", fnb)
		instr.Set(exec.Func(nil, start, end, nil, false))
	})
	p.code.Block(
		exec.CallFn(0),
		exec.Exit,
	)
}

func (p *Compiler) Include(lit string) {

	file, err := strconv.Unquote(lit)
	if err != nil {
		panic("invalid string `" + lit + "`: " + err.Error())
	}

	code := p.code
	instr := code.Reserve()
	p.exits = append(p.exits, func() {
		start := code.Len()
		end := p.Incl(file)
		instr.Set(exec.Block(start, end))
	})
}

func (p *Compiler) Return(e interpreter.Engine) {

	arity := p.popArity()
	p.code.Block(exec.Return(arity))
}

func (p *Compiler) Done() {

	for {
		n := len(p.exits)
		if n == 0 {
			break
		}
		onExit := p.exits[n-1]
		p.exits = p.exits[:n-1]
		onExit()
	}
}

func (p *Compiler) Defer(e interpreter.Engine) {

	src, _ := p.gstk.Pop()
	instr := p.code.Reserve()
	p.exits = append(p.exits, func() {
		start, end := p.cl(e, "expr", src)
		instr.Set(exec.Defer(start, end))
	})
}

func (p *Compiler) Recover() {

	p.code.Block(exec.Recover)
}

func (p *Compiler) cl(e interpreter.Engine, g string, src interface{}) (start, end int) {

	start = p.code.Len()
	if src != nil {
		if err := e.EvalCode(p, g, src); err != nil {
			panic(err)
		}
	}
	end = p.code.Len()
	return
}

// -----------------------------------------------------------------------------
<|endoftext|>


<commit_before>package spscq

import (
	"sync/atomic"

	"github.com/fmstephe/flib/fmath"
	"github.com/fmstephe/flib/fsync/fatomic"
	"github.com/fmstephe/flib/fsync/padded"
)

type ByteQ struct {
	_prebuffer padded.CacheBuffer
	commonQ
	_midbuffer  padded.CacheBuffer
	ringBuffer  []byte
	_postbuffer padded.CacheBuffer
}

func NewByteQ(size int64) *ByteQ {
	ringBuffer := padded.ByteSlice(int(size))
	q := &ByteQ{ringBuffer: ringBuffer, commonQ: newCommonQ(size)}
	return q
}

func (q *ByteQ) WriteBuffer(bufferSize int64) []byte {
	write := q.write.Value
	idx := write & q.mask
	bufferSize = fmath.Min(bufferSize, q.size-idx)
	writeTo := write + bufferSize
	readLimit := writeTo - q.size
	nxt := idx + bufferSize
	if readLimit > q.readCache.Value {
		q.readCache.Value = atomic.LoadInt64(&q.read.Value)
		if readLimit > q.readCache.Value {
			nxt = q.readCache.Value & q.mask
		}
	}
	if idx == nxt {
		q.failedWrites.Value++
	}
	q.writeSize.Value = nxt - idx
	return q.ringBuffer[idx:nxt]
}

func (q *ByteQ) CommitWrite() {
	atomic.AddInt64(&q.write.Value, q.writeSize.Value)
	q.writeSize.Value = 0
}

func (q *ByteQ) CommitWriteLazy() {
	fatomic.LazyStore(&q.write.Value, q.write.Value+q.writeSize.Value)
	q.writeSize.Value = 0
}

func (q *ByteQ) ReadBuffer(bufferSize int64) []byte {
	read := q.read.Value
	idx := read & q.mask
	bufferSize = fmath.Min(bufferSize, q.size-idx)
	readTo := read + bufferSize
	nxt := idx + bufferSize
	if readTo > q.writeCache.Value {
		q.writeCache.Value = atomic.LoadInt64(&q.write.Value)
		if readTo > q.writeCache.Value {
			nxt = q.writeCache.Value & q.mask
		}
	}
	if idx == nxt {
		q.failedReads.Value++
	}
	q.readSize.Value = nxt - idx
	return q.ringBuffer[idx:nxt]
}

func (q *ByteQ) CommitRead() {
	atomic.AddInt64(&q.read.Value, q.readSize.Value)
	q.readSize.Value = 0
}

func (q *ByteQ) CommitReadLazy() {
	fatomic.LazyStore(&q.read.Value, q.read.Value+q.readSize.Value)
	q.readSize.Value = 0
}
<commit_msg>ByteQ now using commonQ.*Buffer/Commit*() methods<commit_after>package spscq

import (
	"github.com/fmstephe/flib/fsync/padded"
)

type ByteQ struct {
	_prebuffer padded.CacheBuffer
	commonQ
	_midbuffer  padded.CacheBuffer
	ringBuffer  []byte
	_postbuffer padded.CacheBuffer
}

func NewByteQ(size int64) *ByteQ {
	ringBuffer := padded.ByteSlice(int(size))
	q := &ByteQ{ringBuffer: ringBuffer, commonQ: newCommonQ(size)}
	return q
}

func (q *ByteQ) WriteBuffer(bufferSize int64) []byte {
	from, to := q.writeBuffer(bufferSize)
	return q.ringBuffer[from:to]
}

func (q *ByteQ) ReadBuffer(bufferSize int64) []byte {
	from, to := q.readBuffer(bufferSize)
	return q.ringBuffer[from:to]
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	lua "github.com/anaminus/gopher-lua"
	"github.com/anaminus/rbxmk"
	"github.com/anaminus/rbxmk/rtypes"
	"github.com/anaminus/snek"
	"github.com/robloxapi/types"
)

const testdata = "testdata"

// Replace scriptArguments[x] with test script.
const replaceIndex = 3

var scriptArguments = [...]string{
	"rbxmk_test",
	"run",
	"--debug",
	"-",
	"true",
	"false",
	"nil",
	"42",
	"3.141592653589793",
	"-1e-8",
	"hello, world!",
	"hello\000world!",
}

type dummyFile struct {
	r    io.Reader
	info *dummyInfo
}

func (d *dummyFile) Name() string               { return "test" }
func (d *dummyFile) Stat() (os.FileInfo, error) { return d.info, nil }
func (d *dummyFile) Read(b []byte) (int, error) { return d.r.Read(b) }
func (d *dummyFile) Write([]byte) (int, error)  { return 0, nil }

type dummyInfo struct {
	name  string
	size  int64
	mode  os.FileMode
	time  time.Time
	isdir bool
}

func (d *dummyInfo) Name() string       { return d.name }
func (d *dummyInfo) Size() int64        { return d.size }
func (d *dummyInfo) Mode() os.FileMode  { return d.mode }
func (d *dummyInfo) ModTime() time.Time { return d.time }
func (d *dummyInfo) IsDir() bool        { return d.isdir }
func (d *dummyInfo) Sys() interface{}   { return d }

func initMain(s rbxmk.State, t *testing.T) {
	T := s.L.CreateTable(0, 2)

	// Pass makes a positive assertion. If the first argument is a non-function,
	// then an error is emitted if the value is falsy. Otherwise, the function
	// is called. If the call errors or returns a falsy value, then an error is
	// emitted. Returning no value counts as truthy. The error is emitted to
	// testing.T, but does not cause a Lua error to be thrown. The second
	// optional argument is a string that describes the assertion, which is
	// included with an emitted error.
	T.RawSetString("Pass", s.L.NewFunction(func(l *lua.LState) int {
		v := l.CheckAny(1)
		msg := l.OptString(2, "expected pass")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err != nil {
				t.Errorf("%s: %s", msg, err.Error())
				return 0
			}
			if s.Count() > n {
				if !s.L.ToBool(n + 1) {
					if m := s.L.ToString(n + 2); m != "" {
						t.Errorf("%s: %s", msg, m)
					} else {
						t.Errorf(msg)
					}
					return 0
				}
			}
		default:
			if !lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// Fail makes a negative assertion. If the first argument is a non-function,
	// then an error is emitted if the value is truthy. Otherwise, the function
	// is called. If the call does not error or returns a truthy value, then an
	// error is emitted. Returning no value counts as falsy. The error is
	// emitted to testing.T, but does not cause a Lua error to be thrown. The
	// second optional argument is a string that describes the assertion, which
	// is included with an emitted error.
	T.RawSetString("Fail", s.L.NewFunction(func(l *lua.LState) int {
		v := l.CheckAny(1)
		msg := l.OptString(2, "expected fail")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err == nil {
				if s.Count() > n {
					if s.L.ToBool(n + 1) {
						if m := s.L.ToString(n + 2); m != "" {
							t.Errorf("%s: %s", msg, m)
						} else {
							t.Errorf(msg)
						}
					}
					return 0
				}
				t.Errorf(msg)
				return 0
			} else if lua.LVAsBool(s.L.GetGlobal("SHOW_ERRORS")) {
				t.Logf("ERROR: %s\n", err)
			}
		default:
			if lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// GC runs the garbage collector.
	T.RawSetString("GC", s.L.NewFunction(func(l *lua.LState) int {
		runtime.GC()
		return 0
	}))

	// UserDataCacheLen returns the number of cached userdata values.
	T.RawSetString("UserDataCacheLen", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.Int(s.UserDataCacheLen()))
	}))

	T.RawSetString("DummySymbol", s.UserDataOf(rtypes.Symbol{Name: "DummySymbol"}, "Symbol"))

	s.L.SetGlobal("T", T)
}

// TestScripts runs each .lua file in testdata as a Lua script. If the first
// line starts with a comment that contains "fail", then the script is expected
// to throw an error. All scripts receive the arguments from scriptArguments.
func TestScripts(t *testing.T) {
	program := snek.NewProgram("", scriptArguments[:])
	program.Register(snek.Def{
		Name: "run",
		New: func() snek.Command {
			return &RunCommand{Init: func(s rbxmk.State) { initMain(s, t) }}
		},
	})

	var files []string
	wd, _ := os.Getwd()
	for _, arg := range os.Args[2:] {
		if strings.HasPrefix(arg, "-test.") {
			continue
		}
		rel, err := filepath.Rel(wd, arg)
		if err != nil {
			rel = arg
		}
		files = append(files, rel)
	}
	if len(files) == 0 {
		err := filepath.Walk(testdata, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() &&
				filepath.Ext(info.Name()) == ".lua" &&
				!strings.HasPrefix(filepath.Base(info.Name()), "_") {
				files = append(files, path)
			}
			return nil
		})
		if err != nil {
			t.Fatalf("error walking testdata: %s", err)
		}
	}
	for _, file := range files {
		t.Run(filepath.ToSlash(file), func(t *testing.T) {
			args := make([]string, len(scriptArguments))
			copy(args, scriptArguments[:])
			args[replaceIndex] = file
			err := program.RunWithInput("run", snek.Input{
				Program:   args[0],
				Arguments: args[2:],
			})
			if err != nil {
				t.Errorf("script %s: %s", file, err)
			}
		})
	}
}
<commit_msg>Add UserDataID test function.<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	lua "github.com/anaminus/gopher-lua"
	"github.com/anaminus/rbxmk"
	"github.com/anaminus/rbxmk/rtypes"
	"github.com/anaminus/snek"
	"github.com/robloxapi/types"
)

const testdata = "testdata"

// Replace scriptArguments[x] with test script.
const replaceIndex = 3

var scriptArguments = [...]string{
	"rbxmk_test",
	"run",
	"--debug",
	"-",
	"true",
	"false",
	"nil",
	"42",
	"3.141592653589793",
	"-1e-8",
	"hello, world!",
	"hello\000world!",
}

type dummyFile struct {
	r    io.Reader
	info *dummyInfo
}

func (d *dummyFile) Name() string               { return "test" }
func (d *dummyFile) Stat() (os.FileInfo, error) { return d.info, nil }
func (d *dummyFile) Read(b []byte) (int, error) { return d.r.Read(b) }
func (d *dummyFile) Write([]byte) (int, error)  { return 0, nil }

type dummyInfo struct {
	name  string
	size  int64
	mode  os.FileMode
	time  time.Time
	isdir bool
}

func (d *dummyInfo) Name() string       { return d.name }
func (d *dummyInfo) Size() int64        { return d.size }
func (d *dummyInfo) Mode() os.FileMode  { return d.mode }
func (d *dummyInfo) ModTime() time.Time { return d.time }
func (d *dummyInfo) IsDir() bool        { return d.isdir }
func (d *dummyInfo) Sys() interface{}   { return d }

func initMain(s rbxmk.State, t *testing.T) {
	T := s.L.CreateTable(0, 2)

	// Pass makes a positive assertion. If the first argument is a non-function,
	// then an error is emitted if the value is falsy. Otherwise, the function
	// is called. If the call errors or returns a falsy value, then an error is
	// emitted. Returning no value counts as truthy. The error is emitted to
	// testing.T, but does not cause a Lua error to be thrown. The second
	// optional argument is a string that describes the assertion, which is
	// included with an emitted error.
	T.RawSetString("Pass", s.L.NewFunction(func(l *lua.LState) int {
		v := l.CheckAny(1)
		msg := l.OptString(2, "expected pass")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err != nil {
				t.Errorf("%s: %s", msg, err.Error())
				return 0
			}
			if s.Count() > n {
				if !s.L.ToBool(n + 1) {
					if m := s.L.ToString(n + 2); m != "" {
						t.Errorf("%s: %s", msg, m)
					} else {
						t.Errorf(msg)
					}
					return 0
				}
			}
		default:
			if !lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// Fail makes a negative assertion. If the first argument is a non-function,
	// then an error is emitted if the value is truthy. Otherwise, the function
	// is called. If the call does not error or returns a truthy value, then an
	// error is emitted. Returning no value counts as falsy. The error is
	// emitted to testing.T, but does not cause a Lua error to be thrown. The
	// second optional argument is a string that describes the assertion, which
	// is included with an emitted error.
	T.RawSetString("Fail", s.L.NewFunction(func(l *lua.LState) int {
		v := l.CheckAny(1)
		msg := l.OptString(2, "expected fail")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err == nil {
				if s.Count() > n {
					if s.L.ToBool(n + 1) {
						if m := s.L.ToString(n + 2); m != "" {
							t.Errorf("%s: %s", msg, m)
						} else {
							t.Errorf(msg)
						}
					}
					return 0
				}
				t.Errorf(msg)
				return 0
			} else if lua.LVAsBool(s.L.GetGlobal("SHOW_ERRORS")) {
				t.Logf("ERROR: %s\n", err)
			}
		default:
			if lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// GC runs the garbage collector.
	T.RawSetString("GC", s.L.NewFunction(func(l *lua.LState) int {
		runtime.GC()
		return 0
	}))

	// UserDataCacheLen returns the number of cached userdata values.
	T.RawSetString("UserDataCacheLen", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.Int(s.UserDataCacheLen()))
	}))

	// UserDataID returns a string that identifies a userdata value.
	T.RawSetString("UserDataID", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.String(fmt.Sprintf("%p", s.CheckUserData(1))))
	}))

	T.RawSetString("DummySymbol", s.UserDataOf(rtypes.Symbol{Name: "DummySymbol"}, "Symbol"))

	s.L.SetGlobal("T", T)
}

// TestScripts runs each .lua file in testdata as a Lua script. If the first
// line starts with a comment that contains "fail", then the script is expected
// to throw an error. All scripts receive the arguments from scriptArguments.
func TestScripts(t *testing.T) {
	program := snek.NewProgram("", scriptArguments[:])
	program.Register(snek.Def{
		Name: "run",
		New: func() snek.Command {
			return &RunCommand{Init: func(s rbxmk.State) { initMain(s, t) }}
		},
	})

	var files []string
	wd, _ := os.Getwd()
	for _, arg := range os.Args[2:] {
		if strings.HasPrefix(arg, "-test.") {
			continue
		}
		rel, err := filepath.Rel(wd, arg)
		if err != nil {
			rel = arg
		}
		files = append(files, rel)
	}
	if len(files) == 0 {
		err := filepath.Walk(testdata, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() &&
				filepath.Ext(info.Name()) == ".lua" &&
				!strings.HasPrefix(filepath.Base(info.Name()), "_") {
				files = append(files, path)
			}
			return nil
		})
		if err != nil {
			t.Fatalf("error walking testdata: %s", err)
		}
	}
	for _, file := range files {
		t.Run(filepath.ToSlash(file), func(t *testing.T) {
			args := make([]string, len(scriptArguments))
			copy(args, scriptArguments[:])
			args[replaceIndex] = file
			err := program.RunWithInput("run", snek.Input{
				Program:   args[0],
				Arguments: args[2:],
			})
			if err != nil {
				t.Errorf("script %s: %s", file, err)
			}
		})
	}
}
<|endoftext|>


<commit_before>package grifts

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/markbates/grift/grift"
	"github.com/pkg/errors"
)

var _ = grift.Desc("release", "Generates a CHANGELOG and creates a new GitHub release based on what is in the version.go file.")
var _ = grift.Add("release", func(c *grift.Context) error {
	v, err := findVersion()
	if err != nil {
		return err
	}

	err = grift.Run("shoulders", c)
	if err != nil {
		return err
	}

	err = grift.Run("deplist", c)
	if err != nil {
		return err
	}

	err = installBin()
	if err != nil {
		return err
	}

	err = localTest()
	if err != nil {
		return err
	}

	err = dockerTest()
	if err != nil {
		return err
	}

	err = tagRelease(v)
	if err != nil {
		return err
	}

	err = runChangelogGenerator(v)
	if err != nil {
		return err
	}

	return commitAndPush(v)
})

func installBin() error {
	cmd := exec.Command("go", "install", "-v", "./buffalo")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func localTest() error {
	cmd := exec.Command("go", "test", "-v", "-race", "./...")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func dockerTest() error {
	cmd := exec.Command("docker", "build", ".")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func tagRelease(v string) error {
	token := os.Getenv("GITHUB_TOKEN")
	if token == "" {
		return errors.New("GITHUB_TOKEN is not set")
	}

	body := map[string]interface{}{
		"tag_name":   fmt.Sprintf("v%s", v),
		"prerelease": false,
	}

	b, err := json.Marshal(&body)
	if err != nil {
		return err
	}

	res, err := http.Post(fmt.Sprintf("https://api.github.com/repos/gobuffalo/buffalo/releases?access_token=%s", token), "application/json", bytes.NewReader(b))
	if err != nil {
		return err
	}

	code := res.StatusCode
	if code < 200 || code >= 300 {
		return fmt.Errorf("got a not successful status code from github! %d", code)
	}

	return nil
}

func runChangelogGenerator(v string) error {
	cmd := exec.Command("github_changelog_generator")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func commitAndPush(v string) error {
	cmd := exec.Command("git", "commit", "CHANGELOG.md", "-m", fmt.Sprintf("Updated changelog for release v%s", v))
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	err := cmd.Run()
	if err != nil {
		return err
	}

	cmd = exec.Command("git", "push", "origin", "master")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func findVersion() (string, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	vfile, err := ioutil.ReadFile(filepath.Join(pwd, "buffalo/cmd/version.go"))
	if err != nil {
		return "", err
	}

	//var Version = "0.4.0"
	re := regexp.MustCompile(`const Version = "(.+)"`)
	matches := re.FindStringSubmatch(string(vfile))
	if len(matches) < 2 {
		return "", errors.New("failed to find the version")
	}
	v := matches[1]
	if strings.Contains(v, "dev") {
		return "", errors.Errorf("version can not be a dev version %s", v)
	}
	return v, nil
}
<commit_msg>re-ordered the release task a bit<commit_after>package grifts

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/markbates/grift/grift"
	"github.com/pkg/errors"
)

var _ = grift.Desc("release", "Generates a CHANGELOG and creates a new GitHub release based on what is in the version.go file.")
var _ = grift.Add("release", func(c *grift.Context) error {
	v, err := findVersion()
	if err != nil {
		return err
	}

	err = installBin()
	if err != nil {
		return err
	}

	err = localTest()
	if err != nil {
		return err
	}

	err = dockerTest()
	if err != nil {
		return err
	}

	grift.Run("shoulders", c)
	grift.Run("deplist", c)

	err = tagRelease(v)
	if err != nil {
		return err
	}

	err = runChangelogGenerator(v)
	if err != nil {
		return err
	}

	return commitAndPush(v)
})

func installBin() error {
	cmd := exec.Command("go", "install", "-v", "./buffalo")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func localTest() error {
	cmd := exec.Command("go", "test", "-v", "-race", "./...")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func dockerTest() error {
	cmd := exec.Command("docker", "build", ".")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func tagRelease(v string) error {
	token := os.Getenv("GITHUB_TOKEN")
	if token == "" {
		return errors.New("GITHUB_TOKEN is not set")
	}

	body := map[string]interface{}{
		"tag_name":   fmt.Sprintf("v%s", v),
		"prerelease": false,
	}

	b, err := json.Marshal(&body)
	if err != nil {
		return err
	}

	res, err := http.Post(fmt.Sprintf("https://api.github.com/repos/gobuffalo/buffalo/releases?access_token=%s", token), "application/json", bytes.NewReader(b))
	if err != nil {
		return err
	}

	code := res.StatusCode
	if code < 200 || code >= 300 {
		return fmt.Errorf("got a not successful status code from github! %d", code)
	}

	return nil
}

func runChangelogGenerator(v string) error {
	cmd := exec.Command("github_changelog_generator")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func commitAndPush(v string) error {
	cmd := exec.Command("git", "commit", "CHANGELOG.md", "-m", fmt.Sprintf("Updated changelog for release v%s", v))
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	err := cmd.Run()
	if err != nil {
		return err
	}

	cmd = exec.Command("git", "push", "origin", "master")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func findVersion() (string, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	vfile, err := ioutil.ReadFile(filepath.Join(pwd, "buffalo/cmd/version.go"))
	if err != nil {
		return "", err
	}

	//var Version = "0.4.0"
	re := regexp.MustCompile(`const Version = "(.+)"`)
	matches := re.FindStringSubmatch(string(vfile))
	if len(matches) < 2 {
		return "", errors.New("failed to find the version")
	}
	v := matches[1]
	if strings.Contains(v, "dev") {
		return "", errors.Errorf("version can not be a dev version %s", v)
	}
	return v, nil
}
<|endoftext|>


<commit_before>package handler

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"strings"

	"github.com/oursky/ourd/oddb"
	"github.com/oursky/ourd/oderr"
	"github.com/oursky/ourd/router"
)

// recordPayload is the input parameter in RecordHandler
type recordPayload router.Payload

func (p *recordPayload) IsWriteAllowed() bool {
	return p.UserInfo == nil
}

// transportRecord override JSON serialization and deserialization of
// oddb.Record
type transportRecord oddb.Record

func (r transportRecord) ID() string {
	return r.Type + "/" + r.Key
}

func (r transportRecord) MarshalJSON() ([]byte, error) {
	// NOTE(limouren): marshalling of type/key is delegated to responseItem
	return json.Marshal(r.Data)
}

func (r *transportRecord) UnmarshalJSON(data []byte) error {
	object := map[string]interface{}{}
	err := json.Unmarshal(data, &object)

	if err != nil {
		return err
	}

	return r.InitFromMap(object)
}

func (r *transportRecord) InitFromMap(m map[string]interface{}) error {
	rawID, ok := m["_id"].(string)
	if !ok {
		return errors.New(`record/json: required field "_id" not found`)
	}
	delete(m, "_id")

	ss := strings.SplitN(rawID, "/", 2)
	if len(ss) == 1 {
		return fmt.Errorf(`record/json: "_id" should be of format '{type}/{id}', got %#v`, rawID)
	}

	recordType, id := ss[0], ss[1]

	// remove reserved keys
	for key := range m {
		if key[0] == '_' {
			delete(m, key)
		}
	}

	r.Key = id
	r.Type = recordType
	r.Data = m

	return nil
}

type responseItem struct {
	id     string
	record *transportRecord
	err    oderr.Error
}

func newResponseItem(record *transportRecord) responseItem {
	return responseItem{
		id:     record.ID(),
		record: record,
	}
}

func newResponseItemErr(id string, err oderr.Error) responseItem {
	return responseItem{
		id:  id,
		err: err,
	}
}

func (item responseItem) MarshalJSON() ([]byte, error) {
	var (
		buf bytes.Buffer
		i   interface{}
	)
	buf.Write([]byte(`{"_id":"`))
	buf.WriteString(item.id)
	buf.Write([]byte(`","_type":"`))
	if item.err != nil {
		buf.Write([]byte(`error",`))
		i = item.err
	} else if item.record != nil {
		buf.Write([]byte(`record",`))
		i = item.record
	} else {
		panic("inconsistent state: both err and record is nil")
	}

	bodyBytes, err := json.Marshal(i)
	if err != nil {
		return nil, err
	}

	if bodyBytes[0] != '{' {
		return nil, fmt.Errorf("first char of embedded json != {: %v", string(bodyBytes))
	} else if bodyBytes[len(bodyBytes)-1] != '}' {
		return nil, fmt.Errorf("last char of embedded json != }: %v", string(bodyBytes))
	}
	buf.Write(bodyBytes[1:])
	return buf.Bytes(), nil
}

/*
RecordSaveHandler is dummy implementation on save/modify Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://192.168.1.89/ <<EOF
{
    "action": "record:save",
    "access_token": "validToken",
    "database_id": "private",
    "records": [{
        "_id": "EA6A3E68-90F3-49B5-B470-5FFDB7A0D4E8",
        "_type": "note",
        "content": "ewdsa"
    }]
}
EOF
*/
func RecordSaveHandler(payload *router.Payload, response *router.Response) {
	if (*recordPayload)(payload).IsWriteAllowed() {
		response.Err = oderr.ErrWriteDenied
		return
	}

	db := payload.Database
	recordMaps, ok := payload.Data["records"].([]interface{})
	if !ok {
		response.Err = oderr.NewRequestInvalidErr(errors.New("expected list of record"))
		return
	}

	length := len(recordMaps)

	records := make([]transportRecord, length, length)
	results := make([]responseItem, length, length)
	for i := range records {
		r := recordMaps[i].(map[string]interface{})
		if err := records[i].InitFromMap(r); err != nil {
			response.Err = oderr.NewRequestInvalidErr(err)
			return
		} else if err := db.Save((*oddb.Record)(&records[i])); err != nil {
			log.WithFields(log.Fields{
				"record": &records[i],
				"err":    err,
			}).Debugln("Failed to save record")

			results[i] = newResponseItemErr(
				records[i].ID(),
				oderr.NewResourceSaveFailureErrWithStringID("record", records[i].ID()),
			)
		} else {
			results[i] = newResponseItem(&records[i])
		}
	}

	response.Result = results
}

/*
RecordFetchHandler is dummy implementation on fetching Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/ <<EOF
{
    "action": "record:fetch",
    "access_token": "validToken",
    "database_id": "private",
    "ids": ["1004", "1005"]
}
EOF
*/
func RecordFetchHandler(payload *router.Payload, response *router.Response) {
	interfaces, ok := payload.Data["ids"].([]interface{})
	if !ok {
		response.Err = oderr.NewRequestInvalidErr(errors.New("expected list of id"))
		return
	}

	length := len(interfaces)
	recordIDs := make([]compRecordID, length, length)
	for i, it := range interfaces {
		rawID, ok := it.(string)
		if !ok {
			response.Err = oderr.NewRequestInvalidErr(errors.New("expected string id"))
			return
		}

		ss := strings.SplitN(rawID, "/", 2)
		if len(ss) == 1 {
			response.Err = oderr.NewRequestInvalidErr(fmt.Errorf("invalid id format: %v", rawID))
			return
		}

		recordIDs[i] = compRecordID{ss[0], ss[1]}
	}

	db := payload.Database

	results := make([]responseItem, length, length)
	for i, recordID := range recordIDs {
		record := transportRecord{}
		if err := db.Get(recordID.id, (*oddb.Record)(&record)); err != nil {
			if err == oddb.ErrRecordNotFound {
				results[i] = newResponseItemErr(
					recordID.ID(),
					oderr.ErrRecordNotFound,
				)
			} else {
				results[i] = newResponseItemErr(
					recordID.ID(),
					oderr.NewResourceFetchFailureErr("record", recordID.ID()),
				)
			}
		} else {
			results[i] = newResponseItem(&record)
		}
	}

	response.Result = results
}

func keyPathFromRaw(rawKeyPath map[string]interface{}) string {
	mapType := rawKeyPath["$type"]
	if mapType != "keypath" {
		panic(fmt.Errorf("got key path's type %v, want \"keypath\"", mapType))
	}

	keypath := rawKeyPath["$val"].(string)
	if keypath == "" {
		panic(errors.New("empty key path value"))
	}

	return keypath
}

func sortFromRaw(rawSort []interface{}, sort *oddb.Sort) {
	keyPathMap, _ := rawSort[0].(map[string]interface{})
	if len(keyPathMap) == 0 {
		panic(errors.New("empty key path in sort descriptor"))
	}
	keyPath := keyPathFromRaw(keyPathMap)

	orderStr, _ := rawSort[1].(string)
	if orderStr == "" {
		panic(errors.New("empty sort order in sort descriptor"))
	}

	var sortOrder oddb.SortOrder
	switch orderStr {
	case "asc":
		sortOrder = oddb.Asc
	case "desc":
		sortOrder = oddb.Desc
	default:
		panic(fmt.Errorf("unknown sort order: %v", orderStr))
	}

	sort.KeyPath = keyPath
	sort.Order = sortOrder
}

func sortsFromRaw(rawSorts []interface{}) []oddb.Sort {
	length := len(rawSorts)
	sorts := make([]oddb.Sort, length, length)

	for i := range rawSorts {
		sortSlice, _ := rawSorts[i].([]interface{})
		if len(sortSlice) != 2 {
			panic(fmt.Errorf("got len(sort descriptor) = %v, want 2", len(sortSlice)))
		}
		sortFromRaw(sortSlice, &sorts[i])
	}

	return sorts
}

func queryFromPayload(payload *router.Payload, query *oddb.Query) (err oderr.Error) {
	defer func() {
		// use panic to escape from inner error
		if r := recover(); r != nil {
			if queryErr, ok := r.(error); ok {
				log.WithField("payload", payload).Debugln("failed to construct query")
				err = oderr.NewFmt(oderr.RequestInvalidErr, "failed to construct query: %v", queryErr.Error())
			} else {
				log.WithField("recovered", r).Errorln("panic recovered while constructing query")
				err = oderr.New(oderr.RequestInvalidErr, "error occurred while constructing query")
			}
		}
	}()

	recordType, _ := payload.Data["record_type"].(string)
	if recordType == "" {
		return oderr.New(oderr.RequestInvalidErr, "recordType cannot be empty")
	}
	query.Type = recordType

	if rawSorts, ok := payload.Data["sort"]; ok {
		if rawSorts, ok := rawSorts.([]interface{}); ok {
			query.Sorts = sortsFromRaw(rawSorts)
		} else {
			return oderr.New(oderr.RequestInvalidErr, "order has to be an array")
		}
	}

	return nil
}

/*
RecordQueryHandler is dummy implementation on fetching Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/ <<EOF
{
    "action": "record:query",
    "access_token": "validToken",
    "database_id": "private",
    "record_type": "note",
    "sort": [
        [{"$val": "noteOrder", "$type": "desc"}, "asc"]
    ]
}
EOF
*/
func RecordQueryHandler(payload *router.Payload, response *router.Response) {
	db := payload.Database

	query := oddb.Query{}
	if err := queryFromPayload(payload, &query); err != nil {
		response.Err = err
		return
	}

	results, err := db.Query(&query)
	if err != nil {
		response.Err = oderr.ErrDatabaseOpenFailed
		return
	}
	defer results.Close()

	records := []transportRecord{}
	for results.Scan() {
		records = append(records, transportRecord(results.Record()))
	}

	if err != nil {
		response.Err = oderr.ErrDatabaseQueryFailed
		return
	}

	response.Result = records
}

/*
RecordDeleteHandler is dummy implementation on delete Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/ <<EOF
{
    "action": "record:delete",
    "access_token": "validToken",
    "database_id": "_private",
    "ids": ["EA6A3E68-90F3-49B5-B470-5FFDB7A0D4E8"]
}
EOF
*/
func RecordDeleteHandler(payload *router.Payload, response *router.Response) {
	if (*recordPayload)(payload).IsWriteAllowed() {
		response.Err = oderr.ErrWriteDenied
		return
	}

	db := payload.Database

	interfaces, ok := payload.Data["ids"].([]interface{})
	if !ok {
		response.Err = oderr.NewRequestInvalidErr(errors.New("expected list of id"))
		return
	}

	length := len(interfaces)
	recordIDs := make([]compRecordID, length, length)
	for i, it := range interfaces {
		rawID, ok := it.(string)
		if !ok {
			response.Err = oderr.NewRequestInvalidErr(errors.New("expected string id"))
			return
		}

		ss := strings.SplitN(rawID, "/", 2)
		if len(ss) == 1 {
			response.Err = oderr.NewRequestInvalidErr(fmt.Errorf("invalid id format: %v", rawID))
			return
		}

		recordIDs[i] = compRecordID{ss[0], ss[1]}
	}

	results := []responseItem{}
	for i, recordID := range recordIDs {
		record := transportRecord{}
		if err := db.Get(recordID.id, (*oddb.Record)(&record)); err != nil {
			if err == oddb.ErrRecordNotFound {
				results[i] = newResponseItemErr(
					recordID.ID(),
					oderr.ErrRecordNotFound,
				)
			} else {
				results[i] = newResponseItemErr(
					record.ID(),
					oderr.NewResourceDeleteFailureErrWithStringID("record", record.ID()),
				)
			}
		} else {
			results[i] = newResponseItem(&record)
		}
	}

	response.Result = results
	return
}

type compRecordID struct {
	kind string
	id   string
}

func (rid compRecordID) ID() string {
	return rid.kind + "/" + rid.id
}
<commit_msg>Record query handler now returns properly json serializable items, #65<commit_after>package handler

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"strings"

	"github.com/oursky/ourd/oddb"
	"github.com/oursky/ourd/oderr"
	"github.com/oursky/ourd/router"
)

// recordPayload is the input parameter in RecordHandler
type recordPayload router.Payload

func (p *recordPayload) IsWriteAllowed() bool {
	return p.UserInfo == nil
}

// transportRecord override JSON serialization and deserialization of
// oddb.Record
type transportRecord oddb.Record

func (r transportRecord) ID() string {
	return r.Type + "/" + r.Key
}

func (r transportRecord) MarshalJSON() ([]byte, error) {
	// NOTE(limouren): marshalling of type/key is delegated to responseItem
	return json.Marshal(r.Data)
}

func (r *transportRecord) UnmarshalJSON(data []byte) error {
	object := map[string]interface{}{}
	err := json.Unmarshal(data, &object)

	if err != nil {
		return err
	}

	return r.InitFromMap(object)
}

func (r *transportRecord) InitFromMap(m map[string]interface{}) error {
	rawID, ok := m["_id"].(string)
	if !ok {
		return errors.New(`record/json: required field "_id" not found`)
	}
	delete(m, "_id")

	ss := strings.SplitN(rawID, "/", 2)
	if len(ss) == 1 {
		return fmt.Errorf(`record/json: "_id" should be of format '{type}/{id}', got %#v`, rawID)
	}

	recordType, id := ss[0], ss[1]

	// remove reserved keys
	for key := range m {
		if key[0] == '_' {
			delete(m, key)
		}
	}

	r.Key = id
	r.Type = recordType
	r.Data = m

	return nil
}

type responseItem struct {
	id     string
	record *transportRecord
	err    oderr.Error
}

func newResponseItem(record *transportRecord) responseItem {
	return responseItem{
		id:     record.ID(),
		record: record,
	}
}

func newResponseItemErr(id string, err oderr.Error) responseItem {
	return responseItem{
		id:  id,
		err: err,
	}
}

func (item responseItem) MarshalJSON() ([]byte, error) {
	var (
		buf bytes.Buffer
		i   interface{}
	)
	buf.Write([]byte(`{"_id":"`))
	buf.WriteString(item.id)
	buf.Write([]byte(`","_type":"`))
	if item.err != nil {
		buf.Write([]byte(`error",`))
		i = item.err
	} else if item.record != nil {
		buf.Write([]byte(`record",`))
		i = item.record
	} else {
		panic("inconsistent state: both err and record is nil")
	}

	bodyBytes, err := json.Marshal(i)
	if err != nil {
		return nil, err
	}

	if bodyBytes[0] != '{' {
		return nil, fmt.Errorf("first char of embedded json != {: %v", string(bodyBytes))
	} else if bodyBytes[len(bodyBytes)-1] != '}' {
		return nil, fmt.Errorf("last char of embedded json != }: %v", string(bodyBytes))
	}
	buf.Write(bodyBytes[1:])
	return buf.Bytes(), nil
}

/*
RecordSaveHandler is dummy implementation on save/modify Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://192.168.1.89/ <<EOF
{
    "action": "record:save",
    "access_token": "validToken",
    "database_id": "private",
    "records": [{
        "_id": "EA6A3E68-90F3-49B5-B470-5FFDB7A0D4E8",
        "_type": "note",
        "content": "ewdsa"
    }]
}
EOF
*/
func RecordSaveHandler(payload *router.Payload, response *router.Response) {
	if (*recordPayload)(payload).IsWriteAllowed() {
		response.Err = oderr.ErrWriteDenied
		return
	}

	db := payload.Database
	recordMaps, ok := payload.Data["records"].([]interface{})
	if !ok {
		response.Err = oderr.NewRequestInvalidErr(errors.New("expected list of record"))
		return
	}

	length := len(recordMaps)

	records := make([]transportRecord, length, length)
	results := make([]responseItem, length, length)
	for i := range records {
		r := recordMaps[i].(map[string]interface{})
		if err := records[i].InitFromMap(r); err != nil {
			response.Err = oderr.NewRequestInvalidErr(err)
			return
		} else if err := db.Save((*oddb.Record)(&records[i])); err != nil {
			log.WithFields(log.Fields{
				"record": &records[i],
				"err":    err,
			}).Debugln("Failed to save record")

			results[i] = newResponseItemErr(
				records[i].ID(),
				oderr.NewResourceSaveFailureErrWithStringID("record", records[i].ID()),
			)
		} else {
			results[i] = newResponseItem(&records[i])
		}
	}

	response.Result = results
}

/*
RecordFetchHandler is dummy implementation on fetching Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/ <<EOF
{
    "action": "record:fetch",
    "access_token": "validToken",
    "database_id": "private",
    "ids": ["1004", "1005"]
}
EOF
*/
func RecordFetchHandler(payload *router.Payload, response *router.Response) {
	interfaces, ok := payload.Data["ids"].([]interface{})
	if !ok {
		response.Err = oderr.NewRequestInvalidErr(errors.New("expected list of id"))
		return
	}

	length := len(interfaces)
	recordIDs := make([]compRecordID, length, length)
	for i, it := range interfaces {
		rawID, ok := it.(string)
		if !ok {
			response.Err = oderr.NewRequestInvalidErr(errors.New("expected string id"))
			return
		}

		ss := strings.SplitN(rawID, "/", 2)
		if len(ss) == 1 {
			response.Err = oderr.NewRequestInvalidErr(fmt.Errorf("invalid id format: %v", rawID))
			return
		}

		recordIDs[i] = compRecordID{ss[0], ss[1]}
	}

	db := payload.Database

	results := make([]responseItem, length, length)
	for i, recordID := range recordIDs {
		record := transportRecord{}
		if err := db.Get(recordID.id, (*oddb.Record)(&record)); err != nil {
			if err == oddb.ErrRecordNotFound {
				results[i] = newResponseItemErr(
					recordID.ID(),
					oderr.ErrRecordNotFound,
				)
			} else {
				results[i] = newResponseItemErr(
					recordID.ID(),
					oderr.NewResourceFetchFailureErr("record", recordID.ID()),
				)
			}
		} else {
			results[i] = newResponseItem(&record)
		}
	}

	response.Result = results
}

func keyPathFromRaw(rawKeyPath map[string]interface{}) string {
	mapType := rawKeyPath["$type"]
	if mapType != "keypath" {
		panic(fmt.Errorf("got key path's type %v, want \"keypath\"", mapType))
	}

	keypath := rawKeyPath["$val"].(string)
	if keypath == "" {
		panic(errors.New("empty key path value"))
	}

	return keypath
}

func sortFromRaw(rawSort []interface{}, sort *oddb.Sort) {
	keyPathMap, _ := rawSort[0].(map[string]interface{})
	if len(keyPathMap) == 0 {
		panic(errors.New("empty key path in sort descriptor"))
	}
	keyPath := keyPathFromRaw(keyPathMap)

	orderStr, _ := rawSort[1].(string)
	if orderStr == "" {
		panic(errors.New("empty sort order in sort descriptor"))
	}

	var sortOrder oddb.SortOrder
	switch orderStr {
	case "asc":
		sortOrder = oddb.Asc
	case "desc":
		sortOrder = oddb.Desc
	default:
		panic(fmt.Errorf("unknown sort order: %v", orderStr))
	}

	sort.KeyPath = keyPath
	sort.Order = sortOrder
}

func sortsFromRaw(rawSorts []interface{}) []oddb.Sort {
	length := len(rawSorts)
	sorts := make([]oddb.Sort, length, length)

	for i := range rawSorts {
		sortSlice, _ := rawSorts[i].([]interface{})
		if len(sortSlice) != 2 {
			panic(fmt.Errorf("got len(sort descriptor) = %v, want 2", len(sortSlice)))
		}
		sortFromRaw(sortSlice, &sorts[i])
	}

	return sorts
}

func queryFromPayload(payload *router.Payload, query *oddb.Query) (err oderr.Error) {
	defer func() {
		// use panic to escape from inner error
		if r := recover(); r != nil {
			if queryErr, ok := r.(error); ok {
				log.WithField("payload", payload).Debugln("failed to construct query")
				err = oderr.NewFmt(oderr.RequestInvalidErr, "failed to construct query: %v", queryErr.Error())
			} else {
				log.WithField("recovered", r).Errorln("panic recovered while constructing query")
				err = oderr.New(oderr.RequestInvalidErr, "error occurred while constructing query")
			}
		}
	}()

	recordType, _ := payload.Data["record_type"].(string)
	if recordType == "" {
		return oderr.New(oderr.RequestInvalidErr, "recordType cannot be empty")
	}
	query.Type = recordType

	if rawSorts, ok := payload.Data["sort"]; ok {
		if rawSorts, ok := rawSorts.([]interface{}); ok {
			query.Sorts = sortsFromRaw(rawSorts)
		} else {
			return oderr.New(oderr.RequestInvalidErr, "order has to be an array")
		}
	}

	return nil
}

/*
RecordQueryHandler is dummy implementation on fetching Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/ <<EOF
{
    "action": "record:query",
    "access_token": "validToken",
    "database_id": "private",
    "record_type": "note",
    "sort": [
        [{"$val": "noteOrder", "$type": "desc"}, "asc"]
    ]
}
EOF
*/
func RecordQueryHandler(payload *router.Payload, response *router.Response) {
	db := payload.Database

	query := oddb.Query{}
	if err := queryFromPayload(payload, &query); err != nil {
		response.Err = err
		return
	}

	results, err := db.Query(&query)
	if err != nil {
		response.Err = oderr.ErrDatabaseOpenFailed
		return
	}
	defer results.Close()

	records := []responseItem{}
	for results.Scan() {
		record := transportRecord(results.Record())
		records = append(records, newResponseItem(&record))
	}

	if err != nil {
		response.Err = oderr.ErrDatabaseQueryFailed
		return
	}

	response.Result = records
}

/*
RecordDeleteHandler is dummy implementation on delete Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/ <<EOF
{
    "action": "record:delete",
    "access_token": "validToken",
    "database_id": "_private",
    "ids": ["EA6A3E68-90F3-49B5-B470-5FFDB7A0D4E8"]
}
EOF
*/
func RecordDeleteHandler(payload *router.Payload, response *router.Response) {
	if (*recordPayload)(payload).IsWriteAllowed() {
		response.Err = oderr.ErrWriteDenied
		return
	}

	db := payload.Database

	interfaces, ok := payload.Data["ids"].([]interface{})
	if !ok {
		response.Err = oderr.NewRequestInvalidErr(errors.New("expected list of id"))
		return
	}

	length := len(interfaces)
	recordIDs := make([]compRecordID, length, length)
	for i, it := range interfaces {
		rawID, ok := it.(string)
		if !ok {
			response.Err = oderr.NewRequestInvalidErr(errors.New("expected string id"))
			return
		}

		ss := strings.SplitN(rawID, "/", 2)
		if len(ss) == 1 {
			response.Err = oderr.NewRequestInvalidErr(fmt.Errorf("invalid id format: %v", rawID))
			return
		}

		recordIDs[i] = compRecordID{ss[0], ss[1]}
	}

	results := []responseItem{}
	for i, recordID := range recordIDs {
		record := transportRecord{}
		if err := db.Get(recordID.id, (*oddb.Record)(&record)); err != nil {
			if err == oddb.ErrRecordNotFound {
				results[i] = newResponseItemErr(
					recordID.ID(),
					oderr.ErrRecordNotFound,
				)
			} else {
				results[i] = newResponseItemErr(
					record.ID(),
					oderr.NewResourceDeleteFailureErrWithStringID("record", record.ID()),
				)
			}
		} else {
			results[i] = newResponseItem(&record)
		}
	}

	response.Result = results
	return
}

type compRecordID struct {
	kind string
	id   string
}

func (rid compRecordID) ID() string {
	return rid.kind + "/" + rid.id
}
<|endoftext|>