<commit_before>//usr/bin/env go run $0 $@; exit
// Copyright 2014-2015 The DevMine authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package repotool is able to fetch information from a source code repository.
// Typically, it can get all commits, their authors and commiters and so on
// and return this information in a JSON object. Alternatively, it is able
// to populate the information into a PostgreSQL database.
// Currently, on the Git VCS is supported.
package main

import (
	"bytes"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/lib/pq"
	mmh3 "github.com/spaolacci/murmur3"

	"github.com/DevMine/repotool/config"
	"github.com/DevMine/repotool/model"
	"github.com/DevMine/repotool/repo"
)

const version = "0.1.0"

func main() {
	flag.Usage = func() {
		fmt.Printf("usage: %s [OPTION(S)] [REPOSITORY PATH]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}

	configPath := flag.String("c", "", "configuration file")
	vflag := flag.Bool("v", false, "print version.")
	jsonflag := flag.Bool("json", true, "json output")
	dbflag := flag.Bool("db", false, "import data into the database")
	flag.Parse()

	if *vflag {
		fmt.Printf("%s - %s\n", filepath.Base(os.Args[0]), version)
		os.Exit(0)
	}

	if len(flag.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "invalid # of arguments")
		flag.Usage()
	}

	if *dbflag && len(*configPath) == 0 {
		fatal(errors.New("a configuration file must be specified when using db option"))
	}

	repoPath := flag.Arg(0)
	repository, err := repo.New(repoPath)
	if err != nil {
		fatal(err)
	}

	fmt.Fprintln(os.Stderr, "fetching repository commits...")
	tic := time.Now()
	err = repository.FetchCommits()
	if err != nil {
		fatal(err)
	}
	toc := time.Now()
	fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))

	if *jsonflag {
		bs, err := json.Marshal(repository)
		if err != nil {
			fatal(err)
		}

		fmt.Println(string(bs))
	}

	if *dbflag {
		cfg, err := config.ReadConfig(*configPath)
		if err != nil {
			fatal(err)
		}

		db, err := OpenDBSession(cfg.Database)
		if err != nil {
			fatal(err)
		}
		defer db.Close()

		fmt.Fprintf(os.Stderr, "inserting %d commits into the database...\n",
			len(repository.GetCommits()))
		tic := time.Now()
		insertRepoData(db, repository)
		toc := time.Now()
		fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))
	}
}

// fatal prints an error on standard error stream and exits.
func fatal(a ...interface{}) {
	fmt.Fprintln(os.Stderr, a...)
	os.Exit(1)
}

// OpenDBSession creates a session to the database.
func OpenDBSession(cfg config.DatabaseConfig) (*sql.DB, error) {
	dbURL := fmt.Sprintf(
		"user='%s' password='%s' host='%s' port=%d dbname='%s' sslmode='%s'",
		cfg.UserName, cfg.Password, cfg.HostName, cfg.Port, cfg.DBName, cfg.SSLMode)

	return sql.Open("postgres", dbURL)
}

// insertRepoData inserts repository data into the database, or updates it
// if it is already there.
func insertRepoData(db *sql.DB, r repo.Repo) {
	if db == nil {
		fatal(errors.New("nil database given"))
	}

	repoID := getRepoID(db, r)
	if repoID == 0 {
		fatal(errors.New("no corresponding repository found in the database: impossible to insert data"))
	}

	for _, c := range r.GetCommits() {
		insertCommit(db, repoID, c)
	}
}

// insertCommit inserts a commit into the database
func insertCommit(db *sql.DB, repoID int, c model.Commit) {
	commitFields := []string{
		"repository_id",
		"author_id",
		"committer_id",
		"hash",
		"vcs_id",
		"message",
		"author_date",
		"commit_date",
		"file_changed_count",
		"insertions_count",
		"deletions_count"}

	authorID := getUserID(db, c.Author.Email)
	committerID := getUserID(db, c.Committer.Email)
	hash := genCommitHash(c)

	query := genInsQuery("commits", commitFields...)
	_, err := db.Exec(query,
		repoID, authorID, committerID, hash,
		c.VCSID, c.Message, c.AuthorDate, c.CommitDate,
		c.FileChangedCount, c.InsertionsCount, c.DeletionsCount)
	if err != nil {
		fatal(err)
	}
}

// genCommitHash generates a hash (mmh3) from commit fields.
// This hash can then be used to uniquely identify a commit.
func genCommitHash(c model.Commit) string {
	h := mmh3.New128()

	io.WriteString(h, c.VCSID)
	io.WriteString(h, c.Message)
	io.WriteString(h, c.Author.Name)
	io.WriteString(h, c.Author.Email)
	io.WriteString(h, c.Committer.Name)
	io.WriteString(h, c.Committer.Email)
	io.WriteString(h, c.AuthorDate.String())
	io.WriteString(h, c.CommitDate.String())
	io.WriteString(h, strconv.FormatInt(int64(c.FileChangedCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.InsertionsCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.DeletionsCount), 10))

	return hex.EncodeToString(h.Sum(nil))
}

// getRepoID returns the repository id of a repo in repositories table.
// If repo is not in the table, then 0 is returned.
func getRepoID(db *sql.DB, r repo.Repo) int {
	if db == nil {
		fatal(errors.New("nil database given"))
	}

	var id int
	// Clone URL is unique
	err := db.QueryRow("SELECT id FROM repositories WHERE clone_url=$1", r.GetCloneURL()).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return 0
	case err != nil:
		fatal(err)
	}
	return id
}

// getUserID attempts to find a user ID given its email address.
// Email addresses are unique, however they may not be provided.
// If no user ID is found, nil is returned, otherwhise the user ID
// is returned.
func getUserID(db *sql.DB, email string) *int {
	if db == nil {
		fatal(errors.New("nil database given"))
	}

	if len(email) == 0 {
		return nil
	}

	var id *int
	err := db.QueryRow("SELECT id FROM users WHERE email=$1", email).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return nil
	case err != nil:
		fatal(err)
	}
	return id
}

// genInsQuery generates a query string for an insertion in the database.
func genInsQuery(tableName string, fields ...string) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("INSERT INTO %s(%s)\n",
		tableName, strings.Join(fields, ",")))
	buf.WriteString("VALUES(")

	for ind := range fields {
		if ind > 0 {
			buf.WriteString(",")
		}

		buf.WriteString(fmt.Sprintf("$%d", ind+1))
	}

	buf.WriteString(")\n")

	return buf.String()
}
<commit_msg>misc: Remove shebang from repotool.<commit_after>// Copyright 2014-2015 The DevMine authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package repotool is able to fetch information from a source code repository.
// Typically, it can get all commits, their authors and commiters and so on
// and return this information in a JSON object. Alternatively, it is able
// to populate the information into a PostgreSQL database.
// Currently, on the Git VCS is supported.
package main

import (
	"bytes"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/lib/pq"
	mmh3 "github.com/spaolacci/murmur3"

	"github.com/DevMine/repotool/config"
	"github.com/DevMine/repotool/model"
	"github.com/DevMine/repotool/repo"
)

const version = "0.1.0"

func main() {
	flag.Usage = func() {
		fmt.Printf("usage: %s [OPTION(S)] [REPOSITORY PATH]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}

	configPath := flag.String("c", "", "configuration file")
	vflag := flag.Bool("v", false, "print version.")
	jsonflag := flag.Bool("json", true, "json output")
	dbflag := flag.Bool("db", false, "import data into the database")
	flag.Parse()

	if *vflag {
		fmt.Printf("%s - %s\n", filepath.Base(os.Args[0]), version)
		os.Exit(0)
	}

	if len(flag.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "invalid # of arguments")
		flag.Usage()
	}

	if *dbflag && len(*configPath) == 0 {
		fatal(errors.New("a configuration file must be specified when using db option"))
	}

	repoPath := flag.Arg(0)
	repository, err := repo.New(repoPath)
	if err != nil {
		fatal(err)
	}

	fmt.Fprintln(os.Stderr, "fetching repository commits...")
	tic := time.Now()
	err = repository.FetchCommits()
	if err != nil {
		fatal(err)
	}
	toc := time.Now()
	fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))

	if *jsonflag {
		bs, err := json.Marshal(repository)
		if err != nil {
			fatal(err)
		}

		fmt.Println(string(bs))
	}

	if *dbflag {
		cfg, err := config.ReadConfig(*configPath)
		if err != nil {
			fatal(err)
		}

		db, err := OpenDBSession(cfg.Database)
		if err != nil {
			fatal(err)
		}
		defer db.Close()

		fmt.Fprintf(os.Stderr, "inserting %d commits into the database...\n",
			len(repository.GetCommits()))
		tic := time.Now()
		insertRepoData(db, repository)
		toc := time.Now()
		fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))
	}
}

// fatal prints an error on standard error stream and exits.
func fatal(a ...interface{}) {
	fmt.Fprintln(os.Stderr, a...)
	os.Exit(1)
}

// OpenDBSession creates a session to the database.
func OpenDBSession(cfg config.DatabaseConfig) (*sql.DB, error) {
	dbURL := fmt.Sprintf(
		"user='%s' password='%s' host='%s' port=%d dbname='%s' sslmode='%s'",
		cfg.UserName, cfg.Password, cfg.HostName, cfg.Port, cfg.DBName, cfg.SSLMode)

	return sql.Open("postgres", dbURL)
}

// insertRepoData inserts repository data into the database, or updates it
// if it is already there.
func insertRepoData(db *sql.DB, r repo.Repo) {
	if db == nil {
		fatal(errors.New("nil database given"))
	}

	repoID := getRepoID(db, r)
	if repoID == 0 {
		fatal(errors.New("no corresponding repository found in the database: impossible to insert data"))
	}

	for _, c := range r.GetCommits() {
		insertCommit(db, repoID, c)
	}
}

// insertCommit inserts a commit into the database
func insertCommit(db *sql.DB, repoID int, c model.Commit) {
	commitFields := []string{
		"repository_id",
		"author_id",
		"committer_id",
		"hash",
		"vcs_id",
		"message",
		"author_date",
		"commit_date",
		"file_changed_count",
		"insertions_count",
		"deletions_count"}

	authorID := getUserID(db, c.Author.Email)
	committerID := getUserID(db, c.Committer.Email)
	hash := genCommitHash(c)

	query := genInsQuery("commits", commitFields...)
	_, err := db.Exec(query,
		repoID, authorID, committerID, hash,
		c.VCSID, c.Message, c.AuthorDate, c.CommitDate,
		c.FileChangedCount, c.InsertionsCount, c.DeletionsCount)
	if err != nil {
		fatal(err)
	}
}

// genCommitHash generates a hash (mmh3) from commit fields.
// This hash can then be used to uniquely identify a commit.
func genCommitHash(c model.Commit) string {
	h := mmh3.New128()

	io.WriteString(h, c.VCSID)
	io.WriteString(h, c.Message)
	io.WriteString(h, c.Author.Name)
	io.WriteString(h, c.Author.Email)
	io.WriteString(h, c.Committer.Name)
	io.WriteString(h, c.Committer.Email)
	io.WriteString(h, c.AuthorDate.String())
	io.WriteString(h, c.CommitDate.String())
	io.WriteString(h, strconv.FormatInt(int64(c.FileChangedCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.InsertionsCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.DeletionsCount), 10))

	return hex.EncodeToString(h.Sum(nil))
}

// getRepoID returns the repository id of a repo in repositories table.
// If repo is not in the table, then 0 is returned.
func getRepoID(db *sql.DB, r repo.Repo) int {
	if db == nil {
		fatal(errors.New("nil database given"))
	}

	var id int
	// Clone URL is unique
	err := db.QueryRow("SELECT id FROM repositories WHERE clone_url=$1", r.GetCloneURL()).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return 0
	case err != nil:
		fatal(err)
	}
	return id
}

// getUserID attempts to find a user ID given its email address.
// Email addresses are unique, however they may not be provided.
// If no user ID is found, nil is returned, otherwhise the user ID
// is returned.
func getUserID(db *sql.DB, email string) *int {
	if db == nil {
		fatal(errors.New("nil database given"))
	}

	if len(email) == 0 {
		return nil
	}

	var id *int
	err := db.QueryRow("SELECT id FROM users WHERE email=$1", email).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return nil
	case err != nil:
		fatal(err)
	}
	return id
}

// genInsQuery generates a query string for an insertion in the database.
func genInsQuery(tableName string, fields ...string) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("INSERT INTO %s(%s)\n",
		tableName, strings.Join(fields, ",")))
	buf.WriteString("VALUES(")

	for ind := range fields {
		if ind > 0 {
			buf.WriteString(",")
		}

		buf.WriteString(fmt.Sprintf("$%d", ind+1))
	}

	buf.WriteString(")\n")

	return buf.String()
}
<|endoftext|>


<commit_before>package resolv

import (
	"fmt"
	"sync"

	"github.com/miekg/dns"
	"golang.org/x/net/context"
)

type Resolver struct {
}

func NewResolver() *Resolver {
	return &Resolver{}
}

func (r *Resolver) Resolve(req *Request) <-chan *Response {
	// Buffered channel to avoid goroutine leaking
	c := make(chan *Response, 1)

	go func() {
		defer func() {
			// Unhandled error
			if err := recover(); err != nil {
				c <- &Response{Err: fmt.Errorf("resolve: %v", err)}
			}
			close(c)
		}()

		// Prepare message
		m := new(dns.Msg)
		m.Id = dns.Id()
		m.RecursionDesired = true
		m.Question = make([]dns.Question, 1)
		m.Question[0] = dns.Question{req.Name, req.Type, req.Class}

		cli := new(dns.Client)
		cli.Net = req.Mode

		// TODO:
		// https://godoc.org/github.com/miekg/dns#Client
		// - Truncated response
		// - Timeout
		// - Retry
		in, rtt, err := cli.Exchange(m, req.Addr)
		if err != nil {
			c <- &Response{Err: err}
			return
		}

		// DNS error
		if in.Rcode != dns.RcodeSuccess {
			err := NewDNSError(
				dns.RcodeToString[in.Rcode],
				req,
			)
			c <- &Response{Err: err}
			return
		}

		c <- &Response{
			Msg: in,
			Rtt: rtt,
		}
	}()

	return c
}

func (r *Resolver) merge(ctx context.Context, cs ...<-chan *Response) <-chan *Response {
	var wg sync.WaitGroup
	out := make(chan *Response)

	// Start an output goroutine for each input channel in cs.  output
	// copies values from c to out until c or done is closed, then calls
	// wg.Done.
	output := func(c <-chan *Response) {
		defer wg.Done()
		for resp := range c {
			select {
			case out <- resp:
			case <-ctx.Done():
				return
			}
		}
	}

	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// Start a goroutine to close out once all the output goroutines are
	// done.  This must start after the wg.Add call.
	go func() {
		wg.Wait()
		close(out)
	}()

	return out
}

func (r *Resolver) ResolveTypes(ctx context.Context, addr string, name string, types []uint16, options ...RequestOption) <-chan *Response {
	cs := []<-chan *Response{}

	for i := 0; i < len(types); i++ {
		req := NewRequest(addr, name, types[i], options...)
		c := r.Resolve(req)
		cs = append(cs, c)
	}

	return r.merge(ctx, cs...)
}

func (r *Resolver) ResolveNames(ctx context.Context, addr string, type_ uint16, names []string, options ...RequestOption) <-chan *Response {
	cs := []<-chan *Response{}

	for i := 0; i < len(names); i++ {
		req := NewRequest(addr, names[i], type_, options...)
		c := r.Resolve(req)
		cs = append(cs, c)
	}

	return r.merge(ctx, cs...)
}
<commit_msg>Move common code to Exchange function<commit_after>package resolv

import (
	"fmt"
	"sync"

	"github.com/miekg/dns"
	"golang.org/x/net/context"
)

type Resolver struct {
}

func NewResolver() *Resolver {
	return &Resolver{}
}

func (r *Resolver) Resolve(req *Request) <-chan *Response {
	// Buffered channel to avoid goroutine leaking
	c := make(chan *Response, 1)

	go func() {
		defer func() {
			// Unhandled error
			if err := recover(); err != nil {
				c <- &Response{Err: fmt.Errorf("resolve: %v", err)}
			}
			close(c)
		}()

		// Prepare message
		m := new(dns.Msg)
		m.Id = dns.Id()
		m.RecursionDesired = true
		m.Question = make([]dns.Question, 1)
		m.Question[0] = dns.Question{req.Name, req.Type, req.Class}

		cli := new(dns.Client)
		cli.Net = req.Mode

		// TODO:
		// https://godoc.org/github.com/miekg/dns#Client
		// - Truncated response
		// - Timeout
		// - Retry
		in, rtt, err := cli.Exchange(m, req.Addr)
		if err != nil {
			c <- &Response{Err: err}
			return
		}

		// DNS error
		if in.Rcode != dns.RcodeSuccess {
			err := NewDNSError(
				dns.RcodeToString[in.Rcode],
				req,
			)
			c <- &Response{Err: err}
			return
		}

		c <- &Response{
			Msg: in,
			Rtt: rtt,
		}
	}()

	return c
}

func (r *Resolver) Exchange(ctx context.Context, reqs ...*Request) <-chan *Response {
	cs := []<-chan *Response{}

	for i := 0; i < len(reqs); i++ {
		c := r.Resolve(reqs[i])
		cs = append(cs, c)
	}

	return r.merge(ctx, cs...)
}

func (r *Resolver) merge(ctx context.Context, cs ...<-chan *Response) <-chan *Response {
	var wg sync.WaitGroup
	out := make(chan *Response)

	// Start an output goroutine for each input channel in cs.  output
	// copies values from c to out until c or done is closed, then calls
	// wg.Done.
	output := func(c <-chan *Response) {
		defer wg.Done()
		for resp := range c {
			select {
			case out <- resp:
			case <-ctx.Done():
				return
			}
		}
	}

	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// Start a goroutine to close out once all the output goroutines are
	// done.  This must start after the wg.Add call.
	go func() {
		wg.Wait()
		close(out)
	}()

	return out
}

func (r *Resolver) ResolveTypes(ctx context.Context, addr string, name string, types []uint16, options ...RequestOption) <-chan *Response {
	reqs := []*Request{}

	for i := 0; i < len(types); i++ {
		req := NewRequest(addr, name, types[i], options...)
		reqs = append(reqs, req)
	}

	return r.Exchange(ctx, reqs...)
}

func (r *Resolver) ResolveNames(ctx context.Context, addr string, type_ uint16, names []string, options ...RequestOption) <-chan *Response {
	reqs := []*Request{}

	for i := 0; i < len(names); i++ {
		req := NewRequest(addr, names[i], type_, options...)
		reqs = append(reqs, req)
	}

	return r.Exchange(ctx, reqs...)
}
<|endoftext|>


<commit_before>package bartender

import (
	"fmt"
	"reflect"

	"github.com/go-kit/kit/endpoint"
)

var errorType reflect.Type

func init() {
	var err error
	errorType = reflect.TypeOf(&err).Elem()
}

func isFunc(v interface{}) bool {
	return reflect.TypeOf(v).Kind() == reflect.Func
}

func isTypeError(t reflect.Type) bool {
	return t.String() == "error"
}

// validEndpoint returns the reflect.Type of the given fn
// and error if the fn is not valid for casting as Endpoint
func validEndpoint(fn interface{}) (fnt reflect.Type, err error) {

	if !isFunc(fn) {
		err = fmt.Errorf("input is not a function: %#v", fn)
		return
	}
	fnt = reflect.TypeOf(fn)

	if n := fnt.NumIn(); n != 1 {
		err = fmt.Errorf("Given function acceptEndpoint functions accepts only 1 argument. Given function has %d", n)
		// TODO: should allow to take context.Context as input, too
		return
	}

	if n := fnt.NumOut(); n != 2 {
		err = fmt.Errorf("Given function returns too many parameters (%d). Should be 2", n)
		// TODO: might accept lesser return parameter (e.g. no error or not output)
		return
	}

	if out2 := fnt.Out(1); !isTypeError(out2) {
		err = fmt.Errorf("Second return parameter is not error")
		return
	}

	return
}

// CanEndpoint test if a function can be cast as Endpoint
func CanEndpoint(fn interface{}) bool {
	if _, err := validEndpoint(fn); err != nil {
		return false
	}
	return true
}

// Endpoint wraps a provided function and returns
// a valid go-kit Endpoint
func Endpoint(fn interface{}) (e endpoint.Endpoint, err error) {

	var fnt reflect.Type
	if fnt, err = validEndpoint(fn); err != nil {
		return
	}

	// type of request
	reqt := fnt.In(0)

	// need to call fn using fnv.Call
	fnv := reflect.ValueOf(fn)

	// reflect on endpoint
	et := reflect.TypeOf(e)
	respt := et.Out(0)
	//errt := et.Out(1)

	// TODO:
	// -> to examine if the input values type
	//    and return error if mismatch
	callFn := func(in []reflect.Value) []reflect.Value {

		// TODO: test if the input variable
		//       match the type of request

		// set the input variable to typed variable
		reqv := reflect.New(reqt)
		reqv.Elem().Set(in[1].Elem())
		out := fnv.Call([]reflect.Value{reqv.Elem()})

		// cast the output variable's address
		// into an empty interface
		respv := reflect.New(respt)
		respv.Elem().Set(out[0])

		return []reflect.Value{respv.Elem(), out[1]}
	}

	makeEndpoint := func(fin, fpout interface{}) {
		// fptr is a pointer to a function.
		// Obtain the function value itself (likely nil) as a reflect.Value
		// so that we can query its type and then set the value.
		fn := reflect.ValueOf(fpout).Elem()

		// Make a function of the right type.
		v := reflect.MakeFunc(fn.Type(), callFn)

		// Assign it to the value fn represents.
		fn.Set(v)
	}

	makeEndpoint(fn, &e)

	return
}
<commit_msg>Refactored endpoint input and output check<commit_after>package bartender

import (
	"fmt"
	"reflect"

	"github.com/go-kit/kit/endpoint"
)

var errorType reflect.Type

func init() {
	var err error
	errorType = reflect.TypeOf(&err).Elem()
}

func isFunc(v interface{}) bool {
	return reflect.TypeOf(v).Kind() == reflect.Func
}

func isTypeError(t reflect.Type) bool {
	return t.String() == "error"
}

func validEndpointIn(fnt *reflect.Type) (err error) {
	if n := (*fnt).NumIn(); n != 1 {
		err = fmt.Errorf("Given function acceptEndpoint functions accepts only 1 argument. Given function has %d", n)
		// TODO: should allow to take context.Context as input, too
		return
	}
	return
}

func validEndpointOut(fnt *reflect.Type) (err error) {
	if n := (*fnt).NumOut(); n != 2 {
		err = fmt.Errorf("Given function returns too many parameters (%d). Should be 2", n)
		// TODO: might accept lesser return parameter (e.g. no error or not output)
		return
	}

	if out2 := (*fnt).Out(1); !isTypeError(out2) {
		err = fmt.Errorf("Second return parameter is not error")
		return
	}
	return
}

// validEndpoint returns the reflect.Type of the given fn
// and error if the fn is not valid for casting as Endpoint
func validEndpoint(fn interface{}) (fnt reflect.Type, err error) {

	if !isFunc(fn) {
		err = fmt.Errorf("input is not a function: %#v", fn)
		return
	}
	fnt = reflect.TypeOf(fn)

	if err = validEndpointIn(&fnt); err != nil {
		return
	}

	if err = validEndpointOut(&fnt); err != nil {
		return
	}

	return
}

// CanEndpoint test if a function can be cast as Endpoint
func CanEndpoint(fn interface{}) bool {
	if _, err := validEndpoint(fn); err != nil {
		return false
	}
	return true
}

// Endpoint wraps a provided function and returns
// a valid go-kit Endpoint
func Endpoint(fn interface{}) (e endpoint.Endpoint, err error) {

	var fnt reflect.Type
	if fnt, err = validEndpoint(fn); err != nil {
		return
	}

	// type of request
	reqt := fnt.In(0)

	// need to call fn using fnv.Call
	fnv := reflect.ValueOf(fn)

	// reflect on endpoint
	et := reflect.TypeOf(e)
	respt := et.Out(0)
	//errt := et.Out(1)

	// TODO:
	// -> to examine if the input values type
	//    and return error if mismatch
	callFn := func(in []reflect.Value) []reflect.Value {

		// TODO: test if the input variable
		//       match the type of request

		// set the input variable to typed variable
		reqv := reflect.New(reqt)
		reqv.Elem().Set(in[1].Elem())
		out := fnv.Call([]reflect.Value{reqv.Elem()})

		// cast the output variable's address
		// into an empty interface
		respv := reflect.New(respt)
		respv.Elem().Set(out[0])

		return []reflect.Value{respv.Elem(), out[1]}
	}

	makeEndpoint := func(fin, fpout interface{}) {
		// fptr is a pointer to a function.
		// Obtain the function value itself (likely nil) as a reflect.Value
		// so that we can query its type and then set the value.
		fn := reflect.ValueOf(fpout).Elem()

		// Make a function of the right type.
		v := reflect.MakeFunc(fn.Type(), callFn)

		// Assign it to the value fn represents.
		fn.Set(v)
	}

	makeEndpoint(fn, &e)

	return
}
<|endoftext|>


<commit_before>/*
Entry methods and functions that access a database are placed here.
*/
package forum

import (
	"database/sql"
	"errors"
	"strings"
)

// Stores an entry to the database and correctly builds its ancestry based
// on its parent's ID.
func (e *Entry) Persist(parentId int64) error {
	//Trim
	e.Title = strings.TrimSpace(e.Title)
	e.Body = strings.TrimSpace(e.Body)

	//Validate
	if e.Body == "" {
		return errors.New("The Body must not be empty or consist solely of whitespace.")
	}

	//Wrap in a transaction
	tx, err := Config.DB.Begin()

	EntryCreateStmt, err := tx.Prepare(queries.EntryCreate)
	defer EntryCreateStmt.Close()
	if err != nil {
		_ = tx.Rollback()
		return errors.New("Error: We had a database problem trying to create your entry.")
	}

	//Note: because pq handles LastInsertId oddly (or not at all?), instead of
	//calling .Exec() then .LastInsertId, we prepare a statement that ends in
	//`RETURNING id` and we .QueryRow().Select() the result
	err = EntryCreateStmt.QueryRow(e.Title, e.Body, e.Url, e.AuthorId).Scan(&e.Id)
	defer EntryCreateStmt.Close()
	if err != nil {
		tx.Rollback()
		return errors.New("Error: your username or email address was already found in the database. Please choose differently.")
	}

	EntryClosureTableCreateStmt, err := tx.Prepare(queries.EntryClosureTableCreate)
	defer EntryClosureTableCreateStmt.Close()
	if err != nil {
		tx.Rollback()
		// return err
		return errors.New("Error: We had a database problem trying to create ancestry information.")
	}

	_, err = EntryClosureTableCreateStmt.Exec(e.Id, parentId)
	if err != nil {
		tx.Rollback()
		return errors.New("Error: We couldn't save the relationship between your comment and its parent comment.")
	}

	tx.Commit()

	return nil
}

//Retrieve one entry by its ID, if it exists. Error if not.
func OneEntry(id int64) (*Entry, error) {
	e := new(Entry)
	var err error = nil

	stmt, err := Config.DB.Prepare(queries.OneEntry)
	defer stmt.Close()
	if err != nil {
		return e, err
	}

	err = stmt.QueryRow(id).Scan(&e.Id, &e.Title, &e.Body, &e.Url, &e.Created, &e.AuthorId, &e.Forum, &e.AuthorHandle, &e.Seconds, &e.Upvotes, &e.Downvotes)
	if err != nil {
		e = new(Entry)
		return e, err
	}

	return e, err
}

// Retrieves all entries that are descendants of the ancestral entry, including the ancestral entry itself
func DescendantEntries(root int64, user User) (*Entry, error) {
	return getEntries(root, "AllDescendants", user)
}

// Retrieves entries that are immediate descendants of the ancestral entry, including the ancestral entry itself
func DepthOneDescendantEntries(root int64, user User) (*Entry, error) {
	return getEntries(root, "DepthOneDescendants", user)
}

func getEntries(root int64, flag string, user User) (*Entry, error) {
	// Store output in a map initially. Get it all in here before you try to build the tree.
	entries := map[int64]*Entry{}                //k: id => v: Entry
	relationships := make([]map[string]int64, 0) //A slice of maps with k: parentId in entries map => v: childId in entries map

	var stmt *sql.Stmt
	var err error
	switch flag {
	case "AllDescendants":
		stmt, err = Config.DB.Prepare(queries.DescendantEntriesChildParent)
		defer stmt.Close()
	case "DepthOneDescendants":
		stmt, err = Config.DB.Prepare(queries.DepthOneDescendantEntriesChildParent)
		defer stmt.Close()
	}
	if err != nil {
		return New(), err
	}

	// Query from that prepared statement
	rows, err := stmt.Query(root, user.GetId())
	if err != nil {
		return New(), err
	}

	// Iterate over the rows
	for rows.Next() {
		var e *Entry = New()
		var ancestor int64
		err = rows.Scan(&ancestor, &e.Id, &e.Title, &e.Body, &e.Url, &e.Created, &e.AuthorId, &e.Forum, &e.AuthorHandle, &e.Seconds, &e.Upvotes, &e.Downvotes, &e.UserVote.Upvote, &e.UserVote.Downvote)
		if err != nil {
			return e, err
		}

		entries[e.Id] = e
		relationships = append(relationships, map[string]int64{"Parent": ancestor, "Child": e.Id})
	}

	//Construct the full Entry:
	for _, rel := range relationships {
		if rel["Parent"] == rel["Child"] {
			continue
		}
		entries[int64(rel["Parent"])].AddChild(entries[int64(rel["Child"])])
	}

	return Arrange(entries[root]), nil
}
<commit_msg>Was wrong, error checking must be before statement closure.<commit_after>/*
Entry methods and functions that access a database are placed here.
*/
package forum

import (
	"database/sql"
	"errors"
	"strings"
)

// Stores an entry to the database and correctly builds its ancestry based
// on its parent's ID.
func (e *Entry) Persist(parentId int64) error {
	//Trim
	e.Title = strings.TrimSpace(e.Title)
	e.Body = strings.TrimSpace(e.Body)

	//Validate
	if e.Body == "" {
		return errors.New("The Body must not be empty or consist solely of whitespace.")
	}

	//Wrap in a transaction
	tx, err := Config.DB.Begin()

	EntryCreateStmt, err := tx.Prepare(queries.EntryCreate)
	if err != nil {
		tx.Rollback()
		return errors.New("Error: We had a database problem trying to create your entry.")
	}
	defer EntryCreateStmt.Close()

	//Note: because pq handles LastInsertId oddly (or not at all?), instead of
	//calling .Exec() then .LastInsertId, we prepare a statement that ends in
	//`RETURNING id` and we .QueryRow().Select() the result
	err = EntryCreateStmt.QueryRow(e.Title, e.Body, e.Url, e.AuthorId).Scan(&e.Id)
	if err != nil {
		tx.Rollback()
		return errors.New("Error: your username or email address was already found in the database. Please choose differently.")
	}
	defer EntryCreateStmt.Close()

	EntryClosureTableCreateStmt, err := tx.Prepare(queries.EntryClosureTableCreate)
	if err != nil {
		tx.Rollback()
		return errors.New("Error: We had a database problem trying to create ancestry information.")
	}
	defer EntryClosureTableCreateStmt.Close()

	_, err = EntryClosureTableCreateStmt.Exec(e.Id, parentId)
	if err != nil {
		tx.Rollback()
		return errors.New("Error: We couldn't save the relationship between your comment and its parent comment.")
	}

	tx.Commit()

	return nil
}

//Retrieve one entry by its ID, if it exists. Error if not.
func OneEntry(id int64) (*Entry, error) {
	e := new(Entry)
	var err error = nil

	stmt, err := Config.DB.Prepare(queries.OneEntry)
	if err != nil {
		return e, err
	}
	defer stmt.Close()

	err = stmt.QueryRow(id).Scan(&e.Id, &e.Title, &e.Body, &e.Url, &e.Created, &e.AuthorId, &e.Forum, &e.AuthorHandle, &e.Seconds, &e.Upvotes, &e.Downvotes)
	if err != nil {
		e = new(Entry)
		return e, err
	}

	return e, err
}

// Retrieves all entries that are descendants of the ancestral entry, including the ancestral entry itself
func DescendantEntries(root int64, user User) (*Entry, error) {
	return getEntries(root, "AllDescendants", user)
}

// Retrieves entries that are immediate descendants of the ancestral entry, including the ancestral entry itself
func DepthOneDescendantEntries(root int64, user User) (*Entry, error) {
	return getEntries(root, "DepthOneDescendants", user)
}

func getEntries(root int64, flag string, user User) (*Entry, error) {
	// Store output in a map initially. Get it all in here before you try to build the tree.
	entries := map[int64]*Entry{}                //k: id => v: Entry
	relationships := make([]map[string]int64, 0) //A slice of maps with k: parentId in entries map => v: childId in entries map

	var stmt *sql.Stmt
	var err error
	switch flag {
	case "AllDescendants":
		stmt, err = Config.DB.Prepare(queries.DescendantEntriesChildParent)
	case "DepthOneDescendants":
		stmt, err = Config.DB.Prepare(queries.DepthOneDescendantEntriesChildParent)
	}
	if err != nil {
		return New(), err
	}
	defer stmt.Close()

	// Query from that prepared statement
	rows, err := stmt.Query(root, user.GetId())
	if err != nil {
		return New(), err
	}
	defer rows.Close()

	// Iterate over the rows
	for rows.Next() {
		var e *Entry = New()
		var ancestor int64
		err = rows.Scan(&ancestor, &e.Id, &e.Title, &e.Body, &e.Url, &e.Created, &e.AuthorId, &e.Forum, &e.AuthorHandle, &e.Seconds, &e.Upvotes, &e.Downvotes, &e.UserVote.Upvote, &e.UserVote.Downvote)
		if err != nil {
			return e, err
		}

		entries[e.Id] = e
		relationships = append(relationships, map[string]int64{"Parent": ancestor, "Child": e.Id})
	}

	//Construct the full Entry:
	for _, rel := range relationships {
		if rel["Parent"] == rel["Child"] {
			continue
		}
		entries[int64(rel["Parent"])].AddChild(entries[int64(rel["Child"])])
	}

	return Arrange(entries[root]), nil
}
<|endoftext|>


<commit_before>package mongodb

import (
	"fmt"
	"strings"
	"time"

	"github.com/elastic/libbeat/common"
	"github.com/elastic/libbeat/logp"
	"github.com/elastic/packetbeat/config"
	"github.com/elastic/packetbeat/procs"
	"github.com/elastic/packetbeat/protos"
	"github.com/elastic/packetbeat/protos/tcp"
)

type Mongodb struct {
	// config
	Ports          []int
	Send_request   bool
	Send_response  bool
	Max_docs       int
	Max_doc_length int

	transactionsMap map[common.HashableTcpTuple]*MongodbTransaction

	results chan common.MapStr
}

func (mongodb *Mongodb) InitDefaults() {
	mongodb.Send_request = false
	mongodb.Send_response = false
	mongodb.Max_docs = 10
	mongodb.Max_doc_length = 5000
}

func (mongodb *Mongodb) setFromConfig(config config.Mongodb) error {
	mongodb.Ports = config.Ports

	if config.Send_request != nil {
		mongodb.Send_request = *config.Send_request
	}
	if config.Send_response != nil {
		mongodb.Send_response = *config.Send_response
	}
	if config.Max_docs != nil {
		mongodb.Max_docs = *config.Max_docs
	}
	if config.Max_doc_length != nil {
		mongodb.Max_doc_length = *config.Max_doc_length
	}
	return nil
}

func (mongodb *Mongodb) GetPorts() []int {
	return mongodb.Ports
}

func (mongodb *Mongodb) Init(test_mode bool, results chan common.MapStr) error {
	logp.Debug("mongodb", "Init a MongoDB protocol parser")

	mongodb.InitDefaults()
	if !test_mode {
		err := mongodb.setFromConfig(config.ConfigSingleton.Protocols.Mongodb)
		if err != nil {
			return err
		}
	}

	mongodb.transactionsMap = make(map[common.HashableTcpTuple]*MongodbTransaction, TransactionsHashSize)
	mongodb.results = results

	return nil
}

func (mongodb *Mongodb) Parse(pkt *protos.Packet, tcptuple *common.TcpTuple, dir uint8,
	private protos.ProtocolData) protos.ProtocolData {

	logp.Debug("mongodb", "Parse method triggered")

	defer logp.Recover("ParseMongodb exception")

	// Either fetch or initialize current data struct for this parser
	priv := mongodbPrivateData{}
	if private != nil {
		var ok bool
		priv, ok = private.(mongodbPrivateData)
		if !ok {
			priv = mongodbPrivateData{}
		}
	}

	if priv.Data[dir] == nil {
		priv.Data[dir] = &MongodbStream{
			tcptuple: tcptuple,
			data:     pkt.Payload,
			message:  &MongodbMessage{Ts: pkt.Ts},
		}
	} else {
		// concatenate bytes
		priv.Data[dir].data = append(priv.Data[dir].data, pkt.Payload...)
		if len(priv.Data[dir].data) > tcp.TCP_MAX_DATA_IN_STREAM {
			logp.Debug("mongodb", "Stream data too large, dropping TCP stream")
			priv.Data[dir] = nil
			return priv
		}
	}

	stream := priv.Data[dir]
	for len(stream.data) > 0 {
		if stream.message == nil {
			stream.message = &MongodbMessage{Ts: pkt.Ts}
		}

		ok, complete := mongodbMessageParser(priv.Data[dir])

		if !ok {
			// drop this tcp stream. Will retry parsing with the next
			// segment in it
			priv.Data[dir] = nil
			logp.Debug("mongodb", "Ignore Mongodb message. Drop tcp stream. Try parsing with the next segment")
			return priv
		}

		if complete {

			logp.Debug("mongodb", "MongoDB message complete")

			// all ok, go to next level
			mongodb.handleMongodb(stream.message, tcptuple, dir)

			// and reset message
			stream.PrepareForNewMessage()
		} else {
			// wait for more data
			logp.Debug("mongodb", "MongoDB wait for more data before parsing message")
			break
		}
	}

	return priv
}

func (mongodb *Mongodb) handleMongodb(m *MongodbMessage, tcptuple *common.TcpTuple,
	dir uint8) {

	m.TcpTuple = *tcptuple
	m.Direction = dir
	m.CmdlineTuple = procs.ProcWatcher.FindProcessesTuple(tcptuple.IpPort())

	if m.IsResponse {
		logp.Debug("mongodb", "MongoDB response message")
		mongodb.receivedMongodbResponse(m)
	} else {
		logp.Debug("mongodb", "MongoDB request message")
		mongodb.receivedMongodbRequest(m)
	}
}

func (mongodb *Mongodb) receivedMongodbRequest(msg *MongodbMessage) {
	// Add it to the HT
	tuple := msg.TcpTuple

	trans := mongodb.transactionsMap[tuple.Hashable()]
	if trans != nil {
		if trans.Mongodb != nil {
			logp.Warn("Two requests without a Response. Dropping old request")
		}
	} else {
		logp.Debug("mongodb", "Initialize new transaction from request")
		trans = &MongodbTransaction{Type: "mongodb", tuple: tuple}
		mongodb.transactionsMap[tuple.Hashable()] = trans
	}

	trans.Mongodb = common.MapStr{}

	trans.event = msg.event

	trans.method = msg.method

	trans.cmdline = msg.CmdlineTuple
	trans.ts = msg.Ts
	trans.Ts = int64(trans.ts.UnixNano() / 1000) // transactions have microseconds resolution
	trans.JsTs = msg.Ts
	trans.Src = common.Endpoint{
		Ip:   msg.TcpTuple.Src_ip.String(),
		Port: msg.TcpTuple.Src_port,
		Proc: string(msg.CmdlineTuple.Src),
	}
	trans.Dst = common.Endpoint{
		Ip:   msg.TcpTuple.Dst_ip.String(),
		Port: msg.TcpTuple.Dst_port,
		Proc: string(msg.CmdlineTuple.Dst),
	}
	if msg.Direction == tcp.TcpDirectionReverse {
		trans.Src, trans.Dst = trans.Dst, trans.Src
	}
	trans.params = msg.params
	trans.resource = msg.resource

	if trans.timer != nil {
		trans.timer.Stop()
	}
	trans.timer = time.AfterFunc(TransactionTimeout, func() { mongodb.expireTransaction(trans) })

}

func (mongodb *Mongodb) expireTransaction(trans *MongodbTransaction) {
	logp.Debug("mongodb", "Expire transaction")
	// remove from map
	delete(mongodb.transactionsMap, trans.tuple.Hashable())
}

func (mongodb *Mongodb) receivedMongodbResponse(msg *MongodbMessage) {

	tuple := msg.TcpTuple
	trans := mongodb.transactionsMap[tuple.Hashable()]
	if trans == nil {
		logp.Warn("Response from unknown transaction. Ignoring.")
		return
	}
	// check if the request was received
	if trans.Mongodb == nil {
		logp.Warn("Response from unknown transaction. Ignoring.")
		return

	}

	// Merge request and response events attributes
	for k, v := range msg.event {
		trans.event[k] = v
	}

	trans.error = msg.error
	trans.documents = msg.documents

	trans.ResponseTime = int32(msg.Ts.Sub(trans.ts).Nanoseconds() / 1e6) // resp_time in milliseconds

	mongodb.publishTransaction(trans)

	logp.Debug("mongodb", "Mongodb transaction completed: %s", trans.Mongodb)

	// remove from map
	delete(mongodb.transactionsMap, trans.tuple.Hashable())
	if trans.timer != nil {
		trans.timer.Stop()
	}
}

func (mongodb *Mongodb) GapInStream(tcptuple *common.TcpTuple, dir uint8,
	private protos.ProtocolData) protos.ProtocolData {

	// TODO

	return private
}

func (mongodb *Mongodb) ReceivedFin(tcptuple *common.TcpTuple, dir uint8,
	private protos.ProtocolData) protos.ProtocolData {

	// TODO
	return private
}

func copy_map_without_key(d map[string]interface{}, key string) map[string]interface{} {
	res := map[string]interface{}{}
	for k, v := range d {
		if k != key {
			res[k] = v
		}
	}
	return res
}

func reconstructQuery(t *MongodbTransaction, full bool) (query string) {
	query = t.resource + "." + t.method + "("
	if len(t.params) > 0 {
		var err error
		var params string
		if !full {
			// remove the actual data.
			// TODO: review if we need to add other commands here
			if t.method == "insert" {
				params, err = doc2str(copy_map_without_key(t.params, "documents"))
			} else if t.method == "update" {
				params, err = doc2str(copy_map_without_key(t.params, "updates"))
			} else if t.method == "findandmodify" {
				params, err = doc2str(copy_map_without_key(t.params, "update"))
			}
		} else {
			params, err = doc2str(t.params)
		}
		if err != nil {
			logp.Debug("mongodb", "Error marshaling params: %v", err)
		} else {
			query += params
		}
	}
	query += ")"
	skip, _ := t.event["numberToSkip"].(int)
	if skip > 0 {
		query += fmt.Sprintf(".skip(%d)", skip)
	}

	limit, _ := t.event["numberToReturn"].(int)
	if limit > 0 && limit < 0x7fffffff {
		query += fmt.Sprintf(".limit(%d)", limit)
	}
	return
}

func (mongodb *Mongodb) publishTransaction(t *MongodbTransaction) {

	if mongodb.results == nil {
		logp.Debug("mongodb", "Try to publish transaction with null results")
		return
	}

	event := common.MapStr{}
	event["type"] = "mongodb"
	if t.error == "" {
		event["status"] = common.OK_STATUS
	} else {
		t.event["error"] = t.error
		event["status"] = common.ERROR_STATUS
	}
	event["mongodb"] = t.event
	event["method"] = t.method
	event["resource"] = t.resource
	event["query"] = reconstructQuery(t, false)
	event["responsetime"] = t.ResponseTime
	event["bytes_in"] = uint64(t.BytesIn)
	event["bytes_out"] = uint64(t.BytesOut)
	event["timestamp"] = common.Time(t.ts)
	event["src"] = &t.Src
	event["dst"] = &t.Dst

	if mongodb.Send_request {
		event["request"] = reconstructQuery(t, true)
	}
	if mongodb.Send_response {
		if len(t.documents) > 0 {
			// response field needs to be a string
			docs := make([]string, 0, len(t.documents))
			for i, doc := range t.documents {
				if mongodb.Max_docs > 0 && i >= mongodb.Max_docs {
					docs = append(docs, "[...]")
					break
				}
				str, err := doc2str(doc)
				if err != nil {
					logp.Warn("Failed to JSON marshal document from Mongo: %v (error: %v)", doc, err)
				} else {
					if mongodb.Max_doc_length > 0 && len(str) > mongodb.Max_doc_length {
						str = str[:mongodb.Max_doc_length] + " ..."
					}
					docs = append(docs, str)
				}
			}
			event["response"] = strings.Join(docs, "\n")
		}
	}

	mongodb.results <- event
}
<commit_msg>post merge fix: GapInStream interface changed slightly<commit_after>package mongodb

import (
	"fmt"
	"strings"
	"time"

	"github.com/elastic/libbeat/common"
	"github.com/elastic/libbeat/logp"
	"github.com/elastic/packetbeat/config"
	"github.com/elastic/packetbeat/procs"
	"github.com/elastic/packetbeat/protos"
	"github.com/elastic/packetbeat/protos/tcp"
)

type Mongodb struct {
	// config
	Ports          []int
	Send_request   bool
	Send_response  bool
	Max_docs       int
	Max_doc_length int

	transactionsMap map[common.HashableTcpTuple]*MongodbTransaction

	results chan common.MapStr
}

func (mongodb *Mongodb) InitDefaults() {
	mongodb.Send_request = false
	mongodb.Send_response = false
	mongodb.Max_docs = 10
	mongodb.Max_doc_length = 5000
}

func (mongodb *Mongodb) setFromConfig(config config.Mongodb) error {
	mongodb.Ports = config.Ports

	if config.Send_request != nil {
		mongodb.Send_request = *config.Send_request
	}
	if config.Send_response != nil {
		mongodb.Send_response = *config.Send_response
	}
	if config.Max_docs != nil {
		mongodb.Max_docs = *config.Max_docs
	}
	if config.Max_doc_length != nil {
		mongodb.Max_doc_length = *config.Max_doc_length
	}
	return nil
}

func (mongodb *Mongodb) GetPorts() []int {
	return mongodb.Ports
}

func (mongodb *Mongodb) Init(test_mode bool, results chan common.MapStr) error {
	logp.Debug("mongodb", "Init a MongoDB protocol parser")

	mongodb.InitDefaults()
	if !test_mode {
		err := mongodb.setFromConfig(config.ConfigSingleton.Protocols.Mongodb)
		if err != nil {
			return err
		}
	}

	mongodb.transactionsMap = make(map[common.HashableTcpTuple]*MongodbTransaction, TransactionsHashSize)
	mongodb.results = results

	return nil
}

func (mongodb *Mongodb) Parse(pkt *protos.Packet, tcptuple *common.TcpTuple, dir uint8,
	private protos.ProtocolData) protos.ProtocolData {

	logp.Debug("mongodb", "Parse method triggered")

	defer logp.Recover("ParseMongodb exception")

	// Either fetch or initialize current data struct for this parser
	priv := mongodbPrivateData{}
	if private != nil {
		var ok bool
		priv, ok = private.(mongodbPrivateData)
		if !ok {
			priv = mongodbPrivateData{}
		}
	}

	if priv.Data[dir] == nil {
		priv.Data[dir] = &MongodbStream{
			tcptuple: tcptuple,
			data:     pkt.Payload,
			message:  &MongodbMessage{Ts: pkt.Ts},
		}
	} else {
		// concatenate bytes
		priv.Data[dir].data = append(priv.Data[dir].data, pkt.Payload...)
		if len(priv.Data[dir].data) > tcp.TCP_MAX_DATA_IN_STREAM {
			logp.Debug("mongodb", "Stream data too large, dropping TCP stream")
			priv.Data[dir] = nil
			return priv
		}
	}

	stream := priv.Data[dir]
	for len(stream.data) > 0 {
		if stream.message == nil {
			stream.message = &MongodbMessage{Ts: pkt.Ts}
		}

		ok, complete := mongodbMessageParser(priv.Data[dir])

		if !ok {
			// drop this tcp stream. Will retry parsing with the next
			// segment in it
			priv.Data[dir] = nil
			logp.Debug("mongodb", "Ignore Mongodb message. Drop tcp stream. Try parsing with the next segment")
			return priv
		}

		if complete {

			logp.Debug("mongodb", "MongoDB message complete")

			// all ok, go to next level
			mongodb.handleMongodb(stream.message, tcptuple, dir)

			// and reset message
			stream.PrepareForNewMessage()
		} else {
			// wait for more data
			logp.Debug("mongodb", "MongoDB wait for more data before parsing message")
			break
		}
	}

	return priv
}

func (mongodb *Mongodb) handleMongodb(m *MongodbMessage, tcptuple *common.TcpTuple,
	dir uint8) {

	m.TcpTuple = *tcptuple
	m.Direction = dir
	m.CmdlineTuple = procs.ProcWatcher.FindProcessesTuple(tcptuple.IpPort())

	if m.IsResponse {
		logp.Debug("mongodb", "MongoDB response message")
		mongodb.receivedMongodbResponse(m)
	} else {
		logp.Debug("mongodb", "MongoDB request message")
		mongodb.receivedMongodbRequest(m)
	}
}

func (mongodb *Mongodb) receivedMongodbRequest(msg *MongodbMessage) {
	// Add it to the HT
	tuple := msg.TcpTuple

	trans := mongodb.transactionsMap[tuple.Hashable()]
	if trans != nil {
		if trans.Mongodb != nil {
			logp.Warn("Two requests without a Response. Dropping old request")
		}
	} else {
		logp.Debug("mongodb", "Initialize new transaction from request")
		trans = &MongodbTransaction{Type: "mongodb", tuple: tuple}
		mongodb.transactionsMap[tuple.Hashable()] = trans
	}

	trans.Mongodb = common.MapStr{}

	trans.event = msg.event

	trans.method = msg.method

	trans.cmdline = msg.CmdlineTuple
	trans.ts = msg.Ts
	trans.Ts = int64(trans.ts.UnixNano() / 1000) // transactions have microseconds resolution
	trans.JsTs = msg.Ts
	trans.Src = common.Endpoint{
		Ip:   msg.TcpTuple.Src_ip.String(),
		Port: msg.TcpTuple.Src_port,
		Proc: string(msg.CmdlineTuple.Src),
	}
	trans.Dst = common.Endpoint{
		Ip:   msg.TcpTuple.Dst_ip.String(),
		Port: msg.TcpTuple.Dst_port,
		Proc: string(msg.CmdlineTuple.Dst),
	}
	if msg.Direction == tcp.TcpDirectionReverse {
		trans.Src, trans.Dst = trans.Dst, trans.Src
	}
	trans.params = msg.params
	trans.resource = msg.resource

	if trans.timer != nil {
		trans.timer.Stop()
	}
	trans.timer = time.AfterFunc(TransactionTimeout, func() { mongodb.expireTransaction(trans) })

}

func (mongodb *Mongodb) expireTransaction(trans *MongodbTransaction) {
	logp.Debug("mongodb", "Expire transaction")
	// remove from map
	delete(mongodb.transactionsMap, trans.tuple.Hashable())
}

func (mongodb *Mongodb) receivedMongodbResponse(msg *MongodbMessage) {

	tuple := msg.TcpTuple
	trans := mongodb.transactionsMap[tuple.Hashable()]
	if trans == nil {
		logp.Warn("Response from unknown transaction. Ignoring.")
		return
	}
	// check if the request was received
	if trans.Mongodb == nil {
		logp.Warn("Response from unknown transaction. Ignoring.")
		return

	}

	// Merge request and response events attributes
	for k, v := range msg.event {
		trans.event[k] = v
	}

	trans.error = msg.error
	trans.documents = msg.documents

	trans.ResponseTime = int32(msg.Ts.Sub(trans.ts).Nanoseconds() / 1e6) // resp_time in milliseconds

	mongodb.publishTransaction(trans)

	logp.Debug("mongodb", "Mongodb transaction completed: %s", trans.Mongodb)

	// remove from map
	delete(mongodb.transactionsMap, trans.tuple.Hashable())
	if trans.timer != nil {
		trans.timer.Stop()
	}
}

func (mongodb *Mongodb) GapInStream(tcptuple *common.TcpTuple, dir uint8,
	nbytes int, private protos.ProtocolData) (priv protos.ProtocolData, drop bool) {

	// TODO

	return private, true
}

func (mongodb *Mongodb) ReceivedFin(tcptuple *common.TcpTuple, dir uint8,
	private protos.ProtocolData) protos.ProtocolData {

	// TODO
	return private
}

func copy_map_without_key(d map[string]interface{}, key string) map[string]interface{} {
	res := map[string]interface{}{}
	for k, v := range d {
		if k != key {
			res[k] = v
		}
	}
	return res
}

func reconstructQuery(t *MongodbTransaction, full bool) (query string) {
	query = t.resource + "." + t.method + "("
	if len(t.params) > 0 {
		var err error
		var params string
		if !full {
			// remove the actual data.
			// TODO: review if we need to add other commands here
			if t.method == "insert" {
				params, err = doc2str(copy_map_without_key(t.params, "documents"))
			} else if t.method == "update" {
				params, err = doc2str(copy_map_without_key(t.params, "updates"))
			} else if t.method == "findandmodify" {
				params, err = doc2str(copy_map_without_key(t.params, "update"))
			}
		} else {
			params, err = doc2str(t.params)
		}
		if err != nil {
			logp.Debug("mongodb", "Error marshaling params: %v", err)
		} else {
			query += params
		}
	}
	query += ")"
	skip, _ := t.event["numberToSkip"].(int)
	if skip > 0 {
		query += fmt.Sprintf(".skip(%d)", skip)
	}

	limit, _ := t.event["numberToReturn"].(int)
	if limit > 0 && limit < 0x7fffffff {
		query += fmt.Sprintf(".limit(%d)", limit)
	}
	return
}

func (mongodb *Mongodb) publishTransaction(t *MongodbTransaction) {

	if mongodb.results == nil {
		logp.Debug("mongodb", "Try to publish transaction with null results")
		return
	}

	event := common.MapStr{}
	event["type"] = "mongodb"
	if t.error == "" {
		event["status"] = common.OK_STATUS
	} else {
		t.event["error"] = t.error
		event["status"] = common.ERROR_STATUS
	}
	event["mongodb"] = t.event
	event["method"] = t.method
	event["resource"] = t.resource
	event["query"] = reconstructQuery(t, false)
	event["responsetime"] = t.ResponseTime
	event["bytes_in"] = uint64(t.BytesIn)
	event["bytes_out"] = uint64(t.BytesOut)
	event["timestamp"] = common.Time(t.ts)
	event["src"] = &t.Src
	event["dst"] = &t.Dst

	if mongodb.Send_request {
		event["request"] = reconstructQuery(t, true)
	}
	if mongodb.Send_response {
		if len(t.documents) > 0 {
			// response field needs to be a string
			docs := make([]string, 0, len(t.documents))
			for i, doc := range t.documents {
				if mongodb.Max_docs > 0 && i >= mongodb.Max_docs {
					docs = append(docs, "[...]")
					break
				}
				str, err := doc2str(doc)
				if err != nil {
					logp.Warn("Failed to JSON marshal document from Mongo: %v (error: %v)", doc, err)
				} else {
					if mongodb.Max_doc_length > 0 && len(str) > mongodb.Max_doc_length {
						str = str[:mongodb.Max_doc_length] + " ..."
					}
					docs = append(docs, str)
				}
			}
			event["response"] = strings.Join(docs, "\n")
		}
	}

	mongodb.results <- event
}
<|endoftext|>


<commit_before>// Go support for Protocol Buffers - Google's data interchange format
//
// Copyright 2015 The Go Authors.  All rights reserved.
// https://github.com/golang/protobuf
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Package grpc outputs gRPC service descriptions in Go code.
// It runs as a plugin for the Go protocol buffer compiler plugin.
// It is linked in to protoc-gen-go.
package grpc

import (
	"fmt"
	"path"
	"strconv"
	"strings"

	pb "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/johanbrandhorst/protobuf/protoc-gen-gopherjs/generator"
)

// generatedCodeVersion indicates a version of the generated code.
// It is incremented whenever an incompatibility between the generated code and
// the grpcweb package is introduced; the generated code references
// a constant, grpcweb.GrpcWebPackageIsVersionN (where N is generatedCodeVersion).
const generatedCodeVersion = 1

// Paths for packages used by code generated in this file,
// relative to the import_prefix of the generator.Generator.
const (
	contextPkgPath = "context"
	grpcPkgPath    = "github.com/johanbrandhorst/protobuf/grpcweb"
)

func init() {
	generator.RegisterPlugin(new(grpc))
}

// grpc is an implementation of the Go protocol buffer compiler's
// plugin architecture.  It generates bindings for gRPC support.
type grpc struct {
	gen *generator.Generator
}

// Name returns the name of this plugin, "grpc".
func (g *grpc) Name() string {
	return "grpc"
}

// The names for packages imported in the generated code.
// They may vary from the final path component of the import path
// if the name is used by other packages.
var (
	contextPkg string
	grpcPkg    string
)

// Init initializes the plugin.
func (g *grpc) Init(gen *generator.Generator) {
	g.gen = gen
	contextPkg = generator.RegisterUniquePackageName("context", nil)
	grpcPkg = generator.RegisterUniquePackageName("grpcweb", nil)
}

// Given a type name defined in a .proto, return its object.
// Also record that we're using it, to guarantee the associated import.
func (g *grpc) objectNamed(name string) generator.Object {
	g.gen.RecordTypeUse(name)
	return g.gen.ObjectNamed(name)
}

// Given a type name defined in a .proto, return its name as we will print it.
func (g *grpc) typeName(str string) string {
	return g.gen.TypeName(g.objectNamed(str))
}

// P forwards to g.gen.P.
func (g *grpc) P(args ...interface{}) { g.gen.P(args...) }

// In forwards to g.gen.In.
func (g *grpc) In() { g.gen.In() }

// Out forwards to g.gen.Out.
func (g *grpc) Out() { g.gen.Out() }

// Generate generates code for the services in the given file.
func (g *grpc) Generate(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}

	g.P("// Reference imports to suppress errors if they are not otherwise used.")
	g.P("var _ ", contextPkg, ".Context")
	g.P("var _ ", grpcPkg, ".Client")
	g.P()

	// Assert version compatibility.
	g.P("// This is a compile-time assertion to ensure that this generated file")
	g.P("// is compatible with the grpcweb package it is being compiled against.")
	g.P("const _ = ", grpcPkg, ".GrpcWebPackageIsVersion", generatedCodeVersion)
	g.P()

	for i, service := range file.FileDescriptorProto.Service {
		g.generateService(file, service, i)
	}
}

// GenerateImports generates the import declaration for this file.
func (g *grpc) GenerateImports(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}
	g.P("import (")
	g.In()
	g.P(contextPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, contextPkgPath)))
	g.P()
	g.P(grpcPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, grpcPkgPath)))
	g.Out()
	g.P(")")
	g.P()
}

// reservedClientName records whether a client name is reserved on the client side.
var reservedClientName = map[string]bool{
// TODO: do we need any in gRPC?
}

func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }

// generateService generates all the code for the named service.
func (g *grpc) generateService(file *generator.FileDescriptor, service *pb.ServiceDescriptorProto, index int) {
	path := fmt.Sprintf("6,%d", index) // 6 means service.

	origServName := service.GetName()
	fullServName := origServName
	if pkg := file.GetPackage(); pkg != "" {
		fullServName = pkg + "." + fullServName
	}
	servName := generator.CamelCase(origServName)

	g.P()
	g.P("// Client API for ", servName, " service")
	g.P()

	// Client interface.
	g.gen.PrintComments(path)
	g.P("type ", servName, "Client interface {")
	g.In()
	for i, method := range service.Method {
		g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
		g.P(g.generateClientSignature(servName, method))
	}
	g.Out()
	g.P("}")
	g.P()

	// Client structure.
	g.P("type ", unexport(servName), "Client struct {")
	g.In()
	g.P("client *", grpcPkg, ".Client")
	g.Out()
	g.P("}")
	g.P()

	// NewClient factory.
	g.P("// New", servName, "Client creates a new gRPC-Web client.")
	g.P("func New", servName, "Client (hostname string, opts ...grpcweb.DialOption) ", servName, "Client {")
	g.In()
	g.P("return &", unexport(servName), "Client{")
	g.In()
	g.P("client: ", grpcPkg, `.NewClient(hostname, "`, fullServName, `", opts...),`)
	g.Out()
	g.P("}")
	g.Out()
	g.P("}")
	g.P()

	var methodIndex, streamIndex int
	serviceDescVar := "_" + servName + "_serviceDesc"
	// Client method implementations.
	for _, method := range service.Method {
		var descExpr string
		if method.GetClientStreaming() {
			g.gen.Fail("Client streaming is not supported by gRPC-Web yet")
		}
		if !method.GetServerStreaming() && !method.GetClientStreaming() {
			// Unary RPC method
			descExpr = fmt.Sprintf("&%s.Methods[%d]", serviceDescVar, methodIndex)
			methodIndex++
		} else {
			// Streaming RPC method
			descExpr = fmt.Sprintf("&%s.Streams[%d]", serviceDescVar, streamIndex)
			streamIndex++
		}
		g.generateClientMethod(servName, fullServName, serviceDescVar, method, descExpr)
	}
}

// generateClientSignature returns the client-side signature for a method.
func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
	origMethName := method.GetName()
	methName := generator.CamelCase(origMethName)
	if reservedClientName[methName] {
		methName += "_"
	}
	reqArg := ", in *" + g.typeName(method.GetInputType())
	if method.GetClientStreaming() {
		reqArg = ""
	}
	respName := "*" + g.typeName(method.GetOutputType())
	if method.GetServerStreaming() || method.GetClientStreaming() {
		respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
	}
	return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
}

func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto, descExpr string) {
	methName := generator.CamelCase(method.GetName())
	outType := g.typeName(method.GetOutputType())
	streamType := unexport(servName) + methName + "Client"

	g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
	g.In()
	g.P("req := in.Marshal()")
	g.P()
	switch {
	case !method.GetServerStreaming() && !method.GetClientStreaming():
		g.P(`resp, err := c.client.RPCCall(ctx, "`, method.GetName(), `", req, opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
		return
	case method.GetServerStreaming():
		g.P(`srv, err := c.client.Stream(ctx, "`, method.GetName(), `", req, opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{")
		g.In()
		g.P("stream: srv,")
		g.Out()
		g.P("}, nil")
		g.Out()
		g.P("}")
		g.P()
	case method.GetClientStreaming():
		g.gen.Fail("Client streaming is not yet supported by gRPC-Web")
	}

	genRecv := method.GetServerStreaming()

	// Stream auxiliary types and methods.
	g.P("type ", servName, "_", methName, "Client interface {")
	if genRecv {
		g.In()
		g.P("Recv() (*", outType, ", error)")
		g.Out()
	}
	g.P("}")
	g.P()

	g.P("type ", streamType, " struct {")
	g.In()
	g.P("stream *", grpcPkg, ".StreamClient")
	g.Out()
	g.P("}")
	g.P()

	if genRecv {
		g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.Recv()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	}
}
<commit_msg>Generate gRPC functions for bidi streams<commit_after>// Go support for Protocol Buffers - Google's data interchange format
//
// Copyright 2015 The Go Authors.  All rights reserved.
// https://github.com/golang/protobuf
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Package grpc outputs gRPC service descriptions in Go code.
// It runs as a plugin for the Go protocol buffer compiler plugin.
// It is linked in to protoc-gen-go.
package grpc

import (
	"fmt"
	"path"
	"strconv"
	"strings"

	pb "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/johanbrandhorst/protobuf/protoc-gen-gopherjs/generator"
)

// generatedCodeVersion indicates a version of the generated code.
// It is incremented whenever an incompatibility between the generated code and
// the grpcweb package is introduced; the generated code references
// a constant, grpcweb.GrpcWebPackageIsVersionN (where N is generatedCodeVersion).
const generatedCodeVersion = 2

// Paths for packages used by code generated in this file,
// relative to the import_prefix of the generator.Generator.
const (
	contextPkgPath = "context"
	grpcPkgPath    = "github.com/johanbrandhorst/protobuf/grpcweb"
)

func init() {
	generator.RegisterPlugin(new(grpc))
}

// grpc is an implementation of the Go protocol buffer compiler's
// plugin architecture.  It generates bindings for gRPC support.
type grpc struct {
	gen *generator.Generator
}

// Name returns the name of this plugin, "grpc".
func (g *grpc) Name() string {
	return "grpc"
}

// The names for packages imported in the generated code.
// They may vary from the final path component of the import path
// if the name is used by other packages.
var (
	contextPkg string
	grpcPkg    string
)

// Init initializes the plugin.
func (g *grpc) Init(gen *generator.Generator) {
	g.gen = gen
	contextPkg = generator.RegisterUniquePackageName("context", nil)
	grpcPkg = generator.RegisterUniquePackageName("grpcweb", nil)
}

// Given a type name defined in a .proto, return its object.
// Also record that we're using it, to guarantee the associated import.
func (g *grpc) objectNamed(name string) generator.Object {
	g.gen.RecordTypeUse(name)
	return g.gen.ObjectNamed(name)
}

// Given a type name defined in a .proto, return its name as we will print it.
func (g *grpc) typeName(str string) string {
	return g.gen.TypeName(g.objectNamed(str))
}

// P forwards to g.gen.P.
func (g *grpc) P(args ...interface{}) { g.gen.P(args...) }

// In forwards to g.gen.In.
func (g *grpc) In() { g.gen.In() }

// Out forwards to g.gen.Out.
func (g *grpc) Out() { g.gen.Out() }

// Generate generates code for the services in the given file.
func (g *grpc) Generate(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}

	g.P("// Reference imports to suppress errors if they are not otherwise used.")
	g.P("var _ ", contextPkg, ".Context")
	g.P("var _ ", grpcPkg, ".Client")
	g.P()

	// Assert version compatibility.
	g.P("// This is a compile-time assertion to ensure that this generated file")
	g.P("// is compatible with the grpcweb package it is being compiled against.")
	g.P("const _ = ", grpcPkg, ".GrpcWebPackageIsVersion", generatedCodeVersion)
	g.P()

	for i, service := range file.FileDescriptorProto.Service {
		g.generateService(file, service, i)
	}
}

// GenerateImports generates the import declaration for this file.
func (g *grpc) GenerateImports(file *generator.FileDescriptor) {
	if len(file.FileDescriptorProto.Service) == 0 {
		return
	}
	g.P("import (")
	g.In()
	g.P(contextPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, contextPkgPath)))
	g.P()
	g.P(grpcPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, grpcPkgPath)))
	g.Out()
	g.P(")")
	g.P()
}

// reservedClientName records whether a client name is reserved on the client side.
var reservedClientName = map[string]bool{
// TODO: do we need any in gRPC?
}

func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }

// generateService generates all the code for the named service.
func (g *grpc) generateService(file *generator.FileDescriptor, service *pb.ServiceDescriptorProto, index int) {
	path := fmt.Sprintf("6,%d", index) // 6 means service.

	origServName := service.GetName()
	fullServName := origServName
	if pkg := file.GetPackage(); pkg != "" {
		fullServName = pkg + "." + fullServName
	}
	servName := generator.CamelCase(origServName)

	g.P()
	g.P("// Client API for ", servName, " service")
	g.P()

	// Client interface.
	g.gen.PrintComments(path)
	g.P("type ", servName, "Client interface {")
	g.In()
	for i, method := range service.Method {
		g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
		g.P(g.generateClientSignature(servName, method))
	}
	g.Out()
	g.P("}")
	g.P()

	// Client structure.
	g.P("type ", unexport(servName), "Client struct {")
	g.In()
	g.P("client *", grpcPkg, ".Client")
	g.Out()
	g.P("}")
	g.P()

	// NewClient factory.
	g.P("// New", servName, "Client creates a new gRPC-Web client.")
	g.P("func New", servName, "Client (hostname string, opts ...grpcweb.DialOption) ", servName, "Client {")
	g.In()
	g.P("return &", unexport(servName), "Client{")
	g.In()
	g.P("client: ", grpcPkg, `.NewClient(hostname, "`, fullServName, `", opts...),`)
	g.Out()
	g.P("}")
	g.Out()
	g.P("}")
	g.P()

	serviceDescVar := "_" + servName + "_serviceDesc"
	// Client method implementations.
	for _, method := range service.Method {
		g.generateClientMethod(servName, fullServName, serviceDescVar, method)
	}
}

// generateClientSignature returns the client-side signature for a method.
func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
	origMethName := method.GetName()
	methName := generator.CamelCase(origMethName)
	if reservedClientName[methName] {
		methName += "_"
	}
	reqArg := ", in *" + g.typeName(method.GetInputType())
	if method.GetClientStreaming() {
		reqArg = ""
	}
	respName := "*" + g.typeName(method.GetOutputType())
	if method.GetServerStreaming() || method.GetClientStreaming() {
		respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
	}
	return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
}

func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto) {
	methName := generator.CamelCase(method.GetName())
	outType := g.typeName(method.GetOutputType())
	inType := g.typeName(method.GetInputType())
	streamType := unexport(servName) + methName + "Client"

	g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
	g.In()
	switch {
	case !method.GetServerStreaming() && !method.GetClientStreaming():
		// Unary
		g.P(`resp, err := c.client.RPCCall(ctx, "`, method.GetName(), `", in.Marshal(), opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
		return
	case method.GetServerStreaming() && !method.GetClientStreaming():
		// Server-side stream
		g.P(`srv, err := c.client.NewServerStream(ctx, "`, method.GetName(), `", in.Marshal(), opts...)`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{")
		g.In()
		g.P("stream: srv,")
		g.Out()
		g.P("}, nil")
		g.Out()
		g.P("}")
		g.P()
	case method.GetClientStreaming():
		// This case covers both client-side streaming and bidi streaming
		g.P(`srv, err := c.client.NewClientStream(ctx, "`, method.GetName(), `")`)
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return &", streamType, "{stream: srv}, nil")
		g.Out()
		g.P("}")
		g.P()
	}

	// Stream auxiliary types and methods.
	g.P("type ", servName, "_", methName, "Client interface {")
	g.In()
	if method.GetClientStreaming() {
		g.P("Send(*", inType, ") error")
	}
	if method.GetServerStreaming() {
		g.P("Recv() (*", outType, ", error)")
	}
	if method.GetClientStreaming() && !method.GetServerStreaming() {
		g.P("CloseAndRecv() (*", outType, ", error)")
	} else if method.GetServerStreaming() && !method.GetClientStreaming() {
		g.P("CloseSend() error")
	}
	g.P("Context() context.Context")
	g.Out()
	g.P("}")
	g.P()

	g.P("type ", streamType, " struct {")
	g.In()
	if method.GetServerStreaming() && !method.GetClientStreaming() {
		// Server-side streaming only
		g.P("stream ", grpcPkg, ".ServerStream")
	} else {
		// Client side and bidi streams
		g.P("stream ", grpcPkg, ".ClientStream")
	}
	g.Out()
	g.P("}")
	g.P()

	if method.GetClientStreaming() {
		g.P("func (x *", streamType, ") Send(req *", inType, ") error {")
		g.In()
		g.P("return x.stream.SendMsg(req.Marshal())")
		g.Out()
		g.P("}")
		g.P()
	}
	if method.GetServerStreaming() {
		g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.Recv()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	}
	if method.GetClientStreaming() && !method.GetServerStreaming() {
		g.P("func (x *", streamType, ") CloseAndRecv() (*", outType, ", error) {")
		g.In()
		g.P("resp, err := x.stream.CloseAndRecv()")
		g.P("if err != nil {")
		g.In()
		g.P("return nil, err")
		g.Out()
		g.P("}")
		g.P()
		g.P("return new(", outType, ").Unmarshal(resp)")
		g.Out()
		g.P("}")
		g.P()
	} else if method.GetServerStreaming() && !method.GetClientStreaming() {
		g.P("func (x *", streamType, ") CloseSend() error {")
		g.In()
		g.P("return x.stream.CloseSend()")
		g.Out()
		g.P("}")
		g.P()
	}
}
<|endoftext|>


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

package common

import (
	"github.com/juju/errors"
	"github.com/juju/utils/series"

	apimetadata "github.com/juju/juju/api/imagemetadata"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/imagemetadata"
	"github.com/juju/juju/environs/simplestreams"
	"github.com/juju/juju/juju"
)

// FindImageMetadata looks for image metadata in state.
// If none are found, we fall back on original image search in simple streams.
func FindImageMetadata(env environs.Environ, imageConstraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	stateMetadata, stateInfo, err := imageMetadataFromState(env, imageConstraint, signedOnly)
	if err != nil && !errors.IsNotFound(err) {
		// look into simple stream if for some reason can't get from state server
		logger.Infof("could not get image metadata from state server: %v", err)
	}

	logger.Debugf("\n\n GOT FROM STATE %d METADATA \n\n", len(stateMetadata))
	// No need to look in data sources if found in state?
	if len(stateMetadata) != 0 {
		return stateMetadata, stateInfo, nil
	}

	// If none are found, fall back to original simple stream impl.
	// Currently, an image metadata worker picks up this metadata periodically (daily),
	// and stores it in state. So potentially, this collection could be different
	// to what is in state.
	dsMetadata, dsInfo, err := imageMetadataFromDataSources(env, imageConstraint, signedOnly)
	if err != nil {
		if !errors.IsNotFound(err) {
			return nil, nil, errors.Trace(err)
		}
	}

	logger.Debugf("\n\n GOT FROM DATA SOURCES %d METADATA \n\n", len(dsMetadata))
	// If still none found, complain
	if len(dsMetadata) == 0 {
		return nil, nil, errors.NotFoundf("image metadata for series %v, arch %v", imageConstraint.Series, imageConstraint.Arches)
	}

	return dsMetadata, dsInfo, nil
}

var metadataAPI = func(env environs.Environ) (*apimetadata.Client, error) {
	api, err := juju.NewAPIFromName(env.Config().Name())
	if err != nil {
		return nil, errors.Annotate(err, "could not connect to api")
	}
	return apimetadata.NewClient(api), nil
}

func imageMetadataFromState(env environs.Environ, ic *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	metadataAPI, err := metadataAPI(env)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	stored, err := metadataAPI.List(ic.Stream, ic.Region, ic.Series, ic.Arches, "", "")
	if err != nil {
		return nil, nil, errors.Annotate(err, "could not list image metadata from state server")
	}

	// Convert to common format.
	images := make([]*imagemetadata.ImageMetadata, len(stored))
	for i, one := range stored {
		m := &imagemetadata.ImageMetadata{
			Storage:    one.RootStorageType,
			Id:         one.ImageId,
			VirtType:   one.VirtType,
			Arch:       one.Arch,
			RegionName: one.Region,
			Stream:     one.Stream,
		}
		m.Version, _ = series.SeriesVersion(one.Series)
		images[i] = m
	}

	info := &simplestreams.ResolveInfo{}
	info.Source = "state server"
	// This is currently ignored for image metadata that is stored in state
	// since when stored, both signed and unsigned metadata are written,
	// but whether it was signed or not is not.
	info.Signed = signedOnly
	return images, info, nil
}

// imageMetadataFromDataSources finds image metadata using
// existing data sources.
func imageMetadataFromDataSources(env environs.Environ, constraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	sources, err := environs.ImageMetadataSources(env)
	if err != nil {
		return nil, nil, err
	}
	return imagemetadata.Fetch(sources, constraint, signedOnly)
}
<commit_msg>My blutanly obvious log statements.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common

import (
	"github.com/juju/errors"
	"github.com/juju/utils/series"

	apimetadata "github.com/juju/juju/api/imagemetadata"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/imagemetadata"
	"github.com/juju/juju/environs/simplestreams"
	"github.com/juju/juju/juju"
)

// FindImageMetadata looks for image metadata in state.
// If none are found, we fall back on original image search in simple streams.
func FindImageMetadata(env environs.Environ, imageConstraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	stateMetadata, stateInfo, err := imageMetadataFromState(env, imageConstraint, signedOnly)
	if err != nil && !errors.IsNotFound(err) {
		// look into simple stream if for some reason can't get from state server
		logger.Infof("could not get image metadata from state server: %v", err)
	}

	logger.Warningf("\n\n GOT FROM STATE %d METADATA \n\n", len(stateMetadata))
	// No need to look in data sources if found in state?
	if len(stateMetadata) != 0 {
		return stateMetadata, stateInfo, nil
	}

	// If none are found, fall back to original simple stream impl.
	// Currently, an image metadata worker picks up this metadata periodically (daily),
	// and stores it in state. So potentially, this collection could be different
	// to what is in state.
	dsMetadata, dsInfo, err := imageMetadataFromDataSources(env, imageConstraint, signedOnly)
	if err != nil {
		if !errors.IsNotFound(err) {
			return nil, nil, errors.Trace(err)
		}
	}

	logger.Warningf("\n\n GOT FROM DATA SOURCES %d METADATA \n\n", len(dsMetadata))
	// If still none found, complain
	if len(dsMetadata) == 0 {
		return nil, nil, errors.NotFoundf("image metadata for series %v, arch %v", imageConstraint.Series, imageConstraint.Arches)
	}

	return dsMetadata, dsInfo, nil
}

var metadataAPI = func(env environs.Environ) (*apimetadata.Client, error) {
	api, err := juju.NewAPIFromName(env.Config().Name())
	if err != nil {
		return nil, errors.Annotate(err, "could not connect to api")
	}
	return apimetadata.NewClient(api), nil
}

func imageMetadataFromState(env environs.Environ, ic *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	metadataAPI, err := metadataAPI(env)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	stored, err := metadataAPI.List(ic.Stream, ic.Region, ic.Series, ic.Arches, "", "")
	if err != nil {
		return nil, nil, errors.Annotate(err, "could not list image metadata from state server")
	}

	// Convert to common format.
	images := make([]*imagemetadata.ImageMetadata, len(stored))
	for i, one := range stored {
		m := &imagemetadata.ImageMetadata{
			Storage:    one.RootStorageType,
			Id:         one.ImageId,
			VirtType:   one.VirtType,
			Arch:       one.Arch,
			RegionName: one.Region,
			Stream:     one.Stream,
		}
		m.Version, _ = series.SeriesVersion(one.Series)
		images[i] = m
	}

	info := &simplestreams.ResolveInfo{}
	info.Source = "state server"
	// This is currently ignored for image metadata that is stored in state
	// since when stored, both signed and unsigned metadata are written,
	// but whether it was signed or not is not.
	info.Signed = signedOnly
	return images, info, nil
}

// imageMetadataFromDataSources finds image metadata using
// existing data sources.
func imageMetadataFromDataSources(env environs.Environ, constraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	sources, err := environs.ImageMetadataSources(env)
	if err != nil {
		return nil, nil, err
	}
	return imagemetadata.Fetch(sources, constraint, signedOnly)
}
<|endoftext|>


<commit_before>package mpb_test

import (
	"bytes"
	"context"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"

	"github.com/vbauerster/mpb/v7"
	"github.com/vbauerster/mpb/v7/decor"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func TestBarCount(t *testing.T) {
	p := mpb.New(mpb.WithOutput(ioutil.Discard))

	check := make(chan struct{})
	b := p.AddBar(100)
	go func() {
		for i := 0; i < 100; i++ {
			if i == 10 {
				close(check)
			}
			b.Increment()
			time.Sleep((time.Duration(rand.Intn(10)+1) * (10 * time.Millisecond)) / 2)
		}
	}()

	<-check
	count := p.BarCount()
	if count != 1 {
		t.Errorf("BarCount want: %q, got: %q\n", 1, count)
	}

	b.Abort(false)
	p.Wait()
}

func TestBarAbort(t *testing.T) {
	n := 2
	p := mpb.New(mpb.WithOutput(ioutil.Discard))
	bars := make([]*mpb.Bar, n)
	for i := 0; i < n; i++ {
		b := p.AddBar(100)
		switch i {
		case n - 1:
			var abortCalledTimes int
			for j := 0; !b.Completed(); j++ {
				if j >= 33 {
					b.Abort(true)
					abortCalledTimes++
				} else {
					b.Increment()
					time.Sleep((time.Duration(rand.Intn(10)+1) * (10 * time.Millisecond)) / 2)
				}
			}
			if abortCalledTimes != 1 {
				t.Errorf("Expected abortCalledTimes: %d, got: %d\n", 1, abortCalledTimes)
			}
			count := p.BarCount()
			if count != 1 {
				t.Errorf("BarCount want: %d, got: %d\n", 1, count)
			}
		default:
			go func() {
				for !b.Completed() {
					b.Increment()
					time.Sleep((time.Duration(rand.Intn(10)+1) * (10 * time.Millisecond)) / 2)
				}
			}()
		}
		bars[i] = b
	}

	bars[0].Abort(false)
	p.Wait()
}

func TestWithContext(t *testing.T) {
	shutdown := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	p := mpb.NewWithContext(ctx, mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	start := make(chan struct{})
	done := make(chan struct{})
	fail := make(chan struct{})
	bar := p.AddBar(0) // never complete bar
	go func() {
		close(start)
		for !bar.Completed() {
			bar.Increment()
			time.Sleep(randomDuration(100 * time.Millisecond))
		}
		close(done)
	}()

	go func() {
		select {
		case <-done:
			p.Wait()
		case <-time.After(150 * time.Millisecond):
			close(fail)
		}
	}()

	<-start
	cancel()
	select {
	case <-shutdown:
	case <-fail:
		t.Error("Progress didn't shutdown")
	}
}

// MaxWidthDistributor shouldn't stuck in the middle while removing or aborting a bar
func TestMaxWidthDistributor(t *testing.T) {

	makeWrapper := func(f func([]chan int), start, end chan struct{}) func([]chan int) {
		return func(column []chan int) {
			start <- struct{}{}
			f(column)
			<-end
		}
	}

	ready := make(chan struct{})
	start := make(chan struct{})
	end := make(chan struct{})
	mpb.MaxWidthDistributor = makeWrapper(mpb.MaxWidthDistributor, start, end)

	total := 80
	numBars := 6
	p := mpb.New(mpb.WithOutput(ioutil.Discard))
	for i := 0; i < numBars; i++ {
		bar := p.AddBar(int64(total),
			mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0),
			mpb.PrependDecorators(
				decor.EwmaETA(decor.ET_STYLE_GO, 60, decor.WCSyncSpace),
			),
		)
		go func() {
			<-ready
			for i := 0; i < total; i++ {
				start := time.Now()
				if id := bar.ID(); id > 1 && i >= 42 {
					if id&1 == 1 {
						bar.Abort(true)
					} else {
						bar.Abort(false)
					}
				}
				time.Sleep((time.Duration(rand.Intn(10)+1) * (50 * time.Millisecond)) / 2)
				bar.IncrInt64(rand.Int63n(5) + 1)
				bar.DecoratorEwmaUpdate(time.Since(start))
			}
		}()
	}

	go func() {
		<-ready
		p.Wait()
		close(start)
	}()

	res := t.Run("maxWidthDistributor", func(t *testing.T) {
		close(ready)
		for v := range start {
			timer := time.NewTimer(100 * time.Millisecond)
			select {
			case end <- v:
				timer.Stop()
			case <-timer.C:
				t.FailNow()
			}
		}
	})

	if !res {
		t.Error("maxWidthDistributor stuck in the middle")
	}
}

func getLastLine(bb []byte) []byte {
	split := bytes.Split(bb, []byte("\n"))
	return split[len(split)-2]
}

func randomDuration(max time.Duration) time.Duration {
	return time.Duration(rand.Intn(10)+1) * max / 10
}
<commit_msg>refactoring progress test<commit_after>package mpb_test

import (
	"bytes"
	"context"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"

	"github.com/vbauerster/mpb/v7"
	"github.com/vbauerster/mpb/v7/decor"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func TestBarCount(t *testing.T) {
	shutdown := make(chan struct{})
	p := mpb.New(mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	check := make(chan struct{})
	b := p.AddBar(100)
	go func() {
		for i := 0; i < 100; i++ {
			if i == 10 {
				close(check)
			}
			b.Increment()
			time.Sleep(randomDuration(100 * time.Millisecond))
		}
	}()

	<-check
	if count := p.BarCount(); count != 1 {
		t.Errorf("BarCount want: %q, got: %q\n", 1, count)
	}

	b.Abort(false)
	go p.Wait()
	select {
	case <-shutdown:
	case <-time.After(150 * time.Millisecond):
		t.Error("Progress didn't shutdown")
	}
}

func TestBarAbort(t *testing.T) {
	shutdown := make(chan struct{})
	p := mpb.New(mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))
	n := 2
	bars := make([]*mpb.Bar, n)
	for i := 0; i < n; i++ {
		b := p.AddBar(100)
		switch i {
		case n - 1:
			var abortCalledTimes int
			for j := 0; !b.Aborted(); j++ {
				if j >= 10 {
					b.Abort(true)
					abortCalledTimes++
				} else {
					b.Increment()
				}
			}
			if abortCalledTimes != 1 {
				t.Errorf("Expected abortCalledTimes: %d, got: %d\n", 1, abortCalledTimes)
			}
			count := p.BarCount()
			if count != 1 {
				t.Errorf("BarCount want: %d, got: %d\n", 1, count)
			}
		default:
			go func() {
				for !b.Completed() {
					b.Increment()
					time.Sleep(randomDuration(100 * time.Millisecond))
				}
			}()
		}
		bars[i] = b
	}

	bars[0].Abort(false)
	go p.Wait()
	select {
	case <-shutdown:
	case <-time.After(150 * time.Millisecond):
		t.Error("Progress didn't shutdown")
	}
}

func TestWithContext(t *testing.T) {
	shutdown := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	p := mpb.NewWithContext(ctx, mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	done := make(chan struct{})
	fail := make(chan struct{})
	bar := p.AddBar(0) // never complete bar
	go func() {
		for !bar.Aborted() {
			time.Sleep(randomDuration(100 * time.Millisecond))
			cancel()
		}
		close(done)
	}()

	go func() {
		select {
		case <-done:
			p.Wait()
		case <-time.After(150 * time.Millisecond):
			close(fail)
		}
	}()

	select {
	case <-shutdown:
	case <-fail:
		t.Error("Progress didn't shutdown")
	}
}

// MaxWidthDistributor shouldn't stuck in the middle while removing or aborting a bar
func TestMaxWidthDistributor(t *testing.T) {

	makeWrapper := func(f func([]chan int), start, end chan struct{}) func([]chan int) {
		return func(column []chan int) {
			start <- struct{}{}
			f(column)
			<-end
		}
	}

	ready := make(chan struct{})
	start := make(chan struct{})
	end := make(chan struct{})
	mpb.MaxWidthDistributor = makeWrapper(mpb.MaxWidthDistributor, start, end)

	total := 100
	numBars := 6
	p := mpb.New(mpb.WithOutput(ioutil.Discard))
	for i := 0; i < numBars; i++ {
		bar := p.AddBar(int64(total),
			mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0),
			mpb.PrependDecorators(decor.EwmaETA(decor.ET_STYLE_GO, 60, decor.WCSyncSpace)),
		)
		go func() {
			<-ready
			for i := 0; i < total; i++ {
				start := time.Now()
				if id := bar.ID(); id > 1 && i >= 32 {
					if id&1 == 1 {
						bar.Abort(true)
					} else {
						bar.Abort(false)
					}
				}
				time.Sleep(randomDuration(100 * time.Millisecond))
				bar.IncrInt64(rand.Int63n(5) + 1)
				bar.DecoratorEwmaUpdate(time.Since(start))
			}
		}()
	}

	go func() {
		<-ready
		p.Wait()
		close(start)
	}()

	res := t.Run("maxWidthDistributor", func(t *testing.T) {
		close(ready)
		for v := range start {
			timer := time.NewTimer(100 * time.Millisecond)
			select {
			case end <- v:
				timer.Stop()
			case <-timer.C:
				t.FailNow()
			}
		}
	})

	if !res {
		t.Error("maxWidthDistributor stuck in the middle")
	}
}

func getLastLine(bb []byte) []byte {
	split := bytes.Split(bb, []byte("\n"))
	return split[len(split)-2]
}

func randomDuration(max time.Duration) time.Duration {
	return time.Duration(rand.Intn(10)+1) * max / 10
}
<|endoftext|>


<commit_before>package metrics

import (
	"net/url"

	"github.com/cerana/cerana/acomm"
	"github.com/cerana/cerana/provider"
	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/host"
	"github.com/shirou/gopsutil/load"
	"github.com/shirou/gopsutil/mem"
	"github.com/shirou/gopsutil/net"
)

// MockMetrics is a mock Metrics provider.
type MockMetrics struct {
	Data *MockMetricsData
}

// MockMetricsData is the in-memory data structure for the MockMetrics.
type MockMetricsData struct {
	CPU      *CPUResult
	Disk     *DiskResult
	Hardware interface{}
	Host     *host.InfoStat
	Memory   *MemoryResult
	Network  *NetworkResult
}

// NewMockMetrics creates a new MockMetrics and populates default data.
func NewMockMetrics() *MockMetrics {
	return &MockMetrics{
		Data: &MockMetricsData{
			CPU: &CPUResult{
				Info: []cpu.InfoStat{
					{
						CPU:        0,
						VendorID:   "",
						Family:     "6",
						Model:      "69",
						Stepping:   1,
						PhysicalID: "0",
						CoreID:     "0",
						Cores:      1,
						ModelName:  "Intel(R) Core(TM) i5-4278U CPU @ 2.60GHz",
						Mhz:        2599.79,
						CacheSize:  3072,
						Flags:      []string{},
					}, {
						CPU:        1,
						VendorID:   "",
						Family:     "6",
						Model:      "69",
						Stepping:   1,
						PhysicalID: "0",
						CoreID:     "1",
						Cores:      1,
						ModelName:  "Intel(R) Core(TM) i5-4278U CPU @ 2.60GHz",
						Mhz:        2599.79,
						CacheSize:  3072,
						Flags:      []string{},
					},
				},
				Load: load.AvgStat{
					Load1:  0.72,
					Load5:  0.22,
					Load15: 0.11,
				},
				Times: []cpu.TimesStat{
					{
						CPU:       "cpu0",
						User:      4973.34,
						System:    3579.08,
						Idle:      824805.64,
						Nice:      0.93,
						Iowait:    921.25,
						Irq:       0,
						Softirq:   674.86,
						Steal:     0,
						Guest:     0,
						GuestNice: 0,
						Stolen:    0,
					}, {
						CPU:       "cpu1",
						User:      5104.66,
						System:    3458.4,
						Idle:      837419.2,
						Nice:      1.29,
						Iowait:    869.55,
						Irq:       0,
						Softirq:   161.14,
						Steal:     0,
						Guest:     0,
						GuestNice: 0,
						Stolen:    0,
					},
				},
			},
			Disk: &DiskResult{
				IO:         make(map[string]disk.IOCountersStat),
				Partitions: []disk.PartitionStat{},
				Usage: []*disk.UsageStat{
					{
						Path:        "/",
						Fstype:      "zfs",
						Total:       100000,
						Free:        50000,
						Used:        50000,
						UsedPercent: 0.5,
					},
				},
			},
			Hardware: struct{}{},
			Host: &host.InfoStat{
				Hostname:       "mockMetrics",
				Uptime:         uint64(123456),
				BootTime:       uint64(123456),
				Procs:          5,
				OS:             "linux",
				Platform:       "cerana",
				PlatformFamily: "nix",
			},
			Memory: &MemoryResult{
				Swap: &mem.SwapMemoryStat{
					Total:       100000,
					Used:        50000,
					Free:        50000,
					UsedPercent: 0.5,
				},
				Virtual: &mem.VirtualMemoryStat{
					Total:       100000,
					Available:   100000,
					Used:        50000,
					Free:        50000,
					UsedPercent: 0.5,
				},
			},
			Network: &NetworkResult{
				IO: []net.IOCountersStat{
					{
						Name:        "fooIface",
						BytesSent:   12345,
						BytesRecv:   54321,
						PacketsSent: 12345,
						PacketsRecv: 54321,
					},
				},
				Interfaces: []net.InterfaceStat{
					{
						Name:         "fooIface",
						MTU:          1500,
						HardwareAddr: "en0",
						Addrs: []net.InterfaceAddr{
							{"123.123.123.123"},
						},
					},
				},
			},
		},
	}
}

// RegisterTasks registes all MockMetric task handlers.
func (m *MockMetrics) RegisterTasks(server *provider.Server) {
	server.RegisterTask("metrics-cpu", m.CPU)
	server.RegisterTask("metrics-disk", m.Disk)
	server.RegisterTask("metrics-host", m.Host)
	server.RegisterTask("metrics-memory", m.Memory)
	server.RegisterTask("metrics-network", m.Network)
}

// CPU returns mock CPU information.
func (m *MockMetrics) CPU(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.CPU, nil, nil
}

// Disk returns mock disk information.
func (m *MockMetrics) Disk(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Disk, nil, nil
}

// Hardware returns mock hardware information.
func (m *MockMetrics) Hardware(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Hardware, nil, nil
}

// Host returns mock host information.
func (m *MockMetrics) Host(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Host, nil, nil
}

// Memory returns mock memory information.
func (m *MockMetrics) Memory(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Memory, nil, nil
}

// Network returns mock network information.
func (m *MockMetrics) Network(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Network, nil, nil
}
<commit_msg>providers/metric: fix "composite literal uses unkeyed fields" go vet error<commit_after>package metrics

import (
	"net/url"

	"github.com/cerana/cerana/acomm"
	"github.com/cerana/cerana/provider"
	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/host"
	"github.com/shirou/gopsutil/load"
	"github.com/shirou/gopsutil/mem"
	"github.com/shirou/gopsutil/net"
)

// MockMetrics is a mock Metrics provider.
type MockMetrics struct {
	Data *MockMetricsData
}

// MockMetricsData is the in-memory data structure for the MockMetrics.
type MockMetricsData struct {
	CPU      *CPUResult
	Disk     *DiskResult
	Hardware interface{}
	Host     *host.InfoStat
	Memory   *MemoryResult
	Network  *NetworkResult
}

// NewMockMetrics creates a new MockMetrics and populates default data.
func NewMockMetrics() *MockMetrics {
	return &MockMetrics{
		Data: &MockMetricsData{
			CPU: &CPUResult{
				Info: []cpu.InfoStat{
					{
						CPU:        0,
						VendorID:   "",
						Family:     "6",
						Model:      "69",
						Stepping:   1,
						PhysicalID: "0",
						CoreID:     "0",
						Cores:      1,
						ModelName:  "Intel(R) Core(TM) i5-4278U CPU @ 2.60GHz",
						Mhz:        2599.79,
						CacheSize:  3072,
						Flags:      []string{},
					}, {
						CPU:        1,
						VendorID:   "",
						Family:     "6",
						Model:      "69",
						Stepping:   1,
						PhysicalID: "0",
						CoreID:     "1",
						Cores:      1,
						ModelName:  "Intel(R) Core(TM) i5-4278U CPU @ 2.60GHz",
						Mhz:        2599.79,
						CacheSize:  3072,
						Flags:      []string{},
					},
				},
				Load: load.AvgStat{
					Load1:  0.72,
					Load5:  0.22,
					Load15: 0.11,
				},
				Times: []cpu.TimesStat{
					{
						CPU:       "cpu0",
						User:      4973.34,
						System:    3579.08,
						Idle:      824805.64,
						Nice:      0.93,
						Iowait:    921.25,
						Irq:       0,
						Softirq:   674.86,
						Steal:     0,
						Guest:     0,
						GuestNice: 0,
						Stolen:    0,
					}, {
						CPU:       "cpu1",
						User:      5104.66,
						System:    3458.4,
						Idle:      837419.2,
						Nice:      1.29,
						Iowait:    869.55,
						Irq:       0,
						Softirq:   161.14,
						Steal:     0,
						Guest:     0,
						GuestNice: 0,
						Stolen:    0,
					},
				},
			},
			Disk: &DiskResult{
				IO:         make(map[string]disk.IOCountersStat),
				Partitions: []disk.PartitionStat{},
				Usage: []*disk.UsageStat{
					{
						Path:        "/",
						Fstype:      "zfs",
						Total:       100000,
						Free:        50000,
						Used:        50000,
						UsedPercent: 0.5,
					},
				},
			},
			Hardware: struct{}{},
			Host: &host.InfoStat{
				Hostname:       "mockMetrics",
				Uptime:         uint64(123456),
				BootTime:       uint64(123456),
				Procs:          5,
				OS:             "linux",
				Platform:       "cerana",
				PlatformFamily: "nix",
			},
			Memory: &MemoryResult{
				Swap: &mem.SwapMemoryStat{
					Total:       100000,
					Used:        50000,
					Free:        50000,
					UsedPercent: 0.5,
				},
				Virtual: &mem.VirtualMemoryStat{
					Total:       100000,
					Available:   100000,
					Used:        50000,
					Free:        50000,
					UsedPercent: 0.5,
				},
			},
			Network: &NetworkResult{
				IO: []net.IOCountersStat{
					{
						Name:        "fooIface",
						BytesSent:   12345,
						BytesRecv:   54321,
						PacketsSent: 12345,
						PacketsRecv: 54321,
					},
				},
				Interfaces: []net.InterfaceStat{
					{
						Name:         "fooIface",
						MTU:          1500,
						HardwareAddr: "en0",
						Addrs: []net.InterfaceAddr{
							{Addr: "123.123.123.123"},
						},
					},
				},
			},
		},
	}
}

// RegisterTasks registes all MockMetric task handlers.
func (m *MockMetrics) RegisterTasks(server *provider.Server) {
	server.RegisterTask("metrics-cpu", m.CPU)
	server.RegisterTask("metrics-disk", m.Disk)
	server.RegisterTask("metrics-host", m.Host)
	server.RegisterTask("metrics-memory", m.Memory)
	server.RegisterTask("metrics-network", m.Network)
}

// CPU returns mock CPU information.
func (m *MockMetrics) CPU(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.CPU, nil, nil
}

// Disk returns mock disk information.
func (m *MockMetrics) Disk(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Disk, nil, nil
}

// Hardware returns mock hardware information.
func (m *MockMetrics) Hardware(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Hardware, nil, nil
}

// Host returns mock host information.
func (m *MockMetrics) Host(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Host, nil, nil
}

// Memory returns mock memory information.
func (m *MockMetrics) Memory(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Memory, nil, nil
}

// Network returns mock network information.
func (m *MockMetrics) Network(req *acomm.Request) (interface{}, *url.URL, error) {
	return m.Data.Network, nil, nil
}
<|endoftext|>