<commit_before>package ionic

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/ion-channel/ionic/organizations"
	"github.com/ion-channel/ionic/pagination"
	"github.com/ion-channel/ionic/requests"
)

// CreateOrganizationOptions represents all the values that can be provided for an organization
// at the time of creation
type CreateOrganizationOptions struct {
	Name string `json:"name"`
}

// CreateOrganization takes a create team options, validates the minimum info is
// present, and makes the calls to create the team. It returns the ID of the created organization
// and any errors it encounters with the API.
func (ic *IonClient) CreateOrganization(opts CreateOrganizationOptions, token string) (*organizations.Organization, error) {
	if opts.Name == "" {
		return nil, fmt.Errorf("name missing from options")
	}

	b, err := json.Marshal(opts)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	b, err = ic.Post(organizations.OrganizationsCreateEndpoint, token, nil, *buff, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create organization: %v", err.Error())
	}

	var org organizations.Organization
	err = json.Unmarshal(b, &org)
	if err != nil {
		return nil, fmt.Errorf("failed to parse organization from response: %v", err.Error())
	}

	return &org, nil
}

// GetOwnOrganizations takes a token and returns a list of organizations the user belongs to.
func (ic *IonClient) GetOwnOrganizations(token string) (*[]organizations.UserOrganizationRole, error) {
	resp, _, err := ic.Get(organizations.OrganizationsGetOwnEndpoint, token, nil, nil, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get own organizations: %v", err.Error())
	}

	var orgs []organizations.UserOrganizationRole
	err = json.Unmarshal(resp, &orgs)
	if err != nil {
		return nil, fmt.Errorf("cannot parse own organizations: %v", err.Error())
	}

	return &orgs, nil
}

// GetOrganization takes an organization id and returns the Ion Channel representation of that organization.
func (ic *IonClient) GetOrganization(id, token string) (*organizations.Organization, error) {
	b, _, err := ic.Get(fmt.Sprintf("%s/%s", organizations.OrganizationsGetEndpoint, id), token, nil, nil, pagination.Pagination{})
	if err != nil {
		return nil, fmt.Errorf("failed to get organization: %v", err.Error())
	}

	var organization organizations.Organization
	err = json.Unmarshal(b, &organization)
	if err != nil {
		return nil, fmt.Errorf("cannot parse organization: %v", err.Error())
	}

	return &organization, nil
}

// GetOrganizations takes one or more IDs and returns those organizations.
func (ic *IonClient) GetOrganizations(ids requests.ByIDs, token string) (*[]organizations.Organization, error) {
	b, err := json.Marshal(ids)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	resp, err := ic.Post(organizations.OrganizationsGetBulkEndpoint, token, nil, *buff, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get organizations: %v", err.Error())
	}

	var orgs []organizations.Organization
	err = json.Unmarshal(resp, &orgs)
	if err != nil {
		return nil, fmt.Errorf("cannot parse organizations: %v", err.Error())
	}

	return &orgs, nil
}

// UpdateOrganization takes an organization ID, and the fields to update, returns the updated organization.
func (ic *IonClient) UpdateOrganization(id string, name string, token string) (*organizations.Organization, error) {
	req := struct {
		Name string `json:"name"`
	}{Name: name}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	resp, err := ic.Put(fmt.Sprintf("%s/%s", organizations.OrganizationsUpdateEndpoint, id), token, nil, *buff, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to update organization: %v", err.Error())
	}

	var org organizations.Organization
	err = json.Unmarshal(resp, &org)
	if err != nil {
		return nil, fmt.Errorf("cannot parse organization: %v", err.Error())
	}

	return &org, nil
}

// DisableOrganization takes an organization ID and returns any errors that occurred.
func (ic *IonClient) DisableOrganization(id string, token string) error {
	_, err := ic.Delete(fmt.Sprintf("%s/%s", organizations.OrganizationsDisableEndpoint, id), token, nil, nil)
	if err != nil {
		return fmt.Errorf("failed to disable organization: %v", err.Error())
	}

	return nil
}

// AddMemberToOrganization takes an organization ID, a user ID, and a role, and returns any errors that occurred.
func (ic *IonClient) AddMemberToOrganization(organizationID string, userID string, role organizations.OrganizationRole, token string) error {
	req := struct {
		UserID string                         `json:"user_id"`
		Role   organizations.OrganizationRole `json:"role"`
	}{
		UserID: userID,
		Role:   role,
	}

	b, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	_, err = ic.Post(fmt.Sprintf("%s/%s", organizations.OrganizationsAddMemberEndpoint, organizationID), token, nil, *buff, nil)
	if err != nil {
		return fmt.Errorf("failed to add member to organization: %v", err.Error())
	}

	return nil
}
<commit_msg>Revise pointer usage in GetOwnOrganizations<commit_after>package ionic

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/ion-channel/ionic/organizations"
	"github.com/ion-channel/ionic/pagination"
	"github.com/ion-channel/ionic/requests"
)

// CreateOrganizationOptions represents all the values that can be provided for an organization
// at the time of creation
type CreateOrganizationOptions struct {
	Name string `json:"name"`
}

// CreateOrganization takes a create team options, validates the minimum info is
// present, and makes the calls to create the team. It returns the ID of the created organization
// and any errors it encounters with the API.
func (ic *IonClient) CreateOrganization(opts CreateOrganizationOptions, token string) (*organizations.Organization, error) {
	if opts.Name == "" {
		return nil, fmt.Errorf("name missing from options")
	}

	b, err := json.Marshal(opts)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	b, err = ic.Post(organizations.OrganizationsCreateEndpoint, token, nil, *buff, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create organization: %v", err.Error())
	}

	var org organizations.Organization
	err = json.Unmarshal(b, &org)
	if err != nil {
		return nil, fmt.Errorf("failed to parse organization from response: %v", err.Error())
	}

	return &org, nil
}

// GetOwnOrganizations takes a token and returns a list of organizations the user belongs to.
func (ic *IonClient) GetOwnOrganizations(token string) (*[]organizations.UserOrganizationRole, error) {
	resp, _, err := ic.Get(organizations.OrganizationsGetOwnEndpoint, token, nil, nil, pagination.Pagination{})
	if err != nil {
		return nil, fmt.Errorf("failed to get own organizations: %v", err.Error())
	}

	var orgs []organizations.UserOrganizationRole
	err = json.Unmarshal(resp, &orgs)
	if err != nil {
		return nil, fmt.Errorf("cannot parse own organizations: %v", err.Error())
	}

	return &orgs, nil
}

// GetOrganization takes an organization id and returns the Ion Channel representation of that organization.
func (ic *IonClient) GetOrganization(id, token string) (*organizations.Organization, error) {
	b, _, err := ic.Get(fmt.Sprintf("%s/%s", organizations.OrganizationsGetEndpoint, id), token, nil, nil, pagination.Pagination{})
	if err != nil {
		return nil, fmt.Errorf("failed to get organization: %v", err.Error())
	}

	var organization organizations.Organization
	err = json.Unmarshal(b, &organization)
	if err != nil {
		return nil, fmt.Errorf("cannot parse organization: %v", err.Error())
	}

	return &organization, nil
}

// GetOrganizations takes one or more IDs and returns those organizations.
func (ic *IonClient) GetOrganizations(ids requests.ByIDs, token string) (*[]organizations.Organization, error) {
	b, err := json.Marshal(ids)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	resp, err := ic.Post(organizations.OrganizationsGetBulkEndpoint, token, nil, *buff, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get organizations: %v", err.Error())
	}

	var orgs []organizations.Organization
	err = json.Unmarshal(resp, &orgs)
	if err != nil {
		return nil, fmt.Errorf("cannot parse organizations: %v", err.Error())
	}

	return &orgs, nil
}

// UpdateOrganization takes an organization ID, and the fields to update, returns the updated organization.
func (ic *IonClient) UpdateOrganization(id string, name string, token string) (*organizations.Organization, error) {
	req := struct {
		Name string `json:"name"`
	}{Name: name}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	resp, err := ic.Put(fmt.Sprintf("%s/%s", organizations.OrganizationsUpdateEndpoint, id), token, nil, *buff, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to update organization: %v", err.Error())
	}

	var org organizations.Organization
	err = json.Unmarshal(resp, &org)
	if err != nil {
		return nil, fmt.Errorf("cannot parse organization: %v", err.Error())
	}

	return &org, nil
}

// DisableOrganization takes an organization ID and returns any errors that occurred.
func (ic *IonClient) DisableOrganization(id string, token string) error {
	_, err := ic.Delete(fmt.Sprintf("%s/%s", organizations.OrganizationsDisableEndpoint, id), token, nil, nil)
	if err != nil {
		return fmt.Errorf("failed to disable organization: %v", err.Error())
	}

	return nil
}

// AddMemberToOrganization takes an organization ID, a user ID, and a role, and returns any errors that occurred.
func (ic *IonClient) AddMemberToOrganization(organizationID string, userID string, role organizations.OrganizationRole, token string) error {
	req := struct {
		UserID string                         `json:"user_id"`
		Role   organizations.OrganizationRole `json:"role"`
	}{
		UserID: userID,
		Role:   role,
	}

	b, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("failed to marshal request body: %v", err.Error())
	}

	buff := bytes.NewBuffer(b)

	_, err = ic.Post(fmt.Sprintf("%s/%s", organizations.OrganizationsAddMemberEndpoint, organizationID), token, nil, *buff, nil)
	if err != nil {
		return fmt.Errorf("failed to add member to organization: %v", err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/jpillora/backoff"
	"golang.org/x/net/context"
)

// Backoff for fetching. It starts by waiting the minimum duration after a
// failed fetch, doubling it each time (with a bitter of jitter) up to max
// duration between requests.
var defaultBackoff = backoff.Backoff{
	Min:    1 * time.Second,
	Max:    5 * time.Minute,
	Jitter: true,
	Factor: 2,
}

type Worker struct {
	query   *Query
	state   *Set
	payload []byte
	client  *http.Client
	metrics *SetMetrics
	log     *log.Logger
	backoff backoff.Backoff
}

func (w *Worker) Fetch(url string) (*Set, error) {
	var (
		t    time.Time
		err  error
		resp *http.Response
	)

	for {
		t = time.Now()

		resp, err = w.client.Post(url, "application/json", bytes.NewBuffer(w.payload))

		// No formal error, but a non-successful status code. Construct an error.
		if err == nil && resp.StatusCode != 200 {
			b, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			err = fmt.Errorf("%s: %s", resp.Status, string(b))
		}

		// No error, break to read the data.
		if err == nil {
			break
		}

		// Backoff on an error.
		w.log.Print(err)
		d := w.backoff.Duration()
		w.log.Printf("Backing off for %s", d)
		time.Sleep(d)
	}

	w.backoff.Reset()

	w.log.Printf("Fetch took %s", time.Now().Sub(t))

	var records []record

	defer resp.Body.Close()

	if err = json.NewDecoder(resp.Body).Decode(&records); err != nil {
		return nil, err
	}

	return NewSet(w.query.Identifier, records)
}

func (w *Worker) Start(cxt context.Context, url string) {
	state, err := w.Fetch(url)

	if err != nil {
		w.log.Printf("Error fetching state: %s", err)
	}

	// Set the initial state.
	w.state = state

	if state != nil {
		w.metrics.Compare(nil, state)
	}

	ticker := time.NewTicker(w.query.Interval)

	for {
		select {
		case <-cxt.Done():
			wg, _ := cxt.Value("wg").(*sync.WaitGroup)
			wg.Done()
			w.log.Printf("Stopping worker")
			return

		case <-ticker.C:
			if state, err = w.Fetch(url); err != nil {
				w.log.Printf("Error fetching state: %s", err)
				break
			}

			// Compare with new state and log metrics.
			w.metrics.Compare(w.state, state)
			w.state = state
		}
	}
}

// StateHandler returns an http.HandlerFunc that writes the state of the query.
func (w *Worker) StateHandler() http.HandlerFunc {
	return func(rp http.ResponseWriter, _ *http.Request) {
		if w.state == nil {
			rp.WriteHeader(http.StatusNoContent)
			return
		}

		if err := json.NewEncoder(rp).Encode(w.state.records); err != nil {
			rp.WriteHeader(http.StatusInternalServerError)
			rp.Write([]byte(err.Error()))
			return
		}

		rp.Header().Set("content-type", "application/json")
	}
}

// NewWorker creates a new worker for a query.
func NewWorker(q *Query) *Worker {
	// Encode the payload once for all subsequent requests.
	payload, err := json.Marshal(map[string]interface{}{
		"driver":     q.Driver,
		"connection": q.Connection,
		"sql":        q.SQL,
		"params":     q.Params,
	})

	if err != nil {
		panic(err)
	}

	return &Worker{
		query:   q,
		metrics: NewSetMetrics(q.Name),
		payload: payload,
		backoff: defaultBackoff,
		log:     log.New(os.Stderr, fmt.Sprintf("[%s] ", q.Name), log.LstdFlags),
		client: &http.Client{
			Timeout: q.Timeout,
		},
	}
}
<commit_msg>Set explicit Accept header for SQL agent service<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/jpillora/backoff"
	"golang.org/x/net/context"
)

// Backoff for fetching. It starts by waiting the minimum duration after a
// failed fetch, doubling it each time (with a bitter of jitter) up to max
// duration between requests.
var defaultBackoff = backoff.Backoff{
	Min:    1 * time.Second,
	Max:    5 * time.Minute,
	Jitter: true,
	Factor: 2,
}

type Worker struct {
	query   *Query
	state   *Set
	payload []byte
	client  *http.Client
	metrics *SetMetrics
	log     *log.Logger
	backoff backoff.Backoff
}

func (w *Worker) Fetch(url string) (*Set, error) {
	var (
		t    time.Time
		err  error
		req  *http.Request
		resp *http.Response
	)

	for {
		t = time.Now()

		req, err = http.NewRequest("POST", url, bytes.NewBuffer(w.payload))

		if err != nil {
			panic(err)
		}

		// Set the content-type of the request body and accept LD-JSON.
		req.Header.Set("content-type", "application/json")
		req.Header.Set("accept", "application/json")

		resp, err = w.client.Do(req)

		// No formal error, but a non-successful status code. Construct an error.
		if err == nil && resp.StatusCode != 200 {
			b, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			err = fmt.Errorf("%s: %s", resp.Status, string(b))
		}

		// No error, break to read the data.
		if err == nil {
			break
		}

		// Backoff on an error.
		w.log.Print(err)
		d := w.backoff.Duration()
		w.log.Printf("Backing off for %s", d)
		time.Sleep(d)
	}

	w.backoff.Reset()

	w.log.Printf("Fetch took %s", time.Now().Sub(t))

	var records []record

	defer resp.Body.Close()

	if err = json.NewDecoder(resp.Body).Decode(&records); err != nil {
		return nil, err
	}

	return NewSet(w.query.Identifier, records)
}

func (w *Worker) Start(cxt context.Context, url string) {
	state, err := w.Fetch(url)

	if err != nil {
		w.log.Printf("Error fetching state: %s", err)
	}

	// Set the initial state.
	w.state = state

	if state != nil {
		w.metrics.Compare(nil, state)
	}

	ticker := time.NewTicker(w.query.Interval)

	for {
		select {
		case <-cxt.Done():
			wg, _ := cxt.Value("wg").(*sync.WaitGroup)
			wg.Done()
			w.log.Printf("Stopping worker")
			return

		case <-ticker.C:
			if state, err = w.Fetch(url); err != nil {
				w.log.Printf("Error fetching state: %s", err)
				break
			}

			// Compare with new state and log metrics.
			w.metrics.Compare(w.state, state)
			w.state = state
		}
	}
}

// StateHandler returns an http.HandlerFunc that writes the state of the query.
func (w *Worker) StateHandler() http.HandlerFunc {
	return func(rp http.ResponseWriter, _ *http.Request) {
		if w.state == nil {
			rp.WriteHeader(http.StatusNoContent)
			return
		}

		if err := json.NewEncoder(rp).Encode(w.state.records); err != nil {
			rp.WriteHeader(http.StatusInternalServerError)
			rp.Write([]byte(err.Error()))
			return
		}

		rp.Header().Set("content-type", "application/json")
	}
}

// NewWorker creates a new worker for a query.
func NewWorker(q *Query) *Worker {
	// Encode the payload once for all subsequent requests.
	payload, err := json.Marshal(map[string]interface{}{
		"driver":     q.Driver,
		"connection": q.Connection,
		"sql":        q.SQL,
		"params":     q.Params,
	})

	if err != nil {
		panic(err)
	}

	return &Worker{
		query:   q,
		metrics: NewSetMetrics(q.Name),
		payload: payload,
		backoff: defaultBackoff,
		log:     log.New(os.Stderr, fmt.Sprintf("[%s] ", q.Name), log.LstdFlags),
		client: &http.Client{
			Timeout: q.Timeout,
		},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"runtime/pprof"
	"time"

	"github.com/dgraph-io/badger/badger"
	"github.com/dgraph-io/badger/value"

	"github.com/dgraph-io/dgraph/store"
	"github.com/dgraph-io/dgraph/x"
)

var (
	flagBench        = flag.String("bench", "", "Run which benchmark?")
	flagDB           = flag.String("db", "", "Which DB: rocksdb, badger")
	flagValueSize    = flag.Int("value_size", 100, "Size of each value.")
	flagBatchSize    = flag.Int("batch_size", 1, "Size of writebatch.")
	flagNumWrites    = flag.Int("writes", 1000000, "Number of key-value pairs to write.")
	flagNumReads     = flag.Int("reads", 1000000, "Number of key-value pairs to read.")
	flagRandSize     = flag.Int("rand_size", 1000000, "Size of rng buffer.")
	flagCpuProfile   = flag.String("cpu_profile", "", "Write cpu profile to file.")
	flagVerbose      = flag.Bool("verbose", false, "Verbose.")
	flagDir          = flag.String("dir", "/tmp/badger_bench", "Where data is temporarily stored.")
	flagDirLowLevels = flag.String("dir_low_levels", "/tmp/badger_bench",
		"Where data is stored for lower levels for Badger.")

	rdbStore *store.Store
	rng      randomGenerator
)

type randomGenerator struct {
	data []byte
	idx  int
}

func (s *randomGenerator) Init() {
	if *flagRandSize <= 0 {
		// Will not precompute the randomness.
		return
	}
	s.data = make([]byte, *flagRandSize)
	n, err := rand.Read(s.data)
	x.Check(err)
	x.AssertTrue(n == *flagRandSize)
	s.idx = 0
}

// Bytes generates len(out) random bytes and writes to out.
func (s *randomGenerator) Bytes(out []byte) {
	if *flagRandSize == 0 {
		n, err := rand.Read(out)
		x.AssertTrue(n == len(out))
		x.Check(err)
		return
	}
	size := len(out)
	if s.idx+size > len(s.data) {
		s.idx = 0
	}
	x.AssertTrue(size == copy(out, s.data[s.idx:s.idx+size]))
	s.idx += size
}

func (s *randomGenerator) Int() int {
	var buf [4]byte
	s.Bytes(buf[:])
	return int(binary.LittleEndian.Uint32(buf[:]))
}

func main() {
	x.Init()
	x.AssertTrue(len(*flagBench) > 0)
	x.AssertTrue(*flagValueSize > 0)
	rng.Init()

	if *flagCpuProfile != "" {
		f, err := os.Create(*flagCpuProfile)
		if err != nil {
			x.Fatalf("Profiler error: %v", err)
		}
		pprof.StartCPUProfile(f)
	}

	var database Database
	switch *flagDB {
	case "badger":
		database = new(BadgerAdapter)
	case "rocksdb":
		database = new(RocksDBAdapter)
	default:
		x.Fatalf("Database invalid: %v", *flagDB)
	}
	database.Init()
	defer database.Close()

	x.AssertTrue(*flagDB == "rocksdb" || *flagDB == "badger")
	switch *flagBench {
	case "writerandom":
		WriteRandom(database)
	case "batchwriterandom":
		BatchWriteRandom(database)
	case "readrandom":
		ReadRandom(database)
	default:
		x.Fatalf("Unknown benchmark: %v", *flagBench)
	}
	if *flagCpuProfile != "" {
		pprof.StopCPUProfile()
	}
}

func report(d time.Duration, n int) string {
	secs := d.Seconds()
	throughput := float64(n*(16+*flagValueSize)) / ((1 << 20) * secs)
	return fmt.Sprintf("%.2fs, %.2fMb/s", secs, throughput)
}

type Database interface {
	Init()
	Close()
	Put(key, val []byte)
	BatchPut(key, val [][]byte)
	Get(key []byte)
}

type RocksDBAdapter struct {
	rdb *store.Store
	dir string
}

func (s *RocksDBAdapter) Init() {
	var err error
	s.dir, err = ioutil.TempDir(*flagDir, "storetest_")
	x.Check(err)
	s.rdb, err = store.NewSyncStore(s.dir)
	x.Check(err)
}

func (s *RocksDBAdapter) Close() {
	//	s.rdb.Close()
}

func (s *RocksDBAdapter) Put(key, val []byte) {
	s.rdb.SetOne(key, val)
}

func (s *RocksDBAdapter) BatchPut(key, val [][]byte) {
	wb := s.rdb.NewWriteBatch()
	x.AssertTrue(len(key) == len(val))
	for i := 0; i < len(key); i++ {
		wb.Put(key[i], val[i])
	}
	x.Check(s.rdb.WriteBatch(wb))
}

func (s *RocksDBAdapter) Get(key []byte) {
	_, err := s.rdb.Get(key)
	x.Check(err)
}

type BadgerAdapter struct {
	db  *badger.DB
	dir string
}

func (s *BadgerAdapter) Init() {
	opt := badger.DBOptions{
		WriteBufferSize:         1 << 20, // Size of each memtable.
		NumLevelZeroTables:      5,
		NumLevelZeroTablesStall: 6,
		LevelOneSize:            5 << 20,
		MaxLevels:               7,
		NumCompactWorkers:       6,
		MaxTableSize:            2 << 20,
		LevelSizeMultiplier:     5,
		Verbose:                 *flagVerbose,
		Dir:                     *flagDir,
		DirLowLevels:            *flagDirLowLevels,
	}
	fmt.Printf("Dir: %s\nDirForLowLevels: %s\n", *flagDir, *flagDirLowLevels)
	s.db = badger.NewDB(opt)
}

func (s *BadgerAdapter) Close() {
}

func (s *BadgerAdapter) Put(key, val []byte) {
	s.db.Put(key, val)
}

func (s *BadgerAdapter) BatchPut(key, val [][]byte) {
	var entries []value.Entry
	x.AssertTrue(len(key) == len(val))
	for i := 0; i < len(key); i++ {
		entries = append(entries, value.Entry{
			Key:   key[i],
			Value: val[i],
		})
	}
	x.Check(s.db.Write(entries))
}

func (s *BadgerAdapter) Get(key []byte) {
	s.db.Get(key)
}

// No batching.
func WriteRandom(database Database) {
	fmt.Println("WriteRandom test")
	timeStart := time.Now()
	timeLog := timeStart
	timeLogI := 0
	// If you use b.N, you might add too few samples and be working only in memory.
	// We need to fix a large number of pairs. This is what LevelDB benchmark does as well.
	for i := 0; i < *flagNumWrites; i++ {
		key := []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumWrites))
		val := make([]byte, *flagValueSize)
		rng.Bytes(val)
		database.Put(key, val)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent: %s\n", float64(i)*100.0/float64(*flagNumWrites),
				report(timeElapsed, i-timeLogI))
			timeLog = time.Now()
			timeLogI = i
		}
	}
	x.Printf("Overall: %s\n", report(time.Since(timeStart), *flagNumWrites))
}

// With batching.
func BatchWriteRandom(database Database) {
	fmt.Println("BatchWriteRandom test")
	timeStart := time.Now()
	timeLog := timeStart
	timeLogI := 0
	keys := make([][]byte, *flagBatchSize)
	vals := make([][]byte, *flagBatchSize)

	for i := 0; i < *flagNumWrites; i++ {
		for j := 0; j < *flagBatchSize; j++ {
			keys[j] = []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumWrites))
			vals[j] = make([]byte, *flagValueSize)
			rng.Bytes(vals[j])
		}
		database.BatchPut(keys, vals)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent: %s\n", float64(i)*100.0/float64(*flagNumWrites),
				report(timeElapsed, (i-timeLogI)*(*flagBatchSize)))
			timeLog = time.Now()
			timeLogI = i
		}
	}
	x.Printf("Overall: %s\n", report(time.Since(timeStart), *flagNumWrites*(*flagBatchSize)))
}

// No batching.
func ReadRandom(database Database) {
	fmt.Println("ReadRandom test")

	keys := make([][]byte, *flagBatchSize)
	vals := make([][]byte, *flagBatchSize)
	// Write some key-value pairs first.
	// TODO: Allow user to just specify a database to open.
	timeLog := time.Now()
	x.Printf("Preparing database")
	for i := 0; i < *flagNumWrites; i++ {
		for j := 0; j < *flagBatchSize; j++ {
			keys[j] = []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumWrites))
			vals[j] = make([]byte, *flagValueSize)
			rng.Bytes(vals[j])
		}
		database.BatchPut(keys, vals)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent written\n", float64(i)*100.0/float64(*flagNumWrites))
			timeLog = time.Now()
		}
	}

	x.Printf("Slight pause\n")
	time.Sleep(time.Second)

	x.Printf("Starting reads")
	timeLogI := 0
	timeStart := time.Now()
	timeLog = timeStart
	for i := 0; i < *flagNumReads; i++ {
		key := []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumReads))
		database.Get(key)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent: %s\n", float64(i)*100.0/float64(*flagNumReads),
				report(timeElapsed, (i-timeLogI)))
			timeLog = time.Now()
			timeLogI = i
		}
	}
	x.Printf("Overall: %s\n", report(time.Since(timeStart), *flagNumReads))
}
<commit_msg>Small fixes.<commit_after>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"runtime/pprof"
	"time"

	"github.com/dgraph-io/badger/badger"
	"github.com/dgraph-io/badger/value"

	"github.com/dgraph-io/dgraph/store"
	"github.com/dgraph-io/dgraph/x"
)

var (
	flagBench      = flag.String("bench", "", "Run which benchmark?")
	flagDB         = flag.String("db", "", "Which DB: rocksdb, badger")
	flagValueSize  = flag.Int("value_size", 100, "Size of each value.")
	flagBatchSize  = flag.Int("batch_size", 1, "Size of writebatch.")
	flagNumWrites  = flag.Int("writes", 1000000, "Number of key-value pairs to write.")
	flagNumReads   = flag.Int("reads", 1000000, "Number of key-value pairs to read.")
	flagRandSize   = flag.Int("rand_size", 1000000, "Size of rng buffer.")
	flagCpuProfile = flag.String("cpu_profile", "", "Write cpu profile to file.")
	flagVerbose    = flag.Bool("verbose", false, "Verbose.")
	flagDir        = flag.String("dir", "/tmp/badger_bench", "Where data is temporarily stored.")

	rdbStore *store.Store
	rng      randomGenerator
)

type randomGenerator struct {
	data []byte
	idx  int
}

func (s *randomGenerator) Init() {
	if *flagRandSize <= 0 {
		// Will not precompute the randomness.
		return
	}
	s.data = make([]byte, *flagRandSize)
	n, err := rand.Read(s.data)
	x.Check(err)
	x.AssertTrue(n == *flagRandSize)
	s.idx = 0
}

// Bytes generates len(out) random bytes and writes to out.
func (s *randomGenerator) Bytes(out []byte) {
	if *flagRandSize == 0 {
		n, err := rand.Read(out)
		x.AssertTrue(n == len(out))
		x.Check(err)
		return
	}
	size := len(out)
	if s.idx+size > len(s.data) {
		s.idx = 0
	}
	x.AssertTrue(size == copy(out, s.data[s.idx:s.idx+size]))
	s.idx += size
}

func (s *randomGenerator) Int() int {
	var buf [4]byte
	s.Bytes(buf[:])
	return int(binary.LittleEndian.Uint32(buf[:]))
}

func main() {
	x.Init()
	x.AssertTrue(len(*flagBench) > 0)
	x.AssertTrue(*flagValueSize > 0)
	rng.Init()

	if *flagCpuProfile != "" {
		f, err := os.Create(*flagCpuProfile)
		if err != nil {
			x.Fatalf("Profiler error: %v", err)
		}
		pprof.StartCPUProfile(f)
	}

	var database Database
	switch *flagDB {
	case "badger":
		database = new(BadgerAdapter)
	case "rocksdb":
		database = new(RocksDBAdapter)
	default:
		x.Fatalf("Database invalid: %v", *flagDB)
	}
	database.Init()
	defer database.Close()

	x.AssertTrue(*flagDB == "rocksdb" || *flagDB == "badger")
	switch *flagBench {
	case "writerandom":
		WriteRandom(database)
	case "batchwriterandom":
		BatchWriteRandom(database)
	case "readrandom":
		ReadRandom(database)
	default:
		x.Fatalf("Unknown benchmark: %v", *flagBench)
	}
	if *flagCpuProfile != "" {
		pprof.StopCPUProfile()
	}
}

func report(d time.Duration, n int) string {
	secs := d.Seconds()
	throughput := float64(n*(16+*flagValueSize)) / ((1 << 20) * secs)
	return fmt.Sprintf("%.2fs, %.2fMb/s", secs, throughput)
}

type Database interface {
	Init()
	Close()
	Put(key, val []byte)
	BatchPut(key, val [][]byte)
	Get(key []byte)
}

type RocksDBAdapter struct {
	rdb *store.Store
	dir string
}

func (s *RocksDBAdapter) Init() {
	var err error
	s.dir, err = ioutil.TempDir(*flagDir, "storetest_")
	x.Check(err)
	s.rdb, err = store.NewSyncStore(s.dir)
	x.Check(err)
}

func (s *RocksDBAdapter) Close() {
	//	s.rdb.Close()
}

func (s *RocksDBAdapter) Put(key, val []byte) {
	s.rdb.SetOne(key, val)
}

func (s *RocksDBAdapter) BatchPut(key, val [][]byte) {
	wb := s.rdb.NewWriteBatch()
	x.AssertTrue(len(key) == len(val))
	for i := 0; i < len(key); i++ {
		wb.Put(key[i], val[i])
	}
	x.Check(s.rdb.WriteBatch(wb))
}

func (s *RocksDBAdapter) Get(key []byte) {
	_, err := s.rdb.Get(key)
	x.Check(err)
}

type BadgerAdapter struct {
	db  *badger.DB
	dir string
}

func (s *BadgerAdapter) Init() {
	opt := badger.DefaultOptions
	opt.Verbose = *flagVerbose
	opt.Dir = *flagDir

	fmt.Printf("Dir: %s\n", *flagDir)
	s.db = badger.NewDB(opt)
}

func (s *BadgerAdapter) Close() {
}

func (s *BadgerAdapter) Put(key, val []byte) {
	s.db.Put(key, val)
}

func (s *BadgerAdapter) BatchPut(key, val [][]byte) {
	var entries []value.Entry
	x.AssertTrue(len(key) == len(val))
	for i := 0; i < len(key); i++ {
		entries = append(entries, value.Entry{
			Key:   key[i],
			Value: val[i],
		})
	}
	x.Check(s.db.Write(entries))
}

func (s *BadgerAdapter) Get(key []byte) {
	s.db.Get(key)
}

// No batching.
func WriteRandom(database Database) {
	fmt.Println("WriteRandom test")
	timeStart := time.Now()
	timeLog := timeStart
	timeLogI := 0
	// If you use b.N, you might add too few samples and be working only in memory.
	// We need to fix a large number of pairs. This is what LevelDB benchmark does as well.
	for i := 0; i < *flagNumWrites; i++ {
		key := []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumWrites))
		val := make([]byte, *flagValueSize)
		rng.Bytes(val)
		database.Put(key, val)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent: %s\n", float64(i)*100.0/float64(*flagNumWrites),
				report(timeElapsed, i-timeLogI))
			timeLog = time.Now()
			timeLogI = i
		}
	}
	x.Printf("Overall: %s\n", report(time.Since(timeStart), *flagNumWrites))
}

// With batching.
func BatchWriteRandom(database Database) {
	fmt.Println("BatchWriteRandom test")
	timeStart := time.Now()
	timeLog := timeStart
	timeLogI := 0
	keys := make([][]byte, *flagBatchSize)
	vals := make([][]byte, *flagBatchSize)

	for i := 0; i < *flagNumWrites; i++ {
		for j := 0; j < *flagBatchSize; j++ {
			keys[j] = []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumWrites))
			vals[j] = make([]byte, *flagValueSize)
			rng.Bytes(vals[j])
		}
		database.BatchPut(keys, vals)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent: %s\n", float64(i)*100.0/float64(*flagNumWrites),
				report(timeElapsed, (i-timeLogI)*(*flagBatchSize)))
			timeLog = time.Now()
			timeLogI = i
		}
	}
	x.Printf("Overall: %s\n", report(time.Since(timeStart), *flagNumWrites*(*flagBatchSize)))
}

// No batching.
func ReadRandom(database Database) {
	fmt.Println("ReadRandom test")

	keys := make([][]byte, *flagBatchSize)
	vals := make([][]byte, *flagBatchSize)
	// Write some key-value pairs first.
	// TODO: Allow user to just specify a database to open.
	timeLog := time.Now()
	x.Printf("Preparing database")
	for i := 0; i < *flagNumWrites; i++ {
		for j := 0; j < *flagBatchSize; j++ {
			keys[j] = []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumWrites))
			vals[j] = make([]byte, *flagValueSize)
			rng.Bytes(vals[j])
		}
		database.BatchPut(keys, vals)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent written\n", float64(i)*100.0/float64(*flagNumWrites))
			timeLog = time.Now()
		}
	}

	x.Printf("Slight pause\n")
	time.Sleep(time.Second)

	x.Printf("Starting reads")
	timeLogI := 0
	timeStart := time.Now()
	timeLog = timeStart
	for i := 0; i < *flagNumReads; i++ {
		key := []byte(fmt.Sprintf("%016d", rng.Int()%*flagNumReads))
		database.Get(key)
		timeElapsed := time.Since(timeLog)
		if timeElapsed > 5*time.Second {
			x.Printf("%.2f percent: %s\n", float64(i)*100.0/float64(*flagNumReads),
				report(timeElapsed, (i-timeLogI)))
			timeLog = time.Now()
			timeLogI = i
		}
	}
	x.Printf("Overall: %s\n", report(time.Since(timeStart), *flagNumReads))
}
<|endoftext|>


<commit_before>package handler

import (
	"net/http"

	"socialapi/config"
	"socialapi/models"
	"socialapi/workers/common/metrics"

	kmetrics "github.com/koding/metrics"

	tigertonic "github.com/rcrowley/go-tigertonic"

	gometrics "github.com/rcrowley/go-metrics"
)

const (
	PostRequest   = "POST"
	GetRequest    = "GET"
	DeleteRequest = "DELETE"
)

var (
	// TODO allowed origins must be configurable
	cors     = tigertonic.NewCORSBuilder().AddAllowedOrigins("*")
	conf     *config.Config
	trackers *metrics.Trackers
)

type Request struct {
	Handler        interface{}
	Endpoint       string
	Type           string
	Name           string
	CollectMetrics bool
	Metrics        *kmetrics.Metrics
}

// todo add prooper logging
func getAccount(r *http.Request) *models.Account {
	cookie, err := r.Cookie("clientId")
	if err != nil {
		return nil
	}

	// if cookie doenst exists return nil
	if cookie.Value == "" {
		return nil
	}

	session, err := models.Cache.Session.ById(cookie.Value)
	if err != nil {
		return nil
	}

	// if session doesnt have username, return nil
	if session.Username == "" {
		return nil
	}

	acc, err := models.Cache.Account.ByNick(session.Username)
	if err != nil {
		return nil
	}

	return acc
}

func Wrapper(r Request) http.Handler {
	handler := r.Handler

	var hHandler http.Handler

	// count the statuses of the requests
	hHandler = buildHandlerWithStatusCount(handler, r)

	hHandler = buildHandlerWithTimeTracking(hHandler, r)

	// create the final handler
	return cors.Build(hHandler)
}

// count the statuses of the requests
func buildHandlerWithStatusCount(handler interface{}, r Request) http.Handler {
	return CountedByStatus(
		tigertonic.Marshaled(handler), r.Name, r.CollectMetrics,
	)
}

// add request time tracking
func buildHandlerWithTimeTracking(handler http.Handler, r Request) http.Handler {
	var registry gometrics.Registry

	if r.Metrics != nil {
		registry = r.Metrics.Registry
	}

	return tigertonic.Timed(
		handler,
		r.Name,
		registry,
	)
}

func BuildHandlerWithContext(handler http.Handler) http.Handler {
	// add context
	return tigertonic.If(
		func(r *http.Request) (http.Header, error) {
			// this is an example
			// set group name to context
			//
			context := &models.Context{
				GroupName: "koding",
				Client: &models.Client{
					Account: getAccount(r),
				},
			}

			*(tigertonic.Context(r).(*models.Context)) = *context
			return nil, nil
		},
		handler,
	)
}

//----------------------------------------------------------
// CounterByStatus
//----------------------------------------------------------

type CounterByStatus struct {
	name           string
	handler        http.Handler
	collectMetrics bool
}

func (c *CounterByStatus) ServeHTTP(w0 http.ResponseWriter, r *http.Request) {
	w := tigertonic.NewTeeHeaderResponseWriter(w0)
	c.handler.ServeHTTP(w, r)

	if w.StatusCode <= 300 {
		c.track()
	}
}

func (c *CounterByStatus) track() {
	// don't log if analytics are disabled for that endpoint
	if !c.collectMetrics {
		return
	}

	// set `conf` and `trackers` if either is not set
	if conf == nil || trackers == nil {
		conf = config.MustGet()

		// don't log if analytics are disabled globally
		if !conf.Mixpanel.Enabled {
			return
		}

		trackers = metrics.InitTrackers(
			metrics.NewMixpanelTracker(conf.Mixpanel.Token),
		)
	}

	trackers.Track(c.name)
}

func CountedByStatus(handler http.Handler, name string, collectMetrics bool) *CounterByStatus {
	return &CounterByStatus{
		name:           name,
		handler:        handler,
		collectMetrics: collectMetrics,
	}
}
<commit_msg>handler: use request for also outgoing connections<commit_after>package handler

import (
	"fmt"
	"net/http"

	"socialapi/config"
	"socialapi/models"
	"socialapi/workers/common/metrics"
	"strings"

	kmetrics "github.com/koding/metrics"

	tigertonic "github.com/rcrowley/go-tigertonic"

	gometrics "github.com/rcrowley/go-metrics"
)

const (
	PostRequest   = "POST"
	GetRequest    = "GET"
	DeleteRequest = "DELETE"
)

var (
	// TODO allowed origins must be configurable
	cors     = tigertonic.NewCORSBuilder().AddAllowedOrigins("*")
	conf     *config.Config
	trackers *metrics.Trackers
)

type Request struct {
	Handler        interface{}
	Endpoint       string
	Type           string
	Name           string
	CollectMetrics bool
	Metrics        *kmetrics.Metrics
	// used for external requests
	Params  map[string]string
	Cookie  string
	Cookies []*http.Cookie
}

// todo add prooper logging
func getAccount(r *http.Request) *models.Account {
	cookie, err := r.Cookie("clientId")
	if err != nil {
		return nil
	}

	// if cookie doenst exists return nil
	if cookie.Value == "" {
		return nil
	}

	session, err := models.Cache.Session.ById(cookie.Value)
	if err != nil {
		return nil
	}

	// if session doesnt have username, return nil
	if session.Username == "" {
		return nil
	}

	acc, err := models.Cache.Account.ByNick(session.Username)
	if err != nil {
		return nil
	}

	return acc
}

func Wrapper(r Request) http.Handler {
	handler := r.Handler

	var hHandler http.Handler

	// count the statuses of the requests
	hHandler = buildHandlerWithStatusCount(handler, r)

	hHandler = buildHandlerWithTimeTracking(hHandler, r)

	// create the final handler
	return cors.Build(hHandler)
}

// count the statuses of the requests
func buildHandlerWithStatusCount(handler interface{}, r Request) http.Handler {
	return CountedByStatus(
		tigertonic.Marshaled(handler), r.Name, r.CollectMetrics,
	)
}

// add request time tracking
func buildHandlerWithTimeTracking(handler http.Handler, r Request) http.Handler {
	var registry gometrics.Registry

	if r.Metrics != nil {
		registry = r.Metrics.Registry
	}

	return tigertonic.Timed(
		handler,
		r.Name,
		registry,
	)
}

func BuildHandlerWithContext(handler http.Handler) http.Handler {
	// add context
	return tigertonic.If(
		func(r *http.Request) (http.Header, error) {
			// this is an example
			// set group name to context
			//
			context := &models.Context{
				GroupName: "koding",
				Client: &models.Client{
					Account: getAccount(r),
				},
			}

			*(tigertonic.Context(r).(*models.Context)) = *context
			return nil, nil
		},
		handler,
	)
}

//----------------------------------------------------------
// CounterByStatus
//----------------------------------------------------------

type CounterByStatus struct {
	name           string
	handler        http.Handler
	collectMetrics bool
}

func (c *CounterByStatus) ServeHTTP(w0 http.ResponseWriter, r *http.Request) {
	w := tigertonic.NewTeeHeaderResponseWriter(w0)
	c.handler.ServeHTTP(w, r)

	if w.StatusCode <= 300 {
		c.track()
	}
}

func (c *CounterByStatus) track() {
	// don't log if analytics are disabled for that endpoint
	if !c.collectMetrics {
		return
	}

	// set `conf` and `trackers` if either is not set
	if conf == nil || trackers == nil {
		conf = config.MustGet()

		// don't log if analytics are disabled globally
		if !conf.Mixpanel.Enabled {
			return
		}

		trackers = metrics.InitTrackers(
			metrics.NewMixpanelTracker(conf.Mixpanel.Token),
		)
	}

	trackers.Track(c.name)
}

func CountedByStatus(handler http.Handler, name string, collectMetrics bool) *CounterByStatus {
	return &CounterByStatus{
		name:           name,
		handler:        handler,
		collectMetrics: collectMetrics,
	}
}

//
func MakeRequest(request *Request) (*http.Response, error) {
	if request.Cookie != "" {
		request.Cookies = parseCookies(request.Cookie)
	}

	request.Endpoint = prepareQueryString(request.Endpoint, request.Params)

	client := new(http.Client)
	req, err := http.NewRequest(request.Type, request.Endpoint, nil)
	if err != nil {
		return nil, err
	}

	// Add cookies
	for _, cookie := range request.Cookies {
		req.AddCookie(cookie)
	}

	return client.Do(req)
}

func prepareQueryString(endpoint string, params map[string]string) string {
	if len(params) == 0 {
		return endpoint
	}

	fullPath := fmt.Sprintf("%s?", endpoint)

	for key, value := range params {
		fullPath = fmt.Sprintf("%s%s=%s&", fullPath, key, value)
	}

	return fullPath[0 : len(fullPath)-1]
}

func parseCookies(cookie string) []*http.Cookie {
	pairs := strings.Split(cookie, "; ")
	cookies := make([]*http.Cookie, 0)

	if len(pairs) == 0 {
		return cookies
	}

	for _, val := range pairs {
		cp := strings.Split(val, "=")
		if len(cp) != 2 {
			continue
		}

		c := new(http.Cookie)
		c.Name = cp[0]
		c.Value = cp[1]

		cookies = append(cookies, c)
	}

	return cookies
}
<|endoftext|>


<commit_before>package proto

import (
	"encoding"
	"fmt"
	"reflect"
	"time"

	"github.com/go-redis/redis/v8/internal/util"
)

// Scan parses bytes `b` to `v` with appropriate type.
func Scan(b []byte, v interface{}) error {
	switch v := v.(type) {
	case nil:
		return fmt.Errorf("redis: Scan(nil)")
	case *string:
		*v = util.BytesToString(b)
		return nil
	case *[]byte:
		*v = b
		return nil
	case *int:
		var err error
		*v, err = util.Atoi(b)
		return err
	case *int8:
		n, err := util.ParseInt(b, 10, 8)
		if err != nil {
			return err
		}
		*v = int8(n)
		return nil
	case *int16:
		n, err := util.ParseInt(b, 10, 16)
		if err != nil {
			return err
		}
		*v = int16(n)
		return nil
	case *int32:
		n, err := util.ParseInt(b, 10, 32)
		if err != nil {
			return err
		}
		*v = int32(n)
		return nil
	case *int64:
		n, err := util.ParseInt(b, 10, 64)
		if err != nil {
			return err
		}
		*v = n
		return nil
	case *uint:
		n, err := util.ParseUint(b, 10, 64)
		if err != nil {
			return err
		}
		*v = uint(n)
		return nil
	case *uint8:
		n, err := util.ParseUint(b, 10, 8)
		if err != nil {
			return err
		}
		*v = uint8(n)
		return nil
	case *uint16:
		n, err := util.ParseUint(b, 10, 16)
		if err != nil {
			return err
		}
		*v = uint16(n)
		return nil
	case *uint32:
		n, err := util.ParseUint(b, 10, 32)
		if err != nil {
			return err
		}
		*v = uint32(n)
		return nil
	case *uint64:
		n, err := util.ParseUint(b, 10, 64)
		if err != nil {
			return err
		}
		*v = n
		return nil
	case *float32:
		n, err := util.ParseFloat(b, 32)
		if err != nil {
			return err
		}
		*v = float32(n)
		return err
	case *float64:
		var err error
		*v, err = util.ParseFloat(b, 64)
		return err
	case *bool:
		*v = len(b) == 1 && b[0] == '1'
		return nil
	case *time.Time:
		var err error
		*v, err = time.Parse(time.RFC3339Nano, util.BytesToString(b))
		return err
	case *time.Duration:
		n, err := util.ParseInt(b, 10, 64)
		if err != nil {
			return err
		}
		*v = time.Duration(n)
		return nil
	case encoding.BinaryUnmarshaler:
		return v.UnmarshalBinary(b)
	default:
		return fmt.Errorf(
			"redis: can't unmarshal %T (consider implementing BinaryUnmarshaler)", v)
	}
}

func ScanSlice(data []string, slice interface{}) error {
	v := reflect.ValueOf(slice)
	if !v.IsValid() {
		return fmt.Errorf("redis: ScanSlice(nil)")
	}
	if v.Kind() != reflect.Ptr {
		return fmt.Errorf("redis: ScanSlice(non-pointer %T)", slice)
	}
	v = v.Elem()
	if v.Kind() != reflect.Slice {
		return fmt.Errorf("redis: ScanSlice(non-slice %T)", slice)
	}

	next := makeSliceNextElemFunc(v)
	for i, s := range data {
		elem := next()
		if err := Scan([]byte(s), elem.Addr().Interface()); err != nil {
			err = fmt.Errorf("redis: ScanSlice index=%d value=%q failed: %w", i, s, err)
			return err
		}
	}

	return nil
}

func makeSliceNextElemFunc(v reflect.Value) func() reflect.Value {
	elemType := v.Type().Elem()

	if elemType.Kind() == reflect.Ptr {
		elemType = elemType.Elem()
		return func() reflect.Value {
			if v.Len() < v.Cap() {
				v.Set(v.Slice(0, v.Len()+1))
				elem := v.Index(v.Len() - 1)
				if elem.IsNil() {
					elem.Set(reflect.New(elemType))
				}
				return elem.Elem()
			}

			elem := reflect.New(elemType)
			v.Set(reflect.Append(v, elem))
			return elem.Elem()
		}
	}

	zero := reflect.Zero(elemType)
	return func() reflect.Value {
		if v.Len() < v.Cap() {
			v.Set(v.Slice(0, v.Len()+1))
			return v.Index(v.Len() - 1)
		}

		v.Set(reflect.Append(v, zero))
		return v.Index(v.Len() - 1)
	}
}
<commit_msg>chore: fix linter<commit_after>package proto

import (
	"encoding"
	"fmt"
	"reflect"
	"time"

	"github.com/go-redis/redis/v8/internal/util"
)

// Scan parses bytes `b` to `v` with appropriate type.
//nolint:gocyclo
func Scan(b []byte, v interface{}) error {
	switch v := v.(type) {
	case nil:
		return fmt.Errorf("redis: Scan(nil)")
	case *string:
		*v = util.BytesToString(b)
		return nil
	case *[]byte:
		*v = b
		return nil
	case *int:
		var err error
		*v, err = util.Atoi(b)
		return err
	case *int8:
		n, err := util.ParseInt(b, 10, 8)
		if err != nil {
			return err
		}
		*v = int8(n)
		return nil
	case *int16:
		n, err := util.ParseInt(b, 10, 16)
		if err != nil {
			return err
		}
		*v = int16(n)
		return nil
	case *int32:
		n, err := util.ParseInt(b, 10, 32)
		if err != nil {
			return err
		}
		*v = int32(n)
		return nil
	case *int64:
		n, err := util.ParseInt(b, 10, 64)
		if err != nil {
			return err
		}
		*v = n
		return nil
	case *uint:
		n, err := util.ParseUint(b, 10, 64)
		if err != nil {
			return err
		}
		*v = uint(n)
		return nil
	case *uint8:
		n, err := util.ParseUint(b, 10, 8)
		if err != nil {
			return err
		}
		*v = uint8(n)
		return nil
	case *uint16:
		n, err := util.ParseUint(b, 10, 16)
		if err != nil {
			return err
		}
		*v = uint16(n)
		return nil
	case *uint32:
		n, err := util.ParseUint(b, 10, 32)
		if err != nil {
			return err
		}
		*v = uint32(n)
		return nil
	case *uint64:
		n, err := util.ParseUint(b, 10, 64)
		if err != nil {
			return err
		}
		*v = n
		return nil
	case *float32:
		n, err := util.ParseFloat(b, 32)
		if err != nil {
			return err
		}
		*v = float32(n)
		return err
	case *float64:
		var err error
		*v, err = util.ParseFloat(b, 64)
		return err
	case *bool:
		*v = len(b) == 1 && b[0] == '1'
		return nil
	case *time.Time:
		var err error
		*v, err = time.Parse(time.RFC3339Nano, util.BytesToString(b))
		return err
	case *time.Duration:
		n, err := util.ParseInt(b, 10, 64)
		if err != nil {
			return err
		}
		*v = time.Duration(n)
		return nil
	case encoding.BinaryUnmarshaler:
		return v.UnmarshalBinary(b)
	default:
		return fmt.Errorf(
			"redis: can't unmarshal %T (consider implementing BinaryUnmarshaler)", v)
	}
}

func ScanSlice(data []string, slice interface{}) error {
	v := reflect.ValueOf(slice)
	if !v.IsValid() {
		return fmt.Errorf("redis: ScanSlice(nil)")
	}
	if v.Kind() != reflect.Ptr {
		return fmt.Errorf("redis: ScanSlice(non-pointer %T)", slice)
	}
	v = v.Elem()
	if v.Kind() != reflect.Slice {
		return fmt.Errorf("redis: ScanSlice(non-slice %T)", slice)
	}

	next := makeSliceNextElemFunc(v)
	for i, s := range data {
		elem := next()
		if err := Scan([]byte(s), elem.Addr().Interface()); err != nil {
			err = fmt.Errorf("redis: ScanSlice index=%d value=%q failed: %w", i, s, err)
			return err
		}
	}

	return nil
}

func makeSliceNextElemFunc(v reflect.Value) func() reflect.Value {
	elemType := v.Type().Elem()

	if elemType.Kind() == reflect.Ptr {
		elemType = elemType.Elem()
		return func() reflect.Value {
			if v.Len() < v.Cap() {
				v.Set(v.Slice(0, v.Len()+1))
				elem := v.Index(v.Len() - 1)
				if elem.IsNil() {
					elem.Set(reflect.New(elemType))
				}
				return elem.Elem()
			}

			elem := reflect.New(elemType)
			v.Set(reflect.Append(v, elem))
			return elem.Elem()
		}
	}

	zero := reflect.Zero(elemType)
	return func() reflect.Value {
		if v.Len() < v.Cap() {
			v.Set(v.Slice(0, v.Len()+1))
			return v.Index(v.Len() - 1)
		}

		v.Set(reflect.Append(v, zero))
		return v.Index(v.Len() - 1)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2021 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package stdscript

import (
	"bytes"
	"testing"

	"github.com/decred/dcrd/txscript/v4"
)

// complexScriptV0 is a version 0 script comprised of half as many opcodes as
// the maximum allowed followed by as many max size data pushes fit without
// exceeding the max allowed script size.
var complexScriptV0 = func() []byte {
	const (
		maxScriptSize        = txscript.MaxScriptSize
		maxScriptElementSize = txscript.MaxScriptElementSize
	)
	var scriptLen int
	builder := txscript.NewScriptBuilder()
	for i := 0; i < txscript.MaxOpsPerScript/2; i++ {
		builder.AddOp(txscript.OP_TRUE)
		scriptLen++
	}
	maxData := bytes.Repeat([]byte{0x02}, maxScriptElementSize)
	for i := 0; i < (maxScriptSize-scriptLen)/maxScriptElementSize; i++ {
		builder.AddData(maxData)
	}
	script, err := builder.Script()
	if err != nil {
		panic(err)
	}
	return script
}()

// makeBenchmarks constructs a slice of tests to use in the benchmarks as
// follows:
// - Start with a version 0 complex non standard script
// - Add all tests for which the provided filter function returns true
func makeBenchmarks(filterFn func(test scriptTest) bool) []scriptTest {
	benches := make([]scriptTest, 0, 5)
	benches = append(benches, scriptTest{
		name:    "v0 complex non standard",
		version: 0,
		script:  complexScriptV0,
	})
	for _, test := range scriptV0Tests {
		if filterFn(test) {
			benches = append(benches, test)
		}
	}
	return benches
}

// benchIsX is a convenience function that runs benchmarks for the entries that
// match the provided filter function using the given script type determination
// function and ensures the result matches the expected one.
func benchIsX(b *testing.B, filterFn func(test scriptTest) bool, isXFn func(scriptVersion uint16, script []byte) bool) {
	b.Helper()

	benches := makeBenchmarks(filterFn)
	for _, bench := range benches {
		want := filterFn(bench)
		b.Run(bench.name, func(b *testing.B) {
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				got := isXFn(bench.version, bench.script)
				if got != want {
					b.Fatalf("%q: unexpected result -- got %v, want %v",
						bench.name, got, want)
				}
			}
		})
	}
}

// BenchmarkIsPubKeyScript benchmarks the performance of analyzing various
// public key scripts to determine if they are p2pk-ecdsa-secp256k1 scripts.
func BenchmarkIsPubKeyScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyEcdsaSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeyScript)
}

// BenchmarkIsPubKeyEd25519Script benchmarks the performance of analyzing
// various public key scripts to determine if they are p2pkh-ed25519 scripts.
func BenchmarkIsPubKeyEd25519Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyEd25519
	}
	benchIsX(b, filterFn, IsPubKeyEd25519Script)
}

// BenchmarkIsPubKeySchnorrSecp256k1Script benchmarks the performance of
// analyzing various public key scripts to determine if they are
// p2pkh-schnorr-secp256k1 scripts.
func BenchmarkIsPubKeySchnorrSecp256k1Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeySchnorrSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeySchnorrSecp256k1Script)
}

// BenchmarkIsPubKeyHashScript benchmarks the performance of analyzing various
// public key scripts to determine if they are p2pkh-ecdsa-secp256k1 scripts.
func BenchmarkIsPubKeyHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyHashEcdsaSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeyHashScript)
}

// BenchmarkIsPubKeyHashEd25519Script benchmarks the performance of analyzing
// various public key scripts to determine if they are p2pkh-ed25519 scripts.
func BenchmarkIsPubKeyHashEd25519Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyHashEd25519
	}
	benchIsX(b, filterFn, IsPubKeyHashEd25519Script)
}

// BenchmarkIsPubKeyHashSchnorrSecp256k1Script benchmarks the performance of
// analyzing various public key scripts to determine if they are
// p2pkh-schnorr-secp256k1 scripts.
func BenchmarkIsPubKeyHashSchnorrSecp256k1Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyHashSchnorrSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeyHashSchnorrSecp256k1Script)
}

// BenchmarkIsScriptHashScript benchmarks the performance of analyzing various
// public key scripts to determine if they are p2sh scripts.
func BenchmarkIsScriptHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STScriptHash
	}
	benchIsX(b, filterFn, IsScriptHashScript)
}

// BenchmarkIsMultiSigScript benchmarks the performance of analyzing various
// public key scripts to determine if they are multisignature scripts.
func BenchmarkIsMultiSigScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STMultiSig && !test.isSig
	}
	benchIsX(b, filterFn, IsMultiSigScript)
}

// BenchmarkIsMultiSigSigScript benchmarks the performance of analyzing various
// signature scripts to determine if they are likely to be multisignature redeem
// scripts.
func BenchmarkIsMultiSigSigScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STMultiSig && test.isSig
	}
	benchIsX(b, filterFn, IsMultiSigSigScript)
}

// BenchmarkMultiSigRedeemScriptFromScriptSigV0 benchmarks the performance of
// extracting the redeem script from various version 0 signature scripts.
func BenchmarkMultiSigRedeemScriptFromScriptSigV0(b *testing.B) {
	for _, test := range scriptV0Tests {
		if test.version != 0 || test.wantType != STMultiSig || !test.isSig {
			continue
		}
		want := test.wantData.([]byte)

		b.Run(test.name, func(b *testing.B) {
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				got := MultiSigRedeemScriptFromScriptSigV0(test.script)
				if !bytes.Equal(got, want) {
					b.Fatalf("%q: unexpected result -- got %x, want %x",
						test.name, got, want)
				}
			}
		})
	}
}

// BenchmarkIsNullDataScript benchmarks the performance of analyzing various
// public key scripts to determine if they are nulldata scripts.
func BenchmarkIsNullDataScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STNullData
	}
	benchIsX(b, filterFn, IsNullDataScript)
}

// BenchmarkIsStakeSubmissionPubKeyHashScript benchmarks the performance of
// analyzing various public key scripts to determine if they are stake
// submission p2pkh-ecdsa-secp256k1 scripts.
func BenchmarkIsStakeSubmissionPubKeyHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeSubmissionPubKeyHash
	}
	benchIsX(b, filterFn, IsStakeSubmissionPubKeyHashScript)
}

// BenchmarkIsStakeSubmissionScriptHashScript benchmarks the performance of
// analyzing various public key scripts to determine if they are stake
// submission p2sh scripts.
func BenchmarkIsStakeSubmissionScriptHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeSubmissionScriptHash
	}
	benchIsX(b, filterFn, IsStakeSubmissionScriptHashScript)
}

// BenchmarkIsStakeGenPubKeyHashScript benchmarks the performance of analyzing
// various public key scripts to determine if they are stake generation
// p2pkh-ecdsa-secp256k1 scripts.
func BenchmarkIsStakeGenPubKeyHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeGenPubKeyHash
	}
	benchIsX(b, filterFn, IsStakeGenPubKeyHashScript)
}
<commit_msg>stdscript: Add v0 stake gen p2sh benchmark.<commit_after>// Copyright (c) 2021 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package stdscript

import (
	"bytes"
	"testing"

	"github.com/decred/dcrd/txscript/v4"
)

// complexScriptV0 is a version 0 script comprised of half as many opcodes as
// the maximum allowed followed by as many max size data pushes fit without
// exceeding the max allowed script size.
var complexScriptV0 = func() []byte {
	const (
		maxScriptSize        = txscript.MaxScriptSize
		maxScriptElementSize = txscript.MaxScriptElementSize
	)
	var scriptLen int
	builder := txscript.NewScriptBuilder()
	for i := 0; i < txscript.MaxOpsPerScript/2; i++ {
		builder.AddOp(txscript.OP_TRUE)
		scriptLen++
	}
	maxData := bytes.Repeat([]byte{0x02}, maxScriptElementSize)
	for i := 0; i < (maxScriptSize-scriptLen)/maxScriptElementSize; i++ {
		builder.AddData(maxData)
	}
	script, err := builder.Script()
	if err != nil {
		panic(err)
	}
	return script
}()

// makeBenchmarks constructs a slice of tests to use in the benchmarks as
// follows:
// - Start with a version 0 complex non standard script
// - Add all tests for which the provided filter function returns true
func makeBenchmarks(filterFn func(test scriptTest) bool) []scriptTest {
	benches := make([]scriptTest, 0, 5)
	benches = append(benches, scriptTest{
		name:    "v0 complex non standard",
		version: 0,
		script:  complexScriptV0,
	})
	for _, test := range scriptV0Tests {
		if filterFn(test) {
			benches = append(benches, test)
		}
	}
	return benches
}

// benchIsX is a convenience function that runs benchmarks for the entries that
// match the provided filter function using the given script type determination
// function and ensures the result matches the expected one.
func benchIsX(b *testing.B, filterFn func(test scriptTest) bool, isXFn func(scriptVersion uint16, script []byte) bool) {
	b.Helper()

	benches := makeBenchmarks(filterFn)
	for _, bench := range benches {
		want := filterFn(bench)
		b.Run(bench.name, func(b *testing.B) {
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				got := isXFn(bench.version, bench.script)
				if got != want {
					b.Fatalf("%q: unexpected result -- got %v, want %v",
						bench.name, got, want)
				}
			}
		})
	}
}

// BenchmarkIsPubKeyScript benchmarks the performance of analyzing various
// public key scripts to determine if they are p2pk-ecdsa-secp256k1 scripts.
func BenchmarkIsPubKeyScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyEcdsaSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeyScript)
}

// BenchmarkIsPubKeyEd25519Script benchmarks the performance of analyzing
// various public key scripts to determine if they are p2pkh-ed25519 scripts.
func BenchmarkIsPubKeyEd25519Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyEd25519
	}
	benchIsX(b, filterFn, IsPubKeyEd25519Script)
}

// BenchmarkIsPubKeySchnorrSecp256k1Script benchmarks the performance of
// analyzing various public key scripts to determine if they are
// p2pkh-schnorr-secp256k1 scripts.
func BenchmarkIsPubKeySchnorrSecp256k1Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeySchnorrSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeySchnorrSecp256k1Script)
}

// BenchmarkIsPubKeyHashScript benchmarks the performance of analyzing various
// public key scripts to determine if they are p2pkh-ecdsa-secp256k1 scripts.
func BenchmarkIsPubKeyHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyHashEcdsaSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeyHashScript)
}

// BenchmarkIsPubKeyHashEd25519Script benchmarks the performance of analyzing
// various public key scripts to determine if they are p2pkh-ed25519 scripts.
func BenchmarkIsPubKeyHashEd25519Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyHashEd25519
	}
	benchIsX(b, filterFn, IsPubKeyHashEd25519Script)
}

// BenchmarkIsPubKeyHashSchnorrSecp256k1Script benchmarks the performance of
// analyzing various public key scripts to determine if they are
// p2pkh-schnorr-secp256k1 scripts.
func BenchmarkIsPubKeyHashSchnorrSecp256k1Script(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STPubKeyHashSchnorrSecp256k1
	}
	benchIsX(b, filterFn, IsPubKeyHashSchnorrSecp256k1Script)
}

// BenchmarkIsScriptHashScript benchmarks the performance of analyzing various
// public key scripts to determine if they are p2sh scripts.
func BenchmarkIsScriptHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STScriptHash
	}
	benchIsX(b, filterFn, IsScriptHashScript)
}

// BenchmarkIsMultiSigScript benchmarks the performance of analyzing various
// public key scripts to determine if they are multisignature scripts.
func BenchmarkIsMultiSigScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STMultiSig && !test.isSig
	}
	benchIsX(b, filterFn, IsMultiSigScript)
}

// BenchmarkIsMultiSigSigScript benchmarks the performance of analyzing various
// signature scripts to determine if they are likely to be multisignature redeem
// scripts.
func BenchmarkIsMultiSigSigScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STMultiSig && test.isSig
	}
	benchIsX(b, filterFn, IsMultiSigSigScript)
}

// BenchmarkMultiSigRedeemScriptFromScriptSigV0 benchmarks the performance of
// extracting the redeem script from various version 0 signature scripts.
func BenchmarkMultiSigRedeemScriptFromScriptSigV0(b *testing.B) {
	for _, test := range scriptV0Tests {
		if test.version != 0 || test.wantType != STMultiSig || !test.isSig {
			continue
		}
		want := test.wantData.([]byte)

		b.Run(test.name, func(b *testing.B) {
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				got := MultiSigRedeemScriptFromScriptSigV0(test.script)
				if !bytes.Equal(got, want) {
					b.Fatalf("%q: unexpected result -- got %x, want %x",
						test.name, got, want)
				}
			}
		})
	}
}

// BenchmarkIsNullDataScript benchmarks the performance of analyzing various
// public key scripts to determine if they are nulldata scripts.
func BenchmarkIsNullDataScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STNullData
	}
	benchIsX(b, filterFn, IsNullDataScript)
}

// BenchmarkIsStakeSubmissionPubKeyHashScript benchmarks the performance of
// analyzing various public key scripts to determine if they are stake
// submission p2pkh-ecdsa-secp256k1 scripts.
func BenchmarkIsStakeSubmissionPubKeyHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeSubmissionPubKeyHash
	}
	benchIsX(b, filterFn, IsStakeSubmissionPubKeyHashScript)
}

// BenchmarkIsStakeSubmissionScriptHashScript benchmarks the performance of
// analyzing various public key scripts to determine if they are stake
// submission p2sh scripts.
func BenchmarkIsStakeSubmissionScriptHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeSubmissionScriptHash
	}
	benchIsX(b, filterFn, IsStakeSubmissionScriptHashScript)
}

// BenchmarkIsStakeGenPubKeyHashScript benchmarks the performance of analyzing
// various public key scripts to determine if they are stake generation
// p2pkh-ecdsa-secp256k1 scripts.
func BenchmarkIsStakeGenPubKeyHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeGenPubKeyHash
	}
	benchIsX(b, filterFn, IsStakeGenPubKeyHashScript)
}

// BenchmarkIsStakeGenScriptHashScript benchmarks the performance of analyzing various
// public key scripts to determine if they are stake generation p2sh scripts.
func BenchmarkIsStakeGenScriptHashScript(b *testing.B) {
	filterFn := func(test scriptTest) bool {
		return test.wantType == STStakeGenScriptHash
	}
	benchIsX(b, filterFn, IsStakeGenScriptHashScript)
}
<|endoftext|>


<commit_before>// deploy2deter is responsible for kicking off the deployment process
// for deterlab. Given a list of hostnames, it will create an overlay
// tree topology, using all but the last node. It will create multiple
// nodes per server and run timestamping processes. The last node is
// reserved for the logging server, which is forwarded to localhost:8080
//
// options are "bf" which specifies the branching factor
//
// 	and "hpn" which specifies the replicaiton factor: hosts per node
//
// Creates the following directory structure in remote:
// exec, timeclient, logserver/...,
// this way it can rsync the remove to each of the destinations
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"

	"github.com/dedis/prifi/coco/test/cliutils"
	"github.com/dedis/prifi/coco/test/config"
	"github.com/dedis/prifi/coco/test/graphs"
)

// bf is the branching factor of the tree that we want to build
var bf int

// hpn is the replication factor of hosts per node: how many hosts do we want per node
var hpn int

var nmsgs int
var debug bool
var rate int
var failures int

var rounds int

func init() {
	flag.IntVar(&bf, "bf", 2, "branching factor: default binary")
	flag.IntVar(&hpn, "hpn", 1, "hosts per node: default 1")
	flag.IntVar(&nmsgs, "nmsgs", 100, "number of messages per round")
	flag.IntVar(&rate, "rate", -1, "number of milliseconds between messages: if rate > 0 then used")
	flag.BoolVar(&debug, "debug", false, "run in debugging mode")
	flag.IntVar(&failures, "failures", 0, "percent showing per node probability of failure")
	flag.IntVar(&rounds, "rounds", 100, "number of rounds to run for")
}

func main() {
	log.SetFlags(log.Lshortfile)
	flag.Parse()
	log.Println("RUNNING DEPLOY2DETER WITH RATE:", rate)
	os.MkdirAll("remote", 0777)
	var wg sync.WaitGroup
	// start building the necessary packages
	packages := []string{"../logserver", "../timeclient", "../exec", "../forkexec", "../deter"}
	for _, p := range packages {
		wg.Add(1)
		if p == "../deter" {
			go func(p string) {
				defer wg.Done()
				// the users node has a 386 FreeBSD architecture
				err := cliutils.Build(p, "386", "freebsd")
				if err != nil {
					log.Fatal(err)
				}
			}(p)
			continue
		}
		go func(p string) {
			defer wg.Done()
			// deter has an amd64, linux architecture
			err := cliutils.Build(p, "amd64", "linux")
			if err != nil {
				log.Fatal(err)
			}
		}(p)
	}
	// killssh processes on users
	cliutils.SshRunStdout("dvisher", "users.isi.deterlab.net", "killall ssh scp deter 2>/dev/null 1>/dev/null")
	// parse the hosts.txt file to create a separate list (and file)
	// of physical nodes and virtual nodes. Such that each host on line i, in phys.txt
	// corresponds to each host on line i, in virt.txt.
	physVirt, err := cliutils.ReadLines("hosts.txt")
	phys := make([]string, 0, len(physVirt)/2)
	virt := make([]string, 0, len(physVirt)/2)
	for i := 0; i < len(physVirt); i += 2 {
		phys = append(phys, physVirt[i])
		virt = append(virt, physVirt[i+1])
	}
	phys = phys[:36]
	virt = virt[:36]
	physOut := strings.Join(phys, "\n")
	virtOut := strings.Join(virt, "\n")

	err = ioutil.WriteFile("remote/phys.txt", []byte(physOut), 0666)
	if err != nil {
		log.Fatal("failed to write physical nodes file", err)
	}

	err = ioutil.WriteFile("remote/virt.txt", []byte(virtOut), 0666)
	if err != nil {
		log.Fatal("failed to write virtual nodes file", err)
	}

	masterLogger := phys[0]
	// slaveLogger1 := phys[1]
	// slaveLogger2 := phys[2]
	virt = virt[3:]
	phys = phys[3:]
	t, hostnames, depth, err := graphs.TreeFromList(virt, hpn, bf)
	log.Println("DEPTH:", depth)
	log.Println("TOTAL HOSTS:", len(hostnames))

	// wait for the build to finish
	wg.Wait()

	// copy the logserver directory to the current directory
	err = exec.Command("rsync", "-au", "../logserver", "remote/").Run()
	if err != nil {
		log.Fatal("error rsyncing logserver directory into remote directory:", err)
	}
	err = exec.Command("rsync", "-au", "remote/phys.txt", "remote/virt.txt", "remote/logserver/").Run()
	if err != nil {
		log.Fatal("error rsyncing phys, virt, and remote/logserver:", err)
	}
	err = os.Rename("logserver", "remote/logserver/logserver")
	if err != nil {
		log.Fatal("error renaming logserver:", err)
	}

	b, err := json.Marshal(t)
	if err != nil {
		log.Fatal("unable to generate tree from list")
	}
	err = ioutil.WriteFile("remote/logserver/cfg.json", b, 0660)
	if err != nil {
		log.Fatal("unable to write configuration file")
	}

	// NOTE: now remote/logserver is ready for transfer
	// it has logserver/ folder, binary, and cfg.json, and phys.txt, virt.txt

	// generate the configuration file from the tree
	cf := config.ConfigFromTree(t, hostnames)
	cfb, err := json.Marshal(cf)
	err = ioutil.WriteFile("remote/cfg.json", cfb, 0666)
	if err != nil {
		log.Fatal(err)
	}

	// scp the files that we need over to the boss node
	files := []string{"timeclient", "exec", "forkexec", "deter", "cfg.json"}
	for _, f := range files {
		cmd := exec.Command("rsync", "-au", f, "remote/")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		err := cmd.Run()
		if err != nil {
			log.Fatal("error unable to rsync file into remote directory:", err)
		}
	}
	err = cliutils.Rsync("dvisher", "users.isi.deterlab.net", "remote", "")
	if err != nil {
		log.Fatal(err)
	}

	// setup port forwarding for viewing log server
	// ssh -L 8080:pcXXX:80 username@users.isi.deterlab.net
	// ssh username@users.deterlab.net -L 8118:somenode.experiment.YourClass.isi.deterlab.net:80
	fmt.Println("setup port forwarding for master logger: ", masterLogger)
	cmd := exec.Command(
		"ssh",
		"-t",
		"-t",
		"dvisher@users.isi.deterlab.net",
		"-L",
		"8080:"+masterLogger+":10000")
	cmd.Start()
	if err != nil {
		log.Fatal("failed to setup portforwarding for logging server")
	}
	log.Println("runnning deter with nmsgs:", nmsgs)
	// run the deter lab boss nodes process
	// it will be responsible for forwarding the files and running the individual
	// timestamping servers
	log.Fatal(cliutils.SshRunStdout("dvisher", "users.isi.deterlab.net",
		"GOMAXPROCS=8 remote/deter -nmsgs="+strconv.Itoa(nmsgs)+
			" -hpn="+strconv.Itoa(hpn)+
			" -bf="+strconv.Itoa(bf)+
			" -rate="+strconv.Itoa(rate)+
			" -rounds="+strconv.Itoa(rounds)+
			" -debug="+strconv.FormatBool(debug)))
}
<commit_msg>fixed assumption about config file<commit_after>// deploy2deter is responsible for kicking off the deployment process
// for deterlab. Given a list of hostnames, it will create an overlay
// tree topology, using all but the last node. It will create multiple
// nodes per server and run timestamping processes. The last node is
// reserved for the logging server, which is forwarded to localhost:8080
//
// options are "bf" which specifies the branching factor
//
// 	and "hpn" which specifies the replicaiton factor: hosts per node
//
// Creates the following directory structure in remote:
// exec, timeclient, logserver/...,
// this way it can rsync the remove to each of the destinations
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"

	"github.com/dedis/prifi/coco/test/cliutils"
	"github.com/dedis/prifi/coco/test/config"
	"github.com/dedis/prifi/coco/test/graphs"
)

// bf is the branching factor of the tree that we want to build
var bf int

// hpn is the replication factor of hosts per node: how many hosts do we want per node
var hpn int

var nmsgs int
var debug bool
var rate int
var failures int

var rounds int

func init() {
	flag.IntVar(&bf, "bf", 2, "branching factor: default binary")
	flag.IntVar(&hpn, "hpn", 1, "hosts per node: default 1")
	flag.IntVar(&nmsgs, "nmsgs", 100, "number of messages per round")
	flag.IntVar(&rate, "rate", -1, "number of milliseconds between messages: if rate > 0 then used")
	flag.BoolVar(&debug, "debug", false, "run in debugging mode")
	flag.IntVar(&failures, "failures", 0, "percent showing per node probability of failure")
	flag.IntVar(&rounds, "rounds", 100, "number of rounds to run for")
}

func main() {
	log.SetFlags(log.Lshortfile)
	flag.Parse()
	log.Println("RUNNING DEPLOY2DETER WITH RATE:", rate)
	os.MkdirAll("remote", 0777)
	var wg sync.WaitGroup
	// start building the necessary packages
	packages := []string{"../logserver", "../timeclient", "../exec", "../forkexec", "../deter"}
	for _, p := range packages {
		wg.Add(1)
		if p == "../deter" {
			go func(p string) {
				defer wg.Done()
				// the users node has a 386 FreeBSD architecture
				err := cliutils.Build(p, "386", "freebsd")
				if err != nil {
					log.Fatal(err)
				}
			}(p)
			continue
		}
		go func(p string) {
			defer wg.Done()
			// deter has an amd64, linux architecture
			err := cliutils.Build(p, "amd64", "linux")
			if err != nil {
				log.Fatal(err)
			}
		}(p)
	}
	// killssh processes on users
	cliutils.SshRunStdout("dvisher", "users.isi.deterlab.net", "killall ssh scp deter 2>/dev/null 1>/dev/null")
	// parse the hosts.txt file to create a separate list (and file)
	// of physical nodes and virtual nodes. Such that each host on line i, in phys.txt
	// corresponds to each host on line i, in virt.txt.
	physVirt, err := cliutils.ReadLines("hosts.txt")
	phys := make([]string, 0, len(physVirt)/2)
	virt := make([]string, 0, len(physVirt)/2)
	for i := 0; i < len(physVirt); i += 2 {
		phys = append(phys, physVirt[i])
		virt = append(virt, physVirt[i+1])
	}
	phys = phys[:36]
	virt = virt[:36]
	physOut := strings.Join(phys, "\n")
	virtOut := strings.Join(virt, "\n")

	err = ioutil.WriteFile("remote/phys.txt", []byte(physOut), 0666)
	if err != nil {
		log.Fatal("failed to write physical nodes file", err)
	}

	err = ioutil.WriteFile("remote/virt.txt", []byte(virtOut), 0666)
	if err != nil {
		log.Fatal("failed to write virtual nodes file", err)
	}

	masterLogger := phys[0]
	// slaveLogger1 := phys[1]
	// slaveLogger2 := phys[2]
	virt = virt[3:]
	phys = phys[3:]
	t, hostnames, depth, err := graphs.TreeFromList(virt, hpn, bf)
	log.Println("DEPTH:", depth)
	log.Println("TOTAL HOSTS:", len(hostnames))

	// wait for the build to finish
	wg.Wait()

	// copy the logserver directory to the current directory
	err = exec.Command("rsync", "-au", "../logserver", "remote/").Run()
	if err != nil {
		log.Fatal("error rsyncing logserver directory into remote directory:", err)
	}
	err = exec.Command("rsync", "-au", "remote/phys.txt", "remote/virt.txt", "remote/logserver/").Run()
	if err != nil {
		log.Fatal("error rsyncing phys, virt, and remote/logserver:", err)
	}
	err = os.Rename("logserver", "remote/logserver/logserver")
	if err != nil {
		log.Fatal("error renaming logserver:", err)
	}

	b, err := json.Marshal(t)
	if err != nil {
		log.Fatal("unable to generate tree from list")
	}
	err = ioutil.WriteFile("remote/logserver/cfg.json", b, 0660)
	if err != nil {
		log.Fatal("unable to write configuration file")
	}

	// NOTE: now remote/logserver is ready for transfer
	// it has logserver/ folder, binary, and cfg.json, and phys.txt, virt.txt

	// generate the configuration file from the tree
	cf := config.ConfigFromTree(t, hostnames)
	cfb, err := json.Marshal(cf)
	err = ioutil.WriteFile("remote/cfg.json", cfb, 0666)
	if err != nil {
		log.Fatal(err)
	}

	// scp the files that we need over to the boss node
	files := []string{"timeclient", "exec", "forkexec", "deter"}
	for _, f := range files {
		cmd := exec.Command("rsync", "-au", f, "remote/")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		err := cmd.Run()
		if err != nil {
			log.Fatal("error unable to rsync file into remote directory:", err)
		}
	}
	err = cliutils.Rsync("dvisher", "users.isi.deterlab.net", "remote", "")
	if err != nil {
		log.Fatal(err)
	}

	// setup port forwarding for viewing log server
	// ssh -L 8080:pcXXX:80 username@users.isi.deterlab.net
	// ssh username@users.deterlab.net -L 8118:somenode.experiment.YourClass.isi.deterlab.net:80
	fmt.Println("setup port forwarding for master logger: ", masterLogger)
	cmd := exec.Command(
		"ssh",
		"-t",
		"-t",
		"dvisher@users.isi.deterlab.net",
		"-L",
		"8080:"+masterLogger+":10000")
	cmd.Start()
	if err != nil {
		log.Fatal("failed to setup portforwarding for logging server")
	}
	log.Println("runnning deter with nmsgs:", nmsgs)
	// run the deter lab boss nodes process
	// it will be responsible for forwarding the files and running the individual
	// timestamping servers
	log.Fatal(cliutils.SshRunStdout("dvisher", "users.isi.deterlab.net",
		"GOMAXPROCS=8 remote/deter -nmsgs="+strconv.Itoa(nmsgs)+
			" -hpn="+strconv.Itoa(hpn)+
			" -bf="+strconv.Itoa(bf)+
			" -rate="+strconv.Itoa(rate)+
			" -rounds="+strconv.Itoa(rounds)+
			" -debug="+strconv.FormatBool(debug)))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"reflect"
	"regexp"
	"strings"
	"sync/atomic"
	"time"
	"unsafe"

	r "github.com/dancannon/gorethink"
	"github.com/jaracil/ei"
	. "github.com/jaracil/nexus/log"
	"github.com/jaracil/nxcli/nxcore"
)

type LoginResponse struct {
	User string                            `json:"user"`
	Tags map[string]map[string]interface{} `json:"tags"`
}

func (nc *NexusConn) handleSysReq(req *JsonRpcReq) {
	switch req.Method {
	case "sys.ping":
		req.Result(ei.M{"ok": true})

	case "sys.version":
		req.Result(ei.M{"version": Version.String()})

	case "sys.watchdog":
		wdt, err := ei.N(req.Params).M("watchdog").Lower().Int64()
		if err == nil {
			if wdt < 10 {
				wdt = 10
			}
			atomic.StoreInt64(&nc.wdog, wdt)
		}
		req.Result(ei.M{"watchdog": nc.wdog})

	case "sys.login":
		var user string
		var mask map[string]map[string]interface{}

		// Auth
		method := ei.N(req.Params).M("method").Lower().StringZ()
		switch method {
		case "", "basic":
			var err int
			user, _, err = nc.BasicAuth(req.Params)
			if err != ErrNoError {
				req.Error(err, "", nil)
				return
			}

		default:
			nic := NewInternalClient()
			defer nic.Close()
			ret, err := nic.TaskPush(fmt.Sprintf("sys.login.%s.login", method), req.Params, time.Second*10, &nxcore.TaskOpts{})
			if err != nil {
				req.Error(ErrPermissionDenied, err.Error(), nil)
				return
			}

			var loginResponse LoginResponse
			b, err := json.Marshal(ret)
			if err != nil {
				req.Error(ErrInternal, "", nil)
				return
			}

			if err := json.Unmarshal(b, &loginResponse); err != nil {
				req.Error(ErrInternal, "", nil)
				return
			}
			user = loginResponse.User
			mask = loginResponse.Tags
		}

		// Check user limits

		ud, err := loadUserData(user)
		if err != ErrNoError {
			req.Error(err, "", nil)
			return
		}

		if !nc.checkUserLimits(ud) {
			req.Error(ErrPermissionDenied, "", nil)
			return
		}

		// LOGGED IN!
		atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&nc.user)), unsafe.Pointer(&UserData{
			User:        ud.User,
			Mask:        mask,
			Tags:        maskTags(ud.Tags, mask),
			MaxSessions: ud.MaxSessions,
			Whitelist:   ud.Whitelist,
			Blacklist:   ud.Blacklist,
		}))
		nc.updateSession()
		hook("user", ud.User, ud.User, ei.M{
			"action":      "login",
			"user":        nc.user.User,
			"mask":        nc.user.Mask,
			"tags":        nc.user.Tags,
			"maxSessions": nc.user.MaxSessions,
			"whitelist":   nc.user.Whitelist,
			"blacklist":   nc.user.Blacklist,
		})
		req.Result(ei.M{"ok": true, "user": nc.user.User, "connid": nc.connId})

	default:
		req.Error(ErrMethodNotFound, "", nil)
	}
}

// Return tags that are equal in both A and B
func maskTags(src map[string]map[string]interface{}, mask map[string]map[string]interface{}) (m map[string]map[string]interface{}) {
	m = make(map[string]map[string]interface{})

	if mask == nil {
		return src
	}

	for prefix, tags := range src {
		if bprefix, ok := mask[prefix]; ok {
			m[prefix] = make(map[string]interface{})
			for k, v := range tags {
				if vb, ok := bprefix[k]; ok && reflect.DeepEqual(v, vb) {
					m[prefix][k] = v
				}
			}
		}
	}
	return
}

func loadUserData(user string) (*UserData, int) {
	return loadUserDataWithTemplates(user, make(map[string]bool))
}

func loadUserDataWithTemplates(user string, loadedTemplates map[string]bool) (*UserData, int) {
	ud := &UserData{}
	cur, err := r.Table("users").Get(strings.ToLower(user)).Run(db)
	if err != nil {
		return nil, ErrInternal
	}
	defer cur.Close()
	err = cur.One(ud)
	if err != nil {
		if err == r.ErrEmptyResult {
			return nil, ErrPermissionDenied
		}
		return nil, ErrInternal
	}
	grabTemplates(ud, loadedTemplates)
	return ud, ErrNoError
}

func grabTemplates(ud *UserData, loadedTemplates map[string]bool) {
	for _, template := range ud.Templates {
		if loadedTemplates[template] {
			continue
		}
		loadedTemplates[template] = true

		tud, err := loadUserDataWithTemplates(template, loadedTemplates)
		if err != ErrNoError {
			continue
		}

		mergeTags(tud, ud)
	}
}

// Bring from src tags that did not exist on dst
func mergeTags(src, dst *UserData) {
	for prefix, tags := range src.Tags {
		for tag, val := range tags {
			if _, ok := dst.Tags[prefix]; !ok {
				dst.Tags[prefix] = make(map[string]interface{})
			}
			if _, ok := dst.Tags[prefix][tag]; !ok {
				dst.Tags[prefix][tag] = val
			}
		}
	}
}

func (nc *NexusConn) checkUserLimits(ud *UserData) bool {
	nci := NewInternalClient()
	defer nci.Close()

	// Max Sessions opened?
	// soft limit because race condition checking sessions
	sessions, err := nci.SessionList(ud.User, -1, 0)
	seslen := 0
	for _, u := range sessions {
		if u.User == ud.User {
			seslen = len(u.Sessions)
			break
		}
	}
	if err != nil || (ud.MaxSessions > 0 && seslen+1 > ud.MaxSessions) {
		Log.Warnf("User %s has too many sessions opened: %d/%d", ud.User, seslen, ud.MaxSessions)
		return false
	}

	remoteaddr := nc.conn.RemoteAddr().String()

	// Whitelisted?
	for _, wr := range ud.Whitelist {
		if match, err := regexp.MatchString(wr, remoteaddr); err == nil && match {
			Log.Warnf("User %s from %s whitelisted by %s", ud.User, remoteaddr, wr)
			return true
		}
	}

	// Blacklisted?
	for _, br := range ud.Blacklist {
		if match, err := regexp.MatchString(br, remoteaddr); err == nil && match {
			Log.Warnf("User %s from %s blacklisted by %s", ud.User, remoteaddr, br)
			return false
		}
	}

	return true
}

func (nc *NexusConn) BasicAuth(params interface{}) (string, map[string]map[string]interface{}, int) {
	user, err := ei.N(params).M("user").Lower().String()
	if err != nil {
		return "", nil, ErrInvalidParams
	}
	pass, err := ei.N(params).M("pass").String()
	if err != nil {
		return "", nil, ErrInvalidParams
	}
	var suser string
	split := strings.Split(user, ">")
	switch len(split) {
	case 1:
	case 2:
		if len(split[0]) > 0 && len(split[1]) > 0 {
			user = split[0]
			suser = split[1]
		} else {
			return "", nil, ErrInvalidParams

		}
	default:
		return "", nil, ErrInvalidParams
	}

	ud, rerr := loadUserData(user)
	if rerr != ErrNoError {
		return "", nil, rerr
	}
	dk, err := HashPass(pass, ud.Salt)
	if err != nil {
		return "", nil, ErrInternal
	}
	if ud.Pass != dk {
		return "", nil, ErrPermissionDenied
	}

	if suser != "" {
		tags := getTags(ud, suser)
		if !(ei.N(tags).M("@admin").BoolZ()) {
			return "", nil, ErrPermissionDenied
		}
		sud, err := loadUserData(suser)
		if err != ErrNoError {
			return "", nil, ErrPermissionDenied
		}
		return sud.User, sud.Tags, ErrNoError
	}

	return ud.User, ud.Tags, ErrNoError
}
<commit_msg>Initialize nil values loaded from rethink<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"reflect"
	"regexp"
	"strings"
	"sync/atomic"
	"time"
	"unsafe"

	r "github.com/dancannon/gorethink"
	"github.com/jaracil/ei"
	. "github.com/jaracil/nexus/log"
	"github.com/jaracil/nxcli/nxcore"
)

type LoginResponse struct {
	User string                            `json:"user"`
	Tags map[string]map[string]interface{} `json:"tags"`
}

func (nc *NexusConn) handleSysReq(req *JsonRpcReq) {
	switch req.Method {
	case "sys.ping":
		req.Result(ei.M{"ok": true})

	case "sys.version":
		req.Result(ei.M{"version": Version.String()})

	case "sys.watchdog":
		wdt, err := ei.N(req.Params).M("watchdog").Lower().Int64()
		if err == nil {
			if wdt < 10 {
				wdt = 10
			}
			atomic.StoreInt64(&nc.wdog, wdt)
		}
		req.Result(ei.M{"watchdog": nc.wdog})

	case "sys.login":
		var user string
		var mask map[string]map[string]interface{}

		// Auth
		method := ei.N(req.Params).M("method").Lower().StringZ()
		switch method {
		case "", "basic":
			var err int
			user, _, err = nc.BasicAuth(req.Params)
			if err != ErrNoError {
				req.Error(err, "", nil)
				return
			}

		default:
			nic := NewInternalClient()
			defer nic.Close()
			ret, err := nic.TaskPush(fmt.Sprintf("sys.login.%s.login", method), req.Params, time.Second*10, &nxcore.TaskOpts{})
			if err != nil {
				req.Error(ErrPermissionDenied, err.Error(), nil)
				return
			}

			var loginResponse LoginResponse
			b, err := json.Marshal(ret)
			if err != nil {
				req.Error(ErrInternal, "", nil)
				return
			}

			if err := json.Unmarshal(b, &loginResponse); err != nil {
				req.Error(ErrInternal, "", nil)
				return
			}
			user = loginResponse.User
			mask = loginResponse.Tags
		}

		// Check user limits

		ud, err := loadUserData(user)
		if err != ErrNoError {
			req.Error(err, "", nil)
			return
		}

		if !nc.checkUserLimits(ud) {
			req.Error(ErrPermissionDenied, "", nil)
			return
		}

		// LOGGED IN!
		atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&nc.user)), unsafe.Pointer(&UserData{
			User:        ud.User,
			Mask:        mask,
			Tags:        maskTags(ud.Tags, mask),
			MaxSessions: ud.MaxSessions,
			Whitelist:   ud.Whitelist,
			Blacklist:   ud.Blacklist,
		}))
		nc.updateSession()
		hook("user", ud.User, ud.User, ei.M{
			"action":      "login",
			"user":        nc.user.User,
			"mask":        nc.user.Mask,
			"tags":        nc.user.Tags,
			"maxSessions": nc.user.MaxSessions,
			"whitelist":   nc.user.Whitelist,
			"blacklist":   nc.user.Blacklist,
		})
		req.Result(ei.M{"ok": true, "user": nc.user.User, "connid": nc.connId})

	default:
		req.Error(ErrMethodNotFound, "", nil)
	}
}

// Return tags that are equal in both A and B
func maskTags(src map[string]map[string]interface{}, mask map[string]map[string]interface{}) (m map[string]map[string]interface{}) {
	m = make(map[string]map[string]interface{})

	if mask == nil {
		return src
	}

	for prefix, tags := range src {
		if bprefix, ok := mask[prefix]; ok {
			m[prefix] = make(map[string]interface{})
			for k, v := range tags {
				if vb, ok := bprefix[k]; ok && reflect.DeepEqual(v, vb) {
					m[prefix][k] = v
				}
			}
		}
	}
	return
}

func loadUserData(user string) (*UserData, int) {
	return loadUserDataWithTemplates(user, make(map[string]bool))
}

func loadUserDataWithTemplates(user string, loadedTemplates map[string]bool) (*UserData, int) {
	ud := &UserData{}
	cur, err := r.Table("users").Get(strings.ToLower(user)).Run(db)
	if err != nil {
		return nil, ErrInternal
	}
	defer cur.Close()
	err = cur.One(ud)
	if err != nil {
		if err == r.ErrEmptyResult {
			return nil, ErrPermissionDenied
		}
		return nil, ErrInternal
	}

	if ud.Tags == nil {
		ud.Tags = make(map[string]map[string]interface{})
	}

	if ud.Mask == nil {
		ud.Mask = make(map[string]map[string]interface{})
	}

	if ud.Blacklist == nil {
		ud.Blacklist = make([]string, 0)
	}

	if ud.Whitelist == nil {
		ud.Whitelist = make([]string, 0)
	}

	if ud.Templates == nil {
		ud.Templates = make([]string, 0)
	}

	grabTemplates(ud, loadedTemplates)
	return ud, ErrNoError
}

func grabTemplates(ud *UserData, loadedTemplates map[string]bool) {
	for _, template := range ud.Templates {
		if loadedTemplates[template] {
			continue
		}
		loadedTemplates[template] = true

		tud, err := loadUserDataWithTemplates(template, loadedTemplates)
		if err != ErrNoError {
			continue
		}

		mergeTags(tud, ud)
	}
}

// Bring from src tags that did not exist on dst
func mergeTags(src, dst *UserData) {
	for prefix, tags := range src.Tags {
		for tag, val := range tags {
			if _, ok := dst.Tags[prefix]; !ok {
				dst.Tags[prefix] = make(map[string]interface{})
			}
			if _, ok := dst.Tags[prefix][tag]; !ok {
				dst.Tags[prefix][tag] = val
			}
		}
	}
}

func (nc *NexusConn) checkUserLimits(ud *UserData) bool {
	nci := NewInternalClient()
	defer nci.Close()

	// Max Sessions opened?
	// soft limit because race condition checking sessions
	sessions, err := nci.SessionList(ud.User, -1, 0)
	seslen := 0
	for _, u := range sessions {
		if u.User == ud.User {
			seslen = len(u.Sessions)
			break
		}
	}
	if err != nil || (ud.MaxSessions > 0 && seslen+1 > ud.MaxSessions) {
		Log.Warnf("User %s has too many sessions opened: %d/%d", ud.User, seslen, ud.MaxSessions)
		return false
	}

	remoteaddr := nc.conn.RemoteAddr().String()

	// Whitelisted?
	for _, wr := range ud.Whitelist {
		if match, err := regexp.MatchString(wr, remoteaddr); err == nil && match {
			Log.Warnf("User %s from %s whitelisted by %s", ud.User, remoteaddr, wr)
			return true
		}
	}

	// Blacklisted?
	for _, br := range ud.Blacklist {
		if match, err := regexp.MatchString(br, remoteaddr); err == nil && match {
			Log.Warnf("User %s from %s blacklisted by %s", ud.User, remoteaddr, br)
			return false
		}
	}

	return true
}

func (nc *NexusConn) BasicAuth(params interface{}) (string, map[string]map[string]interface{}, int) {
	user, err := ei.N(params).M("user").Lower().String()
	if err != nil {
		return "", nil, ErrInvalidParams
	}
	pass, err := ei.N(params).M("pass").String()
	if err != nil {
		return "", nil, ErrInvalidParams
	}
	var suser string
	split := strings.Split(user, ">")
	switch len(split) {
	case 1:
	case 2:
		if len(split[0]) > 0 && len(split[1]) > 0 {
			user = split[0]
			suser = split[1]
		} else {
			return "", nil, ErrInvalidParams

		}
	default:
		return "", nil, ErrInvalidParams
	}

	ud, rerr := loadUserData(user)
	if rerr != ErrNoError {
		return "", nil, rerr
	}
	dk, err := HashPass(pass, ud.Salt)
	if err != nil {
		return "", nil, ErrInternal
	}
	if ud.Pass != dk {
		return "", nil, ErrPermissionDenied
	}

	if suser != "" {
		tags := getTags(ud, suser)
		if !(ei.N(tags).M("@admin").BoolZ()) {
			return "", nil, ErrPermissionDenied
		}
		sud, err := loadUserData(suser)
		if err != ErrNoError {
			return "", nil, ErrPermissionDenied
		}
		return sud.User, sud.Tags, ErrNoError
	}

	return ud.User, ud.Tags, ErrNoError
}
<|endoftext|>


<commit_before>package token

//Comma is true if v is a literal comma.
func (v Value) Comma() bool {
	return v.Kind == Literal && v.Value == ","
}

//Op is true if v is literal and first rune is one of "!@%&|-+=/<>"".
func (v Value) Op() bool {
	if v.Kind != Literal {
		return false
	}
	s := v.Value[0]
	for _, r := range "!@%&|-+=/<>*" {
		if byte(r) == s {
			return true
		}
	}
	return false
}

//Head is true if v is a literal whose canonical value is one of the list of
//valid literals for beginning a statement, as recognized by this SQLite superset.
func (v Value) Head(subquery bool) bool {
	if v.Kind != Literal {
		return false
	}
	lits := headLiterals[:]
	if subquery {
		lits = sqLiterals[:]
	}
	for _, h := range lits {
		if v.Canon == h {
			return true
		}
	}
	return false
}

var sqLiterals = [...]string{
	"IMPORT",
	"SELECT",
	"WITH",
}

var headLiterals = [...]string{
	"USE",
	"DISPLAY",
	"IMPORT",
	"SELECT",
	"INSERT",
	"UPDATE",
	"DELETE",
	"REPLACE",
	"WITH",
	"CREATE",
	"DROP",
	"REINDEX",
	"ALTER",
	"VACUUM",
	"ATTACH",
	"DETACH",
	"PRAGMA",
	"COMMIT",
	"ROLLBACK",
	"SAVEPOINT",
	"RELEASE",

	"ANALYZE", //verboten for now but handled specially for better errors
	"EXPLAIN",
}
<commit_msg>fixed error in sql unparser preventing dot from being handled correctly<commit_after>package token

//Comma is true if v is a literal comma.
func (v Value) Comma() bool {
	return v.Kind == Literal && v.Value == ","
}

//Op is true if v is literal and first rune is one of "!@%&|-+=/<>"".
func (v Value) Op() bool {
	if v.Kind != Literal {
		return false
	}
	s := v.Value[0]
	for _, r := range "!@%&|-+=/<>*." {
		if byte(r) == s {
			return true
		}
	}
	return false
}

//Head is true if v is a literal whose canonical value is one of the list of
//valid literals for beginning a statement, as recognized by this SQLite superset.
func (v Value) Head(subquery bool) bool {
	if v.Kind != Literal {
		return false
	}
	lits := headLiterals[:]
	if subquery {
		lits = sqLiterals[:]
	}
	for _, h := range lits {
		if v.Canon == h {
			return true
		}
	}
	return false
}

var sqLiterals = [...]string{
	"IMPORT",
	"SELECT",
	"WITH",
}

var headLiterals = [...]string{
	"USE",
	"DISPLAY",
	"IMPORT",
	"SELECT",
	"INSERT",
	"UPDATE",
	"DELETE",
	"REPLACE",
	"WITH",
	"CREATE",
	"DROP",
	"REINDEX",
	"ALTER",
	"VACUUM",
	"ATTACH",
	"DETACH",
	"PRAGMA",
	"COMMIT",
	"ROLLBACK",
	"SAVEPOINT",
	"RELEASE",

	"ANALYZE", //verboten for now but handled specially for better errors
	"EXPLAIN",
}
<|endoftext|>