<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"sync"

	"github.com/yi-jiayu/telegram-bot-api"
	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
)

var callbackQueryHandlers = map[string]CallbackQueryHandler{
	"refresh":                       RefreshCallbackHandler,
	"eta":                           EtaCallbackHandler,
	"eta_demo":                      EtaDemoCallbackHandler,
	"new_eta":                       NewEtaHandler,
	"no_show_redundant_eta_command": NoShowRedundantEtaCommandCallbackHandler,
}

// CallbackQueryHandler is a handler for callback queries
type CallbackQueryHandler func(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error

func answerCallbackQuery(bot *BusEtaBot, c tgbotapi.Chattable, answer tgbotapi.CallbackConfig) error {
	errs := make([]error, 0)

	var wg sync.WaitGroup
	wg.Add(2)

	errChan := make(chan error, 2)
	go func() {
		for err := range errChan {
			errs = append(errs, err)
		}
	}()

	go func() {
		defer wg.Done()

		_, err := bot.Telegram.Send(c)
		if err != nil {
			errChan <- err
		}
	}()

	go func() {
		defer wg.Done()

		_, err := bot.Telegram.AnswerCallbackQuery(answer)
		if err != nil {
			errChan <- err
		}
	}()

	wg.Wait()
	close(errChan)

	switch len(errs) {
	case 1:
		return errs[0]
	case 2:
		return errors.New(errs[0].Error() + "\n" + errs[1].Error())
	default:
		return nil
	}
}

func updateEtaMessage(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery, busStopID string, serviceNos []string) error {
	text, err := EtaMessage(ctx, bot, busStopID, serviceNos)
	if err != nil {
		return err
	}

	var reply tgbotapi.EditMessageTextConfig
	if cbq.Message != nil {
		chatID := cbq.Message.Chat.ID
		messageID := cbq.Message.MessageID
		reply = tgbotapi.NewEditMessageText(chatID, messageID, text)
	} else {
		inlineMessageID := cbq.InlineMessageID
		reply = tgbotapi.EditMessageTextConfig{
			BaseEdit: tgbotapi.BaseEdit{
				InlineMessageID: inlineMessageID,
			},
			Text: text,
		}
	}

	callbackData := CallbackData{
		Type:       "refresh",
		BusStopID:  busStopID,
		ServiceNos: serviceNos,
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return err
	}
	callbackDataJSONStr := string(callbackDataJSON)

	reply.ParseMode = "markdown"
	reply.ReplyMarkup = &tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{
			{
				tgbotapi.InlineKeyboardButton{
					Text:         "Refresh",
					CallbackData: &callbackDataJSONStr,
				},
			},
		},
	}

	answer := tgbotapi.NewCallback(cbq.ID, "Etas updated!")

	if cbq.Message != nil {
		go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionRefreshCallback, cbq.Message.Chat.Type)
	} else {
		go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionRefreshCallback, LabelInlineMessage)
	}

	err = answerCallbackQuery(bot, reply, answer)
	return err
}

// RefreshCallbackHandler handles the callback for the Refresh button on an eta message.
func RefreshCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	var data CallbackData
	err := json.Unmarshal([]byte(cbq.Data), &data)
	if err != nil {
		return err
	}

	bsID, sNos := data.BusStopID, data.ServiceNos

	return updateEtaMessage(ctx, bot, cbq, bsID, sNos)
}

// EtaCallbackHandler handles callback queries from eta messages from old versions of the bot for
// backwards-compatibility.
func EtaCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	var data CallbackData
	err := json.Unmarshal([]byte(cbq.Data), &data)
	if err != nil {
		return err
	}

	var bsID string
	var sNos []string
	if data.Argstr != "" {
		bsID, sNos, _ = InferEtaQuery(data.Argstr)
	} else {
		bsID = data.BusStopID
		sNos = data.ServiceNos
	}

	callbackData := CallbackData{
		Type:       "refresh",
		BusStopID:  bsID,
		ServiceNos: sNos,
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return err
	}
	callbackDataJSONStr := string(callbackDataJSON)
	cbq.Data = callbackDataJSONStr

	return updateEtaMessage(ctx, bot, cbq, bsID, sNos)
}

// EtaDemoCallbackHandler handles an eta_demo callback from a start command.
func EtaDemoCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	chatID := cbq.Message.Chat.ID

	text, err := EtaMessage(ctx, bot, "96049", nil)
	if err != nil {
		return err
	}

	callbackData := CallbackData{
		Type:      "refresh",
		BusStopID: "96049",
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return err
	}
	callbackDataJSONStr := string(callbackDataJSON)

	reply := tgbotapi.NewMessage(chatID, text)
	reply.ParseMode = "markdown"
	reply.ReplyMarkup = tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{
			{
				tgbotapi.InlineKeyboardButton{
					Text:         "Refresh",
					CallbackData: &callbackDataJSONStr,
				},
			},
		},
	}

	if !cbq.Message.Chat.IsPrivate() {
		reply.ReplyToMessageID = cbq.Message.MessageID
	}

	answer := tgbotapi.NewCallback(cbq.ID, "")

	go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionEtaDemoCallback, cbq.Message.Chat.Type)

	err = answerCallbackQuery(bot, reply, answer)
	return err
}

// NewEtaHandler sends etas for a bus stop when a user taps "Get etas" on a bus stop location returned from a
// location query
func NewEtaHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	chatID := cbq.Message.Chat.ID

	var data CallbackData
	err := json.Unmarshal([]byte(cbq.Data), &data)
	if err != nil {
		return err
	}

	bsID, sNos := data.BusStopID, data.ServiceNos

	text, err := EtaMessage(ctx, bot, bsID, sNos)
	if err != nil {
		return err
	}

	callbackData := CallbackData{
		Type:       "refresh",
		BusStopID:  bsID,
		ServiceNos: sNos,
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return err
	}
	callbackDataJSONStr := string(callbackDataJSON)

	reply := tgbotapi.NewMessage(chatID, text)
	reply.ParseMode = "markdown"
	reply.ReplyMarkup = tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{
			{
				tgbotapi.InlineKeyboardButton{
					Text:         "Refresh",
					CallbackData: &callbackDataJSONStr,
				},
			},
		},
	}

	if !cbq.Message.Chat.IsPrivate() {
		reply.ReplyToMessageID = cbq.Message.MessageID
	}

	answer := tgbotapi.NewCallback(cbq.ID, "")

	go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionEtaFromLocationCallback, cbq.Message.Chat.Type)

	err = answerCallbackQuery(bot, reply, answer)
	return err
}

func NoShowRedundantEtaCommandCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	chatID := cbq.Message.Chat.ID
	messageID := cbq.Message.MessageID
	userID := cbq.From.ID

	prefs, err := GetUserPreferences(ctx, userID)
	if err != nil {
		return err
	}

	prefs.NoRedundantEtaCommandReminder = true

	err = SetUserPreferences(ctx, userID, &prefs)
	if err != nil {
		return err
	}

	markup := tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{},
	}
	edit := tgbotapi.NewEditMessageReplyMarkup(chatID, messageID, markup)

	answer := tgbotapi.NewCallback(cbq.ID, "Got it!")

	err = answerCallbackQuery(bot, edit, answer)
	return err
}

// callbackErrorHandler is for informing the user about an error while processing a callback query.
func callbackErrorHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery, err error) {
	log.Errorf(ctx, "%v", err)

	text := fmt.Sprintf("Oh no! Something went wrong. \n\nRequest ID: `%s`", appengine.RequestID(ctx))
	answer := tgbotapi.NewCallbackWithAlert(cbq.ID, text)

	_, err = bot.Telegram.AnswerCallbackQuery(answer)
	if err != nil {
		log.Errorf(ctx, "%v", err)
	}
}
<commit_msg>Add context to errors<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/pkg/errors"
	"github.com/yi-jiayu/telegram-bot-api"
	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
)

var callbackQueryHandlers = map[string]CallbackQueryHandler{
	"refresh":                       RefreshCallbackHandler,
	"eta":                           EtaCallbackHandler,
	"eta_demo":                      EtaDemoCallbackHandler,
	"new_eta":                       NewEtaHandler,
	"no_show_redundant_eta_command": NoShowRedundantEtaCommandCallbackHandler,
}

// CallbackQueryHandler is a handler for callback queries
type CallbackQueryHandler func(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error

func answerCallbackQuery(bot *BusEtaBot, c tgbotapi.Chattable, answer tgbotapi.CallbackConfig) error {
	errs := make([]error, 0)

	var wg sync.WaitGroup
	wg.Add(2)

	errChan := make(chan error, 2)
	go func() {
		for err := range errChan {
			errs = append(errs, err)
		}
	}()

	go func() {
		defer wg.Done()

		_, err := bot.Telegram.Send(c)
		if err != nil {
			err = errors.Wrap(err, fmt.Sprintf("%#v", c))
			errChan <- err
		}
	}()

	go func() {
		defer wg.Done()

		_, err := bot.Telegram.AnswerCallbackQuery(answer)
		if err != nil {
			err = errors.Wrap(err, fmt.Sprintf("%#v", answer))
			errChan <- err
		}
	}()

	wg.Wait()
	close(errChan)

	switch len(errs) {
	case 1:
		return errs[0]
	case 2:
		return fmt.Errorf("%v\n%v", errs[0], errs[1])
	default:
		return nil
	}
}

func updateEtaMessage(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery, busStopID string, serviceNos []string) error {
	text, err := EtaMessage(ctx, bot, busStopID, serviceNos)
	if err != nil {
		return err
	}

	var reply tgbotapi.EditMessageTextConfig
	if cbq.Message != nil {
		chatID := cbq.Message.Chat.ID
		messageID := cbq.Message.MessageID
		reply = tgbotapi.NewEditMessageText(chatID, messageID, text)
	} else {
		inlineMessageID := cbq.InlineMessageID
		reply = tgbotapi.EditMessageTextConfig{
			BaseEdit: tgbotapi.BaseEdit{
				InlineMessageID: inlineMessageID,
			},
			Text: text,
		}
	}

	callbackData := CallbackData{
		Type:       "refresh",
		BusStopID:  busStopID,
		ServiceNos: serviceNos,
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return err
	}
	callbackDataJSONStr := string(callbackDataJSON)

	reply.ParseMode = "markdown"
	reply.ReplyMarkup = &tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{
			{
				tgbotapi.InlineKeyboardButton{
					Text:         "Refresh",
					CallbackData: &callbackDataJSONStr,
				},
			},
		},
	}

	answer := tgbotapi.NewCallback(cbq.ID, "Etas updated!")

	if cbq.Message != nil {
		go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionRefreshCallback, cbq.Message.Chat.Type)
	} else {
		go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionRefreshCallback, LabelInlineMessage)
	}

	err = answerCallbackQuery(bot, reply, answer)
	return err
}

// RefreshCallbackHandler handles the callback for the Refresh button on an eta message.
func RefreshCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	var data CallbackData
	err := json.Unmarshal([]byte(cbq.Data), &data)
	if err != nil {
		return err
	}

	bsID, sNos := data.BusStopID, data.ServiceNos

	return updateEtaMessage(ctx, bot, cbq, bsID, sNos)
}

// EtaCallbackHandler handles callback queries from eta messages from old versions of the bot for
// backwards-compatibility.
func EtaCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	var data CallbackData
	err := json.Unmarshal([]byte(cbq.Data), &data)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("error unmarshalling callback data: %s", cbq.Data))
	}

	var bsID string
	var sNos []string
	if data.Argstr != "" {
		bsID, sNos, _ = InferEtaQuery(data.Argstr)
	} else {
		bsID = data.BusStopID
		sNos = data.ServiceNos
	}

	callbackData := CallbackData{
		Type:       "refresh",
		BusStopID:  bsID,
		ServiceNos: sNos,
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return err
	}
	callbackDataJSONStr := string(callbackDataJSON)
	cbq.Data = callbackDataJSONStr

	return updateEtaMessage(ctx, bot, cbq, bsID, sNos)
}

// EtaDemoCallbackHandler handles an eta_demo callback from a start command.
func EtaDemoCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	chatID := cbq.Message.Chat.ID

	text, err := EtaMessage(ctx, bot, "96049", nil)
	if err != nil {
		return err
	}

	callbackData := CallbackData{
		Type:      "refresh",
		BusStopID: "96049",
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("error marshalling callback data: %#v", callbackData))
	}
	callbackDataJSONStr := string(callbackDataJSON)

	reply := tgbotapi.NewMessage(chatID, text)
	reply.ParseMode = "markdown"
	reply.ReplyMarkup = tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{
			{
				tgbotapi.InlineKeyboardButton{
					Text:         "Refresh",
					CallbackData: &callbackDataJSONStr,
				},
			},
		},
	}

	if !cbq.Message.Chat.IsPrivate() {
		reply.ReplyToMessageID = cbq.Message.MessageID
	}

	answer := tgbotapi.NewCallback(cbq.ID, "")

	go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionEtaDemoCallback, cbq.Message.Chat.Type)

	err = answerCallbackQuery(bot, reply, answer)
	return err
}

// NewEtaHandler sends etas for a bus stop when a user taps "Get etas" on a bus stop location returned from a
// location query
func NewEtaHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	chatID := cbq.Message.Chat.ID

	var data CallbackData
	err := json.Unmarshal([]byte(cbq.Data), &data)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("error unmarshalling callback data: %s", cbq.Data))
	}

	bsID, sNos := data.BusStopID, data.ServiceNos

	text, err := EtaMessage(ctx, bot, bsID, sNos)
	if err != nil {
		return err
	}

	callbackData := CallbackData{
		Type:       "refresh",
		BusStopID:  bsID,
		ServiceNos: sNos,
	}

	callbackDataJSON, err := json.Marshal(callbackData)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("error marshalling callback data: %#v", callbackData))
	}
	callbackDataJSONStr := string(callbackDataJSON)

	reply := tgbotapi.NewMessage(chatID, text)
	reply.ParseMode = "markdown"
	reply.ReplyMarkup = tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{
			{
				tgbotapi.InlineKeyboardButton{
					Text:         "Refresh",
					CallbackData: &callbackDataJSONStr,
				},
			},
		},
	}

	if !cbq.Message.Chat.IsPrivate() {
		reply.ReplyToMessageID = cbq.Message.MessageID
	}

	answer := tgbotapi.NewCallback(cbq.ID, "")

	go bot.LogEvent(ctx, cbq.From, CategoryCallback, ActionEtaFromLocationCallback, cbq.Message.Chat.Type)

	err = answerCallbackQuery(bot, reply, answer)
	return err
}

// NoShowRedundantEtaCommandCallbackHandler handles the "Don't show again" callback button on the reminder sent to
// users that the /eta command isn't necessary in a private chat.
func NoShowRedundantEtaCommandCallbackHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery) error {
	chatID := cbq.Message.Chat.ID
	messageID := cbq.Message.MessageID
	userID := cbq.From.ID

	prefs, err := GetUserPreferences(ctx, userID)
	if err != nil {
		return err
	}

	prefs.NoRedundantEtaCommandReminder = true

	err = SetUserPreferences(ctx, userID, &prefs)
	if err != nil {
		return err
	}

	markup := tgbotapi.InlineKeyboardMarkup{
		InlineKeyboard: [][]tgbotapi.InlineKeyboardButton{},
	}
	edit := tgbotapi.NewEditMessageReplyMarkup(chatID, messageID, markup)

	answer := tgbotapi.NewCallback(cbq.ID, "Got it!")

	err = answerCallbackQuery(bot, edit, answer)
	return err
}

// callbackErrorHandler is for informing the user about an error while processing a callback query.
func callbackErrorHandler(ctx context.Context, bot *BusEtaBot, cbq *tgbotapi.CallbackQuery, err error) {
	log.Errorf(ctx, "%v", err)

	text := fmt.Sprintf("Oh no! Something went wrong. \n\nRequest ID: `%s`", appengine.RequestID(ctx))
	answer := tgbotapi.NewCallbackWithAlert(cbq.ID, text)

	_, err = bot.Telegram.AnswerCallbackQuery(answer)
	if err != nil {
		err := errors.Wrap(err, fmt.Sprintf("%#v", answer))
		log.Errorf(ctx, "%v", err)
	}
}
<|endoftext|>


<commit_before>package helpers

import (
	"sync"

	"bytes"

	"io/ioutil"

	"os"

	"path/filepath"

	"errors"

	"fmt"

	"strings"

	"time"

	"strconv"

	"encoding/binary"

	"mime"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/globalsign/mgo/bson"
	"github.com/kennygrant/sanitize"
	"github.com/minio/minio-go"
	uuid "github.com/satori/go.uuid"
)

var (
	minioBucket string
	minioClient *minio.Client
	minioLock   sync.Mutex
)

// TODO: watch cache folder size

type AddFileMetadata struct {
	Filename           string            // the actual file name, can be empty
	ChannelID          string            // the source channel ID, can be empty, but should be set if possible
	UserID             string            // the source user ID, can be empty, but should be set if possible
	AdditionalMetadata map[string]string // additional metadata attached to the object
}

// TODO: prevent duplicates
// Stores a file
// name		: the name of the new object, can be empty to generate an unique name
// data		: the file data
// metadata	: metadata attached to the object
// source	: the source name for the file, for example the module name, can not be empty
// public	: if true file will be available via the website proxy
func AddFile(name string, data []byte, metadata AddFileMetadata, source string, public bool) (objectName string, err error) {
	// check if source is set
	if source == "" {
		return "", errors.New("source can not be empty")
	}
	// check if user uploads are disabled for given userID, if userID is given
	if metadata.UserID != "" && UseruploadsIsDisabled(metadata.UserID) {
		return "", errors.New("uploads are disabled for this user")
	}
	// set new object name
	objectName = name
	if objectName == "" {
		// generate unique filename
		objectName = uuid.NewV4().String()
	}
	// retrieve guildID if channelID is set
	var guildID string
	if metadata.ChannelID != "" {
		channel, err := GetChannel(metadata.ChannelID)
		RelaxLog(err)
		if err == nil {
			guildID = channel.GuildID
		}
	}
	// get filetype
	filetype, _ := SniffMime(data)
	// get filesize
	filesize := binary.Size(data)
	// update metadata
	if metadata.AdditionalMetadata == nil {
		metadata.AdditionalMetadata = make(map[string]string, 0)
	}
	metadata.AdditionalMetadata["filename"] = metadata.Filename
	metadata.AdditionalMetadata["userid"] = metadata.UserID
	metadata.AdditionalMetadata["guildid"] = guildID
	metadata.AdditionalMetadata["channelid"] = metadata.ChannelID
	metadata.AdditionalMetadata["source"] = source
	metadata.AdditionalMetadata["mimetype"] = filetype
	metadata.AdditionalMetadata["filesize"] = strconv.Itoa(filesize)
	metadata.AdditionalMetadata["public"] = "no"
	if public {
		metadata.AdditionalMetadata["public"] = "yes"
	}
	// upload file
	err = uploadFile(objectName, data, metadata.AdditionalMetadata)
	if err != nil {
		return "", err
	}
	// store in database
	err = MDbUpsert(
		models.StorageTable,
		bson.M{"objectname": objectName},
		models.StorageEntry{
			ObjectName:     objectName,
			ObjectNameHash: GetMD5Hash(objectName),
			UploadDate:     time.Now(),
			Filename:       metadata.Filename,
			UserID:         metadata.UserID,
			GuildID:        guildID,
			ChannelID:      metadata.ChannelID,
			Source:         source,
			MimeType:       filetype,
			Filesize:       filesize,
			Public:         public,
			Metadata:       metadata.AdditionalMetadata,
		},
	)
	if err != nil {
		return "", err
	}
	// return new objectName
	return objectName, nil
}

// retrieves a file
// objectName	: the name of the file to retrieve
func RetrieveFile(objectName string) (data []byte, err error) {
	// setup minioClient if not yet done
	if minioClient == nil {
		err = setupMinioClient()
		if err != nil {
			return data, err
		}
	}

	// Increase MongoDB RetrievedCount
	go func() {
		defer Recover()
		err := MDbUpdateQuery(models.StorageTable, bson.M{"objectname": objectName}, bson.M{"$inc": bson.M{"retrievedcount": 1}})
		if err != nil && !strings.Contains(err.Error(), "not found") {
			RelaxLog(err)
		}
	}()

	data = getBucketCache(objectName)
	if data != nil {
		cache.GetLogger().WithField("module", "storage").Infof("retrieving " + objectName + " from minio cache")
		return data, nil
	}

	cache.GetLogger().WithField("module", "storage").Infof("retrieving " + objectName + " from minio storage")

	// retrieve the object
	minioObject, err := minioClient.GetObject(minioBucket, sanitize.BaseName(objectName), minio.GetObjectOptions{})
	if err != nil {
		return data, err
	}

	// read the object into a byte slice
	data, err = ioutil.ReadAll(minioObject)
	if err != nil {
		return data, err
	}

	go func() {
		defer Recover()
		cache.GetLogger().WithField("module", "storage").Infof("caching " + objectName + " into minio cache")
		err := setBucketCache(objectName, data)
		RelaxLog(err)
	}()

	return data, nil
}

// Retrieves a file by the object name md5 hash
// currently supported file sources: custom commands
// hash	: the md5 hash
func RetrieveFileByHash(hash string) (filename, filetype string, data []byte, err error) {
	var entryBucket models.StorageEntry
	err = MdbOne(
		MdbCollection(models.StorageTable).Find(bson.M{"objectnamehash": hash}),
		&entryBucket,
	)
	if err != nil {
		// try fallback to deprecated custom commands object storage
		var entryBucket models.CustomCommandsEntry
		err = MdbOne(
			MdbCollection(models.CustomCommandsTable).Find(bson.M{"storagehash": hash}),
			&entryBucket,
		)
		if err != nil {
			return "", "", nil, errors.New("file not found")
		}
		data, err = RetrieveFile(entryBucket.StorageObjectName)
		RelaxLog(err)
		if err != nil {
			return "", "", nil, err
		}
		return entryBucket.StorageFilename, entryBucket.StorageMimeType, data, nil
	}

	data, err = RetrieveFile(entryBucket.ObjectName)
	RelaxLog(err)
	if err != nil {
		return "", "", nil, err
	}
	return entryBucket.Filename, entryBucket.MimeType, data, nil
}

// Retrieves a file's public url, returns an error if file is not public
// objectName	:  the name of the object
func GetFileLink(objectName string) (url string, err error) {
	var entryBucket models.StorageEntry
	err = MdbOne(
		MdbCollection(models.StorageTable).Find(bson.M{"objectname": objectName}),
		&entryBucket,
	)
	if err != nil {
		return "", err
	}

	if !entryBucket.Public {
		return "", errors.New("this file is not available publicly")
	}

	filename := entryBucket.Filename
	if filename == "" {
		filename = "robyul"
		extensions, err := mime.ExtensionsByType(entryBucket.MimeType)
		if err == nil && extensions != nil && len(extensions) >= 0 {
			filename += extensions[0]
		}
	}

	url = GeneratePublicFileLink(filename, entryBucket.ObjectNameHash)

	return url, nil
}

// Deletes a file
// objectName	: the name of the object
func DeleteFile(objectName string) (err error) {
	// setup minioClient if not yet done
	if minioClient == nil {
		err = setupMinioClient()
		if err != nil {
			return err
		}
	}

	cache.GetLogger().WithField("module", "storage").Infof("deleting " + objectName + " from minio storage")

	go func() {
		defer Recover()
		cache.GetLogger().WithField("module", "storage").Infof("deleting " + objectName + " from minio cache")
		err = deleteBucketCache(objectName)
		RelaxLog(err)
	}()

	// delete the object
	err = minioClient.RemoveObject(minioBucket, sanitize.BaseName(objectName))

	// delete mongo db entry
	go func() {
		defer Recover()
		err := MdbDeleteQuery(models.StorageTable, bson.M{"objectname": objectName})
		if err != nil && !strings.Contains(err.Error(), "not found") {
			RelaxLog(err)
		}
	}()

	return err
}

// Gets a public link for a file
// filename		: the name of the file
// filetype		: the type of the file
func GeneratePublicFileLink(filename, filehash string) (link string) {
	dots := strings.Count(filename, ".")
	if dots > 1 {
		filename = strings.Replace(filename, ".", "-", dots-1)
	}
	return fmt.Sprintf(GetConfig().Path("imageproxy.base_url").Data().(string),
		filehash, filename)
}

// TODO: prevent overwrites
// uploads a file to the minio object storage
// objectName	: the name of the file to upload
// data			: the data for the new object
// metadata		: additional metadata attached to the object
func uploadFile(objectName string, data []byte, metadata map[string]string) (err error) {
	// setup minioClient if not yet done
	if minioClient == nil {
		err = setupMinioClient()
		if err != nil {
			return err
		}
	}

	cache.GetLogger().WithField("module", "storage").Infof("uploading " + objectName + " to minio storage")

	options := minio.PutObjectOptions{}

	// add content type
	filetype, err := SniffMime(data)
	if err == nil {
		options.ContentType = filetype
	}

	// add metadata
	if metadata != nil && len(metadata) > 0 {
		options.UserMetadata = metadata
	}

	// upload the data
	_, err = minioClient.PutObject(minioBucket, sanitize.BaseName(objectName), bytes.NewReader(data), -1, options)
	return err
}

func getBucketCache(objectName string) (data []byte) {
	var err error

	if _, err = os.Stat(getObjectPath(objectName)); os.IsNotExist(err) {
		return nil
	}

	data, err = ioutil.ReadFile(getObjectPath(objectName))
	if err != nil {
		return nil
	}

	return data
}

func setBucketCache(objectName string, data []byte) (err error) {
	if _, err = os.Stat(filepath.Dir(getObjectPath(objectName))); os.IsNotExist(err) {
		err = os.MkdirAll(filepath.Dir(getObjectPath(objectName)), os.ModePerm)
		if err != nil {
			return err
		}
	}

	err = ioutil.WriteFile(getObjectPath(objectName), data, 0644)
	return err
}

func deleteBucketCache(objectName string) (err error) {
	if _, err = os.Stat(getObjectPath(objectName)); os.IsNotExist(err) {
		return nil
	}

	err = os.Remove(getObjectPath(objectName))
	return err
}

func getObjectPath(objectName string) (path string) {
	return GetConfig().Path("cache_folder").Data().(string) + "/minio-" + GetConfig().Path("s3.bucket").Data().(string) + "/" + sanitize.BaseName(objectName)
}

// Initialize the minio client object, and creates the bucket if it doesn't exist yet
func setupMinioClient() (err error) {
	minioLock.Lock()
	minioBucket = GetConfig().Path("s3.bucket").Data().(string)
	minioClient, err = minio.New(
		GetConfig().Path("s3.endpoint").Data().(string),
		GetConfig().Path("s3.access_key").Data().(string),
		GetConfig().Path("s3.secret_secret_key").Data().(string),
		true,
	)
	minioLock.Unlock()

	bucketExists, err := minioClient.BucketExists(minioBucket)
	if err != nil {
		return err
	}

	if !bucketExists {
		err = minioClient.MakeBucket(minioBucket, "ams3")
		if err != nil {
			return err
		}
	}

	return err
}
<commit_msg>[storage] adds cache warm up for public files 🍵<commit_after>package helpers

import (
	"sync"

	"bytes"

	"io/ioutil"

	"os"

	"path/filepath"

	"errors"

	"fmt"

	"strings"

	"time"

	"strconv"

	"encoding/binary"

	"mime"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/globalsign/mgo/bson"
	"github.com/kennygrant/sanitize"
	"github.com/minio/minio-go"
	uuid "github.com/satori/go.uuid"
)

var (
	minioBucket string
	minioClient *minio.Client
	minioLock   sync.Mutex
)

// TODO: watch cache folder size

type AddFileMetadata struct {
	Filename           string            // the actual file name, can be empty
	ChannelID          string            // the source channel ID, can be empty, but should be set if possible
	UserID             string            // the source user ID, can be empty, but should be set if possible
	AdditionalMetadata map[string]string // additional metadata attached to the object
}

// TODO: prevent duplicates
// Stores a file
// name		: the name of the new object, can be empty to generate an unique name
// data		: the file data
// metadata	: metadata attached to the object
// source	: the source name for the file, for example the module name, can not be empty
// public	: if true file will be available via the website proxy
func AddFile(name string, data []byte, metadata AddFileMetadata, source string, public bool) (objectName string, err error) {
	// check if source is set
	if source == "" {
		return "", errors.New("source can not be empty")
	}
	// check if user uploads are disabled for given userID, if userID is given
	if metadata.UserID != "" && UseruploadsIsDisabled(metadata.UserID) {
		return "", errors.New("uploads are disabled for this user")
	}
	// set new object name
	objectName = name
	if objectName == "" {
		// generate unique filename
		objectName = uuid.NewV4().String()
	}
	// retrieve guildID if channelID is set
	var guildID string
	if metadata.ChannelID != "" {
		channel, err := GetChannel(metadata.ChannelID)
		RelaxLog(err)
		if err == nil {
			guildID = channel.GuildID
		}
	}
	// get filetype
	filetype, _ := SniffMime(data)
	// get filesize
	filesize := binary.Size(data)
	// update metadata
	if metadata.AdditionalMetadata == nil {
		metadata.AdditionalMetadata = make(map[string]string, 0)
	}
	metadata.AdditionalMetadata["filename"] = metadata.Filename
	metadata.AdditionalMetadata["userid"] = metadata.UserID
	metadata.AdditionalMetadata["guildid"] = guildID
	metadata.AdditionalMetadata["channelid"] = metadata.ChannelID
	metadata.AdditionalMetadata["source"] = source
	metadata.AdditionalMetadata["mimetype"] = filetype
	metadata.AdditionalMetadata["filesize"] = strconv.Itoa(filesize)
	metadata.AdditionalMetadata["public"] = "no"
	if public {
		metadata.AdditionalMetadata["public"] = "yes"
	}
	// upload file
	err = uploadFile(objectName, data, metadata.AdditionalMetadata)
	if err != nil {
		return "", err
	}
	// store in database
	err = MDbUpsert(
		models.StorageTable,
		bson.M{"objectname": objectName},
		models.StorageEntry{
			ObjectName:     objectName,
			ObjectNameHash: GetMD5Hash(objectName),
			UploadDate:     time.Now(),
			Filename:       metadata.Filename,
			UserID:         metadata.UserID,
			GuildID:        guildID,
			ChannelID:      metadata.ChannelID,
			Source:         source,
			MimeType:       filetype,
			Filesize:       filesize,
			Public:         public,
			Metadata:       metadata.AdditionalMetadata,
		},
	)
	if err != nil {
		return "", err
	}
	// warm up cache for public files
	if public {
		go func() {
			defer Recover()
			link, err := GetFileLink(objectName)
			Relax(err)
			_, err = NetGetUAWithError(link, DEFAULT_UA)
			Relax(err)
		}()
	}
	// return new objectName
	return objectName, nil
}

// retrieves a file
// objectName	: the name of the file to retrieve
func RetrieveFile(objectName string) (data []byte, err error) {
	// setup minioClient if not yet done
	if minioClient == nil {
		err = setupMinioClient()
		if err != nil {
			return data, err
		}
	}

	// Increase MongoDB RetrievedCount
	go func() {
		defer Recover()
		err := MDbUpdateQuery(models.StorageTable, bson.M{"objectname": objectName}, bson.M{"$inc": bson.M{"retrievedcount": 1}})
		if err != nil && !strings.Contains(err.Error(), "not found") {
			RelaxLog(err)
		}
	}()

	data = getBucketCache(objectName)
	if data != nil {
		cache.GetLogger().WithField("module", "storage").Infof("retrieving " + objectName + " from minio cache")
		return data, nil
	}

	cache.GetLogger().WithField("module", "storage").Infof("retrieving " + objectName + " from minio storage")

	// retrieve the object
	minioObject, err := minioClient.GetObject(minioBucket, sanitize.BaseName(objectName), minio.GetObjectOptions{})
	if err != nil {
		return data, err
	}

	// read the object into a byte slice
	data, err = ioutil.ReadAll(minioObject)
	if err != nil {
		return data, err
	}

	go func() {
		defer Recover()
		cache.GetLogger().WithField("module", "storage").Infof("caching " + objectName + " into minio cache")
		err := setBucketCache(objectName, data)
		RelaxLog(err)
	}()

	return data, nil
}

// Retrieves a file by the object name md5 hash
// currently supported file sources: custom commands
// hash	: the md5 hash
func RetrieveFileByHash(hash string) (filename, filetype string, data []byte, err error) {
	var entryBucket models.StorageEntry
	err = MdbOne(
		MdbCollection(models.StorageTable).Find(bson.M{"objectnamehash": hash}),
		&entryBucket,
	)
	if err != nil {
		// try fallback to deprecated custom commands object storage
		var entryBucket models.CustomCommandsEntry
		err = MdbOne(
			MdbCollection(models.CustomCommandsTable).Find(bson.M{"storagehash": hash}),
			&entryBucket,
		)
		if err != nil {
			return "", "", nil, errors.New("file not found")
		}
		data, err = RetrieveFile(entryBucket.StorageObjectName)
		RelaxLog(err)
		if err != nil {
			return "", "", nil, err
		}
		return entryBucket.StorageFilename, entryBucket.StorageMimeType, data, nil
	}

	data, err = RetrieveFile(entryBucket.ObjectName)
	RelaxLog(err)
	if err != nil {
		return "", "", nil, err
	}
	return entryBucket.Filename, entryBucket.MimeType, data, nil
}

// Retrieves a file's public url, returns an error if file is not public
// objectName	:  the name of the object
func GetFileLink(objectName string) (url string, err error) {
	var entryBucket models.StorageEntry
	err = MdbOne(
		MdbCollection(models.StorageTable).Find(bson.M{"objectname": objectName}),
		&entryBucket,
	)
	if err != nil {
		return "", err
	}

	if !entryBucket.Public {
		return "", errors.New("this file is not available publicly")
	}

	filename := entryBucket.Filename
	if filename == "" {
		filename = "robyul"
		extensions, err := mime.ExtensionsByType(entryBucket.MimeType)
		if err == nil && extensions != nil && len(extensions) >= 0 {
			filename += extensions[0]
		}
	}

	url = GeneratePublicFileLink(filename, entryBucket.ObjectNameHash)

	return url, nil
}

// Deletes a file
// objectName	: the name of the object
func DeleteFile(objectName string) (err error) {
	// setup minioClient if not yet done
	if minioClient == nil {
		err = setupMinioClient()
		if err != nil {
			return err
		}
	}

	cache.GetLogger().WithField("module", "storage").Infof("deleting " + objectName + " from minio storage")

	go func() {
		defer Recover()
		cache.GetLogger().WithField("module", "storage").Infof("deleting " + objectName + " from minio cache")
		err = deleteBucketCache(objectName)
		RelaxLog(err)
	}()

	// delete the object
	err = minioClient.RemoveObject(minioBucket, sanitize.BaseName(objectName))

	// delete mongo db entry
	go func() {
		defer Recover()
		err := MdbDeleteQuery(models.StorageTable, bson.M{"objectname": objectName})
		if err != nil && !strings.Contains(err.Error(), "not found") {
			RelaxLog(err)
		}
	}()

	return err
}

// Gets a public link for a file
// filename		: the name of the file
// filetype		: the type of the file
func GeneratePublicFileLink(filename, filehash string) (link string) {
	dots := strings.Count(filename, ".")
	if dots > 1 {
		filename = strings.Replace(filename, ".", "-", dots-1)
	}
	return fmt.Sprintf(GetConfig().Path("imageproxy.base_url").Data().(string),
		filehash, filename)
}

// TODO: prevent overwrites
// uploads a file to the minio object storage
// objectName	: the name of the file to upload
// data			: the data for the new object
// metadata		: additional metadata attached to the object
func uploadFile(objectName string, data []byte, metadata map[string]string) (err error) {
	// setup minioClient if not yet done
	if minioClient == nil {
		err = setupMinioClient()
		if err != nil {
			return err
		}
	}

	cache.GetLogger().WithField("module", "storage").Infof("uploading " + objectName + " to minio storage")

	options := minio.PutObjectOptions{}

	// add content type
	filetype, err := SniffMime(data)
	if err == nil {
		options.ContentType = filetype
	}

	// add metadata
	if metadata != nil && len(metadata) > 0 {
		options.UserMetadata = metadata
	}

	// upload the data
	_, err = minioClient.PutObject(minioBucket, sanitize.BaseName(objectName), bytes.NewReader(data), -1, options)
	return err
}

func getBucketCache(objectName string) (data []byte) {
	var err error

	if _, err = os.Stat(getObjectPath(objectName)); os.IsNotExist(err) {
		return nil
	}

	data, err = ioutil.ReadFile(getObjectPath(objectName))
	if err != nil {
		return nil
	}

	return data
}

func setBucketCache(objectName string, data []byte) (err error) {
	if _, err = os.Stat(filepath.Dir(getObjectPath(objectName))); os.IsNotExist(err) {
		err = os.MkdirAll(filepath.Dir(getObjectPath(objectName)), os.ModePerm)
		if err != nil {
			return err
		}
	}

	err = ioutil.WriteFile(getObjectPath(objectName), data, 0644)
	return err
}

func deleteBucketCache(objectName string) (err error) {
	if _, err = os.Stat(getObjectPath(objectName)); os.IsNotExist(err) {
		return nil
	}

	err = os.Remove(getObjectPath(objectName))
	return err
}

func getObjectPath(objectName string) (path string) {
	return GetConfig().Path("cache_folder").Data().(string) + "/minio-" + GetConfig().Path("s3.bucket").Data().(string) + "/" + sanitize.BaseName(objectName)
}

// Initialize the minio client object, and creates the bucket if it doesn't exist yet
func setupMinioClient() (err error) {
	minioLock.Lock()
	minioBucket = GetConfig().Path("s3.bucket").Data().(string)
	minioClient, err = minio.New(
		GetConfig().Path("s3.endpoint").Data().(string),
		GetConfig().Path("s3.access_key").Data().(string),
		GetConfig().Path("s3.secret_secret_key").Data().(string),
		true,
	)
	minioLock.Unlock()

	bucketExists, err := minioClient.BucketExists(minioBucket)
	if err != nil {
		return err
	}

	if !bucketExists {
		err = minioClient.MakeBucket(minioBucket, "ams3")
		if err != nil {
			return err
		}
	}

	return err
}
<|endoftext|>


<commit_before>package octokit

import (
	"net/http"
	"testing"

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

func TestSearchService_Users(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/users", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("user_search.json"))
	})

	url, err := SearchURL.Expand(map[string]interface{}{
		"type":  "users",
		"query": "dhruvsinghal"})
	assert.NoError(t, err)

	searchResults, result := client.Search(url).Users()

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 2)
	assert.Equal(t, len(searchResults.Items), 2)
	assert.Equal(t, searchResults.Items[0].ID, 3338555)
	assert.Equal(t, searchResults.Items[0].Login, "dhruvsinghal")
	assert.Equal(t, searchResults.Items[1].ID, 9294878)
	assert.Equal(t, searchResults.Items[1].Login, "dhruvsinghal5")
}

func TestSearchService_Issues(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/issues", func(w http.ResponseWriter,
		r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("issue_search.json"))
	})

	url, err := SearchURL.Expand(map[string]interface{}{
		"type":  "issues",
		"query": "color"})
	assert.NoError(t, err)

	searchResults, result := client.Search(url).Issues()

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 180172)
	assert.Equal(t, searchResults.Items[0].Number, 1551)
	assert.Equal(t, searchResults.Items[0].Title, "Colorizer")
	assert.Equal(t, searchResults.Items[1].Number, 3402)
	assert.Equal(t, searchResults.Items[1].Title, "Colorizer")
}

func TestSearchService_Repositories(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/repositories", func(w http.ResponseWriter,
		r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("repository_search.json"))
	})

	url, err := SearchURL.Expand(map[string]interface{}{
		"type":  "repositories",
		"query": "asdfghjk"})
	assert.NoError(t, err)

	searchResults, result := client.Search(url).Repositories()

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 21)
	assert.Equal(t, len(searchResults.Items), 21)
	assert.Equal(t, searchResults.Items[0].ID, 8269299)
	assert.Equal(t, searchResults.Items[0].FullName, "ysai/asdfghjk")
	assert.Equal(t, searchResults.Items[1].ID, 8511889)
	assert.Equal(t, searchResults.Items[1].FullName, "ines949494/ikadasd")
}

func TestSearchService_Code(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/code", func(w http.ResponseWriter,
		r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("code_search.json"))
	})

	url, err := SearchURL.Expand(map[string]interface{}{
		"type":  "code",
		"query": "addClass in:file language:js repo:jquery/jquery"})
	assert.NoError(t, err)

	searchResults, result := client.Search(url).Code()

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 7)
	assert.Equal(t, len(searchResults.Items), 7)
	assert.Equal(t, searchResults.Items[0].Name, "classes.js")
	assert.Equal(t, searchResults.Items[0].Path, "src/attributes/classes.js")
	assert.Equal(t, searchResults.Items[0].SHA,
		"f9dba94f7de43d6b6b7256e05e0d17c4741a4cde")
	assert.Equal(t, string(searchResults.Items[0].URL),
		"https://api.github.com/repositories/167174/contents/src/attributes/classes.js?ref=53aa87f3bf4284763405f3eb8affff296e55ba4f")
	assert.Equal(t, searchResults.Items[0].GitURL,
		"https://api.github.com/repositories/167174/git/blobs/f9dba94f7de43d6b6b7256e05e0d17c4741a4cde")
	assert.Equal(t, searchResults.Items[0].HTMLURL,
		"https://github.com/jquery/jquery/blob/53aa87f3bf4284763405f3eb8affff296e55ba4f/src/attributes/classes.js")
	assert.Equal(t, searchResults.Items[0].Repository.ID, 167174)
	assert.Equal(t, searchResults.Items[0].Repository.FullName,
		"jquery/jquery")
}
<commit_msg>Updated tests<commit_after>package octokit

import (
	"net/http"
	"testing"

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

func TestSearchService_Users(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/users", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("user_search.json"))
	})

	searchResults, result := client.Search(SearchURITemplate).Users(
		M{"query": "dhruvsinghal"})

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 2)
	assert.Equal(t, len(searchResults.Items), 2)
	assert.Equal(t, searchResults.Items[0].ID, 3338555)
	assert.Equal(t, searchResults.Items[0].Login, "dhruvsinghal")
	assert.Equal(t, searchResults.Items[1].ID, 9294878)
	assert.Equal(t, searchResults.Items[1].Login, "dhruvsinghal5")
}

func TestSearchService_Issues(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/issues", func(w http.ResponseWriter,
		r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("issue_search.json"))
	})

	searchResults, result := client.Search(SearchURITemplate).Issues(M{
		"query": "color"})

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 180172)
	assert.Equal(t, searchResults.Items[0].Number, 1551)
	assert.Equal(t, searchResults.Items[0].Title, "Colorizer")
	assert.Equal(t, searchResults.Items[1].Number, 3402)
	assert.Equal(t, searchResults.Items[1].Title, "Colorizer")
}

func TestSearchService_Repositories(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/repositories", func(w http.ResponseWriter,
		r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("repository_search.json"))
	})

	searchResults, result := client.Search(SearchURITemplate).Repositories(M{
		"query": "asdfghjk"})

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 21)
	assert.Equal(t, len(searchResults.Items), 21)
	assert.Equal(t, searchResults.Items[0].ID, 8269299)
	assert.Equal(t, searchResults.Items[0].FullName, "ysai/asdfghjk")
	assert.Equal(t, searchResults.Items[1].ID, 8511889)
	assert.Equal(t, searchResults.Items[1].FullName, "ines949494/ikadasd")
}

func TestSearchService_Code(t *testing.T) {
	setup()
	defer tearDown()

	mux.HandleFunc("/search/code", func(w http.ResponseWriter,
		r *http.Request) {
		testMethod(t, r, "GET")
		respondWithJSON(w, loadFixture("code_search.json"))
	})

	searchResults, result := client.Search(SearchURITemplate).Code(M{
		"query": "addClass in:file language:js repo:jquery/jquery"})

	assert.False(t, result.HasError())
	assert.False(t, searchResults.IncompleteResults)
	assert.Equal(t, searchResults.TotalCount, 7)
	assert.Equal(t, len(searchResults.Items), 7)
	assert.Equal(t, searchResults.Items[0].Name, "classes.js")
	assert.Equal(t, searchResults.Items[0].Path, "src/attributes/classes.js")
	assert.Equal(t, searchResults.Items[0].SHA,
		"f9dba94f7de43d6b6b7256e05e0d17c4741a4cde")
	assert.Equal(t, string(searchResults.Items[0].URL),
		"https://api.github.com/repositories/167174/contents/src/attributes/classes.js?ref=53aa87f3bf4284763405f3eb8affff296e55ba4f")
	assert.Equal(t, searchResults.Items[0].GitURL,
		"https://api.github.com/repositories/167174/git/blobs/f9dba94f7de43d6b6b7256e05e0d17c4741a4cde")
	assert.Equal(t, searchResults.Items[0].HTMLURL,
		"https://github.com/jquery/jquery/blob/53aa87f3bf4284763405f3eb8affff296e55ba4f/src/attributes/classes.js")
	assert.Equal(t, searchResults.Items[0].Repository.ID, 167174)
	assert.Equal(t, searchResults.Items[0].Repository.FullName,
		"jquery/jquery")
}
<|endoftext|>


<commit_before>package file

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/gogap/logrus"
)

func NewHook(file string) (f *FileHook) {
	path := strings.Split(file, "/")
	if len(path) > 1 {
		exec.Command("mkdir", path[0]).Run()
	}
	w := NewFileWriter()
	config := fmt.Sprintf(`{"filename":"%s","maxdays":7}`, file)
	w.Init(config)
	return &FileHook{w}
}

type FileHook struct {
	W LoggerInterface
}

func (hook *FileHook) Fire(entry *logrus.Entry) (err error) {
	message, err := getMessage(entry)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to read entry, %v", err)
		return err
	}
	switch entry.Level {
	case logrus.PanicLevel:
		fallthrough
	case logrus.FatalLevel:
		fallthrough
	case logrus.ErrorLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[ERROR] %s", message), LevelError)
	case logrus.WarnLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[WARN] %s", message), LevelWarn)
	case logrus.InfoLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[INFO] %s", message), LevelInfo)
	case logrus.DebugLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[DEBUG] %s", message), LevelDebug)
	default:
		return nil
	}
	return
}

func (hook *FileHook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}

func getMessage(entry *logrus.Entry) (message string, err error) {
	message = message + fmt.Sprintf("%s\n", entry.Message)
	for k, v := range entry.Data {
		if !strings.HasPrefix(k, "err_") {
			message = message + fmt.Sprintf("%s\n", v)
		}
	}
	message = fmt.Sprintf("%s", entry.Data["err_full"])
	return
}
<commit_msg>update file hook<commit_after>package file

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/gogap/logrus"
)

func NewHook(file string) (f *FileHook) {
	path := strings.Split(file, "/")
	if len(path) > 1 {
		exec.Command("mkdir", path[0]).Run()
	}
	w := NewFileWriter()
	config := fmt.Sprintf(`{"filename":"%s","maxdays":7}`, file)
	w.Init(config)
	return &FileHook{w}
}

type FileHook struct {
	W LoggerInterface
}

func (hook *FileHook) Fire(entry *logrus.Entry) (err error) {
	message, err := getMessage(entry)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to read entry, %v", err)
		return err
	}
	switch entry.Level {
	case logrus.PanicLevel:
		fallthrough
	case logrus.FatalLevel:
		fallthrough
	case logrus.ErrorLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[ERROR] %s", message), LevelError)
	case logrus.WarnLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[WARN] %s", message), LevelWarn)
	case logrus.InfoLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[INFO] %s", message), LevelInfo)
	case logrus.DebugLevel:
		return hook.W.WriteMsg(fmt.Sprintf("[DEBUG] %s", message), LevelDebug)
	default:
		return nil
	}
	return
}

func (hook *FileHook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}

func getMessage(entry *logrus.Entry) (message string, err error) {
	message = message + fmt.Sprintf("%s\n", entry.Message)
	for k, v := range entry.Data {
		if !strings.HasPrefix(k, "err_") {
			message = message + fmt.Sprintf("%s\n", v)
		}
	}
	if full, ok := entry.Data["err_full"]; ok {
		message = message + fmt.Sprintf("%v", full)
	}
	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Jan Broer
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package hosts

import (
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

type hostlist []*hostname

type hostname struct {
	domain string
	ip     net.IP
	ipv6   bool
}

// newHostlist creates a hostlist by parsing a file
func newHostlist(data []byte) *hostlist {
	hostlist := hostlist{}
	for _, v := range strings.Split(string(data), "\n") {
		for _, hostname := range parseLine(v) {
			err := hostlist.add(hostname)
			if err != nil {
				log.Printf("go-dnsmasq: warning: %s", err)
			}
		}
	}
	return &hostlist
}

func (h *hostlist) add(hostnamev *hostname) error {
	hostname := newHostname(hostnamev.domain, hostnamev.ip, hostnamev.ipv6)
	for _, found := range *h {
		if found.domain == hostname.domain && found.ip.Equal(hostname.ip) {
			return fmt.Errorf("Duplicate hostname entry for %s -> %s",
				hostname.domain, hostname.ip)
		}
	}
	*h = append(*h, hostname)
	return nil
}

// newHostname creates a new Hostname struct
func newHostname(domain string, ip net.IP, ipv6 bool) (host *hostname) {
	domain = strings.ToLower(domain)
	host = &hostname{domain, ip, ipv6}
	return
}

// ParseLine parses an individual line in a hostfile, which may contain one
// (un)commented ip and one or more hostnames. For example
//
//	127.0.0.1 localhost mysite1 mysite2
func parseLine(line string) hostlist {
	var hostnames hostlist

	if len(line) == 0 {
		return hostnames
	}

	// Parse leading # for disabled lines
	if line[0:1] == "#" {
		return hostnames
	}

	// Parse other #s for actual comments
	line = strings.Split(line, "#")[0]

	// Replace tabs and multispaces with single spaces throughout
	line = strings.Replace(line, "\t", " ", -1)
	for strings.Contains(line, "  ") {
		line = strings.Replace(line, "  ", " ", -1)
	}

	line = strings.TrimSpace(line)

	// Break line into words
	words := strings.Split(line, " ")
	for idx, word := range words {
		words[idx] = strings.TrimSpace(word)
	}

	// Separate the first bit (the ip) from the other bits (the domains)
	address := words[0]
	domains := words[1:]

	if strings.Contains(address, "%") {
		return hostnames
	}

	ip := net.ParseIP(address)

	var isIPv6 bool

	switch {
	case ip.To4() != nil:
		isIPv6 = false
	case ip.To16() != nil:
		isIPv6 = true
	default:
		log.Printf("go-dnsmasq: notice: invalid IP address found in hostsfile: %s", address)
		return hostnames
	}

	for _, v := range domains {
		hostname := newHostname(v, ip, isIPv6)
		hostnames = append(hostnames, hostname)
	}

	return hostnames
}

// hostsFileMetadata returns metadata about the hosts file.
func hostsFileMetadata(path string) (time.Time, int64, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return time.Time{}, 0, err
	}

	return fi.ModTime(), fi.Size(), nil
}
<commit_msg>Hostfile: Ignore loopback addresses<commit_after>// Copyright (c) 2015 Jan Broer
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package hosts

import (
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

type hostlist []*hostname

type hostname struct {
	domain string
	ip     net.IP
	ipv6   bool
}

// newHostlist creates a hostlist by parsing a file
func newHostlist(data []byte) *hostlist {
	hostlist := hostlist{}
	for _, v := range strings.Split(string(data), "\n") {
		for _, hostname := range parseLine(v) {
			err := hostlist.add(hostname)
			if err != nil {
				log.Printf("go-dnsmasq: warning: %s", err)
			}
		}
	}
	return &hostlist
}

func (h *hostlist) add(hostnamev *hostname) error {
	hostname := newHostname(hostnamev.domain, hostnamev.ip, hostnamev.ipv6)
	for _, found := range *h {
		if found.domain == hostname.domain && found.ip.Equal(hostname.ip) {
			return fmt.Errorf("Duplicate hostname entry for %s -> %s",
				hostname.domain, hostname.ip)
		}
	}
	*h = append(*h, hostname)
	return nil
}

// newHostname creates a new Hostname struct
func newHostname(domain string, ip net.IP, ipv6 bool) (host *hostname) {
	domain = strings.ToLower(domain)
	host = &hostname{domain, ip, ipv6}
	return
}

// ParseLine parses an individual line in a hostfile, which may contain one
// (un)commented ip and one or more hostnames. For example
//
//	127.0.0.1 localhost mysite1 mysite2
func parseLine(line string) hostlist {
	var hostnames hostlist

	if len(line) == 0 {
		return hostnames
	}

	// Parse leading # for disabled lines
	if line[0:1] == "#" {
		return hostnames
	}

	// Parse other #s for actual comments
	line = strings.Split(line, "#")[0]

	// Replace tabs and multispaces with single spaces throughout
	line = strings.Replace(line, "\t", " ", -1)
	for strings.Contains(line, "  ") {
		line = strings.Replace(line, "  ", " ", -1)
	}

	line = strings.TrimSpace(line)

	// Break line into words
	words := strings.Split(line, " ")
	for idx, word := range words {
		words[idx] = strings.TrimSpace(word)
	}

	// Separate the first bit (the ip) from the other bits (the domains)
	address := words[0]
	domains := words[1:]

	if strings.Contains(address, "%") {
		return hostnames
	}

	ip := net.ParseIP(address)

	var isIPv6 bool

	switch {
	case ip.IsLoopback():
		return hostnames
	case ip.IsUnspecified():
		return hostnames
	case ip.To4() != nil:
		isIPv6 = false
	case ip.To16() != nil:
		isIPv6 = true
	default:
		log.Printf("go-dnsmasq: notice: invalid IP address found in hostsfile: %s", address)
		return hostnames
	}

	for _, v := range domains {
		hostname := newHostname(v, ip, isIPv6)
		hostnames = append(hostnames, hostname)
	}

	return hostnames
}

// hostsFileMetadata returns metadata about the hosts file.
func hostsFileMetadata(path string) (time.Time, int64, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return time.Time{}, 0, err
	}

	return fi.ModTime(), fi.Size(), nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"encoding/json"
	"fmt"
	"github.com/conformal/btcjson"
	"time"
	"sync"
)

// Errors
var (
	// Standard JSON-RPC 2.0 errors
	InvalidRequest = btcjson.Error{
		Code: -32600,
		Message: "Invalid request",
	}
	MethodNotFound = btcjson.Error{
		Code: -32601,
		Message: "Method not found",
	}
	InvalidParams = btcjson.Error{
		Code: -32602,
		Message: "Invalid paramaters",
	}
	InternalError = btcjson.Error{
		Code: -32603,
		Message: "Internal error",
	}
	ParseError = btcjson.Error{
		Code: -32700,
		Message: "Parse error",
	}

	// General application defined errors
	MiscError = btcjson.Error{
		Code: -1,
		Message: "Miscellaneous error",
	}
	ForbiddenBySafeMode = btcjson.Error{
		Code: -2,
		Message: "Server is in safe mode, and command is not allowed in safe mode",
	}
	TypeError = btcjson.Error{
		Code: -3,
		Message: "Unexpected type was passed as parameter",
	}
	InvalidAddressOrKey = btcjson.Error{
		Code: -5,
		Message: "Invalid address or key",
	}
	OutOfMemory = btcjson.Error{
		Code: -7,
		Message: "Ran out of memory during operation",
	}
	InvalidParameter = btcjson.Error{
		Code: -8,
		Message: "Invalid, missing or duplicate parameter",
	}
	DatabaseError = btcjson.Error{
		Code: -20,
		Message: "Database error",
	}
	DeserializationError = btcjson.Error{
		Code: -22,
		Message: "Error parsing or validating structure in raw format",
	}

	// Wallet errors
	WalletError = btcjson.Error{
		Code: -4,
		Message: "Unspecified problem with wallet",
	}
	WalletInsufficientFunds = btcjson.Error{
		Code: -6,
		Message: "Not enough funds in wallet or account",
	}
	WalletInvalidAccountName = btcjson.Error{
		Code: -11,
		Message: "Invalid account name",
	}
	WalletKeypoolRanOut = btcjson.Error{
		Code: -12,
		Message: "Keypool ran out, call keypoolrefill first",
	}
	WalletUnlockNeeded = btcjson.Error{
		Code: -13,
		Message: "Enter the wallet passphrase with walletpassphrase first",
	}
	WalletPassphraseIncorrect = btcjson.Error{
		Code: -14,
		Message: "The wallet passphrase entered was incorrect",
	}
	WalletWrongEncState = btcjson.Error{
		Code: -15,
		Message: "Command given in wrong wallet encryption state",
	}
	WalletEncryptionFailed = btcjson.Error{
		Code: -16,
		Message: "Failed to encrypt the wallet",
	}
	WalletAlreadyUnlocked = btcjson.Error{
		Code: -17,
		Message: "Wallet is already unlocked",
	}
)

var (
	// seq holds the btcwallet sequence number for frontend messages
	// which must be sent to and received from btcd.  A Mutex protects
	// against concurrent access.
	seq = struct {
		sync.Mutex
		n uint64
	}{}

	// replyRouter maps uint64 ids to reply channels, so btcd replies can
	// be routed to the correct frontend.
	replyRouter = struct {
		sync.Mutex
		m map[uint64]chan []byte
	}{
		m: make(map[uint64]chan []byte),
	}

	incorrectParameters = errors.New("Incorrect parameters.")
)

// ProcessFrontendMsg checks the message sent from a frontend.  If the
// message method is one that must be handled by btcwallet, the request
// is processed here.  Otherwise, the message is sent to btcd.
func ProcessFrontendMsg(reply chan []byte, msg []byte) {
	cmd, err := btcjson.JSONGetMethod(msg)
	if err != nil {
		log.Error("Unable to parse JSON method from message.")
		return
	}

	switch cmd {
	case "getaddressesbyaccount":
		GetAddressesByAccount(reply, msg)
	case "getnewaddress":
		GetNewAddress(reply, msg)
	case "walletlock":
		WalletLock(reply, msg)
	case "walletpassphrase":
		WalletPassphrase(reply, msg)
	default:
		// btcwallet does not understand method.  Pass to btcd.
		log.Info("Unknown btcwallet method", cmd)

		seq.Lock()
		n := seq.n
		seq.n++
		seq.Unlock()

		var m map[string]interface{}
		json.Unmarshal(msg, &m)
		m["id"] = fmt.Sprintf("btcwallet(%v)-%v", n, m["id"])
		newMsg, err := json.Marshal(m)
		if err != nil {
			log.Info("Error marshalling json: " + err.Error())
		}
		replyRouter.Lock()
		replyRouter.m[n] = reply
		replyRouter.Unlock()
		btcdMsgs <- newMsg
	}
}

// ReplyError creates and marshalls a btcjson.Reply with the error e,
// sending the reply to a reply channel.
func ReplyError(reply chan []byte, id interface{}, e *btcjson.Error) {
	r := btcjson.Reply{
		Error: e,
		Id: &id,
	}
	if mr, err := json.Marshal(r); err != nil {
		reply <- mr
	}
}

// GetAddressesByAccount Gets all addresses for an account.
func GetAddressesByAccount(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	id := v["id"]
	r := btcjson.Reply{
		Id: &id,
	}
	if w := wallets[params[0].(string)]; w != nil {
		r.Result = w.GetActiveAddresses()
	} else {
		r.Result = []interface{}{}
	}
	mr, err := json.Marshal(r)
	if err != nil {
		log.Info("Error marshalling reply: %v", err)
		return
	}
	reply <- mr
}

// GetNewAddress gets or generates a new address for an account.
//
// TODO(jrick): support non-default account wallets.
func GetNewAddress(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	if len(params) == 0 || params[0].(string) == "" {
		if w := wallets[""]; w != nil {
			addr := w.NextUnusedAddress()
			id := v["id"]
			r := btcjson.Reply{
				Result: addr,
				Id:     &id,
			}
			mr, err := json.Marshal(r)
			if err != nil {
				log.Info("Error marshalling reply: %v", err)
				return
			}
			reply <- mr
		}
	}
}

// WalletLock locks the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletLock(reply chan []byte, msg []byte) {
	// TODO(jrick)
}

// WalletPassphrase stores the decryption key for the default account,
// unlocking the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletPassphrase(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	if len(params) != 2 {
		log.Error(incorrectParameters)
		return
	}
	passphrase, ok := params[0].(string)
	if !ok {
		log.Error(incorrectParameters)
		return
	}
	timeout, ok := params[1].(float64)
	if !ok {
		log.Error(incorrectParameters)
		return
	}

	if w := wallets[""]; w != nil {
		w.Unlock([]byte(passphrase))
		go func() {
			time.Sleep(time.Second * time.Duration(int64(timeout)))
			fmt.Println("finally locking")
			w.Lock()
		}()
	}
}
<commit_msg>Send error replies for walletpassphrase<commit_after>/*
 * Copyright (c) 2013 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"encoding/json"
	"fmt"
	"github.com/conformal/btcjson"
	"time"
	"sync"
)

// Errors
var (
	// Standard JSON-RPC 2.0 errors
	InvalidRequest = btcjson.Error{
		Code: -32600,
		Message: "Invalid request",
	}
	MethodNotFound = btcjson.Error{
		Code: -32601,
		Message: "Method not found",
	}
	InvalidParams = btcjson.Error{
		Code: -32602,
		Message: "Invalid paramaters",
	}
	InternalError = btcjson.Error{
		Code: -32603,
		Message: "Internal error",
	}
	ParseError = btcjson.Error{
		Code: -32700,
		Message: "Parse error",
	}

	// General application defined errors
	MiscError = btcjson.Error{
		Code: -1,
		Message: "Miscellaneous error",
	}
	ForbiddenBySafeMode = btcjson.Error{
		Code: -2,
		Message: "Server is in safe mode, and command is not allowed in safe mode",
	}
	TypeError = btcjson.Error{
		Code: -3,
		Message: "Unexpected type was passed as parameter",
	}
	InvalidAddressOrKey = btcjson.Error{
		Code: -5,
		Message: "Invalid address or key",
	}
	OutOfMemory = btcjson.Error{
		Code: -7,
		Message: "Ran out of memory during operation",
	}
	InvalidParameter = btcjson.Error{
		Code: -8,
		Message: "Invalid, missing or duplicate parameter",
	}
	DatabaseError = btcjson.Error{
		Code: -20,
		Message: "Database error",
	}
	DeserializationError = btcjson.Error{
		Code: -22,
		Message: "Error parsing or validating structure in raw format",
	}

	// Wallet errors
	WalletError = btcjson.Error{
		Code: -4,
		Message: "Unspecified problem with wallet",
	}
	WalletInsufficientFunds = btcjson.Error{
		Code: -6,
		Message: "Not enough funds in wallet or account",
	}
	WalletInvalidAccountName = btcjson.Error{
		Code: -11,
		Message: "Invalid account name",
	}
	WalletKeypoolRanOut = btcjson.Error{
		Code: -12,
		Message: "Keypool ran out, call keypoolrefill first",
	}
	WalletUnlockNeeded = btcjson.Error{
		Code: -13,
		Message: "Enter the wallet passphrase with walletpassphrase first",
	}
	WalletPassphraseIncorrect = btcjson.Error{
		Code: -14,
		Message: "The wallet passphrase entered was incorrect",
	}
	WalletWrongEncState = btcjson.Error{
		Code: -15,
		Message: "Command given in wrong wallet encryption state",
	}
	WalletEncryptionFailed = btcjson.Error{
		Code: -16,
		Message: "Failed to encrypt the wallet",
	}
	WalletAlreadyUnlocked = btcjson.Error{
		Code: -17,
		Message: "Wallet is already unlocked",
	}
)

var (
	// seq holds the btcwallet sequence number for frontend messages
	// which must be sent to and received from btcd.  A Mutex protects
	// against concurrent access.
	seq = struct {
		sync.Mutex
		n uint64
	}{}

	// replyRouter maps uint64 ids to reply channels, so btcd replies can
	// be routed to the correct frontend.
	replyRouter = struct {
		sync.Mutex
		m map[uint64]chan []byte
	}{
		m: make(map[uint64]chan []byte),
	}
)

// ProcessFrontendMsg checks the message sent from a frontend.  If the
// message method is one that must be handled by btcwallet, the request
// is processed here.  Otherwise, the message is sent to btcd.
func ProcessFrontendMsg(reply chan []byte, msg []byte) {
	cmd, err := btcjson.JSONGetMethod(msg)
	if err != nil {
		log.Error("Unable to parse JSON method from message.")
		return
	}

	switch cmd {
	case "getaddressesbyaccount":
		GetAddressesByAccount(reply, msg)
	case "getnewaddress":
		GetNewAddress(reply, msg)
	case "walletlock":
		WalletLock(reply, msg)
	case "walletpassphrase":
		WalletPassphrase(reply, msg)
	default:
		// btcwallet does not understand method.  Pass to btcd.
		log.Info("Unknown btcwallet method", cmd)

		seq.Lock()
		n := seq.n
		seq.n++
		seq.Unlock()

		var m map[string]interface{}
		json.Unmarshal(msg, &m)
		m["id"] = fmt.Sprintf("btcwallet(%v)-%v", n, m["id"])
		newMsg, err := json.Marshal(m)
		if err != nil {
			log.Info("Error marshalling json: " + err.Error())
		}
		replyRouter.Lock()
		replyRouter.m[n] = reply
		replyRouter.Unlock()
		btcdMsgs <- newMsg
	}
}

// ReplyError creates and marshalls a btcjson.Reply with the error e,
// sending the reply to a reply channel.
func ReplyError(reply chan []byte, id interface{}, e *btcjson.Error) {
	r := btcjson.Reply{
		Error: e,
		Id: &id,
	}
	if mr, err := json.Marshal(r); err != nil {
		reply <- mr
	}
}

// GetAddressesByAccount Gets all addresses for an account.
func GetAddressesByAccount(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	id := v["id"]
	r := btcjson.Reply{
		Id: &id,
	}
	if w := wallets[params[0].(string)]; w != nil {
		r.Result = w.GetActiveAddresses()
	} else {
		r.Result = []interface{}{}
	}
	mr, err := json.Marshal(r)
	if err != nil {
		log.Info("Error marshalling reply: %v", err)
		return
	}
	reply <- mr
}

// GetNewAddress gets or generates a new address for an account.
//
// TODO(jrick): support non-default account wallets.
func GetNewAddress(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	if len(params) == 0 || params[0].(string) == "" {
		if w := wallets[""]; w != nil {
			addr := w.NextUnusedAddress()
			id := v["id"]
			r := btcjson.Reply{
				Result: addr,
				Id:     &id,
			}
			mr, err := json.Marshal(r)
			if err != nil {
				log.Info("Error marshalling reply: %v", err)
				return
			}
			reply <- mr
		}
	}
}

// WalletLock locks the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletLock(reply chan []byte, msg []byte) {
	// TODO(jrick)
}

// WalletPassphrase stores the decryption key for the default account,
// unlocking the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletPassphrase(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	if len(params) != 2 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}
	passphrase, ok1 := params[0].(string)
	timeout, ok2 := params[1].(float64)
	if !ok1 || !ok2 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}

	if w := wallets[""]; w != nil {
		if err := w.Unlock([]byte(passphrase)); err != nil {
			ReplyError(reply, v["id"], &WalletPassphraseIncorrect)
			return
		}
		go func() {
			time.Sleep(time.Second * time.Duration(int64(timeout)))
			w.Lock()
		}()
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"encoding/json"
	"fmt"
	"github.com/conformal/btcjson"
	"github.com/conformal/btcwallet/wallet"
	"sync"
	"time"
)

// Errors
var (
	// Standard JSON-RPC 2.0 errors
	InvalidRequest = btcjson.Error{
		Code:    -32600,
		Message: "Invalid request",
	}
	MethodNotFound = btcjson.Error{
		Code:    -32601,
		Message: "Method not found",
	}
	InvalidParams = btcjson.Error{
		Code:    -32602,
		Message: "Invalid paramaters",
	}
	InternalError = btcjson.Error{
		Code:    -32603,
		Message: "Internal error",
	}
	ParseError = btcjson.Error{
		Code:    -32700,
		Message: "Parse error",
	}

	// General application defined errors
	MiscError = btcjson.Error{
		Code:    -1,
		Message: "Miscellaneous error",
	}
	ForbiddenBySafeMode = btcjson.Error{
		Code:    -2,
		Message: "Server is in safe mode, and command is not allowed in safe mode",
	}
	TypeError = btcjson.Error{
		Code:    -3,
		Message: "Unexpected type was passed as parameter",
	}
	InvalidAddressOrKey = btcjson.Error{
		Code:    -5,
		Message: "Invalid address or key",
	}
	OutOfMemory = btcjson.Error{
		Code:    -7,
		Message: "Ran out of memory during operation",
	}
	InvalidParameter = btcjson.Error{
		Code:    -8,
		Message: "Invalid, missing or duplicate parameter",
	}
	DatabaseError = btcjson.Error{
		Code:    -20,
		Message: "Database error",
	}
	DeserializationError = btcjson.Error{
		Code:    -22,
		Message: "Error parsing or validating structure in raw format",
	}

	// Wallet errors
	WalletError = btcjson.Error{
		Code:    -4,
		Message: "Unspecified problem with wallet",
	}
	WalletInsufficientFunds = btcjson.Error{
		Code:    -6,
		Message: "Not enough funds in wallet or account",
	}
	WalletInvalidAccountName = btcjson.Error{
		Code:    -11,
		Message: "Invalid account name",
	}
	WalletKeypoolRanOut = btcjson.Error{
		Code:    -12,
		Message: "Keypool ran out, call keypoolrefill first",
	}
	WalletUnlockNeeded = btcjson.Error{
		Code:    -13,
		Message: "Enter the wallet passphrase with walletpassphrase first",
	}
	WalletPassphraseIncorrect = btcjson.Error{
		Code:    -14,
		Message: "The wallet passphrase entered was incorrect",
	}
	WalletWrongEncState = btcjson.Error{
		Code:    -15,
		Message: "Command given in wrong wallet encryption state",
	}
	WalletEncryptionFailed = btcjson.Error{
		Code:    -16,
		Message: "Failed to encrypt the wallet",
	}
	WalletAlreadyUnlocked = btcjson.Error{
		Code:    -17,
		Message: "Wallet is already unlocked",
	}
)

var (
	// seq holds the btcwallet sequence number for frontend messages
	// which must be sent to and received from btcd.  A Mutex protects
	// against concurrent access.
	seq = struct {
		sync.Mutex
		n uint64
	}{}

	// replyRouter maps uint64 ids to reply channels, so btcd replies can
	// be routed to the correct frontend.
	replyRouter = struct {
		sync.Mutex
		m map[uint64]chan []byte
	}{
		m: make(map[uint64]chan []byte),
	}
)

// ProcessFrontendMsg checks the message sent from a frontend.  If the
// message method is one that must be handled by btcwallet, the request
// is processed here.  Otherwise, the message is sent to btcd.
func ProcessFrontendMsg(reply chan []byte, msg []byte) {
	cmd, err := btcjson.JSONGetMethod(msg)
	if err != nil {
		log.Error("Unable to parse JSON method from message.")
		return
	}

	switch cmd {
	// Standard bitcoind methods
	case "getaddressesbyaccount":
		GetAddressesByAccount(reply, msg)
	case "getnewaddress":
		GetNewAddress(reply, msg)
	case "walletlock":
		WalletLock(reply, msg)
	case "walletpassphrase":
		WalletPassphrase(reply, msg)

	// btcwallet extensions
	case "createencryptedwallet":
		CreateEncryptedWallet(reply, msg)
	case "walletislocked":
		WalletIsLocked(reply, msg)

	default:
		// btcwallet does not understand method.  Pass to btcd.
		log.Info("Unknown btcwallet method", cmd)

		seq.Lock()
		n := seq.n
		seq.n++
		seq.Unlock()

		var m map[string]interface{}
		json.Unmarshal(msg, &m)
		m["id"] = fmt.Sprintf("btcwallet(%v)-%v", n, m["id"])
		newMsg, err := json.Marshal(m)
		if err != nil {
			log.Info("Error marshalling json: " + err.Error())
		}
		replyRouter.Lock()
		replyRouter.m[n] = reply
		replyRouter.Unlock()
		btcdMsgs <- newMsg
	}
}

// ReplyError creates and marshalls a btcjson.Reply with the error e,
// sending the reply to a reply channel.
func ReplyError(reply chan []byte, id interface{}, e *btcjson.Error) {
	r := btcjson.Reply{
		Error: e,
		Id:    &id,
	}
	if mr, err := json.Marshal(r); err == nil {
		reply <- mr
	}
}

// ReplySuccess creates and marshalls a btcjson.Reply with the result r,
// sending the reply to a reply channel.
func ReplySuccess(reply chan []byte, id interface{}, result interface{}) {
	r := btcjson.Reply{
		Result: result,
		Id:     &id,
	}
	if mr, err := json.Marshal(r); err == nil {
		reply <- mr
	}
}

// GetAddressesByAccount Gets all addresses for an account.
func GetAddressesByAccount(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})

	var result interface{}
	wallets.RLock()
	w := wallets.m[params[0].(string)]
	wallets.RUnlock()
	if w != nil {
		result = w.Wallet.GetActiveAddresses()
	} else {
		result = []interface{}{}
	}
	ReplySuccess(reply, v["id"], result)
}

// GetNewAddress gets or generates a new address for an account.  If
// the requested wallet does not exist, a JSON error will be returned to
// the client.
func GetNewAddress(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	var wname string
	if len(params) == 0 || params[0].(string) == "" {
		wname = ""
	} else {
		wname = "params[0].(string)"
	}

	wallets.RLock()
	w := wallets.m[wname]
	wallets.RUnlock()
	if w != nil {
		// TODO(jrick): generate new addresses if the address pool is empty.
		addr := w.NextUnusedAddress()
		ReplySuccess(reply, v["id"], addr)
	} else {
		e := WalletInvalidAccountName
		e.Message = fmt.Sprintf("Wallet for account '%s' does not exist.", wname)
		ReplyError(reply, v["id"], &e)
	}
}

// CreateEncryptedWallet creates a new encrypted wallet.  The form of the command is:
//
//  createencryptedwallet [account] [description] [passphrase]
//
// All three parameters are required, and must be of type string.  If
// the wallet specified by account already exists, an invalid account
// name error is returned to the client.
func CreateEncryptedWallet(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	var wname string
	if len(params) != 3 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}
	wname, ok1 := params[0].(string)
	desc, ok2 := params[1].(string)
	pass, ok3 := params[2].(string)
	if !ok1 || !ok2 || !ok3 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}

	// Does this wallet already exist?
	wallets.RLock()
	if w := wallets.m[wname]; w != nil {
		e := WalletInvalidAccountName
		e.Message = "Wallet already exists."
		ReplyError(reply, v["id"], &e)
		return
	}
	wallets.RUnlock()

	w, err := wallet.NewWallet(wname, desc, []byte(pass))
	if err != nil {
		log.Error("Error creating wallet: " + err.Error())
		ReplyError(reply, v["id"], &InternalError)
		return
	}

	wallets.Lock()
	wallets.m[wname] = &BtcWallet{
		Wallet: w,
	}
	wallets.Unlock()
}

// WalletIsLocked returns whether the wallet used by the specified
// account, or default account, is locked.
func WalletIsLocked(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	account := ""
	if len(params) > 0 {
		if acct, ok := params[0].(string); ok {
			account = acct
		} else {
			ReplyError(reply, v["id"], &InvalidParams)
			return
		}
	}
	wallets.RLock()
	w := wallets.m[account]
	wallets.RUnlock()
	if w != nil {
		result := w.IsLocked()
		ReplySuccess(reply, v["id"], result)
	} else {
		ReplyError(reply, v["id"], &WalletInvalidAccountName)
	}
}

// WalletLock locks the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletLock(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	wallets.RLock()
	w := wallets.m[""]
	wallets.RUnlock()
	if w != nil {
		if err := w.Lock(); err != nil {
			ReplyError(reply, v["id"], &WalletWrongEncState)
		} else {
			ReplySuccess(reply, v["id"], nil)
			NotifyWalletLockStateChange(reply, true)
		}
	}
}

// WalletPassphrase stores the decryption key for the default account,
// unlocking the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletPassphrase(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	if len(params) != 2 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}
	passphrase, ok1 := params[0].(string)
	timeout, ok2 := params[1].(float64)
	if !ok1 || !ok2 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}

	wallets.RLock()
	w := wallets.m[""]
	wallets.RUnlock()
	if w != nil {
		if err := w.Unlock([]byte(passphrase)); err != nil {
			ReplyError(reply, v["id"], &WalletPassphraseIncorrect)
			return
		}
		ReplySuccess(reply, v["id"], nil)
		NotifyWalletLockStateChange(reply, false)
		go func() {
			time.Sleep(time.Second * time.Duration(int64(timeout)))
			w.Lock()
			NotifyWalletLockStateChange(reply, true)
		}()
	}
}

// NotifyWalletLockStateChange sends a notification to all frontends
// that the wallet has just been locked or unlocked.
func NotifyWalletLockStateChange(reply chan []byte, locked bool) {
	var id interface{} = "btcwallet:newwalletlockstate"
	m := btcjson.Reply{
		Result: locked,
		Id:     &id,
	}
	msg, _ := json.Marshal(&m)
	frontendNotificationMaster <- msg
}
<commit_msg>Implement 'getbalance' replier (currently replies with 0).<commit_after>/*
 * Copyright (c) 2013 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"encoding/json"
	"fmt"
	"github.com/conformal/btcjson"
	"github.com/conformal/btcwallet/wallet"
	"sync"
	"time"
)

// Errors
var (
	// Standard JSON-RPC 2.0 errors
	InvalidRequest = btcjson.Error{
		Code:    -32600,
		Message: "Invalid request",
	}
	MethodNotFound = btcjson.Error{
		Code:    -32601,
		Message: "Method not found",
	}
	InvalidParams = btcjson.Error{
		Code:    -32602,
		Message: "Invalid paramaters",
	}
	InternalError = btcjson.Error{
		Code:    -32603,
		Message: "Internal error",
	}
	ParseError = btcjson.Error{
		Code:    -32700,
		Message: "Parse error",
	}

	// General application defined errors
	MiscError = btcjson.Error{
		Code:    -1,
		Message: "Miscellaneous error",
	}
	ForbiddenBySafeMode = btcjson.Error{
		Code:    -2,
		Message: "Server is in safe mode, and command is not allowed in safe mode",
	}
	TypeError = btcjson.Error{
		Code:    -3,
		Message: "Unexpected type was passed as parameter",
	}
	InvalidAddressOrKey = btcjson.Error{
		Code:    -5,
		Message: "Invalid address or key",
	}
	OutOfMemory = btcjson.Error{
		Code:    -7,
		Message: "Ran out of memory during operation",
	}
	InvalidParameter = btcjson.Error{
		Code:    -8,
		Message: "Invalid, missing or duplicate parameter",
	}
	DatabaseError = btcjson.Error{
		Code:    -20,
		Message: "Database error",
	}
	DeserializationError = btcjson.Error{
		Code:    -22,
		Message: "Error parsing or validating structure in raw format",
	}

	// Wallet errors
	WalletError = btcjson.Error{
		Code:    -4,
		Message: "Unspecified problem with wallet",
	}
	WalletInsufficientFunds = btcjson.Error{
		Code:    -6,
		Message: "Not enough funds in wallet or account",
	}
	WalletInvalidAccountName = btcjson.Error{
		Code:    -11,
		Message: "Invalid account name",
	}
	WalletKeypoolRanOut = btcjson.Error{
		Code:    -12,
		Message: "Keypool ran out, call keypoolrefill first",
	}
	WalletUnlockNeeded = btcjson.Error{
		Code:    -13,
		Message: "Enter the wallet passphrase with walletpassphrase first",
	}
	WalletPassphraseIncorrect = btcjson.Error{
		Code:    -14,
		Message: "The wallet passphrase entered was incorrect",
	}
	WalletWrongEncState = btcjson.Error{
		Code:    -15,
		Message: "Command given in wrong wallet encryption state",
	}
	WalletEncryptionFailed = btcjson.Error{
		Code:    -16,
		Message: "Failed to encrypt the wallet",
	}
	WalletAlreadyUnlocked = btcjson.Error{
		Code:    -17,
		Message: "Wallet is already unlocked",
	}
)

var (
	// seq holds the btcwallet sequence number for frontend messages
	// which must be sent to and received from btcd.  A Mutex protects
	// against concurrent access.
	seq = struct {
		sync.Mutex
		n uint64
	}{}

	// replyRouter maps uint64 ids to reply channels, so btcd replies can
	// be routed to the correct frontend.
	replyRouter = struct {
		sync.Mutex
		m map[uint64]chan []byte
	}{
		m: make(map[uint64]chan []byte),
	}
)

// ProcessFrontendMsg checks the message sent from a frontend.  If the
// message method is one that must be handled by btcwallet, the request
// is processed here.  Otherwise, the message is sent to btcd.
func ProcessFrontendMsg(reply chan []byte, msg []byte) {
	cmd, err := btcjson.JSONGetMethod(msg)
	if err != nil {
		log.Error("Unable to parse JSON method from message.")
		return
	}

	switch cmd {
	// Standard bitcoind methods
	case "getaddressesbyaccount":
		GetAddressesByAccount(reply, msg)
	case "getbalance":
		GetBalance(reply, msg)
	case "getnewaddress":
		GetNewAddress(reply, msg)
	case "walletlock":
		WalletLock(reply, msg)
	case "walletpassphrase":
		WalletPassphrase(reply, msg)

	// btcwallet extensions
	case "createencryptedwallet":
		CreateEncryptedWallet(reply, msg)
	case "walletislocked":
		WalletIsLocked(reply, msg)

	default:
		// btcwallet does not understand method.  Pass to btcd.
		log.Info("Unknown btcwallet method ", cmd)

		seq.Lock()
		n := seq.n
		seq.n++
		seq.Unlock()

		var m map[string]interface{}
		json.Unmarshal(msg, &m)
		m["id"] = fmt.Sprintf("btcwallet(%v)-%v", n, m["id"])
		newMsg, err := json.Marshal(m)
		if err != nil {
			log.Info("Error marshalling json: " + err.Error())
		}
		replyRouter.Lock()
		replyRouter.m[n] = reply
		replyRouter.Unlock()
		btcdMsgs <- newMsg
	}
}

// ReplyError creates and marshalls a btcjson.Reply with the error e,
// sending the reply to a reply channel.
func ReplyError(reply chan []byte, id interface{}, e *btcjson.Error) {
	r := btcjson.Reply{
		Error: e,
		Id:    &id,
	}
	if mr, err := json.Marshal(r); err == nil {
		reply <- mr
	}
}

// ReplySuccess creates and marshalls a btcjson.Reply with the result r,
// sending the reply to a reply channel.
func ReplySuccess(reply chan []byte, id interface{}, result interface{}) {
	r := btcjson.Reply{
		Result: result,
		Id:     &id,
	}
	if mr, err := json.Marshal(r); err == nil {
		reply <- mr
	}
}

// GetAddressesByAccount replies with all addresses for an account.
func GetAddressesByAccount(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})

	var result interface{}
	wallets.RLock()
	w := wallets.m[params[0].(string)]
	wallets.RUnlock()
	if w != nil {
		result = w.Wallet.GetActiveAddresses()
	} else {
		result = []interface{}{}
	}
	ReplySuccess(reply, v["id"], result)
}

// GetBalance replies with the balance for an account (wallet).  If
// the requested wallet does not exist, a JSON error will be returned to
// the client.
//
// TODO(jrick): Actually calculate correct balance.
func GetBalance(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	var wname string
	if len(params) > 0 {
		if s, ok := params[0].(string); ok {
			wname = s
		} else {
			ReplyError(reply, v["id"], &InvalidParams)
		}
	}

	wallets.RLock()
	w := wallets.m[wname]
	wallets.RUnlock()
	var result interface{}
	if w != nil {
		result = 0 // TODO(jrick)
		ReplySuccess(reply, v["id"], result)
	} else {
		e := WalletInvalidAccountName
		e.Message = fmt.Sprintf("Wallet for account '%s' does not exist.", wname)
		ReplyError(reply, v["id"], &e)
	}
}

// GetNewAddress gets or generates a new address for an account.  If
// the requested wallet does not exist, a JSON error will be returned to
// the client.
func GetNewAddress(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	var wname string
	if len(params) == 0 || params[0].(string) == "" {
		wname = ""
	} else {
		wname = "params[0].(string)"
	}

	wallets.RLock()
	w := wallets.m[wname]
	wallets.RUnlock()
	if w != nil {
		// TODO(jrick): generate new addresses if the address pool is empty.
		addr := w.NextUnusedAddress()
		ReplySuccess(reply, v["id"], addr)
	} else {
		e := WalletInvalidAccountName
		e.Message = fmt.Sprintf("Wallet for account '%s' does not exist.", wname)
		ReplyError(reply, v["id"], &e)
	}
}

// CreateEncryptedWallet creates a new encrypted wallet.  The form of the command is:
//
//  createencryptedwallet [account] [description] [passphrase]
//
// All three parameters are required, and must be of type string.  If
// the wallet specified by account already exists, an invalid account
// name error is returned to the client.
func CreateEncryptedWallet(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	var wname string
	if len(params) != 3 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}
	wname, ok1 := params[0].(string)
	desc, ok2 := params[1].(string)
	pass, ok3 := params[2].(string)
	if !ok1 || !ok2 || !ok3 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}

	// Does this wallet already exist?
	wallets.RLock()
	if w := wallets.m[wname]; w != nil {
		e := WalletInvalidAccountName
		e.Message = "Wallet already exists."
		ReplyError(reply, v["id"], &e)
		return
	}
	wallets.RUnlock()

	w, err := wallet.NewWallet(wname, desc, []byte(pass))
	if err != nil {
		log.Error("Error creating wallet: " + err.Error())
		ReplyError(reply, v["id"], &InternalError)
		return
	}

	wallets.Lock()
	wallets.m[wname] = &BtcWallet{
		Wallet: w,
	}
	wallets.Unlock()
}

// WalletIsLocked returns whether the wallet used by the specified
// account, or default account, is locked.
func WalletIsLocked(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	account := ""
	if len(params) > 0 {
		if acct, ok := params[0].(string); ok {
			account = acct
		} else {
			ReplyError(reply, v["id"], &InvalidParams)
			return
		}
	}
	wallets.RLock()
	w := wallets.m[account]
	wallets.RUnlock()
	if w != nil {
		result := w.IsLocked()
		ReplySuccess(reply, v["id"], result)
	} else {
		ReplyError(reply, v["id"], &WalletInvalidAccountName)
	}
}

// WalletLock locks the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletLock(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	wallets.RLock()
	w := wallets.m[""]
	wallets.RUnlock()
	if w != nil {
		if err := w.Lock(); err != nil {
			ReplyError(reply, v["id"], &WalletWrongEncState)
		} else {
			ReplySuccess(reply, v["id"], nil)
			NotifyWalletLockStateChange(reply, true)
		}
	}
}

// WalletPassphrase stores the decryption key for the default account,
// unlocking the wallet.
//
// TODO(jrick): figure out how multiple wallets/accounts will work
// with this.
func WalletPassphrase(reply chan []byte, msg []byte) {
	var v map[string]interface{}
	json.Unmarshal(msg, &v)
	params := v["params"].([]interface{})
	if len(params) != 2 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}
	passphrase, ok1 := params[0].(string)
	timeout, ok2 := params[1].(float64)
	if !ok1 || !ok2 {
		ReplyError(reply, v["id"], &InvalidParams)
		return
	}

	wallets.RLock()
	w := wallets.m[""]
	wallets.RUnlock()
	if w != nil {
		if err := w.Unlock([]byte(passphrase)); err != nil {
			ReplyError(reply, v["id"], &WalletPassphraseIncorrect)
			return
		}
		ReplySuccess(reply, v["id"], nil)
		NotifyWalletLockStateChange(reply, false)
		go func() {
			time.Sleep(time.Second * time.Duration(int64(timeout)))
			w.Lock()
			NotifyWalletLockStateChange(reply, true)
		}()
	}
}

// NotifyWalletLockStateChange sends a notification to all frontends
// that the wallet has just been locked or unlocked.
func NotifyWalletLockStateChange(reply chan []byte, locked bool) {
	var id interface{} = "btcwallet:newwalletlockstate"
	m := btcjson.Reply{
		Result: locked,
		Id:     &id,
	}
	msg, _ := json.Marshal(&m)
	frontendNotificationMaster <- msg
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Eric Chou <zhssmail@gmail.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

// Package cmdmux is used to parse and route commands of terminal programe.
package cmdmux

import (
	"errors"
	"os"
	"strings"
)

// CmdHanlder is the type of callback function for command.
// if error is nil, then return value int is useful.
type CmdHandler func([]string, interface{}) (int, error)

// CmdMux represents programme's commands.
type CmdMux struct {
	root *cmdNode
}

type cmdNode struct {
	name     string
	subNodes []*cmdNode
	handler  CmdHandler
}

var std = New()

// New creates a CmdMux variable with the default "/" empty command.
func New() *CmdMux {
	n := &cmdNode{name: "/"}
	c := &CmdMux{root: n}
	return c
}

// String() return the string format of CmdMux c.
func (c *CmdMux) String() string {
	return func(n *cmdNode) string {
		var result string
		n.toString("", &result)
		return result
	}(c.root)
}

// HandleFunc registers the handler function for the given command path cmdpath
func (c *CmdMux) HandleFunc(cmdpath string, handler CmdHandler) error {
	return func(n *cmdNode, cmdpath string, handler CmdHandler) error {
		if cmdpath[0] != '/' {
			return errors.New("cmdmux: cmdpath should be absolute")
		}

		if cmdpath == "/" {
			n.handler = handler
			return nil
		}

		cmdStrs := strings.Split(cmdpath, "/")[1:]
		last := len(cmdStrs) - 1
		node := n
		for i, v := range cmdStrs {
			sub := node.hasSubNode(v)
			if sub == nil {
				sub = &cmdNode{name: v}
				if i == last {
					sub.handler = handler
				}
				node.subNodes = append(node.subNodes, sub)
			}
			node = sub
		}

		return nil
	}(c.root, cmdpath, handler)
}

// Execute accepts the os.Args as command and executes it with data
func (c *CmdMux) Execute(data interface{}) (int, error) {
	return func(node *cmdNode, args []string, data interface{}) (int, error) {
		var opts []string
		args = args[1:]
		for i, v := range args {
			sub := node.hasSubNode(v)
			if sub == nil {
				opts = args[i:]
				break
			}
			node = sub
		}

		if node == nil {
			return 0, errors.New("cmdmux: no such cmdNode")
		}

		if node.handler == nil {
			return 0, errors.New("cmdmux: cmdNode does not have a handler")
		}

		return node.handler(opts, data)
	}(c.root, os.Args, data)
}

// HandleFunc registers the handler function for the given command path cmdpath
// in the default CmdMux.
func HandleFunc(cmdpath string, handler CmdHandler) error {
	return std.HandleFunc(cmdpath, handler)
}

// Execute accepts the os.Args as command and executes it with data
// in the default CmdMux
func Execute(data interface{}) (int, error) {
	return std.Execute(data)
}

// String() return the string format of default CmdMux
func String() string {
	return std.String()
}

func (n *cmdNode) hasSubNode(name string) *cmdNode {
	for _, v := range n.subNodes {
		if v.name == name {
			return v
		}
	}
	return nil
}

func (n *cmdNode) toString(prefix string, result *string) {
	switch prefix {
	case "/":
		prefix = prefix + n.name
	case "":
		prefix = "/"
	default:
		prefix = prefix + "/" + n.name
	}

	if len(n.subNodes) == 0 {
		*result = *result + prefix + "\n"
	} else {
		for _, v := range n.subNodes {
			v.toString(prefix, result)
		}
	}
}
<commit_msg>remove closure<commit_after>/*
 * Copyright (C) 2016 Eric Chou <zhssmail@gmail.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

// Package cmdmux is used to parse and route commands of terminal programe.
package cmdmux

import (
	"errors"
	"strings"
)

// CmdHanlder is the type of callback function for command.
// if error is nil, then return value int is useful.
type CmdHandler func([]string, interface{}) (int, error)

// CmdMux represents programme's commands.
type CmdMux struct {
	root *cmdNode
}

type cmdNode struct {
	name     string
	subNodes []*cmdNode
	handler  CmdHandler
}

var std = New()

// New creates a CmdMux variable with the default "/" empty command.
func New() *CmdMux {
	n := &cmdNode{name: "/"}
	c := &CmdMux{root: n}
	return c
}

// String() return the string format of CmdMux c.
func (c *CmdMux) String() string {
	return func(n *cmdNode) string {
		var result string
		n.toString("", &result)
		return result
	}(c.root)
}

// HandleFunc registers the handler function for the given command path cmdpath
func (c *CmdMux) HandleFunc(cmdpath string, handler CmdHandler) error {
	n := c.root
	if cmdpath[0] != '/' {
		return errors.New("cmdmux: cmdpath should be absolute")
	}

	if cmdpath == "/" {
		n.handler = handler
		return nil
	}

	cmdStrs := strings.Split(cmdpath, "/")[1:]
	last := len(cmdStrs) - 1
	node := n
	for i, v := range cmdStrs {
		sub := node.hasSubNode(v)
		if sub == nil {
			sub = &cmdNode{name: v}
			if i == last {
				sub.handler = handler
			}
			node.subNodes = append(node.subNodes, sub)
		}
		node = sub
	}

	return nil
}

// Execute accepts the os.Args as command and executes it with data
func (c *CmdMux) Execute(data interface{}) (int, error) {
	node := c.root
	var opts []string
	args = args[1:]
	for i, v := range args {
		sub := node.hasSubNode(v)
		if sub == nil {
			opts = args[i:]
			break
		}
		node = sub
	}

	if node == nil {
		return 0, errors.New("cmdmux: no such cmdNode")
	}

	if node.handler == nil {
		return 0, errors.New("cmdmux: cmdNode does not have a handler")
	}

	return node.handler(opts, data)
}

// HandleFunc registers the handler function for the given command path cmdpath
// in the default CmdMux.
func HandleFunc(cmdpath string, handler CmdHandler) error {
	return std.HandleFunc(cmdpath, handler)
}

// Execute accepts the os.Args as command and executes it with data
// in the default CmdMux
func Execute(data interface{}) (int, error) {
	return std.Execute(data)
}

// String() return the string format of default CmdMux
func String() string {
	return std.String()
}

func (n *cmdNode) hasSubNode(name string) *cmdNode {
	for _, v := range n.subNodes {
		if v.name == name {
			return v
		}
	}
	return nil
}

func (n *cmdNode) toString(prefix string, result *string) {
	switch prefix {
	case "/":
		prefix = prefix + n.name
	case "":
		prefix = "/"
	default:
		prefix = prefix + "/" + n.name
	}

	if len(n.subNodes) == 0 {
		*result = *result + prefix + "\n"
	} else {
		for _, v := range n.subNodes {
			v.toString(prefix, result)
		}
	}
}
<|endoftext|>