<commit_before>package main

import (
	"code.google.com/p/go.crypto/md4"
	"code.google.com/p/go.crypto/ripemd160"
	"crypto/hmac"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"errors"
	"flag"
	"fmt"
	"hash"
	"hash/adler32"
	"hash/crc32"
	"hash/crc64"
	"hash/fnv"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

var (
	alg     string
	key     string
	salt    string
	hashStr bool
)

func init() {
	flag.StringVar(&alg, "a", "sha256", "algorithm")
	flag.StringVar(&key, "k", "", "key (for hashes that use a key, e.g. HMAC)")
	flag.StringVar(&salt, "salt", "", "salt")
	flag.BoolVar(&hashStr, "s", false, "hash a string")
	flag.Parse()
}

var (
	algDescs = map[string]string{
		"adler32":         "Adler-32 checksum (RFC 1950)",
		"crc32":           "32-bit cyclic redundancy check (CRC-32) checksum (defaults to IEEE polynomial)",
		"crc32ieee":       "CRC-32 using the IEEE polynomial (0xedb88320)",
		"crc32castagnoli": "CRC-32 using the Castagnoli polynomial (0x82f63b78)",
		"crc32koopman":    "CRC-32 using the Koopman polynomial (0xeb31d82e)",
		"crc64":           "64-bit cyclic redundancy check (CRC-64) checksum (defaults to ISO polynomial)",
		"crc64iso":        "CRC-64 using the ISO polynomial (0xD800000000000000)",
		"crc64ecma":       "CRC-64 using the ECMA polynomial (0xC96C5795D7870F42)",
		"fnv":             "FNV-1 non-cryptographic hash (defaults to fnv32)",
		"fnv32":           "32-bit FNV-1",
		"fnv32a":          "32-bit FNV-1a",
		"fnv64":           "64-bit FNV-1",
		"fnv64a":          "64-bit FNV-1a",
		"hmac":            "Keyed-Hash Message Authentication Code (HMAC) (requires -k <key>) (defaults to SHA-256)",
		"hmacmd5":         "HMAC using MD5 (requires -k <key>)",
		"hmacsha1":        "HMAC using SHA-1 (requires -k <key>)",
		"hmacsha256":      "HMAC using SHA-256 (requires -k <key>)",
		"hmacsha512":      "HMAC using SHA-512 (requires -k <key>)",
		"md4":             "MD4 hash (RFC 1320)",
		"md5":             "MD5 hash (RFC 1321)",
		"ripemd160":       "RIPEMD-160 hash",
		"sha1":            "SHA-1 hash (RFC 3174)",
		"sha224":          "SHA-224 hash (FIPS 180-2)",
		"sha256":          "SHA-256 hash (FIPS 180-2)",
		"sha384":          "SHA-384 hash (FIPS 180-2)",
		"sha512":          "SHA-512 hash (FIPS 180-2)",
	}
)

func GetHash(a string) (hash.Hash, error) {
	var h hash.Hash
	switch a {
	case "adler32":
		h = adler32.New()
	case "crc32", "crc32ieee":
		h = crc32.New(crc32.MakeTable(crc32.IEEE))
	case "crc32castagnoli":
		h = crc32.New(crc32.MakeTable(crc32.Castagnoli))
	case "crc32koopman":
		h = crc32.New(crc32.MakeTable(crc32.Koopman))
	case "crc64", "crc64iso":
		h = crc64.New(crc64.MakeTable(crc64.ISO))
	case "crc64ecma":
		h = crc64.New(crc64.MakeTable(crc64.ECMA))
	case "fnv", "fnv32":
		h = fnv.New32()
	case "fnv32a":
		h = fnv.New32a()
	case "fnv64":
		h = fnv.New64()
	case "fnv64a":
		h = fnv.New64a()
	case "hmac", "hmacsha256":
		h = hmac.New(sha256.New, []byte(key))
	case "hmacmd5":
		h = hmac.New(md5.New, []byte(key))
	case "hmacsha1":
		h = hmac.New(sha1.New, []byte(key))
	case "hmacsha512":
		h = hmac.New(sha512.New, []byte(key))
	case "md4":
		h = md4.New()
	case "md5":
		h = md5.New()
	case "ripemd160":
		h = ripemd160.New()
	case "sha1":
		h = sha1.New()
	case "sha224":
		h = sha256.New224()
	case "sha256":
		h = sha256.New()
	case "sha384":
		h = sha512.New384()
	case "sha512":
		h = sha512.New()
	default:
		return nil, errors.New("Invalid algorithm")
	}
	return h, nil
}

func HashString(h hash.Hash, s string) string {
	h.Write([]byte(s))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func HashFile(h hash.Hash, f io.Reader) (string, error) {
	h.Write([]byte(salt))
	_, err := io.Copy(h, f)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func Usage() {
	fmt.Println("Usage:", os.Args[0], "[-a <algorithm>] [-salt <salt>] [-s <string to hash>] / <file(s) to hash>\n")
	fmt.Println("Examples:")
	fmt.Println(" ", os.Args[0], `-a md5 document.txt               `, "Generate MD5 digest of a file")
	fmt.Println(" ", os.Args[0], `-a md5 *                          `, "Generate MD5 digests of all files in folder")
	fmt.Println(" ", os.Args[0], `-a sha1 -s hello world            `, "Generate SHA-1 digest of a string")
	fmt.Println(" ", os.Args[0], `-a sha1 -salt s4lt -s hello world `, "Generate salted SHA-1 digest of a string")
	fmt.Println(" ", os.Args[0], `-a hmacsha1 -k k3y -s untouched   `, "Generate HMAC of a string using SHA-1")
	fmt.Println("")
	fmt.Println("Available algorithms (default is SHA-256):")
	mk := make([]string, len(algDescs))
	i := 0
	for k, _ := range algDescs {
		mk[i] = k
		i++
	}
	sort.Strings(mk)
	for _, v := range mk {
		fmt.Printf("  %-15s %s\n", v, algDescs[v])
	}
	fmt.Println("")
	fmt.Println(`Note: For complex strings, put the string in a file, then run Picugen on
      the file. (Don't add newlines to the file as they will alter the output.)`)
}

func main() {
	if flag.NArg() == 0 {
		Usage()
		return
	}
	alg = strings.ToLower(alg)
	h, err := GetHash(alg)
	if err != nil {
		fmt.Println("Error:", err)
		return
	}
	if hashStr {
		var s string
		for i, word := range flag.Args() {
			if i > 0 {
				s += " "
			}
			s += word
		}
		fmt.Println(HashString(h, salt+s))
	} else {
		for _, globStr := range flag.Args() {
			// Okay to ignore error here since error just means no matches
			paths, _ := filepath.Glob(globStr)
			for _, path := range paths {
				var res string
				f, err := os.Open(path)
				if err != nil {
					res = err.Error()
				} else {
					h, err := HashFile(h, f)
					if err != nil {
						res = err.Error()
					} else {
						res = h
					}
					f.Close()
				}
				fmt.Println(res, "", path)
				h.Reset()
			}
		}
	}
}
<commit_msg>group-same and only-same flags<commit_after>package main

import (
	"code.google.com/p/go.crypto/md4"
	"code.google.com/p/go.crypto/ripemd160"
	"crypto/hmac"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"errors"
	"flag"
	"fmt"
	"hash"
	"hash/adler32"
	"hash/crc32"
	"hash/crc64"
	"hash/fnv"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

var (
	alg       string
	key       string
	salt      string
	hashStr   bool
	groupSame bool
	onlySame  bool
)

func init() {
	flag.StringVar(&alg, "a", "sha256", "algorithm")
	flag.StringVar(&key, "k", "", "key (for hashes that use a key, e.g. HMAC)")
	flag.StringVar(&salt, "salt", "", "salt")
	flag.BoolVar(&hashStr, "s", false, "hash a string")
	flag.BoolVar(&groupSame, "group-same", false, "group files with identical hash digests (buffers output)")
	flag.BoolVar(&onlySame, "only-same", false, "only show files with identical hash digests (enables group-same and buffers output)")
	flag.Parse()
	if onlySame {
		groupSame = true
	}
}

var (
	algDescs = map[string]string{
		"adler32":         "Adler-32 checksum (RFC 1950)",
		"crc32":           "32-bit cyclic redundancy check (CRC-32) checksum (defaults to IEEE polynomial)",
		"crc32ieee":       "CRC-32 using the IEEE polynomial (0xedb88320)",
		"crc32castagnoli": "CRC-32 using the Castagnoli polynomial (0x82f63b78)",
		"crc32koopman":    "CRC-32 using the Koopman polynomial (0xeb31d82e)",
		"crc64":           "64-bit cyclic redundancy check (CRC-64) checksum (defaults to ISO polynomial)",
		"crc64iso":        "CRC-64 using the ISO polynomial (0xD800000000000000)",
		"crc64ecma":       "CRC-64 using the ECMA polynomial (0xC96C5795D7870F42)",
		"fnv":             "FNV-1 non-cryptographic hash (defaults to fnv32)",
		"fnv32":           "32-bit FNV-1",
		"fnv32a":          "32-bit FNV-1a",
		"fnv64":           "64-bit FNV-1",
		"fnv64a":          "64-bit FNV-1a",
		"hmac":            "Keyed-Hash Message Authentication Code (HMAC) (requires -k <key>) (defaults to SHA-256)",
		"hmacmd5":         "HMAC using MD5 (requires -k <key>)",
		"hmacsha1":        "HMAC using SHA-1 (requires -k <key>)",
		"hmacsha256":      "HMAC using SHA-256 (requires -k <key>)",
		"hmacsha512":      "HMAC using SHA-512 (requires -k <key>)",
		"md4":             "MD4 hash (RFC 1320)",
		"md5":             "MD5 hash (RFC 1321)",
		"ripemd160":       "RIPEMD-160 hash",
		"sha1":            "SHA-1 hash (RFC 3174)",
		"sha224":          "SHA-224 hash (FIPS 180-2)",
		"sha256":          "SHA-256 hash (FIPS 180-2)",
		"sha384":          "SHA-384 hash (FIPS 180-2)",
		"sha512":          "SHA-512 hash (FIPS 180-2)",
	}
)

func GetHash(a string) (hash.Hash, error) {
	var h hash.Hash
	switch a {
	case "adler32":
		h = adler32.New()
	case "crc32", "crc32ieee":
		h = crc32.New(crc32.MakeTable(crc32.IEEE))
	case "crc32castagnoli":
		h = crc32.New(crc32.MakeTable(crc32.Castagnoli))
	case "crc32koopman":
		h = crc32.New(crc32.MakeTable(crc32.Koopman))
	case "crc64", "crc64iso":
		h = crc64.New(crc64.MakeTable(crc64.ISO))
	case "crc64ecma":
		h = crc64.New(crc64.MakeTable(crc64.ECMA))
	case "fnv", "fnv32":
		h = fnv.New32()
	case "fnv32a":
		h = fnv.New32a()
	case "fnv64":
		h = fnv.New64()
	case "fnv64a":
		h = fnv.New64a()
	case "hmac", "hmacsha256":
		h = hmac.New(sha256.New, []byte(key))
	case "hmacmd5":
		h = hmac.New(md5.New, []byte(key))
	case "hmacsha1":
		h = hmac.New(sha1.New, []byte(key))
	case "hmacsha512":
		h = hmac.New(sha512.New, []byte(key))
	case "md4":
		h = md4.New()
	case "md5":
		h = md5.New()
	case "ripemd160":
		h = ripemd160.New()
	case "sha1":
		h = sha1.New()
	case "sha224":
		h = sha256.New224()
	case "sha256":
		h = sha256.New()
	case "sha384":
		h = sha512.New384()
	case "sha512":
		h = sha512.New()
	default:
		return nil, errors.New("Invalid algorithm")
	}
	return h, nil
}

func HashString(h hash.Hash, s string) string {
	h.Write([]byte(s))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func HashFile(h hash.Hash, f io.Reader) (string, error) {
	h.Write([]byte(salt))
	_, err := io.Copy(h, f)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func Usage() {
	fmt.Println("Usage:", os.Args[0], "[-a <algorithm>] [-salt <salt>] [-s <string to hash>] / <file(s) to hash>\n")
	fmt.Println("Examples:")
	fmt.Println(" ", os.Args[0], `-a md5 document.txt               `, "Generate MD5 digest of a file")
	fmt.Println(" ", os.Args[0], `-a md5 *                          `, "Generate MD5 digests of all files in folder")
	fmt.Println(" ", os.Args[0], `-a sha1 -s hello world            `, "Generate SHA-1 digest of a string")
	fmt.Println(" ", os.Args[0], `-a sha1 -salt s4lt -s hello world `, "Generate salted SHA-1 digest of a string")
	fmt.Println(" ", os.Args[0], `-a hmacsha1 -k k3y -s untouched   `, "Generate HMAC of a string using SHA-1")
	fmt.Println("")
	fmt.Println("Available algorithms (default is SHA-256):")
	mk := make([]string, len(algDescs))
	i := 0
	for k, _ := range algDescs {
		mk[i] = k
		i++
	}
	sort.Strings(mk)
	for _, v := range mk {
		fmt.Printf("  %-15s %s\n", v, algDescs[v])
	}
	fmt.Println("")
	fmt.Println(`Note: For complex strings, put the string in a file, then run Picugen on
      the file. (Don't add newlines to the file as they will alter the output.)`)
}

func main() {
	if flag.NArg() == 0 {
		Usage()
		return
	}
	args := flag.Args()
	alg = strings.ToLower(alg)
	h, err := GetHash(alg)
	if err != nil {
		fmt.Println("Error:", err)
		return
	}
	if hashStr {
		var s string
		for i, word := range args {
			if i > 0 {
				s += " "
			}
			s += word
		}
		fmt.Println(HashString(h, salt+s))
		return
	}
	var g map[string][]string
	if groupSame {
		// We won't get a lot of identical digests, so reduce allocation
		// by assuming each file/glob will produce a unique digest.
		g = make(map[string][]string, len(args))
	}
	for _, globStr := range flag.Args() {
		// Okay to ignore error here since error just means no matches
		paths, _ := filepath.Glob(globStr)
		for _, path := range paths {
			var res string
			f, err := os.Open(path)
			if err != nil {
				res = err.Error()
			} else {
				h, err := HashFile(h, f)
				if err != nil {
					res = err.Error()
				} else {
					res = h
				}
				f.Close()
			}
			if groupSame {
				g[res] = append(g[res], path)
			} else {
				fmt.Println(res, "", path)
			}
			h.Reset()
		}
	}
	if groupSame {
		for k, v := range g {
			if !onlySame || len(v) > 1 {
				for _, ov := range v {
					fmt.Println(k, "", ov)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 The Oto Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build js

package oto

import (
	"io"
)

const pipeBufSize = 4096

// pipe returns a set of an io.ReadCloser and an io.WriteCloser.
//
// This is basically same as io.Pipe, but is implemented in more effient way under these assumptions:
// 1) These works on a single thread environment so that locks are not required.
// 2) Reading and writing happen equally.
// 3) Closing doesn't have to do anything special.
func pipe() (io.ReadCloser, io.WriteCloser) {
	w := &pipeWriter{provided: make(chan struct{}), consumed: make(chan struct{})}
	r := &pipeReader{w: w}
	return r, w
}

type pipeReader struct {
	w *pipeWriter
}

func (r *pipeReader) Read(buf []byte) (int, error) {
	// If this returns 0, the caller might block forever on browsers.
	// For example, bufio.Reader tries to Read until any byte can be read, but context switch never happens on browsers.
	for len(r.w.buf) == 0 {
		<-r.w.provided
	}
	notify := len(r.w.buf) > pipeBufSize
	n := copy(buf, r.w.buf)
	r.w.buf = r.w.buf[n:]
	if notify {
		r.w.consumed <- struct{}{}
	}
	return n, nil
}

func (r *pipeReader) Close() error {
	return nil
}

type pipeWriter struct {
	buf      []byte
	provided chan struct{}
	consumed chan struct{}
}

func (w *pipeWriter) Write(buf []byte) (int, error) {
	for len(w.buf) > pipeBufSize {
		<-w.consumed
	}
	notify := len(w.buf) == 0
	w.buf = append(w.buf, buf...)
	if notify {
		w.provided <- struct{}{}
	}
	return len(buf), nil
}

func (w *pipeWriter) Close() error {
	return nil
}
<commit_msg>js: Avoid using channels on pipe<commit_after>// Copyright 2019 The Oto Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build js

package oto

import (
	"io"
	"runtime"
)

const pipeBufSize = 4096

// pipe returns a set of an io.ReadCloser and an io.WriteCloser.
//
// This is basically same as io.Pipe, but is implemented in more effient way under these assumptions:
// 1) These works on a single thread environment so that locks are not required.
// 2) Reading and writing happen equally.
// 3) Closing doesn't have to do anything special.
func pipe() (io.ReadCloser, io.WriteCloser) {
	w := &pipeWriter{}
	r := &pipeReader{w: w}
	return r, w
}

type pipeReader struct {
	w *pipeWriter
}

func (r *pipeReader) Read(buf []byte) (int, error) {
	// If this returns 0, the caller might block forever on browsers.
	// For example, bufio.Reader tries to Read until any byte can be read, but context switch never happens on browsers.
	for len(r.w.buf) == 0 {
		runtime.Gosched()
	}
	n := copy(buf, r.w.buf)
	r.w.buf = r.w.buf[n:]
	return n, nil
}

func (r *pipeReader) Close() error {
	return nil
}

type pipeWriter struct {
	buf []byte
}

func (w *pipeWriter) Write(buf []byte) (int, error) {
	for len(w.buf) >= pipeBufSize {
		runtime.Gosched()
	}
	w.buf = append(w.buf, buf...)
	return len(buf), nil
}

func (w *pipeWriter) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"app/counter"
	"math/rand"
	"sort"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/search"
)

type DB struct {
	Context context.Context
}

func (db *DB) FetchUser(userId string) {
	//Step 1. Get Vine user's data from the Vine API
	vineApi := VineRequest{db.Context}
	vineUser, err := vineApi.GetUser(userId)

	if err != nil {
		if err.Error() == ErrUserDoesntExist {
			db.UnqueueUser(userId)
		} else {
			log.Errorf(db.Context, "got error getting user %s from vine: %v", userId, err)
		}
		return
	} else if vineUser == nil {
		log.Errorf(db.Context, "failed fetch on user %v. got err %v", userId, err)
		return
	} else if vineUser.Private == 1 {
		return
	}

	recordKey := datastore.NewKey(db.Context, "UserRecord", "", vineUser.UserId, nil)

	//Step 2. Add user to user search index.
	userIndex := &UserIndex{
		Username:    vineUser.Username,
		Location:    vineUser.Location,
		Description: vineUser.Description,
	}

	if len(vineUser.VanityUrls) != 0 {
		userIndex.VanityUrl = strings.ToLower(vineUser.VanityUrls[0])
	}

	index, err := search.Open("users")
	if err != nil {
		log.Errorf(db.Context, "error opening users search index: %v", err)
	} else {
		index.Put(db.Context, vineUser.UserIdStr, userIndex)
	}

	var newLoops, newPosts int64

	//Step 3. Write records (user {meta, record, data}).
	if userRecord, err := db.GetUserRecord(vineUser.UserId); err == nil {
		var userMeta []*UserMeta

		if vineUser.Username != userRecord.Username {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "username", userRecord.Username, time.Now()})
		}

		if vineUser.Location != userRecord.Location {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "location", userRecord.Location, time.Now()})
		}

		if vineUser.Description != userRecord.Description {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "description", userRecord.Description, time.Now()})
		}

		if (vineUser.Verified != 0) != userRecord.Verified {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "verified", strconv.FormatBool(userRecord.Verified), time.Now()})
		}

		newLoops = vineUser.LoopCount - userRecord.LoopCount
		newPosts = vineUser.AuthoredPostCount - userRecord.PostCount

		var metaKey *datastore.Key
		for _, m := range userMeta {
			metaKey = datastore.NewIncompleteKey(db.Context, "UserMeta", recordKey)
			if key, err := datastore.Put(db.Context, metaKey, m); err != nil {
				log.Errorf(db.Context, "got error storing user meta %s - %v: %v", userId, key, err)
			}
		}
	} else if err == datastore.ErrNoSuchEntity {

		counter.IncrementBy(db.Context, "TotalUsers", 1)
		counter.IncrementBy(db.Context, "24hUsers", 1)

		if vineUser.Verified == 1 {
			counter.IncrementBy(db.Context, "TotalVerified", 1)
		}
		if vineUser.ExplicitContent == 1 {
			counter.IncrementBy(db.Context, "TotalExplicit", 1)
		}
		newLoops = vineUser.LoopCount
		newPosts = vineUser.AuthoredPostCount
	}

	if newLoops > 0 {
		counter.IncrementBy(db.Context, "TotalLoops", newLoops)
		counter.IncrementBy(db.Context, "24hLoops", newLoops)
	}

	if newPosts > 0 {
		counter.IncrementBy(db.Context, "TotalPosts", newPosts)
		counter.IncrementBy(db.Context, "24hPosts", newPosts)
	}

	userRecord := UserRecord{
		UserId:            vineUser.UserIdStr,
		Username:          vineUser.Username,
		Description:       vineUser.Description,
		Location:          vineUser.Location,
		ProfileBackground: vineUser.ProfileBackground,
		AvatarUrl:         vineUser.AvatarUrl,
		FollowerCount:     vineUser.FollowerCount,
		FollowingCount:    vineUser.FollowingCount,
		LoopCount:         vineUser.LoopCount,
		LoopVelocity:      vineUser.LoopVelocity,
		PostCount:         vineUser.AuthoredPostCount,
		RevineCount:       (vineUser.PostCount - vineUser.AuthoredPostCount),
		LikeCount:         vineUser.LikeCount,
		Private:           (vineUser.Private != 0),
		Verified:          (vineUser.Verified != 0),
		Explicit:          (vineUser.ExplicitContent != 0),
	}
	if len(vineUser.VanityUrls) != 0 {
		userRecord.Vanity = strings.ToLower(vineUser.VanityUrls[0])
	}
	if _, err := datastore.Put(db.Context, recordKey, &userRecord); err != nil {
		log.Errorf(db.Context, "got error storing user record %s: %v", userId, err)
	}

	dataKey := datastore.NewIncompleteKey(db.Context, "UserData", recordKey)
	userData := UserData{
		UserId:    vineUser.UserId,
		Recorded:  time.Now(),
		Followers: vineUser.FollowerCount,
		Following: vineUser.FollowingCount,
		Loops:     vineUser.LoopCount,
		Velocity:  vineUser.LoopVelocity,
		Posts:     vineUser.AuthoredPostCount,
		Revines:   (vineUser.PostCount - vineUser.AuthoredPostCount),
		Likes:     vineUser.LikeCount,
	}
	if key, err := datastore.Put(db.Context, dataKey, &userData); err != nil {
		log.Errorf(db.Context, "got error storing user data %s - %v: %v", userId, key, err)
	}
}

func (db *DB) GetRecentUsers(n int, filter ...interface{}) (records []UserRecord, err error) {
	q := datastore.NewQuery("UserRecord").Order("-Discovered")
	if len(filter) > 0 && len(filter)%2 == 0 {
		for i := 0; i < len(filter)/2; i += 2 {
			q = q.Filter(filter[i].(string), filter[i+1])
		}
	}
	q = q.Limit(n)
	_, err = q.GetAll(db.Context, &records)
	return
}

func (db *DB) GetUser(userId int64) (user *UserRecord, err error) {
	user, err = db.GetUserRecord(userId)
	if err != nil {
		log.Infof(db.Context, "error with userRecord")
		return
	}
	user.UserData, err = db.GetUserData(userId)
	if err != nil {
		log.Infof(db.Context, "error with userData")
		return
	}

	user.UserMeta, err = db.GetUserMeta(userId)
	if err != nil {
		log.Infof(db.Context, "error with userMeta")
	}
	return
}

func (db *DB) GetUserRecord(userId int64) (*UserRecord, error) {

	user := UserRecord{}

	recordKey := datastore.NewKey(db.Context, "UserRecord", "", userId, nil)
	err := datastore.Get(db.Context, recordKey, &user)

	if err != nil {
		return nil, err
	} else {
		return &user, nil
	}
}

func (db *DB) GetUserData(userId int64) (userData []*UserData, err error) {

	dataQuery := datastore.NewQuery("UserData").Filter("UserId =", userId).Order("Recorded")

	_, err = dataQuery.GetAll(db.Context, &userData)
	return
}

func (db *DB) GetUserMeta(userId int64) (userMeta []*UserMeta, err error) {

	dataQuery := datastore.NewQuery("UserMeta").Filter("UserId =", userId).Order("Updated")

	_, err = dataQuery.GetAll(db.Context, &userMeta)
	return
}

func (db *DB) GetTotalUsers() (int, error) {

	var metaStats MetaStats

	key := datastore.NewKey(db.Context, "__Stat_Kind_IsRootEntity__", "UserMeta", 0, nil)
	err := datastore.Get(db.Context, key, &metaStats)

	return metaStats.Count, err
}

func (db *DB) GetTop() (data map[string]interface{}) {

	var topOverall, topFollowed, topLooped, topPosts, topRevines []UserRecord

	//top overall
	q := datastore.NewQuery("UserRecord").Order("-FollowerCount").Limit(10)
	q.GetAll(db.Context, &topOverall)

	sort.Sort(ByOverall(topOverall))

	//top followed
	q = datastore.NewQuery("UserRecord").Order("-FollowerCount").Limit(10)
	q.GetAll(db.Context, &topFollowed)

	//top looped
	q = datastore.NewQuery("UserRecord").Order("-LoopCount").Limit(10)
	q.GetAll(db.Context, &topLooped)

	//top posts
	q = datastore.NewQuery("UserRecord").Order("-PostCount").Limit(5)
	q.GetAll(db.Context, &topPosts)

	//top Revines
	q = datastore.NewQuery("UserRecord").Order("-RevineCount").Limit(5)
	q.GetAll(db.Context, &topRevines)

	data = map[string]interface{}{
		"topOverall":  topOverall,
		"topFollowed": topFollowed,
		"topLooped":   topLooped,
		"topPosts":    topPosts,
		"topRevines":  topRevines,
	}
	return
}

func (a ByOverall) Len() int      { return len(a) }
func (a ByOverall) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByOverall) Less(i, j int) bool {
	return a[i].FollowerCount > a[j].FollowerCount && a[i].LoopCount > a[j].LoopCount && a[i].FollowingCount < a[j].FollowingCount
}

func (a ByFollowers) Len() int      { return len(a) }
func (a ByFollowers) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByFollowers) Less(i, j int) bool {
	return a[i].FollowerCount > a[j].FollowerCount
}

func (a ByLoops) Len() int      { return len(a) }
func (a ByLoops) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByLoops) Less(i, j int) bool {
	return a[i].LoopCount > a[j].LoopCount
}

func (a ByPosts) Len() int      { return len(a) }
func (a ByPosts) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByPosts) Less(i, j int) bool {
	return a[i].PostCount > a[j].PostCount
}

func (a ByRevines) Len() int      { return len(a) }
func (a ByRevines) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByRevines) Less(i, j int) bool {
	return a[i].RevineCount > a[j].RevineCount
}

func (db *DB) UnqueueUser(user string) {
	var key *datastore.Key
	vineApi := VineRequest{db.Context}
	if vineApi.IsVanity(user) {
		q := datastore.NewQuery("Queue").Filter("UserID =", user).KeysOnly()
		keys, err := q.GetAll(db.Context, nil)
		if err == nil {
			key = keys[0]
		} else {
			log.Errorf(db.Context, "error removing %v from queue: %v", user, err)
		}
	} else {
		userId, _ := strconv.ParseInt(user, 10, 64)
		key = datastore.NewKey(db.Context, "Queue", "", userId, nil)
	}

	datastore.Delete(db.Context, key)
	log.Infof(db.Context, "%v removed from queue.", user)
}

func RandomKey(a []*datastore.Key) *datastore.Key {
	rand.Seed(time.Now().UTC().UnixNano())
	return a[rand.Intn(len(a))]
}

func RemoveDuplicates(a []string) []string {
	found := make(map[string]bool)
	j := 0
	for i, x := range a {
		if !found[x] {
			found[x] = true
			a[j] = a[i]
			j++
		}
	}
	return a[:j]
}
<commit_msg>Track Discovered date.<commit_after>package main

import (
	"app/counter"
	"math/rand"
	"sort"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/search"
)

type DB struct {
	Context context.Context
}

func (db *DB) FetchUser(userId string) {
	//Step 1. Get Vine user's data from the Vine API
	vineApi := VineRequest{db.Context}
	vineUser, err := vineApi.GetUser(userId)

	if err != nil {
		if err.Error() == ErrUserDoesntExist {
			db.UnqueueUser(userId)
		} else {
			log.Errorf(db.Context, "got error getting user %s from vine: %v", userId, err)
		}
		return
	} else if vineUser == nil {
		log.Errorf(db.Context, "failed fetch on user %v. got err %v", userId, err)
		return
	} else if vineUser.Private == 1 {
		return
	}

	recordKey := datastore.NewKey(db.Context, "UserRecord", "", vineUser.UserId, nil)

	//Step 2. Add user to user search index.
	userIndex := &UserIndex{
		Username:    vineUser.Username,
		Location:    vineUser.Location,
		Description: vineUser.Description,
	}

	if len(vineUser.VanityUrls) != 0 {
		userIndex.VanityUrl = strings.ToLower(vineUser.VanityUrls[0])
	}

	index, err := search.Open("users")
	if err != nil {
		log.Errorf(db.Context, "error opening users search index: %v", err)
	} else {
		index.Put(db.Context, vineUser.UserIdStr, userIndex)
	}

	var newLoops, newPosts int64

	//Step 3. Write records (user {meta, record, data}).
	if userRecord, err := db.GetUserRecord(vineUser.UserId); err == nil {
		var userMeta []*UserMeta

		if vineUser.Username != userRecord.Username {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "username", userRecord.Username, time.Now()})
		}

		if vineUser.Location != userRecord.Location {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "location", userRecord.Location, time.Now()})
		}

		if vineUser.Description != userRecord.Description {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "description", userRecord.Description, time.Now()})
		}

		if (vineUser.Verified != 0) != userRecord.Verified {
			userMeta = append(userMeta, &UserMeta{vineUser.UserId, "verified", strconv.FormatBool(userRecord.Verified), time.Now()})
		}

		newLoops = vineUser.LoopCount - userRecord.LoopCount
		newPosts = vineUser.AuthoredPostCount - userRecord.PostCount

		var metaKey *datastore.Key
		for _, m := range userMeta {
			metaKey = datastore.NewIncompleteKey(db.Context, "UserMeta", recordKey)
			if key, err := datastore.Put(db.Context, metaKey, m); err != nil {
				log.Errorf(db.Context, "got error storing user meta %s - %v: %v", userId, key, err)
			}
		}
	} else if err == datastore.ErrNoSuchEntity {

		counter.IncrementBy(db.Context, "TotalUsers", 1)
		counter.IncrementBy(db.Context, "24hUsers", 1)

		if vineUser.Verified == 1 {
			counter.IncrementBy(db.Context, "TotalVerified", 1)
		}
		if vineUser.ExplicitContent == 1 {
			counter.IncrementBy(db.Context, "TotalExplicit", 1)
		}
		newLoops = vineUser.LoopCount
		newPosts = vineUser.AuthoredPostCount
	}

	if newLoops > 0 {
		counter.IncrementBy(db.Context, "TotalLoops", newLoops)
		counter.IncrementBy(db.Context, "24hLoops", newLoops)
	}

	if newPosts > 0 {
		counter.IncrementBy(db.Context, "TotalPosts", newPosts)
		counter.IncrementBy(db.Context, "24hPosts", newPosts)
	}

	userRecord := UserRecord{
		UserId:            vineUser.UserIdStr,
		Username:          vineUser.Username,
		Description:       vineUser.Description,
		Location:          vineUser.Location,
		ProfileBackground: vineUser.ProfileBackground,
		AvatarUrl:         vineUser.AvatarUrl,
		FollowerCount:     vineUser.FollowerCount,
		FollowingCount:    vineUser.FollowingCount,
		LoopCount:         vineUser.LoopCount,
		LoopVelocity:      vineUser.LoopVelocity,
		PostCount:         vineUser.AuthoredPostCount,
		RevineCount:       (vineUser.PostCount - vineUser.AuthoredPostCount),
		LikeCount:         vineUser.LikeCount,
		Private:           (vineUser.Private != 0),
		Verified:          (vineUser.Verified != 0),
		Explicit:          (vineUser.ExplicitContent != 0),
		Discovered:        time.Now(),
	}
	if len(vineUser.VanityUrls) != 0 {
		userRecord.Vanity = strings.ToLower(vineUser.VanityUrls[0])
	}
	if _, err := datastore.Put(db.Context, recordKey, &userRecord); err != nil {
		log.Errorf(db.Context, "got error storing user record %s: %v", userId, err)
	}

	dataKey := datastore.NewIncompleteKey(db.Context, "UserData", recordKey)
	userData := UserData{
		UserId:    vineUser.UserId,
		Recorded:  time.Now(),
		Followers: vineUser.FollowerCount,
		Following: vineUser.FollowingCount,
		Loops:     vineUser.LoopCount,
		Velocity:  vineUser.LoopVelocity,
		Posts:     vineUser.AuthoredPostCount,
		Revines:   (vineUser.PostCount - vineUser.AuthoredPostCount),
		Likes:     vineUser.LikeCount,
	}
	if key, err := datastore.Put(db.Context, dataKey, &userData); err != nil {
		log.Errorf(db.Context, "got error storing user data %s - %v: %v", userId, key, err)
	}
}

func (db *DB) GetRecentUsers(n int, filter ...interface{}) (records []UserRecord, err error) {
	q := datastore.NewQuery("UserRecord").Order("-Discovered")
	if len(filter) > 0 && len(filter)%2 == 0 {
		for i := 0; i < len(filter)/2; i += 2 {
			q = q.Filter(filter[i].(string), filter[i+1])
		}
	}
	q = q.Limit(n)
	_, err = q.GetAll(db.Context, &records)
	return
}

func (db *DB) GetUser(userId int64) (user *UserRecord, err error) {
	user, err = db.GetUserRecord(userId)
	if err != nil {
		log.Infof(db.Context, "error with userRecord")
		return
	}
	user.UserData, err = db.GetUserData(userId)
	if err != nil {
		log.Infof(db.Context, "error with userData")
		return
	}

	user.UserMeta, err = db.GetUserMeta(userId)
	if err != nil {
		log.Infof(db.Context, "error with userMeta")
	}
	return
}

func (db *DB) GetUserRecord(userId int64) (*UserRecord, error) {

	user := UserRecord{}

	recordKey := datastore.NewKey(db.Context, "UserRecord", "", userId, nil)
	err := datastore.Get(db.Context, recordKey, &user)

	if err != nil {
		return nil, err
	} else {
		return &user, nil
	}
}

func (db *DB) GetUserData(userId int64) (userData []*UserData, err error) {

	dataQuery := datastore.NewQuery("UserData").Filter("UserId =", userId).Order("Recorded")

	_, err = dataQuery.GetAll(db.Context, &userData)
	return
}

func (db *DB) GetUserMeta(userId int64) (userMeta []*UserMeta, err error) {

	dataQuery := datastore.NewQuery("UserMeta").Filter("UserId =", userId).Order("Updated")

	_, err = dataQuery.GetAll(db.Context, &userMeta)
	return
}

func (db *DB) GetTotalUsers() (int, error) {

	var metaStats MetaStats

	key := datastore.NewKey(db.Context, "__Stat_Kind_IsRootEntity__", "UserMeta", 0, nil)
	err := datastore.Get(db.Context, key, &metaStats)

	return metaStats.Count, err
}

func (db *DB) GetTop() (data map[string]interface{}) {

	var topOverall, topFollowed, topLooped, topPosts, topRevines []UserRecord

	//top overall
	q := datastore.NewQuery("UserRecord").Order("-FollowerCount").Limit(10)
	q.GetAll(db.Context, &topOverall)

	sort.Sort(ByOverall(topOverall))

	//top followed
	q = datastore.NewQuery("UserRecord").Order("-FollowerCount").Limit(10)
	q.GetAll(db.Context, &topFollowed)

	//top looped
	q = datastore.NewQuery("UserRecord").Order("-LoopCount").Limit(10)
	q.GetAll(db.Context, &topLooped)

	//top posts
	q = datastore.NewQuery("UserRecord").Order("-PostCount").Limit(5)
	q.GetAll(db.Context, &topPosts)

	//top Revines
	q = datastore.NewQuery("UserRecord").Order("-RevineCount").Limit(5)
	q.GetAll(db.Context, &topRevines)

	data = map[string]interface{}{
		"topOverall":  topOverall,
		"topFollowed": topFollowed,
		"topLooped":   topLooped,
		"topPosts":    topPosts,
		"topRevines":  topRevines,
	}
	return
}

func (a ByOverall) Len() int      { return len(a) }
func (a ByOverall) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByOverall) Less(i, j int) bool {
	return a[i].FollowerCount > a[j].FollowerCount && a[i].LoopCount > a[j].LoopCount && a[i].FollowingCount < a[j].FollowingCount
}

func (a ByFollowers) Len() int      { return len(a) }
func (a ByFollowers) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByFollowers) Less(i, j int) bool {
	return a[i].FollowerCount > a[j].FollowerCount
}

func (a ByLoops) Len() int      { return len(a) }
func (a ByLoops) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByLoops) Less(i, j int) bool {
	return a[i].LoopCount > a[j].LoopCount
}

func (a ByPosts) Len() int      { return len(a) }
func (a ByPosts) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByPosts) Less(i, j int) bool {
	return a[i].PostCount > a[j].PostCount
}

func (a ByRevines) Len() int      { return len(a) }
func (a ByRevines) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByRevines) Less(i, j int) bool {
	return a[i].RevineCount > a[j].RevineCount
}

func (db *DB) UnqueueUser(user string) {
	var key *datastore.Key
	vineApi := VineRequest{db.Context}
	if vineApi.IsVanity(user) {
		q := datastore.NewQuery("Queue").Filter("UserID =", user).KeysOnly()
		keys, err := q.GetAll(db.Context, nil)
		if err == nil {
			key = keys[0]
		} else {
			log.Errorf(db.Context, "error removing %v from queue: %v", user, err)
		}
	} else {
		userId, _ := strconv.ParseInt(user, 10, 64)
		key = datastore.NewKey(db.Context, "Queue", "", userId, nil)
	}

	datastore.Delete(db.Context, key)
	log.Infof(db.Context, "%v removed from queue.", user)
}

func RandomKey(a []*datastore.Key) *datastore.Key {
	rand.Seed(time.Now().UTC().UnixNano())
	return a[rand.Intn(len(a))]
}

func RemoveDuplicates(a []string) []string {
	found := make(map[string]bool)
	j := 0
	for i, x := range a {
		if !found[x] {
			found[x] = true
			a[j] = a[i]
			j++
		}
	}
	return a[:j]
}
<|endoftext|>


<commit_before>// Package tdigest provides a highly accurate mergeable data-structure
// for quantile estimation.
package tdigest

import (
	"fmt"
	"math"
	"math/rand"
)

// TDigest is a quantile summary structure
type TDigest struct {
	summary     *summary
	compression float64
	count       uint32
}

// New creates a new digest.
// The compression parameter rules the threshold in which samples are
// merged together - the more often distinct samples are merged the more
// precision is lost. Compression should be tuned according to your data
// distribution, but a value of 10 is often good enough. A higher
// compression value means holding more centroids in memory, which means
// a bigger serialization payload and higher memory footprint.
func New(compression float64) *TDigest {
	return &TDigest{
		compression: compression,
		summary:     newSummary(estimateCapacity(compression)),
		count:       0,
	}
}

// Quantile returns the desired percentile estimation.
// Values of p must be between 0 and 1 (inclusive), will panic otherwise.
func (t *TDigest) Quantile(q float64) float64 {
	if q < 0 || q > 1 {
		panic("q must be between 0 and 1 (inclusive)")
	}

	if t.summary.Len() == 0 {
		return math.NaN()
	} else if t.summary.Len() == 1 {
		return t.summary.Min().mean
	}

	q *= float64(t.count)
	var total float64
	i := 0

	found := false
	var result float64

	t.summary.Iterate(func(item centroid) bool {
		k := float64(item.count)

		if q < total+k {
			if i == 0 || i+1 == t.summary.Len() {
				result = item.mean
				found = true
				return false
			}
			succ, pred := t.summary.successorAndPredecessorItems(item.mean)
			delta := (succ.mean - pred.mean) / 2
			result = item.mean + ((q-total)/k-0.5)*delta
			found = true
			return false
		}

		i++
		total += k
		return true
	})

	if found {
		return result
	}
	return t.summary.Max().mean
}

// Add registers a new sample in the digest.
// It's the main entry point for the digest and very likely the only
// method to be used for collecting samples. The count parameter is for
// when you are registering a sample that occurred multiple times - the
// most common value for this is 1.
func (t *TDigest) Add(value float64, count uint32) error {

	if count == 0 {
		return fmt.Errorf("Illegal datapoint <value: %.4f, count: %d>", value, count)
	}

	t.count += count

	if t.summary.Len() == 0 {
		t.summary.Add(value, count)
		return nil
	}

	candidates := t.findNearestCentroids(value)

	for len(candidates) > 0 && count > 0 {
		j := rand.Intn(len(candidates))
		chosen := candidates[j]

		quantile := t.computeCentroidQuantile(chosen)

		if float64(chosen.count+count) > t.threshold(quantile) {
			candidates = append(candidates[:j], candidates[j+1:]...)
			continue
		}

		deltaW := math.Min(t.threshold(quantile)-float64(chosen.count), float64(count))
		t.summary.updateAt(chosen.index, value, uint32(deltaW))
		count -= uint32(deltaW)

		candidates = append(candidates[:j], candidates[j+1:]...)
	}

	if count > 0 {
		t.summary.Add(value, count)
	}

	if float64(t.summary.Len()) > 20*t.compression {
		t.Compress()
	}

	return nil
}

// Compress tries to reduce the number of individual centroids stored
// in the digest.
// Compression trades off accuracy for performance and happens
// automatically after a certain amount of distinct samples have been
// stored.
func (t *TDigest) Compress() {
	if t.summary.Len() <= 1 {
		return
	}

	oldTree := t.summary
	t.summary = newSummary(estimateCapacity(t.compression))

	nodes := oldTree.Data()
	shuffle(nodes)

	for _, item := range nodes {
		t.Add(item.mean, item.count)
	}
}

// Merge joins a given digest into itself.
// Merging is useful when you have multiple TDigest instances running
// in separate threads and you want to compute quantiles over all the
// samples. This is particularly important on a scatter-gather/map-reduce
// scenario.
func (t *TDigest) Merge(other *TDigest) {
	if other.summary.Len() == 0 {
		return
	}

	nodes := other.summary.Data()
	shuffle(nodes)

	for _, item := range nodes {
		t.Add(item.mean, item.count)
	}
}

func shuffle(data []centroid) {
	for i := len(data) - 1; i > 1; i-- {
		other := rand.Intn(i + 1)
		tmp := data[other]
		data[other] = data[i]
		data[i] = tmp
	}
}

func estimateCapacity(compression float64) uint {
	return uint(compression) * 10
}

func (t *TDigest) threshold(q float64) float64 {
	return (4 * float64(t.count) * q * (1 - q)) / t.compression
}

func (t *TDigest) computeCentroidQuantile(c *centroid) float64 {
	cumSum := t.summary.sumUntilMean(c.mean)
	return (float64(c.count)/2.0 + float64(cumSum)) / float64(t.count)
}

func (t *TDigest) findNearestCentroids(mean float64) []*centroid {
	ceil, floor := t.summary.ceilingAndFloorItems(mean)

	if !ceil.isValid() && !floor.isValid() {
		panic("findNearestCentroids called on an empty tree")
	}

	if !ceil.isValid() {
		return []*centroid{&floor}
	}

	if !floor.isValid() {
		return []*centroid{&ceil}
	}

	if math.Abs(floor.mean-mean) < math.Abs(ceil.mean-mean) {
		return []*centroid{&floor}
	} else if math.Abs(floor.mean-mean) == math.Abs(ceil.mean-mean) && floor.mean != ceil.mean {
		return []*centroid{&floor, &ceil}
	} else {
		return []*centroid{&ceil}
	}
}
<commit_msg>Add missing zero from tdigest.New docs<commit_after>// Package tdigest provides a highly accurate mergeable data-structure
// for quantile estimation.
package tdigest

import (
	"fmt"
	"math"
	"math/rand"
)

// TDigest is a quantile summary structure
type TDigest struct {
	summary     *summary
	compression float64
	count       uint32
}

// New creates a new digest.
// The compression parameter rules the threshold in which samples are
// merged together - the more often distinct samples are merged the more
// precision is lost. Compression should be tuned according to your data
// distribution, but a value of 100 is often good enough. A higher
// compression value means holding more centroids in memory (thus: better
// precision), which means a bigger serialization payload and higher
// memory footprint.
func New(compression float64) *TDigest {
	return &TDigest{
		compression: compression,
		summary:     newSummary(estimateCapacity(compression)),
		count:       0,
	}
}

// Quantile returns the desired percentile estimation.
// Values of p must be between 0 and 1 (inclusive), will panic otherwise.
func (t *TDigest) Quantile(q float64) float64 {
	if q < 0 || q > 1 {
		panic("q must be between 0 and 1 (inclusive)")
	}

	if t.summary.Len() == 0 {
		return math.NaN()
	} else if t.summary.Len() == 1 {
		return t.summary.Min().mean
	}

	q *= float64(t.count)
	var total float64
	i := 0

	found := false
	var result float64

	t.summary.Iterate(func(item centroid) bool {
		k := float64(item.count)

		if q < total+k {
			if i == 0 || i+1 == t.summary.Len() {
				result = item.mean
				found = true
				return false
			}
			succ, pred := t.summary.successorAndPredecessorItems(item.mean)
			delta := (succ.mean - pred.mean) / 2
			result = item.mean + ((q-total)/k-0.5)*delta
			found = true
			return false
		}

		i++
		total += k
		return true
	})

	if found {
		return result
	}
	return t.summary.Max().mean
}

// Add registers a new sample in the digest.
// It's the main entry point for the digest and very likely the only
// method to be used for collecting samples. The count parameter is for
// when you are registering a sample that occurred multiple times - the
// most common value for this is 1.
func (t *TDigest) Add(value float64, count uint32) error {

	if count == 0 {
		return fmt.Errorf("Illegal datapoint <value: %.4f, count: %d>", value, count)
	}

	t.count += count

	if t.summary.Len() == 0 {
		t.summary.Add(value, count)
		return nil
	}

	candidates := t.findNearestCentroids(value)

	for len(candidates) > 0 && count > 0 {
		j := rand.Intn(len(candidates))
		chosen := candidates[j]

		quantile := t.computeCentroidQuantile(chosen)

		if float64(chosen.count+count) > t.threshold(quantile) {
			candidates = append(candidates[:j], candidates[j+1:]...)
			continue
		}

		deltaW := math.Min(t.threshold(quantile)-float64(chosen.count), float64(count))
		t.summary.updateAt(chosen.index, value, uint32(deltaW))
		count -= uint32(deltaW)

		candidates = append(candidates[:j], candidates[j+1:]...)
	}

	if count > 0 {
		t.summary.Add(value, count)
	}

	if float64(t.summary.Len()) > 20*t.compression {
		t.Compress()
	}

	return nil
}

// Compress tries to reduce the number of individual centroids stored
// in the digest.
// Compression trades off accuracy for performance and happens
// automatically after a certain amount of distinct samples have been
// stored.
func (t *TDigest) Compress() {
	if t.summary.Len() <= 1 {
		return
	}

	oldTree := t.summary
	t.summary = newSummary(estimateCapacity(t.compression))

	nodes := oldTree.Data()
	shuffle(nodes)

	for _, item := range nodes {
		t.Add(item.mean, item.count)
	}
}

// Merge joins a given digest into itself.
// Merging is useful when you have multiple TDigest instances running
// in separate threads and you want to compute quantiles over all the
// samples. This is particularly important on a scatter-gather/map-reduce
// scenario.
func (t *TDigest) Merge(other *TDigest) {
	if other.summary.Len() == 0 {
		return
	}

	nodes := other.summary.Data()
	shuffle(nodes)

	for _, item := range nodes {
		t.Add(item.mean, item.count)
	}
}

func shuffle(data []centroid) {
	for i := len(data) - 1; i > 1; i-- {
		other := rand.Intn(i + 1)
		tmp := data[other]
		data[other] = data[i]
		data[i] = tmp
	}
}

func estimateCapacity(compression float64) uint {
	return uint(compression) * 10
}

func (t *TDigest) threshold(q float64) float64 {
	return (4 * float64(t.count) * q * (1 - q)) / t.compression
}

func (t *TDigest) computeCentroidQuantile(c *centroid) float64 {
	cumSum := t.summary.sumUntilMean(c.mean)
	return (float64(c.count)/2.0 + float64(cumSum)) / float64(t.count)
}

func (t *TDigest) findNearestCentroids(mean float64) []*centroid {
	ceil, floor := t.summary.ceilingAndFloorItems(mean)

	if !ceil.isValid() && !floor.isValid() {
		panic("findNearestCentroids called on an empty tree")
	}

	if !ceil.isValid() {
		return []*centroid{&floor}
	}

	if !floor.isValid() {
		return []*centroid{&ceil}
	}

	if math.Abs(floor.mean-mean) < math.Abs(ceil.mean-mean) {
		return []*centroid{&floor}
	} else if math.Abs(floor.mean-mean) == math.Abs(ceil.mean-mean) && floor.mean != ceil.mean {
		return []*centroid{&floor, &ceil}
	} else {
		return []*centroid{&ceil}
	}
}
<|endoftext|>


<commit_before>/*
 */
package trayhost

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sort"
)

/*
#cgo linux pkg-config: gtk+-2.0
#cgo linux CFLAGS: -DLINUX -I/usr/include/libappindicator-0.1/
#cgo linux LDFLAGS: -ldl
#cgo windows CFLAGS: -DWIN32 -DUNICODE
#cgo darwin CFLAGS: -DDARWIN -x objective-c
#cgo darwin LDFLAGS: -framework Cocoa
#include <stdlib.h>
#include "platform/platform.h"
*/
import "C"

const (
	ICON_PRIMARY     = iota
	ICON_ALTERNATIVE = iota
	ICON_ATTENTION   = iota
)

const (
	WINDOWS       = iota
	OSX           = iota
	GNOME         = iota
	UNITY         = iota
	LINUX_GENERIC = iota
)

type MenuItem struct {
	Title    string
	Disabled bool
	Handler  func()
}

type MenuItemUpdate struct {
	ItemId int
	Item   MenuItem
}

type MenuItems map[int]MenuItem

var isExiting bool = false
var menuItems MenuItems
var UpdateCh = make(chan MenuItemUpdate, 99)
var icons = map[int]string{}
var tmpFiles []string = make([]string, 0, 3)
var clickHandler func()
var Debug bool = false
var curIconId int = -1
var tmpDir string

// Run the host system's event loop
func Initialize(title string, imageData []byte, items MenuItems, tmpDirectory string) (err error) {
	if !Debug {
		log.SetOutput(ioutil.Discard)
	}
	tmpDir = tmpDirectory
	err = SetIconImage(ICON_PRIMARY, imageData)
	if err != nil {
		return
	}
	initialize(title)
	err = SetIcon(ICON_PRIMARY)
	if err != nil {
		return
	}
	setMenu(items)
	return
}

func SetIconImage(iconId int, imageData []byte) (err error) {
	iconPth, err := createTempFile(imageData)
	if err != nil {
		return
	}
	icons[iconId] = iconPth
	return
}

func EnterLoop() {
	go menuUpdateLoop()
	C.native_loop()
	isExiting = true
}

func Exit() {
	close(UpdateCh)
	C.exit_loop()
	cleanup()
}

func SetIcon(iconId int) (err error) {

	if iconId != curIconId {
		iconPth, ok := icons[iconId]
		if !ok {
			err = fmt.Errorf("No icon with icon id %d", iconId)
			return
		}

		log.Printf("Setting icon %s (id: %d)", iconPth, iconId)
		setIcon(iconPth)
		curIconId = iconId
	}
	return
}

func SetClickHandler(handler func()) {
	clickHandler = handler
}

func menuUpdateLoop() {
	runtime.UnlockOSThread()
	for update := range UpdateCh {
		updateMenuItem(update.ItemId, update.Item)
	}
}

func updateMenuItem(id int, item MenuItem) {
	menuItems[id] = item
	setMenuItem(id, item)
}

func setMenu(menu MenuItems) {
	menuItems = menu
	menuItemOrder := make([]int, 0, len(menuItems))
	for key, _ := range menuItems {
		menuItemOrder = append(menuItemOrder, key)
	}
	sort.Ints(menuItemOrder)

	for id := range menuItemOrder {
		item := menuItems[id]
		setMenuItem(id, item)
	}
}

func createTempFile(iconData []byte) (filename string, err error) {
	err = os.MkdirAll(tmpDir, os.ModeDir|os.ModePerm)
	if err != nil {
		return
	}
	file, err := ioutil.TempFile(tmpDir, "trayhosticon")
	if err != nil {
		return
	}
	defer file.Close()
	_, err = file.Write(iconData)
	filename = file.Name()
	tmpFiles = append(tmpFiles, filename)
	return
}

func cleanup() {
	for _, file := range tmpFiles {
		err := os.Remove(file)
		if err != nil {
			log.Printf("Failed to remove tmp file %s: %v\n", file, err)
		} else {
			log.Printf("Tmp file %s removed\n", file)
		}
	}
}

func cbool(b bool) C.int {
	if b {
		return 1
	} else {
		return 0
	}
}
<commit_msg>Remove mkdir<commit_after>/*
 */
package trayhost

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sort"
)

/*
#cgo linux pkg-config: gtk+-2.0
#cgo linux CFLAGS: -DLINUX -I/usr/include/libappindicator-0.1/
#cgo linux LDFLAGS: -ldl
#cgo windows CFLAGS: -DWIN32 -DUNICODE
#cgo darwin CFLAGS: -DDARWIN -x objective-c
#cgo darwin LDFLAGS: -framework Cocoa
#include <stdlib.h>
#include "platform/platform.h"
*/
import "C"

const (
	ICON_PRIMARY     = iota
	ICON_ALTERNATIVE = iota
	ICON_ATTENTION   = iota
)

const (
	WINDOWS       = iota
	OSX           = iota
	GNOME         = iota
	UNITY         = iota
	LINUX_GENERIC = iota
)

type MenuItem struct {
	Title    string
	Disabled bool
	Handler  func()
}

type MenuItemUpdate struct {
	ItemId int
	Item   MenuItem
}

type MenuItems map[int]MenuItem

var isExiting bool = false
var menuItems MenuItems
var UpdateCh = make(chan MenuItemUpdate, 99)
var icons = map[int]string{}
var tmpFiles []string = make([]string, 0, 3)
var clickHandler func()
var Debug bool = false
var curIconId int = -1
var tmpDir string

// Run the host system's event loop
func Initialize(title string, imageData []byte, items MenuItems, tmpDirectory string) (err error) {
	if !Debug {
		log.SetOutput(ioutil.Discard)
	}
	tmpDir = tmpDirectory
	err = SetIconImage(ICON_PRIMARY, imageData)
	if err != nil {
		return
	}
	initialize(title)
	err = SetIcon(ICON_PRIMARY)
	if err != nil {
		return
	}
	setMenu(items)
	return
}

func SetIconImage(iconId int, imageData []byte) (err error) {
	iconPth, err := createTempFile(imageData)
	if err != nil {
		return
	}
	icons[iconId] = iconPth
	return
}

func EnterLoop() {
	go menuUpdateLoop()
	C.native_loop()
	isExiting = true
}

func Exit() {
	close(UpdateCh)
	C.exit_loop()
	cleanup()
}

func SetIcon(iconId int) (err error) {

	if iconId != curIconId {
		iconPth, ok := icons[iconId]
		if !ok {
			err = fmt.Errorf("No icon with icon id %d", iconId)
			return
		}

		log.Printf("Setting icon %s (id: %d)", iconPth, iconId)
		setIcon(iconPth)
		curIconId = iconId
	}
	return
}

func SetClickHandler(handler func()) {
	clickHandler = handler
}

func menuUpdateLoop() {
	runtime.UnlockOSThread()
	for update := range UpdateCh {
		updateMenuItem(update.ItemId, update.Item)
	}
}

func updateMenuItem(id int, item MenuItem) {
	menuItems[id] = item
	setMenuItem(id, item)
}

func setMenu(menu MenuItems) {
	menuItems = menu
	menuItemOrder := make([]int, 0, len(menuItems))
	for key, _ := range menuItems {
		menuItemOrder = append(menuItemOrder, key)
	}
	sort.Ints(menuItemOrder)

	for id := range menuItemOrder {
		item := menuItems[id]
		setMenuItem(id, item)
	}
}

func createTempFile(iconData []byte) (filename string, err error) {
	file, err := ioutil.TempFile(tmpDir, "trayhosticon")
	if err != nil {
		return
	}
	defer file.Close()
	_, err = file.Write(iconData)
	filename = file.Name()
	tmpFiles = append(tmpFiles, filename)
	return
}

func cleanup() {
	for _, file := range tmpFiles {
		err := os.Remove(file)
		if err != nil {
			log.Printf("Failed to remove tmp file %s: %v\n", file, err)
		} else {
			log.Printf("Tmp file %s removed\n", file)
		}
	}
}

func cbool(b bool) C.int {
	if b {
		return 1
	} else {
		return 0
	}
}
<|endoftext|>


<commit_before>package network

import (
	"testing"

	"github.com/mpreath/netcalc/ipv4"
)

func TestGenerateNetwork(t *testing.T) {
	test_cases := []struct {
		dd_address         string
		dd_network_address string
		dd_mask            string
	}{
		{"192.168.1.1", "192.168.1.0", "255.255.255.0"},
	}

	for _, test_case := range test_cases {
		test_network, _ := GenerateNetwork(test_case.dd_address, test_case.dd_mask)

		dd_test_network, _ := ipv4.Itodd(test_network.Address)
		dd_test_mask, _ := ipv4.Itodd(test_network.Mask)

		if dd_test_network != test_case.dd_network_address {
			t.Errorf("generated address (%s) doesn't match spec address (%s)", dd_test_network, test_case.dd_network_address)
		}

		if dd_test_mask != test_case.dd_mask {
			t.Errorf("generated mask (%s) doesn't match spec mask (%s)", dd_test_mask, test_case.dd_mask)
		}
	}
}
<commit_msg>updated network test<commit_after>package network

import (
	"testing"

	"github.com/mpreath/netcalc/ipv4"
)

func TestGenerateNetwork(t *testing.T) {
	test_cases := []struct {
		dd_address         string
		dd_mask            string
		dd_network_address string
	}{
		{"192.168.1.1", "255.255.255.0", "192.168.1.0"},
	}

	for _, test_case := range test_cases {
		test_network, _ := GenerateNetwork(test_case.dd_address, test_case.dd_mask)

		dd_test_network, _ := ipv4.Itodd(test_network.Address)
		dd_test_mask, _ := ipv4.Itodd(test_network.Mask)

		if dd_test_network != test_case.dd_network_address {
			t.Errorf("generated address (%s) doesn't match spec address (%s)", dd_test_network, test_case.dd_network_address)
		}

		if dd_test_mask != test_case.dd_mask {
			t.Errorf("generated mask (%s) doesn't match spec mask (%s)", dd_test_mask, test_case.dd_mask)
		}
	}
}
<|endoftext|>


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

package walk

import (
	"syscall"
	"unsafe"
)

import . "github.com/lxn/go-winapi"

type treeViewItemInfo struct {
	handle       HTREEITEM
	child2Handle map[TreeItem]HTREEITEM
}

type TreeView struct {
	WidgetBase
	model                         TreeModel
	lazyPopulation                bool
	itemsResetEventHandlerHandle  int
	itemChangedEventHandlerHandle int
	item2Info                     map[TreeItem]*treeViewItemInfo
	handle2Item                   map[HTREEITEM]TreeItem
	currItem                      TreeItem
	hIml                          HIMAGELIST
	usingSysIml                   bool
	imageUintptr2Index            map[uintptr]int32
	filePath2IconIndex            map[string]int32
	itemCollapsedPublisher        TreeItemEventPublisher
	itemExpandedPublisher         TreeItemEventPublisher
	currentItemChangedPublisher   EventPublisher
}

func NewTreeView(parent Container) (*TreeView, error) {
	tv := new(TreeView)

	if err := InitChildWidget(
		tv,
		parent,
		"SysTreeView32",
		WS_TABSTOP|WS_VISIBLE|TVS_HASBUTTONS|TVS_HASLINES|TVS_LINESATROOT|TVS_SHOWSELALWAYS,
		WS_EX_CLIENTEDGE); err != nil {
		return nil, err
	}

	succeeded := false
	defer func() {
		if !succeeded {
			tv.Dispose()
		}
	}()

	if err := tv.setTheme("Explorer"); err != nil {
		return nil, err
	}

	succeeded = true

	return tv, nil
}

func (*TreeView) LayoutFlags() LayoutFlags {
	return ShrinkableHorz | ShrinkableVert | GrowableHorz | GrowableVert | GreedyHorz | GreedyVert
}

func (tv *TreeView) SizeHint() Size {
	return tv.dialogBaseUnitsToPixels(Size{100, 100})
}

func (tv *TreeView) Dispose() {
	tv.WidgetBase.Dispose()

	tv.disposeImageListAndCaches()
}

func (tv *TreeView) Model() TreeModel {
	return tv.model
}

func (tv *TreeView) SetModel(model TreeModel) error {
	if tv.model != nil {
		tv.model.ItemsReset().Detach(tv.itemsResetEventHandlerHandle)
		tv.model.ItemChanged().Detach(tv.itemChangedEventHandlerHandle)

		tv.disposeImageListAndCaches()
	}

	tv.model = model

	if model != nil {
		tv.lazyPopulation = model.LazyPopulation()

		tv.itemsResetEventHandlerHandle = model.ItemsReset().Attach(func(parent TreeItem) {
			if parent == nil {
				tv.resetItems()
			} else if tv.item2Info[parent] != nil {
				if err := tv.removeDescendants(parent); err != nil {
					return
				}

				if err := tv.insertChildren(parent); err != nil {
					return
				}
			}
		})

		tv.itemChangedEventHandlerHandle = model.ItemChanged().Attach(func(item TreeItem) {
			if item == nil || tv.item2Info[item] == nil {
				return
			}

			if err := tv.updateItem(item); err != nil {
				return
			}
		})
	}

	return tv.resetItems()
}

func (tv *TreeView) CurrentItem() TreeItem {
	return tv.currItem
}

func (tv *TreeView) SetCurrentItem(item TreeItem) error {
	if item == tv.currItem {
		return nil
	}

	var handle HTREEITEM
	if item != nil {
		if info := tv.item2Info[item]; info == nil {
			return newError("invalid item")
		} else {
			handle = info.handle
		}
	}

	if 0 == tv.SendMessage(TVM_SELECTITEM, TVGN_CARET, uintptr(handle)) {
		return newError("SendMessage(TVM_SELECTITEM) failed")
	}

	tv.currItem = item

	return nil
}

func (tv *TreeView) ItemAt(x, y int) TreeItem {
	hti := TVHITTESTINFO{Pt: POINT{int32(x), int32(y)}}

	tv.SendMessage(TVM_HITTEST, 0, uintptr(unsafe.Pointer(&hti)))

	if item, ok := tv.handle2Item[hti.HItem]; ok {
		return item
	}

	return nil
}

func (tv *TreeView) resetItems() error {
	tv.SetSuspended(true)
	defer tv.SetSuspended(false)

	if err := tv.clearItems(); err != nil {
		return err
	}

	if tv.model == nil {
		return nil
	}

	if err := tv.insertRoots(); err != nil {
		return err
	}

	return nil
}

func (tv *TreeView) clearItems() error {
	if 0 == tv.SendMessage(TVM_DELETEITEM, 0, 0) {
		return newError("SendMessage(TVM_DELETEITEM) failed")
	}

	tv.item2Info = make(map[TreeItem]*treeViewItemInfo)
	tv.handle2Item = make(map[HTREEITEM]TreeItem)

	return nil
}

func (tv *TreeView) insertRoots() error {
	count := tv.model.RootCount()

	for i := 0; i < count; i++ {
		if _, err := tv.insertItem(i, tv.model.RootAt(i)); err != nil {
			return err
		}
	}

	return nil
}

func (tv *TreeView) applyImageListForImage(image interface{}) {
	tv.hIml, tv.usingSysIml, _ = imageListForImage(image)

	tv.SendMessage(TVM_SETIMAGELIST, 0, uintptr(tv.hIml))

	tv.imageUintptr2Index = make(map[uintptr]int32)
	tv.filePath2IconIndex = make(map[string]int32)
}

func (tv *TreeView) disposeImageListAndCaches() {
	if tv.hIml != 0 && !tv.usingSysIml {
		ImageList_Destroy(tv.hIml)
	}
	tv.hIml = 0

	tv.imageUintptr2Index = nil
	tv.filePath2IconIndex = nil
}

func (tv *TreeView) setTVITEMImageInfo(tvi *TVITEM, item TreeItem) {
	if imager, ok := item.(Imager); ok {
		if tv.hIml == 0 {
			tv.applyImageListForImage(imager.Image())
		}

		// FIXME: If not setting TVIF_SELECTEDIMAGE and tvi.ISelectedImage,
		// some default icon will show up, even though we have not asked for it.

		tvi.Mask |= TVIF_IMAGE | TVIF_SELECTEDIMAGE
		tvi.IImage = imageIndexMaybeAdd(
			imager.Image(),
			tv.hIml,
			tv.usingSysIml,
			tv.imageUintptr2Index,
			tv.filePath2IconIndex)

		tvi.ISelectedImage = tvi.IImage
	}
}

func (tv *TreeView) insertItem(index int, item TreeItem) (HTREEITEM, error) {
	var tvins TVINSERTSTRUCT
	tvi := &tvins.Item

	tvi.Mask = TVIF_CHILDREN | TVIF_TEXT
	tvi.PszText = syscall.StringToUTF16Ptr(item.Text())
	tvi.CChildren = I_CHILDRENCALLBACK

	tv.setTVITEMImageInfo(tvi, item)

	parent := item.Parent()

	if parent == nil {
		tvins.HParent = TVI_ROOT
	} else {
		info := tv.item2Info[parent]
		if info == nil {
			return 0, newError("invalid parent")
		}
		tvins.HParent = info.handle
	}

	if index == 0 {
		tvins.HInsertAfter = TVI_LAST
	} else {
		var prevItem TreeItem
		if parent == nil {
			prevItem = tv.model.RootAt(index - 1)
		} else {
			prevItem = parent.ChildAt(index - 1)
		}
		info := tv.item2Info[prevItem]
		if info == nil {
			return 0, newError("invalid prev item")
		}
		tvins.HInsertAfter = info.handle
	}

	hItem := HTREEITEM(tv.SendMessage(TVM_INSERTITEM, 0, uintptr(unsafe.Pointer(&tvins))))
	if hItem == 0 {
		return 0, newError("TVM_INSERTITEM failed")
	}
	tv.item2Info[item] = &treeViewItemInfo{hItem, make(map[TreeItem]HTREEITEM)}
	tv.handle2Item[hItem] = item

	if !tv.lazyPopulation {
		if err := tv.insertChildren(item); err != nil {
			return 0, err
		}
	}

	return hItem, nil
}

func (tv *TreeView) insertChildren(parent TreeItem) error {
	info := tv.item2Info[parent]

	count := parent.ChildCount()
	for i := 0; i < count; i++ {
		child := parent.ChildAt(i)

		if handle, err := tv.insertItem(i, child); err != nil {
			return err
		} else {
			info.child2Handle[child] = handle
		}
	}

	return nil
}

func (tv *TreeView) updateItem(item TreeItem) error {
	tvi := &TVITEM{
		Mask:    TVIF_TEXT,
		HItem:   tv.item2Info[item].handle,
		PszText: syscall.StringToUTF16Ptr(item.Text()),
	}

	tv.setTVITEMImageInfo(tvi, item)

	if 0 == tv.SendMessage(TVM_SETITEM, 0, uintptr(unsafe.Pointer(tvi))) {
		return newError("SendMessage(TVM_SETITEM) failed")
	}

	return nil
}

func (tv *TreeView) removeItem(item TreeItem) error {
	if err := tv.removeDescendants(item); err != nil {
		return err
	}

	info := tv.item2Info[item]
	if info == nil {
		return newError("invalid item")
	}

	if 0 == tv.SendMessage(TVM_DELETEITEM, 0, uintptr(info.handle)) {
		return newError("SendMessage(TVM_DELETEITEM) failed")
	}

	if parentInfo := tv.item2Info[item.Parent()]; parentInfo != nil {
		delete(parentInfo.child2Handle, item)
	}
	delete(tv.item2Info, item)
	delete(tv.handle2Item, info.handle)

	return nil
}

func (tv *TreeView) removeDescendants(parent TreeItem) error {
	for item, _ := range tv.item2Info[parent].child2Handle {
		if err := tv.removeItem(item); err != nil {
			return err
		}
	}

	return nil
}

func (tv *TreeView) ItemCollapsed() *TreeItemEvent {
	return tv.itemCollapsedPublisher.Event()
}

func (tv *TreeView) ItemExpanded() *TreeItemEvent {
	return tv.itemExpandedPublisher.Event()
}

func (tv *TreeView) CurrentItemChanged() *Event {
	return tv.currentItemChangedPublisher.Event()
}

func (tv *TreeView) WndProc(hwnd HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case WM_NOTIFY:
		nmhdr := (*NMHDR)(unsafe.Pointer(lParam))

		switch nmhdr.Code {
		case TVN_GETDISPINFO:
			nmtvdi := (*NMTVDISPINFO)(unsafe.Pointer(lParam))
			item := tv.handle2Item[nmtvdi.Item.HItem]

			if nmtvdi.Item.Mask&TVIF_CHILDREN != 0 {
				nmtvdi.Item.CChildren = int32(item.ChildCount())
			}

		case TVN_ITEMEXPANDING:
			nmtv := (*NMTREEVIEW)(unsafe.Pointer(lParam))
			item := tv.handle2Item[nmtv.ItemNew.HItem]

			if nmtv.Action == TVE_EXPAND && tv.lazyPopulation {
				info := tv.item2Info[item]
				if len(info.child2Handle) == 0 {
					tv.insertChildren(item)
				}
			}

		case TVN_ITEMEXPANDED:
			nmtv := (*NMTREEVIEW)(unsafe.Pointer(lParam))
			item := tv.handle2Item[nmtv.ItemNew.HItem]

			switch nmtv.Action {
			case TVE_COLLAPSE:
				tv.itemCollapsedPublisher.Publish(item)

			case TVE_COLLAPSERESET:

			case TVE_EXPAND:
				tv.itemExpandedPublisher.Publish(item)

			case TVE_EXPANDPARTIAL:

			case TVE_TOGGLE:
			}

		case TVN_SELCHANGED:
			nmtv := (*NMTREEVIEW)(unsafe.Pointer(lParam))

			tv.currItem = tv.handle2Item[nmtv.ItemNew.HItem]

			tv.currentItemChangedPublisher.Publish()
		}
	}

	return tv.WidgetBase.WndProc(hwnd, msg, wParam, lParam)
}
<commit_msg>TreeView: Provide item text via callback<commit_after>// Copyright 2010 The Walk Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package walk

import (
	"syscall"
	"unsafe"
)

import . "github.com/lxn/go-winapi"

type treeViewItemInfo struct {
	handle       HTREEITEM
	child2Handle map[TreeItem]HTREEITEM
}

type TreeView struct {
	WidgetBase
	model                         TreeModel
	lazyPopulation                bool
	itemsResetEventHandlerHandle  int
	itemChangedEventHandlerHandle int
	item2Info                     map[TreeItem]*treeViewItemInfo
	handle2Item                   map[HTREEITEM]TreeItem
	currItem                      TreeItem
	hIml                          HIMAGELIST
	usingSysIml                   bool
	imageUintptr2Index            map[uintptr]int32
	filePath2IconIndex            map[string]int32
	itemCollapsedPublisher        TreeItemEventPublisher
	itemExpandedPublisher         TreeItemEventPublisher
	currentItemChangedPublisher   EventPublisher
}

func NewTreeView(parent Container) (*TreeView, error) {
	tv := new(TreeView)

	if err := InitChildWidget(
		tv,
		parent,
		"SysTreeView32",
		WS_TABSTOP|WS_VISIBLE|TVS_HASBUTTONS|TVS_HASLINES|TVS_LINESATROOT|TVS_SHOWSELALWAYS,
		WS_EX_CLIENTEDGE); err != nil {
		return nil, err
	}

	succeeded := false
	defer func() {
		if !succeeded {
			tv.Dispose()
		}
	}()

	if err := tv.setTheme("Explorer"); err != nil {
		return nil, err
	}

	succeeded = true

	return tv, nil
}

func (*TreeView) LayoutFlags() LayoutFlags {
	return ShrinkableHorz | ShrinkableVert | GrowableHorz | GrowableVert | GreedyHorz | GreedyVert
}

func (tv *TreeView) SizeHint() Size {
	return tv.dialogBaseUnitsToPixels(Size{100, 100})
}

func (tv *TreeView) Dispose() {
	tv.WidgetBase.Dispose()

	tv.disposeImageListAndCaches()
}

func (tv *TreeView) Model() TreeModel {
	return tv.model
}

func (tv *TreeView) SetModel(model TreeModel) error {
	if tv.model != nil {
		tv.model.ItemsReset().Detach(tv.itemsResetEventHandlerHandle)
		tv.model.ItemChanged().Detach(tv.itemChangedEventHandlerHandle)

		tv.disposeImageListAndCaches()
	}

	tv.model = model

	if model != nil {
		tv.lazyPopulation = model.LazyPopulation()

		tv.itemsResetEventHandlerHandle = model.ItemsReset().Attach(func(parent TreeItem) {
			if parent == nil {
				tv.resetItems()
			} else if tv.item2Info[parent] != nil {
				if err := tv.removeDescendants(parent); err != nil {
					return
				}

				if err := tv.insertChildren(parent); err != nil {
					return
				}
			}
		})

		tv.itemChangedEventHandlerHandle = model.ItemChanged().Attach(func(item TreeItem) {
			if item == nil || tv.item2Info[item] == nil {
				return
			}

			if err := tv.updateItem(item); err != nil {
				return
			}
		})
	}

	return tv.resetItems()
}

func (tv *TreeView) CurrentItem() TreeItem {
	return tv.currItem
}

func (tv *TreeView) SetCurrentItem(item TreeItem) error {
	if item == tv.currItem {
		return nil
	}

	var handle HTREEITEM
	if item != nil {
		if info := tv.item2Info[item]; info == nil {
			return newError("invalid item")
		} else {
			handle = info.handle
		}
	}

	if 0 == tv.SendMessage(TVM_SELECTITEM, TVGN_CARET, uintptr(handle)) {
		return newError("SendMessage(TVM_SELECTITEM) failed")
	}

	tv.currItem = item

	return nil
}

func (tv *TreeView) ItemAt(x, y int) TreeItem {
	hti := TVHITTESTINFO{Pt: POINT{int32(x), int32(y)}}

	tv.SendMessage(TVM_HITTEST, 0, uintptr(unsafe.Pointer(&hti)))

	if item, ok := tv.handle2Item[hti.HItem]; ok {
		return item
	}

	return nil
}

func (tv *TreeView) resetItems() error {
	tv.SetSuspended(true)
	defer tv.SetSuspended(false)

	if err := tv.clearItems(); err != nil {
		return err
	}

	if tv.model == nil {
		return nil
	}

	if err := tv.insertRoots(); err != nil {
		return err
	}

	return nil
}

func (tv *TreeView) clearItems() error {
	if 0 == tv.SendMessage(TVM_DELETEITEM, 0, 0) {
		return newError("SendMessage(TVM_DELETEITEM) failed")
	}

	tv.item2Info = make(map[TreeItem]*treeViewItemInfo)
	tv.handle2Item = make(map[HTREEITEM]TreeItem)

	return nil
}

func (tv *TreeView) insertRoots() error {
	count := tv.model.RootCount()

	for i := 0; i < count; i++ {
		if _, err := tv.insertItem(i, tv.model.RootAt(i)); err != nil {
			return err
		}
	}

	return nil
}

func (tv *TreeView) applyImageListForImage(image interface{}) {
	tv.hIml, tv.usingSysIml, _ = imageListForImage(image)

	tv.SendMessage(TVM_SETIMAGELIST, 0, uintptr(tv.hIml))

	tv.imageUintptr2Index = make(map[uintptr]int32)
	tv.filePath2IconIndex = make(map[string]int32)
}

func (tv *TreeView) disposeImageListAndCaches() {
	if tv.hIml != 0 && !tv.usingSysIml {
		ImageList_Destroy(tv.hIml)
	}
	tv.hIml = 0

	tv.imageUintptr2Index = nil
	tv.filePath2IconIndex = nil
}

func (tv *TreeView) setTVITEMImageInfo(tvi *TVITEM, item TreeItem) {
	if imager, ok := item.(Imager); ok {
		if tv.hIml == 0 {
			tv.applyImageListForImage(imager.Image())
		}

		// FIXME: If not setting TVIF_SELECTEDIMAGE and tvi.ISelectedImage,
		// some default icon will show up, even though we have not asked for it.

		tvi.Mask |= TVIF_IMAGE | TVIF_SELECTEDIMAGE
		tvi.IImage = imageIndexMaybeAdd(
			imager.Image(),
			tv.hIml,
			tv.usingSysIml,
			tv.imageUintptr2Index,
			tv.filePath2IconIndex)

		tvi.ISelectedImage = tvi.IImage
	}
}

func (tv *TreeView) insertItem(index int, item TreeItem) (HTREEITEM, error) {
	var tvins TVINSERTSTRUCT
	tvi := &tvins.Item

	tvi.Mask = TVIF_CHILDREN | TVIF_TEXT
	tvi.PszText = LPSTR_TEXTCALLBACK
	tvi.CChildren = I_CHILDRENCALLBACK

	tv.setTVITEMImageInfo(tvi, item)

	parent := item.Parent()

	if parent == nil {
		tvins.HParent = TVI_ROOT
	} else {
		info := tv.item2Info[parent]
		if info == nil {
			return 0, newError("invalid parent")
		}
		tvins.HParent = info.handle
	}

	if index == 0 {
		tvins.HInsertAfter = TVI_LAST
	} else {
		var prevItem TreeItem
		if parent == nil {
			prevItem = tv.model.RootAt(index - 1)
		} else {
			prevItem = parent.ChildAt(index - 1)
		}
		info := tv.item2Info[prevItem]
		if info == nil {
			return 0, newError("invalid prev item")
		}
		tvins.HInsertAfter = info.handle
	}

	hItem := HTREEITEM(tv.SendMessage(TVM_INSERTITEM, 0, uintptr(unsafe.Pointer(&tvins))))
	if hItem == 0 {
		return 0, newError("TVM_INSERTITEM failed")
	}
	tv.item2Info[item] = &treeViewItemInfo{hItem, make(map[TreeItem]HTREEITEM)}
	tv.handle2Item[hItem] = item

	if !tv.lazyPopulation {
		if err := tv.insertChildren(item); err != nil {
			return 0, err
		}
	}

	return hItem, nil
}

func (tv *TreeView) insertChildren(parent TreeItem) error {
	info := tv.item2Info[parent]

	count := parent.ChildCount()
	for i := 0; i < count; i++ {
		child := parent.ChildAt(i)

		if handle, err := tv.insertItem(i, child); err != nil {
			return err
		} else {
			info.child2Handle[child] = handle
		}
	}

	return nil
}

func (tv *TreeView) updateItem(item TreeItem) error {
	tvi := &TVITEM{
		Mask:    TVIF_TEXT,
		HItem:   tv.item2Info[item].handle,
		PszText: LPSTR_TEXTCALLBACK,
	}

	tv.setTVITEMImageInfo(tvi, item)

	if 0 == tv.SendMessage(TVM_SETITEM, 0, uintptr(unsafe.Pointer(tvi))) {
		return newError("SendMessage(TVM_SETITEM) failed")
	}

	return nil
}

func (tv *TreeView) removeItem(item TreeItem) error {
	if err := tv.removeDescendants(item); err != nil {
		return err
	}

	info := tv.item2Info[item]
	if info == nil {
		return newError("invalid item")
	}

	if 0 == tv.SendMessage(TVM_DELETEITEM, 0, uintptr(info.handle)) {
		return newError("SendMessage(TVM_DELETEITEM) failed")
	}

	if parentInfo := tv.item2Info[item.Parent()]; parentInfo != nil {
		delete(parentInfo.child2Handle, item)
	}
	delete(tv.item2Info, item)
	delete(tv.handle2Item, info.handle)

	return nil
}

func (tv *TreeView) removeDescendants(parent TreeItem) error {
	for item, _ := range tv.item2Info[parent].child2Handle {
		if err := tv.removeItem(item); err != nil {
			return err
		}
	}

	return nil
}

func (tv *TreeView) ItemCollapsed() *TreeItemEvent {
	return tv.itemCollapsedPublisher.Event()
}

func (tv *TreeView) ItemExpanded() *TreeItemEvent {
	return tv.itemExpandedPublisher.Event()
}

func (tv *TreeView) CurrentItemChanged() *Event {
	return tv.currentItemChangedPublisher.Event()
}

func (tv *TreeView) WndProc(hwnd HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case WM_NOTIFY:
		nmhdr := (*NMHDR)(unsafe.Pointer(lParam))

		switch nmhdr.Code {
		case TVN_GETDISPINFO:
			nmtvdi := (*NMTVDISPINFO)(unsafe.Pointer(lParam))
			item := tv.handle2Item[nmtvdi.Item.HItem]

			if nmtvdi.Item.Mask&TVIF_TEXT != 0 {
				nmtvdi.Item.PszText = uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(item.Text())))
			}
			if nmtvdi.Item.Mask&TVIF_CHILDREN != 0 {
				nmtvdi.Item.CChildren = int32(item.ChildCount())
			}

		case TVN_ITEMEXPANDING:
			nmtv := (*NMTREEVIEW)(unsafe.Pointer(lParam))
			item := tv.handle2Item[nmtv.ItemNew.HItem]

			if nmtv.Action == TVE_EXPAND && tv.lazyPopulation {
				info := tv.item2Info[item]
				if len(info.child2Handle) == 0 {
					tv.insertChildren(item)
				}
			}

		case TVN_ITEMEXPANDED:
			nmtv := (*NMTREEVIEW)(unsafe.Pointer(lParam))
			item := tv.handle2Item[nmtv.ItemNew.HItem]

			switch nmtv.Action {
			case TVE_COLLAPSE:
				tv.itemCollapsedPublisher.Publish(item)

			case TVE_COLLAPSERESET:

			case TVE_EXPAND:
				tv.itemExpandedPublisher.Publish(item)

			case TVE_EXPANDPARTIAL:

			case TVE_TOGGLE:
			}

		case TVN_SELCHANGED:
			nmtv := (*NMTREEVIEW)(unsafe.Pointer(lParam))

			tv.currItem = tv.handle2Item[nmtv.ItemNew.HItem]

			tv.currentItemChangedPublisher.Publish()
		}
	}

	return tv.WidgetBase.WndProc(hwnd, msg, wParam, lParam)
}
<|endoftext|>


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

package utils_test

import (
	"fmt"

	gc "launchpad.net/gocheck"

	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/testing/testbase"
	"launchpad.net/juju-core/utils"
)

type AptSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&AptSuite{})

func (s *AptSuite) TestOnePackage(c *gc.C) {
	cmdChan, cleanup := testbase.HookCommandOutput(&utils.AptCommandOutput, []byte{}, nil)
	defer cleanup()
	err := utils.AptGetInstall("test-package")
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-get", "--option=Dpkg::Options::=--force-confold",
		"--option=Dpkg::options::=--force-unsafe-io", "--assume-yes", "--quiet",
		"install", "test-package",
	})
	c.Assert(cmd.Env[len(cmd.Env)-1], gc.Equals, "DEBIAN_FRONTEND=noninteractive")
}

func (s *AptSuite) TestAptGetError(c *gc.C) {
	const expected = `E: frobnicator failure detected`
	cmdError := fmt.Errorf("error")
	cmdExpectedError := fmt.Errorf("apt-get failed: error")
	cmdChan, cleanup := testbase.HookCommandOutput(&utils.AptCommandOutput, []byte(expected), cmdError)
	defer cleanup()
	err := utils.AptGetInstall("foo")
	c.Assert(err, gc.DeepEquals, cmdExpectedError)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-get", "--option=Dpkg::Options::=--force-confold",
		"--option=Dpkg::options::=--force-unsafe-io", "--assume-yes", "--quiet",
		"install", "foo",
	})
}

func (s *AptSuite) TestConfigProxyEmpty(c *gc.C) {
	cmdChan, cleanup := testbase.HookCommandOutput(&utils.AptCommandOutput, []byte{}, nil)
	defer cleanup()
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, "")
}

func (s *AptSuite) TestConfigProxyConfigured(c *gc.C) {
	const expected = `Acquire::http::Proxy "10.0.3.1:3142";
Acquire::https::Proxy "false";`
	cmdChan, cleanup := testbase.HookCommandOutput(&utils.AptCommandOutput, []byte(expected), nil)
	defer cleanup()
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, expected)
}

func (s *AptSuite) TestConfigProxyConfiguredFilterOutput(c *gc.C) {
	const (
		output = `CommandLine::AsString "apt-config dump";
Acquire::http::Proxy  "10.0.3.1:3142";
Acquire::https::Proxy "false";`
		expected = `Acquire::http::Proxy  "10.0.3.1:3142";
Acquire::https::Proxy "false";`
	)
	cmdChan, cleanup := testbase.HookCommandOutput(&utils.AptCommandOutput, []byte(output), nil)
	defer cleanup()
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, expected)
}

func (s *AptSuite) TestConfigProxyError(c *gc.C) {
	const expected = `E: frobnicator failure detected`
	cmdError := fmt.Errorf("error")
	cmdExpectedError := fmt.Errorf("apt-config failed: error")
	cmdChan, cleanup := testbase.HookCommandOutput(&utils.AptCommandOutput, []byte(expected), cmdError)
	defer cleanup()
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.DeepEquals, cmdExpectedError)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, "")
}

func (s *AptSuite) patchLsbRelease(c *gc.C, name string) {
	content := fmt.Sprintf("#!/bin/bash --norc\necho %s", name)
	patchExecutable(s, c.MkDir(), "lsb_release", content)
}

func (s *AptSuite) TestIsUbuntu(c *gc.C) {
	s.patchLsbRelease(c, "Ubuntu")
	c.Assert(utils.IsUbuntu(), jc.IsTrue)
}

func (s *AptSuite) TestIsNotUbuntu(c *gc.C) {
	s.patchLsbRelease(c, "Windows NT")
	c.Assert(utils.IsUbuntu(), jc.IsFalse)
}

func (s *AptSuite) patchDpkgQuery(c *gc.C, installed bool) {
	rc := 0
	if !installed {
		rc = 1
	}
	content := fmt.Sprintf("#!/bin/bash --norc\nexit %v", rc)
	patchExecutable(s, c.MkDir(), "dpkg-query", content)
}

func (s *AptSuite) TestIsPackageInstalled(c *gc.C) {
	s.patchDpkgQuery(c, true)
	c.Assert(utils.IsPackageInstalled("foo-bar"), jc.IsTrue)
}

func (s *AptSuite) TestIsPackageNotInstalled(c *gc.C) {
	s.patchDpkgQuery(c, false)
	c.Assert(utils.IsPackageInstalled("foo-bar"), jc.IsFalse)
}
<commit_msg>Use the cleanup suite HookCommandOutput<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils_test

import (
	"fmt"

	gc "launchpad.net/gocheck"

	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/testing/testbase"
	"launchpad.net/juju-core/utils"
)

type AptSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&AptSuite{})

func (s *AptSuite) TestOnePackage(c *gc.C) {
	cmdChan := s.HookCommandOutput(&utils.AptCommandOutput, []byte{}, nil)
	err := utils.AptGetInstall("test-package")
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-get", "--option=Dpkg::Options::=--force-confold",
		"--option=Dpkg::options::=--force-unsafe-io", "--assume-yes", "--quiet",
		"install", "test-package",
	})
	c.Assert(cmd.Env[len(cmd.Env)-1], gc.Equals, "DEBIAN_FRONTEND=noninteractive")
}

func (s *AptSuite) TestAptGetError(c *gc.C) {
	const expected = `E: frobnicator failure detected`
	cmdError := fmt.Errorf("error")
	cmdExpectedError := fmt.Errorf("apt-get failed: error")
	cmdChan := s.HookCommandOutput(&utils.AptCommandOutput, []byte(expected), cmdError)
	err := utils.AptGetInstall("foo")
	c.Assert(err, gc.DeepEquals, cmdExpectedError)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-get", "--option=Dpkg::Options::=--force-confold",
		"--option=Dpkg::options::=--force-unsafe-io", "--assume-yes", "--quiet",
		"install", "foo",
	})
}

func (s *AptSuite) TestConfigProxyEmpty(c *gc.C) {
	cmdChan := s.HookCommandOutput(&utils.AptCommandOutput, []byte{}, nil)
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, "")
}

func (s *AptSuite) TestConfigProxyConfigured(c *gc.C) {
	const expected = `Acquire::http::Proxy "10.0.3.1:3142";
Acquire::https::Proxy "false";`
	cmdChan := s.HookCommandOutput(&utils.AptCommandOutput, []byte(expected), nil)
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, expected)
}

func (s *AptSuite) TestConfigProxyConfiguredFilterOutput(c *gc.C) {
	const (
		output = `CommandLine::AsString "apt-config dump";
Acquire::http::Proxy  "10.0.3.1:3142";
Acquire::https::Proxy "false";`
		expected = `Acquire::http::Proxy  "10.0.3.1:3142";
Acquire::https::Proxy "false";`
	)
	cmdChan := s.HookCommandOutput(&utils.AptCommandOutput, []byte(output), nil)
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.IsNil)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, expected)
}

func (s *AptSuite) TestConfigProxyError(c *gc.C) {
	const expected = `E: frobnicator failure detected`
	cmdError := fmt.Errorf("error")
	cmdExpectedError := fmt.Errorf("apt-config failed: error")
	cmdChan := s.HookCommandOutput(&utils.AptCommandOutput, []byte(expected), cmdError)
	out, err := utils.AptConfigProxy()
	c.Assert(err, gc.DeepEquals, cmdExpectedError)
	cmd := <-cmdChan
	c.Assert(cmd.Args, gc.DeepEquals, []string{
		"apt-config", "dump", "Acquire::http::Proxy",
		"Acquire::https::Proxy", "Acquire::ftp::Proxy",
	})
	c.Assert(out, gc.Equals, "")
}

func (s *AptSuite) patchLsbRelease(c *gc.C, name string) {
	content := fmt.Sprintf("#!/bin/bash --norc\necho %s", name)
	patchExecutable(s, c.MkDir(), "lsb_release", content)
}

func (s *AptSuite) TestIsUbuntu(c *gc.C) {
	s.patchLsbRelease(c, "Ubuntu")
	c.Assert(utils.IsUbuntu(), jc.IsTrue)
}

func (s *AptSuite) TestIsNotUbuntu(c *gc.C) {
	s.patchLsbRelease(c, "Windows NT")
	c.Assert(utils.IsUbuntu(), jc.IsFalse)
}

func (s *AptSuite) patchDpkgQuery(c *gc.C, installed bool) {
	rc := 0
	if !installed {
		rc = 1
	}
	content := fmt.Sprintf("#!/bin/bash --norc\nexit %v", rc)
	patchExecutable(s, c.MkDir(), "dpkg-query", content)
}

func (s *AptSuite) TestIsPackageInstalled(c *gc.C) {
	s.patchDpkgQuery(c, true)
	c.Assert(utils.IsPackageInstalled("foo-bar"), jc.IsTrue)
}

func (s *AptSuite) TestIsPackageNotInstalled(c *gc.C) {
	s.patchDpkgQuery(c, false)
	c.Assert(utils.IsPackageInstalled("foo-bar"), jc.IsFalse)
}
<|endoftext|>


<commit_before>package biz

import (
	"crypto/tls"
	"fmt"
	"net"
	"time"

	"github.com/cuigh/auxo/data/guid"
	"github.com/cuigh/auxo/errors"
	"github.com/cuigh/auxo/log"
	"github.com/cuigh/auxo/net/web"
	"github.com/cuigh/auxo/security/password"
	"github.com/cuigh/swirl/dao"
	"github.com/cuigh/swirl/model"
	"github.com/go-ldap/ldap"
)

var ErrIncorrectAuth = errors.New("login name or password is incorrect")

var User = &userBiz{}

type userBiz struct {
}

func (b *userBiz) GetByID(id string) (user *model.User, err error) {
	do(func(d dao.Interface) {
		user, err = d.UserGetByID(id)
	})
	return
}

func (b *userBiz) GetByName(loginName string) (user *model.User, err error) {
	do(func(d dao.Interface) {
		user, err = d.UserGetByName(loginName)
	})
	return
}

func (b *userBiz) Create(user *model.User, ctxUser web.User) (err error) {
	user.ID = guid.New()
	user.Status = model.UserStatusActive
	user.CreatedAt = time.Now()
	user.UpdatedAt = user.CreatedAt
	if user.Type == model.UserTypeInternal {
		user.Password, user.Salt, err = password.Generate(user.Password)
		if err != nil {
			return
		}
	}

	do(func(d dao.Interface) {
		if err = d.UserCreate(user); err == nil && ctxUser != nil {
			Event.CreateUser(model.EventActionCreate, user.LoginName, user.Name, ctxUser)
		}
	})
	return
}

func (b *userBiz) Update(user *model.User, ctxUser web.User) (err error) {
	do(func(d dao.Interface) {
		user.UpdatedAt = time.Now()
		if err = d.UserUpdate(user); err == nil {
			Event.CreateUser(model.EventActionUpdate, user.LoginName, user.Name, ctxUser)
		}
	})
	return
}

func (b *userBiz) Block(id string) (err error) {
	do(func(d dao.Interface) {
		err = d.UserBlock(id, true)
	})
	return
}

func (b *userBiz) Unblock(id string) (err error) {
	do(func(d dao.Interface) {
		err = d.UserBlock(id, false)
	})
	return
}

func (b *userBiz) Delete(id string) (err error) {
	do(func(d dao.Interface) {
		err = d.UserDelete(id)
	})
	return
}

func (b *userBiz) UpdateInfo(user *model.User) (err error) {
	do(func(d dao.Interface) {
		err = d.ProfileUpdateInfo(user)
	})
	return
}

func (b *userBiz) UpdatePassword(id, oldPwd, newPwd string) (err error) {
	do(func(d dao.Interface) {
		var (
			user      *model.User
			pwd, salt string
		)

		user, err = d.UserGetByID(id)
		if err != nil {
			return
		}

		if !password.Validate(user.Password, oldPwd, user.Salt) {
			err = errors.New("Current password is incorrect")
			return
		}

		pwd, salt, err = password.Generate(newPwd)
		if err != nil {
			return
		}

		err = d.ProfileUpdatePassword(id, pwd, salt)
	})
	return
}

func (b *userBiz) List(args *model.UserListArgs) (users []*model.User, count int, err error) {
	do(func(d dao.Interface) {
		users, count, err = d.UserList(args)
	})
	return
}

func (b *userBiz) Count() (count int, err error) {
	do(func(d dao.Interface) {
		count, err = d.UserCount()
	})
	return
}

func (b *userBiz) Login(name, pwd string) (token string, err error) {
	do(func(d dao.Interface) {
		var (
			user *model.User
		)

		user, err = d.UserGetByName(name)
		if err != nil {
			return
		}

		if user == nil {
			user = &model.User{
				Type:      model.UserTypeLDAP,
				LoginName: name,
			}
			err = b.loginLDAP(d, user, pwd)
		} else {
			if user.Status == model.UserStatusBlocked {
				err = fmt.Errorf("user %s is blocked", name)
				return
			}

			if user.Type == model.UserTypeInternal {
				err = b.loginInternal(user, pwd)
			} else {
				err = b.loginLDAP(d, user, pwd)
			}
		}

		if err != nil {
			return
		}

		session := &model.Session{
			UserID:    user.ID,
			Token:     guid.New(),
			UpdatedAt: time.Now(),
		}
		session.Expires = session.UpdatedAt.Add(time.Hour * 24)
		err = d.SessionUpdate(session)
		if err != nil {
			return
		}
		token = session.Token

		// create event
		Event.CreateAuthentication(model.EventActionLogin, user.ID, user.LoginName, user.Name)
	})
	return
}

func (b *userBiz) loginInternal(user *model.User, pwd string) error {
	if !password.Validate(user.Password, pwd, user.Salt) {
		return ErrIncorrectAuth
	}
	return nil
}

func (b *userBiz) loginLDAP(d dao.Interface, user *model.User, pwd string) error {
	setting, err := Setting.Get()
	if err != nil {
		return err
	}

	if !setting.LDAP.Enabled {
		return ErrIncorrectAuth
	}

	l, err := b.ldapDial(setting)
	if err != nil {
		return err
	}
	defer l.Close()

	// bind
	err = b.ldapBind(setting, l, user, pwd)
	if err != nil {
		log.Get("user").Error("LDAP > Failed to bind: ", err)
		return ErrIncorrectAuth
	}

	// Stop here for an exist user because we only need validate password.
	if user.ID != "" {
		return nil
	}

	// If user wasn't exist, we need create it
	entry, err := b.ldapSearchOne(setting, l, user.LoginName, setting.LDAP.NameAttr, setting.LDAP.EmailAttr)
	if err != nil {
		return err
	}

	user.Email = entry.GetAttributeValue(setting.LDAP.EmailAttr)
	user.Name = entry.GetAttributeValue(setting.LDAP.NameAttr)
	return b.Create(user, nil)
}

func (b *userBiz) ldapDial(setting *model.Setting) (*ldap.Conn, error) {
	host, _, err := net.SplitHostPort(setting.LDAP.Address)
	if err != nil {
		return nil, err
	}

	// TODO: support tls cert and verification
	tc := &tls.Config{
		ServerName:         host,
		InsecureSkipVerify: true,
		Certificates:       nil,
	}

	if setting.LDAP.Security == model.LDAPSecurityTLS {
		return ldap.DialTLS("tcp", setting.LDAP.Address, tc)
	}

	conn, err := ldap.Dial("tcp", setting.LDAP.Address)
	if err != nil {
		return nil, err
	}

	if setting.LDAP.Security == model.LDAPSecurityStartTLS {
		if err = conn.StartTLS(tc); err != nil {
			conn.Close()
			log.Get("user").Error("LDAP > Failed to switch to TLS: ", err)
			return nil, err
		}
	}
	return conn, nil
}

func (b *userBiz) ldapBind(setting *model.Setting, l *ldap.Conn, user *model.User, pwd string) (err error) {
	if setting.LDAP.Authentication == 0 {
		// simple auth
		err = l.Bind(fmt.Sprintf(setting.LDAP.UserDN, user.LoginName), pwd)
	} else {
		// bind auth
		err = l.Bind(setting.LDAP.BindDN, setting.LDAP.BindPassword)
		if err == nil {
			var entry *ldap.Entry
			entry, err = b.ldapSearchOne(setting, l, user.LoginName, "cn")
			if err == nil {
				err = l.Bind(entry.DN, pwd)
			}
		}
	}
	return
}

func (b *userBiz) ldapSearchOne(setting *model.Setting, l *ldap.Conn, name string, attrs ...string) (entry *ldap.Entry, err error) {
	filter := fmt.Sprintf(setting.LDAP.UserFilter, name)
	req := ldap.NewSearchRequest(
		setting.LDAP.BaseDN, ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
		filter, attrs, nil,
	)
	sr, err := l.Search(req)
	if err != nil {
		return nil, err
	}

	if length := len(sr.Entries); length == 0 {
		return nil, errors.New("User not found with filter: " + filter)
	} else if length > 1 {
		return nil, errors.New("Found more than one account when searching user with filter: " + filter)
	}

	return sr.Entries[0], nil
}

// Identify authenticate user
func (b *userBiz) Identify(token string) (user web.User) {
	do(func(d dao.Interface) {
		var (
			roles []*model.Role
			role  *model.Role
		)

		session, err := d.SessionGet(token)
		if err != nil {
			log.Get("user").Errorf("Load session failed: %v", err)
			return
		}
		if session == nil || session.Expires.Before(time.Now()) {
			return
		}

		u, err := d.UserGetByID(session.UserID)
		if err != nil {
			log.Get("user").Errorf("Load user failed: %v", err)
			return
		}
		if u == nil {
			return
		}

		if len(u.Roles) > 0 {
			roles = make([]*model.Role, len(u.Roles))
			for i, id := range u.Roles {
				role, err = d.RoleGet(id)
				if err != nil {
					return
				} else if role != nil {
					roles[i] = role
				}
			}
		}
		user = model.NewAuthUser(u, roles)
	})
	return
}

// Authorize check permission of user
func (b *userBiz) Authorize(user web.User, h web.HandlerInfo) bool {
	if au, ok := user.(*model.AuthUser); ok {
		return au.IsAllowed(h.Name())
	}
	return false
}

//
//func (b *userBiz) Find(key string) string {
//	b.locker.Lock()
//	defer b.locker.Unlock()
//
//	return b.tickets[key]
//}
//
//func (b *userBiz) setTicket(name string) (key string) {
//	b.locker.Lock()
//	defer b.locker.Unlock()
//
//	key = guid.New()
//	b.tickets[key] = name
//	return
//}
<commit_msg>Fix password validation<commit_after>package biz

import (
	"crypto/tls"
	"fmt"
	"net"
	"time"

	"github.com/cuigh/auxo/data/guid"
	"github.com/cuigh/auxo/errors"
	"github.com/cuigh/auxo/log"
	"github.com/cuigh/auxo/net/web"
	"github.com/cuigh/auxo/security/password"
	"github.com/cuigh/swirl/dao"
	"github.com/cuigh/swirl/model"
	"github.com/go-ldap/ldap"
)

var ErrIncorrectAuth = errors.New("login name or password is incorrect")

var User = &userBiz{}

type userBiz struct {
}

func (b *userBiz) GetByID(id string) (user *model.User, err error) {
	do(func(d dao.Interface) {
		user, err = d.UserGetByID(id)
	})
	return
}

func (b *userBiz) GetByName(loginName string) (user *model.User, err error) {
	do(func(d dao.Interface) {
		user, err = d.UserGetByName(loginName)
	})
	return
}

func (b *userBiz) Create(user *model.User, ctxUser web.User) (err error) {
	user.ID = guid.New()
	user.Status = model.UserStatusActive
	user.CreatedAt = time.Now()
	user.UpdatedAt = user.CreatedAt
	if user.Type == model.UserTypeInternal {
		user.Password, user.Salt, err = password.Generate(user.Password)
		if err != nil {
			return
		}
	}

	do(func(d dao.Interface) {
		if err = d.UserCreate(user); err == nil && ctxUser != nil {
			Event.CreateUser(model.EventActionCreate, user.LoginName, user.Name, ctxUser)
		}
	})
	return
}

func (b *userBiz) Update(user *model.User, ctxUser web.User) (err error) {
	do(func(d dao.Interface) {
		user.UpdatedAt = time.Now()
		if err = d.UserUpdate(user); err == nil {
			Event.CreateUser(model.EventActionUpdate, user.LoginName, user.Name, ctxUser)
		}
	})
	return
}

func (b *userBiz) Block(id string) (err error) {
	do(func(d dao.Interface) {
		err = d.UserBlock(id, true)
	})
	return
}

func (b *userBiz) Unblock(id string) (err error) {
	do(func(d dao.Interface) {
		err = d.UserBlock(id, false)
	})
	return
}

func (b *userBiz) Delete(id string) (err error) {
	do(func(d dao.Interface) {
		err = d.UserDelete(id)
	})
	return
}

func (b *userBiz) UpdateInfo(user *model.User) (err error) {
	do(func(d dao.Interface) {
		err = d.ProfileUpdateInfo(user)
	})
	return
}

func (b *userBiz) UpdatePassword(id, oldPwd, newPwd string) (err error) {
	do(func(d dao.Interface) {
		var (
			user      *model.User
			pwd, salt string
		)

		user, err = d.UserGetByID(id)
		if err != nil {
			return
		}

		if !password.Validate(oldPwd, user.Password, user.Salt) {
			err = errors.New("Current password is incorrect")
			return
		}

		pwd, salt, err = password.Generate(newPwd)
		if err != nil {
			return
		}

		err = d.ProfileUpdatePassword(id, pwd, salt)
	})
	return
}

func (b *userBiz) List(args *model.UserListArgs) (users []*model.User, count int, err error) {
	do(func(d dao.Interface) {
		users, count, err = d.UserList(args)
	})
	return
}

func (b *userBiz) Count() (count int, err error) {
	do(func(d dao.Interface) {
		count, err = d.UserCount()
	})
	return
}

func (b *userBiz) Login(name, pwd string) (token string, err error) {
	do(func(d dao.Interface) {
		var (
			user *model.User
		)

		user, err = d.UserGetByName(name)
		if err != nil {
			return
		}

		if user == nil {
			user = &model.User{
				Type:      model.UserTypeLDAP,
				LoginName: name,
			}
			err = b.loginLDAP(d, user, pwd)
		} else {
			if user.Status == model.UserStatusBlocked {
				err = fmt.Errorf("user %s is blocked", name)
				return
			}

			if user.Type == model.UserTypeInternal {
				err = b.loginInternal(user, pwd)
			} else {
				err = b.loginLDAP(d, user, pwd)
			}
		}

		if err != nil {
			return
		}

		session := &model.Session{
			UserID:    user.ID,
			Token:     guid.New(),
			UpdatedAt: time.Now(),
		}
		session.Expires = session.UpdatedAt.Add(time.Hour * 24)
		err = d.SessionUpdate(session)
		if err != nil {
			return
		}
		token = session.Token

		// create event
		Event.CreateAuthentication(model.EventActionLogin, user.ID, user.LoginName, user.Name)
	})
	return
}

func (b *userBiz) loginInternal(user *model.User, pwd string) error {
	if !password.Validate(pwd, user.Password, user.Salt) {
		return ErrIncorrectAuth
	}
	return nil
}

func (b *userBiz) loginLDAP(d dao.Interface, user *model.User, pwd string) error {
	setting, err := Setting.Get()
	if err != nil {
		return err
	}

	if !setting.LDAP.Enabled {
		return ErrIncorrectAuth
	}

	l, err := b.ldapDial(setting)
	if err != nil {
		return err
	}
	defer l.Close()

	// bind
	err = b.ldapBind(setting, l, user, pwd)
	if err != nil {
		log.Get("user").Error("LDAP > Failed to bind: ", err)
		return ErrIncorrectAuth
	}

	// Stop here for an exist user because we only need validate password.
	if user.ID != "" {
		return nil
	}

	// If user wasn't exist, we need create it
	entry, err := b.ldapSearchOne(setting, l, user.LoginName, setting.LDAP.NameAttr, setting.LDAP.EmailAttr)
	if err != nil {
		return err
	}

	user.Email = entry.GetAttributeValue(setting.LDAP.EmailAttr)
	user.Name = entry.GetAttributeValue(setting.LDAP.NameAttr)
	return b.Create(user, nil)
}

func (b *userBiz) ldapDial(setting *model.Setting) (*ldap.Conn, error) {
	host, _, err := net.SplitHostPort(setting.LDAP.Address)
	if err != nil {
		return nil, err
	}

	// TODO: support tls cert and verification
	tc := &tls.Config{
		ServerName:         host,
		InsecureSkipVerify: true,
		Certificates:       nil,
	}

	if setting.LDAP.Security == model.LDAPSecurityTLS {
		return ldap.DialTLS("tcp", setting.LDAP.Address, tc)
	}

	conn, err := ldap.Dial("tcp", setting.LDAP.Address)
	if err != nil {
		return nil, err
	}

	if setting.LDAP.Security == model.LDAPSecurityStartTLS {
		if err = conn.StartTLS(tc); err != nil {
			conn.Close()
			log.Get("user").Error("LDAP > Failed to switch to TLS: ", err)
			return nil, err
		}
	}
	return conn, nil
}

func (b *userBiz) ldapBind(setting *model.Setting, l *ldap.Conn, user *model.User, pwd string) (err error) {
	if setting.LDAP.Authentication == 0 {
		// simple auth
		err = l.Bind(fmt.Sprintf(setting.LDAP.UserDN, user.LoginName), pwd)
	} else {
		// bind auth
		err = l.Bind(setting.LDAP.BindDN, setting.LDAP.BindPassword)
		if err == nil {
			var entry *ldap.Entry
			entry, err = b.ldapSearchOne(setting, l, user.LoginName, "cn")
			if err == nil {
				err = l.Bind(entry.DN, pwd)
			}
		}
	}
	return
}

func (b *userBiz) ldapSearchOne(setting *model.Setting, l *ldap.Conn, name string, attrs ...string) (entry *ldap.Entry, err error) {
	filter := fmt.Sprintf(setting.LDAP.UserFilter, name)
	req := ldap.NewSearchRequest(
		setting.LDAP.BaseDN, ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
		filter, attrs, nil,
	)
	sr, err := l.Search(req)
	if err != nil {
		return nil, err
	}

	if length := len(sr.Entries); length == 0 {
		return nil, errors.New("User not found with filter: " + filter)
	} else if length > 1 {
		return nil, errors.New("Found more than one account when searching user with filter: " + filter)
	}

	return sr.Entries[0], nil
}

// Identify authenticate user
func (b *userBiz) Identify(token string) (user web.User) {
	do(func(d dao.Interface) {
		var (
			roles []*model.Role
			role  *model.Role
		)

		session, err := d.SessionGet(token)
		if err != nil {
			log.Get("user").Errorf("Load session failed: %v", err)
			return
		}
		if session == nil || session.Expires.Before(time.Now()) {
			return
		}

		u, err := d.UserGetByID(session.UserID)
		if err != nil {
			log.Get("user").Errorf("Load user failed: %v", err)
			return
		}
		if u == nil {
			return
		}

		if len(u.Roles) > 0 {
			roles = make([]*model.Role, len(u.Roles))
			for i, id := range u.Roles {
				role, err = d.RoleGet(id)
				if err != nil {
					return
				} else if role != nil {
					roles[i] = role
				}
			}
		}
		user = model.NewAuthUser(u, roles)
	})
	return
}

// Authorize check permission of user
func (b *userBiz) Authorize(user web.User, h web.HandlerInfo) bool {
	if au, ok := user.(*model.AuthUser); ok {
		return au.IsAllowed(h.Name())
	}
	return false
}

//
//func (b *userBiz) Find(key string) string {
//	b.locker.Lock()
//	defer b.locker.Unlock()
//
//	return b.tickets[key]
//}
//
//func (b *userBiz) setTicket(name string) (key string) {
//	b.locker.Lock()
//	defer b.locker.Unlock()
//
//	key = guid.New()
//	b.tickets[key] = name
//	return
//}
<|endoftext|>