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

package libkbfs

import (
	"fmt"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/keybase/kbfs/kbfscodec"
)

// BlockSplitterSimple implements the BlockSplitter interface by using
// a simple max-size algorithm to determine when to split blocks.
type BlockSplitterSimple struct {
	maxSize                 int64
	maxPtrsPerBlock         int
	blockChangeEmbedMaxSize uint64
	maxDirEntriesPerBlock   int
}

func getMaxDirEntriesPerBlock() (int, error) {
	dirEnv := os.Getenv("KEYBASE_BSPLIT_MAX_DIR_ENTRIES")
	if len(dirEnv) > 0 {
		maxDirEntriesPerBlock, err := strconv.Atoi(dirEnv)
		if err != nil {
			return 0, err
		}
		return maxDirEntriesPerBlock, nil
	}
	return 0, nil // disabled by default
}

// NewBlockSplitterSimple creates a new BlockSplittleSimple and
// adjusts the max size to try to match the desired size for file
// blocks, given the overhead of encoding a file block and the
// round-up padding we do.
func NewBlockSplitterSimple(desiredBlockSize int64,
	blockChangeEmbedMaxSize uint64, codec kbfscodec.Codec) (
	*BlockSplitterSimple, error) {
	// If the desired block size is exactly a power of 2, subtract one
	// from it to account for the padding we will do, which rounds up
	// when the encoded size is exactly a power of 2.
	if desiredBlockSize&(desiredBlockSize-1) == 0 {
		desiredBlockSize--
	}

	// Make a FileBlock of the expected size to see what the encoded
	// overhead is.
	block := NewFileBlock().(*FileBlock)
	fullData := make([]byte, desiredBlockSize)
	// Fill in the block with varying data to make sure not to trigger
	// any encoding optimizations.
	for i := range fullData {
		fullData[i] = byte(i)
	}

	maxSize := desiredBlockSize
	var encodedLen int64
	// Iterate until we find the right size (up to a maximum number of
	// attempts), because the overhead is not constant across
	// different Contents lengths (probably due to variable length
	// encoding of the buffer size).
	for i := 0; i < 10; i++ {
		block.Contents = fullData[:maxSize]
		encodedBlock, err := codec.Encode(block)
		if err != nil {
			return nil, err
		}

		encodedLen = int64(len(encodedBlock))
		if encodedLen >= 2*desiredBlockSize {
			return nil, fmt.Errorf("Encoded block of %d bytes is more than "+
				"twice as big as the desired block size %d",
				encodedLen, desiredBlockSize)
		}

		if encodedLen == desiredBlockSize {
			break
		}

		maxSize += (desiredBlockSize - encodedLen)
	}

	if encodedLen != desiredBlockSize {
		return nil, fmt.Errorf("Couldn't converge on a max block size for a "+
			"desired size of %d", desiredBlockSize)
	}

	// Trial and error shows that this magic 75% constant maximizes
	// the number of realistic indirect pointers you can fit into the
	// default block size.  TODO: calculate this number more exactly
	// during initialization for a given `maxSize`.
	maxPtrs := int(.75 * float64(maxSize/int64(bpSize)))
	if maxPtrs < 2 {
		maxPtrs = 2
	}

	maxDirEntriesPerBlock, err := getMaxDirEntriesPerBlock()
	if err != nil {
		return nil, err
	}
	// TODO(KBFS-3306): call `SetMaxDirEntriesByBlockSize()` here.

	return &BlockSplitterSimple{
		maxSize:                 maxSize,
		maxPtrsPerBlock:         maxPtrs,
		blockChangeEmbedMaxSize: blockChangeEmbedMaxSize,
		maxDirEntriesPerBlock:   maxDirEntriesPerBlock,
	}, nil
}

// SetMaxDirEntriesByBlockSize sets the maximum number of directory
// entries per directory block, based on the maximum block size.  If
// the `KEYBASE_BSPLIT_MAX_DIR_ENTRIES` is set, this function does
// nothing.
func (b *BlockSplitterSimple) SetMaxDirEntriesByBlockSize(
	codec kbfscodec.Codec) error {
	dirEnv := os.Getenv("KEYBASE_BSPLIT_MAX_DIR_ENTRIES")
	if len(dirEnv) > 0 {
		// Don't override the environment variable.
		return nil
	}

	block := NewDirBlock().(*DirBlock)
	bigName := strings.Repeat("a", maxNameBytesDefault)
	// Make "typical" DirEntry, though the max dir entry is a bit
	// bigger than this (can contain a variable-length symlink path,
	// for example).
	de := DirEntry{
		BlockInfo: BlockInfo{
			BlockPointer: BlockPointer{
				DirectType: DirectBlock,
			},
		},
		EntryInfo: EntryInfo{
			PrevRevisions: PrevRevisions{
				{Revision: 0, Count: 0},
				{Revision: 1, Count: 1},
				{Revision: 2, Count: 2},
				{Revision: 3, Count: 3},
				{Revision: 4, Count: 4},
			},
		},
	}
	block.Children[bigName] = de
	encodedBlock, err := codec.Encode(block)
	if err != nil {
		return err
	}
	oneEntrySize := int64(len(encodedBlock))
	b.maxDirEntriesPerBlock = int(b.maxSize / oneEntrySize)
	if b.maxDirEntriesPerBlock == 0 {
		b.maxDirEntriesPerBlock = 1
	}
	return nil
}

// CopyUntilSplit implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) CopyUntilSplit(
	block *FileBlock, lastBlock bool, data []byte, off int64) int64 {
	n := int64(len(data))
	currLen := int64(len(block.Contents))
	// lastBlock is irrelevant since we only copy fixed sizes

	toCopy := n
	if currLen < (off + n) {
		moreNeeded := (n + off) - currLen
		// Reduce the number of additional bytes if it will take this block
		// over maxSize.
		if moreNeeded+currLen > b.maxSize {
			moreNeeded = b.maxSize - currLen
			if moreNeeded < 0 {
				// If it is already over maxSize w/o any added bytes,
				// just give up.
				return 0
			}
			// only copy to the end of the block
			toCopy = b.maxSize - off
		}

		if moreNeeded > 0 {
			block.Contents = append(block.Contents, make([]byte, moreNeeded)...)
		}
	}

	// we may have filled out the block above, but we still can't copy anything
	if off > int64(len(block.Contents)) {
		return 0
	}

	copy(block.Contents[off:off+toCopy], data[:toCopy])
	return toCopy
}

// CheckSplit implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) CheckSplit(block *FileBlock) int64 {
	// The split will always be right
	return 0
}

// MaxPtrsPerBlock implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) MaxPtrsPerBlock() int {
	return b.maxPtrsPerBlock
}

// ShouldEmbedBlockChanges implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) ShouldEmbedBlockChanges(
	bc *BlockChanges) bool {
	return bc.SizeEstimate() <= b.blockChangeEmbedMaxSize
}

// SplitDirIfNeeded implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) SplitDirIfNeeded(block *DirBlock) (
	[]*DirBlock, *StringOffset) {
	if block.IsIndirect() {
		panic("SplitDirIfNeeded must be given only a direct block")
	}

	if b.maxDirEntriesPerBlock == 0 ||
		len(block.Children) <= b.maxDirEntriesPerBlock {
		return []*DirBlock{block}, nil
	}

	// Sort the entries and split them down the middle.
	names := make([]string, 0, len(block.Children))
	for name := range block.Children {
		names = append(names, name)
	}

	sort.Strings(names)
	// Delete the second half of the names from the original block,
	// and add to the new block.
	newBlock := NewDirBlock().(*DirBlock)
	startOff := int(len(names) / 2)
	for _, name := range names[int(len(names)/2):] {
		newBlock.Children[name] = block.Children[name]
		delete(block.Children, name)
	}
	newOffset := StringOffset(names[startOff])
	return []*DirBlock{block, newBlock}, &newOffset
}
<commit_msg>bsplitter: remove old TODO<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/keybase/kbfs/kbfscodec"
)

// BlockSplitterSimple implements the BlockSplitter interface by using
// a simple max-size algorithm to determine when to split blocks.
type BlockSplitterSimple struct {
	maxSize                 int64
	maxPtrsPerBlock         int
	blockChangeEmbedMaxSize uint64
	maxDirEntriesPerBlock   int
}

func getMaxDirEntriesPerBlock() (int, error) {
	dirEnv := os.Getenv("KEYBASE_BSPLIT_MAX_DIR_ENTRIES")
	if len(dirEnv) > 0 {
		maxDirEntriesPerBlock, err := strconv.Atoi(dirEnv)
		if err != nil {
			return 0, err
		}
		return maxDirEntriesPerBlock, nil
	}
	return 0, nil // disabled by default
}

// NewBlockSplitterSimple creates a new BlockSplittleSimple and
// adjusts the max size to try to match the desired size for file
// blocks, given the overhead of encoding a file block and the
// round-up padding we do.
func NewBlockSplitterSimple(desiredBlockSize int64,
	blockChangeEmbedMaxSize uint64, codec kbfscodec.Codec) (
	*BlockSplitterSimple, error) {
	// If the desired block size is exactly a power of 2, subtract one
	// from it to account for the padding we will do, which rounds up
	// when the encoded size is exactly a power of 2.
	if desiredBlockSize&(desiredBlockSize-1) == 0 {
		desiredBlockSize--
	}

	// Make a FileBlock of the expected size to see what the encoded
	// overhead is.
	block := NewFileBlock().(*FileBlock)
	fullData := make([]byte, desiredBlockSize)
	// Fill in the block with varying data to make sure not to trigger
	// any encoding optimizations.
	for i := range fullData {
		fullData[i] = byte(i)
	}

	maxSize := desiredBlockSize
	var encodedLen int64
	// Iterate until we find the right size (up to a maximum number of
	// attempts), because the overhead is not constant across
	// different Contents lengths (probably due to variable length
	// encoding of the buffer size).
	for i := 0; i < 10; i++ {
		block.Contents = fullData[:maxSize]
		encodedBlock, err := codec.Encode(block)
		if err != nil {
			return nil, err
		}

		encodedLen = int64(len(encodedBlock))
		if encodedLen >= 2*desiredBlockSize {
			return nil, fmt.Errorf("Encoded block of %d bytes is more than "+
				"twice as big as the desired block size %d",
				encodedLen, desiredBlockSize)
		}

		if encodedLen == desiredBlockSize {
			break
		}

		maxSize += (desiredBlockSize - encodedLen)
	}

	if encodedLen != desiredBlockSize {
		return nil, fmt.Errorf("Couldn't converge on a max block size for a "+
			"desired size of %d", desiredBlockSize)
	}

	// Trial and error shows that this magic 75% constant maximizes
	// the number of realistic indirect pointers you can fit into the
	// default block size.  TODO: calculate this number more exactly
	// during initialization for a given `maxSize`.
	maxPtrs := int(.75 * float64(maxSize/int64(bpSize)))
	if maxPtrs < 2 {
		maxPtrs = 2
	}

	maxDirEntriesPerBlock, err := getMaxDirEntriesPerBlock()
	if err != nil {
		return nil, err
	}

	return &BlockSplitterSimple{
		maxSize:                 maxSize,
		maxPtrsPerBlock:         maxPtrs,
		blockChangeEmbedMaxSize: blockChangeEmbedMaxSize,
		maxDirEntriesPerBlock:   maxDirEntriesPerBlock,
	}, nil
}

// SetMaxDirEntriesByBlockSize sets the maximum number of directory
// entries per directory block, based on the maximum block size.  If
// the `KEYBASE_BSPLIT_MAX_DIR_ENTRIES` is set, this function does
// nothing.
func (b *BlockSplitterSimple) SetMaxDirEntriesByBlockSize(
	codec kbfscodec.Codec) error {
	dirEnv := os.Getenv("KEYBASE_BSPLIT_MAX_DIR_ENTRIES")
	if len(dirEnv) > 0 {
		// Don't override the environment variable.
		return nil
	}

	block := NewDirBlock().(*DirBlock)
	bigName := strings.Repeat("a", maxNameBytesDefault)
	// Make "typical" DirEntry, though the max dir entry is a bit
	// bigger than this (can contain a variable-length symlink path,
	// for example).
	de := DirEntry{
		BlockInfo: BlockInfo{
			BlockPointer: BlockPointer{
				DirectType: DirectBlock,
			},
		},
		EntryInfo: EntryInfo{
			PrevRevisions: PrevRevisions{
				{Revision: 0, Count: 0},
				{Revision: 1, Count: 1},
				{Revision: 2, Count: 2},
				{Revision: 3, Count: 3},
				{Revision: 4, Count: 4},
			},
		},
	}
	block.Children[bigName] = de
	encodedBlock, err := codec.Encode(block)
	if err != nil {
		return err
	}
	oneEntrySize := int64(len(encodedBlock))
	b.maxDirEntriesPerBlock = int(b.maxSize / oneEntrySize)
	if b.maxDirEntriesPerBlock == 0 {
		b.maxDirEntriesPerBlock = 1
	}
	return nil
}

// CopyUntilSplit implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) CopyUntilSplit(
	block *FileBlock, lastBlock bool, data []byte, off int64) int64 {
	n := int64(len(data))
	currLen := int64(len(block.Contents))
	// lastBlock is irrelevant since we only copy fixed sizes

	toCopy := n
	if currLen < (off + n) {
		moreNeeded := (n + off) - currLen
		// Reduce the number of additional bytes if it will take this block
		// over maxSize.
		if moreNeeded+currLen > b.maxSize {
			moreNeeded = b.maxSize - currLen
			if moreNeeded < 0 {
				// If it is already over maxSize w/o any added bytes,
				// just give up.
				return 0
			}
			// only copy to the end of the block
			toCopy = b.maxSize - off
		}

		if moreNeeded > 0 {
			block.Contents = append(block.Contents, make([]byte, moreNeeded)...)
		}
	}

	// we may have filled out the block above, but we still can't copy anything
	if off > int64(len(block.Contents)) {
		return 0
	}

	copy(block.Contents[off:off+toCopy], data[:toCopy])
	return toCopy
}

// CheckSplit implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) CheckSplit(block *FileBlock) int64 {
	// The split will always be right
	return 0
}

// MaxPtrsPerBlock implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) MaxPtrsPerBlock() int {
	return b.maxPtrsPerBlock
}

// ShouldEmbedBlockChanges implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) ShouldEmbedBlockChanges(
	bc *BlockChanges) bool {
	return bc.SizeEstimate() <= b.blockChangeEmbedMaxSize
}

// SplitDirIfNeeded implements the BlockSplitter interface for
// BlockSplitterSimple.
func (b *BlockSplitterSimple) SplitDirIfNeeded(block *DirBlock) (
	[]*DirBlock, *StringOffset) {
	if block.IsIndirect() {
		panic("SplitDirIfNeeded must be given only a direct block")
	}

	if b.maxDirEntriesPerBlock == 0 ||
		len(block.Children) <= b.maxDirEntriesPerBlock {
		return []*DirBlock{block}, nil
	}

	// Sort the entries and split them down the middle.
	names := make([]string, 0, len(block.Children))
	for name := range block.Children {
		names = append(names, name)
	}

	sort.Strings(names)
	// Delete the second half of the names from the original block,
	// and add to the new block.
	newBlock := NewDirBlock().(*DirBlock)
	startOff := int(len(names) / 2)
	for _, name := range names[int(len(names)/2):] {
		newBlock.Children[name] = block.Children[name]
		delete(block.Children, name)
	}
	newOffset := StringOffset(names[startOff])
	return []*DirBlock{block, newBlock}, &newOffset
}
<|endoftext|>


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

package libkbfs

import (
	"path/filepath"
	"sync"
	"time"

	"github.com/keybase/client/go/logger"
	"github.com/keybase/kbfs/kbfsblock"
	"github.com/keybase/kbfs/kbfscrypto"
	"github.com/keybase/kbfs/tlf"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/storage"
	"github.com/syndtr/goleveldb/leveldb/util"
	"golang.org/x/net/context"
)

const (
	// 10 GB maximum storage by default
	defaultDiskBlockCacheMaxBytes uint64 = 10 * (1 << 30)
	evictionConsiderationFactor   uint   = 3
	blockDbFilename               string = "diskCacheBlocks.leveldb"
	lruDbFilename                 string = "diskCacheLRU.leveldb"
)

// diskBlockCacheEntry packages an encoded block and serverHalf into one data
// structure, allowing us to encode it as one set of bytes.
type diskBlockCacheEntry struct {
	Buf        []byte
	ServerHalf kbfscrypto.BlockCryptKeyServerHalf
}

// diskBlockCacheConfig specifies the interfaces that a DiskBlockCacheStandard
// needs to perform its functions. This adheres to the standard libkbfs Config
// API.
type diskBlockCacheConfig interface {
	codecGetter
	logMaker
}

// DiskBlockCacheStandard is the standard implementation for DiskBlockCache.
type DiskBlockCacheStandard struct {
	config   diskBlockCacheConfig
	maxBytes uint64
	log      logger.Logger
	// protects the disk caches from being shutdown while they're being
	// accessed
	lock     sync.RWMutex
	isClosed bool
	blockDb  *leveldb.DB
	lruDb    *leveldb.DB
}

var _ DiskBlockCache = (*DiskBlockCacheStandard)(nil)

// openLevelDB opens or recovers a leveldb.DB with a passed-in storage.Storage
// as its underlying storage layer.
func openLevelDB(stor storage.Storage) (db *leveldb.DB, err error) {
	db, err = leveldb.Open(stor, leveldbOptions)
	if _, isErrCorrupted := err.(*storage.ErrCorrupted); isErrCorrupted {
		return leveldb.Recover(stor, leveldbOptions)
	}
	return db, err
}

// newDiskBlockCacheStandardFromStorage creates a new *DiskBlockCacheStandard
// with the passed-in storage.Storage interfaces as storage layers for each
// cache.
func newDiskBlockCacheStandardFromStorage(config diskBlockCacheConfig,
	blockStorage, lruStorage storage.Storage, maxBytes uint64) (
	*DiskBlockCacheStandard, error) {
	log := config.MakeLogger("KBC")
	blockDb, err := openLevelDB(blockStorage)
	if err != nil {
		return nil, err
	}

	lruDb, err := openLevelDB(lruStorage)
	if err != nil {
		blockDb.Close()
		return nil, err
	}
	return &DiskBlockCacheStandard{
		config:   config,
		maxBytes: maxBytes,
		log:      log,
		blockDb:  blockDb,
		lruDb:    lruDb,
	}, nil
}

// newDiskBlockCacheStandard creates a new *DiskBlockCacheStandard with a
// specified directory on the filesystem as storage.
func newDiskBlockCacheStandard(config diskBlockCacheConfig, dirPath string,
	maxBytes uint64) (*DiskBlockCacheStandard, error) {
	blockDbPath := filepath.Join(dirPath, blockDbFilename)
	blockStorage, err := storage.OpenFile(blockDbPath, false)
	if err != nil {
		return nil, err
	}
	lruDbPath := filepath.Join(dirPath, lruDbFilename)
	lruStorage, err := storage.OpenFile(lruDbPath, false)
	if err != nil {
		blockStorage.Close()
		return nil, err
	}
	return newDiskBlockCacheStandardFromStorage(config, blockStorage,
		lruStorage, maxBytes)
}

// TODO: Fix getSizesLocked(), where leveldb.DB.SizeOf() currently doesn't work
// for the full range. As it is right now, this is a very cheap operation,
// since it just takes the beginning and end offsets from levelDB and returns
// the difference.
func (cache *DiskBlockCacheStandard) getSizesLocked() (blockSize int64,
	lruSize int64, err error) {
	// a util.Range{nil, nil} is supposed to work for the full range based on
	// other methods, but it appears not to work for SizeOf. Need to implement
	// by querying the files in the LevelDB directory instead. Ugh.
	blockSizes, err := cache.blockDb.SizeOf([]util.Range{{}})
	if err != nil {
		return -1, -1, err
	}
	lruSizes, err := cache.lruDb.SizeOf([]util.Range{{}})
	if err != nil {
		return -1, -1, err
	}
	return blockSizes[0], lruSizes[0], nil
}

// compactCachesLocked manually forces both the block cache and LRU cache to
// compact their underlying data.
func (cache *DiskBlockCacheStandard) compactCachesLocked(ctx context.Context) {
	cache.blockDb.CompactRange(util.Range{})
	cache.lruDb.CompactRange(util.Range{})
	bSize, lruSize, err := cache.getSizesLocked()
	cache.log.CDebugf(ctx, "Disk Cache bSize=%d lruSize=%d err=%+v", bSize, lruSize, err)
}

// lruKey generates an LRU cache key from a tlf.ID and a binary-encoded block
// ID.
func (_ *DiskBlockCacheStandard) lruKey(tlfID tlf.ID, blockKey []byte) []byte {
	return append(tlfID.Bytes(), blockKey...)
}

// updateLruLocked updates the LRU time of a block in the LRU cache to
// time.Now().
func (cache *DiskBlockCacheStandard) updateLruLocked(tlfID tlf.ID,
	blockKey []byte) error {
	key := cache.lruKey(tlfID, blockKey)
	val, err := time.Now().MarshalBinary()
	if err != nil {
		return err
	}
	cache.lruDb.Put(key, val, nil)
	return nil
}

// timeFromBytes converts a value from the LRU cache into a time.Time.
func (_ *DiskBlockCacheStandard) timeFromBytes(b []byte) (t time.Time, err error) {
	err = t.UnmarshalBinary(b)
	return t, err
}

// getLru retrieves the LRU time for a block in the cache, or returns
// leveldb.ErrNotFound and a zero-valued time.Time otherwise.
func (cache *DiskBlockCacheStandard) getLru(tlfID tlf.ID,
	blockID kbfsblock.ID) (time.Time, error) {
	lruKey := cache.lruKey(tlfID, blockID.Bytes())
	lruBytes, err := cache.lruDb.Get(lruKey, nil)
	if err != nil {
		return time.Time{}, err
	}
	return cache.timeFromBytes(lruBytes)
}

// decodeBlockCacheEntry decodes a disk block cache entry buffer into an
// encoded block and server half.
func (cache *DiskBlockCacheStandard) decodeBlockCacheEntry(buf []byte) ([]byte,
	kbfscrypto.BlockCryptKeyServerHalf, error) {
	entry := diskBlockCacheEntry{}
	err := cache.config.Codec().Decode(buf, &entry)
	if err != nil {
		return nil, kbfscrypto.BlockCryptKeyServerHalf{}, err
	}
	return entry.Buf, entry.ServerHalf, nil
}

// encodeBlockCacheEntry encodes an encoded block and serverHalf into a single
// buffer.
func (cache *DiskBlockCacheStandard) encodeBlockCacheEntry(buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf) ([]byte, error) {
	entry := diskBlockCacheEntry{
		Buf:        buf,
		ServerHalf: serverHalf,
	}
	return cache.config.Codec().Encode(&entry)
}

// Get implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Get(ctx context.Context, tlfID tlf.ID,
	blockID kbfsblock.ID) (buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf, err error) {
	cache.lock.RLock()
	defer cache.lock.RUnlock()
	if cache.isClosed {
		return nil, kbfscrypto.BlockCryptKeyServerHalf{},
			DiskCacheClosedError{"Get"}
	}
	defer func() {
		cache.log.CDebugf(ctx, "Cache Get id=%s tlf=%s bSize=%d err=%+v", blockID, tlfID, len(buf), err)
	}()
	var entry []byte
	err = runUnlessCanceled(ctx, func() error {
		blockKey := blockID.Bytes()
		buf, err := cache.blockDb.Get(blockKey, nil)
		if err != nil {
			return NoSuchBlockError{blockID}
		}
		err = cache.updateLruLocked(tlfID, blockKey)
		if err != nil {
			return err
		}
		entry = buf
		return nil
	})
	if err != nil {
		return nil, kbfscrypto.BlockCryptKeyServerHalf{}, err
	}
	return cache.decodeBlockCacheEntry(entry)
}

// Put implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Put(ctx context.Context, tlfID tlf.ID,
	blockID kbfsblock.ID, buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf) error {
	cache.lock.RLock()
	defer cache.lock.RUnlock()
	if cache.isClosed {
		return DiskCacheClosedError{"Put"}
	}
	blockLen := len(buf)
	// TODO: accounting
	entry, err := cache.encodeBlockCacheEntry(buf, serverHalf)
	encodeLen := len(entry)
	cache.log.CDebugf(ctx, "Cache Put id=%s tlf=%s bSize=%d entrySize=%d err=%+v", blockID, tlfID, blockLen, encodeLen, err)
	if err != nil {
		return err
	}
	blockKey := blockID.Bytes()
	err = cache.blockDb.Put(blockKey, entry, nil)
	if err != nil {
		return err
	}
	return cache.updateLruLocked(tlfID, blockKey)
}

// Delete implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Delete(ctx context.Context, tlfID tlf.ID,
	blockIDs []kbfsblock.ID) error {
	cache.lock.RLock()
	defer cache.lock.RUnlock()
	if cache.isClosed {
		return DiskCacheClosedError{"Delete"}
	}
	if len(blockIDs) == 0 {
		return nil
	}
	cache.log.CDebugf(ctx, "Cache Delete tlf=%s numBlocks=%d", tlfID, len(blockIDs))
	blockBatch := new(leveldb.Batch)
	lruBatch := new(leveldb.Batch)
	for _, id := range blockIDs {
		blockKey := id.Bytes()
		blockBatch.Delete(blockKey)
		lruKey := cache.lruKey(tlfID, blockKey)
		lruBatch.Delete(lruKey)
	}
	if err := cache.blockDb.Write(blockBatch, nil); err != nil {
		return err
	}
	if err := cache.lruDb.Write(lruBatch, nil); err != nil {
		return err
	}

	cache.compactCachesLocked(ctx)

	return nil
}

// evictLocked evicts a number of blocks from the cache.
func (cache *DiskBlockCacheStandard) evictLocked(ctx context.Context,
	tlfID tlf.ID, numBlocks int) error {
	// Use kbfscrypto.MakeTemporaryID() to create a random hash ID. Then begin
	// an iterator into cache.lruDb.Range(tlfID + b, tlfID + MaxBlockID) and
	// iterate from there to get numBlocks * evictionConsiderationFactor block
	// IDs.  We sort the resulting blocks by value (LRU time) and pick the
	// minimum numBlocks. We then call cache.Delete() on that list of block
	// IDs.
	return nil
}

// Shutdown implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Shutdown() {
	cache.lock.Lock()
	defer cache.lock.Unlock()
	if cache.isClosed {
		return
	}
	cache.isClosed = true
	cache.blockDb.Close()
	cache.lruDb.Close()
}
<commit_msg>disk_block_cache: Lint fixes<commit_after>// Copyright 2017 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"path/filepath"
	"sync"
	"time"

	"github.com/keybase/client/go/logger"
	"github.com/keybase/kbfs/kbfsblock"
	"github.com/keybase/kbfs/kbfscrypto"
	"github.com/keybase/kbfs/tlf"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/storage"
	"github.com/syndtr/goleveldb/leveldb/util"
	"golang.org/x/net/context"
)

const (
	// 10 GB maximum storage by default
	defaultDiskBlockCacheMaxBytes uint64 = 10 * (1 << 30)
	evictionConsiderationFactor   uint   = 3
	blockDbFilename               string = "diskCacheBlocks.leveldb"
	lruDbFilename                 string = "diskCacheLRU.leveldb"
)

// diskBlockCacheEntry packages an encoded block and serverHalf into one data
// structure, allowing us to encode it as one set of bytes.
type diskBlockCacheEntry struct {
	Buf        []byte
	ServerHalf kbfscrypto.BlockCryptKeyServerHalf
}

// diskBlockCacheConfig specifies the interfaces that a DiskBlockCacheStandard
// needs to perform its functions. This adheres to the standard libkbfs Config
// API.
type diskBlockCacheConfig interface {
	codecGetter
	logMaker
}

// DiskBlockCacheStandard is the standard implementation for DiskBlockCache.
type DiskBlockCacheStandard struct {
	config   diskBlockCacheConfig
	maxBytes uint64
	log      logger.Logger
	// protects the disk caches from being shutdown while they're being
	// accessed
	lock     sync.RWMutex
	isClosed bool
	blockDb  *leveldb.DB
	lruDb    *leveldb.DB
}

var _ DiskBlockCache = (*DiskBlockCacheStandard)(nil)

// openLevelDB opens or recovers a leveldb.DB with a passed-in storage.Storage
// as its underlying storage layer.
func openLevelDB(stor storage.Storage) (db *leveldb.DB, err error) {
	db, err = leveldb.Open(stor, leveldbOptions)
	if _, isErrCorrupted := err.(*storage.ErrCorrupted); isErrCorrupted {
		return leveldb.Recover(stor, leveldbOptions)
	}
	return db, err
}

// newDiskBlockCacheStandardFromStorage creates a new *DiskBlockCacheStandard
// with the passed-in storage.Storage interfaces as storage layers for each
// cache.
func newDiskBlockCacheStandardFromStorage(config diskBlockCacheConfig,
	blockStorage, lruStorage storage.Storage, maxBytes uint64) (
	*DiskBlockCacheStandard, error) {
	log := config.MakeLogger("KBC")
	blockDb, err := openLevelDB(blockStorage)
	if err != nil {
		return nil, err
	}

	lruDb, err := openLevelDB(lruStorage)
	if err != nil {
		blockDb.Close()
		return nil, err
	}
	return &DiskBlockCacheStandard{
		config:   config,
		maxBytes: maxBytes,
		log:      log,
		blockDb:  blockDb,
		lruDb:    lruDb,
	}, nil
}

// newDiskBlockCacheStandard creates a new *DiskBlockCacheStandard with a
// specified directory on the filesystem as storage.
func newDiskBlockCacheStandard(config diskBlockCacheConfig, dirPath string,
	maxBytes uint64) (*DiskBlockCacheStandard, error) {
	blockDbPath := filepath.Join(dirPath, blockDbFilename)
	blockStorage, err := storage.OpenFile(blockDbPath, false)
	if err != nil {
		return nil, err
	}
	lruDbPath := filepath.Join(dirPath, lruDbFilename)
	lruStorage, err := storage.OpenFile(lruDbPath, false)
	if err != nil {
		blockStorage.Close()
		return nil, err
	}
	return newDiskBlockCacheStandardFromStorage(config, blockStorage,
		lruStorage, maxBytes)
}

// TODO: Fix getSizesLocked(), where leveldb.DB.SizeOf() currently doesn't work
// for the full range. As it is right now, this is a very cheap operation,
// since it just takes the beginning and end offsets from levelDB and returns
// the difference.
func (cache *DiskBlockCacheStandard) getSizesLocked() (blockSize int64,
	lruSize int64, err error) {
	// a util.Range{nil, nil} is supposed to work for the full range based on
	// other methods, but it appears not to work for SizeOf. Need to implement
	// by querying the files in the LevelDB directory instead. Ugh.
	blockSizes, err := cache.blockDb.SizeOf([]util.Range{{}})
	if err != nil {
		return -1, -1, err
	}
	lruSizes, err := cache.lruDb.SizeOf([]util.Range{{}})
	if err != nil {
		return -1, -1, err
	}
	return blockSizes[0], lruSizes[0], nil
}

// compactCachesLocked manually forces both the block cache and LRU cache to
// compact their underlying data.
func (cache *DiskBlockCacheStandard) compactCachesLocked(ctx context.Context) {
	cache.blockDb.CompactRange(util.Range{})
	cache.lruDb.CompactRange(util.Range{})
	bSize, lruSize, err := cache.getSizesLocked()
	cache.log.CDebugf(ctx, "Disk Cache bSize=%d lruSize=%d err=%+v", bSize, lruSize, err)
}

// lruKey generates an LRU cache key from a tlf.ID and a binary-encoded block
// ID.
func (*DiskBlockCacheStandard) lruKey(tlfID tlf.ID, blockKey []byte) []byte {
	return append(tlfID.Bytes(), blockKey...)
}

// updateLruLocked updates the LRU time of a block in the LRU cache to
// time.Now().
func (cache *DiskBlockCacheStandard) updateLruLocked(tlfID tlf.ID,
	blockKey []byte) error {
	key := cache.lruKey(tlfID, blockKey)
	val, err := time.Now().MarshalBinary()
	if err != nil {
		return err
	}
	cache.lruDb.Put(key, val, nil)
	return nil
}

// timeFromBytes converts a value from the LRU cache into a time.Time.
func (*DiskBlockCacheStandard) timeFromBytes(b []byte) (t time.Time, err error) {
	err = t.UnmarshalBinary(b)
	return t, err
}

// getLru retrieves the LRU time for a block in the cache, or returns
// leveldb.ErrNotFound and a zero-valued time.Time otherwise.
func (cache *DiskBlockCacheStandard) getLru(tlfID tlf.ID,
	blockID kbfsblock.ID) (time.Time, error) {
	lruKey := cache.lruKey(tlfID, blockID.Bytes())
	lruBytes, err := cache.lruDb.Get(lruKey, nil)
	if err != nil {
		return time.Time{}, err
	}
	return cache.timeFromBytes(lruBytes)
}

// decodeBlockCacheEntry decodes a disk block cache entry buffer into an
// encoded block and server half.
func (cache *DiskBlockCacheStandard) decodeBlockCacheEntry(buf []byte) ([]byte,
	kbfscrypto.BlockCryptKeyServerHalf, error) {
	entry := diskBlockCacheEntry{}
	err := cache.config.Codec().Decode(buf, &entry)
	if err != nil {
		return nil, kbfscrypto.BlockCryptKeyServerHalf{}, err
	}
	return entry.Buf, entry.ServerHalf, nil
}

// encodeBlockCacheEntry encodes an encoded block and serverHalf into a single
// buffer.
func (cache *DiskBlockCacheStandard) encodeBlockCacheEntry(buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf) ([]byte, error) {
	entry := diskBlockCacheEntry{
		Buf:        buf,
		ServerHalf: serverHalf,
	}
	return cache.config.Codec().Encode(&entry)
}

// Get implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Get(ctx context.Context, tlfID tlf.ID,
	blockID kbfsblock.ID) (buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf, err error) {
	cache.lock.RLock()
	defer cache.lock.RUnlock()
	if cache.isClosed {
		return nil, kbfscrypto.BlockCryptKeyServerHalf{},
			DiskCacheClosedError{"Get"}
	}
	defer func() {
		cache.log.CDebugf(ctx, "Cache Get id=%s tlf=%s bSize=%d err=%+v", blockID, tlfID, len(buf), err)
	}()
	var entry []byte
	err = runUnlessCanceled(ctx, func() error {
		blockKey := blockID.Bytes()
		buf, err := cache.blockDb.Get(blockKey, nil)
		if err != nil {
			return NoSuchBlockError{blockID}
		}
		err = cache.updateLruLocked(tlfID, blockKey)
		if err != nil {
			return err
		}
		entry = buf
		return nil
	})
	if err != nil {
		return nil, kbfscrypto.BlockCryptKeyServerHalf{}, err
	}
	return cache.decodeBlockCacheEntry(entry)
}

// Put implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Put(ctx context.Context, tlfID tlf.ID,
	blockID kbfsblock.ID, buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf) error {
	cache.lock.RLock()
	defer cache.lock.RUnlock()
	if cache.isClosed {
		return DiskCacheClosedError{"Put"}
	}
	blockLen := len(buf)
	// TODO: accounting
	entry, err := cache.encodeBlockCacheEntry(buf, serverHalf)
	encodeLen := len(entry)
	cache.log.CDebugf(ctx, "Cache Put id=%s tlf=%s bSize=%d entrySize=%d err=%+v", blockID, tlfID, blockLen, encodeLen, err)
	if err != nil {
		return err
	}
	blockKey := blockID.Bytes()
	err = cache.blockDb.Put(blockKey, entry, nil)
	if err != nil {
		return err
	}
	return cache.updateLruLocked(tlfID, blockKey)
}

// Delete implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Delete(ctx context.Context, tlfID tlf.ID,
	blockIDs []kbfsblock.ID) error {
	cache.lock.RLock()
	defer cache.lock.RUnlock()
	if cache.isClosed {
		return DiskCacheClosedError{"Delete"}
	}
	if len(blockIDs) == 0 {
		return nil
	}
	cache.log.CDebugf(ctx, "Cache Delete tlf=%s numBlocks=%d", tlfID, len(blockIDs))
	blockBatch := new(leveldb.Batch)
	lruBatch := new(leveldb.Batch)
	for _, id := range blockIDs {
		blockKey := id.Bytes()
		blockBatch.Delete(blockKey)
		lruKey := cache.lruKey(tlfID, blockKey)
		lruBatch.Delete(lruKey)
	}
	if err := cache.blockDb.Write(blockBatch, nil); err != nil {
		return err
	}
	if err := cache.lruDb.Write(lruBatch, nil); err != nil {
		return err
	}

	cache.compactCachesLocked(ctx)

	return nil
}

// evictLocked evicts a number of blocks from the cache.
func (cache *DiskBlockCacheStandard) evictLocked(ctx context.Context,
	tlfID tlf.ID, numBlocks int) error {
	// Use kbfscrypto.MakeTemporaryID() to create a random hash ID. Then begin
	// an iterator into cache.lruDb.Range(tlfID + b, tlfID + MaxBlockID) and
	// iterate from there to get numBlocks * evictionConsiderationFactor block
	// IDs.  We sort the resulting blocks by value (LRU time) and pick the
	// minimum numBlocks. We then call cache.Delete() on that list of block
	// IDs.
	return nil
}

// Shutdown implements the DiskBlockCache interface for DiskBlockCacheStandard.
func (cache *DiskBlockCacheStandard) Shutdown() {
	cache.lock.Lock()
	defer cache.lock.Unlock()
	if cache.isClosed {
		return
	}
	cache.isClosed = true
	cache.blockDb.Close()
	cache.lruDb.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"log"

	// flagged "github.com/Urban4M/go-flagged"
	flagged ".."
)

//
var setting struct {
	aString string  `usage:"A String."`
	aBool   bool    `usage:"A Bool."`
	aFloat  float64 `usage:"A Float."`
	ints    struct {
		anInt    int    `usage:"An Int."`
		anInt64  int64  `usage:"An Int64."`
		anUint   uint   `usage:"An Uint."`
		anUint64 uint64 `usage:"An Uint64."`
	}
	values struct {
		aString string  `value:"default string" usage:"A String."`
		aBool   bool    `value:"true" usage:"A Bool."`
		aFloat  float64 `value:"123.456" usage:"A Float."`
		ints    struct {
			anInt    int    `value:"-1" usage:"An Int."`
			anInt64  int64  `value:"-2" usage:"An Int64."`
			anUint   uint   `value:"1" usage:"An Uint."`
			anUint64 uint64 `value:"2" usage:"An Uint64."`
		}
	}
	envs struct {
		aString string  `env:"ASTRING" value:"default string" usage:"A String."`
		aBool   bool    `env:"ABOOL" value:"true" usage:"A Bool."`
		aFloat  float64 `env:"AFLOAT" value:"123.456" usage:"A Float."`
		ints    struct {
			anInt    int    `env:"ANINT" value:"-1" usage:"An Int."`
			anInt64  int64  `env:"ANINT64" value:"-2" usage:"An Int64."`
			anUint   uint   `env:"ANUINT" value:"1" usage:"An Uint."`
			anUint64 uint64 `env:"ANUINT64" value:"2" usage:"An Uint64."`
		}
	}
	named struct {
		aString string  `flag:"flag.string" env:"ASTRING" value:"default string" usage:"A String."`
		aBool   bool    `flag:"flag.bool" env:"ABOOL" value:"true" usage:"A Bool."`
		aFloat  float64 `flag:"flag.float" env:"AFLOAT" value:"123.456" usage:"A Float."`
		ints    struct {
			anInt    int    `flag:"flag.int" env:"ANINT" value:"-1" usage:"An Int."`
			anInt64  int64  `flag:"flag.int64" env:"ANINT64" value:"-2" usage:"An Int64."`
			anUint   uint   `flag:"flag.uint" env:"ANUINT" value:"1" usage:"An Uint."`
			anUint64 uint64 `flag:"flag.uint64" env:"ANUINT64" value:"2" usage:"An Uint64."`
		}
	}
	aliased struct {
		x int `flag:"i,_" env:"ALIASED" value:"123" usage:"An Int."`
		y int `flag:"yy" env:"ALIASED" value:"345" usage:"An Int."`
		z int `flag:"_,a" env:"ALIASED" value:"567" usage:"An Int."`
	}
	subbed struct {
		one subsub `usage:"One. "`
		two subsub
	} `usage:"A sub. "`
	_positional struct {
		first  string  `usage:"Just the first"`
		second int     `usage:"Just the second"`
		third  float64 `usage:"Just the third" value:"3.0"`
		fourth uint64  `usage:"Just the fourth" value:"4"`
		fifth  int64   `usage:"Just the fifth"`
		sixth  string  `usage:"Just the sixth" value:"sixth"`
	}
}

type subsub struct {
	aString string  `usage:"A String."`
	aBool   bool    `usage:"A Bool."`
	aFloat  float64 `usage:"A Float."`
}

//
func main() {
	flagged.FlaggedDebugging = true
	// This is just an example of how you can use parts of a struct instead of the entire struct:
	flagged.ParseWithPrefix(&setting.ints, "ok", flagged.Continue)
	flagged.ParseWithPrefix(&setting.named, "weird", flagged.Continue)
	// This is typically all you'll need to call:
	flagged.Parse(&setting)
	flagged.Usage()
	log.Printf("%+v", setting)
}
<commit_msg>Support for named and typed positional parameters. Changed the Usage() format slightly. Added flagged.Usage(). Added Program().<commit_after>package main

import (
	"log"

	flagged "github.com/Urban4M/go-flagged"
)

//
var setting struct {
	aString string  `usage:"A String."`
	aBool   bool    `usage:"A Bool."`
	aFloat  float64 `usage:"A Float."`
	ints    struct {
		anInt    int    `usage:"An Int."`
		anInt64  int64  `usage:"An Int64."`
		anUint   uint   `usage:"An Uint."`
		anUint64 uint64 `usage:"An Uint64."`
	}
	values struct {
		aString string  `value:"default string" usage:"A String."`
		aBool   bool    `value:"true" usage:"A Bool."`
		aFloat  float64 `value:"123.456" usage:"A Float."`
		ints    struct {
			anInt    int    `value:"-1" usage:"An Int."`
			anInt64  int64  `value:"-2" usage:"An Int64."`
			anUint   uint   `value:"1" usage:"An Uint."`
			anUint64 uint64 `value:"2" usage:"An Uint64."`
		}
	}
	envs struct {
		aString string  `env:"ASTRING" value:"default string" usage:"A String."`
		aBool   bool    `env:"ABOOL" value:"true" usage:"A Bool."`
		aFloat  float64 `env:"AFLOAT" value:"123.456" usage:"A Float."`
		ints    struct {
			anInt    int    `env:"ANINT" value:"-1" usage:"An Int."`
			anInt64  int64  `env:"ANINT64" value:"-2" usage:"An Int64."`
			anUint   uint   `env:"ANUINT" value:"1" usage:"An Uint."`
			anUint64 uint64 `env:"ANUINT64" value:"2" usage:"An Uint64."`
		}
	}
	named struct {
		aString string  `flag:"flag.string" env:"ASTRING" value:"default string" usage:"A String."`
		aBool   bool    `flag:"flag.bool" env:"ABOOL" value:"true" usage:"A Bool."`
		aFloat  float64 `flag:"flag.float" env:"AFLOAT" value:"123.456" usage:"A Float."`
		ints    struct {
			anInt    int    `flag:"flag.int" env:"ANINT" value:"-1" usage:"An Int."`
			anInt64  int64  `flag:"flag.int64" env:"ANINT64" value:"-2" usage:"An Int64."`
			anUint   uint   `flag:"flag.uint" env:"ANUINT" value:"1" usage:"An Uint."`
			anUint64 uint64 `flag:"flag.uint64" env:"ANUINT64" value:"2" usage:"An Uint64."`
		}
	}
	aliased struct {
		x int `flag:"i,_" env:"ALIASED" value:"123" usage:"An Int."`
		y int `flag:"yy" env:"ALIASED" value:"345" usage:"An Int."`
		z int `flag:"_,a" env:"ALIASED" value:"567" usage:"An Int."`
	}
	subbed struct {
		one subsub `usage:"One. "`
		two subsub
	} `usage:"A sub. "`
	_positional struct {
		first  string  `usage:"Just the first"`
		second int     `usage:"Just the second"`
		third  float64 `usage:"Just the third" value:"3.0"`
		fourth uint64  `usage:"Just the fourth" value:"4"`
		fifth  int64   `usage:"Just the fifth"`
		sixth  string  `usage:"Just the sixth" value:"sixth"`
	}
}

type subsub struct {
	aString string  `usage:"A String."`
	aBool   bool    `usage:"A Bool."`
	aFloat  float64 `usage:"A Float."`
}

//
func main() {
	flagged.FlaggedDebugging = true
	// This is just an example of how you can use parts of a struct instead of the entire struct:
	flagged.ParseWithPrefix(&setting.ints, "ok", flagged.Continue)
	flagged.ParseWithPrefix(&setting.named, "weird", flagged.Continue)
	// This is typically all you'll need to call:
	flagged.Parse(&setting)
	flagged.Usage()
	log.Printf("%+v", setting)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"flag"
	"libgeo"
)

func main() {
	flag.Parse()

	// Check the number of arguments
	if flag.NArg() < 2 {
		fmt.Printf("usage: main DBFILE IPADDRESS\n")
		return
	}

	// Set the arguments
	dbFile := flag.Arg(0)
	ipAddr := flag.Arg(1)

	// Load the database file, exit on failure
	gi, err := libgeo.Load(dbFile)
	if err != nil {
		fmt.Printf("Error: " + err.String() + "\n")
		return
	}

	// Lookup the IP and display the details if country is found
	loc := gi.GetLocationByIP(ipAddr)
	if loc != nil {
		fmt.Printf("Country: %s (%s)\n", loc.CountryName, loc.CountryCode)
		fmt.Printf("City: %s\n", loc.City)
		fmt.Printf("Region: %s\n", loc.Region)
		fmt.Printf("Postal Code: %s\n", loc.PostalCode)
		fmt.Printf("Latitude: %f\n", loc.Latitude)
		fmt.Printf("Longitude: %f\n", loc.Longitude)
	}
}
<commit_msg>sort printf out<commit_after>package main

import (
	"fmt"
	"flag"
	"libgeo"
)

func main() {
	flag.Parse()

	// Check the number of arguments
	if flag.NArg() < 2 {
		fmt.Printf("usage: main DBFILE IPADDRESS\n")
		return
	}

	// Set the arguments
	dbFile := flag.Arg(0)
	ipAddr := flag.Arg(1)

	// Load the database file, exit on failure
	gi, err := libgeo.Load(dbFile)
	if err != nil {
		fmt.Printf("Error: %s\n", err.String())
		return
	}

	// Lookup the IP and display the details if country is found
	loc := gi.GetLocationByIP(ipAddr)
	if loc != nil {
		fmt.Printf("Country: %s (%s)\n", loc.CountryName, loc.CountryCode)
		fmt.Printf("City: %s\n", loc.City)
		fmt.Printf("Region: %s\n", loc.Region)
		fmt.Printf("Postal Code: %s\n", loc.PostalCode)
		fmt.Printf("Latitude: %f\n", loc.Latitude)
		fmt.Printf("Longitude: %f\n", loc.Longitude)
	}
}
<|endoftext|>


<commit_before>package harvest

import (
	"fmt"
	"time"
)

type ExpensesResponse struct {
	Expenses     []*Expense `json:"expenses"`
	PerPage      int64      `json:"per_page"`
	TotalPages   int64      `json:"total_pages"`
	TotalEntries int64      `json:"total_entries"`
	NextPage     *int64     `json:"next_page"`
	PreviousPage *int64     `json:"previous_page"`
	Page         int64      `json:"page"`
}

type Expense struct {
	ID                int64     `json:"id"`
	Notes             string    `json:"notes"`
	TotalCost         float64   `json:"total_cost"`
	Units             float64   `json:"units"`
	CreatedAt         time.Time `json:"created_at"`
	UpdatedAt         time.Time `json:"updated_at"`
	ProjectID         int64     `json:"project_id"`
	ExpenseCategoryID int64     `json:"expense_category_id"`
	UserID            int64     `json:"user_id"`
	SpentAt           string    `json:"spent_at"`
	IsClosed          bool      `json:"is_closed"`
	InvoiceID         int64     `json:"invoice_id"`
	Billable          bool      `json:"billable"`
	CompanyID         int64     `json:"company_id"`
	HasReceipt        bool      `json:"has_receipt"`
	ReceiptURL        string    `json:"receipt_url"`
	IsLocked          bool      `json:"is_locked"`
	LockedReason      string    `json:"locked_reason"`
}

func (a *API) GetExpense(expenseID int64, args Arguments) (expense *Expense, err error) {
	expense = &Expense{}
	path := fmt.Sprintf("/expenses/%v", expenseID)
	err = a.Get(path, args, expense)
	return expense, err
}

func (a *API) GetExpenses(args Arguments) (expenses []*Expense, err error) {
	expensesResponse := ExpensesResponse{}
	path := fmt.Sprintf("/expenses")
	err = a.Get(path, args, &expensesResponse)
	return expensesResponse.Expenses, err
}
<commit_msg>Adds pagination support to Expenses.<commit_after>package harvest

import (
	"fmt"
	"time"
)

type ExpensesResponse struct {
	PagedResponse
	Expenses []*Expense `json:"expenses"`
}

type Expense struct {
	ID                int64     `json:"id"`
	Notes             string    `json:"notes"`
	TotalCost         float64   `json:"total_cost"`
	Units             float64   `json:"units"`
	CreatedAt         time.Time `json:"created_at"`
	UpdatedAt         time.Time `json:"updated_at"`
	ProjectID         int64     `json:"project_id"`
	ExpenseCategoryID int64     `json:"expense_category_id"`
	UserID            int64     `json:"user_id"`
	SpentAt           string    `json:"spent_at"`
	IsClosed          bool      `json:"is_closed"`
	InvoiceID         int64     `json:"invoice_id"`
	Billable          bool      `json:"billable"`
	CompanyID         int64     `json:"company_id"`
	HasReceipt        bool      `json:"has_receipt"`
	ReceiptURL        string    `json:"receipt_url"`
	IsLocked          bool      `json:"is_locked"`
	LockedReason      string    `json:"locked_reason"`
}

func (a *API) GetExpense(expenseID int64, args Arguments) (expense *Expense, err error) {
	expense = &Expense{}
	path := fmt.Sprintf("/expenses/%v", expenseID)
	err = a.Get(path, args, expense)
	return expense, err
}

func (a *API) GetExpenses(args Arguments) (expenses []*Expense, err error) {
	expenses = make([]*Expense, 0)
	expensesResponse := ExpensesResponse{}
	err = a.GetPaginated("/expenses", args, &expensesResponse, func() {
		for _, e := range expensesResponse.Expenses {
			expenses = append(expenses, e)
		}
		expensesResponse.Expenses = make([]*Expense, 0)
	})
	return expenses, err
}
<|endoftext|>


<commit_before>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charmstore

import (
	"io"
	"io/ioutil"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"gopkg.in/juju/charm.v6-unstable"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/juju/charmrepo.v2-unstable"
	"gopkg.in/juju/charmrepo.v2-unstable/csclient"
)

var logger = loggo.GetLogger("juju.charmstore")

// Client exposes the functionality of the charm store, as provided
// by github.com/juju/charmrepo/csclient.Client.
//
// Note that the following csclient.Client methods are used as well,
// but only in tests:
//  - Put(path string, val interface{}) error
//  - UploadCharm(id *charm.URL, ch charm.Charm) (*charm.URL, error)
//  - UploadCharmWithRevision(id *charm.URL, ch charm.Charm, promulgatedRevision int) error
//  - UploadBundleWithRevision()
type Client interface {
	BaseClient
	// TODO(ericsnow) Move this over to BaseClient once it supports it.
	ResourcesClient
	io.Closer

	// AsRepo returns a charm repo that wraps the client. If the model's
	// UUID is provided then it is associated with all of the repo's
	// interaction with the charm store.
	AsRepo(modelUUID string) Repo
}

// BaseClient exposes the functionality Juju needs which csclient.Client
// provides.
type BaseClient interface {
	// TODO(ericsnow) Replace use of Get with use of more specific API
	// methods? We only use Get() for authorization on the Juju client
	// side.

	// Get makes a GET request to the given path in the charm store. The
	// path must have a leading slash, but must not include the host
	// name or version prefix. The result is parsed as JSON into the
	// given result value, which should be a pointer to the expected
	// data, but may be nil if no result is desired.
	Get(path string, result interface{}) error
}

// ResourcesClient exposes the charm store client functionality for
// dealing with resources.
type ResourcesClient interface {
	// TODO(ericsnow) Just embed resource/charmstore.BaseClient (or vice-versa)?

	// ListResources composes, for each of the identified charms, the
	// list of details for each of the charm's resources. Those details
	// are those associated with the specific charm revision. They
	// include the resource's metadata and revision.
	ListResources(charmURLs []*charm.URL) ([][]charmresource.Resource, error)

	// GetResource returns a reader for the resource's data. That data
	// is streamed from the charm store. The charm's revision, if any,
	// is ignored. If the identified resource is not in the charm store
	// then errors.NotFound is returned.
	GetResource(cURL *charm.URL, resourceName string, revision int) (io.ReadCloser, error)
}

// client adapts csclient.Client to the needs of Juju.
type client struct {
	*csclient.Client
	fakeCharmStoreClient
	io.Closer
}

// NewClient returns a Juju charm store client for the given client
// config.
func NewClient(config *csclient.Params) Client {
	if config == nil {
		config = &csclient.Params{}
	}
	base := csclient.New(*config)
	return WrapBaseClient(base, nil)
}

// WrapBaseClient returns a Juju charm store client that wraps
// the provided client.
func WrapBaseClient(base *csclient.Client, closer io.Closer) Client {
	if closer == nil {
		closer = ioutil.NopCloser(nil)
	}
	return &client{
		Client: base,
		Closer: closer,
	}
}

// AsRepo implements Client.
func (client client) AsRepo(envUUID string) Repo {
	repo := charmrepo.NewCharmStoreFromClient(client.Client)
	if envUUID != "" {
		repo = repo.WithJujuAttrs(map[string]string{
			"environment_uuid": envUUID,
		})
	}
	return repo
}

// LatestCharmInfo returns the most up-to-date information about each
// of the identified charms at their latest revision. The revisions in
// the provided URLs are ignored.
func LatestCharmInfo(client Client, modelUUID string, cURLs []*charm.URL) ([]CharmInfoResult, error) {
	now := time.Now().UTC()

	// We must use charmrepo.CharmStore since csclient.Client does not
	// have the "Latest" method.
	repo := client.AsRepo(modelUUID)

	// Do a bulk call to get the revision info for all charms.
	logger.Infof("retrieving revision information for %d charms", len(cURLs))
	revResults, err := repo.Latest(cURLs...)
	if err != nil {
		err = errors.Annotate(err, "finding charm revision info")
		logger.Infof(err.Error())
		return nil, errors.Trace(err)
	}

	// Do a bulk call to get the latest info for each charm's resources.
	// TODO(ericsnow) Only do this for charms that *have* resources.
	chResources, err := client.ListResources(cURLs)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Extract the results.
	var results []CharmInfoResult
	for i, cURL := range cURLs {
		revResult := revResults[i]
		resources := chResources[i]

		var result CharmInfoResult
		result.URL = cURL
		result.Timestamp = now
		if revResult.Err != nil {
			result.Error = errors.Trace(revResult.Err)
		} else {
			result.LatestRevision = revResult.Revision
			result.LatestResources = resources
		}
		results = append(results, result)
	}
	return results, nil
}

// TODO(ericsnow) Remove the fake once the charm store adds support.

type fakeCharmStoreClient struct{}

// ListResources implements Client as a noop.
func (fakeCharmStoreClient) ListResources(charmURLs []*charm.URL) ([][]charmresource.Resource, error) {
	res := make([][]charmresource.Resource, len(charmURLs))
	return res, nil
}

// GetResource implements Client as a noop.
func (fakeCharmStoreClient) GetResource(cURL *charm.URL, resourceName string, revision int) (io.ReadCloser, error) {
	return nil, errors.NotFoundf("resource %q", resourceName)
}
<commit_msg>Clean up doc comments.<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charmstore

import (
	"io"
	"io/ioutil"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"gopkg.in/juju/charm.v6-unstable"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/juju/charmrepo.v2-unstable"
	"gopkg.in/juju/charmrepo.v2-unstable/csclient"
)

var logger = loggo.GetLogger("juju.charmstore")

// Client provides the common interface for the charm store client needs
// of Juju.
type Client interface {
	BaseClient
	ResourcesClient
	io.Closer

	// AsRepo returns a charm repo that wraps the client. If the model's
	// UUID is provided then it is associated with all of the repo's
	// interaction with the charm store.
	AsRepo(modelUUID string) Repo
}

// BaseClient exposes the functionality of the charm store, as provided
// by github.com/juju/charmrepo/csclient.Client.
//
// Note that the following csclient.Client methods are used as well,
// but only in tests:
//  - Put(path string, val interface{}) error
//  - UploadCharm(id *charm.URL, ch charm.Charm) (*charm.URL, error)
//  - UploadCharmWithRevision(id *charm.URL, ch charm.Charm, promulgatedRevision int) error
//  - UploadBundleWithRevision()
type BaseClient interface {
	// TODO(ericsnow) Embed ResourcesClient here once the charm store
	// supports it.

	// TODO(ericsnow) Replace use of Get with use of more specific API
	// methods? We only use Get() for authorization on the Juju client
	// side.

	// Get makes a GET request to the given path in the charm store. The
	// path must have a leading slash, but must not include the host
	// name or version prefix. The result is parsed as JSON into the
	// given result value, which should be a pointer to the expected
	// data, but may be nil if no result is desired.
	Get(path string, result interface{}) error
}

// ResourcesClient exposes the charm store client functionality for
// dealing with resources.
type ResourcesClient interface {
	// TODO(ericsnow) Just embed resource/charmstore.BaseClient (or vice-versa)?

	// ListResources composes, for each of the identified charms, the
	// list of details for each of the charm's resources. Those details
	// are those associated with the specific charm revision. They
	// include the resource's metadata and revision.
	ListResources(charmURLs []*charm.URL) ([][]charmresource.Resource, error)

	// GetResource returns a reader for the resource's data. That data
	// is streamed from the charm store. The charm's revision, if any,
	// is ignored. If the identified resource is not in the charm store
	// then errors.NotFound is returned.
	GetResource(cURL *charm.URL, resourceName string, revision int) (io.ReadCloser, error)
}

// client adapts csclient.Client to the needs of Juju.
type client struct {
	*csclient.Client
	fakeCharmStoreClient
	io.Closer
}

// NewClient returns a Juju charm store client for the given client
// config.
func NewClient(config *csclient.Params) Client {
	if config == nil {
		config = &csclient.Params{}
	}
	base := csclient.New(*config)
	return WrapBaseClient(base, nil)
}

// WrapBaseClient returns a Juju charm store client that wraps
// the provided client.
func WrapBaseClient(base *csclient.Client, closer io.Closer) Client {
	if closer == nil {
		closer = ioutil.NopCloser(nil)
	}
	return &client{
		Client: base,
		Closer: closer,
	}
}

// AsRepo implements Client.
func (client client) AsRepo(envUUID string) Repo {
	repo := charmrepo.NewCharmStoreFromClient(client.Client)
	if envUUID != "" {
		repo = repo.WithJujuAttrs(map[string]string{
			"environment_uuid": envUUID,
		})
	}
	return repo
}

// LatestCharmInfo returns the most up-to-date information about each
// of the identified charms at their latest revision. The revisions in
// the provided URLs are ignored.
func LatestCharmInfo(client Client, modelUUID string, cURLs []*charm.URL) ([]CharmInfoResult, error) {
	now := time.Now().UTC()

	// We must use charmrepo.CharmStore since csclient.Client does not
	// have the "Latest" method.
	repo := client.AsRepo(modelUUID)

	// Do a bulk call to get the revision info for all charms.
	logger.Infof("retrieving revision information for %d charms", len(cURLs))
	revResults, err := repo.Latest(cURLs...)
	if err != nil {
		err = errors.Annotate(err, "finding charm revision info")
		logger.Infof(err.Error())
		return nil, errors.Trace(err)
	}

	// Do a bulk call to get the latest info for each charm's resources.
	// TODO(ericsnow) Only do this for charms that *have* resources.
	chResources, err := client.ListResources(cURLs)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Extract the results.
	var results []CharmInfoResult
	for i, cURL := range cURLs {
		revResult := revResults[i]
		resources := chResources[i]

		var result CharmInfoResult
		result.URL = cURL
		result.Timestamp = now
		if revResult.Err != nil {
			result.Error = errors.Trace(revResult.Err)
		} else {
			result.LatestRevision = revResult.Revision
			result.LatestResources = resources
		}
		results = append(results, result)
	}
	return results, nil
}

// TODO(ericsnow) Remove the fake once the charm store adds support.

type fakeCharmStoreClient struct{}

// ListResources implements Client as a noop.
func (fakeCharmStoreClient) ListResources(charmURLs []*charm.URL) ([][]charmresource.Resource, error) {
	res := make([][]charmresource.Resource, len(charmURLs))
	return res, nil
}

// GetResource implements Client as a noop.
func (fakeCharmStoreClient) GetResource(cURL *charm.URL, resourceName string, revision int) (io.ReadCloser, error) {
	return nil, errors.NotFoundf("resource %q", resourceName)
}
<|endoftext|>


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

package models

import (
	"fmt"

	"code.gitea.io/gitea/modules/setting"
	api "code.gitea.io/gitea/modules/structs"
	"code.gitea.io/gitea/modules/timeutil"
	"xorm.io/builder"

	"xorm.io/xorm"
)

// Milestone represents a milestone of repository.
type Milestone struct {
	ID              int64 `xorm:"pk autoincr"`
	RepoID          int64 `xorm:"INDEX"`
	Name            string
	Content         string `xorm:"TEXT"`
	RenderedContent string `xorm:"-"`
	IsClosed        bool
	NumIssues       int
	NumClosedIssues int
	NumOpenIssues   int  `xorm:"-"`
	Completeness    int  // Percentage(1-100).
	IsOverdue       bool `xorm:"-"`

	DeadlineString string `xorm:"-"`
	DeadlineUnix   timeutil.TimeStamp
	ClosedDateUnix timeutil.TimeStamp

	TotalTrackedTime int64 `xorm:"-"`
}

// BeforeUpdate is invoked from XORM before updating this object.
func (m *Milestone) BeforeUpdate() {
	if m.NumIssues > 0 {
		m.Completeness = m.NumClosedIssues * 100 / m.NumIssues
	} else {
		m.Completeness = 0
	}
}

// AfterLoad is invoked from XORM after setting the value of a field of
// this object.
func (m *Milestone) AfterLoad() {
	m.NumOpenIssues = m.NumIssues - m.NumClosedIssues
	if m.DeadlineUnix.Year() == 9999 {
		return
	}

	m.DeadlineString = m.DeadlineUnix.Format("2006-01-02")
	if timeutil.TimeStampNow() >= m.DeadlineUnix {
		m.IsOverdue = true
	}
}

// State returns string representation of milestone status.
func (m *Milestone) State() api.StateType {
	if m.IsClosed {
		return api.StateClosed
	}
	return api.StateOpen
}

// APIFormat returns this Milestone in API format.
func (m *Milestone) APIFormat() *api.Milestone {
	apiMilestone := &api.Milestone{
		ID:           m.ID,
		State:        m.State(),
		Title:        m.Name,
		Description:  m.Content,
		OpenIssues:   m.NumOpenIssues,
		ClosedIssues: m.NumClosedIssues,
	}
	if m.IsClosed {
		apiMilestone.Closed = m.ClosedDateUnix.AsTimePtr()
	}
	if m.DeadlineUnix.Year() < 9999 {
		apiMilestone.Deadline = m.DeadlineUnix.AsTimePtr()
	}
	return apiMilestone
}

// NewMilestone creates new milestone of repository.
func NewMilestone(m *Milestone) (err error) {
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	if _, err = sess.Insert(m); err != nil {
		return err
	}

	if _, err = sess.Exec("UPDATE `repository` SET num_milestones = num_milestones + 1 WHERE id = ?", m.RepoID); err != nil {
		return err
	}
	return sess.Commit()
}

func getMilestoneByRepoID(e Engine, repoID, id int64) (*Milestone, error) {
	m := &Milestone{
		ID:     id,
		RepoID: repoID,
	}
	has, err := e.Get(m)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrMilestoneNotExist{id, repoID}
	}
	return m, nil
}

// GetMilestoneByRepoID returns the milestone in a repository.
func GetMilestoneByRepoID(repoID, id int64) (*Milestone, error) {
	return getMilestoneByRepoID(x, repoID, id)
}

// GetMilestoneByID returns the milestone via id .
func GetMilestoneByID(id int64) (*Milestone, error) {
	var m Milestone
	has, err := x.ID(id).Get(&m)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrMilestoneNotExist{id, 0}
	}
	return &m, nil
}

// MilestoneList is a list of milestones offering additional functionality
type MilestoneList []*Milestone

func (milestones MilestoneList) loadTotalTrackedTimes(e Engine) error {
	type totalTimesByMilestone struct {
		MilestoneID int64
		Time        int64
	}
	if len(milestones) == 0 {
		return nil
	}
	var trackedTimes = make(map[int64]int64, len(milestones))

	// Get total tracked time by milestone_id
	rows, err := e.Table("issue").
		Join("INNER", "milestone", "issue.milestone_id = milestone.id").
		Join("LEFT", "tracked_time", "tracked_time.issue_id = issue.id").
		Select("milestone_id, sum(time) as time").
		In("milestone_id", milestones.getMilestoneIDs()).
		GroupBy("milestone_id").
		Rows(new(totalTimesByMilestone))
	if err != nil {
		return err
	}

	defer rows.Close()

	for rows.Next() {
		var totalTime totalTimesByMilestone
		err = rows.Scan(&totalTime)
		if err != nil {
			return err
		}
		trackedTimes[totalTime.MilestoneID] = totalTime.Time
	}

	for _, milestone := range milestones {
		milestone.TotalTrackedTime = trackedTimes[milestone.ID]
	}
	return nil
}

// LoadTotalTrackedTimes loads for every milestone in the list the TotalTrackedTime by a batch request
func (milestones MilestoneList) LoadTotalTrackedTimes() error {
	return milestones.loadTotalTrackedTimes(x)
}

func (milestones MilestoneList) getMilestoneIDs() []int64 {
	var ids = make([]int64, 0, len(milestones))
	for _, ms := range milestones {
		ids = append(ids, ms.ID)
	}
	return ids
}

// GetMilestonesByRepoID returns all opened milestones of a repository.
func GetMilestonesByRepoID(repoID int64, state api.StateType) (MilestoneList, error) {
	sess := x.Where("repo_id = ?", repoID)

	switch state {
	case api.StateClosed:
		sess = sess.And("is_closed = ?", true)

	case api.StateAll:
		break

	case api.StateOpen:
		fallthrough

	default:
		sess = sess.And("is_closed = ?", false)
	}

	miles := make([]*Milestone, 0, 10)
	return miles, sess.Asc("deadline_unix").Asc("id").Find(&miles)
}

// GetMilestones returns a list of milestones of given repository and status.
func GetMilestones(repoID int64, page int, isClosed bool, sortType string) (MilestoneList, error) {
	miles := make([]*Milestone, 0, setting.UI.IssuePagingNum)
	sess := x.Where("repo_id = ? AND is_closed = ?", repoID, isClosed)
	if page > 0 {
		sess = sess.Limit(setting.UI.IssuePagingNum, (page-1)*setting.UI.IssuePagingNum)
	}

	switch sortType {
	case "furthestduedate":
		sess.Desc("deadline_unix")
	case "leastcomplete":
		sess.Asc("completeness")
	case "mostcomplete":
		sess.Desc("completeness")
	case "leastissues":
		sess.Asc("num_issues")
	case "mostissues":
		sess.Desc("num_issues")
	default:
		sess.Asc("deadline_unix")
	}
	return miles, sess.Find(&miles)
}

func updateMilestone(e Engine, m *Milestone) error {
	_, err := e.ID(m.ID).AllCols().
		SetExpr("num_issues", builder.Select("count(*)").From("issue").Where(
			builder.Eq{"milestone_id": m.ID},
		)).
		SetExpr("num_closed_issues", builder.Select("count(*)").From("issue").Where(
			builder.Eq{
				"milestone_id": m.ID,
				"is_closed":    true,
			},
		)).
		Update(m)
	return err
}

// UpdateMilestone updates information of given milestone.
func UpdateMilestone(m *Milestone) error {
	if err := updateMilestone(x, m); err != nil {
		return err
	}

	return updateMilestoneCompleteness(x, m.ID)
}

func updateMilestoneCompleteness(e Engine, milestoneID int64) error {
	_, err := e.Exec("UPDATE `milestone` SET completeness=100*num_closed_issues/(CASE WHEN num_issues > 0 THEN num_issues ELSE 1 END) WHERE id=?",
		milestoneID,
	)
	return err
}

func countRepoMilestones(e Engine, repoID int64) (int64, error) {
	return e.
		Where("repo_id=?", repoID).
		Count(new(Milestone))
}

func countRepoClosedMilestones(e Engine, repoID int64) (int64, error) {
	return e.
		Where("repo_id=? AND is_closed=?", repoID, true).
		Count(new(Milestone))
}

// CountRepoClosedMilestones returns number of closed milestones in given repository.
func CountRepoClosedMilestones(repoID int64) (int64, error) {
	return countRepoClosedMilestones(x, repoID)
}

// MilestoneStats returns number of open and closed milestones of given repository.
func MilestoneStats(repoID int64) (open int64, closed int64, err error) {
	open, err = x.
		Where("repo_id=? AND is_closed=?", repoID, false).
		Count(new(Milestone))
	if err != nil {
		return 0, 0, nil
	}
	closed, err = CountRepoClosedMilestones(repoID)
	return open, closed, err
}

// ChangeMilestoneStatus changes the milestone open/closed status.
func ChangeMilestoneStatus(m *Milestone, isClosed bool) (err error) {
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	m.IsClosed = isClosed
	if _, err := sess.ID(m.ID).Cols("is_closed").Update(m); err != nil {
		return err
	}

	if err := updateRepoMilestoneNum(sess, m.RepoID); err != nil {
		return err
	}

	return sess.Commit()
}

func updateRepoMilestoneNum(e Engine, repoID int64) error {
	_, err := e.Exec("UPDATE `repository` SET num_milestones=(SELECT count(*) FROM milestone WHERE repo_id=?),num_closed_milestones=(SELECT count(*) FROM milestone WHERE repo_id=? AND is_closed=?) WHERE id=?",
		repoID,
		repoID,
		true,
		repoID,
	)
	return err
}

func updateMilestoneTotalNum(e Engine, milestoneID int64) (err error) {
	if _, err = e.Exec("UPDATE `milestone` SET num_issues=(SELECT count(*) FROM issue WHERE milestone_id=?) WHERE id=?",
		milestoneID,
		milestoneID,
	); err != nil {
		return
	}

	return updateMilestoneCompleteness(e, milestoneID)
}

func updateMilestoneClosedNum(e Engine, milestoneID int64) (err error) {
	if _, err = e.Exec("UPDATE `milestone` SET num_closed_issues=(SELECT count(*) FROM issue WHERE milestone_id=? AND is_closed=?) WHERE id=?",
		milestoneID,
		true,
		milestoneID,
	); err != nil {
		return
	}

	return updateMilestoneCompleteness(e, milestoneID)
}

func changeMilestoneAssign(e *xorm.Session, doer *User, issue *Issue, oldMilestoneID int64) error {
	if err := updateIssueCols(e, issue, "milestone_id"); err != nil {
		return err
	}

	if oldMilestoneID > 0 {
		if err := updateMilestoneTotalNum(e, oldMilestoneID); err != nil {
			return err
		}
		if issue.IsClosed {
			if err := updateMilestoneClosedNum(e, oldMilestoneID); err != nil {
				return err
			}
		}
	}

	if issue.MilestoneID > 0 {
		if err := updateMilestoneTotalNum(e, issue.MilestoneID); err != nil {
			return err
		}
		if issue.IsClosed {
			if err := updateMilestoneClosedNum(e, issue.MilestoneID); err != nil {
				return err
			}
		}
	}

	if oldMilestoneID > 0 || issue.MilestoneID > 0 {
		if err := issue.loadRepo(e); err != nil {
			return err
		}

		if _, err := createMilestoneComment(e, doer, issue.Repo, issue, oldMilestoneID, issue.MilestoneID); err != nil {
			return err
		}
	}

	return nil
}

// ChangeMilestoneAssign changes assignment of milestone for issue.
func ChangeMilestoneAssign(issue *Issue, doer *User, oldMilestoneID int64) (err error) {
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	if err = changeMilestoneAssign(sess, doer, issue, oldMilestoneID); err != nil {
		return err
	}

	if err = sess.Commit(); err != nil {
		return fmt.Errorf("Commit: %v", err)
	}
	return nil
}

// DeleteMilestoneByRepoID deletes a milestone from a repository.
func DeleteMilestoneByRepoID(repoID, id int64) error {
	m, err := GetMilestoneByRepoID(repoID, id)
	if err != nil {
		if IsErrMilestoneNotExist(err) {
			return nil
		}
		return err
	}

	repo, err := GetRepositoryByID(m.RepoID)
	if err != nil {
		return err
	}

	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	if _, err = sess.ID(m.ID).Delete(new(Milestone)); err != nil {
		return err
	}

	numMilestones, err := countRepoMilestones(sess, repo.ID)
	if err != nil {
		return err
	}
	numClosedMilestones, err := countRepoClosedMilestones(sess, repo.ID)
	if err != nil {
		return err
	}
	repo.NumMilestones = int(numMilestones)
	repo.NumClosedMilestones = int(numClosedMilestones)

	if _, err = sess.ID(repo.ID).Cols("num_milestones, num_closed_milestones").Update(repo); err != nil {
		return err
	}

	if _, err = sess.Exec("UPDATE `issue` SET milestone_id = 0 WHERE milestone_id = ?", m.ID); err != nil {
		return err
	}
	return sess.Commit()
}
<commit_msg>[Fix] milestone close timestamp (#8728)<commit_after>// Copyright 2017 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package models

import (
	"fmt"

	"code.gitea.io/gitea/modules/setting"
	api "code.gitea.io/gitea/modules/structs"
	"code.gitea.io/gitea/modules/timeutil"
	"xorm.io/builder"

	"xorm.io/xorm"
)

// Milestone represents a milestone of repository.
type Milestone struct {
	ID              int64 `xorm:"pk autoincr"`
	RepoID          int64 `xorm:"INDEX"`
	Name            string
	Content         string `xorm:"TEXT"`
	RenderedContent string `xorm:"-"`
	IsClosed        bool
	NumIssues       int
	NumClosedIssues int
	NumOpenIssues   int  `xorm:"-"`
	Completeness    int  // Percentage(1-100).
	IsOverdue       bool `xorm:"-"`

	DeadlineString string `xorm:"-"`
	DeadlineUnix   timeutil.TimeStamp
	ClosedDateUnix timeutil.TimeStamp

	TotalTrackedTime int64 `xorm:"-"`
}

// BeforeUpdate is invoked from XORM before updating this object.
func (m *Milestone) BeforeUpdate() {
	if m.NumIssues > 0 {
		m.Completeness = m.NumClosedIssues * 100 / m.NumIssues
	} else {
		m.Completeness = 0
	}
}

// AfterLoad is invoked from XORM after setting the value of a field of
// this object.
func (m *Milestone) AfterLoad() {
	m.NumOpenIssues = m.NumIssues - m.NumClosedIssues
	if m.DeadlineUnix.Year() == 9999 {
		return
	}

	m.DeadlineString = m.DeadlineUnix.Format("2006-01-02")
	if timeutil.TimeStampNow() >= m.DeadlineUnix {
		m.IsOverdue = true
	}
}

// State returns string representation of milestone status.
func (m *Milestone) State() api.StateType {
	if m.IsClosed {
		return api.StateClosed
	}
	return api.StateOpen
}

// APIFormat returns this Milestone in API format.
func (m *Milestone) APIFormat() *api.Milestone {
	apiMilestone := &api.Milestone{
		ID:           m.ID,
		State:        m.State(),
		Title:        m.Name,
		Description:  m.Content,
		OpenIssues:   m.NumOpenIssues,
		ClosedIssues: m.NumClosedIssues,
	}
	if m.IsClosed {
		apiMilestone.Closed = m.ClosedDateUnix.AsTimePtr()
	}
	if m.DeadlineUnix.Year() < 9999 {
		apiMilestone.Deadline = m.DeadlineUnix.AsTimePtr()
	}
	return apiMilestone
}

// NewMilestone creates new milestone of repository.
func NewMilestone(m *Milestone) (err error) {
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	if _, err = sess.Insert(m); err != nil {
		return err
	}

	if _, err = sess.Exec("UPDATE `repository` SET num_milestones = num_milestones + 1 WHERE id = ?", m.RepoID); err != nil {
		return err
	}
	return sess.Commit()
}

func getMilestoneByRepoID(e Engine, repoID, id int64) (*Milestone, error) {
	m := &Milestone{
		ID:     id,
		RepoID: repoID,
	}
	has, err := e.Get(m)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrMilestoneNotExist{id, repoID}
	}
	return m, nil
}

// GetMilestoneByRepoID returns the milestone in a repository.
func GetMilestoneByRepoID(repoID, id int64) (*Milestone, error) {
	return getMilestoneByRepoID(x, repoID, id)
}

// GetMilestoneByID returns the milestone via id .
func GetMilestoneByID(id int64) (*Milestone, error) {
	var m Milestone
	has, err := x.ID(id).Get(&m)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrMilestoneNotExist{id, 0}
	}
	return &m, nil
}

// MilestoneList is a list of milestones offering additional functionality
type MilestoneList []*Milestone

func (milestones MilestoneList) loadTotalTrackedTimes(e Engine) error {
	type totalTimesByMilestone struct {
		MilestoneID int64
		Time        int64
	}
	if len(milestones) == 0 {
		return nil
	}
	var trackedTimes = make(map[int64]int64, len(milestones))

	// Get total tracked time by milestone_id
	rows, err := e.Table("issue").
		Join("INNER", "milestone", "issue.milestone_id = milestone.id").
		Join("LEFT", "tracked_time", "tracked_time.issue_id = issue.id").
		Select("milestone_id, sum(time) as time").
		In("milestone_id", milestones.getMilestoneIDs()).
		GroupBy("milestone_id").
		Rows(new(totalTimesByMilestone))
	if err != nil {
		return err
	}

	defer rows.Close()

	for rows.Next() {
		var totalTime totalTimesByMilestone
		err = rows.Scan(&totalTime)
		if err != nil {
			return err
		}
		trackedTimes[totalTime.MilestoneID] = totalTime.Time
	}

	for _, milestone := range milestones {
		milestone.TotalTrackedTime = trackedTimes[milestone.ID]
	}
	return nil
}

// LoadTotalTrackedTimes loads for every milestone in the list the TotalTrackedTime by a batch request
func (milestones MilestoneList) LoadTotalTrackedTimes() error {
	return milestones.loadTotalTrackedTimes(x)
}

func (milestones MilestoneList) getMilestoneIDs() []int64 {
	var ids = make([]int64, 0, len(milestones))
	for _, ms := range milestones {
		ids = append(ids, ms.ID)
	}
	return ids
}

// GetMilestonesByRepoID returns all opened milestones of a repository.
func GetMilestonesByRepoID(repoID int64, state api.StateType) (MilestoneList, error) {
	sess := x.Where("repo_id = ?", repoID)

	switch state {
	case api.StateClosed:
		sess = sess.And("is_closed = ?", true)

	case api.StateAll:
		break

	case api.StateOpen:
		fallthrough

	default:
		sess = sess.And("is_closed = ?", false)
	}

	miles := make([]*Milestone, 0, 10)
	return miles, sess.Asc("deadline_unix").Asc("id").Find(&miles)
}

// GetMilestones returns a list of milestones of given repository and status.
func GetMilestones(repoID int64, page int, isClosed bool, sortType string) (MilestoneList, error) {
	miles := make([]*Milestone, 0, setting.UI.IssuePagingNum)
	sess := x.Where("repo_id = ? AND is_closed = ?", repoID, isClosed)
	if page > 0 {
		sess = sess.Limit(setting.UI.IssuePagingNum, (page-1)*setting.UI.IssuePagingNum)
	}

	switch sortType {
	case "furthestduedate":
		sess.Desc("deadline_unix")
	case "leastcomplete":
		sess.Asc("completeness")
	case "mostcomplete":
		sess.Desc("completeness")
	case "leastissues":
		sess.Asc("num_issues")
	case "mostissues":
		sess.Desc("num_issues")
	default:
		sess.Asc("deadline_unix")
	}
	return miles, sess.Find(&miles)
}

func updateMilestone(e Engine, m *Milestone) error {
	_, err := e.ID(m.ID).AllCols().
		SetExpr("num_issues", builder.Select("count(*)").From("issue").Where(
			builder.Eq{"milestone_id": m.ID},
		)).
		SetExpr("num_closed_issues", builder.Select("count(*)").From("issue").Where(
			builder.Eq{
				"milestone_id": m.ID,
				"is_closed":    true,
			},
		)).
		Update(m)
	return err
}

// UpdateMilestone updates information of given milestone.
func UpdateMilestone(m *Milestone) error {
	if err := updateMilestone(x, m); err != nil {
		return err
	}

	return updateMilestoneCompleteness(x, m.ID)
}

func updateMilestoneCompleteness(e Engine, milestoneID int64) error {
	_, err := e.Exec("UPDATE `milestone` SET completeness=100*num_closed_issues/(CASE WHEN num_issues > 0 THEN num_issues ELSE 1 END) WHERE id=?",
		milestoneID,
	)
	return err
}

func countRepoMilestones(e Engine, repoID int64) (int64, error) {
	return e.
		Where("repo_id=?", repoID).
		Count(new(Milestone))
}

func countRepoClosedMilestones(e Engine, repoID int64) (int64, error) {
	return e.
		Where("repo_id=? AND is_closed=?", repoID, true).
		Count(new(Milestone))
}

// CountRepoClosedMilestones returns number of closed milestones in given repository.
func CountRepoClosedMilestones(repoID int64) (int64, error) {
	return countRepoClosedMilestones(x, repoID)
}

// MilestoneStats returns number of open and closed milestones of given repository.
func MilestoneStats(repoID int64) (open int64, closed int64, err error) {
	open, err = x.
		Where("repo_id=? AND is_closed=?", repoID, false).
		Count(new(Milestone))
	if err != nil {
		return 0, 0, nil
	}
	closed, err = CountRepoClosedMilestones(repoID)
	return open, closed, err
}

// ChangeMilestoneStatus changes the milestone open/closed status.
func ChangeMilestoneStatus(m *Milestone, isClosed bool) (err error) {
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	m.IsClosed = isClosed
	if isClosed {
		m.ClosedDateUnix = timeutil.TimeStampNow()
	}

	if _, err := sess.ID(m.ID).Cols("is_closed", "closed_date_unix").Update(m); err != nil {
		return err
	}

	if err := updateRepoMilestoneNum(sess, m.RepoID); err != nil {
		return err
	}

	return sess.Commit()
}

func updateRepoMilestoneNum(e Engine, repoID int64) error {
	_, err := e.Exec("UPDATE `repository` SET num_milestones=(SELECT count(*) FROM milestone WHERE repo_id=?),num_closed_milestones=(SELECT count(*) FROM milestone WHERE repo_id=? AND is_closed=?) WHERE id=?",
		repoID,
		repoID,
		true,
		repoID,
	)
	return err
}

func updateMilestoneTotalNum(e Engine, milestoneID int64) (err error) {
	if _, err = e.Exec("UPDATE `milestone` SET num_issues=(SELECT count(*) FROM issue WHERE milestone_id=?) WHERE id=?",
		milestoneID,
		milestoneID,
	); err != nil {
		return
	}

	return updateMilestoneCompleteness(e, milestoneID)
}

func updateMilestoneClosedNum(e Engine, milestoneID int64) (err error) {
	if _, err = e.Exec("UPDATE `milestone` SET num_closed_issues=(SELECT count(*) FROM issue WHERE milestone_id=? AND is_closed=?) WHERE id=?",
		milestoneID,
		true,
		milestoneID,
	); err != nil {
		return
	}

	return updateMilestoneCompleteness(e, milestoneID)
}

func changeMilestoneAssign(e *xorm.Session, doer *User, issue *Issue, oldMilestoneID int64) error {
	if err := updateIssueCols(e, issue, "milestone_id"); err != nil {
		return err
	}

	if oldMilestoneID > 0 {
		if err := updateMilestoneTotalNum(e, oldMilestoneID); err != nil {
			return err
		}
		if issue.IsClosed {
			if err := updateMilestoneClosedNum(e, oldMilestoneID); err != nil {
				return err
			}
		}
	}

	if issue.MilestoneID > 0 {
		if err := updateMilestoneTotalNum(e, issue.MilestoneID); err != nil {
			return err
		}
		if issue.IsClosed {
			if err := updateMilestoneClosedNum(e, issue.MilestoneID); err != nil {
				return err
			}
		}
	}

	if oldMilestoneID > 0 || issue.MilestoneID > 0 {
		if err := issue.loadRepo(e); err != nil {
			return err
		}

		if _, err := createMilestoneComment(e, doer, issue.Repo, issue, oldMilestoneID, issue.MilestoneID); err != nil {
			return err
		}
	}

	return nil
}

// ChangeMilestoneAssign changes assignment of milestone for issue.
func ChangeMilestoneAssign(issue *Issue, doer *User, oldMilestoneID int64) (err error) {
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	if err = changeMilestoneAssign(sess, doer, issue, oldMilestoneID); err != nil {
		return err
	}

	if err = sess.Commit(); err != nil {
		return fmt.Errorf("Commit: %v", err)
	}
	return nil
}

// DeleteMilestoneByRepoID deletes a milestone from a repository.
func DeleteMilestoneByRepoID(repoID, id int64) error {
	m, err := GetMilestoneByRepoID(repoID, id)
	if err != nil {
		if IsErrMilestoneNotExist(err) {
			return nil
		}
		return err
	}

	repo, err := GetRepositoryByID(m.RepoID)
	if err != nil {
		return err
	}

	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	if _, err = sess.ID(m.ID).Delete(new(Milestone)); err != nil {
		return err
	}

	numMilestones, err := countRepoMilestones(sess, repo.ID)
	if err != nil {
		return err
	}
	numClosedMilestones, err := countRepoClosedMilestones(sess, repo.ID)
	if err != nil {
		return err
	}
	repo.NumMilestones = int(numMilestones)
	repo.NumClosedMilestones = int(numClosedMilestones)

	if _, err = sess.ID(repo.ID).Cols("num_milestones, num_closed_milestones").Update(repo); err != nil {
		return err
	}

	if _, err = sess.Exec("UPDATE `issue` SET milestone_id = 0 WHERE milestone_id = ?", m.ID); err != nil {
		return err
	}
	return sess.Commit()
}
<|endoftext|>


<commit_before>package wallet

import (
	"bytes"
	"crypto/rand"
	"errors"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

var (
	errAlreadyUnlocked   = errors.New("wallet has already been unlocked")
	errReencrypt         = errors.New("wallet is already encrypted, cannot encrypt again")
	errUnencryptedWallet = errors.New("wallet has not been encrypted yet")

	unlockModifier = types.Specifier{'u', 'n', 'l', 'o', 'c', 'k'}
)

// uidEncryptionKey creates an encryption key that is used to decrypt a
// specific key file.
func uidEncryptionKey(masterKey crypto.TwofishKey, uid UniqueID) crypto.TwofishKey {
	return crypto.TwofishKey(crypto.HashAll(masterKey, uid))
}

// checkMasterKey verifies that the master key is correct.
func (w *Wallet) checkMasterKey(masterKey crypto.TwofishKey) error {
	uk := uidEncryptionKey(masterKey, w.persist.UID)
	verification, err := uk.DecryptBytes(w.persist.EncryptionVerification)
	if err != nil {
		// Most of the time, the failure is an authentication failure.
		return modules.ErrBadEncryptionKey
	}
	expected := make([]byte, encryptionVerificationLen)
	if !bytes.Equal(expected, verification) {
		return modules.ErrBadEncryptionKey
	}
	return nil
}

// initEncryption checks that the provided encryption key is the valid
// encryption key for the wallet. If encryption has not yet been established
// for the wallet, an encryption key is created.
func (w *Wallet) initEncryption(masterKey crypto.TwofishKey) (modules.Seed, error) {
	// Check if the wallet encryption key has already been set.
	if len(w.persist.EncryptionVerification) != 0 {
		return modules.Seed{}, errReencrypt
	}

	// Create a random seed and use it to generate the seed file for the
	// wallet.
	var seed modules.Seed
	_, err := rand.Read(seed[:])
	if err != nil {
		return modules.Seed{}, err
	}

	// If the input key is blank, use the seed to create the master key.
	// Otherwise, use the input key.
	if masterKey == (crypto.TwofishKey{}) {
		masterKey = crypto.TwofishKey(crypto.HashObject(seed))
	}
	err = w.createSeed(masterKey, seed)
	if err != nil {
		return modules.Seed{}, err
	}

	// Establish the encryption verification using the masterKey. After this
	// point, the wallet is encrypted.
	uk := uidEncryptionKey(masterKey, w.persist.UID)
	encryptionBase := make([]byte, encryptionVerificationLen)
	w.persist.EncryptionVerification, err = uk.EncryptBytes(encryptionBase)
	if err != nil {
		return modules.Seed{}, err
	}
	err = w.saveSettings()
	if err != nil {
		return modules.Seed{}, err
	}
	return seed, nil
}

// unlock loads all of the encrypted file structures into wallet memory. Even
// after loading, the structures are kept encrypted, but some data such as
// addresses are decrypted so that the wallet knows what to track.
func (w *Wallet) unlock(masterKey crypto.TwofishKey) error {
	// Wallet should only be unlocked once.
	if w.unlocked {
		return errAlreadyUnlocked
	}

	// Check if the wallet encryption key has already been set.
	if len(w.persist.EncryptionVerification) == 0 {
		return errUnencryptedWallet
	}

	// Initialize the encryption of the wallet.
	err := w.checkMasterKey(masterKey)
	if err != nil {
		return err
	}

	// Load the wallet seed that is used to generate new addresses.
	err = w.initPrimarySeed(masterKey)
	if err != nil {
		return err
	}

	// Load all wallet seeds that are not used to generate new addresses.
	err = w.initAuxiliarySeeds(masterKey)
	if err != nil {
		return err
	}

	// Load all keys that were not generated by a seed.
	err = w.initUnseededKeys(masterKey)
	if err != nil {
		return err
	}
	w.unlocked = true
	return nil
}

// wipeSecrets erases all of the seeds and secret keys in the wallet.
func (w *Wallet) wipeSecrets() {
	// 'for i := range' must be used to prevent copies of secret data from
	// being made.
	for i := range w.keys {
		for j := range w.keys[i].SecretKeys {
			crypto.SecureWipe(w.keys[i].SecretKeys[j][:])
		}
	}
	for i := range w.seeds {
		crypto.SecureWipe(w.seeds[i][:])
	}
	crypto.SecureWipe(w.primarySeed[:])
	w.seeds = w.seeds[:0]
}

// Encrypted returns whether or not the wallet has been encrypted.
func (w *Wallet) Encrypted() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	if build.DEBUG && w.unlocked && len(w.persist.EncryptionVerification) == 0 {
		panic("wallet is both unlocked and unencrypted")
	}
	return len(w.persist.EncryptionVerification) != 0
}

// Encrypt will encrypt the wallet using the input key. Upon encryption, a
// primary seed will be created for the wallet (no seed exists prior to this
// point). If the key is blank, then the hash of the seed that is generated
// will be used as the key. The wallet will still be locked after encryption.
//
// Encrypt can only be called once throughout the life of the wallet, and will
// return an error on subsequent calls (even after restarting the wallet). To
// reset the wallet, the wallet files must be moved to a different directory or
// deleted.
func (w *Wallet) Encrypt(masterKey crypto.TwofishKey) (modules.Seed, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.initEncryption(masterKey)
}

// Unlocked indicates whether the wallet is locked or unlocked.
func (w *Wallet) Unlocked() bool {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.unlocked
}

// Lock will erase all keys from memory and prevent the wallet from spending
// coins until it is unlocked.
func (w *Wallet) Lock() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return modules.ErrLockedWallet
	}
	w.log.Println("INFO: Locking wallet.")

	// Wipe all of the seeds and secret keys, they will be replaced upon
	// calling 'Unlock' again.
	w.wipeSecrets()
	w.unlocked = false
	return nil
}

// Unlock will decrypt the wallet seed and load all of the addresses into
// memory.
func (w *Wallet) Unlock(masterKey crypto.TwofishKey) error {
	w.log.Println("INFO: Unlocking wallet.")

	// Initialize all of the keys in the wallet under a lock. While holding the
	// lock, also grab the subscriber status.
	w.mu.Lock()
	subscribed := w.subscribed
	err := w.unlock(masterKey)
	w.mu.Unlock()
	if err != nil {
		return err
	}

	// Subscribe to the consensus set if this is the first unlock for the
	// wallet object.
	if !subscribed {
		// During rescan, print height every 3 seconds.
		if build.Release != "testing" {
			go func() {
				println("Rescanning consensus set...")
				for range time.Tick(time.Second * 3) {
					w.mu.RLock()
					height := w.consensusSetHeight
					done := w.subscribed
					w.mu.RUnlock()
					if done {
						println("\nDone!")
						break
					}
					print("\rScanned to height ", height, "...")
				}
			}()
		}
		err = w.cs.ConsensusSetSubscribe(w, modules.ConsensusChangeBeginning)
		if err != nil {
			return errors.New("wallet subscription failed: " + err.Error())
		}
		w.tpool.TransactionPoolSubscribe(w)
		w.mu.Lock()
		w.subscribed = true
		w.mu.Unlock()
	}
	return nil
}
<commit_msg>don't mark the wallet as unlocked until rescan finishes<commit_after>package wallet

import (
	"bytes"
	"crypto/rand"
	"errors"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

var (
	errAlreadyUnlocked   = errors.New("wallet has already been unlocked")
	errReencrypt         = errors.New("wallet is already encrypted, cannot encrypt again")
	errUnencryptedWallet = errors.New("wallet has not been encrypted yet")

	unlockModifier = types.Specifier{'u', 'n', 'l', 'o', 'c', 'k'}
)

// uidEncryptionKey creates an encryption key that is used to decrypt a
// specific key file.
func uidEncryptionKey(masterKey crypto.TwofishKey, uid UniqueID) crypto.TwofishKey {
	return crypto.TwofishKey(crypto.HashAll(masterKey, uid))
}

// checkMasterKey verifies that the master key is correct.
func (w *Wallet) checkMasterKey(masterKey crypto.TwofishKey) error {
	uk := uidEncryptionKey(masterKey, w.persist.UID)
	verification, err := uk.DecryptBytes(w.persist.EncryptionVerification)
	if err != nil {
		// Most of the time, the failure is an authentication failure.
		return modules.ErrBadEncryptionKey
	}
	expected := make([]byte, encryptionVerificationLen)
	if !bytes.Equal(expected, verification) {
		return modules.ErrBadEncryptionKey
	}
	return nil
}

// initEncryption checks that the provided encryption key is the valid
// encryption key for the wallet. If encryption has not yet been established
// for the wallet, an encryption key is created.
func (w *Wallet) initEncryption(masterKey crypto.TwofishKey) (modules.Seed, error) {
	// Check if the wallet encryption key has already been set.
	if len(w.persist.EncryptionVerification) != 0 {
		return modules.Seed{}, errReencrypt
	}

	// Create a random seed and use it to generate the seed file for the
	// wallet.
	var seed modules.Seed
	_, err := rand.Read(seed[:])
	if err != nil {
		return modules.Seed{}, err
	}

	// If the input key is blank, use the seed to create the master key.
	// Otherwise, use the input key.
	if masterKey == (crypto.TwofishKey{}) {
		masterKey = crypto.TwofishKey(crypto.HashObject(seed))
	}
	err = w.createSeed(masterKey, seed)
	if err != nil {
		return modules.Seed{}, err
	}

	// Establish the encryption verification using the masterKey. After this
	// point, the wallet is encrypted.
	uk := uidEncryptionKey(masterKey, w.persist.UID)
	encryptionBase := make([]byte, encryptionVerificationLen)
	w.persist.EncryptionVerification, err = uk.EncryptBytes(encryptionBase)
	if err != nil {
		return modules.Seed{}, err
	}
	err = w.saveSettings()
	if err != nil {
		return modules.Seed{}, err
	}
	return seed, nil
}

// unlock loads all of the encrypted file structures into wallet memory. Even
// after loading, the structures are kept encrypted, but some data such as
// addresses are decrypted so that the wallet knows what to track.
func (w *Wallet) managedUnlock(masterKey crypto.TwofishKey) error {
	var subscribed bool
	err := func() error {
		w.mu.Lock()
		defer w.mu.Unlock()
		subscribed = w.subscribed

		// Wallet should only be unlocked once.
		if w.unlocked {
			return errAlreadyUnlocked
		}

		// Check if the wallet encryption key has already been set.
		if len(w.persist.EncryptionVerification) == 0 {
			return errUnencryptedWallet
		}

		// Initialize the encryption of the wallet.
		err := w.checkMasterKey(masterKey)
		if err != nil {
			return err
		}

		// Load the wallet seed that is used to generate new addresses.
		err = w.initPrimarySeed(masterKey)
		if err != nil {
			return err
		}

		// Load all wallet seeds that are not used to generate new addresses.
		err = w.initAuxiliarySeeds(masterKey)
		if err != nil {
			return err
		}

		// Load all keys that were not generated by a seed.
		return w.initUnseededKeys(masterKey)
	}()
	if err != nil {
		return err
	}

	// Subscribe to the consensus set if this is the first unlock for the
	// wallet object.
	if !subscribed {
		// During rescan, print height every 3 seconds.
		if build.Release != "testing" {
			go func() {
				println("Rescanning consensus set...")
				for range time.Tick(time.Second * 3) {
					w.mu.RLock()
					height := w.consensusSetHeight
					done := w.subscribed
					w.mu.RUnlock()
					if done {
						println("\nDone!")
						break
					}
					print("\rScanned to height ", height, "...")
				}
			}()
		}
		err = w.cs.ConsensusSetSubscribe(w, modules.ConsensusChangeBeginning)
		if err != nil {
			return errors.New("wallet subscription failed: " + err.Error())
		}
		w.tpool.TransactionPoolSubscribe(w)
		w.mu.Lock()
		w.subscribed = true
		w.mu.Unlock()
	}

	w.mu.Lock()
	w.unlocked = true
	w.mu.Unlock()
	return nil
}

// wipeSecrets erases all of the seeds and secret keys in the wallet.
func (w *Wallet) wipeSecrets() {
	// 'for i := range' must be used to prevent copies of secret data from
	// being made.
	for i := range w.keys {
		for j := range w.keys[i].SecretKeys {
			crypto.SecureWipe(w.keys[i].SecretKeys[j][:])
		}
	}
	for i := range w.seeds {
		crypto.SecureWipe(w.seeds[i][:])
	}
	crypto.SecureWipe(w.primarySeed[:])
	w.seeds = w.seeds[:0]
}

// Encrypted returns whether or not the wallet has been encrypted.
func (w *Wallet) Encrypted() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	if build.DEBUG && w.unlocked && len(w.persist.EncryptionVerification) == 0 {
		panic("wallet is both unlocked and unencrypted")
	}
	return len(w.persist.EncryptionVerification) != 0
}

// Encrypt will encrypt the wallet using the input key. Upon encryption, a
// primary seed will be created for the wallet (no seed exists prior to this
// point). If the key is blank, then the hash of the seed that is generated
// will be used as the key. The wallet will still be locked after encryption.
//
// Encrypt can only be called once throughout the life of the wallet, and will
// return an error on subsequent calls (even after restarting the wallet). To
// reset the wallet, the wallet files must be moved to a different directory or
// deleted.
func (w *Wallet) Encrypt(masterKey crypto.TwofishKey) (modules.Seed, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.initEncryption(masterKey)
}

// Unlocked indicates whether the wallet is locked or unlocked.
func (w *Wallet) Unlocked() bool {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.unlocked
}

// Lock will erase all keys from memory and prevent the wallet from spending
// coins until it is unlocked.
func (w *Wallet) Lock() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return modules.ErrLockedWallet
	}
	w.log.Println("INFO: Locking wallet.")

	// Wipe all of the seeds and secret keys, they will be replaced upon
	// calling 'Unlock' again.
	w.wipeSecrets()
	w.unlocked = false
	return nil
}

// Unlock will decrypt the wallet seed and load all of the addresses into
// memory.
func (w *Wallet) Unlock(masterKey crypto.TwofishKey) error {
	w.log.Println("INFO: Unlocking wallet.")

	// Initialize all of the keys in the wallet under a lock. While holding the
	// lock, also grab the subscriber status.
	return w.managedUnlock(masterKey)
}
<|endoftext|>