<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"strings"
	"testing"
	"time"

	"github.com/cyberark/conjur-api-go/conjurapi"
	conjur_authn "github.com/cyberark/conjur-api-go/conjurapi/authn"
	. "github.com/smartystreets/goconvey/convey"
)

func RunCommand(name string, arg ...string) (bytes.Buffer, bytes.Buffer, error) {
	cmd := exec.Command(name, arg...)
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	return stdout, stderr, err
}

func WithoutArgs() {
	Convey("Given summon-conjur is run with no arguments", func() {
		_, stderr, err := RunCommand(PackageName)

		Convey("Returns with error", func() {
			So(err, ShouldNotBeNil)
			So(stderr.String(), ShouldEqual, "A variable name or version flag must be given as the first and only argument!")
		})
	})
}

const PackageName = "summon-conjur"

func TestPackage(t *testing.T) {
	ApplianceURL := os.Getenv("CONJUR_APPLIANCE_URL")
	Account := os.Getenv("CONJUR_ACCOUNT")
	Login := os.Getenv("CONJUR_AUTHN_LOGIN")
	APIKey := os.Getenv("CONJUR_AUTHN_API_KEY")

	ApplianceURL_V4 := os.Getenv("CONJUR_V4_APPLIANCE_URL")
	SSLCert_V4 := os.Getenv("CONJUR_V4_SSL_CERTIFICATE")
	Login_V4 := os.Getenv("CONJUR_V4_AUTHN_LOGIN")
	APIKey_V4 := os.Getenv("CONJUR_V4_AUTHN_API_KEY")

	Path := os.Getenv("PATH")

	Convey("Given a compiled summon-conjur package", t, func() {
		Convey("Given no configuration and no authentication information", func() {
			e := ClearEnv()
			defer e.RestoreEnv()
			os.Setenv("PATH", Path)

			WithoutArgs()
		})

		Convey("Given valid configuration", func() {
			e := ClearEnv()
			defer e.RestoreEnv()
			os.Setenv("PATH", Path)

			os.Setenv("CONJUR_APPLIANCE_URL", ApplianceURL)
			os.Setenv("CONJUR_ACCOUNT", Account)

			Convey("Given valid APIKey credentials", func() {
				os.Setenv("CONJUR_AUTHN_LOGIN", Login)
				os.Setenv("CONJUR_AUTHN_API_KEY", APIKey)

				WithoutArgs()

				Convey("Retrieves existent variable's defined value", func() {
					variableIdentifier := "db/password"
					secretValue := fmt.Sprintf("secret-value-%v", rand.Intn(123456))
					policy := fmt.Sprintf(`
- !variable %s
`, variableIdentifier)

					config := conjurapi.Config{
						ApplianceURL: ApplianceURL,
						Account:      Account,
					}
					conjur, _ := conjurapi.NewClientFromKey(config, conjur_authn.LoginPair{Login, APIKey})

					conjur.LoadPolicy(
						"root",
						strings.NewReader(policy),
					)
					defer conjur.LoadPolicy(
						"root",
						strings.NewReader(""),
					)

					conjur.AddSecret(variableIdentifier, secretValue)

					stdout, _, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldBeNil)
					So(stdout.String(), ShouldEqual, secretValue)
				})

				Convey("Returns 404 on non-existent variable", func() {
					variableIdentifier := "non-existent-variable"

					_, stderr, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldNotBeNil)
					So(stderr.String(), ShouldContainSubstring, "404")
				})

				Convey("Given a non-existent Login is set", func() {
					os.Setenv("CONJUR_AUTHN_LOGIN", "non-existent-user")

					Convey("Returns 401", func() {
						variableIdentifier := "existent-or-non-existent-variable"

						_, stderr, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldNotBeNil)
						So(stderr.String(), ShouldContainSubstring, "401")
					})
				})
			})

			Convey("Given valid TokenFile credentials", func() {

				getToken := fmt.Sprintf(`
token=$(curl --data "%s" "$CONJUR_APPLIANCE_URL/authn/$CONJUR_ACCOUNT/%s/authenticate")
echo $token
`, APIKey, Login)
				stdout, _, err := RunCommand("bash", "-c", getToken)

				So(err, ShouldBeNil)
				So(stdout.String(), ShouldContainSubstring, "signature")

				tokenFile, _ := ioutil.TempFile("", "existent-token-file")
				tokenFileName := tokenFile.Name()
				tokenFileContents := stdout.String()
				os.Remove(tokenFileName)
				go func() {
					ioutil.WriteFile(tokenFileName, []byte(tokenFileContents), 0600)
				}()
				defer os.Remove(tokenFileName)

				os.Setenv("CONJUR_AUTHN_TOKEN_FILE", tokenFileName)

				WithoutArgs()

				Convey("Retrieves existent variable's defined value", func() {
					variableIdentifier := "db/password"
					secretValue := fmt.Sprintf("secret-value-%v", rand.Intn(123456))
					policy := fmt.Sprintf(`
- !variable %s
`, variableIdentifier)

					config := conjurapi.Config{
						ApplianceURL: ApplianceURL,
						Account:      Account,
					}
					conjur, _ := conjurapi.NewClientFromKey(config, conjur_authn.LoginPair{Login, APIKey})

					conjur.LoadPolicy(
						"root",
						strings.NewReader(policy),
					)
					defer conjur.LoadPolicy(
						"root",
						strings.NewReader(""),
					)

					conjur.AddSecret(variableIdentifier, secretValue)

					stdout, _, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldBeNil)
					So(stdout.String(), ShouldEqual, secretValue)
				})

				Convey("Returns 404 on non-existent variable", func() {
					variableIdentifier := "non-existent-variable"

					_, stderr, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldNotBeNil)
					So(stderr.String(), ShouldContainSubstring, "404")
				})

				Convey("Given a non-existent TokenFile is set", func() {
					os.Setenv("CONJUR_AUTHN_TOKEN_FILE", "non-existent-token-file")

					Convey("Waits for longer than a second", func() {
						timeout := time.After(1 * time.Second)
						unexpected_response := make(chan int)

						go func() {
							variableIdentifier := "existent-or-non-existent-variable"
							RunCommand(PackageName, variableIdentifier)
							unexpected_response <- 1
						}()

						select {
						case <-unexpected_response:
							So("receive unexpected response", ShouldEqual, "not receive unexpected response")
						case <-timeout:
							So(true, ShouldEqual, true)
						}
					})
				})
			})

			Convey("Given no authentication credentials", func() {

				WithoutArgs()

				Convey("Returns with error on non-existent variable", func() {
					variableIdentifier := "existent-or-non-existent-variable"

					_, stderr, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldNotBeNil)
					So(stderr.String(), ShouldContainSubstring, "at least one authentication strategy")
				})
			})
		})

		Convey("Given a v4 appliance", func() {
			Convey("Given valid configuration", func() {
				e := ClearEnv()
				defer e.RestoreEnv()
				os.Setenv("PATH", Path)

				os.Setenv("CONJUR_MAJOR_VERSION", "4")
				os.Setenv("CONJUR_APPLIANCE_URL", ApplianceURL_V4)
				os.Setenv("CONJUR_ACCOUNT", Account)
				os.Setenv("CONJUR_SSL_CERTIFICATE", SSLCert_V4)

				Convey("Given valid APIKey credentials", func() {
					os.Setenv("CONJUR_AUTHN_LOGIN", Login_V4)
					os.Setenv("CONJUR_AUTHN_API_KEY", APIKey_V4)

					WithoutArgs()

					Convey("Retrieves existent variable's defined value", func() {
						variableIdentifier := "existent-variable-with-defined-value"
						secretValue := "existent-variable-defined-value"

						stdout, _, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldBeNil)
						So(stdout.String(), ShouldEqual, secretValue)
					})

					Convey("Returns 404 on existent-variable-undefined-value", func() {
						variableIdentifier := "existent-variable-with-undefined-value"

						_, stderr, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldNotBeNil)
						So(stderr.String(), ShouldContainSubstring, "404")
					})

					Convey("Returns 404 on non-existent variable", func() {
						variableIdentifier := "non-existent-variable"

						_, stderr, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldNotBeNil)
						So(stderr.String(), ShouldContainSubstring, "404")
					})

				})
			})
		})

	})
}
<commit_msg>Get the rest of the tests working.<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"strings"
	"testing"
	"time"

	"github.com/cyberark/conjur-api-go/conjurapi"
	conjur_authn "github.com/cyberark/conjur-api-go/conjurapi/authn"
	. "github.com/smartystreets/goconvey/convey"
)

func RunCommand(name string, arg ...string) (bytes.Buffer, bytes.Buffer, error) {
	cmd := exec.Command(name, arg...)
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	return stdout, stderr, err
}

func WithoutArgs() {
	Convey("Given summon-conjur is run with no arguments", func() {
		_, stderr, err := RunCommand(PackageName)

		Convey("Returns with error", func() {
			So(err, ShouldNotBeNil)
			So(stderr.String(), ShouldEqual, `Usage of summon-conjur:
  -h, --help
	show help (default: false)
  -V, --version
	show version (default: false)
  -v, --verbose
	be verbose (default: false)
`)
		})
	})
}

const PackageName = "summon-conjur"

func TestPackage(t *testing.T) {
	ApplianceURL := os.Getenv("CONJUR_APPLIANCE_URL")
	Account := os.Getenv("CONJUR_ACCOUNT")
	Login := os.Getenv("CONJUR_AUTHN_LOGIN")
	APIKey := os.Getenv("CONJUR_AUTHN_API_KEY")

	ApplianceURL_V4 := os.Getenv("CONJUR_V4_APPLIANCE_URL")
	SSLCert_V4 := os.Getenv("CONJUR_V4_SSL_CERTIFICATE")
	APIKey_V4 := os.Getenv("CONJUR_V4_AUTHN_API_KEY")

	Path := os.Getenv("PATH")

	Convey("Given a compiled summon-conjur package", t, func() {
		Convey("Given no configuration and no authentication information", func() {
			e := ClearEnv()
			defer e.RestoreEnv()
			os.Setenv("PATH", Path)

			WithoutArgs()
		})

		Convey("Given valid configuration", func() {
			e := ClearEnv()
			defer e.RestoreEnv()
			os.Setenv("PATH", Path)

			os.Setenv("CONJUR_APPLIANCE_URL", ApplianceURL)
			os.Setenv("CONJUR_ACCOUNT", Account)

			Convey("Given valid APIKey credentials", func() {
				os.Setenv("CONJUR_AUTHN_LOGIN", Login)
				os.Setenv("CONJUR_AUTHN_API_KEY", APIKey)

				WithoutArgs()

				Convey("Retrieves existing variable's defined value", func() {
					variableIdentifier := "db/password"
					secretValue := fmt.Sprintf("secret-value-%v", rand.Intn(123456))
					policy := fmt.Sprintf(`
- !variable %s
`, variableIdentifier)

					config := conjurapi.Config{
						ApplianceURL: ApplianceURL,
						Account:      Account,
					}
					conjur, _ := conjurapi.NewClientFromKey(config, conjur_authn.LoginPair{Login, APIKey})

					conjur.LoadPolicy(
						conjurapi.PolicyModePost,
						"root",
						strings.NewReader(policy),
					)
					defer conjur.LoadPolicy(
						conjurapi.PolicyModePut,
						"root",
						strings.NewReader(""),
					)

					conjur.AddSecret(variableIdentifier, secretValue)

					stdout, _, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldBeNil)
					So(stdout.String(), ShouldEqual, secretValue)
				})

				Convey("Returns error on non-existent variable", func() {
					variableIdentifier := "non-existent-variable"

					_, stderr, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldNotBeNil)
					So(stderr.String(), ShouldContainSubstring, "not found")
				})

				Convey("Given a non-existent Login is set", func() {
					os.Setenv("CONJUR_AUTHN_LOGIN", "non-existent-user")

					Convey("Returns 401", func() {
						variableIdentifier := "existent-or-non-existent-variable"

						_, stderr, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldNotBeNil)
						So(stderr.String(), ShouldContainSubstring, "401")
					})
				})
			})

			Convey("Given valid TokenFile credentials", func() {

				getToken := fmt.Sprintf(`
token=$(curl --data "%s" "$CONJUR_APPLIANCE_URL/authn/$CONJUR_ACCOUNT/%s/authenticate")
echo $token
`, APIKey, Login)
				stdout, _, err := RunCommand("bash", "-c", getToken)

				So(err, ShouldBeNil)
				So(stdout.String(), ShouldContainSubstring, "signature")

				tokenFile, _ := ioutil.TempFile("", "existent-token-file")
				tokenFileName := tokenFile.Name()
				tokenFileContents := stdout.String()
				os.Remove(tokenFileName)
				go func() {
					ioutil.WriteFile(tokenFileName, []byte(tokenFileContents), 0600)
				}()
				defer os.Remove(tokenFileName)

				os.Setenv("CONJUR_AUTHN_TOKEN_FILE", tokenFileName)

				WithoutArgs()

				Convey("Retrieves existent variable's defined value", func() {
					variableIdentifier := "db/password"
					secretValue := fmt.Sprintf("secret-value-%v", rand.Intn(123456))
					policy := fmt.Sprintf(`
- !variable %s
`, variableIdentifier)

					config := conjurapi.Config{
						ApplianceURL: ApplianceURL,
						Account:      Account,
					}
					conjur, _ := conjurapi.NewClientFromKey(config, conjur_authn.LoginPair{Login, APIKey})

					conjur.LoadPolicy(
						conjurapi.PolicyModePost,
						"root",
						strings.NewReader(policy),
					)
					defer conjur.LoadPolicy(
						conjurapi.PolicyModePut,
						"root",
						strings.NewReader(""),
					)

					conjur.AddSecret(variableIdentifier, secretValue)

					stdout, _, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldBeNil)
					So(stdout.String(), ShouldEqual, secretValue)
				})

				Convey("Returns error on non-existent variable", func() {
					variableIdentifier := "non-existent-variable"

					_, stderr, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldNotBeNil)
					So(stderr.String(), ShouldContainSubstring, "not found in account")
				})

				Convey("Given a non-existent TokenFile is set", func() {
					os.Setenv("CONJUR_AUTHN_TOKEN_FILE", "non-existent-token-file")

					Convey("Waits for longer than a second", func() {
						timeout := time.After(1 * time.Second)
						unexpected_response := make(chan int)

						go func() {
							variableIdentifier := "existent-or-non-existent-variable"
							RunCommand(PackageName, variableIdentifier)
							unexpected_response <- 1
						}()

						select {
						case <-unexpected_response:
							So("receive unexpected response", ShouldEqual, "not receive unexpected response")
						case <-timeout:
							So(true, ShouldEqual, true)
						}
					})
				})
			})

			Convey("Given no authentication credentials", func() {

				WithoutArgs()

				Convey("Returns with error on non-existent variable", func() {
					variableIdentifier := "existent-or-non-existent-variable"

					_, stderr, err := RunCommand(PackageName, variableIdentifier)

					So(err, ShouldNotBeNil)
					So(stderr.String(), ShouldContainSubstring, "at least one authentication strategy")
				})
			})
		})

		Convey("Given a v4 appliance", func() {
			Convey("Given valid configuration", func() {
				e := ClearEnv()
				defer e.RestoreEnv()
				os.Setenv("PATH", Path)

				os.Setenv("CONJUR_MAJOR_VERSION", "4")
				os.Setenv("CONJUR_APPLIANCE_URL", ApplianceURL_V4)
				os.Setenv("CONJUR_ACCOUNT", Account)
				os.Setenv("CONJUR_AUTHN_LOGIN", Login)
				os.Setenv("CONJUR_SSL_CERTIFICATE", SSLCert_V4)

				Convey("Given valid APIKey credentials", func() {
					os.Setenv("CONJUR_AUTHN_API_KEY", APIKey_V4)

					WithoutArgs()

					Convey("Retrieves existing variable's defined value", func() {
						variableIdentifier := "existent-variable-with-defined-value"
						secretValue := "existent-variable-defined-value"

						stdout, _, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldBeNil)
						So(stdout.String(), ShouldEqual, secretValue)
					})

					Convey("Returns error on existent-variable-undefined-value", func() {
						variableIdentifier := "existent-variable-with-undefined-value"

						_, stderr, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldNotBeNil)
						So(stderr.String(), ShouldContainSubstring, "Not Found")
					})

					Convey("Returns error on non-existent variable", func() {
						variableIdentifier := "non-existent-variable"

						_, stderr, err := RunCommand(PackageName, variableIdentifier)

						So(err, ShouldNotBeNil)
						So(stderr.String(), ShouldContainSubstring, "not found")
					})

				})
			})
		})

	})
}
<|endoftext|>


<commit_before>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package wallet

import (
	"bytes"
	"os"
	"testing"

	"github.com/FactomProject/factom"
)

func TestNewWallet(t *testing.T) {
	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w1, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// check that the seed got written
	w1.lock.RLock()
	seed, err := w1.ldb.Get(seedDBKey, nil)
	if err != nil {
		t.Error(err)
	}
	w1.lock.RUnlock()
	if len(seed) != 64 {
		t.Errorf("stored db seed is the wrong length: %x", seed)
	}
	if bytes.Equal(seed, make([]byte, 64)) {
		t.Errorf("stored db seed is blank")
	}

	if err := w1.Close(); err != nil {
		t.Error(err)
	}

	// try and create a new database where one already exists
	if _, err := NewWallet(dbpath); err == nil {
		t.Errorf("NewWallet did not report error on existing path: %s", dbpath)
	}

	// remove the testing db
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestOpenWallet(t *testing.T) {
	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w1, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}
	w1.Close()

	// make sure we can open the db
	w2, err := OpenWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// check that the seed is there
	w2.lock.RLock()
	seed, err := w2.ldb.Get(seedDBKey, nil)
	if err != nil {
		t.Error(err)
	}
	w2.lock.RUnlock()
	if len(seed) != 64 {
		t.Errorf("stored db seed is the wrong length: %x", seed)
	}
	if bytes.Equal(seed, make([]byte, 64)) {
		t.Errorf("stored db seed is blank")
	}

	if err := w2.Close(); err != nil {
		t.Error(err)
	}

	// remove the testing db
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestPutECAddress(t *testing.T) {
	zSec := "Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG"

	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// write a new ec address to the db
	e, err := factom.GetECAddress(zSec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutECAddress(e); err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetECAddress(e.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestPutFCTAddress(t *testing.T) {
	zSec := "Fs1KWJrpLdfucvmYwN2nWrwepLn8ercpMbzXshd1g8zyhKXLVLWj"

	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// write a new fct address to the db
	f, err := factom.GetFactoidAddress(zSec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutFCTAddress(f); err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetFCTAddress(f.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestGenerateECAddress(t *testing.T) {
	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// Generate a new ec address
	e, err := w.GenerateECAddress()
	if err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetECAddress(e.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestGenerateFCTAddress(t *testing.T) {
	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// Generate a new fct address
	f, err := w.GenerateFCTAddress()
	if err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetFCTAddress(f.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestGetAllAddresses(t *testing.T) {
	e1Sec := "Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG"
	f1Sec := "Fs1KWJrpLdfucvmYwN2nWrwepLn8ercpMbzXshd1g8zyhKXLVLWj"
	e2Sec := "Es4NQHwo8F4Z4oMnVwndtjV1rzZN3t5pP5u5jtdgiR1RA6FH4Tmc"
	f2Sec := "Fs3GFV6GNV6ar4b8eGcQWpGFbFtkNWKfEPdbywmha8ez5p7XMJyk"
	correctLen := 2
	dbpath := os.TempDir() + "/ldb1"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// write a new ec address to the db
	e1, err := factom.GetECAddress(e1Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutECAddress(e1); err != nil {
		t.Error(err)
	}
	e2, err := factom.GetECAddress(e2Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutECAddress(e2); err != nil {
		t.Error(err)
	}

	// write a new fct address to the db
	f1, err := factom.GetFactoidAddress(f1Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutFCTAddress(f1); err != nil {
		t.Error(err)
	}
	f2, err := factom.GetFactoidAddress(f2Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutFCTAddress(f2); err != nil {
		t.Error(err)
	}

	// get all addresses out of db
	fs, es, err := w.GetAllAddresses()
	if err != nil {
		t.Error(err)
	} else if fs == nil {
		t.Errorf("No Factoid address was retrived")
	} else if es == nil {
		t.Errorf("No EC address was retrived")
	}
	// check that all the addresses are there
	if len(fs) != correctLen {
		t.Errorf("Wrong number of factoid addesses were retrived: %v", fs)
	}
	if len(es) != correctLen {
		t.Errorf("Wrong number of ec addesses were retrived: %v", es)
	}

	// print the addresses
	for _, f := range fs {
		t.Logf("%s %s", f, f.SecString())
	}
	for _, e := range es {
		t.Logf("%s %s", e, e.SecString())
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}
<commit_msg>updated tests<commit_after>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package wallet

import (
	"bytes"
	"os"
	"testing"

	"github.com/FactomProject/factom"
)

func TestNewWallet(t *testing.T) {
	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w1, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// check that the seed got written
	w1.lock.RLock()
	seed, err := w1.ldb.Get(seedDBKey, nil)
	if err != nil {
		t.Error(err)
	}
	w1.lock.RUnlock()
	if len(seed) != 64 {
		t.Errorf("stored db seed is the wrong length: %x", seed)
	}
	if bytes.Equal(seed, make([]byte, 64)) {
		t.Errorf("stored db seed is blank")
	}

	if err := w1.Close(); err != nil {
		t.Error(err)
	}

	// try and create a new database where one already exists
	if _, err := NewWallet(dbpath); err == nil {
		t.Errorf("NewWallet did not report error on existing path: %s", dbpath)
	}

	// remove the testing db
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestOpenWallet(t *testing.T) {
	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w1, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}
	w1.Close()

	// make sure we can open the db
	w2, err := OpenWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// check that the seed is there
	w2.lock.RLock()
	seed, err := w2.ldb.Get(seedDBKey, nil)
	if err != nil {
		t.Error(err)
	}
	w2.lock.RUnlock()
	if len(seed) != 64 {
		t.Errorf("stored db seed is the wrong length: %x", seed)
	}
	if bytes.Equal(seed, make([]byte, 64)) {
		t.Errorf("stored db seed is blank")
	}

	if err := w2.Close(); err != nil {
		t.Error(err)
	}

	// remove the testing db
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestPutECAddress(t *testing.T) {
	zSec := "Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG"

	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// write a new ec address to the db
	e, err := factom.GetECAddress(zSec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutECAddress(e); err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetECAddress(e.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestPutFCTAddress(t *testing.T) {
	zSec := "Fs1KWJrpLdfucvmYwN2nWrwepLn8ercpMbzXshd1g8zyhKXLVLWj"

	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// write a new fct address to the db
	f, err := factom.GetFactoidAddress(zSec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutFCTAddress(f); err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetFCTAddress(f.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestGenerateECAddress(t *testing.T) {
	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// Generate a new ec address
	e, err := w.GenerateECAddress()
	if err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetECAddress(e.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestGenerateFCTAddress(t *testing.T) {
	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// Generate a new fct address
	f, err := w.GenerateFCTAddress()
	if err != nil {
		t.Error(err)
	}

	// Check that the address was written into the db
	if _, err := w.GetFCTAddress(f.PubString()); err != nil {
		t.Error(err)
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}

func TestGetAllAddresses(t *testing.T) {
	e1Sec := "Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG"
	f1Sec := "Fs1KWJrpLdfucvmYwN2nWrwepLn8ercpMbzXshd1g8zyhKXLVLWj"
	e2Sec := "Es4NQHwo8F4Z4oMnVwndtjV1rzZN3t5pP5u5jtdgiR1RA6FH4Tmc"
	f2Sec := "Fs3GFV6GNV6ar4b8eGcQWpGFbFtkNWKfEPdbywmha8ez5p7XMJyk"
	correctLen := 2
	dbpath := os.TempDir() + "/test_wallet-01"

	// create a new database
	w, err := NewWallet(dbpath)
	if err != nil {
		t.Error(err)
	}

	// write a new ec address to the db
	e1, err := factom.GetECAddress(e1Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutECAddress(e1); err != nil {
		t.Error(err)
	}
	e2, err := factom.GetECAddress(e2Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutECAddress(e2); err != nil {
		t.Error(err)
	}

	// write a new fct address to the db
	f1, err := factom.GetFactoidAddress(f1Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutFCTAddress(f1); err != nil {
		t.Error(err)
	}
	f2, err := factom.GetFactoidAddress(f2Sec)
	if err != nil {
		t.Error(err)
	}
	if err := w.PutFCTAddress(f2); err != nil {
		t.Error(err)
	}

	// get all addresses out of db
	fs, es, err := w.GetAllAddresses()
	if err != nil {
		t.Error(err)
	} else if fs == nil {
		t.Errorf("No Factoid address was retrived")
	} else if es == nil {
		t.Errorf("No EC address was retrived")
	}
	// check that all the addresses are there
	if len(fs) != correctLen {
		t.Errorf("Wrong number of factoid addesses were retrived: %v", fs)
	}
	if len(es) != correctLen {
		t.Errorf("Wrong number of ec addesses were retrived: %v", es)
	}

	// print the addresses
	for _, f := range fs {
		t.Logf("%s %s", f, f.SecString())
	}
	for _, e := range es {
		t.Logf("%s %s", e, e.SecString())
	}

	// close and remove the testing db
	if err := w.Close(); err != nil {
		t.Error(err)
	}
	if err := os.RemoveAll(dbpath); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The sutil Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package paconn

import (
	"net"
	"fmt"
	"sync"
	"sync/atomic"
	"time"
	"errors"

	"github.com/sdming/gosnow"
	"github.com/golang/protobuf/proto"

	"github.com/shawnfeng/sutil"
	"github.com/shawnfeng/sutil/slog"
	"github.com/shawnfeng/sutil/stime"
	"github.com/shawnfeng/sutil/snetutil"

	"github.com/shawnfeng/sutil/paconn/pb"
)

var msgidgen *gosnow.SnowFlake

const (
	DEFAULT_SEND_TIMEOUT time.Duration = time.Millisecond*200
)

func init() {
	gosnow.Since = stime.Since2014 / 1000
	v, err := gosnow.NewSnowFlake(0)
	if err != nil {
		panic("snowflake init error, msgid can not get!")
	}

	msgidgen = v
}

type ackNotify struct {
	err error
	busstype int32
	result []byte
}

type FunAgOnewaynotify func(*Agent, int32, []byte)
type FunAgTwowaynotify func(*Agent, int32, []byte) (int32, []byte)
type FunAgClose func(*Agent, []byte, error)

type Agent struct {
	id string
	sendLock sync.Mutex
	conn net.Conn
	tuple4 string

	callmsgLock sync.Mutex
	callmsg  map[uint64] chan *ackNotify

	readTimeout time.Duration
	heartIntv time.Duration
	isConn int32

	cbOnewaynotify FunAgOnewaynotify
	cbTwowaynotify FunAgTwowaynotify
	cbClose FunAgClose

}

func (m *Agent) String() string {
	return fmt.Sprintf("%s@%s", m.id, m.tuple4)
}

func (m *Agent) Id() string {
	return m.id
}

func (m *Agent) Close() {
	// 不管错误，即使关闭了，也再关一次，修改isConn状态
	m.conn.Close()
	//err := m.conn.Close()
	//if err != nil {
	//	slog.Warnf("Agent.Close Close err:%s", err)
	//}

	atomic.StoreInt32(&m.isConn, 0)
}

func (m *Agent) Oneway(btype int32, data []byte, timeout time.Duration) error {
	pb := &connproto.ConnProto {
		Type: connproto.ConnProto_CALL.Enum(),
		Busstype: proto.Int32(btype),
		Bussdata: data,

	}

	spb, _ := proto.Marshal(pb)

	return m.send(spb, timeout)
}

func (m *Agent) Twoway(btype int32, data []byte, timeout time.Duration) (int32, []byte, error) {
	//fun := "Agent.Twoway"

	msgid, _ := msgidgen.Next()
	pb := &connproto.ConnProto {
		Type: connproto.ConnProto_CALL.Enum(),
		Msgid: proto.Uint64(msgid),
		Busstype: proto.Int32(btype),
		Bussdata: data,
	}

	spb, _ := proto.Marshal(pb)

	done := make(chan *ackNotify)

	func () {
		m.callmsgLock.Lock()
		defer m.callmsgLock.Unlock()
		m.callmsg[msgid] = done
	}()

	defer func () {
		m.callmsgLock.Lock()
		defer m.callmsgLock.Unlock()
		delete(m.callmsg, msgid)
	}()

	st := stime.NewTimeStat()
	err := m.send(spb, timeout)
	if err != nil {
		return 0, nil, err
	}

	senduse := st.Duration()

	if senduse >= timeout {
		return 0, nil, errors.New(fmt.Sprintf("call send timetout:%d", senduse))
	}

	select {
	case v := <-done:
		return v.busstype, v.result, v.err
	case <-time.After(timeout-senduse):
		m.Close()
		return 0, nil, errors.New("call ack timetout")
	}

}

func (m *Agent) recvACK(pb *connproto.ConnProto) {
	fun := "Agent.recvACK"
	msgid := pb.GetAckmsgid()
	btype := pb.GetBusstype()

	c, ok := func() (chan *ackNotify, bool) {
		m.callmsgLock.Lock()
		defer m.callmsgLock.Unlock()
		cc, o := m.callmsg[msgid]
		return cc, o
	}()

	if ok {
		an := &ackNotify {
			err: nil,
			busstype: btype,
			result: pb.GetBussdata(),
		}
		select {
		case c <-an:
		default:
			slog.Warnf("%s agent:%s msgid:%d no wait notify", fun, m, msgid)
		}
	} else {
		slog.Warnf("%s agent:%s msgid:%d not found", fun, m, msgid)
	}

}

func (m *Agent) recvCALL(pb *connproto.ConnProto) {
	fun := "Agent.recvCALL"
	data := pb.GetBussdata()
	btype := pb.GetBusstype()
	msgid := pb.GetMsgid()

	if msgid != 0 {
		res := make([]byte, 0)
		var rbtype int32 = 0
		if m.cbTwowaynotify != nil {
			rbtype, res = m.cbTwowaynotify(m, btype, data)
		}

		// 需要回执
		ack := &connproto.ConnProto {
			Type: connproto.ConnProto_ACK.Enum(),
			Ackmsgid: proto.Uint64(msgid),
			Busstype: proto.Int32(rbtype),
			Bussdata: res,
		}

		sdata, _ := proto.Marshal(ack)
		err := m.send(sdata, DEFAULT_SEND_TIMEOUT)
		if err != nil {
			slog.Warnf("%s agent:%s ack error:%s", fun, m, err)
		}

	} else {
		if m.cbOnewaynotify != nil {
			m.cbOnewaynotify(m, btype, data)
		}

	}

}

func (m *Agent) proto(data []byte) {
	fun := "Agent.proto"
	pb := &connproto.ConnProto{}
	err := proto.Unmarshal(data, pb)
	if err != nil {
		m.Close()
		slog.Warnf("%s agent:%s unmarshaling error: %s data:%v sd:%s", fun, m, err, data, data)
		return
	}

	slog.Infof("%s a:%s %s", fun, m, pb)

	pb_type := pb.GetType()
	if pb_type == connproto.ConnProto_ACK {
		m.recvACK(pb)

	} else if pb_type == connproto.ConnProto_CALL {
		m.recvCALL(pb)
	} else if pb_type == connproto.ConnProto_HEART {
		m.recvHEART()
	} else {
		m.Close()
		slog.Warnf("%s agent:%s type error: %s data:%v sd:%s", fun, m, err, data, data)
	}

}

func (m *Agent) send(data []byte, timeout time.Duration) error {
	if atomic.LoadInt32(&m.isConn) == 0 {
		return errors.New("connection is not ok")

	}

	s := snetutil.Packdata(data)

	m.sendLock.Lock()
	defer m.sendLock.Unlock()
	m.conn.SetWriteDeadline(time.Now().Add(timeout))
	a, err := m.conn.Write(s)
	//slog.Infof("%s agent:%s Send Write %d rv %d", fun, m, len(s), a)

	if err != nil {
		m.Close()
		return errors.New(fmt.Sprintf("send write error:%s", err))
	}

	if len(s) != a {
		m.Close()
		return errors.New("send write error:len")
	}

	return nil
}

func (m *Agent) sendHEART() {
	fun := "Agent.sendHEART"
	heart := &connproto.ConnProto {
		Type: connproto.ConnProto_HEART.Enum(),
	}

	//slog.Debugf("%s agent:%s msg:%s", fun, m, heart)

	data, _ := proto.Marshal(heart)
	err := m.send(data, DEFAULT_SEND_TIMEOUT)

	if err != nil {
		slog.Warnf("%s agent:%s error:%s", fun, m, err)
	}

}

func (m *Agent) recvHEART() {
	if m.heartIntv <= 0 {
		// 被动心跳
		m.sendHEART();
	}
}

func (m *Agent) heart() {
	fun := "Agent.heart"
	if m.heartIntv > 0 {
		// 主动心跳
		slog.Infof("%s agent:%s heart:%d", fun, m, m.heartIntv)
		ticker := time.NewTicker(m.heartIntv)
		for {
			select {
			case <-ticker.C:
				if atomic.LoadInt32(&m.isConn) != 0 {
					m.sendHEART();
				}
			}
		}

	} else {
		slog.Infof("%s agent:%s noheart:%d", fun, m, m.heartIntv)
	}

}

func (m *Agent) recv() {

	// 是否是read返回错误socket已经关闭，返回时候没有处理的数据，错误信息
	isclose, data, err := snetutil.PackageSplit(m.conn, m.readTimeout, m.proto)

	if !isclose {
		m.Close()
	}

	if m.cbClose != nil {
		// 所有关闭回调放到这里就好了，关闭时候，其他地方Close会走到这里
		m.cbClose(m, data, err)
	}

}

func NewAgent(
	c net.Conn,
	readto time.Duration,
	heart time.Duration,
	onenotify FunAgOnewaynotify,
	twonotify FunAgTwowaynotify,
	close FunAgClose,
) *Agent {
	fun := "NewAgent"

	if readto <= 0 {
		// 15分
		readto = (1000*60*5) * 3
	}

	aid := sutil.GetUniqueMd5()

	a := &Agent {
		id: aid,
		conn: c,
		tuple4: fmt.Sprintf("%s-%s", c.LocalAddr().String(), c.RemoteAddr().String()),
		callmsg: make(map[uint64] chan *ackNotify),
		readTimeout: readto,
		heartIntv: heart,
		isConn: 1,
		cbOnewaynotify: onenotify,
		cbTwowaynotify: twonotify,
		cbClose:close,
	}

	go a.recv()
	go a.heart()

	slog.Infof("%s a:%s", fun, a)

	return a
}

func NewAgentFromAddr(addr string,
	readtimeout time.Duration,
	heart time.Duration,
	onenotify FunAgOnewaynotify,
	twonotify FunAgTwowaynotify,
	close FunAgClose,
) (*Agent, error) {

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	a := NewAgent(
		conn,
		readtimeout,
		heart,
		onenotify,
		twonotify,
		close,
	)

	return a, nil

}

<commit_msg>build fail<commit_after>// Copyright 2014 The sutil Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package paconn

import (
	"net"
	"fmt"
	"sync"
	"sync/atomic"
	"time"
	"errors"

	"github.com/sdming/gosnow"
	"github.com/golang/protobuf/proto"

	"github.com/shawnfeng/sutil"
	"github.com/shawnfeng/sutil/slog"
	"github.com/shawnfeng/sutil/stime"
	"github.com/shawnfeng/sutil/snetutil"

	"github.com/shawnfeng/sutil/paconn/pb"
)

var msgidgen *gosnow.SnowFlake

const (
	DEFAULT_SEND_TIMEOUT time.Duration = time.Millisecond*200
)

func init() {
	gosnow.Since = stime.Since2014 / 1000
	v, err := gosnow.NewSnowFlake(0)
	if err != nil {
		panic("snowflake init error, msgid can not get!")
	}

	msgidgen = v
}

type ackNotify struct {
	err error
	busstype int32
	result []byte
}

type FunAgOnewaynotify func(*Agent, int32, []byte)
type FunAgTwowaynotify func(*Agent, int32, []byte) (int32, []byte)
type FunAgClose func(*Agent, []byte, error)

type Agent struct {
	id string
	sendLock sync.Mutex
	conn net.Conn
	tuple4 string

	callmsgLock sync.Mutex
	callmsg  map[uint64] chan *ackNotify

	readTimeout time.Duration
	heartIntv time.Duration
	isConn int32

	cbOnewaynotify FunAgOnewaynotify
	cbTwowaynotify FunAgTwowaynotify
	cbClose FunAgClose

}

func (m *Agent) String() string {
	return fmt.Sprintf("%s@%s", m.id, m.tuple4)
}

func (m *Agent) Id() string {
	return m.id
}

func (m *Agent) Close() {
	// 不管错误，即使关闭了，也再关一次，修改isConn状态
	m.conn.Close()
	//err := m.conn.Close()
	//if err != nil {
	//	slog.Warnf("Agent.Close Close err:%s", err)
	//}

	atomic.StoreInt32(&m.isConn, 0)
}

func (m *Agent) Oneway(btype int32, data []byte, timeout time.Duration) error {
	pb := &connproto.ConnProto {
		Type: connproto.ConnProto_CALL.Enum(),
		Busstype: proto.Int32(btype),
		Bussdata: data,

	}

	spb, _ := proto.Marshal(pb)

	return m.send(spb, timeout)
}

func (m *Agent) Twoway(btype int32, data []byte, timeout time.Duration) (int32, []byte, error) {
	//fun := "Agent.Twoway"

	msgid, _ := msgidgen.Next()
	pb := &connproto.ConnProto {
		Type: connproto.ConnProto_CALL.Enum(),
		Msgid: proto.Uint64(msgid),
		Busstype: proto.Int32(btype),
		Bussdata: data,
	}

	spb, _ := proto.Marshal(pb)

	done := make(chan *ackNotify)

	func () {
		m.callmsgLock.Lock()
		defer m.callmsgLock.Unlock()
		m.callmsg[msgid] = done
	}()

	defer func () {
		m.callmsgLock.Lock()
		defer m.callmsgLock.Unlock()
		delete(m.callmsg, msgid)
	}()

	st := stime.NewTimeStat()
	err := m.send(spb, timeout)
	if err != nil {
		return 0, nil, err
	}

	senduse := st.Duration()

	if senduse >= timeout {
		return 0, nil, errors.New(fmt.Sprintf("call send timetout:%d", senduse))
	}

	select {
	case v := <-done:
		return v.busstype, v.result, v.err
	case <-time.After(timeout-senduse):
		m.Close()
		return 0, nil, errors.New("call ack timetout")
	}

}

func (m *Agent) recvACK(pb *connproto.ConnProto) {
	fun := "Agent.recvACK"
	msgid := pb.GetAckmsgid()
	btype := pb.GetBusstype()

	c, ok := func() (chan *ackNotify, bool) {
		m.callmsgLock.Lock()
		defer m.callmsgLock.Unlock()
		cc, o := m.callmsg[msgid]
		return cc, o
	}()

	if ok {
		an := &ackNotify {
			err: nil,
			busstype: btype,
			result: pb.GetBussdata(),
		}
		select {
		case c <-an:
		default:
			slog.Warnf("%s agent:%s msgid:%d no wait notify", fun, m, msgid)
		}
	} else {
		slog.Warnf("%s agent:%s msgid:%d not found", fun, m, msgid)
	}

}

func (m *Agent) recvCALL(pb *connproto.ConnProto) {
	fun := "Agent.recvCALL"
	data := pb.GetBussdata()
	btype := pb.GetBusstype()
	msgid := pb.GetMsgid()

	if msgid != 0 {
		res := make([]byte, 0)
		var rbtype int32 = 0
		if m.cbTwowaynotify != nil {
			rbtype, res = m.cbTwowaynotify(m, btype, data)
		}

		// 需要回执
		ack := &connproto.ConnProto {
			Type: connproto.ConnProto_ACK.Enum(),
			Ackmsgid: proto.Uint64(msgid),
			Busstype: proto.Int32(rbtype),
			Bussdata: res,
		}

		sdata, _ := proto.Marshal(ack)
		err := m.send(sdata, DEFAULT_SEND_TIMEOUT)
		if err != nil {
			slog.Warnf("%s agent:%s ack error:%s", fun, m, err)
		}

	} else {
		if m.cbOnewaynotify != nil {
			m.cbOnewaynotify(m, btype, data)
		}

	}

}

func (m *Agent) proto(data []byte) {
	fun := "Agent.proto"
	pb := &connproto.ConnProto{}
	err := proto.Unmarshal(data, pb)
	if err != nil {
		m.Close()
		slog.Warnf("%s agent:%s unmarshaling error: %s data:%v sd:%s", fun, m, err, data, data)
		return
	}

	slog.Infof("%s a:%s %s", fun, m, pb)

	pb_type := pb.GetType()
	if pb_type == connproto.ConnProto_ACK {
		m.recvACK(pb)

	} else if pb_type == connproto.ConnProto_CALL {
		m.recvCALL(pb)
	} else if pb_type == connproto.ConnProto_HEART {
		m.recvHEART()
	} else {
		m.Close()
		slog.Warnf("%s agent:%s type error: %s data:%v sd:%s", fun, m, err, data, data)
	}

}

func (m *Agent) send(data []byte, timeout time.Duration) error {
	if atomic.LoadInt32(&m.isConn) == 0 {
		return errors.New("connection is not ok")

	}

	s := snetutil.Packdata(data)

	m.sendLock.Lock()
	defer m.sendLock.Unlock()
	m.conn.SetWriteDeadline(time.Now().Add(timeout))
	a, err := m.conn.Write(s)
	//slog.Infof("%s agent:%s Send Write %d rv %d", fun, m, len(s), a)

	if err != nil {
		m.Close()
		return errors.New(fmt.Sprintf("send write error:%s", err))
	}

	if len(s) != a {
		m.Close()
		return errors.New("send write error:len")
	}

	return nil
}

func (m *Agent) sendHEART() {
	fun := "Agent.sendHEART"
	heart := &connproto.ConnProto {
		Type: connproto.ConnProto_HEART.Enum(),
	}

	//slog.Debugf("%s agent:%s msg:%s", fun, m, heart)

	data, _ := proto.Marshal(heart)
	err := m.send(data, DEFAULT_SEND_TIMEOUT)

	if err != nil {
		slog.Warnf("%s agent:%s error:%s", fun, m, err)
	}

}

func (m *Agent) recvHEART() {
	if m.heartIntv <= 0 {
		// 被动心跳
		m.sendHEART();
	}
}

func (m *Agent) heart() {
	fun := "Agent.heart"
	if m.heartIntv > 0 {
		// 主动心跳
		slog.Infof("%s agent:%s heart:%d", fun, m, m.heartIntv)
		ticker := time.NewTicker(m.heartIntv)
		for {
			select {
			case <-ticker.C:
				if atomic.LoadInt32(&m.isConn) != 0 {
					m.sendHEART();
				}
			}
		}

	} else {
		slog.Infof("%s agent:%s noheart:%d", fun, m, m.heartIntv)
	}

}

func (m *Agent) recv() {

	// 是否是read返回错误socket已经关闭，返回时候没有处理的数据，错误信息
	isclose, data, err := snetutil.PackageSplit(m.conn, m.readTimeout, m.proto)

	if !isclose {
		m.Close()
	}

	if m.cbClose != nil {
		// 所有关闭回调放到这里就好了，关闭时候，其他地方Close会走到这里
		m.cbClose(m, data, err)
	}

}

func NewAgent(
	c net.Conn,
	readto time.Duration,
	heart time.Duration,
	onenotify FunAgOnewaynotify,
	twonotify FunAgTwowaynotify,
	close FunAgClose,
) *Agent {
	fun := "NewAgent"

	if readto <= 0 {
		// 15分
		readto = (1000*60*5) * 3
	}

	aid, err := sutil.GetUniqueMd5()
	if err != nil {
		slog.Errorf("%s new err:%s", fun, err)
		return nil
	}

	a := &Agent {
		id: aid,
		conn: c,
		tuple4: fmt.Sprintf("%s-%s", c.LocalAddr().String(), c.RemoteAddr().String()),
		callmsg: make(map[uint64] chan *ackNotify),
		readTimeout: readto,
		heartIntv: heart,
		isConn: 1,
		cbOnewaynotify: onenotify,
		cbTwowaynotify: twonotify,
		cbClose:close,
	}

	go a.recv()
	go a.heart()

	slog.Infof("%s a:%s", fun, a)

	return a
}

func NewAgentFromAddr(addr string,
	readtimeout time.Duration,
	heart time.Duration,
	onenotify FunAgOnewaynotify,
	twonotify FunAgTwowaynotify,
	close FunAgClose,
) (*Agent, error) {

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	a := NewAgent(
		conn,
		readtimeout,
		heart,
		onenotify,
		twonotify,
		close,
	)

	return a, nil

}

<|endoftext|>


<commit_before>package helpers

import (
	"database/sql"
	"strings"

	"github.com/pborman/uuid"
	"github.com/stripe/stripe-go"

	g "github.com/ghmeier/bloodlines/gateways"
	"github.com/ghmeier/coinage/gateways"
	"github.com/ghmeier/coinage/models"
	w "github.com/lcollin/warehouse/gateways"
)

/*Plan manages retrieval and manipulating roaster plan information*/
type Plan struct {
	*base
	Stripe    gateways.Stripe
	Warehouse w.Warehouse
}

/*NewPlan initializes and returns a plan with the given gateways*/
func NewPlan(sql g.SQL, stripe gateways.Stripe, warehouse w.Warehouse) *Plan {
	return &Plan{
		base:      &base{sql: sql},
		Stripe:    stripe,
		Warehouse: warehouse,
	}
}

/*Insert creates a new roaster plan in stripe and adds a record to the db*/
func (p *Plan) Insert(roaster *models.Roaster, req *models.PlanRequest) (*models.Plan, error) {
	item, err := p.Warehouse.GetItemByID(req.ItemID)
	if err != nil {
		return nil, err
	}

	planIDs := make([]string, 0)
	var consumerPrice uint64
	for i := 1; i < len(models.Frequencies); i++ {
		stripe, err := p.Stripe.NewPlan(roaster.Secret, item, models.Frequencies[i])
		if err != nil {
			return nil, err
		}

		planIDs = append(planIDs, stripe.ID)
		consumerPrice = stripe.Amount
	}

	plan := models.NewPlan(roaster.ID, req.ItemID, planIDs)
	err = p.sql.Modify("INSERT INTO plan (roasterId,itemId,planIds)VALUES(?,?,?)",
		plan.RoasterID.String(),
		plan.ItemID.String(),
		strings.Join(plan.PlanIDs, ","),
	)
	if err != nil {
		return nil, err
	}

	item.ConsumerPrice = float64(consumerPrice) / 100.0 * (1.00 + p.Stripe.ApplicationFee())
	p.Warehouse.UpdateItem(item)
	return plan, nil
}

/*GetByRoaster returns all plans associated with a roaster*/
func (p *Plan) GetByRoaster(roaster *models.Roaster, offset int, limit int) ([]*models.Plan, error) {
	rows, err := p.sql.Select("SELECT roasterId,itemId,planIds FROM plan WHERE roasterId=? ORDER BY itemId ASC LIMIT ?,?",
		roaster.ID,
		offset,
		limit,
	)
	if err != nil {
		return nil, err
	}

	return p.plan(roaster.Secret, rows)
}

/*Get returns a plan associated with a roaster and itemid*/
func (p *Plan) Get(roaster *models.Roaster, itemID uuid.UUID) (*models.Plan, error) {
	rows, err := p.sql.Select("SELECT roasterId,itemId,planIds FROM plan WHERE roasterId=? AND itemId=?",
		roaster.ID,
		itemID,
	)
	if err != nil {
		return nil, err
	}

	plans, err := p.plan(roaster.Secret, rows)
	if err != nil {
		return nil, err
	}

	return plans[0], err
}

func (p *Plan) plan(secret string, rows *sql.Rows) ([]*models.Plan, error) {
	plans, _ := models.PlanFromSQL(rows)

	return plans, nil
}

func (p *Plan) plans(secret string, ids []string) ([]*stripe.Plan, error) {
	plans := make([]*stripe.Plan, 0)
	var err error
	var plan *stripe.Plan
	for i := 0; i < len(ids); i++ {
		plan, err = p.Stripe.GetPlan(secret, ids[i])
		plans = append(plans, plan)
	}
	return plans, err
}

/*Update is not implemented*/
func (p *Plan) Update(id string, itemID uuid.UUID) (*models.Plan, error) {
	// err := p.sql.Modify("UPDATE plan SET roasterId=?,itemId=?,planIds=? WHERE itemId=?",
	// 	plan.RoasterID,
	// 	plan.ItemID,
	// 	strings.Join(plan.PlanIDs, ","),
	// )
	return nil, nil
}

/*Delete is not implemented*/
func (p *Plan) Delete(id, planID string) error {
	return nil
}
<commit_msg>Remove unused varable<commit_after>package helpers

import (
	"database/sql"
	"strings"

	"github.com/pborman/uuid"
	"github.com/stripe/stripe-go"

	g "github.com/ghmeier/bloodlines/gateways"
	"github.com/ghmeier/coinage/gateways"
	"github.com/ghmeier/coinage/models"
	w "github.com/lcollin/warehouse/gateways"
)

/*Plan manages retrieval and manipulating roaster plan information*/
type Plan struct {
	*base
	Stripe    gateways.Stripe
	Warehouse w.Warehouse
}

/*NewPlan initializes and returns a plan with the given gateways*/
func NewPlan(sql g.SQL, stripe gateways.Stripe, warehouse w.Warehouse) *Plan {
	return &Plan{
		base:      &base{sql: sql},
		Stripe:    stripe,
		Warehouse: warehouse,
	}
}

/*Insert creates a new roaster plan in stripe and adds a record to the db*/
func (p *Plan) Insert(roaster *models.Roaster, req *models.PlanRequest) (*models.Plan, error) {
	item, err := p.Warehouse.GetItemByID(req.ItemID)
	if err != nil {
		return nil, err
	}

	planIDs := make([]string, 0)
	for i := 1; i < len(models.Frequencies); i++ {
		stripe, err := p.Stripe.NewPlan(roaster.Secret, item, models.Frequencies[i])
		if err != nil {
			return nil, err
		}

		planIDs = append(planIDs, stripe.ID)
	}

	plan := models.NewPlan(roaster.ID, req.ItemID, planIDs)
	err = p.sql.Modify("INSERT INTO plan (roasterId,itemId,planIds)VALUES(?,?,?)",
		plan.RoasterID.String(),
		plan.ItemID.String(),
		strings.Join(plan.PlanIDs, ","),
	)
	if err != nil {
		return nil, err
	}

	item.ConsumerPrice = float64(item.ConsumerPrice) / 100.0 * (1.00 + p.Stripe.ApplicationFee())
	p.Warehouse.UpdateItem(item)
	return plan, nil
}

/*GetByRoaster returns all plans associated with a roaster*/
func (p *Plan) GetByRoaster(roaster *models.Roaster, offset int, limit int) ([]*models.Plan, error) {
	rows, err := p.sql.Select("SELECT roasterId,itemId,planIds FROM plan WHERE roasterId=? ORDER BY itemId ASC LIMIT ?,?",
		roaster.ID,
		offset,
		limit,
	)
	if err != nil {
		return nil, err
	}

	return p.plan(roaster.Secret, rows)
}

/*Get returns a plan associated with a roaster and itemid*/
func (p *Plan) Get(roaster *models.Roaster, itemID uuid.UUID) (*models.Plan, error) {
	rows, err := p.sql.Select("SELECT roasterId,itemId,planIds FROM plan WHERE roasterId=? AND itemId=?",
		roaster.ID,
		itemID,
	)
	if err != nil {
		return nil, err
	}

	plans, err := p.plan(roaster.Secret, rows)
	if err != nil {
		return nil, err
	}

	return plans[0], err
}

func (p *Plan) plan(secret string, rows *sql.Rows) ([]*models.Plan, error) {
	plans, _ := models.PlanFromSQL(rows)

	return plans, nil
}

func (p *Plan) plans(secret string, ids []string) ([]*stripe.Plan, error) {
	plans := make([]*stripe.Plan, 0)
	var err error
	var plan *stripe.Plan
	for i := 0; i < len(ids); i++ {
		plan, err = p.Stripe.GetPlan(secret, ids[i])
		plans = append(plans, plan)
	}
	return plans, err
}

/*Update is not implemented*/
func (p *Plan) Update(id string, itemID uuid.UUID) (*models.Plan, error) {
	// err := p.sql.Modify("UPDATE plan SET roasterId=?,itemId=?,planIds=? WHERE itemId=?",
	// 	plan.RoasterID,
	// 	plan.ItemID,
	// 	strings.Join(plan.PlanIDs, ","),
	// )
	return nil, nil
}

/*Delete is not implemented*/
func (p *Plan) Delete(id, planID string) error {
	return nil
}
<|endoftext|>


<commit_before>package service_test

import (
	"fmt"
	"strings"
	"time"

	"github.com/pborman/uuid"

	"github.com/pivotal-cf-experimental/cf-test-helpers/cf"
	"github.com/pivotal-cf-experimental/cf-test-helpers/runner"
	"github.com/pivotal-cf-experimental/cf-test-helpers/services"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("Redis Service", func() {
	var shortTimeout = time.Second * 180
	var longTimeout = time.Minute * 15
	var retryInterval = time.Second * 1
	var appPath = "../assets/cf-redis-example-app"

	var appName string
	var context services.Context

	randomName := func() string {
		return uuid.NewRandom().String()
	}

	appUri := func(appName string) string {
		return "https://" + appName + "." + testConfig.AppsDomain
	}

	assertAppIsRunning := func(appName string) {
		pingUri := appUri(appName) + "/ping"
		fmt.Println("Checking that the app is responding at url: ", pingUri)
		Eventually(runner.Curl(pingUri, "-k"), shortTimeout, retryInterval).Should(
			Say("key not present"),
			`{"FailReason": "Test app deployed but did not respond in time"}`,
		)
		fmt.Println()
	}

	createTestOrgAndSpace := func() {
		apiCmd := []string{"api", testConfig.ApiEndpoint}

		if testConfig.SkipSSLValidation {
			apiCmd = append(apiCmd, "--skip-ssl-validation")
		}

		Eventually(cf.Cf(apiCmd...), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to target Cloud Foundry"}`,
		)

		Eventually(cf.Cf("auth", testConfig.AdminUser, testConfig.AdminPassword), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf auth` with target Cloud Foundry\"}",
		)

		Eventually(cf.Cf("create-org", testConfig.OrgName), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to create CF test org"}`,
		)

		Eventually(cf.Cf("target", "-o", testConfig.OrgName)).Should(
			Exit(0),
			`{"FailReason": "Failed to target test org"}`,
		)

		Eventually(cf.Cf("create-space", testConfig.SpaceName), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to create CF test space"}`,
		)
	}

	createAndBindSecurityGroup := func() {
		Eventually(cf.Cf("auth", testConfig.AdminUser, testConfig.AdminPassword), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf auth` with target Cloud Foundry\"}",
		)

		Eventually(
			cf.Cf("create-security-group", "redis-smoke-tests-sg", securityGroupConfigPath),
			shortTimeout,
		).Should(
			Exit(0),
			`{"FailReason": "Failed to create security group"}`,
		)

		Eventually(
			cf.Cf("bind-security-group", "redis-smoke-tests-sg", testConfig.OrgName, testConfig.SpaceName),
			shortTimeout,
		).Should(
			Exit(0),
			`{"FailReason": "Failed to bind security group to space"}`,
		)
	}

	deleteSecurityGroup := func() {
		Eventually(cf.Cf("auth", testConfig.AdminUser, testConfig.AdminPassword), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf auth` with target Cloud Foundry\"}",
		)

		Eventually(cf.Cf("delete-security-group", "redis-smoke-tests-sg", "-f"), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to remove security group"}`,
		)
	}

	BeforeSuite(func() {
		createTestOrgAndSpace()
		createAndBindSecurityGroup()

		context = services.NewContext(testConfig, "redis-test")
		context.Setup()
	})

	BeforeEach(func() {
		appName = randomName()
		Eventually(
			cf.Cf("push", appName, "-m", "256M", "-p", appPath, "-s", "cflinuxfs2", "-no-start"),
			shortTimeout,
		).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf push` test app\"}",
		)
	})

	AfterEach(func() {
		Eventually(cf.Cf("delete", appName, "-f"), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf delete` test app\"}",
		)
	})

	AfterSuite(func() {
		context.Teardown()
		deleteSecurityGroup()
	})

	AssertLifeCycleBehavior := func(planName string) {
		It(strings.ToUpper(planName)+": create, bind to, write to, read from, unbind, and destroy a service instance", func() {
			serviceInstanceName := randomName()

			createServiceSession := cf.Cf("create-service", redisConfig.ServiceName, planName, serviceInstanceName)
			createServiceSession.Wait(shortTimeout)

			createServiceStdout := createServiceSession.Out

			select {
			case <-createServiceStdout.Detect("FAILED"):
				Eventually(createServiceSession, shortTimeout).Should(
					Say("instance limit for this service has been reached"),
					`{"FailReason": "Failed to bind Redis service instance to test app"}`,
				)
				Eventually(createServiceSession, shortTimeout).Should(Exit(1))
				fmt.Printf("No Plan Instances available for testing %s plan\n", planName)

			case <-createServiceStdout.Detect("OK"):
				Eventually(createServiceSession, shortTimeout).Should(
					Exit(0),
					`{"FailReason": "Failed to create Redis service instance"}`,
				)

				Eventually(cf.Cf("bind-service", appName, serviceInstanceName), shortTimeout).Should(
					Exit(0),
					`{"FailReason": "Failed to bind Redis service instance to test app"}`,
				)

				Eventually(cf.Cf("start", appName), longTimeout).Should(
					Exit(0),
					`{"FailReason": "Failed to start test app"}`,
				)

				assertAppIsRunning(appName)

				uri := appUri(appName) + "/mykey"
				fmt.Println("Posting to url: ", uri)
				Eventually(
					runner.Curl("-d", "data=myvalue", "-X", "PUT", uri, "-k"),
					shortTimeout,
					retryInterval,
				).Should(
					Say("success"),
					fmt.Sprintf(`{"FailReason": "Failed to write to test %s instance"}`, planName),
				)
				fmt.Println()
				fmt.Println("Getting from url: ", uri)

				Eventually(
					runner.Curl(uri, "-k"),
					shortTimeout,
					retryInterval,
				).Should(
					Say("myvalue"),
					fmt.Sprintf(`{"FailReason": "Failed to read from test %s instance"}`, planName),
				)
				fmt.Println()

				Eventually(cf.Cf("unbind-service", appName, serviceInstanceName), shortTimeout).Should(
					Exit(0),
					fmt.Sprintf(`{"FailReason": "Failed to unbind %s instance from test app"}`, planName),
				)

				Eventually(cf.Cf("delete-service", "-f", serviceInstanceName), shortTimeout).Should(
					Exit(0),
					fmt.Sprintf(`{"FailReason": "Failed to delete test %s instance"}`, planName),
				)
			}
			createServiceStdout.CancelDetects()

		})
	}

	Context("for each plan", func() {
		for _, planName := range redisConfig.PlanNames {
			AssertLifeCycleBehavior(planName)
		}
	})
})
<commit_msg>Consistently use shortTimeout<commit_after>package service_test

import (
	"fmt"
	"strings"
	"time"

	"github.com/pborman/uuid"

	"github.com/pivotal-cf-experimental/cf-test-helpers/cf"
	"github.com/pivotal-cf-experimental/cf-test-helpers/runner"
	"github.com/pivotal-cf-experimental/cf-test-helpers/services"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("Redis Service", func() {
	var shortTimeout = time.Minute * 3
	var longTimeout = time.Minute * 15
	var retryInterval = time.Second * 1
	var appPath = "../assets/cf-redis-example-app"

	var appName string
	var context services.Context

	randomName := func() string {
		return uuid.NewRandom().String()
	}

	appUri := func(appName string) string {
		return "https://" + appName + "." + testConfig.AppsDomain
	}

	assertAppIsRunning := func(appName string) {
		pingUri := appUri(appName) + "/ping"
		fmt.Println("Checking that the app is responding at url: ", pingUri)
		Eventually(runner.Curl(pingUri, "-k"), shortTimeout, retryInterval).Should(
			Say("key not present"),
			`{"FailReason": "Test app deployed but did not respond in time"}`,
		)
		fmt.Println()
	}

	createTestOrgAndSpace := func() {
		apiCmd := []string{"api", testConfig.ApiEndpoint}

		if testConfig.SkipSSLValidation {
			apiCmd = append(apiCmd, "--skip-ssl-validation")
		}

		Eventually(cf.Cf(apiCmd...), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to target Cloud Foundry"}`,
		)

		Eventually(cf.Cf("auth", testConfig.AdminUser, testConfig.AdminPassword), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf auth` with target Cloud Foundry\"}",
		)

		Eventually(cf.Cf("create-org", testConfig.OrgName), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to create CF test org"}`,
		)

		Eventually(cf.Cf("target", "-o", testConfig.OrgName), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to target test org"}`,
		)

		Eventually(cf.Cf("create-space", testConfig.SpaceName), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to create CF test space"}`,
		)
	}

	createAndBindSecurityGroup := func() {
		Eventually(cf.Cf("auth", testConfig.AdminUser, testConfig.AdminPassword), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf auth` with target Cloud Foundry\"}",
		)

		Eventually(
			cf.Cf("create-security-group", "redis-smoke-tests-sg", securityGroupConfigPath),
			shortTimeout,
		).Should(
			Exit(0),
			`{"FailReason": "Failed to create security group"}`,
		)

		Eventually(
			cf.Cf("bind-security-group", "redis-smoke-tests-sg", testConfig.OrgName, testConfig.SpaceName),
			shortTimeout,
		).Should(
			Exit(0),
			`{"FailReason": "Failed to bind security group to space"}`,
		)
	}

	deleteSecurityGroup := func() {
		Eventually(cf.Cf("auth", testConfig.AdminUser, testConfig.AdminPassword), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf auth` with target Cloud Foundry\"}",
		)

		Eventually(cf.Cf("delete-security-group", "redis-smoke-tests-sg", "-f"), shortTimeout).Should(
			Exit(0),
			`{"FailReason": "Failed to remove security group"}`,
		)
	}

	BeforeSuite(func() {
		createTestOrgAndSpace()
		createAndBindSecurityGroup()

		context = services.NewContext(testConfig, "redis-test")
		context.Setup()
	})

	BeforeEach(func() {
		appName = randomName()
		Eventually(
			cf.Cf("push", appName, "-m", "256M", "-p", appPath, "-s", "cflinuxfs2", "-no-start"),
			shortTimeout,
		).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf push` test app\"}",
		)
	})

	AfterEach(func() {
		Eventually(cf.Cf("delete", appName, "-f"), shortTimeout).Should(
			Exit(0),
			"{\"FailReason\": \"Failed to `cf delete` test app\"}",
		)
	})

	AfterSuite(func() {
		context.Teardown()
		deleteSecurityGroup()
	})

	AssertLifeCycleBehavior := func(planName string) {
		It(strings.ToUpper(planName)+": create, bind to, write to, read from, unbind, and destroy a service instance", func() {
			serviceInstanceName := randomName()

			createServiceSession := cf.Cf("create-service", redisConfig.ServiceName, planName, serviceInstanceName)
			createServiceSession.Wait(shortTimeout)

			createServiceStdout := createServiceSession.Out

			select {
			case <-createServiceStdout.Detect("FAILED"):
				Eventually(createServiceSession, shortTimeout).Should(
					Say("instance limit for this service has been reached"),
					`{"FailReason": "Failed to bind Redis service instance to test app"}`,
				)
				Eventually(createServiceSession, shortTimeout).Should(Exit(1))
				fmt.Printf("No Plan Instances available for testing %s plan\n", planName)

			case <-createServiceStdout.Detect("OK"):
				Eventually(createServiceSession, shortTimeout).Should(
					Exit(0),
					`{"FailReason": "Failed to create Redis service instance"}`,
				)

				Eventually(cf.Cf("bind-service", appName, serviceInstanceName), shortTimeout).Should(
					Exit(0),
					`{"FailReason": "Failed to bind Redis service instance to test app"}`,
				)

				Eventually(cf.Cf("start", appName), longTimeout).Should(
					Exit(0),
					`{"FailReason": "Failed to start test app"}`,
				)

				assertAppIsRunning(appName)

				uri := appUri(appName) + "/mykey"
				fmt.Println("Posting to url: ", uri)
				Eventually(
					runner.Curl("-d", "data=myvalue", "-X", "PUT", uri, "-k"),
					shortTimeout,
					retryInterval,
				).Should(
					Say("success"),
					fmt.Sprintf(`{"FailReason": "Failed to write to test %s instance"}`, planName),
				)
				fmt.Println()
				fmt.Println("Getting from url: ", uri)

				Eventually(
					runner.Curl(uri, "-k"),
					shortTimeout,
					retryInterval,
				).Should(
					Say("myvalue"),
					fmt.Sprintf(`{"FailReason": "Failed to read from test %s instance"}`, planName),
				)
				fmt.Println()

				Eventually(cf.Cf("unbind-service", appName, serviceInstanceName), shortTimeout).Should(
					Exit(0),
					fmt.Sprintf(`{"FailReason": "Failed to unbind %s instance from test app"}`, planName),
				)

				Eventually(cf.Cf("delete-service", "-f", serviceInstanceName), shortTimeout).Should(
					Exit(0),
					fmt.Sprintf(`{"FailReason": "Failed to delete test %s instance"}`, planName),
				)
			}
			createServiceStdout.CancelDetects()

		})
	}

	Context("for each plan", func() {
		for _, planName := range redisConfig.PlanNames {
			AssertLifeCycleBehavior(planName)
		}
	})
})
<|endoftext|>


<commit_before>package irc_test

import (
	. "github.com/jriddick/geoffrey/irc"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Irc", func() {
	var (
		client *IRC
		cfg    Config
		emp    Config
	)

	BeforeEach(func() {
		cfg = Config{
			Hostname:           "localhost",
			Port:               6697,
			Secure:             false,
			InsecureSkipVerify: false,
		}
		emp = Config{}
		client = NewIRC(cfg)
	})

	AfterEach(func() {
		client.Disconnect()
	})

	It("should be able to create an irc client", func() {
		Expect(NewIRC(cfg)).NotTo(BeNil())
	})

	It("should be able to connect", func() {
		Expect(client.Connect()).NotTo(HaveOccurred())
	})

	It("should not be able to connect with an empty configuration", func() {
		Expect(NewIRC(emp).Connect()).To(HaveOccurred())
	})

	It("should be able to read from the server", func(done Done) {
		By("connecting to the server")
		Expect(client.Connect()).NotTo(HaveOccurred())

		By("getting the reader")
		reader := client.Reader()
		Expect(reader).NotTo(BeNil())

		Expect(<-reader).NotTo(BeNil())
		close(done)
	}, 2)

	It("should be able to register to the server", func(done Done) {
		By("connecting to the server")
		Expect(client.Connect()).NotTo(HaveOccurred())

		By("getting the reader and writer")
		reader := client.Reader()
		Expect(reader).NotTo(BeNil())
		writer := client.Writer()
		Expect(writer).NotTo(BeNil())

		By("sending nick and user")
		writer <- "NICK geoffrey"
		writer <- "USER geoffrey 0 * :geoffrey"

		for {
			line, ok := <-reader
			Expect(ok).To(BeTrue())

			if line.Command == RPL_WELCOME {
				break
			}
		}

		close(done)
	}, 2)
})
<commit_msg>added test case for reconnecting the client<commit_after>package irc_test

import (
	. "github.com/jriddick/geoffrey/irc"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Irc", func() {
	var (
		client *IRC
		cfg    Config
		emp    Config
	)

	BeforeEach(func() {
		cfg = Config{
			Hostname:           "localhost",
			Port:               6697,
			Secure:             false,
			InsecureSkipVerify: false,
		}
		emp = Config{}
		client = NewIRC(cfg)
	})

	AfterEach(func() {
		client.Disconnect()
	})

	It("should be able to create an irc client", func() {
		Expect(NewIRC(cfg)).NotTo(BeNil())
	})

	It("should be able to connect", func() {
		Expect(client.Connect()).NotTo(HaveOccurred())
	})

	It("should not be able to connect with an empty configuration", func() {
		Expect(NewIRC(emp).Connect()).To(HaveOccurred())
	})

	It("should be able to read from the server", func(done Done) {
		By("connecting to the server")
		Expect(client.Connect()).NotTo(HaveOccurred())

		By("getting the reader")
		reader := client.Reader()
		Expect(reader).NotTo(BeNil())

		Expect(<-reader).NotTo(BeNil())
		close(done)
	}, 2)

	It("should be able to register to the server", func(done Done) {
		By("connecting to the server")
		Expect(client.Connect()).NotTo(HaveOccurred())

		By("getting the reader and writer")
		reader := client.Reader()
		Expect(reader).NotTo(BeNil())
		writer := client.Writer()
		Expect(writer).NotTo(BeNil())

		By("sending nick and user")
		writer <- "NICK geoffrey"
		writer <- "USER geoffrey 0 * :geoffrey"

		for {
			line, ok := <-reader
			Expect(ok).To(BeTrue())

			if line.Command == RPL_WELCOME {
				break
			}
		}

		close(done)
	}, 2)

	It("should be able to reconnect to the server", func(done Done) {
		By("connecting to the server")
		Expect(client.Connect()).NotTo(HaveOccurred())

		By("reading from the server")
		Expect(<-client.Reader()).NotTo(BeNil())

		By("reconnecting")
		Expect(client.Reconnect()).NotTo(HaveOccurred())

		By("reading from the server")
		Expect(<-client.Reader()).NotTo(BeNil())

		close(done)
	}, 2)
})
<|endoftext|>


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

package session_test

import (
	"reflect"
	"testing"

	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/testkit"
	"github.com/pingcap/tidb/testkit/external"
	"github.com/stretchr/testify/require"
)

func TestInitMetaTable(t *testing.T) {
	store := testkit.CreateMockStore(t)

	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")
	for _, sql := range session.DDLJobTables {
		tk.MustExec(sql.SQL)
	}

	tbls := map[string]struct{}{
		"tidb_ddl_job":     {},
		"tidb_ddl_reorg":   {},
		"tidb_ddl_history": {},
	}

	for tbl := range tbls {
		metaInMySQL := external.GetTableByName(t, tk, "mysql", tbl).Meta()
		metaInTest := external.GetTableByName(t, tk, "test", tbl).Meta()

		require.Greater(t, metaInMySQL.ID, int64(0))
		require.Greater(t, metaInMySQL.UpdateTS, uint64(0))

		metaInTest.ID = metaInMySQL.ID
		metaInMySQL.UpdateTS = metaInTest.UpdateTS
		require.True(t, reflect.DeepEqual(metaInMySQL, metaInTest))
	}
}
<commit_msg>session: fix data race in the TestInitMetaTable (#37464)<commit_after>// Copyright 2022 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package session_test

import (
	"reflect"
	"testing"

	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/testkit"
	"github.com/pingcap/tidb/testkit/external"
	"github.com/stretchr/testify/require"
)

func TestInitMetaTable(t *testing.T) {
	store := testkit.CreateMockStore(t)

	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")
	for _, sql := range session.DDLJobTables {
		tk.MustExec(sql.SQL)
	}

	tbls := map[string]struct{}{
		"tidb_ddl_job":     {},
		"tidb_ddl_reorg":   {},
		"tidb_ddl_history": {},
	}

	for tbl := range tbls {
		metaInMySQL := external.GetTableByName(t, tk, "mysql", tbl).Meta().Clone()
		metaInTest := external.GetTableByName(t, tk, "test", tbl).Meta().Clone()

		require.Greater(t, metaInMySQL.ID, int64(0))
		require.Greater(t, metaInMySQL.UpdateTS, uint64(0))

		metaInTest.ID = metaInMySQL.ID
		metaInMySQL.UpdateTS = metaInTest.UpdateTS
		require.True(t, reflect.DeepEqual(metaInMySQL, metaInTest))
	}
}
<|endoftext|>


<commit_before>package par1

import (
	"errors"
	"fmt"
	"os"
	"path"
	"path/filepath"

	"github.com/akalin/gopar/fs"
	"github.com/akalin/gopar/hashutil"
	"github.com/klauspost/reedsolomon"
)

// A Decoder keeps track of all information needed to check the
// integrity of a set of data files, and possibly repair any
// missing/corrupt data files from the parity files (.P00, .P01,
// etc.).
type Decoder struct {
	fs       fs.FS
	delegate DecoderDelegate

	indexFile   string
	indexVolume volume

	fileData [][]byte

	shardByteCount int
	parityData     [][]byte
}

// DecoderDelegate holds methods that are called during the decode
// process.
type DecoderDelegate interface {
	OnHeaderLoad(headerInfo string)
	OnFileEntryLoad(i, n int, filename, entryInfo string)
	OnCommentLoad(comment []byte)
	OnDataFileLoad(i, n int, path string, byteCount int, corrupt bool, err error)
	OnDataFileWrite(i, n int, path string, byteCount int, err error)
	OnVolumeFileLoad(i uint64, path string, setHash [16]byte, dataByteCount int, err error)
}

// DoNothingDecoderDelegate is an implementation of DecoderDelegate
// that does nothing for all methods.
type DoNothingDecoderDelegate struct{}

// OnHeaderLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnHeaderLoad(headerInfo string) {}

// OnFileEntryLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnFileEntryLoad(i, n int, filename, entryInfo string) {}

// OnCommentLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnCommentLoad(comment []byte) {}

// OnDataFileLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, corrupt bool, err error) {
}

// OnDataFileWrite implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnDataFileWrite(i, n int, path string, byteCount int, err error) {}

// OnVolumeFileLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnVolumeFileLoad(i uint64, path string, setHash [16]byte, dataByteCount int, err error) {
}

func newDecoder(fs fs.FS, delegate DecoderDelegate, indexFile string) *Decoder {
	return &Decoder{fs, delegate, indexFile, volume{}, nil, 0, nil}
}

// NewDecoder reads the given index file, which usually has a .PAR
// extension.
func NewDecoder(delegate DecoderDelegate, indexFile string) *Decoder {
	return newDecoder(fs.MakeDefaultFS(), delegate, indexFile)
}

func (d *Decoder) getFilePath(entry fileEntry) (string, error) {
	filename := entry.filename
	if filepath.Base(filename) != filename {
		return "", errors.New("bad filename")
	}

	basePath := filepath.Dir(d.indexFile)
	return filepath.Join(basePath, filename), nil
}

func readAndCheckVolume(filesystem fs.FS, path string, checkVolumeFn func(volume) error) (v volume, data []byte, err error) {
	readStream, err := filesystem.GetReadStream(path)
	if err != nil {
		return volume{}, nil, err
	}

	defer func() {
		if readStream != nil {
			fs.CloseCloser(readStream, &err)
		}
	}()

	v, err = readVolume(readStream)
	if err != nil {
		return volume{}, nil, err
	}

	err = checkVolumeFn(v)
	if err != nil {
		return v, nil, err
	}

	data, err = fs.ReadAndClose(readStream)
	readStream = nil
	if err != nil {
		return v, nil, err
	}
	return v, data, nil
}

func checkIndexVolume(indexVolume volume) error {
	if indexVolume.header.VolumeNumber != 0 {
		// TODO: Relax this check.
		return errors.New("expected volume number 0 for index volume")
	}
	return nil
}

// LoadIndexFile reads the data from the index file.
func (d *Decoder) LoadIndexFile() error {
	indexVolume, data, err := readAndCheckVolume(d.fs, d.indexFile, checkIndexVolume)
	d.delegate.OnVolumeFileLoad(0, d.indexFile, indexVolume.header.SetHash, len(data), err)
	if err != nil {
		return err
	}

	d.delegate.OnHeaderLoad(indexVolume.header.String())
	for i, entry := range indexVolume.entries {
		d.delegate.OnFileEntryLoad(i+1, len(indexVolume.entries), entry.filename, entry.header.String())
	}

	// The comment could be in any encoding, so just pass it
	// through as bytes.
	d.delegate.OnCommentLoad(data)

	d.indexVolume = indexVolume
	return nil
}

// LoadFileData loads existing file data into memory.
func (d *Decoder) LoadFileData() error {
	fileData := make([][]byte, 0, len(d.indexVolume.entries))

	for i, entry := range d.indexVolume.entries {
		if !entry.header.Status.savedInVolumeSet() {
			continue
		}

		path, err := d.getFilePath(entry)
		if err != nil {
			return err
		}

		data, corrupt, err := func() ([]byte, bool, error) {
			readStream, err := d.fs.GetReadStream(path)
			if os.IsNotExist(err) {
				return nil, true, err
			} else if err != nil {
				return nil, false, err
			}
			hasher := hashutil.MakeMD5HashCheckerWith16k(entry.header.Hash16k, entry.header.Hash, false)
			data, err := fs.ReadAndClose(hashutil.TeeReadStream(readStream, hasher))
			if _, ok := err.(hashutil.HashMismatchError); ok {
				return nil, true, err
			} else if err != nil {
				return nil, false, err
			}
			return data, false, nil
		}()
		d.delegate.OnDataFileLoad(i+1, len(d.indexVolume.entries), path, len(data), corrupt, err)
		if corrupt {
			fileData = append(fileData, nil)
			continue
		} else if err != nil {
			return err
		}

		fileData = append(fileData, data)
	}

	if len(fileData) == 0 {
		return errors.New("no file data found")
	}

	d.fileData = fileData
	return nil
}

func (d *Decoder) volumePath(volumeNumber uint64) string {
	if volumeNumber == 0 {
		panic("unexpected zero volume number")
	}
	ext := path.Ext(d.indexFile)
	base := d.indexFile[:len(d.indexFile)-len(ext)]
	return base + fmt.Sprintf(".p%02d", volumeNumber)
}

func checkParityVolume(parityVolume volume, expectedSetHash [16]byte, expectedVolumeNumber uint64) error {
	if parityVolume.header.SetHash != expectedSetHash {
		// TODO: Relax this check.
		return errors.New("unexpected set hash for parity volume")
	}

	if parityVolume.header.VolumeNumber != expectedVolumeNumber {
		// TODO: Relax this check.
		return errors.New("unexpected volume number for parity volume")
	}
	return nil
}

// LoadParityData searches for parity volumes and loads them into
// memory.
func (d *Decoder) LoadParityData() error {
	// TODO: Support searching for volume data without relying on
	// filenames.

	// TODO: Count only files saved in volume set.
	fileCount := d.indexVolume.header.FileCount
	maxParityVolumeCount := maxVolumeCount - fileCount
	// TODO: Support more than 99 parity volumes.
	if maxParityVolumeCount > 99 {
		maxParityVolumeCount = 99
	}

	shardByteCount := 0
	parityData := make([][]byte, maxParityVolumeCount)
	var maxI uint64
	for i := uint64(0); i < maxParityVolumeCount; i++ {
		// TODO: Find the file case-insensitively.
		volumeNumber := i + 1
		volumePath := d.volumePath(volumeNumber)
		parityVolume, data, err := func() (parityVolume volume, data []byte, err error) {
			parityVolume, data, err = readAndCheckVolume(d.fs, volumePath, func(parityVolume volume) error {
				return checkParityVolume(parityVolume, d.indexVolume.header.SetHash, volumeNumber)
			})
			if err != nil {
				return parityVolume, data, err
			}

			if len(data) == 0 {
				// TODO: Relax this check.
				return parityVolume, data, errors.New("no parity data in volume")
			}
			if shardByteCount == 0 {
				shardByteCount = len(data)
			} else if len(data) != shardByteCount {
				// TODO: Relax this check.
				return parityVolume, data, errors.New("mismatched parity data byte counts")
			}
			return parityVolume, data, nil
		}()
		d.delegate.OnVolumeFileLoad(volumeNumber, volumePath, parityVolume.header.SetHash, len(data), err)
		if os.IsNotExist(err) {
			continue
		} else if err != nil {
			return err
		}

		parityData[i] = data
		maxI = i
	}

	d.shardByteCount = shardByteCount
	d.parityData = parityData[:maxI+1]
	return nil
}

func (d *Decoder) buildShards() [][]byte {
	shards := make([][]byte, len(d.fileData)+len(d.parityData))
	for i, data := range d.fileData {
		if data == nil {
			continue
		}
		padding := make([]byte, d.shardByteCount-len(data))
		shards[i] = append(data, padding...)
	}

	for i, data := range d.parityData {
		if data == nil {
			continue
		}
		shards[len(d.fileData)+i] = data
	}

	return shards
}

func (d *Decoder) newReedSolomon() (reedsolomon.Encoder, error) {
	return reedsolomon.New(len(d.fileData), len(d.parityData), reedsolomon.WithPAR1Matrix())
}

// FileCounts contains file counts which can be used to deduce whether
// repair is necessary and/or possible.
type FileCounts struct {
	// UsableDataFileCount is the number of data files that are
	// usable, i.e. not missing and not corrupt.
	UsableDataFileCount int
	// UnusableDataFileCount is the number of data files that are
	// unusable, i.e. missing or corrupt.
	UnusableDataFileCount int

	// UsableParityFileCount is the number of parity files that
	// exist, i.e. not missing and not corrupt.
	UsableParityFileCount int
	// UnusableDataFileCount is the number of parity files that
	// are unusable, i.e. missing or corrupt.
	//
	// Note that we can only infer missing parity files from gaps
	// in the counts, e.g. if we only have foo.p01 and foo.p03,
	// then we can infer that foo.p02 is missing, but if we have
	// foo.p01 and foo.p02, then foo.p03 might be missing or it
	// might not have existed in the first place.
	UnusableParityFileCount int
}

// RepairNeeded returns whether repair is needed, i.e. whether
// UnusableDataFileCount is non-zero.
func (fc FileCounts) RepairNeeded() bool {
	return fc.UnusableDataFileCount > 0
}

// RepairPossible returns whether repair is possible i.e. whether
// UsableParityFileCount >= UnusableDataFileCount.
func (fc FileCounts) RepairPossible() bool {
	return fc.UsableParityFileCount >= fc.UnusableDataFileCount
}

// AllFilesUsable returns whether or not all files are usable,
// i.e. whether UnusableDataFileCount == 0 and UnusableParityFileCount
// == 0.
func (fc FileCounts) AllFilesUsable() bool {
	return fc.UnusableDataFileCount == 0 && fc.UnusableParityFileCount == 0
}

// FileCounts returns a FileCounts object for the current file set.
func (d *Decoder) FileCounts() FileCounts {
	usableDataFileCount := 0
	unusableDataFileCount := 0

	for _, data := range d.fileData {
		if data == nil {
			unusableDataFileCount++
		} else {
			usableDataFileCount++
		}
	}

	usableParityFileCount := 0
	unusableParityFileCount := 0

	for _, data := range d.parityData {
		if data == nil {
			unusableParityFileCount++
		} else {
			usableParityFileCount++
		}
	}

	return FileCounts{
		UsableDataFileCount:     usableDataFileCount,
		UnusableDataFileCount:   unusableDataFileCount,
		UsableParityFileCount:   usableParityFileCount,
		UnusableParityFileCount: unusableParityFileCount,
	}
}

// VerifyAllData checks and returns whether all data and parity files
// contain correct data. Note that this is worth calling only when
// there are no unusable data or parity files, since if there are,
// then false is guaranteed to be returned for ok.
func (d *Decoder) VerifyAllData() (ok bool, err error) {
	rs, err := d.newReedSolomon()
	if err != nil {
		return false, err
	}

	shards := d.buildShards()

	return rs.Verify(shards)
}

// Repair tries to repair any missing or corrupt data, using the
// parity volumes. Returns a list of paths to files that were
// successfully repaired (relative to the indexFile passed to
// NewDecoder) in no particular order, which is present even if an
// error is returned. If checkParity is true, extra checking is done
// of the reconstructed parity data.
func (d *Decoder) Repair(checkParity bool) ([]string, error) {
	rs, err := d.newReedSolomon()
	if err != nil {
		return nil, err
	}

	shards := d.buildShards()

	err = rs.Reconstruct(shards)
	if err != nil {
		return nil, err
	}

	if checkParity {
		ok, err := rs.Verify(shards)
		if err != nil {
			return nil, err
		}

		if !ok {
			return nil, errors.New("repair failed")
		}
	}

	var repairedPaths []string

	for i, data := range d.fileData {
		if data != nil {
			continue
		}

		entry := d.indexVolume.entries[i]
		data = shards[i][:entry.header.FileBytes]
		if err := hashutil.CheckMD5Hashes(data, entry.header.Hash16k, entry.header.Hash, true); err != nil {
			return repairedPaths, err
		}

		path, err := d.getFilePath(entry)
		if err != nil {
			return repairedPaths, err
		}

		err = func() error {
			writeStream, err := d.fs.GetWriteStream(path)
			if err != nil {
				return err
			}
			return fs.WriteAndClose(writeStream, data)
		}()
		d.delegate.OnDataFileWrite(i+1, len(d.fileData), path, len(data), err)
		if err != nil {
			return repairedPaths, err
		}

		repairedPaths = append(repairedPaths, path)
		d.fileData[i] = data
	}

	// TODO: Repair missing parity volumes, too, and then make
	// sure d.VerifyAllData() passes.

	return repairedPaths, nil
}
<commit_msg>Pull out parity file loading from par1.Decoder.LoadParityData<commit_after>package par1

import (
	"errors"
	"fmt"
	"os"
	"path"
	"path/filepath"

	"github.com/akalin/gopar/fs"
	"github.com/akalin/gopar/hashutil"
	"github.com/klauspost/reedsolomon"
)

// A Decoder keeps track of all information needed to check the
// integrity of a set of data files, and possibly repair any
// missing/corrupt data files from the parity files (.P00, .P01,
// etc.).
type Decoder struct {
	fs       fs.FS
	delegate DecoderDelegate

	indexFile   string
	indexVolume volume

	fileData [][]byte

	shardByteCount int
	parityData     [][]byte
}

// DecoderDelegate holds methods that are called during the decode
// process.
type DecoderDelegate interface {
	OnHeaderLoad(headerInfo string)
	OnFileEntryLoad(i, n int, filename, entryInfo string)
	OnCommentLoad(comment []byte)
	OnDataFileLoad(i, n int, path string, byteCount int, corrupt bool, err error)
	OnDataFileWrite(i, n int, path string, byteCount int, err error)
	OnVolumeFileLoad(i uint64, path string, setHash [16]byte, dataByteCount int, err error)
}

// DoNothingDecoderDelegate is an implementation of DecoderDelegate
// that does nothing for all methods.
type DoNothingDecoderDelegate struct{}

// OnHeaderLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnHeaderLoad(headerInfo string) {}

// OnFileEntryLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnFileEntryLoad(i, n int, filename, entryInfo string) {}

// OnCommentLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnCommentLoad(comment []byte) {}

// OnDataFileLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, corrupt bool, err error) {
}

// OnDataFileWrite implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnDataFileWrite(i, n int, path string, byteCount int, err error) {}

// OnVolumeFileLoad implements the DecoderDelegate interface.
func (DoNothingDecoderDelegate) OnVolumeFileLoad(i uint64, path string, setHash [16]byte, dataByteCount int, err error) {
}

func newDecoder(fs fs.FS, delegate DecoderDelegate, indexFile string) *Decoder {
	return &Decoder{fs, delegate, indexFile, volume{}, nil, 0, nil}
}

// NewDecoder reads the given index file, which usually has a .PAR
// extension.
func NewDecoder(delegate DecoderDelegate, indexFile string) *Decoder {
	return newDecoder(fs.MakeDefaultFS(), delegate, indexFile)
}

func (d *Decoder) getFilePath(entry fileEntry) (string, error) {
	filename := entry.filename
	if filepath.Base(filename) != filename {
		return "", errors.New("bad filename")
	}

	basePath := filepath.Dir(d.indexFile)
	return filepath.Join(basePath, filename), nil
}

func readAndCheckVolume(filesystem fs.FS, path string, checkVolumeFn func(volume) error) (v volume, data []byte, err error) {
	readStream, err := filesystem.GetReadStream(path)
	if err != nil {
		return volume{}, nil, err
	}

	defer func() {
		if readStream != nil {
			fs.CloseCloser(readStream, &err)
		}
	}()

	v, err = readVolume(readStream)
	if err != nil {
		return volume{}, nil, err
	}

	err = checkVolumeFn(v)
	if err != nil {
		return v, nil, err
	}

	data, err = fs.ReadAndClose(readStream)
	readStream = nil
	if err != nil {
		return v, nil, err
	}
	return v, data, nil
}

func checkIndexVolume(indexVolume volume) error {
	if indexVolume.header.VolumeNumber != 0 {
		// TODO: Relax this check.
		return errors.New("expected volume number 0 for index volume")
	}
	return nil
}

// LoadIndexFile reads the data from the index file.
func (d *Decoder) LoadIndexFile() error {
	indexVolume, data, err := readAndCheckVolume(d.fs, d.indexFile, checkIndexVolume)
	d.delegate.OnVolumeFileLoad(0, d.indexFile, indexVolume.header.SetHash, len(data), err)
	if err != nil {
		return err
	}

	d.delegate.OnHeaderLoad(indexVolume.header.String())
	for i, entry := range indexVolume.entries {
		d.delegate.OnFileEntryLoad(i+1, len(indexVolume.entries), entry.filename, entry.header.String())
	}

	// The comment could be in any encoding, so just pass it
	// through as bytes.
	d.delegate.OnCommentLoad(data)

	d.indexVolume = indexVolume
	return nil
}

// LoadFileData loads existing file data into memory.
func (d *Decoder) LoadFileData() error {
	fileData := make([][]byte, 0, len(d.indexVolume.entries))

	for i, entry := range d.indexVolume.entries {
		if !entry.header.Status.savedInVolumeSet() {
			continue
		}

		path, err := d.getFilePath(entry)
		if err != nil {
			return err
		}

		data, corrupt, err := func() ([]byte, bool, error) {
			readStream, err := d.fs.GetReadStream(path)
			if os.IsNotExist(err) {
				return nil, true, err
			} else if err != nil {
				return nil, false, err
			}
			hasher := hashutil.MakeMD5HashCheckerWith16k(entry.header.Hash16k, entry.header.Hash, false)
			data, err := fs.ReadAndClose(hashutil.TeeReadStream(readStream, hasher))
			if _, ok := err.(hashutil.HashMismatchError); ok {
				return nil, true, err
			} else if err != nil {
				return nil, false, err
			}
			return data, false, nil
		}()
		d.delegate.OnDataFileLoad(i+1, len(d.indexVolume.entries), path, len(data), corrupt, err)
		if corrupt {
			fileData = append(fileData, nil)
			continue
		} else if err != nil {
			return err
		}

		fileData = append(fileData, data)
	}

	if len(fileData) == 0 {
		return errors.New("no file data found")
	}

	d.fileData = fileData
	return nil
}

func (d *Decoder) volumePath(volumeNumber uint64) string {
	if volumeNumber == 0 {
		panic("unexpected zero volume number")
	}
	ext := path.Ext(d.indexFile)
	base := d.indexFile[:len(d.indexFile)-len(ext)]
	return base + fmt.Sprintf(".p%02d", volumeNumber)
}

func checkParityVolume(parityVolume volume, expectedSetHash [16]byte, expectedVolumeNumber uint64) error {
	if parityVolume.header.SetHash != expectedSetHash {
		// TODO: Relax this check.
		return errors.New("unexpected set hash for parity volume")
	}

	if parityVolume.header.VolumeNumber != expectedVolumeNumber {
		// TODO: Relax this check.
		return errors.New("unexpected volume number for parity volume")
	}
	return nil
}

func (d *Decoder) loadParityFile(volumeNumber uint64, volumePath string, shardByteCount *int) (volume, []byte, error) {
	parityVolume, data, err := readAndCheckVolume(d.fs, volumePath, func(parityVolume volume) error {
		return checkParityVolume(parityVolume, d.indexVolume.header.SetHash, volumeNumber)
	})
	if err != nil {
		return parityVolume, data, err
	}

	if len(data) == 0 {
		// TODO: Relax this check.
		return parityVolume, data, errors.New("no parity data in volume")
	}
	if *shardByteCount == 0 {
		*shardByteCount = len(data)
	} else if len(data) != *shardByteCount {
		// TODO: Relax this check.
		return parityVolume, data, errors.New("mismatched parity data byte counts")
	}
	return parityVolume, data, nil
}

// LoadParityData searches for parity volumes and loads them into
// memory.
func (d *Decoder) LoadParityData() error {
	// TODO: Support searching for volume data without relying on
	// filenames.

	// TODO: Count only files saved in volume set.
	fileCount := d.indexVolume.header.FileCount
	maxParityVolumeCount := maxVolumeCount - fileCount
	// TODO: Support more than 99 parity volumes.
	if maxParityVolumeCount > 99 {
		maxParityVolumeCount = 99
	}

	shardByteCount := 0
	parityData := make([][]byte, maxParityVolumeCount)
	var maxI uint64
	for i := uint64(0); i < maxParityVolumeCount; i++ {
		volumeNumber := i + 1
		// TODO: Find the file case-insensitively.
		volumePath := d.volumePath(volumeNumber)
		parityVolume, data, err := d.loadParityFile(volumeNumber, volumePath, &shardByteCount)
		d.delegate.OnVolumeFileLoad(volumeNumber, volumePath, parityVolume.header.SetHash, len(data), err)
		if os.IsNotExist(err) {
			continue
		} else if err != nil {
			return err
		}

		parityData[i] = data
		maxI = i
	}

	d.shardByteCount = shardByteCount
	d.parityData = parityData[:maxI+1]
	return nil
}

func (d *Decoder) buildShards() [][]byte {
	shards := make([][]byte, len(d.fileData)+len(d.parityData))
	for i, data := range d.fileData {
		if data == nil {
			continue
		}
		padding := make([]byte, d.shardByteCount-len(data))
		shards[i] = append(data, padding...)
	}

	for i, data := range d.parityData {
		if data == nil {
			continue
		}
		shards[len(d.fileData)+i] = data
	}

	return shards
}

func (d *Decoder) newReedSolomon() (reedsolomon.Encoder, error) {
	return reedsolomon.New(len(d.fileData), len(d.parityData), reedsolomon.WithPAR1Matrix())
}

// FileCounts contains file counts which can be used to deduce whether
// repair is necessary and/or possible.
type FileCounts struct {
	// UsableDataFileCount is the number of data files that are
	// usable, i.e. not missing and not corrupt.
	UsableDataFileCount int
	// UnusableDataFileCount is the number of data files that are
	// unusable, i.e. missing or corrupt.
	UnusableDataFileCount int

	// UsableParityFileCount is the number of parity files that
	// exist, i.e. not missing and not corrupt.
	UsableParityFileCount int
	// UnusableDataFileCount is the number of parity files that
	// are unusable, i.e. missing or corrupt.
	//
	// Note that we can only infer missing parity files from gaps
	// in the counts, e.g. if we only have foo.p01 and foo.p03,
	// then we can infer that foo.p02 is missing, but if we have
	// foo.p01 and foo.p02, then foo.p03 might be missing or it
	// might not have existed in the first place.
	UnusableParityFileCount int
}

// RepairNeeded returns whether repair is needed, i.e. whether
// UnusableDataFileCount is non-zero.
func (fc FileCounts) RepairNeeded() bool {
	return fc.UnusableDataFileCount > 0
}

// RepairPossible returns whether repair is possible i.e. whether
// UsableParityFileCount >= UnusableDataFileCount.
func (fc FileCounts) RepairPossible() bool {
	return fc.UsableParityFileCount >= fc.UnusableDataFileCount
}

// AllFilesUsable returns whether or not all files are usable,
// i.e. whether UnusableDataFileCount == 0 and UnusableParityFileCount
// == 0.
func (fc FileCounts) AllFilesUsable() bool {
	return fc.UnusableDataFileCount == 0 && fc.UnusableParityFileCount == 0
}

// FileCounts returns a FileCounts object for the current file set.
func (d *Decoder) FileCounts() FileCounts {
	usableDataFileCount := 0
	unusableDataFileCount := 0

	for _, data := range d.fileData {
		if data == nil {
			unusableDataFileCount++
		} else {
			usableDataFileCount++
		}
	}

	usableParityFileCount := 0
	unusableParityFileCount := 0

	for _, data := range d.parityData {
		if data == nil {
			unusableParityFileCount++
		} else {
			usableParityFileCount++
		}
	}

	return FileCounts{
		UsableDataFileCount:     usableDataFileCount,
		UnusableDataFileCount:   unusableDataFileCount,
		UsableParityFileCount:   usableParityFileCount,
		UnusableParityFileCount: unusableParityFileCount,
	}
}

// VerifyAllData checks and returns whether all data and parity files
// contain correct data. Note that this is worth calling only when
// there are no unusable data or parity files, since if there are,
// then false is guaranteed to be returned for ok.
func (d *Decoder) VerifyAllData() (ok bool, err error) {
	rs, err := d.newReedSolomon()
	if err != nil {
		return false, err
	}

	shards := d.buildShards()

	return rs.Verify(shards)
}

// Repair tries to repair any missing or corrupt data, using the
// parity volumes. Returns a list of paths to files that were
// successfully repaired (relative to the indexFile passed to
// NewDecoder) in no particular order, which is present even if an
// error is returned. If checkParity is true, extra checking is done
// of the reconstructed parity data.
func (d *Decoder) Repair(checkParity bool) ([]string, error) {
	rs, err := d.newReedSolomon()
	if err != nil {
		return nil, err
	}

	shards := d.buildShards()

	err = rs.Reconstruct(shards)
	if err != nil {
		return nil, err
	}

	if checkParity {
		ok, err := rs.Verify(shards)
		if err != nil {
			return nil, err
		}

		if !ok {
			return nil, errors.New("repair failed")
		}
	}

	var repairedPaths []string

	for i, data := range d.fileData {
		if data != nil {
			continue
		}

		entry := d.indexVolume.entries[i]
		data = shards[i][:entry.header.FileBytes]
		if err := hashutil.CheckMD5Hashes(data, entry.header.Hash16k, entry.header.Hash, true); err != nil {
			return repairedPaths, err
		}

		path, err := d.getFilePath(entry)
		if err != nil {
			return repairedPaths, err
		}

		err = func() error {
			writeStream, err := d.fs.GetWriteStream(path)
			if err != nil {
				return err
			}
			return fs.WriteAndClose(writeStream, data)
		}()
		d.delegate.OnDataFileWrite(i+1, len(d.fileData), path, len(data), err)
		if err != nil {
			return repairedPaths, err
		}

		repairedPaths = append(repairedPaths, path)
		d.fileData[i] = data
	}

	// TODO: Repair missing parity volumes, too, and then make
	// sure d.VerifyAllData() passes.

	return repairedPaths, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 Pagoda Box Inc.
//
// This Source Code Form is subject to the terms of the Mozilla Public License,
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.

package jobs

//
import (
	"github.com/pagodabox/nanobox-router"
	"github.com/pagodabox/nanobox-server/config"
	"github.com/pagodabox/nanobox-server/util"
)

//
type Startup struct{}

// process on startup
func (j *Startup) Process() {

	// start all the old containers
	containers, _ := util.ListContainers()
	for _, container := range containers {
		err := util.StartContainer(container.ID)
		if err != nil {
			config.Log.Error(err.Error())
		}
	}

	// when we start we need to reset up the routing
	if container, err := util.GetContainer("web1"); err == nil {
		dc, _ := util.InspectContainer(container.ID)

		config.Router.AddTarget("/", "http://"+dc.NetworkSettings.IPAddress+":8080")
		config.Router.Handler = nil
	} else {
		config.Router.Handler = router.NoDeploy{}
	}

	// we also need to set up a ssh tunnel for each running docker container
	for _, container := range containers {
		dc, err := util.InspectContainer(container.ID)
		if err == nil {
			config.Router.AddForward(dc.NetworkSettings.IPAddress + ":22")
		}
	}
}
<commit_msg>remove the forward adding on start<commit_after>// Copyright (c) 2014 Pagoda Box Inc.
//
// This Source Code Form is subject to the terms of the Mozilla Public License,
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.

package jobs

//
import (
	"github.com/pagodabox/nanobox-router"
	"github.com/pagodabox/nanobox-server/config"
	"github.com/pagodabox/nanobox-server/util"
)

//
type Startup struct{}

// process on startup
func (j *Startup) Process() {

	// start all the old containers
	containers, _ := util.ListContainers()
	for _, container := range containers {
		err := util.StartContainer(container.ID)
		if err != nil {
			config.Log.Error(err.Error())
		}
	}

	// when we start we need to reset up the routing
	if container, err := util.GetContainer("web1"); err == nil {
		dc, _ := util.InspectContainer(container.ID)

		config.Router.AddTarget("/", "http://"+dc.NetworkSettings.IPAddress+":8080")
		config.Router.Handler = nil
	} else {
		config.Router.Handler = router.NoDeploy{}
	}

}
<|endoftext|>