<commit_before>package db_test

import (
	"database/sql"
	"strconv"
	"testing"

	"github.com/tomogoma/authms/db"
	testingH "github.com/tomogoma/authms/testing"
	"github.com/tomogoma/go-commons/database/cockroach"
)

var isInit bool

func setup(t *testing.T) cockroach.DSN {
	conf := testingH.ReadConfig(t)
	conf.Database.DB = conf.Database.DB + "_test"
	if !isInit {
		rdb := getDB(t, conf.Database)
		_, err := rdb.Exec("DROP DATABASE IF EXISTS " + conf.Database.DB)
		if err != nil {
			t.Fatalf("Error setting up: deleting db: %v", err)
		}
		isInit = true
	}
	return conf.Database
}

func tearDown(t *testing.T, conf cockroach.DSN) {
	rdb := getDB(t, conf)
	if _, err := rdb.Exec("SET DATABASE=" + conf.DB); err != nil {
		return
	}
	for i := len(db.AllTableNames) - 1; i >= 0; i-- {
		_, err := rdb.Exec("DELETE FROM " + db.AllTableNames[i])
		if err != nil {
			t.Fatalf("Error tearing down: delete %s: %v",
				db.AllTableNames[i], err)
		}
	}
}

func TestNewRoach(t *testing.T) {
	conf := setup(t)
	defer tearDown(t, conf)
	tt := []struct {
		name   string
		opts   []db.Option
		expErr bool
	}{
		{
			name: "valid",
			opts: []db.Option{
				db.WithDBName(conf.DBName()),
				db.WithDSN(conf.FormatDSN()),
			},
			expErr: false,
		},
		{
			name:   "valid (no options)",
			opts:   nil,
			expErr: false,
		},
	}
	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			r := db.NewRoach(tc.opts...)
			if r == nil {
				t.Fatalf("Got nil roach")
			}
		})
	}
}

func TestRoach_InitDBIfNot(t *testing.T) {

	conf := setup(t)
	defer tearDown(t, conf)
	r := newRoach(t, conf)
	rdb := getDB(t, conf)
	if err := r.InitDBIfNot(); err != nil {
		t.Fatalf("Initial init call failed: %v", err)
	}

	tt := []struct {
		name       string
		hasVersion bool
		version    []byte
		expErr     bool
	}{
		{
			name:       "first use",
			hasVersion: false,
			expErr:     false,
		},
		{
			name:       "versions equal",
			hasVersion: true,
			version:    []byte(strconv.Itoa(db.Version)),
			expErr:     false,
		},
		{
			name:       "db version smaller",
			hasVersion: true,
			version:    []byte(strconv.Itoa(db.Version - 1)),
			expErr:     true,
		},
		{
			name:       "db version bigger",
			hasVersion: true,
			version:    []byte(strconv.Itoa(db.Version + 1)),
			expErr:     true,
		},
	}

	cols := db.ColDesc(db.ColKey, db.ColValue, db.ColUpdateDate)
	upsertQ := `UPSERT INTO ` + db.TblConfigurations + ` (` + cols + `)
					VALUES ('db.version', $1, CURRENT_TIMESTAMP)`
	delQ := `DELETE FROM ` + db.TblConfigurations + ` WHERE ` + db.ColKey + `='db.version'`

	for _, tc := range tt {
		if _, err := rdb.Exec(delQ); err != nil {
			t.Fatalf("Error setting up: clear previous config: %v", err)
		}
		if tc.hasVersion {
			if _, err := rdb.Exec(upsertQ, tc.version); err != nil {
				t.Fatalf("Error setting up: insert test config: %v", err)
			}
		}
		t.Run(tc.name, func(t *testing.T) {
			r = newRoach(t, conf)
			err := r.InitDBIfNot()
			if tc.expErr {
				if err == nil {
					t.Fatalf("Expected an error, got nil")
				}
				// set db to have correct version (init error should be cached not queried)
				if _, err := rdb.Exec(upsertQ, []byte(strconv.Itoa(db.Version))); err != nil {
					t.Fatalf("Error setting up: insert test config: %v", err)
				}
				if err := r.InitDBIfNot(); err == nil {
					t.Fatalf("Subsequent init db not returning error")
				}
				return
			}
			if err != nil {
				t.Fatalf("Got an error: %v", err)
			}
			// set db to have incorrect version (isInit flag should be cached, not queried)
			if _, err := rdb.Exec(upsertQ, []byte(strconv.Itoa(db.Version+10))); err != nil {
				t.Fatalf("Error setting up: insert test config: %v", err)
			}
			if err = r.InitDBIfNot(); err != nil {
				t.Fatalf("Subsequent init not working")
			}
		})
	}
}

func newRoach(t *testing.T, conf cockroach.DSN) *db.Roach {
	r := db.NewRoach(
		db.WithDBName(conf.DBName()),
		db.WithDSN(conf.FormatDSN()),
	)
	if r == nil {
		t.Fatalf("Got nil roach")
	}
	return r
}

func getDB(t *testing.T, conf cockroach.DSN) *sql.DB {
	DB, err := cockroach.DBConn(conf)
	if err != nil {
		t.Fatalf("unable to tear down: cockroach.DBConn(): %s", err)
	}
	return DB
}
<commit_msg>improve test speed (setup)<commit_after>package db_test

import (
	"database/sql"
	"strconv"
	"testing"

	"github.com/tomogoma/authms/db"
	testingH "github.com/tomogoma/authms/testing"
	"github.com/tomogoma/go-commons/database/cockroach"
	"github.com/tomogoma/go-commons/errors"
)

var isInit bool

func setup(t *testing.T) cockroach.DSN {
	conf := testingH.ReadConfig(t)
	conf.Database.DB = conf.Database.DB + "_test"
	if !isInit {
		err := delAllTables(getDB(t, conf.Database), conf.Database.DBName())
		if err != nil {
			t.Fatalf("Error setting up: delete all tables: %v", err)
		}
		isInit = true
	}
	return conf.Database
}

func tearDown(t *testing.T, conf cockroach.DSN) {
	if err := delAllTables(getDB(t, conf), conf.DBName()); err != nil {
		t.Fatalf("Error tearing down: delete all tables: %v", err)
	}
}

func delAllTables(rdb *sql.DB, dbName string) error {
	if _, err := rdb.Exec("SET DATABASE=" + dbName); err != nil {
		return nil
	}
	for i := len(db.AllTableNames) - 1; i >= 0; i-- {
		_, err := rdb.Exec("DELETE FROM " + db.AllTableNames[i])
		if err != nil {
			return errors.Newf("delete %s: %v", db.AllTableNames[i], err)
		}
	}
	return nil
}

func TestNewRoach(t *testing.T) {
	conf := setup(t)
	defer tearDown(t, conf)
	tt := []struct {
		name   string
		opts   []db.Option
		expErr bool
	}{
		{
			name: "valid",
			opts: []db.Option{
				db.WithDBName(conf.DBName()),
				db.WithDSN(conf.FormatDSN()),
			},
			expErr: false,
		},
		{
			name:   "valid (no options)",
			opts:   nil,
			expErr: false,
		},
	}
	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			r := db.NewRoach(tc.opts...)
			if r == nil {
				t.Fatalf("Got nil roach")
			}
		})
	}
}

func TestRoach_InitDBIfNot(t *testing.T) {

	conf := setup(t)
	defer tearDown(t, conf)
	r := newRoach(t, conf)
	rdb := getDB(t, conf)
	if err := r.InitDBIfNot(); err != nil {
		t.Fatalf("Initial init call failed: %v", err)
	}

	tt := []struct {
		name       string
		hasVersion bool
		version    []byte
		expErr     bool
	}{
		{
			name:       "first use",
			hasVersion: false,
			expErr:     false,
		},
		{
			name:       "versions equal",
			hasVersion: true,
			version:    []byte(strconv.Itoa(db.Version)),
			expErr:     false,
		},
		{
			name:       "db version smaller",
			hasVersion: true,
			version:    []byte(strconv.Itoa(db.Version - 1)),
			expErr:     true,
		},
		{
			name:       "db version bigger",
			hasVersion: true,
			version:    []byte(strconv.Itoa(db.Version + 1)),
			expErr:     true,
		},
	}

	cols := db.ColDesc(db.ColKey, db.ColValue, db.ColUpdateDate)
	upsertQ := `UPSERT INTO ` + db.TblConfigurations + ` (` + cols + `)
					VALUES ('db.version', $1, CURRENT_TIMESTAMP)`
	delQ := `DELETE FROM ` + db.TblConfigurations + ` WHERE ` + db.ColKey + `='db.version'`

	for _, tc := range tt {
		if _, err := rdb.Exec(delQ); err != nil {
			t.Fatalf("Error setting up: clear previous config: %v", err)
		}
		if tc.hasVersion {
			if _, err := rdb.Exec(upsertQ, tc.version); err != nil {
				t.Fatalf("Error setting up: insert test config: %v", err)
			}
		}
		t.Run(tc.name, func(t *testing.T) {
			r = newRoach(t, conf)
			err := r.InitDBIfNot()
			if tc.expErr {
				if err == nil {
					t.Fatalf("Expected an error, got nil")
				}
				// set db to have correct version (init error should be cached not queried)
				if _, err := rdb.Exec(upsertQ, []byte(strconv.Itoa(db.Version))); err != nil {
					t.Fatalf("Error setting up: insert test config: %v", err)
				}
				if err := r.InitDBIfNot(); err == nil {
					t.Fatalf("Subsequent init db not returning error")
				}
				return
			}
			if err != nil {
				t.Fatalf("Got an error: %v", err)
			}
			// set db to have incorrect version (isInit flag should be cached, not queried)
			if _, err := rdb.Exec(upsertQ, []byte(strconv.Itoa(db.Version+10))); err != nil {
				t.Fatalf("Error setting up: insert test config: %v", err)
			}
			if err = r.InitDBIfNot(); err != nil {
				t.Fatalf("Subsequent init not working")
			}
		})
	}
}

func newRoach(t *testing.T, conf cockroach.DSN) *db.Roach {
	r := db.NewRoach(
		db.WithDBName(conf.DBName()),
		db.WithDSN(conf.FormatDSN()),
	)
	if r == nil {
		t.Fatalf("Got nil roach")
	}
	return r
}

func getDB(t *testing.T, conf cockroach.DSN) *sql.DB {
	DB, err := cockroach.DBConn(conf)
	if err != nil {
		t.Fatalf("unable to tear down: cockroach.DBConn(): %s", err)
	}
	return DB
}
<|endoftext|>


<commit_before>package main

import (
    "flag"
    "fmt"
    "log"
    "os"
    "math/big"
    "time"

    "github.com/tooru/pnbot/prime"
    "github.com/dghubble/go-twitter/twitter"
    "github.com/dghubble/oauth1"
)

const (
    queueSize = 10
    maxRetry = 10
)

func main() {
    log.SetFlags(log.Ldate | log.Ltime | log.Llongfile)

    flags := flag.NewFlagSet("pnbot", flag.ExitOnError)
    consumerKey := flags.String("ck", "", "Twitter Consumer Key")
    consumerSecret := flags.String("cs", "", "Twitter Consumer Secret")
    accessToken := flags.String("at", "", "Twitter Access Token")
    accessSecret := flags.String("as", "", "Twitter Access Secret")
    flags.Parse(os.Args[1:])

    if *consumerKey == "" || *consumerSecret == "" || *accessToken == "" || *accessSecret == "" {
        log.Fatal("Consumer key/secret and Access token/secret required")
    }

    pnbot := NewPNBot(consumerKey, consumerSecret, accessToken, accessSecret)
    pnbot.Start()
}

func NewPNBot(consumerKey *string, consumerSecret *string,
              accessToken *string, accessSecret *string) *PNBot {
    pnbot := &PNBot{
        consumerKey: consumerKey,
        consumerSecret: consumerSecret,
        accessToken: accessToken,
        accessSecret: accessSecret,
        prime: prime.NewPrime(),
        ch: make(chan *big.Int, queueSize),
    }

    return pnbot
}

type PNBot struct {
    consumerKey *string
    consumerSecret *string
    accessToken *string
    accessSecret *string

    prime *prime.Prime
    ch chan *big.Int

    client *twitter.Client

}

func (pnbot *PNBot) Start() error {
    pnbot.client = pnbot.newClient()

    maxPrime, err := pnbot.getMaxPrime()
    if err != nil {
        log.Printf("getMaxPrime: %v\n", err)
        return err
    }

    go pnbot.makePrimes(maxPrime)

    return pnbot.tweetPrimes()
}

func (pnbot *PNBot) newClient() *twitter.Client {
    config := oauth1.NewConfig(*pnbot.consumerKey, *pnbot.consumerSecret)
    token := oauth1.NewToken(*pnbot.accessToken, *pnbot.accessSecret)
    httpClient := config.Client(oauth1.NoContext, token)

    return twitter.NewClient(httpClient)
}

func (pnbot *PNBot) getMaxPrime() (prime *big.Int, err error) {
    tweets, _, err := pnbot.client.Timelines.HomeTimeline(&twitter.HomeTimelineParams{
        SinceID: 0,
    })
    if err != nil {
        return nil, err
    }

    var maxPrime big.Int

    for _, tweet := range tweets {
        if _, ok := maxPrime.SetString(tweet.Text, 10); ok {
            log.Printf("continue from '%s'", tweet.Text)
            return &maxPrime, nil
        }
        log.Printf("skip '%v'", tweet.Text)
    }
    maxPrime.SetInt64(0)
    return &maxPrime, nil
}

func (pnbot *PNBot) makePrimes(maxPrime *big.Int) {
    var prime *big.Int
    var err error

    for {
        prime, err = pnbot.prime.Next()
        if err != nil {
            log.Fatalf("makePrimes: %v\n", err)
            pnbot.ch <- nil
            return
        }

        if prime.Cmp(maxPrime) > 0 {
            break;
        }
        log.Printf("makePrimes: skip %v\n", prime)
    }

    for {
        log.Printf("makePrimes: found %v\n", prime)
        pnbot.ch <- prime

        prime, err = pnbot.prime.Next()
        if err != nil {
            log.Fatalf("makePrimes: %v\n", err)
            pnbot.ch <- nil
            return
        }
    }
}

func (pnbot *PNBot) tweetPrimes() error {
    interval := 1 * time.Second

    for {
        prime := <- pnbot.ch
        if prime == nil {
            log.Fatalf("prime generator was died")
            return fmt.Errorf("prime generator was died")
        }

        text := prime.Text(10)
        retry := 0
        retryInterval := 10 * time.Second
        for {
            _, _, err := pnbot.client.Statuses.Update(text, nil)
            if err == nil {
                break
            }
            if retry >= maxRetry {
                return fmt.Errorf("Too many tweet error: %v\n", err)
            }
            log.Printf("Tweet error[%d/%d]:sleep=%s: %v\n", retry+1, maxRetry, retryInterval, err)

            time.Sleep(interval)
            pnbot.client = pnbot.newClient()
            retry++
            retryInterval *= 2
            continue
        }
        if retry > 0 {
            interval += 1 * time.Second
        }

        log.Printf("tweet %s\n", text)
        time.Sleep(interval)
    }
}
<commit_msg>Modified pnbot.go.<commit_after>package main

import (
    "flag"
    "fmt"
    "log"
    "os"
    "math/big"
    "time"

    "github.com/tooru/pnbot/prime"
    "github.com/dghubble/go-twitter/twitter"
    "github.com/dghubble/oauth1"
)

const (
    queueSize = 10
    maxRetry = 10
)

func main() {
    log.SetFlags(log.Ldate | log.Ltime | log.Llongfile)

    flags := flag.NewFlagSet("pnbot", flag.ExitOnError)
    consumerKey := flags.String("ck", "", "Twitter Consumer Key")
    consumerSecret := flags.String("cs", "", "Twitter Consumer Secret")
    accessToken := flags.String("at", "", "Twitter Access Token")
    accessSecret := flags.String("as", "", "Twitter Access Secret")
    flags.Parse(os.Args[1:])

    if *consumerKey == "" || *consumerSecret == "" || *accessToken == "" || *accessSecret == "" {
        log.Fatal("Consumer key/secret and Access token/secret required")
    }

    pnbot := NewPNBot(consumerKey, consumerSecret, accessToken, accessSecret)
    pnbot.Start()
}

func NewPNBot(consumerKey *string, consumerSecret *string,
              accessToken *string, accessSecret *string) *PNBot {
    pnbot := &PNBot{
        consumerKey: consumerKey,
        consumerSecret: consumerSecret,
        accessToken: accessToken,
        accessSecret: accessSecret,
        prime: prime.NewPrime(),
        ch: make(chan *big.Int, queueSize),
    }

    return pnbot
}

type PNBot struct {
    consumerKey *string
    consumerSecret *string
    accessToken *string
    accessSecret *string

    prime *prime.Prime
    ch chan *big.Int

    client *twitter.Client

}

func (pnbot *PNBot) Start() error {
    pnbot.client = pnbot.newClient()

    maxPrime, err := pnbot.getMaxPrime()
    if err != nil {
        log.Printf("getMaxPrime: %v\n", err)
        return err
    }

    go pnbot.makePrimes(maxPrime)

    return pnbot.tweetPrimes()
}

func (pnbot *PNBot) newClient() *twitter.Client {
    config := oauth1.NewConfig(*pnbot.consumerKey, *pnbot.consumerSecret)
    token := oauth1.NewToken(*pnbot.accessToken, *pnbot.accessSecret)
    httpClient := config.Client(oauth1.NoContext, token)

    return twitter.NewClient(httpClient)
}

func (pnbot *PNBot) getMaxPrime() (prime *big.Int, err error) {
    tweets, _, err := pnbot.client.Timelines.HomeTimeline(&twitter.HomeTimelineParams{
        SinceID: 0,
    })
    if err != nil {
        return nil, err
    }

    var maxPrime big.Int

    for _, tweet := range tweets {
        if _, ok := maxPrime.SetString(tweet.Text, 10); ok {
            log.Printf("continue from '%s'", tweet.Text)
            return &maxPrime, nil
        }
        log.Printf("skip '%v'", tweet.Text)
    }
    maxPrime.SetInt64(0)
    return &maxPrime, nil
}

func (pnbot *PNBot) makePrimes(maxPrime *big.Int) {
    var prime *big.Int
    var err error

    for {
        prime, err = pnbot.prime.Next()
        if err != nil {
            log.Fatalf("makePrimes: %v\n", err)
            pnbot.ch <- nil
            return
        }

        if prime.Cmp(maxPrime) > 0 {
            break;
        }
        log.Printf("makePrimes: skip %v\n", prime)
    }

    for {
        log.Printf("makePrimes: found %v\n", prime)
        pnbot.ch <- prime

        prime, err = pnbot.prime.Next()
        if err != nil {
            log.Fatalf("makePrimes: %v\n", err)
            pnbot.ch <- nil
            return
        }
    }
}

func (pnbot *PNBot) tweetPrimes() error {
    interval := 1 * time.Second

    for {
        prime := <- pnbot.ch
        if prime == nil {
            log.Fatalf("prime generator was died")
            return fmt.Errorf("prime generator was died")
        }

        text := prime.Text(10)
        retry := 0
        retryInterval := 10 * time.Second
        for {
            _, _, err := pnbot.client.Statuses.Update(text, nil)
            if err == nil {
                break
            }
            if retry >= maxRetry {
                return fmt.Errorf("Too many tweet error: %v\n", err)
            }
            log.Printf("Tweet error[%d/%d]:sleep=%s: %v\n", retry+1, maxRetry, retryInterval, err)

            time.Sleep(retryInterval)
            pnbot.client = pnbot.newClient()
            retry++
            retryInterval *= 2
            continue
        }
        if retry > 0 {
            interval += 1 * time.Second
        }

        log.Printf("tweet %s :sleep=%s\n", text, interval)
        time.Sleep(interval)
    }
}
<|endoftext|>


<commit_before>package inmemory

import (
	"fmt"
	"io"
	"path"
	"sort"
	"strings"
	"time"
)

var (
	errExists    = fmt.Errorf("exists")
	errNotExists = fmt.Errorf("notexists")
	errIsNotDir  = fmt.Errorf("notdir")
	errIsDir     = fmt.Errorf("isdir")
)

type node interface {
	name() string
	path() string
	isdir() bool
	modtime() time.Time
}

// dir is the central type for the memory-based  storagedriver. All operations
// are dispatched from a root dir.
type dir struct {
	common

	// TODO(stevvooe): Use sorted slice + search.
	children map[string]node
}

var _ node = &dir{}

func (d *dir) isdir() bool {
	return true
}

// add places the node n into dir d.
func (d *dir) add(n node) {
	if d.children == nil {
		d.children = make(map[string]node)
	}

	d.children[n.name()] = n
	d.mod = time.Now()
}

// find searches for the node, given path q in dir. If the node is found, it
// will be returned. If the node is not found, the closet existing parent. If
// the node is found, the returned (node).path() will match q.
func (d *dir) find(q string) node {
	q = strings.Trim(q, "/")
	i := strings.Index(q, "/")

	if q == "" {
		return d
	}

	if i == 0 {
		panic("shouldn't happen, no root paths")
	}

	var component string
	if i < 0 {
		// No more path components
		component = q
	} else {
		component = q[:i]
	}

	child, ok := d.children[component]
	if !ok {
		// Node was not found. Return p and the current node.
		return d
	}

	if child.isdir() {
		// traverse down!
		q = q[i+1:]
		return child.(*dir).find(q)
	}

	return child
}

func (d *dir) list(p string) ([]string, error) {
	n := d.find(p)

	if n.path() != p {
		return nil, errNotExists
	}

	if !n.isdir() {
		return nil, errIsNotDir
	}

	var children []string
	for _, child := range n.(*dir).children {
		children = append(children, child.path())
	}

	sort.Strings(children)
	return children, nil
}

// mkfile or return the existing one. returns an error if it exists and is a
// directory. Essentially, this is open or create.
func (d *dir) mkfile(p string) (*file, error) {
	n := d.find(p)
	if n.path() == p {
		if n.isdir() {
			return nil, errIsDir
		}

		return n.(*file), nil
	}

	dirpath, filename := path.Split(p)
	// Make any non-existent directories
	n, err := d.mkdirs(dirpath)
	if err != nil {
		return nil, err
	}

	dd := n.(*dir)
	n = &file{
		common: common{
			p:   path.Join(dd.path(), filename),
			mod: time.Now(),
		},
	}

	dd.add(n)
	return n.(*file), nil
}

// mkdirs creates any missing directory entries in p and returns the result.
func (d *dir) mkdirs(p string) (*dir, error) {
	p = normalize(p)

	n := d.find(p)

	if !n.isdir() {
		// Found something there
		return nil, errIsNotDir
	}

	if n.path() == p {
		return n.(*dir), nil
	}

	dd := n.(*dir)

	relative := strings.Trim(strings.TrimPrefix(p, n.path()), "/")

	if relative == "" {
		return dd, nil
	}

	components := strings.Split(relative, "/")
	for _, component := range components {
		d, err := dd.mkdir(component)

		if err != nil {
			// This should actually never happen, since there are no children.
			return nil, err
		}
		dd = d
	}

	return dd, nil
}

// mkdir creates a child directory under d with the given name.
func (d *dir) mkdir(name string) (*dir, error) {
	if name == "" {
		return nil, fmt.Errorf("invalid dirname")
	}

	_, ok := d.children[name]
	if ok {
		return nil, errExists
	}

	child := &dir{
		common: common{
			p:   path.Join(d.path(), name),
			mod: time.Now(),
		},
	}
	d.add(child)
	d.mod = time.Now()

	return child, nil
}

func (d *dir) move(src, dst string) error {
	dstDirname, _ := path.Split(dst)

	dp, err := d.mkdirs(dstDirname)
	if err != nil {
		return err
	}

	srcDirname, srcFilename := path.Split(src)
	sp := d.find(srcDirname)

	if normalize(srcDirname) != normalize(sp.path()) {
		return errNotExists
	}

	spd, ok := sp.(*dir)
	if !ok {
		return errIsNotDir // paranoid.
	}

	s, ok := spd.children[srcFilename]
	if !ok {
		return errNotExists
	}

	delete(spd.children, srcFilename)

	switch n := s.(type) {
	case *dir:
		n.p = dst
	case *file:
		n.p = dst
	}

	dp.add(s)

	return nil
}

func (d *dir) delete(p string) error {
	dirname, filename := path.Split(p)
	parent := d.find(dirname)

	if normalize(dirname) != normalize(parent.path()) {
		return errNotExists
	}

	parentDir, ok := parent.(*dir)
	if !ok {
		return errIsNotDir
	}

	if _, ok := parentDir.children[filename]; !ok {
		return errNotExists
	}

	delete(parentDir.children, filename)
	return nil
}

func (d *dir) String() string {
	return fmt.Sprintf("&dir{path: %v, children: %v}", d.p, d.children)
}

// file stores actual data in the fs tree. It acts like an open, seekable file
// where operations are conducted through ReadAt and WriteAt. Use it with
// SectionReader for the best effect.
type file struct {
	common
	data []byte
}

var _ node = &file{}

func (f *file) isdir() bool {
	return false
}

func (f *file) truncate() {
	f.data = f.data[:0]
}

func (f *file) sectionReader(offset int64) io.Reader {
	return io.NewSectionReader(f, offset, int64(len(f.data))-offset)
}

func (f *file) ReadAt(p []byte, offset int64) (n int, err error) {
	return copy(p, f.data[offset:]), nil
}

func (f *file) WriteAt(p []byte, offset int64) (n int, err error) {
	off := int(offset)
	if cap(f.data) < off+len(p) {
		data := make([]byte, len(f.data), off+len(p))
		copy(data, f.data)
		f.data = data
	}

	f.mod = time.Now()
	f.data = f.data[:off+len(p)]

	return copy(f.data[off:off+len(p)], p), nil
}

func (f *file) String() string {
	return fmt.Sprintf("&file{path: %q}", f.p)
}

// common provides shared fields and methods for node implementations.
type common struct {
	p   string
	mod time.Time
}

func (c *common) name() string {
	_, name := path.Split(c.p)
	return name
}

func (c *common) path() string {
	return c.p
}

func (c *common) modtime() time.Time {
	return c.mod
}

func normalize(p string) string {
	return "/" + strings.Trim(p, "/")
}
<commit_msg>Fix panic in inmemory driver<commit_after>package inmemory

import (
	"fmt"
	"io"
	"path"
	"sort"
	"strings"
	"time"
)

var (
	errExists    = fmt.Errorf("exists")
	errNotExists = fmt.Errorf("notexists")
	errIsNotDir  = fmt.Errorf("notdir")
	errIsDir     = fmt.Errorf("isdir")
)

type node interface {
	name() string
	path() string
	isdir() bool
	modtime() time.Time
}

// dir is the central type for the memory-based  storagedriver. All operations
// are dispatched from a root dir.
type dir struct {
	common

	// TODO(stevvooe): Use sorted slice + search.
	children map[string]node
}

var _ node = &dir{}

func (d *dir) isdir() bool {
	return true
}

// add places the node n into dir d.
func (d *dir) add(n node) {
	if d.children == nil {
		d.children = make(map[string]node)
	}

	d.children[n.name()] = n
	d.mod = time.Now()
}

// find searches for the node, given path q in dir. If the node is found, it
// will be returned. If the node is not found, the closet existing parent. If
// the node is found, the returned (node).path() will match q.
func (d *dir) find(q string) node {
	q = strings.Trim(q, "/")
	i := strings.Index(q, "/")

	if q == "" {
		return d
	}

	if i == 0 {
		panic("shouldn't happen, no root paths")
	}

	var component string
	if i < 0 {
		// No more path components
		component = q
	} else {
		component = q[:i]
	}

	child, ok := d.children[component]
	if !ok {
		// Node was not found. Return p and the current node.
		return d
	}

	if child.isdir() {
		// traverse down!
		q = q[i+1:]
		return child.(*dir).find(q)
	}

	return child
}

func (d *dir) list(p string) ([]string, error) {
	n := d.find(p)

	if n.path() != p {
		return nil, errNotExists
	}

	if !n.isdir() {
		return nil, errIsNotDir
	}

	var children []string
	for _, child := range n.(*dir).children {
		children = append(children, child.path())
	}

	sort.Strings(children)
	return children, nil
}

// mkfile or return the existing one. returns an error if it exists and is a
// directory. Essentially, this is open or create.
func (d *dir) mkfile(p string) (*file, error) {
	n := d.find(p)
	if n.path() == p {
		if n.isdir() {
			return nil, errIsDir
		}

		return n.(*file), nil
	}

	dirpath, filename := path.Split(p)
	// Make any non-existent directories
	n, err := d.mkdirs(dirpath)
	if err != nil {
		return nil, err
	}

	dd := n.(*dir)
	n = &file{
		common: common{
			p:   path.Join(dd.path(), filename),
			mod: time.Now(),
		},
	}

	dd.add(n)
	return n.(*file), nil
}

// mkdirs creates any missing directory entries in p and returns the result.
func (d *dir) mkdirs(p string) (*dir, error) {
	p = normalize(p)

	n := d.find(p)

	if !n.isdir() {
		// Found something there
		return nil, errIsNotDir
	}

	if n.path() == p {
		return n.(*dir), nil
	}

	dd := n.(*dir)

	relative := strings.Trim(strings.TrimPrefix(p, n.path()), "/")

	if relative == "" {
		return dd, nil
	}

	components := strings.Split(relative, "/")
	for _, component := range components {
		d, err := dd.mkdir(component)

		if err != nil {
			// This should actually never happen, since there are no children.
			return nil, err
		}
		dd = d
	}

	return dd, nil
}

// mkdir creates a child directory under d with the given name.
func (d *dir) mkdir(name string) (*dir, error) {
	if name == "" {
		return nil, fmt.Errorf("invalid dirname")
	}

	_, ok := d.children[name]
	if ok {
		return nil, errExists
	}

	child := &dir{
		common: common{
			p:   path.Join(d.path(), name),
			mod: time.Now(),
		},
	}
	d.add(child)
	d.mod = time.Now()

	return child, nil
}

func (d *dir) move(src, dst string) error {
	dstDirname, _ := path.Split(dst)

	dp, err := d.mkdirs(dstDirname)
	if err != nil {
		return err
	}

	srcDirname, srcFilename := path.Split(src)
	sp := d.find(srcDirname)

	if normalize(srcDirname) != normalize(sp.path()) {
		return errNotExists
	}

	spd, ok := sp.(*dir)
	if !ok {
		return errIsNotDir // paranoid.
	}

	s, ok := spd.children[srcFilename]
	if !ok {
		return errNotExists
	}

	delete(spd.children, srcFilename)

	switch n := s.(type) {
	case *dir:
		n.p = dst
	case *file:
		n.p = dst
	}

	dp.add(s)

	return nil
}

func (d *dir) delete(p string) error {
	dirname, filename := path.Split(p)
	parent := d.find(dirname)

	if normalize(dirname) != normalize(parent.path()) {
		return errNotExists
	}

	parentDir, ok := parent.(*dir)
	if !ok {
		return errIsNotDir
	}

	if _, ok := parentDir.children[filename]; !ok {
		return errNotExists
	}

	delete(parentDir.children, filename)
	return nil
}

func (d *dir) String() string {
	return fmt.Sprintf("&dir{path: %v, children: %v}", d.p, d.children)
}

// file stores actual data in the fs tree. It acts like an open, seekable file
// where operations are conducted through ReadAt and WriteAt. Use it with
// SectionReader for the best effect.
type file struct {
	common
	data []byte
}

var _ node = &file{}

func (f *file) isdir() bool {
	return false
}

func (f *file) truncate() {
	f.data = f.data[:0]
}

func (f *file) sectionReader(offset int64) io.Reader {
	return io.NewSectionReader(f, offset, int64(len(f.data))-offset)
}

func (f *file) ReadAt(p []byte, offset int64) (n int, err error) {
	if offset >= int64(len(f.data)) {
		return 0, io.EOF
	}
	return copy(p, f.data[offset:]), nil
}

func (f *file) WriteAt(p []byte, offset int64) (n int, err error) {
	off := int(offset)
	if cap(f.data) < off+len(p) {
		data := make([]byte, len(f.data), off+len(p))
		copy(data, f.data)
		f.data = data
	}

	f.mod = time.Now()
	f.data = f.data[:off+len(p)]

	return copy(f.data[off:off+len(p)], p), nil
}

func (f *file) String() string {
	return fmt.Sprintf("&file{path: %q}", f.p)
}

// common provides shared fields and methods for node implementations.
type common struct {
	p   string
	mod time.Time
}

func (c *common) name() string {
	_, name := path.Split(c.p)
	return name
}

func (c *common) path() string {
	return c.p
}

func (c *common) modtime() time.Time {
	return c.mod
}

func normalize(p string) string {
	return "/" + strings.Trim(p, "/")
}
<|endoftext|>


<commit_before>package database

import (
  "crypto/rand"
  "math/big"
  "testing"

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

// used for testing empty data
var EmptyData []byte = []byte("")
var SingleData []byte = []byte("a")
var DoubleData []byte = []byte("ab")
var ShortData []byte = []byte("abcdefghijklmnopqrstuvwxyz")
var LongData []byte = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
var UnicodeData []byte = []byte("a®Ďƃɕʶ ̂ΆԃЌԵﬗאر݃ݓޤ߅ࡄখஷഖคබໄ၇ꩦႦᄓᎄⷄꬓᏄᑖᣆᚅᛕᜅᜤᝄᝣ‴№⁷✚z")

var AllData [][]byte = [][]byte{
  EmptyData,
  SingleData,
  DoubleData,
  ShortData,
  LongData,
  UnicodeData,
}

// skip the given test if running in short mode
func skipIfShort(t *testing.T) {
  if (testing.Short()) { t.Skip("Skipping test in short mode") }
}

// should be able to decompress the compressed empty array
func TestDecompressMinCompressed(t *testing.T) {
  minCompressed, err := compress([]byte{})
  assert.NoError(t, err)

  decompressed, err := decompress(minCompressed)
  assert.NoError(t, err)

  assert.Equal(t, decompressed, EmptyData);
}

func TestDecompressTooShort(t *testing.T) {
  minCompressed, err := compress([]byte{})
  assert.NoError(t, err)

  minCompressedLength := len(minCompressed)

  // check all possible sizes below the minimum compressed length to ensure that
  // they error.
  for size := len(minCompressed) - 1; size >= 0; size-- {
    _, err := decompress(make([]byte, minCompressedLength))
    assert.Error(t, err)
  }
}

// decompressing invalid data should fail
func TestDecompressInvalid(t *testing.T) {
  // null data is certainly invalid
  data := make([]byte, 50)

  _, err := decompress(data)
  assert.Error(t, err)
}

// compression and decompression are inverse operations, and therefor passing
// input through the compressor and then the decompressor should yield the input
// data once again.
func TestCompressAndDecompress(t *testing.T) {
  for _, data := range AllData {
    compressed, err := compress(data)
    assert.NoError(t, err)

    decompressed, err := decompress(compressed)
    assert.NoError(t, err)

    assert.Equal(t, decompressed, data)
  }
}

// fuzz test compression with lots of random data
func TestFuzzCompressAndDecompress(t *testing.T) {
  skipIfShort(t)

  for i := 0; i < 100000; i++ {
    // create a randomly-sized array, possibly larger than a byte in length
    size, err := rand.Int(rand.Reader, big.NewInt(512))
    assert.NoError(t, err)

    // fill the array with random data
    data := make([]byte, size.Int64())
    _, err = rand.Read(data)
    assert.NoError(t, err)

    // compress, decompress, and compare to the original
    compressed, err := compress(data)
    assert.NoError(t, err)

    decompressed, err := decompress(compressed)
    assert.NoError(t, err)

    assert.Equal(t, decompressed, data)
  }
}
<commit_msg>Add sign/verify tests<commit_after>package database

import (
  "crypto/rand"
  "math/big"
  "testing"

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

// used for testing empty data
var EmptyData []byte = []byte("")
var SingleData []byte = []byte("a")
var DoubleData []byte = []byte("ab")
var ShortData []byte = []byte("abcdefghijklmnopqrstuvwxyz")
var LongData []byte = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
var UnicodeData []byte = []byte("a®Ďƃɕʶ ̂ΆԃЌԵﬗאر݃ݓޤ߅ࡄখஷഖคබໄ၇ꩦႦᄓᎄⷄꬓᏄᑖᣆᚅᛕᜅᜤᝄᝣ‴№⁷✚z")

var AllData [][]byte = [][]byte{
  EmptyData,
  SingleData,
  DoubleData,
  ShortData,
  LongData,
  UnicodeData,
}

// skip the given test if running in short mode
func skipIfShort(t *testing.T) {
  if (testing.Short()) { t.Skip("Skipping test in short mode") }
}

// should be able to decompress the compressed empty array
func TestDecompressMinCompressed(t *testing.T) {
  minCompressed, err := compress([]byte{})
  assert.NoError(t, err)

  decompressed, err := decompress(minCompressed)
  assert.NoError(t, err)

  assert.Equal(t, decompressed, EmptyData);
}

func TestDecompressTooShort(t *testing.T) {
  minCompressed, err := compress([]byte{})
  assert.NoError(t, err)

  minCompressedLength := len(minCompressed)

  // check all possible sizes below the minimum compressed length to ensure that
  // they error.
  for size := len(minCompressed) - 1; size >= 0; size-- {
    _, err := decompress(make([]byte, minCompressedLength))
    assert.Error(t, err)
  }
}

// decompressing invalid data should fail
func TestDecompressInvalid(t *testing.T) {
  // null data is certainly invalid
  data := make([]byte, 50)

  _, err := decompress(data)
  assert.Error(t, err)
}

// compression and decompression are inverse operations, and therefor passing
// input through the compressor and then the decompressor should yield the input
// data once again.
func TestCompressAndDecompress(t *testing.T) {
  for _, data := range AllData {
    compressed, err := compress(data)
    assert.NoError(t, err)

    decompressed, err := decompress(compressed)
    assert.NoError(t, err)

    assert.Equal(t, decompressed, data)
  }
}

// fuzz test compression with lots of random data
func TestFuzzCompressAndDecompress(t *testing.T) {
  skipIfShort(t)

  for i := 0; i < 100000; i++ {
    // create a randomly-sized array, possibly larger than a byte in length
    size, err := rand.Int(rand.Reader, big.NewInt(512))
    assert.NoError(t, err)

    // fill the array with random data
    data := make([]byte, size.Int64())
    _, err = rand.Read(data)
    assert.NoError(t, err)

    // compress, decompress, and compare to the original
    compressed, err := compress(data)
    assert.NoError(t, err)

    decompressed, err := decompress(compressed)
    assert.NoError(t, err)

    assert.Equal(t, decompressed, data)
  }
}

// make sure that signing identical blocks of bytes always produces a signature
// of the reported length.
func TestGetSignatureSize(t *testing.T) {
  for _, data := range AllData {
    signature := getSignature(data)
    assert.Equal(t, len(signature), SignatureSize)
  }
}

// make sure that signing identical (but distinct) blocks of bytes always
// produces the same signature.
func TestGetSignatureIdentical(t *testing.T) {
  for _, data := range AllData {
    copiedData := make([]byte, len(data))
    copy(copiedData, data)

    signature1 := getSignature(data)
    signature2 := getSignature(copiedData)

    assert.Equal(t, signature1, signature2)
  }
}

// make sure that signing identical (but distinct) blocks of bytes always
// produces the same output.
func TestSignIdentical(t *testing.T) {
  for _, data := range AllData {
    copiedData := make([]byte, len(data))
    copy(copiedData, data)

    signed1 := sign(data)
    signed2 := sign(copiedData)

    assert.Equal(t, signed1, signed2)
  }
}

// make sure that signing data has a length equal to the length of the original
// data plus the size of a signature.
func TestSignLength(t *testing.T) {
  for _, data := range AllData {
    signed := sign(data)
    assert.Equal(t, len(signed), len(data) + SignatureSize)
  }
}

// make sure that signed data always starts with the original data
func TestSignStartsWithOriginalData(t *testing.T) {
  for _, data := range AllData {
    signed := sign(data)
    assert.Equal(t, data, signed[:len(data)])
  }
}

// data with an invalid signature shouldn't verify
func TestVerifyInvalid(t *testing.T) {
  for _, data := range AllData {
    // create a signature with all null bytes
    invalidSigned := make([]byte, len(data) + SignatureSize)
    copy(invalidSigned, data)

    _, err := verify(invalidSigned)
    assert.Error(t, err)
  }
}

// data shorter than the signature length shouldn't verify
func TestVerifyShort(t *testing.T) {
  for i := SignatureSize; i >= 0; i-- {
    invalidSigned := make([]byte, i)

    _, err := verify(invalidSigned)
    assert.Error(t, err)
  }
}

// signing and verifying data should work, and the result data should equal the
// input data.
func TestSignAndVerify(t *testing.T) {
  skipIfShort(t)

  for _, data := range AllData {
    signed := sign(data)
    verified, err := verify(signed)

    assert.NoError(t, err)
    assert.Equal(t, data, verified)
  }
}

// test signing and verifying lots of random data
func TestFuzzSignAndVerify(t *testing.T) {
  skipIfShort(t)

  for i := 0; i < 100000; i++ {
    // create a randomly-sized array
    size, err := rand.Int(rand.Reader, big.NewInt(512))
    assert.NoError(t, err)

    // fill the array with random data
    data := make([]byte, size.Int64())
    _, err = rand.Read(data)
    assert.NoError(t, err)

    // sign, verify, and compare to the original
    signed := sign(data)
    verified, err := verify(signed)

    assert.NoError(t, err)
    assert.Equal(t, data, verified)
  }
}

// data should take the same amount of time to determine validity, no matter how
// soon the signatures differ (i.e. should be robust against timing attacks).
func TestVerifyConstantTime(t *testing.T) {
  skipIfShort(t)

  // TODO: build this and statistically verify the reported times
  assert.Fail(t, "Implement this!")
}
<|endoftext|>


<commit_before>package daemon

import (
    "os"
    "fmt"
    "github.com/sevlyar/go-daemon"
    "github.com/Sirupsen/logrus"
    "github.com/BluePecker/JwtAuth/pkg/storage"
    "github.com/BluePecker/JwtAuth/daemon/server"
    "syscall"
)

const (
    TOKEN_TTL = 2 * 3600

    VERSION = "1.0.0"

    ALLOW_LOGIN_NUM = 3
)

type Storage struct {
    Driver string
    Opts   string
}

type TLS struct {
    Key  string
    Cert string
}

type Options struct {
    PidFile  string
    LogFile  string
    LogLevel string
    Port     int
    Host     string
    SockFile string
    Daemon   bool
    Version  bool
    TLS      TLS
    Storage  Storage
    Secret   string
}

type Daemon struct {
    Options  *Options

    shadow   *server.Shadow
    rosiness *server.Rosiness

    StorageE *storage.Engine
}

func Logger(level string) {
    logrus.SetFormatter(&logrus.TextFormatter{
        TimestampFormat: "2006-01-02 15:04:05",
    })
    Level, err := logrus.ParseLevel(level)
    if err != nil {
        logrus.Error(err)
        os.Exit(0)
    }
    logrus.SetLevel(Level)
}

func Version(version bool) {
    if version == true {
        fmt.Printf("JwtAuth version %s.\n", VERSION)
        os.Exit(0)
    }
}

func NewDaemon(background bool, args Options) *Daemon {
    if background {
        ctx := daemon.Context{
            PidFileName: args.PidFile,
            PidFilePerm: 0644,
            LogFilePerm: 0640,
            Umask:       027,
            WorkDir:     "/",
            LogFileName: args.LogFile,
        }
        if process, err := ctx.Reborn(); err == nil {
            defer ctx.Release()
            if process != nil {
                return nil
            }
        } else {
            if err == daemon.ErrWouldBlock {
                logrus.Error("daemon already exists.")
            } else {
                logrus.Errorf("Unable to run: ", err)
            }
            os.Exit(0)
        }
    }
    return &Daemon{Options: &args}
}

func NewStart(args Options) {

    Logger(args.LogLevel)
    Version(args.Version)

    if args.Secret == "" {
        fmt.Println("please specify the key.")
        os.Exit(0)
    }

    if progress := NewDaemon(args.Daemon, args); progress == nil {
        return
    } else {
        if progress == nil {
            return
        }
        if err := progress.Storage(); err != nil {
            logrus.Error(err)
            os.Exit(0)
        }

        stopSignal := make(chan struct{})
        go progress.Shadow(stopSignal)
        go progress.Rosiness(stopSignal)
        daemon.SetSigHandler(func(sig os.Signal) error {
            close(stopSignal)
            return daemon.ErrStop
        }, syscall.SIGTERM, syscall.SIGQUIT)

        if err := daemon.ServeSignals(); err != nil {
            logrus.Error(err)
        }
        logrus.Error("jwt daemon terminated.")
    }
}<commit_msg>fix bug<commit_after>package daemon

import (
    "os"
    "fmt"
    "github.com/sevlyar/go-daemon"
    "github.com/Sirupsen/logrus"
    "github.com/BluePecker/JwtAuth/pkg/storage"
    "github.com/BluePecker/JwtAuth/daemon/server"
    "syscall"
)

const (
    TOKEN_TTL = 2 * 3600

    VERSION = "1.0.0"

    ALLOW_LOGIN_NUM = 3
)

type Storage struct {
    Driver string
    Opts   string
}

type TLS struct {
    Key  string
    Cert string
}

type Options struct {
    PidFile  string
    LogFile  string
    LogLevel string
    Port     int
    Host     string
    SockFile string
    Daemon   bool
    Version  bool
    TLS      TLS
    Storage  Storage
    Secret   string
}

type Daemon struct {
    Options  *Options

    shadow   *server.Shadow
    rosiness *server.Rosiness

    StorageE *storage.Engine
}

func Logger(level string) {
    logrus.SetFormatter(&logrus.TextFormatter{
        TimestampFormat: "2006-01-02 15:04:05",
    })
    Level, err := logrus.ParseLevel(level)
    if err != nil {
        logrus.Error(err)
        os.Exit(0)
    }
    logrus.SetLevel(Level)
}

func Version(version bool) {
    if version == true {
        fmt.Printf("JwtAuth version %s.\n", VERSION)
        os.Exit(0)
    }
}

func NewDaemon(background bool, args Options) *Daemon {
    if background {
        ctx := daemon.Context{
            PidFileName: args.PidFile,
            PidFilePerm: 0644,
            LogFilePerm: 0640,
            Umask:       027,
            WorkDir:     "/",
            LogFileName: args.LogFile,
        }
        if process, err := ctx.Reborn(); err == nil {
            defer ctx.Release()
            if process != nil {
                return nil
            }
        } else {
            if err == daemon.ErrWouldBlock {
                logrus.Error("daemon already exists.")
            } else {
                logrus.Errorf("Unable to run: ", err)
            }
            os.Exit(0)
        }
    }
    return &Daemon{Options: &args}
}

func NewStart(args Options) {

    Logger(args.LogLevel)
    Version(args.Version)

    if args.Secret == "" {
        fmt.Println("please specify the key.")
        os.Exit(0)
    }

    if progress := NewDaemon(args.Daemon, args); progress == nil {
        return
    } else {
        if progress == nil {
            return
        }
        if err := progress.Storage(); err != nil {
            logrus.Error(err)
            os.Exit(0)
        }

        quit := make(chan struct{})
        go progress.Shadow(quit)
        go progress.Rosiness(quit)
        daemon.SetSigHandler(func(sig os.Signal) error {
            close(quit)
            return daemon.ErrStop
        }, syscall.SIGTERM, syscall.SIGQUIT)

        if err := daemon.ServeSignals(); err != nil {
            logrus.Error(err)
        }
        logrus.Error("jwt daemon terminated.")
    }
}<|endoftext|>


<commit_before>/*--------------------------------------------------------*\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: https://hprose.com                     |
|                                                          |
| rpc/plugins/cluster/cluster.go                           |
|                                                          |
| LastModified: Mar 6, 2021                                |
| Author: Ma Bingyao <andot@hprose.com>                    |
|                                                          |
\*________________________________________________________*/

package cluster

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hprose/hprose-golang/v3/rpc/core"
)

// Config for cluster.
type Config struct {
	Retry       int
	Idempotent  bool
	OnSuccess   func(context.Context)
	OnFailure   func(context.Context)
	OnRetry     func(context.Context) time.Duration
	minInterval time.Duration
	maxInterval time.Duration
}

// Option for cluster config.
type Option func(*Config)

// WithRetry returns a retry Option for cluster config.
func WithRetry(retry int) Option {
	return func(c *Config) {
		c.Retry = retry
	}
}

// WithIdempotent returns an idempotent Option for cluster config.
func WithIdempotent(idempotent bool) Option {
	return func(c *Config) {
		c.Idempotent = idempotent
	}
}

// WithMinInterval returns a retry minInterval Option for cluster config.
func WithMinInterval(minInterval time.Duration) Option {
	return func(c *Config) {
		c.minInterval = minInterval
	}
}

// WithMaxInterval returns a retry maxInterval Option for cluster config.
func WithMaxInterval(maxInterval time.Duration) Option {
	return func(c *Config) {
		c.maxInterval = maxInterval
	}
}

// FailoverConfig for cluster.
func FailoverConfig(options ...Option) (config Config) {
	config.Retry = 10
	config.minInterval = time.Millisecond * 500
	config.maxInterval = time.Second * 5
	for _, option := range options {
		option(&config)
	}
	var index uint64
	config.OnFailure = func(ctx context.Context) {
		clientContext := core.GetClientContext(ctx)
		urls := clientContext.Client().URLs
		n := uint64(len(urls))
		if n > 1 {
			if atomic.AddUint64(&index, 1) >= n {
				index = 0
			}
			clientContext.URL = urls[index]
		}
	}
	config.OnRetry = func(ctx context.Context) time.Duration {
		clientContext := core.GetClientContext(ctx)
		retried := clientContext.Items().GetInt("retried") + 1
		clientContext.Items().Set("retried", retried)
		interval := config.minInterval * time.Duration(retried-len(clientContext.Client().URLs))
		if interval > config.maxInterval {
			interval = config.maxInterval
		}
		return interval
	}
	return
}

// FailtryConfig for cluster.
func FailtryConfig(options ...Option) (config Config) {
	config.Retry = 10
	config.minInterval = time.Millisecond * 500
	config.maxInterval = time.Second * 5
	for _, option := range options {
		option(&config)
	}
	config.OnRetry = func(ctx context.Context) time.Duration {
		clientContext := core.GetClientContext(ctx)
		retried := clientContext.Items().GetInt("retried") + 1
		clientContext.Items().Set("retried", retried)
		interval := config.minInterval * time.Duration(retried)
		if interval > config.maxInterval {
			interval = config.maxInterval
		}
		return interval
	}
	return
}

// FailfastConfig for cluster.
func FailfastConfig(onFailure func(context.Context)) (config Config) {
	config.Retry = 0
	config.OnFailure = onFailure
	return
}

// Cluster plugin for hprose.
type Cluster struct {
	Config
}

// New returns a Cluster instance.
func New(config ...Config) (cluster *Cluster) {
	cluster = &Cluster{}
	if len(config) > 0 {
		cluster.Config = config[0]
		if cluster.Retry < 0 {
			cluster.Retry = 10
		}
	} else {
		cluster.Config = FailoverConfig()
	}
	return
}

// Handler for Cluster.
func (c *Cluster) Handler(ctx context.Context, request []byte, next core.NextIOHandler) (response []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = core.NewPanicError(e)
		}
		if err == nil {
			return
		}
		if c.OnFailure != nil {
			c.OnFailure(ctx)
		}
		if c.OnRetry == nil {
			return
		}
		clientContext := core.GetClientContext(ctx)
		idempotent := clientContext.Items().GetBool("idempotent", c.Idempotent)
		retry := clientContext.Items().GetInt("retry", c.Retry)
		retried := clientContext.Items().GetInt("retried")
		if idempotent && (retried < retry) {
			interval := c.OnRetry(ctx)
			if interval > 0 {
				time.Sleep(interval)
			}
			response, err = c.Handler(ctx, request, next)
		}
	}()
	if response, err = next(ctx, request); err == nil && c.OnSuccess != nil {
		c.OnSuccess(ctx)
	}
	return
}

// Forking on Cluster.
func Forking(ctx context.Context, request []byte, next core.NextIOHandler) (response []byte, err error) {
	clientContext := core.GetClientContext(ctx)
	urls := clientContext.Client().URLs
	n := len(urls)
	if n == 0 {
		return next(ctx, request)
	}
	count := int64(n)
	var once sync.Once
	done := make(chan struct{})
	for i := 0; i < n; i++ {
		forkingContext := clientContext.Clone().(*core.ClientContext)
		forkingContext.URL = urls[i]
		go func(ctx context.Context) {
			defer func() {
				if e := recover(); e != nil && atomic.AddInt64(&count, -1) <= 0 {
					once.Do(func() {
						err = core.NewPanicError(e)
						close(done)
					})
				}
			}()
			resp, e := next(ctx, request)
			if e == nil {
				once.Do(func() {
					response = resp
					close(done)
				})
				return
			}
			if atomic.AddInt64(&count, -1) <= 0 {
				once.Do(func() {
					err = e
					close(done)
				})
			}
		}(core.WithContext(ctx, forkingContext))
	}
	<-done
	return
}

// Broadcast on Cluster.
func Broadcast(ctx context.Context, name string, args []interface{}, next core.NextInvokeHandler) (result []interface{}, err error) {
	clientContext := core.GetClientContext(ctx)
	urls := clientContext.Client().URLs
	n := len(urls)
	if n == 0 {
		return next(ctx, name, args)
	}
	var wg sync.WaitGroup
	var once sync.Once
	result = make([]interface{}, n)
	wg.Add(n)
	for i := 0; i < n; i++ {
		forkingContext := clientContext.Clone().(*core.ClientContext)
		forkingContext.URL = urls[i]
		go func(i int, ctx context.Context) {
			defer func() {
				if e := recover(); e != nil {
					once.Do(func() { err = core.NewPanicError(e) })
				}
				wg.Done()
			}()
			var e error
			if result[i], e = next(ctx, name, args); e != nil {
				once.Do(func() { err = e })
			}
		}(i, core.WithContext(ctx, forkingContext))
	}
	wg.Wait()
	return
}
<commit_msg>Improve Cluster.<commit_after>/*--------------------------------------------------------*\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: https://hprose.com                     |
|                                                          |
| rpc/plugins/cluster/cluster.go                           |
|                                                          |
| LastModified: Mar 6, 2021                                |
| Author: Ma Bingyao <andot@hprose.com>                    |
|                                                          |
\*________________________________________________________*/

package cluster

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hprose/hprose-golang/v3/rpc/core"
)

// Config for cluster.
type Config struct {
	Retry       int
	Idempotent  bool
	OnSuccess   func(context.Context)
	OnFailure   func(context.Context)
	OnRetry     func(context.Context) time.Duration
	minInterval time.Duration
	maxInterval time.Duration
}

// Option for cluster config.
type Option func(*Config)

// WithRetry returns a retry Option for cluster config.
func WithRetry(retry int) Option {
	return func(c *Config) {
		c.Retry = retry
	}
}

// WithIdempotent returns an idempotent Option for cluster config.
func WithIdempotent(idempotent bool) Option {
	return func(c *Config) {
		c.Idempotent = idempotent
	}
}

// WithMinInterval returns a retry minInterval Option for cluster config.
func WithMinInterval(minInterval time.Duration) Option {
	return func(c *Config) {
		c.minInterval = minInterval
	}
}

// WithMaxInterval returns a retry maxInterval Option for cluster config.
func WithMaxInterval(maxInterval time.Duration) Option {
	return func(c *Config) {
		c.maxInterval = maxInterval
	}
}

// FailoverConfig for cluster.
func FailoverConfig(options ...Option) (config Config) {
	config.Retry = 10
	config.minInterval = time.Millisecond * 500
	config.maxInterval = time.Second * 5
	for _, option := range options {
		option(&config)
	}
	var index int64
	config.OnFailure = func(ctx context.Context) {
		clientContext := core.GetClientContext(ctx)
		urls := clientContext.Client().URLs
		n := int64(len(urls))
		if n > 1 {
			if atomic.AddInt64(&index, 1) >= n {
				atomic.StoreInt64(&index, 0)
			}
			clientContext.URL = urls[atomic.LoadInt64(&index)]
		}
	}
	config.OnRetry = func(ctx context.Context) time.Duration {
		clientContext := core.GetClientContext(ctx)
		retried := clientContext.Items().GetInt("retried") + 1
		clientContext.Items().Set("retried", retried)
		interval := config.minInterval * time.Duration(retried-len(clientContext.Client().URLs))
		if interval > config.maxInterval {
			interval = config.maxInterval
		}
		return interval
	}
	return
}

// FailtryConfig for cluster.
func FailtryConfig(options ...Option) (config Config) {
	config.Retry = 10
	config.minInterval = time.Millisecond * 500
	config.maxInterval = time.Second * 5
	for _, option := range options {
		option(&config)
	}
	config.OnRetry = func(ctx context.Context) time.Duration {
		clientContext := core.GetClientContext(ctx)
		retried := clientContext.Items().GetInt("retried") + 1
		clientContext.Items().Set("retried", retried)
		interval := config.minInterval * time.Duration(retried)
		if interval > config.maxInterval {
			interval = config.maxInterval
		}
		return interval
	}
	return
}

// FailfastConfig for cluster.
func FailfastConfig(onFailure func(context.Context)) (config Config) {
	config.Retry = 0
	config.OnFailure = onFailure
	return
}

// Cluster plugin for hprose.
type Cluster struct {
	Config
}

// New returns a Cluster instance.
func New(config ...Config) (cluster *Cluster) {
	cluster = &Cluster{}
	if len(config) > 0 {
		cluster.Config = config[0]
		if cluster.Retry < 0 {
			cluster.Retry = 10
		}
	} else {
		cluster.Config = FailoverConfig()
	}
	return
}

// Handler for Cluster.
func (c *Cluster) Handler(ctx context.Context, request []byte, next core.NextIOHandler) (response []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = core.NewPanicError(e)
		}
		if err == nil {
			return
		}
		if c.OnFailure != nil {
			c.OnFailure(ctx)
		}
		if c.OnRetry == nil {
			return
		}
		clientContext := core.GetClientContext(ctx)
		idempotent := clientContext.Items().GetBool("idempotent", c.Idempotent)
		retry := clientContext.Items().GetInt("retry", c.Retry)
		retried := clientContext.Items().GetInt("retried")
		if idempotent && (retried < retry) {
			interval := c.OnRetry(ctx)
			if interval > 0 {
				time.Sleep(interval)
			}
			response, err = c.Handler(ctx, request, next)
		}
	}()
	if response, err = next(ctx, request); err == nil && c.OnSuccess != nil {
		c.OnSuccess(ctx)
	}
	return
}

// Forking on Cluster.
func Forking(ctx context.Context, request []byte, next core.NextIOHandler) (response []byte, err error) {
	clientContext := core.GetClientContext(ctx)
	urls := clientContext.Client().URLs
	n := len(urls)
	if n == 0 {
		return next(ctx, request)
	}
	count := int64(n)
	var once sync.Once
	done := make(chan struct{})
	for i := 0; i < n; i++ {
		forkingContext := clientContext.Clone().(*core.ClientContext)
		forkingContext.URL = urls[i]
		go func(ctx context.Context) {
			defer func() {
				if e := recover(); e != nil && atomic.AddInt64(&count, -1) <= 0 {
					once.Do(func() {
						err = core.NewPanicError(e)
						close(done)
					})
				}
			}()
			resp, e := next(ctx, request)
			if e == nil {
				once.Do(func() {
					response = resp
					close(done)
				})
				return
			}
			if atomic.AddInt64(&count, -1) <= 0 {
				once.Do(func() {
					err = e
					close(done)
				})
			}
		}(core.WithContext(ctx, forkingContext))
	}
	<-done
	return
}

// Broadcast on Cluster.
func Broadcast(ctx context.Context, name string, args []interface{}, next core.NextInvokeHandler) (result []interface{}, err error) {
	clientContext := core.GetClientContext(ctx)
	urls := clientContext.Client().URLs
	n := len(urls)
	if n == 0 {
		return next(ctx, name, args)
	}
	var wg sync.WaitGroup
	var once sync.Once
	result = make([]interface{}, n)
	wg.Add(n)
	for i := 0; i < n; i++ {
		forkingContext := clientContext.Clone().(*core.ClientContext)
		forkingContext.URL = urls[i]
		go func(i int, ctx context.Context) {
			defer func() {
				if e := recover(); e != nil {
					once.Do(func() { err = core.NewPanicError(e) })
				}
				wg.Done()
			}()
			var e error
			if result[i], e = next(ctx, name, args); e != nil {
				once.Do(func() { err = e })
			}
		}(i, core.WithContext(ctx, forkingContext))
	}
	wg.Wait()
	return
}
<|endoftext|>


<commit_before>package daemon

import (
	"encoding/json"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
	"github.com/ondevice/ondevice/api"
	"github.com/ondevice/ondevice/config"
	"github.com/ondevice/ondevice/logg"
	"github.com/ondevice/ondevice/service"
	"github.com/ondevice/ondevice/tunnel"
	"github.com/ondevice/ondevice/util"
)

// ControlSocket -- REST API to control the ondevice daemon (the implementation's in the control package)
type ControlSocket interface {
	Start()
	Stop() error
}

// Daemon -- represents a device's connection to the ondevice.io API server
type Daemon struct {
	tunnel.Connection

	PIDFile    string
	ConfigFile string

	lastPing   time.Time
	wdog       *util.Watchdog
	signalChan chan os.Signal
	lock       lockFile

	IsOnline bool // TODO replace by state enum

	Control      ControlSocket
	OnConnection func(tunnelID string, service string, protocol string)
	OnError      func(error)
}

type pingMsg struct {
	Type string `json:"_type"`
	Ts   int    `json:"ts"`
}

// NewDaemon -- Create a new Daemon instance
func NewDaemon() *Daemon {
	return &Daemon{
		signalChan: make(chan os.Signal, 1),
	}
}

// Run -- run ondevice daemon (and return with the exit code of the command)
func (d *Daemon) Run() int {
	d.lock.Path = d.PIDFile
	if !d.lock.TryLock() {
		logg.Fatal("Couldn't acquire lock file")
		return -1
	}
	defer d.lock.Unlock()

	if d.Control != nil {
		d.Control.Start()
		defer d.Control.Stop()
	}

	signal.Notify(d.signalChan, syscall.SIGTERM)

	// TODO implement a sane way to stop this infinite loop (at least SIGTERM, SIGINT or maybe a unix socket call)
	retryDelay := 10 * time.Second
	for true {
		d, err := Connect()
		if err != nil {
			retryDelay = d.waitBeforeRetry(retryDelay, err)
			continue
		}

		d.Wait()

		// connection was successful -> restart after 10sec
		logg.Warning("lost device connection, reconnecting in 10s")
		retryDelay = 10
		time.Sleep(retryDelay * time.Second)
	}

	return 0
}

// Connect -- Go online
func Connect(auths ...api.Authentication) (*Daemon, util.APIError) {
	params := map[string]string{"key": config.GetDeviceKey()}
	rc := Daemon{}

	if len(auths) == 0 {
		auth, err := api.GetDeviceAuth()
		if err != nil {
			logg.Fatal("Couldn't get device auth: ", err)
		}
		auths = []api.Authentication{auth}
	}

	if err := tunnel.OpenWebsocket(&rc.Connection, "/serve", params, rc.onMessage, auths...); err != nil {
		return nil, err
	}

	rc.wdog = util.NewWatchdog(180*time.Second, rc.onTimeout)
	return &rc, nil
}

// SendConnectionError -- Send an connection error message to the API server)
func (d *Daemon) SendConnectionError(code int, msg string, tunnelID string) {
	logg.Debugf("Sending connection error: %s (code %d)", msg, code)
	data := make(map[string]interface{})
	data["_type"] = "connectError"
	data["tunnelId"] = tunnelID
	data["code"] = code
	data["msg"] = msg
	d.SendJSON(data)
}

func (d *Daemon) announce(service string, protocol string) {
	var data = map[string]string{"_type": "announce", "name": service, "protocol": protocol}
	d.SendJSON(data)
}

func (d *Daemon) onConnect(msg *map[string]interface{}) {
	clientIP := _getString(msg, "clientIp")
	clientUser := _getString(msg, "clientUser")
	protocol := _getString(msg, "protocol")
	svc := _getString(msg, "service")
	brokerURL := _getString(msg, "broker")
	tunnelID := _getString(msg, "tunnelId")

	logg.Infof("Connection request for %s:%s from user %s@%s", protocol, svc, clientUser, clientIP)

	handler := service.GetServiceHandler(svc, protocol)
	if handler == nil {
		// TODO send the error back to the API server
		logg.Error("Coudln't find protocol handler: ", protocol)
		return
	}

	service.Start(handler, tunnelID, brokerURL)
}

func (d *Daemon) onError(msg *map[string]interface{}) {
	code := _getInt(msg, "code")
	message := _getString(msg, "msg")
	var codeName = tunnel.GetErrorCodeName(int(code))

	d.IsOnline = false
	logg.Errorf("Device ERROR: %s - %s ", codeName, message)
}

func (d *Daemon) onHello(msg *map[string]interface{}) {
	logg.Debug("Got hello message: ", msg)
	var devID, key string
	if _contains(msg, "name") {
		// deprecated hello message format (for backwards compatibility) -- 2017-01-19
		devID = _getString(msg, "name")
		key = _getString(msg, "devId")
	} else {
		devID = _getString(msg, "devId")
		key = _getString(msg, "key")
	}

	logg.Infof("Connection established, online as '%s'", devID)
	d.IsOnline = true

	// update the key if changed
	if config.GetDeviceKey() != key {
		logg.Debug("Updating device key: ", key)
		config.SetValue("device", "key", key)
	}

	// update devID
	config.SetValue("device", "dev-id", devID)

	// TODO announce configured services
	d.announce("ssh", "ssh")
}

func (d *Daemon) onMessage(_type int, data []byte) {
	// got message from the API server -> reset watchdog
	d.wdog.Kick()

	if _type == websocket.BinaryMessage {
		logg.Error("Got a binary message over the device websocket: ", string(data))
		return
	}

	msg := new(map[string]interface{})

	if err := json.Unmarshal(data, msg); err != nil {
		logg.Fatalf("Malformed device message: %s", data)
	}

	msgType := _getString(msg, "_type")
	switch msgType {
	case "hello":
		d.onHello(msg)
		break
	case "ping":
		var ping pingMsg
		json.Unmarshal(data, &ping)
		d.onPing(ping)
		break
	case "connect":
		d.onConnect(msg)
	case "error":
		d.onError(msg)
	default:
		logg.Error("Unsupported WS message: ", data)
		break
	}
}

func (d *Daemon) onPing(msg pingMsg) {
	// quick'n'dirty way to see if we're leaking goroutines (e.g. with stray bloking reads)
	logg.Debugf("Got ping message: %+v (active goroutines: %d)", msg, runtime.NumGoroutine())
	d.lastPing = time.Now()
	resp := make(map[string]interface{}, 1)
	resp["_type"] = "pong"
	resp["ts"] = msg.Ts
	d.SendJSON(resp)
}

func (d *Daemon) onTimeout() {
	logg.Warning("Haven't heard from the API server in a while, closing connection...")
	d.IsOnline = false
	d.Close()
	d.wdog.Stop()
}

func (d *Daemon) waitBeforeRetry(retryDelay time.Duration, err util.APIError) time.Duration {
	// only abort here if it's an authentication issue
	if err.Code() == util.AuthenticationError {
		logg.Fatal(err)
	}

	// keep retryDelay between 10 and 120sec
	if retryDelay > 120*time.Second {
		retryDelay = 120 * time.Second
	}
	if retryDelay < 10*time.Second {
		retryDelay = 10 * time.Second
	}
	// ... unless we've been rate-limited
	if err.Code() == util.TooManyRequestsError {
		retryDelay = 600 * time.Second
	}

	logg.Debug("device error: ", err)
	logg.Errorf("device error - retrying in %ds", retryDelay/time.Second)

	// sleep to avoid flooding the servers
	time.Sleep(retryDelay)

	// slowly increase retryDelay with each failed attempt
	return time.Duration(float32(retryDelay) * 1.5)
}

func _contains(m *map[string]interface{}, key string) bool {
	_, ok := (*m)[key]
	return ok
}

func _getInt(m *map[string]interface{}, key string) int64 {
	return (*m)[key].(int64)
}

func _getString(m *map[string]interface{}, key string) string {
	if val, ok := (*m)[key]; ok {
		var rc, ok = val.(string)
		if !ok {
			logg.Warningf("Not a string (key %s): %s", key, val)
			return ""
		}
		return rc
	}
	return ""
}
<commit_msg>daemon: added simple signal handler for graceful shutdowns<commit_after>package daemon

import (
	"encoding/json"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
	"github.com/ondevice/ondevice/api"
	"github.com/ondevice/ondevice/config"
	"github.com/ondevice/ondevice/logg"
	"github.com/ondevice/ondevice/service"
	"github.com/ondevice/ondevice/tunnel"
	"github.com/ondevice/ondevice/util"
)

// ControlSocket -- REST API to control the ondevice daemon (the implementation's in the control package)
type ControlSocket interface {
	Start()
	Stop() error
}

// Daemon -- represents a device's connection to the ondevice.io API server
type Daemon struct {
	tunnel.Connection

	PIDFile    string
	ConfigFile string

	lastPing     time.Time
	wdog         *util.Watchdog
	signalChan   chan os.Signal
	firstSIGTERM time.Time
	lock         lockFile

	IsOnline bool // TODO replace by state enum

	Control      ControlSocket
	OnConnection func(tunnelID string, service string, protocol string)
	OnError      func(error)
}

type pingMsg struct {
	Type string `json:"_type"`
	Ts   int    `json:"ts"`
}

// NewDaemon -- Create a new Daemon instance
func NewDaemon() *Daemon {
	return &Daemon{
		signalChan: make(chan os.Signal, 1),
	}
}

// Run -- run ondevice daemon (and return with the exit code of the command)
func (d *Daemon) Run() int {
	d.lock.Path = d.PIDFile
	if !d.lock.TryLock() {
		logg.Fatal("Couldn't acquire lock file")
		return -1
	}
	defer d.lock.Unlock()

	if d.Control != nil {
		d.Control.Start()
		defer d.Control.Stop()
	}

	go d.signalHandler()
	signal.Notify(d.signalChan, syscall.SIGTERM, syscall.SIGINT)

	// TODO implement a sane way to stop this infinite loop (at least SIGTERM, SIGINT or maybe a unix socket call)
	retryDelay := 10 * time.Second
	for true {
		if err := d.connect(); err != nil {
			retryDelay = d.waitBeforeRetry(retryDelay, err)
		} else {
			d.Wait()

			// connection was successful -> restart after 10sec
			logg.Warning("lost device connection, reconnecting in 10s")
			retryDelay = 10
			time.Sleep(retryDelay * time.Second)
		}
	}

	return 0
}

// Close -- Gracefully stopping this ondevice daemon instance
func (d *Daemon) Close() {
	if d.Control != nil {
		d.Control.Stop()
	}
	d.Connection.Close()
}

// connect -- Go online
func (d *Daemon) connect(auths ...api.Authentication) util.APIError {
	params := map[string]string{"key": config.GetDeviceKey()}

	if len(auths) == 0 {
		auth, err := api.GetDeviceAuth()
		if err != nil {
			logg.Fatal("Couldn't get device auth: ", err)
		}
		auths = []api.Authentication{auth}
	}

	if err := tunnel.OpenWebsocket(&d.Connection, "/serve", params, d.onMessage, auths...); err != nil {
		return err
	}

	d.wdog = util.NewWatchdog(180*time.Second, d.onTimeout)
	return nil
}

// SendConnectionError -- Send an connection error message to the API server)
func (d *Daemon) SendConnectionError(code int, msg string, tunnelID string) {
	logg.Debugf("Sending connection error: %s (code %d)", msg, code)
	data := make(map[string]interface{})
	data["_type"] = "connectError"
	data["tunnelId"] = tunnelID
	data["code"] = code
	data["msg"] = msg
	d.SendJSON(data)
}

func (d *Daemon) announce(service string, protocol string) {
	var data = map[string]string{"_type": "announce", "name": service, "protocol": protocol}
	d.SendJSON(data)
}

func (d *Daemon) onConnect(msg *map[string]interface{}) {
	clientIP := _getString(msg, "clientIp")
	clientUser := _getString(msg, "clientUser")
	protocol := _getString(msg, "protocol")
	svc := _getString(msg, "service")
	brokerURL := _getString(msg, "broker")
	tunnelID := _getString(msg, "tunnelId")

	logg.Infof("Connection request for %s:%s from user %s@%s", protocol, svc, clientUser, clientIP)

	handler := service.GetServiceHandler(svc, protocol)
	if handler == nil {
		// TODO send the error back to the API server
		logg.Error("Coudln't find protocol handler: ", protocol)
		return
	}

	service.Start(handler, tunnelID, brokerURL)
}

func (d *Daemon) onError(msg *map[string]interface{}) {
	code := _getInt(msg, "code")
	message := _getString(msg, "msg")
	var codeName = tunnel.GetErrorCodeName(int(code))

	d.IsOnline = false
	logg.Errorf("Device ERROR: %s - %s ", codeName, message)
}

func (d *Daemon) onHello(msg *map[string]interface{}) {
	logg.Debug("Got hello message: ", msg)
	var devID, key string
	if _contains(msg, "name") {
		// deprecated hello message format (for backwards compatibility) -- 2017-01-19
		devID = _getString(msg, "name")
		key = _getString(msg, "devId")
	} else {
		devID = _getString(msg, "devId")
		key = _getString(msg, "key")
	}

	logg.Infof("Connection established, online as '%s'", devID)
	d.IsOnline = true

	// update the key if changed
	if config.GetDeviceKey() != key {
		logg.Debug("Updating device key: ", key)
		config.SetValue("device", "key", key)
	}

	// update devID
	config.SetValue("device", "dev-id", devID)

	// TODO announce configured services
	d.announce("ssh", "ssh")
}

func (d *Daemon) onMessage(_type int, data []byte) {
	// got message from the API server -> reset watchdog
	d.wdog.Kick()

	if _type == websocket.BinaryMessage {
		logg.Error("Got a binary message over the device websocket: ", string(data))
		return
	}

	msg := new(map[string]interface{})

	if err := json.Unmarshal(data, msg); err != nil {
		logg.Fatalf("Malformed device message: %s", data)
	}

	msgType := _getString(msg, "_type")
	switch msgType {
	case "hello":
		d.onHello(msg)
		break
	case "ping":
		var ping pingMsg
		json.Unmarshal(data, &ping)
		d.onPing(ping)
		break
	case "connect":
		d.onConnect(msg)
	case "error":
		d.onError(msg)
	default:
		logg.Error("Unsupported WS message: ", data)
		break
	}
}

func (d *Daemon) onPing(msg pingMsg) {
	// quick'n'dirty way to see if we're leaking goroutines (e.g. with stray bloking reads)
	logg.Debugf("Got ping message: %+v (active goroutines: %d)", msg, runtime.NumGoroutine())
	d.lastPing = time.Now()
	resp := make(map[string]interface{}, 1)
	resp["_type"] = "pong"
	resp["ts"] = msg.Ts
	d.SendJSON(resp)
}

func (d *Daemon) onTimeout() {
	logg.Warning("Haven't heard from the API server in a while, closing connection...")
	d.IsOnline = false
	d.Close()
	d.wdog.Stop()
}

func (d *Daemon) signalHandler() {
	for true {
		var sig, ok = <-d.signalChan
		if !ok {
			break
		}
		switch sig {
		case syscall.SIGTERM:
			logg.Info("Got SIGTERM, gracefully shutting down...")
		case syscall.SIGINT:
			logg.Info("Got SIGINT, gracefully shutting down...")
		default:
			logg.Warning("Caught unexpected signal: ", sig)
		}

		d.Close()
	}

	logg.Info("Stopping to handle signals")
}

func (d *Daemon) waitBeforeRetry(retryDelay time.Duration, err util.APIError) time.Duration {
	// only abort here if it's an authentication issue
	if err.Code() == util.AuthenticationError {
		logg.Fatal(err)
	}

	// keep retryDelay between 10 and 120sec
	if retryDelay > 120*time.Second {
		retryDelay = 120 * time.Second
	}
	if retryDelay < 10*time.Second {
		retryDelay = 10 * time.Second
	}
	// ... unless we've been rate-limited
	if err.Code() == util.TooManyRequestsError {
		retryDelay = 600 * time.Second
	}

	logg.Debug("device error: ", err)
	logg.Errorf("device error - retrying in %ds", retryDelay/time.Second)

	// sleep to avoid flooding the servers
	time.Sleep(retryDelay)

	// slowly increase retryDelay with each failed attempt
	return time.Duration(float32(retryDelay) * 1.5)
}

func _contains(m *map[string]interface{}, key string) bool {
	_, ok := (*m)[key]
	return ok
}

func _getInt(m *map[string]interface{}, key string) int64 {
	return (*m)[key].(int64)
}

func _getString(m *map[string]interface{}, key string) string {
	if val, ok := (*m)[key]; ok {
		var rc, ok = val.(string)
		if !ok {
			logg.Warningf("Not a string (key %s): %s", key, val)
			return ""
		}
		return rc
	}
	return ""
}
<|endoftext|>


<commit_before>package command

import "testing"

type appArgsTestCases dataTestCase

func runAppArgsTestCase(t *testing.T, cases []appArgsTestCases) {
	for _, c := range cases {
		if got := c.in.IsValid(); got != c.want {
			t.Errorf("case '%s' failed", c.desc)
		}
	}
}

// TestLaunchAppArgsIsValid tests validity cases for launchAppArgs
func TestLaunchAppArgsIsValid(t *testing.T) {
	cases := []appArgsTestCases{
		{
			desc: "a valid launchAppArgs object",
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 1,
					path:      "test/path.out",
				},
				name: "test-app",
			},
			want: true,
		},
		{
			desc: testDescInvalidProjectId,
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 0,
					path:      "test/path.out",
				},
				name: "test-app",
			},
			want: false,
		},
		{
			desc: "invalid path given (none)",
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 1,
					path:      "",
				},
				name: "test-app",
			},
			want: false,
		},
		{
			desc: "invalid name given (none)",
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 1,
					path:      "test/path.out",
				},
				name: "",
			},
			want: false,
		},
	}
	runAppArgsTestCase(t, cases)
}

// TestBaseAppArgsIsValid tests validity cases for baseAppArgs
func TestBaseAppArgsIsValid(t *testing.T) {
	cases := []appArgsTestCases{
		{
			desc: "a valid baseAppArgs object w/ both id & name",
			in: &baseAppArgs{
				projectId: 1,
				name:      "test-app",
				id:        1,
			},
			want: true,
		},
		{
			desc: "a valid baseAppArgs object w/ name",
			in: &baseAppArgs{
				projectId: 1,
				name:      "test-app",
			},
			want: true,
		},
		{
			desc: "a valid baseAppArgs object w/ id",
			in: &baseAppArgs{
				projectId: 1,
				id:        1,
			},
			want: true,
		},
		{
			desc: testDescInvalidProjectId,
			in: &baseAppArgs{
				projectId: 0,
				id:        1,
			},
			want: false,
		},
		{
			desc: "invalid, no id or name",
			in: &baseAppArgs{
				projectId: 1,
			},
			want: false,
		},
	}
	runAppArgsTestCase(t, cases)
}
<commit_msg>test: More refactoring to cleanup test code<commit_after>package command

import "testing"

// TestLaunchAppArgsIsValid tests validity cases for launchAppArgs
func TestLaunchAppArgsIsValid(t *testing.T) {
	cases := []dataTestCase{
		{
			desc: "a valid launchAppArgs object",
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 1,
					path:      "test/path.out",
				},
				name: "test-app",
			},
			want: true,
		},
		{
			desc: testDescInvalidProjectId,
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 0,
					path:      "test/path.out",
				},
				name: "test-app",
			},
			want: false,
		},
		{
			desc: "invalid path given (none)",
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 1,
					path:      "",
				},
				name: "test-app",
			},
			want: false,
		},
		{
			desc: "invalid name given (none)",
			in: &launchAppArgs{
				uploadFileArgs: uploadFileArgs{
					projectId: 1,
					path:      "test/path.out",
				},
				name: "",
			},
			want: false,
		},
	}
	runDataTestCase(t, cases)
}

// TestBaseAppArgsIsValid tests validity cases for baseAppArgs
func TestBaseAppArgsIsValid(t *testing.T) {
	cases := []dataTestCase{
		{
			desc: "a valid baseAppArgs object w/ both id & name",
			in: &baseAppArgs{
				projectId: 1,
				name:      "test-app",
				id:        1,
			},
			want: true,
		},
		{
			desc: "a valid baseAppArgs object w/ name",
			in: &baseAppArgs{
				projectId: 1,
				name:      "test-app",
			},
			want: true,
		},
		{
			desc: "a valid baseAppArgs object w/ id",
			in: &baseAppArgs{
				projectId: 1,
				id:        1,
			},
			want: true,
		},
		{
			desc: testDescInvalidProjectId,
			in: &baseAppArgs{
				projectId: 0,
				id:        1,
			},
			want: false,
		},
		{
			desc: "invalid, no id or name",
			in: &baseAppArgs{
				projectId: 1,
			},
			want: false,
		},
	}
	runDataTestCase(t, cases)
}
<|endoftext|>


<commit_before>package daemon

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/disorganizer/brig/daemon/wire"
	"github.com/disorganizer/brig/fuse"
	"github.com/disorganizer/brig/repo"
	"github.com/disorganizer/brig/transfer"
	"github.com/disorganizer/brig/transfer/moose"
	"github.com/disorganizer/brig/util/protocol"
	"github.com/disorganizer/brig/util/security"
	"github.com/gogo/protobuf/proto"
	"golang.org/x/net/context"
)

const (
	// MaxConnections is the upper limit of clients that may connect to a daemon
	// at the same time. Other client will wait in Accept().
	MaxConnections = 20
)

// This is just here to make the maxConnections prettier.
type allowOneConn struct{}

// Server is a TCP server that executed all commands
// on a single repository. Once the daemon is started, it
// attempts to open the repository, for which a password is needed.
type Server struct {
	// The repo we're working on
	Repo *repo.Repository

	// All mountpoints this daemon is serving:
	Mounts *fuse.MountTable

	// The metadata networkinglayer.
	MetaHost *transfer.Connector

	// signals (external and self triggered) arrive on this channel.
	signals chan os.Signal

	// Root context for this daemon
	ctx context.Context

	// TCP Listener for incoming connections:
	listener net.Listener

	// buffered channel with N places,
	// - every active connection holds one.
	maxConnections chan allowOneConn
}

// Summon creates a new up and running Server instance
func Summon(pwd, repoFolder string, port int) (*Server, error) {
	// Load the on-disk repository:
	log.Infof("Opening repo: %s", repoFolder)
	rep, err := repo.Open(pwd, repoFolder)
	if err != nil {
		log.Error("Could not load repository: ", err)
		return nil, err
	}

	// Listen for incoming connections.
	addr := fmt.Sprintf("localhost:%d", port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error("Error listening:", err.Error())
		return nil, err
	}

	if err != nil {
		return nil, err
	}

	// Close the listener when the application closes.
	log.Info("Listening on ", addr)
	ctx, cancel := context.WithCancel(context.Background())
	layer := moose.NewLayer(rep.IPFS, ctx)

	daemon := &Server{
		Repo:           rep,
		Mounts:         fuse.NewMountTable(rep.OwnStore),
		MetaHost:       transfer.NewConnector(layer, rep),
		signals:        make(chan os.Signal, 1),
		listener:       listener,
		maxConnections: make(chan allowOneConn, MaxConnections),
		ctx:            ctx,
	}

	go daemon.loop(cancel)

	if err := daemon.Connect(); err != nil {
		return nil, err
	}

	return daemon, nil
}

// Serve waits until the Server received a quit reason.
func (d *Server) Serve() {
	<-d.ctx.Done()

	if err := d.listener.Close(); err != nil {
		log.Warningf("daemon-close: cannot close listener: %v", err)
	}

	if err := d.Disconnect(); err != nil {
		log.Warningf("Could not shut down online services: %v", err)
	}

	if err := d.Mounts.Close(); err != nil {
		log.Errorf("daemon-close: error while closing mounts: %v", err)
	}

	if err := d.Repo.Close(); err != nil {
		log.Errorf("daemon-close: unable to close repository: %v", err)
	}
}

// Handle incoming connections:
func (d *Server) loop(cancel context.CancelFunc) {
	// Forward signals to the signals channel:
	signal.Notify(d.signals, os.Interrupt, os.Kill)

	// Reserve at least cap(d.maxConnections)
	for i := 0; i < cap(d.maxConnections); i++ {
		d.maxConnections <- allowOneConn{}
	}

	for {
		select {
		case <-d.signals:
			// Break the Serve() loop
			cancel()
			return
		case <-d.maxConnections:
			// Listen for an incoming connection.
			deadline := time.Now().Add(500 * time.Millisecond)
			err := d.listener.(*net.TCPListener).SetDeadline(deadline)
			if err != nil {
				log.Errorf("BUG: SetDeadline failed: %v", err)
				return
			}

			conn, err := d.listener.Accept()
			if err != nil && err.(*net.OpError).Timeout() {
				d.maxConnections <- allowOneConn{}
				continue
			}

			if err != nil {
				log.Errorf("Error in Accept(): %v", err)
				return
			}

			// Handle connections in a new goroutine.
			go d.handleConnection(d.ctx, conn)
		default:
			log.Infof("Max number of connections hit: %d", cap(d.maxConnections))
			time.Sleep(500 * time.Millisecond)
		}
	}
}

// Handles incoming requests:
func (d *Server) handleConnection(ctx context.Context, conn net.Conn) {
	// Make sure this connection count gets released:
	defer func() {
		if err := conn.Close(); err != nil {
			log.Debugf("daemon-loop: connection drop failed: %v", err)
		}

		d.maxConnections <- allowOneConn{}
	}()

	tnl, err := security.NewEllipticTunnel(conn)
	if err != nil {
		log.Error("Tunnel failed", err)
		return
	}

	p := protocol.NewProtocol(tnl, false)

	// Loop until client disconnect or dies otherwise:
	for {
		msg := &wire.Command{}
		if err := p.Recv(msg); err != nil {
			if err != io.EOF {
				log.Warning("daemon-recv: ", err)
			}
			return
		}

		log.Infof("recv: %s: %v", conn.RemoteAddr().String(), msg)
		d.handleCommand(ctx, msg, p)
	}
}

// Handles the actual incoming commands:
func (d *Server) handleCommand(ctx context.Context, cmd *wire.Command, p *protocol.Protocol) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Figure out which handler to call:
	handlerID := *(cmd.CommandType)
	handler, ok := handlerMap[handlerID]

	if !ok {
		log.Warningf("No handler for ID: %v", handlerID)
		return
	}

	resp, err := handler(d, ctx, cmd)

	// Empty response or error response:
	if resp == nil {
		resp = &wire.Response{}
	}

	resp.ResponseType = cmd.CommandType
	resp.Success = proto.Bool(false)

	if err != nil {
		resp.Error = proto.String(err.Error())
	} else {
		resp.Success = proto.Bool(true)
	}

	// Send the response back to the client:
	fmt.Println(resp)
	if err := p.Send(resp); err != nil {
		log.Warning("Unable to send response back to client: ", err)
	}
}

// Connect tries to connect the client and the ipfs daemon to the outside world.
func (sv *Server) Connect() error {
	if sv.IsOnline() {
		return nil
	}

	// Check if a previous offline mode was there:
	if err := sv.Repo.IPFS.Online(); err != nil {
		return err
	}

	if err := sv.MetaHost.Connect(); err != nil {
		log.Warningf("Unable to connect metadata client: %v", err)
		return err
	}

	log.Infof("Connected to the network")

	return nil
}

// Disconnect shuts down all store services that need an connection
// to the outside.
func (sv *Server) Disconnect() (err error) {
	if !sv.IsOnline() {
		return nil
	}

	log.Debugf("Disconnecting metadata host.")

	// Try to close metadata host, even if ipfs is still running:
	if err = sv.MetaHost.Disconnect(); err != nil {
		log.Warningf("Unable to disconnect metadata host: %v", err)
	}

	log.Debugf("Disconnecting ipfs daemon.")

	if err = sv.Repo.IPFS.Close(); err != nil {
		log.Warningf("Unable to close ipfs node: %v", err)
	}

	return err
}

// IsOnline checks if both meta host and ipfs is up and running.
func (sv *Server) IsOnline() bool {
	return sv.MetaHost.IsInOnlineMode() && sv.Repo.IPFS.IsOnline()
}
<commit_msg>daemon/server: Remove silly fmt.Println()<commit_after>package daemon

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/disorganizer/brig/daemon/wire"
	"github.com/disorganizer/brig/fuse"
	"github.com/disorganizer/brig/repo"
	"github.com/disorganizer/brig/transfer"
	"github.com/disorganizer/brig/transfer/moose"
	"github.com/disorganizer/brig/util/protocol"
	"github.com/disorganizer/brig/util/security"
	"github.com/gogo/protobuf/proto"
	"golang.org/x/net/context"
)

const (
	// MaxConnections is the upper limit of clients that may connect to a daemon
	// at the same time. Other client will wait in Accept().
	MaxConnections = 20
)

// This is just here to make the maxConnections prettier.
type allowOneConn struct{}

// Server is a TCP server that executed all commands
// on a single repository. Once the daemon is started, it
// attempts to open the repository, for which a password is needed.
type Server struct {
	// The repo we're working on
	Repo *repo.Repository

	// All mountpoints this daemon is serving:
	Mounts *fuse.MountTable

	// The metadata networkinglayer.
	MetaHost *transfer.Connector

	// signals (external and self triggered) arrive on this channel.
	signals chan os.Signal

	// Root context for this daemon
	ctx context.Context

	// TCP Listener for incoming connections:
	listener net.Listener

	// buffered channel with N places,
	// - every active connection holds one.
	maxConnections chan allowOneConn
}

// Summon creates a new up and running Server instance
func Summon(pwd, repoFolder string, port int) (*Server, error) {
	// Load the on-disk repository:
	log.Infof("Opening repo: %s", repoFolder)
	rep, err := repo.Open(pwd, repoFolder)
	if err != nil {
		log.Error("Could not load repository: ", err)
		return nil, err
	}

	// Listen for incoming connections.
	addr := fmt.Sprintf("localhost:%d", port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error("Error listening:", err.Error())
		return nil, err
	}

	if err != nil {
		return nil, err
	}

	// Close the listener when the application closes.
	log.Info("Listening on ", addr)
	ctx, cancel := context.WithCancel(context.Background())
	layer := moose.NewLayer(rep.IPFS, ctx)

	daemon := &Server{
		Repo:           rep,
		Mounts:         fuse.NewMountTable(rep.OwnStore),
		MetaHost:       transfer.NewConnector(layer, rep),
		signals:        make(chan os.Signal, 1),
		listener:       listener,
		maxConnections: make(chan allowOneConn, MaxConnections),
		ctx:            ctx,
	}

	go daemon.loop(cancel)

	if err := daemon.Connect(); err != nil {
		return nil, err
	}

	return daemon, nil
}

// Serve waits until the Server received a quit reason.
func (d *Server) Serve() {
	<-d.ctx.Done()

	if err := d.listener.Close(); err != nil {
		log.Warningf("daemon-close: cannot close listener: %v", err)
	}

	if err := d.Disconnect(); err != nil {
		log.Warningf("Could not shut down online services: %v", err)
	}

	if err := d.Mounts.Close(); err != nil {
		log.Errorf("daemon-close: error while closing mounts: %v", err)
	}

	if err := d.Repo.Close(); err != nil {
		log.Errorf("daemon-close: unable to close repository: %v", err)
	}
}

// Handle incoming connections:
func (d *Server) loop(cancel context.CancelFunc) {
	// Forward signals to the signals channel:
	signal.Notify(d.signals, os.Interrupt, os.Kill)

	// Reserve at least cap(d.maxConnections)
	for i := 0; i < cap(d.maxConnections); i++ {
		d.maxConnections <- allowOneConn{}
	}

	for {
		select {
		case <-d.signals:
			// Break the Serve() loop
			cancel()
			return
		case <-d.maxConnections:
			// Listen for an incoming connection.
			deadline := time.Now().Add(500 * time.Millisecond)
			err := d.listener.(*net.TCPListener).SetDeadline(deadline)
			if err != nil {
				log.Errorf("BUG: SetDeadline failed: %v", err)
				return
			}

			conn, err := d.listener.Accept()
			if err != nil && err.(*net.OpError).Timeout() {
				d.maxConnections <- allowOneConn{}
				continue
			}

			if err != nil {
				log.Errorf("Error in Accept(): %v", err)
				return
			}

			// Handle connections in a new goroutine.
			go d.handleConnection(d.ctx, conn)
		default:
			log.Infof("Max number of connections hit: %d", cap(d.maxConnections))
			time.Sleep(500 * time.Millisecond)
		}
	}
}

// Handles incoming requests:
func (d *Server) handleConnection(ctx context.Context, conn net.Conn) {
	// Make sure this connection count gets released:
	defer func() {
		if err := conn.Close(); err != nil {
			log.Debugf("daemon-loop: connection drop failed: %v", err)
		}

		d.maxConnections <- allowOneConn{}
	}()

	tnl, err := security.NewEllipticTunnel(conn)
	if err != nil {
		log.Error("Tunnel failed", err)
		return
	}

	p := protocol.NewProtocol(tnl, false)

	// Loop until client disconnect or dies otherwise:
	for {
		msg := &wire.Command{}
		if err := p.Recv(msg); err != nil {
			if err != io.EOF {
				log.Warning("daemon-recv: ", err)
			}
			return
		}

		log.Infof("recv: %s: %v", conn.RemoteAddr().String(), msg)
		d.handleCommand(ctx, msg, p)
	}
}

// Handles the actual incoming commands:
func (d *Server) handleCommand(ctx context.Context, cmd *wire.Command, p *protocol.Protocol) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Figure out which handler to call:
	handlerID := *(cmd.CommandType)
	handler, ok := handlerMap[handlerID]

	if !ok {
		log.Warningf("No handler for ID: %v", handlerID)
		return
	}

	resp, err := handler(d, ctx, cmd)

	// Empty response or error response:
	if resp == nil {
		resp = &wire.Response{}
	}

	resp.ResponseType = cmd.CommandType
	resp.Success = proto.Bool(false)

	if err != nil {
		resp.Error = proto.String(err.Error())
	} else {
		resp.Success = proto.Bool(true)
	}

	// Send the response back to the client:
	if err := p.Send(resp); err != nil {
		log.Warning("Unable to send response back to client: ", err)
	}
}

// Connect tries to connect the client and the ipfs daemon to the outside world.
func (sv *Server) Connect() error {
	if sv.IsOnline() {
		return nil
	}

	// Check if a previous offline mode was there:
	if err := sv.Repo.IPFS.Online(); err != nil {
		return err
	}

	if err := sv.MetaHost.Connect(); err != nil {
		log.Warningf("Unable to connect metadata client: %v", err)
		return err
	}

	log.Infof("Connected to the network")

	return nil
}

// Disconnect shuts down all store services that need an connection
// to the outside.
func (sv *Server) Disconnect() (err error) {
	if !sv.IsOnline() {
		return nil
	}

	log.Debugf("Disconnecting metadata host.")

	// Try to close metadata host, even if ipfs is still running:
	if err = sv.MetaHost.Disconnect(); err != nil {
		log.Warningf("Unable to disconnect metadata host: %v", err)
	}

	log.Debugf("Disconnecting ipfs daemon.")

	if err = sv.Repo.IPFS.Close(); err != nil {
		log.Warningf("Unable to close ipfs node: %v", err)
	}

	return err
}

// IsOnline checks if both meta host and ipfs is up and running.
func (sv *Server) IsOnline() bool {
	return sv.MetaHost.IsInOnlineMode() && sv.Repo.IPFS.IsOnline()
}
<|endoftext|>


<commit_before>// +build OMIT

package main

import (
	"fmt"
	"math"
)

func main() {
	fmt.Printf("Now you have %g problems.", math.Nextafter(2, 3))
}
<commit_msg>[x/tour] content: update code in imports example<commit_after>// +build OMIT

package main

import (
	"fmt"
	"math"
)

func main() {
	fmt.Printf("Now you have %g problems.", math.Sqrt(7))
}
<|endoftext|>


<commit_before>package state

import (
	"errors"
	"fmt"
	"strings"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state/presence"
	"launchpad.net/juju-core/state/watcher"
)

// Info encapsulates information about cluster of
// servers holding juju state and can be used to make a
// connection to that cluster.
type Info struct {
	// Addrs gives the addresses of the MongoDB servers for the state.
	// Each address should be in the form address:port.
	Addrs []string

	// UseSSH specifies whether MongoDB should be contacted through an
	// SSH tunnel.
	UseSSH bool
}

// Open connects to the server described by the given
// info, waits for it to be initialized, and returns a new State
// representing the environment connected to.
func Open(info *Info) (*State, error) {
	log.Printf("state: opening state; mongo addresses: %q", info.Addrs)
	if len(info.Addrs) == 0 {
		return nil, errors.New("no mongo addresses")
	}
	if !info.UseSSH {
		session, err := mgo.DialWithTimeout(strings.Join(info.Addrs, ","), 1 * time.Minute)
		if err != nil {
			return nil, err
		}
		return newState(session, nil)
	}
	if len(info.Addrs) > 1 {
		return nil, errors.New("ssh connect does not support multiple addresses")
	}
	fwd, session, err := sshDial(info.Addrs[0], "")
	if err != nil {
		return nil, err
	}
	st, err := newState(session, fwd)
	if err != nil {
		return nil, err
	}
	return st, err
}

// Initialize sets up an initial empty state and returns it.
// This needs to be performed only once for a given environment.
func Initialize(info *Info, cfg *config.Config) (*State, error) {
	st, err := Open(info)
	if err != nil {
		return nil, err
	}
	// A valid environment config is used as a signal that the
	// state has already been initalized. If this is the case
	// do nothing.
	_, err = st.EnvironConfig()
	if err == nil {
		return st, nil
	}
	log.Printf("state: storing no-secrets environment configuration")
	err = st.SetEnvironConfig(cfg)
	if err != nil {
		st.Close()
		return nil, err
	}
	return st, nil
}

var indexes = []struct {
	collection string
	key        []string
}{
	// After the first public release, do not remove entries from here
	// without adding them to a list of indexes to drop, to ensure
	// old databases are modified to have the correct indexes.
	{"relations", []string{"endpoints.relationname"}},
	{"relations", []string{"endpoints.servicename"}},
	{"units", []string{"service"}},
	{"units", []string{"principal"}},
	{"units", []string{"machineid"}},
}

// The capped collection used for transaction logs defaults to 200MB.
// It's tweaked in export_test.go to 1MB to avoid the overhead of
// creating and deleting the large file repeatedly.
var (
	logSize      = 200000000
	logSizeTests = 1000000
)

func newState(session *mgo.Session, fwd *sshForwarder) (*State, error) {
	db := session.DB("juju")
	pdb := session.DB("presence")
	st := &State{
		db:             db,
		charms:         db.C("charms"),
		machines:       db.C("machines"),
		relations:      db.C("relations"),
		relationScopes: db.C("relationscopes"),
		services:       db.C("services"),
		settings:       db.C("settings"),
		units:          db.C("units"),
		presence:       pdb.C("presence"),
		fwd:            fwd,
	}
	log := db.C("txns.log")
	info := mgo.CollectionInfo{Capped: true, MaxBytes: logSize}
	// The lack of error code for this error was reported upstream:
	//     https://jira.mongodb.org/browse/SERVER-6992
	if err := log.Create(&info); err != nil && err.Error() != "collection already exists" {
		return nil, fmt.Errorf("cannot create log collection: %v", err)
	}
	st.runner = txn.NewRunner(db.C("txns"))
	st.runner.ChangeLog(db.C("txns.log"))
	st.watcher = watcher.New(db.C("txns.log"))
	st.pwatcher = presence.NewWatcher(pdb.C("presence"))
	for _, item := range indexes {
		index := mgo.Index{Key: item.key}
		if err := db.C(item.collection).EnsureIndex(index); err != nil {
			return nil, fmt.Errorf("cannot create database index: %v", err)
		}
	}
	return st, nil
}

func (st *State) Close() error {
	err1 := st.watcher.Stop()
	err2 := st.pwatcher.Stop()
	st.db.Session.Close()
	for _, err := range []error{err1, err2} {
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>state: use timeout when dialing, even without ssh<commit_after>package state

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state/presence"
	"launchpad.net/juju-core/state/watcher"
)

// Info encapsulates information about cluster of
// servers holding juju state and can be used to make a
// connection to that cluster.
type Info struct {
	// Addrs gives the addresses of the MongoDB servers for the state.
	// Each address should be in the form address:port.
	Addrs []string

	// UseSSH specifies whether MongoDB should be contacted through an
	// SSH tunnel.
	UseSSH bool
}

// Open connects to the server described by the given
// info, waits for it to be initialized, and returns a new State
// representing the environment connected to.
func Open(info *Info) (*State, error) {
	log.Printf("state: opening state; mongo addresses: %q", info.Addrs)
	if len(info.Addrs) == 0 {
		return nil, errors.New("no mongo addresses")
	}
	if !info.UseSSH {
		session, err := mgo.DialWithTimeout(strings.Join(info.Addrs, ","), 10 * time.Minute)
		if err != nil {
			return nil, err
		}
		return newState(session, nil)
	}
	if len(info.Addrs) > 1 {
		return nil, errors.New("ssh connect does not support multiple addresses")
	}
	fwd, session, err := sshDial(info.Addrs[0], "")
	if err != nil {
		return nil, err
	}
	st, err := newState(session, fwd)
	if err != nil {
		return nil, err
	}
	return st, err
}

// Initialize sets up an initial empty state and returns it.
// This needs to be performed only once for a given environment.
func Initialize(info *Info, cfg *config.Config) (*State, error) {
	st, err := Open(info)
	if err != nil {
		return nil, err
	}
	// A valid environment config is used as a signal that the
	// state has already been initalized. If this is the case
	// do nothing.
	_, err = st.EnvironConfig()
	if err == nil {
		return st, nil
	}
	log.Printf("state: storing no-secrets environment configuration")
	err = st.SetEnvironConfig(cfg)
	if err != nil {
		st.Close()
		return nil, err
	}
	return st, nil
}

var indexes = []struct {
	collection string
	key        []string
}{
	// After the first public release, do not remove entries from here
	// without adding them to a list of indexes to drop, to ensure
	// old databases are modified to have the correct indexes.
	{"relations", []string{"endpoints.relationname"}},
	{"relations", []string{"endpoints.servicename"}},
	{"units", []string{"service"}},
	{"units", []string{"principal"}},
	{"units", []string{"machineid"}},
}

// The capped collection used for transaction logs defaults to 200MB.
// It's tweaked in export_test.go to 1MB to avoid the overhead of
// creating and deleting the large file repeatedly.
var (
	logSize      = 200000000
	logSizeTests = 1000000
)

func newState(session *mgo.Session, fwd *sshForwarder) (*State, error) {
	db := session.DB("juju")
	pdb := session.DB("presence")
	st := &State{
		db:             db,
		charms:         db.C("charms"),
		machines:       db.C("machines"),
		relations:      db.C("relations"),
		relationScopes: db.C("relationscopes"),
		services:       db.C("services"),
		settings:       db.C("settings"),
		units:          db.C("units"),
		presence:       pdb.C("presence"),
		fwd:            fwd,
	}
	log := db.C("txns.log")
	info := mgo.CollectionInfo{Capped: true, MaxBytes: logSize}
	// The lack of error code for this error was reported upstream:
	//     https://jira.mongodb.org/browse/SERVER-6992
	if err := log.Create(&info); err != nil && err.Error() != "collection already exists" {
		return nil, fmt.Errorf("cannot create log collection: %v", err)
	}
	st.runner = txn.NewRunner(db.C("txns"))
	st.runner.ChangeLog(db.C("txns.log"))
	st.watcher = watcher.New(db.C("txns.log"))
	st.pwatcher = presence.NewWatcher(pdb.C("presence"))
	for _, item := range indexes {
		index := mgo.Index{Key: item.key}
		if err := db.C(item.collection).EnsureIndex(index); err != nil {
			return nil, fmt.Errorf("cannot create database index: %v", err)
		}
	}
	return st, nil
}

func (st *State) Close() error {
	err1 := st.watcher.Stop()
	err2 := st.pwatcher.Stop()
	st.db.Session.Close()
	for _, err := range []error{err1, err2} {
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>