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

package ogletest

import (
	"fmt"
	"math"
	"reflect"
)

// LessThan returns a matcher that matches integer, floating point, or strings
// values v such that v < x. Comparison is not defined between numeric and
// string types, but is defined between all integer and floating point types.
//
// x must itself be an integer, floating point, or string type; otherwise,
// LessThan will panic.
func LessThan(x interface{}) Matcher {
	v := reflect.ValueOf(x)
	kind := v.Kind()

	switch {
	case isInteger(v):
	case isFloat(v):
	case kind == reflect.String:

	default:
		panic(fmt.Sprintf("LessThan: unexpected kind %v", kind))
	}

	return &lessThanMatcher{v}
}

type lessThanMatcher struct {
	limit reflect.Value
}

func (m *lessThanMatcher) Description() string {
	// Special case: make it clear that strings are strings.
	if (m.limit.Kind() == reflect.String) {
		return fmt.Sprintf("less than \"%s\"", m.limit.String())
	}

	return fmt.Sprintf("less than %v", m.limit.Interface())
}

func compareIntegers(v1, v2 reflect.Value) (res MatchResult, err string) {
	res = MATCH_FALSE
	switch {
	case isSignedInteger(v1) && isSignedInteger(v2):
		if v1.Int() < v2.Int() {
			res = MATCH_TRUE
		}
		return

	case isSignedInteger(v1) && isUnsignedInteger(v2):
		if v1.Int() < 0 || uint64(v1.Int()) < v2.Uint() {
			res = MATCH_TRUE
		}
		return

	case isUnsignedInteger(v1) && isSignedInteger(v2):
		if v1.Uint() <= math.MaxInt64 && int64(v1.Uint()) < v2.Int() {
			res = MATCH_TRUE
		}
		return

	case isUnsignedInteger(v1) && isUnsignedInteger(v2):
		if v1.Uint() < v2.Uint() {
			res = MATCH_TRUE
		}
		return
	}

	panic(fmt.Sprintf("compareIntegers: %v %v", v1, v2))
}

func getFloat(v reflect.Value) float64 {
	switch {
	case isSignedInteger(v):
		return float64(v.Int())

	case isUnsignedInteger(v):
		return float64(v.Uint())

	case isFloat(v):
		return v.Float()
	}

	panic(fmt.Sprintf("getFloat: %v", v))
}

func (m *lessThanMatcher) Matches(c interface{}) (res MatchResult, err string) {
	v1 := reflect.ValueOf(c)
	v2 := m.limit

	res = MATCH_FALSE

	// Handle strings as a special case.
	if v1.Kind() == reflect.String && v2.Kind() == reflect.String {
		if v1.String() < v2.String() {
			res = MATCH_TRUE
		}
		return
	}

	// If we get here, we require that we are dealing with integers or floats.
  v1Legal := isInteger(v1) || isFloat(v1)
  v2Legal := isInteger(v2) || isFloat(v2)
	if (!v1Legal || !v2Legal) {
		res = MATCH_UNDEFINED
		err = "which is not comparable"
		return
	}

	// Handle the various comparison cases.
	switch {
	// Both integers
	case isInteger(v1) && isInteger(v2):
		return compareIntegers(v1, v2)

	// At least one float32
	case v1.Kind() == reflect.Float32 || v2.Kind() == reflect.Float32:
		if float32(getFloat(v1)) < float32(getFloat(v2)) {
			res = MATCH_TRUE
		}
		return

	// At least one float64
	case v1.Kind() == reflect.Float64 || v2.Kind() == reflect.Float64:
		if getFloat(v1) < getFloat(v2) {
			res = MATCH_TRUE
		}
		return
	}

	// We shouldn't get here.
	panic(fmt.Sprintf("lessThanMatcher.Matches: Shouldn't get here: %v %v", v1, v2))
}
<commit_msg>Ran gofmt.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ogletest

import (
	"fmt"
	"math"
	"reflect"
)

// LessThan returns a matcher that matches integer, floating point, or strings
// values v such that v < x. Comparison is not defined between numeric and
// string types, but is defined between all integer and floating point types.
//
// x must itself be an integer, floating point, or string type; otherwise,
// LessThan will panic.
func LessThan(x interface{}) Matcher {
	v := reflect.ValueOf(x)
	kind := v.Kind()

	switch {
	case isInteger(v):
	case isFloat(v):
	case kind == reflect.String:

	default:
		panic(fmt.Sprintf("LessThan: unexpected kind %v", kind))
	}

	return &lessThanMatcher{v}
}

type lessThanMatcher struct {
	limit reflect.Value
}

func (m *lessThanMatcher) Description() string {
	// Special case: make it clear that strings are strings.
	if m.limit.Kind() == reflect.String {
		return fmt.Sprintf("less than \"%s\"", m.limit.String())
	}

	return fmt.Sprintf("less than %v", m.limit.Interface())
}

func compareIntegers(v1, v2 reflect.Value) (res MatchResult, err string) {
	res = MATCH_FALSE
	switch {
	case isSignedInteger(v1) && isSignedInteger(v2):
		if v1.Int() < v2.Int() {
			res = MATCH_TRUE
		}
		return

	case isSignedInteger(v1) && isUnsignedInteger(v2):
		if v1.Int() < 0 || uint64(v1.Int()) < v2.Uint() {
			res = MATCH_TRUE
		}
		return

	case isUnsignedInteger(v1) && isSignedInteger(v2):
		if v1.Uint() <= math.MaxInt64 && int64(v1.Uint()) < v2.Int() {
			res = MATCH_TRUE
		}
		return

	case isUnsignedInteger(v1) && isUnsignedInteger(v2):
		if v1.Uint() < v2.Uint() {
			res = MATCH_TRUE
		}
		return
	}

	panic(fmt.Sprintf("compareIntegers: %v %v", v1, v2))
}

func getFloat(v reflect.Value) float64 {
	switch {
	case isSignedInteger(v):
		return float64(v.Int())

	case isUnsignedInteger(v):
		return float64(v.Uint())

	case isFloat(v):
		return v.Float()
	}

	panic(fmt.Sprintf("getFloat: %v", v))
}

func (m *lessThanMatcher) Matches(c interface{}) (res MatchResult, err string) {
	v1 := reflect.ValueOf(c)
	v2 := m.limit

	res = MATCH_FALSE

	// Handle strings as a special case.
	if v1.Kind() == reflect.String && v2.Kind() == reflect.String {
		if v1.String() < v2.String() {
			res = MATCH_TRUE
		}
		return
	}

	// If we get here, we require that we are dealing with integers or floats.
	v1Legal := isInteger(v1) || isFloat(v1)
	v2Legal := isInteger(v2) || isFloat(v2)
	if !v1Legal || !v2Legal {
		res = MATCH_UNDEFINED
		err = "which is not comparable"
		return
	}

	// Handle the various comparison cases.
	switch {
	// Both integers
	case isInteger(v1) && isInteger(v2):
		return compareIntegers(v1, v2)

	// At least one float32
	case v1.Kind() == reflect.Float32 || v2.Kind() == reflect.Float32:
		if float32(getFloat(v1)) < float32(getFloat(v2)) {
			res = MATCH_TRUE
		}
		return

	// At least one float64
	case v1.Kind() == reflect.Float64 || v2.Kind() == reflect.Float64:
		if getFloat(v1) < getFloat(v2) {
			res = MATCH_TRUE
		}
		return
	}

	// We shouldn't get here.
	panic(fmt.Sprintf("lessThanMatcher.Matches: Shouldn't get here: %v %v", v1, v2))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"

	"github.com/dt/thile/gen"
	"github.com/paperstreet/gohfile/hfile"
)

type Collection struct {
	Name string
	Path string
}

type CollectionSet struct {
	readers map[string]*hfile.Reader
}

func LoadCollections(collections []Collection) (*CollectionSet, error) {
	cs := new(CollectionSet)
	cs.readers = make(map[string]*hfile.Reader)

	for _, config := range collections {

		f, err := os.OpenFile(config.Path, os.O_RDONLY, 0)

		if err != nil {
			return nil, err
		}

		reader, err := hfile.NewReader(f)
		if err != nil {
			return nil, err
		}

		cs.readers[config.Name] = reader
	}

	return cs, nil
}

func (cs *CollectionSet) readerFor(c string) (*hfile.Reader, error) {
	reader, ok := cs.readers[c]
	if !ok {
		return nil, fmt.Errorf("not configured with reader for collection %s", c)
	}
	return reader, nil
}

func (cs *CollectionSet) scannerFor(c string) (*hfile.Scanner, error) {
	reader, err := cs.readerFor(c)
	if err != nil {
		return nil, err
	}
	s := hfile.NewScanner(reader)
	return &s, nil
}

func (cs *CollectionSet) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	if req.PerKeyValueLimit != nil {
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		log.Print("[GetValuesSingle] key: %v", key)
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	res.KeyCount = &found
	return res, nil
}

func (cs *CollectionSet) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	log.Println("[GetValuesMulti]", len(req.SortedKeys))
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	//out := make(map[int32][]byte)
	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *CollectionSet) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetIterator(req *gen.IteratorRequest) (r *gen.IteratorResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return nil, nil
}

func (cs *CollectionSet) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, nil
}
<commit_msg>print keys in hex<commit_after>package main

import (
	"encoding/hex"
	"fmt"
	"log"
	"os"

	"github.com/dt/thile/gen"
	"github.com/paperstreet/gohfile/hfile"
)

type Collection struct {
	Name string
	Path string
}

type CollectionSet struct {
	readers map[string]*hfile.Reader
}

func LoadCollections(collections []Collection) (*CollectionSet, error) {
	cs := new(CollectionSet)
	cs.readers = make(map[string]*hfile.Reader)

	for _, config := range collections {

		f, err := os.OpenFile(config.Path, os.O_RDONLY, 0)

		if err != nil {
			return nil, err
		}

		reader, err := hfile.NewReader(f)
		if err != nil {
			return nil, err
		}

		cs.readers[config.Name] = reader
	}

	return cs, nil
}

func (cs *CollectionSet) readerFor(c string) (*hfile.Reader, error) {
	reader, ok := cs.readers[c]
	if !ok {
		return nil, fmt.Errorf("not configured with reader for collection %s", c)
	}
	return reader, nil
}

func (cs *CollectionSet) scannerFor(c string) (*hfile.Scanner, error) {
	reader, err := cs.readerFor(c)
	if err != nil {
		return nil, err
	}
	s := hfile.NewScanner(reader)
	return &s, nil
}

func (cs *CollectionSet) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	if req.PerKeyValueLimit != nil {
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		pretty := make([]byte, len(key)*2)
		hex.Encode(pretty, key)
		log.Print("[GetValuesSingle] key: %s", pretty)
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	res.KeyCount = &found
	return res, nil
}

func (cs *CollectionSet) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	log.Println("[GetValuesMulti]", len(req.SortedKeys))
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	//out := make(map[int32][]byte)
	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *CollectionSet) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetIterator(req *gen.IteratorRequest) (r *gen.IteratorResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return nil, nil
}

func (cs *CollectionSet) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, nil
}
<|endoftext|>


<commit_before>package modules

import (
	"fmt"
	"github.com/davidscholberg/go-i3barjson"
	"net"
	"strings"
)

// Interface represents the configuration for the network interface block.
type Interface struct {
	BlockConfigBase `yaml:",inline"`
	IfaceName       string `yaml:"interface_name"`
	IPv4            string
	IPv4CIDR        string
	IPv6            string
	IPv6CIDR        string
	IPv6Local       string
}

// UpdateBlock updates the network interface block.
func (c Interface) UpdateBlock(b *i3barjson.Block) {
	var (
		status string
	)

	b.Color = c.Color
	fullTextFmt := fmt.Sprintf("%s%%s", c.Label)

	iface, err := net.InterfaceByName(c.IfaceName)
	if err != nil {
		b.Urgent = true
		b.FullText = fmt.Sprintf(fullTextFmt, err.Error())
		return
	}

	if iface.Flags == net.FlagUp {
		b.Urgent = false
		status = "up"
	} else {
		b.Urgent = true
		status = "down"
	}

	addrs, err := iface.Addrs()
	if err != nil {
		b.Urgent = true
		b.FullText = fmt.Sprintf(fullTextFmt, err.Error())
		return
	}

	for _, addr := range addrs {
		ip, _, err := net.ParseCIDR(addr.String())
		if err != nil {
			b.Urgent = true
			b.FullText = fmt.Sprintf(fullTextFmt, err.Error())
			return
		}

		// Checking for address family
		if ip.To4() != nil {
			fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv4\u003e", ip.String(), -1)
			fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr4\u003e", addr.String(), -1)
		} else {

			if ip.String()[0:4] == "fe80" {
				//setting ipv6 link local
				fullTextFmt = strings.Replace(fullTextFmt, "\u003clocal6\u003e", ip.String(), -1)
			} else {
				fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv6\u003e", ip.String()[0:3], -1)
				fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr6\u003e", addr.String(), -1)
			}
		}
	}

	// setting up/down flag
	fullTextFmt = strings.Replace(fullTextFmt, "\u003cstatus\u003e", status, -1)

	// clearing unset fields i.e. because of ipv6 single-stack
	fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv4\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr4\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv6\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr6\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003clocal6\u003e", "", -1)

	// removing the last %s placeholder from final string
	b.FullText = fmt.Sprintf(fullTextFmt, "")
}
<commit_msg>whoops, that shouldn't have been left in there<commit_after>package modules

import (
	"fmt"
	"github.com/davidscholberg/go-i3barjson"
	"net"
	"strings"
)

// Interface represents the configuration for the network interface block.
type Interface struct {
	BlockConfigBase `yaml:",inline"`
	IfaceName       string `yaml:"interface_name"`
	IPv4            string
	IPv4CIDR        string
	IPv6            string
	IPv6CIDR        string
}

// UpdateBlock updates the network interface block.
func (c Interface) UpdateBlock(b *i3barjson.Block) {
	var (
		status string
	)

	b.Color = c.Color
	fullTextFmt := fmt.Sprintf("%s%%s", c.Label)

	iface, err := net.InterfaceByName(c.IfaceName)
	if err != nil {
		b.Urgent = true
		b.FullText = fmt.Sprintf(fullTextFmt, err.Error())
		return
	}

	if iface.Flags == net.FlagUp {
		b.Urgent = false
		status = "up"
	} else {
		b.Urgent = true
		status = "down"
	}

	addrs, err := iface.Addrs()
	if err != nil {
		b.Urgent = true
		b.FullText = fmt.Sprintf(fullTextFmt, err.Error())
		return
	}

	for _, addr := range addrs {
		ip, _, err := net.ParseCIDR(addr.String())
		if err != nil {
			b.Urgent = true
			b.FullText = fmt.Sprintf(fullTextFmt, err.Error())
			return
		}

		// Checking for address family
		if ip.To4() != nil {
			fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv4\u003e", ip.String(), -1)
			fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr4\u003e", addr.String(), -1)
		} else {

			if ip.String()[0:4] == "fe80" {
				//setting ipv6 link local
				fullTextFmt = strings.Replace(fullTextFmt, "\u003clocal6\u003e", ip.String(), -1)
			} else {
				fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv6\u003e", ip.String(), -1)
				fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr6\u003e", addr.String(), -1)
			}
		}
	}

	// setting up/down flag
	fullTextFmt = strings.Replace(fullTextFmt, "\u003cstatus\u003e", status, -1)

	// clearing unset fields i.e. because of ipv6 single-stack
	fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv4\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr4\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003cipv6\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003ccidr6\u003e", "", -1)
	fullTextFmt = strings.Replace(fullTextFmt, "\u003clocal6\u003e", "", -1)

	// removing the last %s placeholder from final string
	b.FullText = fmt.Sprintf(fullTextFmt, "")
}
<|endoftext|>


<commit_before>package mpmogilefs

import (
	"bytes"
	"testing"
)

func TestParse(t *testing.T) {
	var mogilefs MogileFSPlugin
	stub := `uptime 35235
pending_queries 5
processing_queries 38
bored_queryworkers 3
queries 117
work_queue_for_delete 0
work_sent_to_delete 6
.
`

	mogilefsStats := bytes.NewBufferString(stub)

	stats, err := mogilefs.parseStats(mogilefsStats)
	if err != nil {
		t.Errorf("%v", err)
	}
}

func TestGraphDefinition(t *testing.T) {
	var mogilefs MogileFSPlugin

	graphdef := mogilefs.GraphDefinition()
	if len(graphdef) != 5 {
		t.Errorf("GetTempfilename: %d should be 5", len(graphdef))
	}
}
<commit_msg>Fix: stats declared and not used<commit_after>package mpmogilefs

import (
	"bytes"
	"testing"
)

func TestParse(t *testing.T) {
	var mogilefs MogileFSPlugin
	stub := `uptime 35235
pending_queries 5
processing_queries 38
bored_queryworkers 3
queries 117
work_queue_for_delete 0
work_sent_to_delete 6
.
`

	mogilefsStats := bytes.NewBufferString(stub)

	_, err := mogilefs.parseStats(mogilefsStats)
	if err != nil {
		t.Errorf("%v", err)
	}
}

func TestGraphDefinition(t *testing.T) {
	var mogilefs MogileFSPlugin

	graphdef := mogilefs.GraphDefinition()
	if len(graphdef) != 5 {
		t.Errorf("GetTempfilename: %d should be 5", len(graphdef))
	}
}
<|endoftext|>


<commit_before>package chat

import (
	"context"
	"testing"
	"time"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/clockwork"
	"github.com/stretchr/testify/require"
)

func setupLoaderTest(t *testing.T) (*kbtest.ChatTestContext, *kbtest.ChatMockWorld, *chatListener, chat1.NewConversationRemoteRes) {
	ctx, world, ri, _, baseSender, listener := setupTest(t, 1)

	u := world.GetUsers()[0]
	tc := world.Tcs[u.Username]
	trip := newConvTriple(ctx, t, tc, u.Username)
	firstMessagePlaintext := chat1.MessagePlaintext{
		ClientHeader: chat1.MessageClientHeader{
			Conv:        trip,
			TlfName:     u.Username,
			TlfPublic:   false,
			MessageType: chat1.MessageType_TLFNAME,
		},
		MessageBody: chat1.MessageBody{},
	}
	firstMessageBoxed, _, _, _, _, err := baseSender.Prepare(ctx, firstMessagePlaintext,
		chat1.ConversationMembersType_KBFS, nil)
	require.NoError(t, err)
	res, err := ri.NewConversationRemote2(ctx, chat1.NewConversationRemote2Arg{
		IdTriple:   trip,
		TLFMessage: *firstMessageBoxed,
	})
	require.NoError(t, err)
	return tc, world, listener, res
}

func TestConvLoader(t *testing.T) {
	tc, world, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	select {
	case convID := <-listener.bgConvLoads:
		if !convID.Eq(res.ConvID) {
			t.Errorf("loaded conv id: %s, expected %s", convID, res.ConvID)
		}
	case <-time.After(20 * time.Second):
		t.Fatal("timeout waiting for conversation load")
	}
}

type slowestRemote struct {
	chat1.RemoteInterface
	callCh chan struct{}
}

func makeSlowestRemote() slowestRemote {
	return slowestRemote{
		callCh: make(chan struct{}, 5),
	}
}

func (s slowestRemote) delay(ctx context.Context) {
	s.callCh <- struct{}{}
	select {
	case <-ctx.Done():
	case <-time.After(24 * time.Hour):
	}
}

func (s slowestRemote) GetThreadRemote(ctx context.Context, arg chat1.GetThreadRemoteArg) (res chat1.GetThreadRemoteRes, err error) {
	s.delay(ctx)
	return res, context.Canceled
}

func (s slowestRemote) GetMessagesRemote(ctx context.Context, arg chat1.GetMessagesRemoteArg) (res chat1.GetMessagesRemoteRes, err error) {
	s.delay(ctx)
	return res, context.Canceled
}

func TestConvLoaderSuspend(t *testing.T) {
	tc, world, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	ri := tc.ChatG.ConvSource.(*HybridConversationSource).ri
	slowRi := makeSlowestRemote()
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = func() chat1.RemoteInterface {
		return slowRi
	}
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	select {
	case <-slowRi.callCh:
	case <-time.After(20 * time.Second):
		require.Fail(t, "no remote call")
	}
	require.True(t, tc.Context().ConvLoader.Suspend(context.TODO()))
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	require.True(t, tc.Context().ConvLoader.Suspend(context.TODO()))

	tc.ChatG.ConvSource.(*HybridConversationSource).ri = ri
	require.False(t, tc.Context().ConvLoader.Resume(context.TODO()))
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	require.True(t, tc.Context().ConvLoader.Resume(context.TODO()))
	select {
	case convID := <-listener.bgConvLoads:
		require.Equal(t, res.ConvID, convID)
	case <-time.After(20 * time.Second):
		require.Fail(t, "no event")
	}
}

func TestConvLoaderAppState(t *testing.T) {
	tc, world, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	clock := clockwork.NewFakeClock()
	appStateCh := make(chan struct{})
	tc.ChatG.ConvLoader.(*BackgroundConvLoader).clock = clock
	tc.ChatG.ConvLoader.(*BackgroundConvLoader).appStateCh = appStateCh
	ri := tc.ChatG.ConvSource.(*HybridConversationSource).ri
	slowRi := makeSlowestRemote()
	failDuration := 2 * time.Second
	uid := gregor1.UID(tc.G.Env.GetUID().ToBytes())

	// Test that a foreground with no background doesnt do anything
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = func() chat1.RemoteInterface {
		return slowRi
	}
	tc.ChatG.ConvSource.(*HybridConversationSource).Clear(res.ConvID, uid)
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	clock.BlockUntil(1)
	clock.Advance(200 * time.Millisecond) // Get by small sleep
	select {
	case <-slowRi.callCh:
	case <-time.After(failDuration):
		require.Fail(t, "no remote call")
	}
	require.True(t, tc.Context().ConvLoader.Suspend(context.TODO()))
	tc.G.AppState.Update(keybase1.AppState_FOREGROUND)
	select {
	case <-appStateCh:
	case <-time.After(failDuration):
		require.Fail(t, "no app state")
	}
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = ri
	require.True(t, tc.Context().ConvLoader.Resume(context.TODO()))
	clock.BlockUntil(1)
	clock.Advance(2 * time.Second) // Get by resume wait
	clock.BlockUntil(1)
	clock.Advance(time.Hour) // Get by small sleep
	select {
	case convID := <-listener.bgConvLoads:
		require.Equal(t, res.ConvID, convID)
	case <-time.After(failDuration):
		require.Fail(t, "no event")
	}

	t.Logf("testing foreground/background")
	// Test that background/foreground works
	tc.ChatG.ConvSource.(*HybridConversationSource).Clear(res.ConvID, uid)
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = func() chat1.RemoteInterface {
		return slowRi
	}
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	clock.BlockUntil(1)
	clock.Advance(200 * time.Millisecond) // Get by small sleep
	select {
	case <-slowRi.callCh:
	case <-time.After(failDuration):
		require.Fail(t, "no remote call")
	}
	tc.G.AppState.Update(keybase1.AppState_BACKGROUND)
	select {
	case <-appStateCh:
	case <-time.After(failDuration):
		require.Fail(t, "no app state")
	}
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = ri
	tc.G.AppState.Update(keybase1.AppState_FOREGROUND)
	select {
	case <-appStateCh:
	case <-time.After(failDuration):
		require.Fail(t, "no app state")
	}
	// Need to advance clock
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	clock.BlockUntil(1)
	clock.Advance(10 * time.Second)
	clock.BlockUntil(1)
	clock.Advance(time.Hour) // Get by small sleep
	select {
	case convID := <-listener.bgConvLoads:
		require.Equal(t, res.ConvID, convID)
	case <-time.After(failDuration):
		require.Fail(t, "no event")
	}

}
<commit_msg>use long delay in this test<commit_after>package chat

import (
	"context"
	"testing"
	"time"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/clockwork"
	"github.com/stretchr/testify/require"
)

func setupLoaderTest(t *testing.T) (*kbtest.ChatTestContext, *kbtest.ChatMockWorld, *chatListener, chat1.NewConversationRemoteRes) {
	ctx, world, ri, _, baseSender, listener := setupTest(t, 1)

	u := world.GetUsers()[0]
	tc := world.Tcs[u.Username]
	trip := newConvTriple(ctx, t, tc, u.Username)
	firstMessagePlaintext := chat1.MessagePlaintext{
		ClientHeader: chat1.MessageClientHeader{
			Conv:        trip,
			TlfName:     u.Username,
			TlfPublic:   false,
			MessageType: chat1.MessageType_TLFNAME,
		},
		MessageBody: chat1.MessageBody{},
	}
	firstMessageBoxed, _, _, _, _, err := baseSender.Prepare(ctx, firstMessagePlaintext,
		chat1.ConversationMembersType_KBFS, nil)
	require.NoError(t, err)
	res, err := ri.NewConversationRemote2(ctx, chat1.NewConversationRemote2Arg{
		IdTriple:   trip,
		TLFMessage: *firstMessageBoxed,
	})
	require.NoError(t, err)
	return tc, world, listener, res
}

func TestConvLoader(t *testing.T) {
	tc, world, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	select {
	case convID := <-listener.bgConvLoads:
		if !convID.Eq(res.ConvID) {
			t.Errorf("loaded conv id: %s, expected %s", convID, res.ConvID)
		}
	case <-time.After(20 * time.Second):
		t.Fatal("timeout waiting for conversation load")
	}
}

type slowestRemote struct {
	chat1.RemoteInterface
	callCh chan struct{}
}

func makeSlowestRemote() slowestRemote {
	return slowestRemote{
		callCh: make(chan struct{}, 5),
	}
}

func (s slowestRemote) delay(ctx context.Context) {
	s.callCh <- struct{}{}
	select {
	case <-ctx.Done():
	case <-time.After(24 * time.Hour):
	}
}

func (s slowestRemote) GetThreadRemote(ctx context.Context, arg chat1.GetThreadRemoteArg) (res chat1.GetThreadRemoteRes, err error) {
	s.delay(ctx)
	return res, context.Canceled
}

func (s slowestRemote) GetMessagesRemote(ctx context.Context, arg chat1.GetMessagesRemoteArg) (res chat1.GetMessagesRemoteRes, err error) {
	s.delay(ctx)
	return res, context.Canceled
}

func TestConvLoaderSuspend(t *testing.T) {
	tc, world, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	ri := tc.ChatG.ConvSource.(*HybridConversationSource).ri
	slowRi := makeSlowestRemote()
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = func() chat1.RemoteInterface {
		return slowRi
	}
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	select {
	case <-slowRi.callCh:
	case <-time.After(20 * time.Second):
		require.Fail(t, "no remote call")
	}
	require.True(t, tc.Context().ConvLoader.Suspend(context.TODO()))
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	require.True(t, tc.Context().ConvLoader.Suspend(context.TODO()))

	tc.ChatG.ConvSource.(*HybridConversationSource).ri = ri
	require.False(t, tc.Context().ConvLoader.Resume(context.TODO()))
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	require.True(t, tc.Context().ConvLoader.Resume(context.TODO()))
	select {
	case convID := <-listener.bgConvLoads:
		require.Equal(t, res.ConvID, convID)
	case <-time.After(20 * time.Second):
		require.Fail(t, "no event")
	}
}

func TestConvLoaderAppState(t *testing.T) {
	tc, world, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	clock := clockwork.NewFakeClock()
	appStateCh := make(chan struct{})
	tc.ChatG.ConvLoader.(*BackgroundConvLoader).clock = clock
	tc.ChatG.ConvLoader.(*BackgroundConvLoader).appStateCh = appStateCh
	ri := tc.ChatG.ConvSource.(*HybridConversationSource).ri
	slowRi := makeSlowestRemote()
	failDuration := 20 * time.Second
	uid := gregor1.UID(tc.G.Env.GetUID().ToBytes())

	// Test that a foreground with no background doesnt do anything
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = func() chat1.RemoteInterface {
		return slowRi
	}
	tc.ChatG.ConvSource.(*HybridConversationSource).Clear(res.ConvID, uid)
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	clock.BlockUntil(1)
	clock.Advance(200 * time.Millisecond) // Get by small sleep
	select {
	case <-slowRi.callCh:
	case <-time.After(failDuration):
		require.Fail(t, "no remote call")
	}
	require.True(t, tc.Context().ConvLoader.Suspend(context.TODO()))
	tc.G.AppState.Update(keybase1.AppState_FOREGROUND)
	select {
	case <-appStateCh:
	case <-time.After(failDuration):
		require.Fail(t, "no app state")
	}
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = ri
	require.True(t, tc.Context().ConvLoader.Resume(context.TODO()))
	clock.BlockUntil(1)
	clock.Advance(2 * time.Second) // Get by resume wait
	clock.BlockUntil(1)
	clock.Advance(time.Hour) // Get by small sleep
	select {
	case convID := <-listener.bgConvLoads:
		require.Equal(t, res.ConvID, convID)
	case <-time.After(failDuration):
		require.Fail(t, "no event")
	}

	t.Logf("testing foreground/background")
	// Test that background/foreground works
	tc.ChatG.ConvSource.(*HybridConversationSource).Clear(res.ConvID, uid)
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = func() chat1.RemoteInterface {
		return slowRi
	}
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(), res.ConvID))
	clock.BlockUntil(1)
	clock.Advance(200 * time.Millisecond) // Get by small sleep
	select {
	case <-slowRi.callCh:
	case <-time.After(failDuration):
		require.Fail(t, "no remote call")
	}
	tc.G.AppState.Update(keybase1.AppState_BACKGROUND)
	select {
	case <-appStateCh:
	case <-time.After(failDuration):
		require.Fail(t, "no app state")
	}
	tc.ChatG.ConvSource.(*HybridConversationSource).ri = ri
	tc.G.AppState.Update(keybase1.AppState_FOREGROUND)
	select {
	case <-appStateCh:
	case <-time.After(failDuration):
		require.Fail(t, "no app state")
	}
	// Need to advance clock
	select {
	case <-listener.bgConvLoads:
		require.Fail(t, "no load yet")
	default:
	}
	clock.BlockUntil(1)
	clock.Advance(10 * time.Second)
	clock.BlockUntil(1)
	clock.Advance(time.Hour) // Get by small sleep
	select {
	case convID := <-listener.bgConvLoads:
		require.Equal(t, res.ConvID, convID)
	case <-time.After(failDuration):
		require.Fail(t, "no event")
	}

}
<|endoftext|>


<commit_before>package storage

import (
	"fmt"
	"sort"
	"sync"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/encrypteddb"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	context "golang.org/x/net/context"
)

const reacjiDiskVersion = 1

// If the user has less than 5 favorite reacjis we stuff these defaults in.
var DefaultTopReacjis = []string{":+1:", ":-1:", ":tada:", ":joy:", ":sunglasses:"}

var addReacjiMemCacheHookOnce sync.Once

type ReacjiInternalStorage struct {
	FrequencyMap map[string]int
	SkinTone     keybase1.ReacjiSkinTone
}

func NewReacjiInternalStorage() *ReacjiInternalStorage {
	return &ReacjiInternalStorage{
		FrequencyMap: make(map[string]int),
	}
}

type reacjiMemCacheImpl struct {
	sync.RWMutex

	uid  gregor1.UID
	data *ReacjiInternalStorage
}

func newReacjiMemCacheImpl() *reacjiMemCacheImpl {
	return &reacjiMemCacheImpl{
		data: NewReacjiInternalStorage(),
	}
}

func (i *reacjiMemCacheImpl) Get(uid gregor1.UID) (bool, *ReacjiInternalStorage) {
	i.RLock()
	defer i.RUnlock()
	if !uid.Eq(i.uid) {
		return false, nil
	}
	return true, i.data
}

func (i *reacjiMemCacheImpl) Put(uid gregor1.UID, data *ReacjiInternalStorage) {
	i.Lock()
	defer i.Unlock()
	i.uid = uid
	i.data = data
}

func (i *reacjiMemCacheImpl) OnLogout(m libkb.MetaContext) error {
	i.Lock()
	defer i.Unlock()
	i.data = NewReacjiInternalStorage()
	i.uid = nil
	return nil
}

var reacjiMemCache = newReacjiMemCacheImpl()

type reacjiPair struct {
	name string
	freq int
}

type reacjiDiskEntry struct {
	Version int
	// reacji name -> frequency
	Data ReacjiInternalStorage
}

type ReacjiStore struct {
	sync.Mutex
	utils.DebugLabeler
	encryptedDB *encrypteddb.EncryptedDB
}

// Keeps map counting emoji used in reactions for each user. Used to populate
// the reacji heads up display.
// Data is stored in an encrypted leveldb in the form:
// uid -> {
//                reacjiName: frequency,
//                ":+1:": 5,
//                ...
//         },
func NewReacjiStore(g *globals.Context) *ReacjiStore {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return GetSecretBoxKey(ctx, g.ExternalG(), DefaultSecretUI)
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalChatDb
	}
	// add a logout hook to clear the in-memory cache, but only add it once:
	addReacjiMemCacheHookOnce.Do(func() {
		g.ExternalG().AddLogoutHook(reacjiMemCache, "reacjiMemCache")
	})
	return &ReacjiStore{
		DebugLabeler: utils.NewDebugLabeler(g.GetLog(), "ReacjiStore", false),
		encryptedDB:  encrypteddb.New(g.ExternalG(), dbFn, keyFn),
	}
}

func (s *ReacjiStore) dbKey(uid gregor1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBChatReacji,
		Key: fmt.Sprintf("ri:%s", uid),
	}
}

func (s *ReacjiStore) populateCacheLocked(ctx context.Context, uid gregor1.UID) *ReacjiInternalStorage {
	if found, data := reacjiMemCache.Get(uid); found {
		return data
	}

	// populate the cache after we fetch from disk
	data := NewReacjiInternalStorage()
	defer func() { reacjiMemCache.Put(uid, data) }()

	dbKey := s.dbKey(uid)
	var entry reacjiDiskEntry
	found, err := s.encryptedDB.Get(ctx, dbKey, &entry)
	if err != nil || !found {
		s.Debug(ctx, "reacji map not found on disk")
		return data
	}

	if entry.Version != reacjiDiskVersion {
		// drop the history if our format changed
		if err = s.encryptedDB.Delete(ctx, dbKey); err != nil {
			s.Debug(ctx, "unable to delete cache entry: %v", err)
			return data
		}
	}
	data = &entry.Data
	return data
}

func (s *ReacjiStore) PutReacji(ctx context.Context, uid gregor1.UID, reacji string) error {
	s.Lock()
	defer s.Unlock()

	cache := s.populateCacheLocked(ctx, uid)
	cache.FrequencyMap[reacji]++
	reacjiMemCache.Put(uid, cache)
	dbKey := s.dbKey(uid)
	return s.encryptedDB.Put(ctx, dbKey, reacjiDiskEntry{
		Version: reacjiDiskVersion,
		Data:    *cache,
	})
}

func (s *ReacjiStore) PutSkinTone(ctx context.Context, uid gregor1.UID,
	skinTone keybase1.ReacjiSkinTone) error {
	s.Lock()
	defer s.Unlock()

	cache := s.populateCacheLocked(ctx, uid)
	cache.SkinTone = skinTone
	dbKey := s.dbKey(uid)
	return s.encryptedDB.Put(ctx, dbKey, reacjiDiskEntry{
		Version: reacjiDiskVersion,
		Data:    *cache,
	})
}

func (s *ReacjiStore) GetInternalStore(ctx context.Context, uid gregor1.UID) *ReacjiInternalStorage {
	s.Lock()
	defer s.Unlock()
	return s.populateCacheLocked(ctx, uid)
}

// UserReacjis returns the user's most frequently used reacjis falling back to
// `DefaultTopReacjis` if there is not enough history. Results are ordered by
// frequency and then alphabetically.
func (s *ReacjiStore) UserReacjis(ctx context.Context, uid gregor1.UID) keybase1.UserReacjis {
	s.Lock()
	defer s.Unlock()

	cache := s.populateCacheLocked(ctx, uid)
	pairs := []reacjiPair{}
	for name, freq := range cache.FrequencyMap {
		pairs = append(pairs, reacjiPair{name: name, freq: freq})
	}
	// sort by frequency and then alphabetically
	sort.Slice(pairs, func(i, j int) bool {
		if pairs[i].freq == pairs[j].freq {
			return pairs[i].name < pairs[j].name
		}
		return pairs[i].freq > pairs[j].freq
	})

	reacjis := []string{}
	for _, p := range pairs {
		reacjis = append(reacjis, p.name)
	}
	if len(reacjis) < len(DefaultTopReacjis) {
		reacjis = append(reacjis, DefaultTopReacjis[:len(DefaultTopReacjis)-len(reacjis)]...)
	}
	return keybase1.UserReacjis{
		TopReacjis: reacjis,
		SkinTone:   cache.SkinTone,
	}
}
<commit_msg>bump reacji disk version (#17016)<commit_after>package storage

import (
	"fmt"
	"sort"
	"sync"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/encrypteddb"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	context "golang.org/x/net/context"
)

const reacjiDiskVersion = 2

// If the user has less than 5 favorite reacjis we stuff these defaults in.
var DefaultTopReacjis = []string{":+1:", ":-1:", ":tada:", ":joy:", ":sunglasses:"}

var addReacjiMemCacheHookOnce sync.Once

type ReacjiInternalStorage struct {
	FrequencyMap map[string]int
	SkinTone     keybase1.ReacjiSkinTone
}

func NewReacjiInternalStorage() *ReacjiInternalStorage {
	return &ReacjiInternalStorage{
		FrequencyMap: make(map[string]int),
	}
}

type reacjiMemCacheImpl struct {
	sync.RWMutex

	uid  gregor1.UID
	data *ReacjiInternalStorage
}

func newReacjiMemCacheImpl() *reacjiMemCacheImpl {
	return &reacjiMemCacheImpl{
		data: NewReacjiInternalStorage(),
	}
}

func (i *reacjiMemCacheImpl) Get(uid gregor1.UID) (bool, *ReacjiInternalStorage) {
	i.RLock()
	defer i.RUnlock()
	if !uid.Eq(i.uid) {
		return false, nil
	}
	return true, i.data
}

func (i *reacjiMemCacheImpl) Put(uid gregor1.UID, data *ReacjiInternalStorage) {
	i.Lock()
	defer i.Unlock()
	i.uid = uid
	i.data = data
}

func (i *reacjiMemCacheImpl) OnLogout(m libkb.MetaContext) error {
	i.Lock()
	defer i.Unlock()
	i.data = NewReacjiInternalStorage()
	i.uid = nil
	return nil
}

var reacjiMemCache = newReacjiMemCacheImpl()

type reacjiPair struct {
	name string
	freq int
}

type reacjiDiskEntry struct {
	Version int
	// reacji name -> frequency
	Data ReacjiInternalStorage
}

type ReacjiStore struct {
	sync.Mutex
	utils.DebugLabeler
	encryptedDB *encrypteddb.EncryptedDB
}

// Keeps map counting emoji used in reactions for each user. Used to populate
// the reacji heads up display.
// Data is stored in an encrypted leveldb in the form:
// uid -> {
//                reacjiName: frequency,
//                ":+1:": 5,
//                ...
//         },
func NewReacjiStore(g *globals.Context) *ReacjiStore {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return GetSecretBoxKey(ctx, g.ExternalG(), DefaultSecretUI)
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalChatDb
	}
	// add a logout hook to clear the in-memory cache, but only add it once:
	addReacjiMemCacheHookOnce.Do(func() {
		g.ExternalG().AddLogoutHook(reacjiMemCache, "reacjiMemCache")
	})
	return &ReacjiStore{
		DebugLabeler: utils.NewDebugLabeler(g.GetLog(), "ReacjiStore", false),
		encryptedDB:  encrypteddb.New(g.ExternalG(), dbFn, keyFn),
	}
}

func (s *ReacjiStore) dbKey(uid gregor1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBChatReacji,
		Key: fmt.Sprintf("ri:%s", uid),
	}
}

func (s *ReacjiStore) populateCacheLocked(ctx context.Context, uid gregor1.UID) *ReacjiInternalStorage {
	if found, data := reacjiMemCache.Get(uid); found {
		return data
	}

	// populate the cache after we fetch from disk
	data := NewReacjiInternalStorage()
	defer func() { reacjiMemCache.Put(uid, data) }()

	dbKey := s.dbKey(uid)
	var entry reacjiDiskEntry
	found, err := s.encryptedDB.Get(ctx, dbKey, &entry)
	if err != nil || !found {
		s.Debug(ctx, "reacji map not found on disk")
		return data
	}

	if entry.Version != reacjiDiskVersion {
		// drop the history if our format changed
		if err = s.encryptedDB.Delete(ctx, dbKey); err != nil {
			s.Debug(ctx, "unable to delete cache entry: %v", err)
			return data
		}
	}
	data = &entry.Data
	return data
}

func (s *ReacjiStore) PutReacji(ctx context.Context, uid gregor1.UID, reacji string) error {
	s.Lock()
	defer s.Unlock()

	cache := s.populateCacheLocked(ctx, uid)
	cache.FrequencyMap[reacji]++
	reacjiMemCache.Put(uid, cache)
	dbKey := s.dbKey(uid)
	return s.encryptedDB.Put(ctx, dbKey, reacjiDiskEntry{
		Version: reacjiDiskVersion,
		Data:    *cache,
	})
}

func (s *ReacjiStore) PutSkinTone(ctx context.Context, uid gregor1.UID,
	skinTone keybase1.ReacjiSkinTone) error {
	s.Lock()
	defer s.Unlock()

	cache := s.populateCacheLocked(ctx, uid)
	cache.SkinTone = skinTone
	dbKey := s.dbKey(uid)
	return s.encryptedDB.Put(ctx, dbKey, reacjiDiskEntry{
		Version: reacjiDiskVersion,
		Data:    *cache,
	})
}

func (s *ReacjiStore) GetInternalStore(ctx context.Context, uid gregor1.UID) *ReacjiInternalStorage {
	s.Lock()
	defer s.Unlock()
	return s.populateCacheLocked(ctx, uid)
}

// UserReacjis returns the user's most frequently used reacjis falling back to
// `DefaultTopReacjis` if there is not enough history. Results are ordered by
// frequency and then alphabetically.
func (s *ReacjiStore) UserReacjis(ctx context.Context, uid gregor1.UID) keybase1.UserReacjis {
	s.Lock()
	defer s.Unlock()

	cache := s.populateCacheLocked(ctx, uid)
	pairs := []reacjiPair{}
	for name, freq := range cache.FrequencyMap {
		pairs = append(pairs, reacjiPair{name: name, freq: freq})
	}
	// sort by frequency and then alphabetically
	sort.Slice(pairs, func(i, j int) bool {
		if pairs[i].freq == pairs[j].freq {
			return pairs[i].name < pairs[j].name
		}
		return pairs[i].freq > pairs[j].freq
	})

	reacjis := []string{}
	for _, p := range pairs {
		reacjis = append(reacjis, p.name)
	}
	if len(reacjis) < len(DefaultTopReacjis) {
		reacjis = append(reacjis, DefaultTopReacjis[:len(DefaultTopReacjis)-len(reacjis)]...)
	}
	return keybase1.UserReacjis{
		TopReacjis: reacjis,
		SkinTone:   cache.SkinTone,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package grpcclient contains utility methods for gRPC client implementations
// to use. It also supports plug-in authentication.
package grpcclient

import (
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"vitess.io/vitess/go/vt/grpccommon"
	"vitess.io/vitess/go/vt/vttls"

	"vitess.io/vitess/go/vt/log"
)

// FailFast is a self-documenting type for the grpc.FailFast.
type FailFast bool

// grpcDialOptions is a registry of functions that append grpcDialOption to use when dialing a service
var grpcDialOptions []func(opts []grpc.DialOption) ([]grpc.DialOption, error)

// RegisterGRPCDialOptions registers an implementation of AuthServer.
func RegisterGRPCDialOptions(grpcDialOptionsFunc func(opts []grpc.DialOption) ([]grpc.DialOption, error)) {
	grpcDialOptions = append(grpcDialOptions, grpcDialOptionsFunc)
}

// Dial creates a grpc connection to the given target.
// failFast is a non-optional parameter because callers are required to specify
// what that should be.
func Dial(target string, failFast FailFast, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
	grpccommon.EnableTracingOpt()
	newopts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(
			grpc.MaxCallRecvMsgSize(*grpccommon.MaxMessageSize),
			grpc.MaxCallSendMsgSize(*grpccommon.MaxMessageSize),
			grpc.FailFast(bool(failFast)),
		),
	}
	newopts = append(newopts, opts...)
	var err error
	for _, grpcDialOptionInitializer := range grpcDialOptions {
		newopts, err = grpcDialOptionInitializer(newopts)
		if err != nil {
			log.Fatalf("There was an error initializing client grpc.DialOption: %v", err)
		}
	}
	return grpc.Dial(target, newopts...)
}

// SecureDialOption returns the gRPC dial option to use for the
// given client connection. It is either using TLS, or Insecure if
// nothing is set.
func SecureDialOption(cert, key, ca, name string) (grpc.DialOption, error) {
	// No security options set, just return.
	if (cert == "" || key == "") && ca == "" {
		return grpc.WithInsecure(), nil
	}

	// Load the config.
	config, err := vttls.ClientConfig(cert, key, ca, name)
	if err != nil {
		return nil, err
	}

	// Create the creds server options.
	creds := credentials.NewTLS(config)
	return grpc.WithTransportCredentials(creds), nil
}
<commit_msg>Adds gRPC client knobs to control KeepAlive and WindowSize settings<commit_after>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package grpcclient contains utility methods for gRPC client implementations
// to use. It also supports plug-in authentication.
package grpcclient

import (
	"flag"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/keepalive"

	"vitess.io/vitess/go/vt/grpccommon"
	"vitess.io/vitess/go/vt/vttls"

	"vitess.io/vitess/go/vt/log"
)

var (
	networkTimeout        = flag.Duration("grpc_network_timeout", 0, "keepalive time and timeout for any network operation in gRPC.")
	initialConnWindowSize = flag.Int("grpc_initial_window_size", 0, "grpc initial connection window size")
	initialWindowSize     = flag.Int("grpc_initial_window_size", 0, "grpc initial window size")
)

// FailFast is a self-documenting type for the grpc.FailFast.
type FailFast bool

// grpcDialOptions is a registry of functions that append grpcDialOption to use when dialing a service
var grpcDialOptions []func(opts []grpc.DialOption) ([]grpc.DialOption, error)

// RegisterGRPCDialOptions registers an implementation of AuthServer.
func RegisterGRPCDialOptions(grpcDialOptionsFunc func(opts []grpc.DialOption) ([]grpc.DialOption, error)) {
	grpcDialOptions = append(grpcDialOptions, grpcDialOptionsFunc)
}

// Dial creates a grpc connection to the given target.
// failFast is a non-optional parameter because callers are required to specify
// what that should be.
func Dial(target string, failFast FailFast, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
	grpccommon.EnableTracingOpt()
	newopts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(
			grpc.MaxCallRecvMsgSize(*grpccommon.MaxMessageSize),
			grpc.MaxCallSendMsgSize(*grpccommon.MaxMessageSize),
			grpc.FailFast(bool(failFast)),
		),
	}
	if *networkTimeout != 0 {
		kp := keepalive.ClientParameters{
			// Every *networkTimeout a ping will be send to the server to test the connection.
			Time: *networkTimeout,
			// After this timeout gRPC will close connection and fail all inflight gRPC requests.
			Timeout:             *networkTimeout,
			PermitWithoutStream: true,
		}
		newopts = append(newopts, grpc.WithKeepaliveParams(kp))
	}

	if *initialConnWindowSize != 0 {
		newopts = append(newopts, grpc.WithInitialConnWindowSize(int32(*initialConnWindowSize)))
	}

	if *initialWindowSize != 0 {
		newopts = append(newopts, grpc.WithInitialWindowSize(int32(*initialWindowSize)))
	}

	newopts = append(newopts, opts...)
	var err error
	for _, grpcDialOptionInitializer := range grpcDialOptions {
		newopts, err = grpcDialOptionInitializer(newopts)
		if err != nil {
			log.Fatalf("There was an error initializing client grpc.DialOption: %v", err)
		}
	}
	return grpc.Dial(target, newopts...)
}

// SecureDialOption returns the gRPC dial option to use for the
// given client connection. It is either using TLS, or Insecure if
// nothing is set.
func SecureDialOption(cert, key, ca, name string) (grpc.DialOption, error) {
	// No security options set, just return.
	if (cert == "" || key == "") && ca == "" {
		return grpc.WithInsecure(), nil
	}

	// Load the config.
	config, err := vttls.ClientConfig(cert, key, ca, name)
	if err != nil {
		return nil, err
	}

	// Create the creds server options.
	creds := credentials.NewTLS(config)
	return grpc.WithTransportCredentials(creds), nil
}
<|endoftext|>


<commit_before>package gobusterdir

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"log"
	"strings"
	"text/tabwriter"

	"github.com/OJ/gobuster/v3/libgobuster"
	"github.com/google/uuid"
)

// GobusterDir is the main type to implement the interface
type GobusterDir struct {
	options    *OptionsDir
	globalopts *libgobuster.Options
	http       *libgobuster.HTTPClient
}

// GetRequest issues a GET request to the target and returns
// the status code, length and an error
func (d *GobusterDir) get(url string) (*int, *int64, error) {
	return d.http.Get(url, "", d.options.Cookies)
}

// NewGobusterDir creates a new initialized GobusterDir
func NewGobusterDir(cont context.Context, globalopts *libgobuster.Options, opts *OptionsDir) (*GobusterDir, error) {
	if globalopts == nil {
		return nil, fmt.Errorf("please provide valid global options")
	}

	if opts == nil {
		return nil, fmt.Errorf("please provide valid plugin options")
	}

	g := GobusterDir{
		options:    opts,
		globalopts: globalopts,
	}

	httpOpts := libgobuster.HTTPOptions{
		Proxy:          opts.Proxy,
		FollowRedirect: opts.FollowRedirect,
		InsecureSSL:    opts.InsecureSSL,
		Timeout:        opts.Timeout,
		Username:       opts.Username,
		Password:       opts.Password,
		UserAgent:      opts.UserAgent,
	}

	h, err := libgobuster.NewHTTPClient(cont, &httpOpts)
	if err != nil {
		return nil, err
	}
	g.http = h
	return &g, nil
}

// PreRun is the pre run implementation of gobusterdir
func (d *GobusterDir) PreRun() error {
	// add trailing slash
	if !strings.HasSuffix(d.options.URL, "/") {
		d.options.URL = fmt.Sprintf("%s/", d.options.URL)
	}

	_, _, err := d.get(d.options.URL)
	if err != nil {
		return fmt.Errorf("unable to connect to %s: %v", d.options.URL, err)
	}

	guid := uuid.New()
	url := fmt.Sprintf("%s%s", d.options.URL, guid)
	wildcardResp, _, err := d.get(url)
	if err != nil {
		return err
	}

	if d.options.StatusCodesParsed.Contains(*wildcardResp) {
		log.Printf("[-] Wildcard response found: %s => %d", url, *wildcardResp)
		if !d.options.WildcardForced {
			return fmt.Errorf("To force processing of Wildcard responses, specify the '--wildcard' switch.")
		}
	}

	return nil
}

// Run is the process implementation of gobusterdir
func (d *GobusterDir) Run(word string) ([]libgobuster.Result, error) {
	suffix := ""
	if d.options.UseSlash {
		suffix = "/"
	}

	// Try the DIR first
	url := fmt.Sprintf("%s%s%s", d.options.URL, word, suffix)
	dirResp, dirSize, err := d.get(url)
	if err != nil {
		return nil, err
	}
	var ret []libgobuster.Result
	if dirResp != nil {
		ret = append(ret, libgobuster.Result{
			Entity:     fmt.Sprintf("%s%s", word, suffix),
			StatusCode: *dirResp,
			Size:       dirSize,
		})
	}

	// Follow up with files using each ext.
	for ext := range d.options.ExtensionsParsed.Set {
		file := fmt.Sprintf("%s.%s", word, ext)
		url = fmt.Sprintf("%s%s", d.options.URL, file)
		fileResp, fileSize, err := d.get(url)
		if err != nil {
			return nil, err
		}

		if fileResp != nil {
			ret = append(ret, libgobuster.Result{
				Entity:     file,
				StatusCode: *fileResp,
				Size:       fileSize,
			})
		}
	}
	return ret, nil
}

// ResultToString is the to string implementation of gobusterdir
func (d *GobusterDir) ResultToString(r *libgobuster.Result) (*string, error) {
	buf := &bytes.Buffer{}

	// Prefix if we're in verbose mode
	if d.globalopts.Verbose {
		if d.options.StatusCodesParsed.Contains(r.StatusCode) {
			if _, err := fmt.Fprintf(buf, "Found: "); err != nil {
				return nil, err
			}
		} else {
			if _, err := fmt.Fprintf(buf, "Missed: "); err != nil {
				return nil, err
			}
		}
	}

	if d.options.StatusCodesParsed.Contains(r.StatusCode) || d.globalopts.Verbose {
		if d.options.Expanded {
			if _, err := fmt.Fprintf(buf, "%s", d.options.URL); err != nil {
				return nil, err
			}
		} else {
			if _, err := fmt.Fprintf(buf, "/"); err != nil {
				return nil, err
			}
		}
		if _, err := fmt.Fprintf(buf, "%s", r.Entity); err != nil {
			return nil, err
		}

		if !d.options.NoStatus {
			if _, err := fmt.Fprintf(buf, " (Status: %d)", r.StatusCode); err != nil {
				return nil, err
			}
		}

		if r.Size != nil {
			if _, err := fmt.Fprintf(buf, " [Size: %d]", *r.Size); err != nil {
				return nil, err
			}
		}
		if _, err := fmt.Fprintf(buf, "\n"); err != nil {
			return nil, err
		}
	}
	s := buf.String()
	return &s, nil
}

// GetConfigString returns the string representation of the current config
func (d *GobusterDir) GetConfigString() (string, error) {
	var buffer bytes.Buffer
	bw := bufio.NewWriter(&buffer)
	tw := tabwriter.NewWriter(bw, 0, 5, 3, ' ', 0)
	o := d.options
	if _, err := fmt.Fprintf(tw, "[+] Url:\t%s\n", o.URL); err != nil {
		return "", err
	}
	if _, err := fmt.Fprintf(tw, "[+] Threads:\t%d\n", d.globalopts.Threads); err != nil {
		return "", err
	}

	wordlist := "stdin (pipe)"
	if d.globalopts.Wordlist != "-" {
		wordlist = d.globalopts.Wordlist
	}
	if _, err := fmt.Fprintf(tw, "[+] Wordlist:\t%s\n", wordlist); err != nil {
		return "", err
	}

	if _, err := fmt.Fprintf(tw, "[+] Status codes:\t%s\n", o.StatusCodesParsed.Stringify()); err != nil {
		return "", err
	}

	if o.Proxy != "" {
		if _, err := fmt.Fprintf(tw, "[+] Proxy:\t%s\n", o.Proxy); err != nil {
			return "", err
		}
	}

	if o.Cookies != "" {
		if _, err := fmt.Fprintf(tw, "[+] Cookies:\t%s\n", o.Cookies); err != nil {
			return "", err
		}
	}

	if o.UserAgent != "" {
		if _, err := fmt.Fprintf(tw, "[+] User Agent:\t%s\n", o.UserAgent); err != nil {
			return "", err
		}
	}

	if o.IncludeLength {
		if _, err := fmt.Fprintf(tw, "[+] Show length:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.Username != "" {
		if _, err := fmt.Fprintf(tw, "[+] Auth User:\t%s\n", o.Username); err != nil {
			return "", err
		}
	}

	if o.Extensions != "" {
		if _, err := fmt.Fprintf(tw, "[+] Extensions:\t%s\n", o.ExtensionsParsed.Stringify()); err != nil {
			return "", err
		}
	}

	if o.UseSlash {
		if _, err := fmt.Fprintf(tw, "[+] Add Slash:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.FollowRedirect {
		if _, err := fmt.Fprintf(tw, "[+] Follow Redir:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.Expanded {
		if _, err := fmt.Fprintf(tw, "[+] Expanded:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.NoStatus {
		if _, err := fmt.Fprintf(tw, "[+] No status:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if d.globalopts.Verbose {
		if _, err := fmt.Fprintf(tw, "[+] Verbose:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if _, err := fmt.Fprintf(tw, "[+] Timeout:\t%s\n", o.Timeout.String()); err != nil {
		return "", err
	}

	if err := tw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %v", err)
	}

	if err := bw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %v", err)
	}

	return strings.TrimSpace(buffer.String()), nil
}
<commit_msg>use result status enum<commit_after>package gobusterdir

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"log"
	"strings"
	"text/tabwriter"

	"github.com/OJ/gobuster/v3/libgobuster"
	"github.com/google/uuid"
)

// GobusterDir is the main type to implement the interface
type GobusterDir struct {
	options    *OptionsDir
	globalopts *libgobuster.Options
	http       *libgobuster.HTTPClient
}

// GetRequest issues a GET request to the target and returns
// the status code, length and an error
func (d *GobusterDir) get(url string) (*int, *int64, error) {
	return d.http.Get(url, "", d.options.Cookies)
}

// NewGobusterDir creates a new initialized GobusterDir
func NewGobusterDir(cont context.Context, globalopts *libgobuster.Options, opts *OptionsDir) (*GobusterDir, error) {
	if globalopts == nil {
		return nil, fmt.Errorf("please provide valid global options")
	}

	if opts == nil {
		return nil, fmt.Errorf("please provide valid plugin options")
	}

	g := GobusterDir{
		options:    opts,
		globalopts: globalopts,
	}

	httpOpts := libgobuster.HTTPOptions{
		Proxy:          opts.Proxy,
		FollowRedirect: opts.FollowRedirect,
		InsecureSSL:    opts.InsecureSSL,
		Timeout:        opts.Timeout,
		Username:       opts.Username,
		Password:       opts.Password,
		UserAgent:      opts.UserAgent,
	}

	h, err := libgobuster.NewHTTPClient(cont, &httpOpts)
	if err != nil {
		return nil, err
	}
	g.http = h
	return &g, nil
}

// PreRun is the pre run implementation of gobusterdir
func (d *GobusterDir) PreRun() error {
	// add trailing slash
	if !strings.HasSuffix(d.options.URL, "/") {
		d.options.URL = fmt.Sprintf("%s/", d.options.URL)
	}

	_, _, err := d.get(d.options.URL)
	if err != nil {
		return fmt.Errorf("unable to connect to %s: %v", d.options.URL, err)
	}

	guid := uuid.New()
	url := fmt.Sprintf("%s%s", d.options.URL, guid)
	wildcardResp, _, err := d.get(url)
	if err != nil {
		return err
	}

	if d.options.StatusCodesParsed.Contains(*wildcardResp) {
		log.Printf("[-] Wildcard response found: %s => %d", url, *wildcardResp)
		if !d.options.WildcardForced {
			return fmt.Errorf("To force processing of Wildcard responses, specify the '--wildcard' switch.")
		}
	}

	return nil
}

// Run is the process implementation of gobusterdir
func (d *GobusterDir) Run(word string) ([]libgobuster.Result, error) {
	suffix := ""
	if d.options.UseSlash {
		suffix = "/"
	}

	// Try the DIR first
	url := fmt.Sprintf("%s%s%s", d.options.URL, word, suffix)
	dirResp, dirSize, err := d.get(url)
	if err != nil {
		return nil, err
	}
	var ret []libgobuster.Result
	if dirResp != nil {
		resultStatus := libgobuster.StatusMissed
		if d.options.StatusCodesParsed.Contains(*dirResp) {
			resultStatus = libgobuster.StatusFound
		}

		if resultStatus == libgobuster.StatusFound || d.globalopts.Verbose {
			ret = append(ret, libgobuster.Result{
				Entity:     fmt.Sprintf("%s%s", word, suffix),
				StatusCode: *dirResp,
				Size:       dirSize,
				Status:     resultStatus,
			})
		}
	}

	// Follow up with files using each ext.
	for ext := range d.options.ExtensionsParsed.Set {
		file := fmt.Sprintf("%s.%s", word, ext)
		url = fmt.Sprintf("%s%s", d.options.URL, file)
		fileResp, fileSize, err := d.get(url)
		if err != nil {
			return nil, err
		}

		if fileResp != nil {
			resultStatus := libgobuster.StatusMissed
			if d.options.StatusCodesParsed.Contains(*fileResp) {
				resultStatus = libgobuster.StatusFound
			}

			if resultStatus == libgobuster.StatusFound || d.globalopts.Verbose {
				ret = append(ret, libgobuster.Result{
					Entity:     file,
					StatusCode: *fileResp,
					Size:       fileSize,
					Status:     resultStatus,
				})
			}

		}
	}
	return ret, nil
}

// ResultToString is the to string implementation of gobusterdir
func (d *GobusterDir) ResultToString(r *libgobuster.Result) (*string, error) {
	buf := &bytes.Buffer{}

	// Prefix if we're in verbose mode
	if d.globalopts.Verbose {
		if r.Status == libgobuster.StatusFound {
			if _, err := fmt.Fprintf(buf, "Found: "); err != nil {
				return nil, err
			}
		} else if r.Status == libgobuster.StatusMissed {
			if _, err := fmt.Fprintf(buf, "Missed: "); err != nil {
				return nil, err
			}
		} else {
			return nil, fmt.Errorf("unknown status %d", r.Status)
		}
	}

	if d.options.Expanded {
		if _, err := fmt.Fprintf(buf, "%s", d.options.URL); err != nil {
			return nil, err
		}
	} else {
		if _, err := fmt.Fprintf(buf, "/"); err != nil {
			return nil, err
		}
	}
	if _, err := fmt.Fprintf(buf, "%s", r.Entity); err != nil {
		return nil, err
	}

	if !d.options.NoStatus {
		if _, err := fmt.Fprintf(buf, " (Status: %d)", r.StatusCode); err != nil {
			return nil, err
		}
	}

	if r.Size != nil {
		if _, err := fmt.Fprintf(buf, " [Size: %d]", *r.Size); err != nil {
			return nil, err
		}
	}
	if _, err := fmt.Fprintf(buf, "\n"); err != nil {
		return nil, err
	}

	s := buf.String()
	return &s, nil
}

// GetConfigString returns the string representation of the current config
func (d *GobusterDir) GetConfigString() (string, error) {
	var buffer bytes.Buffer
	bw := bufio.NewWriter(&buffer)
	tw := tabwriter.NewWriter(bw, 0, 5, 3, ' ', 0)
	o := d.options
	if _, err := fmt.Fprintf(tw, "[+] Url:\t%s\n", o.URL); err != nil {
		return "", err
	}
	if _, err := fmt.Fprintf(tw, "[+] Threads:\t%d\n", d.globalopts.Threads); err != nil {
		return "", err
	}

	wordlist := "stdin (pipe)"
	if d.globalopts.Wordlist != "-" {
		wordlist = d.globalopts.Wordlist
	}
	if _, err := fmt.Fprintf(tw, "[+] Wordlist:\t%s\n", wordlist); err != nil {
		return "", err
	}

	if _, err := fmt.Fprintf(tw, "[+] Status codes:\t%s\n", o.StatusCodesParsed.Stringify()); err != nil {
		return "", err
	}

	if o.Proxy != "" {
		if _, err := fmt.Fprintf(tw, "[+] Proxy:\t%s\n", o.Proxy); err != nil {
			return "", err
		}
	}

	if o.Cookies != "" {
		if _, err := fmt.Fprintf(tw, "[+] Cookies:\t%s\n", o.Cookies); err != nil {
			return "", err
		}
	}

	if o.UserAgent != "" {
		if _, err := fmt.Fprintf(tw, "[+] User Agent:\t%s\n", o.UserAgent); err != nil {
			return "", err
		}
	}

	if o.IncludeLength {
		if _, err := fmt.Fprintf(tw, "[+] Show length:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.Username != "" {
		if _, err := fmt.Fprintf(tw, "[+] Auth User:\t%s\n", o.Username); err != nil {
			return "", err
		}
	}

	if o.Extensions != "" {
		if _, err := fmt.Fprintf(tw, "[+] Extensions:\t%s\n", o.ExtensionsParsed.Stringify()); err != nil {
			return "", err
		}
	}

	if o.UseSlash {
		if _, err := fmt.Fprintf(tw, "[+] Add Slash:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.FollowRedirect {
		if _, err := fmt.Fprintf(tw, "[+] Follow Redir:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.Expanded {
		if _, err := fmt.Fprintf(tw, "[+] Expanded:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if o.NoStatus {
		if _, err := fmt.Fprintf(tw, "[+] No status:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if d.globalopts.Verbose {
		if _, err := fmt.Fprintf(tw, "[+] Verbose:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if _, err := fmt.Fprintf(tw, "[+] Timeout:\t%s\n", o.Timeout.String()); err != nil {
		return "", err
	}

	if err := tw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %v", err)
	}

	if err := bw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %v", err)
	}

	return strings.TrimSpace(buffer.String()), nil
}
<|endoftext|>


<commit_before>// Copyright 2013 The Go 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 redirect provides hooks to register HTTP handlers that redirect old
// godoc paths to their new equivalents and assist in accessing the issue
// tracker, wiki, code review system, etc.
package redirect

import (
	"net/http"
	"regexp"
	"strings"
)

// Register registers HTTP handlers that redirect old godoc paths to their new
// equivalents and assist in accessing the issue tracker, wiki, code review
// system, etc. If mux is nil it uses http.DefaultServeMux.
func Register(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	handlePathRedirects(mux, pkgRedirects, "/pkg/")
	handlePathRedirects(mux, cmdRedirects, "/cmd/")
	for prefix, redirect := range prefixHelpers {
		p := "/" + prefix + "/"
		mux.Handle(p, PrefixHandler(p, redirect))
	}
	for path, redirect := range redirects {
		mux.Handle(path, Handler(redirect))
	}
	// NB: /src/pkg (sans trailing slash) is the index of packages.
	mux.HandleFunc("/src/pkg/", srcPkgHandler)
	mux.HandleFunc("/cl/", clHandler)
}

func handlePathRedirects(mux *http.ServeMux, redirects map[string]string, prefix string) {
	for source, target := range redirects {
		h := Handler(prefix + target + "/")
		p := prefix + source
		mux.Handle(p, h)
		mux.Handle(p+"/", h)
	}
}

// Packages that were renamed between r60 and go1.
var pkgRedirects = map[string]string{
	"asn1":              "encoding/asn1",
	"big":               "math/big",
	"cmath":             "math/cmplx",
	"csv":               "encoding/csv",
	"exec":              "os/exec",
	"exp/template/html": "html/template",
	"gob":               "encoding/gob",
	"http":              "net/http",
	"http/cgi":          "net/http/cgi",
	"http/fcgi":         "net/http/fcgi",
	"http/httptest":     "net/http/httptest",
	"http/pprof":        "net/http/pprof",
	"json":              "encoding/json",
	"mail":              "net/mail",
	"rand":              "math/rand",
	"rpc":               "net/rpc",
	"rpc/jsonrpc":       "net/rpc/jsonrpc",
	"scanner":           "text/scanner",
	"smtp":              "net/smtp",
	"tabwriter":         "text/tabwriter",
	"template":          "text/template",
	"template/parse":    "text/template/parse",
	"url":               "net/url",
	"utf16":             "unicode/utf16",
	"utf8":              "unicode/utf8",
	"xml":               "encoding/xml",
}

// Commands that were renamed between r60 and go1.
var cmdRedirects = map[string]string{
	"gofix":     "fix",
	"goinstall": "go",
	"gopack":    "pack",
	"gotest":    "go",
	"govet":     "vet",
	"goyacc":    "yacc",
}

var redirects = map[string]string{
	"/blog":       "/blog/",
	"/build":      "http://build.golang.org",
	"/change":     "https://code.google.com/p/go/source/list",
	"/cl":         "https://gocodereview.appspot.com/",
	"/cmd/godoc/": "http://godoc.org/golang.org/x/tools/cmd/godoc/",
	"/cmd/vet/":   "http://godoc.org/golang.org/x/tools/cmd/vet/",
	"/issue":      "https://code.google.com/p/go/issues",
	"/issue/new":  "https://code.google.com/p/go/issues/entry",
	"/issues":     "https://code.google.com/p/go/issues",
	"/play":       "http://play.golang.org",

	// In Go 1.2 the references page is part of /doc/.
	"/ref": "/doc/#references",
	// This next rule clobbers /ref/spec and /ref/mem.
	// TODO(adg): figure out what to do here, if anything.
	// "/ref/": "/doc/#references",

	// Be nice to people who are looking in the wrong place.
	"/doc/mem":  "/ref/mem",
	"/doc/spec": "/ref/spec",

	"/talks": "http://talks.golang.org",
	"/tour":  "http://tour.golang.org",
	"/wiki":  "https://code.google.com/p/go-wiki/w/list",

	"/doc/articles/c_go_cgo.html":                    "/blog/c-go-cgo",
	"/doc/articles/concurrency_patterns.html":        "/blog/go-concurrency-patterns-timing-out-and",
	"/doc/articles/defer_panic_recover.html":         "/blog/defer-panic-and-recover",
	"/doc/articles/error_handling.html":              "/blog/error-handling-and-go",
	"/doc/articles/gobs_of_data.html":                "/blog/gobs-of-data",
	"/doc/articles/godoc_documenting_go_code.html":   "/blog/godoc-documenting-go-code",
	"/doc/articles/gos_declaration_syntax.html":      "/blog/gos-declaration-syntax",
	"/doc/articles/image_draw.html":                  "/blog/go-imagedraw-package",
	"/doc/articles/image_package.html":               "/blog/go-image-package",
	"/doc/articles/json_and_go.html":                 "/blog/json-and-go",
	"/doc/articles/json_rpc_tale_of_interfaces.html": "/blog/json-rpc-tale-of-interfaces",
	"/doc/articles/laws_of_reflection.html":          "/blog/laws-of-reflection",
	"/doc/articles/slices_usage_and_internals.html":  "/blog/go-slices-usage-and-internals",
	"/doc/go_for_cpp_programmers.html":               "https://code.google.com/p/go-wiki/wiki/GoForCPPProgrammers",
	"/doc/go_tutorial.html":                          "http://tour.golang.org/",
}

var prefixHelpers = map[string]string{
	"change": "https://code.google.com/p/go/source/detail?r=",
	"issue":  "https://code.google.com/p/go/issues/detail?id=",
	"play":   "http://play.golang.org/",
	"talks":  "http://talks.golang.org/",
	"wiki":   "https://code.google.com/p/go-wiki/wiki/",
}

func Handler(target string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, target, http.StatusMovedPermanently)
	})
}

var validId = regexp.MustCompile(`^[A-Za-z0-9-]*$`)

func PrefixHandler(prefix, baseURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if p := r.URL.Path; p == prefix {
			// redirect /prefix/ to /prefix
			http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
			return
		}
		id := r.URL.Path[len(prefix):]
		if !validId.MatchString(id) {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		target := baseURL + id
		http.Redirect(w, r, target, http.StatusFound)
	})
}

// Redirect requests from the old "/src/pkg/foo" to the new "/src/foo".
// See http://golang.org/s/go14nopkg
func srcPkgHandler(w http.ResponseWriter, r *http.Request) {
	r.URL.Path = "/src/" + r.URL.Path[len("/src/pkg/"):]
	http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
}

func clHandler(w http.ResponseWriter, r *http.Request) {
	const prefix = "/cl/"
	if p := r.URL.Path; p == prefix {
		// redirect /prefix/ to /prefix
		http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
		return
	}
	id := r.URL.Path[len(prefix):]
	if !validId.MatchString(id) {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	target := ""
	if strings.HasPrefix(id, "I") {
		target = "https://go-review.googlesource.com/#/q/" + id
	} else {
		target = "https://codereview.appspot.com/" + id
	}
	http.Redirect(w, r, target, http.StatusFound)
}
<commit_msg>x/tools/godoc/redirect: point to new git/github URLs<commit_after>// Copyright 2013 The Go 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 redirect provides hooks to register HTTP handlers that redirect old
// godoc paths to their new equivalents and assist in accessing the issue
// tracker, wiki, code review system, etc.
package redirect

import (
	"net/http"
	"regexp"
	"strings"
)

// Register registers HTTP handlers that redirect old godoc paths to their new
// equivalents and assist in accessing the issue tracker, wiki, code review
// system, etc. If mux is nil it uses http.DefaultServeMux.
func Register(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	handlePathRedirects(mux, pkgRedirects, "/pkg/")
	handlePathRedirects(mux, cmdRedirects, "/cmd/")
	for prefix, redirect := range prefixHelpers {
		p := "/" + prefix + "/"
		mux.Handle(p, PrefixHandler(p, redirect))
	}
	for path, redirect := range redirects {
		mux.Handle(path, Handler(redirect))
	}
	// NB: /src/pkg (sans trailing slash) is the index of packages.
	mux.HandleFunc("/src/pkg/", srcPkgHandler)
	mux.HandleFunc("/cl/", clHandler)
}

func handlePathRedirects(mux *http.ServeMux, redirects map[string]string, prefix string) {
	for source, target := range redirects {
		h := Handler(prefix + target + "/")
		p := prefix + source
		mux.Handle(p, h)
		mux.Handle(p+"/", h)
	}
}

// Packages that were renamed between r60 and go1.
var pkgRedirects = map[string]string{
	"asn1":              "encoding/asn1",
	"big":               "math/big",
	"cmath":             "math/cmplx",
	"csv":               "encoding/csv",
	"exec":              "os/exec",
	"exp/template/html": "html/template",
	"gob":               "encoding/gob",
	"http":              "net/http",
	"http/cgi":          "net/http/cgi",
	"http/fcgi":         "net/http/fcgi",
	"http/httptest":     "net/http/httptest",
	"http/pprof":        "net/http/pprof",
	"json":              "encoding/json",
	"mail":              "net/mail",
	"rand":              "math/rand",
	"rpc":               "net/rpc",
	"rpc/jsonrpc":       "net/rpc/jsonrpc",
	"scanner":           "text/scanner",
	"smtp":              "net/smtp",
	"tabwriter":         "text/tabwriter",
	"template":          "text/template",
	"template/parse":    "text/template/parse",
	"url":               "net/url",
	"utf16":             "unicode/utf16",
	"utf8":              "unicode/utf8",
	"xml":               "encoding/xml",
}

// Commands that were renamed between r60 and go1.
var cmdRedirects = map[string]string{
	"gofix":     "fix",
	"goinstall": "go",
	"gopack":    "pack",
	"gotest":    "go",
	"govet":     "vet",
	"goyacc":    "yacc",
}

var redirects = map[string]string{
	"/blog":       "/blog/",
	"/build":      "http://build.golang.org",
	"/change":     "https://go.googlesource.com/go",
	"/cl":         "https://go-review.googlesource.com",
	"/cmd/godoc/": "http://godoc.org/golang.org/x/tools/cmd/godoc/",
	"/cmd/vet/":   "http://godoc.org/golang.org/x/tools/cmd/vet/",
	"/issue":      "https://github.com/golang/go/issues",
	"/issue/new":  "https://github.com/golang/go/issues/new",
	"/issues":     "https://github.com/golang/go/issues",
	"/play":       "http://play.golang.org",

	// In Go 1.2 the references page is part of /doc/.
	"/ref": "/doc/#references",
	// This next rule clobbers /ref/spec and /ref/mem.
	// TODO(adg): figure out what to do here, if anything.
	// "/ref/": "/doc/#references",

	// Be nice to people who are looking in the wrong place.
	"/doc/mem":  "/ref/mem",
	"/doc/spec": "/ref/spec",

	"/talks": "http://talks.golang.org",
	"/tour":  "http://tour.golang.org",
	"/wiki":  "https://github.com/golang/go/wiki",

	"/doc/articles/c_go_cgo.html":                    "/blog/c-go-cgo",
	"/doc/articles/concurrency_patterns.html":        "/blog/go-concurrency-patterns-timing-out-and",
	"/doc/articles/defer_panic_recover.html":         "/blog/defer-panic-and-recover",
	"/doc/articles/error_handling.html":              "/blog/error-handling-and-go",
	"/doc/articles/gobs_of_data.html":                "/blog/gobs-of-data",
	"/doc/articles/godoc_documenting_go_code.html":   "/blog/godoc-documenting-go-code",
	"/doc/articles/gos_declaration_syntax.html":      "/blog/gos-declaration-syntax",
	"/doc/articles/image_draw.html":                  "/blog/go-imagedraw-package",
	"/doc/articles/image_package.html":               "/blog/go-image-package",
	"/doc/articles/json_and_go.html":                 "/blog/json-and-go",
	"/doc/articles/json_rpc_tale_of_interfaces.html": "/blog/json-rpc-tale-of-interfaces",
	"/doc/articles/laws_of_reflection.html":          "/blog/laws-of-reflection",
	"/doc/articles/slices_usage_and_internals.html":  "/blog/go-slices-usage-and-internals",
	"/doc/go_for_cpp_programmers.html":               "https://code.google.com/p/go-wiki/wiki/GoForCPPProgrammers",
	"/doc/go_tutorial.html":                          "http://tour.golang.org/",
}

var prefixHelpers = map[string]string{
	// TODO(adg): add redirects from known hg hashes to the git equivalents
	// and switch this to point to "https://go.googlesource.com/go/+/".
	// (We can only change this once we know all the new git hashes.)
	"change": "https://code.google.com/p/go/source/detail?r=",

	"issue": "https://github.com/golang/go/issues/",
	"play":  "http://play.golang.org/",
	"talks": "http://talks.golang.org/",
	"wiki":  "https://github.com/golang/go/wiki/",
}

func Handler(target string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, target, http.StatusMovedPermanently)
	})
}

var validId = regexp.MustCompile(`^[A-Za-z0-9-]*$`)

func PrefixHandler(prefix, baseURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if p := r.URL.Path; p == prefix {
			// redirect /prefix/ to /prefix
			http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
			return
		}
		id := r.URL.Path[len(prefix):]
		if !validId.MatchString(id) {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		target := baseURL + id
		http.Redirect(w, r, target, http.StatusFound)
	})
}

// Redirect requests from the old "/src/pkg/foo" to the new "/src/foo".
// See http://golang.org/s/go14nopkg
func srcPkgHandler(w http.ResponseWriter, r *http.Request) {
	r.URL.Path = "/src/" + r.URL.Path[len("/src/pkg/"):]
	http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
}

func clHandler(w http.ResponseWriter, r *http.Request) {
	const prefix = "/cl/"
	if p := r.URL.Path; p == prefix {
		// redirect /prefix/ to /prefix
		http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
		return
	}
	id := r.URL.Path[len(prefix):]
	if !validId.MatchString(id) {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	target := ""
	if strings.HasPrefix(id, "I") {
		target = "https://go-review.googlesource.com/#/q/" + id
	} else {
		target = "https://codereview.appspot.com/" + id
	}
	http.Redirect(w, r, target, http.StatusFound)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"encoding/json"
	"runtime"

	"golang.org/x/net/context"
)

// IntrospectionOptions are the options used when introspecting the Channel.
type IntrospectionOptions struct {
	// IncludeExchanges will include all the IDs in the message exchanges.
	IncludeExchanges bool `json:"includeExchanges"`

	// IncludeEmptyPeers will include peers, even if they have no connections.
	IncludeEmptyPeers bool `json:"includeEmptyPeers"`
}

// RuntimeState is a snapshot of the runtime state for a channel.
type RuntimeState struct {
	// LocalPeer is the local peer information (service name, host-port, etc).
	LocalPeer LocalPeerInfo `json:"localPeer"`

	// SubChannels contains information about any subchannels.
	SubChannels map[string]SubChannelRuntimeState `json:"subChannels"`

	// RootPeers contains information about all the peers on this channel and their connections.
	RootPeers map[string]PeerRuntimeState `json:"rootPeers"`

	// Peers is the list of shared peers for this channel.
	Peers []string `json:"peers"`
}

// GoRuntimeStateOptions are the options used when getting Go runtime state.
type GoRuntimeStateOptions struct {
	// IncludeGoStacks will include all goroutine stacks.
	IncludeGoStacks bool `json:"includeGoStacks"`
}

// GoRuntimeState is a snapshot of runtime stats from the runtime.
type GoRuntimeState struct {
	MemStats      runtime.MemStats `json:"memStats"`
	NumGoroutines int              `json:"numGoRoutines"`
	NumCPU        int              `json:"numCPU"`
	NumCGo        int64            `json:"numCGo"`
	GoStacks      []byte           `json:"goStacks,omitempty"`
}

// SubChannelRuntimeState is the runtime state for a subchannel.
type SubChannelRuntimeState struct {
	Service  string `json:"service"`
	Isolated bool   `json:"isolated"`
	// IsolatedPeers is the list of all isolated peers for this channel.
	IsolatedPeers []string `json:"isolatedPeers,omitempty"`
}

// ConnectionRuntimeState is the runtime state for a single connection.
type ConnectionRuntimeState struct {
	ID               uint32               `json:"id"`
	ConnectionState  string               `json:"connectionState"`
	LocalHostPort    string               `json:"localHostPort"`
	RemoteHostPort   string               `json:"remoteHostPort"`
	InboundExchange  ExchangeRuntimeState `json:"inboundExchange"`
	OutboundExchange ExchangeRuntimeState `json:"outboundExchange"`
}

// ExchangeRuntimeState is the runtime state for a message exchange set.
type ExchangeRuntimeState struct {
	Name      string   `json:"name"`
	Count     int      `json:"count"`
	Exchanges []uint32 `json:"exchanges,omitempty"`
}

// PeerRuntimeState is the runtime state for a single peer.
type PeerRuntimeState struct {
	HostPort            string                   `json:"hostPort"`
	OutboundConnections []ConnectionRuntimeState `json:"outboundConnections"`
	InboundConnections  []ConnectionRuntimeState `json:"inboundConnections"`
}

// IntrospectState returns the RuntimeState for this channel.
// Note: this is purely for debugging and monitoring, and may slow down your Channel.
func (ch *Channel) IntrospectState(opts *IntrospectionOptions) *RuntimeState {
	return &RuntimeState{
		LocalPeer:   ch.PeerInfo(),
		SubChannels: ch.subChannels.IntrospectState(opts),
		RootPeers:   ch.rootPeers().IntrospectState(opts),
		Peers:       ch.Peers().IntrospectList(opts),
	}
}

// IntrospectState returns the runtime state of the peer list.
func (l *PeerList) IntrospectState(opts *IntrospectionOptions) map[string]PeerRuntimeState {
	m := make(map[string]PeerRuntimeState)
	l.mut.RLock()
	for _, peer := range l.peers {
		peerState := peer.IntrospectState(opts)
		if len(peerState.InboundConnections)+len(peerState.OutboundConnections) > 0 || opts.IncludeEmptyPeers {
			m[peer.HostPort()] = peerState
		}
	}

	l.mut.RUnlock()
	return m
}

// IntrospectState returns the runtime state of the subchannels.
func (subChMap *subChannelMap) IntrospectState(opts *IntrospectionOptions) map[string]SubChannelRuntimeState {
	m := make(map[string]SubChannelRuntimeState)
	subChMap.mut.RLock()
	for k, sc := range subChMap.subchannels {
		state := SubChannelRuntimeState{
			Service:  k,
			Isolated: sc.Isolated(),
		}
		if state.Isolated {
			state.IsolatedPeers = sc.Peers().IntrospectList(opts)
		}
		m[k] = state
	}
	subChMap.mut.RUnlock()
	return m
}

// IntrospectState returns the runtime state for this peer.
func (p *Peer) IntrospectState(opts *IntrospectionOptions) PeerRuntimeState {
	p.mut.RLock()

	hostPort := p.hostPort
	inboundConns := make([]ConnectionRuntimeState, len(p.inboundConnections))
	outboundConns := make([]ConnectionRuntimeState, len(p.outboundConnections))

	for i, conn := range p.outboundConnections {
		outboundConns[i] = conn.IntrospectState(opts)
	}

	for i, conn := range p.inboundConnections {
		inboundConns[i] = conn.IntrospectState(opts)
	}

	p.mut.RUnlock()

	return PeerRuntimeState{
		HostPort:            hostPort,
		InboundConnections:  inboundConns,
		OutboundConnections: outboundConns,
	}
}

// IntrospectState returns the runtime state for this connection.
func (c *Connection) IntrospectState(opts *IntrospectionOptions) ConnectionRuntimeState {
	return ConnectionRuntimeState{
		ID:               c.connID,
		ConnectionState:  c.state.String(),
		LocalHostPort:    c.conn.LocalAddr().String(),
		RemoteHostPort:   c.conn.RemoteAddr().String(),
		InboundExchange:  c.inbound.IntrospectState(opts),
		OutboundExchange: c.outbound.IntrospectState(opts),
	}
}

// IntrospectState returns the runtime state for this messsage exchange set.
func (mexset *messageExchangeSet) IntrospectState(opts *IntrospectionOptions) ExchangeRuntimeState {
	mexset.mut.RLock()
	state := ExchangeRuntimeState{
		Name:  mexset.name,
		Count: len(mexset.exchanges),
	}

	if opts.IncludeExchanges {
		state.Exchanges = make([]uint32, 0, len(mexset.exchanges))
		for k := range mexset.exchanges {
			state.Exchanges = append(state.Exchanges, k)
		}
	}

	mexset.mut.RUnlock()

	return state
}

func getStacks() []byte {
	var buf []byte
	for n := 4096; n < 10*1024*1024; n *= 2 {
		buf = make([]byte, n)
		stackLen := runtime.Stack(buf, true /* all */)
		if stackLen < n {
			return buf
		}
	}

	// return the first 10MB of stacks if we have more than 10MB.
	return buf
}
func (ch *Channel) handleIntrospection(arg3 []byte) interface{} {
	var opts IntrospectionOptions
	json.Unmarshal(arg3, &opts)
	return ch.IntrospectState(&opts)
}

// IntrospectList returns the list of peers (host:port) in this peer list.
func (l *PeerList) IntrospectList(opts *IntrospectionOptions) []string {
	var peers []string
	l.mut.RLock()
	for peer := range l.peersByHostPort {
		peers = append(peers, peer)
	}
	l.mut.RUnlock()

	return peers
}

func handleInternalRuntime(arg3 []byte) interface{} {
	var opts GoRuntimeStateOptions
	json.Unmarshal(arg3, &opts)

	state := GoRuntimeState{
		NumGoroutines: runtime.NumGoroutine(),
		NumCPU:        runtime.NumCPU(),
		NumCGo:        runtime.NumCgoCall(),
	}
	runtime.ReadMemStats(&state.MemStats)
	if opts.IncludeGoStacks {
		state.GoStacks = getStacks()
	}

	return state
}

// registerInternal registers the following internal handlers which return runtime state:
//  _gometa_introspect: TChannel internal state.
//  _gometa_runtime: Golang runtime stats.
func (ch *Channel) registerInternal() {
	endpoints := []struct {
		name    string
		handler func([]byte) interface{}
	}{
		{"_gometa_introspect", ch.handleIntrospection},
		{"_gometa_runtime", handleInternalRuntime},
	}

	for _, ep := range endpoints {
		// We need ep in our closure.
		ep := ep
		handler := func(ctx context.Context, call *InboundCall) {
			var arg2, arg3 []byte
			if err := NewArgReader(call.Arg2Reader()).Read(&arg2); err != nil {
				return
			}
			if err := NewArgReader(call.Arg3Reader()).Read(&arg3); err != nil {
				return
			}
			if err := NewArgWriter(call.Response().Arg2Writer()).Write(nil); err != nil {
				return
			}
			NewArgWriter(call.Response().Arg3Writer()).WriteJSON(ep.handler(arg3))
		}
		ch.Register(HandlerFunc(handler), ep.name)
	}
}
<commit_msg>add getConnectionRuntimeState to reduse code duplication<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"encoding/json"
	"runtime"

	"golang.org/x/net/context"
)

// IntrospectionOptions are the options used when introspecting the Channel.
type IntrospectionOptions struct {
	// IncludeExchanges will include all the IDs in the message exchanges.
	IncludeExchanges bool `json:"includeExchanges"`

	// IncludeEmptyPeers will include peers, even if they have no connections.
	IncludeEmptyPeers bool `json:"includeEmptyPeers"`
}

// RuntimeState is a snapshot of the runtime state for a channel.
type RuntimeState struct {
	// LocalPeer is the local peer information (service name, host-port, etc).
	LocalPeer LocalPeerInfo `json:"localPeer"`

	// SubChannels contains information about any subchannels.
	SubChannels map[string]SubChannelRuntimeState `json:"subChannels"`

	// RootPeers contains information about all the peers on this channel and their connections.
	RootPeers map[string]PeerRuntimeState `json:"rootPeers"`

	// Peers is the list of shared peers for this channel.
	Peers []string `json:"peers"`
}

// GoRuntimeStateOptions are the options used when getting Go runtime state.
type GoRuntimeStateOptions struct {
	// IncludeGoStacks will include all goroutine stacks.
	IncludeGoStacks bool `json:"includeGoStacks"`
}

// GoRuntimeState is a snapshot of runtime stats from the runtime.
type GoRuntimeState struct {
	MemStats      runtime.MemStats `json:"memStats"`
	NumGoroutines int              `json:"numGoRoutines"`
	NumCPU        int              `json:"numCPU"`
	NumCGo        int64            `json:"numCGo"`
	GoStacks      []byte           `json:"goStacks,omitempty"`
}

// SubChannelRuntimeState is the runtime state for a subchannel.
type SubChannelRuntimeState struct {
	Service  string `json:"service"`
	Isolated bool   `json:"isolated"`
	// IsolatedPeers is the list of all isolated peers for this channel.
	IsolatedPeers []string `json:"isolatedPeers,omitempty"`
}

// ConnectionRuntimeState is the runtime state for a single connection.
type ConnectionRuntimeState struct {
	ID               uint32               `json:"id"`
	ConnectionState  string               `json:"connectionState"`
	LocalHostPort    string               `json:"localHostPort"`
	RemoteHostPort   string               `json:"remoteHostPort"`
	InboundExchange  ExchangeRuntimeState `json:"inboundExchange"`
	OutboundExchange ExchangeRuntimeState `json:"outboundExchange"`
}

// ExchangeRuntimeState is the runtime state for a message exchange set.
type ExchangeRuntimeState struct {
	Name      string   `json:"name"`
	Count     int      `json:"count"`
	Exchanges []uint32 `json:"exchanges,omitempty"`
}

// PeerRuntimeState is the runtime state for a single peer.
type PeerRuntimeState struct {
	HostPort            string                   `json:"hostPort"`
	OutboundConnections []ConnectionRuntimeState `json:"outboundConnections"`
	InboundConnections  []ConnectionRuntimeState `json:"inboundConnections"`
}

// IntrospectState returns the RuntimeState for this channel.
// Note: this is purely for debugging and monitoring, and may slow down your Channel.
func (ch *Channel) IntrospectState(opts *IntrospectionOptions) *RuntimeState {
	return &RuntimeState{
		LocalPeer:   ch.PeerInfo(),
		SubChannels: ch.subChannels.IntrospectState(opts),
		RootPeers:   ch.rootPeers().IntrospectState(opts),
		Peers:       ch.Peers().IntrospectList(opts),
	}
}

// IntrospectState returns the runtime state of the peer list.
func (l *PeerList) IntrospectState(opts *IntrospectionOptions) map[string]PeerRuntimeState {
	m := make(map[string]PeerRuntimeState)
	l.mut.RLock()
	for _, peer := range l.peers {
		peerState := peer.IntrospectState(opts)
		if len(peerState.InboundConnections)+len(peerState.OutboundConnections) > 0 || opts.IncludeEmptyPeers {
			m[peer.HostPort()] = peerState
		}
	}

	l.mut.RUnlock()
	return m
}

// IntrospectState returns the runtime state of the subchannels.
func (subChMap *subChannelMap) IntrospectState(opts *IntrospectionOptions) map[string]SubChannelRuntimeState {
	m := make(map[string]SubChannelRuntimeState)
	subChMap.mut.RLock()
	for k, sc := range subChMap.subchannels {
		state := SubChannelRuntimeState{
			Service:  k,
			Isolated: sc.Isolated(),
		}
		if state.Isolated {
			state.IsolatedPeers = sc.Peers().IntrospectList(opts)
		}
		m[k] = state
	}
	subChMap.mut.RUnlock()
	return m
}

func getConnectionRuntimeState(conns []*Connection, opts *IntrospectionOptions) []ConnectionRuntimeState {
	connStates := make([]ConnectionRuntimeState, len(conns))

	for i, conn := range conns {
		connStates[i] = conn.IntrospectState(opts)
	}

	return connStates
}

// IntrospectState returns the runtime state for this peer.
func (p *Peer) IntrospectState(opts *IntrospectionOptions) PeerRuntimeState {
	p.mut.RLock()

	hostPort := p.hostPort
	inboundConns := getConnectionRuntimeState(p.inboundConnections, opts)
	outboundConns := getConnectionRuntimeState(p.outboundConnections, opts)

	p.mut.RUnlock()

	return PeerRuntimeState{
		HostPort:            hostPort,
		InboundConnections:  inboundConns,
		OutboundConnections: outboundConns,
	}
}

// IntrospectState returns the runtime state for this connection.
func (c *Connection) IntrospectState(opts *IntrospectionOptions) ConnectionRuntimeState {
	return ConnectionRuntimeState{
		ID:               c.connID,
		ConnectionState:  c.state.String(),
		LocalHostPort:    c.conn.LocalAddr().String(),
		RemoteHostPort:   c.conn.RemoteAddr().String(),
		InboundExchange:  c.inbound.IntrospectState(opts),
		OutboundExchange: c.outbound.IntrospectState(opts),
	}
}

// IntrospectState returns the runtime state for this messsage exchange set.
func (mexset *messageExchangeSet) IntrospectState(opts *IntrospectionOptions) ExchangeRuntimeState {
	mexset.mut.RLock()
	state := ExchangeRuntimeState{
		Name:  mexset.name,
		Count: len(mexset.exchanges),
	}

	if opts.IncludeExchanges {
		state.Exchanges = make([]uint32, 0, len(mexset.exchanges))
		for k := range mexset.exchanges {
			state.Exchanges = append(state.Exchanges, k)
		}
	}

	mexset.mut.RUnlock()

	return state
}

func getStacks() []byte {
	var buf []byte
	for n := 4096; n < 10*1024*1024; n *= 2 {
		buf = make([]byte, n)
		stackLen := runtime.Stack(buf, true /* all */)
		if stackLen < n {
			return buf
		}
	}

	// return the first 10MB of stacks if we have more than 10MB.
	return buf
}
func (ch *Channel) handleIntrospection(arg3 []byte) interface{} {
	var opts IntrospectionOptions
	json.Unmarshal(arg3, &opts)
	return ch.IntrospectState(&opts)
}

// IntrospectList returns the list of peers (host:port) in this peer list.
func (l *PeerList) IntrospectList(opts *IntrospectionOptions) []string {
	var peers []string
	l.mut.RLock()
	for peer := range l.peersByHostPort {
		peers = append(peers, peer)
	}
	l.mut.RUnlock()

	return peers
}

func handleInternalRuntime(arg3 []byte) interface{} {
	var opts GoRuntimeStateOptions
	json.Unmarshal(arg3, &opts)

	state := GoRuntimeState{
		NumGoroutines: runtime.NumGoroutine(),
		NumCPU:        runtime.NumCPU(),
		NumCGo:        runtime.NumCgoCall(),
	}
	runtime.ReadMemStats(&state.MemStats)
	if opts.IncludeGoStacks {
		state.GoStacks = getStacks()
	}

	return state
}

// registerInternal registers the following internal handlers which return runtime state:
//  _gometa_introspect: TChannel internal state.
//  _gometa_runtime: Golang runtime stats.
func (ch *Channel) registerInternal() {
	endpoints := []struct {
		name    string
		handler func([]byte) interface{}
	}{
		{"_gometa_introspect", ch.handleIntrospection},
		{"_gometa_runtime", handleInternalRuntime},
	}

	for _, ep := range endpoints {
		// We need ep in our closure.
		ep := ep
		handler := func(ctx context.Context, call *InboundCall) {
			var arg2, arg3 []byte
			if err := NewArgReader(call.Arg2Reader()).Read(&arg2); err != nil {
				return
			}
			if err := NewArgReader(call.Arg3Reader()).Read(&arg3); err != nil {
				return
			}
			if err := NewArgWriter(call.Response().Arg2Writer()).Write(nil); err != nil {
				return
			}
			NewArgWriter(call.Response().Arg3Writer()).WriteJSON(ep.handler(arg3))
		}
		ch.Register(HandlerFunc(handler), ep.name)
	}
}
<|endoftext|>