<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/hoisie/web"
	"io"
	"os"
	"path/filepath"
)

type jsonFile struct {
	Hash string `json:"hash"`
	Name string `json:"name"`
	URL  string `json:"url"`
	Size int    `json:"size"`
}

type jsonResponse struct {
	Success bool       `json:"success"`
	Files   []jsonFile `json:"files"`
}

//Handles POST upload requests. the updateURL is used to pass messages
//to the urlHandler indicating that the DB should be updated.
func handleUpload(ctx *web.Context, updateURL chan<- string, updateResp <-chan *Response) string {
	//TODO: Implemente limits with settings.ini or something
	//TODO: Verify the max number of files that pomf could upload
	err := ctx.Request.ParseMultipartForm(500 * 1024 * 1024)
	if err != nil {
		return "Error handling form!\n"
	}
	form := ctx.Request.MultipartForm
	//Loop through and append to the response struct
    c := 0
	for _, _ = range form.File["files[]"] {
        c++
    }
    fmt.Println(c)
	resFiles := make([]jsonFile, c)
	for idx, fileHeader := range form.File["files[]"] {
		filename := fileHeader.Filename
		file, err := fileHeader.Open()
		size, err := file.Seek(0, 2)
		if err != nil {
			return "Error parsing file!\n"
		}
		if size > 50*1024*1024 {
			return "File too big!\n"
		}
		//Seek back to beginning
		file.Seek(0, 0)
		if err != nil {
			return err.Error()
		}
		hash := Md5(file)
		//If file doesn't already exist, create it
		if _, err := os.Stat("files/" + hash); os.IsNotExist(err) {
			f, err := os.Create("files/" + hash)
			if err != nil {
				return "Error, file could not be created.\n"
			}
			_, err = file.Seek(0, 0)
			if err != nil {
				return "Error reading the file\n"
			}
			_, err = io.Copy(f, file)
			if err != nil {
				return "Error, file could not be written to.\n"
			}
		}

		ext := filepath.Ext(filename)
		//Send the hash and ext for updating
		updateURL <- ext + ":" + hash
		resp := <-updateResp
		//Even though this is redundant, it might eventually be useful
		if resp.status == "Failure" {
			return resp.message + "\n"
		} else {
			jFile := jsonFile{Hash: hash, Name: filename, URL: resp.message, Size: int(size)}
            fmt.Println("About to append to array")
			resFiles[idx] = jFile
		}
	}
	jsonResp := &jsonResponse{Success: true, Files: resFiles}
	jresp, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Println(err.Error())
	}
	return string(jresp)
}
<commit_msg>Minor improvements<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/hoisie/web"
	"io"
	"os"
	"path/filepath"
)

type jsonFile struct {
	Hash string `json:"hash"`
	Name string `json:"name"`
	URL  string `json:"url"`
	Size int    `json:"size"`
}

type jsonResponse struct {
	Success bool       `json:"success"`
	Files   []jsonFile `json:"files"`
}

func throwErr(arr *[]jsonFile) string{
    jsonResp := &jsonResponse{Success: false, Files: *arr}
	jresp, err := json.Marshal(jsonResp)
	if err != nil {
	    fmt.Println(err.Error())
	}
	return string(jresp)
}

//Handles POST upload requests. the updateURL is used to pass messages
//to the urlHandler indicating that the DB should be updated.
func handleUpload(ctx *web.Context, updateURL chan<- string, updateResp <-chan *Response) string {
	//TODO: Implemente limits with settings.ini or something
	err := ctx.Request.ParseMultipartForm(500 * 1024 * 1024)
	if err != nil {
		return "Error handling form!\n"
	}
	form := ctx.Request.MultipartForm

	//Loop through and append to the response struct
	resFiles := make([]jsonFile, len(form.File["files[]"]))
	for idx, fileHeader := range form.File["files[]"] {
		filename := fileHeader.Filename
		file, err := fileHeader.Open()
		size, err := file.Seek(0, 2)
		if err != nil {
            return throwErr(&resFiles)
		}
		if size > 50*1024*1024 {
            return throwErr(&resFiles)
		}
		//Seek back to beginning
		file.Seek(0, 0)
		if err != nil {
			return err.Error()
		}
		hash := Md5(file)
		ext := filepath.Ext(filename)
		//Send the hash and ext for updating
		updateURL <- ext + ":" + hash
		resp := <-updateResp
		//Even though this is redundant, it might eventually be useful
		if resp.status == "Failure" {
            return throwErr(&resFiles)
		} else {
			jFile := jsonFile{Hash: hash, Name: filename, URL: resp.message, Size: int(size)}
			resFiles[idx] = jFile
		}

		//If file doesn't already exist, create it
		if _, err := os.Stat("files/" + hash); os.IsNotExist(err) {
			f, err := os.Create("files/" + hash)
			if err != nil {
                return throwErr(&resFiles)
			}
			_, err = file.Seek(0, 0)
			if err != nil {
                return throwErr(&resFiles)
			}
			_, err = io.Copy(f, file)
			if err != nil {
                return throwErr(&resFiles)
			}
		}
	}
	jsonResp := &jsonResponse{Success: true, Files: resFiles}
	jresp, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Println(err.Error())
	}
	return string(jresp)
}
<|endoftext|>


<commit_before>package queue

import (
	"testing"
	"sync/atomic"
)

func assert(b bool, s string, t *testing.T) {

	if !b {

		t.Fatal(s)
	}
}

func TestQueue (t *testing.T) {

	var count int32 = 0
	handler := func(val interface{}) {

		atomic.AddInt32(&count, int32(val.(int)))

	}
	q := NewQueue(handler, 5)

	for i := 0; i < 200; i++ {

		q.Push(i)
	}

	q.Wait()
	if count != 19900 {

		t.Fail()
	}

}<commit_msg>removed old assert function<commit_after>package queue

import (
	"testing"
	"sync/atomic"
)

func TestQueue (t *testing.T) {

	var count int32 = 0
	handler := func(val interface{}) {

		atomic.AddInt32(&count, int32(val.(int)))

	}
	q := NewQueue(handler, 5)

	for i := 0; i < 200; i++ {

		q.Push(i)
	}

	q.Wait()
	if count != 19900 {

		t.Fail()
	}

}<|endoftext|>


<commit_before>package web

import (
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"github.com/uswitch/dagr/app"
	"github.com/uswitch/dagr/program"
	"log"
	"net/http"
	"strconv"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func programExecutions(app app.App) http.HandlerFunc {
	return func(w http.ResponseWriter, req *http.Request) {
		conn, err := upgrader.Upgrade(w, req, nil)
		if err != nil {
			log.Println("cannot upgrade to websocket")
			return
		}
		vars := mux.Vars(req)
		programName := vars["program"]
		p := app.FindProgram(programName)
		if p == nil {
			log.Println("no such program:", programName)
			http.NotFound(w, req)
		} else {
			program.ProgramLog(p, "subscribing to executions")
			p.Subscribe(conn)
			go readLoop(p, conn)
		}
	}

}

func handleExecutionMessages(app app.App) http.HandlerFunc {
	return func(w http.ResponseWriter, req *http.Request) {
		conn, err := upgrader.Upgrade(w, req, nil)
		if err != nil {
			log.Println("cannot upgrade to websocket")
			return
		}
		vars := mux.Vars(req)
		executionId := vars["executionId"]
		execution := app.FindExecution(executionId)
		if execution == nil {
			log.Println("no such execution:", executionId)
			http.NotFound(w, req)
		} else {
			program.ExecutionLog(execution, "subscribing to messages")
			execution.Subscribe(conn)
			countSoFarStr := vars["countSoFar"]
			countSoFar, err := strconv.Atoi(countSoFarStr)
			if err != nil {
				log.Println("countSoFar not an integer?", countSoFarStr, err)
			} else {
				messagesCaughtUp := execution.CatchUp(conn, countSoFar)
				if messagesCaughtUp > 0 {
					log.Println("caught up", messagesCaughtUp, "message(s)")
				}
			}

			go readLoop(execution, conn)
		}
	}
}

// read is required (http://www.gorillatoolkit.org/pkg/websocket)
func readLoop(s program.Subscriber, c *websocket.Conn) {
	for {
		_, _, err := c.NextReader()
		if err != nil {
			c.Close()
			s.Unsubscribe(c)
			return
		}
	}
}
<commit_msg>Now prints what the actuall error if we cant upgrade to websocket<commit_after>package web

import (
	"fmt"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"github.com/uswitch/dagr/app"
	"github.com/uswitch/dagr/program"
	"log"
	"net/http"
	"strconv"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func programExecutions(app app.App) http.HandlerFunc {
	return func(w http.ResponseWriter, req *http.Request) {
		conn, err := upgrader.Upgrade(w, req, nil)
		if err != nil {
			log.Println(fmt.Errorf("cannot upgrade to websocket: %v", err))
			return
		}
		vars := mux.Vars(req)
		programName := vars["program"]
		p := app.FindProgram(programName)
		if p == nil {
			log.Println("no such program:", programName)
			http.NotFound(w, req)
		} else {
			program.ProgramLog(p, "subscribing to executions")
			p.Subscribe(conn)
			go readLoop(p, conn)
		}
	}

}

func handleExecutionMessages(app app.App) http.HandlerFunc {
	return func(w http.ResponseWriter, req *http.Request) {
		conn, err := upgrader.Upgrade(w, req, nil)
		if err != nil {
			log.Println(fmt.Errorf("cannot upgrade to websocket: %v", err))
			return
		}
		vars := mux.Vars(req)
		executionId := vars["executionId"]
		execution := app.FindExecution(executionId)
		if execution == nil {
			log.Println("no such execution:", executionId)
			http.NotFound(w, req)
		} else {
			program.ExecutionLog(execution, "subscribing to messages")
			execution.Subscribe(conn)
			countSoFarStr := vars["countSoFar"]
			countSoFar, err := strconv.Atoi(countSoFarStr)
			if err != nil {
				log.Println("countSoFar not an integer?", countSoFarStr, err)
			} else {
				messagesCaughtUp := execution.CatchUp(conn, countSoFar)
				if messagesCaughtUp > 0 {
					log.Println("caught up", messagesCaughtUp, "message(s)")
				}
			}

			go readLoop(execution, conn)
		}
	}
}

// read is required (http://www.gorillatoolkit.org/pkg/websocket)
func readLoop(s program.Subscriber, c *websocket.Conn) {
	for {
		_, _, err := c.NextReader()
		if err != nil {
			c.Close()
			s.Unsubscribe(c)
			return
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Aerospike, Inc.
//
// 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 aerospike

import (
)

const (
	AS_PREDEXP_AND					uint16 = 1
	AS_PREDEXP_OR					uint16 = 2
	AS_PREDEXP_NOT					uint16 = 3

	AS_PREDEXP_INTEGER_VALUE		uint16 = 10
	AS_PREDEXP_STRING_VALUE			uint16 = 11
	AS_PREDEXP_GEOJSON_VALUE		uint16 = 12

	AS_PREDEXP_INTEGER_BIN			uint16 = 100
	AS_PREDEXP_STRING_BIN			uint16 = 101
	AS_PREDEXP_GEOJSON_BIN			uint16 = 102

	AS_PREDEXP_RECSIZE				uint16 = 150
	AS_PREDEXP_LAST_UPDATE			uint16 = 151
	AS_PREDEXP_VOID_TIME			uint16 = 152

	AS_PREDEXP_INTEGER_EQUAL		uint16 = 200
	AS_PREDEXP_INTEGER_UNEQUAL		uint16 = 201
	AS_PREDEXP_INTEGER_GREATER		uint16 = 202
	AS_PREDEXP_INTEGER_GREATEREQ	uint16 = 203
	AS_PREDEXP_INTEGER_LESS			uint16 = 204
	AS_PREDEXP_INTEGER_LESSEQ		uint16 = 205

	AS_PREDEXP_STRING_EQUAL			uint16 = 210
	AS_PREDEXP_STRING_UNEQUAL		uint16 = 211
	AS_PREDEXP_STRING_REGEX			uint16 = 212

	AS_PREDEXP_GEOJSON_WITHIN		uint16 = 220
	AS_PREDEXP_GEOJSON_CONTAINS		uint16 = 221
)

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

type PredExp interface {
	MarshaledSize() int
	Marshal(cmd *baseCommand) error
}

type PredExpBase struct {
}

func (self *PredExpBase) MarshaledSize() int {
	return 2 + 4	// sizeof(TAG) + sizeof(LEN)
}

func (self *PredExpBase) MarshalTL(
	cmd *baseCommand,
	tag uint16,
	len uint32) int {
	return 2 + 4	// sizeof(TAG) + sizeof(LEN)
}

// ---------------- PredExpAnd

type PredExpAnd struct {
	PredExpBase
	nexpr uint16	// number of child expressions
}

func NewPredExpAnd(nexpr uint16) *PredExpAnd {
	return &PredExpAnd{ nexpr: nexpr }
}

func (self *PredExpAnd) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() + 2
}

func (self *PredExpAnd) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_AND, 2)
	cmd.WriteUint16(self.nexpr)
	return nil
}

// ---------------- PredExpOr

type PredExpOr struct {
	PredExpBase
	nexpr uint16	// number of child expressions
}

func NewPredExpOr(nexpr uint16) *PredExpOr {
	return &PredExpOr{ nexpr: nexpr }
}

func (self *PredExpOr) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() + 2
}

func (self *PredExpOr) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_OR, 2)
	cmd.WriteUint16(self.nexpr)
	return nil
}

// ---------------- PredExpNot

type PredExpNot struct {
	PredExpBase
}

func NewPredExpNot() *PredExpNot {
	return &PredExpNot{ }
}

func (self *PredExpNot) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize()
}

func (self *PredExpNot) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_NOT, 0)
	return nil
}

// ---------------- PredExpIntegerValue

type PredExpIntegerValue struct {
	PredExpBase
	val int64
}

func NewPredExpIntegerValue(val int64) *PredExpIntegerValue {
	return &PredExpIntegerValue{ val: val }
}

func (self *PredExpIntegerValue) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() + 8
}

func (self *PredExpIntegerValue) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_INTEGER_VALUE, 8)
	cmd.WriteInt64(self.val)
	return nil
}

// ---------------- PredExpStringValue

type PredExpStringValue struct {
	PredExpBase
	val string
}

func NewPredExpStringValue(val string) *PredExpStringValue {
	return &PredExpStringValue{ val: val }
}

func (self *PredExpStringValue) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() + len(self.val)
}

func (self *PredExpStringValue) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_STRING_VALUE, uint32(len(self.val)))
	cmd.WriteString(self.val)
	return nil
}

// ---------------- PredExpGeoJSONValue

type PredExpGeoJSONValue struct {
	PredExpBase
	val string
}

func NewPredExpGeoJSONValue(val string) *PredExpGeoJSONValue {
	return &PredExpGeoJSONValue{ val: val }
}

func (self *PredExpGeoJSONValue) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() +
		1 +				// flags
		2 + 			// ncells
		len(self.val)	// strlen value
}

func (self *PredExpGeoJSONValue) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_GEOJSON_VALUE, uint32(1 + 2 + len(self.val)))
	cmd.WriteByte(uint8(0))
	cmd.WriteUint16(0)
	cmd.WriteString(self.val)
	return nil
}

// ---------------- PredExp???Bin

type PredExpBin struct {
	PredExpBase
	name string
	tag uint16	// not marshaled
}

func NewPredExpIntegerBin(name string) *PredExpBin {
	return &PredExpBin{ name: name, tag: AS_PREDEXP_INTEGER_BIN, }
}

func NewPredExpStringBin(name string) *PredExpBin {
	return &PredExpBin{ name: name, tag: AS_PREDEXP_STRING_BIN, }
}

func NewPredExpGeoJSONBin(name string) *PredExpBin {
	return &PredExpBin{ name: name, tag: AS_PREDEXP_GEOJSON_BIN, }
}

func (self *PredExpBin) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() + 1 + len(self.name)
}

func (self *PredExpBin) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, self.tag, uint32(1 + len(self.name)))
	cmd.WriteByte(uint8(len(self.name)))
	cmd.WriteString(self.name)
	return nil
}

// ---------------- PredExpMD (RecSize, LastUpdate, VoidTime)

type PredExpMD struct {
	PredExpBase
	tag uint16	// not marshaled
}

func (self *PredExpMD) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize()
}

func (self *PredExpMD) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, self.tag, 0)
	return nil
}

func NewPredExpRecSize() *PredExpMD {
	return &PredExpMD{ tag: AS_PREDEXP_RECSIZE }
}

func NewPredExpLastUpdate() *PredExpMD {
	return &PredExpMD{ tag: AS_PREDEXP_LAST_UPDATE }
}

func NewPredExpVoidTime() *PredExpMD {
	return &PredExpMD{ tag: AS_PREDEXP_VOID_TIME }
}

// ---------------- PredExpCompare

type PredExpCompare struct {
	PredExpBase
	tag uint16	// not marshaled
}

func (self *PredExpCompare) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize()
}

func (self *PredExpCompare) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, self.tag, 0)
	return nil
}

func NewPredExpIntegerEqual() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_INTEGER_EQUAL }
}

func NewPredExpIntegerUnequal() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_INTEGER_UNEQUAL }
}

func NewPredExpIntegerGreater() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_INTEGER_GREATER }
}

func NewPredExpIntegerGreaterEq() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_INTEGER_GREATEREQ }
}

func NewPredExpIntegerLess() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_INTEGER_LESS }
}

func NewPredExpIntegerLessEq() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_INTEGER_LESSEQ }
}

func NewPredExpStringEqual() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_STRING_EQUAL }
}

func NewPredExpStringUnequal() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_STRING_UNEQUAL }
}

func NewPredExpGeoJSONWithin() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_GEOJSON_WITHIN }
}

func NewPredExpGeoJSONContains() *PredExpCompare {
	return &PredExpCompare{ tag: AS_PREDEXP_GEOJSON_CONTAINS }
}

// ---------------- PredExpStringRegex

type PredExpStringRegex struct {
	PredExpBase
	cflags uint32		// cflags
}

func NewPredExpStringRegex(cflags uint32) *PredExpStringRegex {
	return &PredExpStringRegex{ cflags: cflags }
}

func (self *PredExpStringRegex) MarshaledSize() int {
	return self.PredExpBase.MarshaledSize() + 4
}

func (self *PredExpStringRegex) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_STRING_REGEX, 4)
	cmd.WriteUint32(self.cflags)
	return nil
}
<commit_msg>predexp: Restricted the export of predexp methods and data structures.<commit_after>// Copyright 2017 Aerospike, Inc.
//
// 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 aerospike

import (
)

const (
	AS_PREDEXP_AND					uint16 = 1
	AS_PREDEXP_OR					uint16 = 2
	AS_PREDEXP_NOT					uint16 = 3

	AS_PREDEXP_INTEGER_VALUE		uint16 = 10
	AS_PREDEXP_STRING_VALUE			uint16 = 11
	AS_PREDEXP_GEOJSON_VALUE		uint16 = 12

	AS_PREDEXP_INTEGER_BIN			uint16 = 100
	AS_PREDEXP_STRING_BIN			uint16 = 101
	AS_PREDEXP_GEOJSON_BIN			uint16 = 102

	AS_PREDEXP_RECSIZE				uint16 = 150
	AS_PREDEXP_LAST_UPDATE			uint16 = 151
	AS_PREDEXP_VOID_TIME			uint16 = 152

	AS_PREDEXP_INTEGER_EQUAL		uint16 = 200
	AS_PREDEXP_INTEGER_UNEQUAL		uint16 = 201
	AS_PREDEXP_INTEGER_GREATER		uint16 = 202
	AS_PREDEXP_INTEGER_GREATEREQ	uint16 = 203
	AS_PREDEXP_INTEGER_LESS			uint16 = 204
	AS_PREDEXP_INTEGER_LESSEQ		uint16 = 205

	AS_PREDEXP_STRING_EQUAL			uint16 = 210
	AS_PREDEXP_STRING_UNEQUAL		uint16 = 211
	AS_PREDEXP_STRING_REGEX			uint16 = 212

	AS_PREDEXP_GEOJSON_WITHIN		uint16 = 220
	AS_PREDEXP_GEOJSON_CONTAINS		uint16 = 221
)

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

type PredExp interface {
	MarshaledSize() int
	Marshal(cmd *baseCommand) error
}

type predExpBase struct {
}

func (self *predExpBase) MarshaledSize() int {
	return 2 + 4	// sizeof(TAG) + sizeof(LEN)
}

func (self *predExpBase) MarshalTL(
	cmd *baseCommand,
	tag uint16,
	len uint32) int {
	return 2 + 4	// sizeof(TAG) + sizeof(LEN)
}

// ---------------- predExpAnd

type predExpAnd struct {
	predExpBase
	nexpr uint16	// number of child expressions
}

func NewPredExpAnd(nexpr uint16) *predExpAnd {
	return &predExpAnd{ nexpr: nexpr }
}

func (self *predExpAnd) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() + 2
}

func (self *predExpAnd) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_AND, 2)
	cmd.WriteUint16(self.nexpr)
	return nil
}

// ---------------- predExpOr

type predExpOr struct {
	predExpBase
	nexpr uint16	// number of child expressions
}

func NewPredExpOr(nexpr uint16) *predExpOr {
	return &predExpOr{ nexpr: nexpr }
}

func (self *predExpOr) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() + 2
}

func (self *predExpOr) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_OR, 2)
	cmd.WriteUint16(self.nexpr)
	return nil
}

// ---------------- predExpNot

type predExpNot struct {
	predExpBase
}

func NewPredExpNot() *predExpNot {
	return &predExpNot{ }
}

func (self *predExpNot) MarshaledSize() int {
	return self.predExpBase.MarshaledSize()
}

func (self *predExpNot) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_NOT, 0)
	return nil
}

// ---------------- predExpIntegerValue

type predExpIntegerValue struct {
	predExpBase
	val int64
}

func NewPredExpIntegerValue(val int64) *predExpIntegerValue {
	return &predExpIntegerValue{ val: val }
}

func (self *predExpIntegerValue) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() + 8
}

func (self *predExpIntegerValue) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_INTEGER_VALUE, 8)
	cmd.WriteInt64(self.val)
	return nil
}

// ---------------- predExpStringValue

type predExpStringValue struct {
	predExpBase
	val string
}

func NewPredExpStringValue(val string) *predExpStringValue {
	return &predExpStringValue{ val: val }
}

func (self *predExpStringValue) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() + len(self.val)
}

func (self *predExpStringValue) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_STRING_VALUE, uint32(len(self.val)))
	cmd.WriteString(self.val)
	return nil
}

// ---------------- predExpGeoJSONValue

type predExpGeoJSONValue struct {
	predExpBase
	val string
}

func NewPredExpGeoJSONValue(val string) *predExpGeoJSONValue {
	return &predExpGeoJSONValue{ val: val }
}

func (self *predExpGeoJSONValue) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() +
		1 +				// flags
		2 + 			// ncells
		len(self.val)	// strlen value
}

func (self *predExpGeoJSONValue) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_GEOJSON_VALUE, uint32(1 + 2 + len(self.val)))
	cmd.WriteByte(uint8(0))
	cmd.WriteUint16(0)
	cmd.WriteString(self.val)
	return nil
}

// ---------------- predExp???Bin

type predExpBin struct {
	predExpBase
	name string
	tag uint16	// not marshaled
}

func NewPredExpIntegerBin(name string) *predExpBin {
	return &predExpBin{ name: name, tag: AS_PREDEXP_INTEGER_BIN, }
}

func NewPredExpStringBin(name string) *predExpBin {
	return &predExpBin{ name: name, tag: AS_PREDEXP_STRING_BIN, }
}

func NewPredExpGeoJSONBin(name string) *predExpBin {
	return &predExpBin{ name: name, tag: AS_PREDEXP_GEOJSON_BIN, }
}

func (self *predExpBin) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() + 1 + len(self.name)
}

func (self *predExpBin) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, self.tag, uint32(1 + len(self.name)))
	cmd.WriteByte(uint8(len(self.name)))
	cmd.WriteString(self.name)
	return nil
}

// ---------------- predExpMD (RecSize, LastUpdate, VoidTime)

type predExpMD struct {
	predExpBase
	tag uint16	// not marshaled
}

func (self *predExpMD) MarshaledSize() int {
	return self.predExpBase.MarshaledSize()
}

func (self *predExpMD) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, self.tag, 0)
	return nil
}

func NewPredExpRecSize() *predExpMD {
	return &predExpMD{ tag: AS_PREDEXP_RECSIZE }
}

func NewPredExpLastUpdate() *predExpMD {
	return &predExpMD{ tag: AS_PREDEXP_LAST_UPDATE }
}

func NewPredExpVoidTime() *predExpMD {
	return &predExpMD{ tag: AS_PREDEXP_VOID_TIME }
}

// ---------------- predExpCompare

type predExpCompare struct {
	predExpBase
	tag uint16	// not marshaled
}

func (self *predExpCompare) MarshaledSize() int {
	return self.predExpBase.MarshaledSize()
}

func (self *predExpCompare) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, self.tag, 0)
	return nil
}

func NewPredExpIntegerEqual() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_INTEGER_EQUAL }
}

func NewPredExpIntegerUnequal() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_INTEGER_UNEQUAL }
}

func NewPredExpIntegerGreater() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_INTEGER_GREATER }
}

func NewPredExpIntegerGreaterEq() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_INTEGER_GREATEREQ }
}

func NewPredExpIntegerLess() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_INTEGER_LESS }
}

func NewPredExpIntegerLessEq() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_INTEGER_LESSEQ }
}

func NewPredExpStringEqual() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_STRING_EQUAL }
}

func NewPredExpStringUnequal() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_STRING_UNEQUAL }
}

func NewPredExpGeoJSONWithin() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_GEOJSON_WITHIN }
}

func NewPredExpGeoJSONContains() *predExpCompare {
	return &predExpCompare{ tag: AS_PREDEXP_GEOJSON_CONTAINS }
}

// ---------------- predExpStringRegex

type predExpStringRegex struct {
	predExpBase
	cflags uint32		// cflags
}

func NewPredExpStringRegex(cflags uint32) *predExpStringRegex {
	return &predExpStringRegex{ cflags: cflags }
}

func (self *predExpStringRegex) MarshaledSize() int {
	return self.predExpBase.MarshaledSize() + 4
}

func (self *predExpStringRegex) Marshal(cmd *baseCommand) error {
	self.MarshalTL(cmd, AS_PREDEXP_STRING_REGEX, 4)
	cmd.WriteUint32(self.cflags)
	return nil
}
<|endoftext|>


<commit_before>package gorm

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"reflect"
)

func getRealValue(value reflect.Value, field string) interface{} {
	result := reflect.Indirect(value).FieldByName(field).Interface()
	if r, ok := result.(driver.Valuer); ok {
		result, _ = r.Value()
	}
	return result
}

func equalAsString(a interface{}, b interface{}) bool {
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func Preload(scope *Scope) {
	fields := scope.Fields()
	isSlice := scope.IndirectValue().Kind() == reflect.Slice

	if scope.Search.Preload != nil {
		for key, conditions := range scope.Search.Preload {
			for _, field := range fields {
				if field.Name == key && field.Relationship != nil {
					results := makeSlice(field.Struct.Type)
					relation := field.Relationship
					primaryName := scope.PrimaryKeyField().Name
					associationPrimaryKey := scope.New(results).PrimaryKeyField().Name

					switch relation.Kind {
					case "has_one":
						condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
						scope.NewDB().Where(condition, scope.getColumnAsArray(primaryName)).Find(results, conditions...)

						resultValues := reflect.Indirect(reflect.ValueOf(results))
						for i := 0; i < resultValues.Len(); i++ {
							result := resultValues.Index(i)
							if isSlice {
								value := getRealValue(result, relation.ForeignFieldName)
								objects := scope.IndirectValue()
								for j := 0; j < objects.Len(); j++ {
									if equalAsString(getRealValue(objects.Index(j), primaryName), value) {
										reflect.Indirect(objects.Index(j)).FieldByName(field.Name).Set(result)
										break
									}
								}
							} else {
								scope.SetColumn(field, result)
							}
						}
					case "has_many":
						condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
						scope.NewDB().Where(condition, scope.getColumnAsArray(primaryName)).Find(results, conditions...)
						resultValues := reflect.Indirect(reflect.ValueOf(results))
						if isSlice {
							for i := 0; i < resultValues.Len(); i++ {
								result := resultValues.Index(i)
								value := getRealValue(result, relation.ForeignFieldName)
								objects := scope.IndirectValue()
								for j := 0; j < objects.Len(); j++ {
									object := reflect.Indirect(objects.Index(j))
									if equalAsString(getRealValue(object, primaryName), value) {
										f := object.FieldByName(field.Name)
										f.Set(reflect.Append(f, result))
										break
									}
								}
							}
						} else {
							scope.SetColumn(field, resultValues)
						}
					case "belongs_to":
						scope.NewDB().Where(scope.getColumnAsArray(relation.ForeignFieldName)).Find(results, conditions...)
						resultValues := reflect.Indirect(reflect.ValueOf(results))
						for i := 0; i < resultValues.Len(); i++ {
							result := resultValues.Index(i)
							if isSlice {
								value := getRealValue(result, associationPrimaryKey)
								objects := scope.IndirectValue()
								for j := 0; j < objects.Len(); j++ {
									object := reflect.Indirect(objects.Index(j))
									if equalAsString(getRealValue(object, relation.ForeignFieldName), value) {
										object.FieldByName(field.Name).Set(result)
									}
								}
							} else {
								scope.SetColumn(field, result)
							}
						}
					case "many_to_many":
						scope.Err(errors.New("not supported relation"))
					default:
						scope.Err(errors.New("not supported relation"))
					}
					break
				}
			}
		}
	}
}

func makeSlice(typ reflect.Type) interface{} {
	if typ.Kind() == reflect.Slice {
		typ = typ.Elem()
	}
	sliceType := reflect.SliceOf(typ)
	slice := reflect.New(sliceType)
	slice.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
	return slice.Interface()
}

func (scope *Scope) getColumnAsArray(column string) (primaryKeys []interface{}) {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		for i := 0; i < values.Len(); i++ {
			primaryKeys = append(primaryKeys, reflect.Indirect(values.Index(i)).FieldByName(column).Interface())
		}
	case reflect.Struct:
		return []interface{}{values.FieldByName(column).Interface()}
	}
	return
}
<commit_msg>Uniq foreign key for Preload<commit_after>package gorm

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"reflect"
)

func getRealValue(value reflect.Value, field string) interface{} {
	result := reflect.Indirect(value).FieldByName(field).Interface()
	if r, ok := result.(driver.Valuer); ok {
		result, _ = r.Value()
	}
	return result
}

func equalAsString(a interface{}, b interface{}) bool {
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func Preload(scope *Scope) {
	fields := scope.Fields()
	isSlice := scope.IndirectValue().Kind() == reflect.Slice

	if scope.Search.Preload != nil {
		for key, conditions := range scope.Search.Preload {
			for _, field := range fields {
				if field.Name == key && field.Relationship != nil {
					results := makeSlice(field.Struct.Type)
					relation := field.Relationship
					primaryName := scope.PrimaryKeyField().Name
					associationPrimaryKey := scope.New(results).PrimaryKeyField().Name

					switch relation.Kind {
					case "has_one":
						condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
						scope.NewDB().Where(condition, scope.getColumnAsArray(primaryName)).Find(results, conditions...)

						resultValues := reflect.Indirect(reflect.ValueOf(results))
						for i := 0; i < resultValues.Len(); i++ {
							result := resultValues.Index(i)
							if isSlice {
								value := getRealValue(result, relation.ForeignFieldName)
								objects := scope.IndirectValue()
								for j := 0; j < objects.Len(); j++ {
									if equalAsString(getRealValue(objects.Index(j), primaryName), value) {
										reflect.Indirect(objects.Index(j)).FieldByName(field.Name).Set(result)
										break
									}
								}
							} else {
								scope.SetColumn(field, result)
							}
						}
					case "has_many":
						condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
						scope.NewDB().Where(condition, scope.getColumnAsArray(primaryName)).Find(results, conditions...)
						resultValues := reflect.Indirect(reflect.ValueOf(results))
						if isSlice {
							for i := 0; i < resultValues.Len(); i++ {
								result := resultValues.Index(i)
								value := getRealValue(result, relation.ForeignFieldName)
								objects := scope.IndirectValue()
								for j := 0; j < objects.Len(); j++ {
									object := reflect.Indirect(objects.Index(j))
									if equalAsString(getRealValue(object, primaryName), value) {
										f := object.FieldByName(field.Name)
										f.Set(reflect.Append(f, result))
										break
									}
								}
							}
						} else {
							scope.SetColumn(field, resultValues)
						}
					case "belongs_to":
						scope.NewDB().Where(scope.getColumnAsArray(relation.ForeignFieldName)).Find(results, conditions...)
						resultValues := reflect.Indirect(reflect.ValueOf(results))
						for i := 0; i < resultValues.Len(); i++ {
							result := resultValues.Index(i)
							if isSlice {
								value := getRealValue(result, associationPrimaryKey)
								objects := scope.IndirectValue()
								for j := 0; j < objects.Len(); j++ {
									object := reflect.Indirect(objects.Index(j))
									if equalAsString(getRealValue(object, relation.ForeignFieldName), value) {
										object.FieldByName(field.Name).Set(result)
									}
								}
							} else {
								scope.SetColumn(field, result)
							}
						}
					case "many_to_many":
						scope.Err(errors.New("not supported relation"))
					default:
						scope.Err(errors.New("not supported relation"))
					}
					break
				}
			}
		}
	}
}

func makeSlice(typ reflect.Type) interface{} {
	if typ.Kind() == reflect.Slice {
		typ = typ.Elem()
	}
	sliceType := reflect.SliceOf(typ)
	slice := reflect.New(sliceType)
	slice.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
	return slice.Interface()
}

func (scope *Scope) getColumnAsArray(column string) (primaryKeys []interface{}) {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		primaryKeyMap := map[interface{}]bool{}
		for i := 0; i < values.Len(); i++ {
			primaryKeyMap[reflect.Indirect(values.Index(i)).FieldByName(column).Interface()] = true
		}
		for key := range primaryKeyMap {
			primaryKeys = append(primaryKeys, key)
		}
	case reflect.Struct:
		return []interface{}{values.FieldByName(column).Interface()}
	}
	return
}
<|endoftext|>


<commit_before>package pegasus

import (
	"github.com/HearthSim/hs-proto/go"
	"github.com/golang/protobuf/proto"
	"log"
	"time"
)

type Account struct{}

func (v *Account) Init(sess *Session) {
	sess.RegisterUtilHandler(0, 201, OnGetAccountInfo)
	sess.RegisterUtilHandler(0, 205, OnUpdateLogin)
	sess.RegisterUtilHandler(0, 239, OnSetOptions)
	sess.RegisterUtilHandler(0, 240, OnGetOptions)
	sess.RegisterUtilHandler(0, 253, OnGetAchieves)
	sess.RegisterUtilHandler(0, 267, OnCheckAccountLicenses)
	sess.RegisterUtilHandler(1, 276, OnCheckGameLicenses)
	sess.RegisterUtilHandler(0, 305, OnGetAdventureProgress)
}

func OnCheckAccountLicenses(s *Session, body []byte) ([]byte, error) {
	return OnCheckLicenses(true)
}

func OnCheckGameLicenses(s *Session, body []byte) ([]byte, error) {
	return OnCheckLicenses(false)
}

func OnCheckLicenses(accountLevel bool) ([]byte, error) {
	res := hsproto.PegasusUtil_CheckLicensesResponse{}
	res.AccountLevel = proto.Bool(accountLevel)
	res.Success = proto.Bool(true)
	return EncodeUtilResponse(277, &res)
}

func OnUpdateLogin(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_UpdateLogin{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	res := hsproto.PegasusUtil_UpdateLoginComplete{}
	return EncodeUtilResponse(307, &res)
}

func OnGetAccountInfo(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_GetAccountInfo{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	switch req.Request.String() {
	case "CAMPAIGN_INFO":
		res := hsproto.PegasusUtil_ProfileProgress{}
		res.Progress = proto.Int64(6)  // ILLIDAN_COMPLETE
		res.BestForge = proto.Int32(0) // Arena wins
		return EncodeUtilResponse(233, &res)
	case "BOOSTERS":
		res := hsproto.PegasusUtil_BoosterList{}
		return EncodeUtilResponse(224, &res)
	case "FEATURES":
		res := hsproto.PegasusUtil_GuardianVars{}
		res.ShowUserUI = proto.Int32(1)
		return EncodeUtilResponse(264, &res)
	case "MEDAL_INFO":
		res := hsproto.PegasusUtil_MedalInfo{}
		res.SeasonWins = proto.Int32(0)
		res.Stars = proto.Int32(2)
		res.Streak = proto.Int32(0)
		res.StarLevel = proto.Int32(1)
		res.LevelStart = proto.Int32(1)
		res.LevelEnd = proto.Int32(3)
		res.CanLose = proto.Bool(false)
		return EncodeUtilResponse(232, &res)
	case "MEDAL_HISTORY":
		res := hsproto.PegasusUtil_MedalHistory{}
		for i := int32(1); i <= 3; i++ {
			info := &hsproto.PegasusUtil_MedalHistoryInfo{}
			info.When = PegasusDate(time.Date(2015, 8, 1, 7, 0, 0, 0, time.UTC))
			info.Season = proto.Int32(i)
			info.Stars = proto.Int32(0)
			info.StarLevel = proto.Int32(0)
			info.LevelStart = proto.Int32(0)
			info.LevelEnd = proto.Int32(0)
			info.LegendRank = proto.Int32(1)
			res.Medals = append(res.Medals, info)
		}
		return EncodeUtilResponse(234, &res)
	case "NOTICES":
		res := hsproto.PegasusUtil_ProfileNotices{}
		return EncodeUtilResponse(212, &res)
	case "DECK_LIST":
		res := hsproto.PegasusUtil_DeckList{}
		for i := 2; i <= 10; i++ {
			info := &hsproto.PegasusShared_DeckInfo{}
			info.Id = proto.Int64(int64(1000 + i))
			info.Name = proto.String("precon")
			info.CardBack = proto.Int32(0)
			info.Hero = proto.Int32(int32(heroIdToAssetId[i]))
			precon := hsproto.PegasusShared_DeckType_PRECON_DECK
			info.DeckType = &precon
			info.Validity = proto.Uint64(31)
			info.HeroPremium = proto.Int32(0)
			info.CardBackOverride = proto.Bool(false)
			info.HeroOverride = proto.Bool(false)
			res.Decks = append(res.Decks, info)
		}
		return EncodeUtilResponse(202, &res)
	case "COLLECTION":
		res := hsproto.PegasusUtil_Collection{}
		return EncodeUtilResponse(207, &res)
	case "DECK_LIMIT":
		res := hsproto.PegasusUtil_ProfileDeckLimit{}
		res.DeckLimit = proto.Int32(9)
		return EncodeUtilResponse(231, &res)
	case "CARD_VALUES":
		res := hsproto.PegasusUtil_CardValues{}
		res.CardNerfIndex = proto.Int32(0)
		return EncodeUtilResponse(260, &res)
	case "ARCANE_DUST_BALANCE":
		res := hsproto.PegasusUtil_ArcaneDustBalance{}
		res.Balance = proto.Int64(10000)
		return EncodeUtilResponse(262, &res)
	case "GOLD_BALANCE":
		res := hsproto.PegasusUtil_GoldBalance{}
		res.Cap = proto.Int64(999999)
		res.CapWarning = proto.Int64(2000)
		res.CappedBalance = proto.Int64(1234)
		res.BonusBalance = proto.Int64(0)
		return EncodeUtilResponse(278, &res)
	case "HERO_XP":
		res := hsproto.PegasusUtil_HeroXP{}
		for i := 2; i <= 10; i++ {
			info := &hsproto.PegasusUtil_HeroXPInfo{}
			level := 2*i + 5
			maxXp := 60 + level*10
			info.ClassId = proto.Int32(int32(i))
			info.Level = proto.Int32(int32(level))
			info.CurrXp = proto.Int64(int64(maxXp / 2))
			info.MaxXp = proto.Int64(int64(maxXp))
			res.XpInfos = append(res.XpInfos, info)
		}
		return EncodeUtilResponse(283, &res)
	case "NOT_SO_MASSIVE_LOGIN":
		res := hsproto.PegasusUtil_NotSoMassiveLoginReply{}
		return EncodeUtilResponse(300, &res)
	case "REWARD_PROGRESS":
		res := hsproto.PegasusUtil_RewardProgress{}
		nextMonth := time.Date(2015, 8, 1, 7, 0, 0, 0, time.UTC)
		res.SeasonEnd = PegasusDate(nextMonth)
		res.WinsPerGold = proto.Int32(3)
		res.GoldPerReward = proto.Int32(10)
		res.MaxGoldPerDay = proto.Int32(100)
		res.SeasonNumber = proto.Int32(21)
		res.XpSoloLimit = proto.Int32(60)
		res.MaxHeroLevel = proto.Int32(60)
		res.NextQuestCancel = PegasusDate(time.Now().UTC())
		res.EventTimingMod = proto.Float32(0.291667)
		return EncodeUtilResponse(271, &res)
	case "PVP_QUEUE":
		res := hsproto.PegasusUtil_PlayQueue{}
		queue := hsproto.PegasusShared_PlayQueueInfo{}
		gametype := hsproto.PegasusShared_BnetGameType_BGT_NORMAL
		queue.GameType = &gametype
		res.Queue = &queue
		return EncodeUtilResponse(286, &res)

	case "PLAYER_RECORD":
		res := hsproto.PegasusUtil_PlayerRecords{}
		return EncodeUtilResponse(270, &res)
	case "CARD_BACKS":
		res := hsproto.PegasusUtil_CardBacks{}
		res.DefaultCardBack = proto.Int32(13)
		res.CardBacks = []int32{0, 13, 24}
		return EncodeUtilResponse(236, &res)
	case "FAVORITE_HEROES":
		res := hsproto.PegasusUtil_FavoriteHeroesResponse{}
		return EncodeUtilResponse(318, &res)
	case "ACCOUNT_LICENSES":
		res := hsproto.PegasusUtil_AccountLicensesInfoResponse{}
		return EncodeUtilResponse(325, &res)
	case "BOOSTER_TALLY":
		res := hsproto.PegasusUtil_BoosterTallyList{}
		return EncodeUtilResponse(313, &res)
	default:

		return nil, nyi
	}
}

func OnGetAdventureProgress(s *Session, body []byte) ([]byte, error) {
	res := hsproto.PegasusUtil_AdventureProgressResponse{}
	return EncodeUtilResponse(306, &res)
}

func OnSetOptions(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_SetOptions{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	return nil, nil
}

func OnGetOptions(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_GetOptions{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	res := hsproto.PegasusUtil_ClientOptions{}
	res.Options = append(res.Options, &hsproto.PegasusUtil_ClientOption{
		Index:    proto.Int32(1),
		AsUint64: proto.Uint64(0x20FFFF3FFFCCFCFF),
	})
	res.Options = append(res.Options, &hsproto.PegasusUtil_ClientOption{
		Index:    proto.Int32(2),
		AsUint64: proto.Uint64(0xF0BFFFEF3FFF),
	})
	res.Options = append(res.Options, &hsproto.PegasusUtil_ClientOption{
		Index:   proto.Int32(18),
		AsInt64: proto.Int64(0xB765A8C),
	})
	return EncodeUtilResponse(241, &res)
}

func OnGetAchieves(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_GetAchieves{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	res := hsproto.PegasusUtil_Achieves{}
	for i := 1; i <= 9; i++ {
		info := &hsproto.PegasusUtil_Achieve{}
		info.Id = proto.Int32(int32(i))
		info.Progress = proto.Int32(1)
		info.AckProgress = proto.Int32(1)
		info.CompletionCount = proto.Int32(1)
		info.StartedCount = proto.Int32(1)
		info.DateGiven = PegasusDate(time.Now())
		res.List = append(res.List, info)
	}
	return EncodeUtilResponse(252, &res)
}

func PegasusDate(t time.Time) *hsproto.PegasusShared_Date {
	return &hsproto.PegasusShared_Date{
		Year:  proto.Int32(int32(t.Year())),
		Month: proto.Int32(int32(t.Month())),
		Day:   proto.Int32(int32(t.Day())),
		Hours: proto.Int32(int32(t.Hour())),
		Min:   proto.Int32(int32(t.Minute())),
		Sec:   proto.Int32(int32(t.Second())),
	}
}

// A map from TAG_CLASS ids to DBF ids
var heroIdToAssetId = map[int]int{
	2:  274,
	3:  31,
	4:  637,
	5:  671,
	6:  813,
	7:  930,
	8:  1066,
	9:  893,
	10: 7,
}
<commit_msg>Improve FAVORITE_HEROES stub<commit_after>package pegasus

import (
	"github.com/HearthSim/hs-proto/go"
	"github.com/golang/protobuf/proto"
	"log"
	"time"
)

type Account struct{}

func (v *Account) Init(sess *Session) {
	sess.RegisterUtilHandler(0, 201, OnGetAccountInfo)
	sess.RegisterUtilHandler(0, 205, OnUpdateLogin)
	sess.RegisterUtilHandler(0, 239, OnSetOptions)
	sess.RegisterUtilHandler(0, 240, OnGetOptions)
	sess.RegisterUtilHandler(0, 253, OnGetAchieves)
	sess.RegisterUtilHandler(0, 267, OnCheckAccountLicenses)
	sess.RegisterUtilHandler(1, 276, OnCheckGameLicenses)
	sess.RegisterUtilHandler(0, 305, OnGetAdventureProgress)
}

func OnCheckAccountLicenses(s *Session, body []byte) ([]byte, error) {
	return OnCheckLicenses(true)
}

func OnCheckGameLicenses(s *Session, body []byte) ([]byte, error) {
	return OnCheckLicenses(false)
}

func OnCheckLicenses(accountLevel bool) ([]byte, error) {
	res := hsproto.PegasusUtil_CheckLicensesResponse{}
	res.AccountLevel = proto.Bool(accountLevel)
	res.Success = proto.Bool(true)
	return EncodeUtilResponse(277, &res)
}

func OnUpdateLogin(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_UpdateLogin{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	res := hsproto.PegasusUtil_UpdateLoginComplete{}
	return EncodeUtilResponse(307, &res)
}

func OnGetAccountInfo(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_GetAccountInfo{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	switch req.Request.String() {
	case "CAMPAIGN_INFO":
		res := hsproto.PegasusUtil_ProfileProgress{}
		res.Progress = proto.Int64(6)  // ILLIDAN_COMPLETE
		res.BestForge = proto.Int32(0) // Arena wins
		return EncodeUtilResponse(233, &res)
	case "BOOSTERS":
		res := hsproto.PegasusUtil_BoosterList{}
		return EncodeUtilResponse(224, &res)
	case "FEATURES":
		res := hsproto.PegasusUtil_GuardianVars{}
		res.ShowUserUI = proto.Int32(1)
		return EncodeUtilResponse(264, &res)
	case "MEDAL_INFO":
		res := hsproto.PegasusUtil_MedalInfo{}
		res.SeasonWins = proto.Int32(0)
		res.Stars = proto.Int32(2)
		res.Streak = proto.Int32(0)
		res.StarLevel = proto.Int32(1)
		res.LevelStart = proto.Int32(1)
		res.LevelEnd = proto.Int32(3)
		res.CanLose = proto.Bool(false)
		return EncodeUtilResponse(232, &res)
	case "MEDAL_HISTORY":
		res := hsproto.PegasusUtil_MedalHistory{}
		for i := int32(1); i <= 3; i++ {
			info := &hsproto.PegasusUtil_MedalHistoryInfo{}
			info.When = PegasusDate(time.Date(2015, 8, 1, 7, 0, 0, 0, time.UTC))
			info.Season = proto.Int32(i)
			info.Stars = proto.Int32(0)
			info.StarLevel = proto.Int32(0)
			info.LevelStart = proto.Int32(0)
			info.LevelEnd = proto.Int32(0)
			info.LegendRank = proto.Int32(1)
			res.Medals = append(res.Medals, info)
		}
		return EncodeUtilResponse(234, &res)
	case "NOTICES":
		res := hsproto.PegasusUtil_ProfileNotices{}
		return EncodeUtilResponse(212, &res)
	case "DECK_LIST":
		res := hsproto.PegasusUtil_DeckList{}
		for i := 2; i <= 10; i++ {
			info := &hsproto.PegasusShared_DeckInfo{}
			info.Id = proto.Int64(int64(1000 + i))
			info.Name = proto.String("precon")
			info.CardBack = proto.Int32(0)
			info.Hero = proto.Int32(int32(heroIdToAssetId[i]))
			precon := hsproto.PegasusShared_DeckType_PRECON_DECK
			info.DeckType = &precon
			info.Validity = proto.Uint64(31)
			info.HeroPremium = proto.Int32(0)
			info.CardBackOverride = proto.Bool(false)
			info.HeroOverride = proto.Bool(false)
			res.Decks = append(res.Decks, info)
		}
		return EncodeUtilResponse(202, &res)
	case "COLLECTION":
		res := hsproto.PegasusUtil_Collection{}
		return EncodeUtilResponse(207, &res)
	case "DECK_LIMIT":
		res := hsproto.PegasusUtil_ProfileDeckLimit{}
		res.DeckLimit = proto.Int32(9)
		return EncodeUtilResponse(231, &res)
	case "CARD_VALUES":
		res := hsproto.PegasusUtil_CardValues{}
		res.CardNerfIndex = proto.Int32(0)
		return EncodeUtilResponse(260, &res)
	case "ARCANE_DUST_BALANCE":
		res := hsproto.PegasusUtil_ArcaneDustBalance{}
		res.Balance = proto.Int64(10000)
		return EncodeUtilResponse(262, &res)
	case "GOLD_BALANCE":
		res := hsproto.PegasusUtil_GoldBalance{}
		res.Cap = proto.Int64(999999)
		res.CapWarning = proto.Int64(2000)
		res.CappedBalance = proto.Int64(1234)
		res.BonusBalance = proto.Int64(0)
		return EncodeUtilResponse(278, &res)
	case "HERO_XP":
		res := hsproto.PegasusUtil_HeroXP{}
		for i := 2; i <= 10; i++ {
			info := &hsproto.PegasusUtil_HeroXPInfo{}
			level := 2*i + 5
			maxXp := 60 + level*10
			info.ClassId = proto.Int32(int32(i))
			info.Level = proto.Int32(int32(level))
			info.CurrXp = proto.Int64(int64(maxXp / 2))
			info.MaxXp = proto.Int64(int64(maxXp))
			res.XpInfos = append(res.XpInfos, info)
		}
		return EncodeUtilResponse(283, &res)
	case "NOT_SO_MASSIVE_LOGIN":
		res := hsproto.PegasusUtil_NotSoMassiveLoginReply{}
		return EncodeUtilResponse(300, &res)
	case "REWARD_PROGRESS":
		res := hsproto.PegasusUtil_RewardProgress{}
		nextMonth := time.Date(2015, 8, 1, 7, 0, 0, 0, time.UTC)
		res.SeasonEnd = PegasusDate(nextMonth)
		res.WinsPerGold = proto.Int32(3)
		res.GoldPerReward = proto.Int32(10)
		res.MaxGoldPerDay = proto.Int32(100)
		res.SeasonNumber = proto.Int32(21)
		res.XpSoloLimit = proto.Int32(60)
		res.MaxHeroLevel = proto.Int32(60)
		res.NextQuestCancel = PegasusDate(time.Now().UTC())
		res.EventTimingMod = proto.Float32(0.291667)
		return EncodeUtilResponse(271, &res)
	case "PVP_QUEUE":
		res := hsproto.PegasusUtil_PlayQueue{}
		queue := hsproto.PegasusShared_PlayQueueInfo{}
		gametype := hsproto.PegasusShared_BnetGameType_BGT_NORMAL
		queue.GameType = &gametype
		res.Queue = &queue
		return EncodeUtilResponse(286, &res)

	case "PLAYER_RECORD":
		res := hsproto.PegasusUtil_PlayerRecords{}
		return EncodeUtilResponse(270, &res)
	case "CARD_BACKS":
		res := hsproto.PegasusUtil_CardBacks{}
		res.DefaultCardBack = proto.Int32(13)
		res.CardBacks = []int32{0, 13, 24}
		return EncodeUtilResponse(236, &res)
	case "FAVORITE_HEROES":
		res := hsproto.PegasusUtil_FavoriteHeroesResponse{}
		for i := 2; i <= 10; i++ {
			fav := &hsproto.PegasusShared_FavoriteHero{}
			fav.ClassId = proto.Int32(int32(i))
			carddef := &hsproto.PegasusShared_CardDef{}
			carddef.Asset = proto.Int32(int32(heroIdToAssetId[i]))
			fav.Hero = carddef
			res.FavoriteHeroes = append(res.FavoriteHeroes, fav)
		}
		return EncodeUtilResponse(318, &res)
	case "ACCOUNT_LICENSES":
		res := hsproto.PegasusUtil_AccountLicensesInfoResponse{}
		return EncodeUtilResponse(325, &res)
	case "BOOSTER_TALLY":
		res := hsproto.PegasusUtil_BoosterTallyList{}
		return EncodeUtilResponse(313, &res)
	default:

		return nil, nyi
	}
}

func OnGetAdventureProgress(s *Session, body []byte) ([]byte, error) {
	res := hsproto.PegasusUtil_AdventureProgressResponse{}
	return EncodeUtilResponse(306, &res)
}

func OnSetOptions(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_SetOptions{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	return nil, nil
}

func OnGetOptions(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_GetOptions{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	res := hsproto.PegasusUtil_ClientOptions{}
	res.Options = append(res.Options, &hsproto.PegasusUtil_ClientOption{
		Index:    proto.Int32(1),
		AsUint64: proto.Uint64(0x20FFFF3FFFCCFCFF),
	})
	res.Options = append(res.Options, &hsproto.PegasusUtil_ClientOption{
		Index:    proto.Int32(2),
		AsUint64: proto.Uint64(0xF0BFFFEF3FFF),
	})
	res.Options = append(res.Options, &hsproto.PegasusUtil_ClientOption{
		Index:   proto.Int32(18),
		AsInt64: proto.Int64(0xB765A8C),
	})
	return EncodeUtilResponse(241, &res)
}

func OnGetAchieves(s *Session, body []byte) ([]byte, error) {
	req := hsproto.PegasusUtil_GetAchieves{}
	err := proto.Unmarshal(body, &req)
	if err != nil {
		return nil, err
	}
	log.Printf("req = %s", req.String())
	res := hsproto.PegasusUtil_Achieves{}
	for i := 1; i <= 9; i++ {
		info := &hsproto.PegasusUtil_Achieve{}
		info.Id = proto.Int32(int32(i))
		info.Progress = proto.Int32(1)
		info.AckProgress = proto.Int32(1)
		info.CompletionCount = proto.Int32(1)
		info.StartedCount = proto.Int32(1)
		info.DateGiven = PegasusDate(time.Now())
		res.List = append(res.List, info)
	}
	return EncodeUtilResponse(252, &res)
}

func PegasusDate(t time.Time) *hsproto.PegasusShared_Date {
	return &hsproto.PegasusShared_Date{
		Year:  proto.Int32(int32(t.Year())),
		Month: proto.Int32(int32(t.Month())),
		Day:   proto.Int32(int32(t.Day())),
		Hours: proto.Int32(int32(t.Hour())),
		Min:   proto.Int32(int32(t.Minute())),
		Sec:   proto.Int32(int32(t.Second())),
	}
}

// A map from TAG_CLASS ids to DBF ids
var heroIdToAssetId = map[int]int{
	2:  274,
	3:  31,
	4:  637,
	5:  671,
	6:  813,
	7:  930,
	8:  1066,
	9:  893,
	10: 7,
}
<|endoftext|>


<commit_before>package rados

// #cgo LDFLAGS: -lrados
// #include <stdlib.h>
// #include <rados/librados.h>
import "C"

import (
	"bytes"
	"errors"
	"unsafe"
)

var (
	ErrNotConnected = errors.New("RADOS not connected")
)

// ClusterStat represents Ceph cluster statistics.
type ClusterStat struct {
	Kb          uint64
	Kb_used     uint64
	Kb_avail    uint64
	Num_objects uint64
}

// Conn is a connection handle to a Ceph cluster.
type Conn struct {
	cluster   C.rados_t
	connected bool
}

// PingMonitor sends a ping to a monitor and returns the reply.
func (c *Conn) PingMonitor(id string) (string, error) {
	c_id := C.CString(id)
	defer C.free(unsafe.Pointer(c_id))

	var strlen C.size_t
	var strout *C.char

	ret := C.rados_ping_monitor(c.cluster, c_id, &strout, &strlen)
	defer C.rados_buffer_free(strout)

	if ret == 0 {
		reply := C.GoStringN(strout, (C.int)(strlen))
		return reply, nil
	} else {
		return "", RadosError(int(ret))
	}
}

// Connect establishes a connection to a RADOS cluster. It returns an error,
// if any.
func (c *Conn) Connect() error {
	ret := C.rados_connect(c.cluster)
	if ret != 0 {
		return RadosError(int(ret))
	}
	c.connected = true
	return nil
}

// Shutdown disconnects from the cluster.
func (c *Conn) Shutdown() {
	if err := c.ensure_connected(); err != nil {
		return
	}
	freeConn(c)
}

// ReadConfigFile configures the connection using a Ceph configuration file.
func (c *Conn) ReadConfigFile(path string) error {
	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))
	ret := C.rados_conf_read_file(c.cluster, c_path)
	if ret == 0 {
		return nil
	} else {
		return RadosError(int(ret))
	}
}

// ReadDefaultConfigFile configures the connection using a Ceph configuration
// file located at default locations.
func (c *Conn) ReadDefaultConfigFile() error {
	ret := C.rados_conf_read_file(c.cluster, nil)
	if ret == 0 {
		return nil
	} else {
		return RadosError(int(ret))
	}
}

func (c *Conn) OpenIOContext(pool string) (*IOContext, error) {
	c_pool := C.CString(pool)
	defer C.free(unsafe.Pointer(c_pool))
	ioctx := &IOContext{}
	ret := C.rados_ioctx_create(c.cluster, c_pool, &ioctx.ioctx)
	if ret == 0 {
		return ioctx, nil
	} else {
		return nil, RadosError(int(ret))
	}
}

// ListPools returns the names of all existing pools.
func (c *Conn) ListPools() (names []string, err error) {
	buf := make([]byte, 4096)
	for {
		ret := int(C.rados_pool_list(c.cluster,
			(*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
		if ret < 0 {
			return nil, RadosError(int(ret))
		}

		if ret > len(buf) {
			buf = make([]byte, ret)
			continue
		}

		tmp := bytes.SplitAfter(buf[:ret-1], []byte{0})
		for _, s := range tmp {
			if len(s) > 0 {
				name := C.GoString((*C.char)(unsafe.Pointer(&s[0])))
				names = append(names, name)
			}
		}

		return names, nil
	}
}

// SetConfigOption sets the value of the configuration option identified by
// the given name.
func (c *Conn) SetConfigOption(option, value string) error {
	c_opt, c_val := C.CString(option), C.CString(value)
	defer C.free(unsafe.Pointer(c_opt))
	defer C.free(unsafe.Pointer(c_val))
	ret := C.rados_conf_set(c.cluster, c_opt, c_val)
	if ret < 0 {
		return RadosError(int(ret))
	} else {
		return nil
	}
}

// GetConfigOption returns the value of the Ceph configuration option
// identified by the given name.
func (c *Conn) GetConfigOption(name string) (value string, err error) {
	buf := make([]byte, 4096)
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int(C.rados_conf_get(c.cluster, c_name,
		(*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
	// FIXME: ret may be -ENAMETOOLONG if the buffer is not large enough. We
	// can handle this case, but we need a reliable way to test for
	// -ENAMETOOLONG constant. Will the syscall/Errno stuff in Go help?
	if ret == 0 {
		value = C.GoString((*C.char)(unsafe.Pointer(&buf[0])))
		return value, nil
	} else {
		return "", RadosError(ret)
	}
}

// WaitForLatestOSDMap blocks the caller until the latest OSD map has been
// retrieved.
func (c *Conn) WaitForLatestOSDMap() error {
	ret := C.rados_wait_for_latest_osdmap(c.cluster)
	if ret < 0 {
		return RadosError(int(ret))
	} else {
		return nil
	}
}

func (c *Conn) ensure_connected() error {
	if c.connected {
		return nil
	} else {
		return ErrNotConnected
	}
}

// GetClusterStats returns statistics about the cluster associated with the
// connection.
func (c *Conn) GetClusterStats() (stat ClusterStat, err error) {
	if err := c.ensure_connected(); err != nil {
		return ClusterStat{}, err
	}
	c_stat := C.struct_rados_cluster_stat_t{}
	ret := C.rados_cluster_stat(c.cluster, &c_stat)
	if ret < 0 {
		return ClusterStat{}, RadosError(int(ret))
	} else {
		return ClusterStat{
			Kb:          uint64(c_stat.kb),
			Kb_used:     uint64(c_stat.kb_used),
			Kb_avail:    uint64(c_stat.kb_avail),
			Num_objects: uint64(c_stat.num_objects),
		}, nil
	}
}

// ParseCmdLineArgs configures the connection from command line arguments.
func (c *Conn) ParseCmdLineArgs(args []string) error {
	// add an empty element 0 -- Ceph treats the array as the actual contents
	// of argv and skips the first element (the executable name)
	argc := C.int(len(args) + 1)
	argv := make([]*C.char, argc)

	// make the first element a string just in case it is ever examined
	argv[0] = C.CString("placeholder")
	defer C.free(unsafe.Pointer(argv[0]))

	for i, arg := range args {
		argv[i+1] = C.CString(arg)
		defer C.free(unsafe.Pointer(argv[i+1]))
	}

	ret := C.rados_conf_parse_argv(c.cluster, argc, &argv[0])
	if ret < 0 {
		return RadosError(int(ret))
	} else {
		return nil
	}
}

// ParseDefaultConfigEnv configures the connection from the default Ceph
// environment variable(s).
func (c *Conn) ParseDefaultConfigEnv() error {
	ret := C.rados_conf_parse_env(c.cluster, nil)
	if ret == 0 {
		return nil
	} else {
		return RadosError(int(ret))
	}
}

// GetFSID returns the fsid of the cluster as a hexadecimal string. The fsid
// is a unique identifier of an entire Ceph cluster.
func (c *Conn) GetFSID() (fsid string, err error) {
	buf := make([]byte, 37)
	ret := int(C.rados_cluster_fsid(c.cluster,
		(*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
	// FIXME: the success case isn't documented correctly in librados.h
	if ret == 36 {
		fsid = C.GoString((*C.char)(unsafe.Pointer(&buf[0])))
		return fsid, nil
	} else {
		return "", RadosError(int(ret))
	}
}

// GetInstanceID returns a globally unique identifier for the cluster
// connection instance.
func (c *Conn) GetInstanceID() uint64 {
	// FIXME: are there any error cases for this?
	return uint64(C.rados_get_instance_id(c.cluster))
}

// MakePool creates a new pool with default settings.
func (c *Conn) MakePool(name string) error {
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int(C.rados_pool_create(c.cluster, c_name))
	if ret == 0 {
		return nil
	} else {
		return RadosError(ret)
	}
}

// DeletePool deletes a pool and all the data inside the pool.
func (c *Conn) DeletePool(name string) error {
	if err := c.ensure_connected(); err != nil {
		return err
	}
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int(C.rados_pool_delete(c.cluster, c_name))
	if ret == 0 {
		return nil
	} else {
		return RadosError(ret)
	}
}

// GetPoolByName returns the ID of the pool with a given name.
func (c *Conn) GetPoolByName(name string) (int64, error) {
	if err := c.ensure_connected(); err != nil {
		return 0, err
	}
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int64(C.rados_pool_lookup(c.cluster, c_name))
	if ret < 0 {
		return 0, RadosError(ret)
	} else {
		return ret, nil
	}
}

// GetPoolByID returns the name of a pool by a given ID.
func (c *Conn) GetPoolByID(id int64) (string, error) {
	buf := make([]byte, 4096)
	if err := c.ensure_connected(); err != nil {
		return "", err
	}
	c_id := C.int64_t(id)
	ret := int(C.rados_pool_reverse_lookup(c.cluster, c_id, (*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
	if ret < 0 {
		return "", RadosError(ret)
	} else {
		return C.GoString((*C.char)(unsafe.Pointer(&buf[0]))), nil
	}
}

// MonCommand sends a command to one of the monitors
func (c *Conn) MonCommand(args []byte) (buffer []byte, info string, err error) {
	return c.monCommand(args, nil)
}

// MonCommandWithInputBuffer sends a command to one of the monitors, with an input buffer
func (c *Conn) MonCommandWithInputBuffer(args, inputBuffer []byte) (buffer []byte, info string, err error) {
	return c.monCommand(args, inputBuffer)
}

func (c *Conn) monCommand(args, inputBuffer []byte) (buffer []byte, info string, err error) {
	argv := C.CString(string(args))
	defer C.free(unsafe.Pointer(argv))

	var (
		outs, outbuf       *C.char
		outslen, outbuflen C.size_t
	)
	inbuf := C.CString(string(inputBuffer))
	inbufLen := len(inputBuffer)
	defer C.free(unsafe.Pointer(inbuf))

	ret := C.rados_mon_command(c.cluster,
		&argv, 1,
		inbuf,              // bulk input (e.g. crush map)
		C.size_t(inbufLen), // length inbuf
		&outbuf,            // buffer
		&outbuflen,         // buffer length
		&outs,              // status string
		&outslen)

	if outslen > 0 {
		info = C.GoStringN(outs, C.int(outslen))
		C.free(unsafe.Pointer(outs))
	}
	if outbuflen > 0 {
		buffer = C.GoBytes(unsafe.Pointer(outbuf), C.int(outbuflen))
		C.free(unsafe.Pointer(outbuf))
	}
	if ret != 0 {
		err = RadosError(int(ret))
		return nil, info, err
	}

	return
}

// PGCommand sends a command to one of the PGs
func (c *Conn) PGCommand(pgid []byte, args [][]byte) (buffer []byte, info string, err error) {
	return c.pgCommand(pgid, args, nil)
}

// PGCommand sends a command to one of the PGs, with an input buffer
func (c *Conn) PGCommandWithInputBuffer(pgid []byte, args [][]byte, inputBuffer []byte) (buffer []byte, info string, err error) {
	return c.pgCommand(pgid, args, inputBuffer)
}

func (c *Conn) pgCommand(pgid []byte, args [][]byte, inputBuffer []byte) (buffer []byte, info string, err error) {
	name := C.CString(string(pgid))
	defer C.free(unsafe.Pointer(name))

	argc := len(args)
	argv := make([]*C.char, argc)

	for i, arg := range args {
		argv[i] = C.CString(string(arg))
		defer C.free(unsafe.Pointer(argv[i]))
	}

	var (
		outs, outbuf       *C.char
		outslen, outbuflen C.size_t
	)
	inbuf := C.CString(string(inputBuffer))
	inbufLen := len(inputBuffer)
	defer C.free(unsafe.Pointer(inbuf))

	ret := C.rados_pg_command(c.cluster,
		name,
		&argv[0],
		C.size_t(argc),
		inbuf,              // bulk input
		C.size_t(inbufLen), // length inbuf
		&outbuf,            // buffer
		&outbuflen,         // buffer length
		&outs,              // status string
		&outslen)

	if outslen > 0 {
		info = C.GoStringN(outs, C.int(outslen))
		C.free(unsafe.Pointer(outs))
	}
	if outbuflen > 0 {
		buffer = C.GoBytes(unsafe.Pointer(outbuf), C.int(outbuflen))
		C.free(unsafe.Pointer(outbuf))
	}
	if ret != 0 {
		err = RadosError(int(ret))
		return nil, info, err
	}

	return
}
<commit_msg>rados: use getRadosError where obviously applicable<commit_after>package rados

// #cgo LDFLAGS: -lrados
// #include <stdlib.h>
// #include <rados/librados.h>
import "C"

import (
	"bytes"
	"errors"
	"unsafe"
)

var (
	ErrNotConnected = errors.New("RADOS not connected")
)

// ClusterStat represents Ceph cluster statistics.
type ClusterStat struct {
	Kb          uint64
	Kb_used     uint64
	Kb_avail    uint64
	Num_objects uint64
}

// Conn is a connection handle to a Ceph cluster.
type Conn struct {
	cluster   C.rados_t
	connected bool
}

// PingMonitor sends a ping to a monitor and returns the reply.
func (c *Conn) PingMonitor(id string) (string, error) {
	c_id := C.CString(id)
	defer C.free(unsafe.Pointer(c_id))

	var strlen C.size_t
	var strout *C.char

	ret := C.rados_ping_monitor(c.cluster, c_id, &strout, &strlen)
	defer C.rados_buffer_free(strout)

	if ret == 0 {
		reply := C.GoStringN(strout, (C.int)(strlen))
		return reply, nil
	} else {
		return "", RadosError(int(ret))
	}
}

// Connect establishes a connection to a RADOS cluster. It returns an error,
// if any.
func (c *Conn) Connect() error {
	ret := C.rados_connect(c.cluster)
	if ret != 0 {
		return RadosError(int(ret))
	}
	c.connected = true
	return nil
}

// Shutdown disconnects from the cluster.
func (c *Conn) Shutdown() {
	if err := c.ensure_connected(); err != nil {
		return
	}
	freeConn(c)
}

// ReadConfigFile configures the connection using a Ceph configuration file.
func (c *Conn) ReadConfigFile(path string) error {
	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))
	ret := C.rados_conf_read_file(c.cluster, c_path)
	return getRadosError(int(ret))
}

// ReadDefaultConfigFile configures the connection using a Ceph configuration
// file located at default locations.
func (c *Conn) ReadDefaultConfigFile() error {
	ret := C.rados_conf_read_file(c.cluster, nil)
	return getRadosError(int(ret))
}

func (c *Conn) OpenIOContext(pool string) (*IOContext, error) {
	c_pool := C.CString(pool)
	defer C.free(unsafe.Pointer(c_pool))
	ioctx := &IOContext{}
	ret := C.rados_ioctx_create(c.cluster, c_pool, &ioctx.ioctx)
	if ret == 0 {
		return ioctx, nil
	} else {
		return nil, RadosError(int(ret))
	}
}

// ListPools returns the names of all existing pools.
func (c *Conn) ListPools() (names []string, err error) {
	buf := make([]byte, 4096)
	for {
		ret := int(C.rados_pool_list(c.cluster,
			(*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
		if ret < 0 {
			return nil, RadosError(int(ret))
		}

		if ret > len(buf) {
			buf = make([]byte, ret)
			continue
		}

		tmp := bytes.SplitAfter(buf[:ret-1], []byte{0})
		for _, s := range tmp {
			if len(s) > 0 {
				name := C.GoString((*C.char)(unsafe.Pointer(&s[0])))
				names = append(names, name)
			}
		}

		return names, nil
	}
}

// SetConfigOption sets the value of the configuration option identified by
// the given name.
func (c *Conn) SetConfigOption(option, value string) error {
	c_opt, c_val := C.CString(option), C.CString(value)
	defer C.free(unsafe.Pointer(c_opt))
	defer C.free(unsafe.Pointer(c_val))
	ret := C.rados_conf_set(c.cluster, c_opt, c_val)
	if ret < 0 {
		return RadosError(int(ret))
	} else {
		return nil
	}
}

// GetConfigOption returns the value of the Ceph configuration option
// identified by the given name.
func (c *Conn) GetConfigOption(name string) (value string, err error) {
	buf := make([]byte, 4096)
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int(C.rados_conf_get(c.cluster, c_name,
		(*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
	// FIXME: ret may be -ENAMETOOLONG if the buffer is not large enough. We
	// can handle this case, but we need a reliable way to test for
	// -ENAMETOOLONG constant. Will the syscall/Errno stuff in Go help?
	if ret == 0 {
		value = C.GoString((*C.char)(unsafe.Pointer(&buf[0])))
		return value, nil
	} else {
		return "", RadosError(ret)
	}
}

// WaitForLatestOSDMap blocks the caller until the latest OSD map has been
// retrieved.
func (c *Conn) WaitForLatestOSDMap() error {
	ret := C.rados_wait_for_latest_osdmap(c.cluster)
	if ret < 0 {
		return RadosError(int(ret))
	} else {
		return nil
	}
}

func (c *Conn) ensure_connected() error {
	if c.connected {
		return nil
	} else {
		return ErrNotConnected
	}
}

// GetClusterStats returns statistics about the cluster associated with the
// connection.
func (c *Conn) GetClusterStats() (stat ClusterStat, err error) {
	if err := c.ensure_connected(); err != nil {
		return ClusterStat{}, err
	}
	c_stat := C.struct_rados_cluster_stat_t{}
	ret := C.rados_cluster_stat(c.cluster, &c_stat)
	if ret < 0 {
		return ClusterStat{}, RadosError(int(ret))
	} else {
		return ClusterStat{
			Kb:          uint64(c_stat.kb),
			Kb_used:     uint64(c_stat.kb_used),
			Kb_avail:    uint64(c_stat.kb_avail),
			Num_objects: uint64(c_stat.num_objects),
		}, nil
	}
}

// ParseCmdLineArgs configures the connection from command line arguments.
func (c *Conn) ParseCmdLineArgs(args []string) error {
	// add an empty element 0 -- Ceph treats the array as the actual contents
	// of argv and skips the first element (the executable name)
	argc := C.int(len(args) + 1)
	argv := make([]*C.char, argc)

	// make the first element a string just in case it is ever examined
	argv[0] = C.CString("placeholder")
	defer C.free(unsafe.Pointer(argv[0]))

	for i, arg := range args {
		argv[i+1] = C.CString(arg)
		defer C.free(unsafe.Pointer(argv[i+1]))
	}

	ret := C.rados_conf_parse_argv(c.cluster, argc, &argv[0])
	if ret < 0 {
		return RadosError(int(ret))
	} else {
		return nil
	}
}

// ParseDefaultConfigEnv configures the connection from the default Ceph
// environment variable(s).
func (c *Conn) ParseDefaultConfigEnv() error {
	ret := C.rados_conf_parse_env(c.cluster, nil)
	return getRadosError(int(ret))
}

// GetFSID returns the fsid of the cluster as a hexadecimal string. The fsid
// is a unique identifier of an entire Ceph cluster.
func (c *Conn) GetFSID() (fsid string, err error) {
	buf := make([]byte, 37)
	ret := int(C.rados_cluster_fsid(c.cluster,
		(*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
	// FIXME: the success case isn't documented correctly in librados.h
	if ret == 36 {
		fsid = C.GoString((*C.char)(unsafe.Pointer(&buf[0])))
		return fsid, nil
	} else {
		return "", RadosError(int(ret))
	}
}

// GetInstanceID returns a globally unique identifier for the cluster
// connection instance.
func (c *Conn) GetInstanceID() uint64 {
	// FIXME: are there any error cases for this?
	return uint64(C.rados_get_instance_id(c.cluster))
}

// MakePool creates a new pool with default settings.
func (c *Conn) MakePool(name string) error {
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int(C.rados_pool_create(c.cluster, c_name))
	return getRadosError(int(ret))
}

// DeletePool deletes a pool and all the data inside the pool.
func (c *Conn) DeletePool(name string) error {
	if err := c.ensure_connected(); err != nil {
		return err
	}
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int(C.rados_pool_delete(c.cluster, c_name))
	return getRadosError(int(ret))
}

// GetPoolByName returns the ID of the pool with a given name.
func (c *Conn) GetPoolByName(name string) (int64, error) {
	if err := c.ensure_connected(); err != nil {
		return 0, err
	}
	c_name := C.CString(name)
	defer C.free(unsafe.Pointer(c_name))
	ret := int64(C.rados_pool_lookup(c.cluster, c_name))
	if ret < 0 {
		return 0, RadosError(ret)
	} else {
		return ret, nil
	}
}

// GetPoolByID returns the name of a pool by a given ID.
func (c *Conn) GetPoolByID(id int64) (string, error) {
	buf := make([]byte, 4096)
	if err := c.ensure_connected(); err != nil {
		return "", err
	}
	c_id := C.int64_t(id)
	ret := int(C.rados_pool_reverse_lookup(c.cluster, c_id, (*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))))
	if ret < 0 {
		return "", RadosError(ret)
	} else {
		return C.GoString((*C.char)(unsafe.Pointer(&buf[0]))), nil
	}
}

// MonCommand sends a command to one of the monitors
func (c *Conn) MonCommand(args []byte) (buffer []byte, info string, err error) {
	return c.monCommand(args, nil)
}

// MonCommandWithInputBuffer sends a command to one of the monitors, with an input buffer
func (c *Conn) MonCommandWithInputBuffer(args, inputBuffer []byte) (buffer []byte, info string, err error) {
	return c.monCommand(args, inputBuffer)
}

func (c *Conn) monCommand(args, inputBuffer []byte) (buffer []byte, info string, err error) {
	argv := C.CString(string(args))
	defer C.free(unsafe.Pointer(argv))

	var (
		outs, outbuf       *C.char
		outslen, outbuflen C.size_t
	)
	inbuf := C.CString(string(inputBuffer))
	inbufLen := len(inputBuffer)
	defer C.free(unsafe.Pointer(inbuf))

	ret := C.rados_mon_command(c.cluster,
		&argv, 1,
		inbuf,              // bulk input (e.g. crush map)
		C.size_t(inbufLen), // length inbuf
		&outbuf,            // buffer
		&outbuflen,         // buffer length
		&outs,              // status string
		&outslen)

	if outslen > 0 {
		info = C.GoStringN(outs, C.int(outslen))
		C.free(unsafe.Pointer(outs))
	}
	if outbuflen > 0 {
		buffer = C.GoBytes(unsafe.Pointer(outbuf), C.int(outbuflen))
		C.free(unsafe.Pointer(outbuf))
	}
	if ret != 0 {
		err = RadosError(int(ret))
		return nil, info, err
	}

	return
}

// PGCommand sends a command to one of the PGs
func (c *Conn) PGCommand(pgid []byte, args [][]byte) (buffer []byte, info string, err error) {
	return c.pgCommand(pgid, args, nil)
}

// PGCommand sends a command to one of the PGs, with an input buffer
func (c *Conn) PGCommandWithInputBuffer(pgid []byte, args [][]byte, inputBuffer []byte) (buffer []byte, info string, err error) {
	return c.pgCommand(pgid, args, inputBuffer)
}

func (c *Conn) pgCommand(pgid []byte, args [][]byte, inputBuffer []byte) (buffer []byte, info string, err error) {
	name := C.CString(string(pgid))
	defer C.free(unsafe.Pointer(name))

	argc := len(args)
	argv := make([]*C.char, argc)

	for i, arg := range args {
		argv[i] = C.CString(string(arg))
		defer C.free(unsafe.Pointer(argv[i]))
	}

	var (
		outs, outbuf       *C.char
		outslen, outbuflen C.size_t
	)
	inbuf := C.CString(string(inputBuffer))
	inbufLen := len(inputBuffer)
	defer C.free(unsafe.Pointer(inbuf))

	ret := C.rados_pg_command(c.cluster,
		name,
		&argv[0],
		C.size_t(argc),
		inbuf,              // bulk input
		C.size_t(inbufLen), // length inbuf
		&outbuf,            // buffer
		&outbuflen,         // buffer length
		&outs,              // status string
		&outslen)

	if outslen > 0 {
		info = C.GoStringN(outs, C.int(outslen))
		C.free(unsafe.Pointer(outs))
	}
	if outbuflen > 0 {
		buffer = C.GoBytes(unsafe.Pointer(outbuf), C.int(outbuflen))
		C.free(unsafe.Pointer(outbuf))
	}
	if ret != 0 {
		err = RadosError(int(ret))
		return nil, info, err
	}

	return
}
<|endoftext|>


<commit_before>package pt

import (
	"bytes"
	"net"
	"os"
	"sort"
	"testing"
)

func stringIsSafe(s string) bool {
	for _, c := range []byte(s) {
		if c == '\x00' || c == '\n' || c > 127 {
			return false
		}
	}
	return true
}

func TestEscape(t *testing.T) {
	tests := [...]string{
		"",
		"abc",
		"a\nb",
		"a\\b",
		"ab\\",
		"ab\\\n",
		"ab\n\\",
	}

	check := func(input string) {
		output := escape(input)
		if !stringIsSafe(output) {
			t.Errorf("escape(%q) → %q", input, output)
		}
	}
	for _, input := range tests {
		check(input)
	}
	for b := 0; b < 256; b++ {
		// check one-byte string with each byte value 0–255
		check(string([]byte{byte(b)}))
		// check UTF-8 encoding of each character 0–255
		check(string(b))
	}
}

func TestGetManagedTransportVer(t *testing.T) {
	badTests := [...]string{
		"",
		"2",
	}
	goodTests := [...]struct {
		input, expected string
	}{
		{"1", "1"},
		{"1,1", "1"},
		{"1,2", "1"},
		{"2,1", "1"},
	}

	os.Clearenv()
	_, err := getManagedTransportVer()
	if err == nil {
		t.Errorf("empty environment unexpectedly succeeded")
	}

	for _, input := range badTests {
		os.Setenv("TOR_PT_MANAGED_TRANSPORT_VER", input)
		_, err := getManagedTransportVer()
		if err == nil {
			t.Errorf("TOR_PT_MANAGED_TRANSPORT_VER=%q unexpectedly succeeded", input)
		}
	}

	for _, test := range goodTests {
		os.Setenv("TOR_PT_MANAGED_TRANSPORT_VER", test.input)
		output, err := getManagedTransportVer()
		if err != nil {
			t.Errorf("TOR_PT_MANAGED_TRANSPORT_VER=%q unexpectedly returned an error: %s", test.input, err)
		}
		if output != test.expected {
			t.Errorf("TOR_PT_MANAGED_TRANSPORT_VER=%q → %q (expected %q)", test.input, output, test.expected)
		}
	}
}

// return true iff the two slices contain the same elements, possibly in a
// different order.
func stringSetsEqual(a, b []string) bool {
	ac := make([]string, len(a))
	bc := make([]string, len(b))
	copy(ac, a)
	copy(bc, b)
	sort.Strings(ac)
	sort.Strings(bc)
	if len(ac) != len(bc) {
		return false
	}
	for i := 0; i < len(ac); i++ {
		if ac[i] != bc[i] {
			return false
		}
	}
	return true
}

func tcpAddrsEqual(a, b *net.TCPAddr) bool {
	return a.IP.Equal(b.IP) && a.Port == b.Port
}

func TestGetClientTransports(t *testing.T) {
	tests := [...]struct {
		ptServerClientTransports string
		methodNames              []string
		expected                 []string
	}{
		{
			"*",
			[]string{},
			[]string{},
		},
		{
			"*",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "beta", "gamma"},
		},
		{
			"alpha,beta,gamma",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "beta", "gamma"},
		},
		{
			"alpha,beta",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "beta"},
		},
		{
			"alpha",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha"},
		},
		{
			"alpha,beta",
			[]string{"alpha", "beta", "alpha"},
			[]string{"alpha", "beta"},
		},
		// my reading of pt-spec.txt says that "*" has special meaning
		// only when it is the entirety of the environment variable.
		{
			"alpha,*,gamma",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "gamma"},
		},
		{
			"alpha",
			[]string{"beta"},
			[]string{},
		},
	}

	os.Clearenv()
	_, err := getClientTransports([]string{"alpha", "beta", "gamma"})
	if err == nil {
		t.Errorf("empty environment unexpectedly succeeded")
	}

	for _, test := range tests {
		os.Setenv("TOR_PT_CLIENT_TRANSPORTS", test.ptServerClientTransports)
		output, err := getClientTransports(test.methodNames)
		if err != nil {
			t.Errorf("TOR_PT_CLIENT_TRANSPORTS=%q unexpectedly returned an error: %s",
				test.ptServerClientTransports, err)
		}
		if !stringSetsEqual(output, test.expected) {
			t.Errorf("TOR_PT_CLIENT_TRANSPORTS=%q %q → %q (expected %q)",
				test.ptServerClientTransports, test.methodNames, output, test.expected)
		}
	}
}

func TestResolveAddr(t *testing.T) {
	badTests := [...]string{
		"",
		"1.2.3.4",
		"1.2.3.4:",
		"9999",
		":9999",
		"[1:2::3:4]",
		"[1:2::3:4]:",
		"[1::2::3:4]",
		"1:2::3:4::9999",
		"1:2:3:4::9999",
		"localhost:9999",
		"[localhost]:9999",
	}
	goodTests := [...]struct {
		input    string
		expected net.TCPAddr
	}{
		{"1.2.3.4:9999", net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 9999}},
		{"[1:2::3:4]:9999", net.TCPAddr{IP: net.ParseIP("1:2::3:4"), Port: 9999}},
		{"1:2::3:4:9999", net.TCPAddr{IP: net.ParseIP("1:2::3:4"), Port: 9999}},
	}

	for _, input := range badTests {
		output, err := resolveAddr(input)
		if err == nil {
			t.Errorf("%q unexpectedly succeeded: %q", input, output)
		}
	}

	for _, test := range goodTests {
		output, err := resolveAddr(test.input)
		if err != nil {
			t.Errorf("%q unexpectedly returned an error: %s", test.input, err)
		}
		if !tcpAddrsEqual(output, &test.expected) {
			t.Errorf("%q → %q (expected %q)", test.input, output, test.expected)
		}
	}
}

func TestReadAuthCookie(t *testing.T) {
	badTests := [...][]byte{
		[]byte(""),
		// bad header
		[]byte("! Impostor ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDEF"),
		// too short
		[]byte("! Extended ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDE"),
		// too long
		[]byte("! Extended ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDEFX"),
	}
	goodTests := [...][]byte{
		[]byte("! Extended ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDEF"),
	}

	for _, input := range badTests {
		var buf bytes.Buffer
		buf.Write(input)
		_, err := readAuthCookie(&buf)
		if err == nil {
			t.Errorf("%q unexpectedly succeeded", input)
		}
	}

	for _, input := range goodTests {
		var buf bytes.Buffer
		buf.Write(input)
		cookie, err := readAuthCookie(&buf)
		if err != nil {
			t.Errorf("%q unexpectedly returned an error: %s", input, err)
		}
		if !bytes.Equal(cookie, input[32:64]) {
			t.Errorf("%q → %q (expected %q)", input, cookie, input[:32])
		}
	}
}
<commit_msg>Add tests for getServerBindaddrs.<commit_after>package pt

import (
	"bytes"
	"net"
	"os"
	"sort"
	"testing"
)

func stringIsSafe(s string) bool {
	for _, c := range []byte(s) {
		if c == '\x00' || c == '\n' || c > 127 {
			return false
		}
	}
	return true
}

func TestEscape(t *testing.T) {
	tests := [...]string{
		"",
		"abc",
		"a\nb",
		"a\\b",
		"ab\\",
		"ab\\\n",
		"ab\n\\",
	}

	check := func(input string) {
		output := escape(input)
		if !stringIsSafe(output) {
			t.Errorf("escape(%q) → %q", input, output)
		}
	}
	for _, input := range tests {
		check(input)
	}
	for b := 0; b < 256; b++ {
		// check one-byte string with each byte value 0–255
		check(string([]byte{byte(b)}))
		// check UTF-8 encoding of each character 0–255
		check(string(b))
	}
}

func TestGetManagedTransportVer(t *testing.T) {
	badTests := [...]string{
		"",
		"2",
	}
	goodTests := [...]struct {
		input, expected string
	}{
		{"1", "1"},
		{"1,1", "1"},
		{"1,2", "1"},
		{"2,1", "1"},
	}

	os.Clearenv()
	_, err := getManagedTransportVer()
	if err == nil {
		t.Errorf("empty environment unexpectedly succeeded")
	}

	for _, input := range badTests {
		os.Setenv("TOR_PT_MANAGED_TRANSPORT_VER", input)
		_, err := getManagedTransportVer()
		if err == nil {
			t.Errorf("TOR_PT_MANAGED_TRANSPORT_VER=%q unexpectedly succeeded", input)
		}
	}

	for _, test := range goodTests {
		os.Setenv("TOR_PT_MANAGED_TRANSPORT_VER", test.input)
		output, err := getManagedTransportVer()
		if err != nil {
			t.Errorf("TOR_PT_MANAGED_TRANSPORT_VER=%q unexpectedly returned an error: %s", test.input, err)
		}
		if output != test.expected {
			t.Errorf("TOR_PT_MANAGED_TRANSPORT_VER=%q → %q (expected %q)", test.input, output, test.expected)
		}
	}
}

// return true iff the two slices contain the same elements, possibly in a
// different order.
func stringSetsEqual(a, b []string) bool {
	ac := make([]string, len(a))
	bc := make([]string, len(b))
	copy(ac, a)
	copy(bc, b)
	sort.Strings(ac)
	sort.Strings(bc)
	if len(ac) != len(bc) {
		return false
	}
	for i := 0; i < len(ac); i++ {
		if ac[i] != bc[i] {
			return false
		}
	}
	return true
}

func tcpAddrsEqual(a, b *net.TCPAddr) bool {
	return a.IP.Equal(b.IP) && a.Port == b.Port
}

func TestGetClientTransports(t *testing.T) {
	tests := [...]struct {
		ptServerClientTransports string
		methodNames              []string
		expected                 []string
	}{
		{
			"*",
			[]string{},
			[]string{},
		},
		{
			"*",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "beta", "gamma"},
		},
		{
			"alpha,beta,gamma",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "beta", "gamma"},
		},
		{
			"alpha,beta",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "beta"},
		},
		{
			"alpha",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha"},
		},
		{
			"alpha,beta",
			[]string{"alpha", "beta", "alpha"},
			[]string{"alpha", "beta"},
		},
		// my reading of pt-spec.txt says that "*" has special meaning
		// only when it is the entirety of the environment variable.
		{
			"alpha,*,gamma",
			[]string{"alpha", "beta", "gamma"},
			[]string{"alpha", "gamma"},
		},
		{
			"alpha",
			[]string{"beta"},
			[]string{},
		},
	}

	os.Clearenv()
	_, err := getClientTransports([]string{"alpha", "beta", "gamma"})
	if err == nil {
		t.Errorf("empty environment unexpectedly succeeded")
	}

	for _, test := range tests {
		os.Setenv("TOR_PT_CLIENT_TRANSPORTS", test.ptServerClientTransports)
		output, err := getClientTransports(test.methodNames)
		if err != nil {
			t.Errorf("TOR_PT_CLIENT_TRANSPORTS=%q unexpectedly returned an error: %s",
				test.ptServerClientTransports, err)
		}
		if !stringSetsEqual(output, test.expected) {
			t.Errorf("TOR_PT_CLIENT_TRANSPORTS=%q %q → %q (expected %q)",
				test.ptServerClientTransports, test.methodNames, output, test.expected)
		}
	}
}

func TestResolveAddr(t *testing.T) {
	badTests := [...]string{
		"",
		"1.2.3.4",
		"1.2.3.4:",
		"9999",
		":9999",
		"[1:2::3:4]",
		"[1:2::3:4]:",
		"[1::2::3:4]",
		"1:2::3:4::9999",
		"1:2:3:4::9999",
		"localhost:9999",
		"[localhost]:9999",
	}
	goodTests := [...]struct {
		input    string
		expected net.TCPAddr
	}{
		{"1.2.3.4:9999", net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 9999}},
		{"[1:2::3:4]:9999", net.TCPAddr{IP: net.ParseIP("1:2::3:4"), Port: 9999}},
		{"1:2::3:4:9999", net.TCPAddr{IP: net.ParseIP("1:2::3:4"), Port: 9999}},
	}

	for _, input := range badTests {
		output, err := resolveAddr(input)
		if err == nil {
			t.Errorf("%q unexpectedly succeeded: %q", input, output)
		}
	}

	for _, test := range goodTests {
		output, err := resolveAddr(test.input)
		if err != nil {
			t.Errorf("%q unexpectedly returned an error: %s", test.input, err)
		}
		if !tcpAddrsEqual(output, &test.expected) {
			t.Errorf("%q → %q (expected %q)", test.input, output, test.expected)
		}
	}
}

func bindaddrSliceContains(s []Bindaddr, v Bindaddr) bool {
	for _, sv := range s {
		if sv.MethodName == v.MethodName && tcpAddrsEqual(sv.Addr, v.Addr) {
			return true
		}
	}
	return false
}

func bindaddrSetsEqual(a, b []Bindaddr) bool {
	for _, v := range a {
		if !bindaddrSliceContains(b, v) {
			return false
		}
	}
	for _, v := range b {
		if !bindaddrSliceContains(a, v) {
			return false
		}
	}
	return true
}

func TestGetServerBindaddrs(t *testing.T) {
	badTests := [...]struct {
		ptServerBindaddr   string
		ptServerTransports string
		methodNames        []string
	}{
		{
			"xxx",
			"xxx",
			[]string{},
		},
		{
			"alpha-1.2.3.4",
			"alpha",
			[]string{"alpha", "beta", "gamma"},
		},
	}
	goodTests := [...]struct {
		ptServerBindaddr   string
		ptServerTransports string
		methodNames        []string
		expected           []Bindaddr
	}{
		{
			"alpha-1.2.3.4:1111,beta-[1:2::3:4]:2222",
			"alpha,beta,gamma",
			[]string{"alpha", "beta"},
			[]Bindaddr{
				{"alpha", &net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 1111}},
				{"beta", &net.TCPAddr{IP: net.ParseIP("1:2::3:4"), Port: 2222}},
			},
		},
		{
			"alpha-1.2.3.4:1111",
			"xxx",
			[]string{"alpha", "beta", "gamma"},
			[]Bindaddr{},
		},
		{
			"alpha-1.2.3.4:1111",
			"alpha,beta,gamma",
			[]string{},
			[]Bindaddr{},
		},
		{
			"alpha-1.2.3.4:1111,beta-[1:2::3:4]:2222",
			"*",
			[]string{"alpha", "beta"},
			[]Bindaddr{
				{"alpha", &net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 1111}},
				{"beta", &net.TCPAddr{IP: net.ParseIP("1:2::3:4"), Port: 2222}},
			},
		},
	}

	os.Clearenv()
	_, err := getServerBindaddrs([]string{"alpha", "beta", "gamma"})
	if err == nil {
		t.Errorf("empty environment unexpectedly succeeded")
	}

	for _, test := range badTests {
		os.Setenv("TOR_PT_SERVER_BINDADDR", test.ptServerBindaddr)
		os.Setenv("TOR_PT_SERVER_TRANSPORTS", test.ptServerTransports)
		_, err := getServerBindaddrs(test.methodNames)
		if err == nil {
			t.Errorf("TOR_PT_SERVER_BINDADDR=%q TOR_PT_SERVER_TRANSPORTS=%q %q unexpectedly succeeded",
				test.ptServerBindaddr, test.ptServerTransports, test.methodNames)
		}
	}

	for _, test := range goodTests {
		os.Setenv("TOR_PT_SERVER_BINDADDR", test.ptServerBindaddr)
		os.Setenv("TOR_PT_SERVER_TRANSPORTS", test.ptServerTransports)
		output, err := getServerBindaddrs(test.methodNames)
		if err != nil {
			t.Errorf("TOR_PT_SERVER_BINDADDR=%q TOR_PT_SERVER_TRANSPORTS=%q %q unexpectedly returned an error: %s",
				test.ptServerBindaddr, test.ptServerTransports, test.methodNames, err)
		}
		if !bindaddrSetsEqual(output, test.expected) {
			t.Errorf("TOR_PT_SERVER_BINDADDR=%q TOR_PT_SERVER_TRANSPORTS=%q %q → %q (expected %q)",
				test.ptServerBindaddr, test.ptServerTransports, test.methodNames, output, test.expected)
		}
	}
}

func TestReadAuthCookie(t *testing.T) {
	badTests := [...][]byte{
		[]byte(""),
		// bad header
		[]byte("! Impostor ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDEF"),
		// too short
		[]byte("! Extended ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDE"),
		// too long
		[]byte("! Extended ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDEFX"),
	}
	goodTests := [...][]byte{
		[]byte("! Extended ORPort Auth Cookie !\x0a0123456789ABCDEF0123456789ABCDEF"),
	}

	for _, input := range badTests {
		var buf bytes.Buffer
		buf.Write(input)
		_, err := readAuthCookie(&buf)
		if err == nil {
			t.Errorf("%q unexpectedly succeeded", input)
		}
	}

	for _, input := range goodTests {
		var buf bytes.Buffer
		buf.Write(input)
		cookie, err := readAuthCookie(&buf)
		if err != nil {
			t.Errorf("%q unexpectedly returned an error: %s", input, err)
		}
		if !bytes.Equal(cookie, input[32:64]) {
			t.Errorf("%q → %q (expected %q)", input, cookie, input[:32])
		}
	}
}
<|endoftext|>


<commit_before>// Package profile provides a simple way to manage runtime/pprof
// profiling of your Go application.
package profile

import (
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
)

type profile struct {
	// Quiet suppresses informational messages during profiling.
	Quiet bool

	// CPUProfile controls if cpu profiling will be enabled.
	CPUProfile bool

	// MemProfile controls if memory profiling will be enabled.
	MemProfile bool

	// BlockProfile controls if block (contention) profiling will
	// be enabled.
	// It defaults to false.
	BlockProfile bool

	// ProfilePath controls the base path where various profiling
	// files are written. If blank, the base path will be generated
	// by ioutil.TempDir.
	ProfilePath string

	// NoShutdownHook controls whether the profiling package should
	// hook SIGINT to write profiles cleanly.
	NoShutdownHook bool

	// MemProfileRate sent the rate for the memory profile
	memProfileRate int

	closers []func()
}

// NoShutdownHook controls whether the profiling package should
// hook SIGINT to write profiles cleanly.
// Programs with more sophisticated signal handling should set
// this to true and ensure the Stop() function returned from Start()
// is called during shutdown.
func NoShutdownHook(p *profile) { p.NoShutdownHook = true }

// Quiet suppresses informational messages during profiling.
func Quiet(p *profile) { p.Quiet = true }

func (p *profile) NoProfiles() {
	p.CPUProfile = false
	p.MemProfile = false
	p.BlockProfile = false
}

// CPUProfile controls if cpu profiling will be enabled. It disables any previous profiling settings.
func CPUProfile(p *profile) {
	p.NoProfiles()
	p.CPUProfile = true
}

// MemProfile controls if memory profiling will be enabled. It disables any previous profiling settings.
func MemProfile(p *profile) {
	p.NoProfiles()
	p.MemProfile = true
}

// BlockProfile controls if block (contention) profiling will be enabled. It disables any previous profiling settings.
func BlockProfile(p *profile) {
	p.NoProfiles()
	p.BlockProfile = true
}

// resolvePath resolves the profile's path or outputs to a temporarry directory
func resolvePath(path string) (resolvedPath string, err error) {
	if path != "" {
		return path, os.MkdirAll(path, 0777)
	}

	return ioutil.TempDir("", "profile")
}

func (p *profile) Stop() {
	for _, c := range p.closers {
		c()
	}
}

func newProfile() *profile {
	prof := &profile{memProfileRate: 4096}
	CPUProfile(prof)
	return prof
}

// Start starts a new profiling session.
// The caller should call the Stop method on the value returned
// to cleanly stop profiling.
func Start(options ...func(*profile)) interface {
	Stop()
} {
	prof := newProfile()
	for _, option := range options {
		option(prof)
	}

	path, err := resolvePath(prof.ProfilePath)
	if err != nil {
		log.Fatalf("profile: could not create initial output directory: %v", err)
	}

	prof.ProfilePath = path
	if prof.Quiet {
		log.SetOutput(ioutil.Discard)
	}

	switch {
	case prof.CPUProfile:
		fn := filepath.Join(prof.ProfilePath, "cpu.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create cpu profile %q: %v", fn, err)
		}
		log.Printf("profile: cpu profiling enabled, %s", fn)
		pprof.StartCPUProfile(f)
		prof.closers = append(prof.closers, func() {
			pprof.StopCPUProfile()
			f.Close()
		})

	case prof.MemProfile:
		fn := filepath.Join(prof.ProfilePath, "mem.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create memory profile %q: %v", fn, err)
		}
		old := runtime.MemProfileRate
		runtime.MemProfileRate = prof.memProfileRate
		log.Printf("profile: memory profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("heap").WriteTo(f, 0)
			f.Close()
			runtime.MemProfileRate = old
		})

	case prof.BlockProfile:
		fn := filepath.Join(prof.ProfilePath, "block.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create block profile %q: %v", fn, err)
		}
		runtime.SetBlockProfileRate(1)
		log.Printf("profile: block profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("block").WriteTo(f, 0)
			f.Close()
			runtime.SetBlockProfileRate(0)
		})
	}

	if !prof.NoShutdownHook {
		go func() {
			c := make(chan os.Signal, 1)
			signal.Notify(c, os.Interrupt)
			<-c

			log.Println("profile: caught interrupt, stopping profiles")
			prof.Stop()

			os.Exit(0)
		}()
	}

	return prof
}
<commit_msg>Missed the resolved path<commit_after>// Package profile provides a simple way to manage runtime/pprof
// profiling of your Go application.
package profile

import (
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
)

type profile struct {
	// Quiet suppresses informational messages during profiling.
	Quiet bool

	// CPUProfile controls if cpu profiling will be enabled.
	CPUProfile bool

	// MemProfile controls if memory profiling will be enabled.
	MemProfile bool

	// BlockProfile controls if block (contention) profiling will
	// be enabled.
	// It defaults to false.
	BlockProfile bool

	// ProfilePath controls the base path where various profiling
	// files are written. If blank, the base path will be generated
	// by ioutil.TempDir.
	ProfilePath string

	// NoShutdownHook controls whether the profiling package should
	// hook SIGINT to write profiles cleanly.
	NoShutdownHook bool

	// MemProfileRate sent the rate for the memory profile
	memProfileRate int

	closers []func()
}

// NoShutdownHook controls whether the profiling package should
// hook SIGINT to write profiles cleanly.
// Programs with more sophisticated signal handling should set
// this to true and ensure the Stop() function returned from Start()
// is called during shutdown.
func NoShutdownHook(p *profile) { p.NoShutdownHook = true }

// Quiet suppresses informational messages during profiling.
func Quiet(p *profile) { p.Quiet = true }

func (p *profile) NoProfiles() {
	p.CPUProfile = false
	p.MemProfile = false
	p.BlockProfile = false
}

// CPUProfile controls if cpu profiling will be enabled. It disables any previous profiling settings.
func CPUProfile(p *profile) {
	p.NoProfiles()
	p.CPUProfile = true
}

// MemProfile controls if memory profiling will be enabled. It disables any previous profiling settings.
func MemProfile(p *profile) {
	p.NoProfiles()
	p.MemProfile = true
}

// BlockProfile controls if block (contention) profiling will be enabled. It disables any previous profiling settings.
func BlockProfile(p *profile) {
	p.NoProfiles()
	p.BlockProfile = true
}

// resolvePath resolves the profile's path or outputs to a temporarry directory
func resolvePath(path string) (resolvedPath string, err error) {
	if path != "" {
		return path, os.MkdirAll(path, 0777)
	}

	return ioutil.TempDir("", "profile")
}

func (p *profile) Stop() {
	for _, c := range p.closers {
		c()
	}
}

func newProfile() *profile {
	prof := &profile{memProfileRate: 4096}
	CPUProfile(prof)
	return prof
}

// Start starts a new profiling session.
// The caller should call the Stop method on the value returned
// to cleanly stop profiling.
func Start(options ...func(*profile)) interface {
	Stop()
} {
	prof := newProfile()
	for _, option := range options {
		option(prof)
	}

	path, err := resolvePath(prof.ProfilePath)
	if err != nil {
		log.Fatalf("profile: could not create initial output directory: %v", err)
	}

	prof.ProfilePath = path
	if prof.Quiet {
		log.SetOutput(ioutil.Discard)
	}

	switch {
	case prof.CPUProfile:
		fn := filepath.Join(path, "cpu.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create cpu profile %q: %v", fn, err)
		}
		log.Printf("profile: cpu profiling enabled, %s", fn)
		pprof.StartCPUProfile(f)
		prof.closers = append(prof.closers, func() {
			pprof.StopCPUProfile()
			f.Close()
		})

	case prof.MemProfile:
		fn := filepath.Join(path, "mem.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create memory profile %q: %v", fn, err)
		}
		old := runtime.MemProfileRate
		runtime.MemProfileRate = prof.memProfileRate
		log.Printf("profile: memory profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("heap").WriteTo(f, 0)
			f.Close()
			runtime.MemProfileRate = old
		})

	case prof.BlockProfile:
		fn := filepath.Join(path, "block.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create block profile %q: %v", fn, err)
		}
		runtime.SetBlockProfileRate(1)
		log.Printf("profile: block profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("block").WriteTo(f, 0)
			f.Close()
			runtime.SetBlockProfileRate(0)
		})
	}

	if !prof.NoShutdownHook {
		go func() {
			c := make(chan os.Signal, 1)
			signal.Notify(c, os.Interrupt)
			<-c

			log.Println("profile: caught interrupt, stopping profiles")
			prof.Stop()

			os.Exit(0)
		}()
	}

	return prof
}
<|endoftext|>