<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * rpc/methods.go                                         *
 *                                                        *
 * hprose methods for Go.                                 *
 *                                                        *
 * LastModified: Sep 11, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package rpc

import (
	"reflect"
	"strings"
)

// MethodOptions is the options of the published service method
type MethodOptions struct {
	Mode      ResultMode
	Simple    bool
	Oneway    bool
	NameSpace string
}

// Method is the published service method
type Method struct {
	Function reflect.Value
	MethodOptions
}

// Methods is the published service methods
type Methods struct {
	MethodNames   []string
	RemoteMethods map[string]Method
}

// NewMethods is the constructor for Methods
func NewMethods() (methods *Methods) {
	methods = new(Methods)
	methods.MethodNames = make([]string, 0, 64)
	methods.RemoteMethods = make(map[string]Method)
	return
}

// AddFunction publish a func or bound method
// name is the method name
// function is a func or bound method
// options includes Mode, Simple, Oneway and NameSpace
func (methods *Methods) AddFunction(
	name string, function interface{}, options MethodOptions) {
	if name == "" {
		panic("name can't be empty")
	}
	if function == nil {
		panic("function can't be nil")
	}
	f, ok := function.(reflect.Value)
	if !ok {
		f = reflect.ValueOf(function)
	}
	if f.Kind() != reflect.Func {
		panic("function must be func or bound method")
	}
	if options.NameSpace != "" && name != "*" {
		name = options.NameSpace + "_" + name
	}
	methods.MethodNames = append(methods.MethodNames, name)
	methods.RemoteMethods[strings.ToLower(name)] = Method{f, options}
}

// AddFunctions is used for batch publishing service method
func (methods *Methods) AddFunctions(
	names []string, functions []interface{}, options MethodOptions) {
	count := len(names)
	if count != len(functions) {
		panic("names and functions must have the same length")
	}
	for i := 0; i < count; i++ {
		methods.AddFunction(names[i], functions[i], options)
	}
}

// AddMethod is used for publishing a method on the obj with an alias
func (methods *Methods) AddMethod(
	name string, obj interface{}, options MethodOptions, alias ...string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	f := reflect.ValueOf(obj).MethodByName(name)
	if len(alias) == 1 && alias[0] != "" {
		name = alias[0]
	}
	if f.CanInterface() {
		methods.AddFunction(name, f, options)
	}
}

// AddMethods is used for batch publishing methods on the obj with aliases
func (methods *Methods) AddMethods(
	names []string, obj interface{}, options MethodOptions, aliases ...[]string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	count := len(names)
	if len(aliases) == 1 {
		if len(aliases[0]) != count {
			panic("names and aliases must have the same length")
		}
		for i := 0; i < count; i++ {
			methods.AddMethod(names[i], obj, options, aliases[0][i])
		}
		return
	}
	for i := 0; i < count; i++ {
		methods.AddMethod(names[i], obj, options)
	}
}

func (methods *Methods) addInstanceMethods(
	v reflect.Value, t reflect.Type, options MethodOptions) {
	n := t.NumMethod()
	for i := 0; i < n; i++ {
		name := t.Method(i).Name
		method := v.Method(i)
		if method.CanInterface() {
			methods.AddFunction(name, method, options)
		}
	}
}

func getPtrTo(v reflect.Value, t reflect.Type) (reflect.Value, reflect.Type) {
	for ; t.Kind() == reflect.Ptr && !v.IsNil(); v = v.Elem() {
		t = t.Elem()
	}
	return v, t
}

func (methods *Methods) addFuncField(
	v reflect.Value, t reflect.Type, i int, options MethodOptions) {
	f := v.Field(i)
	name := t.Field(i).Name
	if f.CanInterface() && f.IsValid() {
		f, _ = getPtrTo(f, f.Type())
		if !f.IsNil() && f.Kind() == reflect.Func {
			methods.AddFunction(name, f, options)
		}
	}
}

// AddInstanceMethods is used for publishing all the public methods and func fields with options.
func (methods *Methods) AddInstanceMethods(
	obj interface{}, options MethodOptions) {
	if obj == nil {
		panic("obj can't be nil")
	}
	v := reflect.ValueOf(obj)
	t := v.Type()
	methods.addInstanceMethods(v, t, options)
	v, t = getPtrTo(v, t)
	if t.Kind() == reflect.Struct {
		n := t.NumField()
		for i := 0; i < n; i++ {
			methods.addFuncField(v, t, i, options)
		}
	}
}

func (methods *Methods) recursiveAddFuncFields(
	v reflect.Value, t reflect.Type, i int, options MethodOptions) {
	f := v.Field(i)
	fs := t.Field(i)
	name := fs.Name
	if f.CanInterface() && f.IsValid() {
		f, _ = getPtrTo(f, f.Type())
		if !f.IsNil() && f.Kind() == reflect.Func {
			methods.AddFunction(name, f, options)
		} else if f.Kind() == reflect.Struct {
			if fs.Anonymous {
				methods.AddAllMethods(f.Interface(), options)
			} else {
				newOptions := options
				if newOptions.NameSpace == "" {
					newOptions.NameSpace = name
				} else {
					newOptions.NameSpace += "_" + name
				}
				methods.AddAllMethods(f.Interface(), newOptions)
			}
		}
	}
}

// AddAllMethods will publish all methods and non-nil function fields on the
// obj self and on its anonymous or non-anonymous struct fields (or pointer to
// pointer ... to pointer struct fields). This is a recursive operation.
// So it's a pit, if you do not know what you are doing, do not step on.
func (methods *Methods) AddAllMethods(
	obj interface{}, options MethodOptions) {
	if obj == nil {
		panic("obj can't be nil")
	}
	v := reflect.ValueOf(obj)
	t := v.Type()
	methods.addInstanceMethods(v, t, options)
	v, t = getPtrTo(v, t)
	if t.Kind() == reflect.Struct {
		n := t.NumField()
		for i := 0; i < n; i++ {
			methods.recursiveAddFuncFields(v, t, i, options)
		}
	}
}

// MissingMethod is missing method
type MissingMethod func(name string, args []reflect.Value) (result []reflect.Value)

// AddMissingMethod is used for publishing a method,
// all methods not explicitly published will be redirected to this method.
func (methods *Methods) AddMissingMethod(
	method MissingMethod, options MethodOptions) {
	methods.AddFunction("*", method, options)
}
<commit_msg>Update methods.go<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * rpc/methods.go                                         *
 *                                                        *
 * hprose methods for Go.                                 *
 *                                                        *
 * LastModified: Sep 11, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package rpc

import (
	"reflect"
	"strings"
)

// MethodOptions is the options of the published service method
type MethodOptions struct {
	Mode      ResultMode
	Simple    bool
	Oneway    bool
	NameSpace string
}

// serviceMethod is the published service method
type serviceMethod struct {
	Function reflect.Value
	MethodOptions
}

// methods is the published service methods
type serviceMethods struct {
	MethodNames   []string
	RemoteMethods map[string]serviceMethod
}

// newServiceMethods is the constructor for Methods
func newServiceMethods() (methods *serviceMethods) {
	methods = new(serviceMethods)
	methods.MethodNames = make([]string, 0, 64)
	methods.RemoteMethods = make(map[string]serviceMethod)
	return
}

// AddFunction publish a func or bound method
// name is the method name
// function is a func or bound method
// options includes Mode, Simple, Oneway and NameSpace
func (methods *serviceMethods) AddFunction(
	name string, function interface{}, options MethodOptions) {
	if name == "" {
		panic("name can't be empty")
	}
	if function == nil {
		panic("function can't be nil")
	}
	f, ok := function.(reflect.Value)
	if !ok {
		f = reflect.ValueOf(function)
	}
	if f.Kind() != reflect.Func {
		panic("function must be func or bound method")
	}
	if options.NameSpace != "" && name != "*" {
		name = options.NameSpace + "_" + name
	}
	methods.MethodNames = append(methods.MethodNames, name)
	methods.RemoteMethods[strings.ToLower(name)] = serviceMethod{f, options}
}

// AddFunctions is used for batch publishing service method
func (methods *serviceMethods) AddFunctions(
	names []string, functions []interface{}, options MethodOptions) {
	count := len(names)
	if count != len(functions) {
		panic("names and functions must have the same length")
	}
	for i := 0; i < count; i++ {
		methods.AddFunction(names[i], functions[i], options)
	}
}

// AddMethod is used for publishing a method on the obj with an alias
func (methods *serviceMethods) AddMethod(
	name string, obj interface{}, options MethodOptions, alias ...string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	f := reflect.ValueOf(obj).MethodByName(name)
	if len(alias) == 1 && alias[0] != "" {
		name = alias[0]
	}
	if f.CanInterface() {
		methods.AddFunction(name, f, options)
	}
}

// AddMethods is used for batch publishing methods on the obj with aliases
func (methods *serviceMethods) AddMethods(
	names []string, obj interface{}, options MethodOptions, aliases ...[]string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	count := len(names)
	if len(aliases) == 1 {
		if len(aliases[0]) != count {
			panic("names and aliases must have the same length")
		}
		for i := 0; i < count; i++ {
			methods.AddMethod(names[i], obj, options, aliases[0][i])
		}
		return
	}
	for i := 0; i < count; i++ {
		methods.AddMethod(names[i], obj, options)
	}
}

func (methods *serviceMethods) addMethods(
	v reflect.Value, t reflect.Type, options MethodOptions) {
	n := t.NumMethod()
	for i := 0; i < n; i++ {
		name := t.Method(i).Name
		method := v.Method(i)
		if method.CanInterface() {
			methods.AddFunction(name, method, options)
		}
	}
}

func getPtrTo(v reflect.Value, t reflect.Type) (reflect.Value, reflect.Type) {
	for ; t.Kind() == reflect.Ptr && !v.IsNil(); v = v.Elem() {
		t = t.Elem()
	}
	return v, t
}

func (methods *serviceMethods) addFuncField(
	v reflect.Value, t reflect.Type, i int, options MethodOptions) {
	f := v.Field(i)
	name := t.Field(i).Name
	if !f.CanInterface() || !f.IsValid() {
		return
	}
	f, _ = getPtrTo(f, f.Type())
	if !f.IsNil() && f.Kind() == reflect.Func {
		methods.AddFunction(name, f, options)
	}
}

func (methods *serviceMethods) recursiveAddFuncFields(
	v reflect.Value, t reflect.Type, i int, options MethodOptions) {
	f := v.Field(i)
	fs := t.Field(i)
	name := fs.Name
	if !f.CanInterface() || !f.IsValid() {
		return
	}
	f, _ = getPtrTo(f, f.Type())
	if !f.IsNil() && f.Kind() == reflect.Func {
		methods.AddFunction(name, f, options)
		return
	}
	if f.Kind() != reflect.Struct {
		return
	}
	if fs.Anonymous {
		methods.AddAllMethods(f.Interface(), options)
	} else {
		newOptions := options
		if newOptions.NameSpace == "" {
			newOptions.NameSpace = name
		} else {
			newOptions.NameSpace += "_" + name
		}
		methods.AddAllMethods(f.Interface(), newOptions)
	}
}

type addFuncFunc func(
	v reflect.Value,
	t reflect.Type,
	i int,
	options MethodOptions)

func (methods *serviceMethods) addInstanceMethods(
	obj interface{}, options MethodOptions, addFunc addFuncFunc) {
	if obj == nil {
		panic("obj can't be nil")
	}
	v := reflect.ValueOf(obj)
	t := v.Type()
	methods.addMethods(v, t, options)
	v, t = getPtrTo(v, t)
	if t.Kind() == reflect.Struct {
		n := t.NumField()
		for i := 0; i < n; i++ {
			addFunc(v, t, i, options)
		}
	}
}

// AddInstanceMethods is used for publishing all the public methods and func fields with options.
func (methods *serviceMethods) AddInstanceMethods(
	obj interface{}, options MethodOptions) {
	methods.addInstanceMethods(obj, options, methods.addFuncField)
}

// AddAllMethods will publish all methods and non-nil function fields on the
// obj self and on its anonymous or non-anonymous struct fields (or pointer to
// pointer ... to pointer struct fields). This is a recursive operation.
// So it's a pit, if you do not know what you are doing, do not step on.
func (methods *serviceMethods) AddAllMethods(
	obj interface{}, options MethodOptions) {
	methods.addInstanceMethods(obj, options, methods.recursiveAddFuncFields)
}

// MissingMethod is missing method
type MissingMethod func(name string, args []reflect.Value) (result []reflect.Value)

// AddMissingMethod is used for publishing a method,
// all methods not explicitly published will be redirected to this method.
func (methods *serviceMethods) AddMissingMethod(
	method MissingMethod, options MethodOptions) {
	methods.AddFunction("*", method, options)
}
<|endoftext|>


<commit_before>package soaap

import (
	"encoding/gob"
	"errors"
	"fmt"
	"math"
	"os"
	"strings"
)

type CallGraph struct {
	Nodes map[string]GraphNode
	Calls map[Call]int
}

//
// Create a new, empty CallGraph with enough capacity to hold some calls.
//
func NewCallGraph() CallGraph {
	return CallGraph{
		make(map[string]GraphNode),
		make(map[Call]int),
	}
}

//
// Load a CallGraph from a binary-encoded file.
//
func LoadGraph(f *os.File, report func(string)) (CallGraph, error) {
	var graph CallGraph
	err := gob.NewDecoder(f).Decode(&graph)

	return graph, err
}

func (cg *CallGraph) AddCall(caller string, callee string) {
	cg.Calls[Call{caller, callee}] += 1
}

//
// Save a CallGraph to an os.File using a binary encoding.
//
func (cg *CallGraph) Save(f *os.File) error {
	return gob.NewEncoder(f).Encode(cg)
}

//
// Simplify a CallGraph by collapsing call chains and dropping any
// unreferenced calls.
//
func (cg *CallGraph) Simplify() {
}

func (cg *CallGraph) Union(g CallGraph) error {
	for id, node := range g.Nodes {
		// If we already have a GraphNode with this identifier,
		// merge the two descriptions and tag sets.
		if n, have := cg.Nodes[id]; have {
			if n.Name != node.Name {
				return errors.New(fmt.Sprintf(
					"Nodes in CallGraph union have"+
						" same identifier ('%s') but"+
						" different names ('%s' vs '%s')",
					id, n.Name, node.Name))
			}

			if n.Description != node.Description {
				node.Description =
					n.Description + "\\n" + node.Description
			}

			for tag := range n.Tags {
				node.Tags[tag] = true
			}
		}

		cg.Nodes[id] = node
	}

	for call, count := range g.Calls {
		cg.Calls[call] += count
	}

	return nil
}

//
// A node in a call graph.
//
// This is derived from a call site or other program location, but can have
// an arbitrary name and description appropriate to a particular analysis.
//
type GraphNode struct {
	Name        string
	Description string

	// A vulnerability (current or previous) is known at this location.
	CVE []CVE

	// The name of this node's sandbox (or the empty string if unsandboxed).
	Sandbox string

	// The name of the sandbox(es) that own the data being accessed.
	Owners []string

	Tags map[string]bool
}

//
// Construct a GraphViz Dot description of a GraphNode.
//
// This applies SOAAP-specific styling depending on a node's tags.
//
func (n GraphNode) Dot() string {
	attrs := map[string]interface{}{
		"label": n.Description,
		"style": "filled",
	}

	switch true {
	case len(n.CVE) > 0 && n.Sandbox != "":
		// A vulnerability has been mitigated through sandboxing!
		attrs["fillcolor"] = "#ffff66cc"
		attrs["shape"] = "octagon"

	case len(n.CVE) > 0:
		// A vulnerability exists/existed outside a sandbox.
		attrs["fillcolor"] = "#ff9999cc"
		attrs["shape"] = "doubleoctagon"

	case len(n.Owners) > 0:
		// Sandbox-private data was accessed outside the sandbox.
		attrs["fillcolor"] = "#ff99cccc"
		attrs["shape"] = "invhouse"

	case n.Sandbox != "":
		attrs["fillcolor"] = "#99ff9999"
		attrs["style"] = "dashed,filled"

	default:
		attrs["fillcolor"] = "#cccccccc"
	}

	return fmt.Sprintf("\"%s\" %s;", n.Name, dotAttrs(attrs))
}

func (n GraphNode) HasTag(tag string) bool {
	_, present := n.Tags[tag]
	return present
}

type Call struct {
	// Identifier of the caller.
	Caller string

	// Identifier of the callee.
	Callee string
}

// Output GraphViz for a Call.
func (c Call) Dot(graph CallGraph, weight int) string {
	caller := graph.Nodes[c.Caller]
	callee := graph.Nodes[c.Callee]

	attrs := map[string]interface{}{
		"penwidth": weight,
		"weight":   weight,
	}

	return fmt.Sprintf("\"%s\" -> \"%s\" %s;\n",
		caller.Name, callee.Name, dotAttrs(attrs))
}

//
// A function that extracts a CallGraph from SOAAP Results.
//
type graphFn func(results Results, progress func(string)) CallGraph

var graphExtractors map[string]graphFn = map[string]graphFn{
	"privaccess": PrivAccessGraph,
	"vuln":       VulnGraph,
}

func GraphAnalyses() []string {
	keys := make([]string, len(graphExtractors))

	i := 0
	for k, _ := range graphExtractors {
		keys[i] = k
		i++
	}

	return keys
}

type callSiteLabeler func(CallSite) (string, GraphNode)

//
// Construct a callgraph from SOAAP's vulnerability analysis.
//
func VulnGraph(results Results, progress func(string)) CallGraph {
	graph := NewCallGraph()

	for _, v := range results.Vulnerabilities {
		trace := results.Traces[v.Trace]
		top := true

		fn := func(cs CallSite) (string, GraphNode) {
			key := cs.String() + " " + v.Sandbox

			desc := cs.Function
			if v.Sandbox != "" {
				desc += "\\n<<" + v.Sandbox + ">>"
			}

			var node GraphNode
			node.Name = cs.String() + "_" + v.Sandbox
			node.Description = desc
			node.Sandbox = v.Sandbox

			if top {
				node.CVE = v.CVE
				top = false
			}

			return key, node
		}

		graph.Union(trace.graph(results.Traces, fn))
	}

	return graph
}

//
// Construct a callgraph of sandbox-private data accesses outside of sandboxes.
//
func PrivAccessGraph(results Results, progress func(string)) CallGraph {
	graph := NewCallGraph()
	accesses := results.PrivateAccess
	total := len(accesses)
	chunk := int(math.Ceil(math.Pow(10, math.Log10(float64(total)/500))))

	go progress(fmt.Sprintf("Processing %d private accesses", total))

	count := 0
	for _, a := range accesses {
		trace := results.Traces[a.Trace]

		fn := func(cs CallSite) (string, GraphNode) {
			sandboxes := strings.Join(a.Sandboxes, ",")
			key := cs.String() + " " + sandboxes

			desc := cs.Function
			if sandboxes != "" {
				desc += "\\n<<" + sandboxes + ">>"
			}

			var node GraphNode
			node.Name = cs.String() + "_" + sandboxes
			node.Description = desc
			node.Owners = a.Sandboxes

			return key, node
		}

		graph.Union(trace.graph(results.Traces, fn))

		count++
		if count%chunk == 0 {
			go progress(
				fmt.Sprintf("Processed %d/%d accesses",
					count, total))
		}
	}

	return graph
}

//
// Graph a single CallTrace, using a callSiteLabeler function to convert
// CallSite instances into graph nodes with identifiers, tags, etc.,
// appropriate to the analysis we're performing.
//
func (t CallTrace) graph(traces []CallTrace, nm callSiteLabeler) CallGraph {
	graph := NewCallGraph()
	var callee string

	t.Foreach(traces, func(cs CallSite) {
		identifier, node := nm(cs)
		graph.Nodes[identifier] = node

		caller := identifier

		if callee != "" {
			graph.AddCall(caller, callee)
		}

		callee = caller
	})

	return graph
}

//
// Format a map as a GraphViz attribute list.
//
func dotAttrs(attrs map[string]interface{}) string {
	fields := make([]string, len(attrs))

	i := 0
	for k, v := range attrs {
		switch v.(type) {
		case string:
			v = fmt.Sprintf("\"%s\"", v)
		}

		fields[i] = fmt.Sprintf("\"%s\" = %v", k, v)
		i++
	}

	return fmt.Sprintf("[ %s ]", strings.Join(fields, ", "))
}
<commit_msg>Show call location information.<commit_after>package soaap

import (
	"encoding/gob"
	"errors"
	"fmt"
	"math"
	"os"
	"strings"
)

type CallGraph struct {
	Nodes map[string]GraphNode
	Calls map[Call]int
}

//
// Create a new, empty CallGraph with enough capacity to hold some calls.
//
func NewCallGraph() CallGraph {
	return CallGraph{
		make(map[string]GraphNode),
		make(map[Call]int),
	}
}

//
// Load a CallGraph from a binary-encoded file.
//
func LoadGraph(f *os.File, report func(string)) (CallGraph, error) {
	var graph CallGraph
	err := gob.NewDecoder(f).Decode(&graph)

	return graph, err
}

func (cg *CallGraph) AddCall(caller string, callee string) {
	cg.Calls[Call{caller, callee}] += 1
}

//
// Save a CallGraph to an os.File using a binary encoding.
//
func (cg *CallGraph) Save(f *os.File) error {
	return gob.NewEncoder(f).Encode(cg)
}

//
// Simplify a CallGraph by collapsing call chains and dropping any
// unreferenced calls.
//
func (cg *CallGraph) Simplify() {
}

func (cg *CallGraph) Union(g CallGraph) error {
	for id, node := range g.Nodes {
		// If we already have a GraphNode with this identifier,
		// merge the two descriptions and tag sets.
		if n, have := cg.Nodes[id]; have {
			if n.Name != node.Name {
				return errors.New(fmt.Sprintf(
					"Nodes in CallGraph union have"+
						" same identifier ('%s') but"+
						" different names ('%s' vs '%s')",
					id, n.Name, node.Name))
			}

			if n.Description != node.Description {
				node.Description =
					n.Description + "\\n" + node.Description
			}

			for tag := range n.Tags {
				node.Tags[tag] = true
			}
		}

		cg.Nodes[id] = node
	}

	for call, count := range g.Calls {
		cg.Calls[call] += count
	}

	return nil
}

//
// A node in a call graph.
//
// This is derived from a call site or other program location, but can have
// an arbitrary name and description appropriate to a particular analysis.
//
type GraphNode struct {
	Name        string
	Description string
	Location    SourceLocation

	// A vulnerability (current or previous) is known at this location.
	CVE []CVE

	// The name of this node's sandbox (or the empty string if unsandboxed).
	Sandbox string

	// The name of the sandbox(es) that own the data being accessed.
	Owners []string

	Tags map[string]bool
}

//
// Construct a GraphViz Dot description of a GraphNode.
//
// This applies SOAAP-specific styling depending on a node's tags.
//
func (n GraphNode) Dot() string {
	attrs := map[string]interface{}{
		"label": n.Description,
		"style": "filled",
	}

	switch true {
	case len(n.CVE) > 0 && n.Sandbox != "":
		// A vulnerability has been mitigated through sandboxing!
		attrs["fillcolor"] = "#ffff66cc"
		attrs["shape"] = "octagon"

	case len(n.CVE) > 0:
		// A vulnerability exists/existed outside a sandbox.
		attrs["fillcolor"] = "#ff9999cc"
		attrs["shape"] = "doubleoctagon"

	case len(n.Owners) > 0:
		// Sandbox-private data was accessed outside the sandbox.
		attrs["fillcolor"] = "#ff99cccc"
		attrs["shape"] = "invhouse"

	case n.Sandbox != "":
		attrs["fillcolor"] = "#99ff9999"
		attrs["style"] = "dashed,filled"

	default:
		attrs["fillcolor"] = "#cccccccc"
	}

	return fmt.Sprintf("\"%s\" %s;", n.Name, dotAttrs(attrs))
}

func (n GraphNode) HasTag(tag string) bool {
	_, present := n.Tags[tag]
	return present
}

type Call struct {
	// Identifier of the caller.
	Caller string

	// Identifier of the callee.
	Callee string
}

// Output GraphViz for a Call.
func (c Call) Dot(graph CallGraph, weight int) string {
	caller := graph.Nodes[c.Caller]
	callee := graph.Nodes[c.Callee]

	attrs := map[string]interface{}{
		"label":    caller.Location.String(),
		"penwidth": weight,
		"weight":   weight,
	}

	return fmt.Sprintf("\"%s\" -> \"%s\" %s;\n",
		caller.Name, callee.Name, dotAttrs(attrs))
}

//
// A function that extracts a CallGraph from SOAAP Results.
//
type graphFn func(results Results, progress func(string)) CallGraph

var graphExtractors map[string]graphFn = map[string]graphFn{
	"privaccess": PrivAccessGraph,
	"vuln":       VulnGraph,
}

func GraphAnalyses() []string {
	keys := make([]string, len(graphExtractors))

	i := 0
	for k, _ := range graphExtractors {
		keys[i] = k
		i++
	}

	return keys
}

type callSiteLabeler func(CallSite) (string, GraphNode)

//
// Construct a callgraph from SOAAP's vulnerability analysis.
//
func VulnGraph(results Results, progress func(string)) CallGraph {
	graph := NewCallGraph()

	for _, v := range results.Vulnerabilities {
		trace := results.Traces[v.Trace]
		top := true

		fn := func(cs CallSite) (string, GraphNode) {
			key := cs.String() + " " + v.Sandbox

			desc := cs.Function
			if v.Sandbox != "" {
				desc += "\\n<<" + v.Sandbox + ">>"
			}

			var node GraphNode
			node.Name = cs.String() + "_" + v.Sandbox
			node.Description = desc
			node.Location = cs.Location
			node.Sandbox = v.Sandbox

			if top {
				node.CVE = v.CVE
				top = false
			}

			return key, node
		}

		graph.Union(trace.graph(results.Traces, fn))
	}

	return graph
}

//
// Construct a callgraph of sandbox-private data accesses outside of sandboxes.
//
func PrivAccessGraph(results Results, progress func(string)) CallGraph {
	graph := NewCallGraph()
	accesses := results.PrivateAccess
	total := len(accesses)
	chunk := int(math.Ceil(math.Pow(10, math.Log10(float64(total)/500))))

	go progress(fmt.Sprintf("Processing %d private accesses", total))

	count := 0
	for _, a := range accesses {
		trace := results.Traces[a.Trace]

		fn := func(cs CallSite) (string, GraphNode) {
			sandboxes := strings.Join(a.Sandboxes, ",")
			key := cs.String() + " " + sandboxes

			desc := cs.Function
			if sandboxes != "" {
				desc += "\\n<<" + sandboxes + ">>"
			}

			var node GraphNode
			node.Name = cs.String() + "_" + sandboxes
			node.Description = desc
			node.Location = cs.Location
			node.Owners = a.Sandboxes

			return key, node
		}

		graph.Union(trace.graph(results.Traces, fn))

		count++
		if count%chunk == 0 {
			go progress(
				fmt.Sprintf("Processed %d/%d accesses",
					count, total))
		}
	}

	return graph
}

//
// Graph a single CallTrace, using a callSiteLabeler function to convert
// CallSite instances into graph nodes with identifiers, tags, etc.,
// appropriate to the analysis we're performing.
//
func (t CallTrace) graph(traces []CallTrace, nm callSiteLabeler) CallGraph {
	graph := NewCallGraph()
	var callee string

	t.Foreach(traces, func(cs CallSite) {
		identifier, node := nm(cs)
		graph.Nodes[identifier] = node

		caller := identifier

		if callee != "" {
			graph.AddCall(caller, callee)
		}

		callee = caller
	})

	return graph
}

//
// Format a map as a GraphViz attribute list.
//
func dotAttrs(attrs map[string]interface{}) string {
	fields := make([]string, len(attrs))

	i := 0
	for k, v := range attrs {
		switch v.(type) {
		case string:
			v = fmt.Sprintf("\"%s\"", v)
		}

		fields[i] = fmt.Sprintf("\"%s\" = %v", k, v)
		i++
	}

	return fmt.Sprintf("[ %s ]", strings.Join(fields, ", "))
}
<|endoftext|>


<commit_before>package handler

import (
	"encoding/json"
	"html/template"
	"net/http"

	"github.com/graphql-go/graphql"
)

// page is the page data structure of the rendered GraphiQL page
type graphiqlPage struct {
	GraphiqlVersion string
	QueryString     string
	ResultString    string
	VariablesString string
	OperationName   string
}

// renderGraphiQL renders the GraphiQL GUI
func renderGraphiQL(w http.ResponseWriter, params graphql.Params) {
	t := template.New("GraphiQL")
	t, err := t.Parse(graphiqlTemplate)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Create variables string
	vars, err := json.MarshalIndent(params.VariableValues, "", "  ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	varsString := string(vars)
	if varsString == "null" {
		varsString = ""
	}

	// Create result string
	result, err := json.MarshalIndent(graphql.Do(params), "", "  ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	resString := string(result)

	p := graphiqlPage{
		GraphiqlVersion: graphiqlVersion,
		QueryString:     params.RequestString,
		ResultString:    resString,
		VariablesString: varsString,
		OperationName:   params.OperationName,
	}

	err = t.ExecuteTemplate(w, "index", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	return
}

// graphiqlVersion is the current version of GraphiQL
const graphiqlVersion = "0.11.3"

// tmpl is the page template to render GraphiQL
const graphiqlTemplate = `
{{ define "index" }}
<!--
The request to this GraphQL server provided the header "Accept: text/html"
and as a result has been presented GraphiQL - an in-browser IDE for
exploring GraphQL.

If you wish to receive JSON, provide the header "Accept: application/json" or
add "&raw" to the end of the URL within a browser.
-->
<!DOCTYPE html>
<html>
<head>
  <meta charset="utf-8" />
  <title>GraphiQL</title>
  <meta name="robots" content="noindex" />
  <style>
    html, body {
      height: 100%;
      margin: 0;
      overflow: hidden;
      width: 100%;
    }
  </style>
  <link href="//cdn.jsdelivr.net/npm/graphiql@{{ .GraphiqlVersion }}/graphiql.css" rel="stylesheet" />
  <script src="//cdn.jsdelivr.net/fetch/0.9.0/fetch.min.js"></script>
  <script src="//cdn.jsdelivr.net/react/15.4.2/react.min.js"></script>
  <script src="//cdn.jsdelivr.net/react/15.4.2/react-dom.min.js"></script>
  <script src="//cdn.jsdelivr.net/npm/graphiql@{{ .GraphiqlVersion }}/graphiql.min.js"></script>
</head>
<body>
  <script>
    // Collect the URL parameters
    var parameters = {};
    window.location.search.substr(1).split('&').forEach(function (entry) {
      var eq = entry.indexOf('=');
      if (eq >= 0) {
        parameters[decodeURIComponent(entry.slice(0, eq))] =
          decodeURIComponent(entry.slice(eq + 1));
      }
    });

    // Produce a Location query string from a parameter object.
    function locationQuery(params) {
      return '?' + Object.keys(params).filter(function (key) {
        return Boolean(params[key]);
      }).map(function (key) {
        return encodeURIComponent(key) + '=' +
          encodeURIComponent(params[key]);
      }).join('&');
    }

    // Derive a fetch URL from the current URL, sans the GraphQL parameters.
    var graphqlParamNames = {
      query: true,
      variables: true,
      operationName: true
    };

    var otherParams = {};
    for (var k in parameters) {
      if (parameters.hasOwnProperty(k) && graphqlParamNames[k] !== true) {
        otherParams[k] = parameters[k];
      }
    }
    var fetchURL = locationQuery(otherParams);

    // Defines a GraphQL fetcher using the fetch API.
    function graphQLFetcher(graphQLParams) {
      return fetch(fetchURL, {
        method: 'post',
        headers: {
          'Accept': 'application/json',
          'Content-Type': 'application/json'
        },
        body: JSON.stringify(graphQLParams),
        credentials: 'include',
      }).then(function (response) {
        return response.text();
      }).then(function (responseBody) {
        try {
          return JSON.parse(responseBody);
        } catch (error) {
          return responseBody;
        }
      });
    }

    // When the query and variables string is edited, update the URL bar so
    // that it can be easily shared.
    function onEditQuery(newQuery) {
      parameters.query = newQuery;
      updateURL();
    }

    function onEditVariables(newVariables) {
      parameters.variables = newVariables;
      updateURL();
    }

    function onEditOperationName(newOperationName) {
      parameters.operationName = newOperationName;
      updateURL();
    }

    function updateURL() {
      history.replaceState(null, null, locationQuery(parameters));
    }

    // Render <GraphiQL /> into the body.
    ReactDOM.render(
      React.createElement(GraphiQL, {
        fetcher: graphQLFetcher,
        onEditQuery: onEditQuery,
        onEditVariables: onEditVariables,
        onEditOperationName: onEditOperationName,
        query: {{ .QueryString }},
        response: {{ .ResultString }},
        variables: {{ .VariablesString }},
        operationName: {{ .OperationName }},
      }),
      document.body
    );
  </script>
</body>
</html>
{{ end }}
`
<commit_msg>Account for empty request<commit_after>package handler

import (
	"encoding/json"
	"html/template"
	"net/http"

	"github.com/graphql-go/graphql"
)

// page is the page data structure of the rendered GraphiQL page
type graphiqlPage struct {
	GraphiqlVersion string
	QueryString     string
	ResultString    string
	VariablesString string
	OperationName   string
}

// renderGraphiQL renders the GraphiQL GUI
func renderGraphiQL(w http.ResponseWriter, params graphql.Params) {
	t := template.New("GraphiQL")
	t, err := t.Parse(graphiqlTemplate)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Create variables string
	vars, err := json.MarshalIndent(params.VariableValues, "", "  ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	varsString := string(vars)
	if varsString == "null" {
		varsString = ""
	}

	// Create result string
	var resString string
	if params.RequestString == "" {
		resString = ""
	} else {
		result, err := json.MarshalIndent(graphql.Do(params), "", "  ")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		resString = string(result)
	}

	p := graphiqlPage{
		GraphiqlVersion: graphiqlVersion,
		QueryString:     params.RequestString,
		ResultString:    resString,
		VariablesString: varsString,
		OperationName:   params.OperationName,
	}

	err = t.ExecuteTemplate(w, "index", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	return
}

// graphiqlVersion is the current version of GraphiQL
const graphiqlVersion = "0.11.3"

// tmpl is the page template to render GraphiQL
const graphiqlTemplate = `
{{ define "index" }}
<!--
The request to this GraphQL server provided the header "Accept: text/html"
and as a result has been presented GraphiQL - an in-browser IDE for
exploring GraphQL.

If you wish to receive JSON, provide the header "Accept: application/json" or
add "&raw" to the end of the URL within a browser.
-->
<!DOCTYPE html>
<html>
<head>
  <meta charset="utf-8" />
  <title>GraphiQL</title>
  <meta name="robots" content="noindex" />
  <style>
    html, body {
      height: 100%;
      margin: 0;
      overflow: hidden;
      width: 100%;
    }
  </style>
  <link href="//cdn.jsdelivr.net/npm/graphiql@{{ .GraphiqlVersion }}/graphiql.css" rel="stylesheet" />
  <script src="//cdn.jsdelivr.net/fetch/0.9.0/fetch.min.js"></script>
  <script src="//cdn.jsdelivr.net/react/15.4.2/react.min.js"></script>
  <script src="//cdn.jsdelivr.net/react/15.4.2/react-dom.min.js"></script>
  <script src="//cdn.jsdelivr.net/npm/graphiql@{{ .GraphiqlVersion }}/graphiql.min.js"></script>
</head>
<body>
  <script>
    // Collect the URL parameters
    var parameters = {};
    window.location.search.substr(1).split('&').forEach(function (entry) {
      var eq = entry.indexOf('=');
      if (eq >= 0) {
        parameters[decodeURIComponent(entry.slice(0, eq))] =
          decodeURIComponent(entry.slice(eq + 1));
      }
    });

    // Produce a Location query string from a parameter object.
    function locationQuery(params) {
      return '?' + Object.keys(params).filter(function (key) {
        return Boolean(params[key]);
      }).map(function (key) {
        return encodeURIComponent(key) + '=' +
          encodeURIComponent(params[key]);
      }).join('&');
    }

    // Derive a fetch URL from the current URL, sans the GraphQL parameters.
    var graphqlParamNames = {
      query: true,
      variables: true,
      operationName: true
    };

    var otherParams = {};
    for (var k in parameters) {
      if (parameters.hasOwnProperty(k) && graphqlParamNames[k] !== true) {
        otherParams[k] = parameters[k];
      }
    }
    var fetchURL = locationQuery(otherParams);

    // Defines a GraphQL fetcher using the fetch API.
    function graphQLFetcher(graphQLParams) {
      return fetch(fetchURL, {
        method: 'post',
        headers: {
          'Accept': 'application/json',
          'Content-Type': 'application/json'
        },
        body: JSON.stringify(graphQLParams),
        credentials: 'include',
      }).then(function (response) {
        return response.text();
      }).then(function (responseBody) {
        try {
          return JSON.parse(responseBody);
        } catch (error) {
          return responseBody;
        }
      });
    }

    // When the query and variables string is edited, update the URL bar so
    // that it can be easily shared.
    function onEditQuery(newQuery) {
      parameters.query = newQuery;
      updateURL();
    }

    function onEditVariables(newVariables) {
      parameters.variables = newVariables;
      updateURL();
    }

    function onEditOperationName(newOperationName) {
      parameters.operationName = newOperationName;
      updateURL();
    }

    function updateURL() {
      history.replaceState(null, null, locationQuery(parameters));
    }

    // Render <GraphiQL /> into the body.
    ReactDOM.render(
      React.createElement(GraphiQL, {
        fetcher: graphQLFetcher,
        onEditQuery: onEditQuery,
        onEditVariables: onEditVariables,
        onEditOperationName: onEditOperationName,
        query: {{ .QueryString }},
        response: {{ .ResultString }},
        variables: {{ .VariablesString }},
        operationName: {{ .OperationName }},
      }),
      document.body
    );
  </script>
</body>
</html>
{{ end }}
`
<|endoftext|>


<commit_before>package rsync

import (
	"crypto/sha1"
	"hash"
	"io"

	"bitbucket.org/kardianos/rsync"
)

type OperationTransmitter func(*Operation) error
type OperationReceiver func() (*Operation, error)

type Rsyncer struct {
	rsyncer *rsync.RSync
}

func NewRsyncer() *Rsyncer {
	return &Rsyncer{
		rsyncer: &rsync.RSync{
			UniqueHasher: sha1.New(),
		},
	}
}

func (r *Rsyncer) Signature(base io.Reader) ([]*BlockHash, error) {
	// Create the result.
	var result []*BlockHash

	// Create a signature writer.
	write := func(b rsync.BlockHash) error {
		result = append(result, &BlockHash{
			Index:      b.Index,
			StrongHash: b.StrongHash,
			WeakHash:   b.WeakHash,
		})
		return nil
	}

	// Perform signature generation.
	if err := r.rsyncer.CreateSignature(base, write); err != nil {
		return nil, err
	}

	// Success.
	return result, nil
}

func (r *Rsyncer) Deltafy(target io.Reader, baseSignature []*BlockHash, transmit OperationTransmitter) error {
	// Convert the base signature.
	baseSignatureRsync := make([]rsync.BlockHash, len(baseSignature))
	for i, b := range baseSignature {
		baseSignatureRsync[i] = rsync.BlockHash{
			Index:      b.Index,
			StrongHash: b.StrongHash,
			WeakHash:   b.WeakHash,
		}
	}

	// Create a wrapper operation writer. Re-use the Operation instance to save
	// on allocations. Rsync already re-uses the underlying data buffer, so
	// we're not making things any more dangerous.
	operation := &Operation{}
	write := func(o rsync.Operation) error {
		*operation = Operation{
			Type:          OpType(o.Type),
			BlockIndex:    o.BlockIndex,
			BlockIndexEnd: o.BlockIndexEnd,
			Data:          o.Data,
		}
		return transmit(operation)
	}

	// Perform delta generation.
	return r.rsyncer.CreateDelta(target, baseSignatureRsync, write, nil)
}

func (r *Rsyncer) Patch(destination io.Writer, base io.ReadSeeker, receive OperationReceiver, digest hash.Hash) error {
	// Create channels to communicate with the ApplyDelta Goroutine.
	operations := make(chan rsync.Operation)
	applyErrors := make(chan error, 1)

	// Start the ApplyDelta operation in a separate Goroutine, recording the
	// hash of the received contents.
	go func() {
		applyErrors <- r.rsyncer.ApplyDelta(destination, base, operations, digest)
	}()

	// Receive and feed operations into the Goroutine, watching for errors.
	var applyError, receiveError error
	applyExited := false
	for {
		// Grab the next operation.
		operation, err := receive()
		if err != nil {
			receiveError = err
			break
		}

		// Check if the operation stream is done.
		if operation == nil {
			break
		}

		// Convert the operation.
		operationRsync := rsync.Operation{
			Type:          rsync.OpType(operation.Type),
			BlockIndex:    operation.BlockIndex,
			BlockIndexEnd: operation.BlockIndexEnd,
			Data:          operation.Data,
		}

		// Forward the operation. If there is an error, burn the remaining
		// operations in this stream.
		select {
		case operations <- operationRsync:
		case applyError = <-applyErrors:
			applyExited = true
			break
		}
	}

	// Tell the ApplyDelta Goroutine that operations are complete. It may have
	// exited already if there was an error, in which case this will have no
	// effect.
	close(operations)

	// Ensure that the ApplyDelta Goroutine has completed.
	if !applyExited {
		applyError = <-applyErrors
	}

	// Check for errors.
	if receiveError != nil {
		return receiveError
	} else if applyError != nil {
		return applyError
	}

	// Success.
	return nil
}
<commit_msg>Refactored Rsync type names a bit.<commit_after>package rsync

import (
	"crypto/sha1"
	"hash"
	"io"

	"bitbucket.org/kardianos/rsync"
)

type OperationTransmitter func(*Operation) error
type OperationReceiver func() (*Operation, error)

type Rsync struct {
	rsync *rsync.RSync
}

func New() *Rsync {
	return &Rsync{
		rsync: &rsync.RSync{
			UniqueHasher: sha1.New(),
		},
	}
}

func (r *Rsync) Signature(base io.Reader) ([]*BlockHash, error) {
	// Create the result.
	var result []*BlockHash

	// Create a signature writer.
	write := func(b rsync.BlockHash) error {
		result = append(result, &BlockHash{
			Index:      b.Index,
			StrongHash: b.StrongHash,
			WeakHash:   b.WeakHash,
		})
		return nil
	}

	// Perform signature generation.
	if err := r.rsync.CreateSignature(base, write); err != nil {
		return nil, err
	}

	// Success.
	return result, nil
}

func (r *Rsync) Deltafy(target io.Reader, baseSignature []*BlockHash, transmit OperationTransmitter) error {
	// Convert the base signature.
	baseSignatureRsync := make([]rsync.BlockHash, len(baseSignature))
	for i, b := range baseSignature {
		baseSignatureRsync[i] = rsync.BlockHash{
			Index:      b.Index,
			StrongHash: b.StrongHash,
			WeakHash:   b.WeakHash,
		}
	}

	// Create a wrapper operation writer. Re-use the Operation instance to save
	// on allocations. Rsync already re-uses the underlying data buffer, so
	// we're not making things any more dangerous.
	operation := &Operation{}
	write := func(o rsync.Operation) error {
		*operation = Operation{
			Type:          OpType(o.Type),
			BlockIndex:    o.BlockIndex,
			BlockIndexEnd: o.BlockIndexEnd,
			Data:          o.Data,
		}
		return transmit(operation)
	}

	// Perform delta generation.
	return r.rsync.CreateDelta(target, baseSignatureRsync, write, nil)
}

func (r *Rsync) Patch(destination io.Writer, base io.ReadSeeker, receive OperationReceiver, digest hash.Hash) error {
	// Create channels to communicate with the ApplyDelta Goroutine.
	operations := make(chan rsync.Operation)
	applyErrors := make(chan error, 1)

	// Start the ApplyDelta operation in a separate Goroutine, recording the
	// hash of the received contents.
	go func() {
		applyErrors <- r.rsync.ApplyDelta(destination, base, operations, digest)
	}()

	// Receive and feed operations into the Goroutine, watching for errors.
	var applyError, receiveError error
	applyExited := false
	for {
		// Grab the next operation.
		operation, err := receive()
		if err != nil {
			receiveError = err
			break
		}

		// Check if the operation stream is done.
		if operation == nil {
			break
		}

		// Convert the operation.
		operationRsync := rsync.Operation{
			Type:          rsync.OpType(operation.Type),
			BlockIndex:    operation.BlockIndex,
			BlockIndexEnd: operation.BlockIndexEnd,
			Data:          operation.Data,
		}

		// Forward the operation. If there is an error, burn the remaining
		// operations in this stream.
		select {
		case operations <- operationRsync:
		case applyError = <-applyErrors:
			applyExited = true
			break
		}
	}

	// Tell the ApplyDelta Goroutine that operations are complete. It may have
	// exited already if there was an error, in which case this will have no
	// effect.
	close(operations)

	// Ensure that the ApplyDelta Goroutine has completed.
	if !applyExited {
		applyError = <-applyErrors
	}

	// Check for errors.
	if receiveError != nil {
		return receiveError
	} else if applyError != nil {
		return applyError
	}

	// Success.
	return nil
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package admin

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/gorilla/rpc/v2"

	"github.com/ava-labs/avalanchego/api"
	"github.com/ava-labs/avalanchego/api/server"
	"github.com/ava-labs/avalanchego/chains"
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/engine/common"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/perms"
	"github.com/ava-labs/avalanchego/utils/profiler"

	cjson "github.com/ava-labs/avalanchego/utils/json"
)

const (
	maxAliasLength = 512

	// Name of file that stacktraces are written to
	stacktraceFile = "stacktrace.txt"
)

var errAliasTooLong = errors.New("alias length is too long")

// Admin is the API service for node admin management
type Admin struct {
	log          logging.Logger
	profiler     profiler.Profiler
	chainManager chains.Manager
	httpServer   *server.Server
	logFactory   logging.Factory
}

// NewService returns a new admin API service
func NewService(log logging.Logger, chainManager chains.Manager, httpServer *server.Server, profileDir string, logFactory logging.Factory) (*common.HTTPHandler, error) {
	newServer := rpc.NewServer()
	codec := cjson.NewCodec()
	newServer.RegisterCodec(codec, "application/json")
	newServer.RegisterCodec(codec, "application/json;charset=UTF-8")
	if err := newServer.RegisterService(&Admin{
		log:          log,
		chainManager: chainManager,
		httpServer:   httpServer,
		profiler:     profiler.New(profileDir),
		logFactory:   logFactory,
	}, "admin"); err != nil {
		return nil, err
	}
	return &common.HTTPHandler{Handler: newServer}, nil
}

// StartCPUProfiler starts a cpu profile writing to the specified file
func (service *Admin) StartCPUProfiler(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: StartCPUProfiler called")
	reply.Success = true
	return service.profiler.StartCPUProfiler()
}

// StopCPUProfiler stops the cpu profile
func (service *Admin) StopCPUProfiler(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: StopCPUProfiler called")

	reply.Success = true
	return service.profiler.StopCPUProfiler()
}

// MemoryProfile runs a memory profile writing to the specified file
func (service *Admin) MemoryProfile(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: MemoryProfile called")

	reply.Success = true
	return service.profiler.MemoryProfile()
}

// LockProfile runs a mutex profile writing to the specified file
func (service *Admin) LockProfile(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: LockProfile called")

	reply.Success = true
	return service.profiler.LockProfile()
}

// AliasArgs are the arguments for calling Alias
type AliasArgs struct {
	Endpoint string `json:"endpoint"`
	Alias    string `json:"alias"`
}

// Alias attempts to alias an HTTP endpoint to a new name
func (service *Admin) Alias(_ *http.Request, args *AliasArgs, reply *api.SuccessResponse) error {
	service.log.Info("Admin: Alias called with URL: %s, Alias: %s", args.Endpoint, args.Alias)

	if len(args.Alias) > maxAliasLength {
		return errAliasTooLong
	}

	reply.Success = true
	return service.httpServer.AddAliasesWithReadLock(args.Endpoint, args.Alias)
}

// AliasChainArgs are the arguments for calling AliasChain
type AliasChainArgs struct {
	Chain string `json:"chain"`
	Alias string `json:"alias"`
}

// AliasChain attempts to alias a chain to a new name
func (service *Admin) AliasChain(_ *http.Request, args *AliasChainArgs, reply *api.SuccessResponse) error {
	service.log.Info("Admin: AliasChain called with Chain: %s, Alias: %s", args.Chain, args.Alias)

	if len(args.Alias) > maxAliasLength {
		return errAliasTooLong
	}
	chainID, err := service.chainManager.Lookup(args.Chain)
	if err != nil {
		return err
	}

	if err := service.chainManager.Alias(chainID, args.Alias); err != nil {
		return err
	}

	reply.Success = true
	return service.httpServer.AddAliasesWithReadLock(constants.ChainAliasPrefix+chainID.String(), constants.ChainAliasPrefix+args.Alias)
}

// GetChainAliasesArgs are the arguments for calling GetChainAliases
type GetChainAliasesArgs struct {
	Chain string `json:"chain"`
}

// GetChainAliasesReply are the aliases of the given chain
type GetChainAliasesReply struct {
	Aliases []string `json:"aliases"`
}

// GetChainAliases returns the aliases of the chain
func (service *Admin) GetChainAliases(r *http.Request, args *GetChainAliasesArgs, reply *GetChainAliasesReply) error {
	service.log.Info("Admin: GetChainAliases called with Chain: %s", args.Chain)

	id, err := ids.FromString(args.Chain)
	if err != nil {
		return err
	}

	reply.Aliases = service.chainManager.Aliases(id)
	return nil
}

// Stacktrace returns the current global stacktrace
func (service *Admin) Stacktrace(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: Stacktrace called")

	reply.Success = true
	stacktrace := []byte(logging.Stacktrace{Global: true}.String())
	return perms.WriteFile(stacktraceFile, stacktrace, perms.ReadWrite)
}

// SetLoggerLevelArgs defines the logger name and log&display levels to be set.
type SetLoggerLevelArgs struct {
	LogLevel     string `json:"logLevel"`
	DisplayLevel string `json:"displayLevel"`
	LoggerName   string `json:"loggerName"`
}

// SetLoggerLevelReply success result.
type SetLoggerLevelReply struct {
	Success bool `json:"success"`
}

// SetLoggerLevel sets the log level and display level for logger name.
func (service *Admin) SetLoggerLevel(r *http.Request, args *SetLoggerLevelArgs, reply *SetLoggerLevelReply) error {
	service.log.Info("Admin: SetLogLevels called with LoggerName: %s LogLevel: %s, DisplayLevel: %s", args.LoggerName, args.LogLevel, args.DisplayLevel)
	if len(args.LogLevel) == 0 && len(args.DisplayLevel) == 0 {
		return fmt.Errorf("need to specify either displayLevel or logLevel args")
	}

	var loggerNames []string
	// Empty name means all loggers
	if len(args.LoggerName) > 0 {
		loggerNames = []string{args.LoggerName}
	} else {
		loggerNames = service.logFactory.GetNames()
	}

	for _, name := range loggerNames {
		if len(args.LogLevel) > 0 {
			logLevel, err := logging.ToLevel(args.LogLevel)
			if err != nil {
				return err
			}
			if err := service.logFactory.SetLogLevel(name, logLevel); err != nil {
				return err
			}
		}
		if len(args.DisplayLevel) > 0 {
			displayLevel, err := logging.ToLevel(args.DisplayLevel)
			if err != nil {
				return err
			}
			if err := service.logFactory.SetDisplayLevel(name, displayLevel); err != nil {
				return err
			}
		}
	}
	reply.Success = true
	return nil
}

type LoggerLevels struct {
	LogLevel     string `json:"logLevel"`
	DisplayLevel string `json:"displayLevel"`
}

// GetLoggerLevelArgs defines logger name arg.
type GetLoggerLevelArgs struct {
	LoggerName string `json:"loggerName"`
}

// GetLogLevelReply returns current log levels in a name=>level map.
type GetLoggerLevelReply struct {
	LoggerLevels map[string]LoggerLevels
}

// GetLogLevel get the log level and display level of loggers.
func (service *Admin) GetLoggerLevel(r *http.Request, args *GetLoggerLevelArgs, reply *GetLoggerLevelReply) error {
	service.log.Info("Admin: GetLoggerLevels called with %s", args.LoggerName)
	reply.LoggerLevels = make(map[string]LoggerLevels)
	var loggerNames []string
	// Empty name means all loggers
	if len(args.LoggerName) > 0 {
		loggerNames = []string{args.LoggerName}
	} else {
		loggerNames = service.logFactory.GetNames()
	}

	for _, name := range loggerNames {
		logLevel, err := service.logFactory.GetLogLevel(name)
		if err != nil {
			return err
		}
		displayLevel, err := service.logFactory.GetDisplayLevel(name)
		if err != nil {
			return err
		}
		loggerLevels := LoggerLevels{
			LogLevel:     logLevel.String(),
			DisplayLevel: displayLevel.String(),
		}
		reply.LoggerLevels[name] = loggerLevels
	}
	return nil
}
<commit_msg>Apply suggestions from code review<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package admin

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/gorilla/rpc/v2"

	"github.com/ava-labs/avalanchego/api"
	"github.com/ava-labs/avalanchego/api/server"
	"github.com/ava-labs/avalanchego/chains"
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/engine/common"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/perms"
	"github.com/ava-labs/avalanchego/utils/profiler"

	cjson "github.com/ava-labs/avalanchego/utils/json"
)

const (
	maxAliasLength = 512

	// Name of file that stacktraces are written to
	stacktraceFile = "stacktrace.txt"
)

var errAliasTooLong = errors.New("alias length is too long")

// Admin is the API service for node admin management
type Admin struct {
	log          logging.Logger
	profiler     profiler.Profiler
	chainManager chains.Manager
	httpServer   *server.Server
	logFactory   logging.Factory
}

// NewService returns a new admin API service
func NewService(log logging.Logger, chainManager chains.Manager, httpServer *server.Server, profileDir string, logFactory logging.Factory) (*common.HTTPHandler, error) {
	newServer := rpc.NewServer()
	codec := cjson.NewCodec()
	newServer.RegisterCodec(codec, "application/json")
	newServer.RegisterCodec(codec, "application/json;charset=UTF-8")
	if err := newServer.RegisterService(&Admin{
		log:          log,
		chainManager: chainManager,
		httpServer:   httpServer,
		profiler:     profiler.New(profileDir),
		logFactory:   logFactory,
	}, "admin"); err != nil {
		return nil, err
	}
	return &common.HTTPHandler{Handler: newServer}, nil
}

// StartCPUProfiler starts a cpu profile writing to the specified file
func (service *Admin) StartCPUProfiler(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: StartCPUProfiler called")
	reply.Success = true
	return service.profiler.StartCPUProfiler()
}

// StopCPUProfiler stops the cpu profile
func (service *Admin) StopCPUProfiler(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: StopCPUProfiler called")

	reply.Success = true
	return service.profiler.StopCPUProfiler()
}

// MemoryProfile runs a memory profile writing to the specified file
func (service *Admin) MemoryProfile(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: MemoryProfile called")

	reply.Success = true
	return service.profiler.MemoryProfile()
}

// LockProfile runs a mutex profile writing to the specified file
func (service *Admin) LockProfile(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: LockProfile called")

	reply.Success = true
	return service.profiler.LockProfile()
}

// AliasArgs are the arguments for calling Alias
type AliasArgs struct {
	Endpoint string `json:"endpoint"`
	Alias    string `json:"alias"`
}

// Alias attempts to alias an HTTP endpoint to a new name
func (service *Admin) Alias(_ *http.Request, args *AliasArgs, reply *api.SuccessResponse) error {
	service.log.Info("Admin: Alias called with URL: %s, Alias: %s", args.Endpoint, args.Alias)

	if len(args.Alias) > maxAliasLength {
		return errAliasTooLong
	}

	reply.Success = true
	return service.httpServer.AddAliasesWithReadLock(args.Endpoint, args.Alias)
}

// AliasChainArgs are the arguments for calling AliasChain
type AliasChainArgs struct {
	Chain string `json:"chain"`
	Alias string `json:"alias"`
}

// AliasChain attempts to alias a chain to a new name
func (service *Admin) AliasChain(_ *http.Request, args *AliasChainArgs, reply *api.SuccessResponse) error {
	service.log.Info("Admin: AliasChain called with Chain: %s, Alias: %s", args.Chain, args.Alias)

	if len(args.Alias) > maxAliasLength {
		return errAliasTooLong
	}
	chainID, err := service.chainManager.Lookup(args.Chain)
	if err != nil {
		return err
	}

	if err := service.chainManager.Alias(chainID, args.Alias); err != nil {
		return err
	}

	reply.Success = true
	return service.httpServer.AddAliasesWithReadLock(constants.ChainAliasPrefix+chainID.String(), constants.ChainAliasPrefix+args.Alias)
}

// GetChainAliasesArgs are the arguments for calling GetChainAliases
type GetChainAliasesArgs struct {
	Chain string `json:"chain"`
}

// GetChainAliasesReply are the aliases of the given chain
type GetChainAliasesReply struct {
	Aliases []string `json:"aliases"`
}

// GetChainAliases returns the aliases of the chain
func (service *Admin) GetChainAliases(r *http.Request, args *GetChainAliasesArgs, reply *GetChainAliasesReply) error {
	service.log.Info("Admin: GetChainAliases called with Chain: %s", args.Chain)

	id, err := ids.FromString(args.Chain)
	if err != nil {
		return err
	}

	reply.Aliases = service.chainManager.Aliases(id)
	return nil
}

// Stacktrace returns the current global stacktrace
func (service *Admin) Stacktrace(_ *http.Request, _ *struct{}, reply *api.SuccessResponse) error {
	service.log.Info("Admin: Stacktrace called")

	reply.Success = true
	stacktrace := []byte(logging.Stacktrace{Global: true}.String())
	return perms.WriteFile(stacktraceFile, stacktrace, perms.ReadWrite)
}

// SetLoggerLevelArgs defines the logger name and log&display levels to be set.
type SetLoggerLevelArgs struct {
	LogLevel     string `json:"logLevel"`
	DisplayLevel string `json:"displayLevel"`
	LoggerName   string `json:"loggerName"`
}

// SetLoggerLevelReply success result.
type SetLoggerLevelReply struct {
	Success bool `json:"success"`
}

// SetLoggerLevel sets the log level and display level for logger name.
func (service *Admin) SetLoggerLevel(r *http.Request, args *SetLoggerLevelArgs, reply *api.SuccessResponse) error {
	service.log.Info("Admin: SetLogLevels called with LoggerName: %q LogLevel: %q, DisplayLevel: %q", args.LoggerName, args.LogLevel, args.DisplayLevel)
	if len(args.LogLevel) == 0 && len(args.DisplayLevel) == 0 {
		return errors.New("need to specify either displayLevel or logLevel")
	}

	var loggerNames []string
	// Empty name means all loggers
	if len(args.LoggerName) > 0 {
		loggerNames = []string{args.LoggerName}
	} else {
		loggerNames = service.logFactory.GetNames()
	}

	for _, name := range loggerNames {
		if len(args.LogLevel) > 0 {
			logLevel, err := logging.ToLevel(args.LogLevel)
			if err != nil {
				return err
			}
			if err := service.logFactory.SetLogLevel(name, logLevel); err != nil {
				return err
			}
		}
		if len(args.DisplayLevel) > 0 {
			displayLevel, err := logging.ToLevel(args.DisplayLevel)
			if err != nil {
				return err
			}
			if err := service.logFactory.SetDisplayLevel(name, displayLevel); err != nil {
				return err
			}
		}
	}
	reply.Success = true
	return nil
}

type LoggerLevels struct {
	LogLevel     string `json:"logLevel"`
	DisplayLevel string `json:"displayLevel"`
}

// GetLoggerLevelArgs defines logger name arg.
type GetLoggerLevelArgs struct {
	LoggerName string `json:"loggerName"`
}

// GetLogLevelReply returns current log levels in a name=>level map.
type GetLoggerLevelReply struct {
	LoggerLevels map[string]LoggerLevels `json:"loggerLevels"`
}

// GetLogLevel get the log level and display level of loggers.
func (service *Admin) GetLoggerLevel(r *http.Request, args *GetLoggerLevelArgs, reply *GetLoggerLevelReply) error {
	service.log.Info("Admin: GetLoggerLevels called with %q", args.LoggerName)
	reply.LoggerLevels = make(map[string]LoggerLevels)
	var loggerNames []string
	// Empty name means all loggers
	if len(args.LoggerName) > 0 {
		loggerNames = []string{args.LoggerName}
	} else {
		loggerNames = service.logFactory.GetNames()
	}

	for _, name := range loggerNames {
		logLevel, err := service.logFactory.GetLogLevel(name)
		if err != nil {
			return err
		}
		displayLevel, err := service.logFactory.GetDisplayLevel(name)
		if err != nil {
			return err
		}
		reply.LoggerLevels[name] = LoggerLevels{
			LogLevel:     logLevel.String(),
			DisplayLevel: displayLevel.String(),
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package name_analysis

import (
	"fmt"
	"github.com/williballenthin/Lancelot/artifacts"
	W "github.com/williballenthin/Lancelot/workspace"
)

func check(e error) {
	if e != nil {
		panic(e)
	}
}

type NameAnalysis struct {
	ws *W.Workspace
}

func New(ws *W.Workspace) (*NameAnalysis, error) {
	return &NameAnalysis{
		ws: ws,
	}, nil
}

/** NameAnalysis implements FunctionAnalysis interface **/
func (a *NameAnalysis) AnalyzeFunction(f *artifacts.Function) error {
	f.SetName(fmt.Sprintf("sub_%s", f.Start))
	return nil
}
<commit_msg>name analysis: attempt to use exported names<commit_after>package name_analysis

import (
	"fmt"
	"github.com/williballenthin/Lancelot/artifacts"
	W "github.com/williballenthin/Lancelot/workspace"
)

func check(e error) {
	if e != nil {
		panic(e)
	}
}

type NameAnalysis struct {
	ws *W.Workspace
}

func New(ws *W.Workspace) (*NameAnalysis, error) {
	return &NameAnalysis{
		ws: ws,
	}, nil
}

/** NameAnalysis implements FunctionAnalysis interface **/
func (a *NameAnalysis) AnalyzeFunction(f *artifacts.Function) error {
	sym, e := a.ws.ResolveAddressToSymbol(f.Start)
	if e == nil {
		f.SetName(sym.SymbolName)
	} else {
		f.SetName(fmt.Sprintf("sub_%s", f.Start))
	}
	return nil
}
<|endoftext|>


<commit_before>package api

import (
	"errors"
	"fmt"
	"testing"

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

func TestHandleRequest_ShouldRunControllerFromFactory(t *testing.T) {
	controller := TestController{}
	fac := TestFactory{}

	controller.On("Run").Once().Return(nil, nil)
	fac.On("Create").Return(controller, nil)

	HandleRequest(
		func(Request, Response) (Controller, ControllerDestroyFunc, error) {
			return controller, nil, nil
		},
		TestRequest{},
		TestResponse{},
	)

	controller.AssertExpectations(t)
}

func TestHandleRequest_ShouldCallDestroyFunction(t *testing.T) {
	controller := TestController{}
	fac := TestFactory{}

	controller.On("Run").Once().Return(nil, nil)
	fac.On("Create").Return(controller, nil)
	fac.On("Destroy").Once().Return(nil)

	HandleRequest(
		fac.Create,
		TestRequest{},
		TestResponse{},
	)

	controller.AssertExpectations(t)
}

func TestHandleRequest_ShouldCallDestroyFunctionEvenIfFactoryReturnsError(t *testing.T) {
	controller := TestController{}
	fac := TestFactory{}

	fac.On("Create").Return(nil, errors.New("Some error"))
	fac.On("Destroy").Once().Return(nil)

	HandleRequest(
		fac.Create,
		TestRequest{},
		TestResponse{},
	)

	controller.AssertExpectations(t)
}

type TestController struct {
	mock.Mock
}

func (t TestController) Run() error {
	args := t.Called()
	return args.Error(0)
}

type TestResponse struct {
	mock.Mock
}

func (t TestResponse) String(code int, s string) error {
	args := t.Called(code, s)
	return args.Error(0)
}
func (t TestResponse) JSON(code int, i interface{}) error {
	args := t.Called(code, i)
	return args.Error(0)
}
func (t TestResponse) NoContent(code int) error {
	args := t.Called(code)
	return args.Error(0)
}

type TestRequest struct {
	mock.Mock
}

func (t TestRequest) Param(name string) string {
	args := t.Called(name)
	return args.String(0)
}
func (t TestRequest) Bind(i interface{}) error {
	args := t.Called(i)
	return args.Error(0)
}

type TestFactory struct {
	mock.Mock
}

func (t TestFactory) Create(Request, Response) (Controller, ControllerDestroyFunc, error) {
	args := t.Called()
	return controller(args.Get(0)), t.Destroy, args.Error(1)
}

func controller(obj interface{}) Controller {
	var c Controller
	var ok bool
	if obj == nil {
		return nil
	}
	if c, ok = obj.(Controller); !ok {
		panic(fmt.Sprintf("assert: arguments: Controller failed because object wasn't correct type: %v", obj))
	}
	return c
}

func (t TestFactory) Destroy() error {
	args := t.Called()
	return args.Error(0)
}
<commit_msg>handling function test<commit_after>package api

import (
	"errors"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"net/http"
	typedErrors "github.com/open-gtd/server/api/validation/errors"
)

const controlerValidationError = "Controller validation error"
const someError = "Some error"
const controllerError = "Controller error"

func TestHandleRequest_ShouldRunControllerFromFactory(t *testing.T) {
	controller := testController{}
	fac := TestFactory{}

	request := testRequest{}
	response := testResponse{}

	controller.On("Run").Return(nil)
	fac.On("Create", request, response).Return(controller, nil)
	fac.On("Destroy").Return(nil)

	HandleRequest(
		fac.Create,
		request,
		response,
	)

	controller.AssertExpectations(t)
}

func TestHandleRequest_ShouldCallDestroyFunction(t *testing.T) {
	controller := testController{}
	fac := TestFactory{}

	controller.On("Run").Return(nil)
	fac.On("Create", mock.Anything, mock.Anything).Return(controller, nil)
	fac.On("Destroy").Return(nil)

	request := testRequest{}
	response := testResponse{}
	HandleRequest(
		fac.Create,
		request,
		response,
	)

	controller.AssertExpectations(t)
}

func TestHandleRequest_ShouldCallDestroyFunctionEvenIfFactoryReturnsError(t *testing.T) {
	controller := testController{}
	fac := TestFactory{}

	fac.On("Create", mock.Anything, mock.Anything).Return(nil, errors.New(someError))
	fac.On("Destroy").Return(nil)

	request := testRequest{}
	response := testResponse{}
	HandleRequest(
		fac.Create,
		request,
		response,
	)

	controller.AssertExpectations(t)
}

func TestHandleRequest_ShouldReturnErrorIfFactoryReturnsError(t *testing.T) {
	fac := TestFactory{}

	fac.On("Create", mock.Anything, mock.Anything).Return(nil, errors.New("Some error"))
	fac.On("Destroy").Return(nil)

	request := testRequest{}
	response := testResponse{}
	err := HandleRequest(
		fac.Create,
		request,
		response,
	)

	assert.EqualError(t, err, someError)
}

func TestHandleRequest_ShouldReturnErrorIfControllerRunReturnsError(t *testing.T) {
	controller := testController{}
	fac := TestFactory{}

	controller.On("Run").Return(errors.New(controllerError))
	fac.On("Create", mock.Anything, mock.Anything).Return(controller, nil)
	fac.On("Destroy").Return(nil)

	request := testRequest{}
	response := testResponse{}
	err := HandleRequest(
		fac.Create,
		request,
		response,
	)

	assert.EqualError(t, err, controllerError)
}

func TestHandleRequest_ShouldProduceBadRequestResponseErrorIfControllerRunReturnsError(t *testing.T) {
	controller := testController{}
	fac := TestFactory{}

	controller.On("Run").Return(testValidationError{})
	fac.On("Create", mock.Anything, mock.Anything).Return(controller, nil)
	fac.On("Destroy").Return(nil)

	request := testRequest{}
	response := testResponse{}
	response.On("JSON", http.StatusBadRequest, mock.Anything).
		Run(func(args mock.Arguments) {
		assert.Equal(t, args.Get(1).(MessageResponse).Message, controlerValidationError)
	}).Return(nil)

	HandleRequest(
		fac.Create,
		request,
		response,
	)
}

type testController struct {
	mock.Mock
}

func (t testController) Run() error {
	args := t.Called()
	return args.Error(0)
}

type testResponse struct {
	mock.Mock
}

func (t testResponse) String(code int, s string) error {
	args := t.Called(code, s)
	return args.Error(0)
}
func (t testResponse) JSON(code int, i interface{}) error {
	args := t.Called(code, i)
	return args.Error(0)
}
func (t testResponse) NoContent(code int) error {
	args := t.Called(code)
	return args.Error(0)
}

type testRequest struct {
	mock.Mock
}

func (t testRequest) Param(name string) string {
	args := t.Called(name)
	return args.String(0)
}
func (t testRequest) Bind(i interface{}) error {
	args := t.Called(i)
	return args.Error(0)
}

type TestFactory struct {
	mock.Mock
}

func (t TestFactory) Create(rq Request, rs Response) (Controller, ControllerDestroyFunc, error) {
	args := t.Called(rq, rs)
	return controller(args.Get(0)), t.Destroy, args.Error(1)
}

func controller(obj interface{}) Controller {
	var c Controller
	var ok bool
	if obj == nil {
		return nil
	}
	if c, ok = obj.(Controller); !ok {
		panic(fmt.Sprintf("assert: arguments: Controller failed because object wasn't correct type: %v", obj))
	}
	return c
}

func (t TestFactory) Destroy() error {
	args := t.Called()
	return args.Error(0)
}

type testValidationError struct {}

func (t testValidationError) Type() typedErrors.ErrorType {
	return typedErrors.Validation
}

func (t testValidationError) Error() string {
	return controlerValidationError
}
<|endoftext|>


<commit_before>package ssdp

import (
	"net"
	"strings"
	"sync"
	"testing"
)

func testMaxAge(t *testing.T, s string, expect int) {
	act := extractMaxAge(s, -1)
	if act != expect {
		t.Errorf("max-age for %q should be %d but actually %d", s, expect, act)
	}
}

func TestExtractMaxAge(t *testing.T) {
	// empty
	testMaxAge(t, "", -1)
	// spaces around `=`
	testMaxAge(t, "max-age=100", 100)
	testMaxAge(t, "max-age = 200", 200)
	testMaxAge(t, "max-age= 300", 300)
	testMaxAge(t, "max-age =400", 400)
	// minus
	testMaxAge(t, "max-age=-100", -1)
	// invalid name
	testMaxAge(t, "foo=100", -1)
	// contained valid name
	testMaxAge(t, "foomax-age=100", -1)
	// surrounded
	testMaxAge(t, ";max-age=500;", 500)
	testMaxAge(t, ";max-age=600", 600)
	testMaxAge(t, "max-age=700;", 700)
}

func TestSearch_Request(t *testing.T) {
	searchType := "test:search+request"

	var mu sync.Mutex
	var mm []*SearchMessage
	m := newTestMonitor(searchType, nil, nil, func(m *SearchMessage) {
		mu.Lock()
		mm = append(mm, m)
		mu.Unlock()
	})
	err := m.Start()
	if err != nil {
		t.Fatalf("failed to start Monitor: %s", err)
	}
	defer m.Close()

	srvs, err := Search(searchType, 1, "")
	if err != nil {
		t.Fatalf("failed to Search: %s", err)
	}
	if len(srvs) > 0 {
		t.Errorf("unexpected services: %+v", srvs)
	}

	if len(mm) < 1 {
		t.Fatal("no search detected")
	}
	_, port, err := net.SplitHostPort(mm[0].From.String())
	if err != nil {
		t.Fatalf("failed to split host and port: %s", err)
	}
	port = ":" + port
	for i, m := range mm {
		if m.Type != "test:search+request" {
			t.Errorf("unmatch type#%d:\nwant=%q\n got=%q", i, "test:search+request", m.Type)
		}
		if strings.HasSuffix(port, m.From.String()) {
			t.Errorf("unmatch port#%d:\nwant=%q\n got=%q", i, port, m.From.String())
		}
	}
}
<commit_msg>test: search result<commit_after>package ssdp

import (
	"net"
	"strings"
	"sync"
	"testing"
)

func testMaxAge(t *testing.T, s string, expect int) {
	act := extractMaxAge(s, -1)
	if act != expect {
		t.Errorf("max-age for %q should be %d but actually %d", s, expect, act)
	}
}

func TestExtractMaxAge(t *testing.T) {
	// empty
	testMaxAge(t, "", -1)
	// spaces around `=`
	testMaxAge(t, "max-age=100", 100)
	testMaxAge(t, "max-age = 200", 200)
	testMaxAge(t, "max-age= 300", 300)
	testMaxAge(t, "max-age =400", 400)
	// minus
	testMaxAge(t, "max-age=-100", -1)
	// invalid name
	testMaxAge(t, "foo=100", -1)
	// contained valid name
	testMaxAge(t, "foomax-age=100", -1)
	// surrounded
	testMaxAge(t, ";max-age=500;", 500)
	testMaxAge(t, ";max-age=600", 600)
	testMaxAge(t, "max-age=700;", 700)
}

func TestSearch_Request(t *testing.T) {
	searchType := "test:search+request"

	var mu sync.Mutex
	var mm []*SearchMessage
	m := newTestMonitor(searchType, nil, nil, func(m *SearchMessage) {
		mu.Lock()
		mm = append(mm, m)
		mu.Unlock()
	})
	err := m.Start()
	if err != nil {
		t.Fatalf("failed to start Monitor: %s", err)
	}
	defer m.Close()

	srvs, err := Search(searchType, 1, "")
	if err != nil {
		t.Fatalf("failed to Search: %s", err)
	}
	if len(srvs) > 0 {
		t.Errorf("unexpected services: %+v", srvs)
	}

	if len(mm) < 1 {
		t.Fatal("no search detected")
	}
	_, port, err := net.SplitHostPort(mm[0].From.String())
	if err != nil {
		t.Fatalf("failed to split host and port: %s", err)
	}
	port = ":" + port
	for i, m := range mm {
		if m.Type != "test:search+request" {
			t.Errorf("unmatch type#%d:\nwant=%q\n got=%q", i, "test:search+request", m.Type)
		}
		if strings.HasSuffix(port, m.From.String()) {
			t.Errorf("unmatch port#%d:\nwant=%q\n got=%q", i, port, m.From.String())
		}
	}
}

func TestSearch_Response(t *testing.T) {
	a, err := Advertise("test:search+response", "usn:search+response", "location:search+response", "server:search+response", 600)
	if err != nil {
		t.Fatalf("failed to Advertise: %s", err)
	}
	defer a.Close()

	srvs, err := Search("test:search+response", 1, "")
	if err != nil {
		t.Fatalf("failed to Search: %s", err)
	}
	if len(srvs) == 0 {
		t.Errorf("no services found")
	}

	for i, s := range srvs {
		if s.Type != "test:search+response" {
			t.Errorf("unexpected service#%d type: want=%q got=%q", i, "test:search+response", s.Type)
		}
		if s.USN != "usn:search+response" {
			t.Errorf("unexpected service#%d usn: want=%q got=%q", i, "usn:search+response", s.USN)
		}
		if s.Location != "location:search+response" {
			t.Errorf("unexpected service#%d location: want=%q got=%q", i, "location:search+response", s.Location)
		}
		if s.Server != "server:search+response" {
			t.Errorf("unexpected service#%d server: want=%q got=%q", i, "server:search+response", s.Server)
		}
	}
}
<|endoftext|>


<commit_before>package logs

import (
	"encoding/json"
	"errors"
	"log"
	"strings"
	"time"

	"github.com/appcelerator/amp/pkg/docker"
	"github.com/appcelerator/amp/pkg/elasticsearch"
	"github.com/appcelerator/amp/pkg/labels"
	"github.com/appcelerator/amp/pkg/nats-streaming"
	"github.com/golang/protobuf/proto"
	"github.com/nats-io/go-nats-streaming"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"gopkg.in/olivere/elastic.v5"
)

// Server is used to implement log.LogServer
type Server struct {
	Docker        *docker.Docker
	Es            *elasticsearch.Elasticsearch
	NatsStreaming *ns.NatsStreaming
}

// Get implements logs.LogsServer
func (s *Server) Get(ctx context.Context, in *GetRequest) (*GetReply, error) {
	if err := s.Es.Connect(); err != nil {
		return nil, errors.New("unable to connect to elasticsearch service")
	}
	log.Println("rpc-logs: Get", in.String())

	// Prepares indices
	indices := []string{}
	date := time.Now().UTC()
	for i := 0; i < 2; i++ {
		indices = append(indices, "ampbeat-"+date.Format("2006.01.02"))
		date = date.AddDate(0, 0, -1)
	}

	// Prepare request to elasticsearch
	request := s.Es.GetClient().Search().Index(indices...).IgnoreUnavailable(true)
	request.Type("logs")
	request.Sort("@timestamp", false)
	if in.Size != 0 {
		request.Size(int(in.Size))
	} else {
		request.Size(100)
	}

	masterQuery := elastic.NewBoolQuery()
	if in.Container != "" {
		boolQuery := elastic.NewBoolQuery()
		masterQuery.Filter(
			boolQuery.Should(elastic.NewPrefixQuery("container_id", in.Service)),
			boolQuery.Should(elastic.NewPrefixQuery("container_name", in.Service)),
		)
	}
	if in.Service != "" {
		boolQuery := elastic.NewBoolQuery()
		masterQuery.Filter(
			boolQuery.Should(elastic.NewPrefixQuery("service_id", in.Service)),
			boolQuery.Should(elastic.NewPrefixQuery("service_name", in.Service)),
		)
	}
	if in.Stack != "" {
		masterQuery.Filter(elastic.NewPrefixQuery("stack_name", in.Stack))
	}
	if in.Node != "" {
		masterQuery.Filter(elastic.NewPrefixQuery("node_id", in.Node))
	}
	if in.Message != "" {
		queryString := elastic.NewSimpleQueryStringQuery(in.Message)
		queryString.Field("msg")
		masterQuery.Filter(queryString)
	}
	if !in.Infra {
		masterQuery.MustNot(elastic.NewTermQuery(convertLabelNameToESName(labels.LabelsNameRole), labels.LabelsValuesRoleInfrastructure))
		//For now Tools role is manage as Infrastructure role
		//Later: ToolsRole should be manage with a user premission (admin)
		masterQuery.MustNot(elastic.NewTermQuery(convertLabelNameToESName(labels.LabelsNameRole), labels.LabelsValuesRoleTools))
	}

	// Perform request
	searchResult, err := request.Query(masterQuery).Do(ctx)
	if err != nil {
		return nil, grpc.Errorf(codes.FailedPrecondition, "%v", err)
	}

	// Build reply (from elasticsearch response)
	reply := GetReply{}
	reply.Entries = make([]*LogEntry, len(searchResult.Hits.Hits))
	for i, hit := range searchResult.Hits.Hits {
		entry := &LogEntry{}
		if err := json.Unmarshal(*hit.Source, &entry); err != nil {
			return nil, grpc.Errorf(codes.Internal, "%v", err)
		}
		reply.Entries[i] = entry
	}

	// Reverse entries
	for i, j := 0, len(reply.Entries)-1; i < j; i, j = i+1, j-1 {
		reply.Entries[i], reply.Entries[j] = reply.Entries[j], reply.Entries[i]
	}
	log.Printf("rpc-logs: Get successful, returned %d entries\n", len(reply.Entries))
	return &reply, nil
}

// GetStream implements log.LogServer
func (s *Server) GetStream(in *GetRequest, stream Logs_GetStreamServer) error {
	if err := s.NatsStreaming.Connect(); err != nil {
		return errors.New("unable to connect to nats service")
	}
	log.Println("rpc-logs: GetStream", in.String())

	sub, err := s.NatsStreaming.GetClient().Subscribe(ns.LogsSubject, func(msg *stan.Msg) {
		entry := &LogEntry{}
		if err := proto.Unmarshal(msg.Data, entry); err != nil {
			return
		}
		if filter(entry, in) {
			stream.Send(entry)
		}
	})
	if err != nil {
		sub.Unsubscribe()
		return grpc.Errorf(codes.Internal, "%v", err)
	}
	for {
		select {
		case <-stream.Context().Done():
			sub.Unsubscribe()
			return stream.Context().Err()
		}
	}
}

func filter(entry *LogEntry, in *GetRequest) bool {
	match := true
	if in.Container != "" {
		containerID := strings.ToLower(entry.ContainerId)
		containerName := strings.ToLower(entry.ContainerName)
		match = strings.HasPrefix(containerID, strings.ToLower(in.Container)) || strings.HasPrefix(containerName, strings.ToLower(in.Container))
	}
	if in.Service != "" {
		serviceID := strings.ToLower(entry.ServiceId)
		serviceName := strings.ToLower(entry.ServiceName)
		match = strings.HasPrefix(serviceID, strings.ToLower(in.Service)) || strings.HasPrefix(serviceName, strings.ToLower(in.Service))
	}
	if in.Stack != "" {
		match = strings.HasPrefix(strings.ToLower(entry.StackName), strings.ToLower(in.Stack))
	}
	if in.Node != "" {
		match = strings.HasPrefix(strings.ToLower(entry.NodeId), strings.ToLower(in.Node))
	}
	if in.Message != "" {
		match = strings.Contains(strings.ToLower(entry.Msg), strings.ToLower(in.Message))
	}
	if !in.Infra {
		role := entry.Labels[labels.LabelsNameRole]
		match = role != labels.LabelsValuesRoleInfrastructure && role != labels.LabelsValuesRoleTools
	}
	return match
}

func convertLabelNameToESName(name string) string {
	return "labels." + strings.Replace(name, ".", "-", -1)
}
<commit_msg>Fix logs by container name or container id<commit_after>package logs

import (
	"encoding/json"
	"errors"
	"log"
	"strings"
	"time"

	"github.com/appcelerator/amp/pkg/docker"
	"github.com/appcelerator/amp/pkg/elasticsearch"
	"github.com/appcelerator/amp/pkg/labels"
	"github.com/appcelerator/amp/pkg/nats-streaming"
	"github.com/golang/protobuf/proto"
	"github.com/nats-io/go-nats-streaming"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"gopkg.in/olivere/elastic.v5"
)

// Server is used to implement log.LogServer
type Server struct {
	Docker        *docker.Docker
	Es            *elasticsearch.Elasticsearch
	NatsStreaming *ns.NatsStreaming
}

// Get implements logs.LogsServer
func (s *Server) Get(ctx context.Context, in *GetRequest) (*GetReply, error) {
	if err := s.Es.Connect(); err != nil {
		return nil, errors.New("unable to connect to elasticsearch service")
	}
	log.Println("rpc-logs: Get", in.String())

	// Prepares indices
	indices := []string{}
	date := time.Now().UTC()
	for i := 0; i < 2; i++ {
		indices = append(indices, "ampbeat-"+date.Format("2006.01.02"))
		date = date.AddDate(0, 0, -1)
	}

	// Prepare request to elasticsearch
	request := s.Es.GetClient().Search().Index(indices...).IgnoreUnavailable(true)
	request.Type("logs")
	request.Sort("@timestamp", false)
	if in.Size != 0 {
		request.Size(int(in.Size))
	} else {
		request.Size(100)
	}

	masterQuery := elastic.NewBoolQuery()
	if in.Container != "" {
		boolQuery := elastic.NewBoolQuery()
		masterQuery.Filter(
			boolQuery.Should(elastic.NewPrefixQuery("container_id", in.Container)),
			boolQuery.Should(elastic.NewPrefixQuery("container_name", in.Container)),
		)
	}
	if in.Service != "" {
		boolQuery := elastic.NewBoolQuery()
		masterQuery.Filter(
			boolQuery.Should(elastic.NewPrefixQuery("service_id", in.Service)),
			boolQuery.Should(elastic.NewPrefixQuery("service_name", in.Service)),
		)
	}
	if in.Stack != "" {
		masterQuery.Filter(elastic.NewPrefixQuery("stack_name", in.Stack))
	}
	if in.Node != "" {
		masterQuery.Filter(elastic.NewPrefixQuery("node_id", in.Node))
	}
	if in.Message != "" {
		queryString := elastic.NewSimpleQueryStringQuery(in.Message)
		queryString.Field("msg")
		masterQuery.Filter(queryString)
	}
	if !in.Infra {
		masterQuery.MustNot(elastic.NewTermQuery(convertLabelNameToESName(labels.LabelsNameRole), labels.LabelsValuesRoleInfrastructure))
		//For now Tools role is manage as Infrastructure role
		//Later: ToolsRole should be manage with a user premission (admin)
		masterQuery.MustNot(elastic.NewTermQuery(convertLabelNameToESName(labels.LabelsNameRole), labels.LabelsValuesRoleTools))
	}

	// Perform request
	searchResult, err := request.Query(masterQuery).Do(ctx)
	if err != nil {
		return nil, grpc.Errorf(codes.FailedPrecondition, "%v", err)
	}

	// Build reply (from elasticsearch response)
	reply := GetReply{}
	reply.Entries = make([]*LogEntry, len(searchResult.Hits.Hits))
	for i, hit := range searchResult.Hits.Hits {
		entry := &LogEntry{}
		if err := json.Unmarshal(*hit.Source, &entry); err != nil {
			return nil, grpc.Errorf(codes.Internal, "%v", err)
		}
		reply.Entries[i] = entry
	}

	// Reverse entries
	for i, j := 0, len(reply.Entries)-1; i < j; i, j = i+1, j-1 {
		reply.Entries[i], reply.Entries[j] = reply.Entries[j], reply.Entries[i]
	}
	log.Printf("rpc-logs: Get successful, returned %d entries\n", len(reply.Entries))
	return &reply, nil
}

// GetStream implements log.LogServer
func (s *Server) GetStream(in *GetRequest, stream Logs_GetStreamServer) error {
	if err := s.NatsStreaming.Connect(); err != nil {
		return errors.New("unable to connect to nats service")
	}
	log.Println("rpc-logs: GetStream", in.String())

	sub, err := s.NatsStreaming.GetClient().Subscribe(ns.LogsSubject, func(msg *stan.Msg) {
		entry := &LogEntry{}
		if err := proto.Unmarshal(msg.Data, entry); err != nil {
			return
		}
		if filter(entry, in) {
			stream.Send(entry)
		}
	})
	if err != nil {
		sub.Unsubscribe()
		return grpc.Errorf(codes.Internal, "%v", err)
	}
	for {
		select {
		case <-stream.Context().Done():
			sub.Unsubscribe()
			return stream.Context().Err()
		}
	}
}

func filter(entry *LogEntry, in *GetRequest) bool {
	match := true
	if in.Container != "" {
		containerID := strings.ToLower(entry.ContainerId)
		containerName := strings.ToLower(entry.ContainerName)
		match = strings.HasPrefix(containerID, strings.ToLower(in.Container)) || strings.HasPrefix(containerName, strings.ToLower(in.Container))
	}
	if in.Service != "" {
		serviceID := strings.ToLower(entry.ServiceId)
		serviceName := strings.ToLower(entry.ServiceName)
		match = strings.HasPrefix(serviceID, strings.ToLower(in.Service)) || strings.HasPrefix(serviceName, strings.ToLower(in.Service))
	}
	if in.Stack != "" {
		match = strings.HasPrefix(strings.ToLower(entry.StackName), strings.ToLower(in.Stack))
	}
	if in.Node != "" {
		match = strings.HasPrefix(strings.ToLower(entry.NodeId), strings.ToLower(in.Node))
	}
	if in.Message != "" {
		match = strings.Contains(strings.ToLower(entry.Msg), strings.ToLower(in.Message))
	}
	if !in.Infra {
		role := entry.Labels[labels.LabelsNameRole]
		match = role != labels.LabelsValuesRoleInfrastructure && role != labels.LabelsValuesRoleTools
	}
	return match
}

func convertLabelNameToESName(name string) string {
	return "labels." + strings.Replace(name, ".", "-", -1)
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/NetSys/quilt/api"
	"github.com/NetSys/quilt/api/pb"
	"github.com/NetSys/quilt/db"
	"github.com/NetSys/quilt/engine"
	"github.com/NetSys/quilt/stitch"

	"golang.org/x/net/context"
	"google.golang.org/grpc"

	log "github.com/Sirupsen/logrus"
)

type server struct {
	dbConn db.Conn
}

// Run accepts incoming `quiltctl` connections and responds to them.
func Run(conn db.Conn, listenAddr string) error {
	proto, addr, err := api.ParseListenAddress(listenAddr)
	if err != nil {
		return err
	}

	var sock net.Listener
	apiServer := server{conn}
	for {
		sock, err = net.Listen(proto, addr)

		if err == nil {
			break
		}
		log.WithError(err).Error("Failed to open socket.")

		time.Sleep(30 * time.Second)
	}

	// Cleanup the socket if we're interrupted.
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, os.Interrupt, os.Kill, syscall.SIGTERM)
	go func(c chan os.Signal) {
		sig := <-c
		log.Printf("Caught signal %s: shutting down.\n", sig)
		sock.Close()
		os.Exit(0)
	}(sigc)

	s := grpc.NewServer()
	pb.RegisterAPIServer(s, apiServer)
	s.Serve(sock)

	return nil
}

func (s server) Query(cts context.Context, query *pb.DBQuery) (*pb.QueryReply, error) {
	var rows interface{}
	err := s.dbConn.Transact(func(view db.Database) error {
		switch db.TableType(query.Table) {
		case db.MachineTable:
			rows = view.SelectFromMachine(nil)
		case db.ContainerTable:
			rows = view.SelectFromContainer(nil)
		default:
			return fmt.Errorf("unrecognized table: %s", query.Table)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	json, err := json.Marshal(rows)
	if err != nil {
		return nil, err
	}

	return &pb.QueryReply{TableContents: string(json)}, nil
}

func (s server) Run(cts context.Context, runReq *pb.RunRequest) (*pb.RunReply, error) {
	stitch, err := stitch.New(runReq.Stitch)
	if err != nil {
		return &pb.RunReply{}, err
	}

	err = engine.UpdatePolicy(s.dbConn, stitch)
	if err != nil {
		return &pb.RunReply{}, err
	}

	return &pb.RunReply{}, nil
}
<commit_msg>quilt: Close daemon socket on SIGHUP<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/NetSys/quilt/api"
	"github.com/NetSys/quilt/api/pb"
	"github.com/NetSys/quilt/db"
	"github.com/NetSys/quilt/engine"
	"github.com/NetSys/quilt/stitch"

	"golang.org/x/net/context"
	"google.golang.org/grpc"

	log "github.com/Sirupsen/logrus"
)

type server struct {
	dbConn db.Conn
}

// Run accepts incoming `quiltctl` connections and responds to them.
func Run(conn db.Conn, listenAddr string) error {
	proto, addr, err := api.ParseListenAddress(listenAddr)
	if err != nil {
		return err
	}

	var sock net.Listener
	apiServer := server{conn}
	for {
		sock, err = net.Listen(proto, addr)

		if err == nil {
			break
		}
		log.WithError(err).Error("Failed to open socket.")

		time.Sleep(30 * time.Second)
	}

	// Cleanup the socket if we're interrupted.
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, os.Interrupt, os.Kill, syscall.SIGTERM, syscall.SIGHUP)
	go func(c chan os.Signal) {
		sig := <-c
		log.Printf("Caught signal %s: shutting down.\n", sig)
		sock.Close()
		os.Exit(0)
	}(sigc)

	s := grpc.NewServer()
	pb.RegisterAPIServer(s, apiServer)
	s.Serve(sock)

	return nil
}

func (s server) Query(cts context.Context, query *pb.DBQuery) (*pb.QueryReply, error) {
	var rows interface{}
	err := s.dbConn.Transact(func(view db.Database) error {
		switch db.TableType(query.Table) {
		case db.MachineTable:
			rows = view.SelectFromMachine(nil)
		case db.ContainerTable:
			rows = view.SelectFromContainer(nil)
		default:
			return fmt.Errorf("unrecognized table: %s", query.Table)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	json, err := json.Marshal(rows)
	if err != nil {
		return nil, err
	}

	return &pb.QueryReply{TableContents: string(json)}, nil
}

func (s server) Run(cts context.Context, runReq *pb.RunRequest) (*pb.RunReply, error) {
	stitch, err := stitch.New(runReq.Stitch)
	if err != nil {
		return &pb.RunReply{}, err
	}

	err = engine.UpdatePolicy(s.dbConn, stitch)
	if err != nil {
		return &pb.RunReply{}, err
	}

	return &pb.RunReply{}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/gorilla/sessions"
	"net/http"
	"encoding/json"
	"log"
)

type DBSessionStore struct {
	Options *sessions.Options
	db DBProxy
}

func (s *DBSessionStore) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

func (s *DBSessionStore) New(r *http.Request, name string) (*sessions.Session, error) {
	session := sessions.NewSession(s, name)

	opts := *s.Options
	session.Options = &opts

	session.IsNew = true

	var err error
	cookie, _ := r.Cookie(name)
	if cookie != nil {
		session.ID = cookie.Value
		err = s.load(session)
		session.IsNew = (err != nil)
	}

	return session, err
}

func (s *DBSessionStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
	if session.ID == "" {
		session.ID = randStr(20)
	}

	if err := s.save(session); err != nil {
		return err
	}

	http.SetCookie(w, sessions.NewCookie(session.Name(), session.ID, session.Options))

	return nil
}

func (s *DBSessionStore) load(session *sessions.Session) error {
	data, err := s.db.LoadSession(session.ID, session.Name())
	if err == nil {
		err = json.Unmarshal([]byte(data), &session.Values)

		log.Println("load: ", session.Values)

		if err != nil {
			log.Println(err)
		}
	}

	return err
}

func (s *DBSessionStore) save(session *sessions.Session) error {
	log.Println("save: ", session.Values)

	tmp := make(map[string]interface{})

	for key, value := range session.Values {
		strkey, ok := key.(string)
		if ok {
			tmp[key] = value
		}
	}

	data, _ := json.Marshal(tmp)

	return s.db.SaveSession(session.ID, session.Name(), string(data))
}

func NewDBSessionStore(db DBProxy) *DBSessionStore {
	return &DBSessionStore {
		Options: &sessions.Options {
			Path:   "/",
			MaxAge: 60 * 60 * 24 * 60, /* sec */
		},
		db: db,
	}
}
<commit_msg>Debug<commit_after>package main

import (
	"github.com/gorilla/sessions"
	"net/http"
	"encoding/json"
	"log"
)

type DBSessionStore struct {
	Options *sessions.Options
	db DBProxy
}

func (s *DBSessionStore) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

func (s *DBSessionStore) New(r *http.Request, name string) (*sessions.Session, error) {
	session := sessions.NewSession(s, name)

	opts := *s.Options
	session.Options = &opts

	session.IsNew = true

	var err error
	cookie, _ := r.Cookie(name)
	if cookie != nil {
		session.ID = cookie.Value
		err = s.load(session)
		session.IsNew = (err != nil)
	}

	return session, err
}

func (s *DBSessionStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
	if session.ID == "" {
		session.ID = randStr(20)
	}

	if err := s.save(session); err != nil {
		return err
	}

	http.SetCookie(w, sessions.NewCookie(session.Name(), session.ID, session.Options))

	return nil
}

func (s *DBSessionStore) load(session *sessions.Session) error {
	data, err := s.db.LoadSession(session.ID, session.Name())
	if err == nil {
		err = json.Unmarshal([]byte(data), &session.Values)

		log.Println("load: ", session.Values)

		if err != nil {
			log.Println(err)
		}
	}

	return err
}

func (s *DBSessionStore) save(session *sessions.Session) error {
	log.Println("save: ", session.Values)

	tmp := make(map[string]interface{})

	for key, value := range session.Values {
		strkey, ok := key.(string)
		if ok {
			tmp[strkey] = value
		}
	}

	data, _ := json.Marshal(tmp)

	return s.db.SaveSession(session.ID, session.Name(), string(data))
}

func NewDBSessionStore(db DBProxy) *DBSessionStore {
	return &DBSessionStore {
		Options: &sessions.Options {
			Path:   "/",
			MaxAge: 60 * 60 * 24 * 60, /* sec */
		},
		db: db,
	}
}
<|endoftext|>


<commit_before>package lfs

import (
	"testing"

	"github.com/github/git-lfs/vendor/_nuts/github.com/technoweenie/assert"
)

func TestBatcherSizeMet(t *testing.T) {
	assertAll([]batcherTestCase{
		{2, 4, false},
		{3, 5, false},
		{0, 0, false},
	}, t)
}

func TestBatcherExit(t *testing.T) {
	assertAll([]batcherTestCase{
		{2, 4, true},
		{3, 5, true},
		{0, 0, true},
	}, t)
}

// Type batcherTestCase specifies information about how to run a particular test
// around the type lfs.Batcher.
type batcherTestCase struct {
	BatchSize  int
	ItemCount  int
	ShouldExit bool
}

// Func Batcher makes and retunrs a lfs.Batcher according to the specification
// given in this instance of batcherTestCase. When returned, it is filled with
// the given amount of items, and has exited if it was told to.
func (b batcherTestCase) Batcher() *Batcher {
	batcher := NewBatcher(b.BatchSize)
	for i := 0; i < b.ItemCount; i++ {
		batcher.Add(&Downloadable{})
	}

	if b.ShouldExit {
		batcher.Exit()
	}

	return batcher
}

// Func Batches returns the number of individual batches expected to be
// processed by the batcher under test.
func (b batcherTestCase) Batches() int {
	if b.BatchSize == 0 {
		return 0
	}

	return b.ItemCount / b.BatchSize
}

// Func assertAll processes all test cases, throwing assertion errors if they
// fail.
func assertAll(cases []batcherTestCase, t *testing.T) {
	for _, c := range cases {
		b := c.Batcher()

		items := c.ItemCount
		for i := 0; i < c.Batches(); i++ {
			group := b.Next()
			assert.Equal(t, c.BatchSize, len(group))
			items -= c.BatchSize
		}
	}
}
<commit_msg>アー アアアア アーアー<commit_after>package lfs

import (
	"testing"

	"github.com/github/git-lfs/vendor/_nuts/github.com/technoweenie/assert"
)

func TestBatcherSizeMet(t *testing.T) {
	assertAll([]batcherTestCase{
		{2, 4, false},
		{3, 5, false},
		{0, 0, false},
	}, t)
}

func TestBatcherExit(t *testing.T) {
	assertAll([]batcherTestCase{
		{2, 4, true},
		{3, 5, true},
		{0, 0, true},
	}, t)
}

// batcherTestCase specifies information about how to run a particular test
// around the type lfs.Batcher.
type batcherTestCase struct {
	BatchSize  int
	ItemCount  int
	ShouldExit bool
}

// Batcher makes and retunrs a lfs.Batcher according to the specification given
// in this instance of batcherTestCase. When returned, it is filled with the
// given amount of items, and has exited if it was told to.
func (b batcherTestCase) Batcher() *Batcher {
	batcher := NewBatcher(b.BatchSize)
	for i := 0; i < b.ItemCount; i++ {
		batcher.Add(&Downloadable{})
	}

	if b.ShouldExit {
		batcher.Exit()
	}

	return batcher
}

// Batches returns the number of individual batches expected to be processed by
// the batcher under test.
func (b batcherTestCase) Batches() int {
	if b.BatchSize == 0 {
		return b.ItemCount
	}

	return b.ItemCount / b.BatchSize
}

// assertAll processes all test cases, throwing assertion errors if they
// fail.
func assertAll(cases []batcherTestCase, t *testing.T) {
	for _, c := range cases {
		b := c.Batcher()

		items := c.ItemCount
		for i := 0; i < c.Batches(); i++ {
			group := b.Next()
			assert.Equal(t, c.BatchSize, len(group))
			items -= c.BatchSize
		}
	}
}
<|endoftext|>